Remapper variables during bytecode method inlining by ASM

Question

I am doing an online bytecode method inlining optimization using ASM. My changes are based on the example 3.2.6 Inline Method (http://asm.ow2.org/current/asm-transformations.pdf). The test example (inline callee's calculate(int,int) at Caller::test) is:

public class Caller {
    final Callee _callee;

    public Caller(Callee callee){
        _callee = callee;
    }
    public static void main(String[] args) {
        new Caller(new Callee("xu", "shijie")).test(5, 100);
    }

    public void test(int a, int b){
        int t = a;
        int p = b;
        int r = t+p-_callee.calculate(a, b);
        int m = t-p;
        System.out.println(t);
    }
}
public class Callee {

    final String _a;
    final String _b;
    public Callee(String a, String b){
        _a = a;
        _b = b;
    }

    public int calculate(int t, int p){
        int tmp = _a.length()+_b.length();
        tmp+=t+p;
        return tmp;
    }
}

Based on ASM 5.0 version, my code is:

//MainInliner.java
public class MainInliner extends ClassLoader{

    public byte[] generator(String caller, String callee) throws ClassNotFoundException{
        String resource = callee.replace('.', '/') + ".class";
        InputStream is =  getResourceAsStream(resource);
        byte[] buffer;
        // adapts the class on the fly
        try {
            resource = caller.replace('.', '/')+".class";
            is = getResourceAsStream(resource);
            ClassReader cr = new ClassReader(is);
            ClassWriter cw = new ClassWriter(0);

            ClassVisitor visitor = new BCMerge(Opcodes.ASM5, cw, callee);
            cr.accept(visitor, 0);

            buffer= cw.toByteArray();

        } catch (Exception e) {
            throw new ClassNotFoundException(caller, e);
        }

        // optional: stores the adapted class on disk
        try {
            FileOutputStream fos = new FileOutputStream("/tmp/data.class");
            fos.write(buffer);
            fos.close();
        } catch (IOException e) {}
        return buffer;
     }


      @Override
      protected synchronized Class<?> loadClass(final String name,
                final boolean resolve) throws ClassNotFoundException {
            if (name.startsWith("java.")) {
                System.err.println("Adapt: loading class '" + name
                        + "' without on the fly adaptation");
                return super.loadClass(name, resolve);
            } else {
                System.err.println("Adapt: loading class '" + name
                        + "' with on the fly adaptation");
            }
            String caller = "code.sxu.asm.example.Caller";
            String callee = "code.sxu.asm.example.Callee";
            byte[] b = generator(caller, callee);
            // returns the adapted class
            return defineClass(caller, b, 0, b.length);
        }

        public static void main(final String args[]) throws Exception {
            // loads the application class (in args[0]) with an Adapt class loader
            ClassLoader loader = new MainInliner();
            Class<?> c = loader.loadClass(args[0]);
            Method m = c.getMethod("main", new Class<?>[] { String[].class });

        }
}


class BCMerge extends ClassVisitor{

    String _callee;
    String _caller;

    public BCMerge(int api, ClassVisitor cv, String callee) {
        super(api, cv);
        // TODO Auto-generated constructor stub
        _callee = callee.replace('.', '/');
    }

    public void visit(int version, int access, String name, String signature,
            String superName, String[] interfaces) {
        super.visit(version, access, name, signature, superName, interfaces);
        this._caller = name;
    } 

    public MethodVisitor visitMethod(int access, String name, String desc,
            String signature, String[] exceptions) {

        if(!name.equals("test")){
            return super.visitMethod(access, name, desc, signature, exceptions);
        }
        MethodVisitor mv = cv.visitMethod(access, name, desc, signature, exceptions);

        ClassReader cr = null;
        try {
            cr = new ClassReader(this.getClass().getClassLoader().getResourceAsStream(_callee.replace('.', '/')+".class"));
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }

         ClassNode classNode = new ClassNode();
         cr.accept(classNode, 0);

         MethodNode inlinedMethod = null;
         for(MethodNode node: classNode.methods){
            if(node.name.equals("calculate")){
                inlinedMethod = node;
                break;
            }
        }

        return new MethodCallInliner(access, desc, mv, inlinedMethod, _callee, _caller  );
    }
}

//MethodCallInliner.java
public class MethodCallInliner extends LocalVariablesSorter {

    private final String oldClass;
    private final String newClass;
    private final MethodNode mn;  //Method Visitor wrappers the mv.
    private List blocks = new ArrayList();
    private boolean inlining;

    public MethodCallInliner(int access, String desc, MethodVisitor mv, MethodNode mn, 
            String oldClass, String newClass){
        super(Opcodes.ASM5, access, desc, mv);
        this.oldClass = oldClass;
        this.newClass =  newClass;
        this.mn = mn;
        inlining = false;
    }

    public void visitMethodInsn(int opcode, String owner, String name,
            String desc, boolean itf) {

        System.out.println("opcode:" + opcode + " owner:" + owner + " name:"
                + name + " desc:" + desc);
        if (!canBeInlined(owner, name, desc)) {
            mv.visitMethodInsn(opcode, owner, name, desc, itf);
            return;
        }
        //if it is INVOKEVIRTUAL  ../Callee::calculate(II), then..
        Remapper remapper = new SimpleRemapper(oldClass, newClass);
        Label end = new Label();
        inlining = true;
        mn.instructions.resetLabels();
        mn.accept(new InliningAdapter(this,opcode == Opcodes.INVOKESTATIC ? Opcodes.ACC_STATIC : 0, desc,remapper, end));
        inlining = false;
        super.visitLabel(end);
    }

    private boolean canBeInlined(String owner, String name, String decs){
            if(name.equals("calculate") && owner.equals("code/sxu/asm/example/Callee")){
            return true;
        }
        return false;
    }
}

//InliningAdapter.java
public class InliningAdapter extends RemappingMethodAdapter {

    private final LocalVariablesSorter lvs;
    private final Label end;

    public InliningAdapter(LocalVariablesSorter mv,
            int acc, String desc,Remapper remapper, Label end) {
        super(acc, desc, mv, remapper);
        this.lvs = mv;
        this.end = end;

//      int offset = (acc & Opcodes.ACC_STATIC)!=0 ?0 : 1;
//      Type[] args = Type.getArgumentTypes(desc);
//      for (int i = args.length-1; i >= 0; i--) {
//          super.visitVarInsn(args[i].getOpcode(
//                  Opcodes.ISTORE), i + offset);
//      }
//      if(offset>0) {
//          super.visitVarInsn(Opcodes.ASTORE, 0);
//      }
    }

    public void visitInsn(int opcode) {
        if(opcode==Opcodes.RETURN || opcode == Opcodes.IRETURN) {
            super.visitJumpInsn(Opcodes.GOTO, end);
        } else {
            super.visitInsn(opcode);
        }
    }

    public void visitMaxs(int stack, int locals) {
        System.out.println("visit maxs: "+stack+"  "+locals);
    }

    protected int newLocalMapping(Type type) {
        return lvs.newLocal(type);
    }
}

In the code, both InliningAdapter and MethodCallInliner extends LocalVariablesSorter, which renumbers local variables. And the inline references coping body of Callee::calculate() at the Caller::test::invokevirtual(Callee::calculate) call site.

The bytecodes for Caller::test(), Callee::calculate, and generated::test are:

 //Caller::test()
 public void test(int, int);
    flags: ACC_PUBLIC
    Code:
      stack=4, locals=7, args_size=3
         0: iload_1       
         1: istore_3      
         2: iload_2       
         3: istore        4
         5: iload_3       
         6: iload         4
         8: iadd          
         9: aload_0       
        10: getfield      #13                 // Field _callee:Lcode/sxu/asm/example/Callee;
        13: iload_1       
        14: iload_2       
        15: invokevirtual #39                 // Method code/sxu/asm/example/Callee.calculate:(II)I   //Copy calculate's body here
        18: isub          
        19: istore        5
        21: iload_3       
        22: iload         4
        24: isub          
        25: istore        6
        27: getstatic     #43                 // Field java/lang/System.out:Ljava/io/PrintStream;
        30: iload_3       
        31: invokevirtual #49                 // Method java/io/PrintStream.println:(I)V
        34: getstatic     #43                 // Field java/lang/System.out:Ljava/io/PrintStream;
        37: ldc           #55                 // String 1..........
        39: invokevirtual #57                 // Method java/io/PrintStream.println:(Ljava/lang/String;)V
        42: return        

//Callee::calculate()

 public int calculate(int, int);
    flags: ACC_PUBLIC
    Code:
      stack=3, locals=4, args_size=3
         0: aload_0       
         1: getfield      #14                 // Field _a:Ljava/lang/String;
         4: invokevirtual #26                 // Method java/lang/String.length:()I
         7: aload_0       
         8: getfield      #16                 // Field _b:Ljava/lang/String;
        11: invokevirtual #26                 // Method java/lang/String.length:()I
        14: iadd          
        15: istore_3      
        16: iload_3       
        17: iload_1       
        18: iload_2       
        19: iadd          
        20: iadd          
        21: istore_3      
        22: iload_3       
        23: ireturn       

 //data.class
  public void test(int, int);
    flags: ACC_PUBLIC
    Code:
      stack=4, locals=8, args_size=3
         0: iload_1       
         1: istore_3      
         2: iload_2       
         3: istore        4
         5: iload_3       
         6: iload         4
         8: iadd          
         9: aload_0       
        10: getfield      #14                 // Field _callee:Lcode/sxu/asm/example/Callee;
        13: iload_1       
        14: iload_2       
        15: aload_0       
        16: getfield      #40                 // Field _a:Ljava/lang/String;
        19: invokevirtual #46                 // Method java/lang/String.length:()I
        22: aload_0       
        23: getfield      #49                 // Field _b:Ljava/lang/String;
        26: invokevirtual #46                 // Method java/lang/String.length:()I
        29: iadd          
        30: istore        6
        32: iload         6
        34: iload_1       
        35: iload_2       
        36: iadd          
        37: iadd          
        38: istore        6
        40: iload         6
        42: goto          45
        45: isub          
        46: istore        6
        48: iload_3       
        49: iload         4
        51: isub          
        52: istore        7
        54: getstatic     #59                 // Field java/lang/System.out:Ljava/io/PrintStream;
        57: iload_3       
        58: invokevirtual #65                 // Method java/io/PrintStream.println:(I)V
        61: getstatic     #59                 // Field java/lang/System.out:Ljava/io/PrintStream;
        64: ldc           #67                 // String 1..........
        66: invokevirtual #70                 // Method java/io/PrintStream.println:(Ljava/lang/String;)V
        69: return        

The javap result on the data.class shows that the body of Callee::calculate has been inserted to the right place (Caller::test line::15). However, there are two main problems:

• The top three stack objects before invokevirtual

  Callee::calculate(line 15) 9: aload_0
  10: getfield #14 // Field _callee:Lcode/sxu/asm/example/Callee; 13: iload_1
  14: iload_2

  should not be on the stack after inline.

• The variable number 0 in the copied body (Callee::calculate()) should mapped to the right number

• The variable number are not correct. First, the variable numbers of the copied body of Callee::calculate in the data.class (from line 15 to line 42) shoud begin 5 (instead of 0). Second, the variable numbers after Callee::calculate() should be renumbered by the rule: a) not change if it is between (0,4]; b) Renumber if it is conflict with the number in the copied body of Callee::calculate()

I went to check base class LocalVariablesSorter's implementation. The problem seems to be at its construction:

protected LocalVariablesSorter(final int api, final int access,
            final String desc, final MethodVisitor mv) {
        super(api, mv);
        Type[] args = Type.getArgumentTypes(desc);
        nextLocal = (Opcodes.ACC_STATIC & access) == 0 ? 1 : 0;
        for (int i = 0; i < args.length; i++) {
            nextLocal += args[i].getSize();
        }
        firstLocal = nextLocal; 
    }
    private int[] mapping = new int[40];

To me, it seems that the firstLocal are always starts at 1+ args.length() (In this case it is 3). This class also provide private int remap(final int var, final Type type), which creates new local variables and keeps the mapping (from existing variable number to the new index) in the mapping array.

---

My problem is that how to use LocalVariablesSorter and inline the bytecode method (Callee::calculate) correctly. Any suggestion for efficient inline is welcome.

For parameters on the stack before inline (before line 15). My idea is to store them as new created local variables which would be referred by the copied body of Callee::calculate. For example, add: astore 5 after 10: getfield #14 // Field _callee:Lcode/sxu/asm/example/Callee;

and add mapping[0]=5+1 in the LocalVariablesSorter

But the main problem is that users are not allowed to update LocalVariablesSorter::mapping (from old variable number to the new variable in the mapping array) because mapping array is private and the only place for its update is in the method:

   private int remap(final int var, final Type type) {
        if (var + type.getSize() <= firstLocal) { 
        //Variable index will never be modified if it is less than firstLocal. 0 < 3. Nothing i can do for ALOAD 0. 
            return var;
        }
        int key = 2 * var + type.getSize() - 1;
        int size = mapping.length;
        if (key >= size) {
                 .....
        }
        int value = mapping[key];
        if (value == 0) {
            value = newLocalMapping(type);
            setLocalType(value, type);
            mapping[key] = value + 1;
        } else {
            value--;
        }
        if (value != var) {
            changed = true;
        }
        return value;
    }

---

Update1: The data.class after uncomment constructor of the InliningAdapter:

     0: iload_1       
     1: istore_3      
     2: iload_2       
     3: istore        4
     5: iload_3       
     6: iload         4
     8: iadd          
     9: aload_0       
    10: getfield      #14                 // Field _callee:Lcode/sxu/asm/example/Callee;
    13: iload_1       
    14: iload_2       
    15: istore_2        //should be istore 5
    16: istore_1        //should be istore 6
    17: astore_0        //should be astore 7
    18: aload_0        
    19: getfield      #40                 // Field _a:Ljava/lang/String;
    22: invokevirtual #46                 // Method java/lang/String.length:()I
    25: aload_0       
    26: getfield      #49                 // Field _b:Ljava/lang/String;
    29: invokevirtual #46                 // Method java/lang/String.length:()I
    32: iadd          
    33: istore        6
    35: iload         6
    37: iload_1       
    38: iload_2       
    39: iadd          
    40: iadd          
    41: istore        6
    43: iload         6
    45: goto          48
    48: isub          
    49: istore        6
    51: iload_3       
    52: iload         4
    54: isub          
    55: istore        7
    57: getstatic     #59 

The new stored three variables (15,16,17) should be numbered as 5,6,7, rather than 2,1,0, and the mapping in *store/*load in the inlined code should be like

       0 => 7
       1 => 6 
       2 => 5
       3 => 8
       4 => 9  ...

These mapping should be in the array:LocalVariablesSorter::mapping which is updated by the method LocalVariablesSorter::remap(). However, it seems not possible for i to insert them into mapping array.

There are two kinds of remappering should be done:

• Remapping inside of inlined code (From line 18 t0 45) and variable index starts at 5. The max index is k
• Remapping after inlined code (From line 46 to the end), and any variable index should be remapped (new index starts at k+1) if it the original index is greater than 5

Answer 1

As @Holger already suggested, start by uncommenting the lines in InliningAdapter.

1. To tackle the main problems you listed: The LocalVariablesSorter (extended by InliningAdapter) thinks that the arguments are already stored in the local variables at fixed locations - this is indeed the normal situation when entering a method. So it does not map those at all (see first line in LocalVariablesSorter.remap() - firstLocal is calculated in constructor). However in this case we get the arguments on the stack instead and need to allocate the local variables manually. The solution is to tell LocalVariablesSorter that there are no parameters already stored in local variables (make firstLocal = 0). Then it will treat any reference to them as new variables and allocate new local variables for them. This we can achieve by fooling LocalVariablesSorter to think that there are no arguments and that the method is static (even if it really isn't). So we change the first line in InliningAdapter from

       super(acc, desc, mv, remapper);
   

   to

       super(acc | Opcodes.ACC_STATIC, "()V", mv, remapper);
   

   Now the variables 0,1,2,... get remapped to 5,6,7,... or similar (doesn't really matter what they are, the LocalVariablesSorter of Caller (i.e. the MethodCallInliner instance) takes care of allocating them).

2. There is another problem also that you map the Callee class to Caller by having InliningAdapter extend RemappingMethodAdaptor - however I guess you want to keep the _a and _b variables stored in the Callee instance.

   • If my guess is correct then you should actually not remap the references from Callee to Caller. You can just make InliningAdapter extend LocalVariableSorter instead and get rid of the remapper.
   • If my guess is incorrect then I guess you will probably need to embed the variables of Callee into Caller as well, in which case you should keep the RemappingMethodAdaptor you have.

3. When debugging the inlined code the linenumbers from the Callee will not make sense since the code was inlined into the Caller class. So all linenumbers from Callee should probably be replaced with the line number of the line in Caller where the inlined call occured. Unfortunately in Java you cannot specify different source code files on a line-by-line basis (like you can in C for example). So you would override visitLineNumber() in InliningAdapter using something like this (inlinedLine would be passed to constructor of InliningAdapter):

   @Override
   public void visitLineNumber(int line, Label start) {
       super.visitLineNumber(inlinedLine, start);
   }
   

   .. or perhaps skip the super call altogether, I'm not 100% sure about this.