Merge pull request #1159 from palaviv/freeze-stdlib

palaviv · web-flow · commit 6ca979e7be01 · 2019-07-24T20:41:35.000+03:00
Allow stdlib freeze
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -17,6 +17,7 @@ path = "./benchmarks/bench.rs"
 [features]
 default = ["rustpython-vm/use-proc-macro-hack"]
 flame-it = ["rustpython-vm/flame-it", "flame", "flamescope"]
+freeze-stdlib = ["rustpython-vm/freeze-stdlib"]
 
 [dependencies]
 log="0.4.1"
diff --git a/derive/Cargo.toml b/derive/Cargo.toml
@@ -21,3 +21,4 @@ rustpython-compiler = { path = "../compiler", version = "0.1.0" }
 rustpython-bytecode = { path = "../bytecode", version = "0.1.0" }
 bincode =  "1.1"
 proc-macro-hack = { version = "0.5", optional = true }
+maplit = "1.0"
diff --git a/derive/src/compile_bytecode.rs b/derive/src/compile_bytecode.rs
@@ -19,15 +19,17 @@ use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::quote;
 use rustpython_bytecode::bytecode::CodeObject;
 use rustpython_compiler::compile;
+use std::collections::HashMap;
 use std::env;
 use std::fs;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use syn::parse::{Parse, ParseStream, Result as ParseResult};
-use syn::{self, parse2, Lit, LitByteStr, Meta, Token};
+use syn::{self, parse2, Lit, LitByteStr, LitStr, Meta, Token};
 
 enum CompilationSourceKind {
     File(PathBuf),
     SourceCode(String),
+    Dir(PathBuf),
 }
 
 struct CompilationSource {
@@ -36,14 +38,22 @@ struct CompilationSource {
 }
 
 impl CompilationSource {
-    fn compile(self, mode: &compile::Mode, module_name: String) -> Result<CodeObject, Diagnostic> {
-        let compile = |source| {
-            compile::compile(source, mode, module_name, 0).map_err(|err| {
-                Diagnostic::spans_error(self.span, format!("Compile error: {}", err))
-            })
-        };
-
-        match &self.kind {
+    fn compile_string(
+        &self,
+        source: &str,
+        mode: &compile::Mode,
+        module_name: String,
+    ) -> Result<CodeObject, Diagnostic> {
+        compile::compile(source, mode, module_name, 0)
+            .map_err(|err| Diagnostic::spans_error(self.span, format!("Compile error: {}", err)))
+    }
+
+    fn compile(
+        &self,
+        mode: &compile::Mode,
+        module_name: String,
+    ) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+        Ok(match &self.kind {
             CompilationSourceKind::File(rel_path) => {
                 let mut path = PathBuf::from(
                     env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR is not present"),
@@ -55,10 +65,59 @@ impl CompilationSource {
                         format!("Error reading file {:?}: {}", path, err),
                     )
                 })?;
-                compile(&source)
+                hashmap! {module_name.clone() => self.compile_string(&source, mode, module_name.clone())?}
+            }
+            CompilationSourceKind::SourceCode(code) => {
+                hashmap! {module_name.clone() => self.compile_string(code, mode, module_name.clone())?}
+            }
+            CompilationSourceKind::Dir(rel_path) => {
+                let mut path = PathBuf::from(
+                    env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR is not present"),
+                );
+                path.push(rel_path);
+                self.compile_dir(&path, String::new(), mode)?
+            }
+        })
+    }
+
+    fn compile_dir(
+        &self,
+        path: &Path,
+        parent: String,
+        mode: &compile::Mode,
+    ) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+        let mut code_map = HashMap::new();
+        let paths = fs::read_dir(&path).map_err(|err| {
+            Diagnostic::spans_error(self.span, format!("Error listing dir {:?}: {}", path, err))
+        })?;
+        for path in paths {
+            let path = path.map_err(|err| {
+                Diagnostic::spans_error(self.span, format!("Failed to list file: {}", err))
+            })?;
+            let path = path.path();
+            let file_name = path.file_name().unwrap().to_str().unwrap();
+            if path.is_dir() {
+                code_map.extend(self.compile_dir(
+                    &path,
+                    format!("{}{}.", parent, file_name),
+                    mode,
+                )?);
+            } else if file_name.ends_with(".py") {
+                let source = fs::read_to_string(&path).map_err(|err| {
+                    Diagnostic::spans_error(
+                        self.span,
+                        format!("Error reading file {:?}: {}", path, err),
+                    )
+                })?;
+                let file_name_splitte: Vec<&str> = file_name.splitn(2, '.').collect();
+                let module_name = format!("{}{}", parent, file_name_splitte[0]);
+                code_map.insert(
+                    module_name.clone(),
+                    self.compile_string(&source, mode, module_name)?,
+                );
             }
-            CompilationSourceKind::SourceCode(code) => compile(code),
         }
+        Ok(code_map)
     }
 }
 
@@ -69,7 +128,7 @@ struct PyCompileInput {
 }
 
 impl PyCompileInput {
-    fn compile(&self) -> Result<CodeObject, Diagnostic> {
+    fn compile(&self) -> Result<HashMap<String, CodeObject>, Diagnostic> {
         let mut module_name = None;
         let mut mode = None;
         let mut source: Option<CompilationSource> = None;
@@ -122,6 +181,16 @@ impl PyCompileInput {
                         kind: CompilationSourceKind::File(path),
                         span: extract_spans(&name_value).unwrap(),
                     });
+                } else if name_value.ident == "dir" {
+                    assert_source_empty(&source)?;
+                    let path = match &name_value.lit {
+                        Lit::Str(s) => PathBuf::from(s.value()),
+                        _ => bail_span!(name_value.lit, "source must be a string"),
+                    };
+                    source = Some(CompilationSource {
+                        kind: CompilationSourceKind::Dir(path),
+                        span: extract_spans(&name_value).unwrap(),
+                    });
                 }
             }
         }
@@ -154,16 +223,23 @@ impl Parse for PyCompileInput {
 pub fn impl_py_compile_bytecode(input: TokenStream2) -> Result<TokenStream2, Diagnostic> {
     let input: PyCompileInput = parse2(input)?;
 
-    let code_obj = input.compile()?;
+    let code_map = input.compile()?;
 
-    let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
-    let bytes = LitByteStr::new(&bytes, Span::call_site());
+    let modules = code_map.iter().map(|(module_name, code_obj)| {
+        let module_name = LitStr::new(&module_name, Span::call_site());
+        let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
+        let bytes = LitByteStr::new(&bytes, Span::call_site());
+        quote! {  #module_name.into() => bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes)
+                .expect("Deserializing CodeObject failed") }
+    });
 
     let output = quote! {
         ({
             use ::rustpython_vm::__exports::bincode;
-            bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes)
-                .expect("Deserializing CodeObject failed")
+            use ::rustpython_vm::__exports::hashmap;
+            hashmap! {
+                #(#modules),*
+            }
         })
     };
 
diff --git a/derive/src/lib.rs b/derive/src/lib.rs
@@ -4,6 +4,9 @@
 
 extern crate proc_macro;
 
+#[macro_use]
+extern crate maplit;
+
 #[macro_use]
 mod error;
 mod compile_bytecode;
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -13,6 +13,7 @@ default = ["rustpython-parser", "rustpython-compiler", "use-proc-macro-hack"]
 vm-tracing-logging = []
 flame-it = ["flame", "flamer"]
 use-proc-macro-hack = ["proc-macro-hack", "rustpython-derive/proc-macro-hack"]
+freeze-stdlib = []
 
 [dependencies]
 # Crypto:
diff --git a/vm/src/frozen.rs b/vm/src/frozen.rs
@@ -2,18 +2,24 @@ use crate::bytecode::CodeObject;
 use std::collections::HashMap;
 
 pub fn get_module_inits() -> HashMap<String, CodeObject> {
-    hashmap! {
-        "__hello__".into() => py_compile_bytecode!(
-            source = "initialized = True; print(\"Hello world!\")\n",
-            module_name = "__hello__",
-        ),
-        "_frozen_importlib".into() => py_compile_bytecode!(
-            file = "Lib/_bootstrap.py",
-            module_name = "_frozen_importlib",
-        ),
-        "_frozen_importlib_external".into() => py_compile_bytecode!(
-            file = "Lib/_bootstrap_external.py",
-            module_name = "_frozen_importlib_external",
-        ),
+    let mut modules = HashMap::new();
+    modules.extend(py_compile_bytecode!(
+        source = "initialized = True; print(\"Hello world!\")\n",
+        module_name = "__hello__",
+    ));
+    modules.extend(py_compile_bytecode!(
+        file = "Lib/_bootstrap.py",
+        module_name = "_frozen_importlib",
+    ));
+    modules.extend(py_compile_bytecode!(
+        file = "Lib/_bootstrap_external.py",
+        module_name = "_frozen_importlib_external",
+    ));
+
+    #[cfg(feature = "freeze-stdlib")]
+    {
+        modules.extend(py_compile_bytecode!(dir = "../Lib/",));
     }
+
+    modules
 }
diff --git a/vm/src/lib.rs b/vm/src/lib.rs
@@ -80,4 +80,5 @@ pub use rustpython_bytecode::*;
 #[doc(hidden)]
 pub mod __exports {
     pub use bincode;
+    pub use maplit::hashmap;
 }
diff --git a/wasm/lib/src/browser_module.rs b/wasm/lib/src/browser_module.rs
@@ -381,8 +381,8 @@ pub fn setup_browser_module(vm: &VirtualMachine) {
     vm.stdlib_inits
         .borrow_mut()
         .insert("_browser".to_string(), Box::new(make_module));
-    vm.frozen.borrow_mut().insert(
-        "browser".to_string(),
-        py_compile_bytecode!(file = "src/browser.py"),
-    );
+    vm.frozen.borrow_mut().extend(py_compile_bytecode!(
+        file = "src/browser.py",
+        module_name = "browser",
+    ));
 }

Original file line number	Diff line number	Diff line change
`@@ -80,4 +80,5 @@ pub use rustpython_bytecode::*;`
`80`	`80`	`#[doc(hidden)]`
`81`	`81`	`pub mod __exports {`
`82`	`82`	`pub use bincode;`
	`83`	`+ pub use maplit::hashmap;`
`83`	`84`	`}`