Improve internals of command control flowHEADmain

3 files changed, 267 insertions, 82 deletions

diff --git a/main/command.go b/main/command.go
index bbbb036..832a236 100644
--- a/main/command.go
+++ b/main/command.go
@@ -67,6 +67,7 @@ func (cmd AppendNextCommand) String() string {
 
 type SubstituteNextCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteNextCommand) exec(state *ProgramState) {
 	item, err := state.Peek()
@@ -77,14 +78,14 @@ func (cmd SubstituteNextCommand) exec(state *ProgramState) {
 	newValue, notOk := runSubex(cmd.subex, []walk.Value{item.Value})
 
 	if notOk {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
+		state.pc++
 		state.Read()
 		state.prevStart = item.PrevStart
 		state.start = item.Start
 		state.end = item.End
 		state.nextEnd = item.NextEnd
-		state.pc += 2
 		state.value = newValue
 	}
 }
@@ -94,6 +95,7 @@ func (cmd SubstituteNextCommand) String() string {
 
 type SubstituteAppendNextCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteAppendNextCommand) exec(state *ProgramState) {
 	item, err := state.Peek()
@@ -104,14 +106,14 @@ func (cmd SubstituteAppendNextCommand) exec(state *ProgramState) {
 	newValue, notOk := runSubex(cmd.subex, []walk.Value{item.Value})
 
 	if notOk {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
 		state.Read()
 		state.prevStart = item.PrevStart
 		state.start = item.Start
 		state.end = item.End
 		state.nextEnd = item.NextEnd
-		state.pc += 2
+		state.pc++
 		state.value = append(state.value, newValue...)
 	}
 }
@@ -140,15 +142,16 @@ func (cmd MergeCommand) String() string {
 
 type FullMergeCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd FullMergeCommand) exec(state *ProgramState) {
 	_, notOk := runSubex(cmd.subex, state.value)
 	if notOk {
-		state.pc++
+		state.pc += cmd.elseJump
 		return
 	}
 	if !state.start {
-		state.pc += 2
+		state.pc++
 		return
 	}
 
@@ -170,7 +173,7 @@ func (cmd FullMergeCommand) exec(state *ProgramState) {
 			state.start = item.Start
 			state.end = item.End
 			state.nextEnd = item.NextEnd
-			state.pc += 2
+			state.pc++
 			return
 		}
 	}
@@ -198,13 +201,14 @@ func runSubex(state subex.Transducer, in []walk.Value) ([]walk.Value, bool) {
 
 type SubstituteValueCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteValueCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.value = newValue
 	}
 }
@@ -212,54 +216,72 @@ func (cmd SubstituteValueCommand) String() string {
 	return "s/.../"
 }
 
-type IsStartCommand struct {}
+type IsStartCommand struct {
+	elseJump int
+}
 func (cmd IsStartCommand) exec(state *ProgramState) {
 	if state.start {
-		state.pc += 2
+		state.pc++
 	} else {
-		state.pc += 1
+		state.pc += cmd.elseJump
 	}
 }
 func (cmd IsStartCommand) String() string {
 	return "a"
 }
 
-type IsPrevStartCommand struct {}
+type IsPrevStartCommand struct {
+	elseJump int
+}
 func (cmd IsPrevStartCommand) exec(state *ProgramState) {
 	if state.prevStart {
-		state.pc += 2
+		state.pc++
 	} else {
-		state.pc += 1
+		state.pc += cmd.elseJump
 	}
 }
 func (cmd IsPrevStartCommand) String() string {
 	return "A"
 }
 
-type IsEndCommand struct {}
+type IsEndCommand struct {
+	elseJump int
+}
 func (cmd IsEndCommand) exec(state *ProgramState) {
 	if state.end {
-		state.pc += 2
+		state.pc++
 	} else {
-		state.pc += 1
+		state.pc += cmd.elseJump
 	}
 }
 func (cmd IsEndCommand) String() string {
 	return "e"
 }
 
-type IsNextEndCommand struct {}
+type IsNextEndCommand struct {
+	elseJump int
+}
 func (cmd IsNextEndCommand) exec(state *ProgramState) {
 	if state.nextEnd {
-		state.pc += 2
+		state.pc++
 	} else {
-		state.pc += 1
+		state.pc += cmd.elseJump
 	}
 }
 func (cmd IsNextEndCommand) String() string {
 	return "E"
 }
 
+type LabelCommand struct {
+	label rune
+}
+func (cmd LabelCommand) exec(state *ProgramState) {
+	state.pc++
+}
+func (cmd LabelCommand) String() string {
+	return fmt.Sprintf(":%c", cmd.label)
+}
+
 type NoopCommand struct {}
 func (cmd NoopCommand) exec(state *ProgramState) {
 	state.pc++
@@ -290,13 +312,14 @@ func (cmd AppendXRegCommand) String() string {
 
 type SubstituteToXRegCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteToXRegCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.xreg = newValue
 	}
 }
@@ -306,13 +329,14 @@ func (cmd SubstituteToXRegCommand) String() string {
 
 type SubstituteAppendXRegCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteAppendXRegCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.xreg = append(state.xreg, newValue...)
 	}
 }
@@ -342,13 +366,14 @@ func (cmd AppendYRegCommand) String() string {
 
 type SubstituteToYRegCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteToYRegCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.yreg = newValue
 	}
 }
@@ -358,13 +383,14 @@ func (cmd SubstituteToYRegCommand) String() string {
 
 type SubstituteAppendYRegCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteAppendYRegCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.yreg = append(state.xreg, newValue...)
 	}
 }
@@ -394,13 +420,14 @@ func (cmd AppendZRegCommand) String() string {
 
 type SubstituteToZRegCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteToZRegCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.zreg = newValue
 	}
 }
@@ -410,13 +437,14 @@ func (cmd SubstituteToZRegCommand) String() string {
 
 type SubstituteAppendZRegCommand struct {
 	subex subex.Transducer
+	elseJump int
 }
 func (cmd SubstituteAppendZRegCommand) exec(state *ProgramState) {
 	newValue, err := runSubex(cmd.subex, state.value)
 	if err {
-		state.pc++
+		state.pc += cmd.elseJump
 	} else {
-		state.pc += 2
+		state.pc++
 		state.zreg = append(state.xreg, newValue...)
 	}
 }
@@ -424,6 +452,16 @@ func (cmd SubstituteAppendZRegCommand) String() string {
 	return "Z/.../"
 }
 
+type RelativeJumpCommand struct {
+	destination int
+}
+func (cmd RelativeJumpCommand) exec(state *ProgramState) {
+	state.pc += cmd.destination
+}
+func (cmd RelativeJumpCommand) String() string {
+	return fmt.Sprintf("b+%v", cmd.destination)
+}
+
 type JumpCommand struct {
 	destination int
 }
diff --git a/main/lex.go b/main/lex.go
index da517cc..0bcdaec 100644
--- a/main/lex.go
+++ b/main/lex.go
@@ -180,7 +180,7 @@ func lexCommand(l *lexer) stateFunc {
 	case '}':
 		l.emit(TokenRBrace)
 		return lexCommand
-	case 's', 'S', 'M':
+	case 's', 'S', 'M', 'r':
 		l.emit(TokenCommand)
 		return lexSubstitution
 	case 'x', 'X', 'y', 'Y', 'z', 'Z', 'n', 'N':
diff --git a/main/parse.go b/main/parse.go
index 36bd3ee..d0a0255 100644
--- a/main/parse.go
+++ b/main/parse.go
@@ -10,7 +10,6 @@ import (
 type parser struct {
 	tokenStream chan Token
 	rewinds []Token
-	labels map[rune]int
 }
 func (p *parser) next() Token {
 	var token Token
@@ -53,146 +52,288 @@ func (p *parser) parseSubex() subex.SubexAST {
 	return subexAST
 }
 
-func (p *parser) parseBasicCommand(commands []Command, commandChar rune) []Command {
+func (p *parser) parseBasicCommand(commandChar rune) []Command {
 	switch commandChar {
 		case 'p':
-			return append(commands, PrintValueCommand{})
+			return []Command {PrintValueCommand{}}
 		case 'd':
-			return append(commands, DeleteValueCommand{})
+			return []Command {DeleteValueCommand{}}
 		case 'n':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, NextCommand{})
+				return []Command {NextCommand{}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteNextCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteNextCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'N':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, AppendNextCommand{})
+				return []Command {AppendNextCommand{}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteAppendNextCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteAppendNextCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'm':
-			return append(commands, MergeCommand{})
+			return []Command {MergeCommand {}}
 		case 'M':
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, FullMergeCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					FullMergeCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 's':
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteValueCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteValueCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
+		case 'r':
+			ast := p.parseSubex()
+			subex := subex.CompileTransducer(ast)
+			return []Command {
+				SubstituteValueCommand {
+					subex: subex,
+					elseJump: 2,
+				},
+				RelativeJumpCommand {
+					destination: -1,
+				},
+			}
 		case 'o':
-			return append(commands, NoopCommand{})
+			return []Command {NoopCommand {}}
 		case 'x':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, SwapXRegCommand{})
+				return []Command {SwapXRegCommand {}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteToXRegCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteToXRegCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'X':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, AppendXRegCommand{})
+				return []Command {AppendXRegCommand {}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteAppendXRegCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteAppendXRegCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'y':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, SwapYRegCommand{})
+				return []Command {SwapYRegCommand {}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteToYRegCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteToYRegCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'Y':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, AppendYRegCommand{})
+				return []Command {AppendYRegCommand {}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteAppendYRegCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteAppendYRegCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'z':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, SwapZRegCommand{})
+				return []Command {SwapZRegCommand {}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteToZRegCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteToZRegCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'Z':
 			delim := p.peek()
 			if delim.typ != TokenSubstituteDelimiter {
-				return append(commands, AppendZRegCommand{})
+				return []Command {AppendZRegCommand {}}
 			}
 			ast := p.parseSubex()
 			subex := subex.CompileTransducer(ast)
-			return append(commands, SubstituteAppendZRegCommand {subex}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					SubstituteAppendZRegCommand {
+						subex: subex,
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'a':
-			return append(commands, IsStartCommand{}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					IsStartCommand {
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'A':
-			return append(commands, IsPrevStartCommand{}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					IsPrevStartCommand {
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'e':
-			return append(commands, IsEndCommand{}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					IsEndCommand {
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case 'E':
-			return append(commands, IsNextEndCommand{}, JumpCommand {len(commands) + 3})
+			elseBranch := p.parseCommand()
+			return append(
+				[]Command {
+					IsNextEndCommand {
+						elseJump: len(elseBranch) + 1,
+					},
+				},
+				elseBranch...,
+			)
 		case ':':
 			labelToken := p.next()
 			if labelToken.typ != TokenLabel {
 				panic("Missing branch label")
 			}
 			label, _ := utf8.DecodeRuneInString(labelToken.val)
-			p.labels[label] = len(commands)
-			return commands
+			return []Command {
+				LabelCommand {
+					label: label,
+				},
+			}
 		case 'b':
 			labelToken := p.next()
 			if labelToken.typ != TokenLabel {
 				panic("Missing branch label")
 			}
 			label, _ := utf8.DecodeRuneInString(labelToken.val)
-			return append(commands, BranchPlaceholderCommand {label})
+			return []Command {
+				BranchPlaceholderCommand {
+					label: label,
+				},
+			}
 		default:
 			panic("Invalid command")
 	}
 }
 
-func (p *parser) parseCommand(commands []Command) []Command {
+func (p *parser) parseCommand() []Command {
 	token := p.next()
 	switch token.typ {
 		case TokenLBrace:
-			jumpToBlockCommand := &JumpCommand{0}
-			commands = append(commands, JumpCommand {len(commands) + 2}, jumpToBlockCommand)
-			commands = p.parseCommands(commands)
+			children := p.parseCommandSequence()
 			if p.next().typ != TokenRBrace {
 				panic("Missing matching }")
 			}
-			jumpToBlockCommand.destination = len(commands)
-			return commands
+			return children
+		case TokenRBrace, TokenEOF:
+			p.rewind(token)
+			return nil
 		case TokenCommand:
 			commandChar, _, err := strings.NewReader(token.val).ReadRune()
 			if err != nil {
 				panic("Error reading a command character!?")
 			}
-			return p.parseBasicCommand(commands, commandChar)
+			return p.parseBasicCommand(commandChar)
 		default:
 			panic("Invalid token, expected command")
 	}
 }
 
-func (p *parser) parseCommands(commands []Command) []Command {
+func (p *parser) parseCommandSequence() []Command {
+	var commands []Command
 	for {
 		nextToken := p.peek()
 		if nextToken.typ == TokenEOF || nextToken.typ == TokenRBrace {
 			return commands
 		}
-		commands = p.parseCommand(commands)
+		commands = append(commands, p.parseCommand()...)
 	}
 }
 
@@ -200,17 +341,23 @@ func Parse(tokens chan Token) []Command {
 	p := parser {
 		tokenStream: tokens,
 		rewinds: nil,
-		labels: make(map[rune]int),
 	}
-	program := p.parseCommands(nil)
+	program := p.parseCommandSequence()
+	labels := make(map[rune]int)
+	for i, command := range program {
+		switch label := command.(type) {
+		case LabelCommand:
+			labels[label.label] = i
+		}
+	}
 	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}
+		case BranchPlaceholderCommand:
+			destination, exists := labels[branch.label]
+			if !exists {
+				panic("Tried to branch to a label that doesn't exist")
+			}
+			program[i] = JumpCommand {destination}
 		}
 	}
 	return program