Optimize CodeGen for executable jar

compiler/ballerina-lang/src/main/java/io/ballerina/projects/DeadFunctionVisitor.java

@@ -0,0 +1,50 @@
+/*
+ *  Copyright (c) 2023, WSO2 LLC. (http://www.wso2.org) All Rights Reserved.
+ *
+ *  WSO2 LLC. 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 io.ballerina.projects;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+
+import java.util.HashSet;
+
+/**
+ * This visitor replaces unused functions with null
+ */
+public class DeadFunctionVisitor extends ClassVisitor {
+
+    private final HashSet<String> deadFunctionNames;
+
+    public DeadFunctionVisitor(int api, ClassWriter cw, HashSet<String> deadFunctionNames) {
+
+        super(api, cw);
+        this.deadFunctionNames = deadFunctionNames;
+    }
+
+    @Override
+    public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
+
+        MethodVisitor methodVisitor;
+        if (!name.contains("init") && deadFunctionNames.contains(name)) {
+            methodVisitor = null;
+        } else {
+            methodVisitor = super.visitMethod(access, name, desc, signature, exceptions);
+        }
+        return methodVisitor;
+    }
+}

compiler/ballerina-lang/src/main/java/io/ballerina/projects/JBallerinaBackend.java

@@ -30,7 +30,6 @@
 import org.apache.commons.compress.archivers.jar.JarArchiveEntry;
 import org.apache.commons.compress.archivers.zip.ZipArchiveEntryPredicate;
 import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream;
-import org.apache.commons.compress.archivers.zip.ZipFile;
 import org.apache.commons.io.FilenameUtils;
 import org.ballerinalang.maven.Dependency;
 import org.ballerinalang.maven.MavenResolver;
@@ -40,6 +39,8 @@
 import org.wso2.ballerinalang.compiler.bir.codegen.CodeGenerator;
 import org.wso2.ballerinalang.compiler.bir.codegen.interop.InteropValidator;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.ObservabilitySymbolCollectorRunner;
+import org.wso2.ballerinalang.compiler.semantics.model.symbols.BInvokableSymbol;
+import org.wso2.ballerinalang.compiler.semantics.model.symbols.BPackageSymbol;
 import org.wso2.ballerinalang.compiler.spi.ObservabilitySymbolCollector;
 import org.wso2.ballerinalang.compiler.tree.BLangPackage;
 import org.wso2.ballerinalang.compiler.util.CompilerContext;
@@ -453,7 +454,10 @@ private Manifest createManifest() {
     private void copyJar(ZipArchiveOutputStream outStream, JarLibrary jarLibrary,
             HashMap<String, JarLibrary> copiedEntries, HashMap<String, StringBuilder> services) throws IOException {
 
-        ZipFile zipFile = new ZipFile(jarLibrary.path().toFile());
+        String pkgName = jarLibrary.packageName().get();
+        HashMap<String, HashSet<String>> classWiseDeadFunctionMap = getDeadFunctionList(pkgName, this.compilerContext);
+
+        OptimisedZipFile zipFile = new OptimisedZipFile(jarLibrary.path().toFile(),classWiseDeadFunctionMap);
         ZipArchiveEntryPredicate predicate = entry -> {
             String entryName = entry.getName();
             if (entryName.equals("META-INF/MANIFEST.MF")) {
@@ -501,10 +505,31 @@ private void copyJar(ZipArchiveOutputStream outStream, JarLibrary jarLibrary,
 
         // Transfers selected entries from this zip file to the output stream, while preserving its compression and
         // all the other original attributes.
-        zipFile.copyRawEntries(outStream, predicate);
+        zipFile.copyOptimisedRawEntries(outStream, predicate);
+//        zipFile.copyRawEntries(outStream, predicate);
         zipFile.close();
     }
 
+    private static HashMap<String, HashSet<String>> getDeadFunctionList(String pkgName, CompilerContext compilerContext) {
+        org.wso2.ballerinalang.compiler.PackageCache packageCache = org.wso2.ballerinalang.compiler.PackageCache.getInstance(compilerContext);
+
+        BPackageSymbol pkgSymbol = packageCache.getSymbol(pkgName);
+        int majorVersion = packageCache.getSymbol(pkgName).descriptor.version().value().major();
+
+        HashMap<String, HashSet<String>> classWiseDeadFunctionMap = new HashMap<>();
+
+        for (BInvokableSymbol deadFunction:pkgSymbol.deadFunctions) {
+            if (deadFunction.source != null) {
+                String className = deadFunction.source.replace(".bal", ".class");
+                String classLocationName = pkgName + "/" + majorVersion + "/" + className;
+                classWiseDeadFunctionMap.putIfAbsent(classLocationName, new HashSet<>());
+                classWiseDeadFunctionMap.get(classLocationName).add(deadFunction.originalName.value);
+            }
+        }
+
+        return classWiseDeadFunctionMap;
+    }
+
     private static boolean isCopiedEntry(String entryName, HashMap<String, JarLibrary> copiedEntries) {
         return copiedEntries.keySet().contains(entryName);
     }

compiler/ballerina-lang/src/main/java/io/ballerina/projects/ModuleContext.java

@@ -31,6 +31,7 @@
 import org.ballerinalang.model.elements.Flag;
 import org.ballerinalang.model.elements.PackageID;
 import org.wso2.ballerinalang.compiler.BIRPackageSymbolEnter;
+import org.wso2.ballerinalang.compiler.bir.model.BIRNode;
 import org.wso2.ballerinalang.compiler.bir.writer.BIRBinaryWriter;
 import org.wso2.ballerinalang.compiler.diagnostic.BLangDiagnosticLocation;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.SymbolEnter;
@@ -44,6 +45,7 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -54,6 +56,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Predicate;
 
 import static org.ballerinalang.model.tree.SourceKind.REGULAR_SOURCE;
 import static org.ballerinalang.model.tree.SourceKind.TEST_SOURCE;
@@ -298,7 +301,11 @@ ModuleCompilationState currentCompilationState() {
         } else if (this.project().kind() == ProjectKind.BUILD_PROJECT
                 && !this.project.buildOptions().enableCache()) {
             moduleCompState = ModuleCompilationState.LOADED_FROM_SOURCES;
-        } else {
+        }
+//        else if (!this.moduleDescriptor.packageName().toString().contains("lang") && !this.moduleDescriptor.packageName().toString().contains("jballerina")) {
+//            moduleCompState = ModuleCompilationState.LOADED_FROM_SOURCES;
+//        }
+        else {
             moduleCompState = ModuleCompilationState.LOADED_FROM_CACHE;
         }
         return moduleCompState;
@@ -476,6 +483,11 @@ static void generateCodeInternal(ModuleContext moduleContext,
             return;
         }
 
+        // Eliminates dead BIR function nodes from BUILD_PROJECT
+        if (moduleContext.project.kind().equals(ProjectKind.BUILD_PROJECT)) {
+//            eliminateDeadFunctionsFromBIR(moduleContext);
+        }
+
         // Generate and write the thin JAR to the file system
         compilerBackend.performCodeGen(moduleContext, moduleContext.compilationCache);
 
@@ -540,6 +552,41 @@ static void loadBirBytesInternal(ModuleContext moduleContext) {
         moduleContext.birBytes = moduleContext.compilationCache.getBir(moduleContext.moduleName());
     }
 
+    /**
+     * Removes the dead BIRFunction nodes from BIR of the package Symbol of the given context
+     * @param moduleContext Module context holding the bir
+     */
+    public static void eliminateDeadFunctionsFromBIR(ModuleContext moduleContext) {
+        HashSet<String> deadFunctionNames = new HashSet<>();
+        moduleContext.bLangPackage().symbol.deadFunctions.forEach((deadSymbol) -> {
+            deadFunctionNames.add(deadSymbol.getName().value);
+        });
+
+        // Filtering out the unused topLevel BIRFunctions
+        List<BIRNode.BIRFunction> deadBIRFunctions = new ArrayList<>();
+        Predicate<BIRNode.BIRFunction> isTopLevelDeadFunction = birFunction -> deadFunctionNames.contains(birFunction.originalName.getValue());
+        for (BIRNode.BIRFunction birFunction : moduleContext.bLangPackage().symbol.bir.functions) {
+            if (isTopLevelDeadFunction.test(birFunction)) {
+                deadBIRFunctions.add(birFunction);
+            }
+        }
+        moduleContext.bLangPackage().symbol.bir.functions.removeAll(deadBIRFunctions);
+
+        // Filtering out the unused classLevel BIRFunctions
+        List<BIRNode.BIRTypeDefinition> typeDefs = moduleContext.bLangPackage.symbol.bir.typeDefs;
+        for (BIRNode.BIRTypeDefinition typeDefinition : typeDefs) {
+            Predicate<BIRNode.BIRFunction> isClassLevelDeadFunction = birFunction -> deadFunctionNames.contains(typeDefinition.name.getValue() + "." + birFunction.originalName.getValue());
+            List<BIRNode.BIRFunction> classLevelDeadFunctions = new ArrayList<>();
+            for (BIRNode.BIRFunction birFunction : typeDefinition.attachedFuncs) {
+                if (isClassLevelDeadFunction.test(birFunction)) {
+                    classLevelDeadFunctions.add(birFunction);
+                }
+            }
+            typeDefinition.attachedFuncs.removeAll(classLevelDeadFunctions);
+        }
+    }
+
+
     static void resolveDependenciesFromBALAInternal(ModuleContext moduleContext) {
         // TODO implement
     }

compiler/ballerina-lang/src/main/java/io/ballerina/projects/OptimisedZipFile.java

@@ -0,0 +1,76 @@
+/*
+ *  Copyright (c) 2023, WSO2 LLC. (http://www.wso2.org) All Rights Reserved.
+ *
+ *  WSO2 LLC. 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 io.ballerina.projects;
+
+import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
+import org.apache.commons.compress.archivers.zip.ZipArchiveEntryPredicate;
+import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream;
+import org.apache.commons.compress.archivers.zip.ZipFile;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.Opcodes;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.HashSet;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Eliminates the unused methods from the output stream of the given zip file
+ *
+ */
+public class OptimisedZipFile extends ZipFile {
+    private final HashMap<String, HashSet<String>> deadFunctionMap;
+
+    public OptimisedZipFile(File file, HashMap<String, HashSet<String>> deadFunctionMap) throws IOException {
+        super(file);
+        this.deadFunctionMap = deadFunctionMap;
+    }
+
+    public void copyOptimisedRawEntries(final ZipArchiveOutputStream target, final ZipArchiveEntryPredicate predicate) throws IOException {
+        final Enumeration<ZipArchiveEntry> src = getEntriesInPhysicalOrder();
+        while (src.hasMoreElements()) {
+            final ZipArchiveEntry entry = src.nextElement();
+            if (predicate.test( entry)) {
+                if (!entry.getName().contains("init.class") && deadFunctionMap.containsKey(entry.getName())) {
+                    target.putArchiveEntry(entry);
+                    target.write(getOptimizedEntry(entry));
+                    target.closeArchiveEntry();
+                } else {
+                    target.addRawArchiveEntry(entry, getRawInputStream(entry));
+                }
+            }
+        }
+    }
+
+    public byte[] getOptimizedEntry(ZipArchiveEntry originalEntry) throws IOException {
+        InputStream inputStream = getInputStream(originalEntry);
+
+        ClassReader classReader = new ClassReader(requireNonNull(inputStream));
+        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
+        ClassVisitor classVisitor = new DeadFunctionVisitor(Opcodes.ASM7, classWriter, deadFunctionMap.get(originalEntry.getName()));
+        classReader.accept(classVisitor, 0);
+
+        return classWriter.toByteArray();
+    }
+}

compiler/ballerina-lang/src/main/java/io/ballerina/projects/internal/CompilerPhaseRunner.java

@@ -28,6 +28,7 @@
 import org.wso2.ballerinalang.compiler.semantics.analyzer.CodeAnalyzer;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.CompilerPluginRunner;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.DataflowAnalyzer;
+import org.wso2.ballerinalang.compiler.semantics.analyzer.DeadCodeAnalyzer;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.DocumentationAnalyzer;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.IsolationAnalyzer;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.ObservabilitySymbolCollectorRunner;
@@ -71,6 +72,7 @@ public class CompilerPhaseRunner {
     private final CompilerPhase compilerPhase;
     private final DataflowAnalyzer dataflowAnalyzer;
     private final IsolationAnalyzer isolationAnalyzer;
+    private final DeadCodeAnalyzer deadCodeAnalyzer;
     private boolean isToolingCompilation;
 
 
@@ -102,6 +104,7 @@ private CompilerPhaseRunner(CompilerContext context) {
         this.compilerPhase = this.options.getCompilerPhase();
         this.dataflowAnalyzer = DataflowAnalyzer.getInstance(context);
         this.isolationAnalyzer = IsolationAnalyzer.getInstance(context);
+        this.deadCodeAnalyzer = DeadCodeAnalyzer.getInstance(context);
         this.isToolingCompilation = this.options.isSet(TOOLING_COMPILATION)
                 && Boolean.parseBoolean(this.options.get(TOOLING_COMPILATION));
     }
@@ -122,12 +125,17 @@ public void performTypeCheckPhases(BLangPackage pkgNode) {
         }
 
         dataflowAnalyze(pkgNode);
-        if (this.stopCompilation(pkgNode, CompilerPhase.DOCUMENTATION_ANALYZE)) {
+        if (this.stopCompilation(pkgNode, CompilerPhase.ISOLATION_ANALYZE)) {
             return;
         }
 
         isolationAnalyze(pkgNode);
-        if (this.stopCompilation(pkgNode, CompilerPhase.ISOLATION_ANALYZE)) {
+        if (this.stopCompilation(pkgNode, CompilerPhase.DEAD_CODE_ANALYZE)) {
+            return;
+        }
+
+        deadCodeAnalyze(pkgNode);
+        if (this.stopCompilation(pkgNode, CompilerPhase.DOCUMENTATION_ANALYZE)) {
             return;
         }
 
@@ -223,6 +231,9 @@ public BLangPackage birGen(BLangPackage pkgNode) {
     private BLangPackage birEmit(BLangPackage pkgNode) {
         return this.birEmitter.emit(pkgNode);
     }
+    private BLangPackage deadCodeAnalyze(BLangPackage pkgNode) {
+        return this.deadCodeAnalyzer.analyze(pkgNode);
+    }
 
     private boolean stopCompilation(BLangPackage pkgNode, CompilerPhase nextPhase) {
         if (compilerPhase.compareTo(nextPhase) < 0) {

compiler/ballerina-lang/src/main/java/org/ballerinalang/compiler/CompilerPhase.java

@@ -40,6 +40,7 @@ public enum CompilerPhase {
 
     CONSTANT_PROPAGATION("constantPropagation"),
 
+    DEAD_CODE_ANALYZE("deadCodeAnalyze"),
     COMPILER_PLUGIN("compilerPlugin"),
 
     DESUGAR("desugar"),
@@ -50,6 +51,7 @@ public enum CompilerPhase {
 
     CODE_GEN("codeGen");
 
+
     private String value;
 
     CompilerPhase(String value) {
@@ -68,6 +70,8 @@ public static CompilerPhase fromValue(String value) {
                 return DOCUMENTATION_ANALYZE;
             case "constantPropagation":
                 return CONSTANT_PROPAGATION;
+            case "deadCodeAnalyze":
+                return DEAD_CODE_ANALYZE;
             case "compilerPlugin":
                 return COMPILER_PLUGIN;
             case "desugar":