commit 2ce3b577fb44a28c8021bda739fc1f65e8797fbf
author Uday Bondhugula <bondhugula@google.com> Mon Aug 13 11:30:36 2018 -0700
committer jpienaar <jpienaar@google.com> Fri Mar 29 12:58:45 2019 -0700
tree 223e3c7ed0a02645f89f0c0e0a03ea31bf9e1075
parent c03c695772f409f9ee0a55229983c98cb77d5403

Add MLFunction::getReturnStmt.

PiperOrigin-RevId: 208514441

include/mlir/IR/MLFunction.h

diff --git a/include/mlir/IR/MLFunction.h b/include/mlir/IR/MLFunction.h
index 248828c..ed12afc 100644
--- a/include/mlir/IR/MLFunction.h
+++ b/include/mlir/IR/MLFunction.h
@@ -80,6 +80,10 @@
 
   ~MLFunction();
 
+  // Return the 'return' statement of this MLFunction.
+  const OperationStmt *getReturnStmt() const;
+  OperationStmt *getReturnStmt();
+
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(const Function *func) {
     return func->getKind() == Function::Kind::MLFunc;

lib/IR/Function.cpp

diff --git a/lib/IR/Function.cpp b/lib/IR/Function.cpp
index 4719b35..1c89de8 100644
--- a/lib/IR/Function.cpp
+++ b/lib/IR/Function.cpp
@@ -152,3 +152,11 @@
   this->~MLFunction();
   free(this);
 }
+
+const OperationStmt *MLFunction::getReturnStmt() const {
+  return cast<OperationStmt>(&back());
+}
+
+OperationStmt *MLFunction::getReturnStmt() {
+  return cast<OperationStmt>(&back());
+}