Merge pull request #68 from JuliaSymbolics/s/ctx

Pass around a context object from simplify to matchers and rewriters

Author: shashi

src/matchers.jl

@@ -22,14 +22,23 @@ struct Segment{F}
     predicate::F
 end
 
-ismatch(s::Segment, t) = s.predicate(t)
-
 Segment(s) = Segment(s, alwaystrue)
 
 Base.show(io::IO, s::Segment) = (print(io, "~~"); print(io, s.name))
 
 makesegment(s::Symbol, keys) = (push!(keys, s); Segment(s))
 
+"""
+A wrapper for slot and segment predicates which allows them to
+take two arguments: the value and a Context
+"""
+struct Contextual{F}
+    f::F
+end
+(c::Contextual)(args...) = c.f(args...)
+
+ctxcall(f, x, ctx) = f isa Contextual ? f(x, ctx) : f(x)
+
 function makesegment(s::Expr, keys)
     if !(s.head == :(::))
         error("Syntax for specifying a segment is ~~x::\$predicate, where predicate is a boolean function")
@@ -91,6 +100,16 @@ function makeconsequent(expr)
                 return Expr(:call, map(makeconsequent, expr.args)...)
             end
         else
+            if expr.head == :macrocall
+                if expr.args[1] === Symbol("@ctx")
+                    if length(filter(x->!(x isa LineNumberNode), expr.args)) != 1
+                        error("@ctx takes no arguments. try (@ctx)")
+                    end
+                    return :__CTX__
+                else
+                    return esc(expr)
+                end
+            end
             return Expr(expr.head, map(makeconsequent, expr.args)...)
         end
     else
@@ -106,21 +125,21 @@ end
 # 3. Callback: takes arguments Dictionary × Number of elements matched
 #
 function matcher(val::Any)
-    function literal_matcher(data, bindings, next)
+    function literal_matcher(next, data, bindings, ctx)
         !isempty(data) && isequal(car(data), val) ? next(bindings, 1) : nothing
     end
 end
 
 function matcher(slot::Slot)
-    function slot_matcher(data, bindings, next)
+    function slot_matcher(next, data, bindings, ctx)
         isempty(data) && return
         val = get(bindings, slot.name, nothing)
         if val !== nothing
             if isequal(val, car(data))
                 return next(bindings, 1)
             end
         else
-            if slot.predicate(car(data))
+            if ctxcall(slot.predicate, car(data), ctx)
                 next(assoc(bindings, slot.name, car(data)), 1)
             end
         end
@@ -156,7 +175,7 @@ function trymatchexpr(data, value, n)
 end
 
 function matcher(segment::Segment)
-    function segment_matcher(data, bindings, success)
+    function segment_matcher(success, data, bindings, ctx)
         val = get(bindings, segment.name, nothing)
 
         if val !== nothing
@@ -170,7 +189,7 @@ function matcher(segment::Segment)
             for i=length(data):-1:0
                 subexpr = take_n(data, i)
 
-                if segment.predicate(subexpr)
+                if ctxcall(segment.predicate, subexpr, ctx)
                     res = success(assoc(bindings, segment.name, subexpr), i)
                     if res !== nothing
                         break
@@ -185,7 +204,7 @@ end
 
 function matcher(term::Term)
     matchers = (matcher(operation(term)), map(matcher, arguments(term))...,)
-    function term_matcher(data, bindings, success)
+    function term_matcher(success, data, bindings, ctx)
 
         isempty(data) && return nothing
         !(car(data) isa Term) && return nothing
@@ -197,8 +216,9 @@ function matcher(term::Term)
                 end
                 return nothing
             end
-            res = car(matchers′)(term, bindings′,
-                                 (b, n) -> loop(drop_n(term, n), b, cdr(matchers′)))
+            car(matchers′)(term, bindings′, ctx) do b, n
+                loop(drop_n(term, n), b, cdr(matchers′))
+            end
         end
 
         loop(car(data), bindings, matchers) # Try to eat exactly one term

src/rule_dsl.jl

@@ -28,14 +28,15 @@ function Base.show(io::IO, r::Rule)
 end
 
 const EMPTY_DICT = ImmutableDict{Symbol, Any}(:____, nothing)
+struct EmptyCtx end
 
-function (r::Rule)(term)
-    match_function = r.matcher
+function (r::Rule)(term, ctx=EmptyCtx())
     rhs = r.rhs
 
-    match_function((term,),
-                   EMPTY_DICT,
-                   (d, n) -> n === 1 ? (@timer "RHS" rhs(d)) : nothing)
+    r.matcher((term,), EMPTY_DICT, ctx) do bindings, n
+        # n == 1 means that exactly one term of the input (term,) was matched
+        n === 1 ? (@timer "RHS" rhs(bindings, ctx)) : nothing
+    end
 end
 
 """
@@ -152,7 +153,7 @@ macro rule(expr)
         Rule($(QuoteNode(expr)),
              lhs_pattern,
              matcher(lhs_pattern),
-             __MATCHES__ -> $(makeconsequent(rhs)),
+             (__MATCHES__, __CTX__) -> $(makeconsequent(rhs)),
              rule_depth($lhs_term))
     end
 end
@@ -177,7 +178,7 @@ end
 
 Base.show(io::IO, acr::ACRule) = print(io, "ACRule(", acr.rule, ")")
 
-function (acr::ACRule)(term)
+function (acr::ACRule)(term, ctx=EmptyCtx())
     r = Rule(acr)
     if !(term isa Term)
         r(term)
@@ -187,12 +188,12 @@ function (acr::ACRule)(term)
         if f != operation(r.lhs) # Maybe offer a fallback if m.term errors. 
             return nothing
         end
-        
+
         T = symtype(term)
         args = arguments(term)
-        
+
         for inds in permutations(eachindex(args), acr.arity)
-            result = r(Term{T}(f, args[inds]))
+            result = r(Term{T}(f, args[inds]), ctx)
             if !isnothing(result)
                 return Term{T}(f, [result, (args[i] for i in eachindex(args) if i ∉ inds)...])
             end
@@ -233,7 +234,7 @@ struct RuleRewriteError
     expr
 end
 
-function (r::RuleSet)(term; depth=typemax(Int), applyall=false, recurse=true)
+function (r::RuleSet)(term, context=EmptyCtx(); depth=typemax(Int), applyall=false, recurse=true)
     rules = r.rules
     term = to_symbolic(term)
     # simplify the subexpressions
@@ -243,21 +244,21 @@ function (r::RuleSet)(term; depth=typemax(Int), applyall=false, recurse=true)
     if term isa Symbolic
         if term isa Term && recurse
             expr = Term{symtype(term)}(operation(term),
-                                       map(t -> r(t, depth=depth-1), arguments(term)))
+                                       map(t -> r(t, context, depth=depth-1), arguments(term)))
         else
             expr = term
         end
         for i in 1:length(rules)
             expr′ = try
-                @timer(repr(rules[i]), rules[i](expr))
+                @timer(repr(rules[i]), rules[i](expr, context))
             catch err
                 throw(RuleRewriteError(rules[i], expr))
             end
             if expr′ === nothing
                 # this rule doesn't apply
                 continue
             else
-                expr = r(expr′, depth=getdepth(rules[i]))# levels touched
+                expr = r(expr′, context, depth=getdepth(rules[i]))# levels touched
                 applyall || return expr
             end
         end
@@ -269,17 +270,15 @@ end
 
 getdepth(::RuleSet) = typemax(Int)
 
-function fixpoint(f, x; kwargs...)
-    x1 = f(x; kwargs...)
+function fixpoint(f, x, ctx; kwargs...)
+    x1 = f(x, ctx; kwargs...)
     while !isequal(x1, x)
         x = x1
-        x1 = f(x; kwargs...)
+        x1 = f(x, ctx; kwargs...)
     end
     return x1
 end
 
-fixpoint(f; kwargs...) = x -> fixpoint(f, x; kwargs...)
-
 @noinline function Base.showerror(io::IO, err::RuleRewriteError)
     msg = "Failed to apply rule $(err.rule) on expression "
     msg *= sprint(io->showraw(io, err.expr))

src/simplify.jl

@@ -11,14 +11,17 @@ Applies them once if `fixpoint=false`.
 The `applyall` and `recurse` keywords are forwarded to the enclosed
 `RuleSet`.
 """
-function simplify(x, rules=SIMPLIFY_RULES; fixpoint=true, applyall=true, recurse=true)
+function simplify(x, ctx=EmptyCtx(); rules=SIMPLIFY_RULES, fixpoint=true, applyall=true, recurse=true)
     if fixpoint
-        SymbolicUtils.fixpoint(rules; recurse=recurse, applyall=recurse)(x)
+        SymbolicUtils.fixpoint(rules, x, ctx; recurse=recurse, applyall=recurse)
     else
-        rules(x; recurse=recurse, applyall=recurse)
+        rules(x, ctx; recurse=recurse, applyall=recurse)
     end
 end
 
+
+Base.@deprecate simplify(x, rules::RuleSet; kwargs...)  simplify(x, rules=rules; kwargs...)
+
 """
     substitute(expr, dict)
 

test/basics.jl

@@ -1,4 +1,4 @@
-using SymbolicUtils: Sym, FnType, Term, symtype
+using SymbolicUtils: Sym, FnType, Term, symtype, Contextual, EmptyCtx
 using SymbolicUtils
 using Test
 
@@ -68,6 +68,19 @@ end
     @test_throws ErrorException t(2)
 end
 
+@testset "Contexts" begin
+    @syms a b c
+
+    @test @rule(~x::Contextual((x, ctx) -> ctx==EmptyCtx()) => "yes")(1) == "yes"
+    @test @rule(~x::Contextual((x, ctx) -> haskey(ctx, x)) => true)(a, Dict(a=>1))
+    @test @rule(~x::Contextual((x, ctx) -> haskey(ctx, x)) => true)(b, Dict(a=>1)) === nothing
+    @test_throws UndefVarError @rule(~x => __CTX__)(a, "test")
+    @test @rule(~x => @ctx)(a, "test") == "test"
+    @test @rule(~x::Contextual((x, ctx) -> haskey(ctx, x)) => (@ctx)[~x])(a, Dict(a=>1)) === 1
+    @test @rule(~x::Contextual((x, ctx) -> haskey(ctx, x)) => (@ctx)[~x])(b, Dict(a=>1)) === nothing
+    @test simplify(a+a, "test", rules=RuleSet([@rule ~x => @ctx])) == "test"
+end
+
 @testset "substitute" begin
     @syms a b
     @test substitute(a, Dict(a=>1)) == 1

test/rulesets.jl

@@ -15,7 +15,7 @@ using SymbolicUtils: fixpoint, getdepth
     @test rset(ex) == (((2 * w) + (2 * w)) + (2 * α)) + (2 * β)
     @test rset(ex) == simplify(ex, rset; fixpoint=false, applyall=false) 
 
-    @test fixpoint(rset, ex) == ((2 * (2 * w)) + (2 * α)) + (2 * β)
+    @test fixpoint(rset, ex, "ctx") == ((2 * (2 * w)) + (2 * α)) + (2 * β)
 end
 
 @testset "Numeric" begin