Carrier type for Choose and Offer; lift Choice<N> restriction

dialectic-compiler/src/cfg.rs

@@ -47,10 +47,13 @@ pub enum Ir {
     /// continuation, which is stored in the "next" pointer of the node. The scope resolution pass
     /// "lowers" this next pointer continuation into the arms of the `Choose`, and so after scope
     /// resolution all `Choose` nodes' next pointers should be `None`.
-    Choose(Vec<Option<Index>>),
-    /// Like `Choose`, `Offer` nodes have a list of choices, and after scope resolution have no
-    /// continuation.
-    Offer(Vec<Option<Index>>),
+    ///
+    /// The `Choose` may also have an optional "carrier" type. If none, this is defaulted to be a
+    /// `Choice<N>` corresponding to the `N` number of branches of the `Choose`.
+    Choose(Vec<Option<Index>>, Option<Type>),
+    /// Like `Choose`, `Offer` nodes have a list of choices and an optional carrier type, and after
+    /// scope resolution have no continuation.
+    Offer(Vec<Option<Index>>, Option<Type>),
     /// `Split` nodes have a transmit-only half and a receive-only half. The nodes' semantics are
     /// similar to `Call`.
     Split {
@@ -236,7 +239,7 @@ impl Cfg {
                     follow(None, *tx_only);
                     follow(None, *rx_only);
                 }
-                Ir::Choose(choices) | Ir::Offer(choices) => {
+                Ir::Choose(choices, _) | Ir::Offer(choices, _) => {
                     // Inline the current implicit continuation into the next pointer of the node,
                     // if it is `Some`
                     follow(implicit_cont, *next);
@@ -285,7 +288,7 @@ impl Cfg {
                 // `Offer` (which have had their explicit continuations erased and lowered into the
                 // arms) then we need to assign the scoped implicit continuation, if there is one,
                 // as the new explicit continuation
-                None if !matches!(expr, Ir::Choose(_) | Ir::Offer(_)) => *next = implicit_cont,
+                None if !matches!(expr, Ir::Choose(..) | Ir::Offer(..)) => *next = implicit_cont,
                 // This will only be reached if there is no explicit continuation and the node is a
                 // `Choose` or `Offer`, in which case we want to do nothing.
                 _ => {}
@@ -321,7 +324,7 @@ impl Cfg {
                         stack.extend(node.next);
                     }
                 }
-                Ir::Choose(choices) | Ir::Offer(choices) => {
+                Ir::Choose(choices, _) | Ir::Offer(choices, _) => {
                     stack.extend(choices.iter().filter_map(Option::as_ref));
 
                     assert!(node.next.is_none(), "at this point in the compiler, all continuations of \
@@ -479,23 +482,23 @@ impl Cfg {
                         cont: Rc::new(cont),
                     }
                 }
-                Ir::Choose(choices) => {
+                Ir::Choose(choices, carrier_type) => {
                     let targets = choices
                         .iter()
                         .map(|&choice| generate_inner(cfg, errors, loop_env, choice))
                         .collect();
                     debug_assert!(node.next.is_none(), "non-`Done` continuation for `Choose`");
 
-                    Target::Choose(targets)
+                    Target::Choose(targets, carrier_type.clone())
                 }
-                Ir::Offer(choices) => {
+                Ir::Offer(choices, carrier_type) => {
                     let targets = choices
                         .iter()
                         .map(|&choice| generate_inner(cfg, errors, loop_env, choice))
                         .collect();
                     debug_assert!(node.next.is_none(), "non-`Done` continuation for `Offer`");
 
-                    Target::Offer(targets)
+                    Target::Offer(targets, carrier_type.clone())
                 }
                 Ir::Loop(body) => {
                     loop_env.push(node_index);

dialectic-compiler/src/flow.rs

@@ -287,7 +287,7 @@ fn preconditions(cfg: &Cfg, constraint: Constraint) -> Dnf {
             // passability of an `Offer` or `Choose` typically does not matter because as control
             // flow analysis must be run after scope resolution, they should not have an implicit
             // continuation.
-            Ir::Offer(choices) | Ir::Choose(choices) => Dnf(choices
+            Ir::Offer(choices, _) | Ir::Choose(choices, _) => Dnf(choices
                 .iter()
                 .filter_map(Option::as_ref)
                 .map(|&c| vec![Constraint::Passable(c)])
@@ -325,7 +325,7 @@ fn preconditions(cfg: &Cfg, constraint: Constraint) -> Dnf {
                 Dnf::only_if(conj)
             }
             // `Choose`/`Offer` are haltable only if any of their choices are haltable.
-            Ir::Choose(choices) | Ir::Offer(choices) => Dnf(choices
+            Ir::Choose(choices, _) | Ir::Offer(choices, _) => Dnf(choices
                 .iter()
                 // If any of the choices are `Done`, we want to emit an empty `Vec` instead,
                 // denoting that this constraint is trivially satisfiable.
@@ -371,7 +371,7 @@ fn preconditions(cfg: &Cfg, constraint: Constraint) -> Dnf {
                 Dnf::only_if(conj)
             }
             // `Choose` and `Offer` break only if any arm breaks.
-            Ir::Choose(choices) | Ir::Offer(choices) => Dnf(choices
+            Ir::Choose(choices, _) | Ir::Offer(choices, _) => Dnf(choices
                 .iter()
                 .filter_map(Option::as_ref)
                 .chain(node.next.as_ref())

dialectic-compiler/src/lib.rs

@@ -393,20 +393,20 @@ mod tests {
         let continue1 = cfg.singleton(Ir::Continue(client));
         cfg[recv].next = Some(continue1);
         let choose_opts = vec![Some(send), Some(recv)];
-        let choose = cfg.singleton(Ir::Choose(choose_opts));
+        let choose = cfg.singleton(Ir::Choose(choose_opts, None));
 
         cfg[client_tally].expr = Ir::Loop(Some(choose));
 
         let break0 = cfg.singleton(Ir::Break(client));
         let send = cfg.send("Operation");
         cfg[send].next = Some(client_tally);
         let choose_opts = vec![Some(break0), Some(send)];
-        let choose = cfg.singleton(Ir::Choose(choose_opts));
+        let choose = cfg.singleton(Ir::Choose(choose_opts, None));
 
         cfg[client].expr = Ir::Loop(Some(choose));
 
         let s = format!("{}", cfg.generate_target(Some(client)).unwrap());
-        assert_eq!(s, "Loop<Choose<(Done, Send<Operation, Loop<Choose<(Send<i64, Continue<0>>, Recv<i64, Continue<1>>)>>>)>>");
+        assert_eq!(s, "Loop<Choose<(Done, Send<Operation, Loop<Choose<(Send<i64, Continue<0>>, Recv<i64, Continue<1>>), Choice<2>>>>), Choice<2>>>");
     }
 
     #[test]
@@ -421,14 +421,14 @@ mod tests {
         let continue0 = cfg.singleton(Ir::Continue(client));
         cfg[call].next = Some(continue0);
         let choose_opts = vec![Some(break0), Some(send)];
-        let choose = cfg.singleton(Ir::Choose(choose_opts));
+        let choose = cfg.singleton(Ir::Choose(choose_opts, None));
 
         cfg[client].expr = Ir::Loop(Some(choose));
 
         let s = format!("{}", cfg.generate_target(Some(client)).unwrap());
         assert_eq!(
             s,
-            "Loop<Choose<(Done, Send<Operation, Call<ClientTally, Continue<0>>>)>>"
+            "Loop<Choose<(Done, Send<Operation, Call<ClientTally, Continue<0>>>), Choice<2>>>"
         );
     }
 }

dialectic-compiler/src/parse.rs

@@ -228,6 +228,12 @@ impl Parse for Spanned<Syntax> {
             let kw_span = input.parse::<kw::choose>()?.span();
             let choose_span = kw_span.join(input.span()).unwrap_or(kw_span);
 
+            let carrier_type = if input.peek(token::Brace) {
+                None
+            } else {
+                Some(input.parse()?)
+            };
+
             let content;
             braced!(content in input);
             let mut choice_arms = Vec::new();
@@ -251,14 +257,20 @@ impl Parse for Spanned<Syntax> {
             }
 
             Ok(Spanned {
-                inner: Syntax::Choose(arm_asts),
+                inner: Syntax::Choose(arm_asts, carrier_type),
                 span: choose_span,
             })
         } else if lookahead.peek(kw::offer) {
             // Ast::Offer: offer { 0 => <Ast>, 1 => <Ast>, ... }
             let kw_span = input.parse::<kw::offer>()?.span();
             let offer_span = kw_span.join(input.span()).unwrap_or(kw_span);
 
+            let carrier_type = if input.peek(token::Brace) {
+                None
+            } else {
+                Some(input.parse()?)
+            };
+
             let content;
             braced!(content in input);
             let mut choice_arms = Vec::new();
@@ -282,7 +294,7 @@ impl Parse for Spanned<Syntax> {
             }
 
             Ok(Spanned {
-                inner: Syntax::Offer(arm_asts),
+                inner: Syntax::Offer(arm_asts, carrier_type),
                 span: offer_span,
             })
         } else if lookahead.peek(kw::split) {

dialectic-compiler/src/syntax.rs

@@ -42,9 +42,9 @@ pub enum Syntax {
     /// Syntax: `call T` or `call { ... }`.
     Call(Box<Spanned<Syntax>>),
     /// Syntax: `choose { 0 => ..., ... }`.
-    Choose(Vec<Spanned<Syntax>>),
+    Choose(Vec<Spanned<Syntax>>, Option<Type>),
     /// Syntax: `offer { 0 => ..., ... }`.
-    Offer(Vec<Spanned<Syntax>>),
+    Offer(Vec<Spanned<Syntax>>, Option<Type>),
     /// Syntax: `split { -> ..., <- ... }`.
     Split {
         /// The transmit-only half.
@@ -172,19 +172,19 @@ fn to_cfg<'a>(
             let rx_only = to_cfg(rx_only, cfg, env).0;
             Ir::Split { tx_only, rx_only }
         }
-        Choose(choices) => {
+        Choose(choices, carrier_type) => {
             let choice_nodes = choices
                 .iter()
                 .map(|choice| to_cfg(choice, cfg, env).0)
                 .collect();
-            Ir::Choose(choice_nodes)
+            Ir::Choose(choice_nodes, carrier_type.clone())
         }
-        Offer(choices) => {
+        Offer(choices, carrier_type) => {
             let choice_nodes = choices
                 .iter()
                 .map(|choice| to_cfg(choice, cfg, env).0)
                 .collect();
-            Ir::Offer(choice_nodes)
+            Ir::Offer(choice_nodes, carrier_type.clone())
         }
         Continue(label) => {
             return convert_jump_to_cfg(label, CompileError::ContinueOutsideLoop, Ir::Continue)
@@ -306,13 +306,13 @@ impl Spanned<Syntax> {
                 let rx = rx_only.to_token_stream_with(add_optional);
                 quote_spanned! {sp=> split { -> #tx, <- #rx, } }
             }
-            Choose(choices) => {
+            Choose(choices, carrier_type) => {
                 let arms = choice_arms_to_tokens(&mut add_optional, choices);
-                quote_spanned! {sp=> choose { #arms } }
+                quote_spanned! {sp=> choose #carrier_type { #arms } }
             }
-            Offer(choices) => {
+            Offer(choices, carrier_type) => {
                 let arms = choice_arms_to_tokens(&mut add_optional, choices);
-                quote_spanned! {sp=> offer { #arms } }
+                quote_spanned! {sp=> offer #carrier_type { #arms } }
             }
             Loop(None, body) => {
                 let body_tokens = body.to_token_stream_with(add_optional);
@@ -354,7 +354,7 @@ impl Spanned<Syntax> {
                         // rather than required.
                         let ends_with_block = matches!(
                             &stmt.inner,
-                            Block(_) | Split { .. } | Offer(_) | Choose(_) | Loop(_, _),
+                            Block(_) | Split { .. } | Offer(..) | Choose(..) | Loop(_, _),
                         );
 
                         if !(is_call_of_block || ends_with_block) || add_optional() {
@@ -438,8 +438,14 @@ mod tests {
                     "type" => Type(parse_quote!(())),
                     other => unreachable!("{}", other),
                 },
-                "choose" => Choose(Arbitrary::arbitrary(g)),
-                "offer" => Offer(Arbitrary::arbitrary(g)),
+                "choose" => Choose(
+                    Arbitrary::arbitrary(g),
+                    Some(parse_quote!(())).filter(|_| Arbitrary::arbitrary(g)),
+                ),
+                "offer" => Offer(
+                    Arbitrary::arbitrary(g),
+                    Some(parse_quote!(())).filter(|_| Arbitrary::arbitrary(g)),
+                ),
                 "split" => Split {
                     tx_only: Arbitrary::arbitrary(g),
                     rx_only: Arbitrary::arbitrary(g),
@@ -475,8 +481,8 @@ mod tests {
                         })
                     })
                     .collect(),
-                Choose(choices) => choices.shrink().map(Choose).collect(),
-                Offer(choices) => choices.shrink().map(Offer).collect(),
+                Choose(choices, _) => choices.shrink().map(|cs| Choose(cs, None)).collect(),
+                Offer(choices, _) => choices.shrink().map(|cs| Offer(cs, None)).collect(),
                 Loop(label, body) => body
                     .shrink()
                     .map(|body_shrunk| Loop(label.clone(), body_shrunk))

dialectic-compiler/src/target.rs

@@ -4,7 +4,7 @@ use {
     proc_macro2::TokenStream,
     quote::{quote_spanned, ToTokens},
     std::{fmt, rc::Rc},
-    syn::{Path, Type},
+    syn::{parse_quote, Path, Type},
 };
 
 use crate::Spanned;
@@ -24,10 +24,10 @@ pub enum Target {
     Recv(Type, Rc<Spanned<Target>>),
     /// Session type: `Send<T, P>`.
     Send(Type, Rc<Spanned<Target>>),
-    /// Session type: `Choose<(P, ...)>`.
-    Choose(Vec<Spanned<Target>>),
-    /// Session type: `Offer<(P, ...)>`.
-    Offer(Vec<Spanned<Target>>),
+    /// Session type: `Choose<(P, ...), Carrier>`.
+    Choose(Vec<Spanned<Target>>, Option<Type>),
+    /// Session type: `Offer<(P, ...), Carrier>`.
+    Offer(Vec<Spanned<Target>>, Option<Type>),
     /// Session type: `Loop<...>`.
     Loop(Rc<Spanned<Target>>),
     /// Session type: `Continue<N>`.
@@ -64,33 +64,55 @@ impl fmt::Display for Target {
             } => write!(f, "Split<{}, {}, {}>", s, p, q)?,
             Call(s, p) => write!(f, "Call<{}, {}>", s, p)?,
             Then(s, p) => write!(f, "<{} as Then<{}>>::Combined", s, p)?,
-            Choose(cs) => {
+            Choose(cs, carrier_type) => {
                 let count = cs.len();
+
                 write!(f, "Choose<(")?;
+
                 for (i, c) in cs.iter().enumerate() {
                     write!(f, "{}", c)?;
                     if i + 1 < count {
                         write!(f, ", ")?;
                     }
                 }
+
                 if count == 1 {
                     write!(f, ",")?;
                 }
-                write!(f, ")>")?;
+
+                write!(f, "), ")?;
+
+                match carrier_type {
+                    Some(carrier) => write!(f, "CustomChoice<{}>", carrier.to_token_stream())?,
+                    None => write!(f, "Choice<{}>", count)?,
+                }
+
+                write!(f, ">")?;
             }
-            Offer(cs) => {
+            Offer(cs, carrier_type) => {
                 let count = cs.len();
+
                 write!(f, "Offer<(")?;
+
                 for (i, c) in cs.iter().enumerate() {
                     write!(f, "{}", c)?;
                     if i + 1 < count {
                         write!(f, ", ")?;
                     }
                 }
+
                 if count == 1 {
                     write!(f, ",")?;
                 }
-                write!(f, ")>")?;
+
+                write!(f, "), ")?;
+
+                match carrier_type {
+                    Some(carrier) => write!(f, "CustomChoice<{}>", carrier.to_token_stream())?,
+                    None => write!(f, "Choice<{}>", count)?,
+                }
+
+                write!(f, ">")?;
             }
             Continue(n) => {
                 write!(f, "Continue<{}>", n)?;
@@ -154,17 +176,33 @@ impl Spanned<Target> {
                 quote_spanned!(span=> <#s as #dialectic_crate::types::Then<#p>>::Combined)
                     .to_tokens(tokens);
             }
-            Choose(cs) => {
+            Choose(cs, carrier_type) => {
+                let carrier: syn::Type = match carrier_type {
+                    Some(ty) => parse_quote!(#dialectic_crate::backend::CustomChoice<#ty>),
+                    None => {
+                        let n = cs.len();
+                        parse_quote!(#dialectic_crate::backend::Choice<#n>)
+                    }
+                };
                 let cs = cs
                     .iter()
                     .map(|c| c.to_token_stream_with_crate_name(dialectic_crate));
-                quote_spanned!(span=> #dialectic_crate::types::Choose<(#(#cs,)*)>).to_tokens(tokens)
+                quote_spanned!(span=> #dialectic_crate::types::Choose<(#(#cs,)*), #carrier>)
+                    .to_tokens(tokens)
             }
-            Offer(cs) => {
+            Offer(cs, carrier_type) => {
+                let carrier: syn::Type = match carrier_type {
+                    Some(ty) => parse_quote!(#dialectic_crate::backend::CustomChoice<#ty>),
+                    None => {
+                        let n = cs.len();
+                        parse_quote!(#dialectic_crate::backend::Choice<#n>)
+                    }
+                };
                 let cs = cs
                     .iter()
                     .map(|c| c.to_token_stream_with_crate_name(dialectic_crate));
-                quote_spanned!(span=> #dialectic_crate::types::Offer<(#(#cs,)*)>).to_tokens(tokens)
+                quote_spanned!(span=> #dialectic_crate::types::Offer<(#(#cs,)*), #carrier>)
+                    .to_tokens(tokens)
             }
             Continue(n) => {
                 quote_spanned!(span=> #dialectic_crate::types::Continue<#n>).to_tokens(tokens)