Use the cache less often

Project.toml

@@ -1,7 +1,7 @@
 name = "Functors"
 uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 authors = ["Mike J Innes <[email protected]>"]
-version = "0.3.0"
+version = "0.4.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/Functors.jl

@@ -5,7 +5,6 @@ export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect
 include("functor.jl")
 include("base.jl")
 
-
 ###
 ### Docstrings for basic functionality
 ###
@@ -132,19 +131,21 @@ Any[23, (45,), (x = 6//7, y = ())]
 [8, 9]
 (a = nothing, b = nothing, c = nothing)
 
-julia> twice = [1, 2];
+julia> twice = [1, 2];  # println only acts once on this
 
 julia> fmap(println, (i = twice, ii = 34, iii = [5, 6], iv = (twice, 34), v = 34.0))
 [1, 2]
 34
 [5, 6]
+34
 34.0
 (i = nothing, ii = nothing, iii = nothing, iv = (nothing, nothing), v = nothing)
 ```
 
-If the same node (same according to `===`) appears more than once,
-it will only be handled once, and only be transformed once with `f`.
-Thus the result will also have this relationship.
+Mutable objects which appear more than once are only handled once (by caching `f(x)` in an `IdDict`).
+Thus the relationship `x.i === x.iv[1]` will be preserved.
+An immutable object which appears twice is not stored in the cache, thus `f(34)` will be called twice,
+and the results will agree only if `f` is pure.
 
 By default, `Tuple`s, `NamedTuple`s, and some other container-like types in Base have
 children to recurse into. Arrays of numbers do not.

src/functor.jl

@@ -2,10 +2,10 @@
 functor(T, x) = (), _ -> x
 functor(x) = functor(typeof(x), x)
 
-functor(::Type{<:Tuple}, x) = x, y -> y
+functor(::Type{<:Tuple}, x) = x, identity
 functor(::Type{<:NamedTuple{L}}, x) where L = NamedTuple{L}(map(s -> getproperty(x, s), L)), identity
 
-functor(::Type{<:AbstractArray}, x) = x, y -> y
+functor(::Type{<:AbstractArray}, x) = x, identity
 functor(::Type{<:AbstractArray{<:Number}}, x) = (), _ -> x
 
 function makefunctor(m::Module, T, fs = fieldnames(T))
@@ -39,13 +39,32 @@ function _default_walk(f, x)
   re(map(f, func))
 end
 
-struct NoKeyword end
+usecache(::AbstractDict, x) = isleaf(x) ? anymutable(x) : ismutable(x)
+usecache(::Nothing, x) = false
 
-function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = IdDict(), prune = NoKeyword())
-  haskey(cache, x) && return prune isa NoKeyword ? cache[x] : prune
-  cache[x] = exclude(x) ? f(x) : walk(x -> fmap(f, x; exclude=exclude, walk=walk, cache=cache, prune=prune), x)
+@generated function anymutable(x::T) where {T}
+  ismutabletype(T) && return true
+  subs =  [:(anymutable(getfield(x, $f))) for f in QuoteNode.(fieldnames(T))]
+  return Expr(:(||), subs...)
 end
 
+struct NoKeyword end
+
+function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = anymutable(x) ? IdDict() : nothing, prune = NoKeyword())
+  if usecache(cache, x) && haskey(cache, x)
+    return prune isa NoKeyword ? cache[x] : prune
+  end
+  ret = if exclude(x)
+    f(x)
+  else
+    walk(x -> fmap(f, x; exclude, walk, cache, prune), x)
+  end
+  if usecache(cache, x)
+    cache[x] = ret
+  end
+  ret
+end  
+
 ###
 ### Extras
 ###
@@ -69,9 +88,19 @@ end
 ### Vararg forms
 ###
 
-function fmap(f, x, ys...; exclude = isleaf, walk = _default_walk, cache = IdDict(), prune = NoKeyword())
-  haskey(cache, x) && return prune isa NoKeyword ? cache[x] : prune
-  cache[x] = exclude(x) ? f(x, ys...) : walk((xy...,) -> fmap(f, xy...; exclude=exclude, walk=walk, cache=cache, prune=prune), x, ys...)
+function fmap(f, x, ys...; exclude = isleaf, walk = _default_walk, cache = anymutable(x) ? IdDict() : nothing, prune = NoKeyword())
+  if usecache(cache, x) && haskey(cache, x)
+    return prune isa NoKeyword ? cache[x] : prune
+  end
+  ret = if exclude(x)
+    f(x, ys...)
+  else
+    walk((xy...,) -> fmap(f, xy...; exclude, walk, cache, prune), x, ys...)
+  end
+  if usecache(cache, x)
+    cache[x] = ret
+  end
+  ret
 end
 
 function _default_walk(f, x, ys...)
@@ -108,3 +137,19 @@ end
 macro flexiblefunctor(args...)
   flexiblefunctorm(args...)
 end
+
+###
+### Compat
+###
+
+if VERSION < v"1.7"
+  # # Copied verbatim from Base, except omitting the macro:
+  # function ismutabletype(@nospecialize t)
+  #     # @_total_meta
+  #     t = Base.unwrap_unionall(t)
+  #     return isa(t, DataType) && t.name.flags & 0x2 == 0x2
+  # end
+
+  # That doesn't work, but this does:
+  ismutabletype(@nospecialize t) = t.mutable
+end

test/basics.jl

@@ -1,5 +1,5 @@
 
-using Functors: functor
+using Functors: functor, usecache
 
 struct Foo; x; y; end
 @functor Foo
@@ -14,6 +14,7 @@ struct NoChildren2; x; y; end
 
 struct NoChild{T}; x::T; end
 
+
 ###
 ### Basic functionality
 ###
@@ -47,28 +48,66 @@ end
   @test (model′.x, model′.y, model′.z) == (1, 4, 3)
 end
 
-@testset "cache" begin
+@testset "Sharing" begin
   shared = [1,2,3]
   m1 = Foo(shared, Foo([1,2,3], Foo(shared, [1,2,3])))
   m1f = fmap(float, m1)
   @test m1f.x === m1f.y.y.x
   @test m1f.x !== m1f.y.x
   m1p = fmapstructure(identity, m1; prune = nothing)
   @test m1p == (x = [1, 2, 3], y = (x = [1, 2, 3], y = (x = nothing, y = [1, 2, 3])))
+  m1no = fmap(float, m1; cache = nothing)  # disable the cache by hand
+  @test m1no.x !== m1no.y.y.x
 
-  # A non-leaf node can also be repeated:
+  # Here "4" is not shared, because Foo isn't leaf:
   m2 = Foo(Foo(shared, 4), Foo(shared, 4))
   @test m2.x === m2.y
   m2f = fmap(float, m2)
   @test m2f.x.x === m2f.y.x
   m2p = fmapstructure(identity, m2; prune = Bar(0))
-  @test m2p == (x = (x = [1, 2, 3], y = 4), y = Bar(0))
+  @test m2p == (x = (x = [1, 2, 3], y = 4), y = (x = Bar{Int64}(0), y = 4))
 
   # Repeated isbits types should not automatically be regarded as shared:
   m3 = Foo(Foo(shared, 1:3), Foo(1:3, shared))
   m3p = fmapstructure(identity, m3; prune = 0)
   @test m3p.y.y == 0
-  @test_broken m3p.y.x == 1:3
+  @test m3p.y.x == 1:3
+
+  # All-isbits trees need not create a cache at all:
+  m4 = (x=1, y=(2, 3), z=4:5)
+  @test isbits(fmap(float, m4))
+  @test_skip 0 == @allocated fmap(float, m4)  # true, but fails in tests
+
+  # Shared mutable containers are preserved, even if all children are isbits:
+  ref = Ref(1)
+  m5 = (x = ref, y = ref, z = Ref(1))
+  m5f = fmap(x -> x/2, m5)
+  @test m5f.x === m5f.y
+  @test m5f.x !== m5f.z
+
+  @testset "usecache" begin
+    d = IdDict()
+
+    # Leaf types:
+    @test usecache(d, [1,2])
+    @test !usecache(d, 4.0)
+    @test usecache(d, NoChild([1,2]))
+    @test !usecache(d, NoChild((3,4)))
+
+    # Not leaf:
+    @test usecache(d, Ref(3))  # mutable container
+    @test !usecache(d, (5, 6.0))
+    @test !usecache(d, (a = 2pi, b = missing))
+
+    @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)
+  end
 end
 
 @testset "functor(typeof(x), y) from @functor" begin