1 files changed, 66 insertions, 23 deletions

diff --git a/src/disasm.cpp b/src/disasm.cpp
index d88a27e..a20cb92 100644
--- a/src/disasm.cpp
+++ b/src/disasm.cpp
@@ -23,6 +23,24 @@ void DisasmNode::AddReferencedBy(const uint32_t address_from, const ReferenceTyp
     this->last_ref_by = node;
 }
 
+void DisasmNode::RemoveReferencedBy(const uint32_t address_from)
+{
+    ReferenceRecord *ref{this->ref_by};
+    while (ref) {
+        ReferenceRecord *prev = ref;
+        ref = ref->next;
+        if (prev->address == address_from) {
+            if (this->ref_by == prev) {
+                this->ref_by = prev->next;
+            }
+            delete prev;
+        }
+    }
+    if (nullptr == this->ref_by) {
+        this->last_ref_by = nullptr;
+    }
+}
+
 DisasmNode::~DisasmNode()
 {
     ReferenceRecord *ref{this->ref_by};
@@ -49,7 +67,7 @@ DisasmNode &DisasmMap::insertNode(uint32_t address, NodeType type)
         // traced data support is yet to come.
         if (IsInstruction(type) && !IsInstruction(node->type)) {
             if (0 == (node->size & 1) && 0 == (node->address & 1)) {
-                *const_cast<NodeType*>(&node->type) = type;
+                node->type = type;
                 // Make sure it is OpCode::kNone so it will be properly disassembled
                 node->op = Op{};
             }
@@ -85,7 +103,7 @@ DisasmNode &DisasmMap::insertNodeQuickPeek(uint32_t address, NodeType type)
         // traced data support is yet to come.
         if (IsInstruction(type) && !IsInstruction(node->type)) {
             if (0 == (node->size & 1) && 0 == (node->address & 1)) {
-                *const_cast<NodeType*>(&node->type) = type;
+                node->type = type;
                 // Make sure it is OpCode::kNone so it will be properly disassembled
                 node->op = Op{};
             }
@@ -102,32 +120,57 @@ DisasmNode &DisasmMap::insertNodeQuickPeek(uint32_t address, NodeType type)
     return *node;
 }
 
-DisasmNode *DisasmMap::mergeNodes(DisasmNode *primary, DisasmNode *secondary)
+static bool WithinRange(uint32_t address, uint32_t range_address, uint32_t range_size)
+{
+    return address >= range_address && address < range_address + range_size;
+}
+
+DisasmNode *DisasmMap::mergeNodeOverlappingSpace(DisasmNode *primary, DisasmNode *secondary)
 {
     ASSERT(primary->address < secondary->address);
     ASSERT(primary->address + primary->size >= secondary->address);
-    ASSERT(primary->address + primary->size >= secondary->address + secondary->size);
-    ReferenceNode *rnode{secondary->ref_by};
-    while (rnode) {
-        for (size_t i = 0; i < rnode->refs_count; i--) {
-            primary->AddReferencedBy(rnode->refs[i].address, rnode->refs[i].type);
+    ReferenceRecord *ref{secondary->ref_by};
+    while (ref) {
+        ReferenceRecord *prev = ref;
+        ref = ref->next;
+        if (WithinRange(prev->address, primary->address, primary->size)) {
+            primary->AddReferencedBy(prev->address, prev->type);
+            if (secondary->ref_by == prev) {
+                secondary->ref_by = prev->next;
+            }
+            delete prev;
         }
-        ReferenceNode *prev = rnode;
-        rnode = rnode->next;
-        delete prev;
     }
-    if (secondary.ref_kinds & kRef1Mask) {
-        DisasmNode *node = _map[secondary.ref1_addr];
+    if (nullptr == secondary->ref_by) {
+        secondary->last_ref_by = nullptr;
+    }
+    if (secondary->ref_kinds & kRef1Mask) {
+        DisasmNode *node = _map[secondary->ref1_addr];
         ASSERT(node);
-        node->RemoveReferencedBy(secondary.ref1_addr);
+        node->RemoveReferencedBy(secondary->ref1_addr);
     }
-    if (secondary.ref_kinds & kRef2Mask) {
-        DisasmNode *node = _map[secondary.ref2_addr];
+    if (secondary->ref_kinds & kRef2Mask) {
+        DisasmNode *node = _map[secondary->ref2_addr];
         ASSERT(node);
-        node->RemoveReferencedBy(secondary.ref2_addr);
+        node->RemoveReferencedBy(secondary->ref2_addr);
+    }
+    // Spread primary across the overlapping space
+    const size_t begin = secondary->address;
+    const size_t end = primary->address + primary->size;
+    for (size_t address = begin; address < end; address++) {
+        ASSERT(_map[address] == secondary);
+        _map[address] = primary;
+    }
+    if (secondary->address + secondary->size > end) {
+        // It is still alive, just shrunk
+        secondary->size = secondary->address + secondary->size - end;
+        secondary->address = end;
+        // Make sure it is OpCode::kNone so it will be properly disassembled
+        node->op = Op{}
+    } else {
+        // It is completely swallowed by the primary node
+        delete secondary;
     }
-    secondary->ref_by = secondary->last_ref_by = nullptr;
-    delete secondary;
     return primary;
 }
 
@@ -344,12 +387,12 @@ void DisasmMap::disasmQuickPeek(const DataView &code, const Settings &s)
         // Spread across the size and merge if an intersection encountered
         const size_t address = node->address;
         const auto size = node->size;
-        for (size_t i = 0; i < size; i++) {
-            auto *const ptr = _map[node->address + i];
+        for (size_t i = address; i < address + size; i++) {
+            auto *const ptr = _map[i];
             if (ptr != nullptr && ptr != node) {
-                node = mergeNodes(node, ptr);
+                mergeNodeOverlappingSpace(node, ptr);
             }
-            _map[address + i] = node;
+            _map[i] = node;
         }
         at += node->size;
         // NOTE: There is not much information about a reference passed further,