Merge pull request #18232 from uschindler/array_invoke_dynamic

painless: Array load/store and length with invokedynamic

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

@@ -26,9 +26,12 @@ import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.lang.invoke.MethodHandles.Lookup;
-import java.lang.reflect.Array;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * Support for dynamic type (def).
@@ -36,18 +39,13 @@ import java.util.Map;
  * Dynamic types can invoke methods, load/store fields, and be passed as parameters to operators without 
  * compile-time type information. 
  * <p>
- * Dynamic methods, loads, and stores involve locating the appropriate field or method depending
- * on the receiver's class. For these, we emit an {@code invokedynamic} instruction that, for each new 
- * type encountered will query a corresponding {@code lookupXXX} method to retrieve the appropriate method.
- * In most cases, the {@code lookupXXX} methods here will only be called once for a given call site, because 
+ * Dynamic methods, loads, stores, and array/list/map load/stores involve locating the appropriate field
+ * or method depending on the receiver's class. For these, we emit an {@code invokedynamic} instruction that,
+ * for each new  type encountered will query a corresponding {@code lookupXXX} method to retrieve the appropriate
+ * method. In most cases, the {@code lookupXXX} methods here will only be called once for a given call site, because 
  * caching ({@link DynamicCallSite}) generally works: usually all objects at any call site will be consistently 
  * the same type (or just a few types).  In extreme cases, if there is type explosion, they may be called every 
  * single time, but simplicity is still more valuable than performance in this code.
- * <p>
- * Dynamic array loads and stores and operator functions (e.g. {@code +}) are called directly
- * with {@code invokestatic}. Because these features cannot be overloaded in painless, they are hardcoded 
- * decision trees based on the only types that are possible. This keeps overhead low, and seems to be as fast
- * on average as the more adaptive methodhandle caching. 
  */
 public class Def {
 
@@ -96,8 +94,6 @@ public class Def {
                                            "for class [" + receiverClass.getCanonicalName() + "].");
     }
 
-    /** pointer to Array.getLength(Object) */
-    private static final MethodHandle ARRAY_LENGTH;
     /** pointer to Map.get(Object) */
     private static final MethodHandle MAP_GET;
     /** pointer to Map.put(Object,Object) */
@@ -109,9 +105,6 @@ public class Def {
     static {
         Lookup lookup = MethodHandles.publicLookup();
         try {
-            // TODO: maybe specialize handles for different array types. this may be slower, but simple :)
-            ARRAY_LENGTH = lookup.findStatic(Array.class, "getLength",
-                                             MethodType.methodType(int.class, Object.class));
             MAP_GET      = lookup.findVirtual(Map.class, "get",
                                              MethodType.methodType(Object.class, Object.class));
             MAP_PUT      = lookup.findVirtual(Map.class, "put",
@@ -125,6 +118,54 @@ public class Def {
         }
     }
     
+    // TODO: Once Java has a factory for those in java.lang.invoke.MethodHandles, use it:
+
+    /** Helper class for isolating MethodHandles and methods to get the length of arrays
+     * (to emulate a "arraystore" byteoode using MethodHandles).
+     * This should really be a method in {@link MethodHandles} class!
+     */
+    private static final class ArrayLengthHelper {
+      private ArrayLengthHelper() {}
+    
+      private static final Lookup PRIV_LOOKUP = MethodHandles.lookup();
+      private static final Map<Class<?>,MethodHandle> ARRAY_TYPE_MH_MAPPING = Collections.unmodifiableMap(
+        Stream.of(boolean[].class, byte[].class, short[].class, int[].class, long[].class,
+        char[].class, float[].class, double[].class, Object[].class)
+          .collect(Collectors.toMap(Function.identity(), type -> {
+            try {
+              return PRIV_LOOKUP.findStatic(PRIV_LOOKUP.lookupClass(), "getArrayLength", MethodType.methodType(int.class, type));
+            } catch (ReflectiveOperationException e) {
+              throw new AssertionError(e);
+            }
+          }))
+      );
+      private static final MethodHandle OBJECT_ARRAY_MH = ARRAY_TYPE_MH_MAPPING.get(Object[].class);
+     
+      static int getArrayLength(boolean[] array) { return array.length; }
+      static int getArrayLength(byte[] array) { return array.length; }
+      static int getArrayLength(short[] array) { return array.length; }
+      static int getArrayLength(int[] array) { return array.length; }
+      static int getArrayLength(long[] array) { return array.length; }
+      static int getArrayLength(char[] array) { return array.length; }
+      static int getArrayLength(float[] array) { return array.length; }
+      static int getArrayLength(double[] array) { return array.length; }
+      static int getArrayLength(Object[] array) { return array.length; }
+      
+      public static MethodHandle arrayLengthGetter(Class<?> arrayType) {
+        if (!arrayType.isArray()) {
+          throw new IllegalArgumentException("type must be an array");
+        }
+        return (ARRAY_TYPE_MH_MAPPING.containsKey(arrayType)) ?
+            ARRAY_TYPE_MH_MAPPING.get(arrayType) :
+            OBJECT_ARRAY_MH.asType(OBJECT_ARRAY_MH.type().changeParameterType(0, arrayType));
+      }
+    }
+    
+    /** Returns an array length getter MethodHandle for the given array type */
+    public static MethodHandle arrayLengthGetter(Class<?> arrayType) {
+      return ArrayLengthHelper.arrayLengthGetter(arrayType);
+    }
+    
     /** 
      * Looks up handle for a dynamic field getter (field load)
      * <p>
@@ -177,7 +218,7 @@ public class Def {
         // special case: arrays, maps, and lists
         if (receiverClass.isArray() && "length".equals(name)) {
             // arrays expose .length as a read-only getter
-            return ARRAY_LENGTH;
+            return arrayLengthGetter(receiverClass);
         } else if (Map.class.isAssignableFrom(receiverClass)) {
             // maps allow access like mymap.key
             // wire 'key' as a parameter, its a constant in painless
@@ -266,57 +307,46 @@ public class Def {
                                            "for class [" + receiverClass.getCanonicalName() + "].");
     }
 
-    // NOTE: below methods are not cached, instead invoked directly because they are performant.
-
     /**
-     * Performs an actual array store.
-     * @param array array object
-     * @param index map key, array index (integer), or list index (integer)
-     * @param value value to store in the array.
+     * Returns a method handle to do an array store.
+     * @param receiverClass Class of the array to store the value in
+     * @return a MethodHandle that accepts the receiver as first argument, the index as second argument,
+     *   and the value to set as 3rd argument. Return value is undefined and should be ignored.
      */
-    @SuppressWarnings({ "unchecked", "rawtypes" })
-    public static void arrayStore(final Object array, Object index, Object value) {
-        if (array instanceof Map) {
-            ((Map)array).put(index, value);
-        } else if (array.getClass().isArray()) {
-            try {
-                Array.set(array, (int)index, value);
-            } catch (final Throwable throwable) {
-                throw new IllegalArgumentException("Error storing value [" + value + "] " +
-                                                   "in array class [" + array.getClass().getCanonicalName() + "].", throwable);
-            }
-        } else if (array instanceof List) {
-            ((List)array).set((int)index, value);
-        } else {
-            throw new IllegalArgumentException("Attempting to address a non-array type " +
-                                               "[" + array.getClass().getCanonicalName() + "] as an array.");
+    static MethodHandle lookupArrayStore(Class<?> receiverClass) {
+        if (receiverClass.isArray()) {
+            return MethodHandles.arrayElementSetter(receiverClass);
+        } else if (Map.class.isAssignableFrom(receiverClass)) {
+            // maps allow access like mymap[key]
+            return MAP_PUT;
+        } else if (List.class.isAssignableFrom(receiverClass)) {
+            return LIST_SET;
         }
+        throw new IllegalArgumentException("Attempting to address a non-array type " +
+                                           "[" + receiverClass.getCanonicalName() + "] as an array.");
     }
-    
+   
     /**
-     * Performs an actual array load.
-     * @param array array object
-     * @param index map key, array index (integer), or list index (integer)
+     * Returns a method handle to do an array load.
+     * @param receiverClass Class of the array to load the value from
+     * @return a MethodHandle that accepts the receiver as first argument, the index as second argument.
+     *   It returns the loaded value.
      */
-    @SuppressWarnings("rawtypes")
-    public static Object arrayLoad(final Object array, Object index) {
-        if (array instanceof Map) {
-            return ((Map)array).get(index);
-        } else if (array.getClass().isArray()) {
-            try {
-                return Array.get(array, (int)index);
-            } catch (final Throwable throwable) {
-                throw new IllegalArgumentException("Error loading value from " +
-                                                   "array class [" + array.getClass().getCanonicalName() + "].", throwable);
-            }
-        } else if (array instanceof List) {
-            return ((List)array).get((int)index);
-        } else {
-            throw new IllegalArgumentException("Attempting to address a non-array type " +
-                                               "[" + array.getClass().getCanonicalName() + "] as an array.");
+    static MethodHandle lookupArrayLoad(Class<?> receiverClass) {
+        if (receiverClass.isArray()) {
+            return MethodHandles.arrayElementGetter(receiverClass);
+        } else if (Map.class.isAssignableFrom(receiverClass)) {
+            // maps allow access like mymap[key]
+            return MAP_GET;
+        } else if (List.class.isAssignableFrom(receiverClass)) {
+            return LIST_GET;
         }
+        throw new IllegalArgumentException("Attempting to address a non-array type " +
+                                           "[" + receiverClass.getCanonicalName() + "] as an array.");
     }
 
+    // NOTE: below methods are not cached, instead invoked directly because they are performant.
+
     public static Object not(final Object unary) {
         if (unary instanceof Double || unary instanceof Float || unary instanceof Long) {
             return ~((Number)unary).longValue();

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

@@ -48,6 +48,10 @@ public final class DynamicCallSite {
     static final int LOAD = 1;
     /** static bootstrap parameter indicating a dynamic store (setter), e.g. foo.bar = baz */
     static final int STORE = 2;
+    /** static bootstrap parameter indicating a dynamic array load, e.g. baz = foo[bar] */
+    static final int ARRAY_LOAD = 3;
+    /** static bootstrap parameter indicating a dynamic array store, e.g. foo[bar] = baz */
+    static final int ARRAY_STORE = 4;
     
     static class InliningCacheCallSite extends MutableCallSite {
         /** maximum number of types before we go megamorphic */
@@ -104,6 +108,10 @@ public final class DynamicCallSite {
                 return Def.lookupGetter(clazz, name, Definition.INSTANCE);
             case STORE:
                 return Def.lookupSetter(clazz, name, Definition.INSTANCE);
+            case ARRAY_LOAD: 
+                return Def.lookupArrayLoad(clazz);
+            case ARRAY_STORE:
+                return Def.lookupArrayStore(clazz);
             default: throw new AssertionError();
         }
     }

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

@@ -58,10 +58,14 @@ class WriterConstants {
     final static Handle DEF_BOOTSTRAP_HANDLE = new Handle(Opcodes.H_INVOKESTATIC, Type.getInternalName(DynamicCallSite.class),
                                                           "bootstrap", WriterConstants.DEF_BOOTSTRAP_TYPE.toMethodDescriptorString());
 
-    final static Method DEF_ARRAY_STORE = getAsmMethod(
-        void.class, "arrayStore", Object.class, Object.class, Object.class);
-    final static Method DEF_ARRAY_LOAD = getAsmMethod(
-        Object.class, "arrayLoad", Object.class, Object.class);
+    final static String DEF_DYNAMIC_LOAD_FIELD_DESC = MethodType.methodType(Object.class, Object.class)
+        .toMethodDescriptorString();
+    final static String DEF_DYNAMIC_STORE_FIELD_DESC = MethodType.methodType(void.class, Object.class, Object.class)
+        .toMethodDescriptorString();
+    final static String DEF_DYNAMIC_ARRAY_LOAD_DESC = MethodType.methodType(Object.class, Object.class, Object.class)
+        .toMethodDescriptorString();
+    final static String DEF_DYNAMIC_ARRAY_STORE_DESC = MethodType.methodType(void.class, Object.class, Object.class, Object.class)
+        .toMethodDescriptorString();
 
     final static Method DEF_NOT_CALL = getAsmMethod(Object.class, "not", Object.class);
     final static Method DEF_NEG_CALL = getAsmMethod(Object.class, "neg", Object.class);

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

@@ -49,8 +49,6 @@ import static org.elasticsearch.painless.PainlessParser.DIV;
 import static org.elasticsearch.painless.PainlessParser.MUL;
 import static org.elasticsearch.painless.PainlessParser.REM;
 import static org.elasticsearch.painless.PainlessParser.SUB;
-import static org.elasticsearch.painless.WriterConstants.DEF_ARRAY_LOAD;
-import static org.elasticsearch.painless.WriterConstants.DEF_ARRAY_STORE;
 import static org.elasticsearch.painless.WriterConstants.TOBYTEEXACT_INT;
 import static org.elasticsearch.painless.WriterConstants.TOBYTEEXACT_LONG;
 import static org.elasticsearch.painless.WriterConstants.TOBYTEWOOVERFLOW_DOUBLE;
@@ -468,10 +466,10 @@ class WriterExternal {
 
     private void writeLoadStoreField(final ParserRuleContext source, final boolean store, final String name) {
         if (store) {
-            execute.visitInvokeDynamicInsn(name, "(Ljava/lang/Object;Ljava/lang/Object;)V", 
+            execute.visitInvokeDynamicInsn(name, WriterConstants.DEF_DYNAMIC_STORE_FIELD_DESC, 
                                            WriterConstants.DEF_BOOTSTRAP_HANDLE, new Object[] { DynamicCallSite.STORE });
         } else {
-            execute.visitInvokeDynamicInsn(name, "(Ljava/lang/Object;)Ljava/lang/Object;", 
+            execute.visitInvokeDynamicInsn(name, WriterConstants.DEF_DYNAMIC_LOAD_FIELD_DESC, 
                                            WriterConstants.DEF_BOOTSTRAP_HANDLE, new Object[] { DynamicCallSite.LOAD });
         }
     }
@@ -483,9 +481,11 @@ class WriterExternal {
 
         if (type.sort == Sort.DEF) {
             if (store) {
-                execute.invokeStatic(definition.defobjType.type, DEF_ARRAY_STORE);
+                execute.visitInvokeDynamicInsn("arrayStore", WriterConstants.DEF_DYNAMIC_ARRAY_STORE_DESC, 
+                                           WriterConstants.DEF_BOOTSTRAP_HANDLE, new Object[] { DynamicCallSite.ARRAY_STORE });
             } else {
-                execute.invokeStatic(definition.defobjType.type, DEF_ARRAY_LOAD);
+                execute.visitInvokeDynamicInsn("arrayLoad", WriterConstants.DEF_DYNAMIC_ARRAY_LOAD_DESC, 
+                                           WriterConstants.DEF_BOOTSTRAP_HANDLE, new Object[] { DynamicCallSite.ARRAY_LOAD });
             }
         } else {
             if (store) {

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

@@ -0,0 +1,62 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless;
+
+/** Tests for or operator across all types */
+public class ArrayTests extends ScriptTestCase {
+
+    public void testArrayLengthHelper() throws Throwable {
+        assertArrayLength(2, new int[2]);
+        assertArrayLength(3, new long[3]);
+        assertArrayLength(4, new byte[4]);
+        assertArrayLength(5, new float[5]);
+        assertArrayLength(6, new double[6]);
+        assertArrayLength(7, new char[7]);
+        assertArrayLength(8, new short[8]);
+        assertArrayLength(9, new Object[9]);
+        assertArrayLength(10, new Integer[10]);
+        assertArrayLength(11, new String[11][2]);
+    }
+    
+    private void assertArrayLength(int length, Object array) throws Throwable {
+        assertEquals(length, (int) Def.arrayLengthGetter(array.getClass()).invoke(array));
+    }
+
+    public void testArrayLoadStoreInt() {
+        assertEquals(5, exec("def x = new int[5]; return x.length"));
+        assertEquals(5, exec("def x = new int[4]; x[0] = 5; return x[0];"));
+    }
+    
+    public void testArrayLoadStoreString() {
+        assertEquals(5, exec("def x = new String[5]; return x.length"));
+        assertEquals("foobar", exec("def x = new String[4]; x[0] = 'foobar'; return x[0];"));
+    }
+    
+    public void testArrayLoadStoreDef() {
+        assertEquals(5, exec("def x = new def[5]; return x.length"));
+        assertEquals(5, exec("def x = new def[4]; x[0] = 5; return x[0];"));
+    }
+    
+    public void testForLoop() {
+        assertEquals(999*1000/2, exec("def a = new int[1000]; for (int x = 0; x < a.length; x++) { a[x] = x; } "+
+            "int total = 0; for (int x = 0; x < a.length; x++) { total += a[x]; } return total;"));
+    }
+    
+}

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

@@ -62,12 +62,6 @@ public class BasicAPITests extends ScriptTestCase {
         assertEquals(5, exec("def x = new ArrayList(); x.add(3); x[0] = 5; return x[0];"));
     }
     
-    /** Test loads and stores with a list */
-    public void testArrayLoadStore() {
-        assertEquals(5, exec("def x = new int[5]; return x.length"));
-        assertEquals(5, exec("def x = new int[4]; x[0] = 5; return x[0];"));
-    }
-    
     /** Test shortcut for getters with isXXXX */
     public void testListEmpty() {
         assertEquals(true, exec("def x = new ArrayList(); return x.empty;"));