WFL functions can now be defined in Lua

src/formula/formula.cpp

@@ -248,6 +248,13 @@ formula::formula(const tk::token* i1, const tk::token* i2, function_symbol_table
 	}
 }
 
+formula::formula(expression_ptr expr, function_symbol_table* symbols) :
+	expr_(expr),
+	str_(expr->str()),
+	symbols_(symbols)
+{
+}
+
 formula_ptr formula::create_optional_formula(const std::string& str, function_symbol_table* symbols)
 {
 	if(str.empty()) {

src/formula/formula.hpp

@@ -20,6 +20,11 @@
 #include "formula/tokenizer.hpp"
 #include "formula/variant.hpp"
 #include <memory>
+#include <set>
+
+namespace lua_formula_bridge {
+	class fwrapper;
+}
 
 namespace wfl
 {
@@ -34,7 +39,8 @@ namespace tk = tokenizer;
 class formula
 {
 public:
-	formula(const std::string& str, function_symbol_table* symbols = nullptr);
+	explicit formula(const std::string& str, function_symbol_table* symbols = nullptr);
+	explicit formula(expression_ptr expr, function_symbol_table* symbols = nullptr);
 	formula(const tk::token* i1, const tk::token* i2, function_symbol_table* symbols = nullptr);
 
 	static variant evaluate(
@@ -73,7 +79,8 @@ public:
 	const std::string& str() const { return str_; }
 
 	static const char* const id_chars;
-
+	static const std::set<std::string> keywords; // defined in formula/tokenizer.cpp
+	function_symbol_table* get_functions() {return symbols_;}
 private:
 	variant execute(const formula_callable& variables, formula_debugger* fdb = nullptr) const;
 	variant execute(formula_debugger* fdb) const;
@@ -86,6 +93,7 @@ private:
 	function_symbol_table* symbols_;
 
 	friend class formula_debugger;
+	friend class lua_formula_bridge::fwrapper;
 };
 
 struct formula_error : public game::error

src/formula/function.hpp

@@ -91,8 +91,8 @@ public:
 	virtual std::string str() const = 0;
 
 private:
-	virtual variant execute(const formula_callable& variables, formula_debugger* fdb = nullptr) const = 0;
-
+	virtual variant execute(const formula_callable& variables, formula_debugger *fdb = nullptr) const = 0;
+	formula_ptr parent_;
 	const std::string name_;
 	friend class formula_debugger;
 };

src/formula/tokenizer.cpp

@@ -14,12 +14,19 @@
 */
 
 #include "formula/tokenizer.hpp"
+#include "formula/formula.hpp" // for formula::keywords
 
 #include <locale>
 #include <sstream>
 
 namespace wfl
 {
+
+// This is defined here to make it easy to remember to update it if a new one is added
+const std::set<std::string> formula::keywords = {
+	"d", "or", "def", "and", "not", "wfl", "fai", "where", "wflend", "faiend", "functions",
+};
+
 namespace tokenizer
 {
 

src/scripting/lua_formula_bridge.cpp

@@ -14,15 +14,19 @@
 
 #include "scripting/lua_formula_bridge.hpp"
 
+#include <set>
+
 #include "game_board.hpp"
 #include "scripting/lua_unit.hpp"
 #include "scripting/lua_common.hpp"
+#include "scripting/lua_kernel_base.hpp"
 #include "scripting/lua_team.hpp"
 #include "scripting/lua_unit_attacks.hpp"
 #include "scripting/lua_unit_type.hpp"
 #include "lua/wrapper_lauxlib.h"
 #include "formula/callable_objects.hpp"
 #include "formula/formula.hpp"
+#include "formula/function.hpp"
 #include "variable.hpp"
 
 #include "resources.hpp"
@@ -30,17 +34,33 @@
 #include "units/unit.hpp"
 
 static const char formulaKey[] = "formula";
+static const char formfcntbKey[] = "formula function table";
 
 using namespace wfl;
 
 void luaW_pushfaivariant(lua_State* L, variant val);
 variant luaW_tofaivariant(lua_State* L, int i);
 
+typedef function_symbol_table* functionstb_ptr;
+
+std::set<functionstb_ptr> lua_owned_fcntb;
+
 class lua_callable : public formula_callable {
 	lua_State* mState;
 	int table_i;
 public:
-	lua_callable(lua_State* L, int i) : mState(L), table_i(lua_absindex(L,i)) {}
+	bool is_list() const {
+		for(lua_pushnil(mState); lua_next(mState, table_i); lua_pop(mState, 1)) {
+			if(!lua_isnumber(mState, -1)) {
+				lua_pop(mState, 2);
+				return false;
+			}
+		}
+		return true;
+	}
+	lua_callable(lua_State* L, int i) : mState(L), table_i(lua_absindex(L,i)) {
+		luaL_checktype(L, i, LUA_TTABLE);
+	}
 	variant get_value(const std::string& key) const {
 		if(key == "__list") {
 			std::vector<variant> values;
@@ -76,7 +96,7 @@ public:
 			lua_pop(mState, 1);
 			if(is_valid_key) {
 				std::string key = lua_tostring(mState, -2);
-				if(key.find_first_not_of(formula::id_chars) != std::string::npos) {
+				if(key.find_first_not_of(formula::id_chars) == std::string::npos) {
 					add_input(inputs, key);
 				}
 			}
@@ -119,6 +139,75 @@ public:
 	}
 };
 
+class lua_formula_function : public function_expression {
+	lua_kernel_base& kernel;
+	const void* registry_index;
+	variant execute(const formula_callable& variables, formula_debugger *fdb) const {
+		lua_State* L = kernel.get_state();
+		lua_rawgetp(L, LUA_REGISTRYINDEX, registry_index);
+		luaW_pushfaivariant(L, variant(variables.fake_ptr()));
+		for(auto arg : args()) {
+			if(fdb) {
+				arg.reset(new wrapper_formula(arg));
+			}
+			new(L) lua_formula_bridge::fwrapper(arg);
+			luaL_setmetatable(L, formulaKey);
+		}
+		bool success = luaW_pcall(L, args().size() + 1, 2);
+		variant result;
+		if(!success || lua_isnoneornil(L, -2)) {
+			return result;
+		} else if(lua_isnoneornil(L, -1)) {
+			// Use normal guessing to determine the type
+			result = luaW_tofaivariant(L, -2);
+		} else {
+			std::string type = lua_tostring(L, -1);
+			if(type == "string") {
+				result = variant(luaL_checkstring(L, -2));
+			} else if(type == "integer") {
+				result = variant(luaL_checkinteger(L, -2));
+			} else if(type == "decimal") {
+				result = variant(luaL_checknumber(L, -2), variant::DECIMAL_VARIANT);
+			} else if(type == "list") {
+				result = lua_callable(L, -2).get_value("__list");
+			} else if(type == "map") {
+				result = lua_callable(L, -2).get_value("__map");
+			} else if(type == "unit") {
+				unit& u = luaW_checkunit(L, -2);
+				result = variant(std::make_shared<unit_callable>(u));
+			} else if(type == "loc") {
+				map_location loc = luaW_checklocation(L, -2);
+				result = variant(std::make_shared<location_callable>(loc));
+			}
+		}
+		return result;
+	}
+public:
+	lua_formula_function(const std::string name, const args_list& args, int min_args, int max_args, const void* ridx, lua_kernel_base& k)
+		: function_expression(name, args, min_args, max_args)
+		, kernel(k)
+		, registry_index(ridx)
+	{
+	}
+};
+
+class lua_formula_function_defn : public formula_function {
+	int min_args, max_args;
+	lua_kernel_base& kernel;
+public:
+	lua_formula_function_defn(const std::string name, int min, int max, lua_kernel_base& k)
+		: formula_function(name)
+		, min_args(min)
+		, max_args(max)
+		, kernel(k)
+	{
+	}
+	virtual function_expression_ptr generate_function_expression(const std::vector<expression_ptr>& args) const {
+		return function_expression_ptr(new lua_formula_function(name_, args, min_args, max_args, this, kernel));
+	}
+	virtual ~lua_formula_function_defn() {}
+};
+
 void luaW_pushfaivariant(lua_State* L, variant val) {
 	if(val.is_int()) {
 		lua_pushinteger(L, val.as_int());
@@ -188,6 +277,7 @@ variant luaW_tofaivariant(lua_State* L, int i) {
 		case LUA_TBOOLEAN:
 			return variant(lua_tointeger(L, i));
 		case LUA_TNUMBER:
+			if(lua_isinteger(L, i)) return variant(lua_tointeger(L, i));
 			return variant(lua_tonumber(L, i), variant::DECIMAL_VARIANT);
 		case LUA_TSTRING:
 			return variant(lua_tostring(L, i));
@@ -242,7 +332,7 @@ lua_formula_bridge::fpointer luaW_check_formula(lua_State* L, int idx, bool allo
 
 /**
  * Evaluates a formula in the formula engine.
- * - Arg 1: Formula string.
+ * - Arg 1: Formula string or object.
  * - Arg 2: optional context; can be a unit or a Lua table.
  * - Ret 1: Result of the formula.
  */
@@ -273,7 +363,25 @@ int lua_formula_bridge::intf_compile_formula(lua_State* L)
 	if(!lua_isstring(L, 1)) {
 		luaW_type_error(L, 1, "string");
 	}
-	new(L) fwrapper(lua_tostring(L, 1));
+	functionstb_ptr fcns = nullptr;
+	if(luaL_testudata(L, 2, formfcntbKey)) {
+		fcns = *static_cast<functionstb_ptr*>(lua_touserdata(L, 2));
+	} else if(lua_istable(L, 2)) {
+		// Create a functions table on the Lua stack
+		fcns = *new(L) functionstb_ptr(new function_symbol_table);
+		luaL_setmetatable(L, formfcntbKey);
+		// Set it to be collected
+		lua_owned_fcntb.insert(fcns);
+		// Iterate through the Lua table and register each entry as a function
+		int i_fcntb = lua_absindex(L, -1);
+		for(lua_pushnil(L); lua_next(L, i_fcntb - 1); lua_pop(L, 1)) {
+			// Just forward to the metafunction, which handles any error checking
+			lua_pushvalue(L, i_fcntb + 1); // The key (hopefully a valid string)
+			lua_pushvalue(L, i_fcntb + 2); // The value (hopefully a valid function)
+			lua_settable(L, i_fcntb); // Invokes the __newindex metafunction
+		}
+	}
+	new(L) fwrapper(lua_tostring(L, 1), fcns);
 	luaL_setmetatable(L, formulaKey);
 	return 1;
 }
@@ -283,6 +391,11 @@ lua_formula_bridge::fwrapper::fwrapper(const std::string& code, function_symbol_
 {
 }
 
+lua_formula_bridge::fwrapper::fwrapper(std::shared_ptr<formula_expression> expr)
+	: expr_ptr(expr)
+{
+}
+
 std::string lua_formula_bridge::fwrapper::str() const
 {
 	if(formula_ptr) {
@@ -291,10 +404,30 @@ std::string lua_formula_bridge::fwrapper::str() const
 	return "";
 }
 
+std::shared_ptr<formula> lua_formula_bridge::fwrapper::to_formula() const
+{
+	if(formula_ptr) {
+		return formula_ptr;
+	}
+	return nullptr;
+}
+
+std::shared_ptr<formula_expression> lua_formula_bridge::fwrapper::to_expr() const
+{
+	if(formula_ptr) {
+		return formula_ptr->expr_;
+	} else if(expr_ptr) {
+		return expr_ptr;
+	}
+	return nullptr;
+}
+
 variant lua_formula_bridge::fwrapper::evaluate(const formula_callable& variables, formula_debugger* fdb) const
 {
 	if(formula_ptr) {
 		return formula_ptr->evaluate(variables, fdb);
+	} else if(expr_ptr) {
+		return expr_ptr->evaluate(variables, fdb);
 	}
 	return variant();
 }
@@ -306,6 +439,27 @@ static int impl_formula_collect(lua_State* L)
 	return 0;
 }
 
+static int impl_formula_get(lua_State* L)
+{
+	lua_formula_bridge::fwrapper* form = static_cast<lua_formula_bridge::fwrapper*>(lua_touserdata(L, 1));
+	std::string key = luaL_checkstring(L, 2);
+	if(key == "functions") {
+		if(std::shared_ptr<formula> f = form->to_formula()) {
+			// This should not happen, but just in case...
+			if(f->get_functions() == nullptr) return 0;
+			new(L) functionstb_ptr(f->get_functions());
+			luaL_setmetatable(L, formfcntbKey);
+			return 1;
+		} else {
+			// The fwrapper is either invalid or wraps just a formula_expression, not a full formula.
+			// This form is only available from being passed to a function, I think.
+			lua_pushstring(L, "can't get functions table from argument formula");
+			return lua_error(L);
+		}
+	}
+	return 0;
+}
+
 static int impl_formula_tostring(lua_State* L)
 {
 	lua_formula_bridge::fwrapper* form = static_cast<lua_formula_bridge::fwrapper*>(lua_touserdata(L, 1));
@@ -314,6 +468,117 @@ static int impl_formula_tostring(lua_State* L)
 	return 1;
 }
 
+static int impl_fcntb_collect(lua_State* L)
+{
+	functionstb_ptr tb = *static_cast<functionstb_ptr*>(lua_touserdata(L, 1));
+	if(lua_owned_fcntb.count(tb)) {
+		lua_owned_fcntb.erase(tb);
+		tb->~function_symbol_table();
+	}
+	return 0;
+}
+
+static int cfun_fcntb_call(lua_State* L)
+{
+	functionstb_ptr tb = *static_cast<functionstb_ptr*>(lua_touserdata(L, lua_upvalueindex(1)));
+	std::string name = luaL_checkstring(L, lua_upvalueindex(2));
+
+	std::vector<expression_ptr> args;
+	for(int i = 1; i <= lua_gettop(L); i++) {
+		if(luaL_testudata(L, i, formulaKey)) {
+			lua_formula_bridge::fwrapper* form = static_cast<lua_formula_bridge::fwrapper*>(lua_touserdata(L, i));
+			args.push_back(form->to_expr());
+		} else {
+			variant v = luaW_tofaivariant(L, i);
+			std::string var;
+			if(!v.is_callable()) {
+				var = v.serialize_to_string();
+			} else if(auto lc = v.try_convert<lua_callable>()) {
+				std::vector<formula_input> inputs;
+				lc->get_inputs(inputs);
+				if(inputs.size() > 2) {
+					// Treat it as a map since it has string keys
+					variant vmap = lc->get_value("__map");
+					var = vmap.serialize_to_string();
+				} else {
+					// It could be a list or a map; check all keys
+					if(lc->is_list()) {
+						variant vlist = lc->get_value("__list");
+						var = vlist.serialize_to_string();
+					} else {
+						variant vmap = lc->get_value("__map");
+						var = vmap.serialize_to_string();
+					}
+				}
+			} else {
+				// Just try to serialize it. It might actually work,
+				// or it might throw an error.
+				// (For example, location objects are serializable.)
+				var = v.serialize_to_string();
+			}
+			if(!var.empty()) {
+				// It's not an object, so push code for generating the value
+				lua_formula_bridge::fwrapper form(var, tb);
+				args.push_back(form.to_expr());
+			}
+		}
+	}
+
+	expression_ptr fcn = tb->create_function(name, args);
+	new(L) lua_formula_bridge::fwrapper(fcn);
+	luaL_setmetatable(L, formulaKey);
+	return 1;
+}
+
+static int impl_fcntb_get(lua_State* L)
+{
+	if(!luaL_testudata(L, 1, formfcntbKey)) {
+		return luaW_type_error(L, 2, "formula function table");
+	}
+	if(!lua_isstring(L, 2)) {
+		return luaW_type_error(L, 2, lua_typename(L, lua_type(L, 2)));
+	}
+	// This closure simply consumes the two arguments to __index
+	lua_pushcclosure(L, cfun_fcntb_call, 2);
+	return 1;
+}
+
+static int impl_fcntb_set(lua_State* L)
+{
+	functionstb_ptr tb = *static_cast<functionstb_ptr*>(lua_touserdata(L, 1));
+	std::string fcn_name = luaL_checkstring(L, 2);
+	if(lua_iscfunction(L, 3)) {
+		return luaL_argerror(L, 3, "expected Lua function, not C function");
+	} else if(!lua_isfunction(L, 3)) {
+		return luaW_type_error(L, 3, lua_typename(L, LUA_TFUNCTION));
+	}
+	if(formula::keywords.count(fcn_name)) {
+		return luaL_argerror(L, 2, "can't create a function whose name is a formula keyword");
+	}
+	// Make sure the function is at the top
+	// It probably is already, though
+	lua_settop(L, 3);
+	lua_pushvalue(L, -1);
+	lua_Debug fcn_info;
+	lua_getinfo(L, ">u", &fcn_info);
+	if(fcn_info.nparams == 0) {
+		return luaL_argerror(L, 3, "expected Lua function taking at least one named argument");
+	}
+	int min_args = 1, max_args = fcn_info.isvararg ? -1 : fcn_info.nparams;
+	for(int i = 2; i <= fcn_info.nparams; i++) {
+		std::string var_name = lua_getlocal(L, nullptr, i);
+		if(var_name.find("_opt") != std::string::npos) {
+			min_args--;
+		}
+	}
+	lua_kernel_base& kernel = lua_kernel_base::get_lua_kernel<lua_kernel_base>(L);
+	formula_function_ptr fcn(new lua_formula_function_defn(fcn_name, min_args - 1, max_args - 1, kernel));
+	lua_pushvalue(L, 3);
+	lua_rawsetp(L, LUA_REGISTRYINDEX, fcn.get());
+	tb->add_function(fcn_name, std::move(fcn));
+	return 0;
+}
+
 std::string lua_formula_bridge::register_metatables(lua_State* L)
 {
 	luaL_newmetatable(L, formulaKey);
@@ -323,8 +588,20 @@ std::string lua_formula_bridge::register_metatables(lua_State* L)
 	lua_setfield(L, -2, "__tostring");
 	lua_pushcfunction(L, intf_eval_formula);
 	lua_setfield(L, -2, "__call");
+	lua_pushcfunction(L, impl_formula_get);
+	lua_setfield(L, -2, "__index");
 	lua_pushstring(L, "formula");
 	lua_setfield(L, -2, "__metatable");
 
-	return "Adding formula metatable...\n";
+	luaL_newmetatable(L, formfcntbKey);
+	lua_pushcfunction(L, impl_fcntb_collect);
+	lua_setfield(L, -2, "__gc");
+	lua_pushcfunction(L, impl_fcntb_get);
+	lua_setfield(L, -2, "__index");
+	lua_pushcfunction(L, impl_fcntb_set);
+	lua_setfield(L, -2, "__newindex");
+	lua_pushstring(L, "formula functions");
+	lua_setfield(L, -2, "__metatable");
+
+	return "Adding formula metatable...\nAdding formula function table metatable";
 }

src/scripting/lua_formula_bridge.hpp

@@ -25,6 +25,7 @@ namespace wfl {
 	class function_symbol_table;
 	class formula_debugger;
 	class formula_callable;
+	class formula_expression;
 	class variant;
 }
 
@@ -36,10 +37,14 @@ namespace lua_formula_bridge {
 
 	class fwrapper {
 		std::shared_ptr<wfl::formula> formula_ptr;
+		std::shared_ptr<wfl::formula_expression> expr_ptr;
 	public:
 		fwrapper(const std::string& code, wfl::function_symbol_table* functions = nullptr);
+		fwrapper(std::shared_ptr<wfl::formula_expression> expr);
 		std::string str() const;
 		wfl::variant evaluate(const wfl::formula_callable& variables, wfl::formula_debugger* fdb = nullptr) const;
+		std::shared_ptr<wfl::formula> to_formula() const;
+		std::shared_ptr<wfl::formula_expression> to_expr() const;
 	};
 
 	using fpointer = std::unique_ptr<fwrapper, std::function<void(fwrapper*)>>;