Support resize/erase/insert for containers, and allow any index in BitArray.

library/LuaTypes.cpp

@@ -147,12 +147,12 @@ void df::pointer_identity::lua_write(lua_State *state, int fname_idx, void *ptr,
     lua_write(state, fname_idx, ptr, target, val_index);
 }
 
-int container_identity::lua_item_count(lua_State *state, void *ptr)
+int container_identity::lua_item_count(lua_State *state, void *ptr, CountMode mode)
 {
     if (lua_isnumber(state, UPVAL_ITEM_COUNT))
         return lua_tointeger(state, UPVAL_ITEM_COUNT);
     else
-        return item_count(ptr);
+        return item_count(ptr, mode);
 }
 
 void container_identity::lua_item_reference(lua_State *state, int fname_idx, void *ptr, int idx)
@@ -176,6 +176,31 @@ void container_identity::lua_item_write(lua_State *state, int fname_idx, void *p
     id->lua_write(state, fname_idx, pitem, val_index);
 }
 
+bool container_identity::lua_insert(lua_State *state, int fname_idx, void *ptr, int idx, int val_index)
+{
+    auto id = (type_identity*)lua_touserdata(state, UPVAL_ITEM_ID);
+
+    char tmp[32];
+    void *pitem = &tmp;
+
+    if (id->isPrimitive())
+    {
+        if (id->isConstructed())
+            luaL_error(state, "Temporaries of type %s not supported", id->getFullName().c_str());
+
+        assert(id->byte_size() <= sizeof(tmp));
+        id->lua_write(state, fname_idx, pitem, val_index);
+    }
+    else
+    {
+        pitem = get_object_internal(state, id, val_index, false);
+        if (!pitem)
+            field_error(state, fname_idx, "incompatible object type", "insert");
+    }
+
+    return insert(ptr, idx, pitem);
+}
+
 void ptr_container_identity::lua_item_reference(lua_State *state, int fname_idx, void *ptr, int idx)
 {
     auto id = (type_identity*)lua_touserdata(state, UPVAL_ITEM_ID);
@@ -197,6 +222,16 @@ void ptr_container_identity::lua_item_write(lua_State *state, int fname_idx, voi
     df::pointer_identity::lua_write(state, fname_idx, pitem, id, val_index);
 }
 
+bool ptr_container_identity::lua_insert(lua_State *state, int fname_idx, void *ptr, int idx, int val_index)
+{
+    auto id = (type_identity*)lua_touserdata(state, UPVAL_ITEM_ID);
+
+    void *pitem = NULL;
+    df::pointer_identity::lua_write(state, fname_idx, &pitem, id, val_index);
+
+    return insert(ptr, idx, pitem);
+}
+
 void bit_container_identity::lua_item_reference(lua_State *state, int, void *, int)
 {
     lua_pushnil(state);
@@ -238,6 +273,15 @@ static void *find_field(lua_State *state, int index, const char *mode)
     return p;
 }
 
+static uint8_t *check_method_call(lua_State *state, int min_args, int max_args)
+{
+    int argc = lua_gettop(state)-1;
+    if (argc < min_args || argc > max_args)
+        field_error(state, UPVAL_METHOD_NAME, "wrong argument count", "call");
+
+    return get_object_addr(state, 1, UPVAL_METHOD_NAME, "call");
+}
+
 static void GetAdHocMetatable(lua_State *state, const struct_field_info *field);
 
 static void read_field(lua_State *state, const struct_field_info *field, void *ptr)
@@ -456,7 +500,7 @@ static int meta_container_len(lua_State *state)
 {
     uint8_t *ptr = get_object_addr(state, 1, 0, "get length");
     auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
-    int len = id->lua_item_count(state, ptr);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_LEN);
     lua_pushinteger(state, len);
     return 1;
 }
@@ -497,7 +541,7 @@ static int check_container_index(lua_State *state, int len,
         field_error(state, fidx, "invalid index", mode);
 
     int idx = lua_tointeger(state, iidx);
-    if (idx < 0 || idx >= len)
+    if (idx < 0 || (idx >= len && len >= 0))
         field_error(state, fidx, "index out of bounds", mode);
 
     return idx;
@@ -514,7 +558,7 @@ static int meta_container_index(lua_State *state)
         return 1;
 
     auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
-    int len = id->lua_item_count(state, ptr);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_READ);
     int idx = check_container_index(state, len, 2, iidx, "read");
     id->lua_item_read(state, 2, ptr, idx);
     return 1;
@@ -531,7 +575,7 @@ static int meta_container_field_reference(lua_State *state)
     int iidx = lookup_container_field(state, 2, "reference");
 
     auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
-    int len = id->lua_item_count(state, ptr);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_WRITE);
     int idx = check_container_index(state, len, 2, iidx, "reference");
     id->lua_item_reference(state, 2, ptr, idx);
     return 1;
@@ -546,12 +590,60 @@ static int meta_container_newindex(lua_State *state)
     int iidx = lookup_container_field(state, 2, "write");
 
     auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
-    int len = id->lua_item_count(state, ptr);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_WRITE);
     int idx = check_container_index(state, len, 2, iidx, "write");
     id->lua_item_write(state, 2, ptr, idx, 3);
     return 0;
 }
 
+/**
+ * Method: resize container
+ */
+static int method_container_resize(lua_State *state)
+{
+    uint8_t *ptr = check_method_call(state, 1, 1);
+
+    auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
+    int idx = check_container_index(state, -1, UPVAL_METHOD_NAME, 2, "call");
+
+    if (!id->resize(ptr, idx))
+        field_error(state, UPVAL_METHOD_NAME, "not supported", "call");
+    return 0;
+}
+
+/**
+ * Method: erase item from container
+ */
+static int method_container_erase(lua_State *state)
+{
+    uint8_t *ptr = check_method_call(state, 1, 1);
+
+    auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_LEN);
+    int idx = check_container_index(state, len, UPVAL_METHOD_NAME, 2, "call");
+
+    if (!id->erase(ptr, idx))
+        field_error(state, UPVAL_METHOD_NAME, "not supported", "call");
+    return 0;
+}
+
+/**
+ * Method: insert item into container
+ */
+static int method_container_insert(lua_State *state)
+{
+    uint8_t *ptr = check_method_call(state, 2, 2);
+
+    auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_LEN);
+    if (len >= 0) len++;
+    int idx = check_container_index(state, len, UPVAL_METHOD_NAME, 2, "call");
+
+    if (!id->lua_insert(state, UPVAL_METHOD_NAME, ptr, idx, 3))
+        field_error(state, UPVAL_METHOD_NAME, "not supported", "call");
+    return 0;
+}
+
 /**
  * Metamethod: __len for bitfields.
  */
@@ -718,6 +810,17 @@ static void MakePrimitiveMetatable(lua_State *state, type_identity *type)
     SetStructMethod(state, base+1, base+2, meta_primitive_newindex, "__newindex");
 }
 
+static void AddContainerMethodFun(lua_State *state, int meta_idx, int field_idx,
+                                  lua_CFunction function, const char *name,
+                                  type_identity *container, type_identity *item, int count)
+{
+    lua_pushfstring(state, "%s()", name);
+    SetContainerMethod(state, meta_idx, lua_gettop(state), function, name, container, item, count);
+    lua_pop(state, 1);
+
+    EnableMetaField(state, field_idx, name);
+}
+
 /**
  * Make a container-style object metatable.
  */
@@ -750,6 +853,10 @@ static void MakeContainerMetatable(lua_State *state, container_identity *type,
 
     SetContainerMethod(state, base+1, base+2, meta_container_field_reference, "_field", type, item, count);
 
+    AddContainerMethodFun(state, base+1, base+2, method_container_resize, "resize", type, item, count);
+    AddContainerMethodFun(state, base+1, base+2, method_container_erase, "erase", type, item, count);
+    AddContainerMethodFun(state, base+1, base+2, method_container_insert, "insert", type, item, count);
+
     AttachEnumKeys(state, base+1, base+2, ienum);
 }
 

library/LuaWrapper.cpp

@@ -593,7 +593,10 @@ static int meta_new(lua_State *state)
 
     void *ptr = id->allocate();
     if (!ptr)
-        return 0;
+    {
+        lua_pushnil(state);
+        return 1;
+    }
 
     if (lua_isuserdata(state, 1))
     {

library/include/BitArray.h

@@ -184,10 +184,15 @@ namespace DFHack
             memcpy(m_data, other.m_data,m_size*sizeof(T));
         }
 
+        typedef T value_type;
+
         T *data() { return m_data; }
         const T *data() const { return m_data; }
         unsigned size() const { return m_size; }
 
+        T *begin() { return m_data; }
+        T *end() { return m_data+m_size; }
+
         T& operator[] (unsigned i) { return m_data[i]; }
         const T& operator[] (unsigned i) const { return m_data[i]; }
 
@@ -217,5 +222,15 @@ namespace DFHack
             memcpy(data(), other.data(), sizeof(T)*size());
             return *this;
         }
+
+        void erase(T *ptr) {
+            memmove(ptr, ptr+1, sizeof(T)*(m_size - (ptr - m_data))); m_size--;
+        }
+        void insert(T *ptr, const T &item) {
+            int idx = ptr - m_data;
+            resize(m_size+1);
+            memmove(m_data + idx + 1, m_data + idx, sizeof(T)*(m_size - idx - 1));
+            m_data[idx] = item;
+        }
     };
 }

library/include/DataDefs.h

@@ -91,6 +91,11 @@ namespace DFHack
         virtual void lua_write(lua_State *state, int fname_idx, void *ptr, int val_index) = 0;
         virtual void build_metatable(lua_State *state);
 
+        // lua_read doesn't just return a reference to the object
+        virtual bool isPrimitive() { return true; }
+        // needs constructor/destructor
+        virtual bool isConstructed() { return false; }
+        // inherits from container_identity
         virtual bool isContainer() { return false; }
 
         void *allocate();
@@ -104,6 +109,9 @@ namespace DFHack
         constructed_identity(size_t size, TAllocateFn alloc)
             : type_identity(size), allocator(alloc) {};
 
+        virtual bool isPrimitive() { return false; }
+        virtual bool isConstructed() { return true; }
+
         virtual bool can_allocate() { return (allocator != NULL); }
         virtual void *do_allocate() { return allocator(NULL,NULL); }
         virtual void do_copy(void *tgt, const void *src) { allocator(tgt,src); }
@@ -161,6 +169,8 @@ namespace DFHack
 
         virtual identity_type type() { return IDTYPE_BITFIELD; }
 
+        virtual bool isConstructed() { return false; }
+
         int getNumBits() { return num_bits; }
         const bitfield_item_info *getBits() { return bits; }
 
@@ -195,6 +205,9 @@ namespace DFHack
 
         type_identity *getBaseType() { return base_type; }
 
+        virtual bool isPrimitive() { return true; }
+        virtual bool isConstructed() { return false; }
+
         virtual void lua_read(lua_State *state, int fname_idx, void *ptr);
         virtual void lua_write(lua_State *state, int fname_idx, void *ptr, int val_index);
     };

library/include/DataIdentity.h

@@ -84,14 +84,24 @@ namespace DFHack
 
         virtual std::string getFullName(type_identity *item);
 
-        int lua_item_count(lua_State *state, void *ptr);
+        enum CountMode {
+            COUNT_LEN, COUNT_READ, COUNT_WRITE
+        };
+
+        int lua_item_count(lua_State *state, void *ptr, CountMode cnt);
 
         virtual void lua_item_reference(lua_State *state, int fname_idx, void *ptr, int idx);
         virtual void lua_item_read(lua_State *state, int fname_idx, void *ptr, int idx);
         virtual void lua_item_write(lua_State *state, int fname_idx, void *ptr, int idx, int val_index);
 
+        virtual bool resize(void *ptr, int size) { return false; }
+        virtual bool erase(void *ptr, int index) { return false; }
+        virtual bool insert(void *ptr, int index, void *pitem) { return false; }
+
+        virtual bool lua_insert(lua_State *state, int fname_idx, void *ptr, int idx, int val_index);
+
     protected:
-        virtual int item_count(void *ptr) = 0;
+        virtual int item_count(void *ptr, CountMode cnt) = 0;
         virtual void *item_pointer(type_identity *item, void *ptr, int idx) = 0;
     };
 
@@ -108,6 +118,8 @@ namespace DFHack
         virtual void lua_item_reference(lua_State *state, int fname_idx, void *ptr, int idx);
         virtual void lua_item_read(lua_State *state, int fname_idx, void *ptr, int idx);
         virtual void lua_item_write(lua_State *state, int fname_idx, void *ptr, int idx, int val_index);
+
+        virtual bool lua_insert(lua_State *state, int fname_idx, void *ptr, int idx, int val_index);
     };
 
     class DFHACK_EXPORT bit_container_identity : public container_identity {
@@ -175,25 +187,33 @@ namespace df
         virtual void lua_write(lua_State *state, int fname_idx, void *ptr, int val_index);
     };
 
-    class stl_string_identity : public primitive_identity {
+    class stl_string_identity : public DFHack::constructed_identity {
     public:
-        stl_string_identity() : primitive_identity(sizeof(std::string)) {};
+        stl_string_identity()
+            : constructed_identity(sizeof(std::string), &allocator_fn<std::string>)
+        {};
 
         std::string getFullName() { return "string"; }
 
+        virtual DFHack::identity_type type() { return DFHack::IDTYPE_PRIMITIVE; }
+
+        virtual bool isPrimitive() { return true; }
+
         virtual void lua_read(lua_State *state, int fname_idx, void *ptr);
         virtual void lua_write(lua_State *state, int fname_idx, void *ptr, int val_index);
     };
 
     class stl_ptr_vector_identity : public ptr_container_identity {
     public:
+        typedef std::vector<void*> container;
+
         /*
          * This class assumes that std::vector<T*> is equivalent
          * in layout and behavior to std::vector<void*> for any T.
          */
 
         stl_ptr_vector_identity(type_identity *item = NULL, enum_identity *ienum = NULL)
-            : ptr_container_identity(sizeof(std::vector<void*>),allocator_fn<std::vector<void*> >,item, ienum)
+            : ptr_container_identity(sizeof(container),allocator_fn<container>,item, ienum)
         {};
 
         std::string getFullName(type_identity *item) {
@@ -202,12 +222,27 @@ namespace df
 
         virtual DFHack::identity_type type() { return DFHack::IDTYPE_STL_PTR_VECTOR; }
 
+        virtual bool resize(void *ptr, int size) {
+            (*(container*)ptr).resize(size);
+            return true;
+        }
+        virtual bool erase(void *ptr, int size) {
+            auto &ct = *(container*)ptr;
+            ct.erase(ct.begin()+size);
+            return true;
+        }
+        virtual bool insert(void *ptr, int idx, void *item) {
+            auto &ct = *(container*)ptr;
+            ct.insert(ct.begin()+idx, item);
+            return true;
+        }
+
     protected:
-        virtual int item_count(void *ptr) {
+        virtual int item_count(void *ptr, CountMode) {
-            return ((std::vector<void*>*)ptr)->size();
+            return ((container*)ptr)->size();
         };
         virtual void *item_pointer(type_identity *, void *ptr, int idx) {
-            return &(*(std::vector<void*>*)ptr)[idx];
+            return &(*(container*)ptr)[idx];
         }
     };
 
@@ -231,7 +266,7 @@ namespace df
         static buffer_container_identity base_instance;
 
     protected:
-        virtual int item_count(void *ptr) { return size; }
+        virtual int item_count(void *ptr, CountMode) { return size; }
         virtual void *item_pointer(type_identity *item, void *ptr, int idx) {
             return ((uint8_t*)ptr) + idx * item->byte_size();
         }
@@ -250,8 +285,23 @@ namespace df
             return name + container_identity::getFullName(item);
         }
 
+        virtual bool resize(void *ptr, int size) {
+            (*(T*)ptr).resize(size);
+            return true;
+        }
+        virtual bool erase(void *ptr, int size) {
+            auto &ct = *(T*)ptr;
+            ct.erase(ct.begin()+size);
+            return true;
+        }
+        virtual bool insert(void *ptr, int idx, void *item) {
+            auto &ct = *(T*)ptr;
+            ct.insert(ct.begin()+idx, *(typename T::value_type*)item);
+            return true;
+        }
+
     protected:
-        virtual int item_count(void *ptr) { return ((T*)ptr)->size(); }
+        virtual int item_count(void *ptr, CountMode) { return ((T*)ptr)->size(); }
         virtual void *item_pointer(type_identity *item, void *ptr, int idx) {
             return &(*(T*)ptr)[idx];
         }
@@ -274,8 +324,15 @@ namespace df
             return "BitArray<>";
         }
 
+        virtual bool resize(void *ptr, int size) {
+            ((container*)ptr)->resize(size);
+            return true;
+        }
+
     protected:
-        virtual int item_count(void *ptr) { return ((container*)ptr)->size * 8; }
+        virtual int item_count(void *ptr, CountMode cnt) {
+            return cnt == COUNT_LEN ? ((container*)ptr)->size * 8 : -1;
+        }
         virtual bool get_item(void *ptr, int idx) {
             return ((container*)ptr)->is_set(idx);
         }
@@ -296,8 +353,15 @@ namespace df
             return "vector" + bit_container_identity::getFullName(item);
         }
 
+        virtual bool resize(void *ptr, int size) {
+            (*(container*)ptr).resize(size);
+            return true;
+        }
+
     protected:
-        virtual int item_count(void *ptr) { return ((container*)ptr)->size(); }
+        virtual int item_count(void *ptr, CountMode) {
+            return ((container*)ptr)->size();
+        }
         virtual bool get_item(void *ptr, int idx) {
             return (*(container*)ptr)[idx];
         }

library/include/LuaWrapper.h

@@ -85,6 +85,7 @@ namespace DFHack { namespace LuaWrapper {
  * Fields that are actually in UPVAL_METATABLE are marked with NULL light udata.
  */
 #define UPVAL_FIELDTABLE lua_upvalueindex(3)
+#define UPVAL_METHOD_NAME lua_upvalueindex(3)
 
 /*
  * Only for containers: light udata with container identity.