tket 0.22.0

//! Resource flow logic for tracking how resources move through operations.
//!
//! This module defines the [`ResourceFlow`] trait that specifies how resources
//! flow through operations, along with a default implementation.

use crate::resource::types::ResourceId;
use derive_more::derive::{Display, Error};
use hugr::HugrView;
use hugr::ops::OpType;
use hugr::types::Type;
use itertools::{EitherOrBoth, Itertools};

/// Error type for unsupported operations in ResourceFlow implementations.
#[derive(Debug, Display, Clone, PartialEq, Error)]
#[display("Unsupported operation: {_0}")]
pub struct UnsupportedOp(#[error(not(source))] OpType);

/// Trait for specifying how resources flow through operations.
///
/// This trait allows different implementations to define how linear resources
/// are mapped from inputs to outputs through various operation types.
///
/// This trait is dyn-compatible.
pub trait ResourceFlow<H: HugrView> {
    /// Map resource IDs from operation inputs to outputs.
    ///
    /// Given a `node`, the `hugr` in which it is defined and `inputs`, a list
    /// of the linear inputs to `node` given as resources, return the list
    /// of resources for the linear outputs of `node`.
    ///
    /// More specifically, `inputs` is a vector of the same length as the number
    /// of inputs in the signature of `node` and such that
    ///  - `inputs[i]` is Some(resource_id) if the i-th input port is a linear
    ///    type, and
    ///  - `inputs[i]` is `None` otherwise, i.e. if the i-th input port is a
    ///    copyable type.
    ///
    /// The returned vector `outputs` must satisfy
    ///  - if `outputs[i]` is `Some(resource_id)`, then the i-th output port is
    ///    a linear type and `resource_id` is an ID passed in `inputs`,
    ///  - if the i-th output port is a copyable type, then `outputs[i]` is
    ///    `None`.
    ///
    /// If the i-th output port is linear and `outputs[i]` is set to None, then
    /// a fresh resource ID will be created and assigned to that port.
    ///
    /// # Arguments
    /// * `node` - The operation node
    /// * `hugr` - The `hugr` in which `node` is defined
    /// * `inputs` - Resource IDs for each input port (None for non-linear
    ///   types)
    fn map_resources(
        &self,
        node: H::Node,
        hugr: &H,
        inputs: &[Option<ResourceId>],
    ) -> Result<Vec<Option<ResourceId>>, UnsupportedOp>;

    /// Convert this ResourceFlow into a boxed trait object.
    fn into_boxed<'a>(self) -> Box<dyn 'a + ResourceFlow<H>>
    where
        Self: 'a + Sized,
    {
        Box::new(self)
    }
}

impl<H: HugrView> ResourceFlow<H> for Box<dyn '_ + ResourceFlow<H>> {
    fn map_resources(
        &self,
        node: H::Node,
        hugr: &H,
        inputs: &[Option<ResourceId>],
    ) -> Result<Vec<Option<ResourceId>>, UnsupportedOp> {
        self.as_ref().map_resources(node, hugr, inputs)
    }
}

/// Default implementation of ResourceFlow.
///
/// This implementation considers that an operation is resource-preserving if
/// for all port indices i, either
///  - the i-th input and i-th output are both linear and are of the same type,
///  - or the i-th input and i-th output are both copyable,
///  - or one of the i-th input/output is copyable and the other does not exist.
///
/// For resource-preserving operations, linear inputs are then mapped to the
/// corresponding output. (All outputs with no corresponding input, must be
/// copyable.)
///
/// If on the other hand an operation is not resource preserving, all input
/// resources are discarded and all outputs will be given fresh resource IDs.
#[derive(Debug, Clone, Default)]
pub struct DefaultResourceFlow;

impl DefaultResourceFlow {
    /// Determine if an operation is resource-preserving based on input/output
    /// types.
    fn is_resource_preserving(
        input_types: impl IntoIterator<Item = Type>,
        output_types: impl IntoIterator<Item = Type>,
    ) -> bool {
        // An operation is resource-preserving if for each i, if input[i] or
        // output[i] is linear, then type(input[i]) == type(output[i])

        for io_ty in input_types.into_iter().zip_longest(output_types) {
            let (input_ty, output_ty) = match io_ty {
                EitherOrBoth::Both(input_ty, output_ty) => (input_ty, output_ty),
                EitherOrBoth::Left(ty) | EitherOrBoth::Right(ty) => {
                    if !ty.copyable() {
                        // linear type on one side, nothing on the other
                        return false;
                    }
                    continue;
                }
            };

            if !input_ty.copyable() || !output_ty.copyable() {
                // If input/output is linear, both must be the same type
                if input_ty != output_ty {
                    return false;
                }
            }
        }

        true
    }
}

impl<H: HugrView> ResourceFlow<H> for DefaultResourceFlow {
    fn map_resources(
        &self,
        node: H::Node,
        hugr: &H,
        inputs: &[Option<ResourceId>],
    ) -> Result<Vec<Option<ResourceId>>, UnsupportedOp> {
        let op = hugr.get_optype(node);

        debug_assert_eq!(
            inputs.len(),
            op.value_input_count(),
            "Input resource array length must match operation input count"
        );

        if Self::is_resource_preserving(
            hugr.in_value_types(node).map(|(_, ty)| ty),
            hugr.out_value_types(node).map(|(_, ty)| ty),
        ) {
            Ok(retain_linear_types(
                inputs.to_vec(),
                hugr.out_value_types(node).map(|(_, ty)| ty),
                op.value_output_count(),
            ))
        } else {
            // Not resource-preserving: all linear outputs are new resources (None)
            Ok(vec![None; op.value_output_count()])
        }
    }
}

fn retain_linear_types(
    mut resources: Vec<Option<ResourceId>>,
    types: impl IntoIterator<Item = Type>,
    output_count: usize,
) -> Vec<Option<ResourceId>> {
    resources.resize(output_count, None);
    for (ty, resource) in types.into_iter().zip(resources.iter_mut()) {
        if ty.copyable() {
            *resource = None;
        }
    }
    resources
}