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::{
11    builder::FlowState,
12    ir::{CollectionKind, HydroIrMetadata},
13};
14
15#[expect(missing_docs, reason = "TODO")]
16#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Debug, Hash, Serialize, Deserialize)]
17pub enum LocationId {
18    Process(usize),
19    Cluster(usize),
20    Atomic(
21        /// The tick that the atomic region is associated with.
22        Box<LocationId>,
23    ),
24    Tick(usize, Box<LocationId>),
25}
26
27#[expect(missing_docs, reason = "TODO")]
28impl LocationId {
29    pub fn root(&self) -> &LocationId {
30        match self {
31            LocationId::Process(_) => self,
32            LocationId::Cluster(_) => self,
33            LocationId::Atomic(tick) => tick.root(),
34            LocationId::Tick(_, id) => id.root(),
35        }
36    }
37
38    pub fn is_root(&self) -> bool {
39        match self {
40            LocationId::Process(_) | LocationId::Cluster(_) => true,
41            LocationId::Atomic(_) => false,
42            LocationId::Tick(_, _) => false,
43        }
44    }
45
46    pub fn is_top_level(&self) -> bool {
47        match self {
48            LocationId::Process(_) | LocationId::Cluster(_) => true,
49            LocationId::Atomic(_) => true,
50            LocationId::Tick(_, _) => false,
51        }
52    }
53
54    pub fn raw_id(&self) -> usize {
55        match self {
56            LocationId::Process(id) => *id,
57            LocationId::Cluster(id) => *id,
58            LocationId::Atomic(_) => panic!("cannot get raw id for atomic"),
59            LocationId::Tick(_, _) => panic!("cannot get raw id for tick"),
60        }
61    }
62
63    pub fn swap_root(&mut self, new_root: LocationId) {
64        match self {
65            LocationId::Tick(_, id) => {
66                id.swap_root(new_root);
67            }
68            LocationId::Atomic(tick) => {
69                tick.swap_root(new_root);
70            }
71            _ => {
72                assert!(new_root.is_root());
73                *self = new_root;
74            }
75        }
76    }
77}
78
79#[cfg(stageleft_runtime)]
80pub(crate) trait DynLocation: Clone {
81    fn id(&self) -> LocationId;
82
83    fn flow_state(&self) -> &FlowState;
84    fn is_top_level() -> bool;
85
86    fn new_node_metadata(&self, collection_kind: CollectionKind) -> HydroIrMetadata {
87        use crate::compile::ir::HydroIrOpMetadata;
88        use crate::compile::ir::backtrace::Backtrace;
89
90        HydroIrMetadata {
91            location_kind: self.id(),
92            collection_kind,
93            cardinality: None,
94            tag: None,
95            op: HydroIrOpMetadata {
96                backtrace: Backtrace::get_backtrace(2),
97                cpu_usage: None,
98                network_recv_cpu_usage: None,
99                id: None,
100            },
101        }
102    }
103}