Performance optimizations and correctness fix for [=[...]=]-style

strings. Also, require % now properly inserts lua code.

3 files changed, 121 insertions, 86 deletions

diff --git a/nomsu.lua b/nomsu.lua
index cf057be..c103777 100644
--- a/nomsu.lua
+++ b/nomsu.lua
@@ -219,10 +219,8 @@ do
         defs = self.defs
       }
       local where_defs_go = (getmetatable(self.defs) or { }).__newindex or self.defs
-      for _index_0 = 1, #signature do
-        local _des_0 = signature[_index_0]
-        local stub, arg_names, escaped_args
-        stub, arg_names, escaped_args = _des_0[1], _des_0[2], _des_0[3]
+      for sig_i = 1, #signature do
+        local stub, arg_names, escaped_args = unpack(signature[sig_i])
         self:assert(stub, "NO STUB FOUND: " .. tostring(repr(signature)))
         if self.debug then
           self:writeln(tostring(colored.bright("DEFINING RULE:")) .. " " .. tostring(colored.underscore(colored.magenta(repr(stub)))) .. " " .. tostring(colored.bright("WITH ARGS")) .. " " .. tostring(colored.dim(repr(arg_names))))
@@ -545,6 +543,43 @@ do
       end
       return ret, lua_code, vars
     end,
+    run_file = function(self, filename, vars)
+      if vars == nil then
+        vars = { }
+      end
+      if filename:match(".*%.lua") then
+        return dofile(filename)(self, vars)
+      end
+      if filename:match(".*%.nom") then
+        if not self.skip_precompiled then
+          local file = io.open(filename:gsub("%.nom", ".lua"), "r")
+          if file then
+            local lua_code = file:read("*a")
+            file:close()
+            return self:run_lua(lua_code, vars)
+          end
+        end
+        local file = file or io.open(filename)
+        if not file then
+          self:error("File does not exist: " .. tostring(filename))
+        end
+        local nomsu_code = file:read('*a')
+        file:close()
+        return self:run(nomsu_code, filename)
+      else
+        return self:error("Invalid filetype for " .. tostring(filename))
+      end
+    end,
+    require_file = function(self, filename, vars)
+      if vars == nil then
+        vars = { }
+      end
+      local loaded = self.defs["#loaded_files"]
+      if not loaded[filename] then
+        loaded[filename] = self:run_file(filename, vars) or true
+      end
+      return loaded[filename]
+    end,
     run_lua = function(self, lua_code, vars)
       if vars == nil then
         vars = { }
@@ -1344,41 +1379,14 @@ end]]):format(lua_code))
           expr = repr(self:source_code(self:tree_to_value(vars.level)))
         }
       end)
-      local run_file
-      run_file = function(self, vars)
-        if vars.filename:match(".*%.lua") then
-          return dofile(vars.filename)(self, vars)
-        end
-        if vars.filename:match(".*%.nom") then
-          if not self.skip_precompiled then
-            local file = io.open(vars.filename:gsub("%.nom", ".lua"), "r")
-            if file then
-              return self:run_lua(file:read("*a"), vars)
-            end
-          end
-          local file = file or io.open(vars.filename)
-          if not file then
-            self:error("File does not exist: " .. tostring(vars.filename))
-          end
-          local contents = file:read('*a')
-          file:close()
-          return self:run(contents, vars.filename)
-        else
-          return self:error("Invalid filetype for " .. tostring(vars.filename))
-        end
-      end
-      self:def("run file %filename", run_file)
+      self:def("run file %filename", function(self, vars)
+        return self:run_file(vars.filename, vars)
+      end)
       return self:defmacro("require %filename", function(self, vars)
         local filename = self:tree_to_value(vars.filename)
-        local loaded = self.defs["#loaded_files"]
-        if not loaded[filename] then
-          loaded[filename] = run_file(self, {
-            filename = filename
-          }) or true
-        end
-        local _ = loaded[filename]
+        self:require_file(filename, vars)
         return {
-          statements = ""
+          statements = "nomsu:require_file(" .. tostring(repr(filename)) .. ");"
         }
       end)
     end
diff --git a/nomsu.moon b/nomsu.moon
index 49922fa..999c4ca 100755
--- a/nomsu.moon
+++ b/nomsu.moon
@@ -177,7 +177,8 @@ class NomsuCompiler
         @@def_number += 1
         def = {:fn, :src, :is_macro, aliases:{}, def_number:@@def_number, defs:@defs}
         where_defs_go = (getmetatable(@defs) or {}).__newindex or @defs
-        for {stub, arg_names, escaped_args} in *signature
+        for sig_i=1,#signature
+            stub, arg_names, escaped_args = unpack(signature[sig_i])
             @assert stub, "NO STUB FOUND: #{repr signature}"
             if @debug then @writeln "#{colored.bright "DEFINING RULE:"} #{colored.underscore colored.magenta repr(stub)} #{colored.bright "WITH ARGS"} #{colored.dim repr(arg_names)}"
             for i=1,#arg_names-1 do for j=i+1,#arg_names
@@ -354,6 +355,31 @@ class NomsuCompiler
             output_file\write(lua_code)
         return ret, lua_code, vars
 
+    run_file: (filename, vars={})=>
+        if filename\match(".*%.lua")
+            return dofile(filename)(@, vars)
+        if filename\match(".*%.nom")
+            if not @skip_precompiled -- Look for precompiled version
+                file = io.open(filename\gsub("%.nom", ".lua"), "r")
+                if file
+                    lua_code = file\read("*a")
+                    file\close!
+                    return @run_lua(lua_code, vars)
+            file = file or io.open(filename)
+            if not file
+                @error "File does not exist: #{filename}"
+            nomsu_code = file\read('*a')
+            file\close!
+            return @run(nomsu_code, filename)
+        else
+            @error "Invalid filetype for #{filename}"
+    
+    require_file: (filename, vars={})=>
+        loaded = @defs["#loaded_files"]
+        if not loaded[filename]
+            loaded[filename] = @run_file(filename, vars) or true
+        return loaded[filename]
+
     run_lua: (lua_code, vars={})=>
         load_lua_fn, err = load([[
 return function(nomsu, vars)
@@ -852,31 +878,13 @@ end]]\format(lua_code))
         @defmacro "__src__ %level", (vars)=>
             expr: repr(@source_code(@tree_to_value(vars.level)))
 
-        run_file = (vars)=>
-            if vars.filename\match(".*%.lua")
-                return dofile(vars.filename)(@, vars)
-            if vars.filename\match(".*%.nom")
-                if not @skip_precompiled -- Look for precompiled version
-                    file = io.open(vars.filename\gsub("%.nom", ".lua"), "r")
-                    if file
-                        return @run_lua(file\read("*a"), vars)
-                file = file or io.open(vars.filename)
-                if not file
-                    @error "File does not exist: #{vars.filename}"
-                contents = file\read('*a')
-                file\close!
-                return @run(contents, vars.filename)
-            else
-                @error "Invalid filetype for #{vars.filename}"
-        @def "run file %filename", run_file
+        @def "run file %filename", (vars)=>
+            @run_file(vars.filename, vars)
+
         @defmacro "require %filename", (vars)=>
             filename = @tree_to_value(vars.filename)
-            loaded = @defs["#loaded_files"]
-            if not loaded[filename]
-                loaded[filename] = run_file(self, {:filename}) or true
-            loaded[filename]
-            return statements:""
-
+            @require_file(filename, vars)
+            return statements:"nomsu:require_file(#{repr filename});"
 
 if arg
     export colors
diff --git a/utils.lua b/utils.lua
index 6e4bfc2..1e77965 100644
--- a/utils.lua
+++ b/utils.lua
@@ -1,6 +1,7 @@
 --
 -- A collection of helper utility functions
 --
+local lpeg, re = require("lpeg"), require("re")
 
 local function is_list(t)
     if type(t) ~= 'table' then
@@ -24,48 +25,66 @@ local function size(t)
     return n
 end
 
-local function repr(x, depth)
+local _quote_state = {}
+local max = math.max
+local _quote_patt = re.compile("(({'\n' / '\"' / \"'\" / '\\'}->mark_char) / (']' ({'='*}->mark_eq) (']' / !.)) / .)*",
+    {mark_char=function(q)
+        if q == "\n" or q == "\\" then
+            if _quote_state.min_eq == nil then
+                _quote_state.min_eq = 0
+            end
+        elseif q == "'" then
+            _quote_state["'"] = false
+        elseif q == '"' then
+            _quote_state['"'] = false
+        end
+    end,
+    mark_eq=function(eq)
+        _quote_state.min_eq = max(_quote_state.min_eq or 0, #eq+1)
+    end})
+local function repr(x)
     -- Create a string representation of the object that is close to the lua code that will
     -- reproduce the object (similar to Python's "repr" function)
-    depth = depth or math.huge
-    if depth == 0 then
-        return tostring(x)
-    end
     local x_type = type(x)
     if x_type == 'table' then
         if getmetatable(x) then
             -- If this object has a weird metatable, then don't pretend like it's a regular table
             return tostring(x)
-        elseif is_list(x) then
-            local ret = {}
-            for i=1,#x do
-                ret[i] = repr(x[i], depth-1)
-            end
-            return "{"..table.concat(ret, ", ").."}"
         else
             local ret = {}
+            local i = 1
             for k, v in pairs(x) do
-                ret[#ret+1] = "["..repr(k,depth-1).."]= "..repr(v,depth-1)
+                if k == i then
+                    ret[#ret+1] = repr(x[i])
+                    i = i + 1
+                elseif type(k) == 'string' and k:match("[_a-zA-Z][_a-zA-Z0-9]*") then
+                    ret[#ret+1] = k.."= "..repr(v)
+                else
+                    ret[#ret+1] = "["..repr(k).."]= "..repr(v)
+                end
             end
             return "{"..table.concat(ret, ", ").."}"
         end
     elseif x_type == 'string' then
         if x == "\n" then
             return "'\\n'"
-        elseif not x:find([["]]) and not x:find("\n") and not x:find("\\") then
-            return "\"" .. x .. "\""
-        elseif not x:find([[']]) and not x:find("\n") and not x:find("\\") then
+        end
+        _quote_state = {}
+        _quote_patt:match(x)
+        if _quote_state["'"] ~= false and _quote_state.min_eq == nil then
             return "\'" .. x .. "\'"
+        elseif _quote_state['"'] ~= false and _quote_state.min_eq == nil then
+            return "\"" .. x .. "\""
         else
-            for i = 0, math.huge do
-                local eq = ("="):rep(i)
-                if not x:find("%]"..eq.."%]") and x:sub(-#eq-1, -1) ~= "]"..eq then
-                    if x:sub(1, 1) == "\n" then
-                        return "["..eq.."[\n"..x.."]"..eq.."]"
-                    else
-                        return "["..eq.."["..x.."]"..eq.."]"
-                    end
-                end
+            local eq = ("="):rep(_quote_state.min_eq or 0)
+            -- Need to add parens around ([=[...]=]) so lua's parser doesn't get confused
+            -- by stuff like x[[=[...]=]], which should obviously parse as x[ ([=[...]=]) ],
+            -- but instead parses as x( [[=[...]=]] ), i.e. a function call whose argument
+            -- is a string that starts with "=[" and ends with "]="
+            if x:sub(1, 1) == "\n" then
+                return "(["..eq.."[\n"..x.."]"..eq.."])"
+            else
+                return "(["..eq.."["..x.."]"..eq.."])"
             end
         end
     else
@@ -73,11 +92,11 @@ local function repr(x, depth)
     end
 end
 
-local function stringify(x, depth)
+local function stringify(x)
     if type(x) == 'string' then
         return x
     else
-        return repr(x, depth)
+        return repr(x)
     end
 end