//! Optimization driver using ISLE rewrite rules on an egraph.

use crate::egraph::Analysis;
use crate::egraph::FuncEGraph;
pub use crate::egraph::{Node, NodeCtx};
use crate::ir::condcodes;
pub use crate::ir::condcodes::{FloatCC, IntCC};
pub use crate::ir::immediates::{Ieee32, Ieee64, Imm64, Offset32, Uimm32, Uimm64, Uimm8};
pub use crate::ir::types::*;
pub use crate::ir::{
    dynamic_to_fixed, AtomicRmwOp, Block, Constant, DynamicStackSlot, FuncRef, GlobalValue, Heap,
    HeapImm, Immediate, InstructionImms, JumpTable, MemFlags, Opcode, StackSlot, Table, TrapCode,
    Type, Value,
};
use crate::isle_common_prelude_methods;
use crate::machinst::isle::*;
use crate::trace;
pub use cranelift_egraph::{Id, NewOrExisting, NodeIter};
use cranelift_entity::{EntityList, EntityRef};
use smallvec::SmallVec;
use std::marker::PhantomData;

pub type IdArray = EntityList<Id>;
#[allow(dead_code)]
pub type Unit = ();
pub type Range = (usize, usize);

pub type ConstructorVec<T> = SmallVec<[T; 8]>;

mod generated_code;
use generated_code::ContextIter;

struct IsleContext<'a, 'b> {
    egraph: &'a mut FuncEGraph<'b>,
}

const REWRITE_LIMIT: usize = 5;

pub fn optimize_eclass<'a>(id: Id, egraph: &mut FuncEGraph<'a>) -> Id {
    trace!("running rules on eclass {}", id.index());
    egraph.stats.rewrite_rule_invoked += 1;

    if egraph.rewrite_depth > REWRITE_LIMIT {
        egraph.stats.rewrite_depth_limit += 1;
        return id;
    }
    egraph.rewrite_depth += 1;

    // Find all possible rewrites and union them in, returning the
    // union.
    let mut ctx = IsleContext { egraph };
    let optimized_ids = generated_code::constructor_simplify(&mut ctx, id);
    let mut union_id = id;
    if let Some(mut ids) = optimized_ids {
        while let Some(new_id) = ids.next(&mut ctx) {
            if ctx.egraph.subsume_ids.contains(&new_id) {
                trace!(" -> eclass {} subsumes {}", new_id, id);
                ctx.egraph.stats.node_subsume += 1;
                // Merge in the unionfind so canonicalization still
                // works, but take *only* the subsuming ID, and break
                // now.
                ctx.egraph.egraph.unionfind.union(union_id, new_id);
                union_id = new_id;
                break;
            }
            ctx.egraph.stats.node_union += 1;
            let old_union_id = union_id;
            union_id = ctx
                .egraph
                .egraph
                .union(&ctx.egraph.node_ctx, union_id, new_id);
            trace!(
                " -> union eclass {} with {} to get {}",
                new_id,
                old_union_id,
                union_id
            );
        }
    }
    trace!(" -> optimize {} got {}", id, union_id);
    ctx.egraph.rewrite_depth -= 1;
    union_id
}

pub(crate) fn store_to_load<'a>(id: Id, egraph: &mut FuncEGraph<'a>) -> Id {
    // Note that we only examine the latest enode in the eclass: opts
    // are invoked for every new enode added to an eclass, so
    // traversing the whole eclass would be redundant.
    let load_key = egraph.egraph.classes[id].get_node().unwrap();
    if let Node::Load {
        op:
            InstructionImms::Load {
                opcode: Opcode::Load,
                offset: load_offset,
                ..
            },
        ty: load_ty,
        addr: load_addr,
        mem_state,
        ..
    } = load_key.node(&egraph.egraph.nodes)
    {
        if let Some(store_inst) = mem_state.as_store() {
            trace!(" -> got load op for id {}", id);
            if let Some((store_ty, store_id)) = egraph.store_nodes.get(&store_inst) {
                trace!(" -> got store id: {} ty: {}", store_id, store_ty);
                let store_key = egraph.egraph.classes[*store_id].get_node().unwrap();
                if let Node::Inst {
                    op:
                        InstructionImms::Store {
                            opcode: Opcode::Store,
                            offset: store_offset,
                            ..
                        },
                    args: store_args,
                    ..
                } = store_key.node(&egraph.egraph.nodes)
                {
                    let store_args = store_args.as_slice(&egraph.node_ctx.args);
                    let store_data = store_args[0];
                    let store_addr = store_args[1];
                    if *load_offset == *store_offset
                        && *load_ty == *store_ty
                        && egraph.egraph.unionfind.equiv_id_mut(*load_addr, store_addr)
                    {
                        trace!(" -> same offset, type, address; forwarding");
                        egraph.stats.store_to_load_forward += 1;
                        return store_data;
                    }
                }
            }
        }
    }

    id
}

struct NodesEtorIter<'a, 'b>
where
    'b: 'a,
{
    root: Id,
    iter: NodeIter<NodeCtx, Analysis>,
    _phantom1: PhantomData<&'a ()>,
    _phantom2: PhantomData<&'b ()>,
}

impl<'a, 'b> generated_code::ContextIter for NodesEtorIter<'a, 'b>
where
    'b: 'a,
{
    type Context = IsleContext<'a, 'b>;
    type Output = (Type, InstructionImms, IdArray);

    fn next(&mut self, ctx: &mut IsleContext<'a, 'b>) -> Option<Self::Output> {
        while let Some(node) = self.iter.next(&ctx.egraph.egraph) {
            trace!("iter from root {}: node {:?}", self.root, node);
            match node {
                Node::Pure {
                    op,
                    args,
                    ty,
                    arity,
                }
                | Node::Inst {
                    op,
                    args,
                    ty,
                    arity,
                    ..
                } if *arity == 1 => {
                    return Some((*ty, op.clone(), args.clone()));
                }
                _ => {}
            }
        }
        None
    }
}

impl<'a, 'b> generated_code::Context for IsleContext<'a, 'b> {
    isle_common_prelude_methods!();

    fn eclass_type(&mut self, eclass: Id) -> Option<Type> {
        let mut iter = self.egraph.egraph.enodes(eclass);
        while let Some(node) = iter.next(&self.egraph.egraph) {
            match node {
                &Node::Pure { ty, arity, .. } | &Node::Inst { ty, arity, .. } if arity == 1 => {
                    return Some(ty);
                }
                &Node::Load { ty, .. } => return Some(ty),
                &Node::Result { ty, .. } => return Some(ty),
                &Node::Param { ty, .. } => return Some(ty),
                _ => {}
            }
        }
        None
    }

    fn at_loop_level(&mut self, eclass: Id) -> (u8, Id) {
        (
            self.egraph.egraph.analysis_value(eclass).loop_level.level() as u8,
            eclass,
        )
    }

    type enodes_etor_iter = NodesEtorIter<'a, 'b>;

    fn enodes_etor(&mut self, eclass: Id) -> Option<NodesEtorIter<'a, 'b>> {
        Some(NodesEtorIter {
            root: eclass,
            iter: self.egraph.egraph.enodes(eclass),
            _phantom1: PhantomData,
            _phantom2: PhantomData,
        })
    }

    fn pure_enode_ctor(&mut self, ty: Type, op: &InstructionImms, args: IdArray) -> Id {
        let op = op.clone();
        match self.egraph.egraph.add(
            Node::Pure {
                op,
                args,
                ty,
                arity: 1,
            },
            &mut self.egraph.node_ctx,
        ) {
            NewOrExisting::New(id) => {
                self.egraph.stats.node_created += 1;
                self.egraph.stats.node_pure += 1;
                self.egraph.stats.node_ctor_created += 1;
                optimize_eclass(id, self.egraph)
            }
            NewOrExisting::Existing(id) => {
                self.egraph.stats.node_ctor_deduped += 1;
                id
            }
        }
    }

    fn id_array_0_etor(&mut self, arg0: IdArray) -> Option<()> {
        let values = arg0.as_slice(&self.egraph.node_ctx.args);
        if values.len() == 0 {
            Some(())
        } else {
            None
        }
    }

    fn id_array_0_ctor(&mut self) -> IdArray {
        EntityList::default()
    }

    fn id_array_1_etor(&mut self, arg0: IdArray) -> Option<Id> {
        let values = arg0.as_slice(&self.egraph.node_ctx.args);
        if values.len() == 1 {
            Some(values[0])
        } else {
            None
        }
    }

    fn id_array_1_ctor(&mut self, arg0: Id) -> IdArray {
        EntityList::from_iter([arg0].into_iter(), &mut self.egraph.node_ctx.args)
    }

    fn id_array_2_etor(&mut self, arg0: IdArray) -> Option<(Id, Id)> {
        let values = arg0.as_slice(&self.egraph.node_ctx.args);
        if values.len() == 2 {
            Some((values[0], values[1]))
        } else {
            None
        }
    }

    fn id_array_2_ctor(&mut self, arg0: Id, arg1: Id) -> IdArray {
        EntityList::from_iter([arg0, arg1].into_iter(), &mut self.egraph.node_ctx.args)
    }

    fn id_array_3_etor(&mut self, arg0: IdArray) -> Option<(Id, Id, Id)> {
        let values = arg0.as_slice(&self.egraph.node_ctx.args);
        if values.len() == 3 {
            Some((values[0], values[1], values[2]))
        } else {
            None
        }
    }

    fn id_array_3_ctor(&mut self, arg0: Id, arg1: Id, arg2: Id) -> IdArray {
        EntityList::from_iter(
            [arg0, arg1, arg2].into_iter(),
            &mut self.egraph.node_ctx.args,
        )
    }

    fn remat(&mut self, id: Id) -> Id {
        trace!("remat: {}", id);
        self.egraph.remat_ids.insert(id);
        id
    }

    fn subsume(&mut self, id: Id) -> Id {
        trace!("subsume: {}", id);
        self.egraph.subsume_ids.insert(id);
        id
    }
}