fix perf issue loadmodel! (#2241)

* fix perf issue loadmodel!

* reinstate filter

* cleanup

* cleanup

Author: CarloLucibello

src/loading.jl

@@ -1,19 +1,23 @@
-loadleaf!(dst, src, err) = dst
-loadleaf!(dst::AbstractArray, src, err) =
+loadleaf!(dst, src) = dst
+loadleaf!(dst::AbstractArray, src) =
   error("Tried to copy $src into an array destination; this is not allowed.")
-loadleaf!(dst, src::AbstractArray, err) =
+loadleaf!(dst, src::AbstractArray) =
   error("Tried to copy an array to $dst; this is not allowed.")
-function loadleaf!(dst::AbstractArray, src::Bool, err)
+
+function loadleaf!(dst::AbstractArray, src::Bool)
   if iszero(src)
     dst .= src
   else
     error("Cannot copy boolean parameter == true to non-zero parameter.")
   end
   return dst
 end
-loadleaf!(dst::Bool, src::AbstractArray, err) = iszero(dst) ? dst :
+
+loadleaf!(dst::Bool, src::AbstractArray) = iszero(dst) ? dst :
   error("Cannot copy non-zero parameter to boolean parameter == true.")
-function loadleaf!(dst::AbstractArray, src::AbstractArray, err)
+
+function loadleaf!(dst::AbstractArray, src::AbstractArray)
+  err = DimensionMismatch("Tried to load size $(size(src)) array into size $(size(dst))")
   (size(dst) == size(src)) || throw(err)
   copyto!(dst, src)
 end
@@ -82,20 +86,19 @@ Likewise, copying a `src` value of `false` to any `dst` array is valid,
 but copying a `src` value of `true` will error.
 """
 function loadmodel!(dst, src; filter = _ -> true, cache = Base.IdSet())
-  ldsts = _filter_children(filter, functor(dst)[1])
-  lsrcs = _filter_children(filter, functor(src)[1])
+  ldsts = _filter_children(filter, Functors.children(dst))
+  lsrcs = _filter_children(filter, Functors.children(src))
   (keys(ldsts) == keys(lsrcs)) ||
-    throw(ArgumentError("Tried to load $src into $dst but the structures do not match."))
+    throw(ArgumentError("Tried to load $(keys(lsrcs)) into $(keys(ldsts)) but the structures do not match."))
 
-  err = DimensionMismatch("Tried to load $src into $dst but the parameter sizes do not match.")
   foreach(ldsts, lsrcs) do ldst, lsrc
     if ldst in cache # we already loaded this parameter before
       _tie_check(ldst, lsrc) && return ldst
     elseif Functors.isleaf(ldst) # our first time loading this leaf
       push!(cache, ldst)
-      loadleaf!(ldst, lsrc, err)
+      loadleaf!(ldst, lsrc)
     else # this isn't a leaf
-      loadmodel!(ldst, lsrc; filter = filter, cache = cache)
+      loadmodel!(ldst, lsrc; filter, cache)
     end
   end