Restore Result compatibility on safe VM leak-fix base

doeff/_vendor.py

@@ -288,8 +288,60 @@ def __hash__(self) -> int:
 FrozenDict = frozendict
 
 
+class Result(Generic[T_co]):
+    """Legacy Result base kept for older imports and persisted annotations."""
+
+    __slots__ = ()
+
+    def is_ok(self) -> bool:
+        return isinstance(self, Ok)
+
+    def is_err(self) -> bool:
+        return isinstance(self, Err)
+
+    def ok(self) -> T_co | None:
+        if isinstance(self, Ok):
+            return self.value
+        return None
+
+    def err(self) -> Any | None:
+        if isinstance(self, Err):
+            return self.error
+        return None
+
+    def expect(self, message: str) -> T_co:
+        if isinstance(self, Ok):
+            return self.value
+        if message:
+            raise RuntimeError(f"{message}: {self.error}") from self.error
+        raise self.error
+
+    def unwrap(self) -> T_co:
+        if isinstance(self, Ok):
+            return self.value
+        raise self.error
+
+    def unwrap_err(self) -> Any:
+        if isinstance(self, Err):
+            return self.error
+        raise RuntimeError("Called unwrap_err on Ok value")
+
+    def map(self, func: Callable[[T_co], U]) -> "Result[U]":
+        if isinstance(self, Ok):
+            return Ok(func(self.value))
+        return self
+
+    def map_err(self, func: Callable[[Any], Any]) -> "Result[T_co]":
+        if isinstance(self, Err):
+            return Err(func(self.error))
+        return self
+
+    def __bool__(self) -> bool:
+        return self.is_ok()
+
+
 @dataclass(frozen=True)
-class Ok(Generic[T_co]):
+class Ok(Result[T_co], Generic[T_co]):
     """Legacy Result shim kept for persisted pickles created before the Rust move."""
 
     value: T_co
@@ -305,7 +357,7 @@ def __bool__(self) -> bool:
 
 
 @dataclass(frozen=True)
-class Err:
+class Err(Result[NoReturn]):
     """Legacy error shim kept for persisted pickles created before the Rust move."""
 
     error: Any
@@ -327,6 +379,7 @@ def __bool__(self) -> bool:
     "Maybe",
     "Nothing",
     "Ok",
+    "Result",
     "Some",
     "TraceError",
     "WGraph",

packages/doeff-vm-core/src/arena.rs

@@ -84,13 +84,22 @@ impl FiberArena {
     }
 
     pub fn capacity(&self) -> usize {
+        self.fibers.capacity()
+    }
+
+    pub fn slot_count(&self) -> usize {
         self.fibers.len()
     }
 
     pub fn clear(&mut self) {
         self.fibers.clear();
         self.free_list.clear();
     }
+
+    pub fn shrink_to_fit(&mut self) {
+        self.fibers.shrink_to_fit();
+        self.free_list.shrink_to_fit();
+    }
 }
 
 impl Default for FiberArena {

packages/doeff-vm-core/src/continuation.rs

@@ -8,14 +8,15 @@ use pyo3::types::{PyDict, PyList};
 
 use crate::frame::CallMetadata;
 use crate::frame::Frame;
+use crate::ids::VarId;
 use crate::ids::{ContId, DispatchId, Marker, ScopeId, SegmentId};
 use crate::kleisli::KleisliRef;
+use crate::memory_stats;
 use crate::py_key::HashedPyKey;
 use crate::py_shared::PyShared;
 use crate::segment::Segment;
 use crate::step::PyException;
 use crate::value::Value;
-use crate::ids::VarId;
 
 #[pyclass(name = "K")]
 pub struct PyK {
@@ -62,7 +63,7 @@ pub(crate) struct DispatchHandlerHint {
     pub(crate) prompt_seg_id: SegmentId,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct Continuation {
     pub cont_id: ContId,
     dispatch_id: Option<DispatchId>,
@@ -80,27 +81,15 @@ pub struct Continuation {
 
 impl Continuation {
     fn captured_segment_snapshot(segment: &Segment) -> Box<Segment> {
-        Box::new(Segment {
-            scope_id: segment.scope_id,
-            persistent_epoch: segment.persistent_epoch,
-            marker: segment.marker,
-            frames: segment.frames.clone(),
-            parent: segment.parent,
-            state_store: segment.state_store.clone(),
-            writer_log: segment.writer_log.clone(),
-            kind: segment.kind.clone(),
-            pending_error_context: segment.pending_error_context.clone(),
-            throw_parent: segment.throw_parent.clone(),
-            interceptor_eval_depth: segment.interceptor_eval_depth,
-            interceptor_skip_stack: segment.interceptor_skip_stack.clone(),
-        })
+        Box::new(segment.clone())
     }
 
     pub fn capture(
         segment: &Segment,
         segment_id: SegmentId,
         dispatch_id: Option<DispatchId>,
     ) -> Self {
+        memory_stats::register_continuation();
         Continuation {
             cont_id: ContId::fresh(),
             dispatch_id,
@@ -122,6 +111,7 @@ impl Continuation {
         segment_id: SegmentId,
         dispatch_id: Option<DispatchId>,
     ) -> Self {
+        memory_stats::register_continuation();
         Continuation {
             cont_id,
             dispatch_id,
@@ -148,6 +138,7 @@ impl Continuation {
         outside_scope: Option<SegmentId>,
     ) -> Self {
         let handler_count = handlers.len();
+        memory_stats::register_continuation();
         Continuation {
             cont_id: ContId::fresh(),
             dispatch_id: None,
@@ -190,6 +181,7 @@ impl Continuation {
         metadata: Option<CallMetadata>,
         outside_scope: Option<SegmentId>,
     ) -> Self {
+        memory_stats::register_continuation();
         Continuation {
             cont_id: ContId::fresh(),
             dispatch_id: None,
@@ -325,19 +317,10 @@ impl Continuation {
     }
 
     pub(crate) fn clone_for_dispatch(&self, dispatch_id: Option<DispatchId>) -> Self {
-        Continuation {
-            cont_id: ContId::fresh(),
-            dispatch_id,
-            resume_dispatch_id: self.resume_dispatch_id,
-            dispatch_handler_hint: self.dispatch_handler_hint,
-            segment_id: self.segment_id,
-            segment_snapshot: self.segment_snapshot.clone(),
-            scope_parent_snapshot: self.scope_parent_snapshot,
-            scope_bindings_snapshot: self.scope_bindings_snapshot.clone(),
-            var_overrides_snapshot: self.var_overrides_snapshot.clone(),
-            unstarted: self.unstarted.clone(),
-            parent: self.parent.clone(),
-        }
+        let mut cloned = self.clone();
+        cloned.cont_id = ContId::fresh();
+        cloned.dispatch_id = dispatch_id;
+        cloned
     }
 
     pub(crate) fn refresh_persistent_segment_state(
@@ -432,15 +415,15 @@ impl Continuation {
     }
 
     pub(crate) fn into_unstarted_parts(
-        self,
+        mut self,
     ) -> Option<(
         PyShared,
         Vec<KleisliRef>,
         Vec<Option<PyShared>>,
         Option<CallMetadata>,
         Option<SegmentId>,
     )> {
-        self.unstarted.map(
+        self.unstarted.take().map(
             |UnstartedContinuation {
                  program,
                  handlers,
@@ -486,13 +469,39 @@ impl Continuation {
     }
 }
 
+impl Clone for Continuation {
+    fn clone(&self) -> Self {
+        memory_stats::register_continuation();
+        Continuation {
+            cont_id: self.cont_id,
+            dispatch_id: self.dispatch_id,
+            resume_dispatch_id: self.resume_dispatch_id,
+            dispatch_handler_hint: self.dispatch_handler_hint,
+            segment_id: self.segment_id,
+            segment_snapshot: self.segment_snapshot.clone(),
+            scope_parent_snapshot: self.scope_parent_snapshot,
+            scope_bindings_snapshot: self.scope_bindings_snapshot.clone(),
+            var_overrides_snapshot: self.var_overrides_snapshot.clone(),
+            unstarted: self.unstarted.clone(),
+            parent: self.parent.clone(),
+        }
+    }
+}
+
+impl Drop for Continuation {
+    fn drop(&mut self) {
+        memory_stats::unregister_continuation();
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::do_ctrl::{DoCtrl, InterceptMode};
     use crate::error::VMError;
     use crate::ids::Marker;
     use crate::kleisli::{Kleisli, KleisliDebugInfo};
+    use crate::memory_stats::live_object_counts;
     use crate::segment::SegmentKind;
     use crate::value::Value;
 
@@ -551,11 +560,15 @@ mod tests {
     #[test]
     fn test_unstarted_continuation_has_no_segment_snapshot() {
         Python::attach(|py| {
+            let baseline = live_object_counts().live_continuations;
             let cont = Continuation::create_unstarted(PyShared::new(py.None()), Vec::new());
             assert!(!cont.is_started());
             assert!(cont.segment_id().is_none());
             assert!(cont.segment().is_none());
             assert!(cont.frames().is_none());
+            assert_eq!(live_object_counts().live_continuations, baseline + 1);
+            drop(cont);
+            assert_eq!(live_object_counts().live_continuations, baseline);
         });
     }
 
@@ -642,4 +655,21 @@ mod tests {
         assert_eq!(snapshot.writer_log, vec![Value::Int(20)]);
         assert_eq!(snapshot.persistent_epoch, 2);
     }
+
+    #[test]
+    fn test_continuation_live_count_tracks_clone_lifetime() {
+        let baseline = live_object_counts().live_continuations;
+        let (seg, seg_id) = make_test_segment();
+        let cont = Continuation::capture(&seg, seg_id, None);
+        assert_eq!(live_object_counts().live_continuations, baseline + 1);
+
+        let cont_clone = cont.clone();
+        assert_eq!(live_object_counts().live_continuations, baseline + 2);
+
+        drop(cont_clone);
+        assert_eq!(live_object_counts().live_continuations, baseline + 1);
+
+        drop(cont);
+        assert_eq!(live_object_counts().live_continuations, baseline);
+    }
 }

packages/doeff-vm-core/src/debug_state.rs

@@ -46,6 +46,10 @@ impl DebugState {
         self.trace_events.clear();
     }
 
+    pub(crate) fn shrink_to_fit(&mut self) {
+        self.trace_events.shrink_to_fit();
+    }
+
     pub(crate) fn trace_events(&self) -> &[TraceEvent] {
         &self.trace_events
     }

packages/doeff-vm-core/src/dispatch_observer.rs

@@ -33,6 +33,27 @@ impl DispatchObserver {
         self.segment_dispatch_ids.clear();
     }
 
+    pub(crate) fn shrink_to_fit(&mut self) {
+        self.dispatches.shrink_to_fit();
+        self.segment_dispatch_ids.shrink_to_fit();
+    }
+
+    pub(crate) fn dispatch_count(&self) -> usize {
+        self.dispatches.len()
+    }
+
+    pub(crate) fn segment_binding_count(&self) -> usize {
+        self.segment_dispatch_ids.len()
+    }
+
+    pub(crate) fn dispatch_capacity(&self) -> usize {
+        self.dispatches.capacity()
+    }
+
+    pub(crate) fn segment_binding_capacity(&self) -> usize {
+        self.segment_dispatch_ids.capacity()
+    }
+
     pub(crate) fn start_dispatch(
         &mut self,
         dispatch_id: DispatchId,

packages/doeff-vm-core/src/frame.rs

@@ -270,9 +270,7 @@ mod tests {
 
     #[test]
     fn test_program_frame_is_program() {
-        let stream = std::sync::Arc::new(std::sync::Mutex::new(
-            Box::new(DummyStream) as Box<dyn IRStream>
-        ));
+        let stream = IRStreamRef::new(Box::new(DummyStream) as Box<dyn IRStream>);
         let frame = Frame::program(stream, None);
         assert!(frame.is_program());
     }