wdext 0.1.0

// SPDX-FileCopyrightText: 2026 takubokudori
// SPDX-License-Identifier: MIT OR Apache-2.0
use std::slice;

use super::raw::bindings::{
    ConstBufferView, GuestAddress, MemoryBuffer, MemoryBufferWithRanges,
    MemoryRange,
};

/// Memory returned from TTD-owned storage by `QueryMemoryRange`.
#[derive(Debug)]
pub struct QueriedMemoryRange<'a> {
    address: GuestAddress,
    memory: &'a [u8],
    sequence: super::raw::bindings::SequenceId,
}

impl<'a> QueriedMemoryRange<'a> {
    /// # Safety
    ///
    /// `result.Memory` must describe TTD-owned storage that remains valid for
    /// `'a`. The caller must arrange for `'a` to end before another memory
    /// query invalidates that storage.
    pub(crate) unsafe fn from_ttd_owned(result: MemoryRange) -> Self {
        let memory = if result.Memory.Size == 0 {
            &[]
        } else {
            assert!(
                !result.Memory.BaseAddress.is_null(),
                "TTD returned a null memory-range pointer with a non-zero size"
            );
            assert!(
                result.Memory.Size <= isize::MAX as usize,
                "TTD returned a memory range too large for a Rust slice"
            );

            unsafe {
                slice::from_raw_parts(
                    result.Memory.BaseAddress.cast::<u8>(),
                    result.Memory.Size,
                )
            }
        };

        Self {
            address: result.Address,
            memory,
            sequence: result.Sequence,
        }
    }

    /// Returns the guest address corresponding to the first returned byte.
    #[inline]
    pub fn address(&self) -> GuestAddress { self.address }

    /// Returns the bytes stored internally by TTD.
    #[inline]
    pub fn memory(&self) -> &'a [u8] { self.memory }

    /// Returns the trace sequence from which the bytes were obtained.
    #[inline]
    pub fn sequence(&self) -> super::raw::bindings::SequenceId { self.sequence }
}

/// Memory returned by a TTD query into caller-owned storage.
///
/// The returned slice borrows the buffer supplied to the query, preventing the
/// native result pointer from outliving that storage.
#[derive(Debug)]
pub struct QueriedMemory<'a> {
    address: GuestAddress,
    memory: &'a [u8],
}

impl<'a> QueriedMemory<'a> {
    #[inline]
    pub(crate) fn from_raw(buffer: &'a [u8], result: MemoryBuffer) -> Self {
        Self {
            address: result.Address,
            memory: validated_memory_slice(buffer, result.Memory),
        }
    }

    /// Returns the guest address corresponding to the first returned byte.
    #[inline]
    pub fn address(&self) -> GuestAddress { self.address }

    /// Returns the bytes written by TTD.
    #[inline]
    pub fn memory(&self) -> &'a [u8] { self.memory }
}

/// Memory and source ranges returned by a TTD query into caller-owned storage.
///
/// Both slices borrow the buffers supplied to the query. This prevents the
/// native `Memory` pointer and `RangeCount` from being used after those buffers
/// cease to be valid.
#[derive(Debug)]
pub struct QueriedMemoryWithRanges<'a, 'r> {
    address: GuestAddress,
    memory: &'a [u8],
    ranges: &'r [MemoryRange],
}

impl<'a, 'r> QueriedMemoryWithRanges<'a, 'r> {
    #[inline]
    pub(crate) fn from_raw(
        buffer: &'a [u8],
        ranges: &'r [MemoryRange],
        result: MemoryBufferWithRanges,
    ) -> Self {
        assert!(
            result.RangeCount <= ranges.len(),
            "TTD returned more memory ranges than the supplied range buffer"
        );

        Self {
            address: result.Address,
            memory: validated_memory_slice(buffer, result.Memory),
            ranges: &ranges[..result.RangeCount],
        }
    }

    /// Returns the guest address corresponding to the first returned byte.
    #[inline]
    pub fn address(&self) -> GuestAddress { self.address }

    /// Returns the bytes written by TTD.
    #[inline]
    pub fn memory(&self) -> &'a [u8] { self.memory }

    /// Returns the prefix of the supplied range buffer populated by TTD.
    #[inline]
    pub fn ranges(&self) -> &'r [MemoryRange] { self.ranges }
}

fn validated_memory_slice(buffer: &[u8], memory: ConstBufferView) -> &[u8] {
    if memory.Size == 0 {
        return &buffer[..0];
    }

    assert!(
        !memory.BaseAddress.is_null(),
        "TTD returned a null memory pointer with a non-zero size"
    );

    let buffer_start = buffer.as_ptr() as usize;
    let buffer_end = buffer_start
        .checked_add(buffer.len())
        .expect("host buffer address overflow");
    let memory_start = memory.BaseAddress as usize;
    let memory_end = memory_start
        .checked_add(memory.Size)
        .expect("TTD memory view address overflow");

    assert!(
        memory_start >= buffer_start && memory_end <= buffer_end,
        "TTD returned a memory view outside the supplied buffer"
    );

    // SAFETY: the bounds checks above prove that TTD's returned view is fully
    // contained in `buffer`. The returned slice is tied to `buffer`'s lifetime.
    unsafe {
        slice::from_raw_parts(memory.BaseAddress.cast::<u8>(), memory.Size)
    }
}