From 150a7f455bc71c5060d99f89e5fc0250a823b09e Mon Sep 17 00:00:00 2001
From: Jeff Davis <jeffda@vmware.com>
Date: Tue, 15 Jun 2021 16:37:48 -0500
Subject: [PATCH] fix(filter): correct filter logic for maps and scalars

Signed-off-by: Jeff Davis <jeffda@vmware.com>
---
 pkg/yamlpath/lexer.go      |  7 ++++++-
 pkg/yamlpath/lexer_test.go |  2 +-
 pkg/yamlpath/path.go       | 34 ++++++++++++++++++++++++++++++----
 pkg/yamlpath/path_test.go  |  9 +++++++++
 4 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/pkg/yamlpath/lexer.go b/pkg/yamlpath/lexer.go
index 8a37471..67fe197 100644
--- a/pkg/yamlpath/lexer.go
+++ b/pkg/yamlpath/lexer.go
@@ -52,6 +52,7 @@ const (
 	lexemePropertyName
 	lexemeBracketPropertyName
 	lexemeArraySubscriptPropertyName
+	lexemeRecursiveFilterBegin
 	lexemeEOF // lexing complete
 )
 
@@ -596,7 +597,11 @@ func lexSubPath(l *lexer) stateFn {
 		return lexOptionalArrayIndex
 
 	case l.consumed(filterBegin):
-		l.emit(lexemeFilterBegin)
+		if l.lastEmittedLexemeType == lexemeRecursiveDescent {
+			l.emit(lexemeRecursiveFilterBegin)
+		} else {
+			l.emit(lexemeFilterBegin)
+		}
 		l.push(lexFilterEnd)
 		return lexFilterExprInitial
 
diff --git a/pkg/yamlpath/lexer_test.go b/pkg/yamlpath/lexer_test.go
index d037964..42cd06f 100644
--- a/pkg/yamlpath/lexer_test.go
+++ b/pkg/yamlpath/lexer_test.go
@@ -822,7 +822,7 @@ func TestLexer(t *testing.T) {
 			expected: []lexeme{
 				{typ: lexemeRoot, val: "$"},
 				{typ: lexemeRecursiveDescent, val: ".."},
-				{typ: lexemeFilterBegin, val: "[?("},
+				{typ: lexemeRecursiveFilterBegin, val: "[?("},
 				{typ: lexemeFilterAt, val: "@"},
 				{typ: lexemeDotChild, val: ".child"},
 				{typ: lexemeFilterEnd, val: ")]"},
diff --git a/pkg/yamlpath/path.go b/pkg/yamlpath/path.go
index f5072b8..fb0fd02 100644
--- a/pkg/yamlpath/path.go
+++ b/pkg/yamlpath/path.go
@@ -116,7 +116,12 @@ func newPath(l *lexer) (*Path, error) {
 		subscript := strings.TrimSuffix(strings.TrimPrefix(lx.val, "["), "]")
 		return arraySubscriptThen(subscript, subPath), nil
 
-	case lexemeFilterBegin:
+	case lexemeFilterBegin, lexemeRecursiveFilterBegin:
+		var recursive bool
+
+		if lx.typ == lexemeRecursiveFilterBegin {
+			recursive = true
+		}
 		filterLexemes := []lexeme{}
 		filterNestingLevel := 1
 	f:
@@ -144,6 +149,9 @@ func newPath(l *lexer) (*Path, error) {
 		if err != nil {
 			return nil, err
 		}
+		if recursive {
+			return recursiveFilterThen(filterLexemes, subPath), nil
+		}
 		return filterThen(filterLexemes, subPath), nil
 	case lexemePropertyName:
 		subPath, err := newPath(l)
@@ -427,10 +435,28 @@ func filterThen(filterLexemes []lexeme, p *Path) *Path {
 	filter := newFilter(newFilterNode(filterLexemes))
 	return new(func(node, root *yaml.Node) yit.Iterator {
 		its := []yit.Iterator{}
-		for _, c := range node.Content {
-			if filter(c, root) {
-				its = append(its, compose(yit.FromNode(c), p, root))
+		if node.Kind == yaml.SequenceNode {
+			for _, c := range node.Content {
+				if filter(c, root) {
+					its = append(its, compose(yit.FromNode(c), p, root))
+				}
 			}
+		} else {
+			if filter(node, root) {
+				its = append(its, compose(yit.FromNode(node), p, root))
+			}
+		}
+		return yit.FromIterators(its...)
+	})
+}
+
+func recursiveFilterThen(filterLexemes []lexeme, p *Path) *Path {
+	filter := newFilter(newFilterNode(filterLexemes))
+	return new(func(node, root *yaml.Node) yit.Iterator {
+		its := []yit.Iterator{}
+
+		if filter(node, root) {
+			its = append(its, compose(yit.FromNode(node), p, root))
 		}
 		return yit.FromIterators(its...)
 	})
diff --git a/pkg/yamlpath/path_test.go b/pkg/yamlpath/path_test.go
index b051b8b..77152dc 100644
--- a/pkg/yamlpath/path_test.go
+++ b/pkg/yamlpath/path_test.go
@@ -969,6 +969,15 @@ price: 12.99
 			},
 			expectedPathErr: "",
 		},
+		{
+			name: "map filter",
+			path: `$.store.bicycle[?(@.color == "red")]`,
+			expectedStrings: []string{
+				`color: red
+price: 19.95
+`,
+			},
+		},
 	}
 
 	focussed := false