Support class defs in enclosed func

ballerina-platform · Mar 5, 2025 · 6e5e203 · 6e5e203
1 parent 11f2e55
commit 6e5e203
Show file tree

Hide file tree

Showing 17 changed files with 231 additions and 7 deletions.
diff --git a/...generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/EnclosedNodeFinder.java b/...generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/EnclosedNodeFinder.java
@@ -18,15 +18,19 @@
 
 package io.ballerina.flowmodelgenerator.core;
 
+import io.ballerina.compiler.syntax.tree.ClassDefinitionNode;
 import io.ballerina.compiler.syntax.tree.FunctionDefinitionNode;
 import io.ballerina.compiler.syntax.tree.ModulePartNode;
 import io.ballerina.compiler.syntax.tree.NonTerminalNode;
+import io.ballerina.compiler.syntax.tree.ServiceDeclarationNode;
 import io.ballerina.projects.Document;
 import io.ballerina.tools.text.LinePosition;
 import io.ballerina.tools.text.LineRange;
 import io.ballerina.tools.text.TextRange;
 import org.ballerinalang.langserver.common.utils.PositionUtil;
 
+import java.util.function.Predicate;
+
 /**
  * Enclosed node finder.
  *
@@ -36,20 +40,31 @@ public class EnclosedNodeFinder {
 
     private final Document document;
     private final LinePosition position;
+    private final boolean findClass;
 
-    public EnclosedNodeFinder(Document document, LinePosition position) {
+    public EnclosedNodeFinder(Document document, LinePosition position, boolean findClass) {
         this.document = document;
         this.position = position;
+        this.findClass = findClass;
     }
 
     public LineRange findEnclosedNode() {
         ModulePartNode modulePartNode = document.syntaxTree().rootNode();
         int positionOffset = PositionUtil.getPositionOffset(PositionUtil.toPosition(position), document.syntaxTree());
         TextRange textRange = TextRange.from(positionOffset, 1);
         NonTerminalNode nonTerminalNode = modulePartNode.findNode(textRange);
-        while (nonTerminalNode != null && !(nonTerminalNode instanceof FunctionDefinitionNode)) {
+
+        Predicate<NonTerminalNode> nodeMatcher;
+        if (findClass) {
+            nodeMatcher = node -> !(node instanceof ClassDefinitionNode) && !(node instanceof ServiceDeclarationNode);
+        } else {
+            nodeMatcher = node -> !(node instanceof FunctionDefinitionNode);
+        }
+
+        while (nonTerminalNode != null && nodeMatcher.test(nonTerminalNode)) {
             nonTerminalNode = nonTerminalNode.parent();
         }
+
         if (nonTerminalNode == null) {
             throw new RuntimeException("No enclosed node found");
         }

diff --git a/...on/src/main/java/io/ballerina/flowmodelgenerator/extension/FlowModelGeneratorService.java b/...on/src/main/java/io/ballerina/flowmodelgenerator/extension/FlowModelGeneratorService.java
@@ -445,7 +445,7 @@ public CompletableFuture<EnclosedFuncDefResponse> getEnclosedFunctionDef(Enclose
                 if (document.isEmpty()) {
                     return response;
                 }
-                EnclosedNodeFinder enclosedNodeFinder = new EnclosedNodeFinder(document.get(), request.position());
+                EnclosedNodeFinder enclosedNodeFinder = new EnclosedNodeFinder(document.get(), request.position(), request.findClass());
                 LineRange enclosedRange = enclosedNodeFinder.findEnclosedNode();
                 response.setFilePath(project.sourceRoot().resolve(enclosedRange.fileName()).toString());
                 response.setStartLine(enclosedRange.startLine());

diff --git a/...c/main/java/io/ballerina/flowmodelgenerator/extension/request/EnclosedFuncDefRequest.java b/...c/main/java/io/ballerina/flowmodelgenerator/extension/request/EnclosedFuncDefRequest.java
@@ -25,8 +25,9 @@
  *
  * @param filePath The file path of the source file.
  * @param position The position of the source.
+ * @param findClass Flag to determine whether to search for class definition instead of function definition.
  *
  * @since 2.0.0
  */
-public record EnclosedFuncDefRequest(String filePath, LinePosition position) {
+public record EnclosedFuncDefRequest(String filePath, LinePosition position, boolean findClass) {
 }
diff --git a/...n/src/test/java/io/ballerina/flowmodelgenerator/extension/GetEnclosedFunctionDefTest.java b/...n/src/test/java/io/ballerina/flowmodelgenerator/extension/GetEnclosedFunctionDefTest.java
@@ -44,7 +44,8 @@ public void test(Path config) throws IOException {
         TestConfig testConfig = gson.fromJson(Files.newBufferedReader(configJsonPath), TestConfig.class);
 
         String filePath = getSourcePath(testConfig.filePath());
-        EnclosedFuncDefRequest request = new EnclosedFuncDefRequest(filePath, testConfig.position());
+        EnclosedFuncDefRequest request =
+                new EnclosedFuncDefRequest(filePath, testConfig.position(), testConfig.findClass());
         JsonObject response = getResponse(request);
         JsonObject acutalJsonObj = testConfig.response();
         acutalJsonObj.add("filePath", new JsonPrimitive(
@@ -54,7 +55,7 @@ public void test(Path config) throws IOException {
                     .replace(sourceDir.toString() + "/", "");
             response.addProperty("filePath", pathToReplace);
             TestConfig updatedConfig = new TestConfig(testConfig.description(), testConfig.filePath(),
-                    testConfig.position(), response);
+                    testConfig.position(), testConfig.findClass(), response);
 //            updateConfig(configJsonPath, updatedConfig);
             Assert.fail(String.format("Failed test: '%s' (%s)", testConfig.description(), configJsonPath));
         }
@@ -81,8 +82,11 @@ protected String getApiName() {
      * @param description The description of the test.
      * @param filePath    The path to the source file that contains the nodes to be deleted.
      * @param position    The position of the node to be deleted.
+     * @param findClass   Flag to determine whether to search for class definition instead of function definition.
      * @param response    The expected response.
      */
-    private record TestConfig(String description, String filePath, LinePosition position, JsonObject response) {
+    private record TestConfig(String description, String filePath, LinePosition position, boolean findClass,
+                              JsonObject response) {
+
     }
 }
diff --git a/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_1.json b/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_1.json
@@ -5,6 +5,7 @@
     "line": 3,
     "offset": 8
   },
+  "findClass": false,
   "response": {
     "filePath" : "proj_1/functions.bal",
     "startLine" : {

diff --git a/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_2.json b/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_2.json
@@ -5,6 +5,7 @@
     "line": 9,
     "offset": 0
   },
+  "findClass": false,
   "response": {
     "filePath": "proj_1/functions.bal",
     "startLine": {

diff --git a/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_3.json b/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_3.json
@@ -5,6 +5,7 @@
     "line": 15,
     "offset": 4
   },
+  "findClass": false,
   "response": {
     "filePath": "proj_1/functions.bal",
     "startLine": {

diff --git a/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_4.json b/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_4.json
@@ -5,6 +5,7 @@
     "line": 5,
     "offset": 24
   },
+  "findClass": false,
   "response": {
     "filePath": "proj_1/services.bal",
     "startLine": {

diff --git a/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_5.json b/...r-ls-extension/src/test/resources/get_enclosed_func_def/config/current_package_ref_5.json
@@ -5,6 +5,7 @@
     "line": 17,
     "offset": 16
   },
+  "findClass": false,
   "response": {
     "filePath": "proj_1/services.bal",
     "startLine": {

diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service1.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service1.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/graphql_service.bal",
+  "position": {
+    "line": 17,
+    "offset": 25
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/graphql_service.bal",
+    "startLine": {
+      "line": 12,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 30,
+      "offset": 1
+    }
+  }
+}
diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service2.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service2.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/graphql_service.bal",
+  "position": {
+    "line": 28,
+    "offset": 18
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/graphql_service.bal",
+    "startLine": {
+      "line": 12,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 30,
+      "offset": 1
+    }
+  }
+}
diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service3.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service3.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/graphql_service.bal",
+  "position": {
+    "line": 14,
+    "offset": 26
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/graphql_service.bal",
+    "startLine": {
+      "line": 12,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 30,
+      "offset": 1
+    }
+  }
+}
diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service4.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service4.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/graphql_service.bal",
+  "position": {
+    "line": 43,
+    "offset": 22
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/graphql_service.bal",
+    "startLine": {
+      "line": 38,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 58,
+      "offset": 1
+    }
+  }
+}
diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service5.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service5.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/graphql_service.bal",
+  "position": {
+    "line": 56,
+    "offset": 35
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/graphql_service.bal",
+    "startLine": {
+      "line": 38,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 58,
+      "offset": 1
+    }
+  }
+}
diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service6.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service6.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/services.bal",
+  "position": {
+    "line": 5,
+    "offset": 24
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/services.bal",
+    "startLine": {
+      "line": 0,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 21,
+      "offset": 1
+    }
+  }
+}
diff --git a/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service7.json b/...odel-generator-ls-extension/src/test/resources/get_enclosed_func_def/config/service7.json
@@ -0,0 +1,20 @@
+{
+  "description": "",
+  "filePath": "proj_1/services.bal",
+  "position": {
+    "line": 17,
+    "offset": 31
+  },
+  "findClass": true,
+  "response": {
+    "filePath": "proj_1/services.bal",
+    "startLine": {
+      "line": 0,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 21,
+      "offset": 1
+    }
+  }
+}
diff --git a/...r-ls-extension/src/test/resources/get_enclosed_func_def/source/proj_1/graphql_service.bal b/...r-ls-extension/src/test/resources/get_enclosed_func_def/source/proj_1/graphql_service.bal
@@ -0,0 +1,59 @@
+import ballerina/graphql;
+
+public type AbstractProfile record {|
+    string name;
+    int age;
+    Address address;
+|};
+
+type UserName record {|
+    string value;
+|};
+
+service class UserProfile {
+    private final string name;
+    private final int age;
+
+    function init(string name, int age) {
+        self.name = name;
+        self.age = age;
+    }
+
+    resource function get name() returns string {
+        return self.name;
+    }
+    resource function get age() returns int {
+        return self.age;
+    }
+    resource function get isAdult() returns boolean {
+        return self.age > 21;
+    }
+}
+
+public type Address record {|
+    string number;
+    string street;
+    string city;
+|};
+
+service /graphql on new graphql:Listener(9090) {
+    resource function get profile21() returns AbstractProfile {
+        return {
+            name: "Walter White",
+            age: 51,
+            address: {
+                number: "308",
+                street: "Negra Arroyo Lane",
+                city: "Albuquerque"
+            }
+        };
+    }
+
+    resource function get profile() returns UserProfile {
+        return new ("Walter White", 51);
+    }
+
+    remote function updateName(UserName userName) returns UserProfile {
+        return new (userName.value, 51);
+    }
+}