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}