Merge pull request #39602 from IMS94/bug_fixes

Fix LS high priority bugs

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/codeaction/CodeActionUtil.java

@@ -86,6 +86,7 @@
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import static org.ballerinalang.langserver.common.utils.CommonUtil.LINE_SEPARATOR;
 
@@ -382,44 +383,45 @@ public static List<TextEdit> getTypeGuardCodeActionEdits(String varName, Range r
         String spaces = StringUtils.repeat(' ', range.getStart().getCharacter());
         String padding = LINE_SEPARATOR + LINE_SEPARATOR + spaces;
 
-        boolean hasError = CodeActionUtil.hasErrorMemberType(unionType);
+        List<TypeSymbol> errorMembers = unionType.memberTypeDescriptors().stream()
+                .filter(typeSymbol -> CommonUtil.getRawType(typeSymbol).typeKind() == TypeDescKind.ERROR)
+                .collect(Collectors.toList());
 
         List<TypeSymbol> members = new ArrayList<>(unionType.memberTypeDescriptors());
-        long errorTypesCount = unionType.memberTypeDescriptors().stream()
-                .map(CommonUtil::getRawType)
-                .filter(t -> t.typeKind() == TypeDescKind.ERROR)
-                .count();
+
         if (members.size() == 1) {
             // Skip type guard
             return edits;
         }
-        boolean transitiveBinaryUnion = unionType.memberTypeDescriptors().size() - errorTypesCount == 1;
+        boolean transitiveBinaryUnion = unionType.memberTypeDescriptors().size() - errorMembers.size() == 1;
         if (transitiveBinaryUnion) {
-            members.removeIf(s -> s.typeKind() == TypeDescKind.ERROR);
+            members.removeIf(typeSymbol -> CommonUtil.getRawType(typeSymbol).typeKind() == TypeDescKind.ERROR);
         }
-        // Check is binary union type with error type
-        if ((unionType.memberTypeDescriptors().size() == 2 || transitiveBinaryUnion) && hasError) {
+
+        ImportsAcceptor importsAcceptor = new ImportsAcceptor(context);
+        // Check if a union type with error and one non-error type
+        if ((unionType.memberTypeDescriptors().size() == 2 || transitiveBinaryUnion) && !errorMembers.isEmpty()) {
             members.forEach(bType -> {
                 if (bType.typeKind() == TypeDescKind.NIL) {
                     // if (foo() is error) {...}
-                    String newText = generateIfElseText(varName, spaces, padding, Collections.singletonList("error"));
+                    String type = errorMembers.get(0).signature();
+                    type = FunctionGenerator.processModuleIDsInText(importsAcceptor, type, context);
+                    String newText = generateIfElseText(varName, spaces, padding, List.of(type));
                     edits.add(new TextEdit(newTextRange, newText));
                 } else {
                     // if (foo() is int) {...} else {...}
                     String type = CodeActionUtil.getPossibleType(bType, edits, context).orElseThrow();
-                    String newText = generateIfElseText(varName, spaces, padding, Collections.singletonList(type));
+                    String newText = generateIfElseText(varName, spaces, padding, List.of(type));
                     edits.add(new TextEdit(newTextRange, newText));
                 }
             });
         } else {
-            boolean addErrorTypeAtEnd;
+            boolean addErrorTypeAtEnd = false;
             List<TypeSymbol> tMembers = new ArrayList<>((unionType).memberTypeDescriptors());
-            if (errorTypesCount > 1) {
+            if (errorMembers.size() > 1) {
                 // merge all error types into generic `error` type
-                tMembers.removeIf(s -> s.typeKind() == TypeDescKind.ERROR);
+                tMembers.removeIf(typeSymbol -> CommonUtil.getRawType(typeSymbol).typeKind() == TypeDescKind.ERROR);
                 addErrorTypeAtEnd = true;
-            } else {
-                addErrorTypeAtEnd = false;
             }
             List<String> memberTypes = new ArrayList<>();
             for (TypeSymbol tMember : tMembers) {
@@ -430,6 +432,8 @@ public static List<TextEdit> getTypeGuardCodeActionEdits(String varName, Range r
             }
             edits.add(new TextEdit(newTextRange, generateIfElseText(varName, spaces, padding, memberTypes)));
         }
+
+        edits.addAll(importsAcceptor.getNewImportTextEdits());
         return edits;
     }
 
@@ -511,9 +515,17 @@ public static List<TextEdit> getAddCheckTextEdits(Position pos, NonTerminalNode
             }
         }
 
+        // If we are in a method call expression and the expression part already doesn't have a brace, we have to add
+        // braces to prevent the "check" being added to the entire method call expression.
+        if (matchedNode.kind() != SyntaxKind.BRACED_EXPRESSION && 
+                matchedNode.parent().kind() == SyntaxKind.METHOD_CALL) {
+            Position endPos = PositionUtil.toPosition(matchedNode.lineRange().endLine());
+            edits.add(new TextEdit(new Range(pos, pos), "("));
+            edits.add(new TextEdit(new Range(endPos, endPos), ")"));
+        }
+
         // Add `check` expression text edit
-        Position insertPos = new Position(pos.getLine(), pos.getCharacter());
-        edits.add(new TextEdit(new Range(insertPos, insertPos), "check "));
+        edits.add(new TextEdit(new Range(pos, pos), "check "));
 
         // Add parent function return change text edits
         if (!returnText.isEmpty()) {

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

@@ -53,7 +53,11 @@
 public class AddCheckCodeAction extends TypeCastCodeAction {
 
     public static final String NAME = "Add Check";
-    public static final Set<String> DIAGNOSTIC_CODES = Set.of("BCE2652", "BCE2066", "BCE2068", "BCE2800", "BCE3998");
+    private static final String DIAGNOSTIC_CODE_3998 = "BCE3998";
+    private static final String DIAGNOSTIC_CODE_2068 = "BCE2068";
+    private static final String DIAGNOSTIC_CODE_2800 = "BCE2800";
+    public static final Set<String> DIAGNOSTIC_CODES = Set.of("BCE2652", "BCE2066",
+            DIAGNOSTIC_CODE_2068, DIAGNOSTIC_CODE_2800, DIAGNOSTIC_CODE_3998);
 
     public AddCheckCodeAction() {
         super();
@@ -78,14 +82,14 @@ public List<CodeAction> getCodeActions(Diagnostic diagnostic, DiagBasedPositionD
         }
 
         Optional<TypeSymbol> foundType;
-        if ("BCE2068".equals(diagnostic.diagnosticInfo().code())) {
+        if (DIAGNOSTIC_CODE_2068.equals(diagnostic.diagnosticInfo().code())) {
             foundType = positionDetails.diagnosticProperty(
                     CodeActionUtil.getDiagPropertyFilterFunction(
                             DiagBasedPositionDetails.DIAG_PROP_INCOMPATIBLE_TYPES_FOUND_SYMBOL_INDEX));
-        } else if ("BCE2800".equals(diagnostic.diagnosticInfo().code())) {
+        } else if (DIAGNOSTIC_CODE_2800.equals(diagnostic.diagnosticInfo().code())) {
             foundType = positionDetails.diagnosticProperty(
                     DiagBasedPositionDetails.DIAG_PROP_INCOMPATIBLE_TYPES_FOR_ITERABLE_FOUND_SYMBOL_INDEX);
-        } else if ("BCE3998".equals(diagnostic.diagnosticInfo().code())) {
+        } else if (DIAGNOSTIC_CODE_3998.equals(diagnostic.diagnosticInfo().code())) {
 
             foundType = context.currentSemanticModel()
                     .flatMap(semanticModel -> semanticModel.typeOf(expressionNode.get()));
@@ -109,7 +113,7 @@ public List<CodeAction> getCodeActions(Diagnostic diagnostic, DiagBasedPositionD
         if (expressionNode.get().kind() == SyntaxKind.BRACED_EXPRESSION) {
             BracedExpressionNode bracedExpressionNode = (BracedExpressionNode) expressionNode.get();
             pos = PositionUtil.toRange(bracedExpressionNode.expression().location().lineRange()).getStart();
-        } else if ("BCE3998".equals(diagnostic.diagnosticInfo().code())) {
+        } else if (DIAGNOSTIC_CODE_3998.equals(diagnostic.diagnosticInfo().code())) {
             // In the case of "BCE3998", we have to consider the position as the position of the initializer 
             // because the diagnostic range is provided for the variable declaration statement instead of the 
             // initializer expression

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/common/utils/CommonUtil.java

@@ -437,9 +437,14 @@ public static List<Token> getQualifiersOfNode(BallerinaCompletionContext context
                 List<Token> qualsAtCursor = getQualifiersAtCursor(context);
                 Set<SyntaxKind> foundQuals = qualifiers.stream().map(Node::kind).collect(Collectors.toSet());
                 context.getNodeAtCursor().leadingInvalidTokens().stream()
-                        .filter(token -> QUALIFIER_KINDS.contains(token.kind())
-                                && !foundQuals.contains(token.kind())).forEach(qualifiers::add);
-                qualifiers.addAll(qualsAtCursor);
+                        .filter(token -> QUALIFIER_KINDS.contains(token.kind()))
+                        .filter(token -> !foundQuals.contains(token.kind()))
+                        .forEach(qualifiers::add);
+                // Avoid duplicating the token at cursor.
+                qualsAtCursor.stream()
+                        .filter(token -> qualifiers.stream()
+                                .noneMatch(qual -> qual.textRange().equals(token.textRange())))
+                        .forEach(qualifiers::add);
                 return qualifiers;
             default:
         }

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/completions/providers/context/ModulePartNodeContext.java

@@ -34,6 +34,7 @@
 import org.ballerinalang.langserver.completions.util.QNameRefCompletionUtil;
 import org.ballerinalang.langserver.completions.util.Snippet;
 import org.ballerinalang.langserver.completions.util.SortingUtil;
+import org.eclipse.lsp4j.Position;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -87,7 +88,7 @@ public List<LSCompletionItem> getCompletions(BallerinaCompletionContext context,
                 currently the qualifier can be isolated/transactional/client.
             */
             completionItems.addAll(this.getCompletionItemsOnQualifiers(node, context));
-            Optional<Token> lastQualifier = CommonUtil.getLastQualifier(context, node);
+            Optional<Token> lastQualifier = getLastQualifier(context, node);
             if (lastQualifier.isPresent() && lastQualifier.get().kind() == SyntaxKind.CONFIGURABLE_KEYWORD) {
                 resolvedContext = ResolvedContext.CONFIGURABLE_QUALIFIER;
             }
@@ -160,6 +161,52 @@ protected List<LSCompletionItem> getCompletionItemsOnQualifiers(Node node, Balle
         return completionItems;
     }
 
+    /**
+     * Check if the cursor is after the qualifiers of the node.
+     *
+     * @param context completion context
+     * @param node    node.
+     * @return True if the cursor is after the qualifiers of the node.
+     */
+    @Override
+    protected boolean onSuggestionsAfterQualifiers(BallerinaCompletionContext context, Node node) {
+        int cursor = context.getCursorPositionInTree();
+        Position cursorPos = context.getCursorPosition();
+        // Get the qualifiers in the same line as the cursor
+        List<Token> qualifiers = CommonUtil.getQualifiersOfNode(context, node)
+                .stream()
+                .filter(qualifier -> qualifier.lineRange().endLine().line() == cursorPos.getLine())
+                .collect(Collectors.toList());
+        if (qualifiers.isEmpty()) {
+            return false;
+        }
+
+        // If cursor is after the last qualifier, no problem
+        Token lastQualifier = qualifiers.get(qualifiers.size() - 1);
+        return lastQualifier.textRange().endOffset() < cursor;
+    }
+
+    /**
+     * Get the last qualifier of the node, but in the same line as the cursor. We have to consider the cursor's
+     * line here due to parser's behavior at the module context.
+     *
+     * @param context Completion context
+     * @param node    Node
+     * @return Optional last qualifier's token
+     */
+    private Optional<Token> getLastQualifier(BallerinaCompletionContext context, Node node) {
+        Position cursorPos = context.getCursorPosition();
+        // Get the qualifiers in the same line as the cursor
+        List<Token> qualifiers = CommonUtil.getQualifiersOfNode(context, node)
+                .stream()
+                .filter(qualifier -> qualifier.lineRange().endLine().line() == cursorPos.getLine())
+                .collect(Collectors.toList());
+        if (qualifiers.isEmpty()) {
+            return Optional.empty();
+        }
+        return Optional.of(qualifiers.get(qualifiers.size() - 1));
+    }
+
     private List<LSCompletionItem> getModulePartContextItems(BallerinaCompletionContext context) {
         NonTerminalNode nodeAtCursor = context.getNodeAtCursor();
         List<LSCompletionItem> completionItems = new ArrayList<>();

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/completions/providers/context/QueryPipelineNodeContext.java

@@ -138,7 +138,7 @@ private boolean onSuggestDirectionKeywords(BallerinaCompletionContext context, Q
                 closestNode = next;
             }
         }
-        if (closestNode.kind() != SyntaxKind.ORDER_BY_CLAUSE) {
+        if (closestNode == null || closestNode.kind() != SyntaxKind.ORDER_BY_CLAUSE) {
             return false;
         }
         SeparatedNodeList<OrderKeyNode> orderKeyNodes = ((OrderByClauseNode) closestNode).orderKey();

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/completions/util/SortingUtil.java

@@ -54,7 +54,6 @@
 import org.ballerinalang.langserver.completions.TypeCompletionItem;
 import org.eclipse.lsp4j.CompletionItemKind;
 
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
@@ -685,17 +684,14 @@ public static void sortCompletionsAfterConfigurableQualifier(CompletionContext c
             return;
         }
 
-        List<String> anyDataSubTypeLabels = Arrays.asList("boolean", "int", "float",
-                "decimal", "string", "xml", "map", "table");
+        TypeSymbol anydataType = context.currentSemanticModel().get().types().ANYDATA;
         completionItems.forEach(lsCompletionItem -> {
             String sortText;
-            if (lsCompletionItem.getCompletionItem().getKind() == CompletionItemKind.Unit &&
-                    lsCompletionItem.getType() == SYMBOL) {
+            if (lsCompletionItem.getType() == SYMBOL) {
                 Optional<Symbol> symbol = ((SymbolCompletionItem) lsCompletionItem).getSymbol();
-                if (symbol.isPresent() && symbol.get() instanceof ModuleSymbol &&
-                        CommonUtil.isLangLib(((ModuleSymbol) symbol.get()).id()) &&
-                        anyDataSubTypeLabels.contains(lsCompletionItem.getCompletionItem().getLabel())
-                ) {
+                if (symbol.isPresent() &&
+                        symbol.get() instanceof TypeSymbol &&
+                        ((TypeSymbol) symbol.get()).subtypeOf(anydataType)) {
                     sortText = SortingUtil.genSortText(1);
                 } else {
                     sortText = SortingUtil.genSortText(3);

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

@@ -52,7 +52,9 @@ public Object[][] dataProvider() {
                 {"add_check_codeaction_wait_action_config1.json"},
                 {"add_check_in_local_var1.json"},
                 {"add_check_in_local_var2.json"},
-                {"add_check_in_module_var.json"}
+                {"add_check_in_module_var.json"},
+                {"add_check_on_field_access_config1.json"},
+                {"add_check_on_field_access_config2.json"}
         };
     }
 

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

@@ -51,6 +51,7 @@ public Object[][] dataProvider() {
                 {"typeGuardVariableCodeAction1.json"},
                 {"typeGuardVariableCodeAction3.json"},
                 {"typeGuardWithTuple1.json"},
+                {"typeGuardErrorType1.json"},
         };
     }
 

language-server/modules/langserver-core/src/test/resources/codeaction/add-check/config/add_check_on_field_access_config1.json

@@ -0,0 +1,54 @@
+{
+  "position": {
+    "line": 7,
+    "character": 20
+  },
+  "source": "add_check_on_field_access_source1.bal",
+  "expected": [
+    {
+      "title": "Add 'check' error",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 7,
+              "character": 17
+            },
+            "end": {
+              "line": 7,
+              "character": 17
+            }
+          },
+          "newText": "("
+        },
+        {
+          "range": {
+            "start": {
+              "line": 7,
+              "character": 23
+            },
+            "end": {
+              "line": 7,
+              "character": 23
+            }
+          },
+          "newText": ")"
+        },
+        {
+          "range": {
+            "start": {
+              "line": 7,
+              "character": 17
+            },
+            "end": {
+              "line": 7,
+              "character": 17
+            }
+          },
+          "newText": "check "
+        }
+      ]
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-check/config/add_check_on_field_access_config2.json

@@ -0,0 +1,28 @@
+{
+  "position": {
+    "line": 7,
+    "character": 20
+  },
+  "source": "add_check_on_field_access_source2.bal",
+  "expected": [
+    {
+      "title": "Add 'check' error",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 7,
+              "character": 18
+            },
+            "end": {
+              "line": 7,
+              "character": 18
+            }
+          },
+          "newText": "check "
+        }
+      ]
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-check/source/add_check_on_field_access_source1.bal

@@ -0,0 +1,9 @@
+type Data record {|
+    json|error field1;
+|};
+
+
+function toString(Data data) returns error? {
+    json|error field1 = data.field1;
+    string str = field1.toString();
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-check/source/add_check_on_field_access_source2.bal

@@ -0,0 +1,9 @@
+type Data record {|
+    json|error field1;
+|};
+
+
+function toString(Data data) returns error? {
+    json|error field1 = data.field1;
+    string str = (field1).toString();
+}