From 9d52ed49ba3267059ca680c112b8ad8c43b9239d Mon Sep 17 00:00:00 2001
From: ThePhD <phdofthehouse@gmail.com>
Date: Mon, 19 Sep 2016 23:37:00 -0400
Subject: [PATCH] balance that stack, goddamnit, BALANCE IT

---
 docs/source/api/api-top.rst | 54 ++++++++++++++++++-------------------
 sol/stack_field.hpp         |  9 ++++++-
 test_tables.cpp             | 36 +++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 28 deletions(-)

diff --git a/docs/source/api/api-top.rst b/docs/source/api/api-top.rst
index 631a2633..f9b0378b 100644
--- a/docs/source/api/api-top.rst
+++ b/docs/source/api/api-top.rst
@@ -9,38 +9,38 @@ Browse the various function and classes :doc:`Sol<../index>` utilizes to make yo
    :name: apitoc
    :maxdepth: 2
 
-   compatibility
-   coroutine
-   c_call
-   error
-   function
-   protected_function
-   object
-   reference
-   stack_reference
-   make_reference
-   overload
-   protect
-   readonly
-   var
-   resolve
-   as_function
-   property
-   proxy
-   stack
-   optional
    state
    table
+   proxy
    as_table
-   metatable_key
-   this_state
-   thread
-   tie
-   types
-   user
    usertype
    simple_usertype
-   userdata
    usertype_memory
    unique_usertype_traits
+   tie
+   function
+   protected_function
+   coroutine
+   error
+   object
+   userdata
+   reference
+   thread
+   stack_reference
+   make_reference
+   optional
+   this_state
    variadic_args
+   overload
+   property
+   var
+   protect
+   readonly
+   as_function
+   c_call
+   resolve
+   stack
+   user
+   compatibility
+   types
+   metatable_key
diff --git a/sol/stack_field.hpp b/sol/stack_field.hpp
index eb7bab64..ca52f3db 100644
--- a/sol/stack_field.hpp
+++ b/sol/stack_field.hpp
@@ -234,9 +234,15 @@ namespace sol {
 				apply<false>(std::index_sequence<I1, I...>(), L, std::forward<Keys>(keys), std::forward<Value>(value), -1);
 			}
 
+			template <bool g, std::size_t I0, std::size_t... I, typename Keys, typename Value>
+			void top_apply(std::index_sequence<I0, I...>, lua_State* L, Keys&& keys, Value&& value, int tableindex) {
+				apply<g>(std::index_sequence<I0, I...>(), L, std::forward<Keys>(keys), std::forward<Value>(value), tableindex);
+				lua_pop(L, static_cast<int>(sizeof...(I)));
+			}
+
 			template <typename Keys, typename Value>
 			void set(lua_State* L, Keys&& keys, Value&& value, int tableindex = -3) {
-				apply<b>(std::make_index_sequence<sizeof...(Args)>(), L, std::forward<Keys>(keys), std::forward<Value>(value), tableindex);
+				top_apply<b>(std::make_index_sequence<sizeof...(Args)>(), L, std::forward<Keys>(keys), std::forward<Value>(value), tableindex);
 			}
 		};
 
@@ -246,6 +252,7 @@ namespace sol {
 			void set(lua_State* L, Keys&& keys, Value&& value, int tableindex = -1) {
 				get_field<b, raw>(L, detail::forward_get<0>(keys), tableindex);
 				set_field<false, raw>(L, detail::forward_get<1>(keys), std::forward<Value>(value));
+				lua_pop(L, 1);
 			}
 		};
 	} // stack
diff --git a/test_tables.cpp b/test_tables.cpp
index bb5a8cbf..cb8f291e 100644
--- a/test_tables.cpp
+++ b/test_tables.cpp
@@ -77,6 +77,42 @@ TEST_CASE("tables/as-enum-classes", "Making sure enums can be put in and gotten
 	REQUIRE(dir == direction::up);
 }
 
+TEST_CASE("tables/cleanup", "make sure tables leave the stack balanced") {
+	sol::state lua;
+	lua.open_libraries();
+
+	auto f = [] { return 5; };
+	for (int i = 0; i < 30; i++) {
+		std::string name = std::string("init") + std::to_string(i);
+		int top = lua_gettop(lua);
+		lua[name] = f;
+		int aftertop = lua_gettop(lua);
+		REQUIRE(aftertop == top);
+		int val = lua[name]();
+		REQUIRE(val == 5);
+	}
+}
+
+TEST_CASE("tables/nested-cleanup", "make sure tables leave the stack balanced") {
+	sol::state lua;
+	lua.open_libraries();
+
+	lua.script("A={}");
+	auto f = [] { return 5; };
+	for (int i = 0; i < 30; i++) {
+		std::string name = std::string("init") + std::to_string(i);
+		int top = lua_gettop(lua);
+		auto A = lua["A"];
+		int beforetop = lua_gettop(lua);
+		REQUIRE(beforetop == top);
+		A[name] = f;
+		int aftertop = lua_gettop(lua);
+		REQUIRE(aftertop == top);
+		int val = A[name]();
+		REQUIRE(val == 5);
+	}
+}
+
 TEST_CASE("tables/new_enum", "Making sure enums can be put in and gotten out as values") {
 	enum class direction {
 		up,