Add support for macro replacement using Lua functions

candran.can

@@ -150,24 +150,29 @@ function candran.preprocess(input, options={})
 		-- parse identifier
 		local iast, ierr = parser.parsemacroidentifier(identifier, options.chunkname)
 		if not iast then
-			return error("in macro identifier: %s":format(ierr))
+			return error("in macro identifier: %s":format(tostring(ierr)))
 		end
 		-- parse replacement value
+		if type(replacement) == "string" then
 			local rast, rerr = parser.parse(replacement, options.chunkname)
 			if not rast then
-			return error("in macro replacement: %s":format(rerr))
+				return error("in macro replacement: %s":format(tostring(rerr)))
 			end
 			-- when giving a single value as a replacement, bypass the implicit push
 			if #rast == 1 and rast[1].tag == "Push" and rast[1].implicit then
 				rast = rast[1][1]
 			end
+			replacement = rast
+		elseif type(replacement) ~= "function" then
+			error("bad argument #2 to 'define' (string or function expected)")
+		end
 		-- add macros
 		if iast.tag == "MacroFunction" then
-			macros.functions[iast[1][1]] = { args = iast[2], replacement = rast }
+			macros.functions[iast[1][1]] = { args = iast[2], replacement = replacement }
 		elseif iast.tag == "Id" then
-			macros.variables[iast[1]] = rast
+			macros.variables[iast[1]] = replacement
 		else
-			error("invalid macro type %s":format(iast.tag))
+			error("invalid macro type %s":format(tostring(iast.tag)))
 		end
 	end
 

compiler/lua54.can

@@ -2,6 +2,8 @@ local util = require("candran.util")
 
 local targetName = "Lua 5.4"
 
+local unpack = unpack or table.unpack
+
 return function(code, ast, options, macros={functions={}, variables={}})
 	--- Line mapping
 	local lastInputPos = 1 -- last token position in the input code
@@ -740,6 +742,15 @@ return function(code, ast, options, macros={functions={}, variables={}})
 			elseif t[1].tag == "Id" and not nomacro.functions[t[1][1]] and macros.functions[t[1][1]] then
 				local macro = macros.functions[t[1][1]]
 				local replacement = macro.replacement
+				local r
+				nomacro.functions[t[1][1]] = true
+				if type(replacement) == "function" then
+					local args = {}
+					for i=2, #t do
+						table.insert(args, lua(t[i]))
+					end
+					r = replacement(unpack(args))
+				else
 					local macroargs = util.merge(peek("macroargs"))
 					for i, arg in ipairs(macro.args) do
 						if arg.tag == "Dots" then
@@ -754,10 +765,10 @@ return function(code, ast, options, macros={functions={}, variables={}})
 						end
 					end
 					push("macroargs", macroargs)
-				nomacro.functions[t[1][1]] = true
-				local r = lua(replacement)
-				nomacro.functions[t[1][1]] = nil
+					r = lua(replacement)
 					pop("macroargs")
+				end
+				nomacro.functions[t[1][1]] = nil
 				return r
 			elseif t[1].tag == "MethodStub" then -- method call
 				if t[1][1].tag == "String" or t[1][1].tag == "Table" then
@@ -793,21 +804,23 @@ return function(code, ast, options, macros={functions={}, variables={}})
 		end,
 		-- Id{ <string> }
 		Id = (t)
+			local r = t[1]
 			local macroargs = peek("macroargs")
 			if not nomacro.variables[t[1]] then
+				nomacro.variables[t[1]] = true
 				if macroargs and macroargs[t[1]] then -- replace with macro argument
-					nomacro.variables[t[1]] = true
-					local r = lua(macroargs[t[1]])
-					nomacro.variables[t[1]] = nil
-					return r
+					r = lua(macroargs[t[1]])
 				elseif macros.variables[t[1]] ~= nil then -- replace with macro variable
-					nomacro.variables[t[1]] = true
-					local r = lua(macros.variables[t[1]])
+					local macro = macros.variables[t[1]]
+					if type(macro) == "function" then
+						r = macro()
+					else
+						r = lua(macro)
+					end
+				end
 				nomacro.variables[t[1]] = nil
+			end
 			return r
-				end
-			end
-			return t[1]
 		end,
 		-- AttributeId{ <string> <string>? }
 		AttributeId = (t)

test/test.lua

@@ -199,6 +199,17 @@ x = f(x)
 return x
 ]], 42)
 
+test("preprocessor macro replace variable with function", [[
+#define("a", function() return "42" end)
+return a
+]], 42)
+
+test("preprocessor macro replace function with function", [[
+#define("test(x)", function(x) return ("%s = 42"):format(x) end)
+test(hello)
+return hello
+]], 42)
+
 ----------------------
 -- SYNTAX ADDITIONS --
 ----------------------