fix: added visitor functions for if statement

1 files changed, 11 insertions, 0 deletions

diff --git a/lib/c.py b/lib/c.py
index 0830dc5..c963b5b 100644
--- a/lib/c.py
+++ b/lib/c.py
@@ -855,6 +855,13 @@ def _(blk: Block, func):
         func(stmt)
 
 @Visit.register
+def _(jmp: If, func):
+    func(jmp.cond)
+    func(jmp.then)
+    if jmp.orelse is not None:
+        func(jmp.orelse)
+
+@Visit.register
 def _(loop: For, func):
     func(loop.init)
     func(loop.cond)
@@ -889,6 +896,10 @@ def _(loop: For, func):
     return For(func(loop.init), func(loop.cond), func(loop.step), func(loop.body))
 
 @Make.register
+def _(jmp: If, func):
+    return If(func(jmp.cond), func(jmp.then), func(loop.orelse) if loop.orelse is not None else None)
+
+@Make.register
 def _(ret: Return, func):
     return Return(func(ret.val))