zhc_ir 0.8.3

use crate::{AsOpId, AsOpRef, Formatted, Op};
use std::{hash::Hash, ops::Deref};
use zhc_utils::{Dumpable, FastSet, iter::CollectInSmallVec};

use super::{Depth, Dialect, IR, OpId, Signature, State, ValId, val_ref::ValRef};

/// Borrowed view of an operation within an [`IR`].
///
/// Provides access to operation metadata, arguments, return values, and
/// dependency traversal methods. The reference is tied to the lifetime of the
/// [`IR`] it was obtained from. Derefs to [`OpId`].
#[derive(Debug, Clone)]
pub struct OpRef<'ir, D: Dialect> {
    pub(super) id: OpId,
    pub(super) ir: &'ir IR<D>,
    pub(super) operation: &'ir D::InstructionSet,
    pub(super) signature: &'ir Signature<D::TypeSystem>,
    pub(super) args: &'ir [ValId],
    pub(super) returns: &'ir [ValId],
    pub(super) state: &'ir State,
    pub(super) depth: &'ir Depth,
    pub(super) comment: &'ir Option<String>,
}

impl<'ir, D: Dialect> Hash for OpRef<'ir, D> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.id.hash(state)
    }
}

impl<'ir, D: Dialect> PartialEq for OpRef<'ir, D> {
    fn eq(&self, other: &Self) -> bool {
        std::ptr::eq(self.ir, other.ir) && self.id == other.id
    }
}

impl<'ir, D: Dialect> Eq for OpRef<'ir, D> {}

impl<'ir, D: Dialect> Deref for OpRef<'ir, D> {
    type Target = OpId;

    fn deref(&self) -> &Self::Target {
        &self.id
    }
}

impl<D: Dialect> AsOpId for OpRef<'_, D> {
    fn op_id(&self) -> OpId {
        self.id
    }
}

impl<D: Dialect> AsOpId for &OpRef<'_, D> {
    fn op_id(&self) -> OpId {
        self.id
    }
}

impl<D: Dialect> AsOpRef for OpRef<'_, D> {
    type Dialect = D;
    fn op_ref(&self) -> OpRef<'_, D> {
        self.clone()
    }
}

impl<D: Dialect> AsOpRef for &OpRef<'_, D> {
    type Dialect = D;
    fn op_ref(&self) -> OpRef<'_, D> {
        (*self).clone()
    }
}

impl<'ir, D: Dialect> OpRef<'ir, D> {
    /// Returns an iterator over the operation's argument values without state checking.
    pub(super) fn raw_get_args_iter(&self) -> impl Iterator<Item = ValRef<'ir, D>> + use<'ir, D> {
        self.args.iter().map(|valid| self.ir.raw_get_val(*valid))
    }

    /// Returns an iterator over the operation's return values without state checking.
    pub(super) fn raw_get_returns_iter(
        &self,
    ) -> impl Iterator<Item = ValRef<'ir, D>> + use<'ir, D> {
        self.returns.iter().map(|valid| self.ir.raw_get_val(*valid))
    }
}

impl<'ir, D: Dialect> OpRef<'ir, D> {
    /// Checks if the operation is active.
    pub fn is_active(&self) -> bool {
        self.state.is_active()
    }

    /// Checks if the operation is inactive.
    pub fn is_inactive(&self) -> bool {
        self.state.is_inactive()
    }

    /// Checks if the operation is an input operation.
    ///
    /// An input operation is one that takes no arguments.
    pub fn is_input(&self) -> bool {
        self.signature.get_args_arity() == 0
    }

    /// Checks if the operation is an effect operation.
    ///
    /// An effect operation is one that produces no return values.
    pub fn is_effect(&self) -> bool {
        self.signature.get_returns_arity() == 0
    }

    /// Returns the unique identifier of the operation.
    pub fn get_id(&self) -> OpId {
        self.id
    }

    /// Returns a copy of the instruction's dialect-specific data.
    pub fn get_instruction(&self) -> D::InstructionSet {
        self.operation.clone()
    }

    /// Returns the depth of the operation relative to the IR inputs.
    pub fn get_depth(&self) -> Depth {
        *self.depth
    }

    /// Returns the optional comment attached to this operation.
    pub fn get_comment(&self) -> Option<&str> {
        self.comment.as_deref()
    }

    /// Returns an iterator over the operation's argument values.
    pub fn get_args_iter(&self) -> impl Iterator<Item = ValRef<'ir, D>> + use<'ir, D> {
        self.args.iter().map(|valid| self.ir.get_val(*valid))
    }

    /// Returns the argument value IDs as a slice.
    pub fn get_arg_valids(&self) -> &[ValId] {
        self.args
    }

    /// Returns the number of argument vals.
    pub fn get_args_arity(&self) -> usize {
        self.signature.get_args_arity()
    }

    /// Returns an iterator over the operation's return values.
    pub fn get_returns_iter(&self) -> impl Iterator<Item = ValRef<'ir, D>> + use<'ir, D> {
        self.returns.iter().map(|valid| self.ir.get_val(*valid))
    }

    /// Returns the return value IDs as a slice.
    pub fn get_return_valids(&self) -> &[ValId] {
        self.returns
    }

    /// Returns the number of return vals.
    pub fn get_return_arity(&self) -> usize {
        self.signature.get_returns_arity()
    }

    /// Returns an iterator over the direct users of the current operation.
    ///
    /// Users are deduplicated, meaning that if an operation uses multiple
    /// return values from this operation, it will appear only once in the
    /// iterator.
    pub fn get_users_iter(&self) -> impl Iterator<Item = OpRef<'ir, D>> + use<'ir, D> {
        let mut raw_users = self
            .get_returns_iter()
            .flat_map(|r| r.get_users_iter().map(|a| a.get_id()))
            .collect::<Vec<OpId>>();
        raw_users.sort_unstable();
        raw_users.dedup();
        raw_users.into_iter().map(|a| self.ir.get_op(a))
    }

    /// Checks if the operation has any users.
    pub fn has_users(&self) -> bool {
        self.get_returns_iter().any(|r| r.has_users())
    }

    /// Returns an iterator over the direct predecessors of the current operation.
    ///
    /// Predecessors are deduplicated, meaning that if a predecessor produces
    /// multiple return values used by this operation, it will appear only once
    /// in the iterator.
    pub fn get_predecessors_iter(&self) -> impl Iterator<Item = OpRef<'ir, D>> + use<'ir, D> {
        let mut raw_predecessors = self
            .get_args_iter()
            .map(|r| r.get_origin().opref.get_id())
            .collect::<Vec<_>>();
        raw_predecessors.sort_unstable();
        raw_predecessors.dedup();
        raw_predecessors.into_iter().map(|a| self.ir.get_op(a))
    }

    /// Returns an iterator over all operations that can reach the current operation.
    ///
    /// Performs a backward traversal through the operation graph, collecting all
    /// operations that directly or indirectly produce values used by this operation.
    /// Operations are deduplicated in the result set.
    pub fn get_reaching_iter(&self) -> impl Iterator<Item = OpRef<'ir, D>> + use<'ir, D> {
        let mut output = FastSet::new();
        let mut worklist = vec![self.clone()];
        while let Some(val) = worklist.pop() {
            for op in val.get_args_iter().map(|a| a.get_origin().opref) {
                output.insert(op.clone());
                worklist.push(op);
            }
        }
        output.into_iter()
    }

    /// Returns an iterator over all operations that can reach this operation, including itself.
    ///
    /// Equivalent to [`get_reaching_iter`](Self::get_reaching_iter) with `self` appended.
    pub fn get_inc_reaching_iter(&self) -> impl Iterator<Item = OpRef<'ir, D>> + use<'ir, D> {
        self.get_reaching_iter()
            .chain(std::iter::once(self.to_owned()))
    }

    /// Returns an iterator over all operations that can be reached from the current operation.
    ///
    /// Performs a forward traversal through the operation graph, collecting all
    /// operations that directly or indirectly use values produced by this operation.
    /// Operations are deduplicated in the result set.
    pub fn get_reached_iter(&self) -> impl Iterator<Item = OpRef<'ir, D>> + use<'ir, D> {
        let mut output = FastSet::new();
        let mut worklist = vec![self.clone()];
        while let Some(val) = worklist.pop() {
            for op in val.get_returns_iter().flat_map(|a| a.get_users_iter()) {
                output.insert(op.clone());
                worklist.push(op);
            }
        }
        output.into_iter()
    }

    /// Checks if this operation can reach the specified `other` operation.
    ///
    /// Returns true if this operation produces values that are directly or
    /// indirectly used by `other`, or if this operation and `other` are the
    /// same operation.
    pub fn reaches<'o>(&self, other: &OpRef<'o, D>) -> bool {
        if self == other {
            return true;
        }
        // We try to leverage the depth to make the reachability analysis faster.
        if self.get_depth() >= other.get_depth() {
            // The other can not be reached for sure -> Its depth would be strictly larger.
            return false;
        }
        self.get_users_iter()
            .any(|a| a.get_id() == other.get_id() || a.reaches(other))
    }

    /// Creates a configurable formatter for this operation.
    pub fn format(&self) -> Formatted<'_, Self> {
        Formatted::new(self)
    }

    pub fn as_op(&self) -> Op<D> {
        Op {
            instruction: self.operation.clone(),
            signature: self.signature.clone(),
            args: self.args.iter().cloned().cosvec(),
            returns: self.returns.iter().cloned().cosvec(),
            state: *self.state,
            depth: *self.depth,
            comment: self.comment.clone(),
        }
    }
}

impl<D: Dialect> Dumpable for OpRef<'_, D> {
    fn dump_to_string(&self) -> String {
        self.format().dump_to_string()
    }
}