From 929657bed4e1300537a88401de9e3b14f2e1dc72 Mon Sep 17 00:00:00 2001
From: Alexander Gavrilov <angavrilov@gmail.com>
Date: Wed, 28 Mar 2012 15:28:42 +0800
Subject: [PATCH] Disable pointer auto-vivification unless new is specified.

Since it is essentially allocating non-gc managed objects,
it can lead to memory leaks and shouldn't happen invisibly.
Also support using the 'assign' key to request assign()
from another object before processing the current map.
---
 library/LuaTypes.cpp   | 13 +++++++++-
 library/LuaWrapper.cpp | 58 +++++++++++++++++++++++++++++++++---------
 2 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/library/LuaTypes.cpp b/library/LuaTypes.cpp
index 1e33ed713..79f515d40 100644
--- a/library/LuaTypes.cpp
+++ b/library/LuaTypes.cpp
@@ -168,10 +168,16 @@ static void autovivify_ptr(lua_State *state, int fname_idx, void **pptr,
 {
     lua_getfield(state, val_index, "new");
 
-    if (!lua_isnil(state, -1))
+    // false or nil => bail out
+    if (!lua_toboolean(state, -1))
+        field_error(state, fname_idx, "null and autovivify not requested", "write");
+
+    // not 'true' => call df.new()
+    if (!lua_isboolean(state, -1))
     {
         int top = lua_gettop(state);
 
+        // Verify new points to a reasonable type of object
         type_identity *suggested = get_object_identity(state, top, "autovivify", true, true);
 
         if (!is_type_compatible(state, target, 0, suggested, top+1, false))
@@ -179,17 +185,22 @@ static void autovivify_ptr(lua_State *state, int fname_idx, void **pptr,
 
         lua_pop(state, 1);
 
+        // Invoke df.new()
         lua_getfield(state, LUA_REGISTRYINDEX, DFHACK_NEW_NAME);
         lua_swap(state);
         lua_call(state, 1, 1);
 
+        // Retrieve the pointer
         void *nval = get_object_internal(state, target, top, false);
 
+        // shouldn't happen: this means suggested type is compatible,
+        // but its new() result isn't for some reason.
         if (!nval)
             field_error(state, fname_idx, "inconsistent autovivify type", "write");
 
         *pptr = nval;
     }
+    // otherwise use the target type
     else
     {
         if (!target)
diff --git a/library/LuaWrapper.cpp b/library/LuaWrapper.cpp
index 11db14875..36b24a403 100644
--- a/library/LuaWrapper.cpp
+++ b/library/LuaWrapper.cpp
@@ -616,20 +616,32 @@ static void invoke_resize(lua_State *state, int table, lua_Integer size)
     lua_call(state, 2, 0);
 }
 
-static void copy_table(lua_State *state, int dest, int src, int skipkey)
+static void copy_table(lua_State *state, int dest, int src, int skipbase)
 {
+    // stack: (skipbase) skipkey skipkey |
+
+    int top = lua_gettop(state);
+
     lua_pushnil(state);
 
     while (lua_next(state, src))
     {
-        if (lua_equal(state, -2, skipkey))
-            lua_pop(state, 1);
-        else
+        for (int i = skipbase+1; i <= top; i++)
+        {
+            if (lua_rawequal(state, -2, i))
+            {
+                lua_pop(state, 1);
+                goto next_outer;
+            }
+        }
+
         {
             lua_pushvalue(state, -2);
             lua_swap(state);
             lua_settable(state, dest);
         }
+
+    next_outer:;
     }
 }
 
@@ -655,18 +667,40 @@ static int meta_assign(lua_State *state)
     {
         type_identity *id = get_object_identity(state, 1, "df.assign()", false);
 
+        int base = lua_gettop(state);
+
+        // x:assign{ assign = foo } => x:assign(foo)
+        bool has_assign = false;
+
+        lua_pushstring(state, "assign");
+        lua_dup(state);
+        lua_rawget(state, 2);
+
+        if (!lua_isnil(state,-1))
+        {
+            has_assign = true;
+            lua_getfield(state, LUA_REGISTRYINDEX, DFHACK_ASSIGN_NAME);
+            lua_pushvalue(state, 1);
+            lua_pushvalue(state, base+2);
+            lua_call(state, 2, 0);
+        }
+
+        lua_pop(state, 1);
+
+        // new is used by autovivification and should be skipped
+        lua_pushstring(state, "new");
+
         if (id->isContainer())
         {
+            // check resize field
             lua_pushstring(state, "resize");
-            int resize_str = lua_gettop(state);
-
             lua_dup(state);
             lua_rawget(state, 2);
 
-            if (lua_isnil(state,-1))
+            if (lua_isnil(state,-1) && !has_assign)
             {
                 /*
-                 * nil or missing resize field => 1-based lua array
+                 * no assign && nil or missing resize field => 1-based lua array
                  */
                 int size = lua_objlen(state, 2);
 
@@ -682,7 +716,7 @@ static int meta_assign(lua_State *state)
             }
             else
             {
-                if (lua_isboolean(state, -1))
+                if (lua_isboolean(state, -1) || lua_isnil(state, -1))
                 {
                     // resize=false => just assign
                     // resize=true => find the largest index
@@ -711,13 +745,13 @@ static int meta_assign(lua_State *state)
                 }
 
                 lua_pop(state, 1);
-                copy_table(state, 1, 2, resize_str);
+                copy_table(state, 1, 2, base);
             }
         }
         else
         {
-            lua_pushstring(state, "new");
-            copy_table(state, 1, 2, lua_gettop(state));
+
+            copy_table(state, 1, 2, base);
         }
     }