Fixes JSONWriter to work with implicit data structures

2 files changed, 321 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:
diff --git a/json/write_test.go b/json/write_test.go
new file mode 100644
index 0000000..60ad609
--- /dev/null
+++ b/json/write_test.go
@@ -0,0 +1,183 @@
+package json
+
+import (
+	"bufio"
+	"main/walk"
+	"strings"
+	"testing"
+	"encoding/json"
+)
+
+type writeTester struct {
+	writer *JSONWriter
+	output *strings.Builder
+	t *testing.T
+}
+
+func (t writeTester) write(value walk.Value, path ...interface{}) writeTester {
+	var pathValues []walk.Value
+	for _, segment := range path {
+		switch s := segment.(type) {
+		case int:
+			pathValues = append(pathValues, walk.NumberScalar(s))
+		case string:
+			pathValues = append(pathValues, walk.StringStructure(s))
+		default:
+			panic("Invalid path segment type")
+		}
+	}
+
+	t.writer.Write(walk.WalkItem {
+		Value: []walk.Value{value},
+		Path: pathValues,
+	})
+
+	return t
+}
+
+func (t writeTester) expect(expected interface{}) writeTester {
+	t.writer.AssertDone()
+	output := t.output.String()
+	var actual interface{}
+	err := json.Unmarshal([]byte(output), &actual)
+	if err != nil {
+		t.t.Log("Produced invalid JSON:")
+		t.t.Log(output)
+		t.t.FailNow()
+	}
+
+	expectedBytes, err1 := json.Marshal(expected)
+	actualBytes, err2 := json.Marshal(actual)
+
+	if err1 != nil || err2 != nil {
+		panic("Error marshalling")
+	}
+
+	expectedString := string(expectedBytes)
+	actualString := string(actualBytes)
+
+	if expectedString != actualString {
+		t.t.Log("Expected:")
+		t.t.Log(expectedString)
+		t.t.Log("Found:")
+		t.t.Log(actualString)
+		t.t.FailNow()
+	}
+
+	return t
+}
+
+func tester(t *testing.T) writeTester {
+	var output strings.Builder
+	return writeTester {
+		writer: NewJSONWriter(bufio.NewWriter(&output)),
+		output: &output,
+		t: t,
+	}
+}
+
+func TestImplicitStructures(t *testing.T) {
+	tester(t).write(
+		walk.NullScalar{},
+		0, "test", 0,
+	).expect(
+		[]interface{}{
+			map[string]interface{}{
+				"test": []interface{}{
+					nil,
+				},
+			},
+		},
+	)
+}
+
+func TestExplicitMap(t *testing.T) {
+	tester(t).write(
+		make(walk.MapStructure),
+	).write(
+		walk.NullScalar{},
+		"test",
+	).expect(
+		map[string]interface{}{
+			"test": nil,
+		},
+	)
+}
+
+func TestExplicitNested(t *testing.T) {
+	tester(t).write(
+		make(walk.MapStructure),
+	).write(
+		make(walk.MapStructure),
+		"first",
+	).write(
+		make(walk.MapStructure),
+		"first", "second",
+	).write(
+		walk.StringStructure("test"),
+		"first", "second", "third",
+	).expect(
+		map[string]interface{}{
+			"first": map[string]interface{}{
+				"second": map[string]interface{}{
+					"third": "test",
+				},
+			},
+		},
+	)
+}
+
+func TestArrayOfMaps(t *testing.T) {
+	tester(t).write(
+		walk.ArrayStructure{},
+	).write(
+		make(walk.MapStructure),
+		0,
+	).write(
+		walk.NumberScalar(0),
+		0, "number",
+	).write(
+		make(walk.MapStructure),
+		1,
+	).write(
+		walk.NumberScalar(1),
+		1, "nested", "number",
+	).write(
+		make(walk.MapStructure),
+		2,
+	).write(
+		walk.NumberScalar(2),
+		2, "number",
+	).expect(
+		[]interface{}{
+			map[string]interface{}{
+				"number": 0,
+			},
+			map[string]interface{}{
+				"nested": map[string]interface{}{
+					"number": 1,
+				},
+			},
+			map[string]interface{}{
+				"number": 2,
+			},
+		},
+	)
+}
+
+func TestStructures1(t *testing.T) {
+	tester(t).write(
+		make(walk.MapStructure),
+	).write(
+		make(walk.MapStructure),
+		"map",
+	).write(
+		walk.ArrayStructure{},
+		"array",
+	).expect(
+		map[string]interface{}{
+			"map": map[string]interface{}{},
+			"array": []interface{}{},
+		},
+	)
+}