From ae44f8661a4e04fcccc3a1dea4e1705753bc510c Mon Sep 17 00:00:00 2001
From: Taeber Rapczak <taeber@rapczak.com>
Date: Sat, 16 Apr 2022 23:25:05 -0400
Subject: [PATCH] Fix jumps to yet-to-be defined labels

fixes #2
---
 lib.go      | 15 +++++++++------
 lib_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/lib.go b/lib.go
index a358836..418ef00 100644
--- a/lib.go
+++ b/lib.go
@@ -37,17 +37,20 @@ func Assemble(dst io.Writer, src io.Reader, headless bool) (written uint, err er
 	err = nil
 
 	for lbl := range s.References {
-		if lbl[0] == '>' || lbl[0] == '<' {
-			num, ok := s.Constants[lbl]
-			if !ok {
+		num, ok := s.Constants[lbl]
+		if !ok {
+			if lbl[0] == '>' || lbl[0] == '<' {
 				err = s.errorf("unknown label: %s", lbl)
 				return
 			}
-
+		} else {
 			for _, ref := range s.References[lbl] {
-				s.Memory[ref.Address] += uint8(num)
+				if num <= 0xFF {
+					s.Memory[ref.Address] += uint8(num)
+				} else {
+					binary.LittleEndian.PutUint16(s.Memory[ref.Address:], num)
+				}
 			}
-
 			continue
 		}
 
diff --git a/lib_test.go b/lib_test.go
index cdba927..b0e776a 100644
--- a/lib_test.go
+++ b/lib_test.go
@@ -285,6 +285,32 @@ EXIT  RTS
 	}
 }
 
+func TestJumpEqu(t *testing.T) {
+	out := bytes.NewBuffer(nil)
+	prg := strings.NewReader(`
+	ORG $800
+	JSR MAIN
+	JMP EXIT
+EXIT	EQU $3D0
+INIT	EQU $FB2F
+HOME	EQU $FC58
+MAIN	JSR INIT
+	JSR HOME
+	RTS
+	`)
+	_, err := Assemble(out, prg, true)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	expected := []byte("\x20\x06\x08\x4C\xD0\x03\x20\x2F\xFB\x20\x58\xFC\x60")
+	actual := out.Bytes()
+	if !bytes.Equal(expected, actual) {
+		t.Errorf("Expected %v; got %v", expected, actual)
+	}
+}
+
 func test(t *testing.T, assembly, expected string) {
 	s := state{
 		Reader: bufio.NewReader(strings.NewReader(assembly)),