lift-export 0.4.6

use lift_core::context::Context;
use lift_core::values::{DefSite, ValueKey};
use lift_quantum::gates::QuantumGate;
use std::collections::HashMap;
use std::fmt::Write;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum QasmExportError {
    #[error("Unsupported gate for QASM export: {0}")]
    UnsupportedGate(String),
    #[error("Export error: {0}")]
    General(String),
}

#[derive(Debug)]
pub struct QasmExporter;

/// Resolves a value to the physical qubit index that first introduced it,
/// by walking the SSA def chain backwards. A gate's results carry on the
/// *same* physical qubit as the input at the same position (results[i] is
/// what inputs[i] became after the gate acted on it), so tracing an
/// arbitrary operand back through a chain of gates always bottoms out at
/// the entry-block argument that owns that qubit. Resolved indices are
/// memoised in `qubit_index` so later gates resolve in O(1).
fn resolve_qubit_index(
    ctx: &Context,
    qubit_index: &mut HashMap<ValueKey, usize>,
    value: ValueKey,
) -> usize {
    if let Some(&idx) = qubit_index.get(&value) {
        return idx;
    }
    let idx = match ctx.get_value(value).map(|v| v.def.clone()) {
        Some(DefSite::OpResult { op, result_index }) => ctx
            .get_op(op)
            .and_then(|op_data| op_data.inputs.get(result_index as usize).copied())
            .map(|input| resolve_qubit_index(ctx, qubit_index, input))
            .unwrap_or(0),
        Some(DefSite::BlockArg { arg_index, .. }) => arg_index as usize,
        None => 0,
    };
    qubit_index.insert(value, idx);
    idx
}

impl QasmExporter {
    pub fn new() -> Self {
        Self
    }

    pub fn export(&self, ctx: &Context) -> Result<String, QasmExportError> {
        let mut output = String::new();

        let _ = writeln!(output, "OPENQASM 3.0;");
        let _ = writeln!(output, "// Generated by LIFT framework");
        let _ = writeln!(output);

        // OpenQASM describes a single circuit. Export the first function (in
        // declaration order) that contains at least one quantum gate, rather
        // than flattening every function in the module into one register —
        // that used to produce e.g. `qubit[4] q;` out of two unrelated
        // 2-qubit functions with their gates concatenated.
        let target_region = ctx.modules.iter().find_map(|module| {
            module.functions.iter().find_map(|func| {
                let region = ctx.get_region(func.body?)?;
                let has_gate = region.blocks.iter().any(|&bk| {
                    ctx.get_block(bk).is_some_and(|b| {
                        b.ops.iter().any(|&ok| {
                            ctx.get_op(ok).is_some_and(|op| {
                                QuantumGate::from_name(ctx.strings.resolve(op.name)).is_some()
                            })
                        })
                    })
                });
                has_gate.then_some(region)
            })
        });

        let Some(region) = target_region else {
            let _ = writeln!(output, "qubit[1] q;");
            let _ = writeln!(output, "bit[1] c;");
            return Ok(output);
        };

        // Assign each qubit-typed entry-block argument a stable physical
        // index, in declaration order.
        let mut qubit_index: HashMap<ValueKey, usize> = HashMap::new();
        if let Some(&entry) = region.blocks.first() {
            if let Some(block) = ctx.get_block(entry) {
                for &arg in &block.args {
                    if ctx.get_value(arg).is_some_and(|v| ctx.is_qubit_type(v.ty)) {
                        let idx = qubit_index.len();
                        qubit_index.insert(arg, idx);
                    }
                }
            }
        }
        let num_qubits = qubit_index.len().max(1);

        let _ = writeln!(output, "qubit[{}] q;", num_qubits);
        let _ = writeln!(output, "bit[{}] c;", num_qubits);
        let _ = writeln!(output);

        // Walk gates in real circuit order (block.ops), not the ops slotmap:
        // slotmap iteration order drifts once a pass removes an op and a
        // later pass inserts a new one into the freed slot, which used to
        // export gates out of order.
        for &block_key in &region.blocks {
            let Some(block) = ctx.get_block(block_key) else {
                continue;
            };
            for &op_key in &block.ops {
                let Some(op) = ctx.get_op(op_key) else {
                    continue;
                };
                let op_name = ctx.strings.resolve(op.name).to_string();

                if let Some(gate) = QuantumGate::from_name(&op_name) {
                    let q0 = op
                        .inputs
                        .first()
                        .map(|&v| resolve_qubit_index(ctx, &mut qubit_index, v))
                        .unwrap_or(0);
                    let q1 = op
                        .inputs
                        .get(1)
                        .map(|&v| resolve_qubit_index(ctx, &mut qubit_index, v))
                        .unwrap_or(0);
                    let q2 = op
                        .inputs
                        .get(2)
                        .map(|&v| resolve_qubit_index(ctx, &mut qubit_index, v))
                        .unwrap_or(0);
                    let angle = op.attrs.get_float("angle").unwrap_or(0.0);
                    let theta = op.attrs.get_float("theta").unwrap_or(0.0);
                    let phi = op.attrs.get_float("phi").unwrap_or(0.0);
                    let lambda = op.attrs.get_float("lambda").unwrap_or(0.0);

                    match gate {
                        // 1-qubit standard
                        QuantumGate::H => {
                            let _ = writeln!(output, "h q[{}];", q0);
                        }
                        QuantumGate::X => {
                            let _ = writeln!(output, "x q[{}];", q0);
                        }
                        QuantumGate::Y => {
                            let _ = writeln!(output, "y q[{}];", q0);
                        }
                        QuantumGate::Z => {
                            let _ = writeln!(output, "z q[{}];", q0);
                        }
                        QuantumGate::S => {
                            let _ = writeln!(output, "s q[{}];", q0);
                        }
                        QuantumGate::Sdg => {
                            let _ = writeln!(output, "sdg q[{}];", q0);
                        }
                        QuantumGate::T => {
                            let _ = writeln!(output, "t q[{}];", q0);
                        }
                        QuantumGate::Tdg => {
                            let _ = writeln!(output, "tdg q[{}];", q0);
                        }
                        QuantumGate::SX => {
                            let _ = writeln!(output, "sx q[{}];", q0);
                        }
                        // 1-qubit parametric
                        QuantumGate::RX => {
                            let _ = writeln!(output, "rx({}) q[{}];", angle, q0);
                        }
                        QuantumGate::RY => {
                            let _ = writeln!(output, "ry({}) q[{}];", angle, q0);
                        }
                        QuantumGate::RZ => {
                            let _ = writeln!(output, "rz({}) q[{}];", angle, q0);
                        }
                        QuantumGate::P => {
                            let _ = writeln!(output, "p({}) q[{}];", angle, q0);
                        }
                        QuantumGate::U1 => {
                            let _ = writeln!(output, "u1({}) q[{}];", lambda, q0);
                        }
                        QuantumGate::U2 => {
                            let _ = writeln!(output, "u2({}, {}) q[{}];", phi, lambda, q0);
                        }
                        QuantumGate::U3 => {
                            let _ =
                                writeln!(output, "u3({}, {}, {}) q[{}];", theta, phi, lambda, q0);
                        }
                        // 1-qubit fixed-angle
                        QuantumGate::Rx90 => {
                            let _ = writeln!(output, "rx(pi/2) q[{}];", q0);
                        }
                        QuantumGate::Rx180 => {
                            let _ = writeln!(output, "rx(pi) q[{}];", q0);
                        }
                        // 1-qubit special
                        QuantumGate::VirtualRZ => {
                            let _ = writeln!(output, "rz({}) q[{}]; // virtual", angle, q0);
                        }
                        QuantumGate::GlobalPhase => {
                            let _ = writeln!(output, "gphase({});", angle);
                        }
                        // 2-qubit standard
                        QuantumGate::CX => {
                            let _ = writeln!(output, "cx q[{}], q[{}];", q0, q1);
                        }
                        QuantumGate::CZ => {
                            let _ = writeln!(output, "cz q[{}], q[{}];", q0, q1);
                        }
                        QuantumGate::CY => {
                            let _ = writeln!(output, "cy q[{}], q[{}];", q0, q1);
                        }
                        QuantumGate::SWAP => {
                            let _ = writeln!(output, "swap q[{}], q[{}];", q0, q1);
                        }
                        QuantumGate::ISWAP => {
                            let _ = writeln!(output, "iswap q[{}], q[{}];", q0, q1);
                        }
                        QuantumGate::ECR => {
                            let _ = writeln!(output, "ecr q[{}], q[{}];", q0, q1);
                        }
                        QuantumGate::RZX => {
                            let _ = writeln!(output, "rzx({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        // 2-qubit parametric
                        QuantumGate::CP => {
                            let _ = writeln!(output, "cp({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        QuantumGate::CPhase => {
                            let _ = writeln!(output, "cphase({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        QuantumGate::XX => {
                            let _ = writeln!(output, "rxx({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        QuantumGate::YY => {
                            let _ = writeln!(output, "ryy({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        QuantumGate::ZZ => {
                            let _ = writeln!(output, "rzz({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        QuantumGate::XY => {
                            let _ = writeln!(output, "xy({}) q[{}], q[{}];", angle, q0, q1);
                        }
                        QuantumGate::MS => {
                            let _ = writeln!(output, "ms q[{}], q[{}];", q0, q1);
                        }
                        // 3-qubit
                        QuantumGate::CCX => {
                            let _ = writeln!(output, "ccx q[{}], q[{}], q[{}];", q0, q1, q2);
                        }
                        QuantumGate::CSWAP => {
                            let _ = writeln!(output, "cswap q[{}], q[{}], q[{}];", q0, q1, q2);
                        }
                        // Multi-controlled
                        QuantumGate::MCX => {
                            let _ = writeln!(output, "mcx q[{}], q[{}], q[{}];", q0, q1, q2);
                        }
                        QuantumGate::MCZ => {
                            let _ = writeln!(output, "mcz q[{}], q[{}], q[{}];", q0, q1, q2);
                        }
                        // Measurement and control
                        QuantumGate::Measure => {
                            let _ = writeln!(output, "c[{}] = measure q[{}];", q0, q0);
                        }
                        QuantumGate::MeasureAll => {
                            for i in 0..num_qubits {
                                let _ = writeln!(output, "c[{}] = measure q[{}];", i, i);
                            }
                        }
                        QuantumGate::Reset => {
                            let _ = writeln!(output, "reset q[{}];", q0);
                        }
                        QuantumGate::Barrier => {
                            let _ = writeln!(output, "barrier q;");
                        }
                        QuantumGate::Init => {
                            let _ = writeln!(output, "reset q[{}];", q0);
                        }
                        QuantumGate::Delay => {
                            let _ = writeln!(output, "delay[100ns] q[{}];", q0);
                        }
                        // IonQ native
                        QuantumGate::GPI => {
                            let _ = writeln!(output, "gpi({}) q[{}];", angle, q0);
                        }
                        QuantumGate::GPI2 => {
                            let _ = writeln!(output, "gpi2({}) q[{}];", angle, q0);
                        }
                        // Control flow and generic
                        QuantumGate::IfElse => {
                            let _ = writeln!(output, "// if-else control flow");
                        }
                        QuantumGate::ParamGate => {
                            let _ = writeln!(output, "// parameterised gate: {}", op_name);
                        }
                    }

                    // Thread the qubit identity forward: results[i] is the
                    // same physical qubit as inputs[i] after the gate acts.
                    for (i, &result) in op.results.iter().enumerate() {
                        if let Some(&input) = op.inputs.get(i) {
                            let idx = resolve_qubit_index(ctx, &mut qubit_index, input);
                            qubit_index.insert(result, idx);
                        }
                    }
                }
            }
        }

        Ok(output)
    }
}

impl Default for QasmExporter {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use lift_core::attributes::Attributes;
    use lift_core::functions::FunctionData;
    use lift_core::location::Location;

    /// Builds a one-function, one-block module with `num_qubits` qubit
    /// params, wires it into `ctx.modules`, and returns the block plus its
    /// qubit-typed block args (in declaration order) for gate-building.
    fn build_module(
        ctx: &mut Context,
        num_qubits: usize,
    ) -> (lift_core::blocks::BlockKey, Vec<ValueKey>) {
        let qty = ctx.make_qubit_type();
        let region = ctx.create_region();
        let block = ctx.create_block();
        ctx.add_block_to_region(region, block);

        let qubits: Vec<ValueKey> = (0..num_qubits)
            .map(|_| ctx.create_block_arg(block, qty))
            .collect();

        let name = ctx.intern_string("f");
        let mut func = FunctionData::new(name, vec![qty; num_qubits], vec![qty; num_qubits]);
        func.body = Some(region);

        let module_idx = ctx.create_module("m");
        ctx.add_function_to_module(module_idx, func);

        (block, qubits)
    }

    fn gate1(
        ctx: &mut Context,
        block: lift_core::blocks::BlockKey,
        name: &str,
        input: ValueKey,
    ) -> ValueKey {
        let qty = ctx.make_qubit_type();
        let (op, results) = ctx.create_op(
            name,
            "quantum",
            vec![input],
            vec![qty],
            Attributes::new(),
            Location::unknown(),
        );
        ctx.add_op_to_block(block, op);
        results[0]
    }

    fn gate2(
        ctx: &mut Context,
        block: lift_core::blocks::BlockKey,
        name: &str,
        a: ValueKey,
        b: ValueKey,
    ) -> (ValueKey, ValueKey) {
        let qty = ctx.make_qubit_type();
        let (op, results) = ctx.create_op(
            name,
            "quantum",
            vec![a, b],
            vec![qty, qty],
            Attributes::new(),
            Location::unknown(),
        );
        ctx.add_op_to_block(block, op);
        (results[0], results[1])
    }

    /// Regression test for #2: three gates in a row on the same qubit (q2)
    /// used to come out on q[0], q[1], q[2] because the exporter numbered
    /// qubits from a counter instead of following the operands.
    #[test]
    fn test_qubit_index_follows_operands_not_a_counter() {
        let mut ctx = Context::new();
        let (block, qubits) = build_module(&mut ctx, 3);
        let a = gate1(&mut ctx, block, "quantum.x", qubits[2]);
        let b = gate1(&mut ctx, block, "quantum.y", a);
        let _ = gate1(&mut ctx, block, "quantum.z", b);

        let qasm = QasmExporter::new().export(&ctx).unwrap();
        assert!(qasm.contains("qubit[3] q;"), "{qasm}");
        assert!(qasm.contains("x q[2];"), "{qasm}");
        assert!(qasm.contains("y q[2];"), "{qasm}");
        assert!(qasm.contains("z q[2];"), "{qasm}");
        assert!(
            !qasm.contains("x q[0];") && !qasm.contains("y q[1];"),
            "gates on q2 must not spill onto q0/q1:\n{qasm}"
        );
    }

    /// Regression test for #2: control/target were swapped because CX's
    /// second operand (the target) was numbered as if it were the next
    /// counter slot instead of following its own operand chain.
    #[test]
    fn test_cx_control_and_target_are_not_swapped() {
        let mut ctx = Context::new();
        let (block, qubits) = build_module(&mut ctx, 2);
        let h_out = gate1(&mut ctx, block, "quantum.h", qubits[0]);
        let _ = gate2(&mut ctx, block, "quantum.cx", h_out, qubits[1]);

        let qasm = QasmExporter::new().export(&ctx).unwrap();
        assert!(qasm.contains("h q[0];"), "{qasm}");
        assert!(
            qasm.contains("cx q[0], q[1];"),
            "control must stay on q0, target on q1:\n{qasm}"
        );
    }
}