jw3126 / Setfield.jl

Update deeply nested immutable structs.
Other
167 stars 17 forks source link

Use gensym for `lens` variable in `setmacro` #156

Closed torfjelde closed 3 years ago

torfjelde commented 3 years ago

This makes the setmacro a bit more "multiple-use-friendly", e.g.

julia> macro f(expr)
           quote
               $(Setfield.setmacro(identity, expr))
               $(Setfield.setmacro(identity, expr))
           end
       end
@f (macro with 1 method)

julia> @macroexpand @f(x[1] = 1)
quote
    #= REPL[28]:3 =#
    begin
        #= /home/tor/.julia/packages/Setfield/XM37G/src/sugar.jl:181 =#
        var"#42#lens" = (identity)((Setfield.compose)((Setfield.IndexLens)((1,))))
        #= /home/tor/.julia/packages/Setfield/XM37G/src/sugar.jl:182 =#
        var"#41###_#392" = (Setfield.set)(x, var"#42#lens", 1)
    end
    #= REPL[28]:4 =#
    begin
        #= /home/tor/.julia/packages/Setfield/XM37G/src/sugar.jl:181 =#
        var"#42#lens" = (identity)((Setfield.compose)((Setfield.IndexLens)((1,))))
        #= /home/tor/.julia/packages/Setfield/XM37G/src/sugar.jl:182 =#
        var"#43###_#393" = (Setfield.set)(x, var"#42#lens", 1)
    end
end

vs. on this branch

julia> @macroexpand @f(x[1] = 1)
quote
    #= REPL[5]:3 =#
    begin
        #= /home/tor/Projects/public/Setfield.jl/src/sugar.jl:182 =#
        var"#1###lens#257" = (identity)((Setfield.compose)((Setfield.IndexLens)((1,))))
        #= /home/tor/Projects/public/Setfield.jl/src/sugar.jl:183 =#
        var"#2###_#258" = (Setfield.set)(x, var"#1###lens#257", 1)
    end
    #= REPL[5]:4 =#
    begin
        #= /home/tor/Projects/public/Setfield.jl/src/sugar.jl:182 =#
        var"#3###lens#259" = (identity)((Setfield.compose)((Setfield.IndexLens)((1,))))
        #= /home/tor/Projects/public/Setfield.jl/src/sugar.jl:183 =#
        var"#4###_#260" = (Setfield.set)(x, var"#3###lens#259", 1)
    end
end

We could just re-use the variable, but AFAIK this makes things a bit easier for the compiler.

jw3126 commented 3 years ago

Thanks! Out of curiosity, does the current behaviour create a real-world bug or problem for you? If so could an MWE of that be turned into a test?

torfjelde commented 3 years ago

Thanks! Out of curiosity, does the current behaviour create a real-world bug or problem for you? If so could an MWE of that be turned into a test?

It doesn't really create a bug per se, but in our use-case (we want to use it within Turing.jl's @model macro) we might have 100s of these in the same scope. In that case it's preferable to not have a single lens variable used everywhere since this can make type-inference difficult. It will still have the correct behavior though:)

torfjelde commented 3 years ago

Okay so this is the sort of thing I'm talking about:

julia> macro f(expr)
           quote
               function f($(esc(:x)))
                   $(Setfield.setmacro(identity, expr, overwrite=true))
                   $(Setfield.setmacro(identity, expr, overwrite=true))
                   $(Setfield.setmacro(identity, expr, overwrite=true))
                   $(Setfield.setmacro(identity, expr, overwrite=true))
                   $(Setfield.setmacro(identity, expr, overwrite=true))
                   return $(esc(:x))
               end
           end
       end
@f (macro with 1 method)

julia> f = @f(x[end] = 1)
#35#f (generic function with 1 method)

julia> @code_warntype f(rand(2))
Variables
  #self#::Core.Const(var"#35#f")
  x@_2::Vector{Float64}
  #13::var"#13#18"
  #12::var"#12#17"
  #11::var"#11#16"
  #10::var"#10#15"
  #9::var"#9#14"
  lens::Setfield.DynamicIndexLens
  x@_9::Vector{Float64}

Body::Vector{Float64}
1 ─       (x@_9 = x@_2)
│         (#9 = %new(Main.:(var"#9#14")))
│   %3  = #9::Core.Const(var"#9#14"())
│   %4  = (Setfield.DynamicIndexLens)(%3)::Core.Const(Setfield.DynamicIndexLens{var"#9#14"}(var"#9#14"()))
│   %5  = (Setfield.compose)(%4)::Core.Const(Setfield.DynamicIndexLens{var"#9#14"}(var"#9#14"()))
│         (lens = (identity)(%5))
│         (x@_9 = (Setfield.set)(x@_9, lens::Core.Const(Setfield.DynamicIndexLens{var"#9#14"}(var"#9#14"())), 1))
│         (#10 = %new(Main.:(var"#10#15")))
│   %9  = #10::Core.Const(var"#10#15"())
│   %10 = (Setfield.DynamicIndexLens)(%9)::Core.Const(Setfield.DynamicIndexLens{var"#10#15"}(var"#10#15"()))
│   %11 = (Setfield.compose)(%10)::Core.Const(Setfield.DynamicIndexLens{var"#10#15"}(var"#10#15"()))
│         (lens = (identity)(%11))
│         (x@_9 = (Setfield.set)(x@_9, lens::Core.Const(Setfield.DynamicIndexLens{var"#10#15"}(var"#10#15"())), 1))
│         (#11 = %new(Main.:(var"#11#16")))
│   %15 = #11::Core.Const(var"#11#16"())
│   %16 = (Setfield.DynamicIndexLens)(%15)::Core.Const(Setfield.DynamicIndexLens{var"#11#16"}(var"#11#16"()))
│   %17 = (Setfield.compose)(%16)::Core.Const(Setfield.DynamicIndexLens{var"#11#16"}(var"#11#16"()))
│         (lens = (identity)(%17))
│         (x@_9 = (Setfield.set)(x@_9, lens::Core.Const(Setfield.DynamicIndexLens{var"#11#16"}(var"#11#16"())), 1))
│         (#12 = %new(Main.:(var"#12#17")))
│   %21 = #12::Core.Const(var"#12#17"())
│   %22 = (Setfield.DynamicIndexLens)(%21)::Core.Const(Setfield.DynamicIndexLens{var"#12#17"}(var"#12#17"()))
│   %23 = (Setfield.compose)(%22)::Core.Const(Setfield.DynamicIndexLens{var"#12#17"}(var"#12#17"()))
│         (lens = (identity)(%23))
│         (x@_9 = (Setfield.set)(x@_9, lens::Core.Const(Setfield.DynamicIndexLens{var"#12#17"}(var"#12#17"())), 1))
│         (#13 = %new(Main.:(var"#13#18")))
│   %27 = #13::Core.Const(var"#13#18"())
│   %28 = (Setfield.DynamicIndexLens)(%27)::Core.Const(Setfield.DynamicIndexLens{var"#13#18"}(var"#13#18"()))
│   %29 = (Setfield.compose)(%28)::Core.Const(Setfield.DynamicIndexLens{var"#13#18"}(var"#13#18"()))
│         (lens = (identity)(%29))
│         (x@_9 = (Setfield.set)(x@_9, lens::Core.Const(Setfield.DynamicIndexLens{var"#13#18"}(var"#13#18"())), 1))
└──       return x@_9

vs

julia> f = @f(x[end] = 1)
#71#f (generic function with 1 method)

julia> @code_warntype f(rand(2))
Variables
  #self#::Core.Const(var"#71#f")
  x@_2::Vector{Float64}
  #23::var"#23#28"
  #22::var"#22#27"
  #21::var"#21#26"
  #20::var"#20#25"
  #19::var"#19#24"
  ##lens#294::Setfield.DynamicIndexLens{var"#23#28"}
  ##lens#292::Setfield.DynamicIndexLens{var"#22#27"}
  ##lens#290::Setfield.DynamicIndexLens{var"#21#26"}
  ##lens#288::Setfield.DynamicIndexLens{var"#20#25"}
  ##lens#286::Setfield.DynamicIndexLens{var"#19#24"}
  x@_13::Vector{Float64}

Body::Vector{Float64}
1 ─       (x@_13 = x@_2)
│         (#19 = %new(Main.:(var"#19#24")))
│   %3  = #19::Core.Const(var"#19#24"())
│   %4  = (Setfield.DynamicIndexLens)(%3)::Core.Const(Setfield.DynamicIndexLens{var"#19#24"}(var"#19#24"()))
│   %5  = (Setfield.compose)(%4)::Core.Const(Setfield.DynamicIndexLens{var"#19#24"}(var"#19#24"()))
│         (##lens#286 = (identity)(%5))
│         (x@_13 = (Setfield.set)(x@_13, ##lens#286, 1))
│         (#20 = %new(Main.:(var"#20#25")))
│   %9  = #20::Core.Const(var"#20#25"())
│   %10 = (Setfield.DynamicIndexLens)(%9)::Core.Const(Setfield.DynamicIndexLens{var"#20#25"}(var"#20#25"()))
│   %11 = (Setfield.compose)(%10)::Core.Const(Setfield.DynamicIndexLens{var"#20#25"}(var"#20#25"()))
│         (##lens#288 = (identity)(%11))
│         (x@_13 = (Setfield.set)(x@_13, ##lens#288, 1))
│         (#21 = %new(Main.:(var"#21#26")))
│   %15 = #21::Core.Const(var"#21#26"())
│   %16 = (Setfield.DynamicIndexLens)(%15)::Core.Const(Setfield.DynamicIndexLens{var"#21#26"}(var"#21#26"()))
│   %17 = (Setfield.compose)(%16)::Core.Const(Setfield.DynamicIndexLens{var"#21#26"}(var"#21#26"()))
│         (##lens#290 = (identity)(%17))
│         (x@_13 = (Setfield.set)(x@_13, ##lens#290, 1))
│         (#22 = %new(Main.:(var"#22#27")))
│   %21 = #22::Core.Const(var"#22#27"())
│   %22 = (Setfield.DynamicIndexLens)(%21)::Core.Const(Setfield.DynamicIndexLens{var"#22#27"}(var"#22#27"()))
│   %23 = (Setfield.compose)(%22)::Core.Const(Setfield.DynamicIndexLens{var"#22#27"}(var"#22#27"()))
│         (##lens#292 = (identity)(%23))
│         (x@_13 = (Setfield.set)(x@_13, ##lens#292, 1))
│         (#23 = %new(Main.:(var"#23#28")))
│   %27 = #23::Core.Const(var"#23#28"())
│   %28 = (Setfield.DynamicIndexLens)(%27)::Core.Const(Setfield.DynamicIndexLens{var"#23#28"}(var"#23#28"()))
│   %29 = (Setfield.compose)(%28)::Core.Const(Setfield.DynamicIndexLens{var"#23#28"}(var"#23#28"()))
│         (##lens#294 = (identity)(%29))
│         (x@_13 = (Setfield.set)(x@_13, ##lens#294, 1))
└──       return x@_13

Notice how lens is not inferrable. It doesn't lead to type-instability in the return-value of course, but it means that the compiler won't be able to make the same optimizations. Not sure how we would write a test for this though :confused:

jw3126 commented 3 years ago

Do you think there is a meaningful way we could test this? If not we can also merge as is.

torfjelde commented 3 years ago

Do you think there is a meaningful way we could test this? If not we can also merge as is.

Trying to see if there's a way to do it; a sec.

jw3126 commented 3 years ago

BTW if you have the bandwidth to do so, it would nice if you could port these fixes to Accessors.jl but don't feel obligated to do so.

torfjelde commented 3 years ago

Maybe something like

function test_all_inferrable(f, argtypes)
    typed = first(code_typed(f, argtypes))
    code = typed.first
    @test all(T -> !(T isa UnionAll || T === Any), code.slottypes)
end

?

Then on master we have

julia> f = @f(x[end] = 1)
#129#f (generic function with 1 method)

julia> test_all_inferrable(f, (Vector{Float64}, ))
Test Failed at REPL[70]:4
  Expression: all((T->begin
            !(T isa UnionAll || T === Any)
        end), code.slottypes)
ERROR: There was an error during testing

while on this branch

julia> f = @f(x[end] = 1)
#116#f (generic function with 1 method)

julia> test_all_inferrable(f, (Vector{Float64}, ))
Test Passed

Seems a bit hacky but I guess it works?

torfjelde commented 3 years ago

BTW if you have the bandwidth to do so, it would nice if you could port these fixes to Accessors.jl but don't feel obligated to do so.

Am a bit swamped atm, but I'll make an issue if one does not already exist and refrence this PR => if I get time in the future or someone else is inspired it could get done.

torfjelde commented 3 years ago

Do you want me to add the test_all_inferrable @jw3126 ?

jw3126 commented 3 years ago

Do you want me to add the test_all_inferrable @jw3126 ?

Yes that would be good I think.

torfjelde commented 3 years ago

Done :+1:

jw3126 commented 3 years ago

I just realized, that this package is still on travis ci, which is silently not working anymore :sweat_smile: #158

torfjelde commented 3 years ago

Yeah I was just looking at what it would take to add GH actions as I also just realized no CI was running, but you beat me to it:)

torfjelde commented 3 years ago

Btw, just because I haven't said so yes: all the work you've done on Setfields.jl and related is so awesome :heart: We're slowly starting to use it in more and more places within TuringLang. AdvancedHMC.jl already uses it, and now we'll likely start using it DynamicPPL.jl enabling us to do a bunch of neat stuff:) So thank you for your work!

jw3126 commented 3 years ago

Thanks a lot! I am glad lenses get used in PPL!