Merge pull request #3401 from coolreader18/stdlib-arrayiter

youknowone · web-flow · commit 93ef12e55329 · 2021-11-02T18:12:28.000+09:00
Don't create an intermediate map for rustpython_stdlib::stdlib_inits
diff --git a/src/lib.rs b/src/lib.rs
@@ -304,12 +304,7 @@ fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> {
 fn add_stdlib(vm: &mut VirtualMachine) {
     let _ = vm;
     #[cfg(feature = "stdlib")]
-    {
-        let stdlib = rustpython_stdlib::get_module_inits();
-        for (name, init) in stdlib.into_iter() {
-            vm.add_native_module(name, init);
-        }
-    }
+    vm.add_native_modules(rustpython_stdlib::get_module_inits());
 }
 
 /// Create settings by examining command line arguments and environment
diff --git a/stdlib/src/lib.rs b/stdlib/src/lib.rs
@@ -52,26 +52,23 @@ mod termios;
 use rustpython_common as common;
 use rustpython_vm as vm;
 
-use crate::vm::{
-    builtins,
-    stdlib::{StdlibInitFunc, StdlibMap},
-};
+use crate::vm::{builtins, stdlib::StdlibInitFunc};
 use std::borrow::Cow;
 
-pub fn get_module_inits() -> StdlibMap {
+pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInitFunc)> {
     macro_rules! modules {
         {
             $(
                 #[cfg($cfg:meta)]
                 { $( $key:expr => $val:expr),* $(,)? }
             )*
         } => {{
-            let modules = [
+            [
                 $(
                     $(#[cfg($cfg)] (Cow::<'static, str>::from($key), Box::new($val) as StdlibInitFunc),)*
                 )*
-            ];
-            modules.into_iter().collect()
+            ]
+            .into_iter()
         }};
     }
     modules! {
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -386,14 +386,24 @@ impl VirtualMachine {
         self.initialized = true;
     }
 
+    fn state_mut(&mut self) -> &mut PyGlobalState {
+        PyRc::get_mut(&mut self.state)
+            .expect("there should not be multiple threads while a user has a mut ref to a vm")
+    }
+
     /// Can only be used in the initialization closure passed to [`Interpreter::new_with_init`]
     pub fn add_native_module<S>(&mut self, name: S, module: stdlib::StdlibInitFunc)
     where
         S: Into<Cow<'static, str>>,
     {
-        let state = PyRc::get_mut(&mut self.state)
-            .expect("can't add_native_module when there are multiple threads");
-        state.module_inits.insert(name.into(), module);
+        self.state_mut().module_inits.insert(name.into(), module);
+    }
+
+    pub fn add_native_modules<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = (Cow<'static, str>, stdlib::StdlibInitFunc)>,
+    {
+        self.state_mut().module_inits.extend(iter);
     }
 
     /// Can only be used in the initialization closure passed to [`Interpreter::new_with_init`]
@@ -402,9 +412,7 @@ impl VirtualMachine {
         I: IntoIterator<Item = (String, bytecode::FrozenModule)>,
     {
         let frozen = frozen::map_frozen(self, frozen).collect::<Vec<_>>();
-        let state = PyRc::get_mut(&mut self.state)
-            .expect("can't add_frozen when there are multiple threads");
-        state.frozen.extend(frozen);
+        self.state_mut().frozen.extend(frozen);
     }
 
     /// Start a new thread with access to the same interpreter.
