refactor: made condition of append a dynamic function as well

1 files changed, 28 insertions, 29 deletions

diff --git a/lib/c.py b/lib/c.py
index b67ea98..25c97b2 100644
--- a/lib/c.py
+++ b/lib/c.py
@@ -730,64 +730,63 @@ def _(x: StmtExpr, func):
 # expression functions
 
 @singledispatch
-def Filter(x: Expr, kind: type, results: List[Expr]):
+def Filter(x: Expr, cond, results: List[Expr]):
     raise TypeError(f"{type(s)} not supported by Filter operation")
 
 @Filter.register
-def _(x: Empty, kind: type, results: List[Expr]):
-    if isinstance(x, kind):
+def _(x: Empty, cond, results: List[Expr]):
+    if cond(x):
         results.append(x)
 
 @Filter.register(Ident)
 @Filter.register(Literal)
-def _(x, kind:type, results: List[Expr]):
-    if isinstance(x, kind):
+def _(x, cond, results: List[Expr]):
+    if cond(x):
         results.append(x)
-    return
 
 @Filter.register
-def _(x: UnaryOp, kind:type, results: List[Expr]):
-    if isinstance(x, kind):
+def _(x: UnaryOp, cond, results: List[Expr]):
+    if cond(x):
         results.append(x)
-    Filter(op.x, kind, results)
+    Filter(op.x, cond, results)
 
 @Filter.register
-def _(s: Assign, kind:type, results: List[Expr]):
-    if isinstance(s, kind):
+def _(s: Assign, cond, results: List[Expr]):
+    if cond(s):
         results.append(s)
-    Filter(s.lhs, kind, results)
-    Filter(s.rhs, kind, results)
+    Filter(s.lhs, cond, results)
+    Filter(s.rhs, cond, results)
 
 @Filter.register
-def _(v: Deref, kind:type, results: List[Expr]):
-    if isinstance(v, kind):
+def _(v: Deref, cond, results: List[Expr]):
+    if cond(v):
         results.append(v)
-    Filter(v.x, kind, results)
+    Filter(v.x, cond, results)
 
 @Filter.register
-def _(i: Index, kind: type, results: List[Expr]):
-    if isinstance(i, kind):
+def _(i: Index, cond, results: List[Expr]):
+    if cond(i):
         results.append(i)
-    Filter(i.x, kind, results)
-    Filter(i.i, kind, results)
+    Filter(i.x, cond, results)
+    Filter(i.i, cond, results)
 
 @Filter.register
-def _(comma: Comma, kind: type, results: List[Expr]):
-    if isinstance(comma, kind):
+def _(comma: Comma, cond, results: List[Expr]):
+    if cond(comma):
         results.append(comma)
-    Filter(comma.expr[0], kind, results)
-    Filter(comma.expr[1], kind, results)
+    Filter(comma.expr[0], cond, results)
+    Filter(comma.expr[1], cond, results)
 
 @Filter.register
-def _(op: BinaryOp, kind: type, results: List[Expr]):
-    if isinstance(op, kind):
+def _(op: BinaryOp, cond, results: List[Expr]):
+    if cond(op):
         results.append(op)
-    Filter(op.l, kind, results)
-    Filter(op.r, kind, results)
+    Filter(op.l, cond, results)
+    Filter(op.r, cond, results)
 
 def VarsUsed(stmt: Stmt) -> Set[Var]:
     vars = []
-    Visit(loop.body, lambda node: Filter(node, Ident, vars))
+    Visit(loop.body, lambda node: Filter(node, lambda x: isinstance(x, Ident), vars))
     vars = set([v.var for v in vars])
 
     return vars