From 0775772e291a091d3fc60699e457154816156a6d Mon Sep 17 00:00:00 2001
From: Adam Eury <euryadam@gmail.com>
Date: Sun, 14 Jan 2024 23:53:39 -0500
Subject: [PATCH] Add JXNZ instruction.

---
 examples/factorial.g |  4 +--
 gmachine.go          | 11 +++++++-
 gmachine_test.go     | 63 ++++++++++++++++++++++++++++++++++++++++++++
 token/token.go       |  1 +
 token/token_test.go  |  1 +
 5 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/examples/factorial.g b/examples/factorial.g
index 79c6dc6..bc14690 100644
--- a/examples/factorial.g
+++ b/examples/factorial.g
@@ -1,8 +1,8 @@
 ; Tasks:
 ; 1. Implement MULA - done
 ; 2. Implement SETX - done
-; 3. Implement DECX (and DECA, DECY)
-; 4. Implement JXNZ
+; 3. Implement DECX (and DECA, DECY) - done
+; 4. Implement JXNZ - done
 ; 5. Implement CALL and RTRN
 
 .factorial
diff --git a/gmachine.go b/gmachine.go
index 3421b98..083b536 100644
--- a/gmachine.go
+++ b/gmachine.go
@@ -36,6 +36,7 @@ const (
 	OpPSHA
 	OpPOPA
 	OpJUMP
+	OpJXNZ
 )
 
 const (
@@ -79,6 +80,7 @@ var opcodes = map[string]Word{
 	"PSHA": OpPSHA,
 	"POPA": OpPOPA,
 	"JUMP": OpJUMP,
+	"JXNZ": OpJXNZ,
 }
 
 type Word uint64
@@ -173,10 +175,17 @@ func (g *Machine) Run() {
 			g.A = g.Memory[g.S]
 		case OpJUMP:
 			g.P = g.Memory[g.MemOffset+g.P]
+		case OpJXNZ:
+			if g.X != 0 {
+				g.P = g.Memory[g.MemOffset+g.P]
+			} else {
+				g.P++
+			}
 		default:
 			g.E = ExceptionIllegalInstruction
 			return
 		}
+
 	}
 }
 
@@ -294,7 +303,7 @@ func assembleOpcodeStatement(stmt *ast.OpcodeStatement, program []Word, refs []R
 		}
 		program = append(program, register)
 	case *ast.Identifier:
-		if !slices.Contains([]Word{OpSETA, OpJUMP}, opcode) {
+		if !slices.Contains([]Word{OpSETA, OpJUMP, OpJXNZ}, opcode) {
 			return nil, nil, fmt.Errorf("%w: %s at line %d", ErrInvalidOperand, stmt.TokenLiteral(), stmt.Token.Line)
 		}
 		ref := Ref{
diff --git a/gmachine_test.go b/gmachine_test.go
index bb0d8c3..6b16497 100644
--- a/gmachine_test.go
+++ b/gmachine_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"gmachine/parser"
 	"io"
 	"os"
@@ -424,6 +425,68 @@ func TestJUMPWithInvalidNumber(t *testing.T) {
 	}
 }
 
+func TestJXNZ(t *testing.T) {
+	t.Parallel()
+	g := gmachine.New(nil)
+	var wantA gmachine.Word = 10
+	var wantX gmachine.Word = 0
+	err := assembleAndRunFromString(g, `
+SETA 0
+SETX 10
+.loop
+INCA
+DECX
+JXNZ loop
+HALT
+`)
+	if err != nil {
+		t.Fatal("didn't expect an error", err)
+	}
+	if wantA != g.A {
+		t.Errorf("want A value %d, got %d", wantA, g.A)
+	}
+	if wantX != g.X {
+		t.Errorf("want X value %d, got %d", wantX, g.X)
+	}
+}
+
+func TestFactorial(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		factorial int
+		wantA     gmachine.Word
+	}{
+		{1, 1},
+		{2, 2},
+		{3, 6},
+		{4, 24},
+		{5, 120},
+		{6, 720},
+		{7, 5040},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%d!", tt.factorial), func(t *testing.T) {
+			g := gmachine.New(nil)
+			program := fmt.Sprintf(`
+SETA 1
+SETX %d
+.factorial
+MULA X
+DECX
+JXNZ factorial
+HALT
+`, tt.factorial)
+			err := assembleAndRunFromString(g, program)
+			if err != nil {
+				t.Fatal("didn't expect an error", err)
+			}
+			if tt.wantA != g.A {
+				t.Errorf("want A value %d, got %d", tt.wantA, g.A)
+			}
+		})
+	}
+}
+
 func TestAssemble_SkipsComments(t *testing.T) {
 	t.Parallel()
 	want := []gmachine.Word{}
diff --git a/token/token.go b/token/token.go
index e7ea5c8..b773ae2 100644
--- a/token/token.go
+++ b/token/token.go
@@ -35,6 +35,7 @@ var opcodes = map[string]TokenType{
 	"PSHA": OPCODE,
 	"POPA": OPCODE,
 	"JUMP": OPCODE,
+	"JXNZ": OPCODE,
 }
 
 var pragmas = map[string]TokenType{
diff --git a/token/token_test.go b/token/token_test.go
index 85aec78..4412d14 100644
--- a/token/token_test.go
+++ b/token/token_test.go
@@ -27,6 +27,7 @@ func TestLookupIdent(t *testing.T) {
 		{"PSHA", token.OPCODE},
 		{"POPA", token.OPCODE},
 		{"JUMP", token.OPCODE},
+		{"JXNZ", token.OPCODE},
 		{"X", token.REGISTER},
 		{"test", token.IDENT},
 		{".test", token.IDENT},