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

hydro_lang/location/
cluster.rs
