From c8ae97d8f62130134fc5f990ff62306e3bd37728 Mon Sep 17 00:00:00 2001
From: Daan Vanden Bosch <daan.v.d.bosch@gmail.com>
Date: Wed, 7 Aug 2019 13:57:59 +0200
Subject: [PATCH] Fixed bug in assembler.

---
 src/scripting/assembly.ts         |  9 +++--
 src/scripting/disassembly.test.ts | 65 +++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 src/scripting/disassembly.test.ts

diff --git a/src/scripting/assembly.ts b/src/scripting/assembly.ts
index 750c9bb7..13a9260f 100644
--- a/src/scripting/assembly.ts
+++ b/src/scripting/assembly.ts
@@ -57,7 +57,8 @@ class Assembler {
     private errors!: AssemblyError[];
     // Encountered labels.
     private labels!: Set<number>;
-    private section: SegmentType = SegmentType.Instructions;
+    private section!: SegmentType;
+    private first_section_marker = true;
 
     constructor(private assembly: string[], private manual_stack: boolean) {}
 
@@ -73,6 +74,7 @@ class Assembler {
         this.labels = new Set();
         // Need to cast SegmentType.Instructions because of TypeScript bug.
         this.section = SegmentType.Instructions as SegmentType;
+        this.first_section_marker = true;
 
         for (const line of this.assembly) {
             this.tokens = this.lexer.tokenize_line(line);
@@ -190,7 +192,7 @@ class Assembler {
             const arr = new Uint8Array(buf);
 
             arr.set(new Uint8Array(this.segment.data));
-            arr.set(new Uint8Array(bytes));
+            arr.set(new Uint8Array(bytes), this.segment.data.byteLength);
 
             this.segment.data = buf;
         } else {
@@ -349,7 +351,7 @@ class Assembler {
                 break;
         }
 
-        if (this.section === section) {
+        if (this.section === section && !this.first_section_marker) {
             this.add_warning({
                 col,
                 length: len,
@@ -358,6 +360,7 @@ class Assembler {
         }
 
         this.section = section;
+        this.first_section_marker = false;
 
         const next_token = this.tokens.shift();
 
diff --git a/src/scripting/disassembly.test.ts b/src/scripting/disassembly.test.ts
new file mode 100644
index 00000000..0c4c0897
--- /dev/null
+++ b/src/scripting/disassembly.test.ts
@@ -0,0 +1,65 @@
+import { readFileSync } from "fs";
+import { Endianness } from "../data_formats";
+import { prs_decompress } from "../data_formats/compression/prs/decompress";
+import { ArrayBufferCursor } from "../data_formats/cursor/ArrayBufferCursor";
+import { BufferCursor } from "../data_formats/cursor/BufferCursor";
+import { parse_bin, write_bin } from "../data_formats/parsing/quest/bin";
+import { assemble } from "./assembly";
+import { disassemble } from "./disassembly";
+
+// Roundtrip test.
+test("assembling dissambled object code with manual stack management should result in the same object code", () => {
+    const orig_buffer = readFileSync("test/resources/quest27_e.bin");
+    const orig_bytes = prs_decompress(new BufferCursor(orig_buffer, Endianness.Little));
+    const bin = parse_bin(orig_bytes);
+
+    const { object_code, warnings, errors } = assemble(disassemble(bin.object_code, true), true);
+
+    expect(errors).toEqual([]);
+    expect(warnings).toEqual([]);
+
+    bin.object_code.splice(0, bin.object_code.length, ...object_code);
+
+    const test_bytes = new ArrayBufferCursor(write_bin(bin), Endianness.Little);
+
+    orig_bytes.seek_start(0);
+    expect(test_bytes.size).toBe(orig_bytes.size);
+
+    let matching_bytes = 0;
+
+    while (orig_bytes.bytes_left) {
+        const test_byte = test_bytes.u8();
+        const orig_byte = orig_bytes.u8();
+
+        if (test_byte !== orig_byte) {
+            throw new Error(
+                `Byte ${matching_bytes} didn't match, expected ${orig_byte}, got ${test_byte}.`
+            );
+        }
+
+        matching_bytes++;
+    }
+
+    expect(matching_bytes).toBe(orig_bytes.size);
+});
+
+// Roundtrip test.
+test("disassembling assembled assembly code with automatic stack management should result the same assembly code", () => {
+    const orig_buffer = readFileSync("test/resources/quest27_e.bin");
+    const orig_bytes = prs_decompress(new BufferCursor(orig_buffer, Endianness.Little));
+    const orig_asm = disassemble(parse_bin(orig_bytes).object_code);
+
+    const { object_code, warnings, errors } = assemble(orig_asm);
+
+    expect(errors).toEqual([]);
+    expect(warnings).toEqual([]);
+
+    const test_asm = disassemble(object_code);
+    const len = Math.min(orig_asm.length, test_asm.length);
+
+    for (let i = 0; i < len; i++) {
+        expect(test_asm[i]).toBe(orig_asm[i]);
+    }
+
+    expect(test_asm.length).toBe(orig_asm.length);
+});