Add code block to the worker node

worker-model-generator/modules/model-generator-core/src/main/java/io.ballerina.workermodelgenerator.core/FlowBuilder.java

@@ -24,6 +24,7 @@
 import io.ballerina.compiler.syntax.tree.MetadataNode;
 import io.ballerina.compiler.syntax.tree.NamedWorkerDeclarationNode;
 import io.ballerina.compiler.syntax.tree.NodeVisitor;
+import io.ballerina.compiler.syntax.tree.StatementNode;
 import io.ballerina.workermodelgenerator.core.model.Flow;
 import io.ballerina.workermodelgenerator.core.model.FlowJsonBuilder;
 import io.ballerina.workermodelgenerator.core.model.WorkerNode;
@@ -82,7 +83,15 @@ public void visit(NamedWorkerDeclarationNode namedWorkerDeclarationNode) {
 
         // Analyze the body of the worker
         BlockStatementNode blockStatementNode = namedWorkerDeclarationNode.workerBody();
-        blockStatementNode.statements().forEach(statement -> statement.accept(nodeBuilder));
+        StringBuilder codeBlock = new StringBuilder();
+        for (StatementNode statement : blockStatementNode.statements()) {
+            statement.accept(nodeBuilder);
+            if (!nodeBuilder.hasProcessed()) {
+                codeBlock.append(statement.toSourceCode());
+            }
+            nodeBuilder.resetProcessFlag();
+        }
+        nodeBuilder.setCodeBlock(codeBlock.toString());
         addNode(nodeBuilder.build());
     }
 

worker-model-generator/modules/model-generator-core/src/main/java/io.ballerina.workermodelgenerator.core/NodeBuilder.java

@@ -45,13 +45,16 @@ class NodeBuilder extends NodeVisitor implements WorkerNodeJsonBuilder {
     private CanvasPosition canvasPosition;
     private final List<InputPort> inputPorts;
     private final List<OutputPort> outputPorts;
+    private String codeBlock;
 
     // State variables
     private String toWorker;
     private String fromWorker;
     private TypeDescKind type;
     private String name;
     private int portId;
+    private boolean capturedFromWorker;
+    private boolean hasProcessed;
 
     public NodeBuilder(SemanticModel semanticModel) {
         this.inputPorts = new ArrayList<>();
@@ -77,6 +80,7 @@ public void visit(ReceiveActionNode receiveActionNode) {
         if (receiverWorker.kind() == SyntaxKind.SIMPLE_NAME_REFERENCE) {
             this.fromWorker = ((SimpleNameReferenceNode) receiverWorker).name().text();
         }
+        this.capturedFromWorker = true;
     }
 
     @Override
@@ -94,6 +98,7 @@ private void analyzeSendAction(SimpleNameReferenceNode receiverNode, ExpressionN
         Optional<TypeSymbol> typeSymbol = this.semanticModel.typeOf(expressionNode);
         this.type = typeSymbol.isPresent() ? typeSymbol.get().typeKind() : TypeDescKind.NONE;
         this.addOutputPort(String.valueOf(++this.portId), this.type, this.toWorker);
+        this.hasProcessed = true;
     }
 
     @Override
@@ -105,6 +110,10 @@ public void visit(VariableDeclarationNode variableDeclarationNode) {
         }
         initializer.get().accept(this);
 
+        if (!this.capturedFromWorker) {
+            return;
+        }
+
         // Find the parameter name
         TypedBindingPatternNode typedBindingPatternNode = variableDeclarationNode.typedBindingPattern();
         typedBindingPatternNode.bindingPattern().accept(this);
@@ -115,13 +124,23 @@ public void visit(VariableDeclarationNode variableDeclarationNode) {
                 TypeDescKind.NONE;
 
         this.addInputPort(String.valueOf(++this.portId), this.type, this.name, this.fromWorker);
+        this.capturedFromWorker = false;
+        this.hasProcessed = true;
     }
 
     @Override
     public void visit(CaptureBindingPatternNode captureBindingPatternNode) {
         this.name = captureBindingPatternNode.variableName().text();
     }
 
+    public boolean hasProcessed() {
+        return this.hasProcessed;
+    }
+
+    public void resetProcessFlag() {
+        this.hasProcessed = false;
+    }
+
     @Override
     public void setName(String id) {
         this.id = id;
@@ -152,8 +171,13 @@ public void setCanvasPosition(int x, int y) {
         this.canvasPosition = new CanvasPosition(x, y);
     }
 
+    @Override
+    public void setCodeBlock(String codeBlock) {
+        this.codeBlock = codeBlock;
+    }
+
     @Override
     public WorkerNode build() {
-        return new WorkerNode(id, templateId, codeLocation, canvasPosition, inputPorts, outputPorts);
+        return new WorkerNode(id, templateId, codeLocation, canvasPosition, inputPorts, outputPorts, codeBlock);
     }
 }

worker-model-generator/modules/model-generator-core/src/main/java/io.ballerina.workermodelgenerator.core/model/WorkerNode.java

@@ -29,9 +29,10 @@
  * @param canvasPosition position of the node in the canvas
  * @param inputPorts     input ports of the node
  * @param outputPorts    output ports of the node
+ * @param codeBlock      editable code block of the node
  * @since 2201.9.0
  */
-public record WorkerNode(String name, String templateId, CodeLocation codeLocation,
-                         CanvasPosition canvasPosition, List<InputPort> inputPorts, List<OutputPort> outputPorts) {
+public record WorkerNode(String name, String templateId, CodeLocation codeLocation, CanvasPosition canvasPosition,
+                         List<InputPort> inputPorts, List<OutputPort> outputPorts, String codeBlock) {
 
 }

worker-model-generator/modules/model-generator-core/src/main/java/io.ballerina.workermodelgenerator.core/model/WorkerNodeJsonBuilder.java

@@ -77,6 +77,13 @@ public interface WorkerNodeJsonBuilder {
      */
     void addOutputPort(String id, TypeDescKind type, String receiver);
 
+    /**
+     * Sets the code block of the node.
+     *
+     * @param codeBlock code block of the node
+     */
+    void setCodeBlock(String codeBlock);
+
     /**
      * Builds the node.
      *

worker-model-generator/modules/model-generator-ls-extension/src/test/resources/config/code_block.json

@@ -0,0 +1,80 @@
+{
+  "start": {
+    "line": 1,
+    "offset": 0
+  },
+  "end": {
+    "line": 30,
+    "offset": 1
+  },
+  "source": "code_block.bal",
+  "description": "Tests a flow with a code block",
+  "flow": {
+    "id": "2",
+    "name": "main/function",
+    "fileName": "code_block.bal",
+    "nodes": [
+      {
+        "name": "A",
+        "templateId": "block",
+        "codeLocation": {
+          "start": {
+            "line": 7,
+            "offset": 4
+          },
+          "end": {
+            "line": 15,
+            "offset": 5
+          }
+        },
+        "canvasPosition": {
+          "x": 11,
+          "y": 32
+        },
+        "inputPorts": [],
+        "outputPorts": [
+          {
+            "id": "1",
+            "type": "INT",
+            "receiver": "B"
+          }
+        ],
+        "codeBlock": ""
+      },
+      {
+        "name": "B",
+        "templateId": "block",
+        "codeLocation": {
+          "start": {
+            "line": 17,
+            "offset": 4
+          },
+          "end": {
+            "line": 28,
+            "offset": 5
+          }
+        },
+        "canvasPosition": {
+          "x": 32,
+          "y": 63
+        },
+        "inputPorts": [
+          {
+            "id": "1",
+            "type": "INT",
+            "name": "x",
+            "sender": "A"
+          }
+        ],
+        "outputPorts": [
+          {
+            "id": "2",
+            "type": "INT",
+            "receiver": "function"
+          }
+        ],
+        "codeBlock": "        int a = 32 + x;\n        int b = a % 12;\n"
+      }
+    ]
+  }
+}

worker-model-generator/modules/model-generator-ls-extension/src/test/resources/config/simple_flow.json

@@ -43,7 +43,8 @@
             "type": "INT",
             "receiver": "C"
           }
-        ]
+        ],
+        "codeBlock": ""
       },
       {
         "name": "B",
@@ -76,7 +77,8 @@
             "type": "INT",
             "receiver": "function"
           }
-        ]
+        ],
+        "codeBlock": ""
       },
       {
         "name": "C",
@@ -109,7 +111,8 @@
             "type": "INT",
             "receiver": "function"
           }
-        ]
+        ],
+        "codeBlock": ""
       }
     ]
   }

worker-model-generator/modules/model-generator-ls-extension/src/test/resources/source/code_block.bal

@@ -0,0 +1,32 @@
+
+@display {
+    label: "Flow",
+    id: "2",
+    name: "main/function"
+}
+public function main() {
+    @display {
+        label: "Node",
+        templateId: "block",
+        xCord: 11,
+        yCord: 32
+    }
+    worker A {
+        12 -> B;
+    }
+
+    @display {
+        label: "Node",
+        templateId: "block",
+        xCord: 32,
+        yCord: 63
+    }
+    worker B {
+        int x = <- A;
+        int a = 32 + x;
+        int b = a % 12;
+        b -> function;
+    }
+
+    int result = <- B;
+}