some cleanup

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

@@ -235,15 +235,21 @@ public final class DefBootstrap {
         /**
          * Does a slow lookup for the operator
          */
-        private MethodHandle lookup(int flavor, String name, Object[] args) throws Throwable {
+        private MethodHandle lookup(Object[] args) throws Throwable {
             switch(flavor) {
                 case UNARY_OPERATOR:
                 case SHIFT_OPERATOR:
+                    if ((flags & OPERATOR_COMPOUND_ASSIGNMENT) != 0) {
+                        return lookupGeneric(); // XXX: optimize better.
+                    }
                     // shifts are treated as unary, as java allows long arguments without a cast (but bits are ignored)
                     return DefMath.lookupUnary(args[0].getClass(), name);
                 case BINARY_OPERATOR:
+                    if ((flags & OPERATOR_COMPOUND_ASSIGNMENT) != 0) {
+                        return lookupGeneric(); // XXX: optimize better.
+                    }
                     if (args[0] == null || args[1] == null) {
-                        return DefMath.lookupGeneric(name); // can handle nulls, if supported
+                        return lookupGeneric(); // can handle nulls, if supported
                     } else {
                         return DefMath.lookupBinary(args[0].getClass(), args[1].getClass(), name);
                     }
@@ -251,6 +257,23 @@ 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;
+                // 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;
+            }
+        }
+        
         /**
          * Called when a new type is encountered (or, when we have encountered more than {@code MAX_DEPTH}
          * types at this call site and given up on caching).
@@ -259,13 +282,13 @@ public final class DefBootstrap {
         Object fallback(Object[] args) throws Throwable {
             if (initialized) {
                 // caching defeated
-                MethodHandle generic = DefMath.lookupGeneric(name);
+                MethodHandle generic = lookupGeneric();
                 setTarget(generic.asType(type()));
                 return generic.invokeWithArguments(args);
             }
             
             final MethodType type = type();
-            final MethodHandle target = lookup(flavor, name, args).asType(type);
+            final MethodHandle target = lookup(args).asType(type);
 
             final MethodHandle test;
             if (flavor == BINARY_OPERATOR || flavor == SHIFT_OPERATOR) {
@@ -301,7 +324,7 @@ public final class DefBootstrap {
             // very special cases, where even the receiver can be null (see JLS rules for string concat)
             // we wrap + with an NPE catcher, and use our generic method in that case.
             if (flavor == BINARY_OPERATOR && (flags & OPERATOR_ALLOWS_NULL) != 0) {
-                MethodHandle handler = MethodHandles.dropArguments(DefMath.lookupGeneric(name).asType(type()), 
+                MethodHandle handler = MethodHandles.dropArguments(lookupGeneric().asType(type()), 
                                                                    0, 
                                                                    NullPointerException.class);
                 guard = MethodHandles.catchException(guard, NullPointerException.class, handler);
@@ -321,7 +344,6 @@ public final class DefBootstrap {
             return leftObject.getClass() == clazz;
         }
 
-        
         /**
          * guard method for inline caching: checks the first argument is the same
          * as the cached first argument.
@@ -338,9 +360,23 @@ 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();
@@ -351,6 +387,8 @@ 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) {
@@ -416,6 +454,10 @@ public final class DefBootstrap {
                     // we just don't need it anywhere else.
                     throw new BootstrapMethodError("This parameter is only supported for BINARY_OPERATORs");
                 }
+                if ((flags & OPERATOR_COMPOUND_ASSIGNMENT) != 0 && flavor != BINARY_OPERATOR) {
+                    // we just don't need it anywhere else.
+                    throw new BootstrapMethodError("This parameter is only supported for BINARY_OPERATORs");
+                }
                 return new MIC(name, type, flavor, flags);
             default:
                 throw new BootstrapMethodError("Illegal static bootstrap parameter for flavor: " + flavor);

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

@@ -263,48 +263,54 @@ public final class MethodWriter extends GeneratorAdapter {
     }
 
     /** Writes a dynamic binary instruction: returnType, lhs, and rhs can be different */
-    public void writeDynamicBinaryInstruction(Location location, Type returnType, Type lhs, Type rhs, Operation operation) {
+    public void writeDynamicBinaryInstruction(Location location, Type returnType, Type lhs, Type rhs, Operation operation, boolean compoundAssignment) {
         org.objectweb.asm.Type methodType = org.objectweb.asm.Type.getMethodType(returnType.type, lhs.type, rhs.type);
         String descriptor = methodType.getDescriptor();
         
+        int flags = 0;
+        if (compoundAssignment) {
+            flags |= DefBootstrap.OPERATOR_COMPOUND_ASSIGNMENT;
+        }
         switch (operation) {
             case MUL:
-                invokeDynamic("mul", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0); 
+                invokeDynamic("mul", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags); 
                 break;
             case DIV:
-                invokeDynamic("div", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0); 
+                invokeDynamic("div", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags); 
                 break;
             case REM:
-                invokeDynamic("rem", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0); 
+                invokeDynamic("rem", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags); 
                 break;
             case ADD:
                 // if either side is primitive, then the + operator should always throw NPE on null,
                 // so we don't need a special NPE guard.
                 // otherwise, we need to allow nulls for possible string concatenation.
                 boolean hasPrimitiveArg = lhs.clazz.isPrimitive() || rhs.clazz.isPrimitive();
-                int flags = hasPrimitiveArg ? 0 : DefBootstrap.OPERATOR_ALLOWS_NULL;
+                if (!hasPrimitiveArg) {
+                    flags |= DefBootstrap.OPERATOR_ALLOWS_NULL;
+                }
                 invokeDynamic("add", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags);
                 break;
             case SUB:
-                invokeDynamic("sub", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0); 
+                invokeDynamic("sub", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags); 
                 break;
             case LSH:
-                invokeDynamic("lsh", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.SHIFT_OPERATOR, 0); 
+                invokeDynamic("lsh", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.SHIFT_OPERATOR, flags);
                 break;
             case USH:
-                invokeDynamic("ush", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.SHIFT_OPERATOR, 0); 
+                invokeDynamic("ush", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.SHIFT_OPERATOR, flags); 
                 break;
             case RSH:
-                invokeDynamic("rsh", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.SHIFT_OPERATOR, 0); 
+                invokeDynamic("rsh", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.SHIFT_OPERATOR, flags); 
                 break;
             case BWAND: 
-                invokeDynamic("and", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0);
+                invokeDynamic("and", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags);
                 break;
             case XOR:   
-                invokeDynamic("xor", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0);
+                invokeDynamic("xor", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags);
                 break;
             case BWOR:  
-                invokeDynamic("or", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, 0);
+                invokeDynamic("or", descriptor, DEF_BOOTSTRAP_HANDLE, DefBootstrap.BINARY_OPERATOR, flags);
                 break;
             default:
                 throw location.createError(new IllegalStateException("Illegal tree structure."));

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EBinary.java

@@ -612,7 +612,7 @@ public final class EBinary extends AExpression {
             right.write(writer);
 
             if (promote.sort == Sort.DEF || (shiftDistance != null && shiftDistance.sort == Sort.DEF)) {
-                writer.writeDynamicBinaryInstruction(location, actual, left.actual, right.actual, operation);
+                writer.writeDynamicBinaryInstruction(location, actual, left.actual, right.actual, operation, false);
             } else {
                 writer.writeBinaryInstruction(location, actual, operation);
             }

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EChain.java

@@ -351,7 +351,7 @@ public final class EChain extends AExpression {
                     // write the operation instruction for compound assignment
                     if (promote.sort == Sort.DEF) {
                         writer.writeDynamicBinaryInstruction(location, promote, 
-                            Definition.DEF_TYPE, Definition.DEF_TYPE, operation);
+                            Definition.DEF_TYPE, Definition.DEF_TYPE, operation, true);
                     } else {
                         writer.writeBinaryInstruction(location, promote, operation);
                     }

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

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