From 81ee0f0349435fbccd7fc50c2573ade63316591e Mon Sep 17 00:00:00 2001
From: Celtic Minstrel <celtic.minstrel.ca@some.place>
Date: Sat, 23 Nov 2019 14:15:30 -0500
Subject: [PATCH] Make AI aspect handling in Lua more robust

There should never again be a bug where aspects of one type return nil instead of their value.

If support for a new type isn't explicitly added either in lua_object or in impl_ai_aspect_get, it'll raise a runtime error instead of returning nil.
---
 src/ai/composite/aspect.hpp | 10 +++++
 src/ai/lua/core.cpp         | 23 ++---------
 src/ai/lua/lua_object.hpp   | 79 +++++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 20 deletions(-)

diff --git a/src/ai/composite/aspect.hpp b/src/ai/composite/aspect.hpp
index 1333a61fb5fe..dcaf2507a20b 100644
--- a/src/ai/composite/aspect.hpp
+++ b/src/ai/composite/aspect.hpp
@@ -52,6 +52,9 @@ class aspect : public readonly_context_proxy, public events::observer, public co
 	virtual std::shared_ptr<wfl::variant> get_variant_ptr() const = 0;
 
 
+	virtual void get_lua(lua_State* L) const = 0;
+
+
 	virtual void recalculate() const = 0;
 
 
@@ -152,6 +155,13 @@ class typesafe_aspect : public aspect {
 		return value_variant_;
 	}
 
+	void get_lua(lua_State* L) const {
+		if(auto p = get_ptr()) {
+			lua_object<T> obj(get());
+			obj.push(L);
+		} else lua_pushnil(L);
+	}
+
 	virtual void recalculate() const = 0;
 
 
diff --git a/src/ai/lua/core.cpp b/src/ai/lua/core.cpp
index 3da2664e6217..75ae58bd2118 100644
--- a/src/ai/lua/core.cpp
+++ b/src/ai/lua/core.cpp
@@ -775,24 +775,8 @@ static int impl_ai_aspect_get(lua_State* L)
 		return 0;
 	}
 
-	typedef std::vector<std::string> string_list;
-	if(typesafe_aspect<bool>* aspect_as_bool = try_aspect_as<bool>(iter->second)) {
-		lua_pushboolean(L, aspect_as_bool->get());
-	} else if(typesafe_aspect<int>* aspect_as_int = try_aspect_as<int>(iter->second)) {
-		lua_pushinteger(L, aspect_as_int->get());
-	} else if(typesafe_aspect<double>* aspect_as_double = try_aspect_as<double>(iter->second)) {
-		lua_pushnumber(L, aspect_as_double->get());
-	} else if(typesafe_aspect<std::string>* aspect_as_string = try_aspect_as<std::string>(iter->second)) {
-		lua_pushstring(L, aspect_as_string->get().c_str());
-	} else if(typesafe_aspect<config>* aspect_as_config = try_aspect_as<config>(iter->second)) {
-		luaW_pushconfig(L, aspect_as_config->get());
-	} else if(typesafe_aspect<string_list>* aspect_as_string_list = try_aspect_as<string_list>(iter->second)) {
-		lua_push(L, aspect_as_string_list->get());
-	} else if(typesafe_aspect<terrain_filter>* aspect_as_terrain_filter = try_aspect_as<terrain_filter>(iter->second)) {
-		std::set<map_location> result;
-		aspect_as_terrain_filter->get().get_locations(result);
-		lua_push(L, result);
-	} else if(typesafe_aspect<attacks_vector>* aspect_as_attacks_vector = try_aspect_as<attacks_vector>(iter->second)) {
+	// A few aspects require special delicate handling...
+	if(typesafe_aspect<attacks_vector>* aspect_as_attacks_vector = try_aspect_as<attacks_vector>(iter->second)) {
 		using ai_default_rca::aspect_attacks_base;
 		aspect_attacks_base* real_aspect = dynamic_cast<aspect_attacks_base*>(aspect_as_attacks_vector);
 		while(real_aspect == nullptr) {
@@ -826,7 +810,6 @@ static int impl_ai_aspect_get(lua_State* L)
 			lua_rawseti(L, -2, i + 1);
 		}
 		lua_setfield(L, -2, "enemy");
-		return 1;
 	} else if(typesafe_aspect<unit_advancements_aspect>* aspect_as_unit_advancements_aspects = try_aspect_as<unit_advancements_aspect>(iter->second)) {
 		const unit_advancements_aspect& val = aspect_as_unit_advancements_aspects->get();
 		int my_side = get_engine(L).get_readonly_context().get_side();
@@ -841,7 +824,7 @@ static int impl_ai_aspect_get(lua_State* L)
 			lua_settable(L, -3);
 		}
 	} else {
-		lua_pushnil(L);
+		iter->second->get_lua(L);
 	}
 	return 1;
 }
diff --git a/src/ai/lua/lua_object.hpp b/src/ai/lua/lua_object.hpp
index 53c639443754..9a80075c4132 100644
--- a/src/ai/lua/lua_object.hpp
+++ b/src/ai/lua/lua_object.hpp
@@ -59,6 +59,12 @@ class lua_object : public lua_object_base
 	{
 		// empty
 	}
+	
+	lua_object(const T& init)
+		: value_(std::make_shared<T>(init))
+	{
+		// empty
+	}
 
 	std::shared_ptr<T> get()
 	{
@@ -69,6 +75,11 @@ class lua_object : public lua_object_base
 	{
 		this->value_ = to_type(L, lua_absindex(L, n));
 	}
+	
+	void push(lua_State* L)
+	{
+		from_type(L, this->value_);
+	}
 
 protected:
 
@@ -77,6 +88,13 @@ class lua_object : public lua_object_base
 	{
 		return std::shared_ptr<T>();
 	}
+	
+	// A group of functions that deal with the translations of values back to Lua
+	void from_type(lua_State* L, std::shared_ptr<T>)
+	{
+		lua_pushliteral(L, "Unsupported AI aspect type for Lua!");
+		lua_error(L);
+	}
 
 	std::shared_ptr<T> value_;
 };
@@ -87,24 +105,52 @@ inline std::shared_ptr<double> lua_object<double>::to_type(lua_State *L, int n)
 	return std::make_shared<double>(lua_tonumber(L, n));
 }
 
+template <>
+inline void lua_object<double>::from_type(lua_State *L, std::shared_ptr<double> value)
+{
+	if(value) lua_pushnumber(L, *value);
+	else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr<std::string> lua_object<std::string>::to_type(lua_State *L, int n)
 {
 	return std::make_shared<std::string>(lua_tostring(L, n));
 }
 
+template <>
+inline void lua_object<std::string>::from_type(lua_State *L, std::shared_ptr<std::string> value)
+{
+	if(value) lua_pushlstring(L, value->c_str(), value->size());
+	else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr<bool> lua_object<bool>::to_type(lua_State *L, int n)
 {
 	return std::make_shared<bool>(luaW_toboolean(L, n));
 }
 
+template <>
+inline void lua_object<bool>::from_type(lua_State *L, std::shared_ptr<bool> value)
+{
+	if(value) lua_pushboolean(L, *value);
+	else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr<int> lua_object<int>::to_type(lua_State *L, int n)
 {
 	return std::make_shared<int>(static_cast<int>(lua_tointeger(L, n)));
 }
 
+template <>
+inline void lua_object<int>::from_type(lua_State *L, std::shared_ptr<int> value)
+{
+	if(value) lua_pushnumber(L, *value);
+	else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr< std::vector<std::string> > lua_object< std::vector<std::string> >::to_type(lua_State *L, int n)
 {
@@ -122,6 +168,18 @@ inline std::shared_ptr< std::vector<std::string> > lua_object< std::vector<std::
 	return v;
 }
 
+template <>
+inline void lua_object< std::vector<std::string> >::from_type(lua_State *L, std::shared_ptr< std::vector<std::string> > value)
+{
+	if(value) {
+		lua_createtable(L, value->size(), 0);
+		for(const std::string& str : *value) {
+			lua_pushlstring(L, str.c_str(), str.size());
+			lua_rawseti(L, -1, lua_rawlen(L, -2) + 1);
+		}
+	} else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr<config> lua_object<config>::to_type(lua_State *L, int n)
 {
@@ -130,6 +188,13 @@ inline std::shared_ptr<config> lua_object<config>::to_type(lua_State *L, int n)
 	return cfg;
 }
 
+template <>
+inline void lua_object<config>::from_type(lua_State *L, std::shared_ptr<config> value)
+{
+	if(value) luaW_pushconfig(L, *value);
+	else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr<terrain_filter> lua_object<terrain_filter>::to_type(lua_State *L, int n)
 {
@@ -143,6 +208,20 @@ inline std::shared_ptr<terrain_filter> lua_object<terrain_filter>::to_type(lua_S
 	return tf;
 }
 
+template <>
+inline void lua_object<terrain_filter>::from_type(lua_State *L, std::shared_ptr<terrain_filter> value)
+{
+	if(value) {
+		std::set<map_location> locs;
+		value->get_locations(locs);
+		lua_createtable(L, locs.size(), 0);
+		for(const map_location& loc : locs) {
+			luaW_pushlocation(L, loc);
+			lua_rawseti(L, -1, lua_rawlen(L, -2) + 1);
+		}
+	} else lua_pushnil(L);
+}
+
 template <>
 inline std::shared_ptr<std::vector<target> > lua_object< std::vector<target> >::to_type(lua_State *L, int n)
 {