jingle 0.5.4

use crate::analysis::cpa::lattice::JoinSemiLattice;
use crate::analysis::pcode_store::{PcodeOpRef, PcodeStore};
use crate::modeling::machine::cpu::concrete::ConcretePcodeAddress;
use jingle_sleigh::PcodeOperation;
use std::borrow::Borrow;
use std::cmp::Ordering;
use std::fmt::{Debug, Display, Formatter, Result as FmtResult};
use std::ops::{Add, AddAssign};

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum MergeOutcome {
    NoOp,
    Merged,
}

impl MergeOutcome {
    pub fn merged(&self) -> bool {
        matches!(self, MergeOutcome::Merged)
    }
}

impl Add for MergeOutcome {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        match (self, rhs) {
            (Self::NoOp, Self::NoOp) => Self::NoOp,
            _ => Self::Merged,
        }
    }
}

impl AddAssign for MergeOutcome {
    fn add_assign(&mut self, rhs: Self) {
        (*self) = *self + rhs
    }
}

/// Object-safe trait for cloneable iterators stored inside `DynIterBox`.
///
/// Any iterator type that is also `Clone` will implement this via the
/// blanket impl below; the `clone_box` method allows cloning through the
/// trait object so the boxed iterator can be cloned cheaply.
pub trait CloneableIter<'a, T>: Iterator<Item = T> {
    fn clone_box(&self) -> Box<dyn CloneableIter<'a, T> + 'a>;
}

impl<'a, T, I> CloneableIter<'a, T> for I
where
    I: Iterator<Item = T> + Clone + 'a,
{
    fn clone_box(&self) -> Box<dyn CloneableIter<'a, T> + 'a> {
        Box::new(self.clone())
    }
}

/// Newtype wrapper around a boxed cloneable iterator.
///
/// This provides `Clone` and `Iterator` implementations for the boxed trait
/// object, so upstream types (like `Successor::IntoIter`) can be `Clone`,
/// which is required by combinators such as `iproduct`.
pub struct DynIterBox<'a, T>(Box<dyn CloneableIter<'a, T> + 'a>);

impl<'a, T> DynIterBox<'a, T> {
    pub fn new<I>(iter: I) -> Self
    where
        I: Iterator<Item = T> + Clone + 'a,
    {
        DynIterBox(Box::new(iter))
    }
}

impl<'a, T> Clone for DynIterBox<'a, T> {
    fn clone(&self) -> Self {
        DynIterBox(self.0.clone_box())
    }
}

impl<'a, T> Iterator for DynIterBox<'a, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.0.next()
    }

    // Forward size_hint if available
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}

impl<'a, T> Debug for DynIterBox<'a, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        // We can't show the inner iterator state meaningfully, so provide a simple marker.
        f.write_str("DynIterBox { .. }")
    }
}

/// Iterator wrapper returned by `transfer` methods.
///
/// This stores a boxed, cloneable iterator so the `Successor` itself can be
/// `Clone` without forcing collection of items into a `Vec`.
///
/// This structure, [DynIterBox], and [ClonableIter] are my quick best-effort at
/// balancing the following goals:
///
/// * [AbstractState::transfer] returns an iterator.
/// * [AbstractState::transfer] impls can return heterogenous iterator types easily
///   (e.g. one path returns [std::iter::empty] and another returns [std::iter::once])
/// * [AbstractState] Implementations don't have to juggle an associated type or generic
///   arguments for the transfer function.
/// * The iterators produced by [AbstractState::transfer] can be cloned for usage with `iproduct`
///
/// There's definitely some trade-offs with this approach and this might change in the future.
pub struct Successor<'a, T>(DynIterBox<'a, T>);

impl<'a, T> Successor<'a, T> {
    /// Create an empty successor iterator.
    pub fn empty() -> Self
    where
        T: 'a,
    {
        // std::iter::Empty implements Clone, so it satisfies CloneableIter via the blanket impl.
        Self(DynIterBox::new(std::iter::empty::<T>()))
    }
}

impl<'a, T> Clone for Successor<'a, T> {
    fn clone(&self) -> Self {
        Successor(self.0.clone())
    }
}

impl<'a, T: 'a> IntoIterator for Successor<'a, T> {
    type Item = T;
    type IntoIter = DynIterBox<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.0
    }
}

/// Construct a `Successor` from any iterator that implements `Clone`.
///
/// We require the iterator to be `Clone` so that the boxed trait object can be
/// cloned cheaply. This avoids collecting into a `Vec`.
impl<'a, T: 'a, I> From<I> for Successor<'a, T>
where
    I: Iterator<Item = T> + Clone + 'a,
{
    fn from(value: I) -> Self {
        Self(DynIterBox::new(value))
    }
}

/// Core trait for abstract states used by the CPA.
pub trait AbstractState: JoinSemiLattice + Clone + Debug + Display {
    /// Merge `other` into `self`. Mutate `self` and return whether merging occurred.
    /// The mutated `self` MUST be >= than it was before.
    fn merge(&mut self, other: &Self) -> MergeOutcome;

    /// Default cartesian merge using `join`.
    fn merge_join(&mut self, new_state: &Self) -> MergeOutcome {
        if self == new_state {
            MergeOutcome::NoOp
        } else {
            self.join(new_state);
            MergeOutcome::Merged
        }
    }

    /// Default separate merge (no-op).
    fn merge_sep(&mut self, _: &Self) -> MergeOutcome {
        MergeOutcome::NoOp
    }

    /// Stop predicate: is `self` covered by any of `states`?
    fn stop<'a, T: Iterator<Item = &'a Self>>(&'a self, states: T) -> bool;

    /// Default stop predicate using piecewise ordering.
    fn stop_sep<'a, T: Iterator<Item = &'a Self>>(&'a self, mut states: T) -> bool {
        states.any(|s| {
            matches!(
                PartialOrd::partial_cmp(self, s),
                Some(Ordering::Less) | Some(Ordering::Equal)
            )
        })
    }

    /// Transfer function: return successor abstract states for a pcode operation.
    ///
    /// The returned `Successor` must be constructed from an iterator that
    /// implements `Clone`. This lets CPA clients clone the successor sequence
    /// cheaply when needed.
    fn transfer<'a, B: Borrow<PcodeOperation>>(&'a self, opcode: B) -> Successor<'a, Self>;
}

/// States that know their program location.
pub trait LocationState: AbstractState {
    /// Retrieve the p-code operation for this location from a pcode store.
    ///
    /// Relaxed lifetime: the returned operation borrows from the provided
    /// `PcodeStore` reference only; the state (`&self`) is not required to be
    /// borrowed for the same lifetime. This allows callers to obtain an op
    /// tied to the store borrow without forcing the state to outlive that borrow.
    fn get_operation<'op, T: PcodeStore<'op> + ?Sized>(&self, t: &'op T)
    -> Option<PcodeOpRef<'op>>;
    fn get_location(&self) -> Option<ConcretePcodeAddress>;
}