Merge pull request #42310 from nipunayf/fix-31630

Add a new code action to wrap an isolated variable access with `lock` statement

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

@@ -0,0 +1,141 @@
+/*
+ *  Copyright (c) 2024, WSO2 LLC. (https://www.wso2.com)
+ *
+ *  WSO2 LLC. licenses this file to you under the Apache License,
+ *  Version 2.0 (the "License"); you may not use this file except
+ *  in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing,
+ *  software distributed under the License is distributed on an
+ *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ *  KIND, either express or implied. See the License for the
+ *  specific language governing permissions and limitations
+ *  under the License.
+ */
+
+package org.ballerinalang.langserver.codeaction.providers;
+
+import io.ballerina.compiler.syntax.tree.AssignmentStatementNode;
+import io.ballerina.compiler.syntax.tree.Node;
+import io.ballerina.compiler.syntax.tree.StatementNode;
+import io.ballerina.compiler.syntax.tree.SyntaxKind;
+import io.ballerina.tools.diagnostics.Diagnostic;
+import io.ballerina.tools.text.LinePosition;
+import io.ballerina.tools.text.TextDocument;
+import io.ballerina.tools.text.TextRange;
+import org.ballerinalang.annotation.JavaSPIService;
+import org.ballerinalang.langserver.codeaction.CodeActionNodeValidator;
+import org.ballerinalang.langserver.codeaction.CodeActionUtil;
+import org.ballerinalang.langserver.common.constants.CommandConstants;
+import org.ballerinalang.langserver.common.utils.CommonUtil;
+import org.ballerinalang.langserver.common.utils.PositionUtil;
+import org.ballerinalang.langserver.commons.CodeActionContext;
+import org.ballerinalang.langserver.commons.codeaction.spi.DiagBasedPositionDetails;
+import org.ballerinalang.langserver.commons.codeaction.spi.DiagnosticBasedCodeActionProvider;
+import org.eclipse.lsp4j.CodeAction;
+import org.eclipse.lsp4j.CodeActionKind;
+import org.eclipse.lsp4j.Position;
+import org.eclipse.lsp4j.Range;
+import org.eclipse.lsp4j.TextEdit;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+
+/**
+ * A code action to wrap an isolated variable with a lock.
+ *
+ * @since 2201.9.0
+ */
+@JavaSPIService("org.ballerinalang.langserver.commons.codeaction.spi.LSCodeActionProvider")
+public class AddLockCodeAction implements DiagnosticBasedCodeActionProvider {
+
+    private static final String NAME = "Add lock";
+    private static final Set<String> DIAGNOSTIC_CODES = Set.of("BCE3957", "BCE3962");
+
+    @Override
+    public boolean validate(Diagnostic diagnostic, DiagBasedPositionDetails positionDetails,
+                            CodeActionContext context) {
+        return DIAGNOSTIC_CODES.contains(diagnostic.diagnosticInfo().code())
+                && CodeActionNodeValidator.validate(context.nodeAtRange());
+    }
+
+    @Override
+    public List<CodeAction> getCodeActions(Diagnostic diagnostic, DiagBasedPositionDetails positionDetails,
+                                           CodeActionContext context) {
+        // Determine the enclosing statement node of the isolated variable
+        Optional<StatementNode> statementNode = getMatchingStatementNode(positionDetails.matchedNode());
+        if (statementNode.isEmpty()) {
+            return Collections.emptyList();
+        }
+
+        // Check if there are multiple isolated variables within a single statement
+        List<Diagnostic> diagnostics = context.diagnostics(context.filePath());
+        if (diagnostics.size() > 1 && hasMultipleIsolationVars(statementNode.get(), diagnostic, diagnostics)) {
+            return Collections.emptyList();
+        }
+
+        // Generate and return the text edit for the lock statement
+        TextEdit surroundWithLockEditText = getTextEdit(statementNode.get());
+        return Collections.singletonList(CodeActionUtil.createCodeAction(
+                CommandConstants.SURROUND_WITH_LOCK,
+                List.of(surroundWithLockEditText),
+                context.fileUri(),
+                CodeActionKind.QuickFix)
+        );
+    }
+
+    private static Optional<StatementNode> getMatchingStatementNode(Node matchedNode) {
+        Node parentNode = matchedNode.parent();
+        while (parentNode != null && !(parentNode instanceof StatementNode)) {
+            // Lock statement does not support async calls
+            if (parentNode.kind() == SyntaxKind.START_ACTION) {
+                return Optional.empty();
+            }
+            parentNode = parentNode.parent();
+        }
+
+        // Check if the lock statement contains any async calls
+        if (parentNode != null && parentNode.kind() == SyntaxKind.ASSIGNMENT_STATEMENT) {
+            SyntaxKind kind = ((AssignmentStatementNode) parentNode).expression().kind();
+            if (kind == SyntaxKind.RECEIVE_ACTION || kind == SyntaxKind.START_ACTION) {
+                return Optional.empty();
+            }
+        }
+
+        return Optional.ofNullable((StatementNode) parentNode);
+    }
+
+    private static TextEdit getTextEdit(Node node) {
+        TextDocument textDocument = node.syntaxTree().textDocument();
+        TextRange textRange = node.textRangeWithMinutiae();
+        LinePosition startLinePosition = textDocument.linePositionFrom(textRange.startOffset());
+        LinePosition endLinePosition = textDocument.linePositionFrom(textRange.endOffset());
+        Position startPosition = PositionUtil.toPosition(startLinePosition);
+        Position endPosition = PositionUtil.toPosition(endLinePosition);
+
+        String spaces = " ".repeat(node.lineRange().startLine().offset());
+        String statement = node.toSourceCode();
+        String indentedStatement = statement.substring(0, statement.length() - 1).replace("\n", "\n\t") + "\n";
+
+        String editText =
+                spaces + "lock {" + CommonUtil.LINE_SEPARATOR + "\t" + indentedStatement + spaces + "}" + "\n";
+        return new TextEdit(new Range(startPosition, endPosition), editText);
+    }
+
+    private static boolean hasMultipleIsolationVars(StatementNode statementNode, Diagnostic currentDiagnostic,
+                                                    List<Diagnostic> diagnostics) {
+        return diagnostics.stream().anyMatch(diagnostic -> !currentDiagnostic.equals(diagnostic) &&
+                DIAGNOSTIC_CODES.contains(diagnostic.diagnosticInfo().code()) &&
+                PositionUtil.isWithinLineRange(diagnostic.location().lineRange(), statementNode.lineRange()));
+    }
+
+    @Override
+    public String getName() {
+        return NAME;
+    }
+}

language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/common/constants/CommandConstants.java

@@ -158,6 +158,8 @@ public class CommandConstants {
 
     public static final String SURROUND_WITH_DO_ON_FAIL = "Surround with do/on-fail";
 
+    public static final String SURROUND_WITH_LOCK = "Surround with lock";
+
     public static final String CONVERT_MODULE_VAR_TO_LISTENER_DECLARATION =
             "Convert module variable '%s' to listener declaration";
 

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

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (http://wso2.com)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.ballerinalang.langserver.codeaction;
+
+import org.ballerinalang.langserver.commons.workspace.WorkspaceDocumentException;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+/**
+ * Test class to test the add lock code action.
+ *
+ * @since 2201.9.0
+ */
+public class AddLockCodeActionTest extends AbstractCodeActionTest {
+
+    @Test(dataProvider = "codeaction-data-provider")
+    @Override
+    public void test(String config) throws IOException, WorkspaceDocumentException {
+        super.test(config);
+    }
+
+    @DataProvider(name = "codeaction-data-provider")
+    @Override
+    public Object[][] dataProvider() {
+        return new Object[][]{
+                {"add_lock1.json"},
+                {"add_lock2.json"},
+                {"add_lock3.json"},
+                {"add_lock4.json"},
+                {"add_lock5.json"},
+                {"add_lock6.json"},
+                {"add_lock7.json"},
+                {"add_lock8.json"},
+                {"add_lock9.json"},
+                {"add_lock10.json"},
+                {"add_lock11.json"},
+                {"add_lock12.json"},
+                {"add_lock13.json"},
+                {"add_lock14.json"},
+                {"add_lock15.json"},
+                {"add_lock16.json"},
+                {"add_lock17.json"},
+                {"add_lock18.json"},
+                {"add_lock19.json"},
+                {"add_lock20.json"},
+                {"add_lock21.json"},
+                {"add_lock22.json"},
+                {"add_lock23.json"},
+                {"add_lock24.json"},
+                {"add_lock25.json"},
+                {"add_lock26.json"}
+        };
+    }
+
+    @Test(dataProvider = "negative-data-provider")
+    @Override
+    public void negativeTest(String config) throws IOException, WorkspaceDocumentException {
+        super.negativeTest(config);
+    }
+
+    @DataProvider(name = "negative-data-provider")
+    public Object[][] negativeDataProvider() {
+        return new Object[][]{
+                {"add_lock_negative1.json"},
+                {"add_lock_negative2.json"},
+                {"add_lock_negative3.json"},
+                {"add_lock_negative4.json"},
+                {"add_lock_negative5.json"},
+                {"add_lock_negative6.json"},
+                {"add_lock_negative7.json"}
+        };
+    }
+
+    @Override
+    public String getResourceDir() {
+        return "add-lock";
+    }
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-lock/config/add_lock1.json

@@ -0,0 +1,30 @@
+{
+  "position": {
+    "line": 3,
+    "character": 9
+  },
+  "source": "add_lock1.bal",
+  "description": "Wrap in a lock statement for a simple isolated variable",
+  "expected": [
+    {
+      "title": "Surround with lock",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 3,
+              "character": 0
+            },
+            "end": {
+              "line": 4,
+              "character": 0
+            }
+          },
+          "newText": "    lock {\n\t    _ = i + 2;\n    }\n"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-lock/config/add_lock10.json

@@ -0,0 +1,30 @@
+{
+  "position": {
+    "line": 4,
+    "character": 20
+  },
+  "source": "add_lock10.bal",
+  "description": "Wrap in a lock statement for returning a mutable storage",
+  "expected": [
+    {
+      "title": "Surround with lock",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 4,
+              "character": 0
+            },
+            "end": {
+              "line": 5,
+              "character": 0
+            }
+          },
+          "newText": "        lock {\n\t        return self.mp.get(\"a\");\n        }\n"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-lock/config/add_lock11.json

@@ -0,0 +1,30 @@
+{
+  "position": {
+    "line": 4,
+    "character": 20
+  },
+  "source": "add_lock11.bal",
+  "description": "Wrap in a lock statement for returning a mutable storage",
+  "expected": [
+    {
+      "title": "Surround with lock",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 4,
+              "character": 0
+            },
+            "end": {
+              "line": 5,
+              "character": 0
+            }
+          },
+          "newText": "        lock {\n\t        return self.mp.get(\"a\");\n        }\n"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}

language-server/modules/langserver-core/src/test/resources/codeaction/add-lock/config/add_lock12.json

@@ -0,0 +1,30 @@
+{
+  "position": {
+    "line": 4,
+    "character": 20
+  },
+  "source": "add_lock12.bal",
+  "description": "Wrap in a lock statement for returning a mutable storage",
+  "expected": [
+    {
+      "title": "Surround with lock",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 4,
+              "character": 0
+            },
+            "end": {
+              "line": 5,
+              "character": 0
+            }
+          },
+          "newText": "        lock {\n\t        return self.arr[0];\n        }\n"
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}