tket 0.21.1

//! A pass that normalizes the structure of Guppy-generated circuits into something that can be optimized by tket.

use crate::passes::composable::WithScope;
use crate::passes::const_fold::{ConstFoldError, ConstantFoldPass};
use crate::passes::dead_funcs::RemoveDeadFuncsError;
use crate::passes::inline_funcs::{InlineFuncsError, InlineFuncsHeuristic, InlineFunctionsPass};
use crate::passes::modifier_resolver::{ModifierResolverErrors, ModifierResolverPass};
use crate::passes::normalize_cfgs::{NormalizeCFGError, NormalizeCFGPass};
use crate::passes::redundant_order_edges::RedundantOrderEdgesPass;
use crate::passes::untuple::{UntupleError, UntuplePass};
use crate::passes::{ComposablePass, InlineDFGsPass, PassScope, RemoveDeadFuncsPass};
use hugr::Node;
use hugr::hugr::HugrError;
use hugr::hugr::hugrmut::HugrMut;
use hugr::hugr::patch::inline_dfg::InlineDFGError;

use crate::passes::BorrowSquashPass;

/// Normalize the structure of a Guppy-generated circuit into something that can be optimized by tket.
///
/// This is a mixture of global optimization passes, and operations that optimize the entrypoint.
#[derive(Clone, Debug)]
pub struct NormalizeGuppy {
    /// Whether to resolve modifier operations.
    resolve_modifiers: bool,
    /// Whether to simplify CFG control flow.
    simplify_cfgs: bool,
    /// Whether to remove tuple/untuple operations.
    untuple: bool,
    /// Whether to constant fold the program.
    constant_fold: bool,
    /// Whether to remove dead functions.
    dead_funcs: bool,
    /// Whether to inline function calls (converting to DFGs).
    inline_funcs: Option<InlineFuncsHeuristic>,
    /// Whether to inline DFG operations.
    inline_dfgs: bool,
    /// Whether to squash BorrowArray borrow/return ops
    squash_borrows: bool,
    /// Whether to remove redundant order edges.
    remove_redundant_order_edges: bool,

    /// Where to apply the pass.
    ///
    /// Configurable via [`WithScope::with_scope`].
    scope: PassScope,
}

impl NormalizeGuppy {
    /// Set whether to resolve modifier operations.
    pub fn resolve_modifiers(&mut self, resolve_modifiers: bool) -> &mut Self {
        self.resolve_modifiers = resolve_modifiers;
        self
    }
    /// Set whether to simplify CFG control flow.
    pub fn simplify_cfgs(&mut self, simplify_cfgs: bool) -> &mut Self {
        self.simplify_cfgs = simplify_cfgs;
        self
    }
    /// Set whether to remove tuple/untuple operations.
    pub fn remove_tuple_untuple(&mut self, untuple: bool) -> &mut Self {
        self.untuple = untuple;
        self
    }
    /// Set whether to constant fold the program.
    pub fn constant_folding(&mut self, constant_fold: bool) -> &mut Self {
        self.constant_fold = constant_fold;
        self
    }
    /// Set whether to remove dead functions.
    pub fn remove_dead_funcs(&mut self, dead_funcs: bool) -> &mut Self {
        self.dead_funcs = dead_funcs;
        self
    }
    /// Set whether to inline DFG operations.
    pub fn inline_dfgs(&mut self, inline: bool) -> &mut Self {
        self.inline_dfgs = inline;
        self
    }
    /// Set whether to inline Function calls. (Does not include [Self::inline_dfgs]
    /// but generates DFGs so the latter is strongly recommended.)
    pub fn inline_funcs(&mut self, inline: Option<InlineFuncsHeuristic>) -> &mut Self {
        self.inline_funcs = inline;
        self
    }
    /// Set whether to squash BorrowArray borrow/return ops
    pub fn squash_borrows(&mut self, squash: bool) -> &mut Self {
        self.squash_borrows = squash;
        self
    }
    /// Set whether to remove redundant order edges.
    pub fn remove_redundant_order_edges(&mut self, remove: bool) -> &mut Self {
        self.remove_redundant_order_edges = remove;
        self
    }
}

impl Default for NormalizeGuppy {
    fn default() -> Self {
        Self {
            resolve_modifiers: true,
            inline_funcs: Some(InlineFuncsHeuristic::default()),
            simplify_cfgs: true,
            constant_fold: true,
            untuple: true,
            dead_funcs: true,
            inline_dfgs: true,
            squash_borrows: true,
            remove_redundant_order_edges: true,
            scope: PassScope::default(),
        }
    }
}

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

impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for NormalizeGuppy {
    type Error = NormalizeGuppyErrors;
    type Result = ();

    fn run(&self, hugr: &mut H) -> Result<Self::Result, Self::Error> {
        // Simplify CFGs first, as (until we start removing statically-impossible branches)
        // nothing else affects CFG structure or creates new opportunities for this.
        // (Possibly also this may assist modifier resolution??)
        if self.simplify_cfgs {
            NormalizeCFGPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)?;
        }
        // Run modifier resolution
        if self.resolve_modifiers {
            ModifierResolverPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)?;
        }
        // Inline function calls creates many opportunities for optimization by other
        // passes by producing copies that can be optimized in a specific context
        if let Some(inline_funcs) = &self.inline_funcs {
            InlineFunctionsPass::default_with_scope(self.scope.clone())
                .with_heuristic(inline_funcs.clone())
                .run(hugr)
                .map_err(NormalizeGuppyErrors::InlineFuncs)?;
        }
        // Clean up after inlining - only to improve compilation speed, not affected by
        // anything else until we start removing untaken branches.
        if self.dead_funcs {
            RemoveDeadFuncsPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)?;
        }
        // Function inlining produces lots of DFGs, so merge those into their surrounds
        if self.inline_dfgs {
            InlineDFGsPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)
                .unwrap_or_else(|e| match e {})
        }
        // This should sort out argument marshalling for function calls (esp. inlined ones)
        if self.untuple {
            UntuplePass::default_with_scope(self.scope.clone()).run(hugr)?;
        }
        // Should propagate through untuple, so could do earlier, but not clear earlier
        // would be any advantage. Must be before BorrowSquash as that needs constant indices.
        if self.constant_fold {
            ConstantFoldPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)?;
        }
        // Potentially, could (need to) do fixpoint here with untuple,
        // as both create opportunities for the other
        if self.squash_borrows {
            BorrowSquashPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)
                .unwrap_or_else(|e| match e {});
        }
        // Remove redundant order edges once all other structural rewrites have been applied.
        if self.remove_redundant_order_edges {
            RedundantOrderEdgesPass::default()
                .with_scope(self.scope.clone())
                .run(hugr)
                .map_err(NormalizeGuppyErrors::RedundantOrderEdges)?;
        }

        Ok(())
    }
}

/// Errors that can occur during the guppy-program normalization process.
#[derive(derive_more::Error, Debug, derive_more::Display, derive_more::From)]
#[non_exhaustive]
pub enum NormalizeGuppyErrors {
    /// Error while resolving modifier operations.
    ModifierResolver(ModifierResolverErrors),
    /// Error while simplifying CFG control flow.
    SimplifyCFG(NormalizeCFGError),
    /// Error while removing tuple/untuple operations.
    Untuple(UntupleError),
    /// Error while constant folding.
    ConstantFold(ConstFoldError),
    /// Error while removing dead functions.
    DeadFuncs(RemoveDeadFuncsError),
    /// Error while inlining DFG operations.
    InlineDFGs(InlineDFGError),
    /// Error while inlining function calls.
    InlineFuncs(InlineFuncsError),
    /// Error while removing redundant order edges.
    #[from(ignore)]
    RedundantOrderEdges(HugrError),
}

#[cfg(test)]
mod test {
    use hugr::builder::{Dataflow, DataflowHugr, FunctionBuilder};
    use hugr::extension::prelude::qb_t;
    use hugr::types::Signature;

    use crate::TketOp;

    use super::*;

    /// Running the pass with all options disabled should still work (and do nothing).
    #[test]
    fn test_guppy_pass_noop() {
        let mut b = FunctionBuilder::new("main", Signature::new_endo(vec![qb_t()])).unwrap();
        let [q] = b.input_wires_arr();
        let [q] = b.add_dataflow_op(TketOp::H, [q]).unwrap().outputs_arr();
        let hugr = b.finish_hugr_with_outputs([q]).unwrap();

        let mut hugr2 = hugr.clone();
        NormalizeGuppy::default()
            .resolve_modifiers(false)
            .simplify_cfgs(false)
            .remove_tuple_untuple(false)
            .constant_folding(false)
            .remove_dead_funcs(false)
            .inline_dfgs(false)
            .remove_redundant_order_edges(false)
            .run(&mut hugr2)
            .unwrap();

        assert_eq!(hugr2, hugr);
    }
}