vyre-wgpu 0.1.0

//! GPU-to-host readback helpers for the DFA scan engine.
//!
//! Wraps [`wgpu::Buffer`] map/read/unmap into synchronous helpers
//! returning plain bytes (or `u32` slots) so the dispatch path does
//! not have to reason about mapped state or channel lifetimes. Every
//! read polls a dedicated submission index and returns a typed
//! [`Error::Dfa`] on driver-level failure.

use crate::engine::dfa::MAX_DFA_MATCHES;
use std::sync::mpsc;
use vyre::error::{Error, Result};

pub(crate) fn read_one_u32(
    device: &wgpu::Device,
    buffer: &wgpu::Buffer,
    label: &str,
    submission: wgpu::SubmissionIndex,
) -> Result<u32> {
    let bytes = read_buffer(
        device,
        buffer,
        u64::try_from(std::mem::size_of::<u32>()).map_err(|source| Error::Dfa {
            message: format!("u32 size cannot fit u64: {source}. Fix: run on a supported target."),
        })?,
        label,
        submission,
    )?;
    Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
}

pub(crate) fn read_matches(
    device: &wgpu::Device,
    buffer: &wgpu::Buffer,
    match_count: u32,
    submission: wgpu::SubmissionIndex,
) -> Result<Vec<vyre::Match>> {
    if match_count > MAX_DFA_MATCHES {
        return Err(Error::Dfa {
            message: format!(
                "DFA readback match_count {match_count} exceeds {MAX_DFA_MATCHES}. Fix: reject this malformed GPU readback."
            ),
        });
    }
    let bytes = read_buffer(
        device,
        buffer,
        u64::from(match_count) * 12,
        "matches",
        submission,
    )?;
    let mut matches =
        Vec::with_capacity(usize::try_from(match_count).map_err(|source| Error::Dfa {
            message: format!(
                "DFA match_count {match_count} cannot fit usize: {source}. Fix: lower max_matches."
            ),
        })?);
    for fields in bytes.chunks_exact(12) {
        matches.push(vyre::Match::new(
            u32::from_le_bytes([fields[0], fields[1], fields[2], fields[3]]),
            u32::from_le_bytes([fields[4], fields[5], fields[6], fields[7]]),
            u32::from_le_bytes([fields[8], fields[9], fields[10], fields[11]]),
        ));
    }
    Ok(matches)
}

pub fn read_buffer(
    device: &wgpu::Device,
    buffer: &wgpu::Buffer,
    byte_len: u64,
    label: &str,
    submission: wgpu::SubmissionIndex,
) -> Result<Vec<u8>> {
    let slice = buffer.slice(0..byte_len);
    let (sender, receiver) = mpsc::channel();
    let readback_label = label.to_string();
    slice.map_async(wgpu::MapMode::Read, move |result| {
        if let Err(send_err) = sender.send(result) {
            tracing::warn!(
                ?send_err,
                readback = %readback_label,
                "DFA readback receiver dropped before map_async result delivery"
            );
        }
    });
    match device.poll(wgpu::Maintain::wait_for(submission)) {
        wgpu::MaintainResult::Ok | wgpu::MaintainResult::SubmissionQueueEmpty => {}
    }
    receiver
        .recv()
        .map_err(|error| Error::Dfa {
            message: format!(
                "failed to receive DFA {label} readback: {error}. Fix: keep the GPU device alive through readback."
            ),
        })?
        .map_err(|error| Error::Dfa {
            message: format!(
                "failed to map DFA {label} readback: {error:?}. Fix: use MAP_READ and COPY_DST readback buffers."
            ),
        })?;

    let mapped = slice.get_mapped_range();
    let output = mapped.to_vec();
    drop(mapped);
    buffer.unmap();
    Ok(output)
}