Inference and @batchlens

[jw3126]

using Kaleido
function doit()
    @batchlens begin
        _.b.c
    end
end
l = doit()
obj = (a=1, b=(c=2,))
@show get(obj, l)
@code_warntype doit()

does not infer:

get(obj, l) = (2,)
Variables
  #self#::Core.Compiler.Const(doit, false)

Body::Setfield.ComposedLens{_A,FlatLens{(0x0000000000000001,)}} where _A
1 ─ %1 = Core.apply_type(Setfield.PropertyLens, :b)::Core.Compiler.Const(Setfield.PropertyLens{:b}, false)
│   %2 = (%1)()::Core.Compiler.Const((@lens _.b), false)
│   %3 = Core.apply_type(Setfield.PropertyLens, :c)::Core.Compiler.Const(Setfield.PropertyLens{:c}, false)
│   %4 = (%3)()::Core.Compiler.Const((@lens _.c), false)
│   %5 = (Setfield.compose)(%2, %4)::Core.Compiler.Const((@lens _.b.c), false)
│   %6 = (identity)(%5)::Core.Compiler.Const((@lens _.b.c), false)
│   %7 = (Kaleido.batch)(%6)::Setfield.ComposedLens{_A,FlatLens{(0x0000000000000001,)}} where _A
└──      return %7

Interestingly the more complicated

function doit()
    @batchlens begin
        _.b.c
        _.a
    end
end

does infer.

[tkf]

OK, this is interesting. Thanks for the analysis. Yes, @batchlens needs more care. I use MultiLens most of the time (especially in performance-sensitive code) as there is complicated code behind batchlens in order to reduce calls to the constructors...

[jw3126]

I wanted to implement batch @set stealing your @batchlens work. Sadly it does almost never infer.

[tkf]

Maybe I should bite the bullet and use @generated...

[jw3126]

BTW here is the @batchset I am toying with. It is relatively easy to implement on top of your work.

using Kaleido, Setfield, MacroTools
using Setfield: parse_obj_lens
using Setfield: get_update_op
using Setfield: _UpdateOp

struct _Constant{T}
    value::T
end
(f::_Constant)(x) = f.value


function compute_f_modify(ex::Expr, val)
    @assert ex.head isa Symbol
    if ex.head == :(=)
        return :($_Constant($val))
    else
        op = get_update_op(ex.head)
        return :($_UpdateOp($op, $val))
    end
end
apply(f, arg) = f(arg)

function batchsetmacro(obj, body)
    obj = esc(obj)
    @assert Meta.isexpr(body, :block)
    body = MacroTools.striplines(body)
    
    lenses = map(body.args) do ex
        ref, val = ex.args
        _, lens = parse_obj_lens(ref)
        lens
    end
    
    fs = map(body.args) do ex
        ref, val = ex.args
        compute_f_modify(ex, val)
    end
    fs = Expr(:tuple, fs...)
    
    quote
        lens = $batch($(lenses...))
        $modify($obj, lens) do vals
            map($apply, $fs, vals)
        end
    end
    
    # # This macro does not infer well. 
    # # One reason is, that batchlens has inference problems
    # # another reason might be that Setfield.modify
    # # causes trouble, since it needs to be passed a closure
    # # to circumvent the latter, the code below "inlines" modify
    #
    # new_vals = map(enumerate(body.args)) do (i,ex)
    #     ref, val = ex.args
    #     if ex.head == :(=)
    #         return val
    #     else
    #         op = get_update_op(ex.head)
    #         return :($op(old[$i], $val))
    #     end
    # end
    # new = Expr(:tuple, new_vals...)
    # quote
    #     lens = $(batch)($(lenses...))
    #     old = $get($obj, lens)
    #     new = $new
    #     $set($obj, lens, new)
    # end
end

macro batchset(obj, body)
    batchsetmacro(obj, body)
end

obj = (a=1, b=(c=1, d=2))
two = 2.0

@batchset obj begin
    _.a = two
    _.b.c += 19
    _.b.d -= 22
end