Set determine type for query and conditional expressions when the expected types are incompatible

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/QueryTypeChecker.java

@@ -183,6 +183,7 @@ public void checkQueryType(BLangQueryExpr queryExpr, TypeChecker.AnalyzerData da
         if (finalClause.getKind() == NodeKind.SELECT) {
             actualType = resolveQueryType(commonAnalyzerData.queryEnvs.peek(),
                     ((BLangSelectClause) finalClause).expression, data.expType, queryExpr, clauses, data);
+            queryExpr.setDeterminedType(actualType);
             actualType = (actualType == symTable.semanticError) ? actualType : types.checkType(queryExpr.pos,
                     actualType, data.expType, DiagnosticErrorCode.INCOMPATIBLE_TYPES);
         } else {
@@ -196,6 +197,7 @@ public void checkQueryType(BLangQueryExpr queryExpr, TypeChecker.AnalyzerData da
             if (completionType != null) {
                 queryType = BUnionType.create(null, queryType, completionType);
             }
+            queryExpr.setDeterminedType(queryType);
             actualType = types.checkType(finalClauseExpr.pos, queryType, data.expType,
                     DiagnosticErrorCode.INCOMPATIBLE_TYPES);
         }

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/TypeChecker.java

@@ -5013,30 +5013,37 @@ private boolean isReferencingNonWorker(BLangExpression expression, AnalyzerData
         return true;
     }
 
-
     public void visit(BLangTernaryExpr ternaryExpr, AnalyzerData data) {
         BType condExprType = checkExpr(ternaryExpr.expr, this.symTable.booleanType, data);
 
         SymbolEnv thenEnv = typeNarrower.evaluateTruth(ternaryExpr.expr, ternaryExpr.thenExpr, data.env);
+        BType thenActualType = silentTypeCheckExpr(ternaryExpr.thenExpr, symTable.noType, data);
         BType thenType = checkExpr(ternaryExpr.thenExpr, thenEnv, data.expType, data);
 
         SymbolEnv elseEnv = typeNarrower.evaluateFalsity(ternaryExpr.expr, ternaryExpr.elseExpr, data.env, false);
+        BType elseActualType = silentTypeCheckExpr(ternaryExpr.elseExpr, symTable.noType, data);
         BType elseType = checkExpr(ternaryExpr.elseExpr, elseEnv, data.expType, data);
 
         if (condExprType == symTable.semanticError || thenType == symTable.semanticError ||
                 elseType == symTable.semanticError) {
             data.resultType = symTable.semanticError;
         } else if (data.expType == symTable.noType) {
-            if (types.isAssignable(elseType, thenType)) {
-                data.resultType = thenType;
-            } else if (types.isAssignable(thenType, elseType)) {
-                data.resultType = elseType;
-            } else {
-                data.resultType = BUnionType.create(null, thenType, elseType);
-            }
+            data.resultType = getConditionalExprType(thenType, elseType);
         } else {
             data.resultType = data.expType;
         }
+
+        ternaryExpr.setDeterminedType(getConditionalExprType(thenActualType, elseActualType));
+    }
+
+    private BType getConditionalExprType(BType lhsType, BType rhsType) {
+        if (types.isAssignable(rhsType, lhsType)) {
+            return lhsType;
+        }
+        if (types.isAssignable(lhsType, rhsType)) {
+            return rhsType;
+        }
+        return BUnionType.create(null, lhsType, rhsType);
     }
 
     public void visit(BLangWaitExpr waitExpr, AnalyzerData data) {
@@ -5286,24 +5293,21 @@ private BType getXMLConstituents(BType bType) {
 
     public void visit(BLangElvisExpr elvisExpr, AnalyzerData data) {
         BType lhsType = checkExpr(elvisExpr.lhsExpr, data);
-        BType actualType = lhsType == symTable.semanticError ?
+        BType lhsActualType = lhsType == symTable.semanticError ?
                 symTable.semanticError : validateElvisExprLhsExpr(elvisExpr, lhsType);
+        BType rhsActualType = silentTypeCheckExpr(elvisExpr.rhsExpr, symTable.noType, data);
         BType rhsReturnType = checkExpr(elvisExpr.rhsExpr, data.expType, data);
-        BType lhsReturnType = types.checkType(elvisExpr.lhsExpr.pos, actualType, data.expType,
+        BType lhsReturnType = types.checkType(elvisExpr.lhsExpr.pos, lhsActualType, data.expType,
                 DiagnosticErrorCode.INCOMPATIBLE_TYPES);
         if (rhsReturnType == symTable.semanticError || lhsReturnType == symTable.semanticError) {
             data.resultType = symTable.semanticError;
         } else if (data.expType == symTable.noType) {
-            if (types.isAssignable(rhsReturnType, lhsReturnType)) {
-                data.resultType = lhsReturnType;
-            } else if (types.isAssignable(lhsReturnType, rhsReturnType)) {
-                data.resultType = rhsReturnType;
-            } else {
-                data.resultType = BUnionType.create(null, lhsReturnType, rhsReturnType);
-            }
+            data.resultType = getConditionalExprType(lhsReturnType, rhsReturnType);
         } else {
             data.resultType = data.expType;
         }
+
+        elvisExpr.setDeterminedType(getConditionalExprType(lhsActualType, rhsActualType));
     }
 
     @Override

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/codeaction/providers/changetype/FixReturnTypeCodeAction.java

@@ -69,14 +69,37 @@
             return false;
         }
 
-        //Suggest the code action only if the immediate parent of the matched node is either of return statement,
-        //check expression, check action.
+        // Suggest code action if the node is aligned with a check keyword.
         NonTerminalNode matchedNode = positionDetails.matchedNode();
-        if (matchedNode.parent() != null && matchedNode.parent().kind() != SyntaxKind.RETURN_STATEMENT &&
-                matchedNode.kind() != SyntaxKind.CHECK_EXPRESSION && 
-                matchedNode.kind() != SyntaxKind.CHECK_ACTION) {
+        if (matchedNode.kind() == SyntaxKind.CHECK_ACTION || matchedNode.kind() == SyntaxKind.CHECK_EXPRESSION) {
+            return CodeActionNodeValidator.validate(context.nodeAtRange());
+        }
+
+        // Suggest code action if the node is an expression inside a return statement.
+        NonTerminalNode parentNode = matchedNode.parent();
+        if (parentNode == null) {
+            return false;
+        }
+        if (parentNode.kind() == SyntaxKind.RETURN_KEYWORD) {
+            return CodeActionNodeValidator.validate(context.nodeAtRange());
+        }
+
+        // Suggest code action if the node is an expression nested into a return statement.
+        NonTerminalNode grandParentNode = parentNode.parent();
+        if (grandParentNode == null) {
             return false;
         }
+        if (grandParentNode.kind() == SyntaxKind.RETURN_STATEMENT) {
+            return CodeActionNodeValidator.validate(context.nodeAtRange());
+        }
+
+        // Suggest code action if the node is an expression inside a collect clause.
+        if (matchedNode.kind() == SyntaxKind.COLLECT_CLAUSE) {
+            NonTerminalNode ancestorNode = grandParentNode.parent();
+            if (ancestorNode == null || ancestorNode.kind() != SyntaxKind.RETURN_STATEMENT) {
+                return false;
+            }
+        }
 
         return CodeActionNodeValidator.validate(context.nodeAtRange());
     }

language-server/modules/langserver-core/src/test/java/org/ballerinalang/langserver/codeaction/FixReturnTypeTest.java

@@ -67,7 +67,19 @@ public Object[][] dataProvider() {
                 {"fixReturnTypeInUnionContext3.json"},
                 {"fixReturnTypeInUnionContext4.json"},
                 {"fixReturnTypeInCommitAction.json"},
-                {"fixReturnTypeInMain1.json"}
+                {"fixReturnTypeInMain1.json"},
+                {"fixReturnTypeForConditionalExpr1.json"},
+                {"fixReturnTypeForConditionalExpr2.json"},
+                {"fixReturnTypeForConditionalExpr3.json"},
+                {"fixReturnTypeForConditionalExpr4.json"},
+                {"fixReturnTypeForConditionalExpr5.json"},
+                {"fixReturnTypeForConditionalExpr6.json"},
+                {"fixReturnTypeForConditionalExpr7.json"},
+                {"fixReturnTypeForConditionalExpr8.json"},
+                {"fixReturnTypeForQueryExpr1.json"},
+                {"fixReturnTypeForQueryExpr2.json"},
+                {"fixReturnTypeForQueryExpr3.json"},
+                {"fixReturnTypeForQueryExpr4.json"}
         };
     }
 

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr1.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 6,
+    "character": 39
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'boolean|int'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 5,
+              "character": 76
+            },
+            "end": {
+              "line": 5,
+              "character": 76
+            }
+          },
+          "newText": " returns boolean|int"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr2.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 6,
+    "character": 51
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'boolean|int'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 5,
+              "character": 76
+            },
+            "end": {
+              "line": 5,
+              "character": 76
+            }
+          },
+          "newText": " returns boolean|int"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr3.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 11,
+    "character": 72
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'module1:TestRecord1|int'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 9,
+              "character": 76
+            },
+            "end": {
+              "line": 9,
+              "character": 76
+            }
+          },
+          "newText": " returns module1:TestRecord1|int"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr4.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 11,
+    "character": 81
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'module1:TestRecord1|int'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 9,
+              "character": 76
+            },
+            "end": {
+              "line": 9,
+              "character": 76
+            }
+          },
+          "newText": " returns module1:TestRecord1|int"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr5.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 18,
+    "character": 18
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'string|boolean'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 17,
+              "character": 50
+            },
+            "end": {
+              "line": 17,
+              "character": 50
+            }
+          },
+          "newText": " returns string|boolean"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr6.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 18,
+    "character": 28
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'string|boolean'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 17,
+              "character": 50
+            },
+            "end": {
+              "line": 17,
+              "character": 50
+            }
+          },
+          "newText": " returns string|boolean"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr7.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 23,
+    "character": 74
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'string|int|boolean|module1:TestRecord1|float'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 21,
+              "character": 91
+            },
+            "end": {
+              "line": 21,
+              "character": 91
+            }
+          },
+          "newText": " returns string|int|boolean|module1:TestRecord1|float"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/fix-return-type/config/fixReturnTypeForConditionalExpr8.json

@@ -0,0 +1,29 @@
+{
+  "position": {
+    "line": 23,
+    "character": 86
+  },
+  "source": "fixReturnTypeForConditionalExpr.bal",
+  "expected": [
+    {
+      "title": "Change return type to 'string|int|boolean|module1:TestRecord1|float'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 21,
+              "character": 91
+            },
+            "end": {
+              "line": 21,
+              "character": 91
+            }
+          },
+          "newText": " returns string|int|boolean|module1:TestRecord1|float"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}