Merge pull request #41400 from gimantha/2201.8.x-oom-fix

[2201.8.x] Remove unncessary CapturedClosureEnvs in Lambda functions
ballerina-platform · Oct 5, 2023 · 775a621 · 775a621
2 parents 059c473 + ddbe706
commit 775a621
Show file tree

Hide file tree

Showing 9 changed files with 23 additions and 10 deletions.
diff --git a/compiler/ballerina-lang/src/main/java/io/ballerina/projects/DocumentContext.java b/compiler/ballerina-lang/src/main/java/io/ballerina/projects/DocumentContext.java
@@ -58,7 +58,7 @@ class DocumentContext {
     private NodeCloner nodeCloner;
     private final DocumentId documentId;
     private final String name;
-    private final String content;
+    private String content;
     private boolean disableSyntaxTree = false;
 
     private DocumentContext(DocumentId documentId, String name, String content, boolean disableSyntaxTree) {
@@ -187,4 +187,11 @@ private void reportSyntaxDiagnostics(PackageID pkgID, SyntaxTree tree, BLangDiag
     DocumentContext duplicate() {
         return new DocumentContext(this.documentId, this.name, syntaxTree().toSourceCode(), false);
     }
+
+    void shrink() {
+        if (this.compilationUnit != null) {
+            this.compilationUnit.topLevelNodes.clear();
+        }
+        this.content = null;
+    }
 }
diff --git a/compiler/ballerina-lang/src/main/java/io/ballerina/projects/JBallerinaBackend.java b/compiler/ballerina-lang/src/main/java/io/ballerina/projects/JBallerinaBackend.java
@@ -176,6 +176,9 @@ private void performCodeGen() {
                 moduleDiagnostics.add(
                         new PackageDiagnostic(diagnostic, moduleContext.descriptor(), moduleContext.project()));
             }
+
+            //TODO: remove this once ballerina-lang#41407 is fixed
+            ModuleContext.shrinkDocuments(moduleContext);
         }
         // add compilation diagnostics
         diagnostics.addAll(moduleDiagnostics);

diff --git a/compiler/ballerina-lang/src/main/java/io/ballerina/projects/ModuleContext.java b/compiler/ballerina-lang/src/main/java/io/ballerina/projects/ModuleContext.java
@@ -560,6 +560,11 @@ 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);
+    }
+
     Optional<MdDocumentContext> moduleMdContext() {
         return Optional.ofNullable(this.moduleMdContext);
     }

diff --git a/...llerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/CodeGenerator.java b/...llerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/CodeGenerator.java
@@ -98,11 +98,11 @@ private CompiledJarFile generate(BPackageSymbol packageSymbol) {
         HashMap<String, String> originalIdentifierMap = JvmDesugarPhase.encodeModuleIdentifiers(packageSymbol.bir);
 
         // TODO Get-rid of the following assignment
-        packageSymbol.compiledJarFile = jvmPackageGen.generate(packageSymbol.bir, true);
+        CompiledJarFile compiledJarFile = jvmPackageGen.generate(packageSymbol.bir, true);
 
         //Revert encoding identifier names
         JvmDesugarPhase.replaceEncodedModuleIdentifiers(packageSymbol.bir, originalIdentifierMap);
-        return packageSymbol.compiledJarFile;
+        return compiledJarFile;
     }
 
     private void populateExternalMap(JvmPackageGen jvmPackageGen) {

diff --git a/...llerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/AnnotationDesugar.java b/...llerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/AnnotationDesugar.java
@@ -840,7 +840,7 @@ private BLangLambdaFunction addReturnAndDefineLambda(BLangFunction function, BLa
 
         BInvokableSymbol lambdaFunctionSymbol = createInvokableSymbol(function, pkgID, owner);
         BLangLambdaFunction lambdaFunction = desugar.createLambdaFunction(function, lambdaFunctionSymbol, env);
-        lambdaFunction.capturedClosureEnv = env.createClone();
+        lambdaFunction.capturedClosureEnv = env;
 
         pkgNode.functions.add(function);
         pkgNode.topLevelNodes.add(function);

diff --git a/.../ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/ClosureDesugar.java b/.../ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/ClosureDesugar.java
@@ -1415,6 +1415,7 @@ public void visit(BLangLambdaFunction bLangLambdaFunction) {
             boolean isWorker = bLangLambdaFunction.function.flagSet.contains(Flag.WORKER);
             bLangLambdaFunction.enclMapSymbols = collectClosureMapSymbols(symbolEnv, enclInvokable, isWorker);
         }
+        bLangLambdaFunction.capturedClosureEnv = null;
         result = bLangLambdaFunction;
     }
 

diff --git a/...allerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/ClosureGenerator.java b/...allerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/ClosureGenerator.java
@@ -551,7 +551,7 @@ private BLangExpression createClosureForDefaultValue(String closureName, String
         BLangReturn returnStmt = ASTBuilderUtil.createReturnStmt(function.pos, (BLangBlockFunctionBody) function.body);
         returnStmt.expr = varNode.expr;
         BLangLambdaFunction lambdaFunction = createLambdaFunction(function);
-        lambdaFunction.capturedClosureEnv = env.createClone();
+        lambdaFunction.capturedClosureEnv = env;
         BInvokableSymbol varSymbol = createSimpleVariable(function, lambdaFunction, false);
         env.enclPkg.symbol.scope.define(function.symbol.name, function.symbol);
         env.enclPkg.functions.add(function);

diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/Desugar.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/Desugar.java
@@ -6222,7 +6222,7 @@ private BLangNode rewriteObjectMemberAccessAsField(BLangFieldBasedAccess fieldAc
 
         BLangLambdaFunction lambdaFunction = (BLangLambdaFunction) TreeBuilder.createLambdaFunctionNode();
         lambdaFunction.function = func;
-        lambdaFunction.capturedClosureEnv = env.createClone();
+        lambdaFunction.capturedClosureEnv = env;
         env.enclPkg.functions.add(func);
         env.enclPkg.topLevelNodes.add(func);
         //env.enclPkg.lambdaFunctions.add(lambdaFunction);
@@ -7809,6 +7809,7 @@ public void visit(BLangLambdaFunction bLangLambdaFunction) {
             funcSymbol.addAnnotation(this.strandAnnotAttachement.annotationAttachmentSymbol);
             funcSymbol.schedulerPolicy = SchedulerPolicy.ANY;
         }
+        bLangLambdaFunction.capturedClosureEnv = null;
         result = bLangLambdaFunction;
     }
 

diff --git a/...src/main/java/org/wso2/ballerinalang/compiler/semantics/model/symbols/BPackageSymbol.java b/...src/main/java/org/wso2/ballerinalang/compiler/semantics/model/symbols/BPackageSymbol.java
@@ -23,7 +23,6 @@
 import org.ballerinalang.model.symbols.SymbolKind;
 import org.ballerinalang.model.symbols.SymbolOrigin;
 import org.ballerinalang.repository.CompiledPackage;
-import org.wso2.ballerinalang.compiler.CompiledJarFile;
 import org.wso2.ballerinalang.compiler.bir.model.BIRNode;
 import org.wso2.ballerinalang.compiler.semantics.model.types.BPackageType;
 import org.wso2.ballerinalang.compiler.util.Name;
@@ -56,9 +55,6 @@ public class BPackageSymbol extends BTypeSymbol {
     public BIRNode.BIRPackage bir;   // TODO try to remove this
     public BIRPackageFile birPackageFile;
 
-    // kep code generated jar binary content in memory
-    public CompiledJarFile compiledJarFile;
-
     // TODO Refactor following two flags
     public boolean entryPointExists = false;