prism-q 0.35.0

//! Independent-subsystem decomposition.
//!
//! Runs each independent block on its own backend, then merges classical bits
//! and factored probabilities into one outcome.

use crate::circuit::Circuit;
use crate::error::Result;

use super::{
    BackendKind, ExecutionPlan, FactoredBlock, Probabilities, ProbabilityRoute, RunMetadata,
    RunOutcome, SimOptions,
    dispatch::{BackendPlan, TemporalCliffordPlan},
    ensure_saves_recorded, execute_circuit, expand_for_backend, fuse_for_backend,
    plan_probability_route, plan_temporal_clifford, resolve, run_route, run_temporal_clifford,
};

pub(super) const MIN_DECOMPOSITION_QUBITS: usize = 8;

pub(super) fn should_decompose(components: &[Vec<usize>], total_qubits: usize) -> bool {
    let max_block = components.iter().map(|c| c.len()).max().unwrap_or(0);
    max_block + 3 <= total_qubits
}

#[cfg(feature = "parallel")]
const MAX_BLOCK_QUBITS_FOR_PAR: usize = crate::backend::PARALLEL_THRESHOLD_QUBITS;

/// Run block `i` of `k` on `run(i, mix_seed(seed, i))`, across Rayon workers
/// when every block sits below the parallel floor.
fn run_blocks_maybe_par(
    _kind: &BackendKind,
    _components: &[Vec<usize>],
    seed: u64,
    k: usize,
    run: impl Fn(usize, u64) -> Result<RunOutcome> + Sync,
) -> Vec<Result<RunOutcome>> {
    #[cfg(feature = "parallel")]
    {
        let all_small = _components
            .iter()
            .all(|c| c.len() < MAX_BLOCK_QUBITS_FOR_PAR);
        #[cfg(feature = "gpu")]
        let all_small = all_small
            && _components
                .iter()
                .all(|c| !super::dispatch::may_resolve_to_gpu(_kind, c.len()));
        if all_small && k >= 2 {
            use rayon::prelude::*;
            crate::backend::init_thread_pool();
            return (0..k)
                .into_par_iter()
                .map(|i| run(i, super::mix_seed(seed, i)))
                .collect();
        }
    }
    (0..k).map(|i| run(i, super::mix_seed(seed, i))).collect()
}

fn run_subcircuit(
    kind: &BackendKind,
    sub_circuit: &Circuit,
    block_seed: u64,
    opts: &SimOptions,
) -> Result<RunOutcome> {
    let block_kind = if matches!(kind, BackendKind::Auto) {
        BackendKind::Auto
    } else {
        kind.clone()
    };
    super::run_with_internal(block_kind, sub_circuit, block_seed, *opts)
}

pub(super) fn run_decomposed(
    kind: &BackendKind,
    components: &[Vec<usize>],
    circuit: &Circuit,
    seed: u64,
    opts: &SimOptions,
) -> Result<RunOutcome> {
    let partitions = circuit.partition_subcircuits(components);
    let k = partitions.len();

    let block_opts = if circuit.num_qubits > 64 {
        SimOptions::classical_only()
    } else {
        *opts
    };
    let results: Vec<Result<RunOutcome>> =
        run_blocks_maybe_par(kind, components, seed, k, |i, block_seed| {
            run_subcircuit(kind, &partitions[i].0, block_seed, &block_opts)
        });

    merge_decomposed_results(
        results,
        components,
        &partitions,
        circuit.num_classical_bits,
        circuit.num_qubits,
        opts,
    )
}

/// One block's route, settled once for a loop that runs the block per seed.
enum BlockRoute {
    TemporalClifford(TemporalCliffordPlan),
    Backend { plan: BackendPlan, fused: Circuit },
    Route(ProbabilityRoute),
}

impl BlockRoute {
    fn prepare(kind: &BackendKind, block: &Circuit) -> Self {
        let has_partial_independence = match plan_probability_route(kind, block) {
            ProbabilityRoute::TemporalClifford {
                has_partial_independence,
            } => match plan_temporal_clifford(kind, block) {
                Some(plan) => return BlockRoute::TemporalClifford(plan),
                None => has_partial_independence,
            },
            ProbabilityRoute::Direct {
                has_partial_independence,
            } => has_partial_independence,
            route => return BlockRoute::Route(route),
        };
        match resolve(kind, block, has_partial_independence) {
            ExecutionPlan::Backend(plan) => {
                let probe = plan.build(0);
                let expanded = expand_for_backend(&*probe, block);
                let fused = fuse_for_backend(&*probe, &expanded).into_owned();
                BlockRoute::Backend { plan, fused }
            }
            _ => BlockRoute::Route(ProbabilityRoute::Direct {
                has_partial_independence,
            }),
        }
    }

    /// What `run_with_internal` returns for `block` on an auto `kind`.
    fn run(
        &self,
        kind: &BackendKind,
        block: &Circuit,
        seed: u64,
        opts: &SimOptions,
    ) -> Result<RunOutcome> {
        let outcome = match self {
            BlockRoute::TemporalClifford(plan) => {
                run_temporal_clifford(plan, seed, opts.probabilities)?
            }
            BlockRoute::Backend { plan, fused } => {
                let mut backend = plan.build(seed);
                execute_circuit(&mut *backend, fused, opts)?
            }
            BlockRoute::Route(route) => run_route(kind, block, seed, *opts, route)?,
        };
        ensure_saves_recorded(block, &outcome)?;
        Ok(outcome)
    }
}

/// The blocks of a decomposed circuit with each block's route settled once, so
/// a shot loop on an auto kind runs [`run_decomposed`]'s blocks per seed without
/// partitioning and planning them per shot.
pub(super) struct PreparedBlocks {
    components: Vec<Vec<usize>>,
    partitions: Vec<(Circuit, Vec<usize>, Vec<usize>)>,
    routes: Vec<BlockRoute>,
}

impl PreparedBlocks {
    pub(super) fn new(kind: &BackendKind, components: Vec<Vec<usize>>, circuit: &Circuit) -> Self {
        let partitions = circuit.partition_subcircuits(&components);
        let routes = partitions
            .iter()
            .map(|(block, _, _)| BlockRoute::prepare(kind, block))
            .collect();
        Self {
            components,
            partitions,
            routes,
        }
    }

    /// [`run_decomposed`] of `circuit` on `seed`, for the `kind` and circuit the
    /// blocks were prepared from.
    pub(super) fn run(
        &self,
        kind: &BackendKind,
        circuit: &Circuit,
        seed: u64,
        opts: &SimOptions,
    ) -> Result<RunOutcome> {
        let block_opts = if circuit.num_qubits > 64 {
            SimOptions::classical_only()
        } else {
            *opts
        };
        let results = run_blocks_maybe_par(
            kind,
            &self.components,
            seed,
            self.partitions.len(),
            |i, block_seed| {
                self.routes[i].run(kind, &self.partitions[i].0, block_seed, &block_opts)
            },
        );
        merge_decomposed_results(
            results,
            &self.components,
            &self.partitions,
            circuit.num_classical_bits,
            circuit.num_qubits,
            opts,
        )
    }
}

fn merge_decomposed_results(
    results: Vec<Result<RunOutcome>>,
    components: &[Vec<usize>],
    partitions: &[(Circuit, Vec<usize>, Vec<usize>)],
    num_classical_bits: usize,
    num_qubits: usize,
    opts: &SimOptions,
) -> Result<RunOutcome> {
    let mut factored_blocks: Vec<FactoredBlock> = Vec::new();
    let mut merged_classical = vec![false; num_classical_bits];
    let mut block_metadata: Vec<RunMetadata> = Vec::with_capacity(partitions.len());

    for (i, result) in results.into_iter().enumerate() {
        let result = result?;
        block_metadata.push(result.metadata);
        let (_, ref qubit_map, ref classical_map) = partitions[i];

        for (local_idx, &global_idx) in classical_map.iter().enumerate() {
            merged_classical[global_idx] = result.classical_bits[local_idx];
        }

        if opts.probabilities
            && num_qubits <= 64
            && let Some(probs) = result.probabilities
        {
            let dense = probs.to_vec();
            let mut mask = 0u64;
            for &global_qubit in qubit_map {
                mask |= 1u64 << global_qubit;
            }
            factored_blocks.push(FactoredBlock { probs: dense, mask });
        }
    }

    let probabilities = if opts.probabilities && factored_blocks.len() == components.len() {
        Some(Probabilities::Factored {
            blocks: factored_blocks,
            total_qubits: num_qubits,
        })
    } else {
        None
    };

    Ok(RunOutcome {
        classical_bits: merged_classical,
        probabilities,
        metadata: RunMetadata::decomposed(block_metadata),
        saves: Vec::new(),
    })
}

pub(super) fn run_decomposed_prefused(
    block_plans: &[BackendPlan],
    components: &[Vec<usize>],
    partitions: &[(Circuit, Vec<usize>, Vec<usize>)],
    fused_blocks: &[std::borrow::Cow<'_, Circuit>],
    seed: u64,
    opts: &SimOptions,
    original_circuit: &Circuit,
) -> Result<RunOutcome> {
    let k = partitions.len();
    let block_opts = if original_circuit.num_qubits > 64 {
        SimOptions::classical_only()
    } else {
        *opts
    };
    let results: Vec<Result<RunOutcome>> = (0..k)
        .map(|i| {
            let block_seed = super::mix_seed(seed, i);
            let mut backend = block_plans[i].build(block_seed);
            execute_circuit(&mut *backend, &fused_blocks[i], &block_opts)
        })
        .collect();
    merge_decomposed_results(
        results,
        components,
        partitions,
        original_circuit.num_classical_bits,
        original_circuit.num_qubits,
        opts,
    )
}

/// Combine per-block probability vectors via Kronecker product.
///
/// Two-pass O(2^N) algorithm:
/// 1. In-place Kronecker product: single allocation, reverse-iteration expansion
/// 2. Bit permutation to map natural (block-sequential) bit positions
///    to global qubit positions (parallelized at ≥2^14 states)
pub(crate) fn merge_probabilities(
    blocks: &[(Vec<f64>, Vec<usize>)],
    total_qubits: usize,
) -> Vec<f64> {
    let n_states = 1usize << total_qubits;

    let mut result = vec![0.0f64; n_states];
    result[0] = 1.0;
    let mut cur_len = 1usize;

    for (probs, _) in blocks {
        let block_len = probs.len();
        let new_len = cur_len * block_len;
        for i in (0..cur_len).rev() {
            let r = result[i];
            for j in 0..block_len {
                result[i * block_len + j] = r * probs[j];
            }
        }
        cur_len = new_len;
    }
    debug_assert_eq!(cur_len, n_states);

    let mut natural_to_global = Vec::with_capacity(total_qubits);
    for (_probs, qubits) in blocks.iter().rev() {
        natural_to_global.extend_from_slice(qubits);
    }

    if natural_to_global.iter().enumerate().all(|(i, &g)| i == g) {
        return result;
    }

    let mut perm = vec![0usize; n_states];
    for (nat_bit, &global_bit) in natural_to_global.iter().enumerate() {
        let half = 1usize << nat_bit;
        for i in 0..half {
            perm[i | half] = perm[i] | (1 << global_bit);
        }
    }

    let mut permuted = vec![0.0f64; n_states];

    #[cfg(feature = "parallel")]
    {
        const MIN_PAR_PERM: usize = 1 << 14;
        if n_states >= MIN_PAR_PERM {
            use rayon::prelude::*;
            result
                .par_iter()
                .zip(perm.par_iter())
                .for_each(|(&prob, &global_idx)| {
                    // SAFETY: perm is a bijection so each global_idx is written exactly once.
                    // No two threads write to the same index.
                    unsafe {
                        let ptr = permuted.as_ptr() as *mut f64;
                        *ptr.add(global_idx) = prob;
                    }
                });
            return permuted;
        }
    }

    for (nat_idx, &prob) in result.iter().enumerate() {
        permuted[perm[nat_idx]] = prob;
    }
    permuted
}