Optimize fieldpath.Set#Leaves

Author: kevindelgado

fieldpath/set.go

@@ -206,30 +206,32 @@ func (s *Set) WithPrefix(pe PathElement) *Set {
 	return subset
 }
 
+func (s *Set) leavesPrefix(prefix Path, set *Set) {
+	for _, child := range s.Children.members {
+		child.set.leavesPrefix(append(prefix, child.pathElement), set)
+	}
+
+	for _, mem := range s.Members.members {
+		isChild := false
+		currentPath := append(prefix, mem)
+		for _, child := range s.Children.members {
+			if mem.Equals(child.pathElement) {
+				isChild = true
+			}
+		}
+		if !isChild {
+			// any members that are not also children are leaves
+			set.Insert(currentPath)
+		}
+	}
+}
+
 // Leaves returns a set containing only the leaf paths
 // of a set.
 func (s *Set) Leaves() *Set {
-	paths := []Path{}
-	s.Iterate(func(p Path) {
-		parentPath := Path(p[0 : len(p)-1])
-		n := 0
-		for _, path := range paths {
-			// perform preorder DFS on every path in the set,
-			// if the parent of the current path exists in paths
-			// then the parent must not be a leaf and should be
-			// overwritten by the current path
-			if !parentPath.Equals(path) {
-				paths[n] = path
-				n++
-			}
-		}
-		paths = paths[:n]
-		// save a copy of current path or else it can be overwritten
-		pathCopy := Path(make([]PathElement, len(p)))
-		copy(pathCopy, p)
-		paths = append(paths, pathCopy)
-	})
-	return NewSet(paths...)
+	out := &Set{}
+	s.leavesPrefix(Path{}, out)
+	return out
 }
 
 // setNode is a pair of PathElement / Set, for the purpose of expressing

fieldpath/set_test.go

@@ -496,23 +496,56 @@ func TestSetLeaves(t *testing.T) {
 				_P("root", "l1", "l2", "l3"),
 			),
 		}, {
-			name: "nested with lists",
+			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"),
 			),