From 1b8ff35b6258c8328cdcdc89b41a54e9bee710af Mon Sep 17 00:00:00 2001
From: Tristan Morgan <tristanmorgan@users.noreply.github.com>
Date: Thu, 30 May 2024 14:29:41 +1000
Subject: [PATCH] Switch to switch and collapse nested loops.

---
 parser/execute_test.go     |  2 +-
 parser/instruction.go      |  1 +
 parser/instruction_test.go |  2 ++
 parser/tokenise.go         | 13 +++++++++----
 parser/tokenise_test.go    | 18 ++----------------
 5 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/parser/execute_test.go b/parser/execute_test.go
index 0376591..1cb6288 100644
--- a/parser/execute_test.go
+++ b/parser/execute_test.go
@@ -58,7 +58,7 @@ type bufferWriter struct {
 
 func (writer *bufferWriter) Write(p []byte) (int, error) {
 	writer.buf = append(writer.buf, p...)
-	return 0, nil
+	return len(p), nil
 }
 
 func BenchmarkExecute(b *testing.B) {
diff --git a/parser/instruction.go b/parser/instruction.go
index f32392a..bad42fa 100644
--- a/parser/instruction.go
+++ b/parser/instruction.go
@@ -54,5 +54,6 @@ func (inst Instruction) IsZeroOp() bool {
 	return inst.operator == opJmpNz ||
 		inst.operator == opNoop ||
 		inst.operator == opMove ||
+		inst.operator == opSkip ||
 		(inst.operator == opSetVal && inst.operand == 0)
 }
diff --git a/parser/instruction_test.go b/parser/instruction_test.go
index 5f72b63..420c450 100644
--- a/parser/instruction_test.go
+++ b/parser/instruction_test.go
@@ -58,6 +58,7 @@ func TestIsZeroOp(t *testing.T) {
 		{opSetVal, -1},
 		{opSetVal, 0},
 		{opOut, 1},
+		{opSkip, 1},
 		{opIn, 1},
 		{opJmpZ, 0},
 		{opJmpNz, 0},
@@ -69,6 +70,7 @@ func TestIsZeroOp(t *testing.T) {
 		false,
 		true,
 		false,
+		true,
 		false,
 		false,
 		true,
diff --git a/parser/tokenise.go b/parser/tokenise.go
index 6957a65..7e4bd1c 100644
--- a/parser/tokenise.go
+++ b/parser/tokenise.go
@@ -51,23 +51,28 @@ func Tokenise(input io.ByteReader) (program []Instruction, err error) {
 			jmpStack = jmpStack[:len(jmpStack)-1]
 			program[pc].operand = jmpPc
 			program[jmpPc].operand = pc
-			if program[jmpPc-1].IsZeroOp() {
+			switch {
+			case program[jmpPc-1].IsZeroOp():
 				pc = jmpPc
 				program = program[:pc]
 				pc--
-			} else if pc-jmpPc == 2 && program[pc-1].SameOp(NewInstruction('+')) {
+			case pc-jmpPc == 2 &&
+				(program[pc-1].SameOp(NewInstruction('+')) ||
+					program[pc-1].operator == opSetVal):
 				pc = jmpPc
 				if program[jmpPc-1].SameOp(NewInstruction('+')) {
 					pc--
 				}
 				program = program[:pc]
 				program = append(program, Instruction{opSetVal, 0})
-			} else if pc-jmpPc == 2 && program[pc-1].SameOp(NewInstruction('>')) {
+			case pc-jmpPc == 2 &&
+				(program[pc-1].SameOp(NewInstruction('>')) ||
+					program[pc-1].operator == opSkip):
 				offset := program[pc-1].operand
 				pc = jmpPc
 				program = program[:pc]
 				program = append(program, Instruction{opSkip, offset})
-			} else if pc-jmpPc == 5 { // looking for opMulVal and opMove
+			case pc-jmpPc == 5: // looking for opMulVal and opMove
 				var factor, offset = 0, 0
 				if program[pc-4].Complement(NewInstruction('+')) &&
 					program[pc-3].Complement(program[pc-1]) &&
diff --git a/parser/tokenise_test.go b/parser/tokenise_test.go
index 0c1f34d..eadb189 100644
--- a/parser/tokenise_test.go
+++ b/parser/tokenise_test.go
@@ -99,25 +99,11 @@ func TestTokenise(t *testing.T) {
 		},
 		{
 			"op_nested",
-			">[[[[[[[,]]]]]]][comment.]",
+			">[[[[[[[>]]]]]]][comment.]",
 			[]Instruction{
 				{opNoop, 0},
 				{opAddDp, 1},
-				{opJmpZ, 16},
-				{opJmpZ, 15},
-				{opJmpZ, 14},
-				{opJmpZ, 13},
-				{opJmpZ, 12},
-				{opJmpZ, 11},
-				{opJmpZ, 10},
-				{opIn, 1},
-				{opJmpNz, 8},
-				{opJmpNz, 7},
-				{opJmpNz, 6},
-				{opJmpNz, 5},
-				{opJmpNz, 4},
-				{opJmpNz, 3},
-				{opJmpNz, 2},
+				{opSkip, 1},
 			},
 		},
 	}