4 files changed, 65 insertions, 17 deletions

diff --git a/src/grammar.ron b/src/grammar.ron
index 109556561..8e8a71010 100644
--- a/src/grammar.ron
+++ b/src/grammar.ron
@@ -70,6 +70,7 @@ Grammar(
         "ref",
         "let",
         "move",
+        "return",
     ],
     contextual_keywords: [
         "auto",
@@ -138,6 +139,8 @@ Grammar(
         "STRUCT_LIT",
         "STRUCT_LIT_FIELD",
         "IF_EXPR",
+        "BLOCK_EXPR",
+        "RETURN_EXPR",
 
         "EXTERN_BLOCK_EXPR",
         "ENUM_VARIANT",
@@ -155,7 +158,6 @@ Grammar(
         "ABI",
         "NAME",
         "NAME_REF",
-        "BLOCK_EXPR",
 
         "LET_STMT",
         "EXPR_STMT",
diff --git a/src/grammar/expressions.rs b/src/grammar/expressions.rs
index 5c59843a4..423e1a95a 100644
--- a/src/grammar/expressions.rs
+++ b/src/grammar/expressions.rs
@@ -13,18 +13,20 @@ use super::*;
 //     let _ = b"e";
 //     let _ = br"f";
 // }
+const LITERAL_FIRST: TokenSet =
+    token_set![TRUE_KW, FALSE_KW, INT_NUMBER, FLOAT_NUMBER, BYTE, CHAR,
+               STRING, RAW_STRING, BYTE_STRING, RAW_BYTE_STRING];
+
 pub(super) fn literal(p: &mut Parser) -> Option<CompletedMarker> {
-    match p.current() {
-        TRUE_KW | FALSE_KW | INT_NUMBER | FLOAT_NUMBER | BYTE | CHAR | STRING | RAW_STRING
-        | BYTE_STRING | RAW_BYTE_STRING => {
-            let m = p.start();
-            p.bump();
-            Some(m.complete(p, LITERAL))
-        }
-        _ => None,
+    if !LITERAL_FIRST.contains(p.current()) {
+        return None;
     }
+    let m = p.start();
+    p.bump();
+    Some(m.complete(p, LITERAL))
 }
 
+const EXPR_FIRST: TokenSet = PREFIX_EXPR_FIRST;
 pub(super) fn expr(p: &mut Parser) {
     let mut lhs = match prefix_expr(p) {
         Some(lhs) => lhs,
@@ -80,6 +82,11 @@ fn let_stmt(p: &mut Parser) {
     m.complete(p, LET_STMT);
 }
 
+const PREFIX_EXPR_FIRST: TokenSet =
+    token_set_union![
+        token_set![AMPERSAND, STAR, EXCL],
+        ATOM_EXPR_FIRST,
+    ];
 fn prefix_expr(p: &mut Parser) -> Option<CompletedMarker> {
     let done = match p.current() {
         AMPERSAND => ref_expr(p),
@@ -128,6 +135,11 @@ fn not_expr(p: &mut Parser) -> CompletedMarker {
     m.complete(p, NOT_EXPR)
 }
 
+const ATOM_EXPR_FIRST: TokenSet =
+    token_set_union![
+        LITERAL_FIRST,
+        token_set![L_PAREN, PIPE, MOVE_KW, IF_KW, UNSAFE_KW, L_CURLY, RETURN_KW],
+    ];
 fn atom_expr(p: &mut Parser) -> Option<CompletedMarker> {
     match literal(p) {
         Some(m) => return Some(m),
@@ -144,6 +156,7 @@ fn atom_expr(p: &mut Parser) -> Option<CompletedMarker> {
         IF_KW => if_expr(p),
         UNSAFE_KW if la == L_CURLY => block_expr(p),
         L_CURLY => block_expr(p),
+        RETURN_KW => return_expr(p),
         _ => {
             p.err_and_bump("expected expression");
             return None;
@@ -237,6 +250,21 @@ fn block_expr(p: &mut Parser) -> CompletedMarker {
     m.complete(p, BLOCK_EXPR)
 }
 
+// test return_expr
+// fn foo() {
+//     return;
+//     return 92;
+// }
+fn return_expr(p: &mut Parser) -> CompletedMarker {
+    assert!(p.at(RETURN_KW));
+    let m = p.start();
+    p.bump();
+    if EXPR_FIRST.contains(p.current()) {
+        expr(p);
+    }
+    m.complete(p, RETURN_EXPR)
+}
+
 // test call_expr
 // fn foo() {
 //     let _ = f();
diff --git a/src/parser_api.rs b/src/parser_api.rs
index 5a0da32c9..d12f773b2 100644
--- a/src/parser_api.rs
+++ b/src/parser_api.rs
@@ -3,22 +3,22 @@ use {
     SyntaxKind::{self, ERROR},
 };
 
-pub(crate) struct TokenSet {
-    pub tokens: &'static [SyntaxKind],
+#[derive(Clone, Copy)]
+pub(crate) struct TokenSet(pub(crate) u128);
+fn mask(kind: SyntaxKind) -> u128 {
+    1u128 << (kind as usize)
 }
 
 impl TokenSet {
     pub fn contains(&self, kind: SyntaxKind) -> bool {
-        self.tokens.contains(&kind)
+        self.0 & mask(kind) != 0
     }
 }
 
 #[macro_export]
 macro_rules! token_set {
     ($($t:ident),*) => {
-        TokenSet {
-            tokens: &[$($t),*],
-        }
+        TokenSet($(1u128 << ($t as usize))|*)
     };
 
     ($($t:ident),* ,) => {
@@ -26,6 +26,17 @@ macro_rules! token_set {
     };
 }
 
+#[macro_export]
+macro_rules! token_set_union {
+    ($($ts:expr),*) => {
+        TokenSet($($ts.0)|*)
+    };
+
+    ($($ts:expr),* ,) => {
+        token_set_union!($($ts),*)
+    };
+}
+
 /// `Parser` struct provides the low-level API for
 /// navigating through the stream of tokens and
 /// constructing the parse tree. The actual parsing
diff --git a/src/syntax_kinds/generated.rs b/src/syntax_kinds/generated.rs
index 5f26cbf0a..50aa2b580 100644
--- a/src/syntax_kinds/generated.rs
+++ b/src/syntax_kinds/generated.rs
@@ -71,6 +71,7 @@ pub enum SyntaxKind {
     REF_KW,
     LET_KW,
     MOVE_KW,
+    RETURN_KW,
     AUTO_KW,
     DEFAULT_KW,
     UNION_KW,
@@ -129,6 +130,8 @@ pub enum SyntaxKind {
     STRUCT_LIT,
     STRUCT_LIT_FIELD,
     IF_EXPR,
+    BLOCK_EXPR,
+    RETURN_EXPR,
     EXTERN_BLOCK_EXPR,
     ENUM_VARIANT,
     NAMED_FIELD,
@@ -145,7 +148,6 @@ pub enum SyntaxKind {
     ABI,
     NAME,
     NAME_REF,
-    BLOCK_EXPR,
     LET_STMT,
     EXPR_STMT,
     TYPE_PARAM_LIST,
@@ -202,6 +204,7 @@ impl SyntaxKind {
             | REF_KW
             | LET_KW
             | MOVE_KW
+            | RETURN_KW
             | AUTO_KW
             | DEFAULT_KW
             | UNION_KW
@@ -278,6 +281,7 @@ impl SyntaxKind {
             REF_KW => &SyntaxInfo { name: "REF_KW" },
             LET_KW => &SyntaxInfo { name: "LET_KW" },
             MOVE_KW => &SyntaxInfo { name: "MOVE_KW" },
+            RETURN_KW => &SyntaxInfo { name: "RETURN_KW" },
             AUTO_KW => &SyntaxInfo { name: "AUTO_KW" },
             DEFAULT_KW => &SyntaxInfo { name: "DEFAULT_KW" },
             UNION_KW => &SyntaxInfo { name: "UNION_KW" },
@@ -336,6 +340,8 @@ impl SyntaxKind {
             STRUCT_LIT => &SyntaxInfo { name: "STRUCT_LIT" },
             STRUCT_LIT_FIELD => &SyntaxInfo { name: "STRUCT_LIT_FIELD" },
             IF_EXPR => &SyntaxInfo { name: "IF_EXPR" },
+            BLOCK_EXPR => &SyntaxInfo { name: "BLOCK_EXPR" },
+            RETURN_EXPR => &SyntaxInfo { name: "RETURN_EXPR" },
             EXTERN_BLOCK_EXPR => &SyntaxInfo { name: "EXTERN_BLOCK_EXPR" },
             ENUM_VARIANT => &SyntaxInfo { name: "ENUM_VARIANT" },
             NAMED_FIELD => &SyntaxInfo { name: "NAMED_FIELD" },
@@ -352,7 +358,6 @@ impl SyntaxKind {
             ABI => &SyntaxInfo { name: "ABI" },
             NAME => &SyntaxInfo { name: "NAME" },
             NAME_REF => &SyntaxInfo { name: "NAME_REF" },
-            BLOCK_EXPR => &SyntaxInfo { name: "BLOCK_EXPR" },
             LET_STMT => &SyntaxInfo { name: "LET_STMT" },
             EXPR_STMT => &SyntaxInfo { name: "EXPR_STMT" },
             TYPE_PARAM_LIST => &SyntaxInfo { name: "TYPE_PARAM_LIST" },
@@ -403,6 +408,7 @@ impl SyntaxKind {
             "ref" => REF_KW,
             "let" => LET_KW,
             "move" => MOVE_KW,
+            "return" => RETURN_KW,
             _ => return None,
         };
         Some(kw)
@@ -506,6 +512,7 @@ impl SyntaxKind {
             REF_KW => "ref",
             LET_KW => "let",
             MOVE_KW => "move",
+            RETURN_KW => "return",
             AUTO_KW => "auto",
             DEFAULT_KW => "default",
             UNION_KW => "union",