Add RzIterator support to the bindgen and linter

src/binding_class.py

@@ -42,6 +42,7 @@ class Class:
     fields: OrderedDict[str, Field]
     funcs: OrderedDict[str, Func]
     methods: OrderedDict[str, Func]
+    python_methods: OrderedDict[str, List[str]]
 
     constructor: Optional[Func]
     destructor: Optional[Func]
@@ -82,6 +83,7 @@ class Class:
         self.fields = OrderedDict()
         self.funcs = OrderedDict()
         self.methods = OrderedDict()
+        self.python_methods = OrderedDict()
         self.constructor = None
         self.destructor = None
 
@@ -192,6 +194,17 @@ class Class:
         assert rename not in self.methods
         self.methods[rename] = method
 
+    def add_python_method(self, decl: str, *lines: str) -> None:
+        """
+        Add a pure-python method with the given declaration
+
+        Used for protocol methods such as __iter__ and __len__ which
+        have no direct C counterpart. The generated SWIG output wraps
+        these in a %pythoncode block inside the class %extend.
+        """
+        assert decl not in self.python_methods
+        self.python_methods[decl] = [f"def {decl}:"] + [f"    {line}" for line in lines]
+
     def add_prefixed_methods(self, prefix: str) -> None:
         """
         Add C functions with a given prefix, and that take a

src/bindings.py

@@ -51,6 +51,20 @@ def run() -> None:
             func(Header(translation_unit, builder))
 
 
+def bind_ht_iterators(ht: Class, prefix: str) -> None:
+    """
+    Bind the RzIterator-returning helpers shared by every hashtable type
+    (see ht_inc.h). These return an `RzIterator *`, which is bound as an
+    iterable class by ``bind_iterator``.
+
+    ``as_iter``/``as_iter_keys`` yield immutable values/keys; ``as_iter_mut``
+    yields mutable values.
+    """
+    ht.add_method(f"{prefix}as_iter", rename="as_iter")
+    ht.add_method(f"{prefix}as_iter_keys", rename="as_iter_keys")
+    ht.add_method(f"{prefix}as_iter_mut", rename="as_iter_mut")
+
+
 ############
 # GENERICS #
 ############
@@ -137,6 +151,35 @@ def bind_vector(vector_h: Header) -> None:
     rz_pvector.add_python_method("__iter__(self)", "return RzPVectorIterator(self)")
 
 
+@threaded_header("rz_util/rz_iterator.h")
+def bind_iterator(iterator_h: Header) -> None:
+    """
+    RzIterator
+
+    The generic, lazy Rizin iterator (see rz_iterator.h / iterator.c).
+    Many Rizin functions return an `RzIterator *` (e.g. the hashtable
+    `as_iter*` helpers, `rz_set_*_as_iter`, graph traversals, ...).
+
+    `RzIterator` is type-erased in C: `rz_iterator_next` returns a borrowed
+    `void *`, so elements are yielded as opaque pointers. Binding it as a
+    class makes every such return value directly usable from a `for` loop
+    via the `RzIteratorIterator` python helper, and ties its lifetime to
+    `rz_iterator_free` through the destructor.
+    """
+    # rz_iterator_new takes C callbacks and cannot be bound
+    iterator_h.ignore("rz_iterator_new")
+
+    rz_iterator = Class(
+        iterator_h,
+        typedef="RzIterator",
+        # Internal state and callbacks, not useful from the guest language
+        ignore_fields={"cur", "u", "next", "free", "free_u"},
+    )
+    rz_iterator.add_method("rz_iterator_next", rename="next")
+    rz_iterator.add_destructor("rz_iterator_free")
+    rz_iterator.add_python_method("__iter__(self)", "return RzIteratorIterator(self)")
+
+
 ###########
 # HEADERS #
 ###########
@@ -434,7 +477,8 @@ def bind_ht_pp(ht_pp_h: Header) -> None:
     """
     ht_pp
     """
-    Class(ht_pp_h, typedef="HtPP")
+    ht_pp = Class(ht_pp_h, typedef="HtPP")
+    bind_ht_iterators(ht_pp, "ht_pp_")
 
 
 @threaded_header("rz_util/ht_pu.h")
@@ -442,7 +486,8 @@ def bind_ht_pu(ht_pu_h: Header) -> None:
     """
     ht_pu
     """
-    Class(ht_pu_h, typedef="HtPU")
+    ht_pu = Class(ht_pu_h, typedef="HtPU")
+    bind_ht_iterators(ht_pu, "ht_pu_")
 
 
 @threaded_header("rz_util/ht_up.h")
@@ -450,7 +495,8 @@ def bind_ht_up(ht_up_h: Header) -> None:
     """
     ht_up
     """
-    Class(ht_up_h, typedef="HtUP")
+    ht_up = Class(ht_up_h, typedef="HtUP")
+    bind_ht_iterators(ht_up, "ht_up_")
 
 
 @threaded_header("rz_util/ht_uu.h")
@@ -458,7 +504,8 @@ def bind_ht_uu(ht_uu_h: Header) -> None:
     """
     ht_uu
     """
-    Class(ht_uu_h, typedef="HtUU")
+    ht_uu = Class(ht_uu_h, typedef="HtUU")
+    bind_ht_iterators(ht_uu, "ht_uu_")
 
 
 @threaded_header("rz_util/ht_sp.h")
@@ -466,7 +513,8 @@ def bind_ht_sp(ht_sp_h: Header) -> None:
     """
     ht_sp
     """
-    Class(ht_sp_h, typedef="HtSP")
+    ht_sp = Class(ht_sp_h, typedef="HtSP")
+    bind_ht_iterators(ht_sp, "ht_sp_")
 
 
 @threaded_header("rz_util/ht_ss.h")
@@ -474,7 +522,8 @@ def bind_ht_ss(ht_ss_h: Header) -> None:
     """
     ht_ss
     """
-    Class(ht_ss_h, typedef="HtSS")
+    ht_ss = Class(ht_ss_h, typedef="HtSS")
+    bind_ht_iterators(ht_ss, "ht_ss_")
 
 
 @threaded_header("rz_util/ht_su.h")
@@ -482,4 +531,5 @@ def bind_ht_su(ht_su_h: Header) -> None:
     """
     ht_su
     """
-    Class(ht_su_h, typedef="HtSU")
+    ht_su = Class(ht_su_h, typedef="HtSU")
+    bind_ht_iterators(ht_su, "ht_su_")

src/generator_swig.py

@@ -221,7 +221,13 @@ def write_class(writer: Writer, cls: Class) -> None:
             writer.line(f'%rename {cls.struct_name}::{field.name} "";')
 
     # Extension
-    if len(cls.funcs) != 0 or len(cls.methods) != 0:
+    if (
+        len(cls.funcs) != 0
+        or len(cls.methods) != 0
+        or len(cls.python_methods) != 0
+        or cls.constructor
+        or cls.destructor
+    ):
         writer.line(f"%extend {cls.struct_name} {{")
         with writer.indent():
             if cls.constructor:
@@ -235,6 +241,12 @@ def write_class(writer: Writer, cls: Class) -> None:
             for name, method in cls.methods.items():
                 write_func(writer, method, name, FuncKind.METHOD)
 
+            for python_lines in cls.python_methods.values():
+                writer.line("%pythoncode %{")
+                with writer.indent():
+                    writer.line(*python_lines)
+                writer.line("%}")
+
         writer.line("}")
 
 

src/lint.py

@@ -114,6 +114,26 @@ generic_types = {"RzList", "RzListIter", "RzPVector", "RzVector", "RzGraph", "Ht
 skip_files = {"ht_inc.c", "ht_inc.h", "rz_th_ht.h", "thread_hash_table.c"}
 
 
+def check_iterator_comment(comment: str, location: SourceLocation) -> None:
+    """
+    Validate an RzIterator /*<type>*/ comment.
+
+    RzIterator is type-erased and takes a single element-type parameter
+    (unlike the two-parameter RzGraph/HtPP form). The comment is optional in
+    rizin -- the element type is often not known statically -- so RzIterator is
+    deliberately kept out of `generic_types`. When present, the comment must
+    name exactly one type.
+    """
+    inner = comment[1:-1].strip()
+    if not inner:
+        warn(f"Type comment at {stringify_location(location)} for RzIterator is empty")
+    elif "," in inner:
+        warn(
+            f"Type comment at {stringify_location(location)} for "
+            f"RzIterator must name exactly one type. Is: '{comment}'"
+        )
+
+
 def cursor_get_comment(cursor: Cursor, *, packed: bool = False) -> Optional[str]:
     """
     Get /*<type>*/ comment on a cursor
@@ -182,6 +202,12 @@ def cursor_get_comment(cursor: Cursor, *, packed: bool = False) -> Optional[str]
         return None
     comment = comment[2:-2]
 
+    # Normalize surrounding whitespace so both the compact `/*<T>*/` style and
+    # the space-padded `/* <T> */` style (used by RzIterator helpers) are
+    # recognized. Existing annotations have no padding, so this is a no-op for
+    # them.
+    comment = comment.strip()
+
     if not comment.startswith("<") or not comment.endswith(">"):
         if typeref_spelling in generic_types:
             warn(
@@ -227,6 +253,8 @@ def cursor_get_comment(cursor: Cursor, *, packed: bool = False) -> Optional[str]
         "SdbList",
     }:
         pass
+    elif typeref_spelling == "RzIterator":
+        check_iterator_comment(comment, cursor.location)
     else:
         warn(
             f"Type comment at {stringify_location(cursor.location)} "

src/snippets_swig/iterators.py

@@ -34,3 +34,23 @@ class RzPVectorIterator:
         data = self.rzpvector.at(self.index)
         self.index += 1
         return data
+
+
+class RzIteratorIterator:
+    def __init__(self, rziterator):
+        self.iter = rziterator
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.iter is None:
+            raise StopIteration
+        data = self.iter.next()
+        if data is None:
+            # Iteration is done: hand ownership of the RzIterator to SWIG so
+            # rz_iterator_free is called once the last reference is dropped.
+            self.iter.thisown = True
+            self.iter = None
+            raise StopIteration
+        return data