From fc712f94a7ed8554d8d44f4965be367354a7e670 Mon Sep 17 00:00:00 2001
From: L0P0P <grassoemanuele@live.com>
Date: Wed, 26 Jun 2024 10:06:12 +0200
Subject: Using child

---
 src/ast/Python3VisitorImpl.java | 40 ++++++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 18 deletions(-)

(limited to 'src/ast/Python3VisitorImpl.java')

diff --git a/src/ast/Python3VisitorImpl.java b/src/ast/Python3VisitorImpl.java
index 604c8d2..f5b369d 100644
--- a/src/ast/Python3VisitorImpl.java
+++ b/src/ast/Python3VisitorImpl.java
@@ -21,17 +21,19 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ```
      */
     public Node visitRoot(RootContext ctx) {
-        ArrayList<Node> stmts = new ArrayList<Node>();
-        ArrayList<Node> compStmts = new ArrayList<Node>();
-
-        for (Simple_stmtsContext stm : ctx.simple_stmts()) {
-            stmts.add(visit(stm));
-        }
-        for (Compound_stmtContext stm : ctx.compound_stmt()) {
-            compStmts.add(visit(stm));
+        ArrayList<Node> childs = new ArrayList<Node>();
+        
+        for (int i = 0; i < ctx.getChildCount(); i++){
+            var child = ctx.getChild(i);
+            
+            if (child instanceof Simple_stmtsContext) {
+                childs.add(visit((Simple_stmtsContext) child));
+            } else if (child instanceof Compound_stmtContext) {
+                childs.add(visit((Compound_stmtContext) child));
+            }
         }
 
-        return new RootNode(stmts, compStmts);
+        return new RootNode(childs);
     }
 
     /**
@@ -331,17 +333,19 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ```
      */
     public Node visitBlock(BlockContext ctx) {
-        ArrayList<Node> stmts = new ArrayList<Node>();
-        ArrayList<Node> compStmts = new ArrayList<Node>();
-
-        for (Simple_stmtsContext s : ctx.simple_stmts()) {
-            stmts.add(visit(s));
-        }
-        for (Compound_stmtContext s : ctx.compound_stmt()) {
-            compStmts.add(visit(s));
+        ArrayList<Node> childs = new ArrayList<Node>();
+
+        for (int i = 0; i < ctx.getChildCount(); i++){
+            var child = ctx.getChild(i);
+            
+            if (child instanceof Simple_stmtsContext) {
+                childs.add(visit((Simple_stmtsContext) child));
+            } else if (child instanceof Compound_stmtContext) {
+                childs.add(visit((Compound_stmtContext) child));
+            }
         }
 
-        return new BlockNode(stmts, compStmts);
+        return new BlockNode(childs);
     }
 
     /**
-- 
cgit v1.2.3-18-g5258