Fix references in modules

Signed-off-by: Rafael Ravedutti <rafaelravedutti@gmail.com>

src/pairs/ir/loops.py

@@ -76,7 +76,7 @@ class ParticleFor(For):
         return f"ParticleFor<>"
 
     def children(self):
-        return [self.sim.nlocal] + ([] if self.local_only else [self.sim.pbc.npbc])
+        return [self.block, self.sim.nlocal] + ([] if self.local_only else [self.sim.pbc.npbc])
 
 
 class While(ASTNode):

src/pairs/ir/visitor.py

@@ -19,7 +19,14 @@ class Visitor:
         if method is not None:
             method(ast_node)
         else:
-            self.visit_children(ast_node)
+            for b in type(ast_node).__bases__:
+                method = self.get_method(f"visit_{b.__name__}")
+                if method is not None:
+                    method(ast_node)
+                    break
+
+            if method is None:
+                self.visit_children(ast_node)
 
     def visit_children(self, ast_node):
         for c in ast_node.children():

src/pairs/sim/simulation.py

@@ -216,7 +216,9 @@ class Simulation:
             self.kernels
         ])
 
+        self.enter(timestep.block)
         timestep.add(VTKWrite(self, self.vtk_file, timestep.timestep() + 1))
+        self.leave()
 
         body = Block.from_list(self, [
             self.setups,

src/pairs/sim/timestep.py

@@ -28,16 +28,19 @@ class Timestep:
         stmts = item if not isinstance(item, Block) else item.statements()
         stmts_else = None
         ts = self.timestep_loop.iter()
+        self.sim.enter(self.block)
 
         if item_else is not None:
             stmts_else = item_else if not isinstance(item_else, Block) else item_else.statements()
 
         if exec_every > 0:
             self.block.add_statement(
-                Branch(self.sim, BinOp.cmp(ts % exec_every, 0), True if stmts_else is None else False,
+                Branch(self.sim, BinOp.inline(BinOp.cmp(ts % exec_every, 0)), True if stmts_else is None else False,
                 Block(self.sim, stmts), None if stmts_else is None else Block(self.sim, stmts_else)))
         else:
             self.block.add_statement(stmts)
 
+        self.sim.leave()
+
     def as_block(self):
         return Block(self.sim, [self.timestep_loop])