From 574fa08747a2809fe41d74f56e999b8684470e98 Mon Sep 17 00:00:00 2001
From: lethosor <lethosor@gmail.com>
Date: Sat, 12 Aug 2023 22:12:45 -0400
Subject: [PATCH] Add index_enum, ref_target

---
 library/LuaTypes.cpp              | 23 +++++++++++++++++------
 test/structures/struct_fields.lua | 12 ++++++++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/library/LuaTypes.cpp b/library/LuaTypes.cpp
index ac99600a1..407dea265 100644
--- a/library/LuaTypes.cpp
+++ b/library/LuaTypes.cpp
@@ -1365,6 +1365,13 @@ static void IndexFields(lua_State *state, int base, struct_identity *pstruct, in
     }
 }
 
+static void PushTypeIdentity(lua_State *state, const type_identity *id)
+{
+    lua_rawgetp(state, LUA_REGISTRYINDEX, &DFHACK_TYPEID_TABLE_TOKEN);
+    lua_rawgetp(state, -1, id);
+    lua_remove(state, -2);  // TYPEID_TABLE
+}
+
 static void PushFieldInfoSubTable(lua_State *state, const struct_field_info *field)
 {
     if (!field) {
@@ -1384,15 +1391,19 @@ static void PushFieldInfoSubTable(lua_State *state, const struct_field_info *fie
         lua_pushlightuserdata(state, field->type);
         lua_setfield(state, -2, "type_identity");
 
-        lua_pushstring(state, "type");
-        lua_rawgetp(state, LUA_REGISTRYINDEX, &DFHACK_TYPEID_TABLE_TOKEN);
-        lua_rawgetp(state, -1, field->type);
-        lua_remove(state, -2);  // TYPEID_TABLE
-        lua_settable(state, -3);
+        PushTypeIdentity(state, field->type);
+        lua_setfield(state, -2, "type");
     }
 
     if (field->extra) {
-        // TODO: index_enum, ref_target
+        if (field->extra->index_enum) {
+            PushTypeIdentity(state, field->extra->index_enum);
+            lua_setfield(state, -2, "index_enum");
+        }
+        if (field->extra->ref_target) {
+            PushTypeIdentity(state, field->extra->ref_target);
+            lua_setfield(state, -2, "ref_target");
+        }
         if (field->extra->union_tag_field) {
             Lua::TableInsert(state, "union_tag_field", field->extra->union_tag_field);
         }
diff --git a/test/structures/struct_fields.lua b/test/structures/struct_fields.lua
index 53c9ebcc1..d61c71846 100644
--- a/test/structures/struct_fields.lua
+++ b/test/structures/struct_fields.lua
@@ -52,6 +52,18 @@ function test.nonexistent()
     end)
 end
 
+function test.count()
+    expect.eq(df.unit._fields.relationship_ids.count, 10)
+end
+
+function test.index_enum()
+    expect.eq(df.unit._fields.relationship_ids.index_enum, df.unit_relationship_type)
+end
+
+function test.ref_target()
+    expect.eq(df.unit._fields.hist_figure_id.ref_target, df.historical_figure)
+end
+
 function test.readonly()
     expect.error_match(READONLY_MSG, function()
         df.coord._fields.x = 'foo'