Add labels and branches with the : and b commands

1 files changed, 32 insertions, 2 deletions

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
 }