Painless: Fix variable scoping issue in lambdas not including captured variables. (#27571)

modules/lang-painless/src/main/java/org/elasticsearch/painless/antlr/Walker.java

@@ -1076,9 +1076,11 @@ public final class Walker extends PainlessParserBaseVisitor<ANode> {
             }
         }
 
+        FunctionReserved lambdaReserved = (FunctionReserved)reserved.pop();
+        reserved.peek().addUsedVariables(lambdaReserved);
+
         String name = nextLambda();
-        return new ELambda(name, (FunctionReserved)reserved.pop(), location(ctx),
-                           paramTypes, paramNames, statements);
+        return new ELambda(name, lambdaReserved, location(ctx), paramTypes, paramNames, statements);
     }
 
     @Override

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFunction.java

@@ -41,11 +41,13 @@ import java.lang.invoke.MethodType;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 
 import static java.util.Collections.emptyList;
+import static java.util.Collections.unmodifiableSet;
 import static org.elasticsearch.painless.WriterConstants.CLASS_TYPE;
 
 /**
@@ -53,11 +55,22 @@ import static org.elasticsearch.painless.WriterConstants.CLASS_TYPE;
  */
 public final class SFunction extends AStatement {
     public static final class FunctionReserved implements Reserved {
+        private final Set<String> usedVariables = new HashSet<>();
         private int maxLoopCounter = 0;
 
         @Override
         public void markUsedVariable(String name) {
-            // Do nothing.
+            usedVariables.add(name);
+        }
+
+        @Override
+        public Set<String> getUsedVariables() {
+            return unmodifiableSet(usedVariables);
+        }
+
+        @Override
+        public void addUsedVariables(FunctionReserved reserved) {
+            usedVariables.addAll(reserved.getUsedVariables());
         }
 
         @Override

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SSource.java

@@ -32,6 +32,7 @@ import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.ScriptClassInfo;
 import org.elasticsearch.painless.SimpleChecksAdapter;
 import org.elasticsearch.painless.WriterConstants;
+import org.elasticsearch.painless.node.SFunction.FunctionReserved;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.Label;
@@ -89,6 +90,8 @@ public final class SSource extends AStatement {
      */
     public interface Reserved {
         void markUsedVariable(String name);
+        Set<String> getUsedVariables();
+        void addUsedVariables(FunctionReserved reserved);
 
         void setMaxLoopCounter(int max);
         int getMaxLoopCounter();
@@ -103,6 +106,16 @@ public final class SSource extends AStatement {
             usedVariables.add(name);
         }
 
+        @Override
+        public Set<String> getUsedVariables() {
+            return unmodifiableSet(usedVariables);
+        }
+
+        @Override
+        public void addUsedVariables(FunctionReserved reserved) {
+            usedVariables.addAll(reserved.getUsedVariables());
+        }
+
         @Override
         public void setMaxLoopCounter(int max) {
             maxLoopCounter = max;
@@ -112,10 +125,6 @@ public final class SSource extends AStatement {
         public int getMaxLoopCounter() {
             return maxLoopCounter;
         }
-
-        public Set<String> getUsedVariables() {
-            return unmodifiableSet(usedVariables);
-        }
     }
 
     private final ScriptClassInfo scriptClassInfo;

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

@@ -191,4 +191,13 @@ public class FactoryTests extends ScriptTestCase {
         assertEquals("def", script.execute());
         assertEquals("def", script.execute());
     }
+
+    public void testGetterInLambda() {
+        FactoryTestScript.Factory factory =
+            scriptEngine.compile("template_test",
+                "IntSupplier createLambda(IntSupplier s) { return s; } createLambda(() -> params['x'] + test).getAsInt()",
+                FactoryTestScript.CONTEXT, Collections.emptyMap());
+        FactoryTestScript script = factory.newInstance(Collections.singletonMap("x", 1));
+        assertEquals(2, script.execute(1));
+    }
 }