tket 0.19.0

//! Hugr hashing.

use derive_more::{Display, Error};
use fxhash::{FxHashMap, FxHasher64};
use hugr_core::HugrView;
use hugr_core::extension::ExtensionId;
use hugr_core::ops::OpNameRef;
use hugr_core::ops::OpType;
use hugr_core::ops::{OpTag, OpTrait};
use hugr_core::types::type_param::TypeArg;
use petgraph::visit::{self as pg, Walker};
use std::fmt::{self, Write};
use std::hash::{Hash, Hasher};

/// Hugr hashing utilities.
pub trait HugrHash: HugrView {
    /// Compute a hash for the entire HUGR structure.
    ///
    /// This corresponds to the hash of the root node.
    fn hugr_hash(&self) -> Result<u64, HashError> {
        self.region_hash(self.module_root())
    }

    /// Compute a hash for a hugr region.
    ///
    /// If the hugr is a dfg, we compute a hash for each node from its operation
    /// and the hash of the predecessors. The hash of the hugr corresponds to
    /// the hash of its output node.
    /// Otherwise, we compute a generic hash combining the hashes of its children.
    ///
    /// This hash is independent from the children node order.
    fn region_hash(&self, node: Self::Node) -> Result<u64, HashError>;
}

impl<H: HugrView> HugrHash for H {
    fn region_hash(&self, node: Self::Node) -> Result<u64, HashError> {
        let node_op = self.get_optype(node);
        let mut hasher = FxHasher64::default();

        let op: &OpType = self.get_optype(node);
        hash_op_type(op, &mut hasher);

        if OpTag::DataflowParent.is_superset(node_op.tag()) {
            // In this case, we have a dataflow container
            dfg_hash(self, node)?.hash(&mut hasher);
        } else {
            // otherwise, use generic hash
            generic_hugr_hash(self, node)?.hash(&mut hasher);
        }

        Ok(hasher.finish())
    }
}

fn dfg_hash<H: HugrView>(dfg_hugr: &H, node: H::Node) -> Result<u64, HashError> {
    let mut node_hashes = HashState {
        hashes: FxHashMap::default(),
    };

    let [_, output_node] = dfg_hugr.get_io(node).expect("DFG region missing I/O nodes");

    let sg = dfg_hugr.scheduling_graph(node);
    for pg_node in pg::Topo::new(sg.petgraph()).iter(sg.petgraph()) {
        let child = sg.pg_to_node(pg_node);
        let hash = dfg_hash_node(dfg_hugr, child, &mut node_hashes)?;
        if node_hashes.set_hash(child, hash).is_some() {
            panic!("Hash already set for node {node}");
        }
    }

    node_hashes
        .get_hash(output_node)
        .ok_or(HashError::CyclicDFG)
}

fn generic_hugr_hash<H: HugrView>(hugr: &H, node: H::Node) -> Result<u64, HashError> {
    let mut child_hashes = Vec::new();

    for child in hugr.children(node) {
        child_hashes.push(hugr.region_hash(child)?);
    }
    // Combine child hashes in an order-independent way
    child_hashes.sort_unstable();
    Ok(fxhash::hash64(&child_hashes))
}

/// Auxiliary data for circuit hashing.
///
/// Contains previously computed hashes.
#[derive(Clone, Default, Debug)]
struct HashState<H: HugrView> {
    /// Computed node hashes.
    pub hashes: FxHashMap<H::Node, u64>,
}
impl<H: HugrView> HashState<H> {
    /// Return the hash for a node.
    #[inline]
    fn get_hash(&self, node: H::Node) -> Option<u64> {
        self.hashes.get(&node).copied()
    }

    /// Register the hash for a node.
    ///
    /// Returns the previous hash, if it was set.
    #[inline]
    fn set_hash(&mut self, node: H::Node, hash: u64) -> Option<u64> {
        self.hashes.insert(node, hash)
    }
}

/// Hash the identity of an operation (extension, name, and type arguments).
fn hash_extension_identity(
    hasher: &mut impl Hasher,
    extension: &ExtensionId,
    name: &OpNameRef,
    args: &[TypeArg],
) {
    extension.hash(hasher);
    name.hash(hasher);
    args.hash(hasher);
}

/// Hash an operation for use in region hashes.
///
/// Extension and opaque ops are hashed structurally (no JSON formatting).
/// Other ops use their tag and name; graph context distinguishes calls etc.
fn hash_op_type(op: &OpType, hasher: &mut impl Hasher) {
    match op {
        OpType::ExtensionOp(ext) => {
            hash_extension_identity(hasher, ext.extension_id(), ext.unqualified_id(), ext.args());
        }
        OpType::OpaqueOp(op) => {
            hash_extension_identity(hasher, op.extension(), op.unqualified_id(), op.args());
        }
        _ => {
            op.tag().hash(hasher);
            // `NamedOp::name` is crate-private in hugr; hash via `Display` without allocating.
            write!(FmtHasher(hasher), "{op}").expect("hashing op name must not fail");
        }
    }
}

/// [`Write`] adapter that feeds formatted text into a [`Hasher`].
struct FmtHasher<'a, H: Hasher>(&'a mut H);

impl<H: Hasher> Write for FmtHasher<'_, H> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        s.hash(self.0);
        Ok(())
    }
}

/// Compute the hash of a circuit command.
///
/// Uses the hash of the operation and the node hash of its predecessors.
///
/// # Panics
/// - If the command is a container node, or if it is a parametric CustomOp.
/// - If the hash of any of its predecessors has not been set.
fn dfg_hash_node<H: HugrView>(
    hugr: &H,
    node: H::Node,
    state: &mut HashState<H>,
) -> Result<u64, HashError> {
    let mut hasher = FxHasher64::default();

    hugr.region_hash(node)?.hash(&mut hasher);

    // Add each each input neighbour hash, including the connected ports.
    // TODO: Ignore state edges?
    for input in hugr.node_inputs(node) {
        // Combine the hash for each subport, ignoring their order.
        let input_hash = hugr
            .linked_ports(node, input)
            .map(|(pred_node, pred_port)| {
                let pred_node_hash = state.get_hash(pred_node);
                fxhash::hash64(&(pred_node_hash, pred_port, input))
            })
            .fold(0, |total, hash| hash ^ total);
        input_hash.hash(&mut hasher);
    }
    Ok(hasher.finish())
}

/// Errors that can occur while hashing a hugr.
#[derive(Debug, Display, Clone, PartialEq, Eq, Error)]
#[non_exhaustive]
pub enum HashError {
    /// The hugr dfg contains a cycle.
    #[display("The hugr dfg contains a cycle.")]
    CyclicDFG,
}

#[cfg(test)]
mod test {
    use crate::passes::test_utils::build_simple_hugr;
    use crate::passes::test_utils::test_quantum_extension::{cx_gate, h_gate};
    use hugr_core::builder::{Dataflow, DataflowSubContainer};

    use super::*;
    #[test]
    fn hash_equality() {
        let hugr1 = build_simple_hugr(2, |mut f_build| {
            // let wires = f_build.input_wires().map(Some).collect();

            let mut linear = f_build.as_circuit(f_build.input_wires());

            linear
                .append(h_gate(), [0])?
                .append(h_gate(), [1])?
                .append(cx_gate(), [0, 1])?;

            let outs = linear.finish();
            f_build.finish_with_outputs(outs)
        })
        .unwrap();

        let hash1 = hugr1.hugr_hash().unwrap();

        // A circuit built in a different order should have the same hash
        let hugr2 = build_simple_hugr(2, |mut f_build| {
            let mut linear = f_build.as_circuit(f_build.input_wires());

            linear
                .append(h_gate(), [1])?
                .append(h_gate(), [0])?
                .append(cx_gate(), [0, 1])?;

            let outs = linear.finish();
            f_build.finish_with_outputs(outs)
        })
        .unwrap();

        let hash2 = hugr2.hugr_hash().unwrap();

        assert_eq!(hash1, hash2);

        // Inverting the CX control and target should produce a different hash
        let hugr3 = build_simple_hugr(2, |mut f_build| {
            let mut linear = f_build.as_circuit(f_build.input_wires());

            linear
                .append(h_gate(), [1])?
                .append(h_gate(), [0])?
                .append(cx_gate(), [1, 0])?;

            let outs = linear.finish();
            f_build.finish_with_outputs(outs)
        })
        .unwrap();
        let hash3 = hugr3.hugr_hash().unwrap();

        assert_ne!(hash1, hash3);
    }

    #[test]
    fn hash_distinguishes_type_args() {
        use hugr_core::std_extensions::arithmetic::int_ops::IntOpDef;

        let op5: OpType = IntOpDef::iadd.with_log_width(5).into();
        let op6: OpType = IntOpDef::iadd.with_log_width(6).into();

        let hash_of = |op: &OpType| {
            let mut hasher = FxHasher64::default();
            hash_op_type(op, &mut hasher);
            hasher.finish()
        };

        assert_eq!(hash_of(&op5), hash_of(&op5));
        assert_ne!(hash_of(&op5), hash_of(&op6));
    }
}