jsrt/compiler/src/link.rs

//! Link: what each name in each module refers to, and the program's entry fn
//!
//! Every binding at the top of a module (`let`, `const`, `function`, `class`, ...) is a
//! global: there is one of each for the whole program. A module's body reaches its own
//! bindings and its imports as captured vars, which QuickJS calls module decls and module
//! imports. Link gives each decl a global, and follows each import to the global of the
//! binding that it names (see `Link::resolve`). transform/modules.rs then rewrites the code
//!
//! A name that no module declares, such as `String` or `console`, is a global name: it
//! refers to the export of that name of the prelude, `jsrt:prelude`, as names in JS refer
//! to the global object. A module can declare a binding of the same name, which hides it
//!
//! The natives are the exports of `jsrt:native`, which QuickJS sees as a C module (see
//! `libqjs::compile`)
//!
//! # The entry fn
//!
//! ```js
//! function <entry>() {
//!     // Phase 1 (`this` is true): each body only makes its top-level function declarations.
//!     // Phase 2: each body runs. Every module's phase 1 comes first, so that modules that
//!     // import each other see each other's functions
//!     let phase = true;
//!     while (true) {
//!         for (const body of bodies) body.call(phase);
//!         if (!phase) break;
//!         phase = false;
//!     }
//!     __drain_microtasks();
//! }
//! ```
//!
//! Both phases call each body from one call instr. Infer names the objects that a walk makes
//! by its call site, so with two call instrs, each phase would make its own copy of every
//! object at the top of a module
//!
//! The bodies run in ES's order: each module after the modules that it imports, from the
//! prelude, then each module that the command line names. A body that awaits at its top runs
//! until its first await, so only the last body may await (see transform/async.rs)

use std::collections::HashMap;

use crate::{
    flatten::{ENTRY_FN, FlatFn, FlatModule},
    libqjs::parse::{
        Atom, AtomRef, ClosureKind, ClosureVar, ExportInfo, FuncKind, ImmOperand as Imm, Instr,
        JsValue, opcodes::OpCode,
    },
    native::NATIVES,
};

/// The module whose exports are every module's global names
pub const PRELUDE: &str = "jsrt:prelude";

/// The module of the natives
pub const NATIVE: &str = "jsrt:native";

/// What a name refers to
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Target {
    Global(usize), // a binding at the top of a module
    Native(usize), // a native, by its index in `native::NATIVES`
}

pub struct Link {
    pub globals: usize,                                // how many there are
    bodies: HashMap<usize, (String, Vec<Target>)>, // for the fn of each module's body: the module's name, and what each of the body's captured vars refers to
    exports: HashMap<String, HashMap<String, Target>>, // for each module, by name: each export, by name
}

impl Link {
    /// The name of the module whose body is `f`, if it is a body, and what each captured var
    /// of the body refers to
    pub fn body(&self, f: usize) -> Option<(&str, &[Target])> {
        self.bodies
            .get(&f)
            .map(|(name, vars)| (&name[..], &vars[..]))
    }

    /// What the global name `name` refers to: the prelude's export of that name
    pub fn global(&self, name: &str) -> Option<Target> {
        self.export(PRELUDE, name)
    }

    /// The export `name` of the module `module`
    pub fn export(&self, module: &str, name: &str) -> Option<Target> {
        self.exports.get(module)?.get(name).copied()
    }

    /// The global of an export that the compiler needs from the stdlib
    pub fn stdlib(&self, module: &str, name: &str) -> usize {
        match self.export(module, name) {
            Some(Target::Global(global)) => global,
            _ => panic!("the compiler needs `{name}` from `{module}`"),
        }
    }
}

/// Link the modules, and write the entry fn. `roots` are the modules that the command line
/// names, in order
pub fn link(
    modules: &[FlatModule],
    roots: &[String],
    fns: &mut [FlatFn],
    captures: &[Vec<ClosureVar>],
    atoms: &[Atom],
) -> Link {
    let linker = Linker::new(modules, captures, atoms);

    let link = Link {
        globals: linker.globals,
        bodies: linker.bodies(),
        exports: linker.all_exports(),
    };

    let order = linker.order(roots);
    check_top_level_await(&order, modules);

    let drain = link.stdlib("jsrt:loop", "__drain_microtasks");
    fns[ENTRY_FN] = entry(&order, modules, drain);
    link
}

/// A body that awaits returns at its first await, so the bodies after it would run too early.
/// Only the last one may await, until the entry waits for each body's promise
fn check_top_level_await(order: &[usize], modules: &[FlatModule]) {
    for &m in &order[..order.len() - 1] {
        if modules[m].has_top_level_await {
            let name = &modules[m].name;
            unimplemented!("{name}: top-level `await` in a module that is not the last to run");
        }
    }
}

/// The natives, as the exports of `jsrt:native`
fn natives() -> HashMap<String, Target> {
    let natives = NATIVES.iter().enumerate();
    natives
        .map(|(idx, native)| (native.name.to_string(), Target::Native(idx)))
        .collect()
}

struct Linker<'a> {
    modules: &'a [FlatModule],
    captures: &'a [Vec<ClosureVar>],
    atoms: &'a [Atom],
    index: HashMap<String, usize>,     // each module, by name
    decls: Vec<HashMap<usize, usize>>, // for each module: the global of each of its body's decls, by captured var
    globals: usize,
}

impl<'a> Linker<'a> {
    /// The linker, with a global for each decl of each module's body
    fn new(modules: &'a [FlatModule], captures: &'a [Vec<ClosureVar>], atoms: &'a [Atom]) -> Self {
        let index = modules.iter().enumerate();
        let mut linker = Linker {
            modules,
            captures,
            atoms,
            index: index.map(|(m, module)| (module.name.clone(), m)).collect(),
            decls: vec![],
            globals: 0,
        };

        linker.add_globals();
        linker
    }

    /// For each module's body: the module's name, and what each captured var refers to
    fn bodies(&self) -> HashMap<usize, (String, Vec<Target>)> {
        let bodies = self.modules.iter().enumerate();
        bodies
            .map(|(m, module)| (module.body, (module.name.clone(), self.body_vars(m))))
            .collect()
    }

    /// Every module's exports, by the module's name, with the natives' module
    fn all_exports(&self) -> HashMap<String, HashMap<String, Target>> {
        let modules = self.modules.iter().enumerate();
        let exports =
            modules.map(|(m, module)| (module.name.clone(), self.exports(m, &mut vec![])));
        exports.chain([(NATIVE.to_string(), natives())]).collect()
    }
}

impl Linker<'_> {
    /// A global for each decl of each module's body. QuickJS makes a namespace import a decl
    /// too, whose object its own linker makes at run time, so that is rejected here
    fn add_globals(&mut self) {
        for module in self.modules {
            let atoms = self.atoms;
            let namespace = module
                .imports
                .iter()
                .find(|import| import.name.text(atoms) == "*");

            if namespace.is_some() {
                unimplemented!("{}: `import * as` needs namespace objects", module.name);
            }

            let vars = self.captures[module.body].iter().enumerate();
            let decls = vars.filter(|(_, var)| var.flags.closure_kind == ClosureKind::ModuleDecl);
            let mut globals = HashMap::new();

            for (i, _) in decls {
                globals.insert(i, self.globals);
                self.globals += 1;
            }

            self.decls.push(globals);
        }
    }

    /// What each captured var of module `m`'s body refers to
    fn body_vars(&self, m: usize) -> Vec<Target> {
        (0..self.captures[self.modules[m].body].len())
            .map(|i| self.var(m, i, &mut vec![]))
            .collect()
    }

    /// What captured var `i` of module `m`'s body refers to: its global if it is a decl, else
    /// the export that it imports. `seen` guards against imports that go round in a circle
    fn var(&self, m: usize, i: usize, seen: &mut Vec<(usize, usize)>) -> Target {
        let module = &self.modules[m];

        if let Some(&global) = self.decls[m].get(&i) {
            return Target::Global(global);
        }

        assert!(
            !seen.contains(&(m, i)),
            "{}: an import refers back to itself",
            module.name
        );
        seen.push((m, i));

        let import = module.imports.iter().find(|import| import.var_idx == i);
        let import =
            import.unwrap_or_else(|| panic!("{}: a module var with no import", module.name));
        let name = import.name.text(self.atoms);
        let from = self.request(m, import.request);

        match self.resolve(&from, &name, seen) {
            Some(target) => target,
            None => panic!("{}: `{from}` has no export `{name}`", module.name),
        }
    }

    /// ES's ResolveExport: what export `name` of module `module` refers to
    fn resolve(&self, module: &str, name: &str, seen: &mut Vec<(usize, usize)>) -> Option<Target> {
        if module == NATIVE {
            return natives().get(name).copied();
        }

        let Some(&m) = self.index.get(module) else {
            panic!("no module `{module}`");
        };

        for export in &self.modules[m].exports {
            match export {
                ExportInfo::Local { var_idx, name: n } if n.text(self.atoms) == name => {
                    return Some(self.var(m, *var_idx, seen));
                }
                ExportInfo::Indirect {
                    request,
                    local,
                    name: n,
                } if n.text(self.atoms) == name => {
                    return self.resolve(&self.request(m, *request), &local.text(self.atoms), seen);
                }
                _ => {}
            }
        }

        // `export *` passes on every export but `default`
        if name == "default" {
            return None;
        }

        let mut stars = self.modules[m].star_exports.iter();
        stars.find_map(|&request| self.resolve(&self.request(m, request), name, seen))
    }

    /// Every export of module `m`, by name. `seen` guards against `export *` in a circle
    fn exports(&self, m: usize, seen: &mut Vec<usize>) -> HashMap<String, Target> {
        let module = &self.modules[m];
        let mut exports = HashMap::new();

        if seen.contains(&m) {
            return exports;
        }

        seen.push(m);

        for &request in &module.star_exports {
            let from = self.request(m, request);
            let passed = match self.index.get(&from) {
                Some(&other) => self.exports(other, seen),
                None if from == NATIVE => natives(),
                None => panic!("no module `{from}`"),
            };

            for (name, target) in passed {
                if name != "default" {
                    exports.insert(name, target);
                }
            }
        }

        for export in &module.exports {
            let (ExportInfo::Local { name, .. } | ExportInfo::Indirect { name, .. }) = export;
            let name = name.text(self.atoms);
            let target = self.resolve(&module.name, &name, &mut vec![]);
            exports.insert(name, target.unwrap());
        }

        exports
    }

    /// The name of request `r` of module `m`. QuickJS made the same name when it loaded it
    /// (see `normalize`)
    fn request(&self, m: usize, r: usize) -> String {
        let module = &self.modules[m];
        normalize(&module.name, &module.requests[r])
    }

    /// The order that the bodies run in: each module after the modules that it imports, from
    /// the prelude, then each root
    fn order(&self, roots: &[String]) -> Vec<usize> {
        let mut order = vec![];

        for root in [PRELUDE.to_string()].iter().chain(roots) {
            self.visit(self.index[root], &mut order);
        }

        order
    }

    fn visit(&self, m: usize, order: &mut Vec<usize>) {
        if order.contains(&m) {
            return;
        }

        order.push(m);
        let at = order.len() - 1;

        for r in 0..self.modules[m].requests.len() {
            if let Some(&other) = self.index.get(&self.request(m, r)) {
                self.visit(other, order);
            }
        }

        // After its imports. It was marked first, so that a circle of imports ends
        let m = order.remove(at);
        order.push(m);
    }
}

/// The name of the module that `specifier` names, in the module `base`. This is QuickJS's
/// default rule, which QuickJS used to load it: a name that does not start with `.` is as it
/// is. Else it is relative to the folder of `base`, and its leading `./` and `../` are taken
/// out
fn normalize(base: &str, specifier: &str) -> String {
    if !specifier.starts_with('.') {
        return specifier.to_string();
    }

    let mut folder = match base.rfind('/') {
        Some(slash) => base[..slash].to_string(),
        None => String::new(),
    };
    let mut rest = specifier;

    loop {
        if let Some(after) = rest.strip_prefix("./") {
            rest = after;
        } else if let Some(after) = rest.strip_prefix("../") {
            let last = folder.rfind('/').map_or(0, |slash| slash + 1);

            if folder.is_empty() || matches!(&folder[last..], "." | "..") {
                break;
            }

            folder.truncate(last.saturating_sub(1));
            rest = after;
        } else {
            break;
        }
    }

    match folder.is_empty() {
        true => rest.to_string(),
        false => format!("{folder}/{rest}"),
    }
}

/// The entry fn (see the module doc). Local `k` holds the closure of the `k`th body to run,
/// and the last local is the phase
fn entry(order: &[usize], modules: &[FlatModule], drain: usize) -> FlatFn {
    use OpCode::*;

    let n = order.len();
    let phase = Imm::Loc(n as u16);
    let body = |k: usize| Imm::Loc(k as u16);

    // The loop starts after the closures and the phase, and ends after a call of each body
    // and the phase's test
    let head = 2 * n + 2;
    let end = head + 4 * n + 5;

    let make = (0..n).flat_map(|k| [(Fclosure, Imm::Const(k as u32)), (PutLoc, body(k))]);
    let call = (0..n).flat_map(|k| {
        [
            (GetLoc, phase),
            (GetLoc, body(k)),
            (CallMethod, Imm::Npop(0)),
            (Drop, Imm::None),
        ]
    });

    let code: Vec<(OpCode, Imm)> = make
        .chain([(PushTrue, Imm::None), (PutLoc, phase)])
        .chain(call)
        .chain([
            (GetLoc, phase),
            (IfFalse, Imm::Label { target: end }),
            (PushFalse, Imm::None),
            (PutLoc, phase),
            (Goto, Imm::Label { target: head }),
            (GetVar, Imm::Global(drain as u32)),
            (Call0, Imm::None),
            (Drop, Imm::None),
            (ReturnUndef, Imm::None),
        ])
        .collect();

    let code = code.into_iter().enumerate();
    let code = code.map(|(pc, (op, imm))| Instr { pc, op, imm }).collect();
    let bodies = order
        .iter()
        .map(|&m| JsValue::FnByIdx(modules[m].body))
        .collect();

    FlatFn {
        name: AtomRef::Predefined(0),
        kind: FuncKind::Normal,
        arg_count: 0,
        local_count: n + 1,
        cpool: bodies,
        code,
    }
}