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

hydro_lang/location/
dynamic.rs
