1 files changed, 97 insertions, 0 deletions
diff --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp
index fd9fda53e68..13bd0bcd8f3 100644
--- a/mlir/lib/IR/Value.cpp
+++ b/mlir/lib/IR/Value.cpp
@@ -9,8 +9,63 @@
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/StandardTypes.h"
 using namespace mlir;
 
+/// Construct a value.
+Value::Value(detail::BlockArgumentImpl *impl)
+    : ownerAndKind(impl, Kind::BlockArgument) {}
+Value::Value(Operation *op, unsigned resultNo) {
+  assert(op->getNumResults() > resultNo && "invalid result number");
+  if (LLVM_LIKELY(canPackResultInline(resultNo))) {
+    ownerAndKind = {op, static_cast<Kind>(resultNo)};
+    return;
+  }
+
+  // If we can't pack the result directly, we need to represent this as a
+  // trailing result.
+  unsigned trailingResultNo =
+      resultNo - static_cast<unsigned>(Kind::TrailingOpResult);
+  ownerAndKind = {op->getTrailingResult(trailingResultNo),
+                  Kind::TrailingOpResult};
+}
+
+/// Return the type of this value.
+Type Value::getType() const {
+  if (BlockArgument arg = dyn_cast<BlockArgument>())
+    return arg.getType();
+
+  // If this is an operation result, query the parent operation.
+  OpResult result = cast<OpResult>();
+  Operation *owner = result.getOwner();
+  if (owner->hasSingleResult)
+    return owner->resultType;
+  return owner->resultType.cast<TupleType>().getType(result.getResultNumber());
+}
+
+/// Mutate the type of this Value to be of the specified type.
+void Value::setType(Type newType) {
+  if (BlockArgument arg = dyn_cast<BlockArgument>())
+    return arg.setType(newType);
+  OpResult result = cast<OpResult>();
+
+  // If the owner has a single result, simply update it directly.
+  Operation *owner = result.getOwner();
+  if (owner->hasSingleResult) {
+    owner->resultType = newType;
+    return;
+  }
+  unsigned resultNo = result.getResultNumber();
+
+  // Otherwise, rebuild the tuple if the new type is different from the current.
+  auto curTypes = owner->resultType.cast<TupleType>().getTypes();
+  if (curTypes[resultNo] == newType)
+    return;
+  auto newTypes = llvm::to_vector<4>(curTypes);
+  newTypes[resultNo] = newType;
+  owner->resultType = TupleType::get(newTypes, newType.getContext());
+}
+
 /// If this value is the result of an Operation, return the operation that
 /// defines it.
 Operation *Value::getDefiningOp() const {
@@ -84,6 +139,48 @@ bool Value::use_empty() const {
 }
 
 //===----------------------------------------------------------------------===//
+// OpResult
+//===----------------------------------------------------------------------===//
+
+/// Returns the operation that owns this result.
+Operation *OpResult::getOwner() const {
+  // If the result is in-place, the `owner` is the operation.
+  if (LLVM_LIKELY(getKind() != Kind::TrailingOpResult))
+    return reinterpret_cast<Operation *>(ownerAndKind.getPointer());
+
+  // Otherwise, we need to do some arithmetic to get the operation pointer.
+  // Move the trailing owner to the start of the array.
+  auto *trailingIt =
+      static_cast<detail::TrailingOpResult *>(ownerAndKind.getPointer());
+  trailingIt -= trailingIt->trailingResultNumber;
+
+  // This point is the first trailing object after the operation. So all we need
+  // to do here is adjust for the operation size.
+  return reinterpret_cast<Operation *>(trailingIt) - 1;
+}
+
+/// Return the result number of this result.
+unsigned OpResult::getResultNumber() const {
+  // If the result is in-place, we can use the kind directly.
+  if (LLVM_LIKELY(getKind() != Kind::TrailingOpResult))
+    return static_cast<unsigned>(ownerAndKind.getInt());
+  // Otherwise, we add the number of inline results to the trailing owner.
+  auto *trailingIt =
+      static_cast<detail::TrailingOpResult *>(ownerAndKind.getPointer());
+  unsigned trailingNumber = trailingIt->trailingResultNumber;
+  return trailingNumber + static_cast<unsigned>(Kind::TrailingOpResult);
+}
+
+/// Given a number of operation results, returns the number that need to be
+/// stored as trailing.
+unsigned OpResult::getNumTrailing(unsigned numResults) {
+  // If we can pack all of the results, there is no need for additional storage.
+  if (numResults <= static_cast<unsigned>(Kind::TrailingOpResult))
+    return 0;
+  return numResults - static_cast<unsigned>(Kind::TrailingOpResult);
+}
+
+//===----------------------------------------------------------------------===//
 // BlockOperand
 //===----------------------------------------------------------------------===//