support broadcasting (#137)

* support broadcasting in optics
* remove _UpdateOp
* support @set with broadcasting

Author: aplavin

src/functionlenses.jl

@@ -89,6 +89,10 @@ end
 
 set(obj, o::Base.Fix1{typeof(map)}, val) = map((ob, v) -> set(ob, o.x, v), obj, val)
 
+set(obj, o::Base.BroadcastFunction, val) = set.(obj, Ref(o.f), val)
+set(obj, o::Base.Fix1{<:Base.BroadcastFunction}, val) = set.(obj, Base.Fix1.(Ref(o.f.f), o.x), val)
+set(obj, o::Base.Fix2{<:Base.BroadcastFunction}, val) = set.(obj, Base.Fix2.(Ref(o.f.f), o.x), val)
+
 set(obj, o::Base.Fix1{typeof(filter)}, val) = @set obj[findall(o.x, obj)] = val
 modify(f, obj, o::Base.Fix1{typeof(filter)}) = @modify(f, obj[findall(o.x, obj)])
 delete(obj, o::Base.Fix1{typeof(filter)}) = filter(!o.x, obj)

src/sugar.jl

@@ -157,6 +157,21 @@ function lower_index(collection::Symbol, index, dim)
     )
 end
 
+_secondarg(_, x) = x
+
+_esc_and_dot_name_to_broadcasted(f) = esc(f)
+_esc_and_dot_name_to_broadcasted(f::Symbol) =
+    if f == :.
+        # eg, in @set a[:] .= 1
+        # the returned function will be called as func(a, 1)
+        :(Base.BroadcastFunction($_secondarg))
+    elseif startswith(string(f), '.')
+        # eg, in @set a[:] .+= 1 or @optic _ .+ 1
+        :(Base.BroadcastFunction($(esc(Symbol(string(f)[2:end])))))
+    else
+        esc(f)
+    end
+
 function parse_obj_optics(ex)
     dollar_exprs = foldtree([], ex) do exs, x
         x isa Expr && x.head == :$ ?
@@ -206,9 +221,15 @@ function parse_obj_optics(ex)
         obj, frontoptic = parse_obj_optics(front)
         optic = :($PropertyLens{$(QuoteNode(property))}())
     elseif @capture(ex, f_(front_))
+        # regular function call
         obj, frontoptic = parse_obj_optics(front)
-        optic = esc(f) # function optic
+        optic = _esc_and_dot_name_to_broadcasted(f)  # broadcasted operators like .- fall here
+    elseif @capture(ex, f_.(front_))
+        # broadcasted function call (not operator)
+        obj, frontoptic = parse_obj_optics(front)
+        optic = :(Base.BroadcastFunction($(esc(f))))
     elseif @capture(ex, f_(args__))
+        # function call with multiple arguments
         args_contain_under = map(args) do arg
             foldtree((yes, x) -> yes || x === :_, false, arg)
         end
@@ -219,12 +240,13 @@ function parse_obj_optics(ex)
         end
         length(args) == 2 || error("Only 1- and 2-argument functions are supported")
         sum(args_contain_under) == 1 || error("Only a single function argument can be the optic target")
+        f = _esc_and_dot_name_to_broadcasted(f)  # multi-arg broadcasted fall here, no matter if regular function or operator
         if args_contain_under[1]
             obj, frontoptic = parse_obj_optics(args[1])
-            optic = :(Base.Fix2($(esc(f)), $(esc(args[2]))))
+            optic = :(Base.Fix2($f, $(esc(args[2]))))
         elseif args_contain_under[2]
             obj, frontoptic = parse_obj_optics(args[2])
-            optic = :(Base.Fix1($(esc(f)), $(esc(args[1]))))
+            optic = :(Base.Fix1($f, $(esc(args[1]))))
         end
     else
         obj = esc(ex)
@@ -261,12 +283,6 @@ function get_update_op(sym::Symbol)
     Symbol(s[1:end-1])
 end
 
-struct _UpdateOp{OP,V}
-    op::OP
-    val::V
-end
-(u::_UpdateOp)(x) = u.op(x, u.val)
-
 """
     setmacro(optictransform, ex::Expr; overwrite::Bool=false)
 
@@ -300,7 +316,7 @@ function setmacro(optictransform, ex::Expr; overwrite::Bool=false)
         :($set($obj, ($optictransform)($optic), $val))
     else
         op = get_update_op(ex.head)
-        f = :($_UpdateOp($op,$val))
+        f = :($Base.Fix2($(_esc_and_dot_name_to_broadcasted(op)), $val))
         :($modify($f, $obj, ($optictransform)($optic)))
     end
     return _macro_expression_result(obj, ret; overwrite=overwrite)

test/test_functionlenses.jl

@@ -198,6 +198,29 @@ end
     test_getset_laws(diag, [1 2 3; 4 5 6], [1., 2.5], [0, 1])
 end
 
+@testset "broadcast" begin
+    # in optic definiton
+    @test (@optic exp.(_)) === Base.BroadcastFunction(exp)
+    @test (@optic .-_) === Base.BroadcastFunction(-)
+    @test (@optic 1 .+ _) === Base.Fix1(Base.BroadcastFunction(+), 1)
+    A = [[1,2], [1,2,3], [1,2,3,4]]
+    @test (@set first.(A) = 10:12) == [[10, 2], [11, 2, 3], [12, 2, 3, 4]]
+
+    test_getset_laws((@optic first.(_)), [[1,2], [1,2,3]], [3., 4.], [5, 6])
+    test_getset_laws((@optic _ .- 1), [1, 2], [3., 4.], [5, 6])
+    test_getset_laws((@optic 1 .- _), [1, 2], [3., 4.], [5, 6])
+    test_getset_laws((@optic [10, 20] .* first.(_)), [[1,2], [1,2,3]], [3., 4.], [5, 6])
+
+    # in @set update operation
+    @test (@set A .= 1) == [1, 1, 1]
+    @test (@set A[1] .= 1) == [[1,1], [1,2,3], [1,2,3,4]]
+    @test (@set A[1] .= [10, 20]) == [[10,20], [1,2,3], [1,2,3,4]]
+
+    # broadcasting in both optic and set, with a user-defined function
+    @accessor f(x) = x+1
+    @test (@set f.(first.(A)) .= [10, 20, 30]) == [[9,2], [19,2,3], [29,2,3,4]]
+end
+
 @testset "math" begin
     @test 2.0       === @set real(1) = 2.0
     @test 2.0 + 1im === @set real(1+1im) = 2.0