fix: properly merge array values (#8)

Author: jensneuse

update_test.go

@@ -1,7 +1,12 @@
 package astjson
 
 import (
+	"encoding/json"
+	"fmt"
+	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestObjectDelSet(t *testing.T) {
@@ -101,3 +106,109 @@ func TestValueDelSet(t *testing.T) {
 	v.Set("x", MustParse(`[]`))
 	v.SetArrayItem(1, MustParse(`[]`))
 }
+
+func TestValue_AppendArrayItems(t *testing.T) {
+	left := MustParse(`[1,2,3]`)
+	right := MustParse(`[4,5,6]`)
+	left.AppendArrayItems(right)
+	if len(left.GetArray()) != 6 {
+		t.Fatalf("unexpected length; got %d; want %d", len(left.GetArray()), 6)
+	}
+	out := left.MarshalTo(nil)
+	if string(out) != `[1,2,3,4,5,6]` {
+		t.Fatalf("unexpected output; got %q; want %q", out, `[1,2,3,4,5,6]`)
+	}
+}
+
+func TestMergeWithSwap(t *testing.T) {
+	left := MustParse(`{"a":{"b":1,"c":2,"e":[],"f":[1],"h":[1,2,3]}}`)
+	right := MustParse(`{"a":{"b":2,"d":3,"e":[1,2,3],"g":[1],"h":[4,5,6]}}`)
+	out, _ := MergeValues(left, right)
+	assert.Equal(t, `{"a":{"b":2,"c":2,"e":[1,2,3],"f":[1],"h":[4,5,6],"d":3,"g":[1]}}`, out.String())
+}
+
+type RootObject struct {
+	Child *ChildObject `json:"child"`
+}
+
+type ChildObject struct {
+	GrandChild *GrandChildObject `json:"grand_child"`
+}
+
+type GrandChildObject struct {
+	Items []string `json:"items"`
+}
+
+func BenchmarkValue_SetArrayItem(b *testing.B) {
+
+	root := &RootObject{
+		Child: &ChildObject{
+			GrandChild: &GrandChildObject{
+				Items: make([]string, 0),
+			},
+		},
+	}
+
+	l, err := json.Marshal(root)
+	assert.NoError(b, err)
+
+	root.Child.GrandChild.Items = make([]string, 1024*1024)
+
+	for i := 0; i < 1024*1024; i++ {
+		root.Child.GrandChild.Items[i] = strings.Repeat("a", 1024)
+	}
+
+	r, err := json.Marshal(root)
+	assert.NoError(b, err)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		l, _ := ParseBytesWithoutCache(l)
+		r, _ := ParseBytesWithoutCache(r)
+		out, _ := MergeValues(l, r)
+		arr := out.GetArray("child", "grand_child", "items")
+		assert.Len(b, arr, 1024*1024)
+	}
+}
+
+func BenchmarkMergeValuesWithATonOfRecursion(b *testing.B) {
+	b.ReportAllocs()
+
+	left := MustParse(`{"a":{}}`)
+	str := fmt.Sprintf(
+		`{"ba":{"bb":{"bc":{"bd":[%s, %s, %s], "be": {"bf": %s}}}}}`,
+		objectWithRecursion(10),
+		objectWithRecursion(20),
+		objectWithRecursion(2),
+		objectWithRecursion(3))
+
+	expected := fmt.Sprintf(
+		`{"a":{},"ba":{"bb":{"bc":{"bd":[%s,%s,%s],"be":{"bf":%s}}}}}`,
+		objectWithRecursion(10),
+		objectWithRecursion(20),
+		objectWithRecursion(2),
+		objectWithRecursion(3))
+
+	_ = expected
+
+	right := MustParse(str)
+
+	for i := 0; i < b.N; i++ {
+		_, _ = MergeValues(left, right)
+		/*v, _ := MergeValues(left, right)
+		if v.String() != expected {
+			assert.Equal(b, expected, v.String())
+		}*/
+	}
+
+	fmt.Printf("left: %s\n", left.String())
+}
+
+func objectWithRecursion(depth int) string {
+	if depth == 0 {
+		return `{}`
+	}
+	return `{"a":` + objectWithRecursion(depth-1) + `}`
+}

util.go

@@ -42,30 +42,39 @@ func MergeValues(a, b *Value) (*Value, bool) {
 	case TypeObject:
 		ao, _ := a.Object()
 		bo, _ := b.Object()
-		ao.Visit(func(key []byte, l *Value) {
-			sKey := b2s(key)
-			r := bo.Get(sKey)
-			if r == nil {
-				return
+		ao.unescapeKeys()
+		bo.unescapeKeys()
+		for i := range bo.kvs {
+			k := bo.kvs[i].k
+			r := bo.kvs[i].v
+			l := ao.Get(k)
+			if l == nil {
+				ao.Set(k, r)
+				continue
 			}
-			merged, changed := MergeValues(l, r)
+			n, changed := MergeValues(l, r)
 			if changed {
-				ao.Set(b2s(key), merged)
+				ao.Set(k, n)
 			}
-		})
-		bo.Visit(func(key []byte, r *Value) {
-			sKey := b2s(key)
-			if ao.Get(sKey) != nil {
-				return
-			}
-			ao.Set(sKey, r)
-		})
+		}
 		return a, false
 	case TypeArray:
 		aa, _ := a.Array()
 		ba, _ := b.Array()
-		for i := 0; i < len(ba); i++ {
-			a.SetArrayItem(len(aa)+i, ba[i])
+		if len(aa) == 0 {
+			return b, true
+		}
+		if len(ba) == 0 {
+			return a, false
+		}
+		if len(aa) != len(ba) {
+			return b, true
+		}
+		for i := range aa {
+			n, changed := MergeValues(aa[i], ba[i])
+			if changed {
+				aa[i] = n
+			}
 		}
 		return a, false
 	case TypeFalse:
@@ -151,6 +160,13 @@ func ValueIsNonNull(v *Value) bool {
 	return true
 }
 
+func (v *Value) AppendArrayItems(right *Value) {
+	if v.t != TypeArray || right.t != TypeArray {
+		return
+	}
+	v.a = append(v.a, right.a...)
+}
+
 func ValueIsNull(v *Value) bool {
 	return !ValueIsNonNull(v)
 }

util_test.go

@@ -46,7 +46,15 @@ func TestMergeValuesArray(t *testing.T) {
 	merged, changed := MergeValues(a, b)
 	require.Equal(t, false, changed)
 	out := merged.MarshalTo(nil)
-	require.Equal(t, `[1,2,3,4]`, string(out))
+	require.Equal(t, `[3,4]`, string(out))
+}
+
+func TestMergeObjectValuesArray(t *testing.T) {
+	a, b := MustParse(`[{"a":1,"b":2},{"x":1}]`), MustParse(`[{"a":2,"b":3,"c":4},{"y":1}]`)
+	merged, changed := MergeValues(a, b)
+	require.Equal(t, false, changed)
+	out := merged.MarshalTo(nil)
+	require.Equal(t, `[{"a":2,"b":3,"c":4},{"x":1,"y":1}]`, string(out))
 }
 
 func TestMergeValuesNestedObjects(t *testing.T) {