Add new caching behavior to fcollect

Author: darsnack

Project.toml

@@ -12,8 +12,9 @@ julia = "1.6"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Documenter", "Zygote"]
+test = ["Test", "Documenter", "StaticArrays", "Zygote"]

src/walks.jl

@@ -88,7 +88,8 @@ end
 
 struct NoKeyword end
 
-usecache(::AbstractDict, x) = isleaf(x) ? anymutable(x) : ismutable(x)
+usecache(::Union{AbstractDict, AbstractSet}, x) =
+  isleaf(x) ? anymutable(x) : ismutable(x)
 usecache(::Nothing, x) = false
 
 @generated function anymutable(x::T) where {T}
@@ -149,9 +150,11 @@ CollectWalk() = CollectWalk(Base.IdSet(), Any[])
 # (to ensure we get exactly 1 copy of each distinct array), and a usual `Vector`
 # for the results, to preserve traversal order (important downstream!).
 function (walk::CollectWalk)(recurse, x)
-  x in walk.cache && return walk.output
+  if usecache(walk.cache, x) && (x in walk.cache)
+    return walk.output
+  end
   # to exclude, we wrap this walk in ExcludeWalk
-  push!(walk.cache, x)
+  usecache(walk.cache, x) && push!(walk.cache, x)
   push!(walk.output, x)
   map(recurse, children(x))
 

test/basics.jl

@@ -85,9 +85,7 @@ end
   @test m5f.x === m5f.y
   @test m5f.x !== m5f.z
 
-  @testset "usecache" begin
-    d = IdDict()
-
+  @testset "usecache ($d)" for d in [IdDict(), Base.IdSet()]
     # Leaf types:
     @test usecache(d, [1,2])
     @test !usecache(d, 4.0)
@@ -101,9 +99,9 @@ end
 
     @test !usecache(d, (5, [6.0]'))  # contains mutable
     @test !usecache(d, (x = [1,2,3], y = 4))
-    
+
     usecache(d, OneChild3([1,2], 3, nothing))  # mutable isn't a child, do we care?
-    
+
     # No dictionary:
     @test !usecache(nothing, [1,2])
     @test !usecache(nothing, 3)
@@ -173,6 +171,14 @@ end
   m2 = [1, 2, 3]
   m3 = Foo(m1, m2)
   @test all(fcollect(m3) .=== [m3, m1, m2])
+
+  m1 = [1, 2, 3]
+  m2 = SVector{length(m1)}(m1)
+  m2′ = SVector{length(m1)}(m1)
+  m3 = Foo(m1, m1)
+  m4 = Foo(m2, m2′)
+  @test all(fcollect(m3) .=== [m3, m1])
+  @test all(fcollect(m4) .=== [m4, m2, m2′])
 end
 
 ###

test/runtests.jl

@@ -1,6 +1,7 @@
 using Functors, Test
 using Zygote
 using LinearAlgebra
+using StaticArrays
 
 @testset "Functors.jl" begin