simplest isbits usecache

Author: mcabbott

src/Functors.jl

@@ -138,13 +138,15 @@ julia> fmap(println, (i = twice, ii = 34, iii = [5, 6], iv = (twice, 34), v = 34
 [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.
+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.
 
 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,11 +39,15 @@ function _default_walk(f, x)
   re(map(f, func))
 end
 
+usecache(x) = !isbits(x)
+
 struct NoKeyword end
 
-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)
+function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = usecache(x) ? IdDict() : nothing, prune = NoKeyword())
+  usecache(x) && haskey(cache, x) && return prune isa NoKeyword ? cache[x] : prune
+  xnew = exclude(x) ? f(x) : walk(x -> fmap(f, x; exclude=exclude, walk=walk, cache=cache, prune=prune), x)
+  usecache(x) && setindex!(cache, xnew, x)
+  return xnew
 end
 
 ###
@@ -70,8 +74,10 @@ end
 ###
 
 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...)
+  usecache(x) && haskey(cache, x) && return prune isa NoKeyword ? cache[x] : prune
+  xnew = exclude(x) ? f(x, ys...) : walk((xy...,) -> fmap(f, xy...; exclude=exclude, walk=walk, cache=cache, prune=prune), x, ys...)
+  usecache(x) && setindex!(cache, xnew, x)
+  return xnew
 end
 
 function _default_walk(f, x, ys...)

test/basics.jl

@@ -1,5 +1,5 @@
 
-using Functors: functor
+using Functors: functor, usecache
 
 struct Foo; x; y; end
 @functor Foo
@@ -68,7 +68,23 @@ end
   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:
+  @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))
+
+  @testset "usecache" begin
+    @test usecache([1,2])
+    @test usecache(Ref(3))
+
+    @test !usecache(4.0)
+    @test !usecache((5, 6.0))
+    @test !usecache((a = 2pi, b = missing))
+
+    @test usecache(Bar([1,2]))
+    @test !usecache(Bar((3,4)))
+  end
 end
 
 @testset "functor(typeof(x), y) from @functor" begin