Refactor mapping node traversal and optimize RNode.GetAnnotations and RNode.GetLabels

kyaml/yaml/fns.go

@@ -197,36 +197,37 @@ func (c FieldClearer) Filter(rn *RNode) (*RNode, error) {
 		return nil, err
 	}
 
-	for i := 0; i < len(rn.Content()); i += 2 {
-		// if name matches, remove these 2 elements from the list because
-		// they are treated as a fieldName/fieldValue pair.
-		if rn.Content()[i].Value == c.Name {
-			if c.IfEmpty {
-				if len(rn.Content()[i+1].Content) > 0 {
-					continue
-				}
-			}
+	var removed *RNode
+	visitFieldsWhileTrue(rn.Content(), func(key, value *yaml.Node, keyIndex int) bool {
+		if key.Value != c.Name {
+			return true
+		}
 
-			// save the item we are about to remove
-			removed := NewRNode(rn.Content()[i+1])
-			if len(rn.YNode().Content) > i+2 {
-				l := len(rn.YNode().Content)
-				// remove from the middle of the list
-				rn.YNode().Content = rn.Content()[:i]
-				rn.YNode().Content = append(
-					rn.YNode().Content,
-					rn.Content()[i+2:l]...)
-			} else {
-				// remove from the end of the list
-				rn.YNode().Content = rn.Content()[:i]
+		// the name matches: remove these 2 elements from the list because
+		// they are treated as a fieldName/fieldValue pair.
+		if c.IfEmpty {
+			if len(value.Content) > 0 {
+				return true
 			}
+		}
 
-			// return the removed field name and value
-			return removed, nil
+		// save the item we are about to remove
+		removed = NewRNode(value)
+		if len(rn.YNode().Content) > keyIndex+2 {
+			l := len(rn.YNode().Content)
+			// remove from the middle of the list
+			rn.YNode().Content = rn.Content()[:keyIndex]
+			rn.YNode().Content = append(
+				rn.YNode().Content,
+				rn.Content()[keyIndex+2:l]...)
+		} else {
+			// remove from the end of the list
+			rn.YNode().Content = rn.Content()[:keyIndex]
 		}
-	}
-	// nothing removed
-	return nil, nil
+		return false
+	})
+
+	return removed, nil
 }
 
 func MatchElement(field, value string) ElementMatcher {
@@ -402,14 +403,15 @@ func (f FieldMatcher) Filter(rn *RNode) (*RNode, error) {
 		return nil, err
 	}
 
-	for i := 0; i < len(rn.Content()); i = IncrementFieldIndex(i) {
-		isMatchingField := rn.Content()[i].Value == f.Name
-		if isMatchingField {
-			requireMatchFieldValue := f.Value != nil
-			if !requireMatchFieldValue || rn.Content()[i+1].Value == f.Value.YNode().Value {
-				return NewRNode(rn.Content()[i+1]), nil
-			}
+	var returnNode *RNode
+	requireMatchFieldValue := f.Value != nil
+	visitMappingNodeFields(rn.Content(), func(key, value *yaml.Node) {
+		if !requireMatchFieldValue || value.Value == f.Value.YNode().Value {
+			returnNode = NewRNode(value)
 		}
+	}, f.Name)
+	if returnNode != nil {
+		return returnNode, nil
 	}
 
 	if f.Create != nil {
@@ -643,13 +645,19 @@ func (s MapEntrySetter) Filter(rn *RNode) (*RNode, error) {
 	if s.Name == "" {
 		s.Name = GetValue(s.Key)
 	}
-	for i := 0; i < len(rn.Content()); i = IncrementFieldIndex(i) {
-		isMatchingField := rn.Content()[i].Value == s.Name
-		if isMatchingField {
-			rn.Content()[i] = s.Key.YNode()
-			rn.Content()[i+1] = s.Value.YNode()
-			return rn, nil
+
+	content := rn.Content()
+	stillMissing := true
+	visitFieldsWhileTrue(content, func(key, value *yaml.Node, keyIndex int) bool {
+		if key.Value == s.Name {
+			content[keyIndex] = s.Key.YNode()
+			content[keyIndex+1] = s.Value.YNode()
+			stillMissing = false
 		}
+		return stillMissing
+	})
+	if !stillMissing {
+		return rn, nil
 	}
 
 	// create the field
@@ -868,9 +876,3 @@ func SplitIndexNameValue(p string) (string, string, error) {
 	}
 	return parts[0], parts[1], nil
 }
-
-// IncrementFieldIndex increments i to point to the next field name element in
-// a slice of Contents.
-func IncrementFieldIndex(i int) int {
-	return i + 2
-}

kyaml/yaml/rnode.go

@@ -242,11 +242,7 @@ func (rn *RNode) IsTaggedNull() bool {
 // IsNilOrEmpty is true if the node is nil,
 // has no YNode, or has YNode that appears empty.
 func (rn *RNode) IsNilOrEmpty() bool {
-	return rn.IsNil() ||
-		IsYNodeTaggedNull(rn.YNode()) ||
-		IsYNodeEmptyMap(rn.YNode()) ||
-		IsYNodeEmptySeq(rn.YNode()) ||
-		IsYNodeZero(rn.YNode())
+	return rn.IsNil() || IsYNodeNilOrEmpty(rn.YNode())
 }
 
 // IsStringValue is true if the RNode is not nil and is scalar string node
@@ -420,12 +416,11 @@ func (rn *RNode) SetApiVersion(av string) {
 // given field, so this function cannot be used to make distinctions
 // between these cases.
 func (rn *RNode) getMapFieldValue(field string) *yaml.Node {
-	for i := 0; i < len(rn.Content()); i = IncrementFieldIndex(i) {
-		if rn.Content()[i].Value == field {
-			return rn.Content()[i+1]
-		}
-	}
-	return nil
+	var result *yaml.Node
+	visitMappingNodeFields(rn.Content(), func(key, value *yaml.Node) {
+		result = value
+	}, field)
+	return result
 }
 
 // GetName returns the name, or empty string if
@@ -696,9 +691,9 @@ func (rn *RNode) Fields() ([]string, error) {
 		return nil, errors.Wrap(err)
 	}
 	var fields []string
-	for i := 0; i < len(rn.Content()); i += 2 {
-		fields = append(fields, rn.Content()[i].Value)
-	}
+	visitMappingNodeFields(rn.Content(), func(key, value *yaml.Node) {
+		fields = append(fields, key.Value)
+	})
 	return fields, nil
 }
 
@@ -709,13 +704,12 @@ func (rn *RNode) FieldRNodes() ([]*RNode, error) {
 		return nil, errors.Wrap(err)
 	}
 	var fields []*RNode
-	for i := 0; i < len(rn.Content()); i += 2 {
-		yNode := rn.Content()[i]
+	visitMappingNodeFields(rn.Content(), func(key, value *yaml.Node) {
 		// for each key node in the input mapping node contents create equivalent rNode
 		rNode := &RNode{}
-		rNode.SetYNode(yNode)
+		rNode.SetYNode(key)
 		fields = append(fields, rNode)
-	}
+	})
 	return fields, nil
 }
 
@@ -725,13 +719,11 @@ func (rn *RNode) Field(field string) *MapNode {
 	if rn.YNode().Kind != yaml.MappingNode {
 		return nil
 	}
-	for i := 0; i < len(rn.Content()); i = IncrementFieldIndex(i) {
-		isMatchingField := rn.Content()[i].Value == field
-		if isMatchingField {
-			return &MapNode{Key: NewRNode(rn.Content()[i]), Value: NewRNode(rn.Content()[i+1])}
-		}
-	}
-	return nil
+	var result *MapNode
+	visitMappingNodeFields(rn.Content(), func(key, value *yaml.Node) {
+		result = &MapNode{Key: NewRNode(key), Value: NewRNode(value)}
+	}, field)
+	return result
 }
 
 // VisitFields calls fn for each field in the RNode.
@@ -752,6 +744,47 @@ func (rn *RNode) VisitFields(fn func(node *MapNode) error) error {
 	return nil
 }
 
+// visitMappingNodeFields calls fn for fields in the content, in content order.
+// The caller is responsible to ensure the node is a mapping node. If fieldNames
+// are specified, then fn is called only for the fields that match the given
+// fieldNames. fieldNames must contain unique values.
+func visitMappingNodeFields(content []*yaml.Node, fn func(key, value *yaml.Node), fieldNames ...string) {
+	switch len(fieldNames) {
+	case 0: // visit all fields
+		visitFieldsWhileTrue(content, func(key, value *yaml.Node, _ int) bool {
+			fn(key, value)
+			return true
+		})
+	default: // visit specified fields
+		// assumption: fields in content have unique names
+		found := 0
+		visitFieldsWhileTrue(content, func(key, value *yaml.Node, _ int) bool {
+			if key == nil {
+				return true
+			}
+			if !sliceutil.Contains(fieldNames, key.Value) {
+				return true
+			}
+			fn(key, value)
+			found++
+			return found < len(fieldNames)
+		})
+	}
+}
+
+// visitFieldsWhileTrue calls fn for the fields in content, in content order,
+// until either fn returns false or all fields have been visited. The caller
+// should ensure that content is from a mapping node, or fits the same expected
+// pattern (consecutive key/value entries in the slice).
+func visitFieldsWhileTrue(content []*yaml.Node, fn func(key, value *yaml.Node, keyIndex int) bool) {
+	for i := 0; i < len(content); i += 2 {
+		continueVisiting := fn(content[i], content[i+1], i)
+		if !continueVisiting {
+			return
+		}
+	}
+}
+
 // Elements returns the list of elements in the RNode.
 // Returns an error for non-SequenceNodes.
 func (rn *RNode) Elements() ([]*RNode, error) {
@@ -1003,17 +1036,19 @@ func findMergeValues(yn *yaml.Node) ([]*yaml.Node, error) {
 // it fails.
 func getMergeTagValue(yn *yaml.Node) (*yaml.Node, error) {
 	var result *yaml.Node
-	for i := 0; i < len(yn.Content); i += 2 {
-		key := yn.Content[i]
-		value := yn.Content[i+1]
+	var err error
+	visitFieldsWhileTrue(yn.Content, func(key, value *yaml.Node, _ int) bool {
 		if isMerge(key) {
 			if result != nil {
-				return nil, fmt.Errorf("duplicate merge key")
+				err = fmt.Errorf("duplicate merge key")
+				result = nil
+				return false
 			}
 			result = value
 		}
-	}
-	return result, nil
+		return true
+	})
+	return result, err
 }
 
 // removeMergeTags removes all merge tags and returns a ordered list of yaml

kyaml/yaml/rnode_test.go

@@ -2311,6 +2311,28 @@ func TestGetAnnotations(t *testing.T) {
 	}
 }
 
+func BenchmarkGetAnnotations(b *testing.B) {
+	counts := []int{0, 2, 5, 8}
+	for _, count := range counts {
+		appliedAnnotations := make(map[string]string, count)
+		for i := 1; i <= count; i++ {
+			key := fmt.Sprintf("annotation-key-%d", i)
+			value := fmt.Sprintf("annotation-value-%d", i)
+			appliedAnnotations[key] = value
+		}
+		rn := NewRNode(nil)
+		if err := rn.UnmarshalJSON([]byte(deploymentBiggerJson)); err != nil {
+			b.Fatalf("unexpected unmarshaljson err: %v", err)
+		}
+		assert.NoError(b, rn.SetAnnotations(appliedAnnotations))
+		b.Run(fmt.Sprintf("%02d", count), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				_ = rn.GetAnnotations()
+			}
+		})
+	}
+}
+
 func TestGetFieldValueWithDot(t *testing.T) {
 	const input = `
 kind: Pod

kyaml/yaml/types.go

@@ -39,11 +39,20 @@ func IsYNodeEmptyMap(n *yaml.Node) bool {
 	return n != nil && n.Kind == yaml.MappingNode && len(n.Content) == 0
 }
 
-// IsYNodeEmptyMap is true if the Node is a non-nil empty sequence.
+// IsYNodeEmptySeq is true if the Node is a non-nil empty sequence.
 func IsYNodeEmptySeq(n *yaml.Node) bool {
 	return n != nil && n.Kind == yaml.SequenceNode && len(n.Content) == 0
 }
 
+// IsYNodeNilOrEmpty is true if the Node is nil or appears empty.
+func IsYNodeNilOrEmpty(n *yaml.Node) bool {
+	return n == nil ||
+		IsYNodeTaggedNull(n) ||
+		IsYNodeEmptyMap(n) ||
+		IsYNodeEmptySeq(n) ||
+		IsYNodeZero(n)
+}
+
 // IsYNodeEmptyDoc is true if the node is a Document with no content.
 // E.g.: "---\n---"
 func IsYNodeEmptyDoc(n *yaml.Node) bool {