From c18d5beac1621a8a265efb375c90cf6813011ca7 Mon Sep 17 00:00:00 2001
From: Jens Neuse <jens.neuse@gmx.de>
Date: Thu, 5 Dec 2024 20:55:20 +0100
Subject: [PATCH] feat: improve array merge

---
 update.go      | 16 ++++++++++++++++
 update_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 util.go        |  8 ++++++--
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/update.go b/update.go
index c064fde..1e58ee0 100644
--- a/update.go
+++ b/update.go
@@ -108,3 +108,19 @@ func (v *Value) SetArrayItem(idx int, value *Value) {
 	}
 	v.a[idx] = value
 }
+
+// AppendArrayItems appends items from value to the array v.
+func (v *Value) AppendArrayItems(value *Value) {
+	if v == nil || v.t != TypeArray || value == nil || value.t != TypeArray {
+		return
+	}
+	if cap(v.a) < len(v.a)+len(value.a) {
+		a := make([]*Value, len(v.a)+len(value.a))
+		copy(a, v.a)
+		for i := range v.a {
+			v.a[i] = nil
+		}
+		v.a = a[:len(v.a)]
+	}
+	v.a = append(v.a, value.a...)
+}
diff --git a/update_test.go b/update_test.go
index 4cd13c7..57aecc6 100644
--- a/update_test.go
+++ b/update_test.go
@@ -1,7 +1,10 @@
 package astjson
 
 import (
+	"encoding/json"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestObjectDelSet(t *testing.T) {
@@ -101,3 +104,49 @@ 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 BenchmarkValue_SetArrayItem(b *testing.B) {
+	input := []byte(`1`)
+	leftInput := make([]any, 2)
+	for i := 0; i < 2; i++ {
+		err := json.Unmarshal(input, &leftInput[i])
+		assert.NoError(b, err)
+	}
+	rightInput := make([]any, 1024*1024)
+	for i := 0; i < 1024*1024; i++ {
+		err := json.Unmarshal(input, &rightInput[i])
+		assert.NoError(b, err)
+	}
+
+	left, err := json.Marshal(leftInput)
+	assert.NoError(b, err)
+	right, err := json.Marshal(rightInput)
+	assert.NoError(b, err)
+
+	expectedLen := 2 + 1024*1024
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		l, _ := ParseBytesWithoutCache(left)
+		r, _ := ParseBytesWithoutCache(right)
+		out, _ := MergeValues(l, r)
+		if len(out.GetArray()) != expectedLen {
+			b.Fatalf("unexpected length; got %d; want %d", len(out.GetArray()), expectedLen)
+		}
+	}
+}
diff --git a/util.go b/util.go
index 599f470..75591c7 100644
--- a/util.go
+++ b/util.go
@@ -64,9 +64,13 @@ func MergeValues(a, b *Value) (*Value, bool) {
 	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
 		}
+		a.AppendArrayItems(b)
 		return a, false
 	case TypeFalse:
 		if b.Type() == TypeTrue {