Some fxprof-processed-profile API changes

fxprof-processed-profile/Cargo.toml

@@ -2,7 +2,7 @@
 name = "fxprof-processed-profile"
 version = "0.8.1"
 edition = "2021"
-rust-version = "1.60" # needed by bytesize
+rust-version = "1.63" # needed by indexmap
 authors = ["Markus Stange <mstange.moz@gmail.com>"]
 license = "MIT OR Apache-2.0"
 description = "Create profiles in the Firefox Profiler's processed profile JSON format."
@@ -16,6 +16,7 @@ serde = "1.0.204"
 serde_derive = "1.0.188"
 debugid = "0.8.0"
 rustc-hash = "2"
+indexmap = { version = "2.7", features = ["serde"] }
 
 [dev-dependencies]
 assert-json-diff = "2.0.1"

fxprof-processed-profile/README.md

@@ -20,10 +20,10 @@ let process = profile.add_process("App process", 54132, Timestamp::from_millis_s
 let thread = profile.add_thread(process, 54132000, Timestamp::from_millis_since_reference(0.0), true);
 profile.set_thread_name(thread, "Main thread");
 let stack_frames = vec![
-    FrameInfo { frame: Frame::Label(profile.intern_string("Root node")), category_pair: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() },
-    FrameInfo { frame: Frame::Label(profile.intern_string("First callee")), category_pair: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() }
+    FrameInfo { frame: Frame::Label(profile.handle_for_string("Root node")), subcategory: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() },
+    FrameInfo { frame: Frame::Label(profile.handle_for_string("First callee")), subcategory: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() }
 ];
-let stack = profile.intern_stack_frames(thread, stack_frames.into_iter());
+let stack = profile.handle_for_stack_frames(thread, stack_frames.into_iter());
 profile.add_sample(thread, Timestamp::from_millis_since_reference(0.0), stack, CpuDelta::ZERO, 1);
 
 let writer = std::io::BufWriter::new(output_file);

fxprof-processed-profile/src/category.rs

@@ -1,10 +1,38 @@
+use std::hash::Hash;
+
+use indexmap::Equivalent;
 use serde::ser::{Serialize, SerializeMap, Serializer};
 
+use crate::Profile;
+
 use super::category_color::CategoryColor;
+use super::fast_hash_map::FastIndexSet;
+
+/// Implemented by [`Category`], [`Subcategory`], [`CategoryHandle`] and [`SubcategoryHandle`].
+pub trait IntoSubcategoryHandle {
+    /// Returns the corresponding [`SubcategoryHandle`].
+    fn into_subcategory_handle(self, profile: &mut Profile) -> SubcategoryHandle;
+}
+
+/// A profiling category. Has a name and a color.
+///
+/// Used to categorize stack frames and markers in the front-end. The category's
+/// color is used in the activity graph and in the call tree, and in a few other places.
+#[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
+pub struct Category<'a>(pub &'a str, pub CategoryColor);
+
+impl IntoSubcategoryHandle for Category<'_> {
+    fn into_subcategory_handle(self, profile: &mut Profile) -> SubcategoryHandle {
+        let category_handle = profile.handle_for_category(self);
+        category_handle.into()
+    }
+}
 
-/// A profiling category, can be set on stack frames and markers as part of a [`CategoryPairHandle`].
+/// The handle for a [`Category`], obtained from [`Profile::handle_for_category`](crate::Profile::handle_for_category).
 ///
-/// Categories can be created with [`Profile::add_category`](crate::Profile::add_category).
+/// Storing and reusing the handle avoids repeated lookups and can improve performance.
+///
+/// The handle is specific to a [`Profile`] instance and cannot be reused across profiles.
 #[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct CategoryHandle(pub(crate) u16);
 
@@ -13,15 +41,33 @@ impl CategoryHandle {
     pub const OTHER: Self = CategoryHandle(0);
 }
 
+impl IntoSubcategoryHandle for CategoryHandle {
+    fn into_subcategory_handle(self, _profile: &mut Profile) -> SubcategoryHandle {
+        self.into()
+    }
+}
+
 impl Serialize for CategoryHandle {
     fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         self.0.serialize(serializer)
     }
 }
 
-/// A profiling subcategory, can be set on stack frames and markers as part of a [`CategoryPairHandle`].
+/// A named subcategory of a [`Category`], for fine-grained annotation of stack frames.
 ///
-/// Subategories can be created with [`Profile::add_subcategory`](crate::Profile::add_subcategory).
+/// If you don't need named subcategories, you can just pass a [`Category`] or a
+/// [`CategoryHandle`] in any place where an [`IntoSubcategoryHandle`] is expected;
+/// this will give you the category's default subcategory.
+pub struct Subcategory<'a>(pub Category<'a>, pub &'a str);
+
+impl IntoSubcategoryHandle for Subcategory<'_> {
+    fn into_subcategory_handle(self, profile: &mut Profile) -> SubcategoryHandle {
+        let Subcategory(category, subcategory_name) = self;
+        let category_handle = profile.handle_for_category(category);
+        profile.handle_for_subcategory(category_handle, subcategory_name)
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct SubcategoryIndex(pub u16);
 
@@ -30,16 +76,30 @@ impl SubcategoryIndex {
     pub const OTHER: Self = SubcategoryIndex(0);
 }
 
-/// A profiling category pair, consisting of a category and an optional subcategory. Can be set on stack frames and markers.
+/// A handle for a [`Subcategory`], or for the default subcategory of a [`CategoryHandle`].
+///
+/// Used to annotate stack frames.
+///
+/// Every [`CategoryHandle`] can be turned into a [`SubcategoryHandle`] by calling `.into()` -
+/// this will give you the default subcategory of that category.
 ///
-/// Category pairs can be created with [`Profile::add_subcategory`](crate::Profile::add_subcategory)
-/// and from a [`CategoryHandle`].
+/// Subcategory handles for named subcategories can be obtained from
+/// [`Profile::handle_for_subcategory`](crate::Profile::handle_for_subcategory).
+/// Storing and reusing the handle avoids repeated lookups and can improve performance.
+///
+/// The handle is specific to a Profile instance and cannot be reused across profiles.
 #[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
-pub struct CategoryPairHandle(pub(crate) CategoryHandle, pub(crate) SubcategoryIndex);
+pub struct SubcategoryHandle(pub(crate) CategoryHandle, pub(crate) SubcategoryIndex);
 
-impl From<CategoryHandle> for CategoryPairHandle {
+impl IntoSubcategoryHandle for SubcategoryHandle {
+    fn into_subcategory_handle(self, _profile: &mut Profile) -> SubcategoryHandle {
+        self
+    }
+}
+
+impl From<CategoryHandle> for SubcategoryHandle {
     fn from(category: CategoryHandle) -> Self {
-        CategoryPairHandle(category, SubcategoryIndex::OTHER)
+        SubcategoryHandle(category, SubcategoryIndex::OTHER)
     }
 }
 
@@ -48,24 +108,61 @@ impl From<CategoryHandle> for CategoryPairHandle {
 pub struct InternalCategory {
     name: String,
     color: CategoryColor,
-    subcategories: Vec<String>,
+    subcategories: FastIndexSet<String>,
+}
+
+impl Hash for InternalCategory {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.as_category().hash(state)
+    }
+}
+
+impl Equivalent<Category<'_>> for InternalCategory {
+    fn equivalent(&self, key: &Category<'_>) -> bool {
+        &self.as_category() == key
+    }
 }
 
+impl Equivalent<InternalCategory> for Category<'_> {
+    fn equivalent(&self, key: &InternalCategory) -> bool {
+        self == &key.as_category()
+    }
+}
+
+impl PartialEq for InternalCategory {
+    fn eq(&self, other: &Self) -> bool {
+        self.as_category() == other.as_category()
+    }
+}
+
+impl Eq for InternalCategory {}
+
 impl InternalCategory {
-    pub fn new(name: String, color: CategoryColor) -> Self {
-        let subcategories = vec!["Other".to_string()];
+    pub fn new(name: &str, color: CategoryColor) -> Self {
+        let mut subcategories = FastIndexSet::default();
+        subcategories.insert("Other".to_string());
         Self {
-            name,
+            name: name.to_string(),
             color,
             subcategories,
         }
     }
 
-    /// Add a subcategory to this category.
-    pub fn add_subcategory(&mut self, subcategory_name: String) -> SubcategoryIndex {
-        let subcategory_index = SubcategoryIndex(u16::try_from(self.subcategories.len()).unwrap());
-        self.subcategories.push(subcategory_name);
-        subcategory_index
+    /// Get or create a subcategory to this category.
+    pub fn index_for_subcategory(&mut self, subcategory_name: &str) -> SubcategoryIndex {
+        let index = self
+            .subcategories
+            .get_index_of(subcategory_name)
+            .unwrap_or_else(|| {
+                self.subcategories
+                    .insert_full(subcategory_name.to_owned())
+                    .0
+            });
+        SubcategoryIndex(u16::try_from(index).unwrap())
+    }
+
+    pub fn as_category(&self) -> Category<'_> {
+        Category(&self.name, self.color)
     }
 }
 

fxprof-processed-profile/src/fast_hash_map.rs

@@ -1,3 +1,5 @@
+use indexmap::IndexSet;
 use rustc_hash::FxHashMap;
 
 pub type FastHashMap<K, V> = FxHashMap<K, V>;
+pub type FastIndexSet<V> = IndexSet<V, rustc_hash::FxBuildHasher>;

fxprof-processed-profile/src/frame.rs

@@ -1,6 +1,6 @@
 use bitflags::bitflags;
 
-use crate::category::CategoryPairHandle;
+use crate::category::SubcategoryHandle;
 use crate::global_lib_table::LibraryHandle;
 use crate::profile::StringHandle;
 
@@ -46,7 +46,7 @@ pub enum Frame {
     /// has already been resolved to a `LibraryHandle`.
     RelativeAddressFromAdjustedReturnAddress(LibraryHandle, u32),
     /// A string, containing an index returned by
-    /// [`Profile::intern_string`](crate::Profile::intern_string).
+    /// [`Profile::handle_for_string`](crate::Profile::handle_for_string).
     Label(StringHandle),
 }
 
@@ -55,8 +55,8 @@ pub enum Frame {
 pub struct FrameInfo {
     /// The absolute address or label of this frame.
     pub frame: Frame,
-    /// The category pair of this frame.
-    pub category_pair: CategoryPairHandle,
+    /// The subcategory of this frame.
+    pub subcategory: SubcategoryHandle,
     /// The flags of this frame. Use `FrameFlags::empty()` if unsure.
     pub flags: FrameFlags,
 }

fxprof-processed-profile/src/frame_table.rs

@@ -1,6 +1,6 @@
 use serde::ser::{Serialize, SerializeMap, SerializeSeq, Serializer};
 
-use crate::category::{CategoryHandle, CategoryPairHandle, SubcategoryIndex};
+use crate::category::{CategoryHandle, SubcategoryHandle, SubcategoryIndex};
 use crate::fast_hash_map::FastHashMap;
 use crate::frame::FrameFlags;
 use crate::func_table::{FuncIndex, FuncTable};
@@ -84,7 +84,7 @@ impl FrameTable {
                 };
                 let func_index =
                     func_table.index_for_func(location_string_index, resource, frame.flags);
-                let CategoryPairHandle(category, subcategory) = frame.category_pair;
+                let SubcategoryHandle(category, subcategory) = frame.subcategory;
                 addresses.push(address);
                 categories.push(category);
                 subcategories.push(subcategory);
@@ -134,7 +134,7 @@ impl Serialize for SerializableFrameTableAddressColumn<'_> {
 #[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct InternalFrame {
     pub location: InternalFrameLocation,
-    pub category_pair: CategoryPairHandle,
+    pub subcategory: SubcategoryHandle,
     pub flags: FrameFlags,
 }
 

fxprof-processed-profile/src/lib.rs

@@ -19,10 +19,10 @@
 //! let thread = profile.add_thread(process, 54132000, Timestamp::from_millis_since_reference(0.0), true);
 //! profile.set_thread_name(thread, "Main thread");
 //! let stack_frames = vec![
-//!     FrameInfo { frame: Frame::Label(profile.intern_string("Root node")), category_pair: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() },
-//!     FrameInfo { frame: Frame::Label(profile.intern_string("First callee")), category_pair: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() }
+//!     FrameInfo { frame: Frame::Label(profile.handle_for_string("Root node")), subcategory: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() },
+//!     FrameInfo { frame: Frame::Label(profile.handle_for_string("First callee")), subcategory: CategoryHandle::OTHER.into(), flags: FrameFlags::empty() }
 //! ];
-//! let stack = profile.intern_stack_frames(thread, stack_frames.into_iter());
+//! let stack = profile.handle_for_stack_frames(thread, stack_frames.into_iter());
 //! profile.add_sample(thread, Timestamp::from_millis_since_reference(0.0), stack, CpuDelta::ZERO, 1);
 //!
 //! let writer = std::io::BufWriter::new(output_file);
@@ -59,7 +59,9 @@ mod thread;
 mod thread_string_table;
 mod timestamp;
 
-pub use category::{CategoryHandle, CategoryPairHandle};
+pub use category::{
+    Category, CategoryHandle, IntoSubcategoryHandle, Subcategory, SubcategoryHandle,
+};
 pub use category_color::CategoryColor;
 pub use counters::CounterHandle;
 pub use cpu_delta::CpuDelta;

fxprof-processed-profile/src/marker_table.rs

@@ -58,7 +58,6 @@ impl MarkerTable {
         schema: &InternalMarkerSchema,
         marker: T,
         timing: MarkerTiming,
-        category: CategoryHandle,
         thread_string_table: &mut ThreadStringTable,
         global_string_table: &mut GlobalStringTable,
     ) -> MarkerHandle {
@@ -68,7 +67,7 @@ impl MarkerTable {
             MarkerTiming::IntervalStart(s) => (Some(s), None, Phase::IntervalStart),
             MarkerTiming::IntervalEnd(e) => (None, Some(e), Phase::IntervalEnd),
         };
-        self.marker_categories.push(category);
+        self.marker_categories.push(schema.category());
         self.marker_name_string_indexes.push(name_string_index);
         self.marker_starts.push(s);
         self.marker_ends.push(e);