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
9use super::LocationKey;
10use crate::compile::builder::ClockId;
11#[cfg(stageleft_runtime)]
12use crate::compile::{
13    builder::FlowState,
14    ir::{CollectionKind, HydroIrMetadata},
15};
16use crate::location::LocationType;
17
18/// An enumeration representing a location heirarchy, including "virtual" locations (atomic/tick).
19#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Hash, Serialize, Deserialize)]
20pub enum LocationId {
21    /// A process root location (i.e. a single node).
22    Process(LocationKey),
23    /// A cluster root location (i.e. multiple nodes).
24    Cluster(LocationKey),
25    /// An atomic region, within a tick.
26    Atomic(
27        /// The tick that the atomic region is associated with.
28        Box<LocationId>,
29    ),
30    /// A tick within a location.
31    Tick(ClockId, Box<LocationId>),
32}
33
34/// Implement Debug to Display-print the key, reduces snapshot verbosity.
35impl std::fmt::Debug for LocationId {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        match self {
38            LocationId::Process(key) => write!(f, "Process({key})"),
39            LocationId::Cluster(key) => write!(f, "Cluster({key})"),
40            LocationId::Atomic(tick) => write!(f, "Atomic({tick:?})"),
41            LocationId::Tick(tick, id) => write!(f, "Tick({tick}, {id:?})"),
42        }
43    }
44}
45
46impl LocationId {
47    /// The [`LocationType`] of this location ID. `None` if this is not a root location.
48    pub fn location_type(&self) -> Option<LocationType> {
49        match self {
50            LocationId::Process(_) => Some(LocationType::Process),
51            LocationId::Cluster(_) => Some(LocationType::Cluster),
52            _ => None,
53        }
54    }
55}
56
57#[expect(missing_docs, reason = "TODO")]
58impl LocationId {
59    pub fn root(&self) -> &LocationId {
60        match self {
61            LocationId::Process(_) => self,
62            LocationId::Cluster(_) => self,
63            LocationId::Atomic(tick) => tick.root(),
64            LocationId::Tick(_, id) => id.root(),
65        }
66    }
67
68    pub fn is_root(&self) -> bool {
69        match self {
70            LocationId::Process(_) | LocationId::Cluster(_) => true,
71            LocationId::Atomic(_) => false,
72            LocationId::Tick(_, _) => false,
73        }
74    }
75
76    pub fn is_top_level(&self) -> bool {
77        match self {
78            LocationId::Process(_) | LocationId::Cluster(_) => true,
79            LocationId::Atomic(_) => true,
80            LocationId::Tick(_, _) => false,
81        }
82    }
83
84    pub fn key(&self) -> LocationKey {
85        match self {
86            LocationId::Process(id) => *id,
87            LocationId::Cluster(id) => *id,
88            LocationId::Atomic(_) => panic!("cannot get raw id for atomic"),
89            LocationId::Tick(_, _) => panic!("cannot get raw id for tick"),
90        }
91    }
92
93    pub fn swap_root(&mut self, new_root: LocationId) {
94        match self {
95            LocationId::Tick(_, id) => {
96                id.swap_root(new_root);
97            }
98            LocationId::Atomic(tick) => {
99                tick.swap_root(new_root);
100            }
101            _ => {
102                assert!(new_root.is_root());
103                *self = new_root;
104            }
105        }
106    }
107}
108
109#[cfg(stageleft_runtime)]
110pub(crate) trait DynLocation: Clone {
111    fn id(&self) -> LocationId;
112
113    fn flow_state(&self) -> &FlowState;
114    fn is_top_level() -> bool;
115    fn multiversioned(&self) -> bool;
116
117    fn new_node_metadata(&self, collection_kind: CollectionKind) -> HydroIrMetadata {
118        use crate::compile::ir::HydroIrOpMetadata;
119        use crate::compile::ir::backtrace::Backtrace;
120
121        HydroIrMetadata {
122            location_id: self.id(),
123            collection_kind,
124            cardinality: None,
125            tag: None,
126            op: HydroIrOpMetadata {
127                backtrace: Backtrace::get_backtrace(2),
128                cpu_usage: None,
129                network_recv_cpu_usage: None,
130                id: None,
131            },
132        }
133    }
134}
hydro_lang/location/dynamic.rs

hydro_lang/location/
dynamic.rs
