Merge pull request #18258 from rmuir/painless_faster_score

painless: retrieve _score directly from Scorer

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

@@ -94,6 +94,7 @@ class Analyzer extends PainlessParserBaseVisitor<Void> {
         utility.incrementScope();
         utility.addVariable(null, "#this", definition.execType);
         metadata.inputValueSlot = utility.addVariable(null, "input", definition.smapType).slot;
+        metadata.scorerValueSlot = utility.addVariable(null, "#scorer", definition.objectType).slot;
         metadata.loopCounterSlot = utility.addVariable(null, "#loop", definition.intType).slot;
         metadata.scoreValueSlot = utility.addVariable(null, "_score", definition.floatType).slot;
 

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

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.painless;
 
+import org.apache.lucene.search.Scorer;
+
 import java.util.Map;
 
 public abstract class Executable {
@@ -46,5 +48,5 @@ public abstract class Executable {
         return definition;
     }
 
-    public abstract Object execute(Map<String, Object> input);
+    public abstract Object execute(Map<String, Object> input, Scorer scorer);
 }

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

@@ -410,6 +410,12 @@ class Metadata {
      */
     int inputValueSlot = -1;
 
+    /**
+     * Used to determine what slot the Scorer variable is stored in.  This is used in the {@link Writer} to load
+     * _score from it, if _score will be accessed by the script.
+     */
+    int scorerValueSlot = -1;
+
     /**
      * Used to determine what slot the loopCounter variable is stored in.  This is used n the {@link Writer} whenever
      * the loop variable is accessed.

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

@@ -22,7 +22,6 @@ package org.elasticsearch.painless;
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.LeafSearchScript;
-import org.elasticsearch.script.ScoreAccessor;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 
 import java.util.HashMap;
@@ -48,6 +47,12 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     private final LeafSearchLookup lookup;
 
+    /**
+     * Current scorer being used
+     * @see #setScorer(Scorer)
+     */
+    private Scorer scorer;
+
     /**
      * Creates a ScriptImpl for the a previously compiled Painless script.
      * @param executable The previously compiled Painless script.
@@ -84,7 +89,7 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     @Override
     public Object run() {
-        return executable.execute(variables);
+        return executable.execute(variables, scorer);
     }
 
     /**
@@ -120,7 +125,7 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     @Override
     public void setScorer(final Scorer scorer) {
-        variables.put("#score", new ScoreAccessor(scorer));
+        this.scorer = scorer;
     }
 
     /**

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

@@ -72,7 +72,6 @@ import org.elasticsearch.painless.PainlessParser.TryContext;
 import org.elasticsearch.painless.PainlessParser.UnaryContext;
 import org.elasticsearch.painless.PainlessParser.WhileContext;
 import org.objectweb.asm.ClassWriter;
-import org.objectweb.asm.Label;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.commons.GeneratorAdapter;
 
@@ -80,10 +79,6 @@ import static org.elasticsearch.painless.WriterConstants.BASE_CLASS_TYPE;
 import static org.elasticsearch.painless.WriterConstants.CLASS_TYPE;
 import static org.elasticsearch.painless.WriterConstants.CONSTRUCTOR;
 import static org.elasticsearch.painless.WriterConstants.EXECUTE;
-import static org.elasticsearch.painless.WriterConstants.MAP_GET;
-import static org.elasticsearch.painless.WriterConstants.MAP_TYPE;
-import static org.elasticsearch.painless.WriterConstants.SCORE_ACCESSOR_FLOAT;
-import static org.elasticsearch.painless.WriterConstants.SCORE_ACCESSOR_TYPE;
 
 class Writer extends PainlessParserBaseVisitor<Void> {
     static byte[] write(Metadata metadata) {
@@ -155,22 +150,11 @@ class Writer extends PainlessParserBaseVisitor<Void> {
     }
 
     private void writeExecute() {
-        final Label fals = new Label();
-        final Label end = new Label();
-
         if (metadata.scoreValueUsed) {
-            execute.visitVarInsn(Opcodes.ALOAD, metadata.inputValueSlot);
+            // if the _score value is used, we do this once:
-            execute.push("#score");
+            //   float _score = scorer.score();
-            execute.invokeInterface(MAP_TYPE, MAP_GET);
+            execute.visitVarInsn(Opcodes.ALOAD, metadata.scorerValueSlot);
-            execute.dup();
+            execute.invokeVirtual(WriterConstants.SCORER_TYPE, WriterConstants.SCORER_SCORE);
-            execute.ifNull(fals);
-            execute.checkCast(SCORE_ACCESSOR_TYPE);
-            execute.invokeVirtual(SCORE_ACCESSOR_TYPE, SCORE_ACCESSOR_FLOAT);
-            execute.goTo(end);
-            execute.mark(fals);
-            execute.pop();
-            execute.push(0F);
-            execute.mark(end);
             execute.visitVarInsn(Opcodes.FSTORE, metadata.scoreValueSlot);
         }
 

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

@@ -19,7 +19,7 @@
 
 package org.elasticsearch.painless;
 
-import org.elasticsearch.script.ScoreAccessor;
+import org.apache.lucene.search.Scorer;
 import org.objectweb.asm.Handle;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
@@ -37,7 +37,7 @@ class WriterConstants {
     final static Type CLASS_TYPE        = Type.getType("L" + CLASS_NAME.replace(".", "/") + ";");
 
     final static Method CONSTRUCTOR = getAsmMethod(void.class, "<init>", Definition.class, String.class, String.class);
-    final static Method EXECUTE     = getAsmMethod(Object.class, "execute", Map.class);
+    final static Method EXECUTE     = getAsmMethod(Object.class, "execute", Map.class, Scorer.class);
 
     final static Type PAINLESS_ERROR_TYPE = Type.getType(PainlessError.class);
 
@@ -47,12 +47,12 @@ class WriterConstants {
 
     final static Type OBJECT_TYPE = Type.getType(Object.class);
 
+    final static Type SCORER_TYPE = Type.getType(Scorer.class);
+    final static Method SCORER_SCORE = getAsmMethod(float.class, "score");
+
     final static Type MAP_TYPE  = Type.getType(Map.class);
     final static Method MAP_GET = getAsmMethod(Object.class, "get", Object.class);
 
-    final static Type SCORE_ACCESSOR_TYPE    = Type.getType(ScoreAccessor.class);
-    final static Method SCORE_ACCESSOR_FLOAT = getAsmMethod(float.class, "floatValue");
-
     /** dynamic callsite bootstrap signature */
     final static MethodType DEF_BOOTSTRAP_TYPE = MethodType.methodType(CallSite.class, MethodHandles.Lookup.class, 
                                                                        String.class, MethodType.class, int.class);