superacli 1.1.5 → 1.1.7

package/ref-monty/crates/monty/src/bytecode/vm/scheduler.rs

@@ -1,620 +0,0 @@
-//! Task scheduler for async execution.
-//!
-//! This module implements the scheduler for managing concurrent async tasks
-//! and tracking external function calls. The scheduler is always present
-//! (created at VM initialization) to maintain separation of concerns.
-//!
-//! # Task Model
-//!
-//! - Task 0 is the "main task" which uses the VM's stack/frames directly
-//! - Spawned tasks (1+) store their own execution context in the Task struct
-//! - When switching tasks, the scheduler swaps contexts with the VM
-
-use std::collections::VecDeque;
-
-use ahash::{AHashMap, AHashSet};
-
-use crate::{
-    args::ArgValues,
-    asyncio::{CallId, TaskId},
-    exception_private::RunError,
-    heap::{DropWithHeap, HeapId},
-    heap_data::HeapDataMut,
-    parse::CodeRange,
-    value::Value,
-};
-
-/// Task execution state for async scheduling.
-///
-/// Tracks whether a task is ready to run, blocked waiting for something,
-/// or has completed (successfully or with an error).
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
-pub(crate) enum TaskState {
-    /// Task is ready to execute (in the ready queue).
-    Ready,
-    /// Task is blocked waiting for an external call to resolve.
-    BlockedOnCall(CallId),
-    /// Task is blocked waiting for a GatherFuture to complete.
-    BlockedOnGather(HeapId),
-    /// Task completed successfully with a return value.
-    Completed(Value),
-    /// Task failed with an error.
-    Failed(RunError),
-}
-
-/// A single async task with its own execution context.
-///
-/// The main task (task 0) doesn't store its own frames/stack - it uses the VM's
-/// directly. Spawned tasks store their execution context here so they can be
-/// swapped in and out.
-///
-/// # Context Switching
-///
-/// When switching away from a non-main task, its context is saved here.
-/// When switching to it, the context is loaded into the VM.
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
-pub(crate) struct Task {
-    /// Unique identifier for this task.
-    pub id: TaskId,
-    /// Serialized call frames for this task's execution.
-    /// Empty for the main task (which uses VM's frames directly).
-    pub frames: Vec<SerializedTaskFrame>,
-    /// Operand stack for this task.
-    /// Empty for the main task (which uses VM's stack directly).
-    pub stack: Vec<Value>,
-    /// Exception stack for nested except blocks.
-    pub exception_stack: Vec<Value>,
-    /// VM-level instruction_ip (for exception table lookup).
-    pub instruction_ip: usize,
-    /// Coroutine being executed by this task (if any).
-    /// Used to mark the coroutine as Completed when the task finishes.
-    pub coroutine_id: Option<HeapId>,
-    /// GatherFuture this task belongs to (if spawned by gather).
-    /// Used to cancel sibling tasks when this task fails.
-    pub gather_id: Option<HeapId>,
-    /// Index in the gather's results where this task's result should be stored.
-    /// Only set for tasks spawned by gather.
-    pub gather_result_idx: Option<usize>,
-    /// Current execution state.
-    pub state: TaskState,
-    /// CallId that unblocked this task (set when task transitions from Blocked to Ready).
-    /// Used to retrieve the resolved value when the task resumes.
-    pub unblocked_by: Option<CallId>,
-}
-
-/// Serialized call frame for task storage.
-///
-/// Similar to `SerializedFrame` but used within the scheduler for task context.
-/// Cannot store `&Code` references - uses `FunctionId` to look up code on resume.
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
-pub(crate) struct SerializedTaskFrame {
-    /// Which function's code this frame executes (None = module-level).
-    pub function_id: Option<crate::intern::FunctionId>,
-    /// Instruction pointer within this frame's bytecode.
-    pub ip: usize,
-    /// Base index into the VM stack for this frame's locals region.
-    pub stack_base: usize,
-    /// Number of local variable slots (0 for module-level frames).
-    pub locals_count: u16,
-    /// Call site position (for tracebacks).
-    pub call_position: Option<CodeRange>,
-}
-
-impl Task {
-    /// Creates a new task in the Ready state.
-    ///
-    /// # Arguments
-    /// * `id` - Unique task identifier
-    /// * `coroutine_id` - Optional HeapId of the coroutine being executed
-    /// * `gather_id` - Optional HeapId of the GatherFuture this task belongs to
-    pub fn new(
-        id: TaskId,
-        coroutine_id: Option<HeapId>,
-        gather_id: Option<HeapId>,
-        gather_result_idx: Option<usize>,
-    ) -> Self {
-        Self {
-            id,
-            frames: Vec::new(),
-            stack: Vec::new(),
-            exception_stack: Vec::new(),
-            instruction_ip: 0,
-            coroutine_id,
-            gather_id,
-            gather_result_idx,
-            state: TaskState::Ready,
-            unblocked_by: None,
-        }
-    }
-
-    /// Returns true if this task has completed (successfully or with failure).
-    #[inline]
-    pub fn is_finished(&self) -> bool {
-        matches!(self.state, TaskState::Completed(_) | TaskState::Failed(_))
-    }
-}
-
-/// Internal representation of a pending external call.
-///
-/// Stores the data needed to retry or resume an external function call,
-/// along with tracking information for the task that created it.
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
-pub(crate) struct PendingCallData {
-    /// Arguments for the function (includes both positional and keyword args).
-    pub args: ArgValues,
-    /// Task that created this call (for ignoring results if task is cancelled).
-    pub creator_task: TaskId,
-}
-
-/// Scheduler for managing concurrent async tasks and external call tracking.
-///
-/// The scheduler is always present (created at VM initialization) to maintain
-/// separation of concerns. All async-related state lives here:
-/// - Task management (creation, scheduling, completion)
-/// - External call ID allocation and tracking
-/// - Resolution of pending futures
-///
-/// # Main Task
-///
-/// Task 0 is the "main task" which executes using the VM's stack/frames directly.
-/// It's always created at scheduler initialization but doesn't store its own context
-/// (the VM holds it). Spawned tasks (1+) store their context in the Task struct.
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
-pub(crate) struct Scheduler {
-    /// All tasks (main task at index 0, spawned tasks follow).
-    tasks: Vec<Task>,
-    /// Queue of task IDs ready to execute.
-    ready_queue: VecDeque<TaskId>,
-    /// Currently executing task (None only during task switching).
-    current_task: Option<TaskId>,
-    /// Counter for generating new task IDs.
-    next_task_id: u32,
-    /// Counter for external call IDs (always incremented, even for sync resolution).
-    next_call_id: u32,
-    /// Maps CallId -> pending call data for unresolved external calls.
-    /// Populated when host calls `run_pending()`.
-    pending_calls: AHashMap<CallId, PendingCallData>,
-    /// Maps CallId -> resolved Value for futures that have been resolved.
-    /// Entry is removed when the value is consumed by awaiting.
-    resolved: AHashMap<CallId, Value>,
-    /// CallIds that have been awaited (to detect double-await).
-    consumed: AHashSet<CallId>,
-    /// Maps CallId -> (gather_heap_id, result_index) for gathers waiting on external futures.
-    /// When a CallId is resolved, the result is stored in the gather's results at the given index.
-    gather_waiters: AHashMap<CallId, (HeapId, usize)>,
-}
-
-impl Scheduler {
-    /// Creates a new scheduler with the main task (task 0) as current.
-    ///
-    /// The main task uses the VM's stack/frames directly and is always present.
-    /// It starts as the current task (not in the ready queue) since it runs
-    /// immediately without needing to be scheduled.
-    pub fn new() -> Self {
-        let mut main_task = Task::new(TaskId::default(), None, None, None);
-        // Main task starts Running, not Ready (it's the current task, not waiting)
-        main_task.state = TaskState::Ready; // Will be set properly when it blocks
-        Self {
-            tasks: vec![main_task],
-            ready_queue: VecDeque::new(), // Main task is current, not in ready queue
-            current_task: Some(TaskId::default()),
-            next_task_id: 1,
-            next_call_id: 0,
-            pending_calls: AHashMap::new(),
-            resolved: AHashMap::new(),
-            consumed: AHashSet::new(),
-            gather_waiters: AHashMap::new(),
-        }
-    }
-
-    /// Returns the currently executing task ID.
-    ///
-    /// Returns `None` only during task switching operations.
-    #[inline]
-    pub fn current_task_id(&self) -> Option<TaskId> {
-        self.current_task
-    }
-
-    /// Returns a reference to a task by ID.
-    ///
-    /// # Panics
-    /// Panics if the task ID doesn't exist.
-    #[inline]
-    pub fn get_task(&self, task_id: TaskId) -> &Task {
-        &self.tasks[task_id.raw() as usize]
-    }
-
-    /// Returns a mutable reference to a task by ID.
-    ///
-    /// # Panics
-    /// Panics if the task ID doesn't exist.
-    #[inline]
-    pub fn get_task_mut(&mut self, task_id: TaskId) -> &mut Task {
-        &mut self.tasks[task_id.raw() as usize]
-    }
-
-    /// Allocates a new CallId for an external function call.
-    ///
-    /// The counter always increments, even for sync resolution, to keep IDs unique.
-    pub fn allocate_call_id(&mut self) -> CallId {
-        let id = CallId::new(self.next_call_id);
-        self.next_call_id += 1;
-        id
-    }
-
-    /// Sets the next call ID counter.
-    ///
-    /// Used when lazily creating the scheduler to inherit the call ID counter
-    /// from the VM, ensuring call IDs remain unique across the transition.
-    pub fn set_next_call_id(&mut self, id: u32) {
-        self.next_call_id = id;
-    }
-
-    /// Stores pending call data for an external function call.
-    ///
-    /// Called when the host uses async resolution (`run_pending()`).
-    pub fn add_pending_call(&mut self, call_id: CallId, data: PendingCallData) {
-        self.pending_calls.insert(call_id, data);
-    }
-
-    /// Removes a call_id from the pending_calls map.
-    ///
-    /// Called when resolving a gather's external future - the call is no longer
-    /// pending once the result has been stored in the gather's results.
-    pub fn remove_pending_call(&mut self, call_id: CallId) {
-        self.pending_calls.remove(&call_id);
-    }
-
-    /// Returns true if a CallId has already been awaited (consumed).
-    #[inline]
-    pub fn is_consumed(&self, call_id: CallId) -> bool {
-        self.consumed.contains(&call_id)
-    }
-
-    /// Marks a CallId as consumed (awaited).
-    pub fn mark_consumed(&mut self, call_id: CallId) {
-        self.consumed.insert(call_id);
-    }
-
-    /// Registers a gather as waiting on an external future.
-    ///
-    /// When the CallId is resolved, the result will be stored in the gather's results
-    /// at the specified index.
-    pub fn register_gather_for_call(&mut self, call_id: CallId, gather_id: HeapId, result_index: usize) {
-        self.gather_waiters.insert(call_id, (gather_id, result_index));
-    }
-
-    /// Returns gather info if a gather is waiting on this CallId.
-    ///
-    /// Returns (gather_heap_id, result_index) if found, None otherwise.
-    /// Removes the entry from gather_waiters.
-    pub fn take_gather_waiter(&mut self, call_id: CallId) -> Option<(HeapId, usize)> {
-        self.gather_waiters.remove(&call_id)
-    }
-
-    /// Resolves a CallId with a value.
-    ///
-    /// Stores the value for later retrieval when the future is awaited.
-    /// If a task is blocked on this call, it will be unblocked.
-    ///
-    /// Uses `pending_calls` for O(1) lookup of the blocked task instead of
-    /// scanning all tasks.
-    pub fn resolve(&mut self, call_id: CallId, value: Value) {
-        // Get blocked task from pending_calls before removing (O(1) lookup)
-        let blocked_task = self.pending_calls.remove(&call_id).map(|data| data.creator_task);
-
-        // Store the resolved value
-        self.resolved.insert(call_id, value);
-
-        // Unblock the task if found
-        if let Some(task_id) = blocked_task {
-            let task = self.get_task_mut(task_id);
-            if matches!(task.state, TaskState::BlockedOnCall(cid) if cid == call_id) {
-                task.state = TaskState::Ready;
-                task.unblocked_by = Some(call_id);
-                self.ready_queue.push_back(task_id);
-            }
-        }
-    }
-
-    /// Takes the resolved value for a CallId, if available.
-    ///
-    /// Removes the value from the resolved map and returns it.
-    /// Returns `None` if the call hasn't been resolved yet.
-    pub fn take_resolved(&mut self, call_id: CallId) -> Option<Value> {
-        self.resolved.remove(&call_id)
-    }
-
-    /// Takes the resolved value for a task that was unblocked.
-    ///
-    /// If the task has an `unblocked_by` CallId set, takes the resolved value
-    /// for that call and clears the `unblocked_by` field.
-    /// Returns `None` if the task wasn't unblocked by a resolved call.
-    pub fn take_resolved_for_task(&mut self, task_id: TaskId) -> Option<Value> {
-        let task = &mut self.tasks[task_id.raw() as usize];
-        if let Some(call_id) = task.unblocked_by.take() {
-            self.resolved.remove(&call_id)
-        } else {
-            None
-        }
-    }
-
-    /// Marks the current task as blocked on an external call.
-    ///
-    /// The task will be unblocked when `resolve()` is called with the matching CallId.
-    pub fn block_current_on_call(&mut self, call_id: CallId) {
-        if let Some(task_id) = self.current_task {
-            let task = self.get_task_mut(task_id);
-            task.state = TaskState::BlockedOnCall(call_id);
-        }
-    }
-
-    /// Marks the current task as blocked on a GatherFuture.
-    ///
-    /// The task will be unblocked when all gathered tasks complete.
-    pub fn block_current_on_gather(&mut self, gather_id: HeapId) {
-        if let Some(task_id) = self.current_task {
-            let task = self.get_task_mut(task_id);
-            task.state = TaskState::BlockedOnGather(gather_id);
-        }
-    }
-
-    /// Returns all pending (unresolved) CallIds.
-    pub fn pending_call_ids(&self) -> Vec<CallId> {
-        self.pending_calls.keys().copied().collect()
-    }
-
-    /// Removes a task from the ready queue.
-    ///
-    /// Used when handling the main task directly (via `prepare_main_task_after_resolve`)
-    /// instead of through the normal task switching mechanism.
-    pub fn remove_from_ready_queue(&mut self, task_id: TaskId) {
-        self.ready_queue.retain(|&id| id != task_id);
-    }
-
-    /// Spawns a new task from a coroutine.
-    ///
-    /// Creates a new task that will execute the given coroutine when scheduled.
-    /// The task is added to the ready queue.
-    ///
-    /// # Arguments
-    /// * `coroutine_id` - HeapId of the coroutine to execute
-    /// * `gather_id` - Optional HeapId of the GatherFuture this task belongs to
-    /// * `gather_result_idx` - Optional index in the gather's results for this task
-    ///
-    /// # Returns
-    /// The TaskId of the newly created task.
-    pub fn spawn(
-        &mut self,
-        coroutine_id: HeapId,
-        gather_id: Option<HeapId>,
-        gather_result_idx: Option<usize>,
-    ) -> TaskId {
-        let task_id = TaskId::new(self.next_task_id);
-        self.next_task_id += 1;
-
-        let task = Task::new(task_id, Some(coroutine_id), gather_id, gather_result_idx);
-        self.tasks.push(task);
-        self.ready_queue.push_back(task_id);
-
-        task_id
-    }
-
-    /// Gets the next ready task from the queue.
-    ///
-    /// Returns `None` if no tasks are ready.
-    pub fn next_ready_task(&mut self) -> Option<TaskId> {
-        self.ready_queue.pop_front()
-    }
-
-    /// Adds a task back to the ready queue.
-    pub fn make_ready(&mut self, task_id: TaskId) {
-        let task = self.get_task_mut(task_id);
-        task.state = TaskState::Ready;
-        self.ready_queue.push_back(task_id);
-    }
-
-    /// Sets the current task.
-    pub fn set_current_task(&mut self, task_id: Option<TaskId>) {
-        self.current_task = task_id;
-    }
-
-    /// Marks a task as completed with a result value.
-    ///
-    /// If the task is part of a gather, updates the gather's results.
-    /// If this completes the gather, unblocks the waiting task.
-    pub fn complete_task(&mut self, task_id: TaskId, result: Value) {
-        let task = self.get_task_mut(task_id);
-        task.state = TaskState::Completed(result);
-        // Note: gather wake-up logic will be implemented when gather is fully integrated
-    }
-
-    /// Marks a task as failed with an error.
-    ///
-    /// If the task is part of a gather, returns the gather_id so the caller
-    /// can collect siblings from `GatherFuture.task_ids` on the heap.
-    ///
-    /// # Returns
-    /// The gather_id if this task belongs to a gather (for sibling lookup).
-    pub fn fail_task(&mut self, task_id: TaskId, error: RunError) -> Option<HeapId> {
-        let task = self.get_task_mut(task_id);
-        let gather_id = task.gather_id;
-        task.state = TaskState::Failed(error);
-        gather_id
-    }
-
-    /// Cancels a task, cleaning up its resources.
-    ///
-    /// This marks the task as Failed with a cancellation error and cleans up:
-    /// - Stack values
-    /// - Exception stack values
-    /// - Frame cell references
-    /// - Frame namespaces
-    /// - Nested gathers (if the task was blocked on one)
-    /// - Completed task results (if task finished before cancellation)
-    ///
-    /// The caller is responsible for cleaning up the task's coroutine on the heap.
-    ///
-    /// # Arguments
-    /// * `task_id` - ID of the task to cancel
-    /// * `heap` - Heap for dropping values and cell cleanup
-    pub fn cancel_task(
-        &mut self,
-        task_id: TaskId,
-        heap: &mut crate::heap::Heap<impl crate::resource::ResourceTracker>,
-    ) {
-        // If task already finished, clean up its result value and return
-        if self.get_task(task_id).is_finished() {
-            let task = self.get_task_mut(task_id);
-            if let TaskState::Completed(value) = std::mem::replace(&mut task.state, TaskState::Ready) {
-                value.drop_with_heap(heap);
-            }
-            // Note: Failed tasks don't have values to clean up (RunError doesn't contain Values)
-            return;
-        }
-
-        // Remove from ready queue if present (do this before getting mutable task reference)
-        self.ready_queue.retain(|&id| id != task_id);
-
-        // Check if task is blocked on a gather and get the gather info before mutating task
-        let inner_gather_info = {
-            let task = self.get_task(task_id);
-            if let TaskState::BlockedOnGather(gather_id) = task.state {
-                // Get inner gather's task IDs from heap
-                if let crate::heap::HeapData::GatherFuture(gather) = heap.get(gather_id) {
-                    Some((gather_id, gather.task_ids.clone()))
-                } else {
-                    None
-                }
-            } else {
-                None
-            }
-        };
-
-        // Recursively cancel inner gather's tasks first
-        if let Some((inner_gather_id, inner_task_ids)) = inner_gather_info {
-            for inner_task_id in inner_task_ids {
-                self.cancel_task(inner_task_id, heap);
-            }
-
-            // Cleanup the inner GatherFuture - extract data first to avoid borrow conflict
-            let (items, results) = if let HeapDataMut::GatherFuture(gather) = heap.get_mut(inner_gather_id) {
-                (std::mem::take(&mut gather.items), std::mem::take(&mut gather.results))
-            } else {
-                (vec![], vec![])
-            };
-
-            // Now cleanup the extracted data with mutable heap access
-            for item in items {
-                if let crate::asyncio::GatherItem::Coroutine(coro_id) = item {
-                    heap.dec_ref(coro_id);
-                }
-            }
-            for value in results.into_iter().flatten() {
-                value.drop_with_heap(heap);
-            }
-
-            // Dec_ref the gather itself
-            heap.dec_ref(inner_gather_id);
-        }
-
-        // Now get mutable reference to the task for cleanup
-        let task = self.get_task_mut(task_id);
-
-        // Clean up stack values
-        for value in std::mem::take(&mut task.stack) {
-            value.drop_with_heap(heap);
-        }
-
-        // Clean up exception stack values
-        for value in std::mem::take(&mut task.exception_stack) {
-            value.drop_with_heap(heap);
-        }
-
-        // Restore this task's depth contribution before cleanup,
-        // since save_task_context subtracted it.
-        let task_depth = task.frames.len();
-        let global_depth = heap.get_recursion_depth();
-        heap.set_recursion_depth(global_depth + task_depth);
-        task.frames.clear();
-
-        // Mark as failed with a cancellation error
-        task.state = TaskState::Failed(
-            crate::exception_private::SimpleException::new_msg(
-                crate::exception_private::ExcType::RuntimeError,
-                "task was cancelled",
-            )
-            .into(),
-        );
-    }
-
-    /// Fails the task blocked on a specific CallId with an error.
-    ///
-    /// Used when an external function returns an error via `FutureSnapshot::resume`.
-    /// Uses `pending_calls` for O(1) lookup of the blocked task.
-    ///
-    /// # Returns
-    /// A tuple of (task_id, gather_id) if a task was found,
-    /// or None if no task was blocked on this CallId.
-    /// Callers should get siblings from `GatherFuture.task_ids` if gather_id is Some.
-    pub fn fail_for_call(&mut self, call_id: CallId, error: RunError) -> Option<(TaskId, Option<HeapId>)> {
-        // Get blocked task from pending_calls (O(1) lookup)
-        let task_id = self.pending_calls.remove(&call_id)?.creator_task;
-        let gather_id = self.fail_task(task_id, error);
-        Some((task_id, gather_id))
-    }
-
-    /// Returns the task that created a specific pending call.
-    ///
-    /// Used to check if a pending call's creator task has been cancelled.
-    #[inline]
-    pub fn get_pending_call_creator(&self, call_id: CallId) -> Option<TaskId> {
-        self.pending_calls.get(&call_id).map(|data| data.creator_task)
-    }
-
-    /// Returns true if a task has been cancelled or failed.
-    #[inline]
-    pub fn is_task_failed(&self, task_id: TaskId) -> bool {
-        matches!(self.tasks.get(task_id.raw() as usize), Some(task) if matches!(task.state, TaskState::Failed(_)))
-    }
-
-    /// Cleans up all scheduler resources: pending calls, resolved values, task
-    /// stacks/exception stacks, completed results, and task frame cell references.
-    ///
-    /// Each task's `recursion_depth` is restored to the global counter before
-    /// dropping cells, because `save_task_context` subtracted the recursion depth
-    /// and cleanup needs the correct depth to avoid underflow.
-    pub fn cleanup(&mut self, heap: &mut crate::heap::Heap<impl crate::resource::ResourceTracker>) {
-        // Drop pending call arguments
-        for (_, data) in std::mem::take(&mut self.pending_calls) {
-            data.args.drop_with_heap(heap);
-        }
-        // Drop resolved values
-        for (_, value) in std::mem::take(&mut self.resolved) {
-            value.drop_with_heap(heap);
-        }
-        // Drop task stack/exception values and completed results
-        for task in &mut self.tasks {
-            for value in std::mem::take(&mut task.stack) {
-                value.drop_with_heap(heap);
-            }
-            for value in std::mem::take(&mut task.exception_stack) {
-                value.drop_with_heap(heap);
-            }
-            if let TaskState::Completed(value) = std::mem::replace(&mut task.state, TaskState::Ready) {
-                value.drop_with_heap(heap);
-            }
-            // Restore recursion depth and clear frames
-            let task_depth = task.frames.len();
-            let global_depth = heap.get_recursion_depth();
-            heap.set_recursion_depth(global_depth + task_depth);
-            task.frames.clear();
-        }
-    }
-}
-
-impl Default for Scheduler {
-    fn default() -> Self {
-        Self::new()
-    }
-}