From e2ee765a1f9413008bce044e16dab898e6e36f57 Mon Sep 17 00:00:00 2001
From: Taeber Rapczak <taeber@rapczak.com>
Date: Fri, 29 Apr 2022 12:40:05 -0400
Subject: [PATCH] Support arbitrary length arithmetic expressions

Before this change, you could only supply "offsets" in the form of
BASE+OFF or BASE-OFF where BASE was a label or number and OFF was a
number. After looking at the Merlin manual, the assembly actually
supports 4 arithmetic operations: add, subtract, divide, and multiply.
Evaluation is done left to right and parenthesis are disallowed.

Another constraint is length--"the maximum allowable combined
OPERAND+COMMENT length is 64 characters"--but this is not enforced by
a2asm.
---
 lib.go      | 99 +++++++++++++++++++++++++++++++++--------------------
 lib_test.go | 20 +++++++++++
 2 files changed, 81 insertions(+), 38 deletions(-)

diff --git a/lib.go b/lib.go
index 418ef00..08bed65 100644
--- a/lib.go
+++ b/lib.go
@@ -194,16 +194,22 @@ func readMneumonic(line []byte) (mneumonic string, remaining []byte) {
 	return
 }
 
+// readNumber parses text for a number (decimal, hex, or binary) and returns
+// the uint16 representation of it along with the remaining text. If parsing
+// fails, an error is returned and the other returned values should be ignored.
 func readNumber(text []byte) (uint16, []byte, error) {
 	if text[0] == '$' {
 		// Read hex literal.
-		num, err := strconv.ParseUint(string(text[1:]), 16, 16)
-		if err != nil {
-			return 0, text, err
+		var i int
+		for i = 1; i < len(text); i++ {
+			if !isHex(text[i]) {
+				break
+			}
 		}
 
-		var i int
-		for i = 1; i < len(text) && isHex(text[i]); i++ {
+		num, err := strconv.ParseUint(string(text[1:i]), 16, 16)
+		if err != nil {
+			return 0, text, err
 		}
 
 		return uint16(num), text[i:], err
@@ -211,11 +217,6 @@ func readNumber(text []byte) (uint16, []byte, error) {
 
 	if text[0] == '%' {
 		// Read binary literal.
-		num, err := strconv.ParseUint(string(text[1:]), 2, 16)
-		if err != nil {
-			return 0, text, err
-		}
-
 		var i int
 		for i = 1; i < len(text); i++ {
 			if text[i] != '0' && text[i] != '1' {
@@ -223,18 +224,25 @@ func readNumber(text []byte) (uint16, []byte, error) {
 			}
 		}
 
-		return uint16(num), text[i:], err
-	}
-
-	if isDigit(text[0]) {
-		num, err := strconv.ParseUint(string(text), 10, 16)
+		num, err := strconv.ParseUint(string(text[1:i]), 2, 16)
 		if err != nil {
 			return 0, text, err
 		}
 
+		return uint16(num), text[i:], err
+	}
+
+	if isDigit(text[0]) {
 		var i int
+		for i = 0; i < len(text); i++ {
+			if !isDigit(text[i]) {
+				break
+			}
+		}
 
-		for i = 0; i < len(text) && isDigit(text[i]); i++ {
+		num, err := strconv.ParseUint(string(text[0:i]), 10, 16)
+		if err != nil {
+			return 0, text, err
 		}
 
 		return uint16(num), text[i:], err
@@ -247,7 +255,7 @@ func readNumber(text []byte) (uint16, []byte, error) {
 			return uint16(text[1]), text[3:], nil
 		}
 		if text[0] == '"' && text[2] == '"' {
-			// high-ASCII (high-bit on
+			// high-ASCII (high-bit on)
 			return uint16(text[1] | highASCII), text[3:], nil
 		}
 	}
@@ -304,7 +312,7 @@ func parseLine(s *state) (err error) {
 
 	switch mneumonic {
 	case "ORG":
-		s.Address, line, err = readNumber(line)
+		s.Address, _, err = readNumber(line)
 		s.Origin = s.Address
 		return
 
@@ -1313,28 +1321,43 @@ func parseOperandValue(val []byte) (num uint16, ref string, err error) {
 
 	val = val[end:]
 
-	if len(val) == 0 {
-		return
-	}
-
-	var num2 uint16
-	switch val[0] {
-	case ' ':
-	case '+':
-		num2, _, err = readNumber(val[1:])
-		if err != nil {
+	// According to the Merlin manual, the "assembler supports four arithmetic
+	// operations: +, -, /, and *." "All ... operations are done from left to
+	// right (2+3*5 would assemble as 25 and not 17)."
+	for len(val) > 0 {
+		var num2 uint16
+		switch val[0] {
+		case ' ':
+			// Assume the rest of the line is a comment.
 			return
-		}
-		num += num2
-	case '-':
-		num2, _, err = readNumber(val[1:])
-		if err != nil {
+		case '+':
+			num2, val, err = readNumber(val[1:])
+			if err != nil {
+				return
+			}
+			num += num2
+		case '-':
+			num2, val, err = readNumber(val[1:])
+			if err != nil {
+				return
+			}
+			num -= num2
+		case '*':
+			num2, val, err = readNumber(val[1:])
+			if err != nil {
+				return
+			}
+			num *= num2
+		case '/':
+			num2, val, err = readNumber(val[1:])
+			if err != nil {
+				return
+			}
+			num /= num2
+		default:
+			err = fmt.Errorf("invalid arithmetic operator: %c", val[0])
 			return
 		}
-		num -= num2
-	default:
-		err = fmt.Errorf("invalid +/- offset")
-		return
 	}
 
 	return
@@ -1345,7 +1368,7 @@ func (s *state) error(err error) error {
 		return nil
 	}
 
-	return fmt.Errorf("Line %d - %s", s.LineNumber, err.Error())
+	return fmt.Errorf("line %d - %s", s.LineNumber, err.Error())
 }
 
 func (s *state) errorf(format string, a ...interface{}) error {
diff --git a/lib_test.go b/lib_test.go
index b0e776a..c95e7ad 100644
--- a/lib_test.go
+++ b/lib_test.go
@@ -311,6 +311,26 @@ MAIN	JSR INIT
 	}
 }
 
+func TestArithmetic(t *testing.T) {
+	out := bytes.NewBuffer(nil)
+	prg := strings.NewReader(`
+	ORG $800
+MAIN	LDA #MAIN+$A+2*2
+	`)
+
+	_, err := Assemble(out, prg, true)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	expected := []byte("\xA9\x18")
+	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)),