From 3ead1fbf81f1eda8b07c8a4d5daba941fa6265aa Mon Sep 17 00:00:00 2001
From: Yingchi Long <i@lyc.dev>
Date: Mon, 12 Feb 2024 15:47:18 +0800
Subject: [PATCH] libnixf/Parse: parse expr_with

---
 libnixf/include/nixf/Basic/NodeKinds.inc |  1 +
 libnixf/include/nixf/Basic/Nodes/Expr.h  | 17 +++++++++++
 libnixf/src/Parse/ParseExpr.cpp          | 37 ++++++++++++++++++++++++
 libnixf/src/Parse/ParserImpl.h           |  6 ++++
 libnixf/test/Parse/ParseExpr.cpp         | 27 +++++++++++++++++
 5 files changed, 88 insertions(+)

diff --git a/libnixf/include/nixf/Basic/NodeKinds.inc b/libnixf/include/nixf/Basic/NodeKinds.inc
index a76fb4671..1c0bfc507 100644
--- a/libnixf/include/nixf/Basic/NodeKinds.inc
+++ b/libnixf/include/nixf/Basic/NodeKinds.inc
@@ -43,5 +43,6 @@ EXPR(ExprOpHasAttr)
 EXPR(ExprIf)
 EXPR(ExprAssert)
 EXPR(ExprLet)
+EXPR(ExprWith)
 
 #endif // EXPR
diff --git a/libnixf/include/nixf/Basic/Nodes/Expr.h b/libnixf/include/nixf/Basic/Nodes/Expr.h
index 9ade208f4..bb4101180 100644
--- a/libnixf/include/nixf/Basic/Nodes/Expr.h
+++ b/libnixf/include/nixf/Basic/Nodes/Expr.h
@@ -153,4 +153,21 @@ class ExprLet : public Expr {
   }
 };
 
+class ExprWith : public Expr {
+  std::unique_ptr<Expr> With;
+  std::unique_ptr<Expr> E;
+
+public:
+  ExprWith(LexerCursorRange Range, std::unique_ptr<Expr> With,
+           std::unique_ptr<Expr> E)
+      : Expr(NK_ExprWith, Range), With(std::move(With)), E(std::move(E)) {}
+
+  [[nodiscard]] Expr *with() const { return With.get(); }
+  [[nodiscard]] Expr *expr() const { return E.get(); }
+
+  [[nodiscard]] ChildVector children() const override {
+    return {With.get(), E.get()};
+  }
+};
+
 } // namespace nixf
diff --git a/libnixf/src/Parse/ParseExpr.cpp b/libnixf/src/Parse/ParseExpr.cpp
index 9454e9e3b..135ec5df4 100644
--- a/libnixf/src/Parse/ParseExpr.cpp
+++ b/libnixf/src/Parse/ParseExpr.cpp
@@ -97,6 +97,7 @@ std::unique_ptr<Expr> Parser::parseExpr() {
       case tok_comma:    // { a ,
       case tok_id:       // { a b
       case tok_ellipsis: // { a ...
+      case tok_r_curly:
         return parseExprLambda();
       default:
         break;
@@ -116,6 +117,8 @@ std::unique_ptr<Expr> Parser::parseExpr() {
     if (peek(1).kind() != tok_l_curly)
       return parseExprLet();
     break;
+  case tok_kw_with:
+    return parseExprWith();
   default:
     break;
   }
@@ -271,4 +274,38 @@ std::unique_ptr<ExprLet> Parser::parseExprLet() {
                                    std::move(In), std::move(E));
 }
 
+std::unique_ptr<ExprWith> Parser::parseExprWith() {
+  LexerCursor LCur = lCur();
+  Token TokWith = peek();
+  assert(TokWith.kind() == tok_kw_with && "token should be tok_kw_with");
+
+  consume(); // with
+  assert(LastToken && "LastToken should be set after consume()");
+
+  auto SyncSemi = withSync(tok_semi_colon);
+
+  auto With = parseExpr();
+
+  if (!With)
+    diagNullExpr(Diags, LastToken->rCur(), "with expression");
+
+  ExpectResult ExpSemi = expect(tok_semi_colon);
+  if (!ExpSemi.ok()) {
+    ExpSemi.diag().note(Note::NK_ToMachThis, TokWith.range())
+        << std::string(tok::spelling(tok_kw_with));
+    return std::make_unique<ExprWith>(LexerCursorRange{LCur, LastToken->rCur()},
+                                      std::move(With), /*E=*/nullptr);
+  }
+
+  consume(); // ;
+
+  auto E = parseExpr();
+
+  if (!E)
+    diagNullExpr(Diags, LastToken->rCur(), "with body");
+
+  return std::make_unique<ExprWith>(LexerCursorRange{LCur, LastToken->rCur()},
+                                    std::move(With), std::move(E));
+}
+
 } // namespace nixf
diff --git a/libnixf/src/Parse/ParserImpl.h b/libnixf/src/Parse/ParserImpl.h
index 8aafd7d77..dd943be72 100644
--- a/libnixf/src/Parse/ParserImpl.h
+++ b/libnixf/src/Parse/ParserImpl.h
@@ -33,6 +33,7 @@ class ExprLambda;
 class ExprIf;
 class ExprAssert;
 class ExprLet;
+class ExprWith;
 
 namespace detail {
 
@@ -342,6 +343,11 @@ class Parser {
   /// \endcode
   std::unique_ptr<ExprLet> parseExprLet();
 
+  /// \code
+  /// expr_with :  'with' expr ';' expr
+  /// \endcode
+  std::unique_ptr<ExprWith> parseExprWith();
+
   std::unique_ptr<Expr> parse() { return parseExpr(); }
 };
 
diff --git a/libnixf/test/Parse/ParseExpr.cpp b/libnixf/test/Parse/ParseExpr.cpp
index f0a8bcb08..fe1ccfe7b 100644
--- a/libnixf/test/Parse/ParseExpr.cpp
+++ b/libnixf/test/Parse/ParseExpr.cpp
@@ -302,4 +302,31 @@ TEST(Parser, ExprLet_Binds) {
   ASSERT_FALSE(static_cast<ExprLet *>(AST.get())->binds());
 }
 
+TEST(Parser, ExprWith_Ok) {
+  auto Src = R"(with 1; 1)"sv;
+
+  std::vector<Diagnostic> Diags;
+  Parser P(Src, Diags);
+  auto AST = P.parseExpr();
+
+  ASSERT_TRUE(AST);
+  ASSERT_EQ(AST->kind(), Node::NK_ExprWith);
+
+  ASSERT_EQ(Diags.size(), 0);
+}
+
+TEST(Parser, ExprWith_NoExpr) {
+  auto Src = R"(with 1;)"sv;
+
+  std::vector<Diagnostic> Diags;
+  Parser P(Src, Diags);
+  auto AST = P.parseExpr();
+
+  ASSERT_TRUE(AST);
+  ASSERT_EQ(AST->kind(), Node::NK_ExprWith);
+
+  ASSERT_EQ(Diags.size(), 1);
+  ASSERT_EQ(Diags[0].kind(), Diagnostic::DK_Expected);
+}
+
 } // namespace