//! This module lowers closures by making implicit environment captures explicit
//!
//! Replace:
//! ```js
//! let x = 0;
//! let y = true;
//!
//! function fn() {
//!     f(x)
//!     f(y)
//! }
//! ```
//!
//! With:
//! ```js
//! let x = 0;
//! let y = true;
//!
//! function fn() {
//!     f(fn.<var>.x);
//!     f(fn.<var>.y);
//! }
//!
//! fn.<var> = {};
//! fn.<var>.x = x; // by-ref
//! fn.<var>.y = y; // by-ref
//! ```
//!
//! i.e. captured bindings are turned into fields on the function. A closure is an object, so
//! it can have fields of its own. The captures are kept apart from them in one field, `<var>`,
//! so that a write to a field of the closure cannot clobber a capture
//!
//! # Locations
//!
//! Every binding (a local or an arg) is a location. A capture must share the location of the
//! binding it captures: a write by the owner is seen by the closure, and a write by the
//! closure is seen by the owner. Normally a field holds a copy of a value, so a special
//! `define_field_ref` instr is introduced by `c_parse.py`:
//!
//! define_field_ref x                  [obj, value] -> [obj]
//!
//! It defines `obj.x` as the location that `value` was read from, not as a copy of `value`.
//! So the instr right before it is always the read of a location. This module only emits:
//! - get_loc k         a local of the owner
//! - get_arg k         an arg of the owner
//! - get_field y       a capture of the owner, to give to its own child. The owner's `y` is
//!                     the location of a binding further out, so the child shares that
//!
//! Where a location lives (the stack or the heap) is not decided here. `escape.rs` decides
//! that from these instrs. This module only states which fields share which locations
//!
//! # `close_loc`: a new binding for each loop iteration
//!
//! In JS, each iteration of a loop gets its own copy of the bindings that are declared in
//! the loop:
//! ```js
//! for (let i = 0; i < 3; i++) fns.push(() => i);
//! // the closures give 0, 1, 2, not 3, 3, 3
//! ```
//!
//! QuickJS uses one local for `i`, and emits `close_loc i` at each way out of the loop body.
//! After this module, `close_loc k` means: the binding of local `k` moves to a new location,
//! with the value of the old location. The closures that were made before the `close_loc`
//! keep the old location. So in the example, `i++` writes the new location of `i`, and the
//! closure of each iteration keeps the `i` of its iteration
//!
//! The owner's code does not change for this: `get_loc k` and `put_loc k` always use the
//! current location of `k`. This module keeps `close_loc` as it is. Then:
//! - infer can ignore it. The old and new locations of a binding have the same address
//!   (`Addr::Shared`), so their types are joined
//! - `escape.rs` must give each iteration a new location if a closure can see the old one
//!   after the `close_loc`
//!
//! # The replacements
//!
//! In the owner, for `fclosure f`:
//! fclosure f                          [fn]
//! object                              [fn, vars]
//! <read of the location>              [fn, vars, value]       for each capture `x` of `f`
//! define_field_ref x                  [fn, vars]
//! define_field <var>                  [fn]
//!
//! A closure that captures nothing gets no `<var>`: only the `fclosure` stays
//!
//! In the closure, with `env` as the closure itself (arg `ENV_ARG`):
//! get_var_ref x                       get_arg env; get_field <var>; get_field x
//! put_var_ref x                       get_arg env; get_field <var>; swap; put_field x
//! set_var_ref x                       dup; get_arg env; get_field <var>; swap; put_field x
//!
//! close_var_ref x                     get_arg env; get_field <var>; close_field x
//!
//! The `_check` forms are the same: the TDZ is not checked. `close_var_ref` is `close_loc` for
//! a capture: transform/async.rs makes it, when a binding that it moves out of a fn's frame
//! needs a new location for each loop iteration
//!
//! A module binding is a local of the synthetic entry fn, which makes `<eval>` once (see
//! flatten.rs), so there is one location for each module binding
//!
//! In the emitted C, every underlying "plain" function is called as:
//! `fn(env, this, arg_0, arg_1, ... arg_n)`
//! where `env` is the closure itself

use super::{ENV_ARG, FIRST_ARG};
use crate::libqjs::parse::{
    AtomRef, ClosureKind, ClosureVar, ImmOperand as Imm, Instr, JsValue, opcodes::OpCode,
};

/// `<var>`: the field of a closure that holds its captures
pub const VAR_FIELD: AtomRef = AtomRef::Predefined(86);

/// Replace a read or write of a capture, or the making of a closure
///
/// `own` is the captures of the fn that holds `instr`
pub fn try_replace_closure(
    instr: &Instr,
    own: &[ClosureVar],
    captures: &[Vec<ClosureVar>],
    cpool: &[JsValue],
    out: &mut Vec<Instr>,
) -> bool {
    use OpCode::*;

    let mut push = |op, imm| {
        out.push(Instr {
            pc: instr.pc,
            op,
            imm,
        })
    };

    if let Some((i, writes, leaves)) = var_ref_access(instr) {
        if leaves {
            push(Dup, Imm::None);
        }

        push(GetArg, Imm::Arg(ENV_ARG));
        push(GetField, Imm::Atom(VAR_FIELD));

        if writes {
            push(Swap, Imm::None);
            push(PutField, Imm::Atom(own[i].name));
        } else {
            push(GetField, Imm::Atom(own[i].name));
        }

        return true;
    }

    // A new location for the binding that a capture shares (see `close_loc` above)
    if let (CloseVarRef, Imm::VarRef(i)) = (instr.op, instr.imm) {
        push(GetArg, Imm::Arg(ENV_ARG));
        push(GetField, Imm::Atom(VAR_FIELD));
        push(CloseField, Imm::Atom(own[i as usize].name));
        return true;
    }

    if !matches!(instr.op, Fclosure | Fclosure8) {
        return false;
    }

    let JsValue::FnByIdx(child) = cpool[instr.imm.const_idx()] else {
        unreachable!()
    };

    push(Fclosure, instr.imm);

    // A closure with no captures needs no `<var>`
    if captures[child].is_empty() {
        return true;
    }

    push(Object, Imm::None);

    for c in &captures[child] {
        match c.flags.closure_kind {
            ClosureKind::Local => push(GetLoc, Imm::Loc(c.var_idx as u16)),
            ClosureKind::Arg => push(GetArg, Imm::Arg(FIRST_ARG + c.var_idx as u16)),
            ClosureKind::Ref => {
                push(GetArg, Imm::Arg(ENV_ARG));
                push(GetField, Imm::Atom(VAR_FIELD));
                push(GetField, Imm::Atom(own[c.var_idx].name));
            }
            _ => unimplemented!("a capture of kind {:?}", c.flags.closure_kind),
        }

        push(DefineFieldRef, Imm::Atom(c.name));
    }

    push(DefineField, Imm::Atom(VAR_FIELD));
    true
}

/// If `instr` reads or writes capture `i` of its fn: `i`, whether it writes, and whether a
/// write also leaves the value on the stack
fn var_ref_access(instr: &Instr) -> Option<(usize, bool, bool)> {
    use OpCode::*;

    let i = match instr.imm {
        Imm::VarRef(i) => i as usize,
        _ => 0,
    };

    match instr.op {
        GetVarRef0 | GetVarRef1 | GetVarRef2 | GetVarRef3 => {
            Some((instr.op as usize - GetVarRef0 as usize, false, false))
        }
        PutVarRef0 | PutVarRef1 | PutVarRef2 | PutVarRef3 => {
            Some((instr.op as usize - PutVarRef0 as usize, true, false))
        }
        SetVarRef0 | SetVarRef1 | SetVarRef2 | SetVarRef3 => {
            Some((instr.op as usize - SetVarRef0 as usize, true, true))
        }
        GetVarRef | GetVarRefCheck => Some((i, false, false)),
        PutVarRef | PutVarRefCheck | PutVarRefCheckInit => Some((i, true, false)),
        SetVarRef => Some((i, true, true)),
        _ => None,
    }
}