Refactor add elements functionality in datamapping

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

@@ -881,34 +881,57 @@ public String addElement(JsonElement node, String propertyKey, Path filePath, St
         if (source.equals("[]")) {
             return "[" + defaultVal + "]";
         }
-        String fieldsPattern = getFieldsPattern(targetField);
-        if (source.matches("(?s).*" + fieldsPattern + "\\s*:\\s*\\[.*$")) {
-            String[] splits = source.split(".*" + fieldsPattern + "\\s*:\\s*\\[");
-            String lastSplit = splits[1];
-            if (!lastSplit.trim().startsWith("]")) {
-                defaultVal = ", " + defaultVal;
-            }
-            String firstSplit = source.substring(0, source.length() - lastSplit.length());
-            StringBuilder sb = new StringBuilder(firstSplit);
-            int openBraceCount = 0;
-            for (int i = 0; i < lastSplit.length(); i++) {
-                char c = lastSplit.charAt(i);
-                if (c == '[') {
-                    openBraceCount = openBraceCount + 1;
-                } else if (c == ']') {
-                    if (openBraceCount == 0) {
-                        sb.append(defaultVal);
-                        sb.append(lastSplit.substring(i));
-                        break;
-                    } else {
-                        openBraceCount = openBraceCount - 1;
+        VariableDeclarationNode varDeclNode = (VariableDeclarationNode) stNode;
+        if (varDeclNode.initializer().isEmpty()) {
+            return source;
+        }
+        ExpressionNode initializer = varDeclNode.initializer().get();
+        ExpressionNode expr = getArrayExpr(targetField, initializer);
+        if (expr == null || expr.kind() != SyntaxKind.LIST_CONSTRUCTOR) {
+            return source;
+        }
+        ListConstructorExpressionNode listCtrExpr = (ListConstructorExpressionNode) expr;
+        if (!listCtrExpr.expressions().isEmpty()) {
+            defaultVal = ", " + defaultVal;
+        }
+        int pos = listCtrExpr.closeBracket().position() - initializer.position();
+        source = initializer.toSourceCode();
+        return source.substring(0, pos) + defaultVal + source.substring(pos);
+    }
+
+    private ExpressionNode getArrayExpr(String targetField, ExpressionNode expr) {
+        String[] splits = targetField.split("\\.");
+        ExpressionNode currentExpr = expr;
+        for (int i = 1; i < splits.length; i++) {
+            String split = splits[i];
+            if (split.matches("\\d+")) {
+                if (currentExpr.kind() == SyntaxKind.LIST_CONSTRUCTOR) {
+                    ListConstructorExpressionNode listCtrExpr = (ListConstructorExpressionNode) currentExpr;
+                    SeparatedNodeList<Node> expressions = listCtrExpr.expressions();
+                    int size = expressions.size();
+                    int index = Integer.parseInt(split);
+                    if (index >= size) {
+                        return null;
+                    }
+                    currentExpr = (ExpressionNode) expressions.get(index);
+                }
+            } else if (currentExpr.kind() == SyntaxKind.MAPPING_CONSTRUCTOR) {
+                MappingConstructorExpressionNode mappingCtrExpr = (MappingConstructorExpressionNode) currentExpr;
+                for (MappingFieldNode field : mappingCtrExpr.fields()) {
+                    if (field.kind() == SyntaxKind.SPECIFIC_FIELD) {
+                        SpecificFieldNode specificFieldNode = (SpecificFieldNode) field;
+                        if (specificFieldNode.fieldName().toSourceCode().trim().equals(split)) {
+                            Optional<ExpressionNode> optFieldExpr = specificFieldNode.valueExpr();
+                            if (optFieldExpr.isEmpty()) {
+                                return null;
+                            }
+                            currentExpr = optFieldExpr.get();
+                        }
                     }
                 }
-                sb.append(c);
             }
-            return sb.toString();
         }
-        return "";
+        return currentExpr;
     }
 
     private TypeSymbol getTargetType(TypeSymbol typeSymbol, String targetField) {

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/java/io/ballerina/flowmodelgenerator/extension/DataMappingAddElementTest.java

@@ -48,6 +48,9 @@ protected Object[] getConfigsList() {
                 {Path.of("variable5.json")},
                 {Path.of("variable6.json")},
                 {Path.of("variable7.json")},
+                {Path.of("variable8.json")},
+                {Path.of("variable9.json")},
+                {Path.of("variable10.json")},
         };
     }
 
@@ -62,7 +65,7 @@ public void test(Path config) throws IOException {
                         testConfig.diagram(), testConfig.position(), "expression", testConfig.targetField());
         String source = getResponse(request).getAsJsonPrimitive("source").getAsString();
 
-        if (!source.equals(testConfig.output())) {
+        if (!source.replaceAll("\\s+", "").equals(testConfig.output().replaceAll("\\s+", ""))) {
             TestConfig updateConfig = new TestConfig(testConfig.source(), testConfig.description(),
                     testConfig.diagram(), testConfig.propertyKey(), testConfig.position(), testConfig.mappings(),
                     source, testConfig.targetField());

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_add_element/config/variable10.json

@@ -0,0 +1,69 @@
+{
+  "source": "variable5.bal",
+  "description": "Sample diagram node",
+  "diagram": {
+    "id": "31",
+    "metadata": {
+      "label": "Variable",
+      "description": "New variable with type"
+    },
+    "codedata": {
+      "node": "VARIABLE",
+      "lineRange": {
+        "startLine": {
+          "line": 16,
+          "offset": 36
+        },
+        "endLine": {
+          "line": 16,
+          "offset": 36
+        }
+      }
+    },
+    "returning": false,
+    "properties": {
+      "variable": {
+        "metadata": {
+          "label": "Name",
+          "description": "Name of the variable"
+        },
+        "valueType": "IDENTIFIER",
+        "value": "var1",
+        "optional": false,
+        "editable": true,
+        "advanced": false
+      },
+      "type": {
+        "metadata": {
+          "label": "Type",
+          "description": "Type of the variable"
+        },
+        "valueType": "TYPE",
+        "value": "Studentsss",
+        "placeholder": "var",
+        "optional": false,
+        "editable": true,
+        "advanced": false
+      },
+      "expression": {
+        "metadata": {
+          "label": "Expression",
+          "description": "Initialize with value"
+        },
+        "valueType": "EXPRESSION",
+        "value": "{names: [[[\"1\", \"2\", \"\"]], [[\"3\"]]]}",
+        "optional": true,
+        "editable": true,
+        "advanced": false
+      }
+    },
+    "flags": 0
+  },
+  "propertyKey": "expression",
+  "position": {
+    "line": 11,
+    "offset": 12
+  },
+  "output": "{names: [[[\"1\", \"2\", \"\"]], [[\"3\", \"\"]]]}",
+  "targetField": "var1.names.1.0"
+}

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_add_element/config/variable7.json

@@ -62,6 +62,6 @@
     "line": 13,
     "offset": 12
   },
-  "output": "[{username: \"uname\", password: \"pword\", \t\nids: [\"id1\", \"id2\", \"\"]}]",
+  "output": "[\n    {\n        username: \"uname\",\n        password: \"pword\",\n        ids: [\"id1\", \"id2\", \"\"]\n    }\n]",
   "targetField": "credentials.0.ids"
 }

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_add_element/config/variable8.json

@@ -0,0 +1,69 @@
+{
+  "source": "variable5.bal",
+  "description": "Sample diagram node",
+  "diagram": {
+    "id": "31",
+    "metadata": {
+      "label": "Variable",
+      "description": "New variable with type"
+    },
+    "codedata": {
+      "node": "VARIABLE",
+      "lineRange": {
+        "startLine": {
+          "line": 11,
+          "offset": 12
+        },
+        "endLine": {
+          "line": 11,
+          "offset": 12
+        }
+      }
+    },
+    "returning": false,
+    "properties": {
+      "variable": {
+        "metadata": {
+          "label": "Name",
+          "description": "Name of the variable"
+        },
+        "valueType": "IDENTIFIER",
+        "value": "var1",
+        "optional": false,
+        "editable": true,
+        "advanced": false
+      },
+      "type": {
+        "metadata": {
+          "label": "Type",
+          "description": "Type of the variable"
+        },
+        "valueType": "TYPE",
+        "value": "Studentsss",
+        "placeholder": "var",
+        "optional": false,
+        "editable": true,
+        "advanced": false
+      },
+      "expression": {
+        "metadata": {
+          "label": "Expression",
+          "description": "Initialize with value"
+        },
+        "valueType": "EXPRESSION",
+        "value": "{names: [[[\"1\", \"2\", \"\"]], [[\"3\"]]]}",
+        "optional": true,
+        "editable": true,
+        "advanced": false
+      }
+    },
+    "flags": 0
+  },
+  "propertyKey": "expression",
+  "position": {
+    "line": 11,
+    "offset": 12
+  },
+  "output": "{names: [[[\"1\", \"2\", \"\", \"\"]], [[\"3\"]]]}",
+  "targetField": "var1.names.0.0"
+}

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_add_element/config/variable9.json

@@ -0,0 +1,69 @@
+{
+  "source": "variable5.bal",
+  "description": "Sample diagram node",
+  "diagram": {
+    "id": "31",
+    "metadata": {
+      "label": "Variable",
+      "description": "New variable with type"
+    },
+    "codedata": {
+      "node": "VARIABLE",
+      "lineRange": {
+        "startLine": {
+          "line": 16,
+          "offset": 36
+        },
+        "endLine": {
+          "line": 16,
+          "offset": 36
+        }
+      }
+    },
+    "returning": false,
+    "properties": {
+      "variable": {
+        "metadata": {
+          "label": "Name",
+          "description": "Name of the variable"
+        },
+        "valueType": "IDENTIFIER",
+        "value": "var1",
+        "optional": false,
+        "editable": true,
+        "advanced": false
+      },
+      "type": {
+        "metadata": {
+          "label": "Type",
+          "description": "Type of the variable"
+        },
+        "valueType": "TYPE",
+        "value": "Studentsss",
+        "placeholder": "var",
+        "optional": false,
+        "editable": true,
+        "advanced": false
+      },
+      "expression": {
+        "metadata": {
+          "label": "Expression",
+          "description": "Initialize with value"
+        },
+        "valueType": "EXPRESSION",
+        "value": "{names: [[[\"1\", \"2\", \"\"]], [[\"3\"]]]}",
+        "optional": true,
+        "editable": true,
+        "advanced": false
+      }
+    },
+    "flags": 0
+  },
+  "propertyKey": "expression",
+  "position": {
+    "line": 11,
+    "offset": 12
+  },
+  "output": "{names: [[[\"1\", \"2\", \"\"]], [[\"3\"]], []]}",
+  "targetField": "var1.names"
+}

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_add_element/source/variable5.bal

@@ -0,0 +1,18 @@
+import ballerina/http;
+
+type Studentsss record {|
+    string[][][] names;
+|};
+
+const string CONST = "CONST";
+
+service OASServiceType on new http:Listener(9090) {
+
+	resource function get pet() returns int|http:NotFound {
+        do {
+            // Studentsss var1 = {names: [[["1", "2"]], [["3"]]]};
+		} on fail error e {
+			return http:NOT_FOUND;
+		}
+	}
+}