From f2b923248e2692a3d35bdff15285be256cfd69b6 Mon Sep 17 00:00:00 2001
From: Knut Wannheden <knut@moderne.io>
Date: Mon, 18 Nov 2024 09:56:35 +0100
Subject: [PATCH] Allow YAML `JsonPatchMatcher` to match document-level
 sequences (#4680)

---
 .../org/openrewrite/yaml/JsonPathMatcher.java | 15 +++++-----
 .../openrewrite/yaml/JsonPathMatcherTest.java | 30 +++++++++++++++++++
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/JsonPathMatcher.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/JsonPathMatcher.java
index 4adc581ae72..260c9f24974 100644
--- a/rewrite-yaml/src/main/java/org/openrewrite/yaml/JsonPathMatcher.java
+++ b/rewrite-yaml/src/main/java/org/openrewrite/yaml/JsonPathMatcher.java
@@ -175,18 +175,17 @@ protected boolean shouldVisitNextChild(RuleNode node, Object currentResult) {
         public Object visitJsonPath(JsonPathParser.JsonPathContext ctx) {
             MATCH:
             if (ctx.ROOT() != null || ctx.start.getType() == JsonPathLexer.LBRACK) {
+                Tree first = cursorPath.get(0);
+                if (first instanceof Yaml.Document) {
+                    scope = ((Yaml.Document) first).getBlock();
+                    break MATCH;
+                }
                 for (Tree tree : cursorPath) {
                     if (tree instanceof Yaml.Mapping) {
                         scope = tree;
                         break MATCH;
                     }
                 }
-                for (Tree t : cursorPath) {
-                    if (t instanceof Yaml.Document && ((Yaml.Document) t).getBlock() instanceof Yaml.Mapping) {
-                        scope = ((Yaml.Document) t).getBlock();
-                        break MATCH;
-                    }
-                }
                 scope = null;
             }
             return super.visitJsonPath(ctx);
@@ -213,7 +212,7 @@ private JsonPathParser.ExpressionContext getExpressionContext(ParserRuleContext
             if (previous.indexOf(current) - 1 < 0 || "$".equals(previous.get(previous.indexOf(current) - 1).getText())) {
                 List<Object> results = new ArrayList<>();
                 for (Tree path : cursorPath) {
-                    JsonPathMatcher.JsonPathYamlVisitor v = new JsonPathMatcher.JsonPathYamlVisitor(cursorPath, path, null, false);
+                    JsonPathYamlVisitor v = new JsonPathYamlVisitor(cursorPath, path, null, false);
                     for (int i = 1; i < ctx.getChildCount(); i++) {
                         result = v.visit(ctx.getChild(i));
                         if (result != null) {
@@ -224,7 +223,7 @@ private JsonPathParser.ExpressionContext getExpressionContext(ParserRuleContext
                 return results;
                 // Otherwise, the recursive descent is scoped to the previous match. `$.foo..['find-in-foo']`.
             } else {
-                JsonPathMatcher.JsonPathYamlVisitor v = new JsonPathMatcher.JsonPathYamlVisitor(cursorPath, scope, null, true);
+                JsonPathYamlVisitor v = new JsonPathYamlVisitor(cursorPath, scope, null, true);
                 for (int i = 1; i < ctx.getChildCount(); i++) {
                     result = v.visit(ctx.getChild(i));
                     if (result != null) {
diff --git a/rewrite-yaml/src/test/java/org/openrewrite/yaml/JsonPathMatcherTest.java b/rewrite-yaml/src/test/java/org/openrewrite/yaml/JsonPathMatcherTest.java
index 1631ff67c56..74664a4317f 100644
--- a/rewrite-yaml/src/test/java/org/openrewrite/yaml/JsonPathMatcherTest.java
+++ b/rewrite-yaml/src/test/java/org/openrewrite/yaml/JsonPathMatcherTest.java
@@ -483,6 +483,36 @@ void allElements() {
         );
     }
 
+    @Test
+    void documentLevelSequenceWildcard() {
+        assertMatched(
+          "$[*].id",
+          List.of(
+            """
+              - id: 0
+              - id: 1
+              - id: 2
+              """
+          ),
+          List.of("id: 0", "id: 1", "id: 2")
+        );
+    }
+
+    @Test
+    void documentLevelSequenceSingle() {
+        assertMatched(
+          "$[1].id",
+          List.of(
+            """
+              - id: 0
+              - id: 1
+              - id: 2
+              """
+          ),
+          List.of("id: 1")
+        );
+    }
+
     @Test
     void bracketOperatorByNames() {
         assertMatched(