extract small_non_zero_pprof_id

profiling/src/profile/internal/function.rs

@@ -2,7 +2,7 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present Datadog, Inc.
 
 use super::super::{pprof, StringId};
-use super::{Id, Item, PprofItem};
+use super::{small_non_zero_pprof_id, Id, Item, PprofItem};
 use std::fmt::Debug;
 use std::num::NonZeroU32;
 
@@ -39,15 +39,8 @@ pub struct FunctionId(NonZeroU32);
 impl Id for FunctionId {
     type RawId = u64;
 
-    fn from_offset(v: usize) -> Self {
-        let index: u32 = v.try_into().expect("FunctionId to fit into a u32");
-
-        // PProf reserves function 0.
-        // Both this, and the serialization of the table, add 1 to avoid the 0 element
-        let index = index.checked_add(1).expect("FunctionId to fit into a u32");
-        // Safety: the `checked_add(1).expect(...)` guards this from ever being zero.
-        let index = unsafe { NonZeroU32::new_unchecked(index) };
-        Self(index)
+    fn from_offset(offset: usize) -> Self {
+        Self(small_non_zero_pprof_id(offset).expect("FunctionId to fit into a u32"))
     }
 
     fn to_raw_id(&self) -> Self::RawId {

profiling/src/profile/internal/location.rs

@@ -2,8 +2,7 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present Datadog, Inc.
 
 use super::super::{pprof, MappingId};
-use super::{Id, Item, Line, PprofItem};
-use std::fmt::Debug;
+use super::{small_non_zero_pprof_id, Id, Item, Line, PprofItem};
 use std::num::NonZeroU32;
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd, Ord)]
@@ -39,15 +38,8 @@ pub struct LocationId(NonZeroU32);
 impl Id for LocationId {
     type RawId = u64;
 
-    fn from_offset(v: usize) -> Self {
-        let index: u32 = v.try_into().expect("LocationId to fit into a u32");
-
-        // PProf reserves location 0.
-        // Both this, and the serialization of the table, add 1 to avoid the 0 element
-        let index = index.checked_add(1).expect("LocationId to fit into a u32");
-        // Safety: the `checked_add(1).expect(...)` guards this from ever being zero.
-        let index = unsafe { NonZeroU32::new_unchecked(index) };
-        Self(index)
+    fn from_offset(offset: usize) -> Self {
+        Self(small_non_zero_pprof_id(offset).expect("LocationId to fit into a u32"))
     }
 
     fn to_raw_id(&self) -> Self::RawId {

profiling/src/profile/internal/mapping.rs

@@ -1,7 +1,7 @@
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2023-Present Datadog, Inc.
 
-use super::{Id, Item, PprofItem};
+use super::{small_non_zero_pprof_id, Id, Item, PprofItem};
 use crate::profile::{pprof, StringId};
 use std::num::NonZeroU32;
 
@@ -52,15 +52,8 @@ pub struct MappingId(NonZeroU32);
 impl Id for MappingId {
     type RawId = u64;
 
-    fn from_offset(v: usize) -> Self {
-        let index: u32 = v.try_into().expect("MappingId to fit into a u32");
-
-        // PProf reserves location 0.
-        // Both this, and the serialization of the table, add 1 to avoid the 0 element
-        let index = index.checked_add(1).expect("MappingId to fit into a u32");
-        // Safety: the `checked_add(1).expect(...)` guards this from ever being zero.
-        let index = unsafe { NonZeroU32::new_unchecked(index) };
-        Self(index)
+    fn from_offset(offset: usize) -> Self {
+        Self(small_non_zero_pprof_id(offset).expect("MappingId to fit into a u32"))
     }
 
     fn to_raw_id(&self) -> Self::RawId {

profiling/src/profile/internal/mod.rs

@@ -20,6 +20,7 @@ pub use string::*;
 pub use value_type::*;
 
 use std::hash::Hash;
+use std::num::NonZeroU32;
 
 pub trait Id: Copy + Eq + Hash {
     type RawId;
@@ -54,3 +55,11 @@ pub trait PprofItem: Item {
 
     fn to_pprof(&self, id: Self::Id) -> Self::PprofMessage;
 }
+
+#[inline]
+fn small_non_zero_pprof_id(offset: usize) -> Option<NonZeroU32> {
+    let small: u32 = offset.try_into().ok()?;
+    let non_zero = small.checked_add(1)?;
+    // Safety: the `checked_add(1)?` guards this from ever being zero.
+    Some(unsafe { NonZeroU32::new_unchecked(non_zero) })
+}