Fix caching of rust modules.

Author: windelbouwman

vm/src/import.rs

@@ -11,54 +11,47 @@ use crate::pyobject::{ItemProtocol, PyResult};
 use crate::util;
 use crate::vm::VirtualMachine;
 
-fn import_uncached_module(
-    vm: &VirtualMachine,
-    current_path: PathBuf,
-    module_name: &str,
-) -> PyResult {
-    // Check for Rust-native modules
-    if let Some(module) = vm.stdlib_inits.borrow().get(module_name) {
-        return Ok(module(vm).clone());
-    }
+pub fn import_module(vm: &VirtualMachine, current_path: PathBuf, module_name: &str) -> PyResult {
+    // Cached modules:
+    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
 
-    let notfound_error = vm.context().exceptions.module_not_found_error.clone();
-    let import_error = vm.context().exceptions.import_error.clone();
+    // First, see if we already loaded the module:
+    if let Ok(module) = sys_modules.get_item(module_name.to_string(), vm) {
+        Ok(module)
+    } else if let Some(make_module_func) = vm.stdlib_inits.borrow().get(module_name) {
+        let module = make_module_func(vm);
+        sys_modules.set_item(module_name, module.clone(), vm)?;
+        Ok(module)
+    } else {
+        let notfound_error = vm.context().exceptions.module_not_found_error.clone();
+        let import_error = vm.context().exceptions.import_error.clone();
 
-    // Time to search for module in any place:
-    let file_path = find_source(vm, current_path, module_name)
-        .map_err(|e| vm.new_exception(notfound_error.clone(), e))?;
-    let source = util::read_file(file_path.as_path())
-        .map_err(|e| vm.new_exception(import_error.clone(), e.to_string()))?;
-    let code_obj = compile::compile(
-        vm,
-        &source,
-        &compile::Mode::Exec,
-        file_path.to_str().unwrap().to_string(),
-    )
-    .map_err(|err| vm.new_syntax_error(&err))?;
-    // trace!("Code object: {:?}", code_obj);
+        // Time to search for module in any place:
+        let file_path = find_source(vm, current_path, module_name)
+            .map_err(|e| vm.new_exception(notfound_error.clone(), e))?;
+        let source = util::read_file(file_path.as_path())
+            .map_err(|e| vm.new_exception(import_error.clone(), e.to_string()))?;
+        let code_obj = compile::compile(
+            vm,
+            &source,
+            &compile::Mode::Exec,
+            file_path.to_str().unwrap().to_string(),
+        )
+        .map_err(|err| vm.new_syntax_error(&err))?;
+        // trace!("Code object: {:?}", code_obj);
 
-    let attrs = vm.ctx.new_dict();
-    attrs.set_item("__name__", vm.new_str(module_name.to_string()), vm)?;
-    let module = vm.ctx.new_module(module_name, attrs.clone());
+        let attrs = vm.ctx.new_dict();
+        attrs.set_item("__name__", vm.new_str(module_name.to_string()), vm)?;
+        let module = vm.ctx.new_module(module_name, attrs.clone());
 
-    // Store module in cache to prevent infinite loop with mutual importing libs:
-    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
-    sys_modules.set_item(module_name, module.clone(), vm)?;
+        // Store module in cache to prevent infinite loop with mutual importing libs:
+        sys_modules.set_item(module_name, module.clone(), vm)?;
 
-    // Execute main code in module:
-    vm.run_code_obj(code_obj, Scope::new(None, attrs))?;
-    Ok(module)
-}
+        // Execute main code in module:
+        vm.run_code_obj(code_obj, Scope::new(None, attrs))?;
 
-pub fn import_module(vm: &VirtualMachine, current_path: PathBuf, module_name: &str) -> PyResult {
-    // First, see if we already loaded the module:
-    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
-    if let Ok(module) = sys_modules.get_item(module_name.to_string(), vm) {
-        return Ok(module);
+        Ok(module)
     }
-    let module = import_uncached_module(vm, current_path, module_name)?;
-    Ok(module)
 }
 
 fn find_source(vm: &VirtualMachine, current_path: PathBuf, name: &str) -> Result<PathBuf, String> {