make flatten more compiler frendly

1. add inference test
2. add more test and remove `cat_nested`:
   `cat_nested` failed to infer in some cases.
    It has been inserted into `make_makeargs`, thus unneeded.

Author: N5N3

base/broadcast.jl

@@ -314,21 +314,21 @@ some cases.
 """
 function flatten(bc::Broadcasted{Style}) where {Style}
     isflat(bc) && return bc
-    # concatenate the nested arguments into {a, b, c, d}
-    args = cat_nested(bc)
-    # build a function `makeargs` that takes a "flat" argument list and
+    # 1. concatenate the nested arguments into {a, b, c, d}
+    # 2. build a function `makeargs` that takes a "flat" argument list and
     # and creates the appropriate input arguments for `f`, e.g.,
     #          makeargs = (w, x, y, z) -> (w, g(x, y), z)
     #
     # `makeargs` is built recursively and looks a bit like this:
     #     makeargs(w, x, y, z) = (w, makeargs1(x, y, z)...)
     #                          = (w, g(x, y), makeargs2(z)...)
     #                          = (w, g(x, y), z)
-    let makeargs = make_makeargs(()->(), bc.args), f = bc.f
-        newf = @inline function(args::Vararg{Any,N}) where N
-            f(makeargs(args...)...)
-        end
-        return Broadcasted{Style}(newf, args, bc.axes)
+    let (makeargs, args) = make_makeargs((), bc.args), f = bc.f
+        _make(::NTuple{N,Any}) where {N} =
+            @inline function (args::Vararg{Any,N})
+                f(makeargs(args...)...)
+            end
+        return Broadcasted{Style}(_make(args), args, bc.axes)
     end
 end
 
@@ -338,79 +338,47 @@ _isflat(args::NestedTuple) = false
 _isflat(args::Tuple) = _isflat(tail(args))
 _isflat(args::Tuple{}) = true
 
-cat_nested(t::Broadcasted, rest...) = (cat_nested(t.args...)..., cat_nested(rest...)...)
-cat_nested(t::Any, rest...) = (t, cat_nested(rest...)...)
-cat_nested() = ()
-
 """
-    make_makeargs(makeargs_tail::Function, t::Tuple) -> Function
+    make_makeargs(args::Tuple, t::Tuple) -> Function, Tuple
 
 Each element of `t` is one (consecutive) node in a broadcast tree.
-Ignoring `makeargs_tail` for the moment, the job of `make_makeargs` is
-to return a function that takes in flattened argument list and returns a
-tuple (each entry corresponding to an entry in `t`, having evaluated
-the corresponding element in the broadcast tree). As an additional
-complication, the passed in tuple may be longer than the number of leaves
-in the subtree described by `t`. The `makeargs_tail` function should
-be called on such additional arguments (but not the arguments consumed
-by `t`).
+`args` contains the rest arguments on the "right" side of `t`.
+The jobs of `make_makeargs` are:
+    1. append the flattened arguments in `t` at the beginning of `args`.
+    2. return a function that takes in flattened argument list and returns a
+    tuple (each entry corresponding to an entry in `t`, having evaluated
+    the corresponding element in the broadcast tree).
 """
-@inline make_makeargs(makeargs_tail, t::Tuple{}) = makeargs_tail
-@inline function make_makeargs(makeargs_tail, t::Tuple)
-    makeargs = make_makeargs(makeargs_tail, tail(t))
-    (head, tail...)->(head, makeargs(tail...)...)
+@inline function make_makeargs(args, t::Tuple{})
+    _make(::NTuple{N,Any}) where {N} = (args::Vararg{Any,N}) -> args
+    _make(args), args
 end
-function make_makeargs(makeargs_tail, t::Tuple{<:Broadcasted, Vararg{Any}})
+@inline function make_makeargs(args, t::Tuple)
+    makeargs, args′ = make_makeargs(args, tail(t))
+    _make(::NTuple{N,Any}) where {N} =
+        @inline function (head, tail::Vararg{Any,N})
+            (head, makeargs(tail...)...)
+        end
+    _make(args′), (t[1], args′...)
+end
+function make_makeargs(args, t::Tuple{<:Broadcasted,Vararg{Any}})
     bc = t[1]
     # c.f. the same expression in the function on leaf nodes above. Here
     # we recurse into siblings in the broadcast tree.
-    let makeargs_tail = make_makeargs(makeargs_tail, tail(t)),
-            # Here we recurse into children. It would be valid to pass in makeargs_tail
-            # here, and not use it below. However, in that case, our recursion is no
-            # longer purely structural because we're building up one argument (the closure)
-            # while destructuing another.
-            makeargs_head = make_makeargs((args...)->args, bc.args),
-            f = bc.f
-        # Create two functions, one that splits of the first length(bc.args)
-        # elements from the tuple and one that yields the remaining arguments.
-        # N.B. We can't call headargs on `args...` directly because
-        # args is flattened (i.e. our children have not been evaluated
-        # yet).
-        headargs, tailargs = make_headargs(bc.args), make_tailargs(bc.args)
-        return @inline function(args::Vararg{Any,N}) where N
-            args1 = makeargs_head(args...)
-            a, b = headargs(args1...), makeargs_tail(tailargs(args1...)...)
-            (f(a...), b...)
-        end
-    end
-end
-
-@inline function make_headargs(t::Tuple)
-    let headargs = make_headargs(tail(t))
-        return @inline function(head, tail::Vararg{Any,N}) where N
-            (head, headargs(tail...)...)
-        end
+    let (makeargs, args′) = make_makeargs(args, tail(t)), f = bc.f
+        # Here we recurse into children. We can pass in `args′` here,
+        # and get `args″` directly, but it is more compiler frendly to
+        # treat `bc` as a new parent "node".
+        makeargs_head, argsˢ = make_makeargs((), bc.args)
+        args″ = (argsˢ..., args′...)
+        _make(::NTuple{L,Any}, ::NTuple{N,Any}) where {L,N} =
+            @inline function (args::Vararg{Any,N})
+                a, b = Base.IteratorsMD.split(args, Val(L)) # split `args...` directly
+                (f(makeargs_head(a...)...), makeargs(b...)...)
+            end
+        _make(argsˢ, args″), args″
     end
 end
-@inline function make_headargs(::Tuple{})
-    return @inline function(tail::Vararg{Any,N}) where N
-        ()
-    end
-end
-
-@inline function make_tailargs(t::Tuple)
-    let tailargs = make_tailargs(tail(t))
-        return @inline function(head, tail::Vararg{Any,N}) where N
-            tailargs(tail...)
-        end
-    end
-end
-@inline function make_tailargs(::Tuple{})
-    return @inline function(tail::Vararg{Any,N}) where N
-        tail
-    end
-end
-
 ## Broadcasting utilities ##
 
 ## logic for deciding the BroadcastStyle

test/broadcast.jl

@@ -772,13 +772,19 @@ end
 
 # issue #27988: inference of Broadcast.flatten
 using .Broadcast: Broadcasted
-let
+let _cat_nested(bc) = Broadcast.flatten(bc).args
     bc = Broadcasted(+, (Broadcasted(*, (1, 2)), Broadcasted(*, (Broadcasted(*, (3, 4)), 5))))
-    @test @inferred(Broadcast.cat_nested(bc)) == (1,2,3,4,5)
+    @test @inferred(_cat_nested(bc)) == (1,2,3,4,5)
     @test @inferred(Broadcast.materialize(Broadcast.flatten(bc))) == @inferred(Broadcast.materialize(bc)) == 62
     bc = Broadcasted(+, (Broadcasted(*, (1, Broadcasted(/, (2.0, 2.5)))), Broadcasted(*, (Broadcasted(*, (3, 4)), 5))))
-    @test @inferred(Broadcast.cat_nested(bc)) == (1,2.0,2.5,3,4,5)
+    @test @inferred(_cat_nested(bc)) == (1,2.0,2.5,3,4,5)
     @test @inferred(Broadcast.materialize(Broadcast.flatten(bc))) == @inferred(Broadcast.materialize(bc)) == 60.8
+    # 1 .* 1 .- 1 .* 1 .^2 .+ 1 .* 1 .+ 1 .^ 3
+    bc = Broadcasted(+, (Broadcasted(+, (Broadcasted(-, (Broadcasted(*, (1, 1)), Broadcasted(*, (1, Broadcasted(Base.literal_pow, (Ref(^), 1, Ref(Val(2)))))))), Broadcasted(*, (1, 1)))), Broadcasted(Base.literal_pow, (Base.RefValue{typeof(^)}(^), 1, Base.RefValue{Val{3}}(Val{3}())))))
+    @test @inferred(Broadcast.materialize(Broadcast.flatten(bc))) == @inferred(Broadcast.materialize(bc)) == 2
+    # @. 1 + 1 * (1 + 1 + 1 + 1)
+    bc = Broadcasted(+, (1, Broadcasted(*, (1, Broadcasted(+, (1, 1, 1, 1))))))
+    @test @inferred(_cat_nested(bc)) == (1,1,1,1,1,1) # `cat_nested` failed to infer this
 end
 
 let