libdrgn: allow passing multiple type kinds to type finder function

For the next change, we want to look up a name which may have one of
multiple type kinds. Make drgn_type_kind_fn in libdrgn take a bitmask of
kinds instead of a single kind.

We could change the Python bindings to take the same bitmask, or a tuple
of drgn.TypeKind, but either would be a breaking API change. For now,
let's call the type finder function for each kind in the bitmask
instead.

Signed-off-by: Omar Sandoval <osandov@osandov.com>

libdrgn/debug_info.h

@@ -247,8 +247,8 @@ drgn_debug_info_main_language(struct drgn_debug_info *dbinfo,
 			      const struct drgn_language **ret);
 
 /** @ref drgn_type_find_fn() that uses debugging information. */
-struct drgn_error *drgn_debug_info_find_type(enum drgn_type_kind kind,
-					     const char *name, size_t name_len,
+struct drgn_error *drgn_debug_info_find_type(uint64_t kinds, const char *name,
+					     size_t name_len,
 					     const char *filename, void *arg,
 					     struct drgn_qualified_type *ret);
 

libdrgn/drgn.h.in

@@ -581,7 +581,9 @@ bool drgn_filename_matches(const char *haystack, const char *needle);
 /**
  * Callback for finding a type.
  *
- * @param[in] kind Kind of type.
+ * @param[in] kinds Kinds of types to find, as a bitmask of bits shifted by @ref
+ * drgn_type_kind. E.g., `(1 << DRGN_TYPE_STRUCT) | (1 << DRGN_TYPE_CLASS)`
+ * means to find a structure or class type.
  * @param[in] name Name of type (or tag, for structs, unions, and enums). This
  * is @em not null-terminated.
  * @param[in] name_len Length of @p name.
@@ -594,8 +596,8 @@ bool drgn_filename_matches(const char *haystack, const char *needle);
  * considered fatal.
  */
 typedef struct drgn_error *
-(*drgn_type_find_fn)(enum drgn_type_kind kind, const char *name,
-		     size_t name_len, const char *filename, void *arg,
+(*drgn_type_find_fn)(uint64_t kinds, const char *name, size_t name_len,
+		     const char *filename, void *arg,
 		     struct drgn_qualified_type *ret);
 
 /**

libdrgn/dwarf_info.c

@@ -6123,48 +6123,39 @@ find_enclosing_namespace(struct drgn_namespace_dwarf_index *global_namespace,
 	return NULL;
 }
 
-struct drgn_error *drgn_debug_info_find_type(enum drgn_type_kind kind,
-					     const char *name, size_t name_len,
+struct drgn_error *drgn_debug_info_find_type(uint64_t kinds, const char *name,
+					     size_t name_len,
 					     const char *filename, void *arg,
 					     struct drgn_qualified_type *ret)
 {
 	struct drgn_error *err;
 	struct drgn_debug_info *dbinfo = arg;
 
-	enum drgn_dwarf_index_tag tag;
-	switch (kind) {
-	case DRGN_TYPE_INT:
-	case DRGN_TYPE_BOOL:
-	case DRGN_TYPE_FLOAT:
-		tag = DRGN_DWARF_INDEX_base_type;
-		break;
-	case DRGN_TYPE_STRUCT:
-		tag = DRGN_DWARF_INDEX_structure_type;
-		break;
-	case DRGN_TYPE_UNION:
-		tag = DRGN_DWARF_INDEX_union_type;
-		break;
-	case DRGN_TYPE_CLASS:
-		tag = DRGN_DWARF_INDEX_class_type;
-		break;
-	case DRGN_TYPE_ENUM:
-		tag = DRGN_DWARF_INDEX_enumeration_type;
-		break;
-	case DRGN_TYPE_TYPEDEF:
-		tag = DRGN_DWARF_INDEX_typedef;
-		break;
-	default:
-		UNREACHABLE();
-	}
+	enum drgn_dwarf_index_tag tags[6];
+	size_t num_tags = 0;
+	if (kinds & ((1 << DRGN_TYPE_INT)
+		     | (1 << DRGN_TYPE_BOOL)
+		     | (1 << DRGN_TYPE_FLOAT)))
+		tags[num_tags++] = DRGN_DWARF_INDEX_base_type;
+	if (kinds & (1 << DRGN_TYPE_STRUCT))
+		tags[num_tags++] = DRGN_DWARF_INDEX_structure_type;
+	if (kinds & (1 << DRGN_TYPE_UNION))
+		tags[num_tags++] = DRGN_DWARF_INDEX_union_type;
+	if (kinds & (1 << DRGN_TYPE_CLASS))
+		tags[num_tags++] = DRGN_DWARF_INDEX_class_type;
+	if (kinds & (1 << DRGN_TYPE_ENUM))
+		tags[num_tags++] = DRGN_DWARF_INDEX_enumeration_type;
+	if (kinds & (1 << DRGN_TYPE_TYPEDEF))
+		tags[num_tags++] = DRGN_DWARF_INDEX_typedef;
 
 	struct drgn_namespace_dwarf_index *namespace;
 	err = find_enclosing_namespace(&dbinfo->dwarf.global,
 				       &name, &name_len, &namespace);
 	if (err)
 		return err;
 	struct drgn_dwarf_index_iterator it;
-	err = drgn_dwarf_index_iterator_init(&it, namespace, name,
-					     name_len, &tag, 1);
+	err = drgn_dwarf_index_iterator_init(&it, namespace, name, name_len,
+					     tags, num_tags);
 	if (err)
 		return err;
 	Dwarf_Die die;
@@ -6178,7 +6169,7 @@ struct drgn_error *drgn_debug_info_find_type(enum drgn_type_kind kind,
 			 * For base_type, we need to check that the type we
 			 * found was the right kind.
 			 */
-			if (drgn_type_kind(ret->type) == kind)
+			if (kinds & (UINT64_C(1) << drgn_type_kind(ret->type)))
 				return NULL;
 		}
 	}

libdrgn/language_c.c

@@ -2273,16 +2273,15 @@ c_parse_specifier_qualifier_list(struct drgn_program *prog,
 	}
 
 	if (specifier == SPECIFIER_NONE) {
-		enum drgn_type_kind kind;
-
+		uint64_t kinds;
 		if (tag_token == C_TOKEN_STRUCT) {
-			kind = DRGN_TYPE_STRUCT;
+			kinds = 1 << DRGN_TYPE_STRUCT;
 		} else if (tag_token == C_TOKEN_UNION) {
-			kind = DRGN_TYPE_UNION;
+			kinds = 1 << DRGN_TYPE_UNION;
 		} else if (tag_token == C_TOKEN_CLASS) {
-			kind = DRGN_TYPE_CLASS;
+			kinds = 1 << DRGN_TYPE_CLASS;
 		} else if (tag_token == C_TOKEN_ENUM) {
-			kind = DRGN_TYPE_ENUM;
+			kinds = 1 << DRGN_TYPE_ENUM;
 		} else if (identifier) {
 			if (identifier_len == sizeof("size_t") - 1 &&
 			    memcmp(identifier, "size_t",
@@ -2305,14 +2304,14 @@ c_parse_specifier_qualifier_list(struct drgn_program *prog,
 				ret->qualifiers = 0;
 				goto out;
 			} else {
-				kind = DRGN_TYPE_TYPEDEF;
+				kinds = 1 << DRGN_TYPE_TYPEDEF;
 			}
 		} else {
 			return drgn_error_create(DRGN_ERROR_SYNTAX,
 						 "expected type specifier");
 		}
 
-		err = drgn_program_find_type_impl(prog, kind, identifier,
+		err = drgn_program_find_type_impl(prog, kinds, identifier,
 						  identifier_len, filename,
 						  ret);
 		if (err)

libdrgn/python/program.c

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: LGPL-2.1-or-later
 
 #include "drgnpy.h"
+#include "../bitops.h"
 #include "../error.h"
 #include "../hash_table.h"
 #include "../log.h"
@@ -372,49 +373,52 @@ static PyObject *Program_add_memory_segment(Program *self, PyObject *args,
 	Py_RETURN_NONE;
 }
 
-static struct drgn_error *py_type_find_fn(enum drgn_type_kind kind,
-					  const char *name, size_t name_len,
-					  const char *filename, void *arg,
+static struct drgn_error *py_type_find_fn(uint64_t kinds, const char *name,
+					  size_t name_len, const char *filename,
+					  void *arg,
 					  struct drgn_qualified_type *ret)
 {
 	PyGILState_guard();
 
-	_cleanup_pydecref_ PyObject *
-		kind_obj = PyObject_CallFunction(TypeKind_class, "k",
-						 (unsigned long)kind);
-	if (!kind_obj)
-		return drgn_error_from_python();
 	_cleanup_pydecref_ PyObject *name_obj =
 		PyUnicode_FromStringAndSize(name, name_len);
 	if (!name_obj)
 		return drgn_error_from_python();
-	_cleanup_pydecref_ PyObject *type_obj =
-		PyObject_CallFunction(PyTuple_GET_ITEM(arg, 1), "OOs", kind_obj,
-				      name_obj, filename);
-	if (!type_obj)
-		return drgn_error_from_python();
-	if (type_obj == Py_None)
-		return &drgn_not_found;
-	if (!PyObject_TypeCheck(type_obj, &DrgnType_type)) {
-		PyErr_SetString(PyExc_TypeError,
-				"type find callback must return Type or None");
-		return drgn_error_from_python();
-	}
-	/*
-	 * This check is also done in libdrgn, but we need it here because if
-	 * the type isn't from this program, then there's no guarantee that it
-	 * will remain valid after we decrement its reference count.
-	 */
-	if (DrgnType_prog((DrgnType *)type_obj) !=
-	    (Program *)PyTuple_GET_ITEM(arg, 0)) {
-		PyErr_SetString(PyExc_ValueError,
-				"type find callback returned type from wrong program");
-		return drgn_error_from_python();
-	}
 
-	ret->type = ((DrgnType *)type_obj)->type;
-	ret->qualifiers = ((DrgnType *)type_obj)->qualifiers;
-	return NULL;
+	int kind;
+	for_each_bit(kind, kinds) {
+		_cleanup_pydecref_ PyObject *
+			kind_obj = PyObject_CallFunction(TypeKind_class, "i",
+							 kind);
+		if (!kind_obj)
+			return drgn_error_from_python();
+		_cleanup_pydecref_ PyObject *type_obj =
+			PyObject_CallFunction(PyTuple_GET_ITEM(arg, 1), "OOs",
+					      kind_obj, name_obj, filename);
+		if (!type_obj)
+			return drgn_error_from_python();
+		if (type_obj == Py_None)
+			continue;
+		if (!PyObject_TypeCheck(type_obj, &DrgnType_type)) {
+			PyErr_SetString(PyExc_TypeError,
+					"type find callback must return Type or None");
+			return drgn_error_from_python();
+		}
+		// This check is also done in libdrgn, but we need it here
+		// because if the type isn't from this program, then there's no
+		// guarantee that it will remain valid after we decrement its
+		// reference count.
+		if (DrgnType_prog((DrgnType *)type_obj)
+		    != (Program *)PyTuple_GET_ITEM(arg, 0)) {
+			PyErr_SetString(PyExc_ValueError,
+					"type find callback returned type from wrong program");
+			return drgn_error_from_python();
+		}
+		ret->type = ((DrgnType *)type_obj)->type;
+		ret->qualifiers = ((DrgnType *)type_obj)->qualifiers;
+		return NULL;
+	}
+	return &drgn_not_found;
 }
 
 static PyObject *Program_add_type_finder(Program *self, PyObject *args,

libdrgn/type.c

@@ -1326,23 +1326,23 @@ drgn_program_add_type_finder(struct drgn_program *prog, drgn_type_find_fn fn,
 	return NULL;
 }
 
-struct drgn_error *
-drgn_program_find_type_impl(struct drgn_program *prog,
-			    enum drgn_type_kind kind, const char *name,
-			    size_t name_len, const char *filename,
-			    struct drgn_qualified_type *ret)
+struct drgn_error *drgn_program_find_type_impl(struct drgn_program *prog,
+					       uint64_t kinds, const char *name,
+					       size_t name_len,
+					       const char *filename,
+					       struct drgn_qualified_type *ret)
 {
 	struct drgn_type_finder *finder = prog->type_finders;
 	while (finder) {
 		struct drgn_error *err =
-			finder->fn(kind, name, name_len, filename, finder->arg,
+			finder->fn(kinds, name, name_len, filename, finder->arg,
 				   ret);
 		if (!err) {
 			if (drgn_type_program(ret->type) != prog) {
 				return drgn_error_create(DRGN_ERROR_INVALID_ARGUMENT,
 							 "type find callback returned type from wrong program");
 			}
-			if (drgn_type_kind(ret->type) != kind) {
+			if (!(kinds & (UINT64_C(1) << drgn_type_kind(ret->type)))) {
 				return drgn_error_create(DRGN_ERROR_TYPE,
 							 "type find callback returned wrong kind of type");
 			}
@@ -1445,7 +1445,8 @@ drgn_program_find_primitive_type(struct drgn_program *prog,
 
 	spellings = drgn_primitive_type_spellings[type];
 	for (i = 0; spellings[i]; i++) {
-		err = drgn_program_find_type_impl(prog, kind, spellings[i],
+		err = drgn_program_find_type_impl(prog, UINT64_C(1) << kind,
+						  spellings[i],
 						  strlen(spellings[i]), NULL,
 						  &qualified_type);
 		if (!err && drgn_type_primitive(qualified_type.type) == type) {

libdrgn/type.h

@@ -579,11 +579,11 @@ void drgn_program_deinit_types(struct drgn_program *prog);
  * @return @c NULL on success, &@ref drgn_not_found if the type wasn't found,
  * non-@c NULL on other error.
  */
-struct drgn_error *
-drgn_program_find_type_impl(struct drgn_program *prog,
-			    enum drgn_type_kind kind, const char *name,
-			    size_t name_len, const char *filename,
-			    struct drgn_qualified_type *ret);
+struct drgn_error *drgn_program_find_type_impl(struct drgn_program *prog,
+					       uint64_t kinds, const char *name,
+					       size_t name_len,
+					       const char *filename,
+					       struct drgn_qualified_type *ret);
 
 /** Find a primitive type in a @ref drgn_program. */
 struct drgn_error *