refactor: expand unsafe and better document it

Author: morrisonlevi

libdd-profiling/benches/add_samples.rs

@@ -113,7 +113,7 @@ pub fn bench_add_sample_vs_add2(c: &mut Criterion) {
             })
             .unwrap();
         Frame2 {
-            function: unsafe { core::mem::transmute::<SetId<Function>, FunctionId2>(set_id) },
+            function: FunctionId2::from(set_id),
             line_number: f.line_number,
         }
     });
@@ -147,7 +147,11 @@ pub fn bench_add_sample_vs_add2(c: &mut Criterion) {
             for _ in 0..1000 {
                 // Provide an empty iterator for labels conversion path
                 let labels_iter = std::iter::empty::<anyhow::Result<api2::Label>>();
-                black_box(profile.try_add_sample2(&locations, &values, labels_iter, None)).unwrap();
+                // SAFETY: all ids come from the profile's dictionary.
+                black_box(unsafe {
+                    profile.try_add_sample2(&locations, &values, labels_iter, None)
+                })
+                .unwrap();
             }
             black_box(profile.only_for_testing_num_aggregated_samples())
         })

libdd-profiling/src/internal/profile/mod.rs

@@ -173,7 +173,15 @@ impl Profile {
     }
 
     /// Tries to add a sample using `api2` structures.
-    pub fn try_add_sample2<'a, L: ExactSizeIterator<Item = anyhow::Result<api2::Label<'a>>>>(
+    ///
+    /// # Safety
+    ///
+    /// All MappingId2, FunctionId2, and StringId2 values should be coming
+    /// from the same profiles dictionary used by this profile internally.
+    pub unsafe fn try_add_sample2<
+        'a,
+        L: ExactSizeIterator<Item = anyhow::Result<api2::Label<'a>>>,
+    >(
         &mut self,
         locations: &[api2::Location2],
         values: &[i64],
@@ -2784,17 +2792,23 @@ mod api_tests {
 
         // Add sample without labels
         let labels_iter = std::iter::empty::<anyhow::Result<api2::Label>>();
-        profile
-            .try_add_sample2(&locations, &values, labels_iter, None)
-            .expect("add to succeed");
+        // SAFETY: adding ids from the correct ProfilesDictionary.
+        unsafe {
+            profile
+                .try_add_sample2(&locations, &values, labels_iter, None)
+                .expect("add to succeed");
+        }
 
         assert_eq!(profile.only_for_testing_num_aggregated_samples(), 1);
 
         // Add same sample again - should aggregate
         let labels_iter = std::iter::empty::<anyhow::Result<api2::Label>>();
-        profile
-            .try_add_sample2(&locations, &values, labels_iter, None)
-            .expect("add to succeed");
+        // SAFETY: adding ids from the correct ProfilesDictionary.
+        unsafe {
+            profile
+                .try_add_sample2(&locations, &values, labels_iter, None)
+                .expect("add to succeed");
+        }
 
         // Still 1 sample because it aggregated
         assert_eq!(profile.only_for_testing_num_aggregated_samples(), 1);
@@ -2804,19 +2818,25 @@ mod api_tests {
         let label_value = "worker-1";
 
         let labels_iter = std::iter::once(Ok(api2::Label::str(label_key, label_value)));
-        profile
-            .try_add_sample2(&locations, &values, labels_iter, None)
-            .expect("add with label to succeed");
+        // SAFETY: adding ids from the correct ProfilesDictionary.
+        unsafe {
+            profile
+                .try_add_sample2(&locations, &values, labels_iter, None)
+                .expect("add with label to succeed");
+        }
 
         // Should be 2 samples now (different label set)
         assert_eq!(profile.only_for_testing_num_aggregated_samples(), 2);
 
         // Test with numeric label
         let thread_id_key = dict.try_insert_str2("thread_id_num").unwrap();
         let labels_iter = std::iter::once(Ok(api2::Label::num(thread_id_key, 42, "")));
-        profile
-            .try_add_sample2(&locations, &values, labels_iter, None)
-            .expect("add with numeric label to succeed");
+        // SAFETY: adding ids from the correct ProfilesDictionary.
+        unsafe {
+            profile
+                .try_add_sample2(&locations, &values, labels_iter, None)
+                .expect("add with numeric label to succeed");
+        }
 
         // Should be 3 samples now
         assert_eq!(profile.only_for_testing_num_aggregated_samples(), 3);

libdd-profiling/src/internal/profile/profiles_dictionary_translator.rs

@@ -10,14 +10,30 @@ use anyhow::Context;
 use indexmap::map::Entry;
 use std::ptr::NonNull;
 
+/// Translates IDs from a [`ProfilesDictionary`] into the IDs used by the
+/// current Profile's internal collections.
+///
+/// # Safety
+///
+/// All IDs passed to the translate methods (translate_function,
+/// translate_mapping, translate_string) MUST have been created by the same
+/// ProfilesDictionary that this translator wraps.
 pub struct ProfilesDictionaryTranslator {
     pub profiles_dictionary: crate::profiles::collections::Arc<ProfilesDictionary>,
     pub mappings: FxIndexMap<SetId<dt::Mapping>, MappingId>,
     pub functions: FxIndexMap<Option<SetId<dt::Function>>, FunctionId>,
     pub strings: FxIndexMap<StringRef, StringId>,
 }
 
-// SAFETY: the profiles_dictionary keeps the storage for Ids alive.
+// SAFETY: ProfilesDictionaryTranslator is Send because:
+// 1. The profiles_dictionary Arc ensures the underlying storage remains alive and valid for the
+//    lifetime of this translator, and Arc<T> is Send when T is Send + Sync. ProfilesDictionary is
+//    Send + Sync.
+// 2. SetId<T> and StringRef are non-owning handles (thin pointers) to immutable data in the
+//    ProfilesDictionary's concurrent collections, which use arena allocation with stable addresses.
+//    The Arc protects this data, making the pointers safe to send across threads.
+// 3. FxIndexMap<K, V> is Send when K and V are Send. The keys (SetId, StringRef) and values
+//    (MappingId, FunctionId, StringId) are all Copy types that are Send.
 unsafe impl Send for ProfilesDictionaryTranslator {}
 
 impl ProfilesDictionaryTranslator {
@@ -32,7 +48,14 @@ impl ProfilesDictionaryTranslator {
         }
     }
 
-    pub fn translate_function(
+    /// Translates a FunctionId2 from the ProfilesDictionary into a FunctionId
+    /// for this profile's StringTable.
+    ///
+    /// # Safety
+    ///
+    /// The `id2` must have been created by `self.profiles_dictionary`, and
+    /// the strings must also live in the same dictionary.
+    pub unsafe fn translate_function(
         &mut self,
         functions: &mut impl Dedup<Function>,
         string_table: &mut StringTable,
@@ -55,12 +78,19 @@ impl ProfilesDictionaryTranslator {
                     return Ok(*internal);
                 }
 
-                // SAFETY: todo
+                // SAFETY: This is safe if `id2` (the FunctionId2) was created by
+                // `self.profiles_dictionary`, which is a precondition of calling
+                // this method.
                 let function = unsafe { *self.profiles_dictionary.functions().get(set_id) };
-                let function = Function {
-                    name: self.translate_string(string_table, function.name)?,
-                    system_name: self.translate_string(string_table, function.system_name)?,
-                    filename: self.translate_string(string_table, function.file_name)?,
+                // SAFETY: safe if the strings were made by
+                // `self.profiles_dictionary`, which is a precondition of
+                // calling this method.
+                let function = unsafe {
+                    Function {
+                        name: self.translate_string(string_table, function.name)?,
+                        system_name: self.translate_string(string_table, function.system_name)?,
+                        filename: self.translate_string(string_table, function.file_name)?,
+                    }
                 };
                 (Some(set_id), function)
             }
@@ -76,7 +106,14 @@ impl ProfilesDictionaryTranslator {
         Ok(internal_id)
     }
 
-    pub fn translate_mapping(
+    /// Translates a MappingId2 from the ProfilesDictionary into a MappingId
+    /// for this profile's internal collections.
+    ///
+    /// # Safety
+    ///
+    /// The `id2` must have been created by `self.profiles_dictionary`, and
+    /// the strings must also live in the same dictionary.
+    pub unsafe fn translate_mapping(
         &mut self,
         mappings: &mut impl Dedup<Mapping>,
         string_table: &mut StringTable,
@@ -93,7 +130,9 @@ impl ProfilesDictionaryTranslator {
             return Ok(Some(*internal));
         }
 
-        // SAFETY: todo
+        // SAFETY: This is safe if `id2` (the MappingId2) was created by
+        // `self.profiles_dictionary`, which is a precondition of calling
+        // this method.
         let mapping = unsafe { *self.profiles_dictionary.mappings().get(set_id) };
         let internal = Mapping {
             memory_start: mapping.memory_start,
@@ -112,7 +151,14 @@ impl ProfilesDictionaryTranslator {
         Ok(Some(internal_id))
     }
 
-    pub fn translate_string(
+    /// Translates a StringRef from the ProfilesDictionary into a StringId
+    /// for this profile's internal string table.
+    ///
+    /// # Safety
+    ///
+    /// The `str_ref` must have been created by `self.profiles_dictionary`.
+    /// Violating this precondition results in undefined behavior.
+    pub unsafe fn translate_string(
         &mut self,
         string_table: &mut StringTable,
         str_ref: StringRef,
@@ -121,12 +167,11 @@ impl ProfilesDictionaryTranslator {
         match self.strings.entry(str_ref) {
             Entry::Occupied(o) => Ok(*o.get()),
             Entry::Vacant(v) => {
+                // SAFETY: This is safe if `str_ref` was created by
+                // `self.profiles_dictionary`, which is a precondition of calling
+                // this method.
                 let str = unsafe { self.profiles_dictionary.strings().get(str_ref) };
-                // SAFETY: we're keeping these lifetimes in sync. I think.
-                // TODO: BUT longer-term we want to avoid copying them
-                //       entirely, so this should go away.
-                let decouple_str = unsafe { core::mem::transmute::<&str, &str>(str) };
-                let internal_id = string_table.try_intern(decouple_str)?;
+                let internal_id = string_table.try_intern(str)?;
                 v.insert(internal_id);
                 Ok(internal_id)
             }