Add StackFrame.locals() method

The StackFrame's __getitem__() method allows looking up names in the
scope of a stack frame, which is an incredibly useful tool for
debugging. However, the names are not discoverable -- you must already
be looking at the source code or some other source to know what names
can be queried. To fix this, add a locals() method to StackFrame, which
lists names that can be queried in the scope. Since this method is named
locals(), it stops at the function scope and doesn't include globals or
class members.

Signed-off-by: Stephen Brennan <stephen.s.brennan@oracle.com>

_drgn.pyi

@@ -1704,6 +1704,16 @@ class StackFrame:
         :param name: Object name.
         """
         ...
+    def locals(self) -> List[str]:
+        """
+        Get a list of the names of all local objects (local variables, function
+        parameters, local constants, and nested functions) in the scope of this
+        frame.
+
+        Not all names may have present values, but they can be used with the
+        :meth:`[] <.__getitem__>` operator to check.
+        """
+        ...
     def source(self) -> Tuple[str, int, int]:
         """
         Get the source code location of this frame.

libdrgn/drgn.h.in

@@ -2789,6 +2789,31 @@ struct drgn_error *drgn_stack_frame_symbol(struct drgn_stack_trace *trace,
 					   size_t frame,
 					   struct drgn_symbol **ret);
 
+/**
+ * Get the names of local objects in the scope of this frame.
+ *
+ * The array of names must be freed with @ref drgn_stack_frame_locals_destroy().
+ *
+ * @param[out] names_ret Returned array of names. On success, must be freed with
+ * @ref drgn_stack_frame_locals_destroy().
+ * @param[out] count_ret Returned number of names in @p names_ret.
+ * @return @c NULL on success, non-@c NULL on error
+ */
+struct drgn_error *
+drgn_stack_frame_locals(struct drgn_stack_trace *trace, size_t frame,
+			const char ***names_ret, size_t *count_ret);
+
+/**
+ * Free an array of names returned by @ref drgn_stack_frame_locals().
+ *
+ * The individual names from this array are invalid once this function is
+ * called. Any string which will be used later should be copied.
+ *
+ * @param names Array of names returned by @ref drgn_stack_frame_locals().
+ * @param count Count returned by @ref drgn_stack_frame_locals().
+ */
+void drgn_stack_frame_locals_destroy(const char **names, size_t count);
+
 /**
  * Find an object in the scope of a stack frame.
  *

libdrgn/dwarf_info.c

@@ -5314,6 +5314,83 @@ drgn_object_from_dwarf(struct drgn_debug_info *dbinfo,
 					       function_die, regs, ret);
 }
 
+DEFINE_VECTOR(const_char_p_vector, const char *);
+
+static struct drgn_error *add_dwarf_enumerators(Dwarf_Die *enumeration_type,
+						struct const_char_p_vector *vec)
+{
+	Dwarf_Die child;
+	int r = dwarf_child(enumeration_type, &child);
+	while (r == 0) {
+		if (dwarf_tag(&child) == DW_TAG_enumerator) {
+			const char *die_name = dwarf_diename(&child);
+			if (!die_name)
+				continue;
+			if (!const_char_p_vector_append(vec, &die_name))
+				return &drgn_enomem;
+		}
+		r = dwarf_siblingof(&child, &child);
+	}
+	if (r < 0)
+		return drgn_error_libdw();
+	return NULL;
+}
+
+struct drgn_error *drgn_dwarf_scopes_names(Dwarf_Die *scopes,
+					   size_t num_scopes,
+					   const char ***names_ret,
+					   size_t *count_ret)
+{
+	struct drgn_error *err;
+	Dwarf_Die die;
+	struct const_char_p_vector vec = VECTOR_INIT;
+	for (size_t scope = 0; scope < num_scopes; scope++) {
+		if (dwarf_child(&scopes[scope], &die) != 0)
+			continue;
+		do {
+			switch (dwarf_tag(&die)) {
+			case DW_TAG_variable:
+			case DW_TAG_formal_parameter:
+			case DW_TAG_subprogram: {
+				const char *die_name = dwarf_diename(&die);
+				if (!die_name)
+					continue;
+				if (!const_char_p_vector_append(&vec,
+								&die_name)) {
+					err = &drgn_enomem;
+					goto err;
+				}
+				break;
+			}
+			case DW_TAG_enumeration_type: {
+				bool enum_class;
+				if (dwarf_flag_integrate(&die, DW_AT_enum_class,
+							 &enum_class)) {
+					err = drgn_error_libdw();
+					goto err;
+				}
+				if (!enum_class) {
+					err = add_dwarf_enumerators(&die, &vec);
+					if (err)
+						goto err;
+				}
+				break;
+			}
+			default:
+				continue;
+			}
+		} while (dwarf_siblingof(&die, &die) == 0);
+	}
+	const_char_p_vector_shrink_to_fit(&vec);
+	*names_ret = vec.data;
+	*count_ret = vec.size;
+	return NULL;
+
+err:
+	const_char_p_vector_deinit(&vec);
+	return err;
+}
+
 static struct drgn_error *find_dwarf_enumerator(Dwarf_Die *enumeration_type,
 						const char *name,
 						Dwarf_Die *ret)

libdrgn/dwarf_info.h

@@ -231,6 +231,19 @@ struct drgn_error *drgn_find_die_ancestors(Dwarf_Die *die, Dwarf_Die **dies_ret,
 					  size_t *length_ret)
 	__attribute__((__nonnull__(2, 3)));
 
+/**
+ * Get an array of names of `DW_TAG_variable` and `DW_TAG_formal_parameter` DIEs
+ * in local scopes.
+ *
+ * @param[out] names_ret Returned array of names. On success, must be freed with
+ * @c free(). The individual strings should not be freed.
+ * @param[out] count_ret Returned number of names in @p names_ret.
+ */
+struct drgn_error *drgn_dwarf_scopes_names(Dwarf_Die *scopes,
+					   size_t num_scopes,
+					   const char ***names_ret,
+					   size_t *count_ret);
+
 /**
  * Find an object DIE in an array of DWARF scopes.
  *

libdrgn/python/stack_trace.c

@@ -113,6 +113,34 @@ static PyObject *StackFrame_str(StackFrame *self)
 	return ret;
 }
 
+static PyObject *StackFrame_locals(StackFrame *self)
+{
+	struct drgn_error *err;
+	const char **names;
+	size_t count;
+	err = drgn_stack_frame_locals(self->trace->trace, self->i, &names,
+				      &count);
+	if (err)
+		return set_drgn_error(err);
+
+	PyObject *list = PyList_New(count);
+	if (!list) {
+		drgn_stack_frame_locals_destroy(names, count);
+		return NULL;
+	}
+	for (size_t i = 0; i < count; i++) {
+		PyObject *string = PyUnicode_FromString(names[i]);
+		if (!string) {
+			drgn_stack_frame_locals_destroy(names, count);
+			Py_DECREF(list);
+			return NULL;
+		}
+		PyList_SET_ITEM(list, i, string);
+	}
+	drgn_stack_frame_locals_destroy(names, count);
+	return list;
+}
+
 static DrgnObject *StackFrame_subscript(StackFrame *self, PyObject *key)
 {
 	struct drgn_error *err;
@@ -297,6 +325,8 @@ static PyMethodDef StackFrame_methods[] = {
 	 METH_O | METH_COEXIST, drgn_StackFrame___getitem___DOC},
 	{"__contains__", (PyCFunction)StackFrame_contains,
 	 METH_O | METH_COEXIST, drgn_StackFrame___contains___DOC},
+	{"locals", (PyCFunction)StackFrame_locals,
+	 METH_NOARGS, drgn_StackFrame_locals_DOC},
 	{"source", (PyCFunction)StackFrame_source, METH_NOARGS,
 	 drgn_StackFrame_source_DOC},
 	{"symbol", (PyCFunction)StackFrame_symbol, METH_NOARGS,

libdrgn/stack_trace.c

@@ -377,6 +377,27 @@ drgn_stack_frame_symbol(struct drgn_stack_trace *trace, size_t frame,
 	return NULL;
 }
 
+LIBDRGN_PUBLIC struct drgn_error *
+drgn_stack_frame_locals(struct drgn_stack_trace *trace, size_t frame_i,
+			const char ***names_ret, size_t *count_ret)
+{
+	struct drgn_stack_frame *frame = &trace->frames[frame_i];
+	if (frame->function_scope >= frame->num_scopes) {
+		*names_ret = NULL;
+		*count_ret = 0;
+		return NULL;
+	}
+	return drgn_dwarf_scopes_names(frame->scopes + frame->function_scope,
+				       frame->num_scopes - frame->function_scope,
+				       names_ret, count_ret);
+}
+
+LIBDRGN_PUBLIC
+void drgn_stack_frame_locals_destroy(const char **names, size_t count)
+{
+	free(names);
+}
+
 LIBDRGN_PUBLIC struct drgn_error *
 drgn_stack_frame_find_object(struct drgn_stack_trace *trace, size_t frame_i,
 			     const char *name, struct drgn_object *ret)

tests/linux_kernel/kmod/drgn_test.c

@@ -10,6 +10,7 @@
 
 #include <linux/io.h>
 #include <linux/kernel.h>
+#include <linux/kthread.h>
 #include <linux/list.h>
 #include <linux/llist.h>
 #include <linux/mm.h>
@@ -385,6 +386,45 @@ static int drgn_test_slab_init(void)
 	return 0;
 }
 
+// kthread for stack trace
+
+static struct task_struct *drgn_kthread;
+
+static int __attribute__((optimize("O0"))) drgn_kthread_fn(void *arg)
+{
+	int a, b, c;
+
+retry:
+	for (a = 0; a < 100; a++) {
+		for (b = 100; b >= 0; b--) {
+			for (c = 0; c < a; c++) {
+				set_current_state(TASK_UNINTERRUPTIBLE);
+				schedule();
+				if (kthread_should_stop())
+					return 0;
+			}
+		}
+	}
+	goto retry;
+}
+
+static void drgn_test_stack_trace_exit(void)
+{
+	if (drgn_kthread) {
+		kthread_stop(drgn_kthread);
+		drgn_kthread = NULL;
+	}
+}
+
+static int drgn_test_stack_trace_init(void)
+{
+	drgn_kthread = kthread_create(drgn_kthread_fn, (void *)0xF0F0F0F0, "drgn_kthread");
+	if (!drgn_kthread)
+		return -1;
+	wake_up_process(drgn_kthread);
+	return 0;
+}
+
 // Dummy function symbol.
 int drgn_test_function(int x)
 {
@@ -396,6 +436,7 @@ static void drgn_test_exit(void)
 	drgn_test_slab_exit();
 	drgn_test_percpu_exit();
 	drgn_test_mm_exit();
+	drgn_test_stack_trace_exit();
 }
 
 static int __init drgn_test_init(void)
@@ -412,6 +453,9 @@ static int __init drgn_test_init(void)
 		goto out;
 	drgn_test_rbtree_init();
 	ret = drgn_test_slab_init();
+	if (ret)
+		goto out;
+	ret = drgn_test_stack_trace_init();
 out:
 	if (ret)
 		drgn_test_exit();

tests/linux_kernel/test_stack_trace.py

@@ -7,7 +7,12 @@
 from drgn import Object, Program, cast
 from drgn.helpers.linux.pid import find_task
 from tests import assertReprPrettyEqualsStr
-from tests.linux_kernel import LinuxKernelTestCase, fork_and_sigwait, setenv
+from tests.linux_kernel import (
+    LinuxKernelTestCase,
+    fork_and_sigwait,
+    setenv,
+    skip_unless_have_test_kmod,
+)
 from util import NORMALIZED_MACHINE_NAME
 
 
@@ -94,3 +99,17 @@ def test_stack__repr_pretty_(self):
             assertReprPrettyEqualsStr(trace)
             for frame in trace:
                 assertReprPrettyEqualsStr(frame)
+
+    @skip_unless_have_test_kmod
+    def test_stack_locals(self):
+        task = self.prog["drgn_kthread"]
+        stack_trace = self.prog.stack_trace(task)
+        for frame in stack_trace:
+            if frame.symbol().name == "drgn_kthread_fn":
+                self.assertSetEqual(
+                    {"arg", "a", "b", "c"},
+                    set(frame.locals()),
+                )
+                break
+        else:
+            self.fail("Couldn't find drgn_kthread_fn frame")