From b23242029ec02fddb28ba35933fa8d7f9a305cd3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 25 Dec 2025 10:18:15 +0000 Subject: [PATCH 1/4] Initial plan From 3c814450d3b2e16c443231ba67546f6a0f5b152b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 25 Dec 2025 10:37:18 +0000 Subject: [PATCH 2/4] Add module state storage to PyModule Co-authored-by: youknowone <69878+youknowone@users.noreply.github.com> --- crates/vm/src/builtins/module.rs | 67 ++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/crates/vm/src/builtins/module.rs b/crates/vm/src/builtins/module.rs index f8e42b28e0..bb785095ce 100644 --- a/crates/vm/src/builtins/module.rs +++ b/crates/vm/src/builtins/module.rs @@ -3,6 +3,7 @@ use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, builtins::{PyStrInterned, pystr::AsPyStr}, class::PyClassImpl, + common::lock::PyRwLock, convert::ToPyObject, function::{FuncArgs, PyMethodDef}, types::{GetAttr, Initializer, Representable}, @@ -48,6 +49,7 @@ pub struct PyModule { // PyObject *md_dict; pub def: Option<&'static PyModuleDef>, // state: Any + state: PyRwLock>, // weaklist // for logging purposes after md_dict is cleared pub name: Option<&'static PyStrInterned>, @@ -72,6 +74,7 @@ impl PyModule { pub const fn new() -> Self { Self { def: None, + state: PyRwLock::new(None), name: None, } } @@ -79,6 +82,7 @@ impl PyModule { pub const fn from_def(def: &'static PyModuleDef) -> Self { Self { def: Some(def), + state: PyRwLock::new(None), name: Some(def.name), } } @@ -87,6 +91,31 @@ impl PyModule { let doc = module.def.unwrap().doc.map(|doc| doc.to_owned()); module.init_dict(module.name.unwrap(), doc, vm); } + + pub fn get_state(&self) -> Option> { + self.state + .read() + .as_ref() + .and_then(|obj| obj.clone().downcast().ok()) + } + + pub fn get_or_try_init_state(&self, vm: &VirtualMachine, init: F) -> PyResult> + where + T: PyPayload, + F: FnOnce(&VirtualMachine) -> PyResult>, + { + if let Some(state) = self.get_state() { + return Ok(state); + } + + if self.state.read().is_some() { + return Err(vm.new_type_error("module state has incompatible type".to_owned())); + } + + let state = init(vm)?; + *self.state.write() = Some(state.clone().into()); + Ok(state) + } } impl Py { @@ -224,3 +253,41 @@ impl Representable for PyModule { pub(crate) fn init(context: &Context) { PyModule::extend_class(context, context.types.module_type); } + +#[cfg(test)] +mod tests { + use crate::{ + AsObject, + builtins::{PyInt, PyStr}, + vm::Interpreter, + }; + use malachite_bigint::ToBigInt; + + #[test] + fn module_state_is_per_module_and_typed() { + Interpreter::without_stdlib(Default::default()).enter(|vm| { + let m1 = vm.new_module("m1", vm.ctx.new_dict(), None); + let m2 = vm.new_module("m2", vm.ctx.new_dict(), None); + + assert!(m1.get_state::().is_none()); + + let s1 = m1 + .get_or_try_init_state(vm, |vm| Ok(vm.ctx.new_int(1))) + .unwrap(); + let s2 = m2 + .get_or_try_init_state(vm, |vm| Ok(vm.ctx.new_int(2))) + .unwrap(); + + assert_eq!(s1.as_bigint(), &1.to_bigint().unwrap()); + assert_eq!(s2.as_bigint(), &2.to_bigint().unwrap()); + + let s1_again = m1.get_state::().unwrap(); + assert!(s1_again.is(&s1)); + + let err = m1 + .get_or_try_init_state::(vm, |vm| Ok(vm.ctx.new_str("oops"))) + .unwrap_err(); + assert!(err.class().is(vm.ctx.exceptions.type_error)); + }); + } +} From 5ee69bd2192feec326bfeef90fa1b7695f4a0bc6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 25 Dec 2025 10:57:41 +0000 Subject: [PATCH 3/4] Refine module state helpers and docs Co-authored-by: youknowone <69878+youknowone@users.noreply.github.com> --- crates/vm/src/builtins/module.rs | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/crates/vm/src/builtins/module.rs b/crates/vm/src/builtins/module.rs index bb785095ce..dc1686c9ca 100644 --- a/crates/vm/src/builtins/module.rs +++ b/crates/vm/src/builtins/module.rs @@ -70,6 +70,9 @@ pub struct ModuleInitArgs { } impl PyModule { + const STATE_TYPE_ERROR: &'static str = + "module state has already been initialized with a different type"; + #[allow(clippy::new_without_default)] pub const fn new() -> Self { Self { @@ -92,6 +95,8 @@ impl PyModule { module.init_dict(module.name.unwrap(), doc, vm); } + /// Return the stored module state if it exists and matches `T`. Returns `None` when no state + /// has been set or when the stored state is of a different type. pub fn get_state(&self) -> Option> { self.state .read() @@ -99,21 +104,25 @@ impl PyModule { .and_then(|obj| obj.clone().downcast().ok()) } + /// Get or initialize the module state of type `T`, using `init` only when no state exists and + /// creating it while holding a write lock. Raises `TypeError` if an incompatible state is + /// already stored. pub fn get_or_try_init_state(&self, vm: &VirtualMachine, init: F) -> PyResult> where T: PyPayload, F: FnOnce(&VirtualMachine) -> PyResult>, { - if let Some(state) = self.get_state() { - return Ok(state); - } - - if self.state.read().is_some() { - return Err(vm.new_type_error("module state has incompatible type".to_owned())); + let mut lock = self.state.write(); + if let Some(existing) = lock.as_ref() { + return existing + .clone() + .downcast() + .map_err(|_| vm.new_type_error(Self::STATE_TYPE_ERROR)); } let state = init(vm)?; - *self.state.write() = Some(state.clone().into()); + + *lock = Some(state.as_object().to_owned()); Ok(state) } } @@ -261,7 +270,7 @@ mod tests { builtins::{PyInt, PyStr}, vm::Interpreter, }; - use malachite_bigint::ToBigInt; + use malachite_bigint::BigInt; #[test] fn module_state_is_per_module_and_typed() { @@ -278,8 +287,8 @@ mod tests { .get_or_try_init_state(vm, |vm| Ok(vm.ctx.new_int(2))) .unwrap(); - assert_eq!(s1.as_bigint(), &1.to_bigint().unwrap()); - assert_eq!(s2.as_bigint(), &2.to_bigint().unwrap()); + assert_eq!(s1.as_bigint(), &BigInt::from(1)); + assert_eq!(s2.as_bigint(), &BigInt::from(2)); let s1_again = m1.get_state::().unwrap(); assert!(s1_again.is(&s1)); From 71376f43665fe1d3659467ae7e86b62fc95cdda9 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 26 Dec 2025 11:34:59 +0900 Subject: [PATCH 4/4] module {get,set}_state --- crates/vm/src/builtins/module.rs | 62 +++++++++----------------------- 1 file changed, 17 insertions(+), 45 deletions(-) diff --git a/crates/vm/src/builtins/module.rs b/crates/vm/src/builtins/module.rs index dc1686c9ca..d318070816 100644 --- a/crates/vm/src/builtins/module.rs +++ b/crates/vm/src/builtins/module.rs @@ -70,9 +70,6 @@ pub struct ModuleInitArgs { } impl PyModule { - const STATE_TYPE_ERROR: &'static str = - "module state has already been initialized with a different type"; - #[allow(clippy::new_without_default)] pub const fn new() -> Self { Self { @@ -104,26 +101,9 @@ impl PyModule { .and_then(|obj| obj.clone().downcast().ok()) } - /// Get or initialize the module state of type `T`, using `init` only when no state exists and - /// creating it while holding a write lock. Raises `TypeError` if an incompatible state is - /// already stored. - pub fn get_or_try_init_state(&self, vm: &VirtualMachine, init: F) -> PyResult> - where - T: PyPayload, - F: FnOnce(&VirtualMachine) -> PyResult>, - { - let mut lock = self.state.write(); - if let Some(existing) = lock.as_ref() { - return existing - .clone() - .downcast() - .map_err(|_| vm.new_type_error(Self::STATE_TYPE_ERROR)); - } - - let state = init(vm)?; - - *lock = Some(state.as_object().to_owned()); - Ok(state) + /// Set the module state. + pub fn set_state(&self, state: PyObjectRef) { + *self.state.write() = Some(state); } } @@ -265,11 +245,7 @@ pub(crate) fn init(context: &Context) { #[cfg(test)] mod tests { - use crate::{ - AsObject, - builtins::{PyInt, PyStr}, - vm::Interpreter, - }; + use crate::{AsObject, builtins::PyInt, vm::Interpreter}; use malachite_bigint::BigInt; #[test] @@ -280,23 +256,19 @@ mod tests { assert!(m1.get_state::().is_none()); - let s1 = m1 - .get_or_try_init_state(vm, |vm| Ok(vm.ctx.new_int(1))) - .unwrap(); - let s2 = m2 - .get_or_try_init_state(vm, |vm| Ok(vm.ctx.new_int(2))) - .unwrap(); - - assert_eq!(s1.as_bigint(), &BigInt::from(1)); - assert_eq!(s2.as_bigint(), &BigInt::from(2)); - - let s1_again = m1.get_state::().unwrap(); - assert!(s1_again.is(&s1)); - - let err = m1 - .get_or_try_init_state::(vm, |vm| Ok(vm.ctx.new_str("oops"))) - .unwrap_err(); - assert!(err.class().is(vm.ctx.exceptions.type_error)); + let s1 = vm.ctx.new_int(1); + let s2 = vm.ctx.new_int(2); + m1.set_state(s1.as_object().to_owned()); + m2.set_state(s2.as_object().to_owned()); + + assert_eq!( + m1.get_state::().unwrap().as_bigint(), + &BigInt::from(1) + ); + assert_eq!( + m2.get_state::().unwrap().as_bigint(), + &BigInt::from(2) + ); }); } }