Lua frontend to allow changing unit's storage (#1879)

beyond-all-reason · Jan 9, 2025 · 77c4366 · lhog · Jan 9, 2025 · sprunk
1 parent eeae3ce
commit 77c4366
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 0 deletions.
diff --git a/rts/Lua/LuaSyncedCtrl.cpp b/rts/Lua/LuaSyncedCtrl.cpp
@@ -164,6 +164,7 @@ bool LuaSyncedCtrl::PushEntries(lua_State* L)
 
 	REGISTER_LUA_CFUNC(SetUnitCosts);
 	REGISTER_LUA_CFUNC(SetUnitResourcing);
+	REGISTER_LUA_CFUNC(SetUnitStorage);
 	REGISTER_LUA_CFUNC(SetUnitTooltip);
 	REGISTER_LUA_CFUNC(SetUnitHealth);
 	REGISTER_LUA_CFUNC(SetUnitMaxHealth);
@@ -1887,6 +1888,32 @@ static bool SetUnitResourceParam(CUnit* unit, const char* name, float value)
 }
 
 
+static bool SetUnitStorageParam(CUnit* unit, const char* name, float value)
+{
+	// [m|e]
+	//
+	// metal | energy
+
+	SResourcePack newStorage = unit->storage;
+
+	switch (name[0]) {
+		case 'm': {
+			newStorage.metal  = value;
+		} break;
+
+		case 'e': {
+			newStorage.energy = value;
+		} break;
+
+		default: {
+			return false;
+		}
+	}
+	unit->SetStorage(newStorage);
+	return true;
+}
+
+
 /***
  * Unit Resourcing
  * @section unitresourcing
@@ -1933,6 +1960,52 @@ int LuaSyncedCtrl::SetUnitResourcing(lua_State* L)
 }
 
 
+/***
+ * Unit Storage
+ * @section unitstorage
+ */
+
+/***
+ * @function Spring.SetUnitStorage
+ * @number unitID
+ * @string res
+ * @number amount
+ * @treturn nil
+ */
+
+/***
+ * @function Spring.SetUnitStorage
+ * @number unitID
+ * @tparam {[string]=number,...} res keys are: "[m|e]" metal | energy. Values are amounts
+ * @treturn nil
 /*** Set unit position (2D) 
  * @function Spring.SetUnitPosition 
  * 
  * Sets a unit's position in 2D, at terrain height. 
  * 
  * @number unitID 
  * @number x 
  * @number z 
  * @bool[opt=false] floating If true, over water the position is on surface. If false, on seafloor. 
  * @treturn nil 
  */ 
 /*** Set unit position (3D) 
  * @function Spring.SetUnitPosition 
  * 
  * Sets a unit's position in 3D, at an arbitrary height. 
  * 
  * @number unitID 
  * @number x 
  * @number y 
  * @number z 
  * @treturn nil 
  */ 
 /*** Set unit position (2D) 
  * @function Spring.SetUnitPosition 
  * 
  * Sets a unit's position in 2D, at terrain height. 
  * 
  * @number unitID 
  * @number x 
  * @number z 
  * @bool[opt=false] floating If true, over water the position is on surface. If false, on seafloor. 
  * @treturn nil 
  */ 
  
  
 /*** Set unit position (3D) 
  * @function Spring.SetUnitPosition 
  * 
  * Sets a unit's position in 3D, at an arbitrary height. 
  * 
  * @number unitID 
  * @number x 
  * @number y 
  * @number z 
  * @treturn nil 
  */ 
+ */
+int LuaSyncedCtrl::SetUnitStorage(lua_State* L)
+{
+	CUnit* unit = ParseUnit(L, __func__, 1);
+
+	if (unit == nullptr)
+		return 0;
+
+	if (lua_israwstring(L, 2)) {
+		SetUnitStorageParam(unit, lua_tostring(L, 2), luaL_checkfloat(L, 3));
+	} else if (lua_istable(L, 2)) {
+		constexpr int tableIdx = 2;
+
+		for (lua_pushnil(L); lua_next(L, tableIdx) != 0; lua_pop(L, 1)) {
+			if (!lua_israwstring(L, LUA_TABLE_KEY_INDEX) || !lua_isnumber(L, LUA_TABLE_VALUE_INDEX))
+				continue;
+
+			SetUnitStorageParam(unit, lua_tostring(L, LUA_TABLE_KEY_INDEX), lua_tofloat(L, LUA_TABLE_VALUE_INDEX));
+		}
+	}
+	else {
+		luaL_error(L, "Incorrect arguments to SetUnitStorage");
+	}
+
+	return 0;
+}
+
+
 /***
  * @function Spring.SetUnitTooltip
  * @number unitID

diff --git a/rts/Lua/LuaSyncedCtrl.h b/rts/Lua/LuaSyncedCtrl.h
@@ -93,6 +93,7 @@ class LuaSyncedCtrl
 		static int SetUnitAlwaysVisible(lua_State* L);
 		static int SetUnitUseAirLos(lua_State* L);
 		static int SetUnitResourcing(lua_State* L);
+		static int SetUnitStorage(lua_State* L);
 		static int SetUnitMetalExtraction(lua_State* L);
 		static int SetUnitHarvestStorage(lua_State* L);
 		static int SetUnitBuildSpeed(lua_State* L);

diff --git a/rts/Lua/LuaSyncedRead.cpp b/rts/Lua/LuaSyncedRead.cpp
@@ -207,6 +207,7 @@ bool LuaSyncedRead::PushEntries(lua_State* L)
 	REGISTER_LUA_CFUNC(GetUnitIsStunned);
 	REGISTER_LUA_CFUNC(GetUnitIsBeingBuilt);
 	REGISTER_LUA_CFUNC(GetUnitResources);
+	REGISTER_LUA_CFUNC(GetUnitStorage);
 	REGISTER_LUA_CFUNC(GetUnitCosts);
 	REGISTER_LUA_CFUNC(GetUnitCostTable);
 	REGISTER_LUA_CFUNC(GetUnitMetalExtraction);
@@ -4218,6 +4219,24 @@ int LuaSyncedRead::GetUnitResources(lua_State* L)
 	return 4;
 }
 
+/***
+ * @function Spring.GetUnitStorage
+ * @number unitID
+ * @treturn number Unit's metal storage
+ * @treturn number Unit's energy storage
+ */
+int LuaSyncedRead::GetUnitStorage(lua_State* L)
+{
+	const CUnit* unit = ParseAllyUnit(L, __func__, 1);
+
+	if (unit == nullptr)
+		return 0;
+
+	lua_pushnumber(L, unit->storage.metal);
+	lua_pushnumber(L, unit->storage.energy);
+	return 2;
+}
+
 /***
  * @function Spring.GetUnitCosts
  * @number unitID

diff --git a/rts/Lua/LuaSyncedRead.h b/rts/Lua/LuaSyncedRead.h
@@ -113,6 +113,7 @@ class LuaSyncedRead {
 		static int GetUnitCosts(lua_State* L);
 		static int GetUnitCostTable(lua_State* L);
 		static int GetUnitResources(lua_State* L);
+		static int GetUnitStorage(lua_State* L);
 		static int GetUnitMetalExtraction(lua_State* L);
 		static int GetUnitExperience(lua_State* L);
 		static int GetUnitStates(lua_State* L);