1 files changed, 47 insertions, 13 deletions

diff --git a/subex/main.go b/subex/main.go
index 2a5f5ad..4451a00 100644
--- a/subex/main.go
+++ b/subex/main.go
@@ -29,25 +29,58 @@ func CompileTransducer(transducerAst SubexAST) SubexState {
 	return transducerAst.compileWith(SubexNoneState{})
 }
 
+// An immutable stack for outputting to
+type OutputStack interface {
+	pop() ([]walk.Atom, OutputStack)
+	push(atoms []walk.Atom) OutputStack
+}
+
+type OutputStackNil struct {}
+func (stack OutputStackNil) pop() ([]walk.Atom, OutputStack) {
+	panic("Tried to pop from an empty OutputStack")
+}
+func (stack OutputStackNil) push(atoms []walk.Atom) OutputStack {
+	return &OutputStackCons {
+		head: atoms,
+		tail: stack,
+	}
+}
+
+type OutputStackCons struct {
+	head []walk.Atom
+	tail OutputStack
+}
+func (stack OutputStackCons) pop() ([]walk.Atom, OutputStack) {
+	return stack.head, stack.tail
+}
+func (stack OutputStackCons) push(atoms []walk.Atom) OutputStack {
+	return &OutputStackCons {
+		head: atoms,
+		tail: stack,
+	}
+}
+
+func topAppend(outputStack OutputStack, atoms []walk.Atom) OutputStack {
+	head, tail := outputStack.pop()
+	return tail.push(walk.ConcatData(head, atoms))
+}
+
 // One branch of subex execution
 type SubexBranch struct {
 	// Content of slots in this branch
 	store Store
 	// State in this branch
 	state SubexState
-	// Output so far in this branch
-	output []walk.Atom
+	// The output stack. At the end of the program, whatever is on top of this will be output
+	// States may push or pop to the stack as they wish, creating sort of a call stack that allows states to capture the output of other states
+	outputStack OutputStack
 }
 // Read a single character and return all the branches resulting from this branch consuming it
 func (pair SubexBranch) eat(char walk.Atom) []SubexBranch {
-	states := pair.state.eat(pair.store, char)
-	for i := range states {
-		states[i].output = walk.ConcatData(pair.output, states[i].output)
-	}
-	return states
+	return pair.state.eat(pair.store, pair.outputStack, char)
 }
-func (pair SubexBranch) accepting() [][]walk.Atom {
-	return pair.state.accepting(pair.store)
+func (pair SubexBranch) accepting() []OutputStack {
+	return pair.state.accepting(pair.store, pair.outputStack)
 }
 
 func equalStates(left SubexBranch, right SubexBranch) bool {
@@ -73,7 +106,7 @@ func pruneStates(states []SubexBranch) (newStates []SubexBranch) {
 func RunTransducer(transducer SubexState, input <-chan walk.Atom) (output []walk.Atom, err bool) {
 	states := []SubexBranch{{
 		state: transducer,
-		output: nil,
+		outputStack: OutputStackNil{}.push(nil),
 		store: make(Store),
 	}}
 	for piece := range input {
@@ -84,9 +117,10 @@ func RunTransducer(transducer SubexState, input <-chan walk.Atom) (output []walk
 		states = pruneStates(newStates)
 	}
 	for _, state := range states {
-		outputEnds := state.accepting()
-		for _, outputEnd := range outputEnds {
-			return walk.ConcatData(state.output, outputEnd), false
+		acceptingStacks := state.accepting()
+		for _, stack := range acceptingStacks {
+			output, _ := stack.pop()
+			return output, false
 		}
 	}
 	return nil, true