hydro_lang/location/cluster/
mod.rs

Help
1use std::fmt::{Debug, Formatter};
2use std::marker::PhantomData;
3
4use proc_macro2::{Span, TokenStream};
5use quote::quote;
6use stageleft::runtime_support::FreeVariableWithContext;
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) -> (Option<TokenStream>, Option<TokenStream>)
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            None,
102            Some(
103                quote! { unsafe { ::std::mem::transmute::<_, &[#root::ClusterId<#c_type>]>(#ident) } },
104            ),
105        )
106    }
107}
108
109impl<'a, C, Ctx> QuotedWithContext<'a, &'a Vec<ClusterId<C>>, Ctx> for ClusterIds<'a, C> {}
110
111pub trait IsCluster {
112    type Tag;
113}
114
115impl<C> IsCluster for Cluster<'_, C> {
116    type Tag = C;
117}
118
119/// A free variable representing the cluster's own ID. When spliced in
120/// a quoted snippet that will run on a cluster, this turns into a [`ClusterId`].
121pub static CLUSTER_SELF_ID: ClusterSelfId = ClusterSelfId { _private: &() };
122
123#[derive(Clone, Copy)]
124pub struct ClusterSelfId<'a> {
125    _private: &'a (),
126}
127
128impl<'a, L: Location<'a>> FreeVariableWithContext<L> for ClusterSelfId<'a>
129where
130    <L as Location<'a>>::Root: IsCluster,
131{
132    type O = ClusterId<<<L as Location<'a>>::Root as IsCluster>::Tag>;
133
134    fn to_tokens(self, ctx: &L) -> (Option<TokenStream>, Option<TokenStream>)
135    where
136        Self: Sized,
137    {
138        let cluster_id = if let LocationId::Cluster(id) = ctx.root().id() {
139            id
140        } else {
141            unreachable!()
142        };
143
144        let ident = syn::Ident::new(
145            &format!("__hydro_lang_cluster_self_id_{}", cluster_id),
146            Span::call_site(),
147        );
148        let root = get_this_crate();
149        let c_type: syn::Type = quote_type::<<<L as Location<'a>>::Root as IsCluster>::Tag>();
150        (
151            None,
152            Some(quote! { #root::ClusterId::<#c_type>::from_raw(#ident) }),
153        )
154    }
155}
156
157impl<'a, L: Location<'a>>
158    QuotedWithContext<'a, ClusterId<<<L as Location<'a>>::Root as IsCluster>::Tag>, L>
159    for ClusterSelfId<'a>
160where
161    <L as Location<'a>>::Root: IsCluster,
162{
163}
hydro_lang/location/cluster/mod.rs

hydro_lang/location/cluster/
mod.rs
