Lua: Support the extra wesnoth.match_unit arguments also in wesnoth.get_units

changelog

@@ -17,6 +17,7 @@ Version 1.13.8+dev:
      though some of its contents are missing.
    * New extra argument to wesnoth.match_location and wesnoth.get_locations
      allows specifying the teleport_unit.
+   * Support the extra wesnoth.match_unit arguments also in wesnoth.get_units
  * Multiplayer:
    * When set to remember your password, Wesnoth now encrypts it. It is still
      possible to obtain the password from preferences, but it's no longer as

src/scripting/game_lua_kernel.cpp

@@ -547,6 +547,8 @@ int game_lua_kernel::intf_get_displayed_unit(lua_State *L)
 /**
  * Gets all the units matching a given filter.
  * - Arg 1: optional table containing a filter
+ * - Arg 2: optional location (to find all units that would match on that location)
+                  OR unit (to find all units that would match adjacent to that unit)
  * - Ret 1: table containing full userdata with __index pointing to
  *          impl_unit_get and __newindex pointing to impl_unit_set.
  */
@@ -554,15 +556,34 @@ int game_lua_kernel::intf_get_units(lua_State *L)
 {
 	vconfig filter = luaW_checkvconfig(L, 1, true);
 
+	// note that if filter is null, this yields a null filter matching everything (and doing no work)
+	filter_context & fc = game_state_;
+	unit_filter filt(filter, &fc);
+	std::vector<const unit*> units;
+
+	if(unit* u_adj = luaW_tounit(L, 2)) {
+		if(!u_adj) {
+			return luaL_argerror(L, 2, "unit not found");
+		}
+		units = filt.all_matches_with_unit(*u_adj);
+	} else if(!lua_isnoneornil(L, 2)) {
+		map_location loc;
+		luaW_tolocation(L, 2, loc);
+		if(!loc.valid()) {
+			return luaL_argerror(L, 2, "invalid location");
+		}
+		units = filt.all_matches_at(loc);
+	} else {
+		units = filt.all_matches_on_map();
+	}
+
 	// Go through all the units while keeping the following stack:
 	// 1: return table, 2: userdata
 	lua_settop(L, 0);
 	lua_newtable(L);
 	int i = 1;
 
-	// note that if filter is null, this yields a null filter matching everything (and doing no work)
-	filter_context & fc = game_state_;
-	for (const unit * ui : unit_filter(filter, &fc).all_matches_on_map()) {
+	for (const unit * ui : units) {
 		luaW_pushunit(L, ui->underlying_id());
 		lua_rawseti(L, 1, i);
 		++i;

src/units/filter.cpp

@@ -80,7 +80,7 @@ public:
 	virtual bool matches(const unit & /*u*/, const map_location & /*loc*/, const unit *) const {
 		return true;
 	}
-	virtual std::vector<const unit *> all_matches_on_map(unsigned max_matches) const {
+	virtual std::vector<const unit *> all_matches_on_map(unsigned max_matches, const map_location* = nullptr, const unit* = nullptr) const {
 		std::vector<const unit *> ret;
 		for(const unit & u : fc_.get_disp_context().units()) {
 			--max_matches;
@@ -168,7 +168,7 @@ public:
 	}
 
 	virtual bool matches(const unit & u, const map_location & loc, const unit * u2) const;
-	virtual std::vector<const unit *> all_matches_on_map(unsigned max_matches) const;
+	virtual std::vector<const unit *> all_matches_on_map(unsigned max_matches, const map_location* loc = nullptr, const unit* u2 = nullptr) const;
 	virtual unit_const_ptr first_match_on_map() const;
 	config to_config() const {
 		return vcfg.get_config();
@@ -722,10 +722,10 @@ bool basic_unit_filter_impl::internal_matches_filter(const unit & u, const map_l
 	return true;
 }
 
-std::vector<const unit *> basic_unit_filter_impl::all_matches_on_map(unsigned max_matches) const {
+std::vector<const unit *> basic_unit_filter_impl::all_matches_on_map(unsigned max_matches, const map_location* loc, const unit* u2) const {
 	std::vector<const unit *> ret;
 	for (const unit & u : fc_.get_disp_context().units()) {
-		if (matches(u, u.get_location(), nullptr)) {
+		if (matches(u, loc ? *loc : u.get_location(), u2)) {
 			if(max_matches == 0) {
 				return ret;
 			}

src/units/filter.hpp

@@ -39,7 +39,7 @@ struct map_location;
 class unit_filter_abstract_impl {
 public:
 	virtual bool matches(const unit & u, const map_location & loc, const unit * u2 = nullptr) const = 0;
-	virtual std::vector<const unit*> all_matches_on_map(unsigned max_matches) const = 0;
+	virtual std::vector<const unit*> all_matches_on_map(unsigned max_matches, const map_location* loc = nullptr, const unit* u2 = nullptr) const = 0;
 	virtual unit_const_ptr first_match_on_map() const = 0;
 	virtual config to_config() const = 0;
 	virtual bool empty() const {return false;}
@@ -95,6 +95,18 @@ public:
 		return impl_->all_matches_on_map(max_matches_);
 	}
 
+	std::vector<const unit*> all_matches_at(const map_location& loc) const {
+		return impl_->all_matches_on_map(max_matches_, &loc);
+	}
+
+	std::vector<const unit*> all_matches_with_unit(const unit& u) const {
+		return impl_->all_matches_on_map(max_matches_, nullptr, &u);
+	}
+
+	std::vector<const unit*> all_matches_with_unit_at(const unit& u, const map_location& loc) const {
+		return impl_->all_matches_on_map(max_matches_, &loc, &u);
+	}
+
 	unit_const_ptr first_match_on_map() const {
 		return impl_->first_match_on_map();
 	}