1use stageleft::{QuotedWithContext, q};
18
19#[cfg(stageleft_runtime)]
20use super::dynamic::DynLocation;
21use super::{Location, LocationId};
22use crate::compile::builder::{ClockId, FlowState};
23use crate::compile::ir::{HydroNode, HydroSource};
24#[cfg(stageleft_runtime)]
25use crate::forward_handle::{CycleCollection, CycleCollectionWithInitial};
26use crate::forward_handle::{TickCycle, TickCycleHandle};
27#[cfg(feature = "tokio")]
28use crate::live_collections::Singleton;
29use crate::live_collections::boundedness::Bounded;
30use crate::live_collections::optional::Optional;
31use crate::live_collections::stream::{ExactlyOnce, Stream, TotalOrder};
32use crate::location::TopLevel;
33#[cfg(feature = "tokio")]
34use crate::nondet::NonDet;
35use crate::nondet::nondet;
36
37#[derive(Clone)]
48pub struct Atomic<Loc> {
49 pub(crate) tick: Tick<Loc>,
50}
51
52impl<L: DynLocation> DynLocation for Atomic<L> {
53 fn dyn_id(&self) -> LocationId {
54 LocationId::Atomic(Box::new(self.tick.dyn_id()))
55 }
56
57 fn flow_state(&self) -> &FlowState {
58 self.tick.flow_state()
59 }
60
61 fn is_top_level() -> bool {
62 L::is_top_level()
63 }
64
65 fn multiversioned(&self) -> bool {
66 self.tick.multiversioned()
67 }
68
69 fn cluster_consistency() -> Option<super::dynamic::ClusterConsistency> {
70 L::cluster_consistency()
71 }
72}
73
74impl<'a, L> Location<'a> for Atomic<L>
75where
76 L: Location<'a>,
77{
78 type Root = L::Root;
79
80 type DropConsistency = Atomic<L::DropConsistency>;
81
82 fn consistency() -> Option<super::dynamic::ClusterConsistency> {
83 L::consistency()
84 }
85
86 fn root(&self) -> Self::Root {
87 self.tick.root()
88 }
89
90 fn drop_consistency(&self) -> Self::DropConsistency {
91 Atomic {
92 tick: self.tick.drop_consistency(),
93 }
94 }
95
96 fn from_drop_consistency(l2: Self::DropConsistency) -> Self {
97 Atomic {
98 tick: Tick::from_drop_consistency(l2.tick),
99 }
100 }
101}
102
103pub trait DeferTick {
110 fn defer_tick(self) -> Self;
112}
113
114#[derive(Clone)]
116pub struct Tick<L> {
117 pub(crate) id: ClockId,
118 pub(crate) l: L,
120}
121
122impl<L: DynLocation> DynLocation for Tick<L> {
123 fn dyn_id(&self) -> LocationId {
124 LocationId::Tick(self.id, Box::new(self.l.dyn_id()))
125 }
126
127 fn flow_state(&self) -> &FlowState {
128 self.l.flow_state()
129 }
130
131 fn is_top_level() -> bool {
132 false
133 }
134
135 fn multiversioned(&self) -> bool {
136 self.l.multiversioned()
137 }
138
139 fn cluster_consistency() -> Option<super::dynamic::ClusterConsistency> {
140 L::cluster_consistency()
141 }
142}
143
144impl<'a, L> Location<'a> for Tick<L>
145where
146 L: Location<'a>,
147{
148 type Root = L::Root;
149
150 type DropConsistency = Tick<L::DropConsistency>;
151
152 fn consistency() -> Option<super::dynamic::ClusterConsistency> {
153 L::consistency()
154 }
155
156 fn root(&self) -> Self::Root {
157 self.l.root()
158 }
159
160 fn drop_consistency(&self) -> Self::DropConsistency {
161 Tick {
162 id: self.id,
163 l: self.l.drop_consistency(),
164 }
165 }
166
167 fn from_drop_consistency(l2: Self::DropConsistency) -> Self {
168 Tick {
169 id: l2.id,
170 l: L::from_drop_consistency(l2.l),
171 }
172 }
173}
174
175impl<'a, L> Tick<L>
176where
177 L: Location<'a>,
178{
179 pub fn outer(&self) -> &L {
184 &self.l
185 }
186
187 pub fn spin_batch(
193 &self,
194 batch_size: impl QuotedWithContext<
195 'a,
196 usize,
197 crate::live_collections::OperatorContext<
198 L,
199 crate::live_collections::boundedness::Unbounded,
200 >,
201 > + Copy
202 + 'a,
203 ) -> Stream<(), Self, Bounded, TotalOrder, ExactlyOnce>
204 where
205 L: TopLevel<'a>,
206 {
207 let out = self
208 .l
209 .spin()
210 .flat_map_ordered(q!(move |_| 0..batch_size))
211 .map(q!(|_| ()));
212
213 let inner = out.batch(self, nondet!());
214 Stream::new(self.clone(), inner.ir_node.replace(HydroNode::Placeholder))
215 }
216
217 pub fn none<T>(&self) -> Optional<T, Self, Bounded> {
236 let e = q!([]);
237 let e = QuotedWithContext::<'a, [(); 0], Self>::splice_typed_ctx(e, self);
238
239 let unit_optional: Optional<(), Self, Bounded> = Optional::new(
240 self.clone(),
241 HydroNode::Source {
242 source: HydroSource::Iter(e.into()),
243 metadata: self.new_node_metadata(Optional::<(), Self, Bounded>::collection_kind()),
244 },
245 );
246
247 unit_optional.map(q!(|_| unreachable!())) }
249
250 pub fn optional_first_tick<T: Clone>(
276 &self,
277 e: impl QuotedWithContext<'a, T, Tick<L>>,
278 ) -> Optional<T, Self, Bounded> {
279 let e = e.splice_untyped_ctx(self);
280
281 Optional::new(
282 self.clone(),
283 HydroNode::SingletonSource {
284 value: e.into(),
285 first_tick_only: true,
286 metadata: self.new_node_metadata(Optional::<T, Self, Bounded>::collection_kind()),
287 },
288 )
289 }
290
291 #[cfg(feature = "tokio")]
299 pub fn current_tick_instant(
300 &self,
301 _nondet: NonDet,
302 ) -> Singleton<tokio::time::Instant, Tick<L::DropConsistency>, Bounded>
303 where
304 Self: Sized,
305 {
306 self.singleton(q!(tokio::time::Instant::now()))
308 }
309
310 #[expect(
319 private_bounds,
320 reason = "only Hydro collections can implement ReceiverComplete"
321 )]
322 pub fn cycle<S, L2: Location<'a, DropConsistency = Tick<L::DropConsistency>>>(
323 &self,
324 ) -> (TickCycleHandle<'a, S>, S)
325 where
326 S: CycleCollection<'a, TickCycle, Location = L2> + DeferTick,
327 {
328 let cycle_id = self.flow_state().borrow_mut().next_cycle_id();
329 (
330 TickCycleHandle::new(cycle_id, Location::id(self)),
331 S::create_source(cycle_id, self.clone().with_consistency_of()).defer_tick(),
332 )
333 }
334
335 #[expect(
342 private_bounds,
343 reason = "only Hydro collections can implement ReceiverComplete"
344 )]
345 pub fn cycle_with_initial<S, L2: Location<'a, DropConsistency = Tick<L::DropConsistency>>>(
346 &self,
347 initial: S,
348 ) -> (TickCycleHandle<'a, S>, S)
349 where
350 S: CycleCollectionWithInitial<'a, TickCycle, Location = L2>,
351 {
352 let cycle_id = self.flow_state().borrow_mut().next_cycle_id();
353 (
354 TickCycleHandle::new(cycle_id, Location::id(self)),
355 S::create_source_with_initial(cycle_id, initial, self.clone().with_consistency_of()),
357 )
358 }
359}
360
361#[cfg(test)]
362mod tests {
363 #[cfg(feature = "sim")]
364 use stageleft::q;
365
366 #[cfg(feature = "sim")]
367 use crate::live_collections::sliced::sliced;
368 #[cfg(feature = "sim")]
369 use crate::location::Location;
370 #[cfg(feature = "sim")]
371 use crate::nondet::nondet;
372 #[cfg(feature = "sim")]
373 use crate::prelude::FlowBuilder;
374
375 #[cfg(feature = "sim")]
376 #[test]
377 fn sim_atomic_stream() {
378 let mut flow = FlowBuilder::new();
379 let node = flow.process::<()>();
380
381 let (write_send, write_req) = node.sim_input();
382 let (read_send, read_req) = node.sim_input::<(), _, _>();
383
384 let atomic_write = write_req.atomic();
385 let current_state = atomic_write.clone().fold(
386 q!(|| 0),
387 q!(|state: &mut i32, v: i32| {
388 *state += v;
389 }),
390 );
391
392 let write_ack_recv = atomic_write.end_atomic().sim_output();
393 let read_response_recv = sliced! {
394 let batch_of_req = use::batch(read_req, nondet!());
395 let latest_singleton = use::atomic(current_state, nondet!());
396 batch_of_req.cross_singleton(latest_singleton)
397 }
398 .sim_output();
399
400 let sim_compiled = flow.sim().compiled();
401 let instances = sim_compiled.exhaustive(async || {
402 write_send.send(1);
403 write_ack_recv.assert_yields([1]).await;
404 read_send.send(());
405 assert!(read_response_recv.next().await.1 >= 1);
406 });
407
408 assert_eq!(instances, 1);
409
410 let instances_read_before_write = sim_compiled.exhaustive(async || {
411 write_send.send(1);
412 read_send.send(());
413 write_ack_recv.assert_yields([1]).await;
414 let _ = read_response_recv.next().await;
415 });
416
417 assert_eq!(instances_read_before_write, 3); }
419
420 #[cfg(feature = "sim")]
421 #[test]
422 #[should_panic]
423 fn sim_non_atomic_stream() {
424 let mut flow = FlowBuilder::new();
426 let node = flow.process::<()>();
427
428 let (write_send, write_req) = node.sim_input();
429 let (read_send, read_req) = node.sim_input::<(), _, _>();
430
431 let current_state = write_req.clone().fold(
432 q!(|| 0),
433 q!(|state: &mut i32, v: i32| {
434 *state += v;
435 }),
436 );
437
438 let write_ack_recv = write_req.sim_output();
439
440 let read_response_recv = sliced! {
441 let batch_of_req = use::batch(read_req, nondet!());
442 let latest_singleton = use::snapshot(current_state, nondet!());
443 batch_of_req.cross_singleton(latest_singleton)
444 }
445 .sim_output();
446
447 flow.sim().exhaustive(async || {
448 write_send.send(1);
449 write_ack_recv.assert_yields([1]).await;
450 read_send.send(());
451
452 let (_, v) = read_response_recv.next().await;
453 assert_eq!(v, 1);
454 });
455 }
456}