Optimize override_mode in assembler, add more test cases

3 files changed, 40 insertions, 7 deletions

diff --git a/atk16_asm/asm_pass3.py b/atk16_asm/asm_pass3.py
index a3f9537..7156518 100644
--- a/atk16_asm/asm_pass3.py
+++ b/atk16_asm/asm_pass3.py
@@ -77,20 +77,26 @@ def pass_3(result2: Result2) -> Result3:
         symbols[args[0]] = eval_expr(symbols, args[1])
         continue
       case _:
-        if override_mode:
-          # TODO this is a bit inefficient, but works
-          result_lines = list(filter(lambda l: l.address != address, result_lines))
-
         dbg_original_text = " ".join(line.original_parts)
         dbg_expanded_text = " ".join(line.parts)
-        save_dbg_source(address, line.src_file, line.line_num, dbg_expanded_text, dbg_original_text)
-        result_lines.append(Result3Line(
+
+        res_line = Result3Line(
           line_num=line.line_num,
           src_file=line.src_file,
           parts=line.parts,
           address=address,
           original_parts=line.original_parts,
-        ))
+        )
+
+        if override_mode:
+          for i in range(len(result_lines)):
+            if result_lines[i].address == address:
+              result_lines[i] = res_line # overwrite the existing line
+              break
+        else:
+          result_lines.append(res_line)
+
+        save_dbg_source(address, line.src_file, line.line_num, dbg_expanded_text, dbg_original_text)
         address += 1
 
   result_lines.sort(key=lambda l: l.address)
diff --git a/test/e2e/std_bump_alloc/run_bump_alloc_overlap.atk16 b/test/e2e/std_bump_alloc/run_bump_alloc_overlap.atk16
new file mode 100644
index 0000000..8dcc6fa
--- /dev/null
+++ b/test/e2e/std_bump_alloc/run_bump_alloc_overlap.atk16
@@ -0,0 +1,14 @@
+@include %bootstrap
+@include %std_bump_alloc
+
+@label n
+  0x8000
+
+@label main
+  calli bump_reset
+
+  ldi n RA          ; allocate too many words
+  ldr RA RA
+  calli bump_alloc
+
+  hlt
diff --git a/test/e2e/std_bump_alloc/test_bump_alloc.py b/test/e2e/std_bump_alloc/test_bump_alloc.py
index 717a122..b2274f8 100644
--- a/test/e2e/std_bump_alloc/test_bump_alloc.py
+++ b/test/e2e/std_bump_alloc/test_bump_alloc.py
@@ -20,3 +20,16 @@ def test_bump_alloc():
 
   returned_value = machine.rg.value
   assert returned_value == heap_next + 1
+
+def test_bump_alloc_overlap():
+  machine = assemble_and_run_until_halted(
+    "test/e2e/std_bump_alloc/run_bump_alloc_overlap.atk16"
+  )
+
+  heap_next_p = 0xE7EF
+  heap_next = machine.mem_read(heap_next_p)
+  assert heap_next == heap_next_p - 1 # alloc failed because of overlap
+                                      # so pointer should not have moved
+
+  returned_value = machine.rg.value
+  assert returned_value == 0          # returns 0 on failure