From 77c4366b069efdca893477f4a5840bb3a7458650 Mon Sep 17 00:00:00 2001
From: Richard Nai <richardnai6@gmail.com>
Date: Thu, 9 Jan 2025 05:31:29 -0800
Subject: [PATCH] Lua frontend to allow changing unit's storage (#1879)

---
 rts/Lua/LuaSyncedCtrl.cpp | 73 +++++++++++++++++++++++++++++++++++++++
 rts/Lua/LuaSyncedCtrl.h   |  1 +
 rts/Lua/LuaSyncedRead.cpp | 19 ++++++++++
 rts/Lua/LuaSyncedRead.h   |  1 +
 4 files changed, 94 insertions(+)

diff --git a/rts/Lua/LuaSyncedCtrl.cpp b/rts/Lua/LuaSyncedCtrl.cpp
index 4f719447a6..9d6e46dabe 100644
--- 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
+ */
+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
index cb50d9a981..94ef50fd37 100644
--- 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
index e642404011..39f4ff501d 100644
--- 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
index 1716d8a256..efb469daaf 100644
--- 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);