Lua: implement key lookup and ipairs for complex enums

library/LuaWrapper.cpp

@@ -1349,6 +1349,58 @@ static int wtype_next_item(lua_State *state)
     return 1;
 }
 
+/*
+ * Complex enums
+ *
+ * upvalues for all of these:
+ *  1: key table? unsure, taken from wtype stuff
+ *  2: enum_identity::ComplexData
+ */
+
+static bool complex_enum_next_item_helper(lua_State *L, int64_t &item)
+{
+    const auto *complex = (enum_identity::ComplexData*)lua_touserdata(L, lua_upvalueindex(2));
+    auto it = complex->value_index_map.find(item);
+    if (it != complex->value_index_map.end())
+    {
+        size_t index = it->second;
+        if (index >= complex->size() - 1)
+            return false;
+
+        item = complex->index_value_map[index + 1];
+        return true;
+    }
+    return false;
+}
+
+static int complex_enum_inext(lua_State *L)
+{
+    int64_t first_item = ((enum_identity::ComplexData*)lua_touserdata(L, lua_upvalueindex(2)))->index_value_map[0];
+    int64_t i = (lua_isuserdata(L, 2)) ? first_item : luaL_checkint(L, 2);
+    if (complex_enum_next_item_helper(L, i))
+    {
+        lua_pushinteger(L, i);
+        lua_rawgeti(L, lua_upvalueindex(1), i);
+        return 2;
+    }
+    else
+    {
+        lua_pushnil(L);
+        return 1;
+    }
+}
+
+static int complex_enum_ipairs(lua_State *L)
+{
+    lua_pushvalue(L, lua_upvalueindex(1));
+    lua_pushvalue(L, lua_upvalueindex(2));
+    lua_pushcclosure(L, complex_enum_inext, 2);
+    lua_pushnil(L);
+    lua_pushlightuserdata(L, (void*)1);
+    return 3;
+}
+
+
 static void RenderTypeChildren(lua_State *state, const std::vector<compound_identity*> &children);
 
 void LuaWrapper::AssociateId(lua_State *state, int table, int val, const char *name)
@@ -1371,17 +1423,39 @@ static void FillEnumKeys(lua_State *state, int ix_meta, int ftable, enum_identit
     int base = lua_gettop(state);
     lua_newtable(state);
 
+    auto *complex = eid->getComplex();
+
     // For enums, set mapping between keys and values
+    if (complex)
+    {
+        for (size_t i = 0; i < complex->size(); i++)
+        {
+            if (keys[i])
+                AssociateId(state, base+1, complex->index_value_map[i], keys[i]);
+        }
+    }
+    else
+    {
         for (int64_t i = eid->getFirstItem(), j = 0; i <= eid->getLastItem(); i++, j++)
         {
             if (keys[j])
                 AssociateId(state, base+1, i, keys[j]);
         }
+    }
 
+    if (complex)
+    {
+        lua_pushvalue(state, base + 1);
+        lua_pushlightuserdata(state, (void*)complex);
+        lua_pushcclosure(state, complex_enum_ipairs, 2);
+        lua_setfield(state, ix_meta, "__ipairs");
+    }
+    else
+    {
         if (eid->getFirstItem() <= eid->getLastItem())
         {
-        lua_pushvalue(state, base+1);
+            lua_pushvalue(state, base + 1);
-        lua_pushinteger(state, eid->getFirstItem()-1);
+            lua_pushinteger(state, eid->getFirstItem() - 1);
             lua_pushinteger(state, eid->getLastItem());
             lua_pushcclosure(state, wtype_ipairs, 3);
             lua_setfield(state, ix_meta, "__ipairs");
@@ -1397,6 +1471,7 @@ static void FillEnumKeys(lua_State *state, int ix_meta, int ftable, enum_identit
             lua_pushinteger(state, eid->getLastItem());
             lua_setfield(state, ftable, "_last_item");
         }
+    }
 
     SaveInTable(state, eid, &DFHACK_ENUM_TABLE_TOKEN);
 

library/include/DataDefs.h

@@ -231,6 +231,7 @@ namespace DFHack
         int64_t getLastItem() { return last_item_value; }
         int getCount() { return count; }
         const char *const *getKeys() { return keys; }
+        const ComplexData *getComplex() { return complex; }
 
         type_identity *getBaseType() { return base_type; }
         const void *getAttrs() { return attrs; }