Add support for serialization of large caches

Project.toml

@@ -1,13 +1,20 @@
 name = "LRUCache"
 uuid = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
-version = "1.6.0"
+version = "1.6.1"
 
 [compat]
-julia = "1"
+julia = "≥ 1.9"
 
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
 [targets]
-test = ["Test", "Random"]
+test = ["Test", "Random", "Serialization"]
+
+[weakdeps]
+Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[extensions]
+SerializationExt = ["Serialization"]

ext/SerializationExt.jl

@@ -0,0 +1,83 @@
+module SerializationExt
+export serialize, deserialize
+using LRUCache
+using Serialization
+
+# Serialization of large LRUs causes a stack overflow error, so we 
+# create a custom serializer that represents LinkedNodes as Ints
+function Serialization.serialize(s::AbstractSerializer, lru::LRU{K, V}) where {K, V}
+    # Create a mapping from memory address to id
+    node_map = Dict{Ptr, Int}()
+    # Create mapping for first node
+    id = 1
+    first_node = node = lru.keyset.first
+    node_map[pointer_from_objref(node)] = id
+    # Go through the rest of the nodes in the cycle and create a mapping
+    node = node.next
+    while node != first_node
+        id += 1
+        node_map[pointer_from_objref(node)] = id
+        node = node.next
+    end
+    @assert id == length(lru) == lru.keyset.length == length(lru.dict)
+    # By this point, the first node has id 1 and the last node has id length(lru)
+    # so when deserializing, we can infer the order by the id
+    # Create the dict with ids instead of nodes
+    dict = Dict{K, Tuple{V, Int, Int}}()
+    for (key, (value, node, s)) in lru.dict
+        id = node_map[pointer_from_objref(node)]
+        dict[key] = (value, id, s)
+    end
+    Serialization.writetag(s.io, Serialization.OBJECT_TAG)
+    Serialization.serialize(s, typeof(lru))
+    Serialization.serialize(s, dict)
+    Serialization.serialize(s, lru.currentsize)
+    Serialization.serialize(s, lru.maxsize)
+    Serialization.serialize(s, lru.hits)
+    Serialization.serialize(s, lru.misses)
+    Serialization.serialize(s, lru.lock)
+    Serialization.serialize(s, lru.by)
+    Serialization.serialize(s, lru.finalizer)
+end
+
+function Serialization.deserialize(s::AbstractSerializer, ::Type{LRU{K, V}}) where {K, V}
+    dict_with_ids = Serialization.deserialize(s)
+    currentsize = Serialization.deserialize(s)
+    maxsize = Serialization.deserialize(s)
+    hits = Serialization.deserialize(s)
+    misses = Serialization.deserialize(s)
+    lock = Serialization.deserialize(s)
+    by = Serialization.deserialize(s)
+    finalizer = Serialization.deserialize(s)
+    # Create a new keyset and mapping from id to node
+    n_nodes = length(dict_with_ids)
+    nodes = Vector{LRUCache.LinkedNode{K}}(undef, n_nodes)
+    dict = Dict{K, Tuple{V, LRUCache.LinkedNode{K}, Int}}()
+    # Create the nodes, but don't link them yet
+    for (key, (value, id, s)) in dict_with_ids
+        nodes[id] = LRUCache.LinkedNode{K}(key)
+        dict[key] = (value, nodes[id], s)
+    end
+    # Link the nodes
+    for (idx, node) in enumerate(nodes)
+        node.next = nodes[idx % n_nodes + 1]
+        node.prev = nodes[idx == 1 ? n_nodes : idx - 1]
+    end
+    # Create keyset with first node and n_nodes
+    keyset = LRUCache.CyclicOrderedSet{K}()
+    keyset.first = nodes[1]
+    keyset.length = n_nodes
+    # Create LRU
+    lru = LRU{K,V}(maxsize=maxsize)
+    lru.dict = dict
+    lru.keyset = keyset
+    lru.currentsize = currentsize
+    lru.hits = hits
+    lru.misses = misses
+    lru.lock = lock
+    lru.by = by
+    lru.finalizer = finalizer
+    lru
+end
+
+end

test/runtests.jl

@@ -297,3 +297,4 @@ end
 end
 
 include("originaltests.jl")
+include("serializationtests.jl")

test/serializationtests.jl

@@ -0,0 +1,39 @@
+using Serialization
+@testset "Serialize and Deserialize" begin
+
+    cache = LRU{Int, Int}(maxsize=100_000)
+
+    # Populate the cache with dummy data
+    for i in 1:1_000_000
+        cache[i] = i+1
+    end
+    serialize("cache.jls", cache)
+    deserialized_cache = deserialize("cache.jls")
+    rm("cache.jls")
+
+    # Check that the cache is the same
+    @test cache.maxsize == deserialized_cache.maxsize
+    @test cache.currentsize == deserialized_cache.currentsize
+    @test cache.hits == deserialized_cache.hits
+    @test cache.misses == deserialized_cache.misses
+    @test cache.by == deserialized_cache.by
+    @test cache.finalizer == deserialized_cache.finalizer
+    @test cache.keyset.length == deserialized_cache.keyset.length
+    @test length(cache.dict) == length(deserialized_cache.dict)
+    # Check that the cache has the same keyset
+    c_node = cache.keyset.first
+    d_node = deserialized_cache.keyset.first
+    for i in 1:cache.keyset.length
+        c_node.val == d_node.val || @test false
+        c_node = c_node.next
+        d_node = d_node.next
+    end
+    # Check that the cache keys, values, and sizes are the same
+    for (key, (c_value, c_node, c_s)) in cache.dict
+        d_value, d_node, d_s = deserialized_cache.dict[key]
+        c_value == d_value || @test false
+        c_node.val == d_node.val || @test false
+        c_s == d_s || @test false
+    end
+end
+