move the functions around

Author: mcabbott

src/functor.jl

@@ -34,60 +34,14 @@ 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) = children(x) === ()
 
 children(x) = functor(x)[1]
 
-function functor_tuple(f, x::Tuple, dx::Tuple)
-  map(x, dx) do x, x̄
-    _default_walk(f, x, x̄)
-  end
-end
-functor_tuple(f, x, dx) = f(x, dx)
-functor_tuple(f, x, ::Nothing) = x
-
-# @functor Chain
-# Chain -> func = (layers = (Dense,Dense),), gs -> (layers...)
-function _default_walk(f, x, dx)
-  func, re = functor(x)
-  map(func, dx) do x, x̄
-    # functor_tuple(f, x, x̄)
-    f(x, x̄)
-  end |> re
-end
-
 function _default_walk(f, x)
   func, re = functor(x)
   re(map(f, func))
 end
-_default_walk(f, ::Nothing, ::Nothing) = nothing
 
 function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = IdDict())
   haskey(cache, x) && return cache[x]
@@ -97,6 +51,10 @@ function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = IdDict())
   return y
 end
 
+###
+### Extras
+###
+
 fmapstructure(f, x; kwargs...) = fmap(f, x; walk = (f, x) -> map(f, children(x)), kwargs...)
 
 function fcollect(x; output = [], cache = Base.IdSet(), exclude = v -> false)
@@ -112,10 +70,59 @@ function fcollect(x; output = [], cache = Base.IdSet(), exclude = v -> false)
     return output
 end
 
-# Allow gradients and other constructs that match the structure of the functor
-# to allow for `map` style computations and return a modified version of the struct.
-# This way we can use `fmap` to update the params with their gradients
+###
+### Vararg forms
+###
+
 function fmap(f, x, dx...; cache = IdDict())
   haskey(cache, x) && return cache[x]
   cache[x] = isleaf(x) ? f(x, dx...) : _default_walk((x...) -> fmap(f, x..., cache = cache), x, dx...)
 end
+
+function functor_tuple(f, x::Tuple, dx::Tuple)
+  map(x, dx) do x, x̄
+    _default_walk(f, x, x̄)
+  end
+end
+functor_tuple(f, x, dx) = f(x, dx)
+functor_tuple(f, x, ::Nothing) = x
+
+function _default_walk(f, x, dx)
+  func, re = functor(x)
+  map(func, dx) do x, x̄
+    # functor_tuple(f, x, x̄)
+    f(x, x̄)
+  end |> re
+end
+_default_walk(f, ::Nothing, ::Nothing) = nothing
+
+###
+### FlexibleFunctors.jl
+###
+
+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