From 51e3db66f8a1cfcf8346b50e07459d9c9e5aa038 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ingo=20M=C3=BCller?= <ingomueller@google.com>
Date: Fri, 31 May 2024 08:18:29 +0000
Subject: [PATCH] [Substrait] Add folder for `project` op with empty
 `expressions`.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Ingo Müller <ingomueller@google.com>
---
 .../Dialect/Substrait/IR/SubstraitOps.td      |  1 +
 lib/Dialect/Substrait/IR/Substrait.cpp        | 11 +++++++++++
 test/Dialect/Substrait/canonicalize.mlir      | 19 +++++++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/include/structured/Dialect/Substrait/IR/SubstraitOps.td b/include/structured/Dialect/Substrait/IR/SubstraitOps.td
index 198e5f8c94df..7862f7cdd349 100644
--- a/include/structured/Dialect/Substrait/IR/SubstraitOps.td
+++ b/include/structured/Dialect/Substrait/IR/SubstraitOps.td
@@ -395,6 +395,7 @@ def Substrait_ProjectOp : Substrait_RelOp<"project", [
     $input attr-dict `:` type($input) `->` type($result) $expressions
   }];
   let hasRegionVerifier = 1;
+  let hasFolder = 1;
   let extraClassDefinition = [{
     /// Implement OpAsmOpInterface.
     ::llvm::StringRef $cppClass::getDefaultDialect() {
diff --git a/lib/Dialect/Substrait/IR/Substrait.cpp b/lib/Dialect/Substrait/IR/Substrait.cpp
index ed7d1c4020b2..3925d0a4c8bf 100644
--- a/lib/Dialect/Substrait/IR/Substrait.cpp
+++ b/lib/Dialect/Substrait/IR/Substrait.cpp
@@ -339,6 +339,17 @@ LogicalResult PlanRelOp::verifyRegions() {
   return verifyNamedStruct(getOperation(), fieldNames, tupleType);
 }
 
+OpFoldResult ProjectOp::fold(FoldAdaptor adaptor) {
+  Operation *terminator = adaptor.getExpressions().front().getTerminator();
+
+  // If the region does not yield any values, the the `project` has no effect.
+  if (terminator->getNumOperands() == 0) {
+    return getInput();
+  }
+
+  return {};
+}
+
 LogicalResult ProjectOp::verifyRegions() {
   // Verify that the expression block has a matching argument type.
   auto inputTupleType = llvm::cast<TupleType>(getInput().getType());
diff --git a/test/Dialect/Substrait/canonicalize.mlir b/test/Dialect/Substrait/canonicalize.mlir
index cf1e748d3e20..79ed18a8e725 100644
--- a/test/Dialect/Substrait/canonicalize.mlir
+++ b/test/Dialect/Substrait/canonicalize.mlir
@@ -72,3 +72,22 @@ substrait.plan version 0 : 42 : 1 {
     yield %5 : tuple<si1, si32>
   }
 }
+
+// -----
+
+// Check that empty `project` folded.
+
+// CHECK-LABEL: substrait.plan
+// CHECK-NEXT:    relation
+// CHECK-NEXT:      %[[V0:.*]] = named_table
+// CHECK-NEXT:      yield %[[V0]]
+
+substrait.plan version 0 : 42 : 1 {
+  relation {
+    %0 = named_table @t1 as ["a"] : tuple<si32>
+    %1 = project %0 : tuple<si32> -> tuple<si32> {
+    ^bb0(%arg0: tuple<si32>):
+    }
+    yield %1 : tuple<si32>
+  }
+}