Preventing broadcast from creating anonymous functions on lowering for a custom type

I want to implement my own type that behaves similarly to an Array, in particular participates in the same arithmetic expressions. Except, instead of eager evaluation, the arithmetic operations should just return an expression tree. Since, for example, .* or .^ are different semantically from * and ^ (i.e. it’s not just a performance optimization), the difference between dotted and non-dotted operators must be preserved in the tree.

In Julia 0.5 it was pretty straightforward: you would just add methods to * and .* separately. Now, as far as I understand, you are supposed to add methods to the Base.Broadcast machinery (although it does not seem to be documented, so I am not sure if it is the proper way). I managed to get it working to some extent, but encountered a problem. Let me illustrate:

struct A
    x
end

Base.Broadcast.containertype(::A) = A
Base.Broadcast.promote_containertype(::Type{A}, ::Type{A}) = A
Base.Broadcast.promote_containertype(::Type{A}, x) = A
Base.Broadcast.promote_containertype(x, ::Type{A}) = A
Base.Broadcast.promote_containertype(::Type{A}, ::Type{Array}) = A
Base.Broadcast.promote_containertype(::Type{Array}, ::Type{A}) = A

function Base.Broadcast.broadcast_c(op, ::Type{A}, a...)
    [Symbol(".", op), a...]
end


a = A("a")
b = A("b")
c = A("c")
s = 5

println("a .+ b: $(a .+ b)")
println("a .+ s: $(a .+ s)")
println("a .+ [1 2 3]: $(a .+ [1 2 3])")

println("a .+ 1: $(a .+ 1)")
println("a .+ b .+ c: $(a .+ b .+ c)")

The first three expressions work fine:

a .+ b: Any[:.+, A("a"), A("b")]
a .+ s: Any[:.+, A("a"), 5]
a .+ [1 2 3]: Any[:.+, A("a"), [1 2 3]]

For the last two, on the other hand, there is a problem. As far as I understand, during lowering a broadcasted operator with a literal, or a sequence of several broadcasted operators is transformed right away into an anonymous function (x -> x + 1 and (a, b, c) -> a + b + c), so broadcast_c receives this function as op instead of +:

a .+ 1: Any[Symbol(".#1"), A("a")]
a .+ b .+ c: Any[Symbol(".#3"), A("a"), A("b"), A("c")]

This means that the information about the actual operation is now hidden in this anonymous function, and I cannot analyze/transform the resulting expression tree, which I would like to do eventually. Is there a way to stop this from happening? For the last two cases, I would like to get something like

a .+ 1: Any[:.+, A("a"), 1]
a .+ b .+ c: Any[:.+, A("a"), A("b"), A("c")]

or

a .+ 1: Any[:.+, A("a"), 1]
a .+ b .+ c: Any[:.+, Any[:.+, A("a"), A("b")], A("c")]

https://github.com/JuliaLang/julia/issues/22060

https://github.com/JuliaLang/julia/pull/22063

1 Like

Perfect, thanks!