MLIR: Fix SKIP <expr> and LIMIT <expr>

query/ir/codegen/DBProgramGenerator.cpp

@@ -74,6 +74,8 @@
 #include "Symbol.h"
 #include "SymbolChain.h"
 
+#include "columns/BinaryOperators.h"
+
 #include "BioAssert.h"
 #include "FatalException.h"
 #include "TuringException.h"
@@ -285,6 +287,91 @@ void flattenConjuncts(const Expr* expr, std::vector<const Expr*>& conjuncts) {
     conjuncts.push_back(expr);
 }
 
+int64_t applyConstantUnary(UnaryOperator op, int64_t operand) {
+    switch (op) {
+        case UnaryOperator::Plus:
+            return operand;
+        break;
+
+        case UnaryOperator::Minus:
+            return -operand;
+        break;
+
+        default:
+            throw TuringException(fmt::format("Unsupported unary operator in SKIP/LIMIT expression: {}",
+                                              UnaryOperatorDescription::value(op)));
+        break;
+    }
+}
+
+int64_t applyConstantBinary(BinaryOperator op, int64_t lhs, int64_t rhs) {
+    switch (op) {
+        case BinaryOperator::Add:
+            return Add {}(lhs, rhs);
+        break;
+
+        case BinaryOperator::Sub:
+            return Sub {}(lhs, rhs);
+        break;
+
+        case BinaryOperator::Mult:
+            return Mul {}(lhs, rhs);
+        break;
+
+        case BinaryOperator::Div:
+            return Div {}(lhs, rhs);
+        break;
+
+        case BinaryOperator::Mod:
+            return Mod {}(lhs, rhs);
+        break;
+
+        default:
+            throw TuringException(fmt::format("Unsupported operator in SKIP/LIMIT expression: {}",
+                                              BinaryOperatorDescription::value(op)));
+        break;
+    }
+}
+
+int64_t evaluateConstantInteger(const Expr* expr) {
+    const Expr::Kind kind = expr->getKind();
+
+    switch (kind) {
+        case Expr::Kind::LITERAL: {
+            const LiteralExpr* literalExpr = static_cast<const LiteralExpr*>(expr);
+            const Literal* literal = literalExpr->getLiteral();
+
+            if (literal->getKind() != Literal::Kind::INTEGER) {
+                throw TuringException("SKIP/LIMIT expression must evaluate to an integer");
+            }
+
+            const IntegerLiteral* integerLiteral = static_cast<const IntegerLiteral*>(literal);
+            return integerLiteral->getValue();
+        }
+        break;
+
+        case Expr::Kind::UNARY: {
+            const UnaryExpr* unaryExpr = static_cast<const UnaryExpr*>(expr);
+            const int64_t operand = evaluateConstantInteger(unaryExpr->getSubExpr());
+            return applyConstantUnary(unaryExpr->getOperator(), operand);
+        }
+        break;
+
+        case Expr::Kind::BINARY: {
+            const BinaryExpr* binaryExpr = static_cast<const BinaryExpr*>(expr);
+            const int64_t lhs = evaluateConstantInteger(binaryExpr->getLHS());
+            const int64_t rhs = evaluateConstantInteger(binaryExpr->getRHS());
+            return applyConstantBinary(binaryExpr->getOperator(), lhs, rhs);
+        }
+        break;
+
+        default:
+            throw TuringException(fmt::format("Unsupported expression in SKIP/LIMIT: {}",
+                                              ExprKindDescription::value(kind)));
+        break;
+    }
+}
+
 }
 
 DBProgramGenerator::DBProgramGenerator(mlir::ModuleOp* mainModule)
@@ -1333,9 +1420,13 @@ void DBProgramGenerator::generateOutput(const CypherAST* ast) {
 
     if (proj->hasSkip()) {
         const Expr* skipExpr = proj->getSkip()->getExpr();
-        const LiteralExpr* litExpr = static_cast<const LiteralExpr*>(skipExpr);
-        const IntegerLiteral* intLit = static_cast<const IntegerLiteral*>(litExpr->getLiteral());
-        const uint64_t skipCount = static_cast<uint64_t>(intLit->getValue());
+        const int64_t skipValue = evaluateConstantInteger(skipExpr);
+
+        if (skipValue < 0) {
+            throw TuringException("SKIP expression must be a non-negative integer");
+        }
+
+        const uint64_t skipCount = static_cast<uint64_t>(skipValue);
 
         llvm::SmallVector<size_t> skippedItems;
         llvm::SmallVector<mlir::Value> skipped;
@@ -1356,9 +1447,13 @@ void DBProgramGenerator::generateOutput(const CypherAST* ast) {
 
     if (proj->hasLimit()) {
         const Expr* limitExpr = proj->getLimit()->getExpr();
-        const LiteralExpr* litExpr = static_cast<const LiteralExpr*>(limitExpr);
-        const IntegerLiteral* intLit = static_cast<const IntegerLiteral*>(litExpr->getLiteral());
-        const uint64_t limitCount = static_cast<uint64_t>(intLit->getValue());
+        const int64_t limitValue = evaluateConstantInteger(limitExpr);
+
+        if (limitValue < 0) {
+            throw TuringException("LIMIT expression must be a non-negative integer");
+        }
+
+        const uint64_t limitCount = static_cast<uint64_t>(limitValue);
 
         llvm::SmallVector<size_t> limitedItems;
         llvm::SmallVector<mlir::Value> limited;

test/query/ir/LimitSkipTest.cpp

@@ -246,6 +246,21 @@ TEST_F(LimitSkipTest, matchedConstantSkipPastTheMatchedRowsEmitsNothing) {
     expectRows("MATCH (n) RETURN 5 SKIP 20", {});
 }
 
+TEST_F(LimitSkipTest, matchSkipConstantExpressionSkipsTheFoldedCount) {
+    const ValueRows expected = {{6}};
+    expectRows("MATCH (n) RETURN n SKIP 5 + 1 LIMIT 1", expected);
+}
+
+TEST_F(LimitSkipTest, matchLimitConstantExpressionKeepsTheFoldedCount) {
+    const ValueRows expected = {{1}, {2}};
+    expectRows("MATCH (n) RETURN n SKIP 1 LIMIT 1 + 1", expected);
+}
+
+TEST_F(LimitSkipTest, matchedConstantSkipMultiplicationKeepsTheSurvivingRows) {
+    const ValueRows expected = {{5}, {5}};
+    expectRows("MATCH (n) RETURN 5 SKIP 8 * 2", expected);
+}
+
 int main(int argc, char** argv) {
     return turing::test::turingTestMain(argc, argv);
 }