Implement select {} statement.

compiler/src/ast.rs

@@ -543,8 +543,7 @@ pub enum Expr {
         span: Span,
     },
     Select {
-        arms: Vec<Arm>,
-        ty: Type, // TODO remove this
+        arms: Vec<SelectArm>,
         span: Span,
     },
     Defer {
@@ -773,7 +772,7 @@ impl Expr {
             Self::Const { ty, .. } => ty,
             Self::Debug { ty, .. } => ty,
             Self::Spawn { ty, .. } => ty,
-            Self::Select { ty, .. } => ty,
+            Self::Select { .. } => Type::discard(),
             Self::Defer { ty, .. } => ty,
             Self::Reference { ty, .. } => ty,
             Self::Index { ty, .. } => ty,
@@ -819,7 +818,7 @@ impl Expr {
             Self::Const { ref mut ty, .. } => *ty = new_ty,
             Self::Debug { ref mut ty, .. } => *ty = new_ty,
             Self::Spawn { ref mut ty, .. } => *ty = new_ty,
-            Self::Select { ref mut ty, .. } => *ty = new_ty,
+            Self::Select { .. } => (),
             Self::Defer { ref mut ty, .. } => *ty = new_ty,
             Self::Reference { ref mut ty, .. } => *ty = new_ty,
             Self::Index { ref mut ty, .. } => *ty = new_ty,
@@ -1029,3 +1028,16 @@ pub struct InterfaceSuperTrait {
     pub span: Span,
     pub ty: Type,
 }
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct SelectArm {
+    pub pat: SelectArmPat,
+    pub expr: Expr,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum SelectArmPat {
+    Recv(Binding, Expr),
+    Send(Expr),
+    Wildcard,
+}

compiler/src/codegen.rs

@@ -4,7 +4,7 @@ use crate::{
     ast::{
         Arm, Binding, Constructor, DebugKind, EnumDefinition, EnumFieldDef, Expr, File, FileId,
         Function, FunctionKind, Generic, InterfaceSuperTrait, Literal, Loop, LoopFlow, Operator,
-        Pat, Span, StructDefinition, StructField, UnOp,
+        Pat, SelectArm, SelectArmPat, Span, StructDefinition, StructField, UnOp,
     },
     global_state::Module,
     infer,
@@ -375,7 +375,7 @@ impl Codegen {
             Expr::CheckType { .. } => EmitResult::empty(),
             Expr::Paren { expr, .. } => self.emit_paren(expr),
             Expr::Spawn { expr, .. } => self.emit_spawn(expr),
-            Expr::Select { arms: _, .. } => todo!(),
+            Expr::Select { arms, .. } => self.emit_select(arms),
             Expr::Defer { expr, .. } => self.emit_defer(expr),
             Expr::Reference { expr, .. } => self.emit_reference(expr),
             Expr::Index { expr, index, .. } => self.emit_index(expr, index),
@@ -1510,6 +1510,71 @@ if {is_matching} != 2 {{
         }
     }
 
+    fn emit_select(&mut self, arms: &[SelectArm]) -> EmitResult {
+        let mut out = emitter();
+        let mut inner = emitter();
+
+        for a in arms {
+            match &a.pat {
+                SelectArmPat::Recv(binding, expr) => match binding.pat {
+                    Pat::Type { ref ident, .. } => {
+                        self.scope.add_binding(ident.clone(), ident.clone());
+
+                        let new_expr = self.emit_select_case(&expr, &mut out);
+                        inner.emit(format!("case {ident} := {new_expr}:",));
+                    }
+
+                    _ => unreachable!(),
+                },
+
+                SelectArmPat::Send(expr) => {
+                    let new_expr = self.emit_select_case(&expr, &mut out);
+                    inner.emit(format!("case {new_expr}:",));
+                }
+
+                SelectArmPat::Wildcard => {
+                    inner.emit("default:".to_string());
+                }
+            }
+
+            let body = self.emit_expr(Ctx::Discard.to_mode(), &ensure_wrapped_in_block(&a.expr));
+            inner.emit(body.output);
+        }
+
+        out.emit("select {".to_string());
+        out.emit(inner.render());
+        out.emit("}".to_string());
+
+        out.no_value()
+    }
+
+    // Unwrap ch.Recv() or ch.Send(x)
+    // and emit corresponding <- ch or ch <- x
+    fn emit_select_case(&mut self, expr: &Expr, out: &mut Emitter) -> String {
+        match expr {
+            Expr::Call { func, args, .. } => match **func {
+                Expr::FieldAccess {
+                    ref field,
+                    expr: ref method_receiver,
+                    ..
+                } => {
+                    let new_expr = self.emit_local(&method_receiver, out);
+
+                    if field == "Recv" {
+                        format!("<- {new_expr}")
+                    } else {
+                        let new_arg = self.emit_local(&args[0], out);
+                        format!("{new_expr} <- {new_arg}")
+                    }
+                }
+
+                _ => unreachable!(),
+            },
+
+            _ => unreachable!(),
+        }
+    }
+
     fn create_make_function(&mut self, def: &EnumDefinition) -> String {
         let mut out = emitter();
 
@@ -2037,7 +2102,7 @@ func {name} {generic_params} ({params}) {ret} {{
                 format!("func ({args}) {ret}")
             }
 
-            Type::Var(_) => format!("TO_TYPE_FAILED"),
+            Type::Var(_) => format!("any"),
             // Type::Var(_) => panic!("unexpected Type::Var in to_type"),
         }
     }

compiler/src/infer.rs

@@ -1,7 +1,8 @@
 use crate::ast::{
     Arm, Binding, Constructor, EnumDefinition, EnumFieldDef, Expr, File, Function, FunctionKind,
-    Generic, InterfaceSuperTrait, Literal, Loop, NewtypeDefinition, Operator, Pat, PkgImport, Span,
-    StructDefinition, StructField, StructFieldDef, StructFieldPat, TypeAliasDef, TypeAst, UnOp,
+    Generic, InterfaceSuperTrait, Literal, Loop, NewtypeDefinition, Operator, Pat, PkgImport,
+    SelectArm, SelectArmPat, Span, StructDefinition, StructField, StructFieldDef, StructFieldPat,
+    TypeAliasDef, TypeAst, UnOp,
 };
 use crate::error::{ArityError, Error, UnificationError};
 use crate::exhaustive;
@@ -232,9 +233,6 @@ impl Infer {
     pub fn type_any(&self) -> Type {
         self.builtin_type("any")
     }
-    pub fn type_channel_op(&self, typ: Type) -> Type {
-        self.builtin_type("ChannelOp").swap_arg(0, typ)
-    }
 
     pub fn infer_expr(&mut self, expr: Expr, expected: &Type) -> Expr {
         match expr {
@@ -743,16 +741,6 @@ impl Infer {
                 let subject_ty = self.fresh_ty_var();
                 let new_subject = self.infer_expr(*subject, &subject_ty);
 
-                // if it's a `match select!()`, then turn it into a Select expr
-                if matches!(new_subject, Expr::Select { .. }) {
-                    let select = Expr::Select {
-                        arms,
-                        span,
-                        ty: Type::dummy(),
-                    };
-                    return self.infer_expr(select, expected);
-                }
-
                 let new_arms = arms
                     .into_iter()
                     .map(|a| {
@@ -1575,41 +1563,51 @@ has no field or method:
                 }
             }
 
-            Expr::Select { arms, ty: _, span } => {
+            Expr::Select { arms, span, .. } => {
                 // Each arm is inferred separately
-                // Arm pat must be of type ChannelOp
-                // Return type is unit
                 let new_arms = arms
                     .into_iter()
                     .map(|a| {
                         // NEW SCOPE
                         self.begin_scope();
 
-                        let var = self.fresh_ty_var();
-                        let channel_op = self.type_channel_op(var);
+                        let new_pat = match a.pat {
+                            SelectArmPat::Recv(binding, expr) => {
+                                let new_expr = self.ensure_channel_call("Recv", expr);
+
+                                let ty = self.to_type(&binding.ann, &span);
+                                let new_binding = Binding {
+                                    pat: self.infer_pat(binding.pat, Type::discard()),
+                                    ty: ty.clone(),
+                                    ..binding
+                                };
+
+                                SelectArmPat::Recv(new_binding, new_expr)
+                            }
+
+                            SelectArmPat::Send(expr) => {
+                                let new_expr = self.ensure_channel_call("Send", expr);
+                                SelectArmPat::Send(new_expr)
+                            }
 
-                        // each arm.pat should unify with ChannelOp
-                        let new_pat = self.infer_pat(a.pat, channel_op);
+                            SelectArmPat::Wildcard => SelectArmPat::Wildcard,
+                        };
 
                         // the resulting type from arm.expr can be discarded
                         let new_expr = self.infer_expr(a.expr, &Type::discard());
 
                         // EXIT SCOPE
                         self.exit_scope();
 
-                        Arm {
+                        SelectArm {
                             pat: new_pat,
                             expr: new_expr,
                         }
                     })
                     .collect();
 
-                let ret = self.type_unit();
-                self.add_constraint(expected, &ret, &span);
-
                 Expr::Select {
                     arms: new_arms,
-                    ty: ret,
                     span,
                 }
             }
@@ -3304,6 +3302,54 @@ has no field or method:
             span: span.clone(),
         })
     }
+
+    // In a select {}, we're looking for:
+    //   - x.Recv()
+    //   - x.Send()
+    // anything else is invalid
+    fn ensure_channel_call(&mut self, method: &str, expr: Expr) -> Expr {
+        let new_expr = self.infer_expr(expr, &Type::discard());
+
+        match &new_expr {
+            Expr::Call { func, .. } => match **func {
+                Expr::FieldAccess {
+                    ref field,
+                    ref span,
+                    expr: ref method_receiver,
+                    ..
+                } => {
+                    if field != method {
+                        self.generic_error(
+                            format!("was expecting {method}, got {field}"),
+                            span.clone(),
+                        );
+                    }
+
+                    let expected = if method == "Recv" {
+                        self.builtin_type("Receiver")
+                    } else {
+                        self.builtin_type("Sender")
+                    }
+                    .swap_arg(0, self.type_any());
+
+                    self.add_constraint(&expected, &method_receiver.get_type(), &span);
+                }
+
+                _ => {
+                    self.generic_error("invalid call in select".to_string(), func.get_span());
+                }
+            },
+
+            _ => {
+                self.generic_error(
+                    "invalid expression in select".to_string(),
+                    new_expr.get_span(),
+                );
+            }
+        }
+
+        new_expr
+    }
 }
 
 #[derive(Debug, Clone)]

compiler/src/lexer.rs

@@ -138,6 +138,7 @@ pub enum TokenKind {
     Loop,
     Break,
     Continue,
+    Select,
 
     Support, // @ensure, @rawgo
     Error,
@@ -174,6 +175,7 @@ fn get_keyword(s: &str) -> Option<TokenKind> {
         "loop" => TokenKind::Loop,
         "break" => TokenKind::Break,
         "continue" => TokenKind::Continue,
+        "select" => TokenKind::Select,
         _ => TokenKind::Error,
     };