adopt a better rule, using anymutable

Author: mcabbott

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,7 +131,7 @@ 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]
@@ -143,10 +142,10 @@ julia> fmap(println, (i = twice, ii = 34, iii = [5, 6], iv = (twice, 34), v = 34
 (i = nothing, ii = nothing, iii = nothing, iv = (nothing, nothing), v = nothing)
 ```
 
-If the same object 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.
-Here "same" means `===` for non-`isbits` types. The same number (e.g. `34 === 34`) at
-different nodes is taken to be a coincidence, and `f` applies twice.
+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

@@ -39,26 +39,39 @@ function _default_walk(f, x)
   re(map(f, func))
 end
 
-usecache(x) = !isbits(x)
-usecache(x::Union{String, Symbol}) = false
+usecache(::AbstractDict, x) = isleaf(x) ? anymutable(x) : ismutable(x)
+usecache(::Nothing, x) = false
+
+# function _anymutable(x::T) where {T}
+#   ismutable(x) && return true
+#   fs = fieldnames(T)
+#   isempty(fs) && return false
+#   return any(f -> anymutable(getfield(x, f)), fs)
+# end
+@generated function anymutable(x::T) where {T}
+  ismutabletype(T) && return true
+  fs = fieldnames(T)
+  isempty(fs) && return false
+  subs =  [:(anymutable(getfield(x, $f))) for f in QuoteNode.(fs)]
+  return :(|($(subs...)))
+end
 
 struct NoKeyword end
 
-function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = usecache(x) ? IdDict() : nothing, prune = NoKeyword())
-  if exclude(x)
-    if usecache(x)
-      if haskey(cache, x)
-        prune isa NoKeyword ? cache[x] : prune
-      else
-        cache[x] = f(x)
-      end
-    else
-      f(x)
-    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 = exclude, walk = walk, cache = cache, prune = prune), x)
+    walk(x -> fmap(f, x; exclude, walk, cache, prune), x)
   end
-end
+  if usecache(cache, x)
+    cache[x] = ret
+  end
+  ret
+end  
 
 ###
 ### Extras
@@ -83,20 +96,19 @@ end
 ### Vararg forms
 ###
 
-function fmap(f, x, ys...; exclude = isleaf, walk = _default_walk, cache = usecache(x) ? IdDict() : nothing, prune = NoKeyword())
-  if exclude(x)
-    if usecache(x)
-      if haskey(cache, x)
-        prune isa NoKeyword ? cache[x] : prune
-      else
-        cache[x] = f(x, ys...)
-      end
-    else
-      f(x, ys...)
-    end
+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 = exclude, walk = walk, cache = cache, prune = prune), x, ys...)
+    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...)
@@ -133,3 +145,17 @@ end
 macro flexiblefunctor(args...)
   flexiblefunctorm(args...)
 end
+
+###
+### Compat
+###
+
+if VERSION < v"1.7"
+  # Copied verbatim from Base:
+  function ismutabletype(@nospecialize t)
+      @_total_meta
+      t = unwrap_unionall(t)
+      # TODO: what to do for `Union`?
+      return isa(t, DataType) && t.name.flags & 0x2 == 0x2
+  end
+end

test/basics.jl

@@ -48,16 +48,18 @@ 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
 
-  # The cache applies only to leaf nodes, so that "4" is not shared:
+  # 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)
@@ -72,22 +74,39 @@ end
   @test m3p.y.x == 1:3
 
   # All-isbits trees need not create a cache at all:
-  @test isbits(fmap(float, (x=1, y=(2, 3), z=4:5)))
-  @test_skip 0 == @allocated fmap(float, (x=1, y=(2, 3), z=4:5))
+  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
-    # Leaf types:
-    @test usecache([1,2])
-    @test !usecache(4.0)
-    @test usecache(NoChild([1,2]))
-    @test !usecache(NoChild((3,4)))
+    d = IdDict()
 
-    # Not leaf by default, but `exclude` can change that:
-    @test usecache(Ref(3))
-    @test !usecache((5, 6.0))
-    @test !usecache((a = 2pi, b = missing))
-
-    @test usecache((x = [1,2,3], y = 4))
+    # 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