//! One walk: the types that one request gives the locals, args and regs of its fn
//!
//! The walk goes over the reachable blocks until nothing changes. Each instr has a rule,
//! which gives the type of its result from the types of its operands. The result joins into
//! the instr's dsts in memory. A rule with an `Unknown` operand gives nothing: no value has
//! reached the instr yet

use super::{
    InferRequest, TypeCache, binding_addr,
    memory::{Addr, InstrRef, ObjAddr},
    types::{
        BaseType, EMPTY_STRING, LENGTH_FIELD, Shape, Type, decide_pair, is_base, join, join_all,
        loose_eq, pairs, strict_eq, truthiness,
    },
};
use crate::{
    libqjs::parse::{JsValue, opcodes::OpCode},
    native,
    regify::{BasicBlock, Operand, RegFn, RegInstr},
    transform::{FIRST_ARG, THIS_ARG},
};

/// Walk `w`. Gives what it returns, and the blocks that can run
pub fn walk_fn(w: usize, fns: &[RegFn], cache: &mut TypeCache) -> (Type, Vec<bool>) {
    let req = cache.walks[w].req.clone();
    let f = &fns[req.fn_idx];

    let walker = Walker {
        w,
        f,
        fns,
        captured: f.captured_slots(),
        req,
        cache,
        ret: Type::Unknown,
        reachable: vec![false; f.blocks.len()],
    };

    walker.run()
}

struct Walker<'a> {
    w: usize,
    req: InferRequest,
    f: &'a RegFn,
    fns: &'a [RegFn],
    captured: Vec<Operand>, // the slots of `f` that a closure shares
    cache: &'a mut TypeCache,
    ret: Type,
    reachable: Vec<bool>,
}

impl Walker<'_> {
    fn addr(&self, slot: Operand) -> Addr {
        binding_addr(self.w, &self.req, &self.captured, slot)
    }

    /// The address of the objects that the instr at `idx` makes in this walk
    fn here(&self, idx: usize) -> ObjAddr {
        ObjAddr {
            made_at: InstrRef {
                fn_idx: self.req.fn_idx,
                instr_idx: idx,
            },
            in_call: self.req.in_call,
        }
    }

    fn run(mut self) -> (Type, Vec<bool>) {
        // The args get their values from the request
        for (i, ty) in self.req.arg_types.iter().enumerate() {
            self.cache.memory.join(self.addr(Operand::Arg(i)), ty);
        }

        self.reachable[0] = true;

        // A return only changes once its block is reachable, and a read only gives more
        // once memory changes, so the blocks and the count of changes are enough to watch
        loop {
            let before = (self.cache.progress(), self.reachable.clone());
            let f = self.f;

            for (i, block) in f.blocks.iter().enumerate() {
                if self.reachable[i] {
                    self.walk_block(block);
                }
            }

            if (self.cache.progress(), &self.reachable) == (before.0, &before.1) {
                return (self.ret, self.reachable);
            }
        }
    }

    fn walk_block(&mut self, block: &BasicBlock) {
        let f = self.f;
        let code = f.code[block.start..block.end].iter().chain(&block.postfix);

        // A postfix copy has no index in the code, but its rule never needs one
        for (offset, instr) in code.enumerate() {
            self.step(instr, block.start + offset);
        }

        self.mark_successors(block);
    }

    fn step(&mut self, instr: &RegInstr, idx: usize) {
        use OpCode::*;

        // Jumps are the block's to walk (see `mark_successors()`). A `close_loc` or a
        // `close_field` gives a binding a new location, which has the same address (see
        // transform/closure.rs)
        if matches!(
            instr.op,
            Nop | CloseLoc
                | CloseField
                | Goto
                | Goto8
                | Goto16
                | IfFalse
                | IfFalse8
                | IfTrue
                | IfTrue8
        ) {
            return;
        }

        let args: Vec<Type> = instr
            .args
            .iter()
            .map(|&slot| self.cache.memory.get(self.addr(slot)).clone())
            .collect();

        // The first operand of `define_field_ref` is a location, not a value, and a binding
        // has no value before its first write
        let values = if instr.op == DefineFieldRef {
            &args[1..]
        } else {
            &args[..]
        };

        if values.contains(&Type::Unknown) {
            return;
        }

        let ty = self.rule(instr, idx, &args);

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

        for &dst in &instr.dsts {
            self.cache.memory.join(self.addr(dst), &ty);
        }
    }

    /// The type that `instr` gives, from the types of its operands `args`. An instr that
    /// gives nothing gives `Unknown`. Operands that an instr pops come top first (see
    /// regify), so a binary operator has its right side in `args[0]`
    fn rule(&mut self, instr: &RegInstr, idx: usize, args: &[Type]) -> Type {
        use OpCode::*;

        let memory = &self.cache.memory;

        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 => self.constant(instr),
            Undefined => Type::Undefined,
            Null => Type::Null,

            // Arithmetic. Every number is a double, and an int is a double that is known to
            // fit in an int32 (see emit). So an operator that can leave that range, or give -0
            // or NaN, gives a float. `+` concatenates if either side becomes a string, which a
            // string or an object does
            Add => join_all(pairs(&args[1], &args[0]).map(|(l, r)| {
                let text = |t: &Type| matches!(t, Type::Str(_) | Type::Ref(_) | Type::Native(_));

                match text(l) || text(r) {
                    true => Type::Str(None),
                    false => Type::Float,
                }
            })),
            Sub | Mul | Div | Mod | Pow | Neg | Inc | Dec | PostInc | PostDec => Type::Float,
            // `+x` is x as a number: an int stays an int
            Plus => join_all(args[0].members().iter().map(|t| match t {
                Type::Int => Type::Int,
                _ => Type::Float,
            })),
            // The bitwise operators give an int32, but `>>>` gives a uint32
            Not | And | Or | Xor | Shl | Sar => Type::Int,
            Shr => Type::Float,
            Lt | Lte | Gt | Gte => Type::Bool(None),

            // Tests. A known result lets the walk skip a branch that cannot run
            StrictEq | StrictNeq | Eq | Neq => {
                let eq = match instr.op {
                    StrictEq | StrictNeq => decide_pair(&args[0], &args[1], strict_eq),
                    _ => decide_pair(&args[0], &args[1], loose_eq),
                };

                Type::Bool(eq.map(|eq| eq == matches!(instr.op, StrictEq | Eq)))
            }
            Lnot => Type::Bool(args[0].decide(truthiness).map(|t| !t)),
            Typeof => Type::Str(args[0].decide(|t| memory.typeof_name(t))),
            IsUndefined | TypeofIsUndefined => {
                Type::Bool(args[0].decide(|t| is_base(t, &[BaseType::Undefined])))
            }
            IsNull => Type::Bool(args[0].decide(|t| is_base(t, &[BaseType::Null]))),
            IsUndefinedOrNull => {
                Type::Bool(args[0].decide(|t| is_base(t, &[BaseType::Undefined, BaseType::Null])))
            }
            TypeofIsFunction => Type::Bool(args[0].decide(|t| memory.is_function(t))),

            // Moves. A key stays as it is: a number is an index
            SetLoc | SetName | ToPropkey => args[0].clone(),

            // Objects
            Object => self.make(idx, Type::Shape(Shape::Obj(vec![]))),
            Fclosure => {
                let JsValue::FnByIdx(fn_idx) = self.f.cpool[instr.imm.const_idx()] else {
                    unreachable!()
                };

                self.make(
                    idx,
                    Type::Closure {
                        fn_idx,
                        fields: vec![],
                    },
                )
            }
            ArrayFrom => {
                let array = self.make(idx, Type::Shape(Shape::Array(Box::new(Type::Unknown))));

                for value in args {
                    self.cache.memory.write_elem(self.here(idx), value);
                }

                array
            }
            GetArrayEl => self.get_array_el(&args[1], &args[0]),
            PutArrayEl => self.put_array_el(&args[2], &args[1], &args[0]),
            DefineField | PutField => self.write_field(instr, &args[1], &args[0]),
            DefineFieldRef => self.share_field(instr, args),
            GetField => memory.read_field(&args[0], instr.imm.atom()),
            // A string's length is its own. Anything else reads a `length` field
            GetLength => join_all(args[0].members().iter().map(|t| match t {
                Type::Str(_) => Type::Int,
                t => memory.read_field(t, LENGTH_FIELD),
            })),

            // Calls
            Call | CallMethod | CallConstructor => self.call(instr, idx, args),

            // Returns
            Return | ReturnAsync => args[0].clone(),
            ReturnUndef => Type::Undefined,

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

    /// The object that the instr at `idx` makes in this walk, made with `entry` the first time
    fn make(&mut self, idx: usize, entry: Type) -> Type {
        self.cache.memory.make(self.here(idx), entry);
        Type::Ref(self.here(idx))
    }

    /// A constant from the fn's constant pool
    fn constant(&self, instr: &RegInstr) -> Type {
        match &self.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(idx) => Type::Native(*idx),
            c => unimplemented!("the constant {c:?}"),
        }
    }

    /// `a[i]` on an array, or `s[i]` on a string: an index is a number. Other objects need
    /// computed field names, which are not supported. Undefined and null throw
    fn get_array_el(&self, obj: &Type, index: &Type) -> Type {
        join_all(obj.members().iter().map(|obj| match obj {
            Type::Ref(obj) if self.cache.memory.is_array(*obj) => {
                check_index(index);
                self.cache.memory.read_elem(*obj)
            }
            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"),
        }))
    }

    /// `a[i] = v` on an array. A string cannot be written (module code is strict), and
    /// undefined and null throw
    fn put_array_el(&mut self, obj: &Type, index: &Type, value: &Type) -> Type {
        for obj in obj.members() {
            match obj {
                // A write past the end leaves holes before it, which hold undefined
                Type::Ref(obj) if self.cache.memory.is_array(*obj) => {
                    check_index(index);
                    self.cache.memory.write_elem(*obj, value);
                    self.cache.memory.write_elem(*obj, &Type::Undefined);
                }
                Type::Undefined | Type::Null | Type::Str(_) => {}
                obj => unimplemented!("`o[k] = v` on {obj:?}: computed field names"),
            }
        }

        Type::Unknown
    }

    /// `define_field` and `put_field`. The write goes into the object's field, which every
    /// holder of the object sees. A primitive has no fields to write, and undefined and null
    /// throw
    fn write_field(&mut self, instr: &RegInstr, obj: &Type, value: &Type) -> Type {
        for obj in obj.members() {
            if let Type::Ref(obj) = obj {
                let defined = instr.op == OpCode::DefineField;
                self.cache
                    .memory
                    .write_field(*obj, instr.imm.atom(), value, defined);
            }
        }

        Type::Unknown
    }

    /// `define_field_ref`: the field becomes a location, which regify gives as the first
    /// operand. That is a binding, or the reg of an object whose field of the same name is
    /// the location (see regify/walk.rs)
    fn share_field(&mut self, instr: &RegInstr, args: &[Type]) -> Type {
        let name = instr.imm.atom();

        let location = match instr.args[0] {
            Operand::Reg(_) => {
                let holders = args[0].members().iter().map(|holder| match holder {
                    Type::Ref(holder) => self.cache.memory.field(*holder, name).map(|f| f.addr),
                    holder => unreachable!("a capture held by {holder:?}"),
                });

                let mut locations: Vec<Addr> = holders.flatten().collect();
                locations.dedup();
                assert!(
                    locations.len() <= 1,
                    "a capture of one of several locations"
                );
                locations.pop()
            }
            slot => Some(self.addr(slot)),
        };

        // A holder's field can be missing until the walk that makes it gets there
        if let Some(location) = location {
            for obj in args[1].members() {
                if let Type::Ref(obj) = obj {
                    self.cache.memory.share_field(*obj, name, location);
                }
            }
        }

        Type::Unknown
    }

    /// Operands are the callee's params: the callee (its env), `this`, then the args (see
    /// regify). A call of each member of a sum passes that member as the env. A callee that
    /// is not a fn throws, so it gives nothing
    ///
    /// `new` gives what the constructor returns if that is an object, else the new object
    /// itself, `this`
    fn call(&mut self, instr: &RegInstr, idx: usize, args: &[Type]) -> Type {
        let in_call = Some(InstrRef {
            fn_idx: self.req.fn_idx,
            instr_idx: idx,
        });

        let rets: Vec<Type> = (args[0].members().iter())
            .map(|callee| {
                // A native takes the args only, without its env and `this`
                if let Type::Native(idx) = callee {
                    return native::ret_type(*idx, &args[FIRST_ARG as usize..]);
                }

                let Type::Ref(obj) = callee else {
                    return Type::Unknown;
                };
                let Some(fn_idx) = self.cache.memory.closure_fn(*obj) else {
                    return Type::Unknown;
                };

                let mut call_args = args.to_vec();
                call_args[0] = callee.clone();
                self.cache.call(fn_idx, in_call, &call_args, self.fns)
            })
            .collect();

        let ret = join_all(rets);

        if instr.op != OpCode::CallConstructor {
            return ret;
        }

        join_all(ret.members().iter().map(|ret| match ret {
            Type::Ref(_) => ret.clone(),
            _ => args[THIS_ARG as usize].clone(),
        }))
    }

    /// Mark the blocks that `block` 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
    fn mark_successors(&mut self, block: &BasicBlock) {
        use OpCode::*;

        let last = &self.f.code[block.end - 1];

        let jumps = match last.op {
            IfFalse | IfFalse8 | IfTrue | IfTrue8 => {
                match self.cache.memory.get(self.addr(last.args[0])) {
                    Type::Unknown => return,
                    cond => cond
                        .decide(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)) {
                self.reachable[succ] = true;
            }
        }
    }
}

/// 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");
        }
    }
}