Impl deep graph walk and distinct jump following option

This commit changes the default behavior which technically breaks
backwards compatibility. It would be a concern if somebody else besides
me is using this tool, but I doubt it would be the case, LOL.

1 files changed, 101 insertions, 63 deletions

diff --git a/main.cpp b/main.cpp
index 01e1895..a6f73b3 100644
--- a/main.cpp
+++ b/main.cpp
@@ -28,12 +28,11 @@ class DisasmMap {
     const DisasmMapType _type;
     DisasmNode *_map[kDisasmMapSizeElements]{};
     constexpr DisasmNode *findNodeByAddress(uint32_t address) const;
-    DisasmNode *insertTracedNode(uint32_t address, TracedNodeType);
-    void insertReferencedBy(
+    DisasmNode &insertNode(uint32_t address, NodeType);
+    DisasmNode &insertReferencedBy(
             const uint32_t by_addr,
             const uint32_t ref_addr,
-            const TracedNodeType type,
-            const DataView &code,
+            const NodeType type,
             const ReferenceType ref_type);
     constexpr bool canBeAllocated(const DisasmNode& node) const;
 public:
@@ -41,14 +40,12 @@ public:
     {
         return findNodeByAddress(address);
     };
-    // Returns true if node inserted, false if node already exist and has not
-    // been changed
-    bool InsertTracedNode(uint32_t address, TracedNodeType type)
+    void InsertNode(uint32_t address, NodeType type)
     {
         assert(_type == DisasmMapType::kTraced);
-        return nullptr != insertTracedNode(address, type);
+        insertNode(address, type);
     }
-    void Disasm(const DataView &code, const Settings &);
+    void Disasm(const DataView &code, const Settings &, size_t from=0, bool nested=false);
     DisasmMap(DisasmMapType type): _type(type) {}
     ~DisasmMap();
 };
@@ -65,45 +62,35 @@ static constexpr uint32_t AlignInstructionAddress(const uint32_t address)
     return address & ~1UL;
 }
 
-DisasmNode *DisasmMap::insertTracedNode(const uint32_t address, const TracedNodeType type)
+DisasmNode &DisasmMap::insertNode(const uint32_t address, const NodeType type)
 {
     auto *node = findNodeByAddress(address);
     if (node) {
         // Instruction nodes take precedence over data nodes. If a node that
         // was previously accessed only as data now turns out to be an
         // instruction, then it must become an instruction node.
-        if (type == TracedNodeType::kInstruction && node->type != TracedNodeType::kInstruction) {
-            *const_cast<TracedNodeType*>(&node->type) = type;
+        if (IsInstruction(type) && !IsInstruction(node->type)) {
+            *const_cast<NodeType*>(&node->type) = type;
             // Make sure it is OpCode::kNone so it will be properly disassembled
             node->op = Op{};
         }
-        return node;
+        return *node;
     }
     node = new DisasmNode(DisasmNode{type, AlignInstructionAddress(address)});
     assert(node);
     _map[address / kInstructionSizeStepBytes] = node;
-    return node;
+    return *node;
 }
 
-void DisasmMap::insertReferencedBy(
+DisasmNode &DisasmMap::insertReferencedBy(
         const uint32_t by_addr,
         const uint32_t ref_addr,
-        const TracedNodeType type,
-        const DataView &code,
+        const NodeType type,
         const ReferenceType ref_type)
 {
-    auto *const ref_node = insertTracedNode(ref_addr, type);
-    const auto size = ref_node->Disasm(code);
-    assert(size >= kInstructionSizeStepBytes);
-    if (canBeAllocated(*ref_node)) {
-        // Spread across the size
-        for (size_t o = kInstructionSizeStepBytes; o < size; o++) {
-            _map[(ref_node->address + o) / kInstructionSizeStepBytes] = ref_node;
-        }
-    } else {
-        ref_node->DisasmAsRaw(code);
-    }
-    ref_node->AddReferencedBy(by_addr, ref_type);
+    auto &ref_node = insertNode(ref_addr, type);
+    ref_node.AddReferencedBy(by_addr, ref_type);
+    return ref_node;
 }
 
 constexpr bool DisasmMap::canBeAllocated(const DisasmNode& node) const
@@ -141,31 +128,55 @@ static constexpr ReferenceType ReferenceTypeFromRefKindMask2(const RefKindMask r
                 : ReferenceType::kBranch;
 }
 
-void DisasmMap::Disasm(const DataView &code, const Settings &s)
+static constexpr bool IsNextLikelyAnInstruction(const Op &op)
+{
+    return (op.opcode != OpCode::kNone &&
+            op.opcode != OpCode::kRaw &&
+            !IsBRA(op) &&
+            op.opcode != OpCode::kJMP &&
+            op.opcode != OpCode::kRTS &&
+            op.opcode != OpCode::kRTE &&
+            op.opcode != OpCode::kSTOP);
+}
+
+void DisasmMap::Disasm(
+        const DataView &code, const Settings &s, size_t at, bool nested)
 {
-    DisasmNode *node;
-    for (size_t i = 0; i < Min(kRomSizeBytes, code.size);) {
+    // Some of logic of this function is covered by integration tests in
+    // `test_walk_and_follow_jumps.bash`.
+    bool inside_code_span = nested;
+    while (at < Min(kRomSizeBytes, code.size)) {
+        DisasmNode *node;
         if (_type == DisasmMapType::kTraced) {
-            node = _map[i / kInstructionSizeStepBytes];
+            node = _map[at / kInstructionSizeStepBytes];
             if (!node) {
-                i += kInstructionSizeStepBytes;
-                continue;
+                if (inside_code_span) {
+                    node = &insertNode(at, NodeType::kTracedInstruction);
+                } else {
+                    at += kInstructionSizeStepBytes;
+                    continue;
+                }
             }
         } else {
-            node = insertTracedNode(i, TracedNodeType::kInstruction);
+            node = &insertNode(at, NodeType::kTracedInstruction);
         }
-        const auto size = node->Disasm(code);
-        assert(size >= kInstructionSizeStepBytes);
-        if (canBeAllocated(*node)) {
-            // Spread across the size
-            for (size_t o = kInstructionSizeStepBytes; o < size; o++) {
-                _map[(node->address + o) / kInstructionSizeStepBytes] = node;
+        if (node->op.opcode == OpCode::kNone || inside_code_span) {
+            const auto size = node->Disasm(code);
+            assert(size >= kInstructionSizeStepBytes);
+            if (canBeAllocated(*node)) {
+                // Spread across the size
+                for (size_t o = kInstructionSizeStepBytes; o < size; o++) {
+                    _map[(node->address + o) / kInstructionSizeStepBytes] = node;
+                }
+            } else {
+                node->DisasmAsRaw(code);
             }
-        } else {
-            node->DisasmAsRaw(code);
         }
-        // FIXME implement deep graph walk for DisasmMapType::kTraced case.
-        //
+        inside_code_span = s.walk && IsNextLikelyAnInstruction(node->op);
+        if (nested && !inside_code_span) {
+            return;
+        }
+        at += node->size;
         // NOTE: There is not much information about a reference passed further,
         // so just don't add a reference of immediate if s.imm_labels is false
         // enabled.
@@ -174,20 +185,35 @@ void DisasmMap::Disasm(const DataView &code, const Settings &s)
             : (node->ref_kinds & kRef1Mask);
         const bool has_code_ref1 = node->ref1_addr < code.size && has_ref1;
         if (has_code_ref1) {
-            const TracedNodeType type = (node->ref_kinds & (kRef1ReadMask | kRef1WriteMask))
-                ? TracedNodeType::kData : TracedNodeType::kInstruction;
+            const NodeType type = (node->ref_kinds & (kRef1ReadMask | kRef1WriteMask))
+                ? NodeType::kData : NodeType::kRefInstruction;
             const auto ref_type = ReferenceTypeFromRefKindMask1(node->ref_kinds);
-            insertReferencedBy(node->address, node->ref1_addr, type, code, ref_type);
+            auto &ref_node = insertReferencedBy(
+                    node->address, node->ref1_addr, type, ref_type);
+            if (ref_node.op.opcode == OpCode::kNone) {
+                if (s.follow_jumps) {
+                    Disasm(code, s, ref_node.address, true);
+                } else {
+                    ref_node.DisasmAsRaw(code);
+                }
+            }
         }
         const bool has_ref2 = (node->ref_kinds & kRef2Mask);
         const bool has_code_ref2 = (has_ref2 && node->ref2_addr < code.size);
         if (has_code_ref2) {
-            const TracedNodeType type = (node->ref_kinds & (kRef2ReadMask | kRef2WriteMask))
-                ? TracedNodeType::kData : TracedNodeType::kInstruction;
+            const NodeType type = (node->ref_kinds & (kRef2ReadMask | kRef2WriteMask))
+                ? NodeType::kData : NodeType::kRefInstruction;
             const auto ref_type = ReferenceTypeFromRefKindMask2(node->ref_kinds);
-            insertReferencedBy(node->address, node->ref2_addr, type, code, ref_type);
+            auto &ref_node = insertReferencedBy(
+                    node->address, node->ref2_addr, type, ref_type);
+            if (ref_node.op.opcode == OpCode::kNone) {
+                if (s.follow_jumps) {
+                    Disasm(code, s, ref_node.address, true);
+                } else {
+                    ref_node.DisasmAsRaw(code);
+                }
+            }
         }
-        i += node->size;
     }
 }
 
@@ -376,7 +402,6 @@ static void RenderNodeDisassembly(
             arg.SNPrint(arg_str, kArgsBufferSize);
             fprintf(output, ", %s", arg_str);
         }
-        fprintf(output, "\n");
     } else {
         const bool with_ref = node.ref_kinds && s.labels && (s.abs_labels || s.rel_labels);
         const auto *ref1 = (node.ref_kinds & kRef1Mask)
@@ -493,7 +518,7 @@ static void ParseTraceData(DisasmMap &disasm_map, const DataView &trace_data)
                 exit(1);
             } else {
                 // Valid value
-                disasm_map.InsertTracedNode(address, TracedNodeType::kInstruction);
+                disasm_map.InsertNode(address, NodeType::kTracedInstruction);
             }
             if (startptr != endptr) {
                 i += endptr - startptr - 1;
@@ -617,6 +642,8 @@ static bool ApplyFeature(Settings& s, const char *feature_arg)
         { &Settings::xrefs_from, "xrefs-from" },
         { &Settings::xrefs_to, "xrefs-to" },
         { &Settings::imm_hex, "imm-hex" },
+        { &Settings::follow_jumps, "follow-jumps" },
+        { &Settings::walk, "walk" },
     };
     constexpr size_t sizeof_no_prefix = (sizeof "no-") - 1;
     const bool disable = FeatureStringHasPrefixNo(feature_arg);
@@ -670,6 +697,9 @@ static void PrintUsage(FILE *s, const char *argv0)
     "  xrefs-from            Print xrefs comments above all places that have xrefs.\n"
     "  xrefs-to              Print xrefs comments after all branch instructions.\n"
     "  imm-hex               Print all immediate values as hexadecimal numbers.\n"
+    "  follow-jumps          Follow jumps to statically known locations.\n"
+    "  walk                  Try best to detect further instructions following known\n"
+    "                        traced locations without overcommitting.\n"
     , argv0);
 }
 
@@ -753,7 +783,11 @@ int main(int, char* argv[])
     FILE *output_stream = stdout;
     FILE *trace_stream = nullptr;
     if (input_file_name) {
-        input_stream = fopen(input_file_name, "r");
+        if (0 == strcmp(input_file_name, "-")) {
+            input_stream = stdin;
+        } else {
+            input_stream = fopen(input_file_name, "r");
+        }
         if (input_stream == nullptr) {
             const int err = errno;
             fprintf(stderr, "main: fopen(\"%s\", \"r\"): Error (%d): \"%s\"\n", input_file_name, err, strerror(err));
@@ -774,7 +808,15 @@ int main(int, char* argv[])
         }
     }
     if (trace_file_name) {
-        trace_stream = fopen(trace_file_name, "r");
+        if (0 == strcmp(trace_file_name, "-")) {
+            if (input_stream == stdin) {
+                fprintf(stderr, "error: trace stream and input stream cannot be both stdin\n");
+                return EXIT_FAILURE;
+            }
+            trace_stream = stdin;
+        } else {
+            trace_stream = fopen(trace_file_name, "r");
+        }
         if (trace_stream == nullptr) {
             const int err = errno;
             fprintf(stderr, "main: fopen(\"%s\", \"r\"): Error (%d): \"%s\"\n", trace_file_name, err, strerror(err));
@@ -788,11 +830,7 @@ int main(int, char* argv[])
     if (trace_stream != nullptr) {
         fclose(trace_stream);
     }
-    if (output_stream != stdout) {
-        fclose(output_stream);
-    }
-    if (input_stream != stdin) {
-        fclose(input_stream);
-    }
+    fclose(output_stream);
+    fclose(input_stream);
     return ret;
 }