From c0380bd19813951d054f1e90fd0df015b4acf04c Mon Sep 17 00:00:00 2001
From: KallDrexx <me@mshapiro.net>
Date: Sat, 11 Oct 2025 23:11:37 -0400
Subject: [PATCH] Allow for sub address instructions.

When a wrap around scenario is detected during single function tracing,
if the address prior to the "entry point" is a single byte, then we do
not have any space to add the required jump call.

This fixes that by adding the concept of sub address instructions. This
allows adding instructions at runtime that get sorted correctly against
the real instructions from the ROM.

This not only solves the wrapping issue, but also allows for adding hooks
at runtime.
---
 .../Decompilation/DecompiledFunction.cs             |  6 ++++--
 .../Decompilation/FunctionDecompiler.cs             | 13 +++++++------
 NESDecompiler.Core/Disassembly/Disassembler.cs      |  8 ++++++++
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/NESDecompiler.Core/Decompilation/DecompiledFunction.cs b/NESDecompiler.Core/Decompilation/DecompiledFunction.cs
index a8d6eee..6192e7a 100644
--- a/NESDecompiler.Core/Decompilation/DecompiledFunction.cs
+++ b/NESDecompiler.Core/Decompilation/DecompiledFunction.cs
@@ -39,11 +39,13 @@ public DecompiledFunction(
         // the first instruction.
         var initialInstructions = instructions
             .Where(x => x.CPUAddress >= address)
-            .OrderBy(x => x.CPUAddress);
+            .OrderBy(x => x.CPUAddress)
+            .ThenBy(x => x.SubAddressOrder);
 
         var trailingInstructions = instructions
             .Where(x => x.CPUAddress < address)
-            .OrderBy(x => x.CPUAddress);
+            .OrderBy(x => x.CPUAddress)
+            .ThenBy(x => x.SubAddressOrder); // real instructions before virtual ones
 
         OrderedInstructions = initialInstructions.Concat(trailingInstructions).ToArray();
     }
diff --git a/NESDecompiler.Core/Decompilation/FunctionDecompiler.cs b/NESDecompiler.Core/Decompilation/FunctionDecompiler.cs
index f26fc9f..9208de6 100644
--- a/NESDecompiler.Core/Decompilation/FunctionDecompiler.cs
+++ b/NESDecompiler.Core/Decompilation/FunctionDecompiler.cs
@@ -27,10 +27,10 @@ public static DecompiledFunction Decompile(ushort functionAddress, IReadOnlyList
                     // the function entrance. This usually happens when there is a jump/branch to right before the
                     // entrypoint, usually due to decompiling in the middle of a loop. To fix this, we need to add
                     // a jump back to the function entrypoint
-                    if (functionAddress == 0x00 || seenInstructions.Contains((ushort)(functionAddress - 1)))
+                    if (functionAddress == 0x00)
                     {
-                        var message = $"Function 0x{functionAddress:X4} wraps around, but there's not enough " +
-                                      $"space to add a jump back to the entrypoint";
+                        const string message = "Wrap around instruction detected for a function at 0000, but that " +
+                                               "doesn't make sense";
 
                         throw new InvalidOperationException(message);
                     }
@@ -44,6 +44,9 @@ public static DecompiledFunction Decompile(ushort functionAddress, IReadOnlyList
                         CPUAddress = (ushort)(nextAddress - 1),
                         Bytes = [0x4C, (byte)addressLow, (byte)addressHigh],
                         TargetAddress = functionAddress,
+
+                        // Make sure they appear after any instruction that already occupies that address
+                        SubAddressOrder = 1,
                     };
 
                     instructions.Add(jumpInstruction);
@@ -109,9 +112,7 @@ private static DisassembledInstruction GetNextInstruction(ushort address, IReadO
         var info = InstructionSet.GetInstruction(bytes[0]);
         if (!info.IsValid)
         {
-            var message = $"Attempted to get instruction at address 0x{address:X4}, but byte 0x{bytes[0]:X4} " +
-                          $"is not a valid/known opcode";
-
+            var message = $"Opcode 0x{bytes[0]:X2} at address 0x{address:X4} is not a known instruction";
             throw new InvalidOperationException(message);
         }
 
diff --git a/NESDecompiler.Core/Disassembly/Disassembler.cs b/NESDecompiler.Core/Disassembly/Disassembler.cs
index ee98afb..f1efd26 100644
--- a/NESDecompiler.Core/Disassembly/Disassembler.cs
+++ b/NESDecompiler.Core/Disassembly/Disassembler.cs
@@ -72,6 +72,14 @@ public class DisassembledInstruction
         /// </summary>
         public bool IsJump => Info.Mnemonic == "JMP" || Info.Mnemonic == "JSR";
 
+        /// <summary>
+        /// Determines the order of this instruction within a single address space. This is mostly
+        /// needed in the cases that additional instructions are needed to be added in the same
+        /// address location at runtime. Can be used to add runtime hooks or to work around
+        /// decompilation issues. Should be 0 for all native instructions from a ROM.
+        /// </summary>
+        public byte SubAddressOrder { get; set; }
+
         /// <summary>
         /// Returns a string representation of this instruction
         /// </summary>