Cleanup of replaced vars, now uses tree_map()

2 files changed, 86 insertions, 61 deletions

diff --git a/nomsu.lua b/nomsu.lua
index 096bf9e..90188ad 100644
--- a/nomsu.lua
+++ b/nomsu.lua
@@ -864,9 +864,6 @@ do
       if not tree.type then
         error("Invalid tree: " .. tostring(repr(tree)), 0)
       end
-      if not (self.tree_metadata[tree]) then
-        error("??? tree: " .. tostring(repr(tree)), 0)
-      end
       local _exp_0 = tree.type
       if "File" == _exp_0 then
         if #tree.value == 1 then
@@ -1240,19 +1237,27 @@ do
       end
       return concat(bits, "\n")
     end,
-    tree_with_replaced_vars = function(self, tree, replacements)
+    tree_map = function(self, tree, fn)
       if type(tree) ~= 'table' then
         return tree
       end
+      local replacement = fn(tree)
+      if replacement ~= nil then
+        return replacement
+      end
       local _exp_0 = tree.type
-      if "Var" == _exp_0 then
-        local id = self:var_to_lua_identifier(tree.value)
-        if replacements[id] ~= nil then
-          tree = replacements[id]
+      if "File" == _exp_0 or "Nomsu" == _exp_0 or "Block" == _exp_0 or "List" == _exp_0 or "FunctionCall" == _exp_0 or "Text" == _exp_0 then
+        local new_values, is_changed = { }, false
+        for i, old_value in ipairs(tree.value) do
+          local new_value = type(old_value) ~= "string" and self:tree_map(old_value, fn) or nil
+          if new_value ~= nil and new_value ~= old_value then
+            is_changed = true
+            new_values[i] = new_value
+          else
+            new_values[i] = old_value
+          end
         end
-      elseif "File" == _exp_0 or "Nomsu" == _exp_0 or "Block" == _exp_0 or "List" == _exp_0 or "FunctionCall" == _exp_0 or "Text" == _exp_0 then
-        local new_value = self:tree_with_replaced_vars(tree.value, replacements)
-        if new_value ~= tree.value then
+        if is_changed then
           local new_tree
           do
             local _tbl_0 = { }
@@ -1262,22 +1267,25 @@ do
             new_tree = _tbl_0
           end
           self.tree_metadata[new_tree] = self.tree_metadata[tree]
-          tree = new_tree
-          tree.value = new_value
+          new_tree.value = new_values
+          return new_tree
         end
       elseif "Dict" == _exp_0 then
-        local dirty = false
-        replacements = { }
+        local new_values, is_changed = { }, false
         for i, e in ipairs(tree.value) do
-          local new_key = self:tree_with_replaced_vars(e.dict_key, replacements)
-          local new_value = self:tree_with_replaced_vars(e.dict_value, replacements)
-          dirty = dirty or (new_key ~= e.dict_key or new_value ~= e.dict_value)
-          replacements[i] = {
-            dict_key = new_key,
-            dict_value = new_value
-          }
+          local new_key = self:tree_map(e.dict_key, fn)
+          local new_value = self:tree_map(e.dict_value, fn)
+          if (new_key ~= nil and new_key ~= e.dict_key) or (new_value ~= nil and new_value ~= e.dict_value) then
+            is_changed = true
+            new_values[i] = {
+              dict_key = new_key,
+              dict_value = new_value
+            }
+          else
+            new_values[i] = e
+          end
         end
-        if dirty then
+        if is_changed then
           local new_tree
           do
             local _tbl_0 = { }
@@ -1287,22 +1295,24 @@ do
             new_tree = _tbl_0
           end
           self.tree_metadata[new_tree] = self.tree_metadata[tree]
-          tree = new_tree
-          tree.value = replacements
+          new_tree.value = new_values
+          return new_tree
         end
       elseif nil == _exp_0 then
-        local new_values = { }
-        local any_different = false
-        for k, v in pairs(tree) do
-          new_values[k] = self:tree_with_replaced_vars(v, replacements)
-          any_different = any_different or (new_values[k] ~= tree[k])
-        end
-        if any_different then
-          tree = new_values
-        end
+        error("Invalid tree: " .. tostring(repr(tree)))
       end
       return tree
     end,
+    tree_with_replaced_vars = function(self, tree, replacements)
+      return self:tree_map(tree, function(t)
+        if t.type == "Var" then
+          local id = self:var_to_lua_identifier(t.value)
+          if replacements[id] ~= nil then
+            return replacements[id]
+          end
+        end
+      end)
+    end,
     tree_to_stub = cached(function(self, tree)
       if tree.type ~= "FunctionCall" then
         error("Tried to get stub from non-functioncall tree: " .. tostring(tree.type), 0)
@@ -1663,6 +1673,8 @@ if arg and debug.getinfo(2).func ~= require then
         end)
         if ok and ret ~= nil then
           print("= " .. repr(ret))
+        elseif not ok then
+          print(colored.bright(colored.red(ret)))
         end
       end
     end
diff --git a/nomsu.moon b/nomsu.moon
index 10e57c4..df6d296 100755
--- a/nomsu.moon
+++ b/nomsu.moon
@@ -640,8 +640,6 @@ class NomsuCompiler
         assert tree, "No tree provided."
         if not tree.type
             error("Invalid tree: #{repr(tree)}", 0)
-        unless @tree_metadata[tree]
-            error "??? tree: #{repr tree}", 0
         switch tree.type
             when "File"
                 if #tree.value == 1
@@ -859,45 +857,58 @@ class NomsuCompiler
         insert bits, close
         return concat(bits)
 
-    tree_with_replaced_vars: (tree, replacements)=>
+    tree_map: (tree, fn)=>
+        -- Return a new tree with fn mapped to each node. If fn provides a replacement,
+        -- use that and stop recursing, otherwise recurse.
         if type(tree) != 'table' then return tree
+        replacement = fn(tree)
+        if replacement != nil
+            return replacement
         switch tree.type
-            when "Var"
-                id = @var_to_lua_identifier tree.value
-                if replacements[id] ~= nil
-                    tree = replacements[id]
             when "File", "Nomsu", "Block", "List", "FunctionCall", "Text"
-                new_value = @tree_with_replaced_vars tree.value, replacements
-                if new_value != tree.value
+                new_values, is_changed = {}, false
+                for i,old_value in ipairs(tree.value)
+                    new_value = type(old_value) != "string" and @tree_map(old_value, fn) or nil
+                    if new_value != nil and new_value != old_value
+                        is_changed = true
+                        new_values[i] = new_value
+                    else
+                        new_values[i] = old_value
+                if is_changed
                     new_tree = {k,v for k,v in pairs(tree)}
                     -- TODO: Maybe generate new metadata?
                     @tree_metadata[new_tree] = @tree_metadata[tree]
-                    tree = new_tree
-                    tree.value = new_value
+                    new_tree.value = new_values
+                    return new_tree
+                
             when "Dict"
-                dirty = false
-                replacements = {}
+                new_values, is_changed = {}, false
                 for i,e in ipairs tree.value
-                    new_key = @tree_with_replaced_vars e.dict_key, replacements
-                    new_value = @tree_with_replaced_vars e.dict_value, replacements
-                    dirty or= new_key != e.dict_key or new_value != e.dict_value
-                    replacements[i] = {dict_key:new_key, dict_value:new_value}
-                if dirty
+                    new_key = @tree_map(e.dict_key, fn)
+                    new_value = @tree_map(e.dict_value, fn)
+                    if (new_key != nil and new_key != e.dict_key) or (new_value != nil and new_value != e.dict_value)
+                        is_changed = true
+                        new_values[i] = {dict_key:new_key, dict_value:new_value}
+                    else
+                        new_values[i] = e
+                if is_changed
                     new_tree = {k,v for k,v in pairs(tree)}
                     -- TODO: Maybe generate new metadata?
                     @tree_metadata[new_tree] = @tree_metadata[tree]
-                    tree = new_tree
-                    tree.value = replacements
+                    new_tree.value = new_values
+                    return new_tree
             when nil -- Raw table, probably from one of the .value of a multi-value tree (e.g. List)
-                new_values = {}
-                any_different = false
-                for k,v in pairs tree
-                    new_values[k] = @tree_with_replaced_vars v, replacements
-                    any_different or= (new_values[k] != tree[k])
-                if any_different
-                    tree = new_values
+                error("Invalid tree: #{repr tree}")
+
         return tree
 
+    tree_with_replaced_vars: (tree, replacements)=>
+        return @tree_map tree, (t)->
+            if t.type == "Var"
+                id = @var_to_lua_identifier t.value
+                if replacements[id] != nil
+                    return replacements[id]
+
     tree_to_stub: cached (tree)=>
         if tree.type != "FunctionCall" then error "Tried to get stub from non-functioncall tree: #{tree.type}", 0
         return concat([(t.type == "Word" and t.value or "%") for t in *tree.value], " ")
@@ -1065,6 +1076,8 @@ if arg and debug.getinfo(2).func != require
                 ok, ret = pcall(-> nomsu\run(buff, "stdin"))
                 if ok and ret != nil
                     print "= "..repr(ret)
+                elseif not ok
+                    print colored.bright colored.red ret
     
     err_hand = (error_message)->
         -- TODO: write properly to stderr