RustPython/crates/wasm/src/vm_class.rs

use crate::{
    browser_module,
    convert::{self, PyResultExt},
    js_module, wasm_builtins,
};
use alloc::rc::{Rc, Weak};
use core::cell::RefCell;
use js_sys::{Object, TypeError};
use rustpython_vm::{
    Interpreter, PyObjectRef, PyRef, PyResult, Settings, VirtualMachine, builtins::PyWeak,
    compiler::Mode, function::ArgMapping, scope::Scope,
};
use std::collections::HashMap;
use wasm_bindgen::prelude::*;

pub(crate) struct StoredVirtualMachine {
    pub interp: Interpreter,
    pub scope: Scope,
    /// you can put a Rc in here, keep it as a Weak, and it'll be held only for
    /// as long as the StoredVM is alive
    held_objects: RefCell<Vec<PyObjectRef>>,
}

#[pymodule]
mod _window {
    use super::{js_module, wasm_builtins};
    use rustpython_vm::{Py, PyPayload, PyResult, VirtualMachine, builtins::PyModule};

    #[expect(clippy::unnecessary_wraps, reason = "Needs to comply with a signature")]
    pub(crate) fn module_exec(vm: &VirtualMachine, module: &Py<PyModule>) -> PyResult<()> {
        __module_exec(vm, module);
        extend_module!(vm, module, {
            "window" => js_module::PyJsValue::new(wasm_builtins::window()).into_ref(&vm.ctx),
        });
        Ok(())
    }
}

impl StoredVirtualMachine {
    fn new(id: String, inject_browser_module: bool) -> Self {
        let mut settings = Settings::default();
        settings.allow_external_library = false;

        let mut builder = Interpreter::builder(settings);

        #[cfg(feature = "freeze-stdlib")]
        {
            let defs = rustpython_stdlib::stdlib_module_defs(&builder.ctx);
            builder = builder
                .add_native_modules(&defs)
                .add_frozen_modules(rustpython_pylib::FROZEN_STDLIB);
        }

        // Add wasm-specific modules
        let js_def = js_module::module_def(&builder.ctx);
        builder = builder.add_native_module(js_def);

        if inject_browser_module {
            let window_def = _window::module_def(&builder.ctx);
            let browser_def = browser_module::module_def(&builder.ctx);
            builder = builder
                .add_native_modules(&[window_def, browser_def])
                .add_frozen_modules(rustpython_vm::py_freeze!(dir = "../Lib"));
        }

        let interp = builder
            .init_hook(move |vm| {
                vm.wasm_id = Some(id);
            })
            .build();

        let scope = interp.enter(|vm| vm.new_scope_with_builtins());

        Self {
            interp,
            scope,
            held_objects: RefCell::new(Vec::new()),
        }
    }
}

// It's fine that it's thread local, since WASM doesn't even have threads yet. thread_local!
// probably gets compiled down to a normal-ish static variable, like Atomic* types do:
// https://rustwasm.github.io/2018/10/24/multithreading-rust-and-wasm.html#atomic-instructions
thread_local! {
    static STORED_VMS: RefCell<HashMap<String, Rc<StoredVirtualMachine>>> = RefCell::default();
}

pub fn get_vm_id(vm: &VirtualMachine) -> &str {
    vm.wasm_id
        .as_ref()
        .expect("VirtualMachine inside of WASM crate should have wasm_id set")
}
pub(crate) fn stored_vm_from_wasm(wasm_vm: &WASMVirtualMachine) -> Rc<StoredVirtualMachine> {
    STORED_VMS.with_borrow(|vms| {
        vms.get(&wasm_vm.id)
            .expect("VirtualMachine is not valid")
            .clone()
    })
}
pub(crate) fn weak_vm(vm: &VirtualMachine) -> Weak<StoredVirtualMachine> {
    let id = get_vm_id(vm);
    STORED_VMS.with_borrow(|vms| Rc::downgrade(vms.get(id).expect("VirtualMachine is not valid")))
}

#[derive(Clone, Copy)]
#[wasm_bindgen(js_name = vmStore)]
pub struct VMStore;

#[wasm_bindgen(js_class = vmStore)]
impl VMStore {
    #[must_use]
    pub fn init(id: String, inject_browser_module: Option<bool>) -> WASMVirtualMachine {
        STORED_VMS.with_borrow_mut(|vms| {
            if !vms.contains_key(&id) {
                let stored_vm =
                    StoredVirtualMachine::new(id.clone(), inject_browser_module.unwrap_or(true));
                vms.insert(id.clone(), Rc::new(stored_vm));
            }
        });
        WASMVirtualMachine { id }
    }

    pub(crate) fn _get(id: String) -> Option<WASMVirtualMachine> {
        STORED_VMS.with_borrow(|vms| vms.contains_key(&id).then_some(WASMVirtualMachine { id }))
    }

    #[must_use]
    pub fn get(id: String) -> JsValue {
        match Self::_get(id) {
            Some(wasm_vm) => wasm_vm.into(),
            None => JsValue::UNDEFINED,
        }
    }

    pub fn destroy(id: String) {
        STORED_VMS.with_borrow_mut(|vms| {
            if let Some(stored_vm) = vms.remove(&id) {
                // for f in stored_vm.drop_handlers.iter() {
                //     f();
                // }
                // deallocate the VM
                drop(stored_vm);
            }
        });
    }

    #[must_use]
    pub fn ids() -> Vec<JsValue> {
        STORED_VMS.with_borrow(|vms| vms.keys().map(|k| k.into()).collect())
    }
}

#[wasm_bindgen(js_name = VirtualMachine)]
#[derive(Clone)]
pub struct WASMVirtualMachine {
    pub(crate) id: String,
}

#[wasm_bindgen(js_class = VirtualMachine)]
impl WASMVirtualMachine {
    pub(crate) fn with_unchecked<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&StoredVirtualMachine) -> R,
    {
        let stored_vm = STORED_VMS.with_borrow_mut(|vms| vms.get_mut(&self.id).unwrap().clone());
        f(&stored_vm)
    }

    pub(crate) fn with<F, R>(&self, f: F) -> Result<R, JsValue>
    where
        F: FnOnce(&StoredVirtualMachine) -> R,
    {
        self.assert_valid()?;
        Ok(self.with_unchecked(f))
    }

    pub(crate) fn with_vm<F, R>(&self, f: F) -> Result<R, JsValue>
    where
        F: FnOnce(&VirtualMachine, &StoredVirtualMachine) -> R,
    {
        self.with(|stored| stored.interp.enter(|vm| f(vm, stored)))
    }

    #[must_use]
    pub fn valid(&self) -> bool {
        STORED_VMS.with_borrow(|vms| vms.contains_key(&self.id))
    }

    pub(crate) fn push_held_rc(
        &self,
        obj: PyObjectRef,
    ) -> Result<PyResult<PyRef<PyWeak>>, JsValue> {
        self.with_vm(|vm, stored_vm| {
            let weak = obj.downgrade(None, vm)?;
            stored_vm.held_objects.borrow_mut().push(obj);
            Ok(weak)
        })
    }

    pub fn assert_valid(&self) -> Result<(), JsValue> {
        if self.valid() {
            Ok(())
        } else {
            Err(TypeError::new(
                "Invalid VirtualMachine, this VM was destroyed while this reference was still held",
            )
            .into())
        }
    }

    pub fn destroy(&self) -> Result<(), JsValue> {
        self.assert_valid()?;
        VMStore::destroy(self.id.clone());
        Ok(())
    }

    #[wasm_bindgen(js_name = addToScope)]
    pub fn add_to_scope(&self, name: String, value: JsValue) -> Result<(), JsValue> {
        self.with_vm(move |vm, StoredVirtualMachine { scope, .. }| {
            let value = convert::js_to_py(vm, value);
            scope.globals.set_item(&name, value, vm).into_js(vm)
        })?
    }

    #[wasm_bindgen(js_name = setStdout)]
    pub fn set_stdout(&self, stdout: JsValue) -> Result<(), JsValue> {
        self.with_vm(|vm, _| {
            fn error() -> JsValue {
                TypeError::new("Unknown stdout option, please pass a function or 'console'").into()
            }
            use wasm_builtins::make_stdout_object;
            let stdout: PyObjectRef = if let Some(s) = stdout.as_string() {
                match s.as_str() {
                    "console" => make_stdout_object(vm, wasm_builtins::sys_stdout_write_console),
                    _ => return Err(error()),
                }
            } else if stdout.is_function() {
                let func = js_sys::Function::from(stdout);
                make_stdout_object(vm, move |data, vm| {
                    func.call1(&JsValue::UNDEFINED, &data.into())
                        .map_err(|err| convert::js_py_typeerror(vm, err))?;
                    Ok(())
                })
            } else if stdout.is_null() {
                make_stdout_object(vm, |_, _| Ok(()))
            } else if stdout.is_undefined() {
                make_stdout_object(vm, wasm_builtins::sys_stdout_write_console)
            } else {
                return Err(error());
            };
            vm.sys_module.set_attr("stdout", stdout, vm).unwrap();
            Ok(())
        })?
    }

    #[wasm_bindgen(js_name = injectModule)]
    pub fn inject_module(
        &self,
        name: String,
        source: &str,
        imports: Option<Object>,
    ) -> Result<(), JsValue> {
        self.with_vm(|vm, _| {
            let code = vm
                .compile(source, Mode::Exec, name.clone())
                .map_err(convert::syntax_err)?;
            let attrs = vm.ctx.new_dict();
            attrs
                .set_item("__name__", vm.new_pyobj(name.as_str()), vm)
                .into_js(vm)?;

            if let Some(imports) = imports {
                for entry in convert::object_entries(&imports) {
                    let (key, value) = entry?;
                    let key: String = Object::from(key).to_string().into();
                    attrs
                        .set_item(key.as_str(), convert::js_to_py(vm, value), vm)
                        .into_js(vm)?;
                }
            }

            vm.run_code_obj(
                code,
                Scope::new(
                    Some(ArgMapping::from_dict_exact(attrs.clone())),
                    attrs.clone(),
                ),
            )
            .into_js(vm)?;

            let module = vm.new_module(&name, attrs, None);

            let sys_modules = vm.sys_module.get_attr("modules", vm).into_js(vm)?;
            sys_modules.set_item(&name, module.into(), vm).into_js(vm)?;

            Ok(())
        })?
    }

    #[wasm_bindgen(js_name = injectJSModule)]
    pub fn inject_js_module(&self, name: String, module: Object) -> Result<(), JsValue> {
        self.with_vm(|vm, _| {
            let py_module = vm.new_module(&name, vm.ctx.new_dict(), None);
            for entry in convert::object_entries(&module) {
                let (key, value) = entry?;
                let key = Object::from(key).to_string();
                extend_module!(vm, &py_module, {
                    String::from(key) => convert::js_to_py(vm, value),
                });
            }

            let sys_modules = vm.sys_module.get_attr("modules", vm).into_js(vm)?;
            sys_modules
                .set_item(&name, py_module.into(), vm)
                .into_js(vm)?;

            Ok(())
        })?
    }

    pub(crate) fn run(
        &self,
        source: &str,
        mode: Mode,
        source_path: Option<String>,
    ) -> Result<JsValue, JsValue> {
        self.with_vm(|vm, StoredVirtualMachine { scope, .. }| {
            let source_path = source_path.unwrap_or_else(|| "<wasm>".to_owned());
            let code = vm.compile(source, mode, source_path);
            let code = code.map_err(convert::syntax_err)?;
            let result = vm.run_code_obj(code, scope.clone());
            convert::pyresult_to_js_result(vm, result)
        })?
    }

    pub fn exec(&self, source: &str, source_path: Option<String>) -> Result<JsValue, JsValue> {
        self.run(source, Mode::Exec, source_path)
    }

    pub fn eval(&self, source: &str, source_path: Option<String>) -> Result<JsValue, JsValue> {
        self.run(source, Mode::Eval, source_path)
    }

    #[wasm_bindgen(js_name = execSingle)]
    pub fn exec_single(
        &self,
        source: &str,
        source_path: Option<String>,
    ) -> Result<JsValue, JsValue> {
        self.run(source, Mode::Single, source_path)
    }
}