wxtla 0.3.1

//! Concurrent, host-agnostic read abstractions and helpers for parser backends.

use std::{
  collections::HashMap,
  fs::File,
  sync::{Arc, Mutex, RwLock},
  time::Instant,
};

use super::{Error, Result};

/// Raw byte-level access to a file, block device, or virtual blob.
///
/// This trait is intentionally path-agnostic. Host path discovery and related
/// resource resolution belong in adapter layers above the parser core.
pub trait ByteSource: Send + Sync {
  /// Read bytes starting at `offset` into `buf`.
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize>;

  /// Return the total size of this source in bytes.
  fn size(&self) -> Result<u64>;

  /// Describe the backend's read behavior.
  fn capabilities(&self) -> ByteSourceCapabilities {
    ByteSourceCapabilities::default()
  }

  /// Return a stable label for tracing and diagnostics.
  fn telemetry_name(&self) -> &'static str {
    std::any::type_name::<Self>()
  }

  /// Read exactly `buf.len()` bytes from `offset`.
  fn read_exact_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
    let mut total_read = 0usize;
    while total_read < buf.len() {
      let chunk_offset = offset
        .checked_add(total_read as u64)
        .ok_or_else(|| Error::invalid_range("data source offset overflow"))?;
      let read = self.read_at(chunk_offset, &mut buf[total_read..])?;
      if read == 0 {
        return Err(Error::UnexpectedEof {
          offset,
          expected: buf.len(),
          actual: total_read,
        });
      }
      total_read += read;
    }
    Ok(())
  }

  /// Read `len` bytes from `offset` into a new buffer.
  fn read_bytes_at(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
    let mut buf = vec![0u8; len];
    self.read_exact_at(offset, &mut buf)?;
    Ok(buf)
  }

  /// Materialize the full source into memory.
  fn read_all(&self) -> Result<Vec<u8>> {
    let size = usize::try_from(self.size()?)
      .map_err(|_| Error::invalid_range("data source is too large to read into memory"))?;
    let mut buf = vec![0u8; size];
    let mut offset = 0usize;
    while offset < size {
      let read = self.read_at(offset as u64, &mut buf[offset..])?;
      if read == 0 {
        break;
      }
      offset += read;
    }
    buf.truncate(offset);
    Ok(buf)
  }
}

/// In-memory data source backed by immutable bytes.
pub struct BytesDataSource {
  bytes: Arc<[u8]>,
}

impl BytesDataSource {
  /// Create an in-memory data source from owned or shared bytes.
  pub fn new(bytes: impl Into<Arc<[u8]>>) -> Self {
    Self {
      bytes: bytes.into(),
    }
  }
}

impl ByteSource for BytesDataSource {
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
    let Ok(offset) = usize::try_from(offset) else {
      return Ok(0);
    };
    if offset >= self.bytes.len() || buf.is_empty() {
      return Ok(0);
    }

    let available = (self.bytes.len() - offset).min(buf.len());
    buf[..available].copy_from_slice(&self.bytes[offset..offset + available]);
    Ok(available)
  }

  fn size(&self) -> Result<u64> {
    Ok(self.bytes.len() as u64)
  }

  fn capabilities(&self) -> ByteSourceCapabilities {
    ByteSourceCapabilities::concurrent(ByteSourceSeekCost::Cheap)
  }

  fn telemetry_name(&self) -> &'static str {
    "core.bytes_data_source"
  }
}

/// OS file-backed data source with random-access reads.
pub struct FileDataSource {
  file: File,
  size: u64,
}

impl FileDataSource {
  /// Open a file-backed data source from a host path.
  pub fn open(path: impl AsRef<std::path::Path>) -> Result<Self> {
    let file = File::open(path)?;
    let size = file.metadata()?.len();
    Ok(Self { file, size })
  }
}

impl ByteSource for FileDataSource {
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
    read_file_at(&self.file, offset, buf)
  }

  fn size(&self) -> Result<u64> {
    Ok(self.size)
  }

  fn capabilities(&self) -> ByteSourceCapabilities {
    ByteSourceCapabilities::concurrent(ByteSourceSeekCost::Cheap)
  }

  fn telemetry_name(&self) -> &'static str {
    "core.file_data_source"
  }
}

#[cfg(unix)]
fn read_file_at(file: &File, offset: u64, buf: &mut [u8]) -> Result<usize> {
  use std::os::unix::fs::FileExt as _;

  Ok(file.read_at(buf, offset)?)
}

#[cfg(windows)]
fn read_file_at(file: &File, offset: u64, buf: &mut [u8]) -> Result<usize> {
  use std::os::windows::fs::FileExt as _;

  Ok(file.seek_read(buf, offset)?)
}

#[cfg(not(any(unix, windows)))]
fn read_file_at(file: &File, offset: u64, buf: &mut [u8]) -> Result<usize> {
  use std::io::{Read, Seek, SeekFrom};

  let mut clone = file.try_clone()?;
  clone.seek(SeekFrom::Start(offset))?;
  Ok(clone.read(buf)?)
}

/// Whether the backing source can serve reads concurrently.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ByteSourceReadConcurrency {
  /// The backend has not declared its concurrency model.
  Unknown,
  /// Reads are effectively serialized by the backend.
  Serialized,
  /// Reads at different offsets can proceed concurrently.
  Concurrent,
}

/// Relative cost of moving between offsets.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ByteSourceSeekCost {
  /// The backend has not declared its seek characteristics.
  Unknown,
  /// Seeking is cheap enough to treat reads as random-access friendly.
  Cheap,
  /// Seeking is expensive and callers should prefer sequential access.
  Expensive,
}

/// Backend capabilities that inform concurrent readers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ByteSourceCapabilities {
  /// Whether the backend can satisfy reads in parallel.
  pub read_concurrency: ByteSourceReadConcurrency,
  /// Whether frequent offset changes are cheap.
  pub seek_cost: ByteSourceSeekCost,
  /// Optional chunk-size hint for high-throughput callers.
  pub preferred_chunk_size: Option<usize>,
}

impl ByteSourceCapabilities {
  /// Construct a capability descriptor.
  pub const fn new(
    read_concurrency: ByteSourceReadConcurrency, seek_cost: ByteSourceSeekCost,
  ) -> Self {
    Self {
      read_concurrency,
      seek_cost,
      preferred_chunk_size: None,
    }
  }

  /// Construct capabilities for a serialized backend.
  pub const fn serialized(seek_cost: ByteSourceSeekCost) -> Self {
    Self::new(ByteSourceReadConcurrency::Serialized, seek_cost)
  }

  /// Construct capabilities for a concurrent backend.
  pub const fn concurrent(seek_cost: ByteSourceSeekCost) -> Self {
    Self::new(ByteSourceReadConcurrency::Concurrent, seek_cost)
  }

  /// Attach an optional preferred chunk size.
  pub fn with_preferred_chunk_size(mut self, preferred_chunk_size: usize) -> Self {
    self.preferred_chunk_size = Some(preferred_chunk_size);
    self
  }
}

impl Default for ByteSourceCapabilities {
  fn default() -> Self {
    Self::new(
      ByteSourceReadConcurrency::Unknown,
      ByteSourceSeekCost::Unknown,
    )
  }
}

/// Shared statistics handle for observed read activity.
#[derive(Debug, Clone, Default)]
pub struct ByteSourceReadStats {
  inner: Arc<Mutex<DataSourceReadStatsState>>,
}

/// Immutable read statistics snapshot.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ByteSourceReadStatsSnapshot {
  /// Number of read requests issued.
  pub read_count: u64,
  /// Total bytes returned across all reads.
  pub read_bytes: u64,
  /// Average read size in bytes.
  pub average_read_size: u64,
  /// Sum of absolute gaps between consecutive read requests.
  pub request_offset_distance_bytes: u64,
  /// Average absolute gap between consecutive read requests.
  pub average_offset_distance_bytes: u64,
  /// Largest single read size seen.
  pub max_read_size: usize,
  /// Largest absolute gap between consecutive read requests.
  pub max_offset_distance_bytes: u64,
  /// Total time spent in reads, in microseconds.
  pub total_read_micros: u128,
  /// Average per-read time, in microseconds.
  pub average_read_micros: u128,
}

#[derive(Debug, Default)]
struct DataSourceReadStatsState {
  read_count: u64,
  read_bytes: u64,
  request_offset_distance_bytes: u64,
  max_read_size: usize,
  max_offset_distance_bytes: u64,
  total_read_micros: u128,
  last_offset: Option<u64>,
  last_len: usize,
}

impl ByteSourceReadStats {
  fn record_read(&self, offset: u64, len: usize, started_at: Instant) {
    let mut state = self
      .inner
      .lock()
      .unwrap_or_else(|poisoned| poisoned.into_inner());
    state.read_count = state.read_count.saturating_add(1);
    state.read_bytes = state.read_bytes.saturating_add(len as u64);
    state.max_read_size = state.max_read_size.max(len);
    state.total_read_micros = state
      .total_read_micros
      .saturating_add(started_at.elapsed().as_micros());

    if let Some(last_offset) = state.last_offset {
      let last_end = last_offset.saturating_add(state.last_len as u64);
      let distance = offset.abs_diff(last_end);
      state.request_offset_distance_bytes =
        state.request_offset_distance_bytes.saturating_add(distance);
      state.max_offset_distance_bytes = state.max_offset_distance_bytes.max(distance);
    }

    state.last_offset = Some(offset);
    state.last_len = len;
  }

  /// Capture the current statistics snapshot.
  pub fn snapshot(&self) -> ByteSourceReadStatsSnapshot {
    let state = self
      .inner
      .lock()
      .unwrap_or_else(|poisoned| poisoned.into_inner());
    let average_read_size = state.read_bytes.checked_div(state.read_count).unwrap_or(0);
    let average_offset_distance_bytes = if state.read_count <= 1 {
      0
    } else {
      state.request_offset_distance_bytes / (state.read_count - 1)
    };
    let average_read_micros = if state.read_count == 0 {
      0
    } else {
      state.total_read_micros / u128::from(state.read_count)
    };

    ByteSourceReadStatsSnapshot {
      read_count: state.read_count,
      read_bytes: state.read_bytes,
      average_read_size,
      request_offset_distance_bytes: state.request_offset_distance_bytes,
      average_offset_distance_bytes,
      max_read_size: state.max_read_size,
      max_offset_distance_bytes: state.max_offset_distance_bytes,
      total_read_micros: state.total_read_micros,
      average_read_micros,
    }
  }
}

/// Wrapper that records read statistics while delegating to another source.
pub struct ObservedDataSource {
  inner: Arc<dyn ByteSource>,
  stats: ByteSourceReadStats,
}

impl ObservedDataSource {
  /// Wrap a source with read-observation metrics.
  pub fn new(inner: Arc<dyn ByteSource>) -> Self {
    Self {
      inner,
      stats: ByteSourceReadStats::default(),
    }
  }

  /// Access the shared statistics handle.
  pub fn stats(&self) -> ByteSourceReadStats {
    self.stats.clone()
  }
}

impl ByteSource for ObservedDataSource {
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
    let started_at = Instant::now();
    let read = self.inner.read_at(offset, buf)?;
    self.stats.record_read(offset, read, started_at);
    Ok(read)
  }

  fn size(&self) -> Result<u64> {
    self.inner.size()
  }

  fn capabilities(&self) -> ByteSourceCapabilities {
    self.inner.capabilities()
  }

  fn telemetry_name(&self) -> &'static str {
    self.inner.telemetry_name()
  }
}

/// Thin wrapper that turns an `Arc<dyn ByteSource>` back into a `ByteSource`.
pub struct SharedDataSource {
  inner: Arc<dyn ByteSource>,
}

impl SharedDataSource {
  /// Wrap a shared source for trait-object handoff.
  pub fn new(inner: Arc<dyn ByteSource>) -> Self {
    Self { inner }
  }
}

impl ByteSource for SharedDataSource {
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
    self.inner.read_at(offset, buf)
  }

  fn size(&self) -> Result<u64> {
    self.inner.size()
  }

  fn capabilities(&self) -> ByteSourceCapabilities {
    self.inner.capabilities()
  }

  fn telemetry_name(&self) -> &'static str {
    self.inner.telemetry_name()
  }
}

/// Windowed view into a sub-range of another source.
pub struct SliceDataSource {
  inner: Arc<dyn ByteSource>,
  base_offset: u64,
  size: u64,
}

impl SliceDataSource {
  /// Create a slice backed by `inner[base_offset..base_offset + size]`.
  pub fn new(inner: Arc<dyn ByteSource>, base_offset: u64, size: u64) -> Self {
    Self {
      inner,
      base_offset,
      size,
    }
  }
}

impl ByteSource for SliceDataSource {
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
    if offset >= self.size || buf.is_empty() {
      return Ok(0);
    }

    let available = usize::try_from(self.size - offset)
      .unwrap_or(usize::MAX)
      .min(buf.len());
    let absolute_offset = self
      .base_offset
      .checked_add(offset)
      .ok_or_else(|| Error::invalid_range("slice data source offset overflow"))?;
    self.inner.read_at(absolute_offset, &mut buf[..available])
  }

  fn size(&self) -> Result<u64> {
    Ok(self.size)
  }

  fn capabilities(&self) -> ByteSourceCapabilities {
    self.inner.capabilities()
  }

  fn telemetry_name(&self) -> &'static str {
    self.inner.telemetry_name()
  }
}

const PROBE_CACHE_WINDOW_SIZE: usize = 4096;
const PROBE_CACHE_LIMIT: u64 = 64 * 1024;

/// Small-window cache for repeated probe reads near the start of a source.
pub struct ProbeCachedDataSource<'a> {
  inner: &'a dyn ByteSource,
  windows: RwLock<HashMap<u64, Arc<[u8]>>>,
}

impl<'a> ProbeCachedDataSource<'a> {
  /// Wrap a source with a probe-oriented cache.
  pub fn new(inner: &'a dyn ByteSource) -> Self {
    Self {
      inner,
      windows: RwLock::new(HashMap::new()),
    }
  }

  fn cacheable(offset: u64, len: usize) -> bool {
    if len == 0 {
      return false;
    }
    let Some(end) = offset.checked_add(len as u64) else {
      return false;
    };
    end <= PROBE_CACHE_LIMIT
  }

  fn read_window(&self, window_offset: u64) -> Result<Arc<[u8]>> {
    if let Some(window) = self
      .windows
      .read()
      .unwrap_or_else(|poisoned| poisoned.into_inner())
      .get(&window_offset)
      .cloned()
    {
      return Ok(window);
    }

    let mut data = vec![0u8; PROBE_CACHE_WINDOW_SIZE];
    let read = self.inner.read_at(window_offset, &mut data)?;
    data.truncate(read);
    let window: Arc<[u8]> = data.into();

    let mut cache = self
      .windows
      .write()
      .unwrap_or_else(|poisoned| poisoned.into_inner());
    let entry = cache.entry(window_offset).or_insert_with(|| window.clone());
    Ok(entry.clone())
  }
}

impl ByteSource for ProbeCachedDataSource<'_> {
  fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
    if !Self::cacheable(offset, buf.len()) {
      return self.inner.read_at(offset, buf);
    }

    let mut written = 0usize;
    while written < buf.len() {
      let absolute = offset
        .checked_add(written as u64)
        .ok_or_else(|| Error::invalid_range("probe cache offset overflow"))?;
      let window_offset =
        (absolute / PROBE_CACHE_WINDOW_SIZE as u64) * PROBE_CACHE_WINDOW_SIZE as u64;
      let window = self.read_window(window_offset)?;
      let window_inner = (absolute - window_offset) as usize;
      if window_inner >= window.len() {
        break;
      }
      let available = (window.len() - window_inner).min(buf.len() - written);
      buf[written..written + available]
        .copy_from_slice(&window[window_inner..window_inner + available]);
      written += available;
      if window.len() < PROBE_CACHE_WINDOW_SIZE {
        break;
      }
    }

    Ok(written)
  }

  fn size(&self) -> Result<u64> {
    self.inner.size()
  }

  fn capabilities(&self) -> ByteSourceCapabilities {
    self.inner.capabilities()
  }

  fn telemetry_name(&self) -> &'static str {
    self.inner.telemetry_name()
  }
}

#[cfg(test)]
mod tests {
  use std::sync::atomic::{AtomicUsize, Ordering};

  use super::*;

  struct MemDataSource {
    data: Vec<u8>,
  }

  impl ByteSource for MemDataSource {
    fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
      let offset = offset as usize;
      if offset >= self.data.len() {
        return Ok(0);
      }
      let available = &self.data[offset..];
      let read = buf.len().min(available.len());
      buf[..read].copy_from_slice(&available[..read]);
      Ok(read)
    }

    fn size(&self) -> Result<u64> {
      Ok(self.data.len() as u64)
    }

    fn capabilities(&self) -> ByteSourceCapabilities {
      ByteSourceCapabilities::concurrent(ByteSourceSeekCost::Cheap).with_preferred_chunk_size(4096)
    }
  }

  #[test]
  fn read_all_materializes_the_source() {
    let source = MemDataSource {
      data: b"read-all".to_vec(),
    };

    assert_eq!(source.read_all().unwrap(), b"read-all");
  }

  #[test]
  fn bytes_data_source_reads_shared_memory() {
    let source = BytesDataSource::new(Arc::<[u8]>::from(&b"shared-bytes"[..]));
    let mut buf = [0u8; 6];

    let read = source.read_at(7, &mut buf).unwrap();
    assert_eq!(read, 5);
    assert_eq!(&buf[..read], b"bytes");
  }

  #[test]
  fn observed_data_source_tracks_requested_read_patterns() {
    let source: Arc<dyn ByteSource> = Arc::new(MemDataSource {
      data: b"abcdefghijklmnopqrstuvwxyz".to_vec(),
    });
    let observed = ObservedDataSource::new(source);
    let stats = observed.stats();

    let mut first = [0u8; 4];
    let mut second = [0u8; 2];
    let mut third = [0u8; 3];
    observed.read_at(0, &mut first).unwrap();
    observed.read_at(4, &mut second).unwrap();
    observed.read_at(10, &mut third).unwrap();

    let snapshot = stats.snapshot();
    assert_eq!(snapshot.read_count, 3);
    assert_eq!(snapshot.read_bytes, 9);
    assert_eq!(snapshot.average_read_size, 3);
    assert_eq!(snapshot.request_offset_distance_bytes, 4);
    assert_eq!(snapshot.average_offset_distance_bytes, 2);
    assert_eq!(snapshot.max_read_size, 4);
    assert_eq!(snapshot.max_offset_distance_bytes, 4);
  }

  #[test]
  fn observed_data_source_forwards_capabilities() {
    let source: Arc<dyn ByteSource> = Arc::new(MemDataSource {
      data: b"capabilities".to_vec(),
    });
    let observed = ObservedDataSource::new(source);

    assert_eq!(
      observed.capabilities(),
      ByteSourceCapabilities::concurrent(ByteSourceSeekCost::Cheap).with_preferred_chunk_size(4096)
    );
  }

  #[test]
  fn shared_data_source_forwards_reads() {
    let source: Arc<dyn ByteSource> = Arc::new(MemDataSource {
      data: b"shared".to_vec(),
    });
    let shared = SharedDataSource::new(source);
    let mut buf = [0u8; 3];

    let read = shared.read_at(1, &mut buf).unwrap();
    assert_eq!(read, 3);
    assert_eq!(&buf, b"har");
  }

  #[test]
  fn slice_data_source_reads_from_the_requested_window() {
    let source: Arc<dyn ByteSource> = Arc::new(MemDataSource {
      data: b"abcdefghijklmnopqrstuvwxyz".to_vec(),
    });
    let slice = SliceDataSource::new(source, 5, 7);

    let mut buf = [0u8; 8];
    let read = slice.read_at(0, &mut buf).unwrap();
    assert_eq!(read, 7);
    assert_eq!(&buf[..read], b"fghijkl");

    let read = slice.read_at(4, &mut buf).unwrap();
    assert_eq!(read, 3);
    assert_eq!(&buf[..read], b"jkl");
  }

  #[test]
  fn slice_data_source_forwards_capabilities() {
    let source: Arc<dyn ByteSource> = Arc::new(MemDataSource {
      data: b"capabilities".to_vec(),
    });
    let slice = SliceDataSource::new(source, 2, 5);

    assert_eq!(
      slice.capabilities(),
      ByteSourceCapabilities::concurrent(ByteSourceSeekCost::Cheap).with_preferred_chunk_size(4096)
    );
  }

  #[test]
  fn probe_cached_data_source_reuses_small_probe_windows() {
    struct CountingDataSource {
      data: Vec<u8>,
      reads: Arc<AtomicUsize>,
    }

    impl ByteSource for CountingDataSource {
      fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
        self.reads.fetch_add(1, Ordering::Relaxed);
        let offset = offset as usize;
        if offset >= self.data.len() {
          return Ok(0);
        }
        let read = buf.len().min(self.data.len() - offset);
        buf[..read].copy_from_slice(&self.data[offset..offset + read]);
        Ok(read)
      }

      fn size(&self) -> Result<u64> {
        Ok(self.data.len() as u64)
      }
    }

    let reads = Arc::new(AtomicUsize::new(0));
    let source = CountingDataSource {
      data: (0..128u8).collect(),
      reads: reads.clone(),
    };
    let cached = ProbeCachedDataSource::new(&source);

    let mut first = [0u8; 16];
    let mut second = [0u8; 8];
    cached.read_at(0, &mut first).unwrap();
    cached.read_at(4, &mut second).unwrap();

    assert_eq!(reads.load(Ordering::Relaxed), 1);
    assert_eq!(&first[..4], &[0, 1, 2, 3]);
    assert_eq!(&second[..4], &[4, 5, 6, 7]);
  }
}