rustc_middle/mir/interpret/allocation/

provenance_map.rs

//! Store the provenance for each byte in the range, with a more efficient
//! representation for the common case where PTR_SIZE consecutive bytes have the same provenance.

use std::cmp;
use std::ops::{Range, RangeBounds};

use rustc_abi::{HasDataLayout, Size};
use rustc_data_structures::sorted_map::SortedMap;
use rustc_macros::HashStable;
use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
use tracing::trace;

use super::{AllocRange, CtfeProvenance, Provenance, alloc_range};
use crate::mir::interpret::{AllocError, AllocResult};

/// A pointer fragment represents one byte of a pointer.
/// If the bytes are re-assembled in their original order, the pointer can be used again.
/// Wildcard provenance is allowed to have index 0 everywhere.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
#[derive(HashStable)]
pub struct PointerFrag<Prov> {
    /// The position of this fragment inside the pointer (in `0..8`).
    pub idx: u8,
    /// The provenance of the pointer this is a fragment of.
    pub prov: Prov,
    /// The raw bytes of the pointer this is a fragment of.
    /// This is taken as a direct subslice of the raw allocation data, so we don't have to worry
    /// about endianness. If the pointer size is less than 8, only the first N bytes of this are
    /// ever non-zero.
    pub bytes: [u8; 8],
}

/// Stores the provenance information of pointers stored in memory.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
#[derive(HashStable)]
pub struct ProvenanceMap<Prov = CtfeProvenance> {
    /// `Provenance` in this map applies from the given offset for an entire pointer-size worth of
    /// bytes. Two entries in this map are always at least a pointer size apart.
    ptrs: SortedMap<Size, Prov>,
    /// This stores byte-sized provenance fragments.
    bytes: Option<Box<SortedMap<Size, PointerFrag<Prov>>>>,
}

// These impls are generic over `Prov` since `CtfeProvenance` is only decodable/encodable
// for some particular `D`/`S`.
impl<D: Decoder, Prov: Provenance + Decodable<D>> Decodable<D> for ProvenanceMap<Prov> {
    fn decode(d: &mut D) -> Self {
        // `bytes` is not in the serialized format
        Self { ptrs: Decodable::decode(d), bytes: None }
    }
}
impl<S: Encoder, Prov: Provenance + Encodable<S>> Encodable<S> for ProvenanceMap<Prov> {
    fn encode(&self, s: &mut S) {
        let Self { ptrs, bytes } = self;
        assert!(bytes.is_none()); // interning refuses allocations with pointer fragments
        ptrs.encode(s)
    }
}

impl<Prov> ProvenanceMap<Prov> {
    pub fn new() -> Self {
        ProvenanceMap { ptrs: SortedMap::new(), bytes: None }
    }

    /// The caller must guarantee that the given provenance list is already sorted
    /// by offset and contain no duplicates.
    pub fn from_presorted_ptrs(r: Vec<(Size, Prov)>) -> Self {
        ProvenanceMap { ptrs: SortedMap::from_presorted_elements(r), bytes: None }
    }
}

impl ProvenanceMap {
    /// Give access to the ptr-sized provenances (which can also be thought of as relocations, and
    /// indeed that is how codegen treats them).
    ///
    /// Only use on interned allocations, as other allocations may have per-byte provenance!
    #[inline]
    pub fn ptrs(&self) -> &SortedMap<Size, CtfeProvenance> {
        assert!(self.bytes.is_none(), "`ptrs()` called on non-interned allocation");
        &self.ptrs
    }
}

impl<Prov: Provenance> ProvenanceMap<Prov> {
    fn adjusted_range_ptrs(range: AllocRange, cx: &impl HasDataLayout) -> Range<Size> {
        // We have to go back `pointer_size - 1` bytes, as that one would still overlap with
        // the beginning of this range.
        let adjusted_start = Size::from_bytes(
            range.start.bytes().saturating_sub(cx.data_layout().pointer_size().bytes() - 1),
        );
        adjusted_start..range.end()
    }

    /// Returns all ptr-sized provenance in the given range.
    /// If the range has length 0, returns provenance that crosses the edge between `start-1` and
    /// `start`.
    fn range_ptrs_get(&self, range: AllocRange, cx: &impl HasDataLayout) -> &[(Size, Prov)] {
        self.ptrs.range(Self::adjusted_range_ptrs(range, cx))
    }

    /// `pm.range_ptrs_is_empty(r, cx)` == `pm.range_ptrs_get(r, cx).is_empty()`, but is faster.
    fn range_ptrs_is_empty(&self, range: AllocRange, cx: &impl HasDataLayout) -> bool {
        self.ptrs.range_is_empty(Self::adjusted_range_ptrs(range, cx))
    }

    /// Check if there is ptr-sized provenance at the given index.
    /// Does not mean anything for bytewise provenance! But can be useful as an optimization.
    pub fn get_ptr(&self, offset: Size) -> Option<Prov> {
        self.ptrs.get(&offset).copied()
    }

    /// Returns all byte-wise provenance in the given range.
    fn range_bytes_get(&self, range: AllocRange) -> &[(Size, PointerFrag<Prov>)] {
        if let Some(bytes) = self.bytes.as_ref() {
            bytes.range(range.start..range.end())
        } else {
            &[]
        }
    }

    /// Same as `range_bytes_get(range).is_empty()`, but faster.
    fn range_bytes_is_empty(&self, range: AllocRange) -> bool {
        self.bytes.as_ref().is_none_or(|bytes| bytes.range_is_empty(range.start..range.end()))
    }

    /// Get the provenance of a single byte. Must only be called if there is no
    /// pointer-sized provenance here.
    pub fn get_byte(&self, offset: Size, cx: &impl HasDataLayout) -> Option<&PointerFrag<Prov>> {
        debug_assert!(self.range_ptrs_is_empty(alloc_range(offset, Size::from_bytes(1)), cx));
        self.bytes.as_ref().and_then(|b| b.get(&offset))
    }

    /// Gets the provenances of all bytes (including from pointers) in a range.
    pub fn get_range(
        &self,
        range: AllocRange,
        cx: &impl HasDataLayout,
    ) -> impl Iterator<Item = (AllocRange, Prov)> {
        let ptr_size = cx.data_layout().pointer_size();
        let ptr_provs = self
            .range_ptrs_get(range, cx)
            .iter()
            .map(move |(offset, p)| (alloc_range(*offset, ptr_size), *p));
        let byte_provs = self
            .range_bytes_get(range)
            .iter()
            .map(move |(offset, frag)| (alloc_range(*offset, Size::from_bytes(1)), frag.prov));
        ptr_provs.chain(byte_provs)
    }

    /// Attempt to merge per-byte provenance back into ptr chunks, if the right fragments
    /// sit next to each other. Return `false` if that is not possible due to partial pointers.
    pub fn merge_bytes(&mut self, cx: &impl HasDataLayout) -> bool {
        let Some(bytes) = self.bytes.as_deref_mut() else {
            return true;
        };
        let ptr_size = cx.data_layout().pointer_size();
        while let Some((offset, first_frag)) = bytes.iter().next() {
            let offset = *offset;
            // Check if this fragment starts a pointer.
            let range = offset..offset + ptr_size;
            let frags = bytes.range(range.clone());
            if frags.len() != ptr_size.bytes_usize() {
                // We can't merge this one, no point in trying to merge the rest.
                return false;
            }
            for (idx, (_offset, frag)) in frags.iter().enumerate() {
                if !(frag.prov == first_frag.prov
                    && frag.bytes == first_frag.bytes
                    && frag.idx == idx as u8)
                {
                    return false;
                }
            }
            // Looks like a pointer! Move it over to the ptr provenance map.
            self.ptrs.insert(offset, first_frag.prov);
            bytes.remove_range(range);
        }
        // We managed to convert everything into whole pointers.
        self.bytes = None;
        true
    }

    /// Try to read a pointer from the given location, possibly by loading from many per-byte
    /// provenances.
    pub fn read_ptr(&self, offset: Size, cx: &impl HasDataLayout) -> AllocResult<Option<Prov>> {
        // If there is pointer-sized provenance exactly here, we can just return that.
        if let Some(prov) = self.get_ptr(offset) {
            return Ok(Some(prov));
        }
        // The other easy case is total absence of provenance, that also always works.
        let range = alloc_range(offset, cx.data_layout().pointer_size());
        let no_ptrs = self.range_ptrs_is_empty(range, cx);
        if no_ptrs && self.range_bytes_is_empty(range) {
            return Ok(None);
        }
        // If we get here, we have to check whether we can merge per-byte provenance.
        let prov = 'prov: {
            // If there is any ptr-sized provenance overlapping with this range,
            // this is definitely mixing multiple pointers and we can bail.
            if !no_ptrs {
                break 'prov None;
            }
            // Scan all fragments, and ensure their indices, provenance, and bytes match.
            // However, we have to ignore wildcard fragments for this (this is needed for Miri's
            // native-lib mode). Therefore, we will only know the expected provenance and bytes
            // once we find the first non-wildcard fragment.
            let mut expected = None;
            for idx in Size::ZERO..range.size {
                // Ensure there is provenance here.
                let Some(frag) = self.get_byte(offset + idx, cx) else {
                    break 'prov None;
                };
                // If this is wildcard provenance, ignore this fragment.
                if Some(frag.prov) == Prov::WILDCARD {
                    continue;
                }
                // For non-wildcard fragments, the index must match.
                if u64::from(frag.idx) != idx.bytes() {
                    break 'prov None;
                }
                // If there are expectations registered, check them.
                // If not, record this fragment as setting the expectations.
                match expected {
                    Some(expected) => {
                        if (frag.prov, frag.bytes) != expected {
                            break 'prov None;
                        }
                    }
                    None => {
                        expected = Some((frag.prov, frag.bytes));
                    }
                }
            }
            // The final provenance is the expected one we found along the way, or wildcard if
            // we didn't find any.
            Some(expected.map(|(prov, _addr)| prov).or_else(|| Prov::WILDCARD).unwrap())
        };
        if prov.is_none() && !Prov::OFFSET_IS_ADDR {
            // There are some bytes with provenance here but overall the provenance does not add up.
            // We need `OFFSET_IS_ADDR` to fall back to no-provenance here; without that option, we
            // must error.
            return Err(AllocError::ReadPartialPointer(offset));
        }
        Ok(prov)
    }

    /// Returns whether this allocation has provenance overlapping with the given range.
    ///
    /// Note: this function exists to allow `range_get_provenance` to be private, in order to somewhat
    /// limit access to provenance outside of the `Allocation` abstraction.
    ///
    pub fn range_empty(&self, range: AllocRange, cx: &impl HasDataLayout) -> bool {
        self.range_ptrs_is_empty(range, cx) && self.range_bytes_is_empty(range)
    }

    /// Yields all the provenances stored in this map.
    pub fn provenances(&self) -> impl Iterator<Item = Prov> {
        let bytes = self.bytes.iter().flat_map(|b| b.values().map(|frag| frag.prov));
        self.ptrs.values().copied().chain(bytes)
    }

    pub fn insert_ptr(&mut self, offset: Size, prov: Prov, cx: &impl HasDataLayout) {
        debug_assert!(self.range_empty(alloc_range(offset, cx.data_layout().pointer_size()), cx));
        self.ptrs.insert(offset, prov);
    }

    /// Returns an iterator that yields the fragments of this pointer whose absolute positions are
    /// inside `pos_range`.
    fn ptr_fragments(
        pos_range: impl RangeBounds<Size>,
        ptr_pos: Size,
        prov: Prov,
        data_bytes: &[u8],
        ptr_size: Size,
    ) -> impl Iterator<Item = (Size, PointerFrag<Prov>)> {
        if pos_range.is_empty() {
            return either::Left(std::iter::empty());
        }
        // Read ptr_size many bytes starting at ptr_pos.
        let mut bytes = [0u8; 8];
        (&mut bytes[..ptr_size.bytes_usize()])
            .copy_from_slice(&data_bytes[ptr_pos.bytes_usize()..][..ptr_size.bytes_usize()]);
        // Yield the fragments of this pointer.
        either::Right(
            (ptr_pos..ptr_pos + ptr_size).filter(move |pos| pos_range.contains(pos)).map(
                move |pos| (pos, PointerFrag { idx: (pos - ptr_pos).bytes() as u8, bytes, prov }),
            ),
        )
    }

    /// Removes all provenance inside the given range.
    #[allow(irrefutable_let_patterns)] // these actually make the code more clear
    pub fn clear(&mut self, range: AllocRange, data_bytes: &[u8], cx: &impl HasDataLayout) {
        if range.size == Size::ZERO {
            return;
        }

        let start = range.start;
        let end = range.end();
        // Clear the bytewise part -- this is easy.
        if let Some(bytes) = self.bytes.as_mut() {
            bytes.remove_range(start..end);
        }

        // Find all provenance overlapping the given range.
        let ptrs_range = Self::adjusted_range_ptrs(range, cx);
        if self.ptrs.range_is_empty(ptrs_range.clone()) {
            // No provenance in this range, we are done. This is the common case.
            return;
        }
        let pointer_size = cx.data_layout().pointer_size();

        // This redoes some of the work of `range_is_empty`, but this path is much
        // colder than the early return above, so it's worth it.
        let ptrs = self.ptrs.range(ptrs_range.clone());

        // We need to handle clearing the provenance from parts of a pointer.
        if let &(first, prov) = ptrs.first().unwrap()
            && first < start
        {
            // Insert the remaining part in the bytewise provenance.
            let bytes = self.bytes.get_or_insert_with(Box::default);
            for (pos, frag) in Self::ptr_fragments(..start, first, prov, data_bytes, pointer_size) {
                bytes.insert(pos, frag);
            }
        }
        if let &(last, prov) = ptrs.last().unwrap()
            && last + pointer_size > end
        {
            // Insert the remaining part in the bytewise provenance.
            let bytes = self.bytes.get_or_insert_with(Box::default);
            for (pos, frag) in Self::ptr_fragments(end.., last, prov, data_bytes, pointer_size) {
                bytes.insert(pos, frag);
            }
        }

        // Forget all the provenance.
        // Since provenance do not overlap, we know that removing until `last` (exclusive) is fine,
        // i.e., this will not remove any other provenance just after the ones we care about.
        self.ptrs.remove_range(ptrs_range);
    }

    /// Overwrites all provenance in the given range with wildcard provenance.
    /// Pointers partially overwritten will have their provenances preserved
    /// bytewise on their remaining bytes.
    ///
    /// Provided for usage in Miri and panics otherwise.
    pub fn write_wildcards(
        &mut self,
        cx: &impl HasDataLayout,
        data_bytes: &[u8],
        range: AllocRange,
    ) {
        let wildcard = Prov::WILDCARD.unwrap();

        // Clear existing provenance in this range.
        self.clear(range, data_bytes, cx);

        // Make everything in the range wildcards.
        let bytes = self.bytes.get_or_insert_with(Box::default);
        for offset in range.start..range.end() {
            // The fragment index and bytes do not matter for wildcard provenance.
            bytes.insert(
                offset,
                PointerFrag { prov: wildcard, idx: Default::default(), bytes: Default::default() },
            );
        }
    }
}

/// A partial, owned list of provenance to transfer into another allocation.
///
/// Offsets are relative to the beginning of the copied range.
pub struct ProvenanceCopy<Prov> {
    ptrs: Box<[(Size, Prov)]>,
    bytes: Box<[(Size, PointerFrag<Prov>)]>,
}

impl<Prov: Provenance> ProvenanceMap<Prov> {
    pub fn prepare_copy(
        &self,
        range: AllocRange,
        data_bytes: &[u8],
        cx: &impl HasDataLayout,
    ) -> ProvenanceCopy<Prov> {
        let shift_offset = move |offset| offset - range.start;
        let ptr_size = cx.data_layout().pointer_size();

        // # Pointer-sized provenances
        // Get the provenances that are entirely within this range.
        // (Different from `range_get_ptrs` which asks if they overlap the range.)
        // Only makes sense if we are copying at least one pointer worth of bytes.
        let mut ptrs_box: Box<[_]> = Box::new([]);
        if range.size >= ptr_size {
            let adjusted_end = Size::from_bytes(range.end().bytes() - (ptr_size.bytes() - 1));
            let ptrs = self.ptrs.range(range.start..adjusted_end);
            ptrs_box = ptrs.iter().map(|&(offset, reloc)| (shift_offset(offset), reloc)).collect();
        };

        // # Byte-sized provenances
        // This includes the existing bytewise provenance in the range, and ptr provenance
        // that overlaps with the begin/end of the range.
        let mut bytes_box: Box<[_]> = Box::new([]);
        let begin_overlap = self.range_ptrs_get(alloc_range(range.start, Size::ZERO), cx).first();
        let end_overlap = self.range_ptrs_get(alloc_range(range.end(), Size::ZERO), cx).first();
        // We only need to go here if there is some overlap or some bytewise provenance.
        if begin_overlap.is_some() || end_overlap.is_some() || self.bytes.is_some() {
            let mut bytes: Vec<(Size, PointerFrag<Prov>)> = Vec::new();
            // First, if there is a part of a pointer at the start, add that.
            if let Some(&(pos, prov)) = begin_overlap {
                // For really small copies, make sure we don't run off the end of the range.
                let end = cmp::min(pos + ptr_size, range.end());
                for (pos, frag) in
                    Self::ptr_fragments(range.start..end, pos, prov, data_bytes, ptr_size)
                {
                    bytes.push((shift_offset(pos), frag));
                }
            } else {
                trace!("no start overlapping entry");
            }

            // Then the main part, bytewise provenance from `self.bytes`.
            bytes.extend(
                self.range_bytes_get(range)
                    .iter()
                    .map(|(offset, frag)| (shift_offset(*offset), frag.clone())),
            );

            // And finally possibly parts of a pointer at the end.
            // We only have to go here if this is actually different than the begin_overlap.
            if let Some(&(pos, prov)) = end_overlap
                && begin_overlap.is_none_or(|(begin, _)| *begin != pos)
            {
                // If this was a really small copy, we'd have handled this in begin_overlap.
                assert!(pos >= range.start);
                for (pos, frag) in
                    Self::ptr_fragments(pos..range.end(), pos, prov, data_bytes, ptr_size)
                {
                    let pos = shift_offset(pos);
                    // The last entry, if it exists, has a lower offset than us, so we
                    // can add it at the end and remain sorted.
                    debug_assert!(bytes.last().is_none_or(|bytes_entry| bytes_entry.0 < pos));
                    bytes.push((pos, frag));
                }
            } else {
                trace!("no end overlapping entry");
            }
            trace!("byte provenances: {bytes:?}");

            // And again a buffer for the new list on the target side.
            bytes_box = bytes.into_boxed_slice();
        }

        ProvenanceCopy { ptrs: ptrs_box, bytes: bytes_box }
    }

    /// Applies a provenance copy.
    /// The affected range, as defined in the parameters to `prepare_copy` is expected
    /// to be clear of provenance.
    pub fn apply_copy(&mut self, copy: ProvenanceCopy<Prov>, range: AllocRange, repeat: u64) {
        let shift_offset = |idx: u64, offset: Size| offset + range.start + idx * range.size;
        if !copy.ptrs.is_empty() {
            // We want to call `insert_presorted` only once so that, if possible, the entries
            // after the range we insert are moved back only once.
            let chunk_len = copy.ptrs.len() as u64;
            self.ptrs.insert_presorted((0..chunk_len * repeat).map(|i| {
                let chunk = i / chunk_len;
                let (offset, prov) = copy.ptrs[(i % chunk_len) as usize];
                (shift_offset(chunk, offset), prov)
            }));
        }
        if !copy.bytes.is_empty() {
            let chunk_len = copy.bytes.len() as u64;
            self.bytes.get_or_insert_with(Box::default).insert_presorted(
                (0..chunk_len * repeat).map(|i| {
                    let chunk = i / chunk_len;
                    let (offset, frag) = &copy.bytes[(i % chunk_len) as usize];
                    (shift_offset(chunk, *offset), frag.clone())
                }),
            );
        }
    }
}