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
-		local rast, rerr = parser.parse(replacement, options.chunkname)
+		if type(replacement) == "string" then
-		if not rast then
+			local rast, rerr = parser.parse(replacement, options.chunkname)
-			return error("in macro replacement: %s":format(rerr))
+			if not rast then
-		end
+				return error("in macro replacement: %s":format(tostring(rerr)))
-		-- when giving a single value as a replacement, bypass the implicit push
+			end
-		if #rast == 1 and rast[1].tag == "Push" and rast[1].implicit then
+			-- when giving a single value as a replacement, bypass the implicit push
-			rast = rast[1][1]
+			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,24 +742,33 @@ 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 macroargs = util.merge(peek("macroargs"))
+				local r
-				for i, arg in ipairs(macro.args) do
-					if arg.tag == "Dots" then
-						macroargs["..."] = [for j=i+1, #t do t[j] end]
-					elseif arg.tag == "Id" then
-						if t[i+1] == nil then
-							error("bad argument #%s to macro %s (value expected)":format(i, t[1][1]))
-						end
-						macroargs[arg[1]] = t[i+1]
-					else
-						error("unexpected argument type %s in macro %s":format(arg.tag, t[1][1]))
-					end
-				end
-				push("macroargs", macroargs)
 				nomacro.functions[t[1][1]] = true
-				local r = lua(replacement)
+				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
+							macroargs["..."] = [for j=i+1, #t do t[j] end]
+						elseif arg.tag == "Id" then
+							if t[i+1] == nil then
+								error("bad argument #%s to macro %s (value expected)":format(i, t[1][1]))
+							end
+							macroargs[arg[1]] = t[i+1]
+						else
+							error("unexpected argument type %s in macro %s":format(arg.tag, t[1][1]))
+						end
+					end
+					push("macroargs", macroargs)
+					r = lua(replacement)
+					pop("macroargs")
+				end
 				nomacro.functions[t[1][1]] = nil
-				pop("macroargs")
 				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
+					r = lua(macroargs[t[1]])
-					local r = lua(macroargs[t[1]])
-					nomacro.variables[t[1]] = nil
-					return r
 				elseif macros.variables[t[1]] ~= nil then -- replace with macro variable
-					nomacro.variables[t[1]] = true
+					local macro = macros.variables[t[1]]
-					local r = lua(macros.variables[t[1]])
+					if type(macro) == "function" then
-					nomacro.variables[t[1]] = nil
+						r = macro()
-					return r
+					else
+						r = lua(macro)
+					end
 				end
+				nomacro.variables[t[1]] = nil
 			end
-			return t[1]
+			return r
 		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 --
 ----------------------