Add list indexing and slicing conversion to AST

Reviewers: teon.banek

Reviewed By: teon.banek

Subscribers: pullbot

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

src/query/frontend/ast/ast.hpp

@@ -299,6 +299,54 @@ class GreaterEqualOperator : public BinaryOperator {
   using BinaryOperator::BinaryOperator;
 };
 
+class ListIndexingOperator : public BinaryOperator {
+  friend class AstTreeStorage;
+
+ public:
+  void Accept(TreeVisitorBase &visitor) override {
+    if (visitor.PreVisit(*this)) {
+      visitor.Visit(*this);
+      expression1_->Accept(visitor);
+      expression2_->Accept(visitor);
+      visitor.PostVisit(*this);
+    }
+  }
+
+ protected:
+  using BinaryOperator::BinaryOperator;
+};
+
+class ListSlicingOperator : public Expression {
+  friend class AstTreeStorage;
+
+ public:
+  void Accept(TreeVisitorBase &visitor) override {
+    if (visitor.PreVisit(*this)) {
+      visitor.Visit(*this);
+      list_->Accept(visitor);
+      if (lower_bound_) {
+        lower_bound_->Accept(visitor);
+      }
+      if (upper_bound_) {
+        upper_bound_->Accept(visitor);
+      }
+      visitor.PostVisit(*this);
+    }
+  }
+
+  Expression *list_;
+  Expression *lower_bound_;
+  Expression *upper_bound_;
+
+ protected:
+  ListSlicingOperator(int uid, Expression *list, Expression *lower_bound,
+                      Expression *upper_bound)
+      : Expression(uid),
+        list_(list),
+        lower_bound_(lower_bound),
+        upper_bound_(upper_bound) {}
+};
+
 class NotOperator : public UnaryOperator {
   friend class AstTreeStorage;
 
@@ -389,8 +437,7 @@ class ListLiteral : public BaseLiteral {
   void Accept(TreeVisitorBase &visitor) override {
     if (visitor.PreVisit(*this)) {
       visitor.Visit(*this);
-      for (auto expr_ptr : elements_)
-        expr_ptr->Accept(visitor);
+      for (auto expr_ptr : elements_) expr_ptr->Accept(visitor);
       visitor.PostVisit(*this);
     }
   }

src/query/frontend/ast/ast_visitor.hpp

@@ -38,6 +38,8 @@ class LessOperator;
 class GreaterOperator;
 class LessEqualOperator;
 class GreaterEqualOperator;
+class ListIndexingOperator;
+class ListSlicingOperator;
 class Delete;
 class Where;
 class SetProperty;
@@ -53,10 +55,10 @@ using TreeVisitorBase = ::utils::Visitor<
     AdditionOperator, SubtractionOperator, MultiplicationOperator,
     DivisionOperator, ModOperator, NotEqualOperator, EqualOperator,
     LessOperator, GreaterOperator, LessEqualOperator, GreaterEqualOperator,
-    UnaryPlusOperator, UnaryMinusOperator, IsNullOperator, Identifier,
-    PrimitiveLiteral, ListLiteral, PropertyLookup, Aggregation, Function,
-    Create, Match, Return, With, Pattern, NodeAtom, EdgeAtom, Delete, Where,
-    SetProperty, SetProperties, SetLabels, RemoveProperty, RemoveLabels, Merge,
-    Unwind>;
+    ListIndexingOperator, ListSlicingOperator, UnaryPlusOperator,
+    UnaryMinusOperator, IsNullOperator, Identifier, PrimitiveLiteral,
+    ListLiteral, PropertyLookup, Aggregation, Function, Create, Match, Return,
+    With, Pattern, NodeAtom, EdgeAtom, Delete, Where, SetProperty,
+    SetProperties, SetLabels, RemoveProperty, RemoveLabels, Merge, Unwind>;
 
 }  // namespace query

src/query/frontend/ast/cypher_main_visitor.cpp

@@ -568,21 +568,18 @@ antlrcpp::Any CypherMainVisitor::visitExpression5(
 // Unary minus and plus.
 antlrcpp::Any CypherMainVisitor::visitExpression4(
     CypherParser::Expression4Context *ctx) {
-  return PrefixUnaryOperator(ctx->expression3(), ctx->children,
+  return PrefixUnaryOperator(ctx->expression3a(), ctx->children,
                              {kUnaryPlusTokenId, kUnaryMinusTokenId});
 }
 
-// List indexing, list range...
-antlrcpp::Any CypherMainVisitor::visitExpression3(
-    CypherParser::Expression3Context *ctx) {
-  // If there is only one child we don't need to generate any code in this since
-  // that child is expression2. Other operations are not implemented at the
-  // moment.
+// IS NULL, IS NOT NULL, ...
+antlrcpp::Any CypherMainVisitor::visitExpression3a(
+    CypherParser::Expression3aContext *ctx) {
   // TODO: implement this.
   // This is a hack. Unfortunately, grammar for expression3 contains a lot of
   // different expressions. We should break that production in parts so that
   // we can easily implement its visitor.
-  Expression *expression = ctx->expression2()[0]->accept(this);
+  Expression *expression = ctx->expression3b()[0].accept(this);
   if (ctx->children.size() - ctx->SP().size() == 3U && ctx->IS().size() == 1U &&
       ctx->CYPHERNULL().size() == 1U) {
     return static_cast<Expression *>(
@@ -596,9 +593,44 @@ antlrcpp::Any CypherMainVisitor::visitExpression3(
   if (ctx->children.size() > 1U) {
     throw utils::NotYetImplemented();
   }
+  // If there is only one child we don't need to generate any code in this since
+  // that child is expression2.
   return static_cast<Expression *>(visitChildren(ctx));
 }
 
+antlrcpp::Any CypherMainVisitor::visitExpression3b(
+    CypherParser::Expression3bContext *ctx) {
+  Expression *expression = ctx->expression2()->accept(this);
+  for (auto *list_op : ctx->listIndexingOrSlicing()) {
+    if (list_op->getTokens(kDotsTokenId).size() == 0U) {
+      // If there is no '..' then we need to create list indexing operator.
+      expression = storage_.Create<ListIndexingOperator>(
+          expression, list_op->expression()[0]->accept(this));
+    } else if (!list_op->lower_bound && !list_op->upper_bound) {
+      throw SemanticException(
+          "List slicing operator requires at least one bound.");
+    } else {
+      Expression *lower_bound_ast =
+          list_op->lower_bound
+              ? static_cast<Expression *>(list_op->lower_bound->accept(this))
+              : nullptr;
+      Expression *upper_bound_ast =
+          list_op->upper_bound
+              ? static_cast<Expression *>(list_op->upper_bound->accept(this))
+              : nullptr;
+      expression = storage_.Create<ListSlicingOperator>(
+          expression, lower_bound_ast, upper_bound_ast);
+    }
+  }
+  return expression;
+}
+
+antlrcpp::Any CypherMainVisitor::visitListIndexingOrSlicing(
+    CypherParser::ListIndexingOrSlicingContext *) {
+  debug_assert(false, "Should never be called. See documentation in hpp.");
+  return 0;
+}
+
 antlrcpp::Any CypherMainVisitor::visitExpression2(
     CypherParser::Expression2Context *ctx) {
   if (ctx->nodeLabels().size()) {

src/query/frontend/ast/cypher_main_visitor.hpp

@@ -88,7 +88,8 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
     return operators;
   }
 
-  /* Convert opencypher's n-ary production to ast binary operators.
+  /**
+   * Convert opencypher's n-ary production to ast binary operators.
    *
    * @param _expressions Subexpressions of child for which we construct ast
    * operators, for example expression6 if we want to create ast nodes for
@@ -208,7 +209,8 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
   /**
    * @return vector<Expression*>
    */
-  antlrcpp::Any visitListLiteral(CypherParser::ListLiteralContext *ctx) override;
+  antlrcpp::Any visitListLiteral(
+      CypherParser::ListLiteralContext *ctx) override;
 
   /**
    * @return GraphDbTypes::Property
@@ -357,12 +359,26 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
       CypherParser::Expression4Context *ctx) override;
 
   /**
-   * Element of a list, range of a list...
+   * IS NULL, IS NOT NULL, ...
    *
    * @return Expression*
    */
-  antlrcpp::Any visitExpression3(
-      CypherParser::Expression3Context *ctx) override;
+  antlrcpp::Any visitExpression3a(
+      CypherParser::Expression3aContext *ctx) override;
+
+  /**
+   * List indexing and slicing.
+   *
+   * @return Expression*
+   */
+  antlrcpp::Any visitExpression3b(
+      CypherParser::Expression3bContext *ctx) override;
+
+  /**
+   * Does nothing, everything is done in visitExpression3b.
+   */
+  antlrcpp::Any visitListIndexingOrSlicing(
+      CypherParser::ListIndexingOrSlicingContext *ctx) override;
 
   /**
    * Property lookup, test for node labels existence...

src/query/frontend/opencypher/grammar/Cypher.g4

@@ -151,9 +151,15 @@ expression6 : expression5 ( ( SP? '*' SP? expression5 ) | ( SP? '/' SP? expressi
 
 expression5 : expression4 ( SP? '^' SP? expression4 )* ;
 
-expression4 : ( ( '+' | '-' ) SP? )* expression3 ;
+expression4 : ( ( '+' | '-' ) SP? )* expression3a ;
 
-expression3 : expression2 ( ( SP? '[' expression ']' ) | ( SP? '[' expression? '..' expression? ']' ) | ( ( ( SP? '=~' ) | ( SP IN ) | ( SP STARTS SP WITH ) | ( SP ENDS SP WITH ) | ( SP CONTAINS ) ) SP? expression2 ) | ( SP IS SP CYPHERNULL ) | ( SP IS SP NOT SP CYPHERNULL ) )* ;
+expression3a : expression3b ( ( ( ( SP? '=~' ) | ( SP IN ) | ( SP STARTS SP WITH ) | ( SP ENDS SP WITH ) | ( SP CONTAINS ) ) SP? expression2 ) | ( SP IS SP CYPHERNULL ) | ( SP IS SP NOT SP CYPHERNULL ) )* ;
+
+expression3b : expression2 ( SP? listIndexingOrSlicing )* ;
+
+listIndexingOrSlicing : ( '[' SP? expression SP? ']' )
+                      | ( '[' SP? lower_bound=expression? SP? '..' SP? upper_bound=expression? SP? ']' )
+                      ;
 
 expression2 : atom ( SP? ( propertyLookup | nodeLabels ) )* ;
 

tests/unit/cypher_main_visitor.cpp

@@ -68,7 +68,7 @@ TEST(CypherMainVisitorTest, PropertyLookup) {
             ast_generator.db_accessor_->property("x"));
 }
 
-TEST(CypherMainVisitor, ReturnNoDistinctNoBagSemantics) {
+TEST(CypherMainVisitorTest, ReturnNoDistinctNoBagSemantics) {
   AstGenerator ast_generator("RETURN x");
   auto *query = ast_generator.query_;
   ASSERT_EQ(query->clauses_.size(), 1U);
@@ -80,7 +80,7 @@ TEST(CypherMainVisitor, ReturnNoDistinctNoBagSemantics) {
   ASSERT_FALSE(return_clause->body_.distinct);
 }
 
-TEST(CypherMainVisitor, ReturnDistinct) {
+TEST(CypherMainVisitorTest, ReturnDistinct) {
   AstGenerator ast_generator("RETURN DISTINCT x");
   auto *query = ast_generator.query_;
   ASSERT_EQ(query->clauses_.size(), 1U);
@@ -88,7 +88,7 @@ TEST(CypherMainVisitor, ReturnDistinct) {
   ASSERT_TRUE(return_clause->body_.distinct);
 }
 
-TEST(CypherMainVisitor, ReturnLimit) {
+TEST(CypherMainVisitorTest, ReturnLimit) {
   AstGenerator ast_generator("RETURN x LIMIT 5");
   auto *query = ast_generator.query_;
   ASSERT_EQ(query->clauses_.size(), 1U);
@@ -99,7 +99,7 @@ TEST(CypherMainVisitor, ReturnLimit) {
   ASSERT_EQ(literal->value_.Value<int64_t>(), 5);
 }
 
-TEST(CypherMainVisitor, ReturnSkip) {
+TEST(CypherMainVisitorTest, ReturnSkip) {
   AstGenerator ast_generator("RETURN x SKIP 5");
   auto *query = ast_generator.query_;
   ASSERT_EQ(query->clauses_.size(), 1U);
@@ -110,7 +110,7 @@ TEST(CypherMainVisitor, ReturnSkip) {
   ASSERT_EQ(literal->value_.Value<int64_t>(), 5);
 }
 
-TEST(CypherMainVisitor, ReturnOrderBy) {
+TEST(CypherMainVisitorTest, ReturnOrderBy) {
   AstGenerator ast_generator("RETURN x, y, z ORDER BY z ASC, x, y DESC");
   auto *query = ast_generator.query_;
   ASSERT_EQ(query->clauses_.size(), 1U);
@@ -345,6 +345,38 @@ TEST(CypherMainVisitorTest, ComparisonOperators) {
 
 #undef CHECK_COMPARISON
 
+TEST(CypherMainVisitorTest, ListIndexingOperator) {
+  AstGenerator ast_generator("RETURN [1,2,3] [ 2 ]");
+  auto *query = ast_generator.query_;
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  auto *list_index_op = dynamic_cast<ListIndexingOperator *>(
+      return_clause->body_.named_expressions[0]->expression_);
+  ASSERT_TRUE(list_index_op);
+  auto *list = dynamic_cast<ListLiteral *>(list_index_op->expression1_);
+  EXPECT_TRUE(list);
+  auto *index = dynamic_cast<PrimitiveLiteral *>(list_index_op->expression2_);
+  ASSERT_EQ(index->value_.Value<int64_t>(), 2);
+}
+
+TEST(CypherMainVisitorTest, ListSlicingOperatorNoBounds) {
+  ASSERT_THROW(AstGenerator("RETURN [1,2,3] [ .. ]"), SemanticException);
+}
+
+TEST(CypherMainVisitorTest, ListSlicingOperator) {
+  AstGenerator ast_generator("RETURN [1,2,3] [ .. 2 ]");
+  auto *query = ast_generator.query_;
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  auto *list_slicing_op = dynamic_cast<ListSlicingOperator *>(
+      return_clause->body_.named_expressions[0]->expression_);
+  ASSERT_TRUE(list_slicing_op);
+  auto *list = dynamic_cast<ListLiteral *>(list_slicing_op->list_);
+  EXPECT_TRUE(list);
+  EXPECT_FALSE(list_slicing_op->lower_bound_);
+  auto *upper_bound =
+      dynamic_cast<PrimitiveLiteral *>(list_slicing_op->upper_bound_);
+  EXPECT_EQ(upper_bound->value_.Value<int64_t>(), 2);
+}
+
 TEST(CypherMainVisitorTest, IsNull) {
   AstGenerator ast_generator("RETURN 2 iS NulL");
   auto *query = ast_generator.query_;
@@ -1032,5 +1064,4 @@ TEST(CypherMainVisitorTest, Unwind) {
 TEST(CypherMainVisitorTest, UnwindWithoutAsError) {
   EXPECT_THROW(AstGenerator("UNWIND [1,2,3] RETURN 42"), SyntaxException);
 }
-
 }