From a214e004074e74510ad3d0b959b3298b1721a2d3 Mon Sep 17 00:00:00 2001
From: Ben Lubar <ben.lubar@gmail.com>
Date: Fri, 6 Mar 2020 16:48:33 -0600
Subject: [PATCH] allow union vectors to have tags that are bit vectors if they
 have exactly 2 members

---
 library/DataDefs.cpp                      |  9 +++++
 library/include/DataDefs.h                |  3 ++
 plugins/devel/check-structures-sanity.cpp | 41 +++++++++++++++++++++--
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/library/DataDefs.cpp b/library/DataDefs.cpp
index cd3474619..c071ad7d3 100644
--- a/library/DataDefs.cpp
+++ b/library/DataDefs.cpp
@@ -539,5 +539,14 @@ const struct_field_info *DFHack::find_union_tag(const struct_field_info *fields,
         return tag_candidate;
     }
 
+    auto union_fields = ((struct_identity*)union_field->type)->getFields();
+    if (tag_container_type->getFullName(nullptr) != "vector<bool>" &&
+            union_fields[0].mode != struct_field_info::END &&
+            union_fields[1].mode != struct_field_info::END &&
+            union_fields[2].mode == struct_field_info::END)
+    {
+        return tag_candidate;
+    }
+
     return nullptr;
 }
diff --git a/library/include/DataDefs.h b/library/include/DataDefs.h
index 704b1abf7..03bb7857f 100644
--- a/library/include/DataDefs.h
+++ b/library/include/DataDefs.h
@@ -804,6 +804,9 @@ namespace DFHack {
      *
      * If the union field is a container type, the returned tag field is
      * a container of primitive enum types.
+     *
+     * As a special case, a container-type union can have a tag field that is
+     * a bit vector if it has exactly two members.
      */
     DFHACK_EXPORT const struct_field_info *find_union_tag(const struct_field_info *fields, const struct_field_info *union_field);
 }
diff --git a/plugins/devel/check-structures-sanity.cpp b/plugins/devel/check-structures-sanity.cpp
index b5cd1aa5f..3a454a0e6 100644
--- a/plugins/devel/check-structures-sanity.cpp
+++ b/plugins/devel/check-structures-sanity.cpp
@@ -150,6 +150,7 @@ private:
     bool maybe_queue_union(const ToCheck &, const struct_field_info *, const struct_field_info *);
     void queue_union(const ToCheck &, const ToCheck &);
     void queue_union_vector(const ToCheck &, const ToCheck &);
+    void queue_union_bitvector(const ToCheck &, const ToCheck &);
     void check_dispatch(ToCheck &);
     void check_global(const ToCheck &);
     void check_primitive(const ToCheck &);
@@ -163,7 +164,7 @@ private:
     void check_deque(const ToCheck &, type_identity *);
     void check_dfarray(const ToCheck &, type_identity *);
     void check_bitarray(const ToCheck &);
-    void check_bitvector(const ToCheck &);
+    bool check_bitvector(const ToCheck &);
     void check_struct(const ToCheck &);
     void check_virtual(const ToCheck &);
 public:
@@ -691,6 +692,12 @@ void Checker::queue_union(const ToCheck & item, const ToCheck & tag_item)
 
 void Checker::queue_union_vector(const ToCheck & item, const ToCheck & tag_item)
 {
+    if (tag_item.identity->getFullName(nullptr) == "vector<bool>")
+    {
+        queue_union_bitvector(item, tag_item);
+        return;
+    }
+
     auto union_type = static_cast<union_identity *>(static_cast<container_identity *>(item.identity)->getItemType());
     auto tag_type = static_cast<enum_identity *>(static_cast<container_identity *>(tag_item.identity)->getItemType());
 
@@ -714,6 +721,33 @@ void Checker::queue_union_vector(const ToCheck & item, const ToCheck & tag_item)
     }
 }
 
+void Checker::queue_union_bitvector(const ToCheck & item, const ToCheck & tag_item)
+{
+    auto union_type = static_cast<union_identity *>(static_cast<container_identity *>(item.identity)->getItemType());
+    auto union_count = check_vector_size(item, union_type->byte_size());
+
+    if (!check_bitvector(tag_item))
+    {
+        return;
+    }
+    auto tag_vector = reinterpret_cast<std::vector<bool> *>(tag_item.ptr);
+
+    if (union_count != tag_vector->size())
+    {
+        FAIL("tagged union vector size (" << union_count << ") does not match tag vector (" << join_strings("", tag_item.path) << ") size (" << tag_count << ")");
+    }
+
+    auto union_base = *reinterpret_cast<void **>(item.ptr);
+
+    auto count = union_count < tag_vector->size() ? union_count : tag_vector->size();
+    for (size_t i = 0; i < count; i++, union_base = PTR_ADD(union_base, union_type->byte_size()))
+    {
+        auto item_field = &union_type->getFields()[tag_vector->at(i) ? 1 : 0];
+        ToCheck tagged_union_item(item, stl_sprintf("[%zu].%s", i, item_field->name), union_base, item_field->type);
+        queue_field(std::move(tagged_union_item), item_field);
+    }
+}
+
 void Checker::check_dispatch(ToCheck & item)
 {
     if (reinterpret_cast<uintptr_t>(item.ptr) == UNINIT_PTR)
@@ -1267,7 +1301,7 @@ void Checker::check_bitarray(const ToCheck & item)
     // TODO: check DFHack::BitArray?
 }
 
-void Checker::check_bitvector(const ToCheck & item)
+bool Checker::check_bitvector(const ToCheck & item)
 {
     struct biterator_data
     {
@@ -1285,10 +1319,11 @@ void Checker::check_bitvector(const ToCheck & item)
     if (item.identity->byte_size() != sizeof(bvector_data))
     {
         UNEXPECTED;
-        return;
+        return false;
     }
 
     // TODO: check vector<bool>?
+    return true;
 }
 
 void Checker::check_struct(const ToCheck & item)