Support submodule types in the data mapper

flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/SourceBuilder.java

@@ -20,6 +20,8 @@
 
 import io.ballerina.compiler.api.SemanticModel;
 import io.ballerina.compiler.api.symbols.RecordTypeSymbol;
+import io.ballerina.compiler.api.symbols.Symbol;
+import io.ballerina.compiler.api.symbols.TypeDefinitionSymbol;
 import io.ballerina.compiler.api.symbols.TypeDescKind;
 import io.ballerina.compiler.api.symbols.TypeReferenceTypeSymbol;
 import io.ballerina.compiler.api.symbols.TypeSymbol;
@@ -39,6 +41,8 @@
 import io.ballerina.projects.Document;
 import io.ballerina.projects.Module;
 import io.ballerina.projects.ModuleDescriptor;
+import io.ballerina.projects.ModuleId;
+import io.ballerina.projects.Project;
 import io.ballerina.tools.text.LinePosition;
 import io.ballerina.tools.text.LineRange;
 import org.ballerinalang.formatter.core.FormattingTreeModifier;
@@ -58,12 +62,14 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class SourceBuilder {
 
@@ -274,28 +280,78 @@ public SourceBuilder acceptImport(String org, String module, boolean defaultName
     }
 
     public Optional<String> getExpressionBodyText(String typeName, Map<String, String> imports) {
-        PackageUtil.loadProject(workspaceManager, filePath);
+        Project project = PackageUtil.loadProject(workspaceManager, filePath);
         Document document = FileSystemUtils.getDocument(workspaceManager, filePath);
 
+        // Fetch the module ids of the project
+        String packageName = project.currentPackage().packageName().value();
+        Set<String> subModuleIds = new HashSet<>();
+        Map<String, ModuleId> subModuleIdMap = new HashMap<>();
+        project.currentPackage().moduleIds().forEach(moduleId -> {
+            String subModuleId = moduleId.moduleName();
+            if (subModuleId.equals(packageName)) {
+                return;
+            }
+            subModuleIds.add(subModuleId);
+            subModuleIdMap.put(subModuleId, moduleId);
+        });
+
         // Obtain the symbols of the imports
         Map<String, BLangPackage> packageMap = new HashMap<>();
         if (imports != null) {
             imports.values().forEach(moduleId -> {
                 ModuleInfo moduleInfo = ModuleInfo.from(moduleId);
-                PackageUtil.pullModuleAndNotify(lsClientLogger, moduleInfo).ifPresent(pkg ->
-                        packageMap.put(CommonUtils.getDefaultModulePrefix(pkg.packageName().value()),
-                                PackageUtil.getCompilation(pkg).defaultModuleBLangPackage())
-                );
+                if (!subModuleIds.contains(moduleInfo.packageName())) {
+                    PackageUtil.pullModuleAndNotify(lsClientLogger, moduleInfo).ifPresent(pkg ->
+                            packageMap.put(CommonUtils.getDefaultModulePrefix(pkg.packageName().value()),
+                                    PackageUtil.getCompilation(pkg).defaultModuleBLangPackage())
+                    );
+                }
             });
         }
 
         SemanticModel semanticModel = FileSystemUtils.getSemanticModel(workspaceManager, filePath);
         Optional<TypeSymbol> optionalType =
                 BallerinaCompilerApi.getInstance().getType(semanticModel.types(), document, typeName, packageMap);
-        if (optionalType.isEmpty()) {
-            return Optional.empty();
+        TypeSymbol typeSymbol;
+        if (optionalType.isPresent()) {
+            typeSymbol = optionalType.get();
+        } else {
+            // Failover mechanism: check if the type belongs to a local submodule
+            // TODO: The following implementation is a temporary workaround. We need to explore ways to enhance the
+            //  semantic APIs to properly account for submodules as well.
+            String[] typeNameParts = typeName.split(":");
+            if (typeNameParts.length != 2) {
+                return Optional.empty();
+            }
+
+            String fullModuleName = packageName + "." + typeNameParts[0];
+            ModuleId moduleId = subModuleIdMap.get(fullModuleName);
+            if (moduleId == null) {
+                return Optional.empty();
+            }
+
+            SemanticModel localModuleSemanticModel =
+                    project.currentPackage().getCompilation().getSemanticModel(moduleId);
+            Optional<Symbol> moduleSymbol = localModuleSemanticModel.moduleSymbols()
+                    .stream()
+                    .filter(symbol -> symbol.nameEquals(typeNameParts[1]))
+                    .findFirst();
+
+            if (moduleSymbol.isEmpty()) {
+                return Optional.empty();
+            }
+
+            Symbol symbol = moduleSymbol.get();
+            if (symbol instanceof TypeSymbol tSymbol) {
+                typeSymbol = tSymbol;
+            }
+            else if (symbol instanceof TypeDefinitionSymbol typeDefinitionSymbol) {
+                typeSymbol = typeDefinitionSymbol.typeDescriptor();
+            } else {
+                return Optional.empty();
+            }
         }
-        TypeSymbol typeSymbol = optionalType.get();
 
         if (typeSymbol.typeKind() == TypeDescKind.TYPE_REFERENCE) {
             TypeReferenceTypeSymbol typeDefinitionSymbol = (TypeReferenceTypeSymbol) typeSymbol;

flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/DataMapperDefinitionBuilder.java

@@ -168,8 +168,7 @@ public Map<Path, List<TextEdit>> toSource(SourceBuilder sourceBuilder) {
                     .token().skipFormatting().stepOut()
                     .textEdit();
         }
-        return sourceBuilder
-                .build();
+        return sourceBuilder.build();
     }
 
     private static class ParameterSchemaHolder {

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/to_source/config/data_mapper_definition11.json

@@ -0,0 +1,243 @@
+{
+  "source": "nested/data_mappings.bal",
+  "description": "Test basic local module type usage - function returning local module type",
+  "diagram": {
+    "id": "99311",
+    "metadata": {
+      "label": "Data Mapper",
+      "description": "Define a data mapper with local module type"
+    },
+    "codedata": {
+      "node": "DATA_MAPPER_DEFINITION",
+      "isNew": true
+    },
+    "returning": false,
+    "properties": {
+      "functionName": {
+        "metadata": {
+          "label": "Name",
+          "description": "Name of the function"
+        },
+        "valueType": "IDENTIFIER",
+        "valueTypeConstraint": "Global",
+        "value": "createPerson",
+        "optional": false,
+        "editable": false,
+        "advanced": false,
+        "hidden": false
+      },
+      "type": {
+        "metadata": {
+          "label": "Return Type",
+          "description": "Type of the return value"
+        },
+        "valueType": "TYPE",
+        "valueTypeConstraint": "record",
+        "value": "mod:Person",
+        "imports": {
+          "mod": "wso2/nested.mod:0.1.0"
+        },
+        "optional": false,
+        "editable": true,
+        "advanced": false,
+        "hidden": false
+      },
+      "parameters": {
+        "metadata": {
+          "label": "Inputs",
+          "description": "Input variables of the data mapper function"
+        },
+        "valueType": "REPEATABLE_PROPERTY",
+        "valueTypeConstraint": {
+          "metadata": {
+            "label": "Parameter",
+            "description": "Function parameter"
+          },
+          "valueType": "FIXED_PROPERTY",
+          "value": {
+            "type": {
+              "metadata": {
+                "label": "Type",
+                "description": "Type of the parameter"
+              },
+              "valueType": "TYPE",
+              "valueTypeConstraint": "json",
+              "value": "",
+              "optional": false,
+              "editable": true,
+              "advanced": false,
+              "hidden": false
+            },
+            "variable": {
+              "metadata": {
+                "label": "Name",
+                "description": "Name of the parameter"
+              },
+              "valueType": "IDENTIFIER",
+              "value": "",
+              "optional": false,
+              "editable": true,
+              "advanced": false,
+              "hidden": false
+            }
+          },
+          "optional": false,
+          "editable": false,
+          "advanced": false,
+          "hidden": false
+        },
+        "value": {
+          "name": {
+            "metadata": {
+              "label": "Parameter",
+              "description": "Function parameter"
+            },
+            "valueType": "FIXED_PROPERTY",
+            "value": {
+              "type": {
+                "metadata": {
+                  "label": "Type",
+                  "description": "Type of the parameter"
+                },
+                "valueType": "TYPE",
+                "valueTypeConstraint": "json",
+                "value": "string",
+                "optional": false,
+                "editable": true,
+                "advanced": false,
+                "hidden": false
+              },
+              "variable": {
+                "metadata": {
+                  "label": "Name",
+                  "description": "Name of the parameter"
+                },
+                "valueType": "IDENTIFIER",
+                "value": "name",
+                "optional": false,
+                "editable": true,
+                "advanced": false,
+                "hidden": false
+              }
+            },
+            "optional": false,
+            "editable": false,
+            "advanced": false,
+            "hidden": false
+          },
+          "age": {
+            "metadata": {
+              "label": "Parameter",
+              "description": "Function parameter"
+            },
+            "valueType": "FIXED_PROPERTY",
+            "value": {
+              "type": {
+                "metadata": {
+                  "label": "Type",
+                  "description": "Type of the parameter"
+                },
+                "valueType": "TYPE",
+                "valueTypeConstraint": "json",
+                "value": "int",
+                "optional": false,
+                "editable": true,
+                "advanced": false,
+                "hidden": false
+              },
+              "variable": {
+                "metadata": {
+                  "label": "Name",
+                  "description": "Name of the parameter"
+                },
+                "valueType": "IDENTIFIER",
+                "value": "age",
+                "optional": false,
+                "editable": true,
+                "advanced": false,
+                "hidden": false
+              }
+            },
+            "optional": false,
+            "editable": false,
+            "advanced": false,
+            "hidden": false
+          },
+          "email": {
+            "metadata": {
+              "label": "Parameter",
+              "description": "Function parameter"
+            },
+            "valueType": "FIXED_PROPERTY",
+            "value": {
+              "type": {
+                "metadata": {
+                  "label": "Type",
+                  "description": "Type of the parameter"
+                },
+                "valueType": "TYPE",
+                "valueTypeConstraint": "json",
+                "value": "string",
+                "optional": false,
+                "editable": true,
+                "advanced": false,
+                "hidden": false
+              },
+              "variable": {
+                "metadata": {
+                  "label": "Name",
+                  "description": "Name of the parameter"
+                },
+                "valueType": "IDENTIFIER",
+                "value": "email",
+                "optional": false,
+                "editable": true,
+                "advanced": false,
+                "hidden": false
+              }
+            },
+            "optional": false,
+            "editable": false,
+            "advanced": false,
+            "hidden": false
+          }
+        },
+        "optional": false,
+        "editable": false,
+        "advanced": false,
+        "hidden": false
+      }
+    },
+    "flags": 0
+  },
+  "output": {
+    "nested/data_mappings.bal": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 0
+          },
+          "end": {
+            "line": 0,
+            "character": 0
+          }
+        },
+        "newText": "import nested.mod;"
+      },
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 0
+          },
+          "end": {
+            "line": 0,
+            "character": 0
+          }
+        },
+        "newText": "function createPerson(string name, int age, string email) returns mod:Person => {\n    name: \"\",\n    age: 0,\n    email: \"\"\n};"
+      }
+    ]
+  }
+}