Fixes JSONWriter to work with implicit data structures

1 files changed, 138 insertions, 46 deletions

diff --git a/json/write.go b/json/write.go
index d024a56..9e349be 100644
--- a/json/write.go
+++ b/json/write.go
@@ -69,25 +69,40 @@ func (writer *JSONWriter) indent(level int) {
 }
 
 func (writer *JSONWriter) write(targetPath []walk.Value, value walk.Value) error {
-	diversionPoint := len(writer.path)
+	diversionPoint := 0
 	for diversionPoint < len(writer.path) && diversionPoint < len(targetPath) && segmentEqual(writer.path[diversionPoint], targetPath[diversionPoint]) {
 		diversionPoint += 1
 	}
-	
+
 	switch writer.state {
 	case JSONWriterStateBeforeValue:
 		goto beforeValue
 	case JSONWriterStateAfterValue:
 		goto afterValue
-	case JSONWriterStateInMap:
-		goto inMap
 	case JSONWriterStateInArray:
 		goto inArray
-	default:
-		panic("Invalid JSONWriterState")
+	case JSONWriterStateInMap:
+		goto inMap
 	}
 
 	beforeValue: {
+		if diversionPoint < len(writer.path) {
+			panic("Writing a value before doing a necessary leave")
+		}
+		if diversionPoint < len(targetPath) {
+			segment := targetPath[diversionPoint]
+			switch segment.(type) {
+			case walk.NumberScalar:
+				writer.writer.WriteString("[\n")
+				goto inArray
+			case walk.StringStructure:
+				writer.writer.WriteString("{\n")
+				goto inMap
+			default:
+				panic("Invalid path segment")
+			}
+		}
+
 		switch value := value.(type) {
 		case walk.NullScalar:
 			writer.writer.WriteString("null")
@@ -124,66 +139,143 @@ func (writer *JSONWriter) write(targetPath []walk.Value, value walk.Value) error
 	}
 
 	afterValue: {
-		if len(writer.path) == 0 {
-			writer.writer.WriteRune('\n')
-			goto beforeValue
-		}
-		switch writer.path[len(writer.path) - 1].(type) {
-		case walk.NumberScalar:
-			// TODO: second part of this condition might be redundant
-			if len(writer.path) - 1 <= diversionPoint && len(targetPath) >= len(writer.path) && isNumber(targetPath[len(writer.path) - 1]) {
-				writer.writer.WriteString(",\n")
+		if diversionPoint < len(writer.path) {
+			if diversionPoint == len(writer.path) - 1 && diversionPoint < len(targetPath) {
+				segment := writer.path[diversionPoint]
+				switch segment.(type) {
+				case walk.NumberScalar:
+					_, isNumber := targetPath[diversionPoint].(walk.NumberScalar)
+					if isNumber {
+						writer.writer.WriteString(",\n")
+						writer.path = writer.path[:diversionPoint]
+						goto inArray
+					}
+				case walk.StringStructure:
+					_, isString := targetPath[diversionPoint].(walk.StringStructure)
+					if isString {
+						writer.writer.WriteString(",\n")
+						writer.path = writer.path[:diversionPoint]
+						goto inMap
+					}
+				default:
+					panic("Invalid segment type")
+				}
+			}
+
+			writer.writer.WriteString("\n")
+			switch writer.path[len(writer.path) - 1].(type) {
+			case walk.NumberScalar:
 				writer.path = writer.path[:len(writer.path) - 1]
 				goto inArray
-			} else {
-				writer.writer.WriteString("\n")
-				writer.indent(len(writer.path) - 1)
-				writer.writer.WriteString("]")
-				writer.path = writer.path[:len(writer.path) - 1]
-				goto afterValue
-			}
-		case walk.StringStructure:
-			if len(writer.path) -1 <= diversionPoint && len(targetPath) >= len(writer.path) && isString(targetPath[len(writer.path) - 1]) {
-				writer.writer.WriteString(",\n")
+			case walk.StringStructure:
 				writer.path = writer.path[:len(writer.path) - 1]
 				goto inMap
-			} else {
+			default:
+				panic("Invalid segment type")
+			}
+		}
+
+		// TODO: handle len(writer.path) == 0
+		if diversionPoint < len(targetPath) {
+			if len(writer.path) == 0 {
 				writer.writer.WriteString("\n")
-				writer.indent(len(writer.path) - 1)
-				writer.writer.WriteString("}")
-				writer.path = writer.path[:len(writer.path) - 1]
-				goto afterValue
+				goto beforeValue
+			}
+			segment := writer.path[diversionPoint - 1]
+			writer.writer.WriteString(",")
+			diversionPoint--
+			writer.path = writer.path[:diversionPoint]
+			switch segment.(type) {
+			case walk.NumberScalar:
+				goto inArray
+			case walk.StringStructure:
+				goto inMap
+			default:
+				panic("Invalid segment type")
 			}
+		}
+
+		if len(writer.path) == 0 {
+			writer.writer.WriteString("\n")
+			goto beforeValue
+		}
+		segment := writer.path[diversionPoint - 1]
+		writer.writer.WriteString(",\n")
+		diversionPoint--
+		writer.path = writer.path[:diversionPoint]
+		switch segment.(type) {
+		case walk.NumberScalar:
+			goto inArray
+		case walk.StringStructure:
+			goto inMap
 		default:
-			panic("Invalid path segment type")
+			panic("Invalid segment type")
 		}
 	}
 
-	inMap: {
-		if len(writer.path) <= diversionPoint && len(targetPath) > len(writer.path) && isString(targetPath[len(writer.path)]) {
-			writer.indent(len(writer.path) + 1)
-			writer.writer.WriteString(fmt.Sprintf("%q: ", targetPath[len(writer.path)].(walk.StringStructure)))
-			writer.path = append(writer.path, targetPath[len(writer.path)].(walk.StringStructure))
-			goto beforeValue
-		} else {
-			writer.writer.WriteString("\n}")
+	inArray: {
+		if diversionPoint < len(writer.path) {
+			writer.indent(len(writer.path))
+			writer.writer.WriteString("]")
 			goto afterValue
 		}
+
+		if diversionPoint < len(targetPath) {
+			switch s := targetPath[diversionPoint].(type) {
+			case walk.NumberScalar:
+				writer.path = append(writer.path, s)
+				diversionPoint++
+				writer.indent(len(writer.path))
+				goto beforeValue
+			case walk.StringStructure:
+				writer.indent(len(writer.path))
+				writer.writer.WriteString("]")
+				goto afterValue
+			default:
+				panic("Invalid segment type")
+			}
+		}
+
+		writer.indent(len(writer.path))
+		writer.writer.WriteString("]")
+		goto afterValue
 	}
 
-	inArray: {
-		if len(writer.path) <= diversionPoint && len(targetPath) > len(writer.path) && isNumber(targetPath[len(writer.path)]) {
-			writer.indent(len(writer.path) + 1)
-			writer.path = append(writer.path, walk.NumberScalar(0))
-			goto beforeValue
-		} else {
-			writer.writer.WriteString("\n]")
+	inMap: {
+		if diversionPoint < len(writer.path) {
+			writer.indent(len(writer.path))
+			writer.writer.WriteString("}")
 			goto afterValue
 		}
+
+		if diversionPoint < len(targetPath) {
+			switch s := targetPath[diversionPoint].(type) {
+			case walk.NumberScalar:
+				writer.indent(len(writer.path))
+				writer.writer.WriteString("}")
+				goto afterValue
+			case walk.StringStructure:
+				writer.path = append(writer.path, s)
+				diversionPoint++
+				writer.indent(len(writer.path))
+				writer.writer.WriteString(fmt.Sprintf("%q: ", s))
+				goto beforeValue
+			}
+		}
+
+		writer.indent(len(writer.path))
+		writer.writer.WriteString("}")
+		goto afterValue
 	}
 }
 
 func (writer *JSONWriter) AssertDone() {
+	switch writer.state {
+	case JSONWriterStateInArray:
+		writer.writer.WriteString("]")
+	case JSONWriterStateInMap:
+		writer.writer.WriteString("}")
+	}
 	for i := len(writer.path) - 1; i >= 0; i -= 1 {
 		switch writer.path[i].(type) {
 		case walk.NumberScalar: