Merge pull request #406 from JuliaSymbolics/s/no-bindings

WIP: create_bindings in DestructuredArgs

Author: shashi

Project.toml

@@ -1,7 +1,7 @@
 name = "SymbolicUtils"
 uuid = "d1185830-fcd6-423d-90d6-eec64667417b"
 authors = ["Shashi Gowda"]
-version = "0.18.2"
+version = "0.19"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/code.jl

@@ -15,10 +15,10 @@ import SymbolicUtils: @matchable, Sym, Term, istree, operation, arguments,
 ##== state management ==##
 
 struct NameState
-    symbolify::Dict{Any, Symbol}
+    rewrites::Dict{Any, Any}
 end
-NameState() = NameState(Dict{Any, Symbol}())
-function union_symbolify!(n, ts)
+NameState() = NameState(Dict{Any, Any}())
+function union_rewrites!(n, ts)
     for t in ts
         n[t] = Symbol(string(t))
     end
@@ -34,7 +34,7 @@ function Base.get(st::LazyState)
     s === nothing ? getfield(st, :ref)[] = NameState() : s
 end
 
-@inline Base.getproperty(st::LazyState, f::Symbol) = getproperty(get(st), f)
+@inline Base.getproperty(st::LazyState, f::Symbol) = f==:symbolify ?  getproperty(st, :rewrites) : getproperty(get(st), f)
 
 ##========================##
 
@@ -75,7 +75,10 @@ when `y(t)` is itself the argument of a function rather than `y`.
 
 """
 toexpr(x) = toexpr(x, LazyState())
-toexpr(s::Sym, st) = nameof(s)
+function toexpr(s::Sym, st)
+    s′ = substitute_name(s, st)
+    s′ isa Sym ? nameof(s′) : toexpr(s′, st)
+end
 
 
 @matchable struct Assignment
@@ -102,7 +105,7 @@ toexpr(a::Assignment, st) = :($(toexpr(a.lhs, st)) = $(toexpr(a.rhs, st)))
 function_to_expr(op, args, st) = nothing
 
 function function_to_expr(op::Union{typeof(*),typeof(+)}, O, st)
-    out = get(st.symbolify, O, nothing)
+    out = get(st.rewrites, O, nothing)
     out === nothing || return out
     args = map(Base.Fix2(toexpr, st), arguments(O))
     if length(args) >= 3 && symtype(O) <: Number
@@ -135,18 +138,27 @@ function function_to_expr(::typeof(SymbolicUtils.ifelse), O, st)
     :($(toexpr(args[1], st)) ? $(toexpr(args[2], st)) : $(toexpr(args[3], st)))
 end
 
-function_to_expr(::Sym, O, st) = get(st.symbolify, O, nothing)
+function_to_expr(::Sym, O, st) = get(st.rewrites, O, nothing)
 
 toexpr(O::Expr, st) = O
 
+function substitute_name(O, st)
+    if (issym(O) || istree(O)) && haskey(st.rewrites, O)
+        st.rewrites[O]
+    else
+        O
+    end
+end
+
 function toexpr(O, st)
+    O = substitute_name(O, st)
     !istree(O) && return O
     op = operation(O)
     expr′ = function_to_expr(op, O, st)
     if expr′ !== nothing
         return expr′
     else
-        haskey(st.symbolify, O) && return st.symbolify[O]
+        !istree(O) && return O
         args = arguments(O)
         return Expr(:call, toexpr(op, st), map(x->toexpr(x, st), args)...)
     end
@@ -158,10 +170,15 @@ end
     inds
     name
     inbounds::Bool
+    create_bindings::Bool
 end
 
-function DestructuredArgs(elems, name=gensym("arg"); inds=eachindex(elems), inbounds=false)
-    DestructuredArgs(elems, inds, name, inbounds)
+function DestructuredArgs(elems, name=nothing; inds=eachindex(elems), inbounds=false, create_bindings=true)
+    if name === nothing
+        # I'm sorry if you get a hash collision here lol
+        name = Symbol("##arg#", hash((elems, inds, inbounds, create_bindings)))
+    end
+    DestructuredArgs(elems, inds, name, inbounds, create_bindings)
 end
 
 """
@@ -176,76 +193,105 @@ components. See example in `Func` for more information.
 DestructuredArgs
 
 toexpr(x::DestructuredArgs, st) = toexpr(x.name, st)
-get_symbolify(args::DestructuredArgs) = ()
-function get_symbolify(args::Union{AbstractArray, Tuple})
-    cflatten(map(get_symbolify, args))
+get_rewrites(args::DestructuredArgs) = ()
+function get_rewrites(args::Union{AbstractArray, Tuple})
+    cflatten(map(get_rewrites, args))
 end
-get_symbolify(x) = istree(x) ? (x,) : ()
+get_rewrites(x) = istree(x) ? (x,) : ()
 cflatten(x) = Iterators.flatten(x) |> collect
 
+# Used in Symbolics
+Base.@deprecate_binding get_symbolify get_rewrites
+
 function get_assignments(d::DestructuredArgs, st)
     name = toexpr(d, st)
     map(d.inds, d.elems) do i, a
         ex = (i isa Symbol ? :($name.$i) : :($name[$i]))
-        ex = d.inbounds ? :(@inbounds($ex)) : ex
+        ex = d.inbounds && d.create_bindings ? :(@inbounds($ex)) : ex
         a ← ex
     end
 end
 
 @matchable struct Let
     pairs::Vector{Union{Assignment,DestructuredArgs}} # an iterator of pairs, ordered
     body
+    let_block::Bool
 end
 
 """
-    Let(assignments, body)
+    Let(assignments, body[, let_block])
 
 A Let block.
 
 - `assignments` is a vector of `Assignment`s
 - `body` is the body of the let block
+- `let_block` boolean (default=true) -- do not create a let block if false.
 """
-Let
+Let(assignments, body) = Let(assignments, body, true)
 
 function toexpr(l::Let, st)
     if all(x->x isa Assignment && !(x.lhs isa DestructuredArgs), l.pairs)
         dargs = l.pairs
     else
-        dargs = map(l.pairs) do x
+        assignments = []
+        for x in l.pairs
             if x isa DestructuredArgs
-                get_assignments(x, st)
+                if x.create_bindings
+                    append!(assignments, get_assignments(x, st))
+                else
+                    for a in get_assignments(x, st)
+                        st.rewrites[a.lhs] = a.rhs
+                    end
+                end
             elseif x isa Assignment && x.lhs isa DestructuredArgs
-                [x.lhs.name ← x.rhs, get_assignments(x.lhs, st)...]
+                if x.lhs.create_bindings
+                    push!(assignments, x.lhs.name ← x.rhs)
+                    append!(assignments, get_assignments(x.lhs, st))
+                else
+                    push!(assignments, x.lhs.name ← x.rhs)
+                    for a in get_assignments(x.lhs, st)
+                        st.rewrites[a.lhs] = a.rhs
+                    end
+                end
             else
-                (x,)
+                push!(assignments, x)
             end
-        end |> cflatten
+        end
         # expand and come back
-        return toexpr(Let(dargs, l.body), st)
+        return toexpr(Let(assignments, l.body, l.let_block), st)
     end
 
-    funkyargs = get_symbolify(map(lhs, dargs))
-    union_symbolify!(st.symbolify, funkyargs)
+    funkyargs = get_rewrites(map(lhs, dargs))
+    union_rewrites!(st.rewrites, funkyargs)
 
-    Expr(:let,
-         Expr(:block, map(p->toexpr(p, st), dargs)...),
-         toexpr(l.body, st))
+    bindings = map(p->toexpr(p, st), dargs)
+    l.let_block ? Expr(:let,
+                       Expr(:block, bindings...),
+                       toexpr(l.body, st)) : Expr(:block,
+                                                  bindings...,
+                                                  toexpr(l.body, st))
 end
 
 @matchable struct Func
     args::Vector
     kwargs
     body
+    pre::Vector
 end
 
+Func(args, kwargs, body) = Func(args, kwargs, body, [])
+
 """
-    Func(args, kwargs, body)
+    Func(args, kwargs, body[, pre])
 
 A function.
 
 - `args` is a vector of expressions
 - `kwargs` is a vector of `Assignment`s
 - `body` is the body of the function
+- `pre` a vector of expressions to be prepended to the function body,
+   for example, it could be `[Expr(:meta, :inline), Expr(:meta, :propagate_inbounds)]`
+   to create an `@inline @propagate_inbounds` function definition.
 
 **Special features in `args`**:
 
@@ -291,21 +337,23 @@ Func
 toexpr_kw(f, st) = Expr(:kw, toexpr(f, st).args...)
 
 function toexpr(f::Func, st)
-    funkyargs = get_symbolify(vcat(f.args, map(lhs, f.kwargs)))
-    union_symbolify!(st.symbolify, funkyargs)
+    funkyargs = get_rewrites(vcat(f.args, map(lhs, f.kwargs)))
+    union_rewrites!(st.rewrites, funkyargs)
     dargs = filter(x->x isa DestructuredArgs, f.args)
     if !isempty(dargs)
-        body = Let(dargs, f.body)
+        body = Let(dargs, f.body, false)
     else
         body = f.body
     end
     if isempty(f.kwargs)
         :(function ($(map(x->toexpr(x, st), f.args)...),)
+              $(f.pre...)
               $(toexpr(body, st))
           end)
     else
         :(function ($(map(x->toexpr(x, st), f.args)...),;
                     $(map(x->toexpr_kw(x, st), f.kwargs)...))
+              $(f.pre...)
               $(toexpr(body, st))
           end)
     end

test/code.jl

@@ -22,7 +22,7 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
     @test toexpr(x(t)+y(t)) == :($(+)(x(t), y(t)))
     @test toexpr(x(t)+y(t)+x(t+1)) == :($(+)($(+)(x(t), y(t)), x($(+)(1, t))))
     s = LazyState()
-    Code.union_symbolify!(s.symbolify, [x(t), y(t)])
+    Code.union_rewrites!(s.rewrites, [x(t), y(t)])
     @test toexpr(x(t)+y(t)+x(t+1), s) == :($(+)($(+)(var"x(t)", var"y(t)"), x($(+)(1, t))))
 
     ex = :(let a = 3, b = $(+)(1,a)
@@ -43,13 +43,34 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
                            DestructuredArgs((a, b), :params)], [],
                           x(t+1) + x(t) + a  + b)),
               :(function (state, params)
-                    let x = state[1], var"x(t)" = state[2], a = params[1], b = params[2]
+                    begin
+                        x = state[1]
+                        var"x(t)" = state[2]
+                        a = params[1]
+                        b = params[2]
                         $(+)($(+)($(+)(a, b), var"x(t)"), x($(+)(1, t)))
                     end
                 end))
 
+    test_repr(toexpr(Func([DestructuredArgs([x, x(t)], :state, create_bindings=false),
+                           DestructuredArgs((a, b), :params, create_bindings=false)], [],
+                          x(t+1) + x(t) + a  + b)),
+              :(function (state, params)
+                    begin
+                        $(+)($(+)($(+)(params[1], params[2]), $getindex(state, 2)), state[1]($(+)(1, t)))
+                    end
+                end))
+
+
+    test_repr(toexpr(Func([],[],:(rand()), [Expr(:meta, :inline)])),
+              :(function ()
+                    $(Expr(:meta, :inline))
+                    rand()
+                end))
+
     ex = toexpr(Func([DestructuredArgs([x, x(t)], :state, inbounds=true)], [], x(t+1) + x(t)))
-    for e ∈ ex.args[2].args[3].args[1].args
+    ex = Base.remove_linenums!(ex)
+    for e ∈ ex.args[2].args[1].args[1:2]
         @test e.args[2].head == :macrocall
     end
 
@@ -89,7 +110,9 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
     test_repr(toexpr(Func([DestructuredArgs([a,b],c,inds=[:a, :b])], [],
                           a + b)),
               :(function (c,)
-                    let a = c.a, b = c.b
+                    begin
+                        a = c.a
+                        b = c.b
                         $(+)(a, b)
                     end
                 end))