diff --git a/library/DataDefs.cpp b/library/DataDefs.cpp index ab50eb942..3d802638f 100644 --- a/library/DataDefs.cpp +++ b/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 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, diff --git a/library/LuaWrapper.cpp b/library/LuaWrapper.cpp index 69fdd734e..fe8309736 100644 --- a/library/LuaWrapper.cpp +++ b/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 (idx < id->getFirstItem() || idx > id->getLastItem()) - idx = id->getLastItem()+1; - idx -= id->getFirstItem(); + 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 &children); void LuaWrapper::AssociateId(lua_State *state, int table, int val, const char *name) @@ -1371,24 +1445,36 @@ 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 - for (int64_t i = eid->getFirstItem(), j = 0; i <= eid->getLastItem(); i++, j++) + 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 { - if (keys[j]) - AssociateId(state, base+1, i, keys[j]); + for (int64_t i = eid->getFirstItem(), j = 0; i <= eid->getLastItem(); i++, j++) + { + if (keys[j]) + AssociateId(state, base+1, i, keys[j]); + } } - if (eid->getFirstItem() <= eid->getLastItem()) + if (complex) { - lua_pushvalue(state, base+1); - lua_pushinteger(state, eid->getFirstItem()-1); - lua_pushinteger(state, eid->getLastItem()); - lua_pushcclosure(state, wtype_ipairs, 3); + 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, eid->getFirstItem()); - lua_pushinteger(state, eid->getLastItem()); - lua_pushcclosure(state, wtype_next_item, 2); + 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()); @@ -1396,6 +1482,34 @@ 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, true); + lua_setfield(state, ftable, "_complex"); + } + else + { + if (eid->getFirstItem() <= eid->getLastItem()) + { + 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"); + + lua_pushinteger(state, eid->getFirstItem()); + lua_pushinteger(state, eid->getLastItem()); + lua_pushcclosure(state, wtype_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, false); + lua_setfield(state, ftable, "_complex"); + } } SaveInTable(state, eid, &DFHACK_ENUM_TABLE_TOKEN); diff --git a/library/include/DataDefs.h b/library/include/DataDefs.h index e72a893e5..f964455a8 100644 --- a/library/include/DataDefs.h +++ b/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 value_index_map; + std::vector index_value_map; + ComplexData(std::initializer_list 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 - inline typename df::enum_traits::enum_type next_enum_item(T v) { + inline typename std::enable_if< + !df::enum_traits::is_complex, + typename df::enum_traits::enum_type + >::type next_enum_item(T v, bool wrap = true) + { typedef df::enum_traits 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 + inline typename std::enable_if< + df::enum_traits::is_complex, + typename df::enum_traits::enum_type + >::type next_enum_item(T v, bool wrap = true) + { + typedef df::enum_traits 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 - inline bool is_valid_enum_item(T v) { + inline typename std::enable_if< + !df::enum_traits::is_complex, + bool + >::type is_valid_enum_item(T v) + { return df::enum_traits::is_valid(v); } + template + inline typename std::enable_if< + df::enum_traits::is_complex, + bool + >::type is_valid_enum_item(T v) + { + const auto &complex = df::enum_traits::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 - inline const char *enum_item_raw_key(T val) { + inline typename std::enable_if< + !df::enum_traits::is_complex, + const char * + >::type enum_item_raw_key(T val) { typedef df::enum_traits traits; return traits::is_valid(val) ? traits::key_table[(short)val - traits::first_item_value] : NULL; } + template + inline typename std::enable_if< + df::enum_traits::is_complex, + const char * + >::type enum_item_raw_key(T val) { + typedef df::enum_traits 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(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. diff --git a/library/xml b/library/xml index f3a49dbf4..028c4a303 160000 --- a/library/xml +++ b/library/xml @@ -1 +1 @@ -Subproject commit f3a49dbf40e4ff0d47f9631adc350196252df854 +Subproject commit 028c4a3031d828122d7ad4efdbc0522251542891