Analyse control flow in runPreservingCurrentBreakpoint (#313)

Naive `wasmPc+1` implementation is buggy and leads to bad user experience, where resuming from a breakpoint completely disables breakpoints set on control flow instructions.

`DebuggerTests` is now significantly expanded with new tests added that fail on `main` to show the problem.

```
  ┌────────┬──────────────────────────────────────────────┬───────────────────────────────────┐
  │ Status │                     Test                     │          Failure on main          │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ stepFollowsBrIf                              │ offset 23 instead of 8..<19       │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ stepFollowsBr                                │ offset 26 instead of 8..<22       │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ stepFollowsBrTable                           │ offset 25 instead of 8..<21       │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ stepFollowsReturnCall                        │ runs to completion                │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ runPreservingBrIfBreakpoint                  │ runs to completion                │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ runPreservingBrTableBreakpoint               │ runs to completion                │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ runPreservingRecursiveCallBreakpoint         │ noInstructionMappingAvailable(63) │
  ├────────┼──────────────────────────────────────────────┼───────────────────────────────────┤
  │ FAIL   │ runPreservingRecursiveCallIndirectBreakpoint │ noInstructionMappingAvailable(80) │
  └────────┴──────────────────────────────────────────────┴───────────────────────────────────┘
```

These are all passing now on this branch.

Author: MaxDesiatov

Sources/WasmKit/Execution/Debugger.swift

@@ -60,6 +60,12 @@
         /// was not compiled yet in lazy compilation mode).
         private let functionAddresses: [(address: Int, instanceFunctionIndex: Int)]
 
+        /// Reverse map from a head code slot to its opcode ID, used for resolving
+        /// which control-flow instruction is at a breakpoint site.
+        /// For token threading the head slot is the opcode ID itself; for direct
+        /// threading it is a function pointer that we map back.
+        private let headSlotToOpcodeID: [CodeSlot: OpcodeID]
+
         /// Initializes a new debugger state instance.
         /// - Parameters:
         ///   - module: Wasm module to instantiate.
@@ -92,6 +98,11 @@
             )
             self.threadingModel = store.engine.configuration.threadingModel
             self.endOfExecution = Instruction.endOfExecution.headSlot(threadingModel: threadingModel)
+
+            self.headSlotToOpcodeID = Instruction.buildControlHeadSlotMap(
+                threadingModel: self.threadingModel
+            )
+
             self.state = .instantiated
         }
 
@@ -251,8 +262,7 @@
                 return
             }
 
-            // TODO: analyze actual instruction branching to set the breakpoint correctly.
-            try self.enableBreakpoint(address: breakpoint.wasmPc + 1)
+            try self.setNextInstructionBreakpoints(breakpoint: breakpoint)
             try self.run()
         }
 
@@ -264,8 +274,7 @@
                 return
             }
 
-            // TODO: analyze actual instruction branching to set the breakpoint correctly.
-            try self.enableBreakpoint(address: breakpoint.wasmPc + 1)
+            try self.setNextInstructionBreakpoints(breakpoint: breakpoint)
             try self.run()
             try self.enableBreakpoint(address: breakpoint.wasmPc)
             try self.run()
@@ -346,9 +355,138 @@
             return result
         }
 
+        /// Analyzes the control-flow instruction at the given breakpoint and sets breakpoints
+        /// at all possible next instruction locations.
+        private mutating func setNextInstructionBreakpoints(breakpoint: BreakpointState) throws {
+            let savedHead = self.breakpoints[breakpoint.wasmPc]!
+            let operandPc = breakpoint.iseq.pc.advanced(by: 1)
+            let sp = breakpoint.iseq.sp
+
+            if let opcodeID = headSlotToOpcodeID[savedHead],
+                let targets = Instruction.predictNextPcs(
+                    opcodeID: opcodeID, operandPc: operandPc, sp: sp,
+                    predictor: &self
+                )
+            {
+                // Control instruction with predicted targets.
+                // Empty targets means terminal (unreachable, endOfExecution) — no breakpoints to set.
+                for pc in targets {
+                    if let wasmAddr = self.instance.handle.instructionMapping.findWasm(forIseqAddress: pc) {
+                        try self.enableBreakpoint(address: wasmAddr)
+                    }
+                }
+                return
+            }
+
+            // Non-control instruction: fall back to next Wasm address
+            try self.enableBreakpoint(address: breakpoint.wasmPc + 1)
+        }
+
         deinit {
             self.valueStack.deallocate()
         }
     }
 
+    extension Debugger: NextInstructionPredictor {
+        mutating func predictNext_br(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let offset = Instruction.BrOperand.load(from: &pc)
+            return [pc.advanced(by: Int(offset))]
+        }
+
+        mutating func predictNext_brIf(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.BrIfOperand.load(from: &pc)
+            // Both fall-through (pc after operand) and branch target are possible
+            return [pc, pc.advanced(by: Int(op.offset))]
+        }
+
+        mutating func predictNext_brIfNot(operandPc: Pc, sp: Sp) -> [Pc] {
+            predictNext_brIf(operandPc: operandPc, sp: sp)
+        }
+
+        mutating func predictNext_brTable(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.BrTableOperand.load(from: &pc)
+            return (0..<Int(op.count)).map { pc.advanced(by: Int(op.baseAddress[$0].offset)) }
+        }
+
+        mutating func predictNext_call(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.CallOperand.load(from: &pc)
+            let (iseq, _, _) = op.callee.assumeCompiled()
+            return [iseq.instructions.baseAddress!]
+        }
+
+        mutating func predictNext_compilingCall(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.CallOperand.load(from: &pc)
+            _ = try? op.callee.wasm.ensureCompiled(store: StoreRef(self.store))
+            let (iseq, _, _) = op.callee.assumeCompiled()
+            return [iseq.instructions.baseAddress!]
+        }
+
+        mutating func predictNext_internalCall(operandPc: Pc, sp: Sp) -> [Pc] {
+            predictNext_call(operandPc: operandPc, sp: sp)
+        }
+
+        mutating func predictNext__return(operandPc: Pc, sp: Sp) -> [Pc] {
+            // returnPC is stored at sp[-2]
+            let returnPc = Pc(bitPattern: UInt(sp.advanced(by: -2).pointee))
+            return returnPc.map { [$0] } ?? []
+        }
+
+        /// Resolves a callee function from a table and returns its iseq base address.
+        /// Returns `nil` if the callee is a host function or the resolution fails.
+        private mutating func resolveIndirectCallee(
+            tableIndex: UInt32, index: VReg, sp: Sp
+        ) -> Pc? {
+            let callerInstance = self.instance.handle
+            let table = callerInstance.tables[Int(tableIndex)]
+            let value = sp[index].asAddressOffset(table.limits.isMemory64)
+            let elementIndex = Int(value)
+            guard elementIndex < table.elements.count,
+                case .function(let rawBitPattern?) = table.elements[elementIndex]
+            else { return nil }
+            let function = InternalFunction(bitPattern: rawBitPattern)
+            guard function.isWasm else { return nil }
+            _ = try? function.wasm.ensureCompiled(store: StoreRef(self.store))
+            let (iseq, _, _) = function.assumeCompiled()
+            return iseq.instructions.baseAddress!
+        }
+
+        mutating func predictNext_callIndirect(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.CallIndirectOperand.load(from: &pc)
+            guard let target = resolveIndirectCallee(tableIndex: op.tableIndex, index: op.index, sp: sp) else {
+                return []
+            }
+            return [target]
+        }
+
+        mutating func predictNext_returnCall(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.ReturnCallOperand.load(from: &pc)
+            let callee = op.callee
+            guard callee.isWasm else { return [] }
+            _ = try? callee.wasm.ensureCompiled(store: StoreRef(self.store))
+            let (iseq, _, _) = callee.assumeCompiled()
+            return [iseq.instructions.baseAddress!]
+        }
+
+        mutating func predictNext_returnCallIndirect(operandPc: Pc, sp: Sp) -> [Pc] {
+            var pc = operandPc
+            let op = Instruction.ReturnCallIndirectOperand.load(from: &pc)
+            guard let target = resolveIndirectCallee(tableIndex: op.tableIndex, index: op.index, sp: sp) else {
+                return []
+            }
+            return [target]
+        }
+
+        // Terminal instructions — no successor exists
+        mutating func predictNext_unreachable(operandPc: Pc, sp: Sp) -> [Pc] { [] }
+        mutating func predictNext_endOfExecution(operandPc: Pc, sp: Sp) -> [Pc] { [] }
+        mutating func predictNext_breakpoint(operandPc: Pc, sp: Sp) -> [Pc] { [] }
+    }
+
 #endif

Sources/WasmKit/Execution/Instructions/Instruction.swift

@@ -2258,3 +2258,123 @@ extension Instruction {
     }
 }
 #endif // EngineStats
+
+#if WasmDebuggingSupport
+
+/// A protocol for predicting the next instruction(s) that will execute after a control-flow instruction.
+///
+/// Each `isControl` instruction in VMSpec gets a dedicated method in this protocol.
+/// Adding a new control instruction automatically adds a new protocol requirement,
+/// so any conforming type will fail to compile until it implements the new prediction.
+protocol NextInstructionPredictor: ~Copyable {
+    mutating func predictNext_call(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_compilingCall(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_internalCall(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_callIndirect(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_returnCall(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_returnCallIndirect(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_unreachable(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_br(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_brIf(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_brIfNot(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_brTable(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext__return(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_endOfExecution(operandPc: Pc, sp: Sp) -> [Pc]
+    mutating func predictNext_breakpoint(operandPc: Pc, sp: Sp) -> [Pc]
+}
+
+extension Instruction {
+    /// Dispatches to the appropriate `NextInstructionPredictor` method based on the opcode ID.
+    /// - Returns: The predicted next Pc(s), or `nil` if the opcode is not a control instruction.
+    static func predictNextPcs(
+        opcodeID: OpcodeID, operandPc: Pc, sp: Sp,
+        predictor: inout some NextInstructionPredictor & ~Copyable
+    ) -> [Pc]? {
+        switch opcodeID {
+        case 3: return predictor.predictNext_call(operandPc: operandPc, sp: sp)
+        case 4: return predictor.predictNext_compilingCall(operandPc: operandPc, sp: sp)
+        case 5: return predictor.predictNext_internalCall(operandPc: operandPc, sp: sp)
+        case 6: return predictor.predictNext_callIndirect(operandPc: operandPc, sp: sp)
+        case 8: return predictor.predictNext_returnCall(operandPc: operandPc, sp: sp)
+        case 9: return predictor.predictNext_returnCallIndirect(operandPc: operandPc, sp: sp)
+        case 10: return predictor.predictNext_unreachable(operandPc: operandPc, sp: sp)
+        case 12: return predictor.predictNext_br(operandPc: operandPc, sp: sp)
+        case 13: return predictor.predictNext_brIf(operandPc: operandPc, sp: sp)
+        case 14: return predictor.predictNext_brIfNot(operandPc: operandPc, sp: sp)
+        case 15: return predictor.predictNext_brTable(operandPc: operandPc, sp: sp)
+        case 16: return predictor.predictNext__return(operandPc: operandPc, sp: sp)
+        case 17: return predictor.predictNext_endOfExecution(operandPc: operandPc, sp: sp)
+        case 202: return predictor.predictNext_breakpoint(operandPc: operandPc, sp: sp)
+        default: return nil
+        }
+    }
+
+    /// Builds a map from head code slot to opcode ID for all control-flow instructions.
+    ///
+    /// This is generated so that adding a new `isControl` instruction automatically
+    /// includes it in the map without manual updates.
+    static func buildControlHeadSlotMap(
+        threadingModel: EngineConfiguration.ThreadingModel
+    ) -> [CodeSlot: OpcodeID] {
+        var map = [CodeSlot: OpcodeID]()
+            do {
+                let inst = Instruction.call(.init(rawCallee: UInt64(0), spAddend: VReg(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.compilingCall(.init(rawCallee: UInt64(0), spAddend: VReg(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.internalCall(.init(rawCallee: UInt64(0), spAddend: VReg(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.callIndirect(.init(tableIndex: UInt32(0), rawType: UInt32(0), index: VReg(0), spAddend: VReg(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.returnCall(.init(rawCallee: UInt64(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.returnCallIndirect(.init(tableIndex: UInt32(0), rawType: UInt32(0), index: VReg(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.unreachable
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.br(Instruction.BrOperand(0))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.brIf(.init(condition: LVReg(0), offset: Int32(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.brIfNot(.init(condition: LVReg(0), offset: Int32(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.brTable(.init(rawBaseAddress: UInt64(0), count: UInt16(0), index: VReg(0)))
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction._return
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.endOfExecution
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+            do {
+                let inst = Instruction.breakpoint
+                map[inst.headSlot(threadingModel: threadingModel)] = inst.opcodeID
+            }
+        return map
+    }
+}
+
+#endif // WasmDebuggingSupport