diff --git a/src/kirin/dialects/ilist/rewrite/hint_len.py b/src/kirin/dialects/ilist/rewrite/hint_len.py index b20229ec7..4c535474e 100644 --- a/src/kirin/dialects/ilist/rewrite/hint_len.py +++ b/src/kirin/dialects/ilist/rewrite/hint_len.py @@ -1,6 +1,6 @@ from kirin import ir, types from kirin.analysis import const -from kirin.dialects import py +from kirin.dialects import py, scf from kirin.rewrite.abc import RewriteRule, RewriteResult from kirin.dialects.ilist.stmts import IListType @@ -26,7 +26,11 @@ def _get_collection_len(self, collection: ir.SSAValue): return None def rewrite_Statement(self, node: ir.Statement) -> RewriteResult: - if not isinstance(node, py.Len): + + if not ( + isinstance(node, py.Len) + and not isinstance(node.parent_stmt, (scf.For, scf.IfElse)) + ): return RewriteResult() if (coll_len := self._get_collection_len(node.value)) is None: