hydro_lang/location/

cluster.rs

use std::fmt::{Debug, Formatter};
use std::marker::PhantomData;

use proc_macro2::Span;
use quote::quote;
use stageleft::runtime_support::{FreeVariableWithContext, QuoteTokens};
use stageleft::{QuotedWithContext, quote_type};

use super::dynamic::LocationId;
use super::{Location, MemberId};
use crate::compile::builder::FlowState;
use crate::staging_util::{Invariant, get_this_crate};

pub struct Cluster<'a, ClusterTag> {
    pub(crate) id: usize,
    pub(crate) flow_state: FlowState,
    pub(crate) _phantom: Invariant<'a, ClusterTag>,
}

impl<C> Debug for Cluster<'_, C> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "Cluster({})", self.id)
    }
}

impl<C> Eq for Cluster<'_, C> {}
impl<C> PartialEq for Cluster<'_, C> {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id && FlowState::ptr_eq(&self.flow_state, &other.flow_state)
    }
}

impl<C> Clone for Cluster<'_, C> {
    fn clone(&self) -> Self {
        Cluster {
            id: self.id,
            flow_state: self.flow_state.clone(),
            _phantom: PhantomData,
        }
    }
}

impl<'a, C> super::dynamic::DynLocation for Cluster<'a, C> {
    fn id(&self) -> LocationId {
        LocationId::Cluster(self.id)
    }

    fn flow_state(&self) -> &FlowState {
        &self.flow_state
    }

    fn is_top_level() -> bool {
        true
    }
}

impl<'a, C> Location<'a> for Cluster<'a, C> {
    type Root = Cluster<'a, C>;

    fn root(&self) -> Self::Root {
        self.clone()
    }
}

pub struct ClusterIds<'a, C> {
    pub(crate) id: usize,
    pub(crate) _phantom: Invariant<'a, C>,
}

impl<C> Clone for ClusterIds<'_, C> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<C> Copy for ClusterIds<'_, C> {}

impl<'a, C: 'a, Ctx> FreeVariableWithContext<Ctx> for ClusterIds<'a, C> {
    type O = &'a [MemberId<C>];

    fn to_tokens(self, _ctx: &Ctx) -> QuoteTokens
    where
        Self: Sized,
    {
        let ident = syn::Ident::new(
            &format!("__hydro_lang_cluster_ids_{}", self.id),
            Span::call_site(),
        );
        let root = get_this_crate();
        let c_type = quote_type::<C>();

        QuoteTokens {
            prelude: None,
            expr: Some(
                quote! { unsafe { ::std::mem::transmute::<_, &[#root::__staged::location::MemberId<#c_type>]>(#ident) } },
            ),
        }
    }
}

impl<'a, C, Ctx> QuotedWithContext<'a, &'a [MemberId<C>], Ctx> for ClusterIds<'a, C> {}

pub trait IsCluster {
    type Tag;
}

impl<C> IsCluster for Cluster<'_, C> {
    type Tag = C;
}

/// A free variable representing the cluster's own ID. When spliced in
/// a quoted snippet that will run on a cluster, this turns into a [`MemberId`].
pub static CLUSTER_SELF_ID: ClusterSelfId = ClusterSelfId { _private: &() };

#[derive(Clone, Copy)]
pub struct ClusterSelfId<'a> {
    _private: &'a (),
}

impl<'a, L> FreeVariableWithContext<L> for ClusterSelfId<'a>
where
    L: Location<'a>,
    <L as Location<'a>>::Root: IsCluster,
{
    type O = MemberId<<<L as Location<'a>>::Root as IsCluster>::Tag>;

    fn to_tokens(self, ctx: &L) -> QuoteTokens
    where
        Self: Sized,
    {
        let cluster_id = if let LocationId::Cluster(id) = ctx.root().id() {
            id
        } else {
            unreachable!()
        };

        let ident = syn::Ident::new(
            &format!("__hydro_lang_cluster_self_id_{}", cluster_id),
            Span::call_site(),
        );
        let root = get_this_crate();
        let c_type: syn::Type = quote_type::<<<L as Location<'a>>::Root as IsCluster>::Tag>();

        QuoteTokens {
            prelude: None,
            expr: Some(quote! { #root::location::MemberId::<#c_type>::from_raw(#ident) }),
        }
    }
}

impl<'a, L> QuotedWithContext<'a, MemberId<<<L as Location<'a>>::Root as IsCluster>::Tag>, L>
    for ClusterSelfId<'a>
where
    L: Location<'a>,
    <L as Location<'a>>::Root: IsCluster,
{
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use stageleft::q;

    use super::CLUSTER_SELF_ID;
    use crate::location::{Location, MembershipEvent};
    use crate::nondet::nondet;
    use crate::prelude::FlowBuilder;

    #[test]
    fn sim_cluster_self_id() {
        let flow = FlowBuilder::new();
        let cluster1 = flow.cluster::<()>();
        let cluster2 = flow.cluster::<()>();

        let node = flow.process::<()>();
        let external = flow.external::<()>();

        let out_port = cluster1
            .source_iter(q!(vec![CLUSTER_SELF_ID.raw_id]))
            .send_bincode(&node)
            .values()
            .interleave(
                cluster2
                    .source_iter(q!(vec![CLUSTER_SELF_ID.raw_id]))
                    .send_bincode(&node)
                    .values(),
            )
            .send_bincode_external(&external);

        flow.sim()
            .with_cluster_size(&cluster1, 3)
            .with_cluster_size(&cluster2, 4)
            .exhaustive(async |mut compiled| {
                let out_recv = compiled.connect(&out_port);
                compiled.launch();

                out_recv
                    .assert_yields_only_unordered(vec![0, 1, 2, 0, 1, 2, 3])
                    .await
            });
    }

    #[test]
    fn sim_cluster_with_tick() {
        let flow = FlowBuilder::new();
        let cluster = flow.cluster::<()>();
        let node = flow.process::<()>();
        let external = flow.external::<()>();

        let out_port = cluster
            .source_iter(q!(vec![1, 2, 3]))
            .batch(&cluster.tick(), nondet!(/** test */))
            .count()
            .all_ticks()
            .send_bincode(&node)
            .entries()
            .map(q!(|(id, v)| (id.raw_id, v)))
            .send_bincode_external(&external);

        let count = flow
            .sim()
            .with_cluster_size(&cluster, 2)
            .exhaustive(async |mut compiled| {
                let out_recv = compiled.connect(&out_port);
                compiled.launch();

                let grouped = out_recv.collect_sorted::<Vec<_>>().await.into_iter().fold(
                    HashMap::new(),
                    |mut acc: HashMap<u32, usize>, (id, v)| {
                        *acc.entry(id).or_default() += v;
                        acc
                    },
                );

                assert!(grouped.len() == 2);
                for (_id, v) in grouped {
                    assert!(v == 3);
                }
            });

        assert_eq!(count, 106);
        // not a square because we simulate all interleavings of ticks across 2 cluster members
        // eventually, we should be able to identify that the members are independent (because
        // there are no dataflow cycles) and avoid simulating redundant interleavings
    }

    #[test]
    fn sim_cluster_membership() {
        let flow = FlowBuilder::new();
        let cluster = flow.cluster::<()>();
        let node = flow.process::<()>();
        let external = flow.external::<()>();

        let out_port = node
            .source_cluster_members(&cluster)
            .entries()
            .map(q!(|(id, v)| (id.raw_id, v)))
            .send_bincode_external(&external);

        flow.sim()
            .with_cluster_size(&cluster, 2)
            .exhaustive(async |mut compiled| {
                let out_recv = compiled.connect(&out_port);
                compiled.launch();

                out_recv
                    .assert_yields_only_unordered(vec![
                        (0, MembershipEvent::Joined),
                        (1, MembershipEvent::Joined),
                    ])
                    .await;
            });
    }
}