Lambda inlining optimization

base/build.gradle

@@ -33,6 +33,7 @@ dependencies {
     testImplementation 'io.kotest:kotest-assertions-core-jvm:5.5.4' // for kotest core jvm assertions
     testImplementation 'io.kotest:kotest-property-jvm:5.5.4' // for kotest property test
     testImplementation 'io.mockk:mockk:1.13.2' // for mocking
+    testImplementation 'com.guardsquare:proguard-retrace:7.3.2'
 
     testImplementation(testFixtures("com.guardsquare:proguard-core:9.0.8")) {
         exclude group: 'com.guardsquare', module: 'proguard-core'

base/src/main/java/proguard/Configuration.java

@@ -200,6 +200,11 @@ public class Configuration
      */
     public boolean                      mergeInterfacesAggressively      = false;
 
+    /**
+     * Specifies whether lambdas should be inlined.
+     */
+    public boolean                      lambdaInlining                   = false;
+
     ///////////////////////////////////////////////////////////////////////////
     // Obfuscation options.
     ///////////////////////////////////////////////////////////////////////////

base/src/main/java/proguard/ProGuard.java

@@ -38,6 +38,7 @@
 import proguard.optimize.LineNumberTrimmer;
 import proguard.optimize.Optimizer;
 import proguard.optimize.gson.GsonOptimizer;
+import proguard.optimize.inline.LambdaInliner;
 import proguard.optimize.peephole.LineNumberLinearizer;
 import proguard.pass.PassRunner;
 import proguard.preverify.PreverificationClearer;
@@ -195,6 +196,13 @@ public void execute() throws Exception
                 new ListParser(new NameParser()).parse(configuration.optimizations) :
                 new ConstantMatcher(true);
 
+            // Inline kotlin lambdas passed as arguments to methods as long as they are
+            // created in the arguments or created in a local variable and passed as an argument
+            if (configuration.lambdaInlining)
+            {
+                inlineLambdas();
+            }
+
             if (configuration.optimize &&
                 filter.matches(Optimizer.LIBRARY_GSON))
             {
@@ -462,6 +470,11 @@ private void optimizeGson() throws Exception
         passRunner.run(new GsonOptimizer(configuration), appView);
     }
 
+    private void inlineLambdas() throws Exception
+    {
+        LambdaInliner lambdaInliner = new LambdaInliner();
+        passRunner.run(lambdaInliner, appView);
+    }
 
     /**
      * Performs the optimization step.

base/src/main/java/proguard/optimize/inline/CannotInlineException.java

@@ -0,0 +1,7 @@
+package proguard.optimize.inline;
+
+public class CannotInlineException extends RuntimeException {
+    public CannotInlineException(String reason) {
+        super(reason);
+    }
+}

base/src/main/java/proguard/optimize/inline/CastPatternRemover.java

@@ -0,0 +1,55 @@
+package proguard.optimize.inline;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import proguard.classfile.Clazz;
+import proguard.classfile.Method;
+import proguard.classfile.attribute.CodeAttribute;
+import proguard.classfile.constant.Constant;
+import proguard.classfile.constant.MethodrefConstant;
+import proguard.classfile.constant.visitor.ConstantVisitor;
+import proguard.classfile.editor.CodeAttributeEditor;
+import proguard.classfile.editor.InstructionSequenceBuilder;
+import proguard.classfile.instruction.Instruction;
+import proguard.classfile.instruction.visitor.InstructionVisitor;
+import proguard.classfile.util.InstructionSequenceMatcher;
+import proguard.classfile.visitor.ClassPrinter;
+
+import java.io.PrintWriter;
+import java.io.StringWriter;
+
+public class CastPatternRemover implements InstructionVisitor {
+    private final Logger logger = LogManager.getLogger(this.getClass());
+    private final InstructionSequenceMatcher insSeqMatcher;
+    private final CodeAttributeEditor codeAttributeEditor;
+
+    public CastPatternRemover(CodeAttributeEditor codeAttributeEditor) {
+        InstructionSequenceBuilder ____ = new InstructionSequenceBuilder();
+        Constant[] constants = ____.constants();
+        Instruction[] pattern = ____.invokestatic(InstructionSequenceMatcher.X).areturn().__();
+        this.insSeqMatcher = new InstructionSequenceMatcher(constants, pattern);
+        this.codeAttributeEditor = codeAttributeEditor;
+    }
+
+    @Override
+    public void visitAnyInstruction(Clazz clazz, Method method, CodeAttribute codeAttribute, int offset, Instruction instruction) {
+        instruction.accept(clazz, method, codeAttribute, offset, insSeqMatcher);
+        if (insSeqMatcher.isMatching()) {
+            int constantIndex = insSeqMatcher.matchedConstantIndex(InstructionSequenceMatcher.X);
+            clazz.constantPoolEntryAccept(constantIndex, new ConstantVisitor() {
+                @Override
+                public void visitMethodrefConstant(Clazz clazz, MethodrefConstant methodrefConstant) {
+                    if (methodrefConstant.getName(clazz).equals("valueOf")) {
+                        if (logger.isDebugEnabled()) {
+                            StringWriter stringWriter = new StringWriter();
+                            instruction.accept(clazz, method, codeAttribute, offset, new ClassPrinter(new PrintWriter(stringWriter)));
+                            logger.debug("Removing " + stringWriter);
+                        }
+
+                        codeAttributeEditor.deleteInstruction(insSeqMatcher.matchedInstructionOffset(0));
+                    }
+                }
+            });
+        }
+    }
+}

base/src/main/java/proguard/optimize/inline/CastRemover.java

@@ -0,0 +1,86 @@
+package proguard.optimize.inline;
+
+import proguard.classfile.Clazz;
+import proguard.classfile.Method;
+import proguard.classfile.attribute.CodeAttribute;
+import proguard.classfile.constant.AnyMethodrefConstant;
+import proguard.classfile.constant.NameAndTypeConstant;
+import proguard.classfile.constant.visitor.ConstantVisitor;
+import proguard.classfile.editor.CodeAttributeEditor;
+import proguard.classfile.instruction.ConstantInstruction;
+import proguard.classfile.instruction.Instruction;
+import proguard.classfile.instruction.visitor.InstructionVisitor;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+public class CastRemover implements InstructionVisitor {
+    private final CodeAttributeEditor codeAttributeEditor;
+    private final List<Integer> keepList;
+    private int argIndex;
+
+    // The names of all method taking a boxed type variable and returning the variable with the unboxed type
+    private final Set<String> castingMethodNames = new HashSet<>(Arrays.asList("intValue", "booleanValue", "byteValue", "shortValue", "longValue", "floatValue", "doubleValue", "charValue"));
+
+
+    public CastRemover(CodeAttributeEditor codeAttributeEditor, List<Integer> keepList) {
+        this.codeAttributeEditor = codeAttributeEditor;
+        this.argIndex = 0;
+        this.keepList = keepList;
+    }
+
+    public CastRemover(CodeAttributeEditor codeAttributeEditor) {
+        this(codeAttributeEditor, new ArrayList<>());
+    }
+
+    @Override
+    public void visitConstantInstruction(Clazz clazz, Method method, CodeAttribute codeAttribute, int offset, ConstantInstruction constantInstruction) {
+        /* Casting on a lambda invoke call looks like this :
+         * iload
+         * invokestatic valueOf
+         * invokeinterface invoke   //lambda call
+         * checkast
+         * invokevirtual intValue
+         * istore
+         *
+         * We remove valueOf, checkast and intValue
+         */
+        if (constantInstruction.opcode == Instruction.OP_CHECKCAST) {
+            codeAttributeEditor.deleteInstruction(offset);
+        } else if (constantInstruction.opcode == Instruction.OP_INVOKESTATIC) {
+            if (getInvokedMethodName(clazz, constantInstruction).equals("valueOf")) {
+                // Don't remove valueOf call when the lambda takes an object as argument
+                if (!keepList.contains(argIndex)) {
+                    codeAttributeEditor.deleteInstruction(offset);
+                }
+                argIndex++;
+            }
+        } else if (constantInstruction.opcode == Instruction.OP_INVOKEVIRTUAL) {
+            if (castingMethodNames.contains(getInvokedMethodName(clazz, constantInstruction))) {
+                codeAttributeEditor.deleteInstruction(offset);
+            }
+        }
+    }
+
+    @Override
+    public void visitAnyInstruction(Clazz clazz, Method method, CodeAttribute codeAttribute, int offset, Instruction instruction) {
+    }
+
+    private String getInvokedMethodName(Clazz clazz, ConstantInstruction constantInstruction) {
+        final String[] invokedMethodName = new String[1];
+        clazz.constantPoolEntryAccept(constantInstruction.constantIndex, new ConstantVisitor() {
+            @Override public void visitAnyMethodrefConstant(Clazz clazz, AnyMethodrefConstant anyMethodrefConstant) {
+                clazz.constantPoolEntryAccept(anyMethodrefConstant.u2nameAndTypeIndex, new ConstantVisitor() {
+                    @Override
+                    public void visitNameAndTypeConstant(Clazz clazz, NameAndTypeConstant nameAndTypeConstant) {
+                        invokedMethodName[0] = nameAndTypeConstant.getName(clazz);
+                    }
+                });
+            }
+        });
+        return invokedMethodName[0];
+    }
+}

base/src/main/java/proguard/optimize/inline/DescriptorModifier.java

@@ -0,0 +1,52 @@
+package proguard.optimize.inline;
+
+import proguard.classfile.Clazz;
+import proguard.classfile.ProgramClass;
+import proguard.classfile.ProgramMethod;
+import proguard.classfile.editor.AttributeAdder;
+import proguard.classfile.editor.ClassBuilder;
+import proguard.classfile.editor.ClassEditor;
+import proguard.obfuscate.NameFactory;
+import proguard.obfuscate.SimpleNameFactory;
+import proguard.obfuscate.UniqueMemberNameFactory;
+import proguard.optimize.info.ProgramMemberOptimizationInfoSetter;
+
+import java.util.function.Function;
+
+public class DescriptorModifier {
+    private final ClassBuilder classBuilder;
+    private final NameFactory nameFactory;
+    private final Clazz clazz;
+
+    DescriptorModifier(Clazz clazz) {
+        classBuilder = new ClassBuilder((ProgramClass) clazz);
+        nameFactory = new UniqueMemberNameFactory(new SimpleNameFactory(), clazz);
+        this.clazz = clazz;
+    }
+
+    public ProgramMethod modify(ProgramMethod sourceMethod, Function<String, String> modifier) {
+        return modify(sourceMethod, modifier, false);
+    }
+
+    public ProgramMethod modify(ProgramMethod sourceMethod, Function<String, String> modifier, boolean removeOriginal) {
+        ProgramMethod newMethod = classBuilder.addAndReturnMethod(
+                sourceMethod.u2accessFlags,
+                nameFactory.nextName(),
+                modifier.apply(sourceMethod.getDescriptor(clazz))
+        );
+
+        sourceMethod.attributesAccept((ProgramClass) clazz,
+                new AttributeAdder((ProgramClass) clazz, newMethod, true));
+
+        if (removeOriginal)
+            removeOriginal(sourceMethod);
+
+        newMethod.accept(clazz, new ProgramMemberOptimizationInfoSetter());
+        return newMethod;
+    }
+
+    private void removeOriginal(ProgramMethod method) {
+        ClassEditor classEditor = new ClassEditor((ProgramClass) clazz);
+        classEditor.removeMethod(method);
+    }
+}

base/src/main/java/proguard/optimize/inline/FirstLambdaParameterFinder.java

@@ -0,0 +1,21 @@
+package proguard.optimize.inline;
+
+import proguard.classfile.util.InternalTypeEnumeration;
+
+public class FirstLambdaParameterFinder {
+    public static int findFirstLambdaParameter(String descriptor) {
+        return findFirstLambdaParameter(descriptor, true);
+    }
+
+    public static int findFirstLambdaParameter(String descriptor, boolean isStatic) {
+        InternalTypeEnumeration internalTypeEnumeration = new InternalTypeEnumeration(descriptor);
+        int index = 0;
+        while (internalTypeEnumeration.hasMoreTypes()) {
+            if (internalTypeEnumeration.nextType().startsWith("Lkotlin/jvm/functions/Function")) {
+                break;
+            }
+            index ++;
+        }
+        return index + (isStatic ? 0 : 1);
+    }
+}

base/src/main/java/proguard/optimize/inline/InstructionAtOffset.java

@@ -0,0 +1,23 @@
+package proguard.optimize.inline;
+
+import proguard.classfile.instruction.Instruction;
+
+import java.util.Objects;
+
+public final class InstructionAtOffset {
+    private final Instruction instruction;
+    private final int offset;
+
+    InstructionAtOffset(Instruction instruction, int offset) {
+        this.instruction = instruction;
+        this.offset = offset;
+    }
+
+    public Instruction instruction() {
+        return instruction;
+    }
+
+    public int offset() {
+        return offset;
+    }
+}

base/src/main/java/proguard/optimize/inline/IterativeInstructionVisitor.java

@@ -0,0 +1,39 @@
+package proguard.optimize.inline;
+
+import proguard.classfile.Clazz;
+import proguard.classfile.Method;
+import proguard.classfile.attribute.CodeAttribute;
+import proguard.classfile.instruction.Instruction;
+import proguard.classfile.instruction.InstructionFactory;
+import proguard.classfile.instruction.visitor.InstructionVisitor;
+
+public class IterativeInstructionVisitor {
+    private boolean changed;
+
+    /**
+     * This helper method does something similar to codeAttribute.instructionsAccept but has the ability to stop and
+     * restart visiting  the instructions, this is useful for when you change the code during iteration.
+     */
+    public void instructionsAccept(Clazz clazz, Method method, CodeAttribute codeAttribute, InstructionVisitor instructionVisitor) {
+        int offset = 0;
+        while (offset < codeAttribute.u4codeLength) {
+            Instruction instruction = InstructionFactory.create(codeAttribute.code, offset);
+            int instructionLength = instruction.length(offset);
+            changed = false;
+            instruction.accept(clazz, method, codeAttribute, offset, instructionVisitor);
+            if (changed) {
+                offset = 0;
+                continue;
+            }
+            offset += instructionLength;
+        }
+    }
+
+    public boolean codeHasChanged() {
+        return changed;
+    }
+
+    public void setChanged(boolean changed) {
+        this.changed = changed;
+    }
+}