Implement inheritance-aware vmethod interposing.

I.e. overwriting the vmethod in all vtables that use it, not only one.

library/DataDefs.cpp

@@ -218,8 +218,8 @@ virtual_identity::virtual_identity(size_t size, TAllocateFn alloc,
 virtual_identity::~virtual_identity()
 {
     // Remove interpose entries, so that they don't try accessing this object later
-    for (int i = interpose_list.size()-1; i >= 0; i--)
+    while (!interpose_list.empty())
-        interpose_list[i]->remove();
+        interpose_list.begin()->second->on_host_delete(this);
 }
 
 /* Vtable name to identity lookup. */

library/VTableInterpose.cpp

@@ -154,6 +154,73 @@ bool virtual_identity::set_vmethod_ptr(int idx, void *ptr)
     return Core::getInstance().p->patchMemory(&vtable[idx], &ptr, sizeof(void*));
 }
 
+/*
+   VMethod interposing data structures.
+
+   In order to properly support adding and removing hooks,
+   it is necessary to track them. This is what this class
+   is for. The task is further complicated by propagating
+   hooks to child classes that use exactly the same original
+   vmethod implementation.
+
+   Every applied link contains in the saved_chain field a
+   pointer to the next vmethod body that should be called
+   by the hook the link represents. This is the actual
+   control flow structure that needs to be maintained.
+
+   There also are connections between link objects themselves,
+   which constitute the bookkeeping for doing that. Finally,
+   every link is associated with a fixed virtual_identity host,
+   which represents the point in the class hierarchy where
+   the hook is applied.
+
+   When there are no subclasses (i.e. only one host), the
+   structures look like this:
+
+   +--------------+             +------------+
+   | link1        |-next------->| link2      |-next=NULL
+   |s_c: original |<-------prev-|s_c: $link1 |<--+
+   +--------------+             +------------+   |
+                                                 |
+         host->interpose_list[vmethod_idx] ------+
+         vtable: $link2
+
+   The original vtable entry is stored in the saved_chain of the
+   first link. The interpose_list map points to the last one.
+   The hooks are called in order: link2 -> link1 -> original.
+
+   When there are subclasses that use the same vmethod, but don't
+   hook it, the topmost link gets a set of the child_hosts, and
+   the hosts have the link added to their interpose_list:
+
+   +--------------+                    +----------------+
+   | link0 @host0 |<--+-interpose_list-| host1          |
+   |              |-child_hosts-+----->| vtable: $link  |
+   +--------------+   |         |      +----------------+
+                      |         |
+                      |         |      +----------------+
+                      +-interpose_list-| host2          |
+                                +----->| vtable: $link  |
+                                       +----------------+
+
+   When a child defines its own hook, the child_hosts link is
+   severed and replaced with a child_next pointer to the new
+   hook. The hook still points back the chain with prev.
+   All child links to subclasses of host2 are migrated from
+   link1 to link2.
+
+   +--------------+-next=NULL         +--------------+-next=NULL
+   | link1 @host1 |-child_next------->| link2 @host2 |-child_*--->subclasses
+   |              |<-------------prev-|s_c: $link1   |
+   +--------------+<-------+          +--------------+<-------+
+                           |                                  |
+   +--------------+        |          +--------------+        |
+   | host1        |-i_list-+          | host2        |-i_list-+
+   |vtable: $link1|                   |vtable: $link2|
+   +--------------+                   +--------------+
+
+ */
+
 void VMethodInterposeLinkBase::set_chain(void *chain)
 {
     saved_chain = chain;
@@ -162,7 +229,7 @@ void VMethodInterposeLinkBase::set_chain(void *chain)
 
 VMethodInterposeLinkBase::VMethodInterposeLinkBase(virtual_identity *host, int vmethod_idx, void *interpose_method, void *chain_mptr)
     : host(host), vmethod_idx(vmethod_idx), interpose_method(interpose_method), chain_mptr(chain_mptr),
-      saved_chain(NULL), next(NULL), prev(NULL)
+      applied(false), saved_chain(NULL), next(NULL), prev(NULL)
 {
     if (vmethod_idx < 0 || interpose_method == NULL)
     {
@@ -179,6 +246,75 @@ VMethodInterposeLinkBase::~VMethodInterposeLinkBase()
         remove();
 }
 
+VMethodInterposeLinkBase *VMethodInterposeLinkBase::get_first_interpose(virtual_identity *id)
+{
+    auto item = id->interpose_list[vmethod_idx];
+    if (!item)
+        return NULL;
+
+    if (item->host != id)
+        return NULL;
+    while (item->prev && item->prev->host == id)
+        item = item->prev;
+
+    return item;
+}
+
+void VMethodInterposeLinkBase::find_child_hosts(virtual_identity *cur, void *vmptr)
+{
+    auto &children = cur->getChildren();
+
+    for (size_t i = 0; i < children.size(); i++)
+    {
+        auto child = static_cast<virtual_identity*>(children[i]);
+        auto base = get_first_interpose(child);
+
+        if (base)
+        {
+            assert(base->prev == NULL);
+
+            if (base->saved_chain != vmptr)
+                continue;
+
+            child_next.insert(base);
+        }
+        else
+        {
+            void *cptr = child->get_vmethod_ptr(vmethod_idx);
+            if (cptr != vmptr)
+                continue;
+
+            child_hosts.insert(child);
+            find_child_hosts(child, vmptr);
+        }
+    }
+}
+
+void VMethodInterposeLinkBase::on_host_delete(virtual_identity *from)
+{
+    if (from == host)
+    {
+        // When in own host, fully delete
+        remove();
+    }
+    else
+    {
+        // Otherwise, drop the link to that child:
+        assert(child_hosts.count(from) != 0 &&
+               from->interpose_list[vmethod_idx] == this);
+
+        // Find and restore the original vmethod ptr
+        auto last = this;
+        while (last->prev) last = last->prev;
+
+        from->set_vmethod_ptr(vmethod_idx, last->saved_chain);
+
+        // Unlink the chains
+        child_hosts.erase(from);
+        from->interpose_list.erase(vmethod_idx);
+    }
+}
+
 bool VMethodInterposeLinkBase::apply()
 {
     if (is_applied())
@@ -188,33 +324,73 @@ bool VMethodInterposeLinkBase::apply()
 
     // Retrieve the current vtable entry
     void *old_ptr = host->get_vmethod_ptr(vmethod_idx);
-    assert(old_ptr != NULL);
+    VMethodInterposeLinkBase *old_link = host->interpose_list[vmethod_idx];
-
-    // Check if there are other interpose entries for the same slot
-    VMethodInterposeLinkBase *old_link = NULL;
-
-    for (int i = host->interpose_list.size()-1; i >= 0; i--)
-    {
-        if (host->interpose_list[i]->vmethod_idx != vmethod_idx)
-            continue;
 
-        old_link = host->interpose_list[i];
+    assert(old_ptr != NULL && (!old_link || old_link->interpose_method == old_ptr));
-        assert(old_link->next == NULL && old_ptr == old_link->interpose_method);
-        break;
-    }
 
     // Apply the new method ptr
+    set_chain(old_ptr);
+
     if (!host->set_vmethod_ptr(vmethod_idx, interpose_method))
+    {
+        set_chain(NULL);
         return false;
+    }
 
-    set_chain(old_ptr);
+    // Push the current link into the home host
-    host->interpose_list.push_back(this);
+    applied = true;
+    host->interpose_list[vmethod_idx] = this;
+    prev = old_link;
 
-    // Link into the chain if any
+    child_hosts.clear();
-    if (old_link)
+    child_next.clear();
+
+    if (old_link && old_link->host == host)
     {
+        // If the old link is home, just push into the plain chain
+        assert(old_link->next == NULL);
         old_link->next = this;
-        prev = old_link;
+
+        // Child links belong to the topmost local entry
+        child_hosts.swap(old_link->child_hosts);
+        child_next.swap(old_link->child_next);
+    }
+    else
+    {
+        // If creating a new local chain, find children with same vmethod
+        find_child_hosts(host, old_ptr);
+
+        if (old_link)
+        {
+            // Enter the child chain set
+            assert(old_link->child_hosts.count(host));
+            old_link->child_hosts.erase(host);
+            old_link->child_next.insert(this);
+
+            // Subtract our own children from the parent's sets
+            for (auto it = child_next.begin(); it != child_next.end(); ++it)
+                old_link->child_next.erase(*it);
+            for (auto it = child_hosts.begin(); it != child_hosts.end(); ++it)
+                old_link->child_hosts.erase(*it);
+        }
+    }
+
+    // Chain subclass hooks
+    for (auto it = child_next.begin(); it != child_next.end(); ++it)
+    {
+        auto nlink = *it;
+        assert(nlink->saved_chain == old_ptr && nlink->prev == old_link);
+        nlink->set_chain(interpose_method);
+        nlink->prev = this;
+    }
+
+    // Chain passive subclass hosts
+    for (auto it = child_hosts.begin(); it != child_hosts.end(); ++it)
+    {
+        auto nhost = *it;
+        assert(nhost->interpose_list[vmethod_idx] == old_link);
+        nhost->set_vmethod_ptr(vmethod_idx, interpose_method);
+        nhost->interpose_list[vmethod_idx] = this;
     }
 
     return true;
@@ -225,25 +401,57 @@ void VMethodInterposeLinkBase::remove()
     if (!is_applied())
         return;
 
-    // Remove from the list in the identity
+    // Remove the link from prev to this
-    for (int i = host->interpose_list.size()-1; i >= 0; i--)
-        if (host->interpose_list[i] == this)
-            vector_erase_at(host->interpose_list, i);
-
-    // Remove from the chain
     if (prev)
+    {
+        if (prev->host == host)
             prev->next = next;
+        else
+        {
+            prev->child_next.erase(this);
+
+            if (next)
+                prev->child_next.insert(next);
+        }
+    }
 
     if (next)
     {
         next->set_chain(saved_chain);
         next->prev = prev;
+
+        assert(child_next.empty() && child_hosts.empty());
     }
     else
     {
+        // Remove from the list in the identity and vtable
+        host->interpose_list[vmethod_idx] = prev;
         host->set_vmethod_ptr(vmethod_idx, saved_chain);
+
+        for (auto it = child_next.begin(); it != child_next.end(); ++it)
+        {
+            auto nlink = *it;
+            assert(nlink->saved_chain == interpose_method && nlink->prev == this);
+            nlink->set_chain(saved_chain);
+            nlink->prev = prev;
+            if (prev)
+                prev->child_next.insert(nlink);
+        }
+
+        for (auto it = child_hosts.begin(); it != child_hosts.end(); ++it)
+        {
+            auto nhost = *it;
+            assert(nhost->interpose_list[vmethod_idx] == this);
+            nhost->interpose_list[vmethod_idx] = prev;
+            nhost->set_vmethod_ptr(vmethod_idx, saved_chain);
+            if (prev)
+                prev->child_hosts.insert(nhost);
+        }
     }
 
+    applied = false;
     prev = next = NULL;
+    child_next.clear();
+    child_hosts.clear();
     set_chain(NULL);
 }

library/include/DataDefs.h

@@ -303,7 +303,7 @@ namespace DFHack
         void *vtable_ptr;
 
         friend class VMethodInterposeLinkBase;
-        std::vector<VMethodInterposeLinkBase*> interpose_list;
+        std::map<int,VMethodInterposeLinkBase*> interpose_list;
 
     protected:
         virtual void doInit(Core *core);

library/include/VTableInterpose.h

@@ -134,21 +134,31 @@ namespace DFHack
           1) Allow multiple hooks into the same vmethod
           2) Auto-remove hooks when a plugin is unloaded.
         */
+        friend class virtual_identity;
 
         virtual_identity *host; // Class with the vtable
         int vmethod_idx;
         void *interpose_method; // Pointer to the code of the interposing method
         void *chain_mptr;       // Pointer to the chain field below
 
+        bool applied;
         void *saved_chain;      // Previous pointer to the code
         VMethodInterposeLinkBase *next, *prev; // Other hooks for the same method
 
+        // inherited vtable members
+        std::set<virtual_identity*> child_hosts;
+        std::set<VMethodInterposeLinkBase*> child_next;
+
         void set_chain(void *chain);
+        void on_host_delete(virtual_identity *host);
+
+        VMethodInterposeLinkBase *get_first_interpose(virtual_identity *id);
+        void find_child_hosts(virtual_identity *cur, void *vmptr);
     public:
         VMethodInterposeLinkBase(virtual_identity *host, int vmethod_idx, void *interpose_method, void *chain_mptr);
         ~VMethodInterposeLinkBase();
 
-        bool is_applied() { return saved_chain != NULL; }
+        bool is_applied() { return applied; }
         bool apply();
         void remove();
     };