paulistrings 0.1.2

//! The storage seam under `run_layers`, [`PartitionStorage`] and [`PartitionBackend`], and its host-memory implementation [`HostPartition`].

use super::partitioned::layer::{
    apply_layer_partitioned_with_plan, LayerExchangeCounts, PartitionState,
};
use super::plan::PartitionPlan;
use super::transport::Transport;
use crate::channel::prepared::Prepared;
use crate::collectives::Collectives;
use crate::engine::bucketed::LayerScratch;
#[cfg(feature = "phase-timing")]
use crate::engine::stats::PhaseStats;
use crate::engine::trace::PartitionLayerRow;
use crate::pauli_sum::hash::{Gf2Hash, PartitionRows};
use crate::pauli_sum::storage::PauliSum;
use crate::truncation::TruncationPolicy;

/// The policy-independent half of a partition.
// `pub` only because the `pub` methods of `PartitionedSum` and `DistributedSum` bound their backend parameter on it.
#[doc(hidden)]
pub trait PartitionStorage<const W: usize>: Send + Sized {
    fn len(&self) -> usize;
    fn hash(&self) -> &Gf2Hash<W>;
    fn refine(&mut self);
    /// The bucket bits this partition proposes for the layer; a backend may raise the host formula but never lower it.
    fn proposed_bits(
        &self,
        prepared: &Prepared<W>,
        target_bucket_len: usize,
        min_buckets: usize,
    ) -> u8 {
        let _ = prepared;
        crate::pauli_sum::storage::desired_bits(self.len(), target_bucket_len, min_buckets)
            .max(self.hash().bits())
    }
    /// Refine to the host bucket policy without the prepared channel and return `true`, or return `false` for a backend whose proposal needs it.
    fn refine_unprepared(&mut self, target_bucket_len: usize, min_buckets: usize) -> bool {
        let _ = (target_bucket_len, min_buckets);
        false
    }
    /// Move the partition out, leaving an empty one under the same hash so the driver stays consistent if the work is never handed back.
    fn detach(&mut self) -> Self;
    #[cfg(feature = "phase-timing")]
    fn stats(&mut self) -> &mut PhaseStats;
}

/// One layer under policy `T`; an implementor must issue exactly the transport calls the host layer issues, in the same order, or the group falls out of step; the device backend's one addition is the exchange vote (`engine/gpu/export.rs`), legal only because every partition of a group runs the same backend.
pub(crate) trait PartitionBackend<const W: usize, T: ?Sized>: PartitionStorage<W> {
    fn apply_layer<X: Transport>(
        &mut self,
        prepared: &Prepared<W>,
        plan: &PartitionPlan,
        rows: &PartitionRows<W>,
        policy: &T,
        transport: &X,
    ) -> LayerExchangeCounts;
    fn finalize_layer(&mut self, policy: &T, collectives: &dyn Collectives);
    /// Amend the layer's trace row with what the backend counted itself; called only when tracing.
    fn trace_layer(&self, row: &mut PartitionLayerRow) {
        let _ = row;
    }
}

/// A partition in host memory with its retained scratch.
// `pub` only because it is the default backend parameter of `PartitionedSum` and `DistributedSum`.
#[doc(hidden)]
#[derive(Debug)]
pub struct HostPartition<const W: usize> {
    pub(crate) sum: PauliSum<W>,
    pub(super) state: PartitionState<W>,
}

impl<const W: usize> HostPartition<W> {
    pub(crate) fn new(sum: PauliSum<W>) -> Self {
        Self::with_layer_scratch(sum, LayerScratch::default())
    }

    pub(crate) fn with_layer_scratch(sum: PauliSum<W>, layer: LayerScratch<W>) -> Self {
        Self { sum, state: PartitionState { layer, ..PartitionState::default() } }
    }

    pub(crate) fn into_parts(self) -> (PauliSum<W>, LayerScratch<W>) {
        (self.sum, self.state.layer)
    }
}

impl<const W: usize> PartitionStorage<W> for HostPartition<W> {
    fn len(&self) -> usize {
        self.sum.len()
    }

    fn hash(&self) -> &Gf2Hash<W> {
        self.sum.hash()
    }

    fn refine(&mut self) {
        self.sum.refine();
    }

    fn refine_unprepared(&mut self, target_bucket_len: usize, min_buckets: usize) -> bool {
        self.sum.rebucket(target_bucket_len, min_buckets);
        true
    }

    fn detach(&mut self) -> Self {
        let placeholder = PauliSum::empty_with_hash(self.sum.num_qubits(), self.sum.hash().clone());
        Self {
            sum: std::mem::replace(&mut self.sum, placeholder),
            state: std::mem::take(&mut self.state),
        }
    }

    #[cfg(feature = "phase-timing")]
    fn stats(&mut self) -> &mut PhaseStats {
        &mut self.state.layer.stats
    }
}

impl<const W: usize, T> PartitionBackend<W, T> for HostPartition<W>
where
    T: TruncationPolicy<W> + ?Sized,
{
    fn apply_layer<X: Transport>(
        &mut self,
        prepared: &Prepared<W>,
        plan: &PartitionPlan,
        rows: &PartitionRows<W>,
        policy: &T,
        transport: &X,
    ) -> LayerExchangeCounts {
        apply_layer_partitioned_with_plan(
            &mut self.sum,
            prepared,
            plan,
            rows,
            policy,
            &mut self.state,
            transport,
        )
    }

    fn finalize_layer(&mut self, policy: &T, collectives: &dyn Collectives) {
        policy.finalize_layer_partitioned(&mut self.sum, collectives);
    }
}

#[cfg(test)]
mod tests;