python: add TypeKindSet

C type finders are passed a bitset of type kinds, but Python type finders currently get called for each type kind. An upcoming update to the type finder interface will fix this, but we need a set of TypeKinds, and we'd rather not construct a real Python set for it. Instead, add a TypeKindSet bitset that satisfies the collections.abc.Set interface. Signed-off-by: Omar Sandoval <[email protected]>
osandov · Jun 5, 2024 · 9b73b44 · 9b73b44
1 parent f20b41c
commit 9b73b44
Show file tree

Hide file tree

Showing 9 changed files with 709 additions and 0 deletions.
diff --git a/_drgn.pyi b/_drgn.pyi
@@ -7,6 +7,7 @@ libdrgn bindings
 Don't use this module directly. Instead, use the drgn package.
 """
 
+import collections.abc
 import enum
 import os
 import sys
@@ -2329,6 +2330,26 @@ class TypeKind(enum.Enum):
     FUNCTION = ...
     """Function type."""
 
+class TypeKindSet(collections.abc.Set[TypeKind]):
+    """
+    Immutable set of :class:`TypeKind`\\ s.
+
+    >>> kinds = TypeKindSet({TypeKind.STRUCT, TypeKind.CLASS})
+    >>> TypeKind.STRUCT in kinds
+    True
+    >>> TypeKind.INT in kinds
+    False
+    >>> for kind in kinds:
+    ...     print(kind)
+    ...
+    TypeKind.STRUCT
+    TypeKind.CLASS
+    """
+
+    def __contains__(self, __x: object) -> bool: ...
+    def __iter__(self) -> Iterator[TypeKind]: ...
+    def __len__(self) -> int: ...
+
 class PrimitiveType(enum.Enum):
     """A ``PrimitiveType`` represents a primitive type known to drgn."""
 

diff --git a/docs/api_reference.rst b/docs/api_reference.rst
@@ -130,6 +130,7 @@ Types
 .. drgndoc:: TypeParameter
 .. drgndoc:: TypeTemplateParameter
 .. drgndoc:: TypeKind
+.. drgndoc:: TypeKindSet
 .. drgndoc:: PrimitiveType
 .. drgndoc:: Qualifiers
 .. drgndoc:: offsetof

diff --git a/drgn/__init__.py b/drgn/__init__.py
@@ -75,6 +75,7 @@
     Type,
     TypeEnumerator,
     TypeKind,
+    TypeKindSet,
     TypeMember,
     TypeParameter,
     TypeTemplateParameter,
@@ -129,6 +130,7 @@
     "Type",
     "TypeEnumerator",
     "TypeKind",
+    "TypeKindSet",
     "TypeMember",
     "TypeParameter",
     "TypeTemplateParameter",

diff --git a/libdrgn/Makefile.am b/libdrgn/Makefile.am
@@ -172,6 +172,7 @@ _drgn_la_SOURCES = python/constants.c \
 		   python/test.c \
 		   python/thread.c \
 		   python/type.c \
+		   python/type_kind_set.c \
 		   python/util.c
 
 _drgn_la_CFLAGS = $(AM_CFLAGS) -fvisibility=hidden

diff --git a/libdrgn/bitops.h b/libdrgn/bitops.h
@@ -87,6 +87,18 @@
 #define ilog2_impl(arg, suffix, x)	\
 	((8 * sizeof(0u##suffix) - 1) ^ __builtin_clz##suffix(x))
 
+/**
+ * Bit population count.
+ *
+ * Return the number of 1-bits in @p x.
+ *
+ * ```
+ * popcount(8) == 1
+ * popcount(3) == 2
+ * ```
+ */
+#define popcount(x) generic_bitop(x, PP_UNIQUE(_x), builtin_bitop_impl, popcount)
+
 #define builtin_bitop_impl(arg, suffix, x) __builtin_##arg##suffix(x)
 #define generic_bitop(x, unique_x, impl, impl_arg) ({			\
 	__auto_type unique_x = (x);					\

diff --git a/libdrgn/python/drgnpy.h b/libdrgn/python/drgnpy.h
@@ -196,6 +196,16 @@ typedef struct {
 	union drgn_lazy_object *lazy_obj;
 } LazyObject;
 
+typedef struct {
+	PyObject_HEAD
+	uint64_t kinds;
+} TypeKindSet;
+
+typedef struct {
+	PyObject_HEAD
+	uint64_t mask;
+} TypeKindSetIterator;
+
 typedef struct {
 	LazyObject lazy_obj;
 	PyObject *name;
@@ -236,6 +246,8 @@ extern PyTypeObject Symbol_type;
 extern PyTypeObject Thread_type;
 extern PyTypeObject ThreadIterator_type;
 extern PyTypeObject TypeEnumerator_type;
+extern PyTypeObject TypeKindSet_type;
+extern PyTypeObject TypeKindSetIterator_type;
 extern PyTypeObject TypeMember_type;
 extern PyTypeObject TypeParameter_type;
 extern PyTypeObject TypeTemplateParameter_type;
@@ -259,6 +271,9 @@ PyObject *Language_wrap(const struct drgn_language *language);
 int language_converter(PyObject *o, void *p);
 int add_languages(void);
 
+PyObject *TypeKindSet_wrap(uint64_t mask);
+int init_type_kind_set(void);
+
 static inline DrgnObject *DrgnObject_alloc(Program *prog)
 {
 	DrgnObject *ret = call_tp_alloc(DrgnObject);

diff --git a/libdrgn/python/main.c b/libdrgn/python/main.c
@@ -277,6 +277,9 @@ DRGNPY_PUBLIC PyMODINIT_FUNC PyInit__drgn(void)
 	    add_type(m, &Thread_type) ||
 	    add_type(m, &ThreadIterator_type) ||
 	    add_type(m, &TypeEnumerator_type) ||
+	    add_type(m, &TypeKindSet_type) ||
+	    PyType_Ready(&TypeKindSetIterator_type) ||
+	    init_type_kind_set() ||
 	    add_type(m, &TypeMember_type) ||
 	    add_type(m, &TypeParameter_type) ||
 	    add_type(m, &TypeTemplateParameter_type) ||