Create NestedSet to populate set branches along with the leaves

Author: apelisse

fieldpath/set.go

@@ -19,6 +19,8 @@ package fieldpath
 import (
 	"sort"
 	"strings"
+
+	"sigs.k8s.io/structured-merge-diff/v4/schema"
 )
 
 // Set identifies a set of fields.
@@ -110,6 +112,30 @@ func (s *Set) RecursiveDifference(s2 *Set) *Set {
 	}
 }
 
+// EnsureNamedFieldsAreMembers returns a Set that contains all the
+// fields in s, as well as all the named fields that are typically not
+// included. For example, a set made of "a.b.c" will end-up also owning
+// "a" if it's a named fields but not "a.b" if it's a map.
+func (s *Set) EnsureNamedFieldsAreMembers(sc *schema.Schema, tr schema.TypeRef) *Set {
+	members := PathElementSet{
+		members: make(sortedPathElements, 0, s.Members.Size()+len(s.Children.members)),
+	}
+	atom, _ := sc.Resolve(tr)
+	members.members = append(members.members, s.Members.members...)
+	for _, node := range s.Children.members {
+		// Only insert named fields.
+		if node.pathElement.FieldName != nil && atom.Map != nil {
+			if _, has := atom.Map.FindField(*node.pathElement.FieldName); has {
+				members.Insert(node.pathElement)
+			}
+		}
+	}
+	return &Set{
+		Members:  members,
+		Children: *s.Children.EnsureNamedFieldsAreMembers(sc, tr),
+	}
+}
+
 // Size returns the number of members of the set.
 func (s *Set) Size() int {
 	return s.Members.Size() + s.Children.Size()
@@ -391,6 +417,31 @@ func (s *SetNodeMap) RecursiveDifference(s2 *Set) *SetNodeMap {
 	return out
 }
 
+// EnsureNamedFieldsAreMembers returns a set that contains all the named fields along with the leaves.
+func (s *SetNodeMap) EnsureNamedFieldsAreMembers(sc *schema.Schema, tr schema.TypeRef) *SetNodeMap {
+	out := make(sortedSetNode, 0, s.Size())
+	atom, _ := sc.Resolve(tr)
+	for _, member := range s.members {
+		tr := schema.TypeRef{}
+		if member.pathElement.FieldName != nil && atom.Map != nil {
+			tr = atom.Map.ElementType
+			if sf, ok := atom.Map.FindField(*member.pathElement.FieldName); ok {
+				tr = sf.Type
+			}
+		} else if member.pathElement.Key != nil && atom.List != nil {
+			tr = atom.List.ElementType
+		}
+		out = append(out, setNode{
+			pathElement: member.pathElement,
+			set:         member.set.EnsureNamedFieldsAreMembers(sc, tr),
+		})
+	}
+
+	return &SetNodeMap{
+		members: out,
+	}
+}
+
 // Iterate calls f for each PathElement in the set.
 func (s *SetNodeMap) Iterate(f func(PathElement)) {
 	for _, n := range s.members {

fieldpath/set_test.go

@@ -21,6 +21,9 @@ import (
 	"fmt"
 	"math/rand"
 	"testing"
+
+	"gopkg.in/yaml.v2"
+	"sigs.k8s.io/structured-merge-diff/v4/schema"
 )
 
 type randomPathAlphabet []PathElement
@@ -541,6 +544,86 @@ func TestSetDifference(t *testing.T) {
 	}
 }
 
+var nestedSchema = func() (*schema.Schema, schema.TypeRef) {
+	sc := &schema.Schema{}
+	name := "type"
+	err := yaml.Unmarshal([]byte(`types:
+- name: type
+  map:
+    elementType:
+      namedType: type
+    fields:
+      - name: named
+        type:
+          namedType: type
+      - name: list
+        type:
+          list:
+            elementRelationShip: associative
+            keys: ["name"]
+            elementType:
+              namedType: type
+      - name: value
+        type:
+          scalar: numeric
+`), &sc)
+	if err != nil {
+		panic(err)
+	}
+	return sc, schema.TypeRef{NamedType: &name}
+}
+
+var _P = MakePathOrDie
+
+func TestEnsureNamedFieldsAreMembers(t *testing.T) {
+	table := []struct {
+		set, expected *Set
+	}{
+		{
+			set: NewSet(_P("named", "named", "value")),
+			expected: NewSet(
+				_P("named", "named", "value"),
+				_P("named", "named"),
+				_P("named"),
+			),
+		},
+		{
+			set: NewSet(_P("named", "a", "named", "value"), _P("a", "named", "value"), _P("a", "b", "value")),
+			expected: NewSet(
+				_P("named", "a", "named", "value"),
+				_P("named", "a", "named"),
+				_P("named"),
+				_P("a", "named", "value"),
+				_P("a", "named"),
+				_P("a", "b", "value"),
+			),
+		},
+		{
+			set: NewSet(_P("named", "list", KeyByFields("name", "a"), "named", "a", "value")),
+			expected: NewSet(
+				_P("named", "list", KeyByFields("name", "a"), "named", "a", "value"),
+				_P("named", "list", KeyByFields("name", "a"), "named"),
+				_P("named", "list"),
+				_P("named"),
+			),
+		},
+	}
+
+	for _, test := range table {
+		t.Run(fmt.Sprintf("%v", test.set), func(t *testing.T) {
+			got := test.set.EnsureNamedFieldsAreMembers(nestedSchema())
+			if !got.Equals(test.expected) {
+				t.Errorf("expected %v, got %v (missing: %v/superfluous: %v)",
+					test.expected,
+					got,
+					test.expected.Difference(got),
+					got.Difference(test.expected),
+				)
+			}
+		})
+	}
+}
+
 func TestSetNodeMapIterate(t *testing.T) {
 	set := &SetNodeMap{}
 	toAdd := 5