pr: Add generic operation folders.

trailofbits · Jan 15, 2025 · 2dfb27c · 2dfb27c
1 parent f6a8b40
commit 2dfb27c
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 11 deletions.
diff --git a/include/vast/Dialect/Parser/Ops.td b/include/vast/Dialect/Parser/Ops.td
@@ -19,6 +19,8 @@ def Parser_Source
 {
   let summary = "Source of parsed data.";
 
+  let hasFolder = 1;
+
   let assemblyFormat = [{
     $arguments attr-dict `:` functional-type($arguments, $result)
   }];
@@ -31,6 +33,8 @@ def Paser_Sink
 {
   let summary = "Sink of parsed data.";
 
+  let hasFolder = 1;
+
   let assemblyFormat = [{
     $arguments attr-dict `:` functional-type($arguments, $result)
   }];
@@ -43,6 +47,8 @@ def Parser_Parse
 {
   let summary = "Parsing operation data.";
 
+  let hasFolder = 1;
+
   let assemblyFormat = [{
     $arguments attr-dict `:` functional-type($arguments, $result)
   }];
@@ -69,20 +75,24 @@ def Parse_MaybeParse
 {
   let summary = "Maybe parsing operation data.";
 
+  let hasFolder = 1;
+
   let assemblyFormat = [{
     $arguments attr-dict `:` functional-type($arguments, $result)
   }];
 }
 
 def Parse_Cast
   : Parser_Op< "cast" >
-  , Arguments< (ins Parser_AnyDataType:$arguments) >
+  , Arguments< (ins Parser_AnyDataType:$operand) >
   , Results< (outs Parser_AnyDataType:$result) >
 {
   let summary = "Casting operation.";
 
+  let hasFolder = 1;
+
   let assemblyFormat = [{
-    $arguments attr-dict `:` functional-type($arguments, $result)
+    $operand attr-dict `:` functional-type($operand, $result)
   }];
 }
 

diff --git a/lib/vast/Dialect/Parser/Ops.cpp b/lib/vast/Dialect/Parser/Ops.cpp
@@ -15,20 +15,53 @@ using namespace vast::pr;
 
 namespace vast::pr {
 
-    using fold_result_t = ::llvm::SmallVectorImpl< ::mlir::OpFoldResult >;
+    using fold_result = ::mlir::OpFoldResult;
+    using fold_results = ::llvm::SmallVectorImpl< fold_result >;
 
-    logical_result NoParse::fold(FoldAdaptor adaptor, fold_result_t &results) {
-        auto change = mlir::failure();
-        auto op = getOperation();
+    template< typename op_t >
+    logical_result forward_same_operation(
+        op_t op, auto adaptor, fold_results &results
+    ) {
+        if (op.getNumOperands() == 1 && op.getNumResults() == 1) {
+            if (auto operand = op.getOperand(0); mlir::isa< op_t >(operand.getDefiningOp())) {
+                if (operand.getType() == op->getOpResult(0).getType()) {
+                    results.push_back(operand);
+                    return mlir::success();
+                }
+            }
+        }
+
+        return mlir::failure();
+    }
+
+    logical_result Source::fold(FoldAdaptor adaptor, fold_results &results) {
+        return forward_same_operation(*this, adaptor, results);
+    }
+
+    logical_result Sink::fold(FoldAdaptor adaptor, fold_results &results) {
+        return forward_same_operation(*this, adaptor, results);
+    }
+
+    logical_result Parse::fold(FoldAdaptor adaptor, fold_results &results) {
+        return forward_same_operation(*this, adaptor, results);
+    }
+
+    logical_result NoParse::fold(FoldAdaptor adaptor, fold_results &results) {
+        return forward_same_operation(*this, adaptor, results);
+    }
+
+    logical_result MaybeParse::fold(FoldAdaptor adaptor, fold_results &results) {
+        return forward_same_operation(*this, adaptor, results);
+    }
 
-        for (auto [idx, operand] : llvm::reverse(llvm::enumerate(getOperands()))) {
-            if (auto noparse = mlir::dyn_cast< NoParse >(operand.getDefiningOp())) {
-                op->eraseOperand(idx);
-                change = mlir::success();
+    fold_result Cast::fold(FoldAdaptor adaptor) {
+        if (auto operand = getOperand(); mlir::isa< Cast >(operand.getDefiningOp())) {
+            if (operand.getType() == getType()) {
+                return operand;
             }
         }
 
-        return change;
+        return {};
     }
 
 } // namespace vast::pr