yq/pkg/yqlib/operator_traverse_path.go

package yqlib

import (
	"container/list"
	"fmt"
	"strconv"

	"github.com/elliotchance/orderedmap"
	yaml "gopkg.in/yaml.v3"
)

type traversePreferences struct {
	DontFollowAlias      bool
	IncludeMapKeys       bool
	DontAutoCreate       bool // by default, we automatically create entries on the fly.
	DontIncludeMapValues bool
	OptionalTraverse     bool // e.g. .adf?
}

func splat(d *dataTreeNavigator, context Context, prefs traversePreferences) (Context, error) {
	return traverseNodesWithArrayIndices(context, make([]*yaml.Node, 0), prefs)
}

func traversePathOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
	log.Debugf("-- traversePathOperator")
	var matches = list.New()

	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
		newNodes, err := traverse(d, context, el.Value.(*CandidateNode), expressionNode.Operation)
		if err != nil {
			return Context{}, err
		}
		matches.PushBackList(newNodes)
	}

	return context.ChildContext(matches), nil
}

func traverse(d *dataTreeNavigator, context Context, matchingNode *CandidateNode, operation *Operation) (*list.List, error) {
	log.Debug("Traversing %v", NodeToString(matchingNode))
	value := matchingNode.Node

	if value.Tag == "!!null" && operation.Value != "[]" {
		log.Debugf("Guessing kind")
		// we must ahve added this automatically, lets guess what it should be now
		switch operation.Value.(type) {
		case int, int64:
			log.Debugf("probably an array")
			value.Kind = yaml.SequenceNode
		default:
			log.Debugf("probably a map")
			value.Kind = yaml.MappingNode
		}
		value.Tag = ""
	}

	switch value.Kind {
	case yaml.MappingNode:
		log.Debug("its a map with %v entries", len(value.Content)/2)
		return traverseMap(context, matchingNode, operation.StringValue, operation.Preferences.(traversePreferences), false)

	case yaml.SequenceNode:
		log.Debug("its a sequence of %v things!", len(value.Content))
		return traverseArray(matchingNode, operation, operation.Preferences.(traversePreferences))

	case yaml.AliasNode:
		log.Debug("its an alias!")
		matchingNode.Node = matchingNode.Node.Alias
		return traverse(d, context, matchingNode, operation)
	case yaml.DocumentNode:
		log.Debug("digging into doc node")

		return traverse(d, context, matchingNode.CreateChild(nil, matchingNode.Node.Content[0]), operation)
	default:
		return list.New(), nil
	}
}

func traverseArrayOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {

	//lhs may update the variable context, we should pass that into the RHS
	// BUT we still return the original context back (see jq)
	// https://stedolan.github.io/jq/manual/#Variable/SymbolicBindingOperator:...as$identifier|...

	lhs, err := d.GetMatchingNodes(context, expressionNode.Lhs)
	if err != nil {
		return Context{}, err
	}

	// rhs is a collect expression that will yield indexes to retreive of the arrays

	rhs, err := d.GetMatchingNodes(context.ReadOnlyClone(), expressionNode.Rhs)

	if err != nil {
		return Context{}, err
	}
	prefs := traversePreferences{}

	if expressionNode.Rhs.Rhs != nil && expressionNode.Rhs.Rhs.Operation.Preferences != nil {
		prefs = expressionNode.Rhs.Rhs.Operation.Preferences.(traversePreferences)
	}
	var indicesToTraverse = rhs.MatchingNodes.Front().Value.(*CandidateNode).Node.Content

	//now we traverse the result of the lhs against the indices we found
	result, err := traverseNodesWithArrayIndices(lhs, indicesToTraverse, prefs)
	if err != nil {
		return Context{}, err
	}
	return context.ChildContext(result.MatchingNodes), nil
}

func traverseNodesWithArrayIndices(context Context, indicesToTraverse []*yaml.Node, prefs traversePreferences) (Context, error) {
	var matchingNodeMap = list.New()
	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
		candidate := el.Value.(*CandidateNode)
		newNodes, err := traverseArrayIndices(context, candidate, indicesToTraverse, prefs)
		if err != nil {
			return Context{}, err
		}
		matchingNodeMap.PushBackList(newNodes)
	}

	return context.ChildContext(matchingNodeMap), nil
}

func traverseArrayIndices(context Context, matchingNode *CandidateNode, indicesToTraverse []*yaml.Node, prefs traversePreferences) (*list.List, error) { // call this if doc / alias like the other traverse
	node := matchingNode.Node
	if node.Tag == "!!null" {
		log.Debugf("OperatorArrayTraverse got a null - turning it into an empty array")
		// auto vivification, make it into an empty array
		node.Tag = ""
		node.Kind = yaml.SequenceNode
	}

	if node.Kind == yaml.AliasNode {
		matchingNode.Node = node.Alias
		return traverseArrayIndices(context, matchingNode, indicesToTraverse, prefs)
	} else if node.Kind == yaml.SequenceNode {
		return traverseArrayWithIndices(matchingNode, indicesToTraverse, prefs)
	} else if node.Kind == yaml.MappingNode {
		return traverseMapWithIndices(context, matchingNode, indicesToTraverse, prefs)
	} else if node.Kind == yaml.DocumentNode {
		return traverseArrayIndices(context, matchingNode.CreateChild(nil, matchingNode.Node.Content[0]), indicesToTraverse, prefs)
	}
	log.Debugf("OperatorArrayTraverse skipping %v as its a %v", matchingNode, node.Tag)
	return list.New(), nil
}

func traverseMapWithIndices(context Context, candidate *CandidateNode, indices []*yaml.Node, prefs traversePreferences) (*list.List, error) {
	if len(indices) == 0 {
		return traverseMap(context, candidate, "", prefs, true)
	}

	var matchingNodeMap = list.New()

	for _, indexNode := range indices {
		log.Debug("traverseMapWithIndices: %v", indexNode.Value)
		newNodes, err := traverseMap(context, candidate, indexNode.Value, prefs, false)
		if err != nil {
			return nil, err
		}
		matchingNodeMap.PushBackList(newNodes)
	}

	return matchingNodeMap, nil
}

func traverseArrayWithIndices(candidate *CandidateNode, indices []*yaml.Node, prefs traversePreferences) (*list.List, error) {
	log.Debug("traverseArrayWithIndices")
	var newMatches = list.New()
	node := unwrapDoc(candidate.Node)
	if len(indices) == 0 {
		log.Debug("splatting")
		var index int64
		for index = 0; index < int64(len(node.Content)); index = index + 1 {

			newMatches.PushBack(candidate.CreateChild(index, node.Content[index]))
		}
		return newMatches, nil

	}

	for _, indexNode := range indices {
		log.Debug("traverseArrayWithIndices: '%v'", indexNode.Value)
		index, err := strconv.ParseInt(indexNode.Value, 10, 64)
		if err != nil && prefs.OptionalTraverse {
			continue
		}
		if err != nil {
			return nil, fmt.Errorf("Cannot index array with '%v' (%v)", indexNode.Value, err)
		}
		indexToUse := index
		contentLength := int64(len(node.Content))
		for contentLength <= index {
			node.Content = append(node.Content, &yaml.Node{Tag: "!!null", Kind: yaml.ScalarNode, Value: "null"})
			contentLength = int64(len(node.Content))
		}

		if indexToUse < 0 {
			indexToUse = contentLength + indexToUse
		}

		if indexToUse < 0 {
			return nil, fmt.Errorf("Index [%v] out of range, array size is %v", index, contentLength)
		}

		newMatches.PushBack(candidate.CreateChild(index, node.Content[indexToUse]))
	}
	return newMatches, nil
}

func keyMatches(key *yaml.Node, wantedKey string) bool {
	return matchKey(key.Value, wantedKey)
}

func traverseMap(context Context, matchingNode *CandidateNode, key string, prefs traversePreferences, splat bool) (*list.List, error) {
	var newMatches = orderedmap.NewOrderedMap()
	err := doTraverseMap(newMatches, matchingNode, key, prefs, splat)

	if err != nil {
		return nil, err
	}

	if !prefs.DontAutoCreate && !context.DontAutoCreate && newMatches.Len() == 0 {
		//no matches, create one automagically
		valueNode := &yaml.Node{Tag: "!!null", Kind: yaml.ScalarNode, Value: "null"}
		keyNode := &yaml.Node{Kind: yaml.ScalarNode, Value: key}
		node := matchingNode.Node
		node.Content = append(node.Content, keyNode, valueNode)

		if prefs.IncludeMapKeys {
			log.Debug("including key")
			candidateNode := matchingNode.CreateChild(key, keyNode)
			candidateNode.IsMapKey = true
			newMatches.Set(fmt.Sprintf("keyOf-%v", candidateNode.GetKey()), candidateNode)
		}
		if !prefs.DontIncludeMapValues {
			log.Debug("including value")
			candidateNode := matchingNode.CreateChild(key, valueNode)
			newMatches.Set(candidateNode.GetKey(), candidateNode)
		}
	}

	results := list.New()
	i := 0
	for el := newMatches.Front(); el != nil; el = el.Next() {
		results.PushBack(el.Value)
		i++
	}
	return results, nil
}

func doTraverseMap(newMatches *orderedmap.OrderedMap, candidate *CandidateNode, wantedKey string, prefs traversePreferences, splat bool) error {
	// value.Content is a concatenated array of key, value,
	// so keys are in the even indexes, values in odd.
	// merge aliases are defined first, but we only want to traverse them
	// if we don't find a match directly on this node first.

	node := candidate.Node

	var contents = node.Content
	for index := 0; index < len(contents); index = index + 2 {
		key := contents[index]
		value := contents[index+1]

		log.Debug("checking %v (%v)", key.Value, key.Tag)
		//skip the 'merge' tag, find a direct match first
		if key.Tag == "!!merge" && !prefs.DontFollowAlias {
			log.Debug("Merge anchor")
			err := traverseMergeAnchor(newMatches, candidate, value, wantedKey, prefs, splat)
			if err != nil {
				return err
			}
		} else if splat || keyMatches(key, wantedKey) {
			log.Debug("MATCHED")
			if prefs.IncludeMapKeys {
				log.Debug("including key")
				candidateNode := candidate.CreateChild(key.Value, key)
				candidateNode.IsMapKey = true
				newMatches.Set(fmt.Sprintf("keyOf-%v", candidateNode.GetKey()), candidateNode)
			}
			if !prefs.DontIncludeMapValues {
				log.Debug("including value")
				candidateNode := candidate.CreateChild(key.Value, value)
				newMatches.Set(candidateNode.GetKey(), candidateNode)
			}
		}
	}

	return nil
}

func traverseMergeAnchor(newMatches *orderedmap.OrderedMap, originalCandidate *CandidateNode, value *yaml.Node, wantedKey string, prefs traversePreferences, splat bool) error {
	switch value.Kind {
	case yaml.AliasNode:
		candidateNode := originalCandidate.CreateChild(nil, value.Alias)
		return doTraverseMap(newMatches, candidateNode, wantedKey, prefs, splat)
	case yaml.SequenceNode:
		for _, childValue := range value.Content {
			err := traverseMergeAnchor(newMatches, originalCandidate, childValue, wantedKey, prefs, splat)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func traverseArray(candidate *CandidateNode, operation *Operation, prefs traversePreferences) (*list.List, error) {
	log.Debug("operation Value %v", operation.Value)
	indices := []*yaml.Node{{Value: operation.StringValue}}
	return traverseArrayWithIndices(candidate, indices, prefs)
}