Fix Lua unit types table and add variations subtable to each type

src/scripting/game_lua_kernel.cpp

@@ -4219,6 +4219,9 @@ game_lua_kernel::game_lua_kernel(CVideo * video, game_state & gs, play_controlle
 	// Create the vconfig metatable.
 	cmd_log_ << lua_common::register_vconfig_metatable(L);
 
+	// Create the unit_types table
+	cmd_log_ << lua_unit_type::register_table(L);
+
 	// Create the ai elements table.
 	cmd_log_ << "Adding ai elements table...\n";
 
@@ -4343,19 +4346,6 @@ void game_lua_kernel::initialize(const config& level)
 		cmd_log_ << "Added wesnoth.sides\n";
 	}
 
-	// Create the unit_types table.
-	cmd_log_ << "Adding unit_types table...\n";
-
-	lua_settop(L, 0);
-	lua_getglobal(L, "wesnoth");
-	lua_newtable(L);
-	for (const unit_type_data::unit_type_map::value_type &ut : unit_types.types())
-	{
-		luaW_pushunittype(L, ut.first);
-		lua_setfield(L, -2, ut.first.c_str());
-	}
-	lua_setfield(L, -2, "unit_types");
-
 	//Create the races table.
 	cmd_log_ << "Adding races table...\n";
 

src/scripting/lua_unit_type.cpp

@@ -30,7 +30,8 @@
  */
 
 // Registry key
-static const char * UnitType = "unit type";
+static const char UnitType[] = "unit type";
+static const char UnitTypeTable[] = "unit types";
 
 /**
  * Gets some data on a unit type (__index metamethod).
@@ -40,12 +41,8 @@ static const char * UnitType = "unit type";
  */
 static int impl_unit_type_get(lua_State *L)
 {
+	const unit_type& ut = luaW_checkunittype(L, 1);
 	char const *m = luaL_checkstring(L, 2);
-	lua_pushstring(L, "id");
-	lua_rawget(L, 1);
-	const unit_type *utp = unit_types.find(lua_tostring(L, -1));
-	if (!utp) return luaL_argerror(L, 1, "unknown unit type");
-	unit_type const &ut = *utp;
 
 	// Find the corresponding attribute.
 	return_tstring_attrib("name", ut.type_name());
@@ -77,9 +74,86 @@ static int impl_unit_type_get(lua_State *L)
 		push_unit_attacks_table(L, 1);
 		return 1;
 	}
+	if(strcmp(m, "variations") == 0) {
+		lua_createtable(L, 0, 1);
+		lua_pushvalue(L, 1);
+		lua_setfield(L, -2, "base");
+		luaL_setmetatable(L, UnitTypeTable);
+		return 1;
+	}
 	return 0;
 }
 
+static int impl_unit_type_equal(lua_State* L)
+{
+	const unit_type& ut1 = luaW_checkunittype(L, 1);
+	if(const unit_type* ut2 = luaW_tounittype(L, 2)) {
+		lua_pushboolean(L, &ut1 == ut2);
+	} else {
+		lua_pushboolean(L, false);
+	}
+	return 1;
+}
+
+static int impl_unit_type_lookup(lua_State* L)
+{
+	lua_pushstring(L, "base");
+	lua_rawget(L, 1);
+	std::string id = luaL_checkstring(L, 2);
+	const unit_type* ut;
+	if(const unit_type* base = luaW_tounittype(L, -1)) {
+		if(id == "male" || id == "female") {
+			ut = &base->get_gender_unit_type(id);
+		} else {
+			ut = &base->get_variation(id);
+		}
+	} else {
+		ut = unit_types.find(id);
+	}
+	luaW_pushunittype(L, ut);
+	return 1;
+}
+
+static int impl_unit_type_new(lua_State* L)
+{
+	// This could someday become a hook to construct new unit types on the fly?
+	// For now though, we just want to avoid the __index callback not being called
+	// because keys got set in the table.
+	lua_pushstring(L, "unit_types table is read-only");
+	return lua_error(L);
+}
+
+static int impl_unit_type_next(lua_State* L)
+{
+	lua_pushstring(L, "base");
+	lua_rawget(L, 1);
+	const unit_type* base = luaW_tounittype(L, -1);
+	auto unit_map = base ? base->variation_types() : unit_types.types();
+	decltype(unit_map)::const_iterator it;
+	if(lua_isnoneornil(L, 2)) {
+		it = unit_map.begin();
+	} else {
+		it = unit_map.find(luaL_checkstring(L, 2));
+		if(it == unit_map.end()) {
+			return 0;
+		}
+		++it;
+	}
+	if (it == unit_map.end()) {
+		return 0;
+	}
+	lua_pushlstring(L, it->first.c_str(), it->first.size());
+	luaW_pushunittype(L, &it->second);
+	return 2;
+}
+
+static int impl_unit_type_pairs(lua_State* L) {
+	lua_pushcfunction(L, &impl_unit_type_next);
+	lua_pushvalue(L, -2);
+	lua_pushnil(L);
+	return 3;
+}
+
 namespace lua_unit_type {
 	std::string register_metatable(lua_State * L)
 	{
@@ -87,17 +161,50 @@ namespace lua_unit_type {
 
 		lua_pushcfunction(L, impl_unit_type_get);
 		lua_setfield(L, -2, "__index");
-		lua_pushstring(L, "unit type");
+		lua_pushcfunction(L, impl_unit_type_equal);
+		lua_setfield(L, -2, "__eq");
+		lua_pushstring(L, UnitType);
 		lua_setfield(L, -2, "__metatable");
 
 		return "Adding unit type metatable...\n";
 	}
+
+	std::string register_table(lua_State* L)
+	{
+		lua_getglobal(L, "wesnoth");
+		lua_newtable(L);
+		luaL_newmetatable(L, UnitTypeTable);
+		lua_pushcfunction(L, impl_unit_type_lookup);
+		lua_setfield(L, -2, "__index");
+		lua_pushcfunction(L, impl_unit_type_new);
+		lua_setfield(L, -2, "__newindex");
+		lua_pushcfunction(L, impl_unit_type_pairs);
+		lua_setfield(L, -2, "__pairs");
+		lua_pushstring(L, UnitTypeTable);
+		lua_setfield(L, -2, "__metatable");
+		lua_setmetatable(L, -2);
+		lua_setfield(L, -2, "unit_types");
+		lua_pop(L, 1);
+
+		return "Adding unit_types table...\n";
+	}
 }
 
-void luaW_pushunittype(lua_State *L, const std::string & id)
+void luaW_pushunittype(lua_State *L, const unit_type* ut)
 {
-	lua_createtable(L, 0, 1);
-	lua_pushstring(L, id.c_str());
-	lua_setfield(L, -2, "id");
+	*static_cast<const unit_type**>(lua_newuserdata(L, sizeof(unit_type*))) = ut;
 	luaL_setmetatable(L, UnitType);
 }
+
+const unit_type* luaW_tounittype(lua_State* L, int idx)
+{
+	if(void* p = luaL_testudata(L, idx, UnitType)) {
+		return *static_cast<const unit_type**>(p);
+	}
+	return nullptr;
+}
+
+const unit_type& luaW_checkunittype(lua_State* L, int idx)
+{
+	return **static_cast<const unit_type**>(luaL_checkudata(L, idx, UnitType));;
+}

src/scripting/lua_unit_type.hpp

@@ -16,6 +16,7 @@
 #define LUA_UNIT_TYPE_HPP_INCLUDED
 
 struct lua_State;
+class unit_type;
 
 #include <string>
 
@@ -25,10 +26,17 @@ struct lua_State;
  */
 namespace lua_unit_type {
 	std::string register_metatable(lua_State *);
+	std::string register_table(lua_State *);
 } //end namespace lua_team
 
 /// Create a lua object containing a reference to a unittype, and a
 /// metatable to access the properties.
-void luaW_pushunittype(lua_State *, const std::string &);
+void luaW_pushunittype(lua_State *, const unit_type*);
+
+/// Test if a stack element is a unit type, and return it if so
+const unit_type* luaW_tounittype(lua_State*, int);
+
+/// Test if a stack element is a unit type, and return it if so
+const unit_type& luaW_checkunittype(lua_State*, int);
 
 #endif

src/units/types.cpp

@@ -100,10 +100,6 @@ unit_type::unit_type(const unit_type& o) :
 {
 	gender_types_[0] = o.gender_types_[0] != nullptr ? new unit_type(*o.gender_types_[0]) : nullptr;
 	gender_types_[1] = o.gender_types_[1] != nullptr ? new unit_type(*o.gender_types_[1]) : nullptr;
-
-	for(variations_map::const_iterator i = o.variations_.begin(); i != o.variations_.end(); ++i) {
-		variations_[i->first] = new unit_type(*i->second);
-	}
 }
 
 
@@ -167,10 +163,6 @@ unit_type::~unit_type()
 {
 	delete gender_types_[0];
 	delete gender_types_[1];
-
-	for(variations_map::iterator i = variations_.begin(); i != variations_.end(); ++i) {
-		delete i->second;
-	}
 }
 
 /**
@@ -211,7 +203,7 @@ void unit_type::build_full(const movement_type_map &mv_types,
 
 	// Propagate the build to the variations.
 	for (variations_map::value_type & variation : variations_) {
-		variation.second->build_full(mv_types, races, traits);
+		variation.second.build_full(mv_types, races, traits);
 	}
 
 	// Deprecation messages, only seen when unit is parsed for the first time.
@@ -351,11 +343,16 @@ void unit_type::build_help_index(const movement_type_map &mv_types,
 		const std::string& var_id = var_cfg["variation_id"].empty() ?
 				var_cfg["variation_name"] : var_cfg["variation_id"];
 
-		unit_type *ut = new unit_type(var_cfg, id_);
-		ut->debug_id_ = debug_id_ + " [" + var_id + "]";
-		ut->base_id_ = base_id_;  // In case this is not id_.
-		ut->build_help_index(mv_types, races, traits);
-		variations_.insert(std::make_pair(var_id, ut));
+		variations_map::iterator ut;
+		bool success;
+		std::tie(ut, success) = variations_.emplace(var_id, unit_type(var_cfg, id_));
+		if(success) {
+			ut->second.debug_id_ = debug_id_ + " [" + var_id + "]";
+			ut->second.base_id_ = base_id_;  // In case this is not id_.
+			ut->second.build_help_index(mv_types, races, traits);
+		} else {
+			ERR_CF << "Skipping duplicate unit variation ID: " << var_id << "\n";
+		}
 	}
 
 	hide_help_= cfg_["hide_help"].to_bool();
@@ -464,7 +461,7 @@ const unit_type& unit_type::get_variation(const std::string& id) const
 {
 	const variations_map::const_iterator i = variations_.find(id);
 	if(i != variations_.end()) {
-		return *i->second;
+		return i->second;
 	} else {
 		return *this;
 	}
@@ -626,7 +623,7 @@ void unit_type::add_advancement(const unit_type &to_unit,int xp)
 		for(variations_map::iterator v=variations_.begin();
 			v!=variations_.end(); ++v) {
 			LOG_CONFIG << "variation advancement: ";
-			v->second->add_advancement(to_unit,xp);
+			v->second.add_advancement(to_unit,xp);
 		}
 	}
 }
@@ -749,8 +746,7 @@ bool unit_type::has_variation(const std::string& variation_id) const
 bool unit_type::show_variations_in_help() const
 {
 	for (const variations_map::value_type &val : variations_) {
-		assert(val.second != nullptr);
-		if (!val.second->hide_help()) {
+		if (!val.second.hide_help()) {
 			return true;
 		}
 	}

src/units/types.hpp

@@ -73,6 +73,8 @@ private: // These will be called by build().
 	/// Load the most needed data into an empty unit_type (build to CREATE).
 	void build_created(const movement_type_map &movement_types,
 		const race_map &races, const config::const_child_itors &traits);
+
+	typedef std::map<std::string,unit_type> variations_map;
 public:
 	/// Performs a build of this to the indicated stage.
 	void build(BUILD_STATUS status, const movement_type_map &movement_types,
@@ -197,6 +199,7 @@ public:
 	/// to the HELP_INDEXED status.
 	const std::vector<unit_race::GENDER>& genders() const { return genders_; }
 	std::vector<std::string> variations() const;
+	const variations_map& variation_types() const {return variations_; }
 
 	/**
 	 * @param variation_id		The id of the variation we search for.
@@ -269,7 +272,6 @@ private:
 
 	unit_type* gender_types_[2];
 
-	typedef std::map<std::string,unit_type*> variations_map;
 	variations_map variations_;
 	std::string default_variation_;
 	std::string variation_name_;