vendor/measureme/src/stringtable.rs - toolchain/rustc - Git at Google

 //! A string table implementation with a tree-like encoding.
 //!
 //! Each entry in the table represents a string and is encoded as a list of
 //! components where each component can either be
 //!
 //! 1. a string _value_ that contains actual UTF-8 string content,
 //! 2. a string _ID_ that contains a reference to another entry, or
 //! 3. a terminator tag which marks the end of a component list.
 //!
 //! The string _content_ of an entry is defined as the concatenation of the
 //! content of its components. The content of a string value is its actual
 //! UTF-8 bytes. The content of a string ID is the contents of the entry
 //! it references.
 //!
 //! The byte-level encoding of component lists uses the structure of UTF-8 in
 //! order to save space:
 //!
 //! - A valid UTF-8 codepoint never starts with the byte `0xFE`. We make use
 //!   of this fact by letting all string ID components start with this `0xFE`
 //!   prefix. Thus when we parse the contents of a value we know to stop if
 //!   we encounter this byte.
 //!
 //! - A valid UTF-8 string cannot contain the `0xFF` byte. Thus we can safely
 //!   use `0xFF` as our component list terminator.
 //!
 //! The sample composite string ["abc", ID(42), "def", TERMINATOR] would thus be
 //! encoded as:
 //!
 //! ```ignore
 //!     ['a', 'b' , 'c', 254, 42, 0, 0, 0, 'd', 'e', 'f', 255]
 //!                      ^^^^^^^^^^^^^^^^                 ^^^
 //!                 string ID with 0xFE prefix      terminator (0xFF)
 //! ```
 //!
 //! As you can see string IDs are encoded in little endian format.
 //!
 //! ----------------------------------------------------------------------------
 //!
 //! Each string in the table is referred to via a `StringId`. `StringId`s may
 //! be generated in two ways:
 //!
 //!   1. Calling `StringTableBuilder::alloc()` which returns the `StringId` for
 //!      the allocated string.
 //!   2. Calling `StringId::new_virtual()` to create a "virtual" `StringId` that
 //!      later can be mapped to an actual string via
 //!      `StringTableBuilder::map_virtual_to_concrete_string()`.
 //!
 //! String IDs allow you to deduplicate strings by allocating a string
 //! once and then referring to it by id over and over. This is a useful trick
 //! for strings which are recorded many times and it can significantly reduce
 //! the size of profile trace files.
 //!
 //! `StringId`s are partitioned according to type:
 //!
 //! > [0 .. MAX_VIRTUAL_STRING_ID, METADATA_STRING_ID, .. ]
 //!
 //! From `0` to `MAX_VIRTUAL_STRING_ID` are the allowed values for virtual strings.
 //! After `MAX_VIRTUAL_STRING_ID`, there is one string id (`METADATA_STRING_ID`)
 //! which is used internally by `measureme` to record additional metadata about
 //! the profiling session. After `METADATA_STRING_ID` are all other `StringId`
 //! values.

 use crate::file_header::{
     write_file_header, FILE_MAGIC_STRINGTABLE_DATA, FILE_MAGIC_STRINGTABLE_INDEX,
 };
 use crate::serialization::Addr;
 use crate::serialization::SerializationSink;
 use std::{error::Error, sync::Arc};

 /// A `StringId` is used to identify a string in the `StringTable`. It is
 /// either a regular `StringId`, meaning that it contains the absolute address
 /// of a string within the string table data. Or it is "virtual", which means
 /// that the address it points to is resolved via the string table index data,
 /// that maps virtual `StringId`s to addresses.
 #[derive(Clone, Copy, Eq, PartialEq, Debug, Hash)]
 #[repr(C)]
 pub struct StringId(u64);

 impl StringId {
     pub const INVALID: StringId = StringId(INVALID_STRING_ID);

     #[inline]
     pub fn new(id: impl Into<u64>) -> StringId {
         StringId(id.into())
     }

     #[inline]
     pub fn new_virtual(id: impl Into<u64>) -> StringId {
         let id = id.into();
         assert!(id <= MAX_USER_VIRTUAL_STRING_ID);
         StringId(id)
     }

     #[inline]
     pub fn is_virtual(self) -> bool {
         self.0 <= METADATA_STRING_ID
     }

     #[inline]
     pub fn as_u64(self) -> u64 {
         self.0
     }

     #[inline]
     pub fn from_addr(addr: Addr) -> StringId {
         let id = addr.0.checked_add(FIRST_REGULAR_STRING_ID).unwrap();
         StringId::new(id)
     }

     #[inline]
     pub fn to_addr(self) -> Addr {
         Addr(self.0.checked_sub(FIRST_REGULAR_STRING_ID).unwrap())
     }
 }

 // See module-level documentation for more information on the encoding.
 pub const TERMINATOR: u8 = 0xFF;
 pub const STRING_REF_TAG: u8 = 0xFE;
 pub const STRING_REF_ENCODED_SIZE: usize = 9;

 /// The maximum id value a virtual string may be.
 const MAX_USER_VIRTUAL_STRING_ID: u64 = 100_000_000;

 /// The id of the profile metadata string entry.
 pub const METADATA_STRING_ID: u64 = MAX_USER_VIRTUAL_STRING_ID + 1;

 /// Some random string ID that we make sure cannot be generated or assigned to.
 const INVALID_STRING_ID: u64 = METADATA_STRING_ID + 1;

 pub const FIRST_REGULAR_STRING_ID: u64 = INVALID_STRING_ID + 1;

 /// Write-only version of the string table
 pub struct StringTableBuilder {
     data_sink: Arc<SerializationSink>,
     index_sink: Arc<SerializationSink>,
 }

 /// Anything that implements `SerializableString` can be written to a
 /// `StringTable`.
 pub trait SerializableString {
     fn serialized_size(&self) -> usize;
     fn serialize(&self, bytes: &mut [u8]);
 }

 // A single string is encoded as `[UTF-8 bytes][TERMINATOR]`
 impl SerializableString for str {
     #[inline]
     fn serialized_size(&self) -> usize {
         self.len() + // actual bytes
         1 // terminator
     }

     #[inline]
     fn serialize(&self, bytes: &mut [u8]) {
         let last_byte_index = bytes.len() - 1;
         bytes[0..last_byte_index].copy_from_slice(self.as_bytes());
         bytes[last_byte_index] = TERMINATOR;
     }
 }

 /// A single component of a string. Used for building composite table entries.
 pub enum StringComponent<'s> {
     Value(&'s str),
     Ref(StringId),
 }

 impl<'s> StringComponent<'s> {
     #[inline]
     fn serialized_size(&self) -> usize {
         match *self {
             StringComponent::Value(s) => s.len(),
             StringComponent::Ref(_) => STRING_REF_ENCODED_SIZE,
         }
     }

     #[inline]
     fn serialize<'b>(&self, bytes: &'b mut [u8]) -> &'b mut [u8] {
         match *self {
             StringComponent::Value(s) => {
                 bytes[..s.len()].copy_from_slice(s.as_bytes());
                 &mut bytes[s.len()..]
             }
             StringComponent::Ref(string_id) => {
                 // The code below assumes we use a 9-byte encoding for string
                 // refs, where the first byte is STRING_REF_TAG and the
                 // following 8 bytes are a little-endian u64 string ID value.
                 assert!(STRING_REF_ENCODED_SIZE == 9);

                 bytes[0] = STRING_REF_TAG;
                 bytes[1..9].copy_from_slice(&string_id.0.to_le_bytes());
                 &mut bytes[9..]
             }
         }
     }
 }

 impl<'a> SerializableString for [StringComponent<'a>] {
     #[inline]
     fn serialized_size(&self) -> usize {
         self.iter().map(|c| c.serialized_size()).sum::<usize>() + // size of components
         1 // terminator
     }

     #[inline]
     fn serialize(&self, mut bytes: &mut [u8]) {
         assert!(bytes.len() == self.serialized_size());
         for component in self.iter() {
             bytes = component.serialize(bytes);
         }

         // Assert that we used the exact number of bytes we anticipated.
         assert!(bytes.len() == 1);
         bytes[0] = TERMINATOR;
     }
 }

 macro_rules! impl_serializable_string_for_fixed_size {
     ($n:expr) => {
         impl<'a> SerializableString for [StringComponent<'a>; $n] {
             #[inline(always)]
             fn serialized_size(&self) -> usize {
                 (&self[..]).serialized_size()
             }

             #[inline(always)]
             fn serialize(&self, bytes: &mut [u8]) {
                 (&self[..]).serialize(bytes);
             }
         }
     };
 }

 impl_serializable_string_for_fixed_size!(0);
 impl_serializable_string_for_fixed_size!(1);
 impl_serializable_string_for_fixed_size!(2);
 impl_serializable_string_for_fixed_size!(3);
 impl_serializable_string_for_fixed_size!(4);
 impl_serializable_string_for_fixed_size!(5);
 impl_serializable_string_for_fixed_size!(6);
 impl_serializable_string_for_fixed_size!(7);
 impl_serializable_string_for_fixed_size!(8);
 impl_serializable_string_for_fixed_size!(9);
 impl_serializable_string_for_fixed_size!(10);
 impl_serializable_string_for_fixed_size!(11);
 impl_serializable_string_for_fixed_size!(12);
 impl_serializable_string_for_fixed_size!(13);
 impl_serializable_string_for_fixed_size!(14);
 impl_serializable_string_for_fixed_size!(15);
 impl_serializable_string_for_fixed_size!(16);

 fn serialize_index_entry(sink: &SerializationSink, id: StringId, addr: Addr) {
     sink.write_atomic(16, |bytes| {
         bytes[0..8].copy_from_slice(&id.0.to_le_bytes());
         bytes[8..16].copy_from_slice(&addr.0.to_le_bytes());
     });
 }

 impl StringTableBuilder {
     pub fn new(
         data_sink: Arc<SerializationSink>,
         index_sink: Arc<SerializationSink>,
     ) -> Result<StringTableBuilder, Box<dyn Error + Send + Sync>> {
         // The first thing in every stream we generate must be the stream header.
         write_file_header(&mut data_sink.as_std_write(), FILE_MAGIC_STRINGTABLE_DATA)?;
         write_file_header(&mut index_sink.as_std_write(), FILE_MAGIC_STRINGTABLE_INDEX)?;

         Ok(StringTableBuilder {
             data_sink,
             index_sink,
         })
     }

     /// Creates a mapping so that `virtual_id` will resolve to the contents of
     /// `concrete_id` when reading the string table.
     pub fn map_virtual_to_concrete_string(&self, virtual_id: StringId, concrete_id: StringId) {
         // This assertion does not use `is_virtual` on purpose because that
         // would also allow to overwrite `METADATA_STRING_ID`.
         assert!(virtual_id.0 <= MAX_USER_VIRTUAL_STRING_ID);
         serialize_index_entry(&*self.index_sink, virtual_id, concrete_id.to_addr());
     }

     pub fn bulk_map_virtual_to_single_concrete_string<I>(
         &self,
         virtual_ids: I,
         concrete_id: StringId,
     ) where
         I: Iterator<Item = StringId> + ExactSizeIterator,
     {
         // TODO: Index data encoding could have a special bulk mode that assigns
         //       multiple StringIds to the same addr, so we don't have to repeat
         //       the `concrete_id` over and over.

         type MappingEntry = [u64; 2];
         assert!(std::mem::size_of::<MappingEntry>() == 16);

         let to_addr_le = concrete_id.to_addr().0.to_le();

         let serialized: Vec<MappingEntry> = virtual_ids
             .map(|from| {
                 let id = from.0;
                 assert!(id <= MAX_USER_VIRTUAL_STRING_ID);
                 [id.to_le(), to_addr_le]
             })
             .collect();

         let num_bytes = serialized.len() * std::mem::size_of::<MappingEntry>();
         let byte_ptr = serialized.as_ptr() as *const u8;

         let bytes = unsafe { std::slice::from_raw_parts(byte_ptr, num_bytes) };

         self.index_sink.write_bytes_atomic(bytes);
     }

     pub fn alloc_metadata<STR: SerializableString + ?Sized>(&self, s: &STR) {
         let concrete_id = self.alloc(s);
         let virtual_id = StringId(METADATA_STRING_ID);
         assert!(virtual_id.is_virtual());
         serialize_index_entry(&*self.index_sink, virtual_id, concrete_id.to_addr());
     }

     pub fn alloc<STR: SerializableString + ?Sized>(&self, s: &STR) -> StringId {
         let size_in_bytes = s.serialized_size();
         let addr = self.data_sink.write_atomic(size_in_bytes, |mem| {
             s.serialize(mem);
         });

         StringId::from_addr(addr)
     }
 }
	//! A string table implementation with a tree-like encoding.
	//!
	//! Each entry in the table represents a string and is encoded as a list of
	//! components where each component can either be
	//!
	//! 1. a string _value_ that contains actual UTF-8 string content,
	//! 2. a string _ID_ that contains a reference to another entry, or
	//! 3. a terminator tag which marks the end of a component list.
	//!
	//! The string _content_ of an entry is defined as the concatenation of the
	//! content of its components. The content of a string value is its actual
	//! UTF-8 bytes. The content of a string ID is the contents of the entry
	//! it references.
	//!
	//! The byte-level encoding of component lists uses the structure of UTF-8 in
	//! order to save space:
	//!
	//! - A valid UTF-8 codepoint never starts with the byte `0xFE`. We make use
	//! of this fact by letting all string ID components start with this `0xFE`
	//! prefix. Thus when we parse the contents of a value we know to stop if
	//! we encounter this byte.
	//!
	//! - A valid UTF-8 string cannot contain the `0xFF` byte. Thus we can safely
	//! use `0xFF` as our component list terminator.
	//!
	//! The sample composite string ["abc", ID(42), "def", TERMINATOR] would thus be
	//! encoded as:
	//!
	//! ```ignore
	//! ['a', 'b' , 'c', 254, 42, 0, 0, 0, 'd', 'e', 'f', 255]
	//! ^^^^^^^^^^^^^^^^ ^^^
	//! string ID with 0xFE prefix terminator (0xFF)
	//! ```
	//!
	//! As you can see string IDs are encoded in little endian format.
	//!
	//! ----------------------------------------------------------------------------
	//!
	//! Each string in the table is referred to via a `StringId`. `StringId`s may
	//! be generated in two ways:
	//!
	//! 1. Calling `StringTableBuilder::alloc()` which returns the `StringId` for
	//! the allocated string.
	//! 2. Calling `StringId::new_virtual()` to create a "virtual" `StringId` that
	//! later can be mapped to an actual string via
	//! `StringTableBuilder::map_virtual_to_concrete_string()`.
	//!
	//! String IDs allow you to deduplicate strings by allocating a string
	//! once and then referring to it by id over and over. This is a useful trick
	//! for strings which are recorded many times and it can significantly reduce
	//! the size of profile trace files.
	//!
	//! `StringId`s are partitioned according to type:
	//!
	//! > [0 .. MAX_VIRTUAL_STRING_ID, METADATA_STRING_ID, .. ]
	//!
	//! From `0` to `MAX_VIRTUAL_STRING_ID` are the allowed values for virtual strings.
	//! After `MAX_VIRTUAL_STRING_ID`, there is one string id (`METADATA_STRING_ID`)
	//! which is used internally by `measureme` to record additional metadata about
	//! the profiling session. After `METADATA_STRING_ID` are all other `StringId`
	//! values.

	use crate::file_header::{
	write_file_header, FILE_MAGIC_STRINGTABLE_DATA, FILE_MAGIC_STRINGTABLE_INDEX,
	};
	use crate::serialization::Addr;
	use crate::serialization::SerializationSink;
	use std::{error::Error, sync::Arc};

	/// A `StringId` is used to identify a string in the `StringTable`. It is
	/// either a regular `StringId`, meaning that it contains the absolute address
	/// of a string within the string table data. Or it is "virtual", which means
	/// that the address it points to is resolved via the string table index data,
	/// that maps virtual `StringId`s to addresses.
	#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash)]
	#[repr(C)]
	pub struct StringId(u64);

	impl StringId {
	pub const INVALID: StringId = StringId(INVALID_STRING_ID);

	#[inline]
	pub fn new(id: impl Into<u64>) -> StringId {
	StringId(id.into())
	}

	#[inline]
	pub fn new_virtual(id: impl Into<u64>) -> StringId {
	let id = id.into();
	assert!(id <= MAX_USER_VIRTUAL_STRING_ID);
	StringId(id)
	}

	#[inline]
	pub fn is_virtual(self) -> bool {
	self.0 <= METADATA_STRING_ID
	}

	#[inline]
	pub fn as_u64(self) -> u64 {
	self.0
	}

	#[inline]
	pub fn from_addr(addr: Addr) -> StringId {
	let id = addr.0.checked_add(FIRST_REGULAR_STRING_ID).unwrap();
	StringId::new(id)
	}

	#[inline]
	pub fn to_addr(self) -> Addr {
	Addr(self.0.checked_sub(FIRST_REGULAR_STRING_ID).unwrap())
	}
	}

	// See module-level documentation for more information on the encoding.
	pub const TERMINATOR: u8 = 0xFF;
	pub const STRING_REF_TAG: u8 = 0xFE;
	pub const STRING_REF_ENCODED_SIZE: usize = 9;

	/// The maximum id value a virtual string may be.
	const MAX_USER_VIRTUAL_STRING_ID: u64 = 100_000_000;

	/// The id of the profile metadata string entry.
	pub const METADATA_STRING_ID: u64 = MAX_USER_VIRTUAL_STRING_ID + 1;

	/// Some random string ID that we make sure cannot be generated or assigned to.
	const INVALID_STRING_ID: u64 = METADATA_STRING_ID + 1;

	pub const FIRST_REGULAR_STRING_ID: u64 = INVALID_STRING_ID + 1;

	/// Write-only version of the string table
	pub struct StringTableBuilder {
	data_sink: Arc<SerializationSink>,
	index_sink: Arc<SerializationSink>,
	}

	/// Anything that implements `SerializableString` can be written to a
	/// `StringTable`.
	pub trait SerializableString {
	fn serialized_size(&self) -> usize;
	fn serialize(&self, bytes: &mut [u8]);
	}

	// A single string is encoded as `[UTF-8 bytes][TERMINATOR]`
	impl SerializableString for str {
	#[inline]
	fn serialized_size(&self) -> usize {
	self.len() + // actual bytes
	1 // terminator
	}

	#[inline]
	fn serialize(&self, bytes: &mut [u8]) {
	let last_byte_index = bytes.len() - 1;
	bytes[0..last_byte_index].copy_from_slice(self.as_bytes());
	bytes[last_byte_index] = TERMINATOR;
	}
	}

	/// A single component of a string. Used for building composite table entries.
	pub enum StringComponent<'s> {
	Value(&'s str),
	Ref(StringId),
	}

	impl<'s> StringComponent<'s> {
	#[inline]
	fn serialized_size(&self) -> usize {
	match *self {
	StringComponent::Value(s) => s.len(),
	StringComponent::Ref(_) => STRING_REF_ENCODED_SIZE,
	}
	}

	#[inline]
	fn serialize<'b>(&self, bytes: &'b mut [u8]) -> &'b mut [u8] {
	match *self {
	StringComponent::Value(s) => {
	bytes[..s.len()].copy_from_slice(s.as_bytes());
	&mut bytes[s.len()..]
	}
	StringComponent::Ref(string_id) => {
	// The code below assumes we use a 9-byte encoding for string
	// refs, where the first byte is STRING_REF_TAG and the
	// following 8 bytes are a little-endian u64 string ID value.
	assert!(STRING_REF_ENCODED_SIZE == 9);

	bytes[0] = STRING_REF_TAG;
	bytes[1..9].copy_from_slice(&string_id.0.to_le_bytes());
	&mut bytes[9..]
	}
	}
	}
	}

	impl<'a> SerializableString for [StringComponent<'a>] {
	#[inline]
	fn serialized_size(&self) -> usize {
	self.iter().map(\|c\| c.serialized_size()).sum::<usize>() + // size of components
	1 // terminator
	}

	#[inline]
	fn serialize(&self, mut bytes: &mut [u8]) {
	assert!(bytes.len() == self.serialized_size());
	for component in self.iter() {
	bytes = component.serialize(bytes);
	}

	// Assert that we used the exact number of bytes we anticipated.
	assert!(bytes.len() == 1);
	bytes[0] = TERMINATOR;
	}
	}

	macro_rules! impl_serializable_string_for_fixed_size {
	($n:expr) => {
	impl<'a> SerializableString for [StringComponent<'a>; $n] {
	#[inline(always)]
	fn serialized_size(&self) -> usize {
	(&self[..]).serialized_size()
	}

	#[inline(always)]
	fn serialize(&self, bytes: &mut [u8]) {
	(&self[..]).serialize(bytes);
	}
	}
	};
	}

	impl_serializable_string_for_fixed_size!(0);
	impl_serializable_string_for_fixed_size!(1);
	impl_serializable_string_for_fixed_size!(2);
	impl_serializable_string_for_fixed_size!(3);
	impl_serializable_string_for_fixed_size!(4);
	impl_serializable_string_for_fixed_size!(5);
	impl_serializable_string_for_fixed_size!(6);
	impl_serializable_string_for_fixed_size!(7);
	impl_serializable_string_for_fixed_size!(8);
	impl_serializable_string_for_fixed_size!(9);
	impl_serializable_string_for_fixed_size!(10);
	impl_serializable_string_for_fixed_size!(11);
	impl_serializable_string_for_fixed_size!(12);
	impl_serializable_string_for_fixed_size!(13);
	impl_serializable_string_for_fixed_size!(14);
	impl_serializable_string_for_fixed_size!(15);
	impl_serializable_string_for_fixed_size!(16);

	fn serialize_index_entry(sink: &SerializationSink, id: StringId, addr: Addr) {
	sink.write_atomic(16, \|bytes\| {
	bytes[0..8].copy_from_slice(&id.0.to_le_bytes());
	bytes[8..16].copy_from_slice(&addr.0.to_le_bytes());
	});
	}

	impl StringTableBuilder {
	pub fn new(
	data_sink: Arc<SerializationSink>,
	index_sink: Arc<SerializationSink>,
	) -> Result<StringTableBuilder, Box<dyn Error + Send + Sync>> {
	// The first thing in every stream we generate must be the stream header.
	write_file_header(&mut data_sink.as_std_write(), FILE_MAGIC_STRINGTABLE_DATA)?;
	write_file_header(&mut index_sink.as_std_write(), FILE_MAGIC_STRINGTABLE_INDEX)?;

	Ok(StringTableBuilder {
	data_sink,
	index_sink,
	})
	}

	/// Creates a mapping so that `virtual_id` will resolve to the contents of
	/// `concrete_id` when reading the string table.
	pub fn map_virtual_to_concrete_string(&self, virtual_id: StringId, concrete_id: StringId) {
	// This assertion does not use `is_virtual` on purpose because that
	// would also allow to overwrite `METADATA_STRING_ID`.
	assert!(virtual_id.0 <= MAX_USER_VIRTUAL_STRING_ID);
	serialize_index_entry(&*self.index_sink, virtual_id, concrete_id.to_addr());
	}

	pub fn bulk_map_virtual_to_single_concrete_string<I>(
	&self,
	virtual_ids: I,
	concrete_id: StringId,
	) where
	I: Iterator<Item = StringId> + ExactSizeIterator,
	{
	// TODO: Index data encoding could have a special bulk mode that assigns
	// multiple StringIds to the same addr, so we don't have to repeat
	// the `concrete_id` over and over.

	type MappingEntry = [u64; 2];
	assert!(std::mem::size_of::<MappingEntry>() == 16);

	let to_addr_le = concrete_id.to_addr().0.to_le();

	let serialized: Vec<MappingEntry> = virtual_ids
	.map(\|from\| {
	let id = from.0;
	assert!(id <= MAX_USER_VIRTUAL_STRING_ID);
	[id.to_le(), to_addr_le]
	})
	.collect();

	let num_bytes = serialized.len() * std::mem::size_of::<MappingEntry>();
	let byte_ptr = serialized.as_ptr() as *const u8;

	let bytes = unsafe { std::slice::from_raw_parts(byte_ptr, num_bytes) };

	self.index_sink.write_bytes_atomic(bytes);
	}

	pub fn alloc_metadata<STR: SerializableString + ?Sized>(&self, s: &STR) {
	let concrete_id = self.alloc(s);
	let virtual_id = StringId(METADATA_STRING_ID);
	assert!(virtual_id.is_virtual());
	serialize_index_entry(&*self.index_sink, virtual_id, concrete_id.to_addr());
	}

	pub fn alloc<STR: SerializableString + ?Sized>(&self, s: &STR) -> StringId {
	let size_in_bytes = s.serialized_size();
	let addr = self.data_sink.write_atomic(size_in_bytes, \|mem\| {
	s.serialize(mem);
	});

	StringId::from_addr(addr)
	}
	}