Skip to main content

vibe_code/
vibe.rs

1//! The user-facing "vibe" interface.
2//!
3//! This module provides the dead-simple API for users who just want to run
4//! code in parallel without thinking about the underlying complexity. It
5//! defines the `VibeSystem`, the `Job` handle, and the `collect` utility.
6
7use crate::task::{Priority, TaskHandle};
8use crate::utils::timer_init;
9use crate::utils::{VibeRng, elapsed_ns};
10use crate::vibe_code::UltraVibeSystem;
11use std::sync::{Arc, mpsc};
12thread_local! {
13    static LOCAL_RNG: std::cell::RefCell<VibeRng> = std::cell::RefCell::new(VibeRng::new(elapsed_ns()));
14}
15
16/// A handle to a job that is running in the background.
17pub struct Job<T> {
18    handle: TaskHandle<T>,
19}
20
21impl<T> Job<T> {
22    /// Waits for the job to finish and returns its result.
23    ///
24    /// # Panics
25    /// Panics if the job failed (e.g., the user's function panicked) or if the
26    /// system was shut down before the job could complete.
27    pub fn get(self) -> T {
28        match self.handle.recv_result() {
29            Ok(Ok(result)) => result,
30            Ok(Err(_)) => panic!("❌ Your job failed! Check your function for bugs."),
31            Err(_) => panic!("❌ Job was cancelled - did you shut down the system?"),
32        }
33    }
34
35    /// Checks if the job is finished without blocking.
36    ///
37    /// Returns `Some(result)` if the job is done, or `None` if it's still running.
38    ///
39    /// # Panics
40    /// Panics if the job failed.
41    pub fn peek(self) -> Option<T> {
42        match self.handle.try_recv_result() {
43            Ok(Ok(result)) => Some(result),
44            Ok(Err(_)) => panic!("❌ Your job failed! Check your function for bugs."),
45            Err(mpsc::TryRecvError::Empty) => None,
46            Err(mpsc::TryRecvError::Disconnected) => panic!("❌ Job was cancelled"),
47        }
48    }
49
50    /// Returns `true` if the job has finished (either completed or failed).
51    pub fn is_done(&self) -> bool {
52        matches!(
53            self.handle.try_recv_result(),
54            Ok(_) | Err(mpsc::TryRecvError::Disconnected)
55        )
56    }
57}
58
59/// The main entry point for the simple parallel execution system.
60pub struct VibeSystem {
61    /// A shared pointer to the internal task scheduling and execution engine.
62    inner: Arc<UltraVibeSystem>,
63}
64
65impl VibeSystem {
66    /// Creates a new `VibeSystem` with sensible default settings.
67    ///
68    /// This initializes the background worker pools and load balancer.
69    pub fn new() -> Self {
70        let _ = timer_init();
71        Self {
72            inner: Arc::new(
73                UltraVibeSystem::builder()
74                    .with_nodes(80)
75                    .with_super_nodes(40)
76                    .build(),
77            ),
78        }
79    }
80
81    /// Runs a function with input data in parallel.
82    ///
83    /// This method takes a function and its input data, submits it to the system
84    /// for background execution, and immediately returns a `Job` handle.
85    ///
86    /// # Example
87    /// `let job = system.run(process_data, my_data);`
88    pub fn run<F, T, R>(&self, func: F, data: T) -> Job<R>
89    where
90        F: FnOnce(T) -> R + Send + 'static,
91        T: Send + 'static,
92        R: Send + 'static,
93    {
94        LOCAL_RNG.with(|rng| {
95            let work = move || Ok(func(data));
96
97            match self
98                .inner
99                .submit_cpu_task(Priority::Normal, 10, work, &mut rng.borrow_mut())
100            {
101                Ok(handle) => Job { handle },
102                Err(_) => panic!("🔥 System overloaded! Too many jobs running at once."),
103            }
104        })
105    }
106
107    /// Runs a function with no input data in parallel.
108    ///
109    /// This is a convenience wrapper around `run` for functions that take no arguments.
110    ///
111    /// # Example
112    /// `let job = system.go(|| expensive_calculation());`
113    pub fn go<F, R>(&self, func: F) -> Job<R>
114    where
115        F: FnOnce() -> R + Send + 'static,
116        R: Send + 'static,
117    {
118        self.run(|_| func(), ())
119    }
120}
121
122impl Default for VibeSystem {
123    fn default() -> Self {
124        Self::new()
125    }
126}
127
128/// Waits for a vector of jobs to finish and collects their results in order.
129pub fn collect<T>(jobs: Vec<Job<T>>) -> Vec<T> {
130    jobs.into_iter().map(|job| job.get()).collect()
131}