//! Infers a type for every value (see `docs/model.md`, section 4)
//!
//! A request is a fn, the call place that calls it, and the types of its args.
//! Its result is the type of each of its values. An object is named by the
//! instr that makes it and the call place of the request that runs the instr,
//! and each name has one cell for the whole program. Both come from the
//! program text, so there are finitely many requests and names, and infer
//! ends. A walk reads cells, protos and the returns of other requests, which
//! only grow, so infer walks the program in rounds until one changes nothing

pub mod cells;
pub mod types;

use std::collections::{BTreeMap, HashMap};

use crate::{
    libqjs::parse::{AtomRef, JsValue, opcodes::OpCode},
    native::{self, NATIVES, ProtoOf},
    regify::RegFn,
    transform::{FIRST_ARG, THIS_ARG},
};

use cells::{Cell, Code, LENGTH_FIELD, ObjKind, Protos, read_value};
use types::{BaseType, EMPTY_STRING, Type, is_function, join, loose_eq, strict_eq, test, truthiness, typeof_name};

/// A call place: a fn and the index of its call instr
pub type Place = (usize, usize);

/// The call place of the entry fn, which the host calls
pub const HOST: Place = (usize::MAX, 0);

/// The call place of the fn that the host calls when an IO op completes (see
/// `Code::HostLink`)
pub const HOST_COMPLETE: Place = (usize::MAX, 1);

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct InferRequest {
    pub fn_id: usize,
    pub place: Place,
    pub arg_types: Vec<Type>, // one for each param
}

#[derive(Debug)]
pub struct InferResult {
    pub val_types: Vec<Type>,
    pub instr_types: Vec<Type>, // the type each instr gives, before it is joined into its dst
    pub ret_type: Type,
    pub reachable: Vec<bool>,   // the blocks that can run
}

/// A request and its result
#[derive(Debug)]
pub struct Walk {
    pub req: InferRequest,
    ret: Type,                       // what it returned in the rounds so far: a recursive call gives that
    walking: bool,
    pub result: Option<InferResult>, // in the current round, once it is done
}

#[derive(Debug)]
pub struct TypeCache {
    pub walks: Vec<Walk>,
    ids: HashMap<InferRequest, usize>,              // the walk of each request
    objects: BTreeMap<(usize, usize, Place), Cell>, // the cell of each name: the fn and instr that make the objects, and the call place
    builtins: Vec<Cell>, // the cells of the natives, then of the proto links
    pub protos: Protos,  // the prototype of each `ProtoOf`, as the stdlib links it
    pub host_complete: Type, // what the host calls when an IO op completes, as the stdlib links it
    version: usize,      // goes up on every change of a cell, a return or a proto, and on every new cell or walk
}

impl TypeCache {
    /// The request of a call of `fn_id` from `place` with args of these types.
    /// An arg it does not pass is undefined, and an arg it passes beyond its
    /// params is never read. Only an arg's kind crosses a call, not a known
    /// value, else every literal passed would make a new walk
    fn request(fn_id: usize, place: Place, arg_types: &[Type], fns: &[RegFn]) -> InferRequest {
        let params = (0..fns[fn_id].arg_count).map(|i| arg_types.get(i).map_or(Type::Undefined, Type::widen));
        InferRequest { fn_id, place, arg_types: params.collect() }
    }

    /// The walk of a call that infer walked
    pub fn find(&self, fn_id: usize, place: Place, arg_types: &[Type], fns: &[RegFn]) -> usize {
        self.ids[&Self::request(fn_id, place, arg_types, fns)]
    }

    /// The cell of a built-in fn
    pub fn builtin(&self, code: Code) -> &Cell {
        match code {
            Code::Native(id) => &self.builtins[id],
            Code::ProtoLink(kind) => &self.builtins[NATIVES.len() + kind as usize],
            Code::HostLink => &self.builtins[NATIVES.len() + ProtoOf::ALL.len()],
            Code::Fn(_) => unreachable!(),
        }
    }

    /// The cells of the built-in fns
    pub fn builtins(&self) -> &[Cell] {
        &self.builtins
    }

    /// Every cell: the built-in ones, then one for each name
    pub fn cells(&self) -> impl Iterator<Item = &Cell> {
        self.builtins.iter().chain(self.objects.values())
    }

    /// The cell of the objects of `name`, made the first time
    fn object(&mut self, name: (usize, usize, Place), kind: ObjKind) -> Cell {
        if let Some(cell) = self.objects.get(&name) {
            return cell.clone();
        }

        self.version += 1;
        self.objects.entry(name).or_insert(Cell::new(kind)).clone()
    }

    /// Count a change, if there was one
    fn changed(&mut self, changed: bool) {
        self.version += changed as usize;
    }

    /// Infer a call of `fn_id` from `place`, walking it unless its request is
    /// done in this round. Gives what it returns
    fn call(&mut self, fn_id: usize, place: Place, arg_types: &[Type], fns: &[RegFn]) -> Type {
        let req = Self::request(fn_id, place, arg_types, fns);

        let w = match self.ids.get(&req) {
            Some(&w) => w,
            None => {
                self.version += 1;
                self.ids.insert(req.clone(), self.walks.len());
                self.walks.push(Walk { req, ret: Type::Unknown, walking: false, result: None });
                self.walks.len() - 1
            }
        };

        match &self.walks[w] {
            Walk { walking: true, ret, .. } => return ret.clone(),
            Walk { result: Some(res), .. } => return res.ret_type.clone(),
            _ => {}
        }

        self.walks[w].walking = true;
        let res = walk_fn(w, fns, self);
        let ret = res.ret_type.clone();

        let walk = &mut self.walks[w];
        let joined = join(&walk.ret, &ret);
        let grew = joined != walk.ret;
        walk.walking = false;
        walk.ret = joined;
        walk.result = Some(res);
        self.changed(grew);
        ret
    }
}

/// The join of `f` over the members of `ty`
fn map(ty: &Type, mut f: impl FnMut(&Type) -> Type) -> Type {
    ty.members().iter().fold(Type::Unknown, |all, member| join(&all, &f(member)))
}

/// What `f` gives for every member of `ty`, if all members give the same
fn agree<T: PartialEq>(ty: &Type, f: impl Fn(&Type) -> Option<T>) -> Option<T> {
    let mut members = ty.members().iter().map(f);
    let first = members.next()??;
    members.all(|t| t.as_ref() == Some(&first)).then_some(first)
}

/// What `f` gives for every pair of members of `a` and `b`, if all pairs give
/// the same
fn agree2<T: PartialEq>(a: &Type, b: &Type, f: impl Fn(&Type, &Type) -> Option<T>) -> Option<T> {
    let pairs = a.members().iter().flat_map(|x| b.members().iter().map(move |y| (x, y)));
    let mut results = pairs.map(|(x, y)| f(x, y));
    let first = results.next()??;
    results.all(|t| t.as_ref() == Some(&first)).then_some(first)
}

/// A call of `code` from `place` with operands of these types (see the call
/// rule in `walk_fn()`). Gives what it returns
fn call_code(cache: &mut TypeCache, code: Code, arg_types: &[Type], fns: &[RegFn], place: Place) -> Type {
    let args = &arg_types[FIRST_ARG as usize..];

    match code {
        Code::Fn(fn_id) => cache.call(fn_id, place, arg_types, fns),
        Code::Native(id) => native::ret_type(id, args),
        Code::ProtoLink(kind) => {
            let linked = &cache.protos[kind as usize];
            let joined = join(linked, args.first().unwrap_or(&Type::Undefined));
            cache.changed(joined != *linked);
            cache.protos[kind as usize] = joined;
            Type::Undefined
        }
        Code::HostLink => {
            let joined = join(&cache.host_complete, args.first().unwrap_or(&Type::Undefined));
            cache.changed(joined != cache.host_complete);
            cache.host_complete = joined;
            Type::Undefined
        }
    }
}

/// The calls that the host makes: the entry, with no args, then the fn that
/// it calls when an IO op completes, with the op's handle. Each member of its
/// type is its own env, and `this` is undefined
fn host_calls(cache: &mut TypeCache, fns: &[RegFn]) {
    // The synthetic program entry is appended last by `flatten`
    cache.call(fns.len() - 1, HOST, &[], fns);

    for member in cache.host_complete.clone().members() {
        if let Type::Obj(cell) = member
            && let Some(Code::Fn(fn_id)) = cell.code()
        {
            cache.call(fn_id, HOST_COMPLETE, &[member.clone(), Type::Undefined, Type::Int], fns);
        }
    }
}

/// An array index must be a number. Any other key is a computed field name,
/// which is not supported
fn check_index(index: &Type) {
    for member in index.members() {
        if !matches!(member, Type::Int | Type::Float) {
            unimplemented!("`a[k]` with a key of type {member:?}: computed field names");
        }
    }
}

/// Walk the program from the entry, in rounds, until a round changes
/// nothing. The results are those of the last round
pub fn infer_types(fns: &[RegFn]) -> TypeCache {
    let mut cache = TypeCache {
        walks: vec![],
        ids: HashMap::new(),
        objects: BTreeMap::new(),
        builtins: vec![],
        protos: std::array::from_fn(|_| Type::Unknown),
        host_complete: Type::Unknown,
        version: 0,
    };

    let codes = (0..NATIVES.len()).map(Code::Native).chain(ProtoOf::ALL.map(Code::ProtoLink)).chain([Code::HostLink]);
    cache.builtins = codes.map(|code| Cell::new(ObjKind::Fn(code))).collect();

    loop {
        let version = cache.version;
        cache.walks.iter_mut().for_each(|w| w.result = None);

        host_calls(&mut cache, fns);

        if cache.version == version {
            return cache;
        }
    }
}

fn walk_fn(w: usize, reg_fns: &[RegFn], cache: &mut TypeCache) -> InferResult {
    use OpCode::*;

    let (fn_id, place) = (cache.walks[w].req.fn_id, cache.walks[w].req.place);
    let f = &reg_fns[fn_id];
    let mut val_types = vec![Type::Unknown; f.val_count];
    let mut instr_types = vec![Type::Unknown; f.code.len()];
    let mut ret_type = Type::Unknown;

    // arg `i` is value id `var_count + i` (see regify)
    for (i, ty) in cache.walks[w].req.arg_types.iter().enumerate() {
        val_types[f.var_count + i] = ty.clone();
    }

    // The blocks that can run, given what the types say about each condition
    let mut reachable = vec![false; f.blocks.len()];
    reachable[0] = true;

    // Scratch for calls, kept across instrs
    let mut arg_types: Vec<Type> = vec![];

    // Walk instrs till nothing changes. A return only changes `ret_type` once
    // its block is reachable, and a field read only gives more once a cell
    // changes, so the values, the blocks and the version are enough to watch
    loop {
        let version = cache.version;
        let mut changed = false;

        for (block_idx, block) in f.blocks.iter().enumerate() {
            if !reachable[block_idx] {
                continue;
            }

            let code = &f.code[block.start..block.end];
            let code = code.iter().chain(&block.postfix);

            for (offset, instr) in code.enumerate() {
                // Jumps are the block's to walk (see below). The others give
                // nothing
                if matches!(instr.op, Nop | Drop | Goto | Goto8 | Goto16 | IfFalse | IfFalse8 | IfTrue | IfTrue8) {
                    continue;
                }

                // A rule with an `Unknown` operand gives `Unknown`: no value has
                // reached the instr yet
                if instr.args.iter().any(|&v| val_types[v] == Type::Unknown) {
                    continue;
                }

                // The instr's index in the code (postfix copies have none)
                let idx = block.start + offset;
                let arg = |i: usize| val_types[instr.args[i]].clone();

                // The cell of the objects that this instr makes in this
                // request: the one of its name
                let cell_here = |cache: &mut TypeCache, kind: ObjKind| cache.object((fn_id, idx, place), kind);

                let ty = match instr.op {
                    // ────────────────────────────────────────── Literals ───────────────────────────────────────────────────────
                    PushMinus1 | Push0 | Push1 | Push2 | Push3 | Push4 | Push5 | Push6 | Push7
                    | PushI8 | PushI16 | PushI32 => Type::Int,

                    PushTrue => Type::Bool(Some(true)),
                    PushFalse => Type::Bool(Some(false)),
                    PushAtomValue => Type::Str(Some(instr.imm.atom())),
                    PushEmptyString => Type::Str(Some(EMPTY_STRING)),

                    PushConst | PushConst8 => match f.cpool[instr.imm.const_idx()] {
                        JsValue::Int(_) => Type::Int,
                        JsValue::Float(_) => Type::Float,
                        JsValue::String(_) => Type::Str(None),
                        JsValue::BoolTrue => Type::Bool(Some(true)),
                        JsValue::BoolFalse => Type::Bool(Some(false)),
                        JsValue::Undefined => Type::Undefined,
                        JsValue::NativeFn(id) => Type::Obj(cache.builtin(Code::Native(id)).clone()),
                        JsValue::ProtoLink(kind) => Type::Obj(cache.builtin(Code::ProtoLink(kind)).clone()),
                        JsValue::HostLink => Type::Obj(cache.builtin(Code::HostLink).clone()),
                        _ => unreachable!(),
                    },

                    Undefined => Type::Undefined,
                    Null => Type::Null,

                    // ──────────────────────────────────────── Arithmetic ───────────────────────────────────────────────────────
                    // `+` concatenates if either side becomes a string, which
                    // a string or an object does. Else it adds numbers
                    Add => {
                        let r = arg(0);

                        map(&arg(1), |l| map(&r, |r| match (l, r) {
                            (l, r) if [l, r].iter().any(|t| matches!(t, Type::Str(_) | Type::Obj(_))) => Type::Str(None),
                            (Type::Int, Type::Int) => Type::Int,
                            _ => Type::Float,
                        }))
                    }
                    // Any other operator converts its operands to numbers
                    Sub | Mul | Mod | Pow => {
                        let r = arg(0);

                        map(&arg(1), |l| map(&r, |r| match (l, r) {
                            (Type::Int, Type::Int) => Type::Int,
                            _ => Type::Float,
                        }))
                    }
                    Neg | Plus | Inc | Dec | PostInc | PostDec => map(&arg(0), |t| match t {
                        Type::Int => Type::Int,
                        _ => Type::Float,
                    }),
                    Div => Type::Float,
                    // Bitwise NOT gives an int32
                    Not => Type::Int,

                    Lt | Lte | Gt | Gte => Type::Bool(None),

                    // ─────────────────────────────── Conditions the types may decide ──────────────────────────────────────────
                    // A known value lets infer skip the branch that cannot run
                    // (see `reachable` below)
                    StrictEq | StrictNeq | Eq | Neq => {
                        let eq = match instr.op {
                            StrictEq | StrictNeq => agree2(&arg(0), &arg(1), strict_eq),
                            _ => agree2(&arg(0), &arg(1), loose_eq),
                        };

                        Type::Bool(eq.map(|eq| eq == matches!(instr.op, StrictEq | Eq)))
                    }
                    Lnot => Type::Bool(agree(&arg(0), truthiness).map(|t| !t)),
                    Typeof => match agree(&arg(0), typeof_name) {
                        Some(name) => Type::Str(Some(AtomRef::Predefined(name))),
                        None => Type::Str(None),
                    },
                    IsUndefined | TypeofIsUndefined => Type::Bool(agree(&arg(0), |t| test(t, &[BaseType::Undefined]))),
                    IsNull => Type::Bool(agree(&arg(0), |t| test(t, &[BaseType::Null]))),
                    IsUndefinedOrNull => Type::Bool(agree(&arg(0), |t| test(t, &[BaseType::Undefined, BaseType::Null]))),
                    TypeofIsFunction => Type::Bool(agree(&arg(0), is_function)),

                    // ─────────────────────────────────────── Value Movers ──────────────────────────────────────────────────────
                    SetLoc | SetName => arg(0),

                    // ───────────────────────────────────────── Objects ─────────────────────────────────────────────────────────
                    Object => Type::Obj(cell_here(cache, ObjKind::Plain)),
                    Fclosure | Fclosure8 => {
                        let JsValue::FnById(child) = f.cpool[instr.imm.const_idx()] else {
                            unreachable!()
                        };

                        Type::Obj(cell_here(cache, ObjKind::Fn(Code::Fn(child))))
                    }
                    // The elements are its operands
                    ArrayFrom => {
                        let cell = cell_here(cache, ObjKind::Array { elems: Type::Unknown });

                        for &v in &instr.args {
                            let grew = cell.write_elem(&val_types[v]);
                            cache.changed(grew);
                        }

                        Type::Obj(cell)
                    }
                    // `a[i]` and `a[i] = x`, on arrays, and `s[i]` on strings:
                    // an index is a number. Other objects need computed field
                    // names, which are not supported. A string cannot be
                    // written (module code is strict), and undefined and null
                    // throw
                    GetArrayEl => {
                        let index = arg(0);

                        map(&arg(1), |obj| match obj {
                            Type::Obj(cell) if cell.is_array() => {
                                check_index(&index);
                                cell.read_elem()
                            }
                            Type::Str(_) => {
                                check_index(&index);
                                join(&Type::Str(None), &Type::Undefined)
                            }
                            Type::Undefined | Type::Null => Type::Unknown,
                            obj => unimplemented!("`o[k]` on {obj:?}: computed field names"),
                        })
                    }
                    PutArrayEl => {
                        let (value, index) = (arg(0), arg(1));

                        for obj in arg(2).members() {
                            match obj {
                                Type::Obj(cell) if cell.is_array() => {
                                    check_index(&index);
                                    let grew = cell.write_elem(&value);
                                    cache.changed(grew);
                                }
                                Type::Undefined | Type::Null | Type::Str(_) => {}
                                obj => unimplemented!("`o[k] = v` on {obj:?}: computed field names"),
                            }
                        }

                        Type::Unknown
                    }
                    // A key stays as it is: a number is an index
                    ToPropkey => arg(0),
                    // The write goes into the object's cell, which every holder
                    // of the object sees. A primitive has no fields to write,
                    // and undefined and null throw
                    DefineField | PutField => {
                        let (field, value) = (instr.imm.atom(), arg(0));

                        for obj in arg(1).members() {
                            if let Type::Obj(cell) = obj {
                                let grew = cell.write(field, &value, instr.op == DefineField);
                                cache.changed(grew);
                            }
                        }

                        Type::Unknown
                    }
                    GetField => read_value(&arg(0), instr.imm.atom(), &cache.protos),
                    // A string's length is its own. Anything else reads a
                    // `length` field
                    GetLength => map(&arg(0), |obj| match obj {
                        Type::Str(_) => Type::Int,
                        obj => read_value(obj, LENGTH_FIELD, &cache.protos),
                    }),

                    // ───────────────────────────────────── Function Calls ──────────────────────────────────────────────────────
                    // Operands are the callee's params: the callee (its env),
                    // `this`, then the args (see regify). A callee that is not
                    // a fn throws, so it gives nothing. A call of each member
                    // of a sum passes that member as the env
                    Call | CallMethod | CallConstructor => {
                        let ret = map(&arg(0), |callee| {
                            let Type::Obj(cell) = callee else {
                                return Type::Unknown;
                            };
                            let Some(code) = cell.code() else {
                                return Type::Unknown;
                            };

                            arg_types.clear();
                            arg_types.extend(instr.args.iter().map(|&v| val_types[v].clone()));
                            arg_types[0] = callee.clone();
                            call_code(cache, code, &arg_types, reg_fns, (fn_id, idx))
                        });

                        // `new` gives what the constructor returns if that is an
                        // object, else the new object itself, `this`
                        match instr.op {
                            CallConstructor => map(&ret, |ret| match ret {
                                Type::Obj(_) => ret.clone(),
                                _ => arg(THIS_ARG as usize),
                            }),
                            _ => ret,
                        }
                    }

                    // ──────────────────────────────────────── Returns ──────────────────────────────────────────────────────────
                    Return | ReturnAsync => arg(0),
                    ReturnUndef => Type::Undefined,

                    _ => unreachable!("{instr:?}"),
                };

                // Postfix copies come after the block's own instrs, and have no
                // entry: a copy gives the type of what it copies
                if idx < block.end {
                    instr_types[idx] = join(&instr_types[idx], &ty);
                }

                if matches!(instr.op, Return | ReturnAsync | ReturnUndef) {
                    ret_type = join(&ret_type, &ty);
                    continue;
                }

                for &dst in &instr.dsts {
                    let joined = join(&val_types[dst], &ty);
                    changed |= joined != val_types[dst];
                    val_types[dst] = joined;
                }
            } // instr

            // Mark the blocks this one can go to. A conditional jump whose
            // condition the types decide goes one way only. One whose
            // condition nothing has reached yet goes nowhere yet
            // (A block that only carries copies on an edge has no instrs)
            let jumps = match f.code[block.start..block.end].last() {
                Some(last) if matches!(last.op, IfFalse | IfFalse8 | IfTrue | IfTrue8) => {
                    match &val_types[last.args[0]] {
                        Type::Unknown => continue,
                        cond => agree(cond, truthiness).map(|t| t == matches!(last.op, IfTrue | IfTrue8)),
                    }
                }
                _ => None,
            };

            // succ[0] is the jump target, succ[1] the fall-through
            for (i, &succ) in block.succ.iter().enumerate() {
                if jumps.is_none_or(|jumps| jumps == (i == 0)) && !reachable[succ] {
                    reachable[succ] = true;
                    changed = true;
                }
            }
        } // block

        if !changed && cache.version == version {
            break;
        }
    } // loop

    InferResult {
        val_types,
        instr_types,
        ret_type,
        reachable,
    }
}