Write an alternative expression compiler using MethodHandles

worldedit-core/build.gradle.kts

@@ -26,6 +26,7 @@ dependencies {
     "compile"("com.google.code.gson:gson:2.8.0")
     "compile"("org.slf4j:slf4j-api:1.7.26")
     "compile"("it.unimi.dsi:fastutil:8.2.1")
+    "compile"("org.ow2.asm:asm:6.2")
 
     val antlrVersion = "4.7.2"
     "antlr"("org.antlr:antlr4:$antlrVersion")

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/BreakException.java

@@ -25,9 +25,12 @@
  */
 public class BreakException extends RuntimeException {
 
+    public static final BreakException BREAK = new BreakException(false);
+    public static final BreakException CONTINUE = new BreakException(true);
+
     public final boolean doContinue;
 
-    public BreakException(boolean doContinue) {
+    private BreakException(boolean doContinue) {
         super(doContinue ? "'continue' encountered outside a loop" : "'break' encountered outside a loop",
             null, true, false);
 

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/CompiledExpression.java

@@ -0,0 +1,10 @@
+package com.sk89q.worldedit.internal.expression;
+
+/**
+ * Represents a "compiled" expression.
+ */
+public interface CompiledExpression {
+
+    Double execute(ExecutionData executionData);
+
+}

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/EvaluatingVisitor.java

@@ -238,12 +238,12 @@ public Double visitSimpleForStatement(ExpressionParser.SimpleForStatementContext
 
     @Override
     public Double visitBreakStatement(ExpressionParser.BreakStatementContext ctx) {
-        throw new BreakException(false);
+        throw BreakException.BREAK;
     }
 
     @Override
     public Double visitContinueStatement(ExpressionParser.ContinueStatementContext ctx) {
-        throw new BreakException(true);
+        throw BreakException.CONTINUE;
     }
 
     @Override

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/EvaluationException.java

@@ -19,8 +19,6 @@
 
 package com.sk89q.worldedit.internal.expression;
 
-import com.sk89q.worldedit.internal.expression.ExpressionException;
-
 /**
  * Thrown when there's a problem during expression evaluation.
  */

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/ExecutionData.java

@@ -0,0 +1,32 @@
+package com.sk89q.worldedit.internal.expression;
+
+import com.google.common.collect.SetMultimap;
+
+import java.lang.invoke.MethodHandle;
+
+import static java.util.Objects.requireNonNull;
+
+public class ExecutionData {
+
+    /**
+     * Special execution context for evaluating constant values. As long as no variables are used,
+     * it can be considered constant.
+     */
+    public static final ExecutionData CONSTANT_EVALUATOR = new ExecutionData(null, null);
+
+    private final SlotTable slots;
+    private final SetMultimap<String, MethodHandle> functions;
+
+    public ExecutionData(SlotTable slots, SetMultimap<String, MethodHandle> functions) {
+        this.slots = slots;
+        this.functions = functions;
+    }
+
+    public SlotTable getSlots() {
+        return requireNonNull(slots, "Cannot use variables in a constant");
+    }
+
+    public SetMultimap<String, MethodHandle> getFunctions() {
+        return requireNonNull(functions, "Cannot use functions in a constant");
+    }
+}

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/Expression.java

@@ -26,6 +26,7 @@
 import com.sk89q.worldedit.WorldEdit;
 import com.sk89q.worldedit.antlr.ExpressionLexer;
 import com.sk89q.worldedit.antlr.ExpressionParser;
+import com.sk89q.worldedit.internal.expression.invoke.ExpressionCompiler;
 import com.sk89q.worldedit.session.request.Request;
 import org.antlr.v4.runtime.CharStream;
 import org.antlr.v4.runtime.CharStreams;
@@ -83,9 +84,10 @@
 
     private final SlotTable slots = new SlotTable();
     private final List<String> providedSlots;
-    private ExpressionParser.AllStatementsContext root;
+    private final ExpressionParser.AllStatementsContext root;
     private final SetMultimap<String, MethodHandle> functions = Functions.getFunctionMap();
     private ExpressionEnvironment environment;
+    private final CompiledExpression compiledExpression;
 
     public static Expression compile(String expression, String... variableNames) throws ExpressionException {
         return new Expression(expression, variableNames);
@@ -119,6 +121,7 @@ private Expression(String expression, String... variableNames) throws Expression
             throw new ParserException(parser.getState(), e);
         }
         ParseTreeWalker.DEFAULT.walk(new ExpressionValidator(slots.keySet(), functions), root);
+        this.compiledExpression = new ExpressionCompiler().compileExpression(root, functions);
     }
 
     public double evaluate(double... values) throws EvaluationException {
@@ -177,7 +180,7 @@ private double evaluateRootTimed(int timeout) throws EvaluationException {
     private Double evaluateRoot() throws EvaluationException {
         pushInstance();
         try {
-            return root.accept(new EvaluatingVisitor(slots, functions));
+            return compiledExpression.execute(new ExecutionData(slots, functions));
         } finally {
             popInstance();
         }

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/ExpressionHelper.java

@@ -33,33 +33,41 @@
 
 import static com.sk89q.worldedit.antlr.ExpressionLexer.ID;
 
-class ExpressionHelper {
+public class ExpressionHelper {
 
-    static void check(boolean condition, ParserRuleContext ctx, String message) {
+    public static void check(boolean condition, ParserRuleContext ctx, String message) {
         if (!condition) {
             throw evalException(ctx, message);
         }
     }
 
-    static EvaluationException evalException(ParserRuleContext ctx, String message) {
+    public static int getErrorPosition(Token token) {
+        return token.getCharPositionInLine();
+    }
+
+    public static EvaluationException evalException(ParserRuleContext ctx, String message) {
+        return evalException(ctx.start, message);
+    }
+
+    public static EvaluationException evalException(Token token, String message) {
         return new EvaluationException(
-            ctx.getStart().getCharPositionInLine(),
+            getErrorPosition(token),
             message
         );
     }
 
-    static void checkIterations(int iterations, ParserRuleContext ctx) {
+    public static void checkIterations(int iterations, ParserRuleContext ctx) {
         check(iterations <= 256, ctx, "Loop exceeded 256 iterations");
     }
 
-    static void checkTimeout() {
+    public static void checkTimeout() {
         if (Thread.interrupted()) {
             throw new ExpressionTimeoutException("Calculations exceeded time limit.");
         }
     }
 
-    static MethodHandle resolveFunction(SetMultimap<String, MethodHandle> functions,
-                                        ExpressionParser.FunctionCallContext ctx) {
+    public static MethodHandle resolveFunction(SetMultimap<String, MethodHandle> functions,
+                                               ExpressionParser.FunctionCallContext ctx) {
         String fnName = ctx.name.getText();
         Set<MethodHandle> matchingFns = functions.get(fnName);
         check(!matchingFns.isEmpty(), ctx, "Unknown function '" + fnName + "'");
@@ -92,14 +100,14 @@ static MethodHandle resolveFunction(SetMultimap<String, MethodHandle> functions,
     /**
      * The argument should be wrapped in a {@link LocalSlot.Constant} before being passed.
      */
-    static final String WRAPPED_CONSTANT = "<wrapped constant>";
+    public static final String WRAPPED_CONSTANT = "<wrapped constant>";
 
     /**
      * If this argument needs a handle, returns the name of the handle needed. Otherwise, returns
      * {@code null}. If {@code arg} isn't a valid handle reference, throws.
      */
-    static String getArgumentHandleName(String fnName, MethodType type, int i,
-                                        ParserRuleContext arg) {
+    public static String getArgumentHandleName(String fnName, MethodType type, int i,
+                                               ParserRuleContext arg) {
         // Pass variable handle in for modification?
         Class<?> pType = type.parameterType(i);
         Optional<String> id = tryResolveId(arg);

worldedit-core/src/main/java/com/sk89q/worldedit/internal/expression/Functions.java

@@ -22,6 +22,7 @@
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.collect.Multimaps;
 import com.google.common.collect.SetMultimap;
 import com.google.common.primitives.Doubles;
 import com.sk89q.worldedit.internal.expression.LocalSlot.Variable;
@@ -36,6 +37,8 @@
 import java.lang.invoke.MethodHandles;
 import java.util.concurrent.ThreadLocalRandom;
 
+import static com.google.common.base.Preconditions.checkState;
+import static java.lang.invoke.MethodHandles.filterReturnValue;
 import static java.lang.invoke.MethodType.methodType;
 
 /**
@@ -56,7 +59,35 @@
             throw new IllegalStateException(e);
         }
 
-        return ImmutableSetMultimap.copyOf(map);
+        // clean up all the functions
+        return ImmutableSetMultimap.copyOf(
+            Multimaps.transformValues(map, Functions::clean)
+        );
+    }
+
+    private static final MethodHandle DOUBLE_VALUE;
+
+    static {
+        MethodHandles.Lookup lookup = MethodHandles.lookup();
+        try {
+            DOUBLE_VALUE = lookup.findVirtual(Number.class, "doubleValue",
+                methodType(double.class));
+        } catch (NoSuchMethodException | IllegalAccessException e) {
+            throw new IllegalStateException(e);
+        }
+    }
+
+    private static MethodHandle clean(MethodHandle handle) {
+        // box it all first
+        handle = handle.asType(handle.type().wrap());
+        if (handle.type().returnType() != Double.class) {
+            // Ensure that the handle returns a Double, even if originally a Number
+            checkState(Number.class.isAssignableFrom(handle.type().returnType()),
+                "Function does not return a number");
+            handle = handle.asType(handle.type().changeReturnType(Number.class));
+            handle = filterReturnValue(handle, DOUBLE_VALUE);
+        }
+        return handle;
     }
 
     private static void addMathHandles(