1 files changed, 30 insertions, 30 deletions

diff --git a/subex/main.go b/subex/main.go
index f8d9093..982b585 100644
--- a/subex/main.go
+++ b/subex/main.go
@@ -150,7 +150,8 @@ 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
+	nestingLen int
+	nestingValue bool
 }
 
 func (aux auxiliaryState) cloneStore() auxiliaryState {
@@ -204,16 +205,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
@@ -237,7 +228,7 @@ 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
+	return left.state == right.state && left.aux.nestingLen == right.aux.nestingLen && left.aux.nestingValue == right.aux.nestingValue
 }
 
 // If two branches have the same state, only the first has a chance of being successful
@@ -257,12 +248,15 @@ 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:
+			if state.aux.nestingLen < len(nesting) && state.aux.nestingLen > 0 {
+				state.aux.nestingValue = nesting[state.aux.nestingLen - 1]
+			}
 			curStates = append(curStates, SubexEatBranch{
 				state: s,
 				aux:   state.aux,
@@ -272,13 +266,18 @@ 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 {
-		if state.aux.nesting == nesting {
-			newStates = addStates(newStates, state.eat(input))
-		} else if state.aux.nesting < nesting {
+		if state.aux.nestingLen > len(nesting) {
+			continue
+		}
+
+		if (state.aux.nestingLen == len(nesting) &&
+				(len(nesting) == 0 || state.aux.nestingValue || nesting[len(nesting) - 1])) {
+			newStates = addStates(newStates, state.eat(input), nesting)
+		} else {
 			newStates = append(newStates, state)
 		}
 	}
@@ -286,21 +285,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)
@@ -321,20 +320,21 @@ 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,
+			nestingLen: 0,
+			nestingValue: true,
 		},
-	}})
+	}}, 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 state.aux.nestingLen > 0 {
 			continue
 		}
 		acceptingStacks := state.accepting()