Check the closure for function symbols

ballerina-platform · Dec 15, 2023 · a0d5760 · a0d5760
1 parent fcaa0ef
commit a0d5760
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 10 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
@@ -5876,8 +5876,7 @@ private SymbolEnv findEnclosingInvokableEnv(SymbolEnv env, BLangInvokableNode en
     private void updateClosureVariable(BVarSymbol varSymbol, BLangInvokableNode encInvokable, Location pos) {
         if (!varSymbol.closure) {
             SymbolEnv encInvokableEnv = findEnclosingInvokableEnv(env, encInvokable);
-            BSymbol resolvedSymbol =
-                    symResolver.lookupClosureVarSymbol(encInvokableEnv, varSymbol);
+            BSymbol resolvedSymbol = symResolver.lookupClosureVarSymbol(encInvokableEnv, varSymbol);
             if (resolvedSymbol != symTable.notFoundSymbol) {
                 varSymbol.closure = true;
                 ((BLangFunction) encInvokable).closureVarSymbols.add(new ClosureVarSymbol(varSymbol, pos));

diff --git a/...er/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/QueryDesugar.java b/...er/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/QueryDesugar.java
@@ -1908,7 +1908,7 @@ public void visit(BLangSimpleVarRef bLangSimpleVarRef) {
             result = bLangSimpleVarRef;
             return;
         }
-        if (symbol.kind == SymbolKind.VARIABLE) {
+        if (symbol.kind == SymbolKind.VARIABLE || symbol.kind == SymbolKind.FUNCTION) {
             BVarSymbol originalSymbol = ((BVarSymbol) symbol).originalSymbol;
             if (originalSymbol != null) {
                 symbol = originalSymbol;
@@ -1965,14 +1965,8 @@ public void visit(BLangSimpleVarRef bLangSimpleVarRef) {
                 bLangSimpleVarRef.symbol = symbol;
                 bLangSimpleVarRef.varSymbol = symbol;
             }
-        } else if (resolvedSymbol != symTable.notFoundSymbol) {
+        } else if (!resolvedSymbol.closure && resolvedSymbol != symTable.notFoundSymbol) {
             resolvedSymbol.closure = true;
-            // When there's a type guard, there can be a enclSymbol before type narrowing.
-            // So, we have to mark that as a closure as well.
-            BSymbol enclSymbol = symResolver.lookupClosureVarSymbol(env.enclEnv, symbol);
-            if (enclSymbol != null && enclSymbol != symTable.notFoundSymbol) {
-                enclSymbol.closure = true;
-            }
         }
         result = bLangSimpleVarRef;
     }

diff --git a/...-test/src/test/java/org/ballerinalang/test/expressions/ternary/TernaryExpressionTest.java b/...-test/src/test/java/org/ballerinalang/test/expressions/ternary/TernaryExpressionTest.java
@@ -418,6 +418,16 @@ public void testTernaryWithQueryForTwoVariables() {
         BRunUtil.invoke(compileResult, "testTernaryWithQueryForTwoVariables");
     }
 
+    @Test
+    public void testTernaryWithQueryWithFunctionPointers() {
+        BRunUtil.invoke(compileResult, "testTernaryWithQueryWithFunctionPointers");
+    }
+
+    @Test
+    public void testTernaryWithQueryWithFunctionAsClosure() {
+        BRunUtil.invoke(compileResult, "testTernaryWithQueryWithFunctionAsClosure");
+    }
+
     @Test(description = "Test type narrowing for ternary expression")
     public void testTernaryTypeNarrow() {
         CompileResult compileResult = BCompileUtil.compile("test-src/expressions/ternary/ternary_expr_type_narrow.bal");

diff --git a/tests/jballerina-unit-test/src/test/resources/test-src/expressions/ternary/ternary-expr.bal b/tests/jballerina-unit-test/src/test/resources/test-src/expressions/ternary/ternary-expr.bal
@@ -417,6 +417,42 @@ function testTernaryWithQueryForTwoVariables() {
     assertEquals([2, 2], elseResult);
 }
 
+function testTernaryWithQueryWithFunctionPointers() {
+    int? i = 3;
+    var k = function() returns int? {
+        return i;
+    };
+
+    int|int[] thenResult = i is int ?
+        from var _ in [1, 2]
+        where k() + 2 == 5
+        select 2 : 2;
+    assertEquals([2,2], thenResult);
+
+    int|int[] elseResult = i is () ? 2 :
+        from var _ in [1, 2]
+        where k() + 2 == 5
+        select 2;
+    assertEquals([2,2], elseResult);
+}
+
+function testTernaryWithQueryWithFunctionAsClosure() {
+    ()|function() returns int fn = function() returns int {
+        return 3;
+    };
+    int|int[] thenResult = fn !is () ?
+        from var _ in [1, 2]
+        where fn() + 2 == 5
+        select 2 : 2;
+    assertEquals([2,2], thenResult);
+
+    int|int[] elseResult = fn is () ? 2 :
+        from var _ in [1, 2]
+        where fn() + 2 == 5
+        select 2;
+    assertEquals([2,2], elseResult);
+}
+
 boolean cond2 = true;
 int i1 =10;
 byte j1 = 100;