weirflow 0.1.0

use crate::ifds_gpu::prepare_ifds_csr_borrowed_via;
use crate::ifds_resident_types::{IfdsResidentDispatch, PreparedIfdsCsr, ResidentPreparedIfdsCsr};
use vyre_foundation::ir::Program;

/// Upload an already prepared IFDS CSR into backend-resident buffers.
pub fn upload_prepared_ifds_csr_resident<D>(
    dispatch: &D,
    prepared: &PreparedIfdsCsr,
) -> Result<ResidentPreparedIfdsCsr<D::Resource>, String>
where
    D: IfdsResidentDispatch,
{
    fn padded(bytes: &[u8]) -> &[u8] {
        const EMPTY_U32: [u8; std::mem::size_of::<u32>()] = [0; std::mem::size_of::<u32>()];
        if bytes.is_empty() {
            &EMPTY_U32
        } else {
            bytes
        }
    }

    fn allocate_one<D: IfdsResidentDispatch>(
        dispatch: &D,
        label: &'static str,
        bytes: &[u8],
        allocated: &mut Vec<(&'static str, D::Resource)>,
    ) -> Result<D::Resource, String> {
        let resource = dispatch
            .allocate_resident(bytes.len())
            .map_err(|error| format!("weir IFDS resident allocate {label} failed: {error}"))?;
        allocated.push((label, resource.clone()));
        Ok(resource)
    }

    let mut allocated =
        crate::staging_reserve::reserved_vec(5, "IFDS resident prepared CSR rollback resource")?;
    let uploaded = (|| {
        let pg_nodes_b = padded(&prepared.pg_nodes_b);
        let row_ptr_b = padded(&prepared.row_ptr_b);
        let col_idx_b = padded(&prepared.col_idx_b);
        let pg_edge_kind_mask_b = padded(&prepared.pg_edge_kind_mask_b);
        let pg_node_tags_b = padded(&prepared.pg_node_tags_b);
        let share_pg_node_tags = pg_node_tags_b == pg_nodes_b;
        let pg_nodes = allocate_one(dispatch, "pg_nodes", pg_nodes_b, &mut allocated)?;
        let row_ptr = allocate_one(dispatch, "row_ptr", row_ptr_b, &mut allocated)?;
        let col_idx = allocate_one(dispatch, "col_idx", col_idx_b, &mut allocated)?;
        let pg_edge_kind_mask = allocate_one(
            dispatch,
            "pg_edge_kind_mask",
            pg_edge_kind_mask_b,
            &mut allocated,
        )?;
        let pg_node_tags = if share_pg_node_tags {
            pg_nodes.clone()
        } else {
            allocate_one(dispatch, "pg_node_tags", pg_node_tags_b, &mut allocated)?
        };
        let mut uploads =
            crate::staging_reserve::reserved_vec(5, "IFDS resident prepared CSR upload")?;
        uploads.push((&pg_nodes, pg_nodes_b));
        uploads.push((&row_ptr, row_ptr_b));
        uploads.push((&col_idx, col_idx_b));
        uploads.push((&pg_edge_kind_mask, pg_edge_kind_mask_b));
        if !share_pg_node_tags {
            uploads.push((&pg_node_tags, pg_node_tags_b));
        }
        dispatch
            .upload_resident_many(&uploads)
            .map_err(|error| format!("weir IFDS resident batch CSR upload failed: {error}"))?;
        Ok(ResidentPreparedIfdsCsr {
            shape: prepared.shape,
            node_count: prepared.node_count,
            words: prepared.words,
            layout_hash: prepared.stable_layout_hash(),
            pg_nodes,
            row_ptr,
            col_idx,
            pg_edge_kind_mask,
            pg_node_tags,
            shares_pg_node_tags_with_nodes: share_pg_node_tags,
        })
    })();
    if uploaded.is_err() {
        let _ = super::free_unique_resident_resources(dispatch, allocated);
    }
    uploaded
}

/// Build, pack, and upload IFDS invariant CSR buffers once for resident solves.
#[allow(clippy::too_many_arguments)]
pub fn prepare_ifds_csr_resident_via<F, D>(
    borrowed_dispatch: &F,
    resident_dispatch: &D,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    intra_edges: &[(u32, u32, u32)],
    inter_edges: &[(u32, u32, u32, u32)],
    flow_gen: &[(u32, u32, u32)],
    flow_kill: &[(u32, u32, u32)],
) -> Result<ResidentPreparedIfdsCsr<D::Resource>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
    D: IfdsResidentDispatch,
{
    let prepared = prepare_ifds_csr_borrowed_via(
        borrowed_dispatch,
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
    )?;
    upload_prepared_ifds_csr_resident(resident_dispatch, &prepared)
}

/// Free invariant resident buffers created by [`upload_prepared_ifds_csr_resident`]
/// or [`prepare_ifds_csr_resident_via`].
pub fn free_resident_prepared_ifds_csr<D>(
    dispatch: &D,
    prepared: ResidentPreparedIfdsCsr<D::Resource>,
) -> Result<(), String>
where
    D: IfdsResidentDispatch,
{
    let ResidentPreparedIfdsCsr {
        pg_nodes,
        row_ptr,
        col_idx,
        pg_edge_kind_mask,
        pg_node_tags,
        shares_pg_node_tags_with_nodes,
        ..
    } = prepared;
    if shares_pg_node_tags_with_nodes {
        super::free_unique_resident_resources(
            dispatch,
            [
                ("prepared CSR pg_nodes", pg_nodes),
                ("prepared CSR row_ptr", row_ptr),
                ("prepared CSR col_idx", col_idx),
                ("prepared CSR pg_edge_kind_mask", pg_edge_kind_mask),
            ],
        )
    } else {
        super::free_unique_resident_resources(
            dispatch,
            [
                ("prepared CSR pg_nodes", pg_nodes),
                ("prepared CSR row_ptr", row_ptr),
                ("prepared CSR col_idx", col_idx),
                ("prepared CSR pg_edge_kind_mask", pg_edge_kind_mask),
                ("prepared CSR pg_node_tags", pg_node_tags),
            ],
        )
    }
}