From 183b07dcf19af46053ad8f421c19d39c1fe0f9f1 Mon Sep 17 00:00:00 2001
From: mindula <rowelmindula@gmail.com>
Date: Mon, 24 Jun 2024 10:25:58 +0530
Subject: [PATCH 1/5] Fix NPE in lambda functions with check expressions

---
 .../compiler/semantics/analyzer/CodeAnalyzer.java             | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
index cbcaace973ce..296ccfe65ca3 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
@@ -3084,6 +3084,9 @@ public void visit(BLangLambdaFunction bLangLambdaFunction, AnalyzerData data) {
                                                                                   bLangLambdaFunction.function);
             }
         }
+
+        boolean failureHandled = data.failureHandled;
+        data.failureHandled = false;
         // If this is a worker we are already in a worker action system,
         // if not we need to initiate a worker action system
         if (isWorker) {
@@ -3100,6 +3103,7 @@ public void visit(BLangLambdaFunction bLangLambdaFunction, AnalyzerData data) {
                 this.finalizeCurrentWorkerActionSystem(data);
             }
         }
+        data.failureHandled = failureHandled;
 
         if (isWorker) {
             data.workerActionSystemStack.peek().endWorkerActionStateMachine();

From 49db9b44a33929467ebcb3b1e1d5db7a954155be Mon Sep 17 00:00:00 2001
From: mindula <rowelmindula@gmail.com>
Date: Mon, 24 Jun 2024 10:27:16 +0530
Subject: [PATCH 2/5] Add tests

---
 .../checkedexpr/CheckedExprNegativeTest.java  |  2 ++
 ...ed_error_return_type_mismatch_negative.bal | 30 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/expressions/checkedexpr/CheckedExprNegativeTest.java b/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/expressions/checkedexpr/CheckedExprNegativeTest.java
index 941d47ab1c82..bf12b214621d 100644
--- a/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/expressions/checkedexpr/CheckedExprNegativeTest.java
+++ b/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/expressions/checkedexpr/CheckedExprNegativeTest.java
@@ -96,6 +96,8 @@ public void testCheckedErrorvsReturnTypeMismatch() {
         BAssertUtil.validateError(compile, i++, ERROR_MISMATCH_ERR_MSG, 76, 20);
         BAssertUtil.validateError(compile, i++, ERROR_MISMATCH_ERR_MSG, 76, 31);
         BAssertUtil.validateError(compile, i++, ERROR_MISMATCH_ERR_MSG, 85, 13);
+        BAssertUtil.validateError(compile, i++, ERROR_MISMATCH_ERR_MSG, 92, 28);
+        BAssertUtil.validateError(compile, i++, ERROR_MISMATCH_ERR_MSG, 107, 24);
         Assert.assertEquals(compile.getWarnCount(), 6);
         Assert.assertEquals(compile.getErrorCount(), i - 6);
     }
diff --git a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
index 0d14d288bae0..3ee3b2a18459 100644
--- a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
+++ b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
@@ -84,3 +84,33 @@ function testCheckedExprWithNoErrorType3() {
     int|error i = error("Error");
     int _ = check i;
 }
+
+function testOnFailWithCheckInLambdaFunction() {
+    function _ = function() {
+        do {
+            function _ = function () {
+                string _ = check getVal("");
+            };
+        } on fail {
+        }
+    };
+}
+
+function getVal(string str) returns string|error {
+    return "";
+}
+
+function testOnFailWithCheckInWorkers() {
+    string[] s = [];
+    worker A {
+        s.forEach(function (string k) {
+            string _ = check getVal2();
+        });
+
+    } on fail {
+    }
+}
+
+function getVal2() returns string|error {
+    return "";
+}

From c5220231467f0339b8ef1fd28e828303bcf38010 Mon Sep 17 00:00:00 2001
From: mindula <rowelmindula@gmail.com>
Date: Fri, 19 Jul 2024 13:59:47 +0530
Subject: [PATCH 3/5] Address review comments

---
 .../ballerinalang/compiler/desugar/Desugar.java   |  3 +++
 .../compiler/semantics/analyzer/CodeAnalyzer.java |  6 +++---
 ...hecked_error_return_type_mismatch_negative.bal |  1 -
 .../test-src/statements/dostatement/do-stmt.bal   | 15 +++++++++++++++
 4 files changed, 21 insertions(+), 4 deletions(-)

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
index 31fc3db1cb13..7a71e9575186 100644
--- 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
@@ -8075,7 +8075,10 @@ public void visit(BLangTypeConversionExpr conversionExpr) {
 
     @Override
     public void visit(BLangLambdaFunction bLangLambdaFunction) {
+        BLangOnFailClause onFailClause = this.onFailClause;
+        this.onFailClause = null;
         bLangLambdaFunction.function = rewrite(bLangLambdaFunction.function, bLangLambdaFunction.capturedClosureEnv);
+        this.onFailClause = onFailClause;
         BLangFunction function = bLangLambdaFunction.function;
         // Add `@strand {thread: "any"}` annotation to an isolated named worker declaration in an isolated function.
         if (function.flagSet.contains(Flag.WORKER) && Symbols.isFlagOn(function.symbol.type.flags, Flags.ISOLATED) &&
diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
index 296ccfe65ca3..75b43ccd7341 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
@@ -497,6 +497,8 @@ private void validateParams(BLangFunction funcNode, AnalyzerData data) {
     }
 
     private void visitFunction(BLangFunction funcNode, AnalyzerData data) {
+        boolean failureHandled = data.failureHandled;
+        data.failureHandled = false;
         data.env = SymbolEnv.createFunctionEnv(funcNode, funcNode.symbol.scope, data.env);
         data.returnWithinTransactionCheckStack.push(true);
         data.returnTypes.push(new LinkedHashSet<>());
@@ -523,6 +525,7 @@ private void visitFunction(BLangFunction funcNode, AnalyzerData data) {
         data.returnTypes.pop();
         data.returnWithinTransactionCheckStack.pop();
         data.transactionalFuncCheckStack.pop();
+        data.failureHandled = failureHandled;
     }
 
     private boolean isPublicInvokableNode(BLangInvokableNode invNode) {
@@ -3085,8 +3088,6 @@ public void visit(BLangLambdaFunction bLangLambdaFunction, AnalyzerData data) {
             }
         }
 
-        boolean failureHandled = data.failureHandled;
-        data.failureHandled = false;
         // If this is a worker we are already in a worker action system,
         // if not we need to initiate a worker action system
         if (isWorker) {
@@ -3103,7 +3104,6 @@ public void visit(BLangLambdaFunction bLangLambdaFunction, AnalyzerData data) {
                 this.finalizeCurrentWorkerActionSystem(data);
             }
         }
-        data.failureHandled = failureHandled;
 
         if (isWorker) {
             data.workerActionSystemStack.peek().endWorkerActionStateMachine();
diff --git a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
index 3ee3b2a18459..49eb60336d71 100644
--- a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
+++ b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
@@ -106,7 +106,6 @@ function testOnFailWithCheckInWorkers() {
         s.forEach(function (string k) {
             string _ = check getVal2();
         });
-
     } on fail {
     }
 }
diff --git a/tests/jballerina-unit-test/src/test/resources/test-src/statements/dostatement/do-stmt.bal b/tests/jballerina-unit-test/src/test/resources/test-src/statements/dostatement/do-stmt.bal
index c189f27af33e..f73c6ba6697f 100644
--- a/tests/jballerina-unit-test/src/test/resources/test-src/statements/dostatement/do-stmt.bal
+++ b/tests/jballerina-unit-test/src/test/resources/test-src/statements/dostatement/do-stmt.bal
@@ -388,3 +388,18 @@ function assertEquality(any|error expected, any|error actual) {
     panic error(ASSERTION_ERROR_REASON,
             message = "expected '" + expectedValAsString + "', found '" + actualValAsString + "'");
 }
+
+function testOnFailWithCheckAndReturnType(function s = function() {
+            do {
+                function _ = function() returns error? {
+                    string _ = check getVal("");
+                };
+            } on fail {
+            }
+        }) returns string|error {
+    return "";
+}
+
+function getVal(string s) returns string|error {
+    return "";
+}

From c1082e7603459916f39f0ca57a4ceec43017fdd4 Mon Sep 17 00:00:00 2001
From: mindula <rowelmindula@gmail.com>
Date: Fri, 19 Jul 2024 17:09:06 +0530
Subject: [PATCH 4/5] Move the logic to the visitor of BLangFunction

---
 .../org/wso2/ballerinalang/compiler/desugar/Desugar.java    | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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
index 7a71e9575186..3886a18f8e07 100644
--- 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
@@ -1359,6 +1359,8 @@ public void visit(BLangResourceFunction resourceFunction) {
 
     @Override
     public void visit(BLangFunction funcNode) {
+        BLangOnFailClause onFailClause = this.onFailClause;
+        this.onFailClause = null;
         Map<Name, BLangStatement> prevXmlnsDecls = this.stmtsToBePropagatedToQuery;
         if (!funcNode.flagSet.contains(Flag.QUERY_LAMBDA)) {
             this.stmtsToBePropagatedToQuery = new HashMap<>();
@@ -1385,6 +1387,7 @@ public void visit(BLangFunction funcNode) {
         }
         this.stmtsToBePropagatedToQuery = prevXmlnsDecls;
         result = funcNode;
+        this.onFailClause = onFailClause;
     }
 
     @Override
@@ -8075,10 +8078,7 @@ public void visit(BLangTypeConversionExpr conversionExpr) {
 
     @Override
     public void visit(BLangLambdaFunction bLangLambdaFunction) {
-        BLangOnFailClause onFailClause = this.onFailClause;
-        this.onFailClause = null;
         bLangLambdaFunction.function = rewrite(bLangLambdaFunction.function, bLangLambdaFunction.capturedClosureEnv);
-        this.onFailClause = onFailClause;
         BLangFunction function = bLangLambdaFunction.function;
         // Add `@strand {thread: "any"}` annotation to an isolated named worker declaration in an isolated function.
         if (function.flagSet.contains(Flag.WORKER) && Symbols.isFlagOn(function.symbol.type.flags, Flags.ISOLATED) &&

From 6c684040a27704cabf4ad4c88328b6a82b798054 Mon Sep 17 00:00:00 2001
From: mindula <rowelmindula@gmail.com>
Date: Mon, 22 Jul 2024 11:29:18 +0530
Subject: [PATCH 5/5] Remove unnecessary spaces

---
 .../compiler/semantics/analyzer/CodeAnalyzer.java             | 1 -
 .../checked_error_return_type_mismatch_negative.bal           | 4 ++--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
index 75b43ccd7341..d9a905b03e60 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
@@ -3087,7 +3087,6 @@ public void visit(BLangLambdaFunction bLangLambdaFunction, AnalyzerData data) {
                                                                                   bLangLambdaFunction.function);
             }
         }
-
         // If this is a worker we are already in a worker action system,
         // if not we need to initiate a worker action system
         if (isWorker) {
diff --git a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
index 49eb60336d71..4955294b03df 100644
--- a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
+++ b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/checkedexpr/checked_error_return_type_mismatch_negative.bal
@@ -88,7 +88,7 @@ function testCheckedExprWithNoErrorType3() {
 function testOnFailWithCheckInLambdaFunction() {
     function _ = function() {
         do {
-            function _ = function () {
+            function _ = function() {
                 string _ = check getVal("");
             };
         } on fail {
@@ -103,7 +103,7 @@ function getVal(string str) returns string|error {
 function testOnFailWithCheckInWorkers() {
     string[] s = [];
     worker A {
-        s.forEach(function (string k) {
+        s.forEach(function(string k) {
             string _ = check getVal2();
         });
     } on fail {