Add Leaves method to fieldpath.Set object

Author: kevindelgado

fieldpath/set.go

@@ -206,6 +206,45 @@ func (s *Set) WithPrefix(pe PathElement) *Set {
 	return subset
 }
 
+// Leaves returns a set containing only the leaf paths
+// of a set.
+func (s *Set) Leaves() *Set {
+	leaves := PathElementSet{}
+	im := 0
+	ic := 0
+	prevDiff := 0
+
+	// any members that are not also children are leaves
+	for im < len(s.Members.members) {
+		member := s.Members.members[im]
+		if ic >= len(s.Children.members) {
+			leaves.members = append(leaves.members, member)
+			im++
+			continue
+		}
+
+		diff := member.Compare(s.Children.members[ic].pathElement)
+		if diff-prevDiff != 0 {
+			leaves.members = append(leaves.members, member)
+		}
+		prevDiff = diff
+
+		if diff < 0 {
+			im++
+		} else if diff > 0 {
+			ic++
+		} else {
+			im++
+			ic++
+		}
+	}
+
+	return &Set{
+		Members:  leaves,
+		Children: *s.Children.Leaves(),
+	}
+}
+
 // setNode is a pair of PathElement / Set, for the purpose of expressing
 // nested set membership.
 type setNode struct {
@@ -455,3 +494,17 @@ func (s *SetNodeMap) iteratePrefix(prefix Path, f func(Path)) {
 		n.set.iteratePrefix(append(prefix, pe), f)
 	}
 }
+
+// Leaves returns a SetNodeMap containing
+// only setNodes with leaf PathElements.
+func (s *SetNodeMap) Leaves() *SetNodeMap {
+	out := &SetNodeMap{}
+	out.members = make(sortedSetNode, len(s.members))
+	for i, n := range s.members {
+		out.members[i] = setNode{
+			pathElement: n.pathElement,
+			set:         n.set.Leaves(),
+		}
+	}
+	return out
+}

fieldpath/set_test.go

@@ -140,6 +140,12 @@ func BenchmarkFieldSet(b *testing.B) {
 				randOperand().RecursiveDifference(randOperand())
 			}
 		})
+		b.Run(fmt.Sprintf("leaves-%v", here.size), func(b *testing.B) {
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				randOperand().Leaves()
+			}
+		})
 	}
 }
 
@@ -456,6 +462,106 @@ func TestSetIntersectionDifference(t *testing.T) {
 	})
 }
 
+func TestSetLeaves(t *testing.T) {
+	table := []struct {
+		name     string
+		input    *Set
+		expected *Set
+	}{
+		{
+			name:     "empty set",
+			input:    NewSet(),
+			expected: NewSet(),
+		}, {
+			name: "all leaves",
+			input: NewSet(
+				_P("path1"),
+				_P("path2"),
+				_P("path3"),
+			),
+			expected: NewSet(
+				_P("path1"),
+				_P("path2"),
+				_P("path3"),
+			),
+		}, {
+			name: "only one leaf",
+			input: NewSet(
+				_P("root"),
+				_P("root", "l1"),
+				_P("root", "l1", "l2"),
+				_P("root", "l1", "l2", "l3"),
+			),
+			expected: NewSet(
+				_P("root", "l1", "l2", "l3"),
+			),
+		}, {
+			name: "multiple values, check for overwrite",
+			input: NewSet(
+				_P("root", KeyByFields("name", "a")),
+				_P("root", KeyByFields("name", "a"), "name"),
+				_P("root", KeyByFields("name", "a"), "value", "b"),
+				_P("root", KeyByFields("name", "a"), "value", "c"),
+			),
+			expected: NewSet(
+				_P("root", KeyByFields("name", "a"), "name"),
+				_P("root", KeyByFields("name", "a"), "value", "b"),
+				_P("root", KeyByFields("name", "a"), "value", "c"),
+			),
+		}, {
+			name: "multiple values and nested",
+			input: NewSet(
+				_P("root", KeyByFields("name", "a")),
+				_P("root", KeyByFields("name", "a"), "name"),
+				_P("root", KeyByFields("name", "a"), "value", "b"),
+				_P("root", KeyByFields("name", "a"), "value", "b", "d"),
+				_P("root", KeyByFields("name", "a"), "value", "c"),
+			),
+			expected: NewSet(
+				_P("root", KeyByFields("name", "a"), "name"),
+				_P("root", KeyByFields("name", "a"), "value", "b", "d"),
+				_P("root", KeyByFields("name", "a"), "value", "c"),
+			),
+		}, {
+			name: "all-in-one",
+			input: NewSet(
+				_P("root"),
+				_P("root", KeyByFields("name", "a")),
+				_P("root", KeyByFields("name", "a"), "name"),
+				_P("root", KeyByFields("name", "a"), "value", "b"),
+				_P("root", KeyByFields("name", "a"), "value", "b", "c"),
+				_P("root", KeyByFields("name", "a"), "value", "d"),
+				_P("root", KeyByFields("name", "a"), "value", "e"),
+				_P("root", "x"),
+				_P("root", "x", "y"),
+				_P("root", "x", "z"),
+				_P("root", KeyByFields("name", "p")),
+				_P("root", KeyByFields("name", "p"), "name"),
+				_P("root", KeyByFields("name", "p"), "value", "q"),
+			),
+			expected: NewSet(
+				_P("root", KeyByFields("name", "a"), "name"),
+				_P("root", KeyByFields("name", "a"), "value", "b", "c"),
+				_P("root", KeyByFields("name", "a"), "value", "d"),
+				_P("root", KeyByFields("name", "a"), "value", "e"),
+				_P("root", "x", "y"),
+				_P("root", "x", "z"),
+				_P("root", KeyByFields("name", "p"), "name"),
+				_P("root", KeyByFields("name", "p"), "value", "q"),
+			),
+		},
+	}
+
+	for _, tt := range table {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.input.Leaves(); !tt.expected.Equals(got) {
+				t.Errorf("expected %v, got %v for input %v", tt.expected, got, tt.input)
+			}
+		})
+	}
+
+}
+
 func TestSetDifference(t *testing.T) {
 	table := []struct {
 		name                      string