Merge branch 'master' into dg/grad

Author: DhairyaLGandhi

.gitignore

@@ -1,2 +1,3 @@
 .DS_Store
 /dev/
+Manifest.toml

Manifest.toml

@@ -1,13 +1,9 @@
 name = "Functors"
 uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
-version = "0.2.3"
-
-[deps]
-MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.2.7"
 
 [compat]
-MacroTools = "0.5"
 julia = "1"
 
 [extras]

Project.toml

@@ -17,7 +17,7 @@ Functors.jl provides tools to express a powerful design pattern for dealing with
 
 Functors.jl provides `fmap` to make those things easy, acting as a 'map over parameters':
 
-```julia-repl
+```julia
 julia> using Functors
 
 julia> struct Foo
@@ -36,7 +36,7 @@ Foo(1.0, [1.0, 2.0, 3.0])
 
 It works also with deeply-nested models:
 
-```julia-repl
+```julia
 julia> struct Bar
          x
        end
@@ -52,7 +52,7 @@ Bar(Foo(1.0, [1.0, 2.0, 3.0]))
 
 The workhorse of `fmap` is actually a lower level function, `functor`:
 
-```julia-repl
+```julia
 julia> xs, re = functor(Foo(1, [1, 2, 3]))
 ((x = 1, y = [1, 2, 3]), var"#21#22"())
 
@@ -64,7 +64,7 @@ Foo(1.0, [1.0, 2.0, 3.0])
 
 To include only certain fields, pass a tuple of field names to `@functor`:
 
-```julia-repl
+```julia
 julia> struct Baz
          x
          y
@@ -87,7 +87,7 @@ For a discussion regarding the need for a `cache` in the implementation of `fmap
 
 Use `exclude` for more fine-grained control over whether `fmap` descends into a particular value (the default is `exclude = Functors.isleaf`):
 
-```julia-repl
+```julia
 julia> using CUDA
 
 julia> x = ['a', 'b', 'c'];

README.md

@@ -1,9 +1,8 @@
 module Functors
 
-using MacroTools
-
-export @functor, fmap, fmapstructure, fcollect
+export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect
 
 include("functor.jl")
+include("base.jl")
 
 end # module

src/Functors.jl

@@ -0,0 +1 @@
+@functor Base.RefValue

src/base.jl

@@ -7,6 +7,10 @@ functor(::Type{<:NamedTuple}, x) = x, y -> y
 functor(::Type{<:AbstractArray}, x) = x, y -> y
 functor(::Type{<:AbstractArray{<:Number}}, x) = (), _ -> x
 
+@static if VERSION >= v"1.6"
+  functor(::Type{<:Base.ComposedFunction}, x) = (outer = x.outer, inner = x.inner), y -> Base.ComposedFunction(y.outer, y.inner)
+end
+
 function makefunctor(m::Module, T, fs = fieldnames(T))
   yᵢ = 0
   escargs = map(fieldnames(T)) do f
@@ -20,15 +24,42 @@ function makefunctor(m::Module, T, fs = fieldnames(T))
 end
 
 function functorm(T, fs = nothing)
-  fs == nothing || isexpr(fs, :tuple) || error("@functor T (a, b)")
-  fs = fs == nothing ? [] : [:($(map(QuoteNode, fs.args)...),)]
+  fs === nothing || Meta.isexpr(fs, :tuple) || error("@functor T (a, b)")
+  fs = fs === nothing ? [] : [:($(map(QuoteNode, fs.args)...),)]
   :(makefunctor(@__MODULE__, $(esc(T)), $(fs...)))
 end
 
 macro functor(args...)
   functorm(args...)
 end
 
+function makeflexiblefunctor(m::Module, T, pfield)
+  pfield = QuoteNode(pfield)
+  @eval m begin
+    function $Functors.functor(::Type{<:$T}, x)
+      pfields = getproperty(x, $pfield)
+      function re(y)
+        all_args = map(fn -> getproperty(fn in pfields ? y : x, fn), fieldnames($T))
+        return $T(all_args...)
+      end
+      func = NamedTuple{pfields}(map(p -> getproperty(x, p), pfields))
+      return func, re
+    end
+
+  end
+
+end
+
+function flexiblefunctorm(T, pfield = :params)
+  pfield isa Symbol || error("@flexiblefunctor T param_field")
+  pfield = QuoteNode(pfield)
+  :(makeflexiblefunctor(@__MODULE__, $(esc(T)), $(esc(pfield))))
+end
+
+macro flexiblefunctor(args...)
+  flexiblefunctorm(args...)
+end
+
 """
     isleaf(x)
 
@@ -137,7 +168,8 @@ fmapstructure(f, x; kwargs...) = fmap(f, x; walk = (f, x) -> map(f, children(x))
     fcollect(x; exclude = v -> false)
 
 Traverse `x` by recursing each child of `x` as defined by [`functor`](@ref)
-and collecting the results into a flat array.
+and collecting the results into a flat array, ordered by a breadth-first
+traversal of `x`, respecting the iteration order of `children` calls.
 
 Doesn't recurse inside branches rooted at nodes `v`
 for which `exclude(v) == true`.
@@ -180,13 +212,17 @@ julia> fcollect(m, exclude = v -> Functors.isleaf(v))
  Bar([1, 2, 3])
 ```
 """
-function fcollect(x; cache = [], exclude = v -> false)
-  x in cache && return cache
-  if !exclude(x)
-    push!(cache, x)
-    foreach(y -> fcollect(y; cache = cache, exclude = exclude), children(x))
-  end
-  return cache
+function fcollect(x; output = [], cache = Base.IdSet(), exclude = v -> false)
+    # note: we don't have an `OrderedIdSet`, so we use an `IdSet` for the cache
+    # (to ensure we get exactly 1 copy of each distinct array), and a usual `Vector`
+    # for the results, to preserve traversal order (important downstream!).
+    x in cache && return output
+    if !exclude(x)
+      push!(cache, x)
+      push!(output, x)
+      foreach(y -> fcollect(y; cache=cache, output=output, exclude=exclude), children(x))
+    end
+    return output
 end
 
 # Allow gradients and other constructs that match the structure of the functor

src/functor.jl

@@ -0,0 +1,8 @@
+@testset "Base" begin
+    @testset "RefValue" begin
+        x = Ref(1)
+        p, re = Functors.functor(x)
+        @test p == (x = 1,)
+        @test re(p) isa Base.RefValue{Int}
+    end
+end

test/base.jl

@@ -1,5 +1,3 @@
-using Functors, Test
-
 struct Foo
   x
   y
@@ -23,6 +21,16 @@ struct NoChildren
   y
 end
 
+@static if VERSION >= v"1.6"
+  @testset "ComposedFunction" begin
+    f1 = Foo(1.1, 2.2)
+    f2 = Bar(3.3)
+    @test Functors.functor(f1 ∘ f2)[1] == (outer = f1, inner = f2)
+    @test Functors.functor(f1 ∘ f2)[2]((outer = f1, inner = f2)) == f1 ∘ f2
+    @test fmap(x -> x + 10, f1 ∘ f2) == Foo(11.1, 12.2) ∘ Bar(13.3)
+  end
+end
+
 @testset "Nested" begin
   model = Bar(Foo(1, [1, 2, 3]))
 
@@ -73,6 +81,76 @@ end
   m0 = NoChildren(:a, :b)
   m3 = Foo(m2, m0)
   m4 = Bar(m3)
-  println(fcollect(m4))
   @test all(fcollect(m4) .=== [m4, m3, m2, m1, m0])
+
+  m1 = [1, 2, 3]
+  m2 = [1, 2, 3]
+  m3 = Foo(m1, m2)
+  @test all(fcollect(m3) .=== [m3, m1, m2])
+end
+
+struct FFoo
+  x
+  y
+  p
+end
+@flexiblefunctor FFoo p
+
+struct FBar
+  x
+  p
+end
+@flexiblefunctor FBar p
+
+struct FBaz
+  x
+  y
+  z
+  p
+end
+@flexiblefunctor FBaz p
+
+@testset "Flexible Nested" begin
+  model = FBar(FFoo(1, [1, 2, 3], (:y, )), (:x,))
+
+  model′ = fmap(float, model)
+
+  @test model.x.y == model′.x.y
+  @test model′.x.y isa Vector{Float64}
+end
+
+@testset "Flexible Walk" begin
+  model = FFoo((0, FBar([1, 2, 3], (:x,))), [4, 5], (:x, :y))
+
+  model′ = fmapstructure(identity, model)
+  @test model′ == (; x=(0, (; x=[1, 2, 3])), y=[4, 5])
+
+  model2 = FFoo((0, FBar([1, 2, 3], (:x,))), [4, 5], (:x,))
+
+  model2′ = fmapstructure(identity, model2)
+  @test model2′ == (; x=(0, (; x=[1, 2, 3])))
+end
+
+@testset "Flexible Property list" begin
+  model = FBaz(1, 2, 3, (:x, :z))
+  model′ = fmap(x -> 2x, model)
+
+  @test (model′.x, model′.y, model′.z) == (2, 2, 6)
+end
+
+@testset "Flexible fcollect" begin
+  m1 = 1
+  m2 = [1, 2, 3]
+  m3 = FFoo(m1, m2, (:y, ))
+  m4 = FBar(m3, (:x,))
+  @test all(fcollect(m4) .=== [m4, m3, m2])
+  @test all(fcollect(m4, exclude = x -> x isa Array) .=== [m4, m3])
+  @test all(fcollect(m4, exclude = x -> x isa FFoo) .=== [m4])
+
+  m0 = NoChildren(:a, :b)
+  m1 = [1, 2, 3]
+  m2 = FBar(m1, ())
+  m3 = FFoo(m2, m0, (:x, :y,))
+  m4 = FBar(m3, (:x,))
+  @test all(fcollect(m4) .=== [m4, m3, m2, m0])
 end