From 8aa998035aa94ffd410aa76a599911f8e7a479b1 Mon Sep 17 00:00:00 2001
From: Omar Sandoval <osandov@osandov.com>
Date: Tue, 5 Mar 2024 10:43:03 -0800
Subject: [PATCH] helpers.linux.idr: add idr_for_each_entry()

This is a convenience wrapper around idr_for_each() and a cast.

Signed-off-by: Omar Sandoval <osandov@osandov.com>
---
 drgn/helpers/linux/idr.py              | 23 ++++++++++++++++++++---
 tests/linux_kernel/helpers/test_idr.py | 13 ++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/drgn/helpers/linux/idr.py b/drgn/helpers/linux/idr.py
index dc3ee9fb3..a5ec741bc 100644
--- a/drgn/helpers/linux/idr.py
+++ b/drgn/helpers/linux/idr.py
@@ -11,15 +11,16 @@
 """
 
 import operator
-from typing import Iterator, Tuple
+from typing import Iterator, Tuple, Union
 
 from _drgn import _linux_helper_idr_find
-from drgn import NULL, IntegerLike, Object, cast, sizeof
+from drgn import NULL, IntegerLike, Object, Type, cast, sizeof
 from drgn.helpers.linux.radixtree import radix_tree_for_each
 
 __all__ = (
     "idr_find",
     "idr_for_each",
+    "idr_for_each_entry",
 )
 
 _IDR_BITS = 8
@@ -66,7 +67,7 @@ def idr_find(idr: Object, id: IntegerLike) -> Object:
 
 def idr_for_each(idr: Object) -> Iterator[Tuple[int, Object]]:
     """
-    Iterate over all of the entries in an IDR.
+    Iterate over all of the pointers in an IDR.
 
     :param idr: ``struct idr *``
     :return: Iterator of (index, ``void *``) tuples.
@@ -97,3 +98,19 @@ def aux(p: Object, id: int, n: int) -> Iterator[Tuple[int, Object]]:
             base = 0
         for index, entry in radix_tree_for_each(idr_rt.address_of_()):
             yield index + base, entry
+
+
+def idr_for_each_entry(
+    idr: Object, type: Union[str, Type]
+) -> Iterator[Tuple[int, Object]]:
+    """
+    Iterate over all of the entries with the given type in an IDR.
+
+    :param idr: ``struct idr *``
+    :param type: Entry type.
+    :return: Iterator of (index, ``type *``) tuples.
+    """
+    prog = idr.prog_
+    type = prog.pointer_type(prog.type(type))
+    for index, entry in idr_for_each(idr):
+        yield index, cast(type, entry)
diff --git a/tests/linux_kernel/helpers/test_idr.py b/tests/linux_kernel/helpers/test_idr.py
index dd0345687..cdc0904ec 100644
--- a/tests/linux_kernel/helpers/test_idr.py
+++ b/tests/linux_kernel/helpers/test_idr.py
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: LGPL-2.1-or-later
 
 from drgn import NULL, Object
-from drgn.helpers.linux.idr import idr_find, idr_for_each
+from drgn.helpers.linux.idr import idr_find, idr_for_each, idr_for_each_entry
 from tests.linux_kernel import LinuxKernelTestCase, skip_unless_have_test_kmod
 
 
@@ -69,3 +69,14 @@ def test_idr_for_each_sparse(self):
                 (0xEE, Object(self.prog, "void *", 0x9ABC)),
             ],
         )
+
+    def test_idr_for_each_entry(self):
+        root = self.prog["drgn_test_idr_sparse"].address_of_()
+        self.assertIdentical(
+            list(idr_for_each_entry(root, "int")),
+            [
+                (1, Object(self.prog, "int *", 0x1234)),
+                (0x80, Object(self.prog, "int *", 0x5678)),
+                (0xEE, Object(self.prog, "int *", 0x9ABC)),
+            ],
+        )