Reduce memory taken during compilation

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

@@ -59,7 +59,7 @@ class DocumentContext {
     private final DocumentId documentId;
     private final String name;
     private String content;
-    private boolean disableSyntaxTree = false;
+    private final boolean disableSyntaxTree;
 
     private DocumentContext(DocumentId documentId, String name, String content, boolean disableSyntaxTree) {
         this.documentId = documentId;
@@ -82,40 +82,44 @@ String name() {
     }
 
     SyntaxTree parse() {
-        if (syntaxTree != null) {
-            return syntaxTree;
+        if (this.syntaxTree != null) {
+            return this.syntaxTree;
         }
-        if (!disableSyntaxTree) {
-            syntaxTree = SyntaxTree.from(this.textDocument(), name);
-            return syntaxTree;
+        if (!this.disableSyntaxTree) {
+            this.syntaxTree = SyntaxTree.from(this.textDocument(), this.name);
+            return this.syntaxTree;
         }
-        return SyntaxTree.from(this.textDocument(), name);
+        return SyntaxTree.from(this.textDocument(), this.name);
     }
 
     SyntaxTree syntaxTree() {
         return parse();
     }
 
     TextDocument textDocument() {
-        if (this.textDocument == null) {
+        if (this.textDocument != null) {
+            return this.textDocument;
+        }
+        if (!this.disableSyntaxTree) {
             this.textDocument = TextDocuments.from(this.content);
+            return this.textDocument;
         }
-        return this.textDocument;
+        return TextDocuments.from(this.content);
     }
 
     BLangCompilationUnit compilationUnit(CompilerContext compilerContext, PackageID pkgID, SourceKind sourceKind) {
         BLangDiagnosticLog dlog = BLangDiagnosticLog.getInstance(compilerContext);
-        SyntaxTree syntaxTree = syntaxTree();
-        reportSyntaxDiagnostics(pkgID, syntaxTree, dlog);
+        SyntaxTree synTree = syntaxTree();
+        reportSyntaxDiagnostics(pkgID, synTree, dlog);
 
-        nodeCloner = NodeCloner.getInstance(compilerContext);
-        if (compilationUnit != null) {
-            return nodeCloner.cloneCUnit(compilationUnit);
+        this.nodeCloner = NodeCloner.getInstance(compilerContext);
+        if (this.compilationUnit != null) {
+            return this.nodeCloner.cloneCUnit(this.compilationUnit);
         }
         BLangNodeBuilder bLangNodeBuilder = new BLangNodeBuilder(compilerContext, pkgID, this.name);
-        compilationUnit = (BLangCompilationUnit) bLangNodeBuilder.accept(syntaxTree.rootNode()).get(0);
-        compilationUnit.setSourceKind(sourceKind);
-        return nodeCloner.cloneCUnit(compilationUnit);
+        this.compilationUnit = (BLangCompilationUnit) bLangNodeBuilder.accept(synTree.rootNode()).get(0);
+        this.compilationUnit.setSourceKind(sourceKind);
+        return this.nodeCloner.cloneCUnit(this.compilationUnit);
     }
 
     Set<ModuleLoadRequest> moduleLoadRequests(ModuleDescriptor currentModuleDesc, PackageDependencyScope scope) {
@@ -129,10 +133,10 @@ Set<ModuleLoadRequest> moduleLoadRequests(ModuleDescriptor currentModuleDesc, Pa
 
     private Set<ModuleLoadRequest> getModuleLoadRequests(ModuleDescriptor currentModuleDesc,
                                                          PackageDependencyScope scope) {
-        Set<ModuleLoadRequest> moduleLoadRequests = new LinkedHashSet<>();
+        Set<ModuleLoadRequest> moduleLoadRequestSet = new LinkedHashSet<>();
         ModulePartNode modulePartNode = syntaxTree().rootNode();
         for (ImportDeclarationNode importDcl : modulePartNode.imports()) {
-            moduleLoadRequests.add(getModuleLoadRequest(importDcl, scope));
+            moduleLoadRequestSet.add(getModuleLoadRequest(importDcl, scope));
         }
 
         // TODO This is a temporary solution for SLP6 release
@@ -145,9 +149,9 @@ private Set<ModuleLoadRequest> getModuleLoadRequests(ModuleDescriptor currentMod
             ModuleLoadRequest ballerinaiLoadReq = new ModuleLoadRequest(
                     PackageOrg.from(Names.BALLERINA_INTERNAL_ORG.value),
                     moduleName, scope, DependencyResolutionType.PLATFORM_PROVIDED);
-            moduleLoadRequests.add(ballerinaiLoadReq);
+            moduleLoadRequestSet.add(ballerinaiLoadReq);
         }
-        return moduleLoadRequests;
+        return moduleLoadRequestSet;
     }
 
     private ModuleLoadRequest getModuleLoadRequest(ImportDeclarationNode importDcl, PackageDependencyScope scope) {

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

@@ -177,8 +177,10 @@ private void performCodeGen() {
                         new PackageDiagnostic(diagnostic, moduleContext.descriptor(), moduleContext.project()));
             }
 
-            //TODO: remove this once ballerina-lang#41407 is fixed
             ModuleContext.shrinkDocuments(moduleContext);
+            if (moduleContext.project().kind() == ProjectKind.BALA_PROJECT) {
+                moduleContext.cleanBLangPackage();
+            }
         }
         // add compilation diagnostics
         diagnostics.addAll(moduleDiagnostics);

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

@@ -245,6 +245,10 @@ BLangPackage bLangPackage() {
         return getBLangPackageOrThrow();
     }
 
+    protected void cleanBLangPackage() {
+        this.bLangPackage = null;
+    }
+
     ModuleCompilationState compilationState() {
         return moduleCompState;
     }
@@ -308,10 +312,6 @@ void setCompilationState(ModuleCompilationState moduleCompState) {
         this.moduleCompState = moduleCompState;
     }
 
-    void parse() {
-        currentCompilationState().parse(this);
-    }
-
     void resolveDependencies(DependencyResolution dependencyResolution) {
         Set<ModuleDependency> moduleDependencies = new HashSet<>();
         if (this.project.kind() == ProjectKind.BALA_PROJECT) {
@@ -506,12 +506,8 @@ private static boolean shouldGenerateBir(ModuleContext moduleContext, CompilerCo
         if (Boolean.parseBoolean(compilerOptions.get(CompilerOptionName.DUMP_BIR_FILE))) {
             return true;
         }
-        if (moduleContext.project.kind().equals(ProjectKind.BUILD_PROJECT)
-                && moduleContext.project().buildOptions().enableCache()) {
-            return true;
-        }
-
-        return false;
+        return moduleContext.project.kind().equals(ProjectKind.BUILD_PROJECT)
+                && moduleContext.project().buildOptions().enableCache();
     }
 
     private static ByteArrayOutputStream generateBIR(ModuleContext moduleContext, CompilerContext compilerContext) {
@@ -560,7 +556,6 @@ static void loadPlatformSpecificCodeInternal(ModuleContext moduleContext, Compil
         // TODO implement
     }
 
-    //TODO: should be removed once we properly fix ballerina-lang#41407
     static void shrinkDocuments(ModuleContext moduleContext) {
         moduleContext.srcDocContextMap.values().forEach(DocumentContext::shrink);
     }

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/CodeGenerator.java

@@ -21,6 +21,7 @@
 import org.wso2.ballerinalang.compiler.PackageCache;
 import org.wso2.ballerinalang.compiler.bir.BIRGenUtils;
 import org.wso2.ballerinalang.compiler.bir.codegen.optimizer.LargeMethodOptimizer;
+import org.wso2.ballerinalang.compiler.bir.model.BIRNode;
 import org.wso2.ballerinalang.compiler.diagnostic.BLangDiagnosticLog;
 import org.wso2.ballerinalang.compiler.semantics.analyzer.Types;
 import org.wso2.ballerinalang.compiler.semantics.model.SymbolTable;
@@ -38,14 +39,12 @@
 public class CodeGenerator {
 
     private static final CompilerContext.Key<CodeGenerator> CODE_GEN = new CompilerContext.Key<>();
-    private SymbolTable symbolTable;
-    private PackageCache packageCache;
-    private BLangDiagnosticLog dlog;
-    private Types types;
-    private LargeMethodOptimizer largeMethodOptimizer;
+    private final SymbolTable symbolTable;
+    private final PackageCache packageCache;
+    private final BLangDiagnosticLog dlog;
+    private final Types types;
 
     private CodeGenerator(CompilerContext compilerContext) {
-
         compilerContext.put(CODE_GEN, this);
         this.symbolTable = SymbolTable.getInstance(compilerContext);
         this.packageCache = PackageCache.getInstance(compilerContext);
@@ -54,12 +53,10 @@ private CodeGenerator(CompilerContext compilerContext) {
     }
 
     public static CodeGenerator getInstance(CompilerContext context) {
-
         CodeGenerator codeGenerator = context.get(CODE_GEN);
         if (codeGenerator == null) {
             codeGenerator = new CodeGenerator(context);
         }
-
         return codeGenerator;
     }
 
@@ -75,7 +72,7 @@ public CompiledJarFile generateTestModule(BLangPackage bLangTestablePackage) {
     private CompiledJarFile generate(BPackageSymbol packageSymbol) {
 
         // Split large BIR functions into smaller methods
-        largeMethodOptimizer = new LargeMethodOptimizer(symbolTable);
+        LargeMethodOptimizer largeMethodOptimizer = new LargeMethodOptimizer(symbolTable);
         largeMethodOptimizer.splitLargeBIRFunctions(packageSymbol.bir);
 
         // Desugar BIR to include the observations
@@ -93,9 +90,42 @@ private CompiledJarFile generate(BPackageSymbol packageSymbol) {
 
         // TODO Get-rid of the following assignment
         CompiledJarFile compiledJarFile = jvmPackageGen.generate(packageSymbol.bir, true);
-
+        cleanUpBirPackage(packageSymbol);
         //Revert encoding identifier names
         JvmDesugarPhase.replaceEncodedModuleIdentifiers(packageSymbol.bir, originalIdentifierMap);
         return compiledJarFile;
     }
+
+    private static void cleanUpBirPackage(BPackageSymbol packageSymbol) {
+        packageSymbol.birPackageFile = null;
+        BIRNode.BIRPackage bir = packageSymbol.bir;
+        for (BIRNode.BIRTypeDefinition typeDef : bir.typeDefs) {
+            for (BIRNode.BIRFunction attachedFunc : typeDef.attachedFuncs) {
+                cleanUpBirFunction(attachedFunc);
+            }
+            typeDef.annotAttachments = null;
+        }
+        bir.importedGlobalVarsDummyVarDcls.clear();
+        for (BIRNode.BIRFunction function : bir.functions) {
+            cleanUpBirFunction(function);
+        }
+        bir.annotations.clear();
+        bir.constants.clear();
+        bir.serviceDecls.clear();
+    }
+
+    private static void cleanUpBirFunction(BIRNode.BIRFunction function) {
+        function.receiver = null;
+        function.localVars = null;
+        function.returnVariable = null;
+        function.parameters = null;
+        function.basicBlocks = null;
+        function.errorTable = null;
+        function.workerChannels = null;
+        function.annotAttachments = null;
+        function.returnTypeAnnots = null;
+        function.dependentGlobalVars = null;
+        function.pathParams = null;
+        function.restPathParam = null;
+    }
 }