Merge pull request #343 from peterahrens/pja/capture-macro

Author: shashi

src/SymbolicUtils.jl

@@ -35,7 +35,7 @@ include("rewriters.jl")
 using .Rewriters
 
 using Combinatorics: permutations, combinations
-export @rule, @acrule, RuleSet
+export @rule, @acrule, RuleSet, @capture
 
 # Rule type and @rule macro
 include("rule.jl")

src/rule.jl

@@ -316,6 +316,49 @@ macro rule(expr)
     end
 end
 
+"""
+    @capture ex pattern
+
+Uses a `Rule` object to capture an expression if it matches the `pattern`. Returns `true` and injects
+slot variable match results into the calling scope when the `pattern` matches, otherwise returns false. The
+rule language for specifying the `pattern` is the same in @capture as it is in `@rule`. Contextual matching
+is not yet supported
+
+```julia
+julia> @syms a; ex = a^a;
+
+julia> if @capture ex (~x)^(~x)
+           @show x
+       elseif @capture ex 2(~y)
+           @show y
+       end;
+x = a
+```
+
+See also: [`@rule`](@ref)
+"""
+macro capture(ex, lhs)
+    keys = Symbol[]
+    lhs_term = makepattern(lhs, keys)
+    unique!(keys)
+    bind = Expr(:block, map(key-> :($(esc(key)) = getindex(__MATCHES__, $(QuoteNode(key)))), keys)...)
+    quote
+        $(__source__)
+        lhs_pattern = $(lhs_term)
+        __MATCHES__ = Rule($(QuoteNode(lhs)),
+             lhs_pattern,
+             matcher(lhs_pattern),
+             identity,
+             rule_depth($lhs_term))($(esc(ex)))
+        if __MATCHES__ !== nothing
+            $bind
+            true
+        else
+            false
+        end
+    end
+end
+
 #-----------------------------
 #### Associative Commutative Rules
 

test/rewrite.jl

@@ -42,3 +42,34 @@ end
     @eqtest @rule(+(~~x,~y, ~~x) => (~~x, ~y, ~~x))(term(+,9,8,9,9,8,type=Any)) == ([9,8], 9, [9,8])
     @eqtest @rule(+(~~x,~y,~~x) => (~~x, ~y, ~~x))(term(+,6,type=Any)) == ([], 6, [])
 end
+
+@testset "Capture form" begin
+    ex = a^a
+
+    #note that @test inserts a soft local scope (try-catch) that would gobble
+    #the matches from assignment statements in @capture macro, so we call it
+    #outside the test macro 
+    ret = @capture ex (~x)^(~x)
+    @test ret
+    @test @isdefined x
+    @test x === a
+
+    ex = b^a
+    ret = @capture ex (~y)^(~y)
+    @test !ret
+    @test !(@isdefined y)
+
+    ret = @capture (a + b) (+)(~~z)
+    @test ret
+    @test @isdefined z
+    @test all(z .=== arguments(a + b))
+
+    #a more typical way to use the @capture macro
+
+    f(x) = if @capture x (~w)^(~w)
+        w
+    end
+
+    @eqtest f(b^b) == b
+    @test f(b+b) == nothing
+end