hydro_lang/location/cluster/
mod.rs

1use std::fmt::{Debug, Formatter};
2use std::marker::PhantomData;
3
4use proc_macro2::Span;
5use quote::quote;
6use stageleft::runtime_support::{FreeVariableWithContext, QuoteTokens};
7use stageleft::{QuotedWithContext, quote_type};
8
9use super::{Location, LocationId};
10use crate::builder::FlowState;
11use crate::staging_util::{Invariant, get_this_crate};
12
13pub mod cluster_id;
14pub use cluster_id::ClusterId;
15
16pub struct Cluster<'a, C> {
17    pub(crate) id: usize,
18    pub(crate) flow_state: FlowState,
19    pub(crate) _phantom: Invariant<'a, C>,
20}
21
22impl<C> Debug for Cluster<'_, C> {
23    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
24        write!(f, "Cluster({})", self.id)
25    }
26}
27
28impl<C> Eq for Cluster<'_, C> {}
29impl<C> PartialEq for Cluster<'_, C> {
30    fn eq(&self, other: &Self) -> bool {
31        self.id == other.id && self.flow_state.as_ptr() == other.flow_state.as_ptr()
32    }
33}
34
35impl<'a, C> Cluster<'a, C> {
36    pub fn members(&self) -> ClusterIds<'a, C> {
37        ClusterIds {
38            id: self.id,
39            _phantom: PhantomData,
40        }
41    }
42}
43
44impl<C> Clone for Cluster<'_, C> {
45    fn clone(&self) -> Self {
46        Cluster {
47            id: self.id,
48            flow_state: self.flow_state.clone(),
49            _phantom: PhantomData,
50        }
51    }
52}
53
54impl<'a, C> Location<'a> for Cluster<'a, C> {
55    type Root = Cluster<'a, C>;
56
57    fn root(&self) -> Self::Root {
58        self.clone()
59    }
60
61    fn id(&self) -> LocationId {
62        LocationId::Cluster(self.id)
63    }
64
65    fn flow_state(&self) -> &FlowState {
66        &self.flow_state
67    }
68
69    fn is_top_level() -> bool {
70        true
71    }
72}
73
74pub struct ClusterIds<'a, C> {
75    pub(crate) id: usize,
76    _phantom: Invariant<'a, C>,
77}
78
79impl<C> Clone for ClusterIds<'_, C> {
80    fn clone(&self) -> Self {
81        *self
82    }
83}
84
85impl<C> Copy for ClusterIds<'_, C> {}
86
87impl<'a, C: 'a, Ctx> FreeVariableWithContext<Ctx> for ClusterIds<'a, C> {
88    type O = &'a Vec<ClusterId<C>>;
89
90    fn to_tokens(self, _ctx: &Ctx) -> QuoteTokens
91    where
92        Self: Sized,
93    {
94        let ident = syn::Ident::new(
95            &format!("__hydro_lang_cluster_ids_{}", self.id),
96            Span::call_site(),
97        );
98        let root = get_this_crate();
99        let c_type = quote_type::<C>();
100
101        QuoteTokens {
102            prelude: None,
103            expr: Some(
104                quote! { unsafe { ::std::mem::transmute::<_, &[#root::ClusterId<#c_type>]>(#ident) } },
105            ),
106        }
107    }
108}
109
110impl<'a, C, Ctx> QuotedWithContext<'a, &'a Vec<ClusterId<C>>, Ctx> for ClusterIds<'a, C> {}
111
112pub trait IsCluster {
113    type Tag;
114}
115
116impl<C> IsCluster for Cluster<'_, C> {
117    type Tag = C;
118}
119
120/// A free variable representing the cluster's own ID. When spliced in
121/// a quoted snippet that will run on a cluster, this turns into a [`ClusterId`].
122pub static CLUSTER_SELF_ID: ClusterSelfId = ClusterSelfId { _private: &() };
123
124#[derive(Clone, Copy)]
125pub struct ClusterSelfId<'a> {
126    _private: &'a (),
127}
128
129impl<'a, L: Location<'a>> FreeVariableWithContext<L> for ClusterSelfId<'a>
130where
131    <L as Location<'a>>::Root: IsCluster,
132{
133    type O = ClusterId<<<L as Location<'a>>::Root as IsCluster>::Tag>;
134
135    fn to_tokens(self, ctx: &L) -> QuoteTokens
136    where
137        Self: Sized,
138    {
139        let cluster_id = if let LocationId::Cluster(id) = ctx.root().id() {
140            id
141        } else {
142            unreachable!()
143        };
144
145        let ident = syn::Ident::new(
146            &format!("__hydro_lang_cluster_self_id_{}", cluster_id),
147            Span::call_site(),
148        );
149        let root = get_this_crate();
150        let c_type: syn::Type = quote_type::<<<L as Location<'a>>::Root as IsCluster>::Tag>();
151
152        QuoteTokens {
153            prelude: None,
154            expr: Some(quote! { #root::ClusterId::<#c_type>::from_raw(#ident) }),
155        }
156    }
157}
158
159impl<'a, L: Location<'a>>
160    QuotedWithContext<'a, ClusterId<<<L as Location<'a>>::Root as IsCluster>::Tag>, L>
161    for ClusterSelfId<'a>
162where
163    <L as Location<'a>>::Root: IsCluster,
164{
165}
hydro_lang/location/cluster/mod.rs

hydro_lang/location/cluster/
mod.rs
