weirflow 0.1.0

#![allow(clippy::too_many_arguments)]

use super::super::{DirectResidentIfdsBatch, DirectResidentIfdsGraph, DirectResidentIfdsGraphKey};
use crate::ifds_gpu::IfdsResidentDispatch;
use std::mem;

impl<R> DirectResidentIfdsBatch<R>
where
    R: Clone,
{
    /// Solve one seed set with a precomputed graph key.
    #[allow(clippy::too_many_arguments)]
    pub fn solve_with_graph_key<D>(
        &mut self,
        dispatch: &D,
        graph_key: DirectResidentIfdsGraphKey,
        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)],
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
    ) -> Result<Vec<u32>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut single_result_scratch = mem::take(&mut self.single_result_scratch);
        let solve_result = self.solve_many_with_graph_key_into(
            dispatch,
            graph_key,
            num_procs,
            blocks_per_proc,
            facts_per_proc,
            intra_edges,
            inter_edges,
            flow_gen,
            flow_kill,
            &[seed_facts],
            max_iterations,
            &mut single_result_scratch,
        );
        let result = solve_result
            .and_then(|()| take_single_result(&mut single_result_scratch, "one seed set"));
        self.single_result_scratch = single_result_scratch;
        result
    }

    /// Solve one seed set with a borrowed graph view.
    pub fn solve_with_graph_view<D>(
        &mut self,
        dispatch: &D,
        graph: DirectResidentIfdsGraph<'_>,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
    ) -> Result<Vec<u32>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut single_result_scratch = mem::take(&mut self.single_result_scratch);
        let solve_result = self.solve_many_with_graph_view_into(
            dispatch,
            graph,
            &[seed_facts],
            max_iterations,
            &mut single_result_scratch,
        );
        let result = solve_result.and_then(|()| {
            take_single_result(&mut single_result_scratch, "one graph-view seed set")
        });
        self.single_result_scratch = single_result_scratch;
        result
    }

    /// Solve one seed set against an already cached direct-resident graph key.
    ///
    /// This is the warmed steady-state single-query hot path: it accepts no
    /// edge slices, so repeated calls cannot re-hash, rebuild, or re-upload the
    /// invariant graph.
    pub fn solve_cached_single_with_graph_key<D>(
        &mut self,
        dispatch: &D,
        graph_key: DirectResidentIfdsGraphKey,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
    ) -> Result<Vec<u32>, String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut single_result_scratch = mem::take(&mut self.single_result_scratch);
        let solve_result = self.solve_cached_with_graph_key_into(
            dispatch,
            graph_key,
            &[seed_facts],
            max_iterations,
            &mut single_result_scratch,
        );
        let result = solve_result.and_then(|()| {
            take_single_result(&mut single_result_scratch, "one cached-key seed set")
        });
        self.single_result_scratch = single_result_scratch;
        result
    }

    /// Solve one seed set with a precomputed graph key into caller-owned
    /// result storage.
    #[allow(clippy::too_many_arguments)]
    pub fn solve_with_graph_key_into<D>(
        &mut self,
        dispatch: &D,
        graph_key: DirectResidentIfdsGraphKey,
        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)],
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
        result: &mut Vec<u32>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        graph_key.validate_edges(
            num_procs,
            blocks_per_proc,
            facts_per_proc,
            intra_edges,
            inter_edges,
            flow_gen,
            flow_kill,
        )?;
        self.solve_single_validated_graph_key_into(
            dispatch,
            graph_key,
            intra_edges,
            inter_edges,
            flow_gen,
            flow_kill,
            seed_facts,
            max_iterations,
            result,
        )
    }

    pub(super) fn solve_single_validated_graph_key_into<D>(
        &mut self,
        dispatch: &D,
        graph_key: DirectResidentIfdsGraphKey,
        intra_edges: &[(u32, u32, u32)],
        inter_edges: &[(u32, u32, u32, u32)],
        flow_gen: &[(u32, u32, u32)],
        flow_kill: &[(u32, u32, u32)],
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
        result: &mut Vec<u32>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut single_result_scratch = mem::take(&mut self.single_result_scratch);
        prepare_single_result_scratch(&mut single_result_scratch)?;
        if seed_facts.is_empty() {
            crate::dispatch_decode::require_positive_iterations(
                "weir direct resident IFDS solve single with graph key into",
                max_iterations,
            )?;
            self.single_result_scratch = single_result_scratch;
            result.clear();
            return Ok(());
        }
        result.clear();
        mem::swap(result, &mut single_result_scratch[0]);
        let solve_result = self.solve_many_validated_graph_key_into(
            dispatch,
            graph_key,
            intra_edges,
            inter_edges,
            flow_gen,
            flow_kill,
            &[seed_facts],
            max_iterations,
            &mut single_result_scratch,
        );
        let route_result = solve_result.and_then(|()| {
            route_single_result_into(&mut single_result_scratch, result, "one seed set")
        });
        self.single_result_scratch = single_result_scratch;
        route_result
    }

    /// Solve one seed set with a borrowed graph view into caller-owned result
    /// storage.
    pub fn solve_with_graph_view_into<D>(
        &mut self,
        dispatch: &D,
        graph: DirectResidentIfdsGraph<'_>,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
        result: &mut Vec<u32>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut single_result_scratch = mem::take(&mut self.single_result_scratch);
        prepare_single_result_scratch(&mut single_result_scratch)?;
        result.clear();
        mem::swap(result, &mut single_result_scratch[0]);
        let solve_result = self.solve_many_with_graph_view_into(
            dispatch,
            graph,
            &[seed_facts],
            max_iterations,
            &mut single_result_scratch,
        );
        let route_result = solve_result.and_then(|()| {
            route_single_result_into(
                &mut single_result_scratch,
                result,
                "one graph-view seed set",
            )
        });
        self.single_result_scratch = single_result_scratch;
        route_result
    }

    /// Solve one seed set against an already cached direct-resident graph key
    /// into caller-owned result storage.
    pub fn solve_cached_single_with_graph_key_into<D>(
        &mut self,
        dispatch: &D,
        graph_key: DirectResidentIfdsGraphKey,
        seed_facts: &[(u32, u32, u32)],
        max_iterations: u32,
        result: &mut Vec<u32>,
    ) -> Result<(), String>
    where
        D: IfdsResidentDispatch<Resource = R>,
    {
        let mut single_result_scratch = mem::take(&mut self.single_result_scratch);
        prepare_single_result_scratch(&mut single_result_scratch)?;
        result.clear();
        mem::swap(result, &mut single_result_scratch[0]);
        let solve_result = self.solve_cached_with_graph_key_into(
            dispatch,
            graph_key,
            &[seed_facts],
            max_iterations,
            &mut single_result_scratch,
        );
        let route_result = solve_result.and_then(|()| {
            route_single_result_into(
                &mut single_result_scratch,
                result,
                "one cached-key seed set",
            )
        });
        self.single_result_scratch = single_result_scratch;
        route_result
    }
}

fn prepare_single_result_scratch(single_result_scratch: &mut Vec<Vec<u32>>) -> Result<(), String> {
    crate::staging_reserve::resize_result_rows(
        single_result_scratch,
        1,
        "direct resident IFDS single-result scratch row",
    )
}

fn take_single_result(
    single_result_scratch: &mut [Vec<u32>],
    label: &str,
) -> Result<Vec<u32>, String> {
    if single_result_scratch.len() != 1 {
        return Err(format!(
            "weir direct resident IFDS batch produced {} result slot(s) for {label}. Fix: resident batch result routing is corrupt.",
            single_result_scratch.len()
        ));
    }
    Ok(mem::take(&mut single_result_scratch[0]))
}

fn route_single_result_into(
    single_result_scratch: &mut [Vec<u32>],
    result: &mut Vec<u32>,
    label: &str,
) -> Result<(), String> {
    if single_result_scratch.len() != 1 {
        return Err(format!(
            "weir direct resident IFDS batch produced {} result slot(s) for {label}. Fix: resident batch result routing is corrupt.",
            single_result_scratch.len()
        ));
    }
    mem::swap(result, &mut single_result_scratch[0]);
    Ok(())
}