From 34d7a2835125498ea5f52cd6e743fcb89741dfc6 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Wed, 23 Nov 2022 19:04:50 -0500
Subject: [PATCH] feat: add check keyword and new assignment syntax

---
 crates/lang/src/ast.rs          |  1 +
 crates/lang/src/format.rs       |  1 +
 crates/lang/src/parser.rs       | 18 +++++++++++++++++-
 crates/lang/src/parser/lexer.rs |  1 +
 crates/lang/src/parser/token.rs |  2 ++
 crates/lang/src/tipo/infer.rs   |  1 +
 6 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/crates/lang/src/ast.rs b/crates/lang/src/ast.rs
index af9d94c7..9b3c0251 100644
--- a/crates/lang/src/ast.rs
+++ b/crates/lang/src/ast.rs
@@ -618,6 +618,7 @@ impl<A, B> Pattern<A, B> {
 pub enum AssignmentKind {
     Let,
     Assert,
+    Check,
 }
 
 pub type MultiPattern<PatternConstructor, Type> = Vec<Pattern<PatternConstructor, Type>>;
diff --git a/crates/lang/src/format.rs b/crates/lang/src/format.rs
index ee873131..a34d2d5d 100644
--- a/crates/lang/src/format.rs
+++ b/crates/lang/src/format.rs
@@ -568,6 +568,7 @@ impl<'comments> Formatter<'comments> {
         let keyword = match kind {
             Some(AssignmentKind::Let) => "let ",
             Some(AssignmentKind::Assert) => "assert ",
+            Some(AssignmentKind::Check) => "check ",
             None => "try ",
         };
 
diff --git a/crates/lang/src/parser.rs b/crates/lang/src/parser.rs
index 71a0ffd8..e377156e 100644
--- a/crates/lang/src/parser.rs
+++ b/crates/lang/src/parser.rs
@@ -494,7 +494,7 @@ pub fn expr_parser(
         let assert_parser = just(Token::Assert)
             .ignore_then(pattern_parser())
             .then(just(Token::Colon).ignore_then(type_parser()).or_not())
-            .then_ignore(just(Token::Equal))
+            .then_ignore(just(Token::Is))
             .then(r.clone())
             .map_with_span(
                 |((pattern, annotation), value), span| expr::UntypedExpr::Assignment {
@@ -506,6 +506,21 @@ pub fn expr_parser(
                 },
             );
 
+        let check_parser = just(Token::Check)
+            .ignore_then(pattern_parser())
+            .then(just(Token::Colon).ignore_then(type_parser()).or_not())
+            .then_ignore(just(Token::Is))
+            .then(r.clone())
+            .map_with_span(
+                |((pattern, annotation), value), span| expr::UntypedExpr::Assignment {
+                    location: span,
+                    value: Box::new(value),
+                    pattern,
+                    kind: ast::AssignmentKind::Check,
+                    annotation,
+                },
+            );
+
         let if_parser = just(Token::If)
             .ignore_then(r.clone().then(block_parser.clone()).map_with_span(
                 |(condition, body), span| ast::IfBranch {
@@ -551,6 +566,7 @@ pub fn expr_parser(
             when_parser,
             let_parser,
             assert_parser,
+            check_parser,
             if_parser,
         ));
 
diff --git a/crates/lang/src/parser/lexer.rs b/crates/lang/src/parser/lexer.rs
index 3630b32d..ce9539ea 100644
--- a/crates/lang/src/parser/lexer.rs
+++ b/crates/lang/src/parser/lexer.rs
@@ -63,6 +63,7 @@ pub fn lexer() -> impl Parser<char, Vec<(Token, Span)>, Error = ParseError> {
     let keyword = text::ident().map(|s: String| match s.as_str() {
         "as" => Token::As,
         "assert" => Token::Assert,
+        "check" => Token::Assert,
         "const" => Token::Const,
         "fn" => Token::Fn,
         "if" => Token::If,
diff --git a/crates/lang/src/parser/token.rs b/crates/lang/src/parser/token.rs
index 4a382980..208d6d86 100644
--- a/crates/lang/src/parser/token.rs
+++ b/crates/lang/src/parser/token.rs
@@ -58,6 +58,7 @@ pub enum Token {
     // Keywords (alphabetically):
     As,
     Assert,
+    Check,
     Const,
     Fn,
     If,
@@ -130,6 +131,7 @@ impl fmt::Display for Token {
             Token::EmptyLine => "EMPTYLINE",
             Token::As => "as",
             Token::Assert => "assert",
+            Token::Check => "check",
             Token::When => "when",
             Token::Is => "is",
             Token::Const => "const",
diff --git a/crates/lang/src/tipo/infer.rs b/crates/lang/src/tipo/infer.rs
index 4956e840..3a029b29 100644
--- a/crates/lang/src/tipo/infer.rs
+++ b/crates/lang/src/tipo/infer.rs
@@ -449,6 +449,7 @@ fn str_to_keyword(word: &str) -> Option<Token> {
     match word {
         "as" => Some(Token::As),
         "assert" => Some(Token::Assert),
+        "check" => Some(Token::Check),
         "when" => Some(Token::When),
         "const" => Some(Token::Const),
         "fn" => Some(Token::Fn),