1 files changed, 6 insertions, 2 deletions
diff --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp
index a2cb9910ab8..6418b062dc1 100644
--- a/mlir/lib/IR/Value.cpp
+++ b/mlir/lib/IR/Value.cpp
@@ -35,8 +35,6 @@ Function *Value::getFunction() {
     return cast<BlockArgument>(this)->getFunction();
   case Value::Kind::InstResult:
     return getDefiningInst()->getFunction();
-  case Value::Kind::ForInst:
-    return cast<ForInst>(this)->getFunction();
   }
 }
 
@@ -83,3 +81,9 @@ Function *BlockArgument::getFunction() {
     return owner->getFunction();
   return nullptr;
 }
+
+/// Returns if the current argument is a function argument.
+bool BlockArgument::isFunctionArgument() const {
+  auto *containingFn = getFunction();
+  return containingFn && &containingFn->front() == getOwner();
+}