hl: Add helpers for generic selection expr.

trailofbits · Sep 9, 2024 · bc8f398 · bc8f398
1 parent 9d2b242
commit bc8f398
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/include/vast/Dialect/HighLevel/HighLevelOps.td b/include/vast/Dialect/HighLevel/HighLevelOps.td
@@ -1734,6 +1734,12 @@ def HighLevel_GenericSelectionExpr
     }] >
   ];
 
+  let extraClassDeclaration = [{
+    std::optional< mlir::Region *> getResultRegion();
+    bool isExprPredicate();
+    bool isTypePredicate();
+  }];
+
   let assemblyFormat = [{ attr-dict `match` `:` ($control^)? ($controlType^)? $body `->` type($result)}];
 }
 

diff --git a/lib/vast/Dialect/HighLevel/HighLevelOps.cpp b/lib/vast/Dialect/HighLevel/HighLevelOps.cpp
@@ -875,6 +875,27 @@ namespace vast::hl
         return std::optional(getCompatibleAttr().getValue());
     }
 
+    //
+    // GenericSelectionExpr
+    //
+
+    std::optional< region_t *> GenericSelectionExpr::getResultRegion() {
+        if (auto selected = getSelected()) {
+            auto op_it = getBody().op_begin();
+            std::advance(op_it, selected.value().getZExtValue());
+            if (auto assoc = mlir::dyn_cast< hl::GenericAssocExpr >(*op_it))
+            return { &assoc.getBody() };
+        }
+        return std::nullopt;
+    }
+
+    bool GenericSelectionExpr::isExprPredicate() {
+        return !getControl().empty();
+    }
+
+    bool GenericSelectionExpr::isTypePredicate() {
+        return (*this)->hasAttr("matchType");
+    }
 }
 
 //===----------------------------------------------------------------------===//