Add offsets for dict/dictkeys/dictvalues

Using offsets for dict paves the way for handling changes to the size of PyObject. It's not strictly necessary to make the same change for dictkeys and dictvalues, but I've done it anyway because it makes the implementation simpler (decoupling the functions attempting to use the object from the different structure of that object between versions). I left `PyDictEntry` alone for now since the layout is still the same across all Python 3 versions. Signed-off-by: Matt Wozniski <[email protected]>
bloomberg · Aug 13, 2024 · 4b4e0bd · 4b4e0bd
1 parent 969807a
commit 4b4e0bd
Show file tree

Hide file tree

Showing 4 changed files with 136 additions and 25 deletions.
diff --git a/src/pystack/_pystack/cpython/dict.h b/src/pystack/_pystack/cpython/dict.h
@@ -29,7 +29,7 @@ typedef struct _dictobject
 
 namespace Python3 {
 typedef Py_ssize_t (*dict_lookup_func)(void* mp, PyObject* key, Py_hash_t hash, PyObject** value_addr);
-struct PyDictKeysObject;
+union PyDictKeysObject;
 
 typedef struct
 {
@@ -99,4 +99,14 @@ typedef struct _dictvalues
 
 }  // namespace Python3_13
 
+typedef union {
+    Python3_3::PyDictKeysObject v3_3;
+    Python3_11::PyDictKeysObject v3_11;
+} PyDictKeysObject;
+
+typedef union {
+    Python3::PyDictValuesObject v3_3;
+    Python3_13::PyDictValuesObject v3_13;
+} PyDictValuesObject;
+
 }  // namespace pystack
diff --git a/src/pystack/_pystack/pytypes.cpp b/src/pystack/_pystack/pytypes.cpp
@@ -254,22 +254,21 @@ getDictEntries(
         ssize_t& num_items,
         std::vector<Python3::PyDictKeyEntry>& valid_entries)
 {
-    auto keys_addr = reinterpret_cast<remote_addr_t>(dict.ma_keys);
+    remote_addr_t keys_addr = manager->getField(dict, &py_dict_v::o_ma_keys);
     assert(manager->versionIsAtLeast(3, 0));
     ssize_t dk_size = 0;
     int dk_kind = 0;
 
+    PyDictKeysObject keys;
+    manager->copyMemoryFromProcess(keys_addr, manager->offsets().py_dictkeys.size, &keys);
+    num_items = manager->getField(keys, &py_dictkeys_v::o_dk_nentries);
+    dk_size = manager->getField(keys, &py_dictkeys_v::o_dk_size);
+
     if (manager->versionIsAtLeast(3, 11)) {
-        Python3_11::PyDictKeysObject keys;
-        manager->copyObjectFromProcess(keys_addr, &keys);
-        num_items = keys.dk_nentries;
-        dk_size = 1L << keys.dk_log2_size;
-        dk_kind = keys.dk_kind;
-    } else {
-        Python3_3::PyDictKeysObject keys;
-        manager->copyObjectFromProcess(keys_addr, &keys);
-        num_items = keys.dk_nentries;
-        dk_size = keys.dk_size;
+        // We're reusing the o_dk_size offset for dk_log2_size. Fix up the value.
+        dk_size = 1L << dk_size;
+        // Added in 3.11
+        dk_kind = manager->getField(keys, &py_dictkeys_v::o_dk_kind);
     }
     if (num_items == 0) {
         LOG(DEBUG) << std::hex << std::showbase << "There are no elements in this dict";
@@ -294,13 +293,7 @@ getDictEntries(
         offset = 8 * dk_size;
     }
 
-    offset_t dk_indices_offset = 0;
-    if (manager->versionIsAtLeast(3, 11)) {
-        dk_indices_offset = offsetof(Python3_11::PyDictKeysObject, dk_indices);
-    } else {
-        dk_indices_offset = offsetof(Python3_3::PyDictKeysObject, dk_indices);
-    }
-
+    offset_t dk_indices_offset = manager->getFieldOffset(&py_dictkeys_v::o_dk_indices);
     remote_addr_t entries_addr = keys_addr + dk_indices_offset + offset;
 
     std::vector<Python3::PyDictKeyEntry> raw_entries;
@@ -371,7 +364,7 @@ void
 DictObject::loadFromPython3(remote_addr_t addr)
 {
     Python3::PyDictObject dict;
-    d_manager->copyObjectFromProcess(addr, &dict);
+    d_manager->copyMemoryFromProcess(addr, d_manager->offsets().py_dict.size, &dict);
 
     ssize_t num_items;
     std::vector<Python3::PyDictKeyEntry> valid_entries;
@@ -400,15 +393,12 @@ DictObject::loadFromPython3(remote_addr_t addr)
      *          All dicts sharing same key must have same insertion order.
      */
 
-    auto dictvalues_addr = (remote_addr_t)dict.ma_values;
+    remote_addr_t dictvalues_addr = d_manager->getField(dict, &py_dict_v::o_ma_values);
 
     // Get the values in one copy if we are dealing with a split-table dictionary
     if (dictvalues_addr != 0) {
         d_values.resize(num_items);
-        auto values_offset =
-                (d_manager->versionIsAtLeast(3, 13) ? offsetof(Python3_13::PyDictValuesObject, values)
-                                                    : offsetof(Python3::PyDictValuesObject, values));
-
+        auto values_offset = d_manager->getFieldOffset(&py_dictvalues_v::o_values);
         auto values_addr = dictvalues_addr + values_offset;
         d_manager->copyMemoryFromProcess(values_addr, num_items * sizeof(PyObject*), d_values.data());
     } else {

diff --git a/src/pystack/_pystack/version.cpp b/src/pystack/_pystack/version.cpp
@@ -331,6 +331,53 @@ py_list()
     };
 }
 
+template<class T>
+constexpr py_dict_v
+py_dict()
+{
+    return {
+            sizeof(T),
+            offsetof(T, ma_keys),
+            offsetof(T, ma_values),
+    };
+}
+
+template<class T>
+constexpr py_dictkeys_v
+py_dictkeys()
+{
+    return {
+            sizeof(T),
+            offsetof(T, dk_size),
+            {},
+            offsetof(T, dk_nentries),
+            offsetof(T, dk_indices),
+    };
+}
+
+template<>
+constexpr py_dictkeys_v
+py_dictkeys<Python3_11::PyDictKeysObject>()
+{
+    return {
+            sizeof(Python3_11::PyDictKeysObject),
+            offsetof(Python3_11::PyDictKeysObject, dk_log2_size),
+            offsetof(Python3_11::PyDictKeysObject, dk_kind),
+            offsetof(Python3_11::PyDictKeysObject, dk_nentries),
+            offsetof(Python3_11::PyDictKeysObject, dk_indices),
+    };
+}
+
+template<class T>
+constexpr py_dictvalues_v
+py_dictvalues()
+{
+    return {
+            sizeof(T),
+            offsetof(T, values),
+    };
+}
+
 template<class T>
 constexpr py_float_v
 py_float()
@@ -357,6 +404,9 @@ py_long()
 python_v python_v2 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        {},
+        {},
+        {},
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         {},
@@ -373,6 +423,9 @@ python_v python_v2 = {
 python_v python_v3_3 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -389,6 +442,9 @@ python_v python_v3_3 = {
 python_v python_v3_4 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -405,6 +461,9 @@ python_v python_v3_4 = {
 python_v python_v3_6 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -421,6 +480,9 @@ python_v python_v3_6 = {
 python_v python_v3_7 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -439,6 +501,9 @@ python_v python_v3_7 = {
 python_v python_v3_8 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -457,6 +522,9 @@ python_v python_v3_8 = {
 python_v python_v3_9 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -475,6 +543,9 @@ python_v python_v3_9 = {
 python_v python_v3_10 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_3::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -493,6 +564,9 @@ python_v python_v3_10 = {
 python_v python_v3_11 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_11::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3::PyUnicodeObject>(),
@@ -512,6 +586,9 @@ python_v python_v3_11 = {
 python_v python_v3_12 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_11::PyDictKeysObject>(),
+        py_dictvalues<Python3::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3_12::PyUnicodeObject>(),
@@ -531,6 +608,9 @@ python_v python_v3_12 = {
 python_v python_v3_13 = {
         py_tuple<PyTupleObject>(),
         py_list<PyListObject>(),
+        py_dict<Python3::PyDictObject>(),
+        py_dictkeys<Python3_11::PyDictKeysObject>(),
+        py_dictvalues<Python3_13::PyDictValuesObject>(),
         py_float<PyFloatObject>(),
         py_long<_PyLongObject>(),
         py_unicode<Python3_12::PyUnicodeObject>(),

diff --git a/src/pystack/_pystack/version.h b/src/pystack/_pystack/version.h
@@ -34,6 +34,31 @@ struct py_list_v
     FieldOffset<PyObject**> o_ob_item;
 };
 
+struct py_dict_v
+{
+    typedef Python3::PyDictObject Structure;
+    ssize_t size;
+    FieldOffset<remote_addr_t> o_ma_keys;
+    FieldOffset<remote_addr_t> o_ma_values;
+};
+
+struct py_dictkeys_v
+{
+    typedef PyDictKeysObject Structure;
+    ssize_t size;
+    FieldOffset<Py_ssize_t> o_dk_size;
+    FieldOffset<uint8_t> o_dk_kind;
+    FieldOffset<Py_ssize_t> o_dk_nentries;
+    FieldOffset<char[1]> o_dk_indices;
+};
+
+struct py_dictvalues_v
+{
+    typedef PyDictValuesObject Structure;
+    ssize_t size;
+    FieldOffset<remote_addr_t[1]> o_values;
+};
+
 struct py_float_v
 {
     typedef PyFloatObject Structure;
@@ -214,6 +239,9 @@ struct python_v
 {
     py_tuple_v py_tuple;
     py_list_v py_list;
+    py_dict_v py_dict;
+    py_dictkeys_v py_dictkeys;
+    py_dictvalues_v py_dictvalues;
     py_float_v py_float;
     py_long_v py_long;
     py_unicode_v py_unicode;
@@ -240,6 +268,9 @@ struct python_v
 
 define_python_v_get_specialization(py_tuple);
 define_python_v_get_specialization(py_list);
+define_python_v_get_specialization(py_dict);
+define_python_v_get_specialization(py_dictkeys);
+define_python_v_get_specialization(py_dictvalues);
 define_python_v_get_specialization(py_float);
 define_python_v_get_specialization(py_long);
 define_python_v_get_specialization(py_unicode);