Fixes and simplifications for nomsu codegen.

2 files changed, 71 insertions, 60 deletions

diff --git a/nomsu_compiler.lua b/nomsu_compiler.lua
index 2ab57fe..44fc91c 100644
--- a/nomsu_compiler.lua
+++ b/nomsu_compiler.lua
@@ -874,27 +874,22 @@ do
       return error("Unknown type: " .. tostring(tree.type))
     end
   end
-  NomsuCompiler.tree_to_nomsu = function(self, tree, consumed_comments)
-    if consumed_comments == nil then
-      consumed_comments = nil
+  NomsuCompiler.tree_to_nomsu = function(self, tree, comments)
+    if comments == nil then
+      comments = nil
     end
-    consumed_comments = consumed_comments or { }
-    local pop_comments
-    pop_comments = function(pos, prefix, suffix)
-      if prefix == nil then
-        prefix = ''
-      end
-      if suffix == nil then
-        suffix = ''
-      end
+    if not (comments) then
+      local visited
+      comments, visited = { }, { }
       local find_comments
       find_comments = function(t)
         if t.comments and t.source.filename == tree.source.filename then
           local _list_0 = t.comments
           for _index_0 = 1, #_list_0 do
             local c = _list_0[_index_0]
-            if not (consumed_comments[c]) then
-              coroutine.yield(c)
+            if not (visited[c]) then
+              insert(comments, c)
+              visited[c] = true
             end
           end
         end
@@ -905,14 +900,26 @@ do
           end
         end
       end
+      find_comments(tree)
+      table.sort(comments, function(a, b)
+        return (a.pos > b.pos)
+      end)
+    end
+    local pop_comments
+    pop_comments = function(pos, prefix, suffix)
+      if prefix == nil then
+        prefix = ''
+      end
+      if suffix == nil then
+        suffix = ''
+      end
       local nomsu = NomsuCode(tree.source)
-      for comment in coroutine.wrap(function()
-        return find_comments(tree)
-      end) do
+      for i = #comments, 1, -1 do
+        local comment = comments[i]
         if comment.pos > pos then
           break
         end
-        consumed_comments[comment] = true
+        comments[i] = nil
         nomsu:append("#" .. (gsub(comment.comment, "\n", "\n    ")) .. "\n")
         if comment.comment:match("^\n.") then
           nomsu:append("\n")
@@ -927,26 +934,26 @@ do
     end
     local recurse
     recurse = function(t)
-      return self:tree_to_nomsu(t, consumed_comments)
+      return self:tree_to_nomsu(t, comments)
     end
     local _exp_0 = tree.type
     if "FileChunks" == _exp_0 then
       local nomsu = NomsuCode(tree.source, pop_comments(tree.source.start))
-      for i, chunk in ipairs(tree) do
-        if i > 1 then
+      for chunk_no, chunk in ipairs(tree) do
+        if chunk_no > 1 then
           nomsu:append("\n\n" .. tostring(("~"):rep(80)) .. "\n\n")
         end
         nomsu:append(pop_comments(chunk.source.start))
         if chunk.type == "Block" then
-          for j, line in ipairs(chunk) do
+          for line_no, line in ipairs(chunk) do
             nomsu:append(pop_comments(line.source.start, '\n'))
             local line_nomsu = recurse(line)
             nomsu:append(line_nomsu)
-            if j < #chunk then
+            if line_no < #chunk then
               nomsu:append(line_nomsu:is_multiline() and "\n\n" or "\n")
             end
           end
-          nomsu:append(pop_comments(tree.source.stop, '\n'))
+          nomsu:append(pop_comments(chunk.source.stop, '\n'))
         else
           nomsu:append(recurse(chunk))
         end
@@ -1050,7 +1057,6 @@ do
               if j > 1 then
                 nomsu:append("\n")
               end
-              line = gsub(line, "\\", "\\\\")
               if #line > 1.25 * MAX_LINE then
                 local remainder = line
                 while #remainder > 0 do
@@ -1126,16 +1132,15 @@ do
           nomsu:append(item_nomsu)
         elseif #tostring(item_nomsu) <= MAX_LINE then
           if nomsu:trailing_line_len() > 0 then
-            nomsu:append("\n")
+            nomsu:append("\n", pop_comments(item_nomsu.source.start))
           end
           nomsu:append(item_nomsu)
         else
           item_nomsu = recurse(item)
-          local _exp_1 = item.type
-          if "List" == _exp_1 or "Dict" == _exp_1 or "Text" == _exp_1 or "Block" == _exp_1 or "EscapedNomsu" == _exp_1 then
-            nomsu:append(item_nomsu)
+          if item.type == "Action" then
+            nomsu:append("(..)\n    ", pop_comments(item_nomsu.source.start), item_nomsu)
           else
-            nomsu:append("(..)\n    ", item_nomsu)
+            nomsu:append(item_nomsu)
           end
           if i < #tree then
             nomsu:append("\n")
@@ -1163,16 +1168,15 @@ do
           nomsu:append(item_nomsu)
         elseif #tostring(item_nomsu) <= MAX_LINE then
           if nomsu:trailing_line_len() > 0 then
-            nomsu:append("\n")
+            nomsu:append("\n", pop_comments(item_nomsu.source.start))
           end
           nomsu:append(item_nomsu)
         else
           item_nomsu = recurse(item)
-          local _exp_1 = item.type
-          if "List" == _exp_1 or "Dict" == _exp_1 or "Text" == _exp_1 or "Block" == _exp_1 or "EscapedNomsu" == _exp_1 then
-            nomsu:append(item_nomsu)
+          if item.type == "Action" then
+            nomsu:append("(..)\n    ", pop_comments(item_nomsu.source.start), item_nomsu)
           else
-            nomsu:append("(..)\n    ", item_nomsu)
+            nomsu:append(item_nomsu)
           end
           if i < #tree then
             nomsu:append("\n")
diff --git a/nomsu_compiler.moon b/nomsu_compiler.moon
index 6f9c9d5..d856a83 100644
--- a/nomsu_compiler.moon
+++ b/nomsu_compiler.moon
@@ -592,19 +592,27 @@ with NomsuCompiler
             else
                 error("Unknown type: #{tree.type}")
 
-    .tree_to_nomsu = (tree, consumed_comments=nil)=>
-        consumed_comments or= {}
-        pop_comments = (pos, prefix='', suffix='')->
+    .tree_to_nomsu = (tree, comments=nil)=>
+        unless comments
+            comments, visited = {}, {}
             find_comments = (t)->
                 if t.comments and t.source.filename == tree.source.filename
                     for c in *t.comments
-                        coroutine.yield(c) unless consumed_comments[c]
+                        unless visited[c]
+                            insert(comments, c)
+                            visited[c] = true
                 for x in *t
                     find_comments(x) if AST.is_syntax_tree x
+            find_comments(tree)
+            -- Sort in reversed order so they can be easily popped
+            table.sort(comments, (a,b)->(a.pos > b.pos))
+
+        pop_comments = (pos, prefix='', suffix='')->
             nomsu = NomsuCode(tree.source)
-            for comment in coroutine.wrap(-> find_comments(tree))
+            for i=#comments,1,-1
+                comment = comments[i]
                 break if comment.pos > pos
-                consumed_comments[comment] = true
+                comments[i] = nil
                 nomsu\append("#"..(gsub(comment.comment, "\n", "\n    ")).."\n")
                 if comment.comment\match("^\n.") then nomsu\append("\n") -- for aesthetics
             return '' if #nomsu.bits == 0
@@ -613,21 +621,21 @@ with NomsuCompiler
             return nomsu
 
         -- For concision:
-        recurse = (t)-> @tree_to_nomsu(t, consumed_comments)
+        recurse = (t)-> @tree_to_nomsu(t, comments)
 
         switch tree.type
             when "FileChunks"
                 nomsu = NomsuCode(tree.source, pop_comments(tree.source.start))
-                for i, chunk in ipairs tree
-                    nomsu\append "\n\n#{("~")\rep(80)}\n\n" if i > 1
+                for chunk_no, chunk in ipairs tree
+                    nomsu\append "\n\n#{("~")\rep(80)}\n\n" if chunk_no > 1
                     nomsu\append pop_comments(chunk.source.start)
                     if chunk.type == "Block"
-                        for j, line in ipairs chunk
+                        for line_no, line in ipairs chunk
                             nomsu\append pop_comments(line.source.start, '\n')
                             line_nomsu = recurse(line)
                             nomsu\append line_nomsu
-                            nomsu\append(line_nomsu\is_multiline! and "\n\n" or "\n") if j < #chunk
-                        nomsu\append pop_comments(tree.source.stop, '\n')
+                            nomsu\append(line_nomsu\is_multiline! and "\n\n" or "\n") if line_no < #chunk
+                        nomsu\append pop_comments(chunk.source.stop, '\n')
                     else
                         nomsu\append recurse(chunk)
                 nomsu\append pop_comments(tree.source.stop, '\n')
@@ -710,7 +718,6 @@ with NomsuCompiler
                             bit_lines = files.get_lines(bit)
                             for j, line in ipairs bit_lines
                                 if j > 1 then nomsu\append "\n"
-                                line = gsub(line, "\\", "\\\\")
                                 if #line > 1.25*MAX_LINE
                                     remainder = line
                                     while #remainder > 0
@@ -765,15 +772,15 @@ with NomsuCompiler
                         nomsu\append ", " if nomsu\trailing_line_len! > 0
                         nomsu\append item_nomsu
                     elseif #tostring(item_nomsu) <= MAX_LINE
-                        nomsu\append "\n" if nomsu\trailing_line_len! > 0
+                        if nomsu\trailing_line_len! > 0
+                            nomsu\append "\n", pop_comments(item_nomsu.source.start)
                         nomsu\append item_nomsu
                     else
                         item_nomsu = recurse(item)
-                        switch item.type
-                            when "List", "Dict", "Text", "Block", "EscapedNomsu"
-                                nomsu\append item_nomsu
-                            else
-                                nomsu\append "(..)\n    ", item_nomsu
+                        if item.type == "Action"
+                            nomsu\append "(..)\n    ", pop_comments(item_nomsu.source.start), item_nomsu
+                        else
+                            nomsu\append item_nomsu
                         nomsu\append "\n" if i < #tree
                 nomsu\append pop_comments(tree.source.stop, '\n')
                 return NomsuCode(tree.source, "[..]\n    ", nomsu)
@@ -791,15 +798,15 @@ with NomsuCompiler
                         nomsu\append ", " if nomsu\trailing_line_len! > 0
                         nomsu\append item_nomsu
                     elseif #tostring(item_nomsu) <= MAX_LINE
-                        nomsu\append "\n" if nomsu\trailing_line_len! > 0
+                        if nomsu\trailing_line_len! > 0
+                            nomsu\append "\n", pop_comments(item_nomsu.source.start)
                         nomsu\append item_nomsu
                     else
                         item_nomsu = recurse(item)
-                        switch item.type
-                            when "List", "Dict", "Text", "Block", "EscapedNomsu"
-                                nomsu\append item_nomsu
-                            else
-                                nomsu\append "(..)\n    ", item_nomsu
+                        if item.type == "Action"
+                            nomsu\append "(..)\n    ", pop_comments(item_nomsu.source.start), item_nomsu
+                        else
+                            nomsu\append item_nomsu
                         nomsu\append "\n" if i < #tree
                 nomsu\append pop_comments(tree.source.stop, '\n')
                 return NomsuCode(tree.source, "{..}\n    ", nomsu)