weirflow 0.1.0

use super::retained_scratch::{
    require_parallel_scratch_budget, solve_resident_many_with_retained_scratch_into,
    solve_resident_one_with_retained_scratch_into, take_single_result_slot,
};
use super::ResidentIfdsBatch;
use crate::ifds_gpu::{
    free_resident_ifds_parallel_scratch, IfdsResidentDispatch, PreparedIfdsCsr,
    ResidentPreparedIfdsCsr,
};

impl<R> ResidentIfdsBatch<R>
where
    R: Clone,
{
    pub fn solve_prepared<D>(
        &mut self,
        dispatch: &D,
        prepared: &PreparedIfdsCsr,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
    ) -> Result<Vec<u32>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        if seed_facts.is_empty() {
            crate::dispatch_decode::require_positive_iterations(
                "weir resident IFDS batch prepared solve",
                max_iterations,
            )?;
            return Ok(Vec::new());
        }
        require_parallel_scratch_budget(
            prepared.frontier_words(),
            1,
            usize::try_from(max_iterations).map_err(|error| {
                format!(
                    "weir resident IFDS batch max_iterations does not fit usize: {error}. Fix: use a smaller iteration budget."
                )
            })?,
            seed_facts.len(),
            self.max_parallel_scratch_bytes,
        )?;
        let resident = self.csr_cache.get_or_upload(dispatch, prepared)?;
        solve_resident_many_with_retained_scratch_into(
            dispatch,
            resident,
            &mut self.parallel_scratch,
            &mut self.scratch_allocations,
            &mut self.scratch_reuses,
            &mut self.scratch_frees,
            self.max_parallel_scratch_bytes,
            &mut self.parallel_host_scratch,
            &[seed_facts],
            max_iterations,
            &mut self.single_result_scratch,
        )?;
        self.record_graph_residency_from_cache_stats()?;
        let result =
            take_single_result_slot(&mut self.single_result_scratch, "weir resident IFDS batch")?;
        self.record_resident_dispatch_io(1, seed_facts.len(), result.len())?;
        Ok(result)
    }

    /// Solve one IFDS seed set into caller-owned result storage, uploading
    /// invariant graph buffers only on a resident cache miss.
    pub fn solve_prepared_into<D>(
        &mut self,
        dispatch: &D,
        prepared: &PreparedIfdsCsr,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
        result: &mut Vec<u32>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        if seed_facts.is_empty() {
            crate::dispatch_decode::require_positive_iterations(
                "weir resident IFDS batch prepared solve into",
                max_iterations,
            )?;
            result.clear();
            return Ok(());
        }
        require_parallel_scratch_budget(
            prepared.frontier_words(),
            1,
            usize::try_from(max_iterations).map_err(|error| {
                format!(
                    "weir resident IFDS batch max_iterations does not fit usize: {error}. Fix: use a smaller iteration budget."
                )
            })?,
            seed_facts.len(),
            self.max_parallel_scratch_bytes,
        )?;
        let resident = self.csr_cache.get_or_upload(dispatch, prepared)?;
        solve_resident_one_with_retained_scratch_into(
            dispatch,
            resident,
            &mut self.parallel_scratch,
            &mut self.scratch_allocations,
            &mut self.scratch_reuses,
            &mut self.scratch_frees,
            self.max_parallel_scratch_bytes,
            &mut self.parallel_host_scratch,
            seed_facts,
            max_iterations,
            &mut self.single_result_scratch,
            result,
            "weir resident IFDS batch",
        )?;
        self.record_graph_residency_from_cache_stats()?;
        self.record_resident_dispatch_io(1, seed_facts.len(), result.len())?;
        Ok(())
    }

    /// Solve many IFDS seed sets against a prepared CSR, uploading invariant
    /// graph buffers only on a resident cache miss.
    pub fn solve_prepared_many<D>(
        &mut self,
        dispatch: &D,
        prepared: &PreparedIfdsCsr,
        seed_sets: &[&[(u32, u32, u32)]],
        max_iterations: u32,
    ) -> Result<Vec<Vec<u32>>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut results = crate::staging_reserve::reserved_vec(
            seed_sets.len(),
            "resident IFDS batch result row",
        )?;
        self.solve_prepared_many_into(dispatch, prepared, seed_sets, max_iterations, &mut results)?;
        Ok(results)
    }

    /// Solve many IFDS seed sets into caller-owned result storage while
    /// uploading invariant graph buffers only on a resident cache miss.
    pub fn solve_prepared_many_into<D>(
        &mut self,
        dispatch: &D,
        prepared: &PreparedIfdsCsr,
        seed_sets: &[&[(u32, u32, u32)]],
        max_iterations: u32,
        results: &mut Vec<Vec<u32>>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        if seed_sets.is_empty() {
            results.clear();
            return Ok(());
        }
        crate::dispatch_decode::require_positive_iterations(
            "weir resident IFDS batch prepared many solve",
            max_iterations,
        )?;
        if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
            resize_and_clear_result_slots(results, seed_sets.len())?;
            return Ok(());
        }
        let max_iterations_usize = usize::try_from(max_iterations).map_err(|error| {
            format!(
                "weir resident IFDS batch max_iterations does not fit usize: {error}. Fix: use a smaller iteration budget."
            )
        })?;
        let total_seed_facts = seed_sets.iter().try_fold(0usize, |total, seed_set| {
            total.checked_add(seed_set.len()).ok_or_else(|| {
                "weir resident IFDS batch seed count overflowed usize. Fix: shard the IFDS seed batch."
                    .to_string()
            })
        })?;
        require_parallel_scratch_budget(
            prepared.frontier_words(),
            seed_sets.len(),
            max_iterations_usize,
            total_seed_facts,
            self.max_parallel_scratch_bytes,
        )?;
        let resident = self.csr_cache.get_or_upload(dispatch, prepared)?;
        solve_resident_many_with_retained_scratch_into(
            dispatch,
            resident,
            &mut self.parallel_scratch,
            &mut self.scratch_allocations,
            &mut self.scratch_reuses,
            &mut self.scratch_frees,
            self.max_parallel_scratch_bytes,
            &mut self.parallel_host_scratch,
            seed_sets,
            max_iterations,
            results,
        )?;
        self.record_graph_residency_from_cache_stats()?;
        let result_words = results.iter().try_fold(0usize, |total, result| {
            total.checked_add(result.len()).ok_or_else(|| {
                "weir resident IFDS batch result word count overflowed usize. Fix: shard the IFDS query batch."
                    .to_string()
            })
        })?;
        self.record_resident_dispatch_io(seed_sets.len(), total_seed_facts, result_words)?;
        Ok(())
    }

    /// Solve one IFDS seed set against an already resident CSR while reusing
    /// this facade's resident parallel scratch allocation.
    pub fn solve_resident_prepared<D>(
        &mut self,
        dispatch: &D,
        resident: &ResidentPreparedIfdsCsr<R>,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
    ) -> Result<Vec<u32>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        if seed_facts.is_empty() {
            crate::dispatch_decode::require_positive_iterations(
                "weir resident IFDS batch resident-prepared solve",
                max_iterations,
            )?;
            return Ok(Vec::new());
        }
        solve_resident_many_with_retained_scratch_into(
            dispatch,
            resident,
            &mut self.parallel_scratch,
            &mut self.scratch_allocations,
            &mut self.scratch_reuses,
            &mut self.scratch_frees,
            self.max_parallel_scratch_bytes,
            &mut self.parallel_host_scratch,
            &[seed_facts],
            max_iterations,
            &mut self.single_result_scratch,
        )?;
        let result = take_single_result_slot(
            &mut self.single_result_scratch,
            "weir resident IFDS batch resident-prepared",
        )?;
        self.record_resident_dispatch_io(1, seed_facts.len(), result.len())?;
        Ok(result)
    }

    /// Solve one IFDS seed set against an already resident CSR into
    /// caller-owned result storage while reusing this facade's resident
    /// parallel scratch allocation.
    pub fn solve_resident_prepared_into<D>(
        &mut self,
        dispatch: &D,
        resident: &ResidentPreparedIfdsCsr<R>,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
        result: &mut Vec<u32>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        if seed_facts.is_empty() {
            crate::dispatch_decode::require_positive_iterations(
                "weir resident IFDS batch resident-prepared solve into",
                max_iterations,
            )?;
            result.clear();
            return Ok(());
        }
        solve_resident_one_with_retained_scratch_into(
            dispatch,
            resident,
            &mut self.parallel_scratch,
            &mut self.scratch_allocations,
            &mut self.scratch_reuses,
            &mut self.scratch_frees,
            self.max_parallel_scratch_bytes,
            &mut self.parallel_host_scratch,
            seed_facts,
            max_iterations,
            &mut self.single_result_scratch,
            result,
            "weir resident IFDS batch resident-prepared",
        )?;
        self.record_resident_dispatch_io(1, seed_facts.len(), result.len())?;
        Ok(())
    }

    /// Solve many IFDS seed sets against an already resident CSR while reusing
    /// this facade's resident parallel scratch allocation.
    pub fn solve_resident_prepared_many<D>(
        &mut self,
        dispatch: &D,
        resident: &ResidentPreparedIfdsCsr<R>,
        seed_sets: &[&[(u32, u32, u32)]],
        max_iterations: u32,
    ) -> Result<Vec<Vec<u32>>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut results = crate::staging_reserve::reserved_vec(
            seed_sets.len(),
            "resident IFDS batch resident-prepared result row",
        )?;
        self.solve_resident_prepared_many_into(
            dispatch,
            resident,
            seed_sets,
            max_iterations,
            &mut results,
        )?;
        Ok(results)
    }

    /// Solve many IFDS seed sets against an already resident CSR into
    /// caller-owned result storage while reusing this facade's resident
    /// parallel scratch allocation.
    pub fn solve_resident_prepared_many_into<D>(
        &mut self,
        dispatch: &D,
        resident: &ResidentPreparedIfdsCsr<R>,
        seed_sets: &[&[(u32, u32, u32)]],
        max_iterations: u32,
        results: &mut Vec<Vec<u32>>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        if seed_sets.is_empty() {
            results.clear();
            return Ok(());
        }
        crate::dispatch_decode::require_positive_iterations(
            "weir resident IFDS batch resident-prepared many solve",
            max_iterations,
        )?;
        if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
            resize_and_clear_result_slots(results, seed_sets.len())?;
            return Ok(());
        }
        let total_seed_facts = seed_sets.iter().try_fold(0usize, |total, seed_set| {
            total.checked_add(seed_set.len()).ok_or_else(|| {
                "weir resident IFDS batch seed count overflowed usize. Fix: shard the IFDS seed batch."
                    .to_string()
            })
        })?;
        solve_resident_many_with_retained_scratch_into(
            dispatch,
            resident,
            &mut self.parallel_scratch,
            &mut self.scratch_allocations,
            &mut self.scratch_reuses,
            &mut self.scratch_frees,
            self.max_parallel_scratch_bytes,
            &mut self.parallel_host_scratch,
            seed_sets,
            max_iterations,
            results,
        )?;
        let result_words = results.iter().try_fold(0usize, |total, result| {
            total.checked_add(result.len()).ok_or_else(|| {
                "weir resident IFDS batch resident-prepared result word count overflowed usize. Fix: shard the IFDS query batch."
                    .to_string()
            })
        })?;
        self.record_resident_dispatch_io(seed_sets.len(), total_seed_facts, result_words)?;
        Ok(())
    }

    /// Free every resident CSR and scratch buffer retained by this batch facade.
    pub fn free_all<D>(&mut self, dispatch: &D) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let scratch_result = match self.parallel_scratch.take() {
            Some(scratch) => {
                let result = free_resident_ifds_parallel_scratch(dispatch, scratch);
                if result.is_ok() {
                    self.scratch_frees = increment_scratch_counter(self.scratch_frees)?;
                }
                result
            }
            None => Ok(()),
        };
        let cache_result = self.csr_cache.free_all(dispatch);
        match (scratch_result, cache_result) {
            (Ok(()), Ok(())) => Ok(()),
            (Err(error), _) | (_, Err(error)) => Err(error),
        }
    }
}

fn increment_scratch_counter(counter: u64) -> Result<u64, String> {
    counter.checked_add(1).ok_or_else(|| {
        "weir resident IFDS batch scratch counter overflowed u64. Fix: recreate the batch facade before continuing an unbounded scratch lifecycle."
            .to_string()
    })
}

fn resize_and_clear_result_slots(results: &mut Vec<Vec<u32>>, len: usize) -> Result<(), String> {
    crate::staging_reserve::resize_result_rows(results, len, "resident IFDS batch result row")?;
    crate::staging_reserve::clear_result_rows(results);
    Ok(())
}