Implement __pairs and __ipairs for DF objects.

Structs enumerate fields in memory order in pairs().
Containers & biftields enumerate int indexes in ipairs, and
string keys in pairs (i.e. using index-enum for arrays).

library/LuaTypes.cpp

@@ -389,6 +389,43 @@ static void *find_field(lua_State *state, int index, const char *mode)
     return p;
 }
 
+static int cur_iter_index(lua_State *state, int len, int fidx, int first_idx = -1)
+{
+    int rv;
+
+    if (lua_isnil(state, fidx))
+        rv = first_idx;
+    else
+    {
+        if (lua_isnumber(state, fidx))
+            rv = lua_tointeger(state, fidx);
+        else
+        {
+            lua_pushvalue(state, fidx);
+            lua_rawget(state, UPVAL_FIELDTABLE);
+            if (!lua_isnumber(state, -1))
+                field_error(state, fidx, "index not found", "iterate");
+            rv = lua_tointeger(state, -1);
+            lua_pop(state, 1);
+        }
+
+        if (rv < 0 || rv >= len)
+            field_error(state, fidx, "index out of bounds", "iterate");
+    }
+
+    return rv;
+}
+
+static void iter_idx_to_name(lua_State *state, int idx)
+{
+    lua_pushvalue(state, idx);
+    lua_rawget(state, UPVAL_FIELDTABLE);
+    if (lua_isnil(state, -1))
+        lua_pop(state, 1);
+    else
+        lua_replace(state, idx);
+}
+
 static uint8_t *check_method_call(lua_State *state, int min_args, int max_args)
 {
     int argc = lua_gettop(state)-1;
@@ -592,6 +629,24 @@ static int meta_struct_newindex(lua_State *state)
     return 0;
 }
 
+/**
+ * Metamethod: iterator for structures.
+ */
+static int meta_struct_next(lua_State *state)
+{
+    if (lua_gettop(state) < 2) lua_pushnil(state);
+
+    int len = lua_objlen(state, UPVAL_FIELDTABLE);
+    int idx = cur_iter_index(state, len+1, 2, 0);
+    if (idx == len)
+        return 0;
+
+    lua_rawgeti(state, UPVAL_FIELDTABLE, idx+1);
+    lua_dup(state);
+    lua_gettable(state, 1);
+    return 2;
+}
+
 /**
  * Metamethod: __index for primitives, i.e. simple object references.
  *   Fields point to identity, or NULL for metafields.
@@ -722,6 +777,39 @@ static int meta_container_newindex(lua_State *state)
     return 0;
 }
 
+/**
+ * Metamethod: integer iterator for containers.
+ */
+static int meta_container_nexti(lua_State *state)
+{
+    if (lua_gettop(state) < 2) lua_pushnil(state);
+
+    uint8_t *ptr = get_object_addr(state, 1, 2, "iterate");
+
+    auto id = (container_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
+    int len = id->lua_item_count(state, ptr, container_identity::COUNT_LEN);
+    int idx = cur_iter_index(state, len, 2);
+
+    if (++idx >= len)
+        return 0;
+
+    lua_pushinteger(state, idx);
+    id->lua_item_read(state, 2, ptr, idx);
+    return 2;
+}
+
+/**
+ * Metamethod: name iterator for containers.
+ */
+static int meta_container_next(lua_State *state)
+{
+    if (!meta_container_nexti(state))
+        return 0;
+
+    iter_idx_to_name(state, lua_gettop(state)-1);
+    return 2;
+}
+
 /**
  * Method: resize container
  */
@@ -781,6 +869,17 @@ static int meta_bitfield_len(lua_State *state)
     return 1;
 }
 
+static void read_bitfield(lua_State *state, uint8_t *ptr, bitfield_identity *id, int idx)
+{
+    int size = id->getBits()[idx].size;
+
+    int value = getBitfieldField(ptr, idx, size);
+    if (size <= 1)
+        lua_pushboolean(state, value != 0);
+    else
+        lua_pushinteger(state, value);
+}
+
 /**
  * Metamethod: __index for bitfields.
  */
@@ -803,13 +902,7 @@ static int meta_bitfield_index(lua_State *state)
     }
 
     int idx = check_container_index(state, id->getNumBits(), 2, iidx, "read");
-    int size = id->getBits()[idx].size;
-
-    int value = getBitfieldField(ptr, idx, size);
-    if (size <= 1)
-        lua_pushboolean(state, value != 0);
-    else
-        lua_pushinteger(state, value);
+    read_bitfield(state, ptr, id, idx);
     return 1;
 }
 
@@ -846,6 +939,44 @@ static int meta_bitfield_newindex(lua_State *state)
     return 0;
 }
 
+/**
+ * Metamethod: integer iterator for bitfields.
+ */
+static int meta_bitfield_nexti(lua_State *state)
+{
+    if (lua_gettop(state) < 2) lua_pushnil(state);
+
+    uint8_t *ptr = get_object_addr(state, 1, 2, "iterate");
+
+    auto id = (bitfield_identity*)lua_touserdata(state, UPVAL_CONTAINER_ID);
+    int len = id->getNumBits();
+    int idx = cur_iter_index(state, len, 2);
+
+    if (idx < 0)
+        idx = 0;
+    else
+        idx += std::max(1, (int)id->getBits()[idx].size);
+
+    if (idx >= len)
+        return 0;
+
+    lua_pushinteger(state, idx);
+    read_bitfield(state, ptr, id, idx);
+    return 2;
+}
+
+/**
+ * Metamethod: name iterator for bitfields.
+ */
+static int meta_bitfield_next(lua_State *state)
+{
+    if (!meta_bitfield_nexti(state))
+        return 0;
+
+    iter_idx_to_name(state, lua_gettop(state)-1);
+    return 2;
+}
+
 /**
  * Metamethod: __index for df.global
  */
@@ -906,36 +1037,43 @@ static void AddMethodWrapper(lua_State *state, int meta_idx, int field_idx,
 /**
  * Add fields in the array to the UPVAL_FIELDTABLE candidates on the stack.
  */
-static void IndexFields(lua_State *state, struct_identity *pstruct)
+static void IndexFields(lua_State *state, int base, struct_identity *pstruct)
 {
-    // stack: metatable fieldtable
+    if (pstruct->getParent())
+        IndexFields(state, base, pstruct->getParent());
 
-    int base = lua_gettop(state) - 2;
+    auto fields = pstruct->getFields();
+    if (!fields)
+        return;
 
-    for (struct_identity *p = pstruct; p; p = p->getParent())
+    int cnt = lua_objlen(state, base+3); // field iter table
+
+    for (int i = 0; fields[i].mode != struct_field_info::END; ++i)
     {
-        auto fields = p->getFields();
-        if (!fields)
-            continue;
+        // Qualify conflicting field names with the type
+        std::string name = fields[i].name;
 
-        for (int i = 0; fields[i].mode != struct_field_info::END; ++i)
+        lua_getfield(state, base+2, name.c_str());
+        if (!lua_isnil(state, -1))
+            name = pstruct->getName() + ("." + name);
+        lua_pop(state, 1);
+
+        // Handle the field
+        switch (fields[i].mode)
         {
-            switch (fields[i].mode)
-            {
-            case struct_field_info::OBJ_METHOD:
-                AddMethodWrapper(state, base+1, base+2, fields[i].name,
-                                 (function_identity_base*)fields[i].type);
-                break;
+        case struct_field_info::OBJ_METHOD:
+            AddMethodWrapper(state, base+1, base+2, name.c_str(),
+                             (function_identity_base*)fields[i].type);
+            break;
 
-            case struct_field_info::CLASS_METHOD:
-                break;
+        case struct_field_info::CLASS_METHOD:
+            break;
 
-            default:
-                lua_pushstring(state,fields[i].name);
-                lua_pushlightuserdata(state,(void*)&fields[i]);
-                lua_rawset(state,base+2);
-                break;
-            }
+        default:
+            AssociateId(state, base+3, ++cnt, name.c_str());
+            lua_pushlightuserdata(state, (void*)&fields[i]);
+            lua_setfield(state, base+2, name.c_str());
+            break;
         }
     }
 }
@@ -976,8 +1114,20 @@ static void MakeFieldMetatable(lua_State *state, struct_identity *pstruct,
 
     MakeMetatable(state, pstruct, "struct"); // meta, fields
 
-    IndexFields(state, pstruct);
+    // Index the fields
+    lua_newtable(state);
+
+    IndexFields(state, base, pstruct);
+
+    // Add the iteration metamethods
+    PushStructMethod(state, base+1, base+3, meta_struct_next);
+    SetPairsMethod(state, base+1, "__pairs");
+    lua_pushnil(state);
+    SetPairsMethod(state, base+1, "__ipairs");
+
+    lua_setfield(state, base+1, "_index_table");
 
+    // Add the indexing metamethods
     SetStructMethod(state, base+1, base+2, reader, "__index");
     SetStructMethod(state, base+1, base+2, writer, "__newindex");
 
@@ -995,8 +1145,21 @@ static void MakePrimitiveMetatable(lua_State *state, type_identity *type)
 
     SetPtrMethods(state, base+1, base+2);
 
+    // Index the fields
+    lua_newtable(state);
+
     EnableMetaField(state, base+2, "value", type);
+    AssociateId(state, base+3, 1, "value");
+
+    // Add the iteration metamethods
+    PushStructMethod(state, base+1, base+3, meta_struct_next);
+    SetPairsMethod(state, base+1, "__pairs");
+    lua_pushnil(state);
+    SetPairsMethod(state, base+1, "__ipairs");
 
+    lua_setfield(state, base+1, "_index_table");
+
+    // Add the indexing metamethods
     SetStructMethod(state, base+1, base+2, meta_primitive_index, "__index");
     SetStructMethod(state, base+1, base+2, meta_primitive_newindex, "__newindex");
 }
@@ -1048,7 +1211,15 @@ static void MakeContainerMetatable(lua_State *state, container_identity *type,
     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);
 
+    // push the index table
     AttachEnumKeys(state, base+1, base+2, ienum);
+
+    PushContainerMethod(state, base+1, base+3, meta_container_next, type, item, count);
+    SetPairsMethod(state, base+1, "__pairs");
+    PushContainerMethod(state, base+1, base+3, meta_container_nexti, type, item, count);
+    SetPairsMethod(state, base+1, "__ipairs");
+
+    lua_pop(state, 1);
 }
 
 /*
@@ -1078,6 +1249,13 @@ void bitfield_identity::build_metatable(lua_State *state)
 
     AttachEnumKeys(state, base+1, base+2, this);
 
+    PushContainerMethod(state, base+1, base+3, meta_bitfield_next, this, NULL, -1);
+    SetPairsMethod(state, base+1, "__pairs");
+    PushContainerMethod(state, base+1, base+3, meta_bitfield_nexti, this, NULL, -1);
+    SetPairsMethod(state, base+1, "__ipairs");
+
+    lua_pop(state, 1);
+
     EnableMetaField(state, base+2, "whole", this);
 }
 

library/LuaWrapper.cpp

@@ -853,6 +853,23 @@ static int meta_enum_attr_index(lua_State *state)
     return 1;
 }
 
+static int meta_nodata(lua_State *state)
+{
+    return 0;
+}
+
+/**
+ * Metamethod: __pairs, returning 1st upvalue as iterator
+ */
+static int meta_pairs(lua_State *state)
+{
+    luaL_checkany(state, 1);
+    lua_pushvalue(state, lua_upvalueindex(1));
+    lua_pushvalue(state, 1);
+    lua_pushnil(state);
+    return 3;
+}
+
 /**
  * Make a metatable with most common fields, and an empty table for UPVAL_FIELDTABLE.
  */
@@ -879,7 +896,8 @@ void LuaWrapper::MakeMetatable(lua_State *state, type_identity *type, const char
     lua_pushstring(state, kind);
     lua_setfield(state, base+1, "_kind");
 
-    lua_newtable(state); // fieldtable
+    // Create the field table
+    lua_newtable(state);
 }
 
 /**
@@ -930,25 +948,49 @@ void LuaWrapper::SetPtrMethods(lua_State *state, int meta_idx, int read_idx)
     EnableMetaField(state, read_idx, "_displace");
 }
 
+/**
+ * Add a __pairs/__ipairs metamethod using iterator on the top of stack.
+ */
+void LuaWrapper::SetPairsMethod(lua_State *state, int meta_idx, const char *name)
+{
+    if (lua_isnil(state, -1))
+    {
+        lua_pop(state, 1);
+        lua_pushcfunction(state, meta_nodata);
+    }
+
+    lua_pushcclosure(state, meta_pairs, 1);
+    lua_setfield(state, meta_idx, name);
+}
+
 /**
  * Add a struct-style (3 upvalues) metamethod to the metatable.
  */
-void LuaWrapper::SetStructMethod(lua_State *state, int meta_idx, int ftable_idx,
-                                 lua_CFunction function, const char *name)
+void LuaWrapper::PushStructMethod(lua_State *state, int meta_idx, int ftable_idx,
+                                  lua_CFunction function)
 {
     lua_getfield(state, LUA_REGISTRYINDEX, DFHACK_TYPETABLE_NAME);
     lua_pushvalue(state, meta_idx);
     lua_pushvalue(state, ftable_idx);
     lua_pushcclosure(state, function, 3);
+}
+
+/**
+ * Add a struct-style (3 upvalues) metamethod to the metatable.
+ */
+void LuaWrapper::SetStructMethod(lua_State *state, int meta_idx, int ftable_idx,
+                                 lua_CFunction function, const char *name)
+{
+    PushStructMethod(state, meta_idx, ftable_idx, function);
     lua_setfield(state, meta_idx, name);
 }
 
 /**
  * Add a 6 upvalue metamethod to the metatable.
  */
-void LuaWrapper::SetContainerMethod(lua_State *state, int meta_idx, int ftable_idx,
-                                    lua_CFunction function, const char *name,
-                                    type_identity *container, type_identity *item, int count)
+void LuaWrapper::PushContainerMethod(lua_State *state, int meta_idx, int ftable_idx,
+                                     lua_CFunction function,
+                                     type_identity *container, type_identity *item, int count)
 {
     lua_getfield(state, LUA_REGISTRYINDEX, DFHACK_TYPETABLE_NAME);
     lua_pushvalue(state, meta_idx);
@@ -962,31 +1004,48 @@ void LuaWrapper::SetContainerMethod(lua_State *state, int meta_idx, int ftable_i
         lua_pushinteger(state, count);
 
     lua_pushcclosure(state, function, 6);
+}
+
+/**
+ * Add a 6 upvalue metamethod to the metatable.
+ */
+void LuaWrapper::SetContainerMethod(lua_State *state, int meta_idx, int ftable_idx,
+                                    lua_CFunction function, const char *name,
+                                    type_identity *container, type_identity *item, int count)
+{
+    PushContainerMethod(state, meta_idx, ftable_idx, function, container, item, count);
     lua_setfield(state, meta_idx, name);
 }
 
 /**
  * If ienum refers to a valid enum, attach its keys to UPVAL_FIELDTABLE,
- * and the enum itself to the _enum metafield.
+ * and the enum itself to the _enum metafield. Pushes the key table on the stack
  */
 void LuaWrapper::AttachEnumKeys(lua_State *state, int meta_idx, int ftable_idx, type_identity *ienum)
 {
+    EnableMetaField(state, ftable_idx, "_enum");
+
+    LookupInTable(state, ienum, DFHACK_TYPEID_TABLE_NAME);
+    lua_setfield(state, meta_idx, "_enum");
+
     LookupInTable(state, ienum, DFHACK_ENUM_TABLE_NAME);
 
     if (!lua_isnil(state, -1))
     {
+        lua_dup(state);
         lua_newtable(state);
         lua_swap(state);
         lua_setfield(state, -2, "__index");
         lua_setmetatable(state, ftable_idx);
     }
     else
+    {
         lua_pop(state, 1);
+        lua_getfield(state, LUA_REGISTRYINDEX, DFHACK_EMPTY_TABLE_NAME);
+    }
 
-    LookupInTable(state, ienum, DFHACK_TYPEID_TABLE_NAME);
-    lua_setfield(state, meta_idx, "_enum");
-
-    EnableMetaField(state, ftable_idx, "_enum");
+    lua_dup(state);
+    lua_setfield(state, meta_idx, "_index_table");
 }
 
 static void BuildTypeMetatable(lua_State *state, type_identity *type)
@@ -1004,7 +1063,7 @@ static void BuildTypeMetatable(lua_State *state, type_identity *type)
 
 static void RenderTypeChildren(lua_State *state, const std::vector<compound_identity*> &children);
 
-static void AssociateId(lua_State *state, int table, int val, const char *name)
+void LuaWrapper::AssociateId(lua_State *state, int table, int val, const char *name)
 {
     lua_pushinteger(state, val);
     lua_pushstring(state, name);
@@ -1157,6 +1216,8 @@ static void RenderType(lua_State *state, compound_identity *node)
 
             lua_getfield(state, -1, "__newindex");
             lua_setfield(state, base+2, "__newindex");
+            lua_getfield(state, -1, "__pairs");
+            lua_setfield(state, base+2, "__pairs");
 
             lua_pop(state, 3);
             return;
@@ -1201,6 +1262,9 @@ static void DoAttach(lua_State *state)
     lua_newtable(state);
     lua_setfield(state, LUA_REGISTRYINDEX, DFHACK_ENUM_TABLE_NAME);
 
+    lua_newtable(state);
+    lua_setfield(state, LUA_REGISTRYINDEX, DFHACK_EMPTY_TABLE_NAME);
+
     lua_pushcfunction(state, change_error);
     lua_setfield(state, LUA_REGISTRYINDEX, DFHACK_CHANGEERROR_NAME);
 

library/include/LuaWrapper.h

@@ -69,6 +69,7 @@ namespace DFHack { namespace LuaWrapper {
 #define DFHACK_NEW_NAME "DFHack::New"
 #define DFHACK_ASSIGN_NAME "DFHack::Assign"
 #define DFHACK_DELETE_NAME "DFHack::Delete"
+#define DFHACK_EMPTY_TABLE_NAME "DFHack::EmptyTable"
 
 /*
  * Upvalue: contents of DFHACK_TYPETABLE_NAME
@@ -160,6 +161,8 @@ namespace DFHack { namespace LuaWrapper {
     void SaveInTable(lua_State *state, void *node, const char *tname);
     void SaveTypeInfo(lua_State *state, void *node);
 
+    void AssociateId(lua_State *state, int table, int val, const char *name);
+
     /**
      * Look up the key on the stack in DFHACK_TYPETABLE;
      * if found, put result on the stack and return true.
@@ -178,11 +181,26 @@ namespace DFHack { namespace LuaWrapper {
      * Set metatable properties common to all actual DF object references.
      */
     void SetPtrMethods(lua_State *state, int meta_idx, int read_idx);
+    /**
+     * Add a __pairs/__ipairs metamethod using iterator on the top of stack.
+     */
+    void SetPairsMethod(lua_State *state, int meta_idx, const char *name);
+    /**
+     * Add a struct-style (3 upvalues) metamethod to the stack.
+     */
+    void PushStructMethod(lua_State *state, int meta_idx, int ftable_idx,
+                          lua_CFunction function);
     /**
      * Add a struct-style (3 upvalues) metamethod to the metatable.
      */
     void SetStructMethod(lua_State *state, int meta_idx, int ftable_idx,
                          lua_CFunction function, const char *name);
+    /**
+     * Add a 6 upvalue metamethod to the stack.
+     */
+    void PushContainerMethod(lua_State *state, int meta_idx, int ftable_idx,
+                             lua_CFunction function,
+                             type_identity *container, type_identity *item, int count);
     /**
      * Add a 6 upvalue metamethod to the metatable.
      */
@@ -191,7 +209,7 @@ namespace DFHack { namespace LuaWrapper {
                             type_identity *container, type_identity *item, int count);
     /**
      * If ienum refers to a valid enum, attach its keys to UPVAL_FIELDTABLE,
-     * and the enum itself to the _enum metafield.
+     * and the enum itself to the _enum metafield. Pushes the key table on the stack.
      */
     void AttachEnumKeys(lua_State *state, int meta_idx, int ftable_idx, type_identity *ienum);