Add labels and branches with the : and b commands

3 files changed, 54 insertions, 2 deletions

diff --git a/main/command.go b/main/command.go
index 44fd9eb..63cc3b8 100644
--- a/main/command.go
+++ b/main/command.go
@@ -207,4 +207,16 @@ func (cmd JumpCommand) exec(state *ProgramState) {
 }
 func (cmd JumpCommand) String() string {
 	return fmt.Sprintf("b%v", cmd.destination)
+}
+
+// Placeholder as the branch destination may not have been parsed when the branch command is
+// Should never appear in a program to actually be run
+type BranchPlaceholderCommand struct {
+	label rune
+}
+func (cmd BranchPlaceholderCommand) exec(state *ProgramState) {
+	panic("Tried to execute a BranchPlaceholderCommand!!!")
+}
+func (cmd BranchPlaceholderCommand) String() string {
+	return fmt.Sprintf("b%c", cmd.label)
 }
\ No newline at end of file
diff --git a/main/lex.go b/main/lex.go
index f28244d..198c346 100644
--- a/main/lex.go
+++ b/main/lex.go
@@ -118,6 +118,7 @@ const (
 	TokenCommand // A command character
 	TokenSubstituteDelimiter // usually / but could be something else
 	TokenSubex // A subex
+	TokenLabel // A label
 )
 
 type Token struct {
@@ -182,6 +183,9 @@ func lexCommand(l *lexer) stateFunc {
 		case 's', 'S', 'f', 'F', 'l', 'L', 'a', 'A':
 			l.emit(TokenCommand)
 			return lexSubstitution
+		case ':', 'b':
+			l.emit(TokenCommand)
+			return lexLabel
 	}
 	if isAlpha(r) {
 		l.emit(TokenCommand)
@@ -212,3 +216,9 @@ func lexSubstitution(l *lexer) stateFunc {
 	}
 	return lexCommand
 }
+
+func lexLabel(l *lexer) stateFunc {
+	l.next()
+	l.emit(TokenLabel)
+	return lexCommand
+}
diff --git a/main/parse.go b/main/parse.go
index ef50e81..cbbfb9a 100644
--- a/main/parse.go
+++ b/main/parse.go
@@ -1,14 +1,16 @@
 package main
 
 import (
-	"strings"
 	"fmt"
 	"main/subex"
+	"strings"
+	"unicode/utf8"
 )
 
 type parser struct {
 	tokenStream chan Token
 	rewinds []Token
+	labels map[rune]int
 }
 func (p *parser) next() Token {
 	var token Token
@@ -146,6 +148,21 @@ func (p *parser) parseBasicCommand(commands []Command, commandChar rune) []Comma
 			return append(commands, SwapPathCommand{})
 		case 'K':
 			return append(commands, AppendPathCommand{})
+		case ':':
+			labelToken := p.next()
+			if labelToken.typ != TokenLabel {
+				panic("Missing branch label")
+			}
+			label, _ := utf8.DecodeRuneInString(labelToken.val)
+			p.labels[label] = len(commands)
+			return commands
+		case 'b':
+			labelToken := p.next()
+			if labelToken.typ != TokenLabel {
+				panic("Missing branch label")
+			}
+			label, _ := utf8.DecodeRuneInString(labelToken.val)
+			return append(commands, BranchPlaceholderCommand {label})
 		default:
 			panic("Invalid command")
 	}
@@ -187,6 +204,19 @@ func (p *parser) parseCommands(commands []Command) []Command {
 func Parse(tokens chan Token) []Command {
 	p := parser {
 		tokenStream: tokens,
+		rewinds: nil,
+		labels: make(map[rune]int),
+	}
+	program := p.parseCommands(nil)
+	for i, command := range program {
+		switch branch := command.(type) {
+			case BranchPlaceholderCommand:
+				destination, exists := p.labels[branch.label]
+				if !exists {
+					panic("Tried to branch to a label that doesn't exist")
+				}
+				program[i] = JumpCommand {destination}
+		}
 	}
-	return p.parseCommands(nil)
+	return program
 }