[7.17] Script: track this pointer capture from blocks within lambdas (#82228)

docs/changelog/82228.yaml

@@ -0,0 +1,5 @@
+pr: 82228
+summary: "Script: track this pointer capture from blocks within lambdas"
+area: Infra/Scripting
+type: bug
+issues: []

modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/SemanticScope.java

@@ -198,6 +198,9 @@ public void setUsesInstanceMethod() {
                 return;
             }
             usesInstanceMethod = true;
+            if (parent != null) {
+                parent.setUsesInstanceMethod();
+            }
         }
 
         @Override
@@ -261,6 +264,12 @@ public Class<?> getReturnType() {
         public String getReturnCanonicalTypeName() {
             return parent.getReturnCanonicalTypeName();
         }
+
+        @Override
+        // If the parent scope is a lambda, we want to track this usage, so forward call to parent.
+        public void setUsesInstanceMethod() {
+            parent.setUsesInstanceMethod();
+        }
     }
 
     /**
@@ -356,7 +365,8 @@ public Variable defineVariable(Location location, Class<?> type, String name, bo
 
     public abstract Variable getVariable(Location location, String name);
 
-    // We only want to track instance method use inside of lambdas for "this" injection. It's a noop for other scopes.
+    // We only want to track instance method use inside of lambdas (and blocks inside lambdas) for "this" injection.
+    // It's a noop for other scopes.
     public void setUsesInstanceMethod() {}
 
     public boolean usesInstanceMethod() {

modules/lang-painless/src/test/java/org/elasticsearch/painless/UserFunctionTests.java

@@ -140,4 +140,40 @@ public void testLambdaCapture() {
         assertEquals(org.elasticsearch.core.List.of(100, 1, -100), exec(source, org.elasticsearch.core.Map.of("a", 1), false));
         assertBytecodeExists(source, "public static synthetic lambda$synthetic$0(ILjava/lang/Object;Ljava/lang/Object;)I");
     }
+
+    public void testCallUserMethodFromStatementWithinLambda() {
+        String source = ""
+            + "int test1() { return 1; }"
+            + "void test(Map params) { "
+            + "  int i = 0;"
+            + "  params.forEach("
+            + "      (k, v) -> { if (i == 0) { test1() } else { 20 } }"
+            + "    );"
+            + "}"
+            + "test(params)";
+        assertNull(exec(source, org.elasticsearch.core.Map.of("a", 5), false));
+        assertBytecodeExists(source, "public synthetic lambda$synthetic$0(ILjava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");
+    }
+
+    public void testCallUserMethodFromStatementWithinNestedLambda() {
+        String source = ""
+            + "int test1() { return 1; }"
+            + "void test(Map params) { "
+            + "  int i = 0;"
+            + "  int j = 5;"
+            + "  params.replaceAll( "
+            + "    (n, m) -> {"
+            + "      m.forEach("
+            + "        (k, v) -> { if (i == 0) { test1() } else { 20 } }"
+            + "      );"
+            + "      return ['aaa': j];"
+            + "    }"
+            + "  );"
+            + "}"
+            + "Map myParams = new HashMap(params);"
+            + "test(myParams);"
+            + "myParams['a']['aaa']";
+        assertEquals(5, exec(source, org.elasticsearch.core.Map.of("a", org.elasticsearch.core.Map.of("b", 1)), false));
+        assertBytecodeExists(source, "public synthetic lambda$synthetic$1(IILjava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");
+    }
 }