1use std::fmt;
11use std::hash::Hash;
12use std::marker::PhantomData;
13use std::sync::Arc;
14use std::time::Duration;
15
16use crossbeam_channel::{Receiver, Sender};
17use key_paths_core::{FieldDiff, hash_value, RefKpTrait};
18
19use crate::effect::EffectId;
20use crate::optics::Casepath;
21use crate::runtime::store::{ChangeSet, StoreHub, StoreTask, StoreWork};
22
23#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum TeaStoreError {
26 RuntimeShutdown,
27 SubscribeTimeout,
28 SnapshotTimeout,
29 SnapshotUnavailable,
30}
31
32impl fmt::Display for TeaStoreError {
33 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34 match self {
35 Self::RuntimeShutdown => write!(f, "tea runtime control channel disconnected"),
36 Self::SubscribeTimeout => write!(f, "timed out waiting for snapshot subscription"),
37 Self::SnapshotTimeout => write!(f, "timed out waiting for model snapshot"),
38 Self::SnapshotUnavailable => write!(f, "model snapshot unavailable"),
39 }
40 }
41}
42
43impl std::error::Error for TeaStoreError {}
44
45const DEFAULT_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(5);
46
47pub(crate) enum TeaControl<S> {
49 Subscribe {
50 snapshot_tx: Sender<Arc<S>>,
51 ready_tx: Sender<()>,
52 },
53 RequestSnapshot {
54 reply: Sender<Arc<S>>,
55 },
56}
57
58pub(crate) struct TeaStoreBackend<S, M> {
59 pub hub: Arc<StoreHub<S, M>>,
60 pub control_tx: Sender<TeaControl<S>>,
61}
62
63impl<S, M> Clone for TeaStoreBackend<S, M>
64where
65 S: Send + 'static,
66 M: Send + 'static,
67{
68 fn clone(&self) -> Self {
69 Self {
70 hub: self.hub.clone(),
71 control_tx: self.control_tx.clone(),
72 }
73 }
74}
75
76impl<S, M> StoreWork<S, M> for TeaStoreBackend<S, M>
77where
78 S: Send + 'static,
79 M: Send + 'static,
80{
81 fn hub(&self) -> &Arc<StoreHub<S, M>> {
82 &self.hub
83 }
84
85 fn end_store_work(&self) {
87 let prev = self.hub.in_flight.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
88 if prev == 1 {
89 self.hub.signal_effect_done();
90 }
91 }
92}
93
94impl<S, M> TeaStoreBackend<S, M>
95where
96 S: Send + 'static,
97 M: Send + 'static,
98{
99 pub fn new(hub: Arc<StoreHub<S, M>>, control_tx: Sender<TeaControl<S>>) -> Self {
100 Self { hub, control_tx }
101 }
102
103 pub fn tea_store(&self) -> TeaStore<S, M> {
104 TeaStore {
105 backend: self.clone(),
106 }
107 }
108
109 pub fn view_store(&self) -> TeaViewStore<S, M> {
110 TeaViewStore {
111 backend: self.clone(),
112 }
113 }
114
115 fn request_snapshot(&self, timeout: Duration) -> Result<Arc<S>, TeaStoreError> {
116 let (reply_tx, reply_rx) = crossbeam_channel::bounded(1);
117 self.control_tx
118 .send(TeaControl::RequestSnapshot { reply: reply_tx })
119 .map_err(|_| TeaStoreError::RuntimeShutdown)?;
120 reply_rx
121 .recv_timeout(timeout)
122 .map_err(|_| TeaStoreError::SnapshotTimeout)
123 }
124}
125
126pub struct TeaStore<S, M> {
128 pub(crate) backend: TeaStoreBackend<S, M>,
129}
130
131impl<S, M> Clone for TeaStore<S, M>
132where
133 S: Send + 'static,
134 M: Send + 'static,
135{
136 fn clone(&self) -> Self {
137 Self {
138 backend: self.backend.clone(),
139 }
140 }
141}
142
143impl<S, M> TeaStore<S, M>
144where
145 S: Send + Sync + Clone + 'static,
146 M: Send + 'static,
147{
148 pub fn send(&self, action: M) -> StoreTask {
149 let (tx, rx) = crossbeam_channel::bounded(1);
150 self.backend.hub.effect_done.lock().push_back(tx);
151 self.backend.begin_store_work();
152 let _ = self.backend.hub.sender.send_blocking(action);
153 StoreTask::from_rx(rx)
154 }
155
156 pub fn dispatch(&self, action: M) {
157 let _ = self.send(action);
158 }
159
160 pub fn try_state(&self) -> Result<S, TeaStoreError> {
162 self.backend
163 .request_snapshot(DEFAULT_SNAPSHOT_TIMEOUT)
164 .map(|s| (*s).clone())
165 }
166
167 pub fn state(&self) -> Result<S, TeaStoreError> {
169 self.try_state()
170 }
171
172 pub fn view_store(&self) -> TeaViewStore<S, M> {
174 self.backend.view_store()
175 }
176
177 pub fn cancel(&self, id: EffectId) {
178 if let Some(handle) = self
179 .backend
180 .hub
181 .interpreter
182 .cancel_tokens
183 .lock()
184 .remove(&id)
185 {
186 handle.abort();
187 }
188 }
189
190 pub fn try_subscribe_state(&self) -> Result<TeaStateSubscriber<S>, TeaStoreError>
191 where
192 S: PartialEq,
193 {
194 self.register_snapshot_subscriber()
195 }
196
197 pub fn subscribe_state(&self) -> Result<TeaStateSubscriber<S>, TeaStoreError>
199 where
200 S: PartialEq,
201 {
202 self.try_subscribe_state()
203 }
204
205 pub fn try_subscribe_changes(&self) -> Result<TeaChangeSubscriber<S>, TeaStoreError>
206 where
207 S: FieldDiff + Hash + PartialEq,
208 {
209 let mut snapshot_sub = self.register_snapshot_subscriber()?;
210 let initial = snapshot_sub
211 .latest()
212 .or_else(|| snapshot_sub.wait_next(DEFAULT_SNAPSHOT_TIMEOUT));
213 let Some(initial) = initial else {
214 return Err(TeaStoreError::SnapshotUnavailable);
215 };
216 let (last_hash, last_fields) = snapshot_hashes(&initial);
217 Ok(TeaChangeSubscriber {
218 snapshot_sub,
219 last_hash: Some(last_hash),
220 last_fields,
221 _marker: PhantomData,
222 })
223 }
224
225 pub fn subscribe_changes(&self) -> Result<TeaChangeSubscriber<S>, TeaStoreError>
227 where
228 S: FieldDiff + Hash + PartialEq,
229 {
230 self.try_subscribe_changes()
231 }
232
233 pub fn scope<CS, CM, AK, SK>(
234 &self,
235 state_kp: SK,
236 action_kp: AK,
237 ) -> ScopedTeaStore<S, M, CS, CM, AK, SK>
238 where
239 CS: Clone + PartialEq + Send + Sync + 'static,
240 CM: Clone + Send + 'static,
241 S: 'static,
242 M: 'static,
243 AK: Casepath<M, CM> + Clone + Send + Sync + 'static,
244 SK: RefKpTrait<S, CS> + Clone,
245 {
246 ScopedTeaStore {
247 store: self.clone(),
248 state_kp,
249 action_kp,
250 _marker: PhantomData,
251 }
252 }
253
254 fn register_snapshot_subscriber(&self) -> Result<TeaStateSubscriber<S>, TeaStoreError> {
255 let (snapshot_tx, snapshot_rx) = crossbeam_channel::unbounded();
256 let (ready_tx, ready_rx) = crossbeam_channel::bounded(1);
257 self.backend
258 .control_tx
259 .send(TeaControl::Subscribe {
260 snapshot_tx,
261 ready_tx,
262 })
263 .map_err(|_| TeaStoreError::RuntimeShutdown)?;
264 ready_rx
265 .recv_timeout(DEFAULT_SNAPSHOT_TIMEOUT)
266 .map_err(|_| TeaStoreError::SubscribeTimeout)?;
267 let initial = snapshot_rx
268 .recv_timeout(DEFAULT_SNAPSHOT_TIMEOUT)
269 .map_err(|_| TeaStoreError::SnapshotUnavailable)?;
270 Ok(TeaStateSubscriber {
271 rx: snapshot_rx,
272 last: Some(initial),
273 })
274 }
275}
276
277pub struct TeaViewStore<S, M> {
279 pub(crate) backend: TeaStoreBackend<S, M>,
280}
281
282impl<S, M> Clone for TeaViewStore<S, M>
283where
284 S: Send + 'static,
285 M: Send + 'static,
286{
287 fn clone(&self) -> Self {
288 Self {
289 backend: self.backend.clone(),
290 }
291 }
292}
293
294impl<S, M> TeaViewStore<S, M>
295where
296 S: Send + Sync + Clone + 'static,
297 M: Send + 'static,
298{
299 pub fn try_with_snapshot<R>(&self, f: impl FnOnce(&S) -> R) -> Result<R, TeaStoreError> {
300 let snap = self
301 .backend
302 .request_snapshot(DEFAULT_SNAPSHOT_TIMEOUT)?;
303 Ok(f(snap.as_ref()))
304 }
305
306 pub fn with_snapshot<R>(&self, f: impl FnOnce(&S) -> R) -> Result<R, TeaStoreError> {
307 self.try_with_snapshot(f)
308 }
309
310 pub fn try_load(&self) -> Result<Arc<S>, TeaStoreError> {
311 self.backend.request_snapshot(DEFAULT_SNAPSHOT_TIMEOUT)
312 }
313
314 pub fn load(&self) -> Result<Arc<S>, TeaStoreError> {
315 self.try_load()
316 }
317
318 pub fn try_state(&self) -> Result<S, TeaStoreError> {
319 self.try_with_snapshot(|s| s.clone())
320 }
321
322 pub fn state(&self) -> Result<S, TeaStoreError> {
323 self.try_state()
324 }
325
326 pub fn subscribe_state(&self) -> Result<TeaStateSubscriber<S>, TeaStoreError>
327 where
328 S: PartialEq,
329 {
330 self.backend.tea_store().subscribe_state()
331 }
332
333 pub fn subscribe_changes(&self) -> Result<TeaChangeSubscriber<S>, TeaStoreError>
334 where
335 S: FieldDiff + Hash + PartialEq,
336 {
337 self.backend.tea_store().subscribe_changes()
338 }
339}
340
341pub struct TeaStateSubscriber<S> {
343 rx: Receiver<Arc<S>>,
344 last: Option<Arc<S>>,
345}
346
347impl<S: PartialEq> TeaStateSubscriber<S> {
348 #[allow(clippy::should_implement_trait)]
349 pub fn next(&mut self) -> Option<Arc<S>> {
350 loop {
351 match self.rx.try_recv() {
352 Ok(first) => {
353 let mut snapshot = first;
354 while let Ok(newer) = self.rx.try_recv() {
355 snapshot = newer;
356 }
357 if self
358 .last
359 .as_ref()
360 .is_some_and(|prev| prev.as_ref() == snapshot.as_ref())
361 {
362 continue;
363 }
364 self.last = Some(snapshot.clone());
365 return Some(snapshot);
366 }
367 Err(crossbeam_channel::TryRecvError::Empty) => return None,
368 Err(crossbeam_channel::TryRecvError::Disconnected) => return None,
369 }
370 }
371 }
372
373 pub fn wait_next(&mut self, timeout: Duration) -> Option<Arc<S>> {
374 let deadline = std::time::Instant::now() + timeout;
375 while std::time::Instant::now() < deadline {
376 if let Some(s) = self.next() {
377 return Some(s);
378 }
379 std::thread::sleep(Duration::from_millis(5));
380 }
381 None
382 }
383
384 pub fn latest(&self) -> Option<Arc<S>> {
385 self.last.clone()
386 }
387
388 pub fn with_latest<R>(&self, f: impl FnOnce(&S) -> R) -> Option<R> {
389 self.last.as_ref().map(|s| f(s.as_ref()))
390 }
391}
392
393pub struct TeaChangeSubscriber<S>
394where
395 S: FieldDiff + Hash,
396{
397 snapshot_sub: TeaStateSubscriber<S>,
398 last_hash: Option<u64>,
399 last_fields: Vec<(S::Path, u64)>,
400 _marker: PhantomData<S>,
401}
402
403fn snapshot_hashes<S: FieldDiff + Hash>(snap: &Arc<S>) -> (u64, Vec<(S::Path, u64)>) {
404 let root = hash_value(snap.as_ref());
405 let mut fields = Vec::new();
406 snap.field_hashes(&mut fields);
407 (root, fields)
408}
409
410impl<S> TeaChangeSubscriber<S>
411where
412 S: FieldDiff + Hash + PartialEq,
413{
414 fn diff_fields(&self, fields: &[(S::Path, u64)]) -> Vec<S::Path> {
415 fields
416 .iter()
417 .filter(|(path, hash)| {
418 self.last_fields
419 .iter()
420 .find(|(prev_path, _)| prev_path == path)
421 .is_none_or(|(_, prev_hash)| prev_hash != hash)
422 })
423 .map(|(path, _)| *path)
424 .collect()
425 }
426
427 #[allow(clippy::should_implement_trait)]
428 pub fn next(&mut self) -> Option<ChangeSet<S::Path>> {
429 loop {
430 let snap = self.snapshot_sub.next()?;
431 let (root_hash, fields) = snapshot_hashes(&snap);
432 if Some(root_hash) == self.last_hash {
433 continue;
434 }
435 let paths = self.diff_fields(&fields);
436 self.last_hash = Some(root_hash);
437 self.last_fields = fields;
438 if paths.is_empty() {
439 continue;
440 }
441 return Some(ChangeSet { paths });
442 }
443 }
444
445 pub fn wait_next(&mut self, timeout: Duration) -> Option<ChangeSet<S::Path>> {
446 let deadline = std::time::Instant::now() + timeout;
447 while std::time::Instant::now() < deadline {
448 if let Some(change) = self.next() {
449 return Some(change);
450 }
451 std::thread::sleep(Duration::from_millis(5));
452 }
453 None
454 }
455
456 pub fn latest_snapshot(&self) -> Option<Arc<S>> {
457 self.snapshot_sub.latest()
458 }
459}
460
461pub struct ScopedTeaStore<S: 'static, M: 'static, CS: 'static, CM: 'static, AK: 'static, SK>
463where
464 SK: RefKpTrait<S, CS> + Clone,
465{
466 store: TeaStore<S, M>,
467 state_kp: SK,
468 action_kp: AK,
469 _marker: PhantomData<(M, CM, S, CS)>,
470}
471
472impl<S: 'static, M: 'static, CS: 'static, CM: 'static, AK: 'static, SK> Clone
473 for ScopedTeaStore<S, M, CS, CM, AK, SK>
474where
475 S: Send,
476 M: Send,
477 AK: Clone,
478 SK: RefKpTrait<S, CS> + Clone,
479{
480 fn clone(&self) -> Self {
481 Self {
482 store: self.store.clone(),
483 state_kp: self.state_kp.clone(),
484 action_kp: self.action_kp.clone(),
485 _marker: PhantomData,
486 }
487 }
488}
489
490impl<S: 'static, M: 'static, CS: 'static, CM: 'static, AK: 'static, SK>
491 ScopedTeaStore<S, M, CS, CM, AK, SK>
492where
493 S: Send + Sync + Clone + PartialEq,
494 M: Send,
495 CS: Clone + PartialEq + Send + Sync,
496 CM: Clone + Send,
497 AK: Casepath<M, CM> + Clone + Send + Sync + 'static,
498 SK: RefKpTrait<S, CS> + Clone,
499{
500 pub fn send(&self, action: CM) -> StoreTask {
501 self.store.send(self.action_kp.wrap(action))
502 }
503
504 pub fn dispatch(&self, action: CM) {
505 self.store.dispatch(self.action_kp.wrap(action));
506 }
507
508 pub fn child_state(&self) -> Result<Option<CS>, TeaStoreError> {
509 let parent = self.store.state()?;
510 Ok(self.state_kp.focus(&parent).cloned())
511 }
512
513 pub fn subscribe_state(&self) -> Result<ScopedTeaStateSubscriber<S, CS, SK>, TeaStoreError>
514 where
515 S: PartialEq,
516 SK: Clone,
517 {
518 Ok(ScopedTeaStateSubscriber {
519 inner: self.store.subscribe_state()?,
520 state_kp: self.state_kp.clone(),
521 _marker: PhantomData,
522 })
523 }
524}
525
526pub struct ScopedTeaStateSubscriber<S: 'static, CS: 'static, SK>
527where
528 SK: RefKpTrait<S, CS> + Clone,
529{
530 inner: TeaStateSubscriber<S>,
531 state_kp: SK,
532 _marker: PhantomData<(S, CS)>,
533}
534
535impl<S, CS, SK> ScopedTeaStateSubscriber<S, CS, SK>
536where
537 S: PartialEq,
538 CS: Clone + PartialEq,
539 SK: RefKpTrait<S, CS> + Clone,
540{
541 #[allow(clippy::should_implement_trait)]
542 pub fn next(&mut self) -> Option<CS> {
543 loop {
544 let parent = self.inner.next()?;
545 let Some(child) = self.state_kp.focus(parent.as_ref()).cloned() else {
546 continue;
547 };
548 return Some(child);
549 }
550 }
551
552 pub fn wait_next(&mut self, timeout: Duration) -> Option<CS> {
553 let deadline = std::time::Instant::now() + timeout;
554 while std::time::Instant::now() < deadline {
555 if let Some(s) = self.next() {
556 return Some(s);
557 }
558 std::thread::sleep(Duration::from_millis(5));
559 }
560 None
561 }
562}
563
564#[cfg(test)]
565mod tests {
566 use super::*;
567
568 #[test]
569 fn tea_change_subscriber_diffs_snapshots() {
570 #[derive(Clone, Hash, PartialEq, Debug)]
571 struct App {
572 n: i32,
573 }
574
575 impl FieldDiff for App {
576 type Path = &'static str;
577 fn field_hashes(&self, out: &mut Vec<(Self::Path, u64)>) {
578 out.push(("n", hash_value(&self.n)));
579 }
580 }
581
582 let (tx, rx) = crossbeam_channel::unbounded();
583 let mut sub = TeaChangeSubscriber {
584 snapshot_sub: TeaStateSubscriber {
585 rx,
586 last: Some(Arc::new(App { n: 0 })),
587 },
588 last_hash: Some(hash_value(&App { n: 0 })),
589 last_fields: vec![("n", hash_value(&0))],
590 _marker: PhantomData,
591 };
592 tx.send(Arc::new(App { n: 1 })).ok();
593 let change = sub.next();
594 assert!(change.is_some());
595 if let Some(change) = change {
596 assert_eq!(change.paths, vec!["n"]);
597 }
598 }
599}