3 files changed, 28 insertions, 8 deletions
diff --git a/src/utils/assembler.spec.ts b/src/utils/assembler.spec.ts
index f0daf64..836f388 100644
--- a/src/utils/assembler.spec.ts
+++ b/src/utils/assembler.spec.ts
@@ -14,6 +14,7 @@ describe('AVR assembler', () => {
       bytes: bytes('0b0d'),
       errors: [],
       lines: [{ byteOffset: 0, bytes: '0d0b', line: 1, text: 'ADD r16, r11' }],
+      labels: {},
     });
   });
 
@@ -36,6 +37,7 @@ describe('AVR assembler', () => {
       bytes: new Uint8Array(0),
       errors: [],
       lines: [],
+      labels: {},
     });
   });
 
@@ -44,6 +46,7 @@ describe('AVR assembler', () => {
       bytes: new Uint8Array(0),
       errors: ['Line 0: Rd out of range: 16<>31'],
       lines: [],
+      labels: {},
     });
   });
 
diff --git a/src/utils/assembler.ts b/src/utils/assembler.ts
index 4823937..278e944 100644
--- a/src/utils/assembler.ts
+++ b/src/utils/assembler.ts
@@ -29,7 +29,7 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-interface LabelTable {
+export interface LabelTable {
   [key: string]: number;
 }
 
@@ -938,7 +938,7 @@ function passTwo(lineTable: LineTablePass1[], labels: LabelTable) {
     }
   }
 
-  return { errors: errorTable, bytes: resultTable, lines: lineTable as LineTable[] };
+  return { errors: errorTable, bytes: resultTable, lines: lineTable as LineTable[], labels };
 }
 
 /**
@@ -951,6 +951,7 @@ export function assemble(input: string) {
       bytes: new Uint8Array(0),
       errors: mid.errors,
       lines: [],
+      labels: {},
     };
   }
   return passTwo(mid.lines, mid.labels);
diff --git a/src/utils/test-utils.ts b/src/utils/test-utils.ts
index 4b2efff..e620748 100644
--- a/src/utils/test-utils.ts
+++ b/src/utils/test-utils.ts
@@ -5,38 +5,54 @@ import { avrInstruction } from '../cpu/instruction';
 const BREAK_OPCODE = 0x9598;
 
 export function asmProgram(source: string) {
-  const { bytes, errors, lines } = assemble(source);
+  const { bytes, errors, lines, labels } = assemble(source);
   if (errors.length) {
     throw new Error('Assembly failed: ' + errors);
   }
-  return { program: new Uint16Array(bytes.buffer), lines, instructionCount: lines.length };
+  return { program: new Uint16Array(bytes.buffer), lines, instructionCount: lines.length, labels };
 }
 
+const defaultOnBreak = () => {
+  throw new Error('BREAK instruction encountered');
+};
+
 export class TestProgramRunner {
-  constructor(private readonly cpu: CPU, private readonly onBreak?: (cpu: CPU) => void) {}
+  constructor(
+    private readonly cpu: CPU,
+    private readonly onBreak: (cpu: CPU) => void = defaultOnBreak
+  ) {}
 
   runInstructions(count: number) {
     const { cpu, onBreak } = this;
     for (let i = 0; i < count; i++) {
       if (cpu.progMem[cpu.pc] === BREAK_OPCODE) {
         onBreak?.(cpu);
-        throw new Error('BREAK instruction encountered');
       }
       avrInstruction(cpu);
       cpu.tick();
     }
   }
 
-  runToBreak(maxIterations = 5000) {
+  runUntil(predicate: (cpu: CPU) => boolean, maxIterations = 5000) {
     const { cpu, onBreak } = this;
     for (let i = 0; i < maxIterations; i++) {
       if (cpu.progMem[cpu.pc] === BREAK_OPCODE) {
         onBreak?.(cpu);
+      }
+      if (predicate(cpu)) {
         return;
       }
       avrInstruction(cpu);
       cpu.tick();
     }
-    throw new Error('Program ran for too long without a BREAK instruction');
+    throw new Error('Test program ran for too long, check your predicate');
+  }
+
+  runToBreak() {
+    this.runUntil((cpu) => cpu.progMem[cpu.pc] === BREAK_OPCODE);
+  }
+
+  runToAddress(byteAddr: number) {
+    this.runUntil((cpu) => cpu.pc * 2 === byteAddr);
   }
 }