From 8a14abc445fef7b4c1baf8499afe51886fe26aae Mon Sep 17 00:00:00 2001
From: yanganto <yanganto@gmail.com>
Date: Sat, 30 Nov 2019 16:19:55 +0800
Subject: [PATCH] Initialize hardcore modules and importlib

Initialize hardcore modules and importlib after vm created and
stdlib_inits setup
---
 src/main.rs              | 11 +++++------
 vm/src/eval.rs           |  4 +++-
 vm/src/vm.rs             | 28 ++++++++++++++++++----------
 wasm/lib/src/vm_class.rs |  6 +-----
 4 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 81e44f537..fb5ad9971 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,7 +7,7 @@ extern crate log;
 use clap::{App, AppSettings, Arg, ArgMatches};
 use rustpython_compiler::compile;
 use rustpython_vm::{
-    import, match_class,
+    match_class,
     obj::{objint::PyInt, objtuple::PyTuple, objtype},
     print_exception,
     pyobject::{ItemProtocol, PyResult},
@@ -30,8 +30,9 @@ fn main() {
     let app = App::new("RustPython");
     let matches = parse_arguments(app);
     let settings = create_settings(&matches);
-    let vm = VirtualMachine::new(settings);
+    let mut uninitialized_vm = VirtualMachine::new(settings);
 
+    let vm = uninitialized_vm.initialize(true);
     let res = run_rustpython(&vm, &matches);
 
     #[cfg(feature = "flame-it")]
@@ -324,9 +325,6 @@ fn write_profile(matches: &ArgMatches) -> Result<(), Box<dyn std::error::Error>>
 }
 
 fn run_rustpython(vm: &VirtualMachine, matches: &ArgMatches) -> PyResult<()> {
-    // We only include the standard library bytecode in WASI
-    import::init_importlib(&vm, cfg!(not(target_os = "wasi")))?;
-
     if let Some(paths) = option_env!("BUILDTIME_RUSTPYTHONPATH") {
         let sys_path = vm.get_attribute(vm.sys_module.clone(), "path")?;
         for (i, path) in std::env::split_paths(paths).enumerate() {
@@ -446,7 +444,8 @@ fn run_script(vm: &VirtualMachine, scope: Scope, script_file: &str) -> PyResult<
 
 #[test]
 fn test_run_script() {
-    let vm: VirtualMachine = Default::default();
+    let mut uninitialized_vm: VirtualMachine = Default::default();
+    let vm = uninitialized_vm.initialize(true);
 
     // test file run
     let r = run_script(
diff --git a/vm/src/eval.rs b/vm/src/eval.rs
index 3d998f8b4..000232726 100644
--- a/vm/src/eval.rs
+++ b/vm/src/eval.rs
@@ -22,7 +22,9 @@ mod tests {
     #[test]
     fn test_print_42() {
         let source = String::from("print('Hello world')");
-        let vm = VirtualMachine::default();
+        let mut uninitialized_vm = VirtualMachine::default();
+        let vm = uninitialized_vm.initialize(true);
+
         let vars = vm.new_scope_with_builtins();
         let result = eval(&vm, &source, vars, "<unittest>").expect("this should pass");
         assert!(result.is(&vm.ctx.none()));
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
index 19ad8b48f..cfc6ea0a0 100644
--- a/vm/src/vm.rs
+++ b/vm/src/vm.rs
@@ -147,7 +147,7 @@ impl Default for PySettings {
 impl VirtualMachine {
     /// Create a new `VirtualMachine` structure.
     pub fn new(settings: PySettings) -> VirtualMachine {
-        flame_guard!("init VirtualMachine");
+        flame_guard!("new VirtualMachine");
         let ctx = PyContext::new();
 
         // make a new module without access to the vm; doesn't
@@ -200,19 +200,27 @@ impl VirtualMachine {
             vm.get_none(),
         );
 
-        builtins::make_module(&vm, builtins.clone());
-        sysmodule::make_module(&vm, sysmod, builtins);
-
-        #[cfg(not(target_arch = "wasm32"))]
-        import::import_builtin(&vm, "signal").expect("Couldn't initialize signal module");
-
         vm
     }
 
-    /// `update_sysmodule` should be called when the `vm.stdlib_inits` are changed,
-    /// and then the `builtin_module_names` of `sys` module will be updated.
-    pub fn update_sysmodule(&self) {
+    pub fn initialize(&mut self, external: bool) -> &Self {
+        flame_guard!("init VirtualMachine");
+
+        builtins::make_module(self, self.builtins.clone());
         sysmodule::make_module(self, self.sys_module.clone(), self.builtins.clone());
+
+        #[cfg(not(target_arch = "wasm32"))]
+        import::import_builtin(self, "signal").expect("Couldn't initialize signal module");
+
+        if external {
+            // We only include the standard library bytecode in WASI
+            import::init_importlib(self, cfg!(not(target_os = "wasi")))
+                .expect("Couldn't initialize importlib");
+        } else {
+            import::init_importlib(self, false).expect("Couldn't initialize importlib");
+        }
+
+        self
     }
 
     pub fn run_code_obj(&self, code: PyCodeRef, scope: Scope) -> PyResult {
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
index a1d95d4cb..141862afc 100644
--- a/wasm/lib/src/vm_class.rs
+++ b/wasm/lib/src/vm_class.rs
@@ -7,7 +7,6 @@ use wasm_bindgen::prelude::*;
 
 use rustpython_compiler::compile;
 use rustpython_vm::function::PyFuncArgs;
-use rustpython_vm::import;
 use rustpython_vm::pyobject::{PyObject, PyObjectPayload, PyObjectRef, PyResult, PyValue};
 use rustpython_vm::scope::{NameProtocol, Scope};
 use rustpython_vm::VirtualMachine;
@@ -44,10 +43,7 @@ impl StoredVirtualMachine {
             setup_browser_module(&vm);
         }
 
-        // sys module needs to be updated because js_module and browser_module is setted
-        vm.update_sysmodule();
-
-        import::init_importlib(&vm, false).unwrap();
+        vm.initialize(false);
 
         StoredVirtualMachine {
             vm,