//! Pinned compile.cpp: DFS/parents, available-scalar coalescing and three
//! depth-one CSE passes, then original fusion boundaries and Compiled lowering.
//! Rebuild descriptors, never mutate captured model data.
use super::array::{Array, Dtype};
use super::ops::Operation;
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

struct Node {
    array: Array,
    inputs: Vec<usize>,
}

struct Tape {
    nodes: HashMap<usize, Node>,
    order: Vec<usize>,
    parents: HashMap<usize, Vec<(usize, usize)>>,
    outputs: HashSet<usize>,
}

impl Tape {
    /// compile.cpp::split_one: only parents already in the candidate section
    /// move to the new Broadcast descriptor. Other consumers keep the view.
    fn split_broadcast(&mut self, id: usize, section: &HashSet<usize>) -> Result<usize, String> {
        let n = &self.nodes[&id];
        let inputs = n.inputs.clone();
        let args = inputs
            .iter()
            .map(|id| self.nodes[id].array.clone())
            .collect::<Vec<_>>();
        let a = Array::make_shared_operations(
            n.array.stream().unwrap(),
            &args,
            &[n.array.layout().clone()],
            n.array.operation().unwrap(),
        )?
        .remove(0);
        let split = a.id();
        self.nodes.insert(split, Node { array: a, inputs });
        let mut moved = Vec::new();
        self.parents
            .get_mut(&id)
            .unwrap()
            .retain(|&(parent, index)| {
                if section.contains(&parent) {
                    self.nodes.get_mut(&parent).unwrap().inputs[index] = split;
                    moved.push((parent, index));
                    false
                } else {
                    true
                }
            });
        self.parents.insert(split, moved);
        Ok(split)
    }

    fn collect_fusable(
        &mut self,
        id: usize,
        depth: usize,
        stream: super::stream::Stream,
        shape: &[i32],
        section: &mut HashSet<usize>,
        candidates: &mut HashSet<usize>,
    ) -> Result<(), String> {
        if section.contains(&id) {
            return Ok(());
        }
        let a = &self.nodes[&id].array;
        let op = a.operation();
        if depth >= 11
            || op
                .as_ref()
                .is_none_or(|p| super::fused::operator(p).is_none())
            || a.stream() != Some(stream)
            || (self.outputs.contains(&id) && a.layout().shape() != shape)
        {
            candidates.insert(id);
            return Ok(());
        }
        if depth > 0 && !self.parents[&id].iter().all(|(p, _)| section.contains(p)) {
            if matches!(&**op.as_ref().unwrap(), Operation::Broadcast) && candidates.len() < 24 {
                let split = self.split_broadcast(id, section)?;
                self.collect_fusable(split, depth, stream, shape, section, candidates)?;
            } else {
                candidates.insert(id);
            }
            return Ok(());
        }
        if self.outputs.contains(&id) {
            candidates.insert(id);
        } else {
            candidates.remove(&id);
        }
        if candidates.len() >= 24 {
            return Ok(());
        }
        section.insert(id);
        for input in self.nodes[&id].inputs.clone() {
            self.collect_fusable(input, depth + 1, stream, shape, section, candidates)?;
        }
        Ok(())
    }

    fn fused_tape(&self, root: usize, section: &HashSet<usize>) -> (Vec<usize>, Vec<usize>) {
        let mut inputs = Vec::new();
        let mut seen_inputs = HashSet::new();
        let mut tape = Vec::new();
        let mut seen = HashSet::new();
        let mut stack = vec![(root, false)];
        while let Some((id, expanded)) = stack.pop() {
            if !section.contains(&id) {
                if seen_inputs.insert(id) {
                    inputs.push(id);
                }
            } else if expanded {
                tape.push(id);
            } else if seen.insert(id) {
                stack.push((id, true));
                stack.extend(self.nodes[&id].inputs.iter().rev().map(|&i| (i, false)));
            }
        }
        (inputs, tape)
    }

    fn fuse(&mut self, inputs: &[Array], outputs: &[Array]) -> Result<Vec<Array>, String> {
        let input_ids = inputs.iter().map(Array::id).collect::<HashSet<_>>();
        let mut output_map = outputs
            .iter()
            .map(|a| (a.id(), a.clone()))
            .collect::<HashMap<_, _>>();
        let mut keep = Vec::new();
        let mut global = HashSet::new();
        for root in self.order.clone().into_iter().rev() {
            if global.contains(&root) {
                continue;
            }
            let a = &self.nodes[&root].array;
            let mut section = HashSet::new();
            if a.operation()
                .is_some_and(|p| !matches!(&*p, Operation::Broadcast))
            {
                let stream = a.stream().unwrap();
                let shape = a.layout().shape().to_vec();
                self.collect_fusable(root, 0, stream, &shape, &mut section, &mut HashSet::new())?;
            }
            if section.len() <= 1 {
                keep.push(root);
                continue;
            }
            let (inputs, tape) = self.fused_tape(root, &section);
            let mut old_outputs = Vec::new();
            for &id in &tape[..tape.len() - 1] {
                if self.outputs.contains(&id) {
                    old_outputs.push(id);
                    self.parents
                        .entry(id)
                        .or_default()
                        .retain(|(p, _)| !section.contains(p));
                } else {
                    self.parents.remove(&id);
                }
                global.insert(id);
            }
            old_outputs.push(root);
            let layouts = old_outputs
                .iter()
                .map(|id| self.nodes[id].array.layout().clone())
                .collect::<Vec<_>>();
            if layouts
                .iter()
                .any(|l| l.shape() != layouts.last().unwrap().shape())
            {
                return Err("compilation tried to fuse different output shapes".into());
            }
            let args = inputs
                .iter()
                .map(|id| self.nodes[id].array.clone())
                .collect::<Vec<_>>();
            let constants = args
                .iter()
                .map(|a| {
                    a.layout().size() == 1
                        && a.operation().is_none()
                        && !input_ids.contains(&a.id())
                })
                .collect::<Vec<_>>();
            let slots = inputs
                .iter()
                .chain(&tape)
                .enumerate()
                .map(|(i, &id)| (id, i))
                .collect::<HashMap<_, _>>();
            let statements = tape
                .iter()
                .map(|id| {
                    let node = &self.nodes[id];
                    super::fused::Statement {
                        dtype: node.array.layout().dtype(),
                        operator: super::fused::operator(&node.array.operation().unwrap()).unwrap(),
                        inputs: node.inputs.iter().map(|i| slots[i]).collect(),
                    }
                })
                .collect();
            let expr = super::fused::Expression::new(
                &args,
                statements,
                old_outputs.iter().map(|id| slots[id]).collect(),
                &constants,
            )?;
            let new = Array::make_operations(
                self.nodes[&root].array.stream().unwrap(),
                &args,
                &layouts,
                Operation::Compiled(expr),
            )?;
            keep.push(new.last().unwrap().id());
            for output in &new {
                self.nodes.insert(
                    output.id(),
                    Node {
                        array: output.clone(),
                        inputs: inputs.clone(),
                    },
                );
            }
            for (position, input) in inputs.iter().enumerate() {
                let parents = self.parents.entry(*input).or_default();
                parents.retain(|(p, _)| !section.contains(p));
                parents.extend(new.iter().map(|a| (a.id(), position)));
            }
            for (&old, new) in old_outputs.iter().zip(new) {
                self.merge_one(new.id(), old);
                if let Some(output) = output_map.get_mut(&old) {
                    *output = new;
                }
            }
        }
        keep.reverse();
        self.order = keep;
        Ok(outputs
            .iter()
            .map(|a| output_map[&a.id()].clone())
            .collect())
    }

    fn capture(outputs: &[Array], original: &[Array]) -> Result<Self, String> {
        let mut tape = Self {
            nodes: HashMap::new(),
            order: Vec::new(),
            parents: HashMap::new(),
            outputs: outputs.iter().map(Array::id).collect(),
        };
        let forbidden = original.iter().map(Array::id).collect::<HashSet<_>>();
        let mut seen = HashSet::new();
        let mut stack = outputs
            .iter()
            .rev()
            .map(|a| (a.clone(), false))
            .collect::<Vec<_>>();
        while let Some((a, expanded)) = stack.pop() {
            let id = a.id();
            if forbidden.contains(&id) {
                return Err("compiled function has uncaptured inputs".into());
            }
            if seen.contains(&id) {
                continue;
            }
            if expanded {
                seen.extend(a.outputs().iter().map(Array::id));
                tape.order.push(id);
                continue;
            }
            let inputs = a.inputs().iter().map(Array::id).collect::<Vec<_>>();
            // Reference registers the visited output first, then its siblings,
            // before descending into its inputs (not sorted by array id).
            let siblings = a.siblings();
            for (position, &input) in inputs.iter().enumerate() {
                let parents = tape.parents.entry(input).or_default();
                parents.push((id, position));
                parents.extend(siblings.iter().map(|s| (s.id(), position)));
            }
            for output in a.outputs() {
                tape.nodes.insert(
                    output.id(),
                    Node {
                        array: output,
                        inputs: inputs.clone(),
                    },
                );
            }
            stack.push((a.clone(), true));
            stack.extend(a.inputs().iter().rev().map(|i| (i.clone(), false)));
        }
        Ok(tape)
    }

    fn merge_one(&mut self, dst: usize, src: usize) {
        let Some(parents) = self.parents.remove(&src) else {
            return;
        };
        for &(parent, index) in &parents {
            self.nodes.get_mut(&parent).unwrap().inputs[index] = dst;
        }
        let pairs = self.parents.entry(dst).or_default();
        pairs.extend(parents);
        pairs.retain(|&(parent, _)| parent != src);
    }

    fn merge(&mut self, dst: usize, src: usize) {
        let dst = self.nodes[&dst].array.outputs();
        let src = self.nodes[&src].array.outputs();
        for (dst, src) in dst.iter().zip(src.iter()) {
            self.merge_one(dst.id(), src.id());
        }
    }

    fn scalars(&mut self) -> Result<(), String> {
        let mut scalars = HashMap::<(u64, Dtype), usize>::new();
        let mut keep = Vec::new();
        for id in self.order.clone() {
            let a = &self.nodes[&id].array;
            if a.is_available()? && a.layout().shape().is_empty() && a.has_data() {
                let dtype = a.layout().dtype();
                let mut bits = [0; 8];
                a.buffer().read(a.offset(), &mut bits[..dtype.itemsize()])?;
                let dst = *scalars
                    .entry((u64::from_le_bytes(bits), dtype))
                    .or_insert(id);
                if dst != id {
                    self.merge(dst, id);
                    continue;
                }
            }
            keep.push(id);
        }
        self.order = keep;
        Ok(())
    }

    fn equivalent(&self, a: usize, b: usize) -> bool {
        let a = &self.nodes[&a];
        let b = &self.nodes[&b];
        let (Some(pa), Some(pb)) = (a.array.operation(), b.array.operation()) else {
            return false;
        };
        // The original explicitly excludes identical primitive pointers and
        // does not compare streams at this stage. Fusion checks streams later.
        !Rc::ptr_eq(&pa, &pb)
            && a.inputs == b.inputs
            && equivalent_operation(&pa, &pb, &a.array, &b.array)
    }

    fn merge_parents(&mut self, id: usize, order: &HashMap<usize, usize>) -> bool {
        let Some(parents) = self.parents.get(&id).cloned() else {
            return !self.outputs.contains(&id);
        };
        let mut mask = vec![false; parents.len()];
        let groups = if parents.len() > 100 {
            // Same high-fanout grouping as VecU64Hash, using Rust's native
            // hash table. Binary/unary variants are distinguished by the
            // final equivalence check, including their primitive parameters.
            let mut groups = HashMap::new();
            for (i, &(parent, _)) in parents.iter().enumerate() {
                let node = &self.nodes[&parent];
                let op = node.array.operation().unwrap();
                groups
                    .entry((node.inputs.clone(), std::mem::discriminant(&*op)))
                    .or_insert_with(Vec::new)
                    .push(i);
            }
            groups.into_values().collect::<Vec<_>>()
        } else {
            vec![(0..parents.len()).collect()]
        };
        for group in groups {
            for (i, &left) in group.iter().enumerate() {
                if mask[left] {
                    continue;
                }
                for &right in &group[i + 1..] {
                    if mask[right] {
                        continue;
                    }
                    let (mut dst, mut src) = (left, right);
                    if order.get(&parents[src].0).copied().unwrap_or(0)
                        < order.get(&parents[dst].0).copied().unwrap_or(0)
                    {
                        std::mem::swap(&mut dst, &mut src);
                    }
                    let (dst_id, src_id) = (parents[dst].0, parents[src].0);
                    if src_id != dst_id
                        && !self.outputs.contains(&src_id)
                        && self.equivalent(src_id, dst_id)
                    {
                        self.merge(dst_id, src_id);
                        mask[src] = true;
                    }
                }
            }
        }
        let pairs = self.parents.get_mut(&id).unwrap();
        let mut i = 0;
        pairs.retain(|_| {
            let keep = !mask[i];
            i += 1;
            keep
        });
        false
    }

    fn simplify(&mut self) -> Result<(), String> {
        self.scalars()?;
        // Copy/StopGradient are absent from this inference graph API: their
        // no-op contract is represented by returning the input descriptor.
        let order = self
            .order
            .iter()
            .enumerate()
            .map(|(i, &id)| (id, i))
            .collect();
        for _ in 0..3 {
            let mut keep = Vec::new();
            for id in self.order.clone() {
                let outputs = self.nodes[&id].array.outputs();
                let mut discard = self.merge_parents(id, &order);
                for sibling in outputs.iter().filter(|a| a.id() != id) {
                    discard &= self.merge_parents(sibling.id(), &order);
                }
                if !discard {
                    keep.push(id);
                }
            }
            self.order = keep;
        }
        Ok(())
    }

    fn rebuild(&self, outputs: &[Array]) -> Result<Vec<Array>, String> {
        let mut values = HashMap::<usize, Array>::new();
        for &id in &self.order {
            let node = &self.nodes[&id];
            let a = &node.array;
            let Some(op) = a.operation() else {
                values.insert(id, a.clone());
                continue;
            };
            let args = node
                .inputs
                .iter()
                .map(|id| {
                    values
                        .get(id)
                        .cloned()
                        .ok_or("simplified input absent from tape".into())
                })
                .collect::<Result<Vec<_>, String>>()?;
            let old = a.outputs();
            let layouts = old.iter().map(|o| o.layout().clone()).collect::<Vec<_>>();
            let result = Array::make_shared_operations(
                a.stream().ok_or("primitive has no stream")?,
                &args,
                &layouts,
                op,
            )?;
            for (old, new) in old.iter().zip(result) {
                values.insert(old.id(), new);
            }
        }
        outputs
            .iter()
            .map(|a| {
                values
                    .get(&a.id())
                    .cloned()
                    .or_else(|| a.operation().is_none().then(|| a.clone()))
                    .ok_or("simplified output absent from tape".into())
            })
            .collect()
    }
}

pub(super) fn simplify(outputs: &[Array], original_inputs: &[Array]) -> Result<Vec<Array>, String> {
    let mut tape = Tape::capture(outputs, original_inputs)?;
    tape.simplify()?;
    tape.rebuild(outputs)
}

pub(super) fn compile(
    outputs: &[Array],
    original_inputs: &[Array],
    placeholders: &[Array],
) -> Result<Vec<Array>, String> {
    let mut tape = Tape::capture(outputs, original_inputs)?;
    tape.simplify()?;
    let outputs = tape.fuse(placeholders, outputs)?;
    tape.rebuild(&outputs)
}

fn equivalent_operation(a: &Operation, b: &Operation, x: &Array, y: &Array) -> bool {
    use Operation::*;
    let same_outputs = || {
        let x = x.outputs();
        let y = y.outputs();
        x.len() == y.len()
            && x.iter().zip(&y).all(|(a, b)| {
                a.layout().shape() == b.layout().shape() && a.layout().dtype() == b.layout().dtype()
            })
    };
    match (a, b) {
        (Flatten { start: a, end: b }, Flatten { start: c, end: d }) => a == c && b == d,
        (Unflatten { axis: a, shape: b }, Unflatten { axis: c, shape: d }) => a == c && b == d,
        (Reshape, Reshape) | (Broadcast, Broadcast) | (AsType, AsType) | (View, View) => {
            same_outputs()
        }
        (Transpose(a), Transpose(b))
        | (ExpandDims(a), ExpandDims(b))
        | (Squeeze(a), Squeeze(b))
        | (Sum(a), Sum(b))
        | (Min(a), Min(b))
        | (Max(a), Max(b))
        | (DynamicSliceUpdate(a), DynamicSliceUpdate(b)) => a == b,
        (DynamicSlice(a), DynamicSlice(b)) => a == b && same_outputs(),
        (
            Slice {
                starts: a,
                stops: b,
                steps: c,
            },
            Slice {
                starts: x,
                stops: y,
                steps: z,
            },
        )
        | (
            SliceUpdate {
                starts: a,
                stops: b,
                steps: c,
            },
            SliceUpdate {
                starts: x,
                stops: y,
                steps: z,
            },
        ) => (a, b, c) == (x, y, z),
        (Contiguous(a), Contiguous(b))
        | (Softmax(a), Softmax(b))
        | (SearchSorted(a), SearchSorted(b))
        | (CumsumLastF32(a), CumsumLastF32(b)) => a == b,
        (RmsNorm(a), RmsNorm(b)) => a == b,
        (
            SdpaVector {
                scale: a,
                causal: b,
            },
            SdpaVector {
                scale: x,
                causal: y,
            },
        ) => (a, b) == (x, y),
        (QuantizedLinearBf16 { bits: a, group: b }, QuantizedLinearBf16 { bits: x, group: y })
        | (AffineQuantize { bits: a, group: b }, AffineQuantize { bits: x, group: y })
        | (AffineDequantize { bits: a, group: b }, AffineDequantize { bits: x, group: y })
        | (
            GatherQmmBf16 {
                bits: a, group: b, ..
            },
            GatherQmmBf16 {
                bits: x, group: y, ..
            },
        ) => (a, b) == (x, y),
        (
            Sort {
                axis: a,
                indices: b,
                kth: c,
            },
            Sort {
                axis: x,
                indices: y,
                kth: z,
            },
        ) => (a, b, c) == (x, y, z),
        (
            Split {
                indices: a,
                axis: b,
            },
            Split {
                indices: x,
                axis: y,
            },
        ) => (a, b) == (x, y),
        (Unary(a), Unary(b)) => std::mem::discriminant(a) == std::mem::discriminant(b),
        (Binary(a), Binary(b)) => std::mem::discriminant(a) == std::mem::discriminant(b),
        (Concatenate(a), Concatenate(b))
        | (ArgMax(a), ArgMax(b))
        | (GatherAxis(a), GatherAxis(b))
        | (ScatterAxis(a), ScatterAxis(b))
        | (GatherTake(a), GatherTake(b)) => a == b,
        (ArangeI32(a), ArangeI32(b)) => a == b && same_outputs(),
        (ArangeF32 { start: a, step: b }, ArangeF32 { start: x, step: y }) => {
            (a, b) == (x, y) && same_outputs()
        }
        (ArangeU32, ArangeU32) | (ArangeI64, ArangeI64) => same_outputs(),
        (RandomBits, RandomBits) => same_outputs(),
        (Full, Full)
        | (DenseMatmul, DenseMatmul)
        | (Select, Select)
        | (ScatterAddRow, ScatterAddRow)
        | (AllBoolContiguous, AllBoolContiguous)
        | (LogSumExp, LogSumExp)
        | (CummaxLastI64, CummaxLastI64)
        | (CumsumLastI32, CumsumLastI32)
        | (DepthwiseConv1d, DepthwiseConv1d)
        | (PleConv1d, PleConv1d)
        | (SiluBf16, SiluBf16)
        | (SwiGluBf16, SwiGluBf16)
        | (ComputeGBf16, ComputeGBf16) => true,
        // CustomKernel does NOT override Primitive::is_equivalent in the
        // reference. Identical shader text is not permission to merge it.
        _ => false,
    }
}

#[test]
fn mtplx_compile_simplify_preserves_reference_parent_rules() {
    use super::{
        binary::{Binary, binary},
        compiled::Replay,
        ops,
        stream::{Device, Stream},
        unary::{Unary, unary},
    };
    let stream = Stream::new(0, Device::Gpu);
    let input = Array::unallocated(&[4], Dtype::F32).unwrap();
    let args = std::slice::from_ref(&input);
    let body = |args: &[Array]| {
        let a = unary(&args[0], Unary::Square, stream)?;
        let b = unary(&args[0], Unary::Square, stream)?;
        let c = unary(&a, Unary::Negative, stream)?;
        let d = unary(&b, Unary::Negative, stream)?;
        Ok(vec![binary(&c, &d, Binary::Add, stream)?])
    };
    let raw = Replay::trace(args, body).unwrap().run(args).unwrap();
    let merged = Replay::trace_simplified(args, body)
        .unwrap()
        .run(args)
        .unwrap();
    assert_ne!(raw[0].inputs()[0].id(), raw[0].inputs()[1].id());
    assert_eq!(merged[0].inputs()[0].id(), merged[0].inputs()[1].id());
    let normalizers = Replay::trace_simplified(args, |args| {
        let a = super::normalization::logsumexp_last(&args[0], true, stream)?;
        let b = super::normalization::logsumexp_last(&args[0], true, stream)?;
        Ok(vec![binary(&a, &b, Binary::Add, stream)?])
    })
    .unwrap()
    .run(args)
    .unwrap();
    assert_eq!(
        normalizers[0].inputs()[0].id(),
        normalizers[0].inputs()[1].id()
    );
    // An otherwise equivalent later global output must never disappear.
    let roots = Replay::trace_simplified(args, |args| {
        let a = unary(&args[0], Unary::Square, stream)?;
        let b = unary(&args[0], Unary::Square, stream)?;
        Ok(vec![a, b])
    })
    .unwrap()
    .run(args)
    .unwrap();
    assert_ne!(roots[0].id(), roots[1].id());
    // Distinct raw stop parameters remain distinct even if output values and
    // shapes are equal. compile_simplify compares Slice's stored parameters.
    let slices = Replay::trace_simplified(args, |args| {
        let a = ops::slice(&args[0], &[0], &[3], &[2], stream)?;
        let b = ops::slice(&args[0], &[0], &[4], &[2], stream)?;
        Ok(vec![binary(&a, &b, Binary::Add, stream)?])
    })
    .unwrap()
    .run(args)
    .unwrap();
    assert_ne!(slices[0].inputs()[0].id(), slices[0].inputs()[1].id());
    // Equivalent Split primitives merge all siblings in canonical order.
    let split = Replay::trace_simplified(args, |args| {
        let a = ops::split(&args[0], &[2], 0, stream)?;
        let b = ops::split(&args[0], &[2], 0, stream)?;
        Ok(vec![
            binary(&a[1], &b[1], Binary::Add, stream)?,
            binary(&a[0], &b[0], Binary::Multiply, stream)?,
        ])
    })
    .unwrap()
    .run(args)
    .unwrap();
    for root in &split {
        assert_eq!(root.inputs()[0].id(), root.inputs()[1].id());
    }
    // Same immutable primitive pointer must remain separate (reference rule).
    let shared = Replay::trace_simplified(args, |args| {
        let a = unary(&args[0], Unary::Square, stream)?;
        let b = Array::make_shared_operations(
            stream,
            args,
            &[a.layout().clone()],
            a.operation().unwrap(),
        )?
        .remove(0);
        Ok(vec![binary(&a, &b, Binary::Add, stream)?])
    })
    .unwrap()
    .run(args)
    .unwrap();
    assert_ne!(shared[0].inputs()[0].id(), shared[0].inputs()[1].id());
    // Exercise the original >100-parent grouping path.
    let fanout = Replay::trace_simplified(args, |args| {
        let mut nodes = (0..110)
            .map(|_| unary(&args[0], Unary::Square, stream))
            .collect::<Result<Vec<_>, _>>()?;
        let mut result = nodes.pop().unwrap();
        for node in nodes {
            result = binary(&result, &node, Binary::Add, stream)?;
        }
        Ok(vec![result])
    })
    .unwrap()
    .run(args)
    .unwrap();
    let mut root = fanout[0].clone();
    let square = root.inputs()[1].id();
    for _ in 0..109 {
        assert_eq!(root.inputs()[1].id(), square);
        let next = root.inputs()[0].clone();
        root = next;
    }
    assert_eq!(root.id(), square);
    assert!(Replay::trace_simplified(args, |_| Ok(vec![input.clone()])).is_err());
}

#[test]
#[ignore = "requires Metal buffers and external supervisor; scalar ownership only, no GPU dispatch"]
fn mtplx_compile_scalar_bits_and_custom_kernel_identity() {
    use super::{
        super::gpu::Context,
        configure_sources, ops,
        stream::{Device, Stream},
    };
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let stream = Stream::new(0, Device::Gpu);
    let scalar = |shape: &[i32], dtype: Dtype, bytes: &[u8]| {
        Array::new(shape, dtype, super::scalar_buffer(bytes).unwrap()).unwrap()
    };
    // Available scalar identity is raw bits AND dtype, not float equality:
    // equal NaN payloads merge, signed zero and distinct dtypes do not.
    for (left, right, ld, rd, rank, same) in [
        (
            0x7fc00001u32,
            0x7fc00001u32,
            Dtype::F32,
            Dtype::F32,
            0,
            true,
        ),
        (0, 0x80000000, Dtype::F32, Dtype::F32, 0, false),
        (1, 1, Dtype::U32, Dtype::I32, 0, false),
        (1, 1, Dtype::U32, Dtype::U32, 1, false),
    ] {
        let shape = if rank == 0 { &[][..] } else { &[1][..] };
        let a = scalar(shape, ld, &left.to_le_bytes());
        let b = scalar(shape, rd, &right.to_le_bytes());
        let roots = [
            ops::full(&[2], &a, ld, stream).unwrap(),
            ops::full(&[2], &b, rd, stream).unwrap(),
        ];
        let mut tape = Tape::capture(&roots, &[]).unwrap();
        tape.scalars().unwrap();
        let leaves = tape
            .order
            .iter()
            .filter(|id| tape.nodes[id].array.operation().is_none())
            .count();
        assert_eq!(leaves, if same { 1 } else { 2 });
    }
    // CustomKernel inherits the reference's default false equivalence.
    let input = scalar(&[1], Dtype::F32, &1f32.to_le_bytes());
    let make = || {
        Array::make_operation(stream, &input, input.layout().clone(), Operation::HyperV3R1).unwrap()
    };
    let a = make();
    let b = make();
    assert!(!equivalent_operation(
        &a.operation().unwrap(),
        &b.operation().unwrap(),
        &a,
        &b
    ));
    // Canonical source descriptors survive simplification unchanged.
    let x = ops::full(&[2], &input, Dtype::F32, stream).unwrap();
    let before = x.inputs()[0].id();
    let optimized = simplify(std::slice::from_ref(&x), &[]).unwrap();
    assert_eq!(x.inputs()[0].id(), before);
    assert_ne!(optimized[0].id(), x.id());
}