From aeab20ac0269b34ec4333c932a86564b2e6b0cdb Mon Sep 17 00:00:00 2001
From: Thushara-Piyasekara <thusharap@wso2.com>
Date: Tue, 26 Mar 2024 17:42:38 +0530
Subject: [PATCH] Add boolean check for $moduleStart() invocation

---
 .../bir/codegen/methodgen/MethodGen.java      | 30 +++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/methodgen/MethodGen.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/methodgen/MethodGen.java
index 06c8df104232..4bab31e250e6 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/methodgen/MethodGen.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/methodgen/MethodGen.java
@@ -282,8 +282,9 @@ public void genJMethodForBFunc(BIRFunction func, ClassWriter cw, BIRPackage modu
         MethodVisitor mv = cw.visitMethod(access, funcName, desc, null, null);
         mv.visitCode();
 
-        visitModuleStartFunction(module.packageID, funcName, mv);
+        visitStartFunction(module.packageID, funcName, mv);
         visitModuleInitFunction(module.packageID, funcName, mv);
+        visitModuleStartFunction(module.packageID, funcName, mv);
 
         Label methodStartLabel = new Label();
         mv.visitLabel(methodStartLabel);
@@ -420,8 +421,8 @@ private BType getReturnType(BIRFunction func) {
         return retType;
     }
 
-    private void visitModuleStartFunction(PackageID packageID, String funcName, MethodVisitor mv) {
-        if (!isModuleStartFunction(funcName)) {
+    private void visitStartFunction(PackageID packageID, String funcName, MethodVisitor mv) {
+        if (!isStartFunction(funcName)) {
             return;
         }
         mv.visitInsn(ICONST_1);
@@ -446,6 +447,20 @@ private void visitModuleInitFunction(PackageID packageID, String funcName, Metho
         mv.visitFieldInsn(PUTSTATIC, moduleClass, MODULE_INIT_ATTEMPTED, "Z");
     }
 
+    private void visitModuleStartFunction(PackageID packageID, String funcName, MethodVisitor mv) {
+        if (!isModuleStartFunction(funcName)) {
+            return;
+        }
+
+        String moduleClass = JvmCodeGenUtil.getModuleLevelClassName(packageID, MODULE_INIT_CLASS_NAME);
+        mv.visitFieldInsn(GETSTATIC, moduleClass, MODULE_START_ATTEMPTED, "Z");
+        Label labelIf = new Label();
+        mv.visitJumpInsn(IFEQ, labelIf);
+        mv.visitInsn(ACONST_NULL);
+        mv.visitInsn(ARETURN);
+        mv.visitLabel(labelIf);
+    }
+
     private void setChannelDetailsToStrand(BIRFunction func, int localVarOffset, MethodVisitor mv,
                                            int invocationVarIndex) {
         // these channel info is required to notify datachannels, when there is a panic
@@ -725,7 +740,7 @@ private void processTerminator(MethodVisitor mv, BIRFunction func, BIRPackage mo
                     false);
         }
         //set module start success to true for $_init class
-        if (isModuleStartFunction(funcName) && terminator.kind == InstructionKind.RETURN) {
+        if (isStartFunction(funcName) && terminator.kind == InstructionKind.RETURN) {
             mv.visitInsn(ICONST_1);
             mv.visitFieldInsn(PUTSTATIC, JvmCodeGenUtil.getModuleLevelClassName(module.packageID,
                             MODULE_INIT_CLASS_NAME),
@@ -739,7 +754,7 @@ private boolean isModuleTestInitFunction(BIRFunction func) {
                         .encodeModuleSpecialFuncName(".<testinit>"));
     }
 
-    private boolean isModuleStartFunction(String functionName) {
+    private boolean isStartFunction(String functionName) {
         return functionName
                 .equals(MethodGenUtils.encodeModuleSpecialFuncName(MethodGenUtils.START_FUNCTION_SUFFIX));
     }
@@ -749,6 +764,11 @@ private boolean isModuleInitFunction(String functionName) {
                 .equals(MethodGenUtils.encodeModuleSpecialFuncName(JvmConstants.MODULE_INIT_METHOD));
     }
 
+    private boolean isModuleStartFunction(String functionName) {
+        return functionName
+                .equals(MethodGenUtils.encodeModuleSpecialFuncName(JvmConstants.MODULE_START_METHOD));
+    }
+
     private void genGetFrameOnResumeIndex(int localVarOffset, MethodVisitor mv, String frameName) {
         mv.visitVarInsn(ALOAD, localVarOffset);
         mv.visitInsn(DUP);