Moving some more functionality into the syntax tree class, including

pattern matching and traversal

syntax_tree.lua

@@ -5,6 +5,11 @@ do
 end
 local Source
 Source = require("code_obj").Source
+local List, Dict
+do
+  local _obj_0 = require('containers')
+  List, Dict = _obj_0.List, _obj_0.Dict
+end
 local Files = require('files')
 local unpack = unpack or table.unpack
 local as_lua
@@ -101,9 +106,12 @@ do
       if type(fn) == 'table' then
         local replacements = fn
         fn = function(t)
-          for k, v in pairs(replacements) do
+          if t.type == "Var" then
-            if k == t then
+            do
-              return v
+              local r = replacements[t:as_var()]
+              if r then
+                return r
+              end
             end
           end
         end
@@ -201,17 +209,8 @@ do
       return args
     end,
     get_stub = function(self)
-      if self.type == "MethodCall" then
+      local _exp_0 = self.type
-        return "0, " .. table.concat((function()
+      if "Action" == _exp_0 then
-          local _accum_0 = { }
-          local _len_0 = 1
-          for i = 2, #self do
-            _accum_0[_len_0] = self[i]:get_stub()
-            _len_0 = _len_0 + 1
-          end
-          return _accum_0
-        end)(), "; ")
-      end
         local stub_bits = { }
         local arg_i = 1
         for _index_0 = 1, #self do
@@ -227,6 +226,19 @@ do
           stub_bits[#stub_bits] = nil
         end
         return concat(stub_bits, " ")
+      elseif "MethodCall" == _exp_0 then
+        return "0, " .. table.concat((function()
+          local _accum_0 = { }
+          local _len_0 = 1
+          for i = 2, #self do
+            _accum_0[_len_0] = self[i]:get_stub()
+            _len_0 = _len_0 + 1
+          end
+          return _accum_0
+        end)(), "; ")
+      else
+        return error(tostring(self.type) .. "s do not have stubs")
+      end
     end,
     as_var = function(self)
       assert(self.type == "Var")
@@ -235,6 +247,72 @@ do
       else
         return self[1]:get_stub()
       end
+    end,
+    matching = function(self, patt)
+      if patt.type == "Var" then
+        return {
+          [patt:as_var()] = self
+        }
+      end
+      if patt:get_stub() ~= self:get_stub() then
+        return nil
+      end
+      if #self ~= #patt then
+        return nil
+      end
+      local match = { }
+      for i = 1, #self do
+        local v = self[i]
+        local pv = patt[i]
+        if type(v) ~= type(pv) then
+          return nil
+        end
+        if type(v) ~= 'table' then
+          if not (v == pv) then
+            return nil
+          end
+        else
+          local m = v:matching(pv)
+          if not (m) then
+            return nil
+          end
+          for mk, mv in pairs(m) do
+            if match[mk] and match[mk] ~= mv then
+              return nil
+            end
+            match[mk] = mv
+          end
+        end
+      end
+      return Dict(match)
+    end,
+    _breadth_first = function(self)
+      coroutine.yield(self)
+      for _index_0 = 1, #self do
+        local child = self[_index_0]
+        if getmetatable(child) == SyntaxTree.__base then
+          child:_breadth_first()
+        end
+      end
+    end,
+    breadth_first = function(self)
+      return coroutine.create(function()
+        return self:_breadth_first()
+      end)
+    end,
+    _depth_first = function(self)
+      coroutine.yield(self)
+      for _index_0 = 1, #self do
+        local child = self[_index_0]
+        if getmetatable(child) == SyntaxTree.__base then
+          child:_depth_first()
+        end
+      end
+    end,
+    depth_first = function(self)
+      return coroutine.create(function()
+        return self:_depth_first()
+      end)
     end
   }
   _base_0.__index = _base_0

syntax_tree.moon

@@ -2,6 +2,7 @@
 -- as well as the logic for converting them to Lua code.
 {:insert, :remove, :concat} = table
 {:Source} = require "code_obj"
+{:List, :Dict} = require 'containers'
 Files = require 'files'
 unpack or= table.unpack
 
@@ -14,6 +15,7 @@ as_lua = =>
     return @as_lua! if @as_lua
     error("Not supported: #{@}")
 
+local SyntaxTree
 class SyntaxTree
     __tostring: =>
         bits = [type(b) == 'string' and b\as_lua! or tostring(b) for b in *@]
@@ -56,8 +58,9 @@ class SyntaxTree
         if type(fn) == 'table'
             replacements = fn
             fn = (t)->
-                for k,v in pairs(replacements)
+                if t.type == "Var"
-                    if k == t then return v
+                    if r = replacements[t\as_var!]
+                        return r
 
         replacement = fn(@)
         if replacement == false then return nil
@@ -103,8 +106,8 @@ class SyntaxTree
         return args
 
     get_stub: =>
-        if @type == "MethodCall"
+        switch @type
-            return "0, "..table.concat([@[i]\get_stub! for i=2,#@], "; ")
+            when "Action"
                 stub_bits = {}
                 arg_i = 1
                 for a in *@
@@ -116,6 +119,10 @@ class SyntaxTree
                 while type(stub_bits[#stub_bits]) == 'number'
                     stub_bits[#stub_bits] = nil
                 return concat stub_bits, " "
+            when "MethodCall"
+                return "0, "..table.concat([@[i]\get_stub! for i=2,#@], "; ")
+            else
+                error("#{@type}s do not have stubs")
 
     as_var: =>
         assert(@type == "Var")
@@ -124,6 +131,43 @@ class SyntaxTree
         else
             return @[1]\get_stub!
 
+    matching: (patt)=>
+        if patt.type == "Var"
+            return {[patt\as_var!]:@}
+        return nil if patt\get_stub! != @get_stub!
+        -- TODO: support vararg matches like (\(say 1 2 3), matching \(say *$values))
+        return nil if #@ != #patt
+        match = {}
+        for i=1,#@
+            v = @[i]
+            pv = patt[i]
+            return nil if type(v) != type(pv)
+            if type(v) != 'table'
+                return nil unless v == pv
+            else
+                m = v\matching(pv)
+                return nil unless m
+                for mk,mv in pairs(m)
+                    return nil if match[mk] and match[mk] != mv
+                    match[mk] = mv
+        return Dict(match)
+
+    _breadth_first: =>
+        coroutine.yield @
+        for child in *@
+            if getmetatable(child) == SyntaxTree.__base
+                child\_breadth_first!
+        return
+    breadth_first: => coroutine.create(-> @_breadth_first!)
+
+    _depth_first: =>
+        coroutine.yield @
+        for child in *@
+            if getmetatable(child) == SyntaxTree.__base
+                child\_depth_first!
+        return
+    depth_first: => coroutine.create(-> @_depth_first!)
+
     @is_instance: (t)=>
         type(t) == 'table' and getmetatable(t) == @__base