tket 0.18.0

//! Pass for removing redundant tuple pack->unpack operations.

use std::collections::VecDeque;

use hugr_core::builder::{DFGBuilder, Dataflow, DataflowHugr};
use hugr_core::extension::prelude::{MakeTuple, UnpackTuple};
use hugr_core::hugr::SimpleReplacementError;
use hugr_core::hugr::hugrmut::HugrMut;
use hugr_core::hugr::views::SiblingSubgraph;
use hugr_core::hugr::views::sibling_subgraph::TopoConvexChecker;
use hugr_core::ops::{OpTrait, OpType};
use hugr_core::types::Type;
use hugr_core::{HugrView, Node, PortIndex, SimpleReplacement};
use itertools::Itertools;

use crate::passes::composable::WithScope;
use crate::passes::{ComposablePass, PassScope};

/// A pass that removes unnecessary `MakeTuple` operations immediately followed
/// by `UnpackTuple`s.
///
/// If the tuple output is consumed by other operations, only the `UnpackTuple`s
/// are removed and their outputs are connected to the original values
/// accordingly.
///
/// Currently only unpack operations in the same region as the `MakeTuple` are
/// removed. This may be extended in the future.
///
/// Removes `MakeTuple` operations that are not consumed by any other
/// operations.
///
/// Ignores pack/unpack nodes with order edges.
// TODO: Supporting those requires updating the `SiblingSubgraph` implementation. See <https://github.com/CQCL/hugr/issues/1974>.
#[derive(Debug, Clone, Default)]
pub struct UntuplePass {
    scope: PassScope,
}

#[derive(Debug, derive_more::Display, derive_more::Error, derive_more::From)]
#[non_exhaustive]
/// Errors produced by [`UntuplePass`].
pub enum UntupleError {
    /// Rewriting the circuit failed.
    RewriteError(SimpleReplacementError),
}

/// Result type for the untuple pass.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct UntupleResult {
    /// Number of `MakeTuple` rewrites applied.
    pub rewrites_applied: usize,
}

impl UntuplePass {
    /// Find tuple pack operations followed by tuple unpack operations
    /// and generate rewrites to remove them.
    ///
    /// The returned rewrites are guaranteed to be independent of each other.
    ///
    /// Returns an iterator over the rewrites.
    pub fn all_rewrites<H: HugrView<Node = Node>>(
        &self,
        hugr: &H,
    ) -> Vec<SimpleReplacement<H::Node>> {
        let Some(parent) = self.scope.root(hugr) else {
            return vec![];
        };
        find_rewrites(hugr, parent, self.scope.recursive())
    }
}

fn find_rewrites<H: HugrView>(
    hugr: &H,
    parent: H::Node,
    recursive: bool,
) -> Vec<SimpleReplacement<H::Node>> {
    let mut res = Vec::new();
    let mut children_queue = VecDeque::new();
    children_queue.push_back(parent);

    while let Some(parent) = children_queue.pop_front() {
        // Required to create SimpleReplacements.
        // Reset for each parent as `TopoConvexChecker` is tied to a specific parent node.
        let mut convex_checker: Option<TopoConvexChecker<H>> = None;

        for node in hugr.children(parent) {
            let op = hugr.get_optype(node);
            if let Some(rw) = make_rewrite(hugr, &mut convex_checker, node, op) {
                res.push(rw);
            }
            if recursive && op.is_container() {
                children_queue.push_back(node);
            }
        }
    }
    res
}

impl<H: HugrMut<Node = Node>> ComposablePass<H> for UntuplePass {
    type Error = UntupleError;
    type Result = UntupleResult;

    fn run(&self, hugr: &mut H) -> Result<Self::Result, Self::Error> {
        let rewrites = self.all_rewrites(hugr);
        let rewrites_applied = rewrites.len();
        // The rewrites are independent, so we can always apply them all.
        for rewrite in rewrites {
            hugr.apply_patch(rewrite)?;
        }
        Ok(UntupleResult { rewrites_applied })
    }
}

impl WithScope for UntuplePass {
    fn with_scope(mut self, scope: impl Into<PassScope>) -> Self {
        self.scope = scope.into();
        self
    }
}

/// Returns true if the given optype is a `MakeTuple` operation.
///
/// Boilerplate required due to <https://github.com/CQCL/hugr/issues/1496>
fn is_make_tuple(optype: &OpType) -> bool {
    optype.cast::<MakeTuple>().is_some()
}

/// Returns true if the given optype is an `UnpackTuple` operation.
///
/// Boilerplate required due to <https://github.com/CQCL/hugr/issues/1496>
fn is_unpack_tuple(optype: &OpType) -> bool {
    optype.cast::<UnpackTuple>().is_some()
}

/// If this is a `MakeTuple` operation followed by some number of `UnpackTuple` operations
/// on the same region, return a rewrite to remove them.
///
/// Otherwise, return None.
fn make_rewrite<'h, T: HugrView>(
    hugr: &'h T,
    convex_checker: &mut Option<TopoConvexChecker<'h, T>>,
    node: T::Node,
    op: &OpType,
) -> Option<SimpleReplacement<T::Node>> {
    // Only process MakeTuple operations
    if !is_make_tuple(op) {
        return None;
    }

    let has_order_edges = |node: T::Node| -> bool {
        let op = hugr.get_optype(node);
        let has_input_order = op
            .other_input_port()
            .and_then(|p| hugr.linked_outputs(node, p).next())
            .is_some();
        let has_output_order = op
            .other_output_port()
            .and_then(|p| hugr.linked_inputs(node, p).next())
            .is_some();
        has_input_order || has_output_order
    };

    // If the node has order edges, ignore it.
    if has_order_edges(node) {
        return None;
    }

    let tuple_types = op.dataflow_signature().unwrap().input_types().to_vec();
    let node_parent = hugr.get_parent(node);

    // See if it is followed by a tuple unpack
    let links = hugr
        .linked_inputs(node, 0)
        .map(|(neigh, _)| neigh)
        .collect_vec();

    let unpack_nodes = links
        .iter()
        .filter(|&&neigh| hugr.get_parent(neigh) == node_parent)
        .filter(|&&neigh| is_unpack_tuple(hugr.get_optype(neigh)))
        .filter(|&&neigh| !has_order_edges(neigh))
        .copied()
        .collect_vec();

    // If there are no unpacks but the tuple is being used, there's nothing to do.
    if unpack_nodes.is_empty() && !links.is_empty() {
        return None;
    }

    // Remove all unpack operations, and remove the pack operation if all neighbours are unpacks.
    let num_other_outputs = links.len() - unpack_nodes.len();
    Some(remove_pack_unpack(
        hugr,
        convex_checker,
        &tuple_types,
        node,
        unpack_nodes,
        num_other_outputs,
    ))
}

/// Returns a rewrite to remove a tuple pack operation that's followed by unpack operations,
/// and `other_tuple_links` other operations.
fn remove_pack_unpack<'h, T: HugrView>(
    hugr: &'h T,
    convex_checker: &mut Option<TopoConvexChecker<'h, T>>,
    tuple_types: &[Type],
    pack_node: T::Node,
    unpack_nodes: Vec<T::Node>,
    num_other_outputs: usize,
) -> SimpleReplacement<T::Node> {
    let parent = hugr.get_parent(pack_node).expect("pack_node has no parent");
    let checker = convex_checker.get_or_insert_with(|| TopoConvexChecker::new(hugr, parent));

    let mut nodes = unpack_nodes.clone();
    nodes.push(pack_node);
    let subcirc = SiblingSubgraph::try_from_nodes_with_checker(nodes, hugr, checker).unwrap();
    let subcirc_signature = subcirc.signature(hugr);

    let mut replacement = DFGBuilder::new(subcirc_signature).unwrap();

    // Wire the inputs directly to the unpack outputs
    // We need to list the **connected** output ports from the unpack nodes.
    // SiblingSubgraph ignores disconnected outputs, so we need these when building the replacement.
    let mut replacement_outputs =
        Vec::with_capacity(unpack_nodes.len() * tuple_types.len() + num_other_outputs);
    let replacement_inputs = replacement.input_wires().collect_vec();
    for unpack_node in unpack_nodes {
        for out_port in hugr.node_outputs(unpack_node) {
            if hugr.is_linked(unpack_node, out_port) {
                let input = replacement_inputs[out_port.index()];
                replacement_outputs.push(input);
            }
        }
    }

    // If needed, re-add the tuple pack node and connect its output to the tuple outputs.
    if num_other_outputs > 0 {
        let op = MakeTuple::new(tuple_types.to_vec().into());
        let [tuple] = replacement
            .add_dataflow_op(op, replacement.input_wires())
            .unwrap()
            .outputs_arr();
        replacement_outputs.extend(std::iter::repeat_n(tuple, num_other_outputs));
    }

    // These should never fail, as we are defining the replacement ourselves.
    let replacement = replacement
        .finish_hugr_with_outputs(replacement_outputs)
        .unwrap_or_else(|e| {
            panic!("Failed to create replacement for removing tuple pack/unpack operations. {e}")
        });
    subcirc
        .create_simple_replacement(hugr, replacement)
        .unwrap_or_else(|e| {
            panic!("Failed to create rewrite for removing tuple pack/unpack operations. {e}")
        })
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::passes::composable::WithScope;
    use hugr_core::Hugr;
    use hugr_core::builder::FunctionBuilder;
    use hugr_core::extension::prelude::{UnpackTuple, bool_t, qb_t};
    use hugr_core::ops::handle::NodeHandle;
    use hugr_core::std_extensions::arithmetic::float_types::float64_type;
    use hugr_core::types::Signature;
    use rstest::{fixture, rstest};

    /// A simple pack operation with unused output.
    ///
    /// These can be removed entirely.
    #[fixture]
    fn unused_pack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new(vec![bool_t(), bool_t()], vec![])).unwrap();
        let mut inps = h.input_wires();
        let b1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let _tuple = h.make_tuple([b1, b2]).unwrap();

        h.finish_hugr_with_outputs([]).unwrap()
    }

    /// A simple pack operation followed by an unpack operation.
    ///
    /// These can be removed entirely.
    #[fixture]
    fn simple_pack_unpack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new_endo([qb_t(), bool_t()])).unwrap();
        let mut inps = h.input_wires();
        let qb1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let tuple = h.make_tuple([qb1, b2]).unwrap();

        let op = UnpackTuple::new(vec![qb_t(), bool_t()].into());
        let [qb1, b2] = h.add_dataflow_op(op, [tuple]).unwrap().outputs_arr();

        h.finish_hugr_with_outputs([qb1, b2]).unwrap()
    }

    /// A simple pack/unpack pair with order edges between them.
    ///
    /// In the future we should be able to preserve some order edges, but for now
    /// we just remove everything.
    #[fixture]
    fn ordered_pack_unpack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new_endo(vec![qb_t(), bool_t()])).unwrap();
        let mut inps = h.input_wires();
        let qb1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let tuple = h.make_tuple([qb1, b2]).unwrap();
        h.set_order(&h.input(), &tuple.node());

        let op = UnpackTuple::new(vec![qb_t(), bool_t()].into());
        let untuple = h.add_dataflow_op(op, [tuple]).unwrap();
        let [qb1, b2] = untuple.outputs_arr();
        h.set_order(&tuple.node(), &untuple.node());

        h.set_order(&untuple.node(), &h.output());
        h.finish_hugr_with_outputs([qb1, b2]).unwrap()
    }

    /// A simple pack/unpack pair with an order from the pack node to a downstream node.
    ///
    /// The order edge should be preserved, so we move it to the predecessor of the pack node.
    #[fixture]
    fn outgoing_ordered_pack_unpack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new_endo(vec![qb_t(), bool_t()])).unwrap();
        let mut inps = h.input_wires();
        let qb1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let tuple = h.make_tuple([qb1, b2]).unwrap();

        let op = UnpackTuple::new(vec![qb_t(), bool_t()].into());
        let untuple = h.add_dataflow_op(op, [tuple]).unwrap();
        let [qb1, b2] = untuple.outputs_arr();

        h.set_order(&tuple.node(), &h.output());
        h.finish_hugr_with_outputs([qb1, b2]).unwrap()
    }

    /// A simple pack/unpack pair with an order from a downstream node to the pack node.
    ///
    /// The order edge should be preserved, so we move it to the successor of the unpack node.
    #[fixture]
    fn incoming_ordered_pack_unpack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new_endo(vec![qb_t(), bool_t()])).unwrap();
        let mut inps = h.input_wires();
        let qb1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let tuple = h.make_tuple([qb1, b2]).unwrap();

        let op = UnpackTuple::new(vec![qb_t(), bool_t()].into());
        let untuple = h.add_dataflow_op(op, [tuple]).unwrap();
        let [qb1, b2] = untuple.outputs_arr();

        h.set_order(&h.input(), &untuple.node());
        h.finish_hugr_with_outputs([qb1, b2]).unwrap()
    }

    /// A pack operation followed by three unpack operations from the same tuple.
    ///
    /// These can be removed entirely.
    #[fixture]
    fn multi_unpack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new(
            vec![bool_t(), bool_t()],
            vec![bool_t(), bool_t(), bool_t(), bool_t()],
        ))
        .unwrap();
        let mut inps = h.input_wires();
        let b1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let tuple = h.make_tuple([b1, b2]).unwrap();

        let op = UnpackTuple::new(vec![bool_t(), bool_t()].into());
        let [b1, b2] = h.add_dataflow_op(op, [tuple]).unwrap().outputs_arr();

        let op = UnpackTuple::new(vec![bool_t(), bool_t()].into());
        let [b3, b4] = h.add_dataflow_op(op, [tuple]).unwrap().outputs_arr();

        // The last one's outputs are disconnected.
        // TODO: Adding this causes the test to fail due to a `NonCovex` error.
        //let op = UnpackTuple::new(vec![bool_t(), bool_t()].into());
        //let _ = h.add_dataflow_op(op, [tuple]).unwrap();

        h.finish_hugr_with_outputs([b1, b2, b3, b4]).unwrap()
    }

    /// A pack operation followed by an unpack operation, where the tuple is also returned.
    ///
    /// The unpack operation can be removed, but the pack operation cannot.
    #[fixture]
    fn partial_unpack() -> Hugr {
        let mut h = DFGBuilder::new(Signature::new(
            vec![bool_t(), bool_t()],
            vec![
                bool_t(),
                bool_t(),
                Type::new_tuple(vec![bool_t(), bool_t()]),
            ],
        ))
        .unwrap();
        let mut inps = h.input_wires();
        let b1 = inps.next().unwrap();
        let b2 = inps.next().unwrap();

        let tuple = h.make_tuple([b1, b2]).unwrap();

        let op = UnpackTuple::new(vec![bool_t(), bool_t()].into());
        let [b1, b2] = h.add_dataflow_op(op, [tuple]).unwrap().outputs_arr();

        h.finish_hugr_with_outputs([b1, b2, tuple]).unwrap()
    }

    /// A pack operation followed by an unpack that discards its first output.
    ///
    /// The unpack operation can be removed, but the pack operation cannot.
    ///
    /// This is a minimal error case for <https://github.com/Quantinuum/tket2/issues/1347>.
    #[fixture]
    fn unpack_discard_first() -> Hugr {
        let mut h = FunctionBuilder::new(
            "test",
            Signature::new(vec![bool_t(), float64_type()], vec![float64_type()]),
        )
        .unwrap();
        let [b, f] = h.input_wires_arr();

        let tuple = h.make_tuple([b, f]).unwrap();

        let op = UnpackTuple::new(vec![bool_t(), float64_type()].into());
        let [_b, f] = h.add_dataflow_op(op, [tuple]).unwrap().outputs_arr();

        h.finish_hugr_with_outputs([f]).unwrap()
    }

    #[rstest]
    #[case::unused(unused_pack(), 1, 2)]
    #[case::simple(simple_pack_unpack(), 1, 2)]
    #[case::multi(multi_unpack(), 1, 2)]
    #[case::partial(partial_unpack(), 1, 3)]
    #[case::unpack_discard_first(unpack_discard_first(), 1, 2)]
    // Nodes with order edges are ignored.
    #[case::ordered(ordered_pack_unpack(), 0, 4)]
    #[case::outgoing_ordered(outgoing_ordered_pack_unpack(), 0, 4)]
    #[case::incoming_ordered(incoming_ordered_pack_unpack(), 0, 4)]
    fn test_pack_unpack(
        #[case] mut hugr: Hugr,
        #[case] expected_rewrites: usize,
        #[case] remaining_nodes: usize,
    ) {
        let parent = hugr.entrypoint();
        let pass = UntuplePass::default().with_scope(PassScope::EntrypointFlat);
        let res = pass.run(&mut hugr).unwrap_or_else(|e| panic!("{e}"));
        assert_eq!(res.rewrites_applied, expected_rewrites);
        assert_eq!(hugr.children(parent).count(), remaining_nodes);
    }
}