From d2f9affb4c03f4a99c6ed955a13e1604011b4479 Mon Sep 17 00:00:00 2001
From: Mislav Bradac <mislav.bradac@memgraph.io>
Date: Thu, 4 May 2017 14:49:29 +0200
Subject: [PATCH] Add labels and edge type test nodes to ast

Reviewers: teon.banek

Reviewed By: teon.banek

Differential Revision: https://phabricator.memgraph.io/D343
---
 src/query/frontend/ast/ast.hpp         | 42 ++++++++++++++++++++++++++
 src/query/frontend/ast/ast_visitor.hpp |  9 ++++--
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/query/frontend/ast/ast.hpp b/src/query/frontend/ast/ast.hpp
index 9a6e1d62c..80c673892 100644
--- a/src/query/frontend/ast/ast.hpp
+++ b/src/query/frontend/ast/ast.hpp
@@ -487,6 +487,48 @@ class PropertyLookup : public Expression {
       : Expression(uid), expression_(expression), property_(property) {}
 };
 
+class LabelsTest : public Expression {
+  friend class AstTreeStorage;
+
+ public:
+  void Accept(TreeVisitorBase &visitor) override {
+    if (visitor.PreVisit(*this)) {
+      visitor.Visit(*this);
+      expression_->Accept(visitor);
+      visitor.PostVisit(*this);
+    }
+  }
+
+  Expression *expression_ = nullptr;
+  std::vector<GraphDbTypes::Label> labels_;
+
+ protected:
+  LabelsTest(int uid, Expression *expression,
+             std::vector<GraphDbTypes::Label> labels)
+      : Expression(uid), expression_(expression), labels_(labels) {}
+};
+
+class EdgeTypeTest : public Expression {
+  friend class AstTreeStorage;
+
+ public:
+  void Accept(TreeVisitorBase &visitor) override {
+    if (visitor.PreVisit(*this)) {
+      visitor.Visit(*this);
+      expression_->Accept(visitor);
+      visitor.PostVisit(*this);
+    }
+  }
+
+  Expression *expression_ = nullptr;
+  std::vector<GraphDbTypes::EdgeType> edge_types_;
+
+ protected:
+  EdgeTypeTest(int uid, Expression *expression,
+               std::vector<GraphDbTypes::EdgeType> edge_types)
+      : Expression(uid), expression_(expression), edge_types_(edge_types) {}
+};
+
 class Function : public Expression {
   friend class AstTreeStorage;
 
diff --git a/src/query/frontend/ast/ast_visitor.hpp b/src/query/frontend/ast/ast_visitor.hpp
index c4a6ed134..65d1beacb 100644
--- a/src/query/frontend/ast/ast_visitor.hpp
+++ b/src/query/frontend/ast/ast_visitor.hpp
@@ -9,6 +9,8 @@ class Query;
 class NamedExpression;
 class Identifier;
 class PropertyLookup;
+class LabelsTest;
+class EdgeTypeTest;
 class Aggregation;
 class Function;
 class Create;
@@ -57,8 +59,9 @@ using TreeVisitorBase = ::utils::Visitor<
     LessOperator, GreaterOperator, LessEqualOperator, GreaterEqualOperator,
     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>;
+    ListLiteral, PropertyLookup, LabelsTest, EdgeTypeTest, Aggregation,
+    Function, Create, Match, Return, With, Pattern, NodeAtom, EdgeAtom, Delete,
+    Where, SetProperty, SetProperties, SetLabels, RemoveProperty, RemoveLabels,
+    Merge, Unwind>;
 
 }  // namespace query