tket 0.18.0

#![allow(clippy::type_complexity)]
//! Replace types with other types across the Hugr. See [`ReplaceTypes`] and [Linearizer].
//!
use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;

use handlers::list_const;
use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy, OnMultiDefn};
use hugr_core::std_extensions::collections::array::array_type_def;
use hugr_core::std_extensions::collections::list::list_type_def;
use itertools::Either;
use thiserror::Error;

use hugr_core::builder::{
    BuildError, BuildHandle, Container, Dataflow, DataflowHugr, FunctionBuilder, HugrBuilder,
};
use hugr_core::extension::{ExtensionId, OpDef, SignatureError, TypeDef};
use hugr_core::hugr::hugrmut::HugrMut;
use hugr_core::ops::constant::{OpaqueValue, Sum};
use hugr_core::ops::handle::{DataflowOpID, FuncID, NodeHandle};
use hugr_core::ops::{
    AliasDefn, CFG, Call, CallIndirect, Case, Conditional, Const, DFG, DataflowBlock, ExitBlock,
    ExtensionOp, Input, LoadConstant, LoadFunction, OpTrait, OpType, Output, Tag, TailLoop, Value,
};
use hugr_core::types::{
    ConstTypeError, CustomType, Signature, Transformable, Type, TypeArg, TypeEnum, TypeRow,
    TypeTransformer,
};
use hugr_core::{Direction, Hugr, HugrView, Node, PortIndex, Visibility, Wire};

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

mod linearize;
pub use linearize::{CallbackHandler, DelegatingLinearizer, LinearizeError, Linearizer};

/// A recipe for creating a dataflow Node - as a new child of a [`DataflowParent`]
/// or in order to replace an existing node.
///
/// [`DataflowParent`]: hugr_core::ops::OpTag::DataflowParent
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub enum NodeTemplate {
    /// A single node - so if replacing an existing node, change only the op.
    ///
    /// An error will be raised if the new op has a static inport; use [Self::LinkedHugr] instead.
    SingleOp(OpType),
    /// Defines a sub-Hugr whose entrypoint-subtree to insert, the entrypoint (which must be
    /// a [CFG], [Conditional], [DFG] or [`TailLoop`]) becoming (/replacing) the desired Node.
    // Not a FuncDefn, nor Case/DataflowBlock
    /// Note 1. `CompoundOp` will be of limited use for polymorphic ops (before [monomorphization])
    /// because the new subtree will not be able to use type variables present in the
    /// parent Hugr or previous op.
    ///
    /// Edges incoming to the entrypoint subtree will be disconnected.
    ///
    /// It is **recommended** to use [Self::LinkedHugr] instead.
    ///
    /// [monomorphization]: crate::passes::MonomorphizePass
    CompoundOp(Box<Hugr>),
    /// Defines a sub-Hugr to insert, whose entrypoint becomes (or replaces) the desired Node.
    /// Other children of the Hugr reachable from the entrypoint will also be inserted
    /// according to the specified linking policy.
    LinkedHugr(Box<Hugr>, NameLinkingPolicy),
}

impl NodeTemplate {
    /// Creates a [Self::LinkedHugr] from the given Hugr with the default linking policy.
    pub fn linked_hugr(h: impl Into<Hugr>) -> Self {
        NodeTemplate::LinkedHugr(Box::new(h.into()), NameLinkingPolicy::default())
    }

    /// Creates a [`NodeTemplate::LinkedHugr`] creating a call node to the given
    /// function definition or declaration.
    ///
    /// Returns a [`BuildError::UnexpectedType`] if the given Hugr does not have
    /// a function definition or declaration as entrypoint.
    pub fn call_to_function(
        func_def: Hugr,
        type_args: &[TypeArg],
    ) -> Result<NodeTemplate, BuildError> {
        // Create a replacement hugr for the op nodes: Add a `call` node in the `func_def` hugr and set it as entrypoint.
        let func_op = func_def.entrypoint_optype();
        let func_signature = match func_op {
            OpType::FuncDecl(decl) => decl.signature().clone(),
            OpType::FuncDefn(defn) => defn.signature().clone(),
            _ => {
                return Err(BuildError::UnexpectedType {
                    node: func_def.entrypoint(),
                    op_desc: "function definition or declaration",
                });
            }
        }
        .instantiate(type_args)?;

        // Build a new hugr and insert the function definition into it
        let mut b = FunctionBuilder::new_vis("", func_signature, Visibility::Private).unwrap();
        let func_id = FuncID::<true>::from(
            b.module_root_builder()
                .add_hugr(func_def)
                .inserted_entrypoint,
        );

        // Build a call to the function in the new separate function.
        let call = b.call(&func_id, type_args, b.input_wires()).unwrap();
        let mut call_hugr = b.finish_hugr_with_outputs(call.outputs()).unwrap();
        call_hugr.set_entrypoint(call.node());

        Ok(NodeTemplate::LinkedHugr(
            Box::new(call_hugr),
            NameLinkingPolicy::default().on_multiple_defn(OnMultiDefn::UseTarget),
        ))
    }

    /// Adds this instance to the specified [`HugrMut`] as a new node or subtree under a
    /// given parent, returning the unique new child (of that parent) thus created
    ///
    /// # Panics
    ///
    /// * If `parent` is not in the `hugr`
    ///
    /// [`signature`]: hugr_core::types::PolyFuncType
    /// [`FuncDecl`]: hugr_core::ops::FuncDecl
    /// [`FuncDefn`]: hugr_core::ops::FuncDefn
    pub fn add_hugr(
        self,
        hugr: &mut impl HugrMut<Node = Node>,
        parent: Node,
    ) -> Result<Node, BuildError> {
        match self {
            NodeTemplate::SingleOp(op_type) => Ok(hugr.add_node_with_parent(parent, op_type)),
            NodeTemplate::CompoundOp(new_h) => {
                Ok(hugr.insert_hugr(parent, *new_h).inserted_entrypoint)
            }
            NodeTemplate::LinkedHugr(h, pol) => {
                Ok(hugr.insert_link_hugr(parent, *h, &pol)?.inserted_entrypoint)
            }
        }
    }

    /// Adds this instance to the specified [Dataflow] builder as a new node or subtree
    pub fn add(
        self,
        dfb: &mut impl Dataflow,
        inputs: impl IntoIterator<Item = Wire>,
    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
        match self {
            NodeTemplate::SingleOp(opty) => dfb.add_dataflow_op(opty, inputs),
            NodeTemplate::CompoundOp(h) => dfb.add_hugr_with_wires(*h, inputs),
            NodeTemplate::LinkedHugr(h, pol) => dfb.add_link_hugr_with_wires(*h, &pol, inputs),
        }
    }

    fn replace<H: HugrMut<Node = Node>>(
        self,
        hugr: &mut H,
        n: Node,
        rt: &ReplaceTypes,
        opts: &ReplacementOptions,
    ) -> Result<(), ReplaceTypesError> {
        let ef = |e| ReplaceTypesError::AddTemplateError(n, Box::new(e));
        assert_eq!(hugr.children(n).count(), 0);
        let (new_optype, static_source, static_inport) = match self {
            NodeTemplate::SingleOp(op_type) => {
                if op_type.static_input_port().is_some() {
                    return Err(ef(BuildError::UnexpectedType {
                        node: n,
                        op_desc: "Replacement SingleOp without static input",
                    }));
                }
                (op_type, None, None)
            }
            NodeTemplate::CompoundOp(new_h) => {
                let root = new_h.entrypoint_optype();
                if !matches!(
                    root,
                    OpType::CFG(_) | OpType::DFG(_) | OpType::Conditional(_) | OpType::TailLoop(_)
                )
                //if !root.is_container() || !root.dataflow_signature().is_some() // Using explicit list as per docs
                {
                    return Err(ef(BuildError::UnexpectedType {
                        node: n,
                        op_desc: "Replacement CompoundOp not a container/dataflow node",
                    }));
                }
                assert!(root.static_input_port().is_none());
                let new_entrypoint = hugr.insert_hugr(n, *new_h).inserted_entrypoint;
                let children = hugr.children(new_entrypoint).collect::<Vec<_>>();
                let root_opty = hugr.remove_node(new_entrypoint);
                for ch in children {
                    hugr.set_parent(ch, n);
                }
                (root_opty, None, None)
            }
            NodeTemplate::LinkedHugr(mut h, pol) => {
                // We have to recursively process any children that *might* be linked in,
                // before linking, as otherwise they'll signature conflict with other,
                // already-recursively-processed, functions with which they might be linked.
                let mut containing_func = h.entrypoint();
                while let Some(parent) = h.get_parent(containing_func)
                    && !h.get_optype(parent).is_module()
                {
                    containing_func = parent;
                }

                for ch in h.children(h.module_root()).collect::<Vec<_>>() {
                    if ch != containing_func {
                        rt.process_subtree_opts(&mut h, ch, opts)?;
                    }
                }
                let new_entrypoint = hugr
                    .insert_link_hugr(n, *h, &pol)
                    .map_err(|e| ef(BuildError::from(e)))?
                    .inserted_entrypoint;
                let children = hugr.children(new_entrypoint).collect::<Vec<_>>();
                let static_source = hugr.static_source(new_entrypoint);
                let root_opty = hugr.remove_node(new_entrypoint);
                let static_inport = root_opty.static_input_port();
                for ch in children {
                    hugr.set_parent(ch, n);
                }
                (root_opty, static_source, static_inport)
            }
        };
        *hugr.optype_mut(n) = new_optype;
        if let Some(static_inport) = static_inport {
            hugr.insert_ports(n, Direction::Incoming, static_inport.index(), 1);
            if let Some(static_source) = static_source {
                hugr.connect(static_source, 0, n, static_inport);
            }
        }
        rt.process_subtree_opts(hugr, n, opts)?;
        Ok(())
    }

    fn check_signature(
        &self,
        inputs: &TypeRow,
        outputs: &TypeRow,
    ) -> Result<(), Option<Signature>> {
        let sig = match self {
            NodeTemplate::SingleOp(op_type) => op_type,
            NodeTemplate::CompoundOp(hugr) => hugr.entrypoint_optype(),
            NodeTemplate::LinkedHugr(hugr, _) => hugr.entrypoint_optype(),
        }
        .dataflow_signature();
        if sig.as_deref().map(Signature::io) == Some((inputs, outputs)) {
            Ok(())
        } else {
            Err(sig.map(Cow::into_owned))
        }
    }
}

/// Options for how a replacement (op or type) is processed.
///
/// May be specified by [ReplaceTypes::set_replace_op] or [ReplaceTypes::set_replace_type].
/// Otherwise (the default), replacements are inserted as given (without further processing).
#[derive(Clone, Default, PartialEq, Eq)] // More derives might inhibit future extension
pub struct ReplacementOptions {
    process_recursive: bool,
    linearize_unchanged: bool,
}

impl ReplacementOptions {
    fn recursive() -> Self {
        Self {
            process_recursive: true,
            linearize_unchanged: false,
        }
    }

    /// Specifies whether all nodes within the replacement should have their
    /// output ports linearized.
    pub fn with_linearization(mut self, lin: bool) -> Self {
        self.linearize_unchanged = lin;
        self
    }
}

/// A *lowering* [ComposablePass] that replaces types, ops and constants, i.e. changing
/// node signatures/interfaces.
///
/// The struct configures what types, ops, and constants should be replaced with what,
/// and may be applied to a Hugr via [`Self::run`].
///
/// Parametrized types and ops will be reparameterized taking into account the
/// replacements, but any ops taking/returning the replaced types *not* as a result of
/// parametrization, will also need to be replaced - see [`Self::set_replace_op`].
/// Similarly [Const]s.
///
/// Types that are [Copyable](hugr_core::types::TypeBound::Copyable) may also be replaced
/// with types that are not, see [Linearizer].
///
/// Note that although this pass may be used before [monomorphization], there are some
/// limitations (that do not apply if done after [monomorphization]):
/// * [`NodeTemplate::CompoundOp`] only works for operations that do not use type variables
/// * "Overrides" of specific instantiations of polymorphic types will not be detected if
///   the instantiations are created inside polymorphic functions. For example, suppose
///   we [`Self::set_replace_type`] type `A` with `X`, [`Self::set_replace_parametrized_type`]
///   container `MyList` with `List`, and [`Self::set_replace_type`] `MyList<A>` with
///   `SpecialListOfXs`. If a function `foo` polymorphic over a type variable `T` dealing
///   with `MyList<T>`s, that is called with type argument `A`, then `foo<T>` will be
///   updated to deal with `List<T>`s and the call `foo<A>` updated to `foo<X>`, but this
///   will still result in using `List<X>` rather than `SpecialListOfXs`. (However this
///   would be fine *after* [monomorphization]: the monomorphic definition of `foo_A`
///   would use `SpecialListOfXs`.)
/// * See also limitations noted for [Linearizer].
///
/// [monomorphization]: crate::passes::MonomorphizePass
#[derive(Clone)]
pub struct ReplaceTypes {
    type_map: HashMap<CustomType, (Type, ReplacementOptions)>,
    param_types:
        HashMap<ParametricType, (Arc<dyn Fn(&[TypeArg]) -> Option<Type>>, ReplacementOptions)>,
    linearize: DelegatingLinearizer,
    op_map: HashMap<OpHashWrapper, (NodeTemplate, ReplacementOptions)>,
    param_ops: HashMap<
        ParametricOp,
        (
            Arc<
                dyn Fn(
                    &[TypeArg],
                    &ReplaceTypes,
                ) -> Result<Option<NodeTemplate>, ReplaceTypesError>,
            >,
            ReplacementOptions,
        ),
    >,
    consts: HashMap<
        CustomType,
        Arc<dyn Fn(&OpaqueValue, &ReplaceTypes) -> Result<Value, ReplaceTypesError>>,
    >,
    param_consts: HashMap<
        ParametricType,
        Arc<dyn Fn(&OpaqueValue, &ReplaceTypes) -> Result<Option<Value>, ReplaceTypesError>>,
    >,
    scope: Either<PassScope, Vec<Node>>,
}

impl Default for ReplaceTypes {
    fn default() -> Self {
        let mut res = Self::new_empty();
        res.linearize = DelegatingLinearizer::default();
        res.replace_consts_parametrized(array_type_def(), handlers::array_const);
        res.replace_consts_parametrized(list_type_def(), list_const);
        res
    }
}

impl TypeTransformer for ReplaceTypes {
    type Err = ReplaceTypesError;

    fn apply_custom(&self, ct: &CustomType) -> Result<Option<Type>, Self::Err> {
        let mut ty_and_opts = None;
        if let Some(res) = self.type_map.get(ct) {
            ty_and_opts = Some(res.clone())
        } else if let Some((dest_fn, opts)) = self.param_types.get(&ct.into()) {
            // `ct` has not had args transformed
            let mut nargs = ct.args().to_vec();
            // We don't care if `nargs` are changed, we're just calling `dest_fn`
            nargs
                .iter_mut()
                .try_for_each(|ta| ta.transform(self).map(|_ch| ()))?;
            ty_and_opts = dest_fn(&nargs).map(|ty| (ty, opts.clone()))
        };
        let Some((mut ty, opts)) = ty_and_opts else {
            return Ok(None);
        };
        if opts.process_recursive {
            ty.transform(self)?;
        }
        Ok(Some(ty))
    }
}

/// An error produced by the [`ReplaceTypes`] pass
#[derive(Debug, Error, PartialEq)]
#[non_exhaustive]
#[expect(missing_docs)]
pub enum ReplaceTypesError {
    #[error(transparent)]
    SignatureError(#[from] SignatureError),
    #[error(transparent)]
    ConstError(#[from] ConstTypeError),
    #[error(transparent)]
    LinearizeError(#[from] LinearizeError),
    #[error("Replacement op for {0} could not be added because {1}")]
    AddTemplateError(Node, Box<BuildError>),
}

impl ReplaceTypes {
    /// Makes a new instance. Unlike [`Self::default`], this does not understand
    /// any extension types, even those in the prelude.
    #[must_use]
    pub fn new_empty() -> Self {
        Self {
            type_map: Default::default(),
            param_types: Default::default(),
            linearize: DelegatingLinearizer::new_empty(),
            op_map: Default::default(),
            param_ops: Default::default(),
            consts: Default::default(),
            param_consts: Default::default(),
            // Not really clear what "preserve" means for a pass that changes signatures,
            // but default to running on whole hugr not just entrypoint.
            scope: Either::Left(PassScope::default()),
        }
    }

    /// Configures this instance to replace occurrences of type `src` with `dest`.
    ///
    /// `dest` will be recursively transformed by this [ReplaceTypes] before replacement.
    /// (Cases where a type should be replaced by a type containing an instance of
    /// the first type, must be handled by two separate [ReplaceTypes]'s via a temporary
    /// type. )
    ///
    /// Note that if `src` is an instance of a *parametrized* [`TypeDef`], this takes
    /// precedence over [`Self::set_replace_parametrized_type`] where the `src`s overlap. Thus, this
    /// should only be used on already-*[monomorphize](crate::passes::MonomorphizePass)d* Hugrs, as
    /// substitution (parametric polymorphism) happening later will not respect this replacement.
    ///
    /// If there are any [`LoadConstant`]s of this type, callers should also call [`Self::replace_consts`]
    /// (or [`Self::replace_consts_parametrized`]) as the [`LoadConstant`]s will be reparameterized
    /// (and this will break the edge from [Const] to [`LoadConstant`]).
    ///
    /// Note that if `src` is Copyable and `dest` is Linear, then (besides linearity violations)
    /// [`SignatureError`] will be raised if this leads to an impossible type e.g. ArrayOfCopyables(src).
    /// (This can be overridden by an additional [`Self::set_replace_type`].)
    pub fn set_replace_type(&mut self, src: CustomType, dest: Type) {
        // We could check that 'dest' is copyable, 'src' is linear, or relevant copy and
        // discard functions are registered with the linearizer; but since we can't check
        // that for parametrized types, we'll be consistent and not check here either.
        self.type_map
            .insert(src, (dest, ReplacementOptions::recursive()));
    }

    /// Configures this instance to change occurrences of a parametrized type `src`
    /// via a callback that builds the replacement type given the [`TypeArg`]s.
    ///
    /// Note that the `TypeArgs` will already have been updated (e.g. they may not
    /// fit the bounds of the original type). The callback may return `None` to indicate
    /// no change (in which case the supplied `TypeArgs` will be given to `src`).
    /// The returned type will also be subject to recursive processing by this [ReplaceTypes].
    /// (Cases where a type should be replaced by a type containing an instance of
    /// the first type, must be handled by two separate [ReplaceTypes]'s via a temporary
    /// type. )
    ///
    /// If there are any [`LoadConstant`]s of any of these types, callers should also call
    /// [`Self::replace_consts_parametrized`] (or [`Self::replace_consts`]) as the
    /// [`LoadConstant`]s will be reparameterized (and this will break the edge from [Const] to
    /// [`LoadConstant`]).
    /// See [Self::set_replace_type] for more details (including recursion).
    pub fn set_replace_parametrized_type(
        &mut self,
        src: &TypeDef,
        dest_fn: impl Fn(&[TypeArg]) -> Option<Type> + 'static,
    ) {
        // No way to check that dest_fn never produces a linear type.
        // We could require copy/discard-generators if src is Copyable, or *might be*
        // (depending on arguments - i.e. if src's TypeDefBound is anything other than
        // `TypeDefBound::Explicit(TypeBound::Copyable)`) but that seems an annoying
        // overapproximation. Moreover, these depend upon the *return type* of the Fn.
        // It would be too awkward to require:
        // dest_fn: impl Fn(&TypeArg) -> (Type,
        //                                Fn(&Linearizer) -> NodeTemplate, // copy
        //                                Fn(&Linearizer) -> NodeTemplate)` // discard
        self.param_types.insert(
            src.into(),
            (Arc::new(dest_fn), ReplacementOptions::recursive()),
        );
    }

    /// Allows to configure how to deal with types/wires that were `Copyable`
    /// Allows to configure how to deal with types/wires that were [Copyable]
    /// but have become linear as a result of type-changing. Specifically,
    /// the [Linearizer] is used whenever lowering produces an outport which both
    /// * has a non-[Copyable] type - perhaps a direct substitution, or perhaps e.g.
    ///   as a result of changing the element type of a collection such as an [`array`]
    /// * has other than one connected inport,
    ///
    /// [Copyable]: hugr_core::types::TypeBound::Copyable
    /// [`array`]: hugr_core::std_extensions::collections::array::array_type
    pub fn linearizer_mut(&mut self) -> &mut DelegatingLinearizer {
        &mut self.linearize
    }

    /// Allows use of the linearizer (e.g. in a callback passed to
    /// [Self::set_replace_parametrized_op])
    pub fn get_linearizer(&self) -> &impl Linearizer {
        &self.linearize
    }

    /// Configures this instance to change occurrences of `src` to `dest`.
    ///
    /// The RHS will be recursively processed by this [ReplaceTypes].
    /// (Cases where an op should be replaced by a container including an
    /// instance of the same op, must be handled by two separate [ReplaceTypes]'s
    /// via a temporary op.)
    ///
    /// Note that if `src` is an instance of a *parametrized* [`OpDef`], this takes
    /// precedence over [`Self::set_replace_parametrized_op`] where the `src`s overlap.
    /// Thus, this method should only be used for already-*[monomorphize](crate::passes::MonomorphizePass)d*
    /// Hugrs, as substitution (parametric polymorphism) happening later will not respect
    /// this replacement.
    pub fn set_replace_op(&mut self, src: &ExtensionOp, dest: NodeTemplate) {
        self.op_map.insert(
            OpHashWrapper::from(src),
            (dest, ReplacementOptions::recursive()),
        );
    }

    /// Configures this instance to change occurrences of a parametrized op `src`
    /// via a callback that builds the replacement type given the [`TypeArg`]s.
    /// Note that the `TypeArgs` will already have been updated (e.g. they may not
    /// fit the bounds of the original op); and the returned [NodeTemplate] will be
    /// recursively processed by this [ReplaceTypes]. (Cases where an op should be
    /// replaced by a container including an instance of the same op, must be handled
    /// by two separate [ReplaceTypes]'s via a temporary op.)
    ///
    /// If the Callback returns None, the new typeargs will be applied to the original op.
    pub fn set_replace_parametrized_op(
        &mut self,
        src: &OpDef,
        dest_fn: impl Fn(&[TypeArg], &ReplaceTypes) -> Result<Option<NodeTemplate>, ReplaceTypesError>
        + 'static,
    ) {
        self.param_ops.insert(
            src.into(),
            (Arc::new(dest_fn), ReplacementOptions::recursive()),
        );
    }

    /// Configures this instance to change [Const]s of type `src_ty`, using
    /// a callback that is passed the value of the constant (of that type).
    ///
    /// Note that if `src_ty` is an instance of a *parametrized* [`TypeDef`],
    /// this takes precedence over [`Self::replace_consts_parametrized`] where
    /// the `src_ty`s overlap.
    pub fn replace_consts(
        &mut self,
        src_ty: CustomType,
        const_fn: impl Fn(&OpaqueValue, &ReplaceTypes) -> Result<Value, ReplaceTypesError> + 'static,
    ) {
        self.consts.insert(src_ty, Arc::new(const_fn));
    }

    /// Configures this instance to change [Const]s of all types that are instances
    /// of a parametrized typedef `src_ty`, using a callback that is passed the
    /// value of the constant (the [`OpaqueValue`] contains the [`TypeArg`]s). The
    /// callback may return `None` to indicate no change to the constant.
    pub fn replace_consts_parametrized(
        &mut self,
        src_ty: &TypeDef,
        const_fn: impl Fn(&OpaqueValue, &ReplaceTypes) -> Result<Option<Value>, ReplaceTypesError>
        + 'static,
    ) {
        self.param_consts.insert(src_ty.into(), Arc::new(const_fn));
    }

    /// Set the regions of the Hugr to which this pass should be applied.
    ///
    /// If not set, the pass is applied to the whole Hugr.
    /// Each call overwrites any previous calls to `set_regions` and/or [Self::with_scope].
    pub fn set_regions(&mut self, regions: impl IntoIterator<Item = Node>) {
        self.scope = Either::Right(regions.into_iter().collect());
    }

    fn process_subtree_opts(
        &self,
        hugr: &mut impl HugrMut<Node = Node>,
        root: Node,
        opts: &ReplacementOptions,
    ) -> Result<(), ReplaceTypesError> {
        if opts.process_recursive {
            self.change_subtree(hugr, root, opts.linearize_unchanged)?;
            // change_subtree does not linearize its root, just as change_node
            // does not linearize the node it's called on; our caller does.
        } else if opts.linearize_unchanged {
            let mut descs = hugr.descendants(root);
            assert_eq!(descs.next(), Some(root));
            for n in descs.collect::<Vec<_>>() {
                self.linearize_outputs(hugr, n)?;
            }
        }
        Ok(())
    }

    fn change_subtree(
        &self,
        hugr: &mut impl HugrMut<Node = Node>,
        root: Node,
        linearize_unchanged_ops: bool,
    ) -> Result<bool, ReplaceTypesError> {
        let mut descs = hugr.descendants(root).collect::<Vec<_>>().into_iter();
        assert_eq!(descs.next(), Some(root));
        let mut changed = self.change_node(hugr, root)?;
        // Do not linearize the root's outputs - that's done by the caller if appropriate,
        // as any copy/discard ops would be *outside* the root
        for n in descs {
            if self.change_node(hugr, n)? {
                changed = true;
            } else if !linearize_unchanged_ops {
                continue;
            }
            self.linearize_outputs(hugr, n)?;
        }
        Ok(changed)
    }

    fn change_node(
        &self,
        hugr: &mut impl HugrMut<Node = Node>,
        n: Node,
    ) -> Result<bool, ReplaceTypesError> {
        match hugr.optype_mut(n) {
            OpType::FuncDefn(fd) => fd.signature_mut().body_mut().transform(self),
            OpType::FuncDecl(fd) => fd.signature_mut().body_mut().transform(self),
            OpType::LoadConstant(LoadConstant { datatype: ty })
            | OpType::AliasDefn(AliasDefn { definition: ty, .. }) => ty.transform(self),

            OpType::ExitBlock(ExitBlock { cfg_outputs: types })
            | OpType::Input(Input { types })
            | OpType::Output(Output { types }) => types.transform(self),
            OpType::LoadFunction(LoadFunction {
                func_sig,
                type_args,
                instantiation,
            })
            | OpType::Call(Call {
                func_sig,
                type_args,
                instantiation,
            }) => {
                let change = func_sig.body_mut().transform(self)? | type_args.transform(self)?;
                if change {
                    let new_inst = func_sig
                        .instantiate(type_args)
                        .map_err(ReplaceTypesError::SignatureError)?;
                    *instantiation = new_inst;
                }
                Ok(change)
            }
            OpType::Case(Case { signature })
            | OpType::CFG(CFG { signature })
            | OpType::DFG(DFG { signature })
            | OpType::CallIndirect(CallIndirect { signature }) => signature.transform(self),
            OpType::Tag(Tag { variants, .. }) => variants.transform(self),
            OpType::Conditional(Conditional {
                other_inputs: row1,
                outputs: row2,
                sum_rows,
                ..
            })
            | OpType::DataflowBlock(DataflowBlock {
                inputs: row1,
                other_outputs: row2,
                sum_rows,
                ..
            }) => Ok(row1.transform(self)? | row2.transform(self)? | sum_rows.transform(self)?),
            OpType::TailLoop(TailLoop {
                just_inputs,
                just_outputs,
                rest,
                ..
            }) => Ok(just_inputs.transform(self)?
                | just_outputs.transform(self)?
                | rest.transform(self)?),

            OpType::Const(Const { value, .. }) => self.change_value(value),
            OpType::ExtensionOp(ext_op) => Ok({
                let def = ext_op.def_arc();
                let mut changed = false;
                let replacement = match self.op_map.get(&OpHashWrapper::from(&*ext_op)) {
                    r @ Some(_) => r.cloned(),
                    None => {
                        let mut args = ext_op.args().to_vec();
                        changed = args.transform(self)?;
                        let r2 = match self.param_ops.get(&def.as_ref().into()) {
                            None => None,
                            Some((rep_fn, opts)) => {
                                rep_fn(&args, self)?.map(|nt| (nt, opts.clone()))
                            }
                        };
                        if r2.is_none() && changed {
                            *ext_op = ExtensionOp::new(def.clone(), args)?;
                        }
                        r2
                    }
                };
                if let Some((replacement, opts)) = replacement {
                    replacement.replace(hugr, n, self, &opts)?;
                    true
                } else {
                    changed
                }
            }),

            OpType::OpaqueOp(_) => panic!("OpaqueOp should not be in a Hugr"),

            OpType::AliasDecl(_) | OpType::Module(_) => Ok(false),
            _ => todo!(),
        }
    }

    /// Modifies the specified Value in-place according to current configuration.
    /// Returns whether the value has changed (conservative over-approximation).
    pub fn change_value(&self, value: &mut Value) -> Result<bool, ReplaceTypesError> {
        match value {
            Value::Sum(Sum {
                values, sum_type, ..
            }) => {
                let mut any_change = false;
                for value in values {
                    any_change |= self.change_value(value)?;
                }
                any_change |= sum_type.transform(self)?;
                Ok(any_change)
            }
            Value::Extension { e } => Ok({
                let new_const = match e.get_type().as_type_enum() {
                    TypeEnum::Extension(exty) => match self.consts.get(exty) {
                        Some(const_fn) => Some(const_fn(e, self)),
                        None => self
                            .param_consts
                            .get(&exty.into())
                            .and_then(|const_fn| const_fn(e, self).transpose()),
                    },
                    _ => None,
                };
                if let Some(new_const) = new_const {
                    *value = new_const?;
                    true
                } else {
                    false
                }
            }),
        }
    }

    fn linearize_outputs<H: HugrMut<Node = Node>>(
        &self,
        hugr: &mut H,
        n: H::Node,
    ) -> Result<(), LinearizeError> {
        if let Some(new_sig) = hugr.get_optype(n).dataflow_signature() {
            let new_sig = new_sig.into_owned();
            for outp in new_sig.output_ports() {
                if !new_sig.out_port_type(outp).unwrap().copyable() {
                    let targets = hugr.linked_inputs(n, outp).collect::<Vec<_>>();
                    if targets.len() != 1 {
                        hugr.disconnect(n, outp);
                        let src = Wire::new(n, outp);
                        self.linearize.insert_copy_discard(hugr, src, &targets)?;
                    }
                }
            }
        }
        Ok(())
    }
}

impl<H: HugrMut<Node = Node>> ComposablePass<H> for ReplaceTypes {
    type Error = ReplaceTypesError;
    type Result = bool;

    fn run(&self, hugr: &mut H) -> Result<bool, ReplaceTypesError> {
        let temp: Vec<Node>; // keep alive
        let regions = match &self.scope {
            Either::Left(scope) => {
                temp = Vec::from_iter(scope.root(hugr));
                &temp
            }
            Either::Right(regs) => regs,
        };
        let mut changed = false;
        for region_root in regions {
            changed |= self.change_subtree(hugr, *region_root, false)?;
        }
        Ok(changed)
    }
}

impl WithScope for ReplaceTypes {
    /// Sets the scope within which the pass will operate. Note that this pass ignores
    /// * [PassScope::preserve_interface], as this is a lowering pass: its purpose is to
    ///   change node signatures.
    /// * [PassScope::recursive], as non-recursion generally leads to invalid Hugrs.
    ///
    /// Hence, really only the [PassScope::root] affects the pass.
    fn with_scope(mut self, scope: impl Into<PassScope>) -> Self {
        self.scope = Either::Left(scope.into());
        self
    }
}

pub mod handlers;

#[derive(Clone, Hash, PartialEq, Eq)]
struct OpHashWrapper {
    op_name: String, // Only because SmolStr not in hugr-passes yet
    args: Vec<TypeArg>,
}

impl From<&ExtensionOp> for OpHashWrapper {
    fn from(op: &ExtensionOp) -> Self {
        Self {
            op_name: op.qualified_id().to_string(),
            args: op.args().to_vec(),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct ParametricType(ExtensionId, String);

impl From<&TypeDef> for ParametricType {
    fn from(value: &TypeDef) -> Self {
        Self(value.extension_id().clone(), value.name().to_string())
    }
}

impl From<&CustomType> for ParametricType {
    fn from(value: &CustomType) -> Self {
        Self(value.extension().clone(), value.name().to_string())
    }
}

// Separate from above for clarity
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct ParametricOp(ExtensionId, String);

impl From<&OpDef> for ParametricOp {
    fn from(value: &OpDef) -> Self {
        Self(value.extension_id().clone(), value.name().to_string())
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use crate::passes::replace_types::handlers::generic_array_const;
    use hugr_core::builder::{
        BuildError, Container, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer,
        FunctionBuilder, HugrBuilder, ModuleBuilder, SubContainer, TailLoopBuilder, endo_sig,
        inout_sig,
    };
    use hugr_core::extension::SignatureError;
    use hugr_core::extension::prelude::{
        ConstUsize, Noop, UnwrapBuilder, bool_t, option_type, qb_t, usize_t,
    };
    use hugr_core::extension::simple_op::{MakeOpDef, MakeRegisteredOp};
    use hugr_core::extension::{TypeDefBound, Version, simple_op::MakeExtensionOp};
    use hugr_core::hugr::{IdentList, ValidationError, hugrmut::HugrMut};
    use hugr_core::ops::constant::{CustomConst, OpaqueValue};
    use hugr_core::ops::{self, ExtensionOp, OpTrait, OpType, Tag, Value, handle::NodeHandle};
    use hugr_core::std_extensions::arithmetic::conversions::ConvertOpDef;
    use hugr_core::std_extensions::arithmetic::int_types::{ConstInt, INT_TYPES};
    use hugr_core::std_extensions::collections::array::{
        self, Array, ArrayKind, ArrayOpDef, GenericArrayValue, array_type, array_type_def,
    };
    use hugr_core::std_extensions::collections::borrow_array::{
        BArrayValue, BorrowArray, borrow_array_type,
    };
    use hugr_core::std_extensions::collections::list::{
        ListOp, ListOpInst, ListValue, list_type, list_type_def,
    };
    use hugr_core::types::{
        EdgeKind, PolyFuncType, Signature, SumType, Term, Type, TypeArg, TypeBound, TypeRow,
    };
    use hugr_core::{Direction, Extension, HugrView, Port, Visibility, type_row};
    use itertools::Itertools;
    use rstest::rstest;

    use crate::passes::{ComposablePass, mangle_name};

    use super::{NodeTemplate, ReplaceTypes, handlers::list_const};

    const PACKED_VEC: &str = "PackedVec";
    const READ: &str = "read";

    fn i64_t() -> Type {
        INT_TYPES[6].clone()
    }

    fn read_op(ext: &Arc<Extension>, t: Type) -> ExtensionOp {
        ExtensionOp::new(ext.get_op(READ).unwrap().clone(), [t.into()]).unwrap()
    }

    fn just_elem_type(args: &[TypeArg]) -> &Type {
        let [TypeArg::Runtime(ty)] = args else {
            panic!("Expected just elem type")
        };
        ty
    }

    fn ext() -> Arc<Extension> {
        Extension::new_arc(
            IdentList::new("TestExt").unwrap(),
            Version::new(0, 0, 1),
            |ext, w| {
                let pv_of_var = ext
                    .add_type(
                        PACKED_VEC.into(),
                        vec![TypeBound::Linear.into()],
                        String::new(),
                        TypeDefBound::from_params(vec![0]),
                        w,
                    )
                    .unwrap()
                    .instantiate(vec![Type::new_var_use(0, TypeBound::Copyable).into()])
                    .unwrap();
                ext.add_op(
                    READ.into(),
                    String::new(),
                    PolyFuncType::new(
                        vec![TypeBound::Copyable.into()],
                        Signature::new(
                            vec![pv_of_var.into(), i64_t()],
                            [Type::new_var_use(0, TypeBound::Linear)],
                        ),
                    ),
                    w,
                )
                .unwrap();
                ext.add_op(
                    "lowered_read_bool".into(),
                    String::new(),
                    Signature::new(vec![i64_t(); 2], [bool_t()]),
                    w,
                )
                .unwrap();
            },
        )
    }

    fn lowered_read<T: Container + Dataflow>(
        elem_ty: Type,
        new: impl Fn(Signature) -> Result<T, BuildError>,
    ) -> T {
        let mut dfb = new(Signature::new(
            [list_type(elem_ty.clone()), i64_t()],
            [elem_ty.clone()],
        ))
        .unwrap();
        let [val, idx] = dfb.input_wires_arr();
        let [idx] = dfb
            .add_dataflow_op(ConvertOpDef::itousize.without_log_width(), [idx])
            .unwrap()
            .outputs_arr();
        let [opt] = dfb
            .add_dataflow_op(
                ListOp::get
                    .with_type(elem_ty.clone())
                    .to_extension_op()
                    .unwrap(),
                [val, idx],
            )
            .unwrap()
            .outputs_arr();
        let [res] = dfb
            .build_unwrap_sum(1, option_type([Type::from(elem_ty)]), opt)
            .unwrap();
        dfb.set_outputs([res]).unwrap();
        dfb
    }

    fn lowerer(ext: &Arc<Extension>) -> ReplaceTypes {
        let pv = ext.get_type(PACKED_VEC).unwrap();
        let mut lw = ReplaceTypes::default();
        lw.set_replace_type(pv.instantiate([bool_t().into()]).unwrap(), i64_t());
        lw.set_replace_parametrized_type(
            pv,
            Box::new(|args: &[TypeArg]| Some(list_type(just_elem_type(args).clone()))),
        );
        lw.set_replace_op(
            &read_op(ext, bool_t()),
            NodeTemplate::SingleOp(
                ExtensionOp::new(ext.get_op("lowered_read_bool").unwrap().clone(), [])
                    .unwrap()
                    .into(),
            ),
        );
        lw.set_replace_parametrized_op(ext.get_op(READ).unwrap().as_ref(), |type_args, _| {
            Ok(Some(NodeTemplate::CompoundOp(Box::new(
                lowered_read(just_elem_type(type_args).clone(), DFGBuilder::new)
                    .finish_hugr()
                    .unwrap(),
            ))))
        });
        lw
    }

    #[test]
    fn module_func_cfg_call() {
        let ext = ext();
        let coln = ext.get_type(PACKED_VEC).unwrap();
        let c_int = Type::from(coln.instantiate([i64_t().into()]).unwrap());
        let c_bool = Type::from(coln.instantiate([bool_t().into()]).unwrap());
        let mut mb = ModuleBuilder::new();
        let sig = Signature::new_endo([Type::new_var_use(0, TypeBound::Linear)]);
        let fb = mb
            .define_function("id", PolyFuncType::new([TypeBound::Linear.into()], sig))
            .unwrap();
        let inps = fb.input_wires();
        let id = fb.finish_with_outputs(inps).unwrap();

        let sig = Signature::new([i64_t(), c_int.clone(), c_bool.clone()], [bool_t()]);
        let mut fb = mb.define_function("main", sig).unwrap();
        let [idx, indices, bools] = fb.input_wires_arr();
        let [indices] = fb
            .call(id.handle(), &[c_int.into()], [indices])
            .unwrap()
            .outputs_arr();
        let [idx2] = fb
            .add_dataflow_op(read_op(&ext, i64_t()), [indices, idx])
            .unwrap()
            .outputs_arr();
        let mut cfg = fb
            .cfg_builder(
                [(i64_t(), idx2), (c_bool.clone(), bools)],
                [bool_t()].into(),
            )
            .unwrap();
        let mut entry = cfg.entry_builder([[bool_t()].into()], type_row![]).unwrap();
        let [idx2, bools] = entry.input_wires_arr();
        let [bools] = entry
            .call(id.handle(), &[c_bool.into()], [bools])
            .unwrap()
            .outputs_arr();
        let bool_read_op = entry
            .add_dataflow_op(read_op(&ext, bool_t()), [bools, idx2])
            .unwrap();
        let [tagged] = entry
            .add_dataflow_op(
                OpType::Tag(Tag::new(0, vec![[bool_t()].into()])),
                bool_read_op.outputs(),
            )
            .unwrap()
            .outputs_arr();
        let entry = entry.finish_with_outputs(tagged, []).unwrap();
        cfg.branch(&entry, 0, &cfg.exit_block()).unwrap();
        let cfg = cfg.finish_sub_container().unwrap();
        fb.finish_with_outputs(cfg.outputs()).unwrap();
        let mut h = mb.finish_hugr().unwrap();

        assert!(lowerer(&ext).run(&mut h).unwrap());

        let ext_ops = h
            .entry_descendants()
            .filter_map(|n| h.get_optype(n).as_extension_op());
        assert_eq!(
            ext_ops
                .map(hugr_core::ops::ExtensionOp::unqualified_id)
                .sorted()
                .collect_vec(),
            ["get", "itousize", "lowered_read_bool", "panic",]
        );
    }

    #[test]
    fn dfg_conditional_case() {
        let ext = ext();
        let coln = ext.get_type(PACKED_VEC).unwrap();
        let pv = |t: Type| Type::new_extension(coln.instantiate([t.into()]).unwrap());
        let sum_rows = [[pv(pv(bool_t())), i64_t()].into(), [pv(i64_t())].into()];
        let mut dfb = DFGBuilder::new(inout_sig(
            vec![Type::new_sum(sum_rows.clone()), pv(bool_t()), pv(i64_t())],
            vec![pv(bool_t()), pv(i64_t())],
        ))
        .unwrap();
        let [sum, vb, vi] = dfb.input_wires_arr();
        let mut cb = dfb
            .conditional_builder(
                (sum_rows, sum),
                [(pv(bool_t()), vb), (pv(i64_t()), vi)],
                vec![pv(bool_t()), pv(i64_t())].into(),
            )
            .unwrap();
        let mut case0 = cb.case_builder(0).unwrap();
        let [vvb, i, _, vi0] = case0.input_wires_arr();
        let [vb0] = case0
            .add_dataflow_op(read_op(&ext, pv(bool_t())), [vvb, i])
            .unwrap()
            .outputs_arr();
        case0.finish_with_outputs([vb0, vi0]).unwrap();

        let case1 = cb.case_builder(1).unwrap();
        let [vi, vb1, _vi1] = case1.input_wires_arr();
        case1.finish_with_outputs([vb1, vi]).unwrap();
        let cond = cb.finish_sub_container().unwrap();
        let mut h = dfb.finish_hugr_with_outputs(cond.outputs()).unwrap();

        lowerer(&ext).run(&mut h).unwrap();

        let ext_ops = h
            .entry_descendants()
            .filter_map(|n| h.get_optype(n).as_extension_op())
            .collect_vec();
        assert_eq!(
            ext_ops
                .iter()
                .map(|x| x.unqualified_id())
                .sorted()
                .collect_vec(),
            ["get", "itousize", "panic"]
        );
        // The PackedVec<PackedVec<bool>> becomes a list<i64>
        let array_gets = ext_ops
            .into_iter()
            .filter_map(|e| ListOpInst::from_extension_op(e).ok())
            .collect_vec();
        assert_eq!(array_gets, [ListOp::get.with_type(i64_t())]);
    }

    #[test]
    fn loop_const() {
        let cu = |u| ConstUsize::new(u).into();
        let mut tl = TailLoopBuilder::new(
            [list_type(usize_t())],
            [list_type(bool_t())],
            [list_type(usize_t())],
        )
        .unwrap();
        let [_, bools] = tl.input_wires_arr();
        let st = SumType::new(vec![[list_type(usize_t())]; 2]);
        let pred = tl.add_load_value(
            Value::sum(
                0,
                [ListValue::new(usize_t(), [cu(1), cu(3), cu(3), cu(7)]).into()],
                st,
            )
            .unwrap(),
        );
        tl.set_outputs(pred, [bools]).unwrap();
        let backup = tl.finish_hugr().unwrap();

        let mut lowerer = ReplaceTypes::default();

        // 1. Lower List<T> to BArray<10, T> UNLESS T is usize_t() or i64_t
        lowerer.set_replace_parametrized_type(list_type_def(), |args| {
            let ty = just_elem_type(args);
            (![usize_t(), i64_t()].contains(ty)).then_some(borrow_array_type(10, ty.clone()))
        });
        {
            let mut h = backup.clone();
            assert_eq!(lowerer.run(&mut h), Ok(true));
            let sig = h.signature(h.entrypoint()).unwrap();
            assert_eq!(
                sig.input(),
                &TypeRow::from(vec![list_type(usize_t()), borrow_array_type(10, bool_t())])
            );
            assert_eq!(sig.input(), sig.output());
        }

        // 2. Now we'll also change usize's to i64_t's
        let usize_custom_t = usize_t().as_extension().unwrap().clone();
        lowerer.set_replace_type(usize_custom_t.clone(), i64_t());
        lowerer.replace_consts(usize_custom_t, |opaq, _| {
            Ok(ConstInt::new_u(
                6,
                opaq.value().downcast_ref::<ConstUsize>().unwrap().value(),
            )
            .unwrap()
            .into())
        });
        {
            let mut h = backup.clone();
            assert_eq!(lowerer.run(&mut h), Ok(true));
            let sig = h.signature(h.entrypoint()).unwrap();
            assert_eq!(
                sig.input(),
                &TypeRow::from(vec![list_type(i64_t()), borrow_array_type(10, bool_t())])
            );
            assert_eq!(sig.input(), sig.output());
            // This will have to update inside the Const
            let cst = h
                .entry_descendants()
                .filter_map(|n| h.get_optype(n).as_const())
                .exactly_one()
                .ok()
                .unwrap();
            assert_eq!(cst.get_type(), Type::new_sum(vec![[list_type(i64_t())]; 2]));
        }

        // 3. Lower all List<T> to BArray<4,T>
        let mut h = backup;
        lowerer.set_replace_parametrized_type(
            list_type_def(),
            Box::new(|args: &[TypeArg]| Some(borrow_array_type(4, just_elem_type(args).clone()))),
        );
        lowerer.replace_consts_parametrized(list_type_def(), |opaq, repl| {
            // First recursively transform the contents
            let Some(Value::Extension { e: opaq }) = list_const(opaq, repl)? else {
                panic!("Expected list value to stay a list value");
            };
            let lv = opaq.value().downcast_ref::<ListValue>().unwrap();

            Ok(Some(
                BArrayValue::new(lv.get_element_type().clone(), lv.get_contents().to_vec()).into(),
            ))
        });
        lowerer.run(&mut h).unwrap();

        assert_eq!(
            h.get_optype(pred.node())
                .as_load_constant()
                .map(hugr_core::ops::LoadConstant::constant_type),
            Some(&Type::new_sum(vec![
                [Type::from(borrow_array_type(
                    4,
                    i64_t()
                ))];
                2
            ]))
        );
    }

    #[test]
    fn partial_replace() {
        let e = Extension::new_arc(
            IdentList::new_unchecked("NoBoundsCheck"),
            Version::new(0, 0, 0),
            |e, w| {
                let params = vec![TypeBound::Linear.into()];
                let tv = Type::new_var_use(0, TypeBound::Linear);
                let list_of_var = list_type(tv.clone());
                e.add_op(
                    READ.into(),
                    "Like List::get but without the option".to_string(),
                    PolyFuncType::new(params, Signature::new([list_of_var, usize_t()], [tv])),
                    w,
                )
                .unwrap();
            },
        );
        fn option_contents(ty: &Type) -> Option<Type> {
            let row = ty.as_sum()?.get_variant(1).unwrap().clone();
            let elem = row.into_owned().into_iter().exactly_one().unwrap();
            Some(elem.try_into_type().unwrap())
        }
        let i32_t = || INT_TYPES[5].clone();
        let opt_i32 = Type::from(option_type([i32_t()]));
        let i32_custom_t = i32_t().as_extension().unwrap().clone();
        let mut dfb = DFGBuilder::new(inout_sig(
            vec![list_type(i32_t()), list_type(opt_i32.clone())],
            vec![i32_t(), opt_i32.clone()],
        ))
        .unwrap();
        let [l_i, l_oi] = dfb.input_wires_arr();
        let idx = dfb.add_load_value(ConstUsize::new(2));
        let [i] = dfb
            .add_dataflow_op(read_op(&e, i32_t()), [l_i, idx])
            .unwrap()
            .outputs_arr();
        let [oi] = dfb
            .add_dataflow_op(read_op(&e, opt_i32.clone()), [l_oi, idx])
            .unwrap()
            .outputs_arr();
        let mut h = dfb.finish_hugr_with_outputs([i, oi]).unwrap();

        let mut lowerer = ReplaceTypes::default();
        lowerer.set_replace_type(i32_custom_t, qb_t());
        // Lower list<option<x>> to list<x>
        lowerer.set_replace_parametrized_type(list_type_def(), |args| {
            option_contents(just_elem_type(args)).map(list_type)
        });
        // and read<option<x>> to get<x> - the latter has the expected option<x> return type
        lowerer.set_replace_parametrized_op(e.get_op(READ).unwrap().as_ref(), |args, _| {
            Ok(option_contents(just_elem_type(args)).map(|elem| {
                NodeTemplate::SingleOp(
                    ListOp::get
                        .with_type(elem)
                        .to_extension_op()
                        .unwrap()
                        .into(),
                )
            }))
        });
        assert!(lowerer.run(&mut h).unwrap());
        // list<usz>      -> read<usz>      -> usz just becomes list<qb> -> read<qb> -> qb
        // list<opt<usz>> -> read<opt<usz>> -> opt<usz> becomes list<qb> -> get<qb>  -> opt<qb>
        assert_eq!(
            h.entrypoint_optype().dataflow_signature().unwrap().io(),
            (
                &vec![list_type(qb_t()); 2].into(),
                &vec![qb_t(), option_type([qb_t()]).into()].into()
            )
        );
        assert_eq!(
            h.entry_descendants()
                .filter_map(|n| h.get_optype(n).as_extension_op())
                .map(hugr_core::ops::ExtensionOp::qualified_id)
                .sorted()
                .collect_vec(),
            ["NoBoundsCheck.read", "collections.list.get"]
        );
    }

    #[rstest]
    #[case(&[], Array)]
    #[case(&[3], Array)]
    #[case(&[5,7,11,13,17,19], BorrowArray)]
    fn array_const<AK: ArrayKind>(#[case] vals: &[u64], #[case] _kind: AK)
    where
        GenericArrayValue<AK>: CustomConst,
    {
        let mut dfb =
            DFGBuilder::new(inout_sig(type_row![], [AK::ty(vals.len() as _, usize_t())])).unwrap();
        let c = dfb.add_load_value(GenericArrayValue::<AK>::new(
            usize_t(),
            vals.iter().map(|u| ConstUsize::new(*u).into()),
        ));
        let backup = dfb.finish_hugr_with_outputs([c]).unwrap();

        let mut repl = ReplaceTypes::new_empty();
        let usize_custom_t = usize_t().as_extension().unwrap().clone();
        repl.set_replace_type(usize_custom_t.clone(), INT_TYPES[6].clone());
        repl.replace_consts(usize_custom_t, |cst: &OpaqueValue, _| {
            let cu = cst.value().downcast_ref::<ConstUsize>().unwrap();
            Ok(ConstInt::new_u(6, cu.value())?.into())
        });

        let mut h = backup.clone();
        repl.run(&mut h).unwrap(); // No validation here
        assert!(
            matches!(h.validate(), Err(ValidationError::IncompatiblePorts {from, to, ..})
             if backup.get_optype(from).is_const() && to == c.node())
        );
        repl.replace_consts_parametrized(AK::type_def(), generic_array_const::<AK>);
        let mut h = backup;
        repl.run(&mut h).unwrap();
        h.validate().unwrap();
    }

    #[rstest]
    fn op_to_call_polymorphic() {
        // Note the resulting Hugr has a polymorphic lowered_read function, which would
        // mean (re)running monomorphization *after* ReplaceTypes; usually we would expect
        // monomorphization to happen first so that ReplaceTypes can act upon the concrete types.
        let e = ext();
        let pv = e.get_type(PACKED_VEC).unwrap();
        let inner = pv.instantiate([usize_t().into()]).unwrap();
        let outer = pv
            .instantiate([Type::new_extension(inner.clone()).into()])
            .unwrap();
        let mut dfb = DFGBuilder::new(inout_sig([outer.into(), i64_t()], [usize_t()])).unwrap();
        let read_func = dfb
            .module_root_builder()
            .add_hugr(
                lowered_read(Type::new_var_use(0, TypeBound::Copyable), |sig| {
                    FunctionBuilder::new_vis(
                        "lowered_read",
                        PolyFuncType::new([TypeBound::Copyable.into()], sig),
                        Visibility::Public,
                    )
                })
                .finish_hugr()
                .unwrap(),
            )
            .inserted_entrypoint;
        let [outer, idx] = dfb.input_wires_arr();
        let [inner] = dfb
            .add_dataflow_op(read_op(&e, inner.clone().into()), [outer, idx])
            .unwrap()
            .outputs_arr();
        let res = dfb
            .add_dataflow_op(read_op(&e, usize_t()), [inner, idx])
            .unwrap();
        let mut h = dfb.finish_hugr_with_outputs(res.outputs()).unwrap();
        let read_poly = h
            .get_optype(read_func)
            .as_func_defn()
            .unwrap()
            .signature()
            .clone();

        let mut lw = lowerer(&e);
        lw.set_replace_parametrized_op(e.get_op(READ).unwrap().as_ref(), move |args, _| {
            let mut decl_b = ModuleBuilder::new();
            let decl_node = decl_b.declare("lowered_read", read_poly.clone()).unwrap();
            let mut decl_hugr = decl_b.finish_hugr().unwrap();
            decl_hugr.set_entrypoint(decl_node.node());

            Ok(Some(
                NodeTemplate::call_to_function(decl_hugr, args).unwrap(),
            ))
        });
        lw.run(&mut h).unwrap();
        h.validate().unwrap();

        assert_eq!(h.output_neighbours(read_func).count(), 2);
        assert_eq!(
            h.entry_descendants()
                .find(|n| h.get_optype(*n).is_extension_op()),
            None
        );
        assert_eq!(h.children(h.module_root()).count(), 2); // main + lowered_read
    }

    #[rstest]
    fn op_to_call_monomorphic(#[values(false, true)] i64_to_usize: bool) {
        let e = ext();
        let pv = e.get_type(PACKED_VEC).unwrap();
        let inner = pv.instantiate([usize_t().into()]).unwrap();
        let outer = pv
            .instantiate([Type::new_extension(inner.clone()).into()])
            .unwrap();
        let read_outer = read_op(&e, inner.clone().into());
        let mut dfb = DFGBuilder::new(inout_sig(
            vec![outer.into(), inner.clone().into(), i64_t()],
            vec![usize_t(); 2],
        ))
        .unwrap();

        let [outer, inner, idx] = dfb.input_wires_arr();
        let res1 = dfb
            .add_dataflow_op(read_op(&e, usize_t()), [inner, idx])
            .unwrap();
        let [inner] = dfb
            .add_dataflow_op(read_outer, [outer, idx])
            .unwrap()
            .outputs_arr();
        let res2 = dfb
            .add_dataflow_op(read_op(&e, usize_t()), [inner, idx])
            .unwrap();
        let mut h = dfb
            .finish_hugr_with_outputs(res1.outputs().chain(res2.outputs()))
            .unwrap();

        let mut lw = lowerer(&e);
        lw.set_replace_parametrized_op(e.get_op(READ).unwrap().as_ref(), move |args, _| {
            Ok(Some({
                let [Term::Runtime(ty)] = args else {
                    return Err(SignatureError::InvalidTypeArgs.into());
                };

                let defn_hugr = lowered_read(ty.clone(), |sig| {
                    FunctionBuilder::new_vis(
                        mangle_name("lowered_read", args),
                        sig,
                        Visibility::Public,
                    )
                })
                .finish_hugr()
                .unwrap();

                NodeTemplate::call_to_function(defn_hugr, &[]).unwrap()
            }))
        });
        if i64_to_usize {
            lw.set_replace_type(i64_t().as_extension().unwrap().clone(), usize_t());
            lw.set_replace_op(
                &ConvertOpDef::itousize
                    .without_log_width()
                    .to_extension_op()
                    .unwrap(),
                NodeTemplate::SingleOp(Noop::new(usize_t()).into()),
            );
        }
        lw.run(&mut h).unwrap();
        h.validate().unwrap();

        assert_eq!(
            h.entry_descendants()
                .find(|n| h.get_optype(*n).is_extension_op()),
            None
        );
        assert_eq!(h.children(h.module_root()).count(), 3); // main + lowered_read
        for n in h.children(h.module_root()) {
            let fd = h.get_optype(n).as_func_defn().unwrap();
            let expected_uses_and_vis = if fd.func_name() == "main" {
                (0, Visibility::Private)
            } else {
                let is_array = !fd.signature().body().output[0]
                    .as_extension()
                    .unwrap()
                    .args()
                    .is_empty();
                (2 - (is_array as usize), Visibility::Public)
            };
            assert_eq!(h.output_neighbours(n).count(), expected_uses_and_vis.0);
            assert_eq!(fd.visibility(), &expected_uses_and_vis.1);
        }
    }

    #[test]
    fn regions() {
        let ext = ext();
        let coln = ext.get_type(PACKED_VEC).unwrap();
        let c_u = Type::new_extension(coln.instantiate(&[usize_t().into()]).unwrap());
        let mut h = {
            let db = DFGBuilder::new(endo_sig([c_u.clone()])).unwrap();
            let inps = db.input_wires();
            db.finish_hugr_with_outputs(inps)
        }
        .unwrap();
        let mut lowerer = lowerer(&ext);

        {
            let backup = h.clone();
            lowerer.set_regions(vec![]);
            assert!(!lowerer.run(&mut h).unwrap());
            assert_eq!(h, backup);
        }

        let ep = h.entrypoint();
        lowerer.set_regions(vec![h.entrypoint()]);
        assert!(lowerer.run(&mut h).unwrap());
        let v_u = list_type(usize_t());
        assert_eq!(h.signature(ep).unwrap().as_ref(), &endo_sig([v_u.clone()]));
        assert_eq!(h.num_nodes(), h.num_nodes());
        let [f_in, _] = h.get_io(h.get_parent(ep).unwrap()).unwrap();
        assert_eq!(
            h.validate(),
            Err(ValidationError::IncompatiblePorts {
                from: f_in,
                from_port: Port::new(Direction::Outgoing, 0),
                to: ep,
                to_port: Port::new(Direction::Incoming, 0),
                from_kind: Box::new(EdgeKind::Value(c_u)),
                to_kind: Box::new(EdgeKind::Value(v_u))
            })
        );
    }

    #[test]
    fn compositionality() {
        let ext = ext();
        let mut lowerer = lowerer(&ext);
        // Replace std Array's with 64 elements with PackedVec's
        let ext2 = ext.clone();
        lowerer.set_replace_parametrized_type(array_type_def(), move |args| {
            let [sz, ty] = args else {
                panic!("Expected two args to array")
            };
            (sz == &Term::BoundedNat(64)).then_some(
                ext2.get_type(PACKED_VEC)
                    .unwrap()
                    .instantiate([ty.clone()])
                    .unwrap()
                    .into(),
            )
        });

        // Replacement of `get` is complex because we need to wrap result of read into a Some
        let ext = ext.clone();
        lowerer.set_replace_parametrized_op(
            array::EXTENSION
                .get_op(ArrayOpDef::get.opdef_id().as_str())
                .unwrap()
                .as_ref(),
            move |args, _| {
                let [sz, Term::Runtime(ty)] = args else {
                    panic!("Expected two args to array-get")
                };
                if sz != &Term::BoundedNat(64) {
                    return Ok(None);
                }
                let pv = ext
                    .get_type(PACKED_VEC)
                    .unwrap()
                    .instantiate([ty.clone().into()])
                    .unwrap();

                let mut dfb = DFGBuilder::new(Signature::new(
                    vec![pv.clone().into(), usize_t()],
                    vec![option_type([ty.clone()]).into(), pv.into()],
                ))
                .unwrap();
                let [pvec, idx] = dfb.input_wires_arr();
                let [idx] = dfb
                    .add_dataflow_op(ConvertOpDef::ifromusize.without_log_width(), [idx])
                    .unwrap()
                    .outputs_arr();
                let [elem] = dfb
                    .add_dataflow_op(read_op(&ext, ty.clone()), [pvec, idx])
                    .unwrap()
                    .outputs_arr();
                let [wrapped_elem] = dfb
                    .add_dataflow_op(
                        ops::Tag::new(1, vec![type_row![], [ty.clone()].into()]),
                        [elem],
                    )
                    .unwrap()
                    .outputs_arr();
                Ok(Some(NodeTemplate::CompoundOp(Box::new(
                    dfb.finish_hugr_with_outputs([wrapped_elem, pvec]).unwrap(),
                ))))
            },
        );

        // Arrays of 64 bools should thus be transformed into PackedVec<bool> and then to int64s
        // Arrays of 64 non-bools should thus become PackedVec<T> and thus List<T>
        let a64 = |t| array_type(64, t);
        let opt = |t| Type::from(option_type([t]));
        let mut dfb = DFGBuilder::new(Signature::new(
            vec![a64(bool_t()), a64(usize_t())],
            vec![opt(bool_t()), a64(bool_t()), opt(usize_t()), a64(usize_t())],
        ))
        .unwrap();
        let [bools, usizes] = dfb.input_wires_arr();
        let idx = dfb.add_load_value(ConstUsize::new(5));
        let [b, bools] = dfb
            .add_dataflow_op(ArrayOpDef::get.to_concrete(bool_t(), 64), [bools, idx])
            .unwrap()
            .outputs_arr();
        let [u, usizes] = dfb
            .add_dataflow_op(ArrayOpDef::get.to_concrete(usize_t(), 64), [usizes, idx])
            .unwrap()
            .outputs_arr();
        let mut h = dfb.finish_hugr_with_outputs([b, bools, u, usizes]).unwrap();

        lowerer.run(&mut h).unwrap();

        h.validate().unwrap();
    }
}