[SandboxIR] sandboxir::Use Part 1: operands

This patch adds the Use class which is used by the new member functions
of the User class that provide access to the operands.

Author: vporpo

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -71,6 +71,69 @@ namespace sandboxir {
 class Function;
 class Context;
 class Instruction;
+class User;
+class Value;
+
+/// Represents a Def-use/Use-def edge in SandboxIR.
+/// NOTE: Unlike llvm::Use, this is not an integral part of the use-def chains.
+/// It is also not uniqued and is currently passed by value, so you can have
+/// more than one sandboxir::Use objects for the same use-def edge.
+class Use {
+  llvm::Use *LLVMUse;
+  User *Usr;
+  Context *Ctx;
+
+  /// Don't allow the user to create a sandboxir::Use directly.
+  Use(llvm::Use *LLVMUse, User *Usr, Context &Ctx)
+      : LLVMUse(LLVMUse), Usr(Usr), Ctx(&Ctx) {}
+  Use() : LLVMUse(nullptr), Ctx(nullptr) {}
+
+  friend class User;               // For constructor
+  friend class OperandUseIterator; // For constructor
+
+public:
+  operator Value *() const { return get(); }
+  Value *get() const;
+  class User *getUser() const { return Usr; }
+  unsigned getOperandNo() const;
+  Context *getContext() const { return Ctx; }
+  bool operator==(const Use &Other) const {
+    assert(Ctx == Other.Ctx && "Contexts differ!");
+    return LLVMUse == Other.LLVMUse && Usr == Other.Usr;
+  }
+  bool operator!=(const Use &Other) const { return !(*this == Other); }
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const;
+  void dump() const;
+#endif // NDEBUG
+};
+
+/// Returns the operand edge when dereferenced.
+class OperandUseIterator {
+  Use Use;
+  /// Don't let the user create a non-empty OperandUseIterator.
+  OperandUseIterator(const class Use &Use) : Use(Use) {}
+  friend class User;                                  // For constructor
+#define DEF_INSTR(ID, OPC, CLASS) friend class CLASS; // For constructor
+#include "llvm/SandboxIR/SandboxIRValues.def"
+
+public:
+  using difference_type = std::ptrdiff_t;
+  using value_type = sandboxir::Use;
+  using pointer = value_type *;
+  using reference = value_type &;
+  using iterator_category = std::input_iterator_tag;
+
+  OperandUseIterator() = default;
+  value_type operator*() const;
+  OperandUseIterator &operator++();
+  bool operator==(const OperandUseIterator &Other) const {
+    return Use == Other.Use;
+  }
+  bool operator!=(const OperandUseIterator &Other) const {
+    return !(*this == Other);
+  }
+};
 
 /// A SandboxIR Value has users. This is the base class.
 class Value {
@@ -174,9 +237,61 @@ class User : public Value {
 protected:
   User(ClassID ID, llvm::Value *V, Context &Ctx) : Value(ID, V, Ctx) {}
 
+  /// \Returns the Use edge that corresponds to \p OpIdx.
+  /// Note: This is the default implementation that works for instructions that
+  /// match the underlying LLVM instruction. All others should use a different
+  /// implementation.
+  Use getOperandUseDefault(unsigned OpIdx, bool Verify) const;
+  virtual Use getOperandUseInternal(unsigned OpIdx, bool Verify) const = 0;
+  friend class OperandUseIterator; // for getOperandUseInternal()
+
+  /// The default implementation works only for single-LLVMIR-instruction
+  /// Users and only if they match exactly the LLVM instruction.
+  unsigned getUseOperandNoDefault(const Use &Use) const {
+    return Use.LLVMUse->getOperandNo();
+  }
+  /// \Returns the operand index of \p Use.
+  virtual unsigned getUseOperandNo(const Use &Use) const = 0;
+  friend unsigned Use::getOperandNo() const; // For getUseOperandNo()
+
 public:
   /// For isa/dyn_cast.
   static bool classof(const Value *From);
+  using op_iterator = OperandUseIterator;
+  using const_op_iterator = OperandUseIterator;
+  using op_range = iterator_range<op_iterator>;
+  using const_op_range = iterator_range<const_op_iterator>;
+
+  virtual op_iterator op_begin() {
+    assert(isa<llvm::User>(Val) && "Expect User value!");
+    return op_iterator(getOperandUseInternal(0, /*Verify=*/false));
+  }
+  virtual op_iterator op_end() {
+    assert(isa<llvm::User>(Val) && "Expect User value!");
+    return op_iterator(
+        getOperandUseInternal(getNumOperands(), /*Verify=*/false));
+  }
+  virtual const_op_iterator op_begin() const {
+    return const_cast<User *>(this)->op_begin();
+  }
+  virtual const_op_iterator op_end() const {
+    return const_cast<User *>(this)->op_end();
+  }
+
+  op_range operands() { return make_range<op_iterator>(op_begin(), op_end()); }
+  const_op_range operands() const {
+    return make_range<const_op_iterator>(op_begin(), op_end());
+  }
+  Value *getOperand(unsigned OpIdx) const { return getOperandUse(OpIdx).get(); }
+  /// \Returns the operand edge for \p OpIdx. NOTE: This should also work for
+  /// OpIdx == getNumOperands(), which is used for op_end().
+  Use getOperandUse(unsigned OpIdx) const {
+    return getOperandUseInternal(OpIdx, /*Verify=*/true);
+  }
+  virtual unsigned getNumOperands() const {
+    return isa<llvm::User>(Val) ? cast<llvm::User>(Val)->getNumOperands() : 0;
+  }
+
 #ifndef NDEBUG
   void verify() const override {
     assert(isa<llvm::User>(Val) && "Expected User!");
@@ -195,6 +310,9 @@ class Constant : public sandboxir::User {
   Constant(llvm::Constant *C, sandboxir::Context &SBCtx)
       : sandboxir::User(ClassID::Constant, C, SBCtx) {}
   friend class Context; // For constructor.
+  Use getOperandUseInternal(unsigned OpIdx, bool Verify) const final {
+    return getOperandUseDefault(OpIdx, Verify);
+  }
 
 public:
   /// For isa/dyn_cast.
@@ -203,6 +321,9 @@ class Constant : public sandboxir::User {
            From->getSubclassID() == ClassID::Function;
   }
   sandboxir::Context &getParent() const { return getContext(); }
+  unsigned getUseOperandNo(const Use &Use) const final {
+    return getUseOperandNoDefault(Use);
+  }
 #ifndef NDEBUG
   void verify() const final {
     assert(isa<llvm::Constant>(Val) && "Expected Constant!");
@@ -309,11 +430,17 @@ class OpaqueInst : public sandboxir::Instruction {
   OpaqueInst(ClassID SubclassID, llvm::Instruction *I, sandboxir::Context &Ctx)
       : sandboxir::Instruction(SubclassID, Opcode::Opaque, I, Ctx) {}
   friend class Context; // For constructor.
+  Use getOperandUseInternal(unsigned OpIdx, bool Verify) const final {
+    return getOperandUseDefault(OpIdx, Verify);
+  }
 
 public:
   static bool classof(const sandboxir::Value *From) {
     return From->getSubclassID() == ClassID::Opaque;
   }
+  unsigned getUseOperandNo(const Use &Use) const final {
+    return getUseOperandNoDefault(Use);
+  }
   unsigned getNumOfIRInstrs() const final { return 1u; }
 #ifndef NDEBUG
   void verify() const final {

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -13,6 +13,51 @@
 
 using namespace llvm::sandboxir;
 
+Value *Use::get() const { return Ctx->getValue(LLVMUse->get()); }
+
+unsigned Use::getOperandNo() const { return Usr->getUseOperandNo(*this); }
+
+#ifndef NDEBUG
+void Use::dump(raw_ostream &OS) const {
+  Value *Def = nullptr;
+  if (LLVMUse == nullptr)
+    OS << "<null> LLVM Use! ";
+  else
+    Def = Ctx->getValue(LLVMUse->get());
+  OS << "Def:  ";
+  if (Def == nullptr)
+    OS << "NULL";
+  else
+    OS << *Def;
+  OS << "\n";
+
+  OS << "User: ";
+  if (Usr == nullptr)
+    OS << "NULL";
+  else
+    OS << *Usr;
+  OS << "\n";
+
+  OS << "OperandNo: ";
+  if (Usr == nullptr)
+    OS << "N/A";
+  else
+    OS << getOperandNo();
+  OS << "\n";
+}
+
+void Use::dump() const { dump(dbgs()); }
+#endif // NDEBUG
+
+Use OperandUseIterator::operator*() const { return Use; }
+
+OperandUseIterator &OperandUseIterator::operator++() {
+  assert(Use.LLVMUse != nullptr && "Already at end!");
+  User *User = Use.getUser();
+  Use = User->getOperandUseInternal(Use.getOperandNo() + 1, /*Verify=*/false);
+  return *this;
+}
+
 Value::Value(ClassID SubclassID, llvm::Value *Val, Context &Ctx)
     : SubclassID(SubclassID), Val(Val), Ctx(Ctx) {
 #ifndef NDEBUG
@@ -71,6 +116,17 @@ void Argument::dump() const {
 }
 #endif // NDEBUG
 
+Use User::getOperandUseDefault(unsigned OpIdx, bool Verify) const {
+  assert((!Verify || OpIdx < getNumOperands()) && "Out of bounds!");
+  assert(isa<llvm::User>(Val) && "Non-users have no operands!");
+  llvm::Use *LLVMUse;
+  if (OpIdx != getNumOperands())
+    LLVMUse = &cast<llvm::User>(Val)->getOperandUse(OpIdx);
+  else
+    LLVMUse = cast<llvm::User>(Val)->op_end();
+  return Use(LLVMUse, const_cast<User *>(this), Ctx);
+}
+
 bool User::classof(const Value *From) {
   switch (From->getSubclassID()) {
 #define DEF_VALUE(ID, CLASS)

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -109,6 +109,79 @@ define void @foo(i32 %v1) {
 #endif
 }
 
+TEST_F(SandboxIRTest, Use) {
+  parseIR(C, R"IR(
+define i32 @foo(i32 %v0, i32 %v1) {
+  %add0 = add i32 %v0, %v1
+  ret i32 %add0
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  BasicBlock *LLVMBB = &*LLVMF.begin();
+  auto LLVMBBIt = LLVMBB->begin();
+  Instruction *LLVMI0 = &*LLVMBBIt++;
+
+  auto &F = *Ctx.createFunction(&LLVMF);
+  auto &BB = *F.begin();
+  auto *Arg0 = F.getArg(0);
+  auto *Arg1 = F.getArg(1);
+  auto It = BB.begin();
+  auto *I0 = &*It++;
+  auto *Ret = &*It++;
+
+  SmallVector<sandboxir::Argument *> Args{Arg0, Arg1};
+  unsigned OpIdx = 0;
+  for (sandboxir::Use Use : I0->operands()) {
+    // Check Use.getOperandNo().
+    EXPECT_EQ(Use.getOperandNo(), OpIdx);
+    // Check Use.getUser().
+    EXPECT_EQ(Use.getUser(), I0);
+    // Check Use.getContext().
+    EXPECT_EQ(Use.getContext(), &Ctx);
+    // Check Use.get().
+    sandboxir::Value *Op = Use.get();
+    EXPECT_EQ(Op, Ctx.getValue(LLVMI0->getOperand(OpIdx)));
+    // Check implicit cast to Value.
+    sandboxir::Value *Cast = Use;
+    EXPECT_EQ(Cast, Op);
+    // Check that Use points to the correct operand.
+    EXPECT_EQ(Op, Args[OpIdx]);
+    // Check getOperand().
+    EXPECT_EQ(Op, I0->getOperand(OpIdx));
+    // Check getOperandUse().
+    EXPECT_EQ(Use, I0->getOperandUse(OpIdx));
+    ++OpIdx;
+  }
+  EXPECT_EQ(OpIdx, 2u);
+
+  // Check Use.operator==() and Use.operator!=().
+  sandboxir::Use UseA = I0->getOperandUse(0);
+  sandboxir::Use UseB = I0->getOperandUse(0);
+  EXPECT_TRUE(UseA == UseB);
+  EXPECT_FALSE(UseA != UseB);
+
+  // Check getNumOperands().
+  EXPECT_EQ(I0->getNumOperands(), 2u);
+  EXPECT_EQ(Ret->getNumOperands(), 1u);
+
+  EXPECT_EQ(Ret->getOperand(0), I0);
+
+#ifndef NDEBUG
+  // Check Use.dump()
+  std::string Buff;
+  raw_string_ostream BS(Buff);
+  BS << "\n";
+  I0->getOperandUse(0).dump(BS);
+  EXPECT_EQ(Buff, R"IR(
+Def:  i32 %v0 ; SB1. (Argument)
+User:   %add0 = add i32 %v0, %v1 ; SB4. (Opaque)
+OperandNo: 0
+)IR");
+#endif // NDEBUG
+}
+
 TEST_F(SandboxIRTest, Function) {
   parseIR(C, R"IR(
 define void @foo(i32 %arg0, i32 %arg1) {