get compound assignment working

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

@@ -259,19 +259,10 @@ public final class DefBootstrap {
         }
         
         private MethodHandle lookupGeneric() throws Throwable {
-            MethodHandle generic = DefMath.lookupGeneric(name);
             if ((flags & OPERATOR_COMPOUND_ASSIGNMENT) != 0) {
-                assert flavor == BINARY_OPERATOR || flavor == SHIFT_OPERATOR;
+                return DefMath.lookupGenericWithCast(name);
-                // adapt dynamic cast to the generic method and the callsite's return value
-                MethodHandle cast = DYNAMIC_CAST.asType(MethodType.methodType(type().returnType(), 
-                                                                              generic.type().returnType(),
-                                                                              generic.type().parameterType(0)));
-                // drop the RHS parameter
-                cast = MethodHandles.dropArguments(cast, 2, generic.type().parameterType(1));
-                // combine: f(x,y) -> g(f(x,y), x, y);
-                return MethodHandles.foldArguments(cast, generic);
             } else {
-                return generic;
+                return DefMath.lookupGeneric(name);
             }
         }
         
@@ -361,23 +352,9 @@ public final class DefBootstrap {
             return leftObject.getClass() == left && rightObject.getClass() == right;
         }
         
-        /**
-         * Slow dynamic cast: casts {@code returnValue} to the runtime type of {@code lhs}
-         * based upon inspection. If {@code lhs} is null, no cast takes place.
-         * This is used for the generic fallback case of compound assignment.
-         */
-        static Object dynamicCast(Object returnValue, Object lhs) {
-            if (lhs != null) {
-                return lhs.getClass().cast(returnValue);
-            } else {
-                return returnValue;
-            }
-        }
-        
         private static final MethodHandle CHECK_LHS;
         private static final MethodHandle CHECK_RHS;
         private static final MethodHandle CHECK_BOTH;
-        private static final MethodHandle DYNAMIC_CAST;
         private static final MethodHandle FALLBACK;
         static {
             final Lookup lookup = MethodHandles.lookup();
@@ -388,8 +365,6 @@ public final class DefBootstrap {
                                               MethodType.methodType(boolean.class, Class.class, Class.class, Object.class, Object.class));
                 CHECK_BOTH = lookup.findStatic(lookup.lookupClass(), "checkBoth",
                                               MethodType.methodType(boolean.class, Class.class, Class.class, Object.class, Object.class));
-                DYNAMIC_CAST = lookup.findStatic(lookup.lookupClass(), "dynamicCast", 
-                                              MethodType.methodType(Object.class, Object.class, Object.class));
                 FALLBACK = lookup.findVirtual(lookup.lookupClass(), "fallback",
                                               MethodType.methodType(Object.class, Object[].class));
             } catch (ReflectiveOperationException e) {

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

@@ -1146,4 +1146,70 @@ public class DefMath {
     public static MethodHandle lookupGeneric(String name) {
         return TYPE_OP_MAPPING.get(Object.class).get(name);
     }
+    
+    
+    /**
+     * Slow dynamic cast: casts {@code returnValue} to the runtime type of {@code lhs}
+     * based upon inspection. If {@code lhs} is null, no cast takes place.
+     * This is used for the generic fallback case of compound assignment.
+     */
+    static Object dynamicCast(Object returnValue, Object lhs) {
+        if (lhs != null) {
+            Class<?> c = lhs.getClass();
+            if (c == Integer.class) {
+                return getNumber(returnValue).intValue();
+            } else if (c == Long.class) {
+                return getNumber(returnValue).longValue();
+            } else if (c == Double.class) {
+                return getNumber(returnValue).doubleValue();
+            } else if (c == Float.class) {
+                return getNumber(returnValue).floatValue();
+            } else if (c == Short.class) {
+                return getNumber(returnValue).shortValue();
+            } else if (c == Byte.class) {
+                return getNumber(returnValue).byteValue();
+            } else if (c == Character.class) {
+                return (char) getNumber(returnValue).intValue();
+            }
+            return lhs.getClass().cast(returnValue);
+        } else {
+            return returnValue;
+        }
+    }
+    
+    /** Slowly returns a Number for o. Just for supporting dynamicCast */
+    static Number getNumber(Object o) {
+        if (o instanceof Number) {
+            return (Number)o;
+        } else if (o instanceof Character) {
+            return Integer.valueOf((char)o);
+        } else {
+            throw new ClassCastException("Cannot convert [" + o.getClass() + "] to a Number");
+        }
+    }
+    
+    private static final MethodHandle DYNAMIC_CAST;
+    static {
+        final Lookup lookup = MethodHandles.lookup();
+        try {
+            DYNAMIC_CAST = lookup.findStatic(lookup.lookupClass(), 
+                                             "dynamicCast", 
+                                             MethodType.methodType(Object.class, Object.class, Object.class));
+        } catch (ReflectiveOperationException e) {
+            throw new AssertionError(e);
+        }
+    }
+
+    /** Looks up generic method, with a dynamic cast to the receiver's type. (compound assignment) */
+    public static MethodHandle lookupGenericWithCast(String name) {
+        MethodHandle generic = lookupGeneric(name);
+        // adapt dynamic cast to the generic method
+        MethodHandle cast = DYNAMIC_CAST.asType(MethodType.methodType(generic.type().returnType(), 
+                                                                      generic.type().returnType(),
+                                                                      generic.type().parameterType(0)));
+        // drop the RHS parameter
+        cast = MethodHandles.dropArguments(cast, 2, generic.type().parameterType(1));
+        // combine: f(x,y) -> g(f(x,y), x, y);
+        return MethodHandles.foldArguments(cast, generic);
+    }
 }

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

@@ -431,4 +431,29 @@ public class AdditionTests extends ScriptTestCase {
         assertEquals(15D, exec("def x = 5.0; x += 10; return x;"));
         assertEquals(-5D, exec("def x = 5.0; x += -10; return x;"));
     }
+    
+    public void testDefCompoundAssignmentRHS() {
+        // byte
+        assertEquals((byte) 15, exec("byte x = 5; def y = 10; x += y; return x;"));
+        assertEquals((byte) -5, exec("byte x = 5; def y = -10; x += y; return x;"));
+
+        // short
+        assertEquals((short) 15, exec("short x = 5; def y = 10; x += y; return x;"));
+        assertEquals((short) -5, exec("short x = 5; def y = -10; x += y; return x;"));
+        // char
+        assertEquals((char) 15, exec("char x = 5; def y = 10; x += y; return x;"));
+        assertEquals((char) 5, exec("char x = 10; def y = -5; x += y; return x;"));
+        // int
+        assertEquals(15, exec("int x = 5; def y = 10; x += y; return x;"));
+        assertEquals(-5, exec("int x = 5; def y = -10; x += y; return x;"));
+        // long
+        assertEquals(15L, exec("long x = 5; def y = 10; x += y; return x;"));
+        assertEquals(-5L, exec("long x = 5; def y = -10; x += y; return x;"));
+        // float
+        assertEquals(15F, exec("float x = 5f; def y = 10; x += y; return x;"));
+        assertEquals(-5F, exec("float x = 5f; def y = -10; x += y; return x;"));
+        // double
+        assertEquals(15D, exec("double x = 5.0; def y = 10; x += y; return x;"));
+        assertEquals(-5D, exec("double x = 5.0; def y = -10; x += y; return x;"));
+    }
 }