1 files changed, 34 insertions, 32 deletions

diff --git a/subex/main.go b/subex/main.go
index 86a8d41..32a5cf3 100644
--- a/subex/main.go
+++ b/subex/main.go
@@ -150,7 +150,7 @@ type auxiliaryState struct {
 	outputStack OutputStack
 	// How deeply nested the current execution is inside of the overall value
 	// i.e. starts at zero, is incremented to one when entering an array
-	nesting int
+	nesting []bool
 }
 
 func (aux auxiliaryState) cloneStore() auxiliaryState {
@@ -204,16 +204,6 @@ func (aux auxiliaryState) topAppendRune(runes []rune) auxiliaryState {
 	return aux
 }
 
-func (aux auxiliaryState) incNest() auxiliaryState {
-	aux.nesting++
-	return aux
-}
-
-func (aux auxiliaryState) decNest() auxiliaryState {
-	aux.nesting--
-	return aux
-}
-
 type SubexBranch struct {
 	state SubexState
 	aux   auxiliaryState
@@ -236,8 +226,15 @@ func (pair SubexEatBranch) accepting() []OutputStack {
 }
 
 func equalStates(left SubexEatBranch, right SubexEatBranch) bool {
-	// Only care about if they are the same pointer
-	return left.state == right.state && left.aux.nesting == right.aux.nesting
+	if left.state != right.state || len(left.aux.nesting) != len(right.aux.nesting) {
+		return false
+	}
+	for i, l := range left.aux.nesting {
+		if l != right.aux.nesting[i] {
+			return false
+		}
+	}
+	return true
 }
 
 // If two branches have the same state, only the first has a chance of being successful
@@ -257,11 +254,11 @@ outer:
 	return states[:uniqueStates]
 }
 
-func addStates(curStates []SubexEatBranch, newStates []SubexBranch) []SubexEatBranch {
+func addStates(curStates []SubexEatBranch, newStates []SubexBranch, nesting []bool) []SubexEatBranch {
 	for _, state := range newStates {
 		switch s := state.state.(type) {
 		case SubexEpsilonState:
-			curStates = addStates(curStates, s.epsilon(state.aux))
+			curStates = addStates(curStates, s.epsilon(state.aux), nesting)
 		case SubexEatState:
 			curStates = append(curStates, SubexEatBranch{
 				state: s,
@@ -272,14 +269,19 @@ func addStates(curStates []SubexEatBranch, newStates []SubexBranch) []SubexEatBr
 	return curStates
 }
 
-func processInput(states []SubexEatBranch, input walk.Edible, nesting int) []SubexEatBranch {
+func processInput(states []SubexEatBranch, input walk.Edible, nesting []bool) []SubexEatBranch {
 	newStates := make([]SubexEatBranch, 0, 2)
 
 	for _, state := range states {
-		// TODO: What if nesting is changed by an epsilon state?
-		if state.aux.nesting == nesting {
-			newStates = addStates(newStates, state.eat(input))
-		} else if state.aux.nesting < nesting {
+		if len(state.aux.nesting) > len(nesting) {
+			continue
+		}
+
+		if (len(state.aux.nesting) == len(nesting) &&
+				(len(state.aux.nesting) == 0 || len(nesting) == 0 ||
+				state.aux.nesting[len(nesting) - 1] || nesting[len(nesting) - 1])) {
+			newStates = addStates(newStates, state.eat(input), nesting)
+		} else {
 			newStates = append(newStates, state)
 		}
 	}
@@ -287,21 +289,21 @@ func processInput(states []SubexEatBranch, input walk.Edible, nesting int) []Sub
 	switch input := input.(type) {
 	case walk.StringValue:
 		for _, r := range input {
-			newStates = processInput(newStates, walk.RuneEdible(r), nesting+1)
+			newStates = processInput(newStates, walk.RuneEdible(r), append(nesting, true))
 		}
-		newStates = processInput(newStates, walk.StringEnd, nesting+1)
+		newStates = processInput(newStates, walk.StringEnd, append(nesting, true))
 	case walk.ArrayValue:
 		for _, el := range input {
-			newStates = processInput(newStates, walk.NumberValue(el.Index), nesting+1)
-			newStates = processInput(newStates, el.Value, nesting+1)
+			newStates = processInput(newStates, walk.NumberValue(el.Index), append(nesting, false))
+			newStates = processInput(newStates, el.Value, append(nesting, true))
 		}
-		newStates = processInput(newStates, walk.ArrayEnd, nesting+1)
+		newStates = processInput(newStates, walk.ArrayEnd, append(nesting, true))
 	case walk.MapValue:
 		for _, el := range input {
-			newStates = processInput(newStates, walk.StringValue(el.Key), nesting+1)
-			newStates = processInput(newStates, el.Value, nesting+1)
+			newStates = processInput(newStates, walk.StringValue(el.Key), append(nesting, false))
+			newStates = processInput(newStates, el.Value, append(nesting, true))
 		}
-		newStates = processInput(newStates, walk.MapEnd, nesting+1)
+		newStates = processInput(newStates, walk.MapEnd, append(nesting, true))
 	}
 
 	newStates = pruneStates(newStates)
@@ -322,20 +324,20 @@ func RunTransducer(transducer Transducer, input []walk.Value) (output []walk.Val
 				values: make([][]walk.Value, transducer.storeSize.values),
 				runes:  make([][]rune, transducer.storeSize.runes),
 			},
-			nesting: 0,
+			nesting: nil,
 		},
-	}})
+	}}, nil)
 
 	for _, value := range input {
 		if len(states) == 0 {
 			break
 		}
 
-		states = processInput(states, value, 0)
+		states = processInput(states, value, nil)
 	}
 
 	for _, state := range states {
-		if state.aux.nesting > 0 {
+		if len(state.aux.nesting) > 0 {
 			continue
 		}
 		acceptingStacks := state.accepting()