From d5a80492343cf3f2481abf12da4e49d09d95d6e9 Mon Sep 17 00:00:00 2001 From: youknowdot Date: Wed, 21 Oct 2026 23:05:38 +0900 Subject: [PATCH] Support collision-safe frozen package containers Generate empty frozen packages only for source namespace directories with real frozen descendants. Preserve genuine modules and packages across frozen-library registrations while retaining real-real replacement order. Record the container marker in FrozenModule and its library envelope; legacy module/package encodings remain unchanged. This is closed-tree frozen packaging, not full PEP 420 namespace merging. Add structural codec, mock generator, and registration-priority checks. --- crates/compiler-core/src/frozen.rs | 85 +++++++++- crates/derive-impl/src/compile_bytecode.rs | 172 ++++++++++++++++++++- crates/vm/src/vm/interpreter.rs | 69 ++++++++- 3 files changed, 315 insertions(+), 11 deletions(-) diff --git a/crates/compiler-core/src/frozen.rs b/crates/compiler-core/src/frozen.rs index 81530610d..0ecf55f92 100644 --- a/crates/compiler-core/src/frozen.rs +++ b/crates/compiler-core/src/frozen.rs @@ -7,6 +7,10 @@ use alloc::vec::Vec; pub struct FrozenModule { pub code: FrozenCodeObject, pub package: bool, + /// Empty package generated for a source namespace directory. Genuine + /// modules take precedence over these containers across frozen libraries. + /// When true, `package` must also be true. + pub namespace_container: bool, } #[derive(Copy, Clone)] @@ -97,8 +101,15 @@ fn read_entry<'a>( let len = rdr.read_u32()?; let code_slice = rdr.read_slice_borrow(len)?; let code = FrozenCodeObject { bytes: code_slice }; - let package = rdr.read_u8()? != 0; - Ok((name, FrozenModule { code, package })) + let flags = rdr.read_u8()?; + Ok(( + name, + FrozenModule { + code, + package: flags & 1 != 0, + namespace_container: flags & 2 != 0, + }, + )) } impl FrozenLib> { @@ -125,7 +136,75 @@ fn write_lib<'a, B: AsRef<[u8]>>( } fn write_entry(buf: &mut Vec, name: &str, module: FrozenModule>) { + assert!(!module.namespace_container || module.package); marshal::write_vec(buf, name.as_bytes()); marshal::write_vec(buf, module.code.bytes.as_ref()); - buf.write_u8(module.package as u8); + // Ordinary modules retain the legacy 0/1 encoding. Older readers treat + // container flag 3 as a package, but cannot preserve container priority. + buf.write_u8(module.package as u8 | ((module.namespace_container as u8) << 1)); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn library_flags_preserve_legacy_and_container_metadata() { + // Opaque bytes exercise only the library envelope, not Python bytecode. + for (package, namespace_container, flags) in + [(false, false, 0), (true, false, 1), (true, true, 3)] + { + let module = FrozenModule { + code: FrozenCodeObject { bytes: &[7_u8][..] }, + package, + namespace_container, + }; + let encoded = FrozenLib::encode([("entry", module)]); + assert_eq!(encoded.bytes.last(), Some(&flags)); + if !namespace_container { + let mut legacy = Vec::new(); + marshal::write_len(&mut legacy, 1); + marshal::write_vec(&mut legacy, b"entry"); + marshal::write_vec(&mut legacy, &[7]); + legacy.write_u8(package as u8); + assert_eq!(encoded.bytes, legacy); + } + let mut decoded = encoded.decode(); + let (name, decoded_module) = decoded.next().unwrap(); + assert_eq!(name, "entry"); + assert_eq!(decoded_module.code.bytes, &[7]); + assert_eq!(decoded_module.package, package); + assert_eq!(decoded_module.namespace_container, namespace_container); + assert!(decoded.next().is_none()); + } + let entries = [ + ( + "real", + FrozenModule { + code: FrozenCodeObject { + bytes: &[1_u8, 2][..], + }, + package: false, + namespace_container: false, + }, + ), + ( + "container", + FrozenModule { + code: FrozenCodeObject { bytes: &[3_u8][..] }, + package: true, + namespace_container: true, + }, + ), + ]; + let encoded = FrozenLib::encode(entries); + let decoded: Vec<_> = encoded.decode().collect(); + assert_eq!(decoded.len(), 2); + assert_eq!(decoded[0].0, "real"); + assert_eq!(decoded[0].1.code.bytes, &[1, 2]); + assert!(!decoded[0].1.namespace_container); + assert_eq!(decoded[1].0, "container"); + assert_eq!(decoded[1].1.code.bytes, &[3]); + assert!(decoded[1].1.namespace_container); + } } diff --git a/crates/derive-impl/src/compile_bytecode.rs b/crates/derive-impl/src/compile_bytecode.rs index e08c230e2..a1c95c116 100644 --- a/crates/derive-impl/src/compile_bytecode.rs +++ b/crates/derive-impl/src/compile_bytecode.rs @@ -14,6 +14,7 @@ //! ``` use crate::Diagnostic; +use alloc::collections::BTreeMap; use proc_macro2::{Span, TokenStream}; use quote::quote; use rustpython_compiler_core::{Mode, bytecode::CodeObject, frozen}; @@ -40,6 +41,7 @@ enum CompilationSourceKind { struct CompiledModule { code: CodeObject, package: bool, + namespace_container: bool, } struct CompilationSource { @@ -91,13 +93,38 @@ impl CompilationSource { ) -> Result, Diagnostic> { match &self.kind { CompilationSourceKind::Dir { base, rel_path } => { - self.compile_dir(base, &base.join(rel_path), "", mode, compiler) + let mut containers = BTreeMap::new(); + let mut modules = self.compile_dir( + base, + &base.join(rel_path), + "", + mode, + compiler, + &mut containers, + )?; + // Wait until all genuine modules are known: foo.py must beat a + // generated foo/ container regardless of directory iteration order. + for (name, path) in containers { + if let std::collections::hash_map::Entry::Vacant(entry) = + modules.entry(name.clone()) + { + let code = self + .compile_string("", Mode::Exec, &name, compiler, || path.display())?; + entry.insert(CompiledModule { + code, + package: true, + namespace_container: true, + }); + } + } + Ok(modules) } _ => Ok(core::iter::once(( module_name.to_string(), CompiledModule { code: self.compile_single(mode, module_name, compiler)?, package: false, + namespace_container: false, }, )) .collect()), @@ -141,7 +168,9 @@ impl CompilationSource { parent: &str, mode: Mode, compiler: &dyn Compiler, + containers: &mut BTreeMap, ) -> Result, Diagnostic> { + let mut has_initializer = false; let mut code_map = HashMap::new(); let paths = fs::read_dir(path) .or_else(|e| { @@ -175,8 +204,10 @@ impl CompilationSource { } else { format!("{parent}.{file_name}") }; - code_map.extend(self.compile_dir(base, &path, &nparent, mode, compiler)?); + code_map + .extend(self.compile_dir(base, &path, &nparent, mode, compiler, containers)?); } else if file_name.ends_with(".py") { + has_initializer |= file_name == "__init__.py"; // These fixtures must be parsed at import time to exercise syntax errors. if parent == "test.test_lazy_import.data.badsyntax" { continue; @@ -232,10 +263,14 @@ impl CompilationSource { CompiledModule { code, package: is_init, + namespace_container: false, }, ); } } + if !parent.is_empty() && !has_initializer && !code_map.is_empty() { + containers.insert(parent.to_owned(), path.to_owned()); + } Ok(code_map) } } @@ -387,6 +422,7 @@ pub(crate) fn impl_py_freeze( let v = frozen::FrozenModule { code: frozen::FrozenCodeObject::encode(&v.code), package: v.package, + namespace_container: v.namespace_container, }; (&**k, v) })); @@ -398,3 +434,135 @@ pub(crate) fn impl_py_freeze( Ok(output) } + +#[cfg(test)] +mod tests { + use super::*; + use core::sync::atomic::{AtomicUsize, Ordering}; + use rustpython_compiler_core::bytecode::{CodeFlags, ConstantData, Constants}; + + struct FixtureCompiler; + + impl Compiler for FixtureCompiler { + fn compile( + &self, + source: &str, + _mode: Mode, + module_name: &str, + ) -> Result> { + if source == "invalid fixture" { + return Err(std::io::Error::other("fixture excluded by existing policy").into()); + } + // Structural generator fixture only. No Python source is compiled + // or executed; source_path records which real entry survived. + Ok(CodeObject { + instructions: Vec::new().into(), + locations: Box::new([]), + flags: CodeFlags::empty(), + posonlyarg_count: 0, + arg_count: 0, + kwonlyarg_count: 0, + source_path: source.to_owned(), + first_line_number: None, + max_stackdepth: 0, + obj_name: module_name.to_owned(), + qualname: module_name.to_owned(), + constants: Constants::::new(Box::new([])), + names: Box::new([]), + varnames: Box::new([]), + cellvars: Box::new([]), + freevars: Box::new([]), + localspluskinds: Box::new([]), + linetable: Box::new([]), + exceptiontable: Box::new([]), + }) + } + } + + struct FixtureDir(PathBuf); + + impl FixtureDir { + fn new() -> Self { + static NEXT: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "rustpython-freeze-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + fs::create_dir(&path).unwrap(); + Self(path) + } + + fn write(&self, path: &str, source: &str) { + let path = self.0.join(path); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(path, source).unwrap(); + } + + fn compile(&self) -> HashMap { + let source = CompilationSource { + kind: CompilationSourceKind::Dir { + base: self.0.clone(), + rel_path: PathBuf::new(), + }, + span: (Span::call_site(), Span::call_site()), + }; + source + .compile(Mode::Exec, "", &FixtureCompiler) + .unwrap_or_else(|_| panic!("fixture compilation failed")) + } + } + + impl Drop for FixtureDir { + fn drop(&mut self) { + fs::remove_dir_all(&self.0).unwrap(); + } + } + + #[test] + fn frozen_containers_follow_physical_directories_and_preserve_real_entries() { + for module_first in [true, false] { + let dir = FixtureDir::new(); + if module_first { + dir.write("collision.py", "real module"); + } + dir.write("collision/child.py", "child"); + if !module_first { + dir.write("collision.py", "real module"); + } + dir.write("namespace/nested/child.py", "child"); + dir.write("namespace/regular/__init__.py", "real initializer"); + dir.write("dotted.name/child.py", "child"); + dir.write("data/readme.txt", "data only"); + fs::create_dir(dir.0.join("empty")).unwrap(); + dir.write("excluded/badsyntax_fixture.py", "invalid fixture"); + dir.write("test/encoded_modules/__init__.py", "invalid fixture"); + dir.write("test/encoded_modules/child.py", "child"); + let modules = dir.compile(); + for name in ["namespace", "namespace.nested", "dotted.name"] { + assert!(modules[name].package); + assert!(modules[name].namespace_container); + assert_eq!(modules[name].code.source_path, ""); + } + assert!(!modules["collision"].package); + assert!(!modules["collision"].namespace_container); + assert_eq!(modules["collision"].code.source_path, "real module"); + assert!(modules["namespace.regular"].package); + assert!(!modules["namespace.regular"].namespace_container); + assert_eq!( + modules["namespace.regular"].code.source_path, + "real initializer" + ); + for absent in [ + "", + "dotted", + "empty", + "data", + "excluded", + "test.encoded_modules", + ] { + assert!(!modules.contains_key(absent), "unexpected entry {absent}"); + } + } + } +} diff --git a/crates/vm/src/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs index 324f17574..a951067f7 100644 --- a/crates/vm/src/vm/interpreter.rs +++ b/crates/vm/src/vm/interpreter.rs @@ -15,6 +15,25 @@ use core::sync::atomic::Ordering; type InitFunc = Box; +type FrozenModuleMap = + std::collections::HashMap<&'static str, FrozenModule, rapidhash::quality::RandomState>; + +fn merge_frozen_modules( + modules: &mut FrozenModuleMap, + incoming: impl IntoIterator, +) { + for (name, module) in incoming { + modules + .entry(name) + .and_modify(|current| { + if !module.namespace_container || current.namespace_container { + *current = module; + } + }) + .or_insert(module); + } +} + /// Exit code used when stdout/stderr flush fails during interpreter shutdown. /// Matches CPython's behavior (see cpython/Python/pylifecycle.c). const EXITCODE_FLUSH_FAILURE: u32 = 120; @@ -122,12 +141,8 @@ where other => other, } as usize); - let mut frozen: std::collections::HashMap< - &'static str, - FrozenModule, - rapidhash::quality::RandomState, - > = core_frozen_inits().collect(); - frozen.extend(frozen_modules); + let mut frozen = FrozenModuleMap::default(); + merge_frozen_modules(&mut frozen, core_frozen_inits().chain(frozen_modules)); (config, all_module_defs, frozen, int_max_str_digits) }; @@ -887,6 +902,7 @@ fn core_frozen_inits() -> impl Iterator { FrozenModule { code: hello_code, package: false, + namespace_container: false, }, )); entries.push(( @@ -894,6 +910,7 @@ fn core_frozen_inits() -> impl Iterator { FrozenModule { code: hello_code, package: true, + namespace_container: false, }, )); entries.push(( @@ -901,6 +918,7 @@ fn core_frozen_inits() -> impl Iterator { FrozenModule { code: hello_code, package: false, + namespace_container: false, }, )); entries.push(( @@ -908,6 +926,7 @@ fn core_frozen_inits() -> impl Iterator { FrozenModule { code: hello_code, package: false, + namespace_container: false, }, )); } @@ -921,6 +940,7 @@ fn core_frozen_inits() -> impl Iterator { FrozenModule { code, package: false, + namespace_container: false, }, )); } @@ -934,6 +954,7 @@ fn core_frozen_inits() -> impl Iterator { FrozenModule { code, package: false, + namespace_container: false, }, )); } @@ -950,6 +971,42 @@ mod tests { }; use malachite_bigint::ToBigInt; + #[test] + fn frozen_container_merge_preserves_genuine_entries() { + fn module(bytes: &'static [u8], package: bool, namespace_container: bool) -> FrozenModule { + FrozenModule { + code: crate::frozen::FrozenCodeObject { bytes }, + package, + namespace_container, + } + } + // No bytecode is decoded or executed: this checks registration priority. + let real_a = module(&[1], false, false); + let real_b = module(&[2], true, false); + let container_a = module(&[3], true, true); + let container_b = module(&[4], true, true); + for (first, second, expected) in [ + (real_a, container_a, real_a), + (container_a, real_a, real_a), + (real_b, container_a, real_b), + (container_a, real_b, real_b), + (real_a, real_b, real_b), + (real_b, real_a, real_a), + (container_a, container_b, container_b), + ] { + let mut modules = FrozenModuleMap::default(); + merge_frozen_modules(&mut modules, [("owned", first), ("unrelated", real_a)]); + merge_frozen_modules(&mut modules, [("owned", second)]); + assert_eq!(modules["owned"].code.bytes, expected.code.bytes); + assert_eq!(modules["owned"].package, expected.package); + assert_eq!(modules["unrelated"].code.bytes, real_a.code.bytes); + assert_eq!( + modules["owned"].namespace_container, + expected.namespace_container + ); + } + } + #[test] fn add_py_integers() { Interpreter::without_stdlib(Default::default()).enter(|vm| {