Draw the connection usages for service classes

architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/CodeAnalyzer.java

@@ -38,6 +38,7 @@
 import io.ballerina.compiler.syntax.tree.AssignmentStatementNode;
 import io.ballerina.compiler.syntax.tree.BlockStatementNode;
 import io.ballerina.compiler.syntax.tree.CheckExpressionNode;
+import io.ballerina.compiler.syntax.tree.ClassDefinitionNode;
 import io.ballerina.compiler.syntax.tree.ClientResourceAccessActionNode;
 import io.ballerina.compiler.syntax.tree.CompoundAssignmentStatementNode;
 import io.ballerina.compiler.syntax.tree.DoStatementNode;
@@ -110,20 +111,40 @@ public class CodeAnalyzer extends NodeVisitor {
     private IntermediateModel.ServiceModel currentServiceModel;
     private final Path rootPath;
     private final ConnectionFinder connectionFinder;
+    private IntermediateModel.ServiceClassModel currentServiceClass;
+    private String serviceClassName;
 
     public CodeAnalyzer(SemanticModel semanticModel, IntermediateModel intermediateModel, Path rootPath,
                         ConnectionFinder connectionFinder) {
         this.semanticModel = semanticModel;
         this.intermediateModel = intermediateModel;
         this.rootPath = rootPath;
         this.connectionFinder = connectionFinder;
+        this.currentFunctionModel = null;
+        this.currentServiceModel = null;
+        this.currentServiceClass = null;
     }
 
     @Override
     public void visit(ModulePartNode modulePartNode) {
         modulePartNode.members().forEach(member -> member.accept(this));
     }
 
+    @Override
+    public void visit(ClassDefinitionNode classDefinitionNode) {
+        classDefinitionNode.classTypeQualifiers().stream()
+                .filter(qualifier -> qualifier.kind() == SyntaxKind.SERVICE_KEYWORD)
+                .findAny()
+                .ifPresent(qualifier -> {
+                    serviceClassName = classDefinitionNode.className().text();
+                    currentServiceClass = new IntermediateModel.ServiceClassModel(serviceClassName);
+                    intermediateModel.serviceClassModelMap.put(serviceClassName, currentServiceClass);
+                });
+        classDefinitionNode.members().forEach(member -> member.accept(this));
+        serviceClassName = null;
+        currentServiceClass = null;
+    }
+
     @Override
     public void visit(ServiceDeclarationNode serviceDeclarationNode) {
         Optional<Symbol> serviceSymbol = this.semanticModel.symbol(serviceDeclarationNode);
@@ -217,6 +238,9 @@ public void visit(FunctionDefinitionNode functionDefinitionNode) {
             }
         }
         functionDefinitionNode.functionBody().accept(this);
+        if (currentServiceClass != null) {
+            currentServiceClass.functionModels.add(this.currentFunctionModel);
+        }
         this.currentFunctionModel = null;
     }
 
@@ -327,11 +351,21 @@ public void visit(ImplicitNewExpressionNode implicitNewExpressionNode) {
         implicitNewExpressionNode.parenthesizedArgList()
                 .ifPresent(parenthesizedArgList -> parenthesizedArgList.arguments()
                         .forEach(expr -> expr.accept(this)));
+        if (currentFunctionModel != null) {
+            semanticModel.symbol(implicitNewExpressionNode).ifPresent(symbol -> {
+                if (symbol instanceof ClassSymbol classSymbol) {
+                    classSymbol.getName().ifPresent(name -> currentFunctionModel.usedClasses.add(name));
+                }
+            });
+        }
     }
 
     @Override
     public void visit(ExplicitNewExpressionNode explicitNewExpressionNode) {
         explicitNewExpressionNode.parenthesizedArgList().arguments().forEach(expr -> expr.accept(this));
+        if (currentFunctionModel != null) {
+            currentFunctionModel.usedClasses.add(explicitNewExpressionNode.typeDescriptor().toSourceCode().trim());
+        }
     }
 
     @Override

architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/DesignModelGenerator.java

@@ -213,8 +213,23 @@ private void buildConnectionGraph(IntermediateModel intermediateModel,
                                       IntermediateModel.FunctionModel functionModel,
                                       IntermediateModel.ServiceModel serviceModel) {
         Set<String> connections = new HashSet<>();
+        Set<String> dependentServiceClasses = new HashSet<>();
         if (!functionModel.visited && !functionModel.analyzed) {
             functionModel.visited = true;
+            functionModel.usedClasses.forEach(usedClass -> {
+                IntermediateModel.ServiceClassModel serviceClassModel = intermediateModel.serviceClassModelMap
+                        .get(usedClass);
+                if (serviceClassModel != null) {
+                    serviceClassModel.functionModels.forEach(serviceClassFunctionModel -> {
+                        if (!serviceClassFunctionModel.analyzed) {
+                            buildConnectionGraph(intermediateModel, serviceClassFunctionModel, serviceModel);
+                        }
+                        connections.addAll(serviceClassFunctionModel.allDependentConnections);
+                        connections.addAll(serviceClassFunctionModel.connections);
+                        dependentServiceClasses.addAll(serviceClassFunctionModel.usedClasses);
+                    });
+                }
+            });
             functionModel.dependentFuncs.forEach(dependentFunc -> {
                 IntermediateModel.FunctionModel dependentFunctionModel = intermediateModel.functionModelMap
                         .get(dependentFunc);
@@ -226,6 +241,7 @@ private void buildConnectionGraph(IntermediateModel intermediateModel,
                 }
                 connections.addAll(dependentFunctionModel.allDependentConnections);
                 connections.addAll(dependentFunctionModel.connections);
+                dependentServiceClasses.addAll(dependentFunctionModel.usedClasses);
             });
 
             functionModel.dependentObjFuncs.forEach(dependentObjFunc -> {
@@ -242,10 +258,12 @@ private void buildConnectionGraph(IntermediateModel intermediateModel,
                 }
                 connections.addAll(dependentFunctionModel.allDependentConnections);
                 connections.addAll(dependentFunctionModel.connections);
+                dependentServiceClasses.addAll(dependentFunctionModel.usedClasses);
             });
         }
         functionModel.visited = true;
         functionModel.allDependentConnections.addAll(functionModel.connections);
+        functionModel.usedClasses.addAll(dependentServiceClasses);
         // Also add transitive dependent connections
         for (String connectionUuid : functionModel.connections) {
             Connection connection = intermediateModel.uuidToConnectionMap.get(connectionUuid);

architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/IntermediateModel.java

@@ -41,13 +41,15 @@ public class IntermediateModel {
     protected final Map<String, Listener> listeners;
     protected final Map<String, Connection> connectionMap;
     protected final Map<String, Connection> uuidToConnectionMap;
+    protected final Map<String, ServiceClassModel> serviceClassModelMap;
 
     public IntermediateModel() {
         this.functionModelMap = new HashMap<>();
         this.serviceModelMap = new HashMap<>();
         this.listeners = new HashMap<>();
         this.connectionMap = new HashMap<>();
         this.uuidToConnectionMap = new HashMap<>();
+        this.serviceClassModelMap = new HashMap<>();
     }
 
     public static class ServiceModel {
@@ -83,6 +85,7 @@ public static class FunctionModel {
         protected String path;
         protected String displayName;
         protected final Set<String> connections = new HashSet<>();
+        protected final Set<String> usedClasses = new HashSet<>();
 
         public FunctionModel(String name) {
             this.name = name;
@@ -92,4 +95,14 @@ public FunctionModel(String name) {
             this.allDependentConnections = new HashSet<>();
         }
     }
+
+    public static class ServiceClassModel {
+        protected final String name;
+        protected final Set<FunctionModel> functionModels;
+
+        public ServiceClassModel(String name) {
+            this.name = name;
+            this.functionModels = new HashSet<>();
+        }
+    }
 }

architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/config/connections_for_service_classes.json

@@ -0,0 +1,112 @@
+{
+  "description": "Test connections for service classes",
+  "projectPath": "project_9",
+  "output": {
+    "designModel": {
+      "connections": [
+        {
+          "symbol": "httpClient",
+          "location": {
+            "filePath": "/Users/lakshanweerasinghe/Documents/repos/ballerina-platform/ballerina-language-server/architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/source/project_9/main.bal",
+            "startLine": {
+              "line": 3,
+              "offset": 0
+            },
+            "endLine": {
+              "line": 3,
+              "offset": 67
+            }
+          },
+          "scope": "GLOBAL",
+          "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_http_2.15.4.png",
+          "dependentFunctions": [],
+          "dependentConnection": [],
+          "kind": "Connection",
+          "uuid": "c0a28d82-a2ad-9271-10fc-b0af24197b0d",
+          "enableFlowModel": true,
+          "sortText": "main.bal3"
+        }
+      ],
+      "listeners": [
+        {
+          "symbol": "ANON",
+          "location": {
+            "filePath": "/Users/lakshanweerasinghe/Documents/repos/ballerina-platform/ballerina-language-server/architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/source/project_9/main.bal",
+            "startLine": {
+              "line": 5,
+              "offset": 0
+            },
+            "endLine": {
+              "line": 11,
+              "offset": 1
+            }
+          },
+          "attachedServices": [
+            "7d7ac521-9ef6-0e29-694b-ecdd92287e37"
+          ],
+          "kind": "ANON",
+          "type": "websocket:Listener",
+          "args": [
+            {
+              "key": "'listener",
+              "value": "9090"
+            }
+          ],
+          "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_websocket_2.15.0.png",
+          "uuid": "88c93300-d438-5c61-5125-c885fa3d596f",
+          "enableFlowModel": false,
+          "sortText": "main.bal5"
+        }
+      ],
+      "services": [
+        {
+          "location": {
+            "filePath": "/Users/lakshanweerasinghe/Documents/repos/ballerina-platform/ballerina-language-server/architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/source/project_9/main.bal",
+            "startLine": {
+              "line": 5,
+              "offset": 0
+            },
+            "endLine": {
+              "line": 11,
+              "offset": 1
+            }
+          },
+          "attachedListeners": [
+            "88c93300-d438-5c61-5125-c885fa3d596f"
+          ],
+          "connections": [
+            "c0a28d82-a2ad-9271-10fc-b0af24197b0d"
+          ],
+          "functions": [],
+          "remoteFunctions": [],
+          "resourceFunctions": [
+            {
+              "accessor": "get",
+              "path": ".",
+              "location": {
+                "filePath": "/Users/lakshanweerasinghe/Documents/repos/ballerina-platform/ballerina-language-server/architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/source/project_9/main.bal",
+                "startLine": {
+                  "line": 7,
+                  "offset": 4
+                },
+                "endLine": {
+                  "line": 10,
+                  "offset": 5
+                }
+              },
+              "connections": [
+                "c0a28d82-a2ad-9271-10fc-b0af24197b0d"
+              ]
+            }
+          ],
+          "absolutePath": "/chat ",
+          "type": "websocket:Service",
+          "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_websocket_2.15.0.png",
+          "uuid": "7d7ac521-9ef6-0e29-694b-ecdd92287e37",
+          "enableFlowModel": true,
+          "sortText": "main.bal5"
+        }
+      ]
+    }
+  }
+}

architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/source/project_9/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "wso2"
+name = "service"
+version = "0.1.0"

architecture-model-generator/modules/architecture-model-generator-ls-extension/src/test/resources/get_design_model/source/project_9/main.bal

@@ -0,0 +1,24 @@
+import ballerina/http;
+import ballerina/websocket;
+
+final http:Client httpClient = check new ("http://localhost:8080");
+
+service /chat on new websocket:Listener(9090) {
+
+    resource function get .() returns websocket:Service {
+        // Accept the WebSocket upgrade by returning a `websocket:Service`.
+        return new ChatService();
+    }
+}
+
+service class ChatService {
+    *websocket:Service;
+
+    // This `remote method` is triggered when a new message is received
+    // from a client. It accepts `anydata` as the function argument. The received data
+    // will be converted to the data type stated as the function argument.
+    remote function onMessage(websocket:Caller caller, string chatMessage) returns error? {
+        json _ = check httpClient->post("/chat", chatMessage);
+        check caller->writeMessage("Hello!, How are you?");
+    }
+}