Create NestedSet to populate set branches along with the leaves

Author: apelisse

fieldpath/set.go

@@ -110,6 +110,23 @@ func (s *Set) RecursiveDifference(s2 *Set) *Set {
 	}
 }
 
+// NestedSet returns a Set that contains all the fields in s, as well as
+// all the branches that are typically not included. For example, a set
+// made of "a.b.c" will end-up also owning "a" and "a.b".
+func (s *Set) NestedSet() *Set {
+	members := PathElementSet{
+		members: make(sortedPathElements, 0, s.Members.Size()+len(s.Children.members)),
+	}
+	members.members = append(members.members, s.Members.members...)
+	for _, node := range s.Children.members {
+		members.Insert(node.pathElement)
+	}
+	return &Set{
+		Members:  members,
+		Children: *s.Children.NestedSet(),
+	}
+}
+
 // Size returns the number of members of the set.
 func (s *Set) Size() int {
 	return s.Members.Size() + s.Children.Size()
@@ -391,6 +408,21 @@ func (s *SetNodeMap) RecursiveDifference(s2 *Set) *SetNodeMap {
 	return out
 }
 
+// NestedSet returns a set that contains all the branches on top of the leaves.
+func (s *SetNodeMap) NestedSet() *SetNodeMap {
+	out := make(sortedSetNode, 0, s.Size())
+	for _, member := range s.members {
+		out = append(out, setNode{
+			pathElement: member.pathElement,
+			set:         member.set.NestedSet(),
+		})
+	}
+
+	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

@@ -541,6 +541,46 @@ func TestSetDifference(t *testing.T) {
 	}
 }
 
+func TestSetNested(t *testing.T) {
+	table := []struct {
+		set, expected *Set
+	}{
+		{
+			set: NewSet(MakePathOrDie("a", "b", "c")),
+			expected: NewSet(
+				MakePathOrDie("a", "b", "c"),
+				MakePathOrDie("a", "b"),
+				MakePathOrDie("a"),
+			),
+		},
+		{
+			set: NewSet(MakePathOrDie("a", "b", "c"), MakePathOrDie("a", "d"), MakePathOrDie("a", "c", "e")),
+			expected: NewSet(
+				MakePathOrDie("a", "b", "c"),
+				MakePathOrDie("a", "b"),
+				MakePathOrDie("a", "c", "e"),
+				MakePathOrDie("a", "c"),
+				MakePathOrDie("a", "d"),
+				MakePathOrDie("a"),
+			),
+		},
+	}
+
+	for _, test := range table {
+		t.Run(fmt.Sprintf("%v", test.set), func(t *testing.T) {
+			got := test.set.NestedSet()
+			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