//! This module provides a framework on top of the normal MIR dataflow framework to simplify the
//! implementation of analyses that track the values stored in places of interest.
//!
//! The default methods of [`ValueAnalysis`] (prefixed with `super_` instead of `handle_`)
//! provide some behavior that should be valid for all abstract domains that are based only on the
//! value stored in a certain place. On top of these default rules, an implementation should
//! override some of the `handle_` methods. For an example, see `ConstAnalysis`.
//!
//! An implementation must also provide a [`Map`]. Before the anaylsis begins, all places that
//! should be tracked during the analysis must be registered. The set of tracked places cannot be
//! changed during the analysis.

use std::fmt::{Debug, Formatter};

use rustc_data_structures::fx::FxHashMap;
use rustc_index::vec::IndexVec;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, Ty, TyCtxt};
use rustc_target::abi::VariantIdx;

use crate::{
    fmt::DebugWithContext, lattice::FlatSet, Analysis, AnalysisDomain, CallReturnPlaces,
    JoinSemiLattice, SwitchIntEdgeEffects,
};

pub trait ValueAnalysis<'tcx> {
    /// For each place of interest, the analysis tracks a value of the given type.
    type Value: Clone + JoinSemiLattice + HasBottom + HasTop;

    const NAME: &'static str;

    fn map(&self) -> &Map;

    fn handle_statement(&self, statement: &Statement<'tcx>, state: &mut State<Self::Value>) {
        self.super_statement(statement, state)
    }

    fn super_statement(&self, statement: &Statement<'tcx>, state: &mut State<Self::Value>) {
        match &statement.kind {
            StatementKind::Assign(box (place, rvalue)) => {
                self.handle_assign(*place, rvalue, state);
            }
            StatementKind::SetDiscriminant { .. } => {
                // Could tread this as writing a constant to a pseudo-place.
            }
            StatementKind::CopyNonOverlapping(..) => {
                // FIXME: What to do here?
            }
            StatementKind::StorageLive(..)
            | StatementKind::StorageDead(..)
            | StatementKind::Deinit(_) => {
                // Could perhaps use these.
            }
            StatementKind::Nop
            | StatementKind::Retag(..)
            | StatementKind::FakeRead(..)
            | StatementKind::Coverage(..)
            | StatementKind::AscribeUserType(..) => (),
        }
    }

    fn handle_assign(
        &self,
        target: Place<'tcx>,
        rvalue: &Rvalue<'tcx>,
        state: &mut State<Self::Value>,
    ) {
        self.super_assign(target, rvalue, state)
    }

    fn super_assign(
        &self,
        target: Place<'tcx>,
        rvalue: &Rvalue<'tcx>,
        state: &mut State<Self::Value>,
    ) {
        match rvalue {
            Rvalue::Ref(_, BorrowKind::Shared, place) => {
                let target_deref = self
                    .map()
                    .find(target.as_ref())
                    .and_then(|target| self.map().apply_elem(target, ProjElem::Deref));
                let place = self.map().find(place.as_ref());
                match (target_deref, place) {
                    (Some(target_deref), Some(place)) => {
                        state.assign_idx(target_deref, ValueOrPlace::Place(place), self.map())
                    }
                    _ => (),
                }
            }
            Rvalue::Ref(_, _, place) | Rvalue::AddressOf(_, place) => {
                state.flood(place.as_ref(), self.map(), Self::Value::top());
            }
            _ => {
                let result = self.handle_rvalue(rvalue, state);
                state.assign(target.as_ref(), result, self.map());
            }
        }
    }

    fn handle_rvalue(
        &self,
        rvalue: &Rvalue<'tcx>,
        state: &mut State<Self::Value>,
    ) -> ValueOrPlace<Self::Value> {
        self.super_rvalue(rvalue, state)
    }

    fn super_rvalue(
        &self,
        rvalue: &Rvalue<'tcx>,
        state: &mut State<Self::Value>,
    ) -> ValueOrPlace<Self::Value> {
        match rvalue {
            Rvalue::Use(operand) => self.handle_operand(operand, state),
            Rvalue::CopyForDeref(place) => self.handle_operand(&Operand::Copy(*place), state),
            Rvalue::Ref(..) | Rvalue::AddressOf(..) => {
                bug!("this rvalue must be handled by handle_assign() or super_assign()")
            }
            _ => {
                // FIXME: Check that other Rvalues really have no side-effect.
                ValueOrPlace::Unknown
            }
        }
    }

    fn handle_operand(
        &self,
        operand: &Operand<'tcx>,
        state: &mut State<Self::Value>,
    ) -> ValueOrPlace<Self::Value> {
        self.super_operand(operand, state)
    }

    fn super_operand(
        &self,
        operand: &Operand<'tcx>,
        state: &mut State<Self::Value>,
    ) -> ValueOrPlace<Self::Value> {
        match operand {
            Operand::Constant(box constant) => {
                ValueOrPlace::Value(self.handle_constant(constant, state))
            }
            Operand::Copy(place) | Operand::Move(place) => {
                // Do want want to handle moves different? Could flood place with bottom.
                self.map()
                    .find(place.as_ref())
                    .map(ValueOrPlace::Place)
                    .unwrap_or(ValueOrPlace::Unknown)
            }
        }
    }

    fn handle_constant(
        &self,
        constant: &Constant<'tcx>,
        state: &mut State<Self::Value>,
    ) -> Self::Value {
        self.super_constant(constant, state)
    }

    fn super_constant(
        &self,
        _constant: &Constant<'tcx>,
        _state: &mut State<Self::Value>,
    ) -> Self::Value {
        Self::Value::top()
    }

    fn handle_terminator(&self, terminator: &Terminator<'tcx>, state: &mut State<Self::Value>) {
        self.super_terminator(terminator, state)
    }

    fn super_terminator(&self, _terminator: &Terminator<'tcx>, _state: &mut State<Self::Value>) {}

    fn handle_call_return(
        &self,
        return_places: CallReturnPlaces<'_, 'tcx>,
        state: &mut State<Self::Value>,
    ) {
        self.super_call_return(return_places, state)
    }

    fn super_call_return(
        &self,
        return_places: CallReturnPlaces<'_, 'tcx>,
        state: &mut State<Self::Value>,
    ) {
        return_places.for_each(|place| {
            state.flood(place.as_ref(), self.map(), Self::Value::top());
        })
    }

    fn handle_switch_int(
        &self,
        discr: &Operand<'tcx>,
        apply_edge_effects: &mut impl SwitchIntEdgeEffects<State<Self::Value>>,
    ) {
        self.super_switch_int(discr, apply_edge_effects)
    }

    fn super_switch_int(
        &self,
        _discr: &Operand<'tcx>,
        _apply_edge_effects: &mut impl SwitchIntEdgeEffects<State<Self::Value>>,
    ) {
    }

    fn wrap(self) -> ValueAnalysisWrapper<Self>
    where
        Self: Sized,
    {
        ValueAnalysisWrapper(self)
    }
}

pub struct ValueAnalysisWrapper<T>(pub T);

impl<'tcx, T: ValueAnalysis<'tcx>> AnalysisDomain<'tcx> for ValueAnalysisWrapper<T> {
    type Domain = State<T::Value>;

    type Direction = crate::Forward;

    const NAME: &'static str = T::NAME;

    fn bottom_value(&self, _body: &Body<'tcx>) -> Self::Domain {
        State(IndexVec::from_elem_n(T::Value::bottom(), self.0.map().value_count))
    }

    fn initialize_start_block(&self, body: &Body<'tcx>, state: &mut Self::Domain) {
        for arg in body.args_iter() {
            state.flood(PlaceRef { local: arg, projection: &[] }, self.0.map(), T::Value::top());
        }
    }
}

impl<'tcx, T> Analysis<'tcx> for ValueAnalysisWrapper<T>
where
    T: ValueAnalysis<'tcx>,
{
    fn apply_statement_effect(
        &self,
        state: &mut Self::Domain,
        statement: &Statement<'tcx>,
        _location: Location,
    ) {
        self.0.handle_statement(statement, state);
    }

    fn apply_terminator_effect(
        &self,
        state: &mut Self::Domain,
        terminator: &Terminator<'tcx>,
        _location: Location,
    ) {
        self.0.handle_terminator(terminator, state);
    }

    fn apply_call_return_effect(
        &self,
        state: &mut Self::Domain,
        _block: BasicBlock,
        return_places: crate::CallReturnPlaces<'_, 'tcx>,
    ) {
        self.0.handle_call_return(return_places, state)
    }

    fn apply_switch_int_edge_effects(
        &self,
        _block: BasicBlock,
        discr: &Operand<'tcx>,
        apply_edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
    ) {
        self.0.handle_switch_int(discr, apply_edge_effects)
    }
}

rustc_index::newtype_index!(
    pub struct PlaceIndex {}
);

rustc_index::newtype_index!(
    struct ValueIndex {}
);

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct State<V>(IndexVec<ValueIndex, V>);

impl<V: Clone + HasTop> State<V> {
    pub fn flood_all(&mut self, value: V) {
        self.0.raw.fill(value);
    }

    pub fn flood(&mut self, place: PlaceRef<'_>, map: &Map, value: V) {
        if let Some(root) = map.find(place) {
            self.flood_idx(root, map, value);
        }
    }

    pub fn flood_idx(&mut self, place: PlaceIndex, map: &Map, value: V) {
        map.preorder_invoke(place, &mut |place| {
            if let Some(vi) = map.places[place].value_index {
                self.0[vi] = value.clone();
            }
        });
    }

    pub fn assign_place_idx(&mut self, target: PlaceIndex, source: PlaceIndex, map: &Map) {
        if let Some(target_value) = map.places[target].value_index {
            if let Some(source_value) = map.places[source].value_index {
                self.0[target_value] = self.0[source_value].clone();
            } else {
                self.0[target_value] = V::top();
            }
        }
        for target_child in map.children(target) {
            // Try to find corresponding child in source.
            let projection = map.places[target_child].proj_elem.unwrap();
            if let Some(source_child) = map.projections.get(&(source, projection)) {
                self.assign_place_idx(target_child, *source_child, map);
            } else {
                self.flood_idx(target_child, map, V::top());
            }
        }
    }

    pub fn assign(&mut self, target: PlaceRef<'_>, result: ValueOrPlace<V>, map: &Map) {
        if let Some(target) = map.find(target) {
            self.assign_idx(target, result, map);
        }
    }

    pub fn assign_idx(&mut self, target: PlaceIndex, result: ValueOrPlace<V>, map: &Map) {
        match result {
            ValueOrPlace::Value(value) => {
                // FIXME: What if not all tracked projections are overwritten? Can this happen?
                if let Some(value_index) = map.places[target].value_index {
                    self.0[value_index] = value;
                }
            }
            ValueOrPlace::Place(source) => self.assign_place_idx(target, source, map),
            ValueOrPlace::Unknown => {
                self.flood_idx(target, map, V::top());
            }
        }
    }

    pub fn get(&self, place: PlaceRef<'_>, map: &Map) -> V {
        map.find(place).map(|place| self.get_idx(place, map)).unwrap_or(V::top())
    }

    pub fn get_idx(&self, place: PlaceIndex, map: &Map) -> V {
        map.places[place].value_index.map(|v| self.0[v].clone()).unwrap_or(V::top())
    }
}

impl<V: JoinSemiLattice> JoinSemiLattice for State<V> {
    fn join(&mut self, other: &Self) -> bool {
        self.0.join(&other.0)
    }
}

#[derive(Debug)]
pub struct Map {
    locals: IndexVec<Local, Option<PlaceIndex>>,
    projections: FxHashMap<(PlaceIndex, ProjElem), PlaceIndex>,
    places: IndexVec<PlaceIndex, PlaceInfo>,
    value_count: usize,
}

impl Map {
    pub fn new() -> Self {
        Self {
            locals: IndexVec::new(),
            projections: FxHashMap::default(),
            places: IndexVec::new(),
            value_count: 0,
        }
    }

    /// Register all places with suitable types up to a certain derefence depth (to prevent cycles).
    pub fn register_with_filter<'tcx>(
        &mut self,
        tcx: TyCtxt<'tcx>,
        source: &impl HasLocalDecls<'tcx>,
        max_derefs: u32,
        mut filter: impl FnMut(Ty<'tcx>) -> bool,
    ) {
        let mut projection = Vec::new();
        for (local, decl) in source.local_decls().iter_enumerated() {
            self.register_with_filter_rec(
                tcx,
                max_derefs,
                local,
                &mut projection,
                decl.ty,
                &mut filter,
            );
        }
    }

    fn register_with_filter_rec<'tcx>(
        &mut self,
        tcx: TyCtxt<'tcx>,
        max_derefs: u32,
        local: Local,
        projection: &mut Vec<PlaceElem<'tcx>>,
        ty: Ty<'tcx>,
        filter: &mut impl FnMut(Ty<'tcx>) -> bool,
    ) {
        if filter(ty) {
            self.register(local, projection)
                .expect("projection should only contain convertible elements");
        }
        if max_derefs > 0 {
            if let Some(ty::TypeAndMut { ty, .. }) = ty.builtin_deref(false) {
                projection.push(PlaceElem::Deref);
                self.register_with_filter_rec(tcx, max_derefs - 1, local, projection, ty, filter);
                projection.pop();
            }
        }
        iter_fields(ty, tcx, |variant, field, ty| {
            if let Some(variant) = variant {
                projection.push(PlaceElem::Downcast(None, variant));
            }
            projection.push(PlaceElem::Field(field, ty));
            self.register_with_filter_rec(tcx, max_derefs, local, projection, ty, filter);
            projection.pop();
            if variant.is_some() {
                projection.pop();
            }
        });
    }

    pub fn register<'tcx>(
        &mut self,
        local: Local,
        projection: &[PlaceElem<'tcx>],
    ) -> Result<(), ()> {
        // Get the base index of the local.
        let mut index =
            *self.locals.get_or_insert_with(local, || self.places.push(PlaceInfo::new(None)));

        // Apply the projection.
        for &elem in projection {
            let elem = elem.try_into()?;
            index = *self.projections.entry((index, elem)).or_insert_with(|| {
                // Prepend new child to the linked list.
                let next = self.places.push(PlaceInfo::new(Some(elem)));
                self.places[next].next_sibling = self.places[index].first_child;
                self.places[index].first_child = Some(next);
                next
            });
        }

        // Allocate a value slot if it doesn't have one.
        if self.places[index].value_index.is_none() {
            self.places[index].value_index = Some(self.value_count.into());
            self.value_count += 1;
        }

        Ok(())
    }

    pub fn apply_elem(&self, place: PlaceIndex, elem: ProjElem) -> Option<PlaceIndex> {
        self.projections.get(&(place, elem)).copied()
    }

    pub fn find(&self, place: PlaceRef<'_>) -> Option<PlaceIndex> {
        let mut index = *self.locals.get(place.local)?.as_ref()?;

        for &elem in place.projection {
            index = self.apply_elem(index, elem.try_into().ok()?)?;
        }

        Some(index)
    }

    pub fn children(&self, parent: PlaceIndex) -> impl Iterator<Item = PlaceIndex> + '_ {
        Children::new(self, parent)
    }

    pub fn preorder_invoke(&self, root: PlaceIndex, f: &mut impl FnMut(PlaceIndex)) {
        f(root);
        for child in self.children(root) {
            self.preorder_invoke(child, f);
        }
    }
}

#[derive(Debug)]
struct PlaceInfo {
    next_sibling: Option<PlaceIndex>,
    first_child: Option<PlaceIndex>,
    /// The projection used to go from parent to this node (only None for root).
    proj_elem: Option<ProjElem>,
    value_index: Option<ValueIndex>,
}

impl PlaceInfo {
    fn new(proj_elem: Option<ProjElem>) -> Self {
        Self { next_sibling: None, first_child: None, proj_elem, value_index: None }
    }
}

struct Children<'a> {
    map: &'a Map,
    next: Option<PlaceIndex>,
}

impl<'a> Children<'a> {
    fn new(map: &'a Map, parent: PlaceIndex) -> Self {
        Self { map, next: map.places[parent].first_child }
    }
}

impl<'a> Iterator for Children<'a> {
    type Item = PlaceIndex;

    fn next(&mut self) -> Option<Self::Item> {
        match self.next {
            Some(child) => {
                self.next = self.map.places[child].next_sibling;
                Some(child)
            }
            None => None,
        }
    }
}

// FIXME: See if we can get rid of `Unknown`.
pub enum ValueOrPlace<V> {
    Value(V),
    Place(PlaceIndex),
    Unknown,
}

pub trait HasBottom {
    fn bottom() -> Self;
}

pub trait HasTop {
    fn top() -> Self;
}

impl<V> HasBottom for FlatSet<V> {
    fn bottom() -> Self {
        Self::Bottom
    }
}

impl<V> HasTop for FlatSet<V> {
    fn top() -> Self {
        Self::Top
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum ProjElem {
    Deref,
    Field(Field),
    Downcast(VariantIdx),
}

impl<V, T> TryFrom<ProjectionElem<V, T>> for ProjElem {
    type Error = ();

    fn try_from(value: ProjectionElem<V, T>) -> Result<Self, Self::Error> {
        match value {
            ProjectionElem::Deref => Ok(ProjElem::Deref),
            ProjectionElem::Field(field, _) => Ok(ProjElem::Field(field)),
            ProjectionElem::Downcast(_, variant) => Ok(ProjElem::Downcast(variant)),
            _ => Err(()),
        }
    }
}

fn iter_fields<'tcx>(
    ty: Ty<'tcx>,
    tcx: TyCtxt<'tcx>,
    mut f: impl FnMut(Option<VariantIdx>, Field, Ty<'tcx>),
) {
    match ty.kind() {
        ty::Tuple(list) => {
            for (field, ty) in list.iter().enumerate() {
                f(None, field.into(), ty);
            }
        }
        ty::Adt(def, substs) => {
            for (v_index, v_def) in def.variants().iter_enumerated() {
                for (f_index, f_def) in v_def.fields.iter().enumerate() {
                    let field_ty = f_def.ty(tcx, substs);
                    let field_ty = tcx
                        .try_normalize_erasing_regions(ty::ParamEnv::reveal_all(), field_ty)
                        .unwrap_or(field_ty);
                    f(Some(v_index), f_index.into(), field_ty);
                }
            }
        }
        ty::Closure(_, substs) => {
            iter_fields(substs.as_closure().tupled_upvars_ty(), tcx, f);
        }
        _ => (),
    }
}

fn debug_with_context_rec<V: Debug + Eq>(
    place: PlaceIndex,
    place_str: &str,
    new: &State<V>,
    old: Option<&State<V>>,
    map: &Map,
    f: &mut Formatter<'_>,
) -> std::fmt::Result {
    if let Some(value) = map.places[place].value_index {
        match old {
            None => writeln!(f, "{}: {:?}", place_str, new.0[value])?,
            Some(old) => {
                if new.0[value] != old.0[value] {
                    writeln!(f, "\u{001f}-{}: {:?}", place_str, old.0[value])?;
                    writeln!(f, "\u{001f}+{}: {:?}", place_str, new.0[value])?;
                }
            }
        }
    }

    for child in map.children(place) {
        let info_elem = map.places[child].proj_elem.unwrap();
        let child_place_str = match info_elem {
            ProjElem::Deref => format!("*{}", place_str),
            ProjElem::Field(field) => {
                if place_str.starts_with("*") {
                    format!("({}).{}", place_str, field.index())
                } else {
                    format!("{}.{}", place_str, field.index())
                }
            }
            ProjElem::Downcast(variant) => format!("({} as #{})", place_str, variant.index()),
        };
        debug_with_context_rec(child, &child_place_str, new, old, map, f)?;
    }

    Ok(())
}

fn debug_with_context<V: Debug + Eq>(
    new: &State<V>,
    old: Option<&State<V>>,
    map: &Map,
    f: &mut Formatter<'_>,
) -> std::fmt::Result {
    for (local, place) in map.locals.iter_enumerated() {
        if let Some(place) = place {
            debug_with_context_rec(*place, &format!("{:?}", local), new, old, map, f)?;
        }
    }
    Ok(())
}

impl<'tcx, T> DebugWithContext<ValueAnalysisWrapper<T>> for State<T::Value>
where
    T: ValueAnalysis<'tcx>,
    T::Value: Debug,
{
    fn fmt_with(&self, ctxt: &ValueAnalysisWrapper<T>, f: &mut Formatter<'_>) -> std::fmt::Result {
        debug_with_context(self, None, ctxt.0.map(), f)
    }

    fn fmt_diff_with(
        &self,
        old: &Self,
        ctxt: &ValueAnalysisWrapper<T>,
        f: &mut Formatter<'_>,
    ) -> std::fmt::Result {
        debug_with_context(self, Some(old), ctxt.0.map(), f)
    }
}