lift-core 0.4.8

use crate::blocks::BlockKey;
use crate::context::Context;
use crate::dialect::DialectRegistry;
use crate::operations::OpKey;
use crate::values::ValueKey;
use std::collections::HashSet;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum VerifyError {
    #[error("SSA violation: value {0:?} used but not defined")]
    UndefinedValue(ValueKey),

    #[error("SSA violation: value {0:?} defined more than once")]
    MultipleDefinition(ValueKey),

    #[error("Dominance violation: value {0:?} used before definition in block {1:?}")]
    DominanceViolation(ValueKey, BlockKey),

    #[error("Type mismatch in operation {op:?}: expected {expected}, got {actual}")]
    TypeMismatch {
        op: OpKey,
        expected: String,
        actual: String,
    },

    #[error("Linearity violation: qubit value {0:?} consumed more than once")]
    LinearityViolation(ValueKey),

    #[error("Linearity violation: qubit {0:?} not consumed (leaked)")]
    QubitLeaked(ValueKey),

    #[error("Branch linearity: arms consume different qubit sets at block {0:?}")]
    BranchLinearityMismatch(BlockKey),

    #[error("Dangling reference: {0}")]
    DanglingReference(String),

    #[error("Empty block {0:?} has no terminator")]
    MissingTerminator(BlockKey),

    #[error("Operation {0:?} has no parent block")]
    OrphanedOperation(OpKey),

    #[error("Block {0:?} has no parent region")]
    OrphanedBlock(BlockKey),

    #[error("Invalid operation: {0}")]
    InvalidOperation(String),

    #[error("Semantic error in operation {op:?}: {message}")]
    SemanticError { op: OpKey, message: String },
}

pub struct Verifier<'a> {
    ctx: &'a Context,
    errors: Vec<VerifyError>,
    defined: HashSet<ValueKey>,
    consumed_qubits: HashSet<ValueKey>,
}

impl<'a> Verifier<'a> {
    pub fn new(ctx: &'a Context) -> Self {
        Self {
            ctx,
            errors: Vec::new(),
            defined: HashSet::new(),
            consumed_qubits: HashSet::new(),
        }
    }

    pub fn verify_all(&mut self) -> Result<(), Vec<VerifyError>> {
        self.verify_ssa();
        self.verify_dominance();
        self.verify_well_formedness();
        self.verify_linearity();

        if self.errors.is_empty() {
            Ok(())
        } else {
            Err(std::mem::take(&mut self.errors))
        }
    }

    /// Runs SSA, well-formedness, linearity, and semantic (dialect) checks.
    pub fn verify_all_with_dialects(
        &mut self,
        registry: &DialectRegistry,
    ) -> Result<(), Vec<VerifyError>> {
        self.verify_ssa();
        self.verify_dominance();
        self.verify_well_formedness();
        self.verify_linearity();
        self.verify_semantics(registry);

        if self.errors.is_empty() {
            Ok(())
        } else {
            Err(std::mem::take(&mut self.errors))
        }
    }

    /// Validates each operation against its dialect's signature: the number of
    /// inputs/results and, where the dialect provides it, type compatibility.
    fn verify_semantics(&mut self, registry: &DialectRegistry) {
        for (op_key, op) in &self.ctx.ops {
            let dialect_name = self.ctx.strings.resolve(op.dialect);
            let op_name = self.ctx.strings.resolve(op.name);

            let dialect = match registry.get(dialect_name) {
                Some(d) => d,
                None => {
                    // Unknown dialect: skip semantic checks (core ops are
                    // validated by the core dialect when registered).
                    continue;
                }
            };

            if let Err(msg) = dialect.verify_op(op_name, op.inputs.len(), op.results.len()) {
                self.errors.push(VerifyError::SemanticError {
                    op: op_key,
                    message: msg,
                });
            }
        }
    }

    fn verify_ssa(&mut self) {
        let mut all_defined: HashSet<ValueKey> = HashSet::new();

        // Collect all defined values from block args
        for (block_key, block) in &self.ctx.blocks {
            for &arg_key in &block.args {
                if !all_defined.insert(arg_key) {
                    self.errors.push(VerifyError::MultipleDefinition(arg_key));
                }
            }
            let _ = block_key;
        }

        // Collect all defined values from operation results
        for (_op_key, op) in &self.ctx.ops {
            for &result_key in &op.results {
                if !all_defined.insert(result_key) {
                    self.errors
                        .push(VerifyError::MultipleDefinition(result_key));
                }
            }
        }

        // Verify all uses are defined
        for (_op_key, op) in &self.ctx.ops {
            for &input_key in &op.inputs {
                if !all_defined.contains(&input_key) {
                    self.errors.push(VerifyError::UndefinedValue(input_key));
                }
            }
        }

        self.defined = all_defined;
    }

    /// Checks that every op's inputs are defined before that op runs, within
    /// its own block's program order (block args count as defined from the
    /// start). `verify_ssa` only checks that a used value is defined
    /// *somewhere* in the whole context, with no ordering — accepting a
    /// value consumed by an op that appears before the op that defines it.
    ///
    /// Only values this same block itself defines (its own args, or results
    /// of its own ops) are checked for ordering here: a value owned by a
    /// different block is left to `verify_ssa`'s existence check. Every
    /// function body today is a single block (no branches yet), so that
    /// scope limit does not miss anything reachable in practice, and it
    /// avoids false positives if a future multi-block construct legitimately
    /// threads a value in from an enclosing scope.
    fn verify_dominance(&mut self) {
        use std::collections::HashMap;

        let mut owner_block: HashMap<ValueKey, BlockKey> = HashMap::new();
        for (block_key, block) in &self.ctx.blocks {
            for &arg in &block.args {
                owner_block.insert(arg, block_key);
            }
            for &op_key in &block.ops {
                if let Some(op) = self.ctx.ops.get(op_key) {
                    for &result in &op.results {
                        owner_block.insert(result, block_key);
                    }
                }
            }
        }

        for (block_key, block) in &self.ctx.blocks {
            let mut visible: HashSet<ValueKey> = block.args.iter().copied().collect();
            for &op_key in &block.ops {
                let Some(op) = self.ctx.ops.get(op_key) else {
                    continue;
                };
                for &input in &op.inputs {
                    if owner_block.get(&input) == Some(&block_key) && !visible.contains(&input) {
                        self.errors
                            .push(VerifyError::DominanceViolation(input, block_key));
                    }
                }
                for &result in &op.results {
                    visible.insert(result);
                }
            }
        }
    }

    fn verify_well_formedness(&mut self) {
        // Verify all operation inputs reference valid values
        for (op_key, op) in &self.ctx.ops {
            for &input in &op.inputs {
                if !self.ctx.values.contains_key(input) {
                    self.errors.push(VerifyError::DanglingReference(format!(
                        "Operation {:?} references non-existent value {:?}",
                        op_key, input
                    )));
                }
            }
            for &result in &op.results {
                if !self.ctx.values.contains_key(result) {
                    self.errors.push(VerifyError::DanglingReference(format!(
                        "Operation {:?} references non-existent result {:?}",
                        op_key, result
                    )));
                }
            }
            for &region in &op.regions {
                if !self.ctx.regions.contains_key(region) {
                    self.errors.push(VerifyError::DanglingReference(format!(
                        "Operation {:?} references non-existent region {:?}",
                        op_key, region
                    )));
                }
            }
        }

        // Verify blocks reference valid operations
        for (block_key, block) in &self.ctx.blocks {
            for &op in &block.ops {
                if !self.ctx.ops.contains_key(op) {
                    self.errors.push(VerifyError::DanglingReference(format!(
                        "Block {:?} references non-existent operation {:?}",
                        block_key, op
                    )));
                }
            }
        }

        // Verify regions reference valid blocks
        for (region_key, region) in &self.ctx.regions {
            for &block in &region.blocks {
                if !self.ctx.blocks.contains_key(block) {
                    self.errors.push(VerifyError::DanglingReference(format!(
                        "Region {:?} references non-existent block {:?}",
                        region_key, block
                    )));
                }
            }
        }
    }

    fn verify_linearity(&mut self) {
        let mut consumed: HashSet<ValueKey> = HashSet::new();
        let mut all_qubits: HashSet<ValueKey> = HashSet::new();

        // Identify all qubit values
        for (val_key, val) in &self.ctx.values {
            if self.ctx.is_qubit_type(val.ty) {
                all_qubits.insert(val_key);
            }
        }

        // Check each operation's inputs for qubit consumption
        for (_op_key, op) in &self.ctx.ops {
            let op_name = self.ctx.strings.resolve(op.name);

            for &input in &op.inputs {
                if let Some(val) = self.ctx.values.get(input) {
                    if self.ctx.is_qubit_type(val.ty) && !consumed.insert(input) {
                        self.errors.push(VerifyError::LinearityViolation(input));
                    }
                }
            }

            // quantum.measure and quantum.reset consume the qubit
            // All gate operations produce new qubit values (SSA)
            if op_name == "quantum.measure" {
                // Qubit is consumed, classical bit produced — no new qubit
            }
        }

        self.consumed_qubits = consumed;
    }

    pub fn errors(&self) -> &[VerifyError] {
        &self.errors
    }
}

pub fn verify(ctx: &Context) -> Result<(), Vec<VerifyError>> {
    let mut verifier = Verifier::new(ctx);
    verifier.verify_all()
}

/// Verifies a context including dialect-level semantic checks.
pub fn verify_with_dialects(
    ctx: &Context,
    registry: &DialectRegistry,
) -> Result<(), Vec<VerifyError>> {
    let mut verifier = Verifier::new(ctx);
    verifier.verify_all_with_dialects(registry)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_context_verifies() {
        let ctx = Context::new();
        assert!(verify(&ctx).is_ok());
    }

    #[test]
    fn test_simple_ssa_valid() {
        let mut ctx = Context::new();
        let f32_ty = ctx.make_float_type(32);

        let block = ctx.create_block();
        let arg = ctx.create_block_arg(block, f32_ty);

        let (op, _results) = ctx.create_op(
            "tensor.relu",
            "tensor",
            vec![arg],
            vec![f32_ty],
            crate::attributes::Attributes::new(),
            crate::location::Location::unknown(),
        );
        ctx.add_op_to_block(block, op);

        assert!(verify(&ctx).is_ok());
    }

    /// Regression test: an op consuming a value produced by a *later* op in
    /// the same block (use-before-def / a textbook dominance violation) used
    /// to verify successfully — `verify_ssa` only checked that the value
    /// existed somewhere in the context, never that it was defined before
    /// its use.
    #[test]
    fn test_use_before_def_is_a_dominance_violation() {
        let mut ctx = Context::new();
        let f32_ty = ctx.make_float_type(32);
        let block = ctx.create_block();
        let arg = ctx.create_block_arg(block, f32_ty);

        // Create the op that DEFINES `later_result` first (in slotmap/build
        // order), but only add it to the block AFTER the op that uses it —
        // so in block.ops program order, the use comes before the def.
        let (producer, producer_results) = ctx.create_op(
            "tensor.relu",
            "tensor",
            vec![arg],
            vec![f32_ty],
            crate::attributes::Attributes::new(),
            crate::location::Location::unknown(),
        );
        let (consumer, _) = ctx.create_op(
            "tensor.relu",
            "tensor",
            vec![producer_results[0]],
            vec![f32_ty],
            crate::attributes::Attributes::new(),
            crate::location::Location::unknown(),
        );

        // Program order: consumer, then producer — consumer's input is not
        // yet defined at that point in the block.
        ctx.add_op_to_block(block, consumer);
        ctx.add_op_to_block(block, producer);

        let result = verify(&ctx);
        assert!(result.is_err(), "use-before-def must fail verification");
        let errors = result.unwrap_err();
        assert!(
            errors
                .iter()
                .any(|e| matches!(e, VerifyError::DominanceViolation(_, _))),
            "{:?}",
            errors
        );
    }

    #[test]
    fn test_qubit_linearity_violation() {
        let mut ctx = Context::new();
        let qubit_ty = ctx.make_qubit_type();

        let block = ctx.create_block();
        let q0 = ctx.create_block_arg(block, qubit_ty);

        // First use of q0 — ok
        let (op1, _) = ctx.create_op(
            "quantum.x",
            "quantum",
            vec![q0],
            vec![qubit_ty],
            crate::attributes::Attributes::new(),
            crate::location::Location::unknown(),
        );
        ctx.add_op_to_block(block, op1);

        // Second use of q0 — linearity violation!
        let (op2, _) = ctx.create_op(
            "quantum.h",
            "quantum",
            vec![q0],
            vec![qubit_ty],
            crate::attributes::Attributes::new(),
            crate::location::Location::unknown(),
        );
        ctx.add_op_to_block(block, op2);

        let result = verify(&ctx);
        assert!(result.is_err());
        let errors = result.unwrap_err();
        assert!(errors
            .iter()
            .any(|e| matches!(e, VerifyError::LinearityViolation(_))));
    }

    #[test]
    fn test_semantic_verification_detects_wrong_input_count() {
        use crate::dialect::Dialect;

        #[derive(Debug)]
        struct FakeTensorDialect;
        impl Dialect for FakeTensorDialect {
            fn name(&self) -> &str {
                "tensor"
            }
            fn verify_op(
                &self,
                op_name: &str,
                num_inputs: usize,
                _num_results: usize,
            ) -> Result<(), String> {
                if op_name == "tensor.matmul" && num_inputs != 2 {
                    return Err(format!("matmul expects 2 inputs, got {}", num_inputs));
                }
                Ok(())
            }
        }

        let mut registry = DialectRegistry::new();
        registry.register(Box::new(FakeTensorDialect));

        let mut ctx = Context::new();
        let f32_ty = ctx.make_float_type(32);
        let block = ctx.create_block();
        let a = ctx.create_block_arg(block, f32_ty);
        let b = ctx.create_block_arg(block, f32_ty);
        let c = ctx.create_block_arg(block, f32_ty);

        // matmul with 3 inputs -> should fail semantic check
        let (op, _) = ctx.create_op(
            "tensor.matmul",
            "tensor",
            vec![a, b, c],
            vec![f32_ty],
            crate::attributes::Attributes::new(),
            crate::location::Location::unknown(),
        );
        ctx.add_op_to_block(block, op);

        let result = verify_with_dialects(&ctx, &registry);
        assert!(result.is_err());
        let errors = result.unwrap_err();
        assert!(errors
            .iter()
            .any(|e| matches!(e, VerifyError::SemanticError { .. })));
    }
}