Merge pull request #18264 from rmuir/painless_ctx

Add 'ctx' keyword to painless.

modules/lang-painless/src/main/java/org/elasticsearch/painless/Analyzer.java

@@ -101,7 +101,7 @@ class Analyzer extends PainlessParserBaseVisitor<Void> {
         // scorer parameter passed to the script. internal use only.
         metadata.scorerValueSlot = utility.addVariable(null, "#scorer", definition.objectType).slot;
         // doc parameter passed to the script.
-        // TODO: currently working as a def type, should be smapType...
+        // TODO: currently working as a Map<String,Def>, we can do better?
         metadata.docValueSlot = utility.addVariable(null, "doc", definition.smapType).slot;
         //
         // reserved words implemented as local variables
@@ -110,6 +110,8 @@ class Analyzer extends PainlessParserBaseVisitor<Void> {
         metadata.loopCounterSlot = utility.addVariable(null, "#loop", definition.intType).slot;
         // document's score as a read-only float.
         metadata.scoreValueSlot = utility.addVariable(null, "_score", definition.floatType).slot;
+        // ctx map set by executable scripts as a read-only map.
+        metadata.ctxValueSlot = utility.addVariable(null, "ctx", definition.smapType).slot;
 
         metadata.createStatementMetadata(metadata.root);
         visit(metadata.root);

modules/lang-painless/src/main/java/org/elasticsearch/painless/AnalyzerExternal.java

@@ -449,16 +449,18 @@ class AnalyzerExternal {
             }
 
             // special cases: reserved words
-            if (varenmd.last && ("_score".equals(id) || "doc".equals(id))) {
-                // read-only: don't allow stores
-                if (parentemd.storeExpr != null) {
+            if ("_score".equals(id) || "doc".equals(id) || "ctx".equals(id)) {
+                // read-only: don't allow stores to ourself
+                if (varenmd.last && parentemd.storeExpr != null) {
                     throw new IllegalArgumentException(AnalyzerUtility.error(ctx) + "Variable [" + id + "] is read-only.");
                 }
-            }
-
-            // track if the _score value is ever used, we will invoke Scorer.score() only once if so.
-            if ("_score".equals(id)) {
-                metadata.scoreValueUsed = true;
+                if ("_score".equals(id)) {
+                    // track if the _score value is ever used, we will invoke Scorer.score() only once if so.
+                    metadata.scoreValueUsed = true;
+                } else if ("ctx".equals(id)) {
+                    // track if ctx value is ever used, we will invoke Map.get() only once if so.
+                    metadata.ctxValueUsed = true;
+                }
             }
 
             varenmd.target = variable.slot;

modules/lang-painless/src/main/java/org/elasticsearch/painless/Metadata.java

@@ -429,8 +429,8 @@ class Metadata {
     int scoreValueSlot = -1;
 
     /**
-     * Used to determine if the _score variable is actually used.  This is used in the {@link Analyzer} to update
-     * variable slots at the completion of analysis if _score is not used.
+     * Used to determine if the _score variable is actually used.  This is used to know if we should call
+     * Scorer.score() once and cache into a local variable, and expose NeedsScore interface (to allow query caching)
      */
     boolean scoreValueUsed = false;
     
@@ -439,6 +439,18 @@ class Metadata {
      * the doc variable is accessed.
      */
     int docValueSlot = -1;
+    
+    /**
+     * Used to determine what slot the ctx variable is stored in.  This is used in the {@link Writer} whenever
+     * the ctx variable is accessed.
+     */
+    int ctxValueSlot = -1;
+    
+    /**
+     * Used to determine if the ctx variable is actually used.  This is used to determine if we should call
+     * Map.get once and store into a local variable on startup.
+     */
+    boolean ctxValueUsed = false;
 
     /**
      * Maps the relevant ANTLR node to its metadata.

modules/lang-painless/src/main/java/org/elasticsearch/painless/Writer.java

@@ -157,6 +157,15 @@ class Writer extends PainlessParserBaseVisitor<Void> {
             execute.invokeVirtual(WriterConstants.SCORER_TYPE, WriterConstants.SCORER_SCORE);
             execute.visitVarInsn(Opcodes.FSTORE, metadata.scoreValueSlot);
         }
+        
+        if (metadata.ctxValueUsed) {
+            // if the _ctx value is used, we do this once:
+            //   Map<String,Object> ctx = input.get("ctx");
+            execute.visitVarInsn(Opcodes.ALOAD, metadata.inputValueSlot);
+            execute.push("ctx");
+            execute.invokeInterface(WriterConstants.MAP_TYPE, WriterConstants.MAP_GET);
+            execute.visitVarInsn(Opcodes.ASTORE, metadata.ctxValueSlot);
+        }
 
         execute.push(settings.getMaxLoopCounter());
         execute.visitVarInsn(Opcodes.ISTORE, metadata.loopCounterSlot);

modules/lang-painless/src/main/java/org/elasticsearch/painless/WriterExternal.java

@@ -422,10 +422,6 @@ class WriterExternal {
             throw new IllegalStateException(WriterUtility.error(source) + "Cannot load/store void type.");
         }
 
-        if (!metadata.scoreValueUsed && slot > metadata.scoreValueSlot) {
-            --slot;
-        }
-
         if (store) {
             execute.visitVarInsn(type.type.getOpcode(Opcodes.ISTORE), slot);
         } else {

modules/lang-painless/src/main/java/org/elasticsearch/painless/WriterStatement.java

@@ -328,10 +328,6 @@ class WriterStatement {
         final Sort sort = declvaremd.to.sort;
         int slot = (int)declvaremd.postConst;
 
-        if (!metadata.scoreValueUsed && slot > metadata.scoreValueSlot) {
-            --slot;
-        }
-
         final ExpressionContext exprctx = ctx.expression();
         final boolean initialize = exprctx == null;
 

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

@@ -19,6 +19,9 @@
 
 package org.elasticsearch.painless;
 
+import java.util.Collections;
+import java.util.HashMap;
+
 /** Tests for special reserved words such as _score */
 public class ReservedWordTests extends ScriptTestCase {
     
@@ -46,11 +49,32 @@ public class ReservedWordTests extends ScriptTestCase {
         assertTrue(expected.getMessage().contains("Variable name [doc] already defined"));
     }
     
-    /** check that we can't write to _score, its read-only! */
+    /** check that we can't write to doc, its read-only! */
     public void testDocStore() {
         IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
             exec("doc = 5; return doc;");
         });
         assertTrue(expected.getMessage().contains("Variable [doc] is read-only"));
     }
+    
+    /** check that we can't declare a variable of ctx, its really reserved! */
+    public void testCtxVar() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("int ctx = 5; return ctx;");
+        });
+        assertTrue(expected.getMessage().contains("Variable name [ctx] already defined"));
+    }
+    
+    /** check that we can't write to ctx, its read-only! */
+    public void testCtxStore() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("ctx = 5; return ctx;");
+        });
+        assertTrue(expected.getMessage().contains("Variable [ctx] is read-only"));
+    }
+    
+    /** check that we can modify its contents though */
+    public void testCtxStoreMap() {
+        assertEquals(5, exec("ctx.foo = 5; return ctx.foo;", Collections.singletonMap("ctx", new HashMap<String,Object>())));
+    }
 }

modules/lang-painless/src/test/resources/rest-api-spec/test/plan_a/15_update.yaml

@@ -18,7 +18,7 @@
           script: "1"
           body:
             lang:   painless
-            script: "input.ctx._source.foo = input.bar"
+            script: "ctx._source.foo = input.bar"
             params: { bar: 'xxx' }
 
   - match: { _index:   test_1 }
@@ -41,7 +41,7 @@
           type:   test
           id:     1
           lang:   painless
-          script: "input.ctx._source.foo = 'yyy'"
+          script: "ctx._source.foo = 'yyy'"
 
   - match: { _index:   test_1 }
   - match: { _type:    test   }

modules/lang-painless/src/test/resources/rest-api-spec/test/plan_a/25_script_upsert.yaml

@@ -7,7 +7,7 @@
           type:     test
           id:       1
           body:
-            script: "input.ctx._source.foo = input.bar"
+            script: "ctx._source.foo = input.bar"
             lang: "painless"
             params: { bar: 'xxx' }
             upsert: { foo: baz }
@@ -27,7 +27,7 @@
           type:     test
           id:       1
           body:
-            script: "input.ctx._source.foo = input.bar"
+            script: "ctx._source.foo = input.bar"
             lang: "painless"
             params: { bar: 'xxx' }
             upsert: { foo: baz }
@@ -46,7 +46,7 @@
           type:     test
           id:       2
           body:
-            script: "input.ctx._source.foo = input.bar"
+            script: "ctx._source.foo = input.bar"
             lang: "painless"
             params: { bar: 'xxx' }
             upsert: { foo: baz }