From 70309c9908ba9b90642174c8941150c0be20539e Mon Sep 17 00:00:00 2001
From: Changyu Moon <121847433+window9u@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:28:20 +0800
Subject: [PATCH] Implement `InitialRoot` option for Document attachment (#986)

This commit allows setting initial values for documents with a
predefined structure. It checks if the document has been set up with
initial values and applies them if not, instead of checking for an
empty or initial state.

The attach function now accepts an initialRoot option, which can be
partially applied. Existing keys in the Document are preserved, while
new keys from initialRoot are added. This implementation applies
initialRoot locally after pushpull during Attach.

Usage example:

doc := document.New("key-1")
client1.Attach(ctx, doc, client.WithInitialRoot(map[string]any{
    "k": json.NewCounter(0, crdt.LongCnt),
}))
doc.Update(func(root *json.Object, p *presence.Presence) error {
    root.GetCounter("k").Increase(1)
    return nil
})

Note:
- No type checking is performed for initialRoot elements.
- Concurrent document creation may lead to overwrites due to LWW.
---
 client/client.go                  |  11 ++
 client/options.go                 |  14 +-
 pkg/document/json/object.go       |   7 +
 test/integration/document_test.go | 297 ++++++++++++++++++++++++++++++
 4 files changed, 327 insertions(+), 2 deletions(-)

diff --git a/client/client.go b/client/client.go
index 19bc29549..3f2f082b0 100644
--- a/client/client.go
+++ b/client/client.go
@@ -344,6 +344,17 @@ func (c *Client) Attach(ctx context.Context, doc *document.Document, options ...
 		}
 	}
 
+	if err = doc.Update(func(root *json.Object, p *presence.Presence) error {
+		for k, v := range opts.InitialRoot {
+			if root.Get(k) == nil {
+				root.SetDynamicValue(k, v)
+			}
+		}
+		return nil
+	}); err != nil {
+		return err
+	}
+
 	return nil
 }
 
diff --git a/client/options.go b/client/options.go
index 739264c7d..fcb03b6a9 100644
--- a/client/options.go
+++ b/client/options.go
@@ -92,8 +92,9 @@ type AttachOption func(*AttachOptions)
 // AttachOptions configures how we set up the document.
 type AttachOptions struct {
 	// Presence is the presence of the client.
-	Presence innerpresence.Presence
-	IsManual bool
+	Presence    innerpresence.Presence
+	InitialRoot map[string]any
+	IsManual    bool
 }
 
 // WithPresence configures the presence of the client.
@@ -101,6 +102,15 @@ func WithPresence(presence innerpresence.Presence) AttachOption {
 	return func(o *AttachOptions) { o.Presence = presence }
 }
 
+// WithInitialRoot sets the initial root of the document. Values in the initial
+// root will be discarded if the key already exists in the document. If some
+// keys are not in the document, they will be added.
+func WithInitialRoot(root map[string]any) AttachOption {
+	return func(o *AttachOptions) {
+		o.InitialRoot = root
+	}
+}
+
 // WithManualSync configures the manual sync of the client.
 func WithManualSync() AttachOption {
 	return func(o *AttachOptions) { o.IsManual = true }
diff --git a/pkg/document/json/object.go b/pkg/document/json/object.go
index de31d8da4..24398abd9 100644
--- a/pkg/document/json/object.go
+++ b/pkg/document/json/object.go
@@ -43,6 +43,13 @@ func NewObject(ctx *change.Context, root *crdt.Object) *Object {
 	}
 }
 
+// SetDynamicValue sets a dynamic value for the given key.
+func (p *Object) SetDynamicValue(k string, v any) {
+	p.setInternal(k, func(ticket *time.Ticket) crdt.Element {
+		return toElement(p.context, buildCRDTElement(p.context, v, ticket, newBuildState()))
+	})
+}
+
 // SetNewObject sets a new Object for the given key.
 func (p *Object) SetNewObject(k string, v ...any) *Object {
 	value := p.setInternal(k, func(ticket *time.Ticket) crdt.Element {
diff --git a/test/integration/document_test.go b/test/integration/document_test.go
index 2451723d7..3b82dfe03 100644
--- a/test/integration/document_test.go
+++ b/test/integration/document_test.go
@@ -32,6 +32,7 @@ import (
 
 	"github.com/yorkie-team/yorkie/client"
 	"github.com/yorkie-team/yorkie/pkg/document"
+	"github.com/yorkie-team/yorkie/pkg/document/crdt"
 	"github.com/yorkie-team/yorkie/pkg/document/innerpresence"
 	"github.com/yorkie-team/yorkie/pkg/document/json"
 	"github.com/yorkie-team/yorkie/pkg/document/presence"
@@ -831,3 +832,299 @@ func TestDocumentWithProjects(t *testing.T) {
 		assert.NotEqual(t, 0, len(docs[0].Snapshot))
 	})
 }
+
+func TestDocumentWithInitialRoot(t *testing.T) {
+	clients := activeClients(t, 3)
+	c1, c2, c3 := clients[0], clients[1], clients[2]
+	defer deactivateAndCloseClients(t, clients)
+
+	t.Run("attach with InitialRoot test", func(t *testing.T) {
+		ctx := context.Background()
+		doc1 := document.New(helper.TestDocKey(t))
+
+		// 01. attach and initialize document
+		assert.NoError(t, c1.Attach(ctx, doc1, client.WithInitialRoot(map[string]any{
+			"counter": json.NewCounter(0, crdt.LongCnt),
+			"content": map[string]any{
+				"x": 1,
+				"y": 1,
+			},
+		})))
+		assert.True(t, doc1.IsAttached())
+		assert.Equal(t, `{"content":{"x":1,"y":1},"counter":0}`, doc1.Marshal())
+		assert.NoError(t, c1.Sync(ctx))
+
+		// 02. attach and initialize document with new fields and if key already exists, it will be discarded
+		doc2 := document.New(helper.TestDocKey(t))
+		assert.NoError(t, c2.Attach(ctx, doc2, client.WithInitialRoot(map[string]any{
+			"counter": json.NewCounter(1, crdt.LongCnt),
+			"content": map[string]any{
+				"x": 2,
+				"y": 2,
+			},
+			"new": map[string]any{
+				"k": "v",
+			},
+		})))
+		assert.True(t, doc2.IsAttached())
+		assert.Equal(t, `{"content":{"x":1,"y":1},"counter":0,"new":{"k":"v"}}`, doc2.Marshal())
+	})
+
+	t.Run("attach with InitialRoot after key deletion test", func(t *testing.T) {
+		ctx := context.Background()
+		doc1 := document.New(helper.TestDocKey(t))
+
+		// 01. client1 attach with initialRoot
+		assert.NoError(t, c1.Attach(ctx, doc1, client.WithInitialRoot(map[string]any{
+			"counter": json.NewCounter(1, crdt.LongCnt),
+			"content": map[string]any{
+				"x": 1,
+				"y": 1,
+			},
+		})))
+		assert.True(t, doc1.IsAttached())
+		assert.Equal(t, `{"content":{"x":1,"y":1},"counter":1}`, doc1.Marshal())
+		assert.NoError(t, c1.Sync(ctx))
+
+		// 02. client2 attach with initialRoot and delete elements
+		doc2 := document.New(helper.TestDocKey(t))
+		assert.NoError(t, c2.Attach(ctx, doc2))
+		assert.True(t, doc2.IsAttached())
+		assert.NoError(t, doc2.Update(func(root *json.Object, p *presence.Presence) error {
+			root.Delete("content")
+			root.Delete("counter")
+			return nil
+		}))
+		assert.NoError(t, c2.Sync(ctx))
+
+		// 03. client3 attach with initialRoot and delete elements
+		doc3 := document.New(helper.TestDocKey(t))
+		assert.NoError(t, c3.Attach(ctx, doc3, client.WithInitialRoot(map[string]any{
+			"counter": json.NewCounter(3, crdt.LongCnt),
+			"content": map[string]any{
+				"x": 3,
+				"y": 3,
+			},
+		})))
+		assert.True(t, doc3.IsAttached())
+		assert.Equal(t, `{"content":{"x":3,"y":3},"counter":3}`, doc3.Marshal())
+	})
+
+	t.Run("concurrent attach with InitialRoot test", func(t *testing.T) {
+		ctx := context.Background()
+		doc1 := document.New(helper.TestDocKey(t))
+
+		// 01. user1 attach with initialRoot and client doesn't sync
+		assert.NoError(t, c1.Attach(ctx, doc1, client.WithInitialRoot(map[string]any{
+			"first_writer": "user1",
+		})))
+		assert.True(t, doc1.IsAttached())
+		assert.Equal(t, `{"first_writer":"user1"}`, doc1.Marshal())
+
+		// 02. user2 attach with initialRoot and client doesn't sync
+		doc2 := document.New(helper.TestDocKey(t))
+		assert.NoError(t, c2.Attach(ctx, doc2, client.WithInitialRoot(map[string]any{
+			"first_writer": "user2",
+		})))
+		assert.True(t, doc2.IsAttached())
+		assert.Equal(t, `{"first_writer":"user2"}`, doc2.Marshal())
+
+		// 03. user1 sync first and user2 seconds
+		assert.NoError(t, c1.Sync(ctx))
+		assert.NoError(t, c2.Sync(ctx))
+
+		// 04. user1's local document's first_writer was user1
+		assert.Equal(t, `{"first_writer":"user1"}`, doc1.Marshal())
+		assert.Equal(t, `{"first_writer":"user2"}`, doc2.Marshal())
+
+		// 05. user1's local document's first_writer is overwritten by user2
+		assert.NoError(t, c1.Sync(ctx))
+		assert.Equal(t, `{"first_writer":"user2"}`, doc1.Marshal())
+	})
+
+	t.Run("attach with InitialRoot by same key test", func(t *testing.T) {
+		ctx := context.Background()
+		doc := document.New(helper.TestDocKey(t))
+
+		k1 := "key"
+		k2 := "key"
+		k3 := "key"
+		k4 := "key"
+		k5 := "key"
+		assert.NoError(t, c1.Attach(ctx, doc, client.WithInitialRoot(map[string]any{
+			k1: 1,
+			k2: 2,
+			k3: 3,
+			k4: 4,
+			k5: 5,
+		})))
+		assert.True(t, doc.IsAttached())
+		// The last value is used when the same key is used.
+		assert.Equal(t, `{"key":5}`, doc.Marshal())
+	})
+
+	t.Run("attach with InitialRoot conflict type test", func(t *testing.T) {
+		ctx := context.Background()
+		doc1 := document.New(helper.TestDocKey(t))
+
+		// 01. attach with initialRoot and set counter
+		assert.NoError(t, c1.Attach(ctx, doc1, client.WithInitialRoot(map[string]any{
+			"k": json.NewCounter(1, crdt.LongCnt),
+		})))
+		assert.True(t, doc1.IsAttached())
+		assert.NoError(t, c1.Sync(ctx))
+
+		// 02. attach with initialRoot and set text
+		doc2 := document.New(helper.TestDocKey(t))
+		assert.NoError(t, c2.Attach(ctx, doc2, client.WithInitialRoot(map[string]any{
+			"k": json.NewText(),
+		})))
+		assert.True(t, doc2.IsAttached())
+		assert.NoError(t, c2.Sync(ctx))
+
+		// 03. client2 try to update counter
+		assert.Panics(t, func() { doc2.Root().GetText("k").Edit(0, 1, "a") })
+	})
+
+	t.Run("attach with initialRoot support type test", func(t *testing.T) {
+		type (
+			Myint    int
+			MyStruct struct {
+				M Myint
+			}
+			t1 struct {
+				M string
+			}
+			T1 struct {
+				M string
+			}
+			T2 struct {
+				T1
+				t1
+				M string
+			}
+		)
+		ctx := context.Background()
+		nowTime := time.Now()
+		tests := []struct {
+			caseName     string
+			input        any
+			expectedJSON string
+			expectPanic  bool
+		}{
+			// supported primitive types
+			{"nil", nil, `{"k":null}`, false},
+			{"int", 1, `{"k":1}`, false},
+			{"int32", int32(1), `{"k":1}`, false},
+			{"int64", int64(1), `{"k":1}`, false},
+			{"float32", float32(1.1), `{"k":1.100000}`, false},
+			{"float64", 1.1, `{"k":1.100000}`, false},
+			{"string", "hello", `{"k":"hello"}`, false},
+			{"bool", true, `{"k":true}`, false},
+			{"time", nowTime, `{"k":"` + nowTime.Format(time.RFC3339) + `"}`, false},
+			{"Myint", Myint(1), `{"k":1}`, false},
+
+			// unsupported primitive types
+			{"int8", int8(1), `{}`, true},
+			{"int16", int16(1), `{}`, true},
+			{"uint32", uint32(1), `{}`, true},
+			{"uint64", uint64(1), `{}`, true},
+
+			// supported slice, array types
+			{"int slice", []int{1, 2, 3}, `{"k":[1,2,3]}`, false},
+			{"&int slice", &[]int{1, 2, 3}, `{"k":[1,2,3]}`, false},
+			{"any slice", []any{nil, 1, 1.0, "hello", true, nowTime, []int{1, 2, 3}}, `{"k":[null,1,1.000000,"hello",true,"` + nowTime.Format(time.RFC3339) + `",[1,2,3]]}`, false},
+			{"&any slice", &[]any{nil, 1, 1.0, "hello", true, nowTime, []int{1, 2, 3}}, `{"k":[null,1,1.000000,"hello",true,"` + nowTime.Format(time.RFC3339) + `",[1,2,3]]}`, false},
+			{"int array", [3]int{1, 2, 3}, `{"k":[1,2,3]}`, false},
+			{"&int array", &[3]int{1, 2, 3}, `{"k":[1,2,3]}`, false},
+			{"string array", [3]string{"a", "b", "c"}, `{"k":["a","b","c"]}`, false},
+			{"&string array", &[3]string{"a", "b", "c"}, `{"k":["a","b","c"]}`, false},
+			{"any array", [7]any{nil, 1, 1.0, "hello", true, nowTime, []int{1, 2, 3}}, `{"k":[null,1,1.000000,"hello",true,"` + nowTime.Format(time.RFC3339) + `",[1,2,3]]}`, false},
+			{"&any array", &[7]any{nil, 1, 1.0, "hello", true, nowTime, []int{1, 2, 3}}, `{"k":[null,1,1.000000,"hello",true,"` + nowTime.Format(time.RFC3339) + `",[1,2,3]]}`, false},
+
+			// supported map types
+			{"string:any map", map[string]any{"a": nil, "b": 1, "c": 1.0, "d": "hello", "e": true, "f": nowTime, "g": []int{1, 2, 3}}, `{"k":{"a":null,"b":1,"c":1.000000,"d":"hello","e":true,"f":"` + nowTime.Format(time.RFC3339) + `","g":[1,2,3]}}`, false},
+			{"&string:any map", &map[string]any{"a": nil, "b": 1, "c": 1.0, "d": "hello", "e": true, "f": nowTime, "g": []int{1, 2, 3}}, `{"k":{"a":null,"b":1,"c":1.000000,"d":"hello","e":true,"f":"` + nowTime.Format(time.RFC3339) + `","g":[1,2,3]}}`, false},
+
+			// unsupported map types
+			{"int map", map[int]int{1: 1, 2: 2}, `{}`, true},
+			{"string map", map[string]string{"a": "a", "b": "b"}, `{}`, true},
+			{"int map", map[int]any{1: 1, 2: 2}, `{}`, true},
+
+			// supported JSON types
+			{"json.Text", json.NewText(), `{"k":[]}`, false},
+			{"*json.Text", *json.NewText(), `{"k":[]}`, false},
+			{"json.Tree", json.NewTree(&json.TreeNode{ // 1: tree
+				Type: "doc",
+				Children: []json.TreeNode{{
+					Type: "p", Children: []json.TreeNode{{Type: "text", Value: "ab"}},
+				}},
+			}), `{"k":{"type":"doc","children":[{"type":"p","children":[{"type":"text","value":"ab"}]}]}}`, false},
+			{"*json.Tree", *json.NewTree(&json.TreeNode{
+				Type: "doc",
+				Children: []json.TreeNode{{
+					Type: "p", Children: []json.TreeNode{{Type: "text", Value: "ab"}},
+				}},
+			}), `{"k":{"type":"doc","children":[{"type":"p","children":[{"type":"text","value":"ab"}]}]}}`, false},
+			{"json.Counter", json.NewCounter(1, crdt.LongCnt), `{"k":1}`, false},
+			{"*json.Counter", *json.NewCounter(1, crdt.LongCnt), `{"k":1}`, false},
+
+			// struct types
+			{"struct", MyStruct{M: 1}, `{"k":{"M":1}}`, false},
+			{"struct pointer", &MyStruct{M: 1}, `{"k":{"M":1}}`, false},
+			{"struct slice", []MyStruct{{M: 1}, {M: 2}}, `{"k":[{"M":1},{"M":2}]}`, false},
+			{"struct array", [2]MyStruct{{M: 1}, {M: 2}}, `{"k":[{"M":1},{"M":2}]}`, false},
+			{"anonymous struct", struct{ M string }{M: "hello"}, `{"k":{"M":"hello"}}`, false},
+			{"anonymous struct pointer", &struct{ M string }{M: "hello"}, `{"k":{"M":"hello"}}`, false},
+			{"anonymous struct slice", []struct{ M string }{{M: "a"}, {M: "b"}}, `{"k":[{"M":"a"},{"M":"b"}]}`, false},
+			{"anonymous struct array", [2]struct{ M string }{{M: "a"}, {M: "b"}}, `{"k":[{"M":"a"},{"M":"b"}]}`, false},
+			{"struct with embedded struct", T2{T1: T1{M: "a"}, t1: t1{M: "b"}, M: "c"}, `{"k":{"M":"c","T1":{"M":"a"}}}`, false},
+			{"strut with unexported field", struct {
+				t int
+				s string
+			}{t: 1, s: "hello"}, `{"k":{}}`, false},
+			{"strut with unexported field pointer", &struct {
+				t int
+				s string
+			}{t: 1, s: "hello"}, `{"k":{}}`, false},
+			{"struct with slice", struct{ M []int }{M: []int{1, 2, 3}}, `{"k":{"M":[1,2,3]}}`, false},
+			{"struct with slice pointer", &struct{ M []int }{M: []int{1, 2, 3}}, `{"k":{"M":[1,2,3]}}`, false},
+			{"struct with array", struct{ M [3]int }{M: [3]int{1, 2, 3}}, `{"k":{"M":[1,2,3]}}`, false},
+			{"struct with array pointer", &struct{ M [3]int }{M: [3]int{1, 2, 3}}, `{"k":{"M":[1,2,3]}}`, false},
+			{"struct with struct", struct{ M MyStruct }{M: MyStruct{M: 1}}, `{"k":{"M":{"M":1}}}`, false},
+			{"struct with struct pointer", &struct{ M MyStruct }{M: MyStruct{M: 1}}, `{"k":{"M":{"M":1}}}`, false},
+			{"struct with json types", struct {
+				T    json.Text
+				C    json.Counter
+				Tree json.Tree
+			}{T: *json.NewText(), C: *json.NewCounter(1, crdt.LongCnt), Tree: *json.NewTree(&json.TreeNode{
+				Type: "doc",
+				Children: []json.TreeNode{{
+					Type: "p", Children: []json.TreeNode{{Type: "text", Value: "ab"}},
+				}},
+			})}, `{"k":{"C":1,"T":[],"Tree":{"type":"doc","children":[{"type":"p","children":[{"type":"text","value":"ab"}]}]}}}`, false},
+
+			// unsupported struct types
+			{"struct with unsupported map", struct{ M map[string]int }{M: map[string]int{"a": 1, "b": 2}}, `{}`, true},
+			{"struct with unsupported primitive type", struct{ M int8 }{M: 1}, `{}`, true},
+
+			{"func", func(a int, b int) int { return a + b }, `{}`, true},
+		}
+		for _, tt := range tests {
+			t.Run(tt.caseName, func(t *testing.T) {
+				doc := document.New(helper.TestDocKey(t))
+				val := func() {
+					assert.NoError(t, c1.Attach(ctx, doc, client.WithInitialRoot(map[string]any{
+						"k": tt.input,
+					})))
+				}
+				if tt.expectPanic {
+					assert.PanicsWithValue(t, "unsupported type", val)
+				} else {
+					assert.NotPanics(t, val)
+				}
+				assert.Equal(t, tt.expectedJSON, doc.Marshal())
+			})
+		}
+	})
+}