hydro_lang/location/
dynamic.rs

1//! Definitions for interacting with locations using an untyped interface.
2//!
3//! Under the hood, locations are associated with a [`LocationId`] value that
4//! uniquely identifies the location. Manipulating these values is useful for
5//! observability and transforming the Hydro IR.
6
7use serde::{Deserialize, Serialize};
8
9#[cfg(stageleft_runtime)]
10use crate::compile::{builder::FlowState, ir::HydroIrMetadata};
11
12#[expect(missing_docs, reason = "TODO")]
13#[derive(PartialEq, Eq, Clone, Debug, Hash, Serialize, Deserialize)]
14pub enum LocationId {
15    Process(usize),
16    Cluster(usize),
17    Atomic(
18        /// The tick that the atomic region is associated with.
19        Box<LocationId>,
20    ),
21    Tick(usize, Box<LocationId>),
22}
23
24#[expect(missing_docs, reason = "TODO")]
25impl LocationId {
26    pub fn root(&self) -> &LocationId {
27        match self {
28            LocationId::Process(_) => self,
29            LocationId::Cluster(_) => self,
30            LocationId::Atomic(tick) => tick.root(),
31            LocationId::Tick(_, id) => id.root(),
32        }
33    }
34
35    pub fn is_root(&self) -> bool {
36        match self {
37            LocationId::Process(_) | LocationId::Cluster(_) => true,
38            LocationId::Atomic(_) => false,
39            LocationId::Tick(_, _) => false,
40        }
41    }
42
43    pub fn is_top_level(&self) -> bool {
44        match self {
45            LocationId::Process(_) | LocationId::Cluster(_) => true,
46            LocationId::Atomic(_) => true,
47            LocationId::Tick(_, _) => false,
48        }
49    }
50
51    pub fn raw_id(&self) -> usize {
52        match self {
53            LocationId::Process(id) => *id,
54            LocationId::Cluster(id) => *id,
55            LocationId::Atomic(_) => panic!("cannot get raw id for atomic"),
56            LocationId::Tick(_, _) => panic!("cannot get raw id for tick"),
57        }
58    }
59
60    pub fn swap_root(&mut self, new_root: LocationId) {
61        match self {
62            LocationId::Tick(_, id) => {
63                id.swap_root(new_root);
64            }
65            LocationId::Atomic(tick) => {
66                tick.swap_root(new_root);
67            }
68            _ => {
69                assert!(new_root.is_root());
70                *self = new_root;
71            }
72        }
73    }
74}
75
76#[cfg(stageleft_runtime)]
77pub(crate) trait DynLocation: Clone {
78    fn id(&self) -> LocationId;
79
80    fn flow_state(&self) -> &FlowState;
81    fn is_top_level() -> bool;
82
83    fn new_node_metadata<T>(&self) -> HydroIrMetadata {
84        use stageleft::quote_type;
85
86        use crate::compile::ir::HydroIrOpMetadata;
87        use crate::compile::ir::backtrace::Backtrace;
88
89        HydroIrMetadata {
90            location_kind: self.id(),
91            output_type: Some(quote_type::<T>().into()),
92            cardinality: None,
93            tag: None,
94            op: HydroIrOpMetadata {
95                backtrace: Backtrace::get_backtrace(2),
96                cpu_usage: None,
97                network_recv_cpu_usage: None,
98                id: None,
99            },
100        }
101    }
102}