candran.can

local candran = {
	VERSION = "1.0.0"
}
package.loaded["candran"] = candran

#import("candran.util")
#import("candran.serpent")

#import("compiler.lua54")
#import("compiler.lua53")
#import("compiler.lua52")
#import("compiler.luajit")
#import("compiler.lua51")

#import("candran.can-parser.scope")
#import("candran.can-parser.validator")
#import("candran.can-parser.pp")
#import("candran.can-parser.parser")

local unpack = unpack or table.unpack

--- Default options.
candran.default = {
	target = "lua54",
	indentation = "",
	newline = "\n",
	variablePrefix = "__CAN_",
	mapLines = true,
	chunkname = "nil",
	rewriteErrors = true,
	builtInMacros = true,
	preprocessorEnv = {},
	import = {}
}

-- Autodetect version
if _VERSION == "Lua 5.1" then
	if package.loaded.jit then
		candran.default.target = "luajit"
	else
		candran.default.target = "lua51"
	end
elseif _VERSION == "Lua 5.2" then
	candran.default.target = "lua52"
elseif _VERSION == "Lua 5.3" then
	candran.default.target = "lua53"
end

--- Run the preprocessor
-- @tparam input string input code
-- @tparam options table arguments for the preprocessor. They will be inserted into the preprocessor environement.
-- @treturn[1] output string output code
-- @treturn[1] macros registered macros
-- @treturn[2] nil nil if error
-- @treturn[2] error string error message
function candran.preprocess(input, options={}, _env)
	options = util.merge(candran.default, options)
	local macros = {
		functions = {},
		variables = {}
	}

	-- add auto imports
	for _, mod in ipairs(options.import) do
		input =.. "#import(%q, {loadLocal=false})\n":format(mod)
	end

	-- generate preprocessor code
	local preprocessor = ""
	local i = 0
	local inLongString = false
	local inComment = false
	for line in (input.."\n"):gmatch("(.-\n)") do
		i += 1
		-- Simple multiline comment/string detection
		if inComment then
			inComment = not line:match("%]%]")
		elseif inLongString then
			inLongString = not line:match("%]%]")
		else
			if line:match("[^%-]%[%[") then
				inLongString = true
			elseif line:match("%-%-%[%[") then
				inComment = true
			end
		end
		if not inComment and not inLongString and line:match("^%s*#") and not line:match("^#!") then -- exclude shebang
			preprocessor ..= line:gsub("^%s*#", "")
		else
			local l = line:sub(1, -2)
			if not inLongString and options.mapLines and not l:match("%-%- (.-)%:(%d+)$") then
				preprocessor ..= ("write(%q)"):format(l .. " -- "..options.chunkname..":" .. i) .. "\n"
			else
				preprocessor ..= ("write(%q)"):format(line:sub(1, -2)) .. "\n"
			end
		end
	end
	preprocessor ..= "return output"

	-- make preprocessor environement
	local exportenv = {}
	local env = util.merge(_G, options.preprocessorEnv)
	--- Candran library table
	env.candran = candran
	--- Current preprocessor output
	env.output = ""
	--- Import an external Candran/Lua module into the generated file
	-- Notable options:
	-- * loadLocal (true): true to automatically load the module into a local variable
	-- * loadPackage (true): true to automatically load the module into the loaded packages table
	-- @tparam modpath string module path
	-- @tparam margs table preprocessor options to use when preprocessessing the module
	env.import = function(modpath, margs={})
		local filepath = assert(util.search(modpath, {"can", "lua"}), "No module named \""..modpath.."\"")

		-- open module file
		local f = io.open(filepath)
		if not f then error("can't open the module file to import") end

		margs = util.merge(options, { chunkname = filepath, loadLocal = true, loadPackage = true }, margs)
		margs.import = {} -- no need for recursive import
		local modcontent, modmacros, modenv = assert(candran.preprocess(f:read("*a"), margs))
		macros = util.recmerge(macros, modmacros)
		for k, v in pairs(modenv) do env[k] = v end
		f:close()

		-- get module name (ex: module name of path.to.module is module)
		local modname = modpath:match("[^%.]+$")

		env.write(
			"-- MODULE "..modpath.." --\n"..
			"local function _()\n"..
				modcontent.."\n"..
			"end\n"..
			(margs.loadLocal and ("local %s = _() or %s\n"):format(modname, modname) or "").. -- auto require
			(margs.loadPackage and ("package.loaded[%q] = %s or true\n"):format(modpath, margs.loadLocal and modname or "_()") or "").. -- add to package.loaded
			"-- END OF MODULE "..modpath.." --"
		)
	end
	--- Include another file content in the preprocessor output.
	-- @tparam file string filepath
	env.include = function(file)
		local f = io.open(file)
		if not f then error("can't open the file "..file.." to include") end
		env.write(f:read("*a"))
		f:close()
	end
	--- Write a line in the preprocessor output.
	-- @tparam ... string strings to write (similar to print)
	env.write = function(...)
		env.output ..= table.concat({...}, "\t") .. "\n"
	end
	--- Will be replaced with the content of the variable with the given name, if it exists.
	-- @tparam name string variable name
	env.placeholder = function(name)
		if env[name] then
			env.write(env[name])
		end
	end
	env.define = function(identifier, replacement)
		-- parse identifier
		local iast, ierr = parser.parsemacroidentifier(identifier, options.chunkname)
		if not iast then
			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(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 = replacement }
		elseif iast.tag == "Id" then
			macros.variables[iast[1]] = replacement
		else
			error("invalid macro type %s":format(tostring(iast.tag)))
		end
	end
	env.set = function(identifier, value)
		exportenv[identifier] = value
		env[identifier] = value
	end

	-- default macros
	if options.builtInMacros then
		env.define("__STR__(x)", function(x) return ("%q"):format(x) end)
		local s = require("candran.serpent")
		env.define("__CONSTEXPR__(expr)", function(expr)
			return s.block(assert(candran.load(expr))(), {fatal = true})
		end)
	end

	-- compile & load preprocessor
	local preprocess, err = candran.compile(preprocessor, options)
	if not preprocess then
		return nil, "in preprocessor: "..err
	end

	preprocess, err = util.load(preprocessor, "candran preprocessor", env)
	if not preprocess then
		return nil, "in preprocessor: "..err
	end

	-- execute preprocessor
	local success, output = pcall(preprocess)
	if not success then
		return nil, "in preprocessor: "..output
	end

	return output, macros, exportenv
end

--- Run the compiler
-- @tparam input string input code
-- @tparam options table options for the compiler
-- @tparam macros table defined macros, as returned by the preprocessor
-- @treturn[1] output string output code
-- @treturn[2] nil nil if error
-- @treturn[2] error string error message
function candran.compile(input, options={}, macros)
	options = util.merge(candran.default, options)

	local ast, errmsg = parser.parse(input, options.chunkname)

	if not ast then
		return nil, errmsg
	end

	return require("compiler."..options.target)(input, ast, options, macros)
end

--- Preprocess & compile code
-- @tparam code string input code
-- @tparam options table arguments for the preprocessor and compiler
-- @treturn[1] output string output code
-- @treturn[2] nil nil if error
-- @treturn[2] error string error message
function candran.make(code, options)
	local r, err = candran.preprocess(code, options)
	if r then
		r, err = candran.compile(r, options, err)
		if r then
			return r
		end
	end
	return r, err
end

local errorRewritingActive = false
local codeCache = {}
--- Candran equivalent to the Lua 5.3's loadfile funtion.
-- Will rewrite errors by default.
function candran.loadfile(filepath, env, options)
	local f, err = io.open(filepath)
	if not f then
		return nil, "cannot open %s":format(tostring(err))
	end
	local content = f:read("*a")
	f:close()

	return candran.load(content, filepath, env, options)
end

--- Candran equivalent to the Lua 5.3's load funtion.
-- Will rewrite errors by default.
function candran.load(chunk, chunkname, env, options={})
	options = util.merge({ chunkname = tostring(chunkname or chunk) }, options)

	local code, err = candran.make(chunk, options)
	if not code then
		return code, err
	end

	codeCache[options.chunkname] = code
	local f
	f, err = util.load(code, "=%s(%s)":format(options.chunkname, "compiled candran"), env)

	-- Um. Candran isn't supposed to generate invalid Lua code, so this is a major issue.
	-- This is not going to raise an error because this is supposed to behave similarly to Lua's load function.
	-- But the error message will likely be useless unless you know how Candran works.
	if f == nil then
		return f, "candran unexpectedly generated invalid code: "..err
	end

	if options.rewriteErrors == false then
		return f
	else
		return function(...)
			if not errorRewritingActive then
				errorRewritingActive = true
				local t = { xpcall(f, candran.messageHandler, ...) }
				errorRewritingActive = false
				if t[1] == false then
					error(t[2], 0)
				end
				return unpack(t, 2)
			else
				return f(...)
			end
		end
	end
end

--- Candran equivalent to the Lua 5.3's dofile funtion.
-- Will rewrite errors by default.
function candran.dofile(filename, options)
	local f, err = candran.loadfile(filename, nil, options)

	if f == nil then
		error(err)
	else
		return f()
	end
end

--- Candran error message handler.
-- Use it in xpcall to rewrite stacktraces to display Candran source file lines instead of compiled Lua lines.
function candran.messageHandler(message, noTraceback)
	message = tostring(message)
	if not noTraceback and not message:match("\nstack traceback:\n") then
		message = debug.traceback(message, 2)
	end
	return message:gsub("(\n?%s*)([^\n]-)%:(%d+)%:", function(indentation, source, line)
		line = tonumber(line)

		local originalFile
		local strName = source:match("^(.-)%(compiled candran%)$")
		if strName then
			if codeCache[strName] then
				originalFile = codeCache[strName]
				source = strName
			end
		else
			if fi = io.open(source, "r") then
				originalFile = fi:read("*a")
				fi:close()
			end
		end

		if originalFile then
			local i = 0
			for l in (originalFile.."\n"):gmatch("([^\n]*)\n") do
				i = i +1
				if i == line then
					local extSource, lineMap = l:match(".*%-%- (.-)%:(%d+)$")
					if lineMap then
						if extSource ~= source then
							return indentation .. extSource .. ":" .. lineMap .. "(" .. extSource .. ":" .. line .. "):"
						else
							return indentation .. extSource .. ":" .. lineMap .. "(" .. line .. "):"
						end
					end
					break
				end
			end
		end
	end)
end

--- Candran package searcher function. Use the existing package.path.
function candran.searcher(modpath)
	local filepath = util.search(modpath, {"can"})
	if not filepath then
		if _VERSION == "Lua 5.4" then
			return "no candran file in package.path"
		else
			return "\n\tno candran file in package.path"
		end
	end
	return (modpath) -- 2nd argument is not passed in Lua 5.1, so a closure is required
		local r, s = candran.loadfile(filepath)
		if r then
			return r(modpath, filepath)
		else
			error("error loading candran module '%s' from file '%s':\n\t%s":format(modpath, filepath, tostring(s)), 0)
		end
	end, filepath
end

--- Register the Candran package searcher.
function candran.setup()
	local searchers = if _VERSION == "Lua 5.1" then
		package.loaders
	else
		package.searchers
	end
	-- check if already setup
	for _, s in ipairs(searchers) do
		if s == candran.searcher then
			return candran
		end
	end
	-- setup
	table.insert(searchers, 1, candran.searcher)
	return candran
end

return candran