Merge pull request #1254 from lethosor/complex-enum

Complex enum metadata

library/DataDefs.cpp

@@ -143,11 +143,30 @@ enum_identity::enum_identity(size_t size,
                              type_identity *base_type,
                              int64_t first_item_value, int64_t last_item_value,
                              const char *const *keys,
+                             const ComplexData *complex,
                              const void *attrs, struct_identity *attr_type)
     : compound_identity(size, NULL, scope_parent, dfhack_name),
-      keys(keys), first_item_value(first_item_value), last_item_value(last_item_value),
+      keys(keys), complex(complex),
+      first_item_value(first_item_value), last_item_value(last_item_value),
       base_type(base_type), attrs(attrs), attr_type(attr_type)
 {
+    if (complex) {
+        count = complex->size();
+        last_item_value = complex->index_value_map.back();
+    }
+    else {
+        count = int(last_item_value-first_item_value+1);
+    }
+}
+
+enum_identity::ComplexData::ComplexData(std::initializer_list<int64_t> values)
+{
+    size_t i = 0;
+    for (int64_t value : values) {
+        value_index_map[value] = i;
+        index_value_map.push_back(value);
+        i++;
+    }
 }
 
 struct_identity::struct_identity(size_t size, TAllocateFn alloc,

library/LuaWrapper.cpp

@@ -1025,11 +1025,23 @@ static int meta_enum_attr_index(lua_State *state)
         luaL_error(state, "Invalid index in enum.attrs[]");
 
     auto id = (enum_identity*)lua_touserdata(state, lua_upvalueindex(2));
+    auto *complex = id->getComplex();
 
     int64_t idx = lua_tonumber(state, 2);
+    if (complex)
+    {
+        auto it = complex->value_index_map.find(idx);
+        if (it != complex->value_index_map.end())
+            idx = int64_t(it->second);
+        else
+            idx = id->getLastItem() + 1;
+    }
+    else
+    {
         if (idx < id->getFirstItem() || idx > id->getLastItem())
             idx = id->getLastItem()+1;
         idx -= id->getFirstItem();
+    }
 
     uint8_t *ptr = (uint8_t*)id->getAttrs();
     auto atype = id->getAttrType();
@@ -1349,6 +1361,68 @@ 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, bool wrap = false)
+{
+    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 (!wrap && index >= complex->size() - 1)
+            return false;
+
+        item = complex->index_value_map[(index + 1) % complex->size()];
+        return true;
+    }
+    return false;
+}
+
+static int complex_enum_inext(lua_State *L)
+{
+    bool is_first = lua_isuserdata(L, 2);
+    int64_t i = (is_first)
+        ? ((enum_identity::ComplexData*)lua_touserdata(L, lua_upvalueindex(2)))->index_value_map[0]
+        : luaL_checkint(L, 2);
+    if (is_first || 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_next_item(lua_State *L)
+{
+    int64_t cur = luaL_checkint(L, lua_gettop(L) > 1 ? 2 : 1); // 'self' optional
+    complex_enum_next_item_helper(L, cur, true);
+    lua_pushinteger(L, cur);
+    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 +1445,53 @@ 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");
+
+        lua_pushinteger(state, 0); // unused; to align ComplexData
+        lua_pushlightuserdata(state, (void*)complex);
+        lua_pushcclosure(state, complex_enum_next_item, 2);
+        lua_setfield(state, ftable, "next_item");
+
+        lua_pushinteger(state, eid->getFirstItem());
+        lua_setfield(state, ftable, "_first_item");
+
+        lua_pushinteger(state, eid->getLastItem());
+        lua_setfield(state, ftable, "_last_item");
 
+        lua_pushboolean(state, true);
+        lua_setfield(state, ftable, "_complex");
+    }
+    else
+    {
         if (eid->getFirstItem() <= eid->getLastItem())
         {
-        lua_pushvalue(state, base+1);
-        lua_pushinteger(state, eid->getFirstItem()-1);
+            lua_pushvalue(state, base + 1);
+            lua_pushinteger(state, eid->getFirstItem() - 1);
             lua_pushinteger(state, eid->getLastItem());
             lua_pushcclosure(state, wtype_ipairs, 3);
             lua_setfield(state, ix_meta, "__ipairs");
@@ -1396,6 +1506,10 @@ static void FillEnumKeys(lua_State *state, int ix_meta, int ftable, enum_identit
 
             lua_pushinteger(state, eid->getLastItem());
             lua_setfield(state, ftable, "_last_item");
+
+            lua_pushboolean(state, false);
+            lua_setfield(state, ftable, "_complex");
+        }
     }
 
     SaveInTable(state, eid, &DFHACK_ENUM_TABLE_TOKEN);

library/include/DataDefs.h

@@ -188,9 +188,22 @@ namespace DFHack
     class struct_identity;
 
     class DFHACK_EXPORT enum_identity : public compound_identity {
+    public:
+        struct ComplexData {
+            std::map<int64_t, size_t> value_index_map;
+            std::vector<int64_t> index_value_map;
+            ComplexData(std::initializer_list<int64_t> values);
+            size_t size() const {
+                return index_value_map.size();
+            }
+        };
+
+    private:
         const char *const *keys;
+        const ComplexData *complex;
         int64_t first_item_value;
         int64_t last_item_value;
+        int count;
 
         type_identity *base_type;
 
@@ -209,14 +222,16 @@ namespace DFHack
                       type_identity *base_type,
                       int64_t first_item_value, int64_t last_item_value,
                       const char *const *keys,
+                      const ComplexData *complex,
                       const void *attrs, struct_identity *attr_type);
 
         virtual identity_type type() { return IDTYPE_ENUM; }
 
         int64_t getFirstItem() { return first_item_value; }
         int64_t getLastItem() { return last_item_value; }
-        int getCount() { return int(last_item_value-first_item_value+1); }
+        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; }
@@ -499,33 +514,100 @@ namespace DFHack {
      */
 
     /**
-     * Return the next item in the enum, wrapping to the first one at the end.
+     * Return the next item in the enum, wrapping to the first one at the end if 'wrap' is true (otherwise an invalid item).
      */
     template<class T>
-    inline typename df::enum_traits<T>::enum_type next_enum_item(T v) {
+    inline typename std::enable_if<
+        !df::enum_traits<T>::is_complex,
+        typename df::enum_traits<T>::enum_type
+    >::type next_enum_item(T v, bool wrap = true)
+    {
         typedef df::enum_traits<T> traits;
         typedef typename traits::base_type base_type;
         base_type iv = base_type(v);
-        return (iv < traits::last_item_value) ? T(iv+1) : traits::first_item;
+        if (iv < traits::last_item_value)
+        {
+            return T(iv + 1);
+        }
+        else
+        {
+            if (wrap)
+                return traits::first_item;
+            else
+                return T(traits::last_item_value + 1);
+        }
+    }
+
+    template<class T>
+    inline typename std::enable_if<
+        df::enum_traits<T>::is_complex,
+        typename df::enum_traits<T>::enum_type
+    >::type next_enum_item(T v, bool wrap = true)
+    {
+        typedef df::enum_traits<T> traits;
+        const auto &complex = traits::complex;
+        const auto it = complex.value_index_map.find(v);
+        if (it != complex.value_index_map.end())
+        {
+            if (!wrap && it->second + 1 == complex.size())
+            {
+                return T(traits::last_item_value + 1);
+            }
+            size_t next_index = (it->second + 1) % complex.size();
+            return T(complex.index_value_map[next_index]);
+        }
+        else
+            return T(traits::last_item_value + 1);
     }
 
     /**
      * Check if the value is valid for its enum type.
      */
     template<class T>
-    inline bool is_valid_enum_item(T v) {
+    inline typename std::enable_if<
+        !df::enum_traits<T>::is_complex,
+        bool
+    >::type is_valid_enum_item(T v)
+    {
         return df::enum_traits<T>::is_valid(v);
     }
 
+    template<class T>
+    inline typename std::enable_if<
+        df::enum_traits<T>::is_complex,
+        bool
+    >::type is_valid_enum_item(T v)
+    {
+        const auto &complex = df::enum_traits<T>::complex;
+        return complex.value_index_map.find(v) != complex.value_index_map.end();
+    }
+
     /**
      * Return the enum item key string pointer, or NULL if none.
      */
     template<class T>
-    inline const char *enum_item_raw_key(T val) {
+    inline typename std::enable_if<
+        !df::enum_traits<T>::is_complex,
+        const char *
+    >::type enum_item_raw_key(T val) {
         typedef df::enum_traits<T> traits;
         return traits::is_valid(val) ? traits::key_table[(short)val - traits::first_item_value] : NULL;
     }
 
+    template<class T>
+    inline typename std::enable_if<
+        df::enum_traits<T>::is_complex,
+        const char *
+    >::type enum_item_raw_key(T val) {
+        typedef df::enum_traits<T> traits;
+        const auto &value_index_map = traits::complex.value_index_map;
+        auto it = value_index_map.find(val);
+        if (it != value_index_map.end())
+            return traits::key_table[it->second];
+        else
+            return NULL;
+    }
+
     /**
      * Return the enum item key string pointer, or "?" if none.
      */
@@ -703,7 +785,7 @@ namespace DFHack {
 #define ENUM_NEXT_ITEM(enum,val) \
     (DFHack::next_enum_item<df::enum>(val))
 #define FOR_ENUM_ITEMS(enum,iter) \
-    for(df::enum iter = ENUM_FIRST_ITEM(enum); is_valid_enum_item(iter); iter = df::enum(1+int(iter)))
+    for(df::enum iter = ENUM_FIRST_ITEM(enum); DFHack::is_valid_enum_item(iter); iter = DFHack::next_enum_item(iter, false))
 
 /*
  * Include mandatory generated headers.

library/xml

@@ -1 +1 @@
-Subproject commit f3a49dbf40e4ff0d47f9631adc350196252df854
+Subproject commit 028c4a3031d828122d7ad4efdbc0522251542891