Convert With clause from antlr to high level ast

Reviewers: teon.banek

Reviewed By: teon.banek

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D230

src/query/frontend/ast/ast.hpp

@@ -563,7 +563,7 @@ class With : public Clause {
     visitor.PostVisit(*this);
   }
 
-  bool distinct_{false};
+  bool distinct_ = false;
   std::vector<NamedExpression *> named_expressions_;
   Where *where_ = nullptr;
 

src/query/frontend/ast/cypher_main_visitor.cpp

@@ -69,6 +69,9 @@ antlrcpp::Any CypherMainVisitor::visitClause(CypherParser::ClauseContext *ctx) {
     // Different return type!!!
     return ctx->remove()->accept(this).as<std::vector<Clause *>>();
   }
+  if (ctx->with()) {
+    return static_cast<Clause *>(ctx->with()->accept(this).as<With *>());
+  }
   // TODO: implement other clauses.
   throw NotYetImplemented();
   return 0;
@@ -97,11 +100,14 @@ antlrcpp::Any CypherMainVisitor::visitCreate(CypherParser::CreateContext *ctx) {
 
 antlrcpp::Any CypherMainVisitor::visitCypherReturn(
     CypherParser::CypherReturnContext *ctx) {
+  auto *return_clause = storage_.Create<Return>();
   if (ctx->DISTINCT()) {
     // TODO: implement other clauses.
     throw NotYetImplemented();
   }
-  return visitChildren(ctx);
+  return_clause->named_expressions_ =
+      ctx->returnBody()->accept(this).as<std::vector<NamedExpression *>>();
+  return return_clause;
 }
 
 antlrcpp::Any CypherMainVisitor::visitReturnBody(
@@ -115,15 +121,15 @@ antlrcpp::Any CypherMainVisitor::visitReturnBody(
 
 antlrcpp::Any CypherMainVisitor::visitReturnItems(
     CypherParser::ReturnItemsContext *ctx) {
-  auto *return_clause = storage_.Create<Return>();
   if (ctx->getTokens(kReturnAllTokenId).size()) {
     // TODO: implement other clauses.
     throw NotYetImplemented();
   }
+  std::vector<NamedExpression *> named_expressions;
   for (auto *item : ctx->returnItem()) {
-    return_clause->named_expressions_.push_back(item->accept(this));
+    named_expressions.push_back(item->accept(this));
   }
-  return return_clause;
+  return named_expressions;
 }
 
 antlrcpp::Any CypherMainVisitor::visitReturnItem(
@@ -785,5 +791,19 @@ antlrcpp::Any CypherMainVisitor::visitPropertyExpression(
   // It is guaranteed by grammar that there is at least one propertyLookup.
   return dynamic_cast<PropertyLookup *>(expression);
 }
+
+antlrcpp::Any CypherMainVisitor::visitWith(CypherParser::WithContext *ctx) {
+  auto *with = storage_.Create<With>();
+  if (ctx->DISTINCT()) {
+    // TODO: implement this
+    throw NotYetImplemented();
+  }
+  with->named_expressions_ =
+      ctx->returnBody()->accept(this).as<std::vector<NamedExpression *>>();
+  if (ctx->where()) {
+    with->where_ = ctx->where()->accept(this);
+  }
+  return with;
+}
 }
 }

src/query/frontend/ast/cypher_main_visitor.hpp

@@ -164,13 +164,13 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
   antlrcpp::Any visitReturnBody(CypherParser::ReturnBodyContext *ctx) override;
 
   /**
-   * @return Return*
+   * @return vector<NamedExpression*>
    */
   antlrcpp::Any visitReturnItems(
       CypherParser::ReturnItemsContext *ctx) override;
 
   /**
-   * @return NamedExpression*
+   * @return vector<NamedExpression*>
    */
   antlrcpp::Any visitReturnItem(CypherParser::ReturnItemContext *ctx) override;
 
@@ -437,6 +437,11 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
   antlrcpp::Any visitPropertyExpression(
       CypherParser::PropertyExpressionContext *ctx) override;
 
+  /**
+   * @return With*
+   */
+  antlrcpp::Any visitWith(CypherParser::WithContext *ctx) override;
+
  public:
   Query *query() { return query_; }
   const static std::string kAnonPrefix;

tests/unit/cypher_main_visitor.cpp

@@ -700,4 +700,35 @@ TEST(Visitor, Remove) {
                                      ast_generator.db_accessor_->label("i")));
   }
 }
+
+TEST(Visitor, With) {
+  AstGenerator ast_generator("WITH n AS m");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *with = dynamic_cast<With *>(query->clauses_[0]);
+  ASSERT_TRUE(with);
+  ASSERT_FALSE(with->where_);
+  ASSERT_EQ(with->named_expressions_.size(), 1U);
+  auto *named_expr = with->named_expressions_[0];
+  ASSERT_EQ(named_expr->name_, "m");
+  auto *identifier = dynamic_cast<Identifier *>(named_expr->expression_);
+  ASSERT_EQ(identifier->name_, "n");
+}
+
+TEST(Visitor, WithWhere) {
+  AstGenerator ast_generator("WITH n AS m WHERE k");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *with = dynamic_cast<With *>(query->clauses_[0]);
+  ASSERT_TRUE(with);
+  ASSERT_TRUE(with->where_);
+  auto *identifier = dynamic_cast<Identifier *>(with->where_->expression_);
+  ASSERT_TRUE(identifier);
+  ASSERT_EQ(identifier->name_, "k");
+  ASSERT_EQ(with->named_expressions_.size(), 1U);
+  auto *named_expr = with->named_expressions_[0];
+  ASSERT_EQ(named_expr->name_, "m");
+  auto *identifier2 = dynamic_cast<Identifier *>(named_expr->expression_);
+  ASSERT_EQ(identifier2->name_, "n");
+}
 }