[SandboxIR] sandboxir::Use Part 2: uses/users

This patch implements uses/users for sandboxir::Value, including
functions like getNumUses(), hasNUsesOrMore() etc.

Author: vporpo

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -88,8 +88,10 @@ class Use {
       : LLVMUse(LLVMUse), User(User), Ctx(&Ctx) {}
   Use() : LLVMUse(nullptr), Ctx(nullptr) {}
 
+  friend class Value;              // For constructor
   friend class User;               // For constructor
   friend class OperandUseIterator; // For constructor
+  friend class UserUseIterator;    // For accessing members
 
 public:
   operator Value *() const { return get(); }
@@ -137,6 +139,31 @@ class OperandUseIterator {
   }
 };
 
+/// Returns user edge when dereferenced.
+class UserUseIterator {
+  Use Use;
+  /// Don't let the user create a non-empty UserUseIterator.
+  UserUseIterator(const class Use &Use) : Use(Use) {}
+  friend class Value; // For constructor
+
+public:
+  using difference_type = std::ptrdiff_t;
+  using value_type = class Use;
+  using pointer = value_type *;
+  using reference = value_type &;
+  using iterator_category = std::input_iterator_tag;
+
+  UserUseIterator() = default;
+  value_type operator*() const { return Use; }
+  UserUseIterator &operator++();
+  bool operator==(const UserUseIterator &Other) const {
+    return Use == Other.Use;
+  }
+  bool operator!=(const UserUseIterator &Other) const {
+    return !(*this == Other);
+  }
+};
+
 /// A SandboxIR Value has users. This is the base class.
 class Value {
 public:
@@ -188,9 +215,77 @@ class Value {
   virtual ~Value() = default;
   ClassID getSubclassID() const { return SubclassID; }
 
+  using use_iterator = UserUseIterator;
+  using const_use_iterator = UserUseIterator;
+
+  use_iterator use_begin();
+  const_use_iterator use_begin() const {
+    return const_cast<Value *>(this)->use_begin();
+  }
+  use_iterator use_end() { return use_iterator(Use(nullptr, nullptr, Ctx)); }
+  const_use_iterator use_end() const {
+    return const_cast<Value *>(this)->use_end();
+  }
+
+  iterator_range<use_iterator> uses() {
+    return make_range<use_iterator>(use_begin(), use_end());
+  }
+  iterator_range<const_use_iterator> uses() const {
+    return make_range<const_use_iterator>(use_begin(), use_end());
+  }
+
+  /// Helper for mapped_iterator.
+  struct UseToUser {
+    User *operator()(const Use &Use) const { return &*Use.getUser(); }
+  };
+
+  using user_iterator = mapped_iterator<sandboxir::UserUseIterator, UseToUser>;
+  using const_user_iterator = user_iterator;
+
+  user_iterator user_begin();
+  user_iterator user_end() {
+    return user_iterator(Use(nullptr, nullptr, Ctx), UseToUser());
+  }
+  const_user_iterator user_begin() const {
+    return const_cast<Value *>(this)->user_begin();
+  }
+  const_user_iterator user_end() const {
+    return const_cast<Value *>(this)->user_end();
+  }
+
+  iterator_range<user_iterator> users() {
+    return make_range<user_iterator>(user_begin(), user_end());
+  }
+  iterator_range<const_user_iterator> users() const {
+    return make_range<const_user_iterator>(user_begin(), user_end());
+  }
+  /// \Returns the number of user edges (not necessarily to unique users).
+  /// WARNING: This is a linear-time operation.
+  unsigned getNumUses() const;
+  /// Return true if this value has N uses or more.
+  /// This is logically equivalent to getNumUses() >= N.
+  /// WARNING: This can be expensive, as it is linear to the number of users.
+  bool hasNUsesOrMore(unsigned Num) const {
+    unsigned Cnt = 0;
+    for (auto It = use_begin(), ItE = use_end(); It != ItE; ++It) {
+      if (++Cnt >= Num)
+        return true;
+    }
+    return false;
+  }
+  /// Return true if this Value has exactly N uses.
+  bool hasNUses(unsigned Num) const {
+    unsigned Cnt = 0;
+    for (auto It = use_begin(), ItE = use_end(); It != ItE; ++It) {
+      if (++Cnt > Num)
+        return false;
+    }
+    return Cnt == Num;
+  }
+
   Type *getType() const { return Val->getType(); }
 
-  Context &getContext() const;
+  Context &getContext() const { return Ctx; }
 #ifndef NDEBUG
   /// Should crash if there is something wrong with the instruction.
   virtual void verify() const = 0;

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/SandboxIR/SandboxIR.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/Support/Debug.h"
 #include <sstream>
@@ -58,13 +59,50 @@ OperandUseIterator &OperandUseIterator::operator++() {
   return *this;
 }
 
+UserUseIterator &UserUseIterator::operator++() {
+  llvm::Use *&LLVMUse = Use.LLVMUse;
+  assert(LLVMUse != nullptr && "Already at end!");
+  LLVMUse = LLVMUse->getNext();
+  if (LLVMUse == nullptr) {
+    Use.User = nullptr;
+    return *this;
+  }
+  auto *Ctx = Use.Ctx;
+  auto *LLVMUser = LLVMUse->getUser();
+  Use.User = cast_or_null<sandboxir::User>(Ctx->getValue(LLVMUser));
+  return *this;
+}
+
 Value::Value(ClassID SubclassID, llvm::Value *Val, Context &Ctx)
     : SubclassID(SubclassID), Val(Val), Ctx(Ctx) {
 #ifndef NDEBUG
   UID = Ctx.getNumValues();
 #endif
 }
 
+Value::use_iterator Value::use_begin() {
+  llvm::Use *LLVMUse = nullptr;
+  if (Val->use_begin() != Val->use_end())
+    LLVMUse = &*Val->use_begin();
+  User *User = LLVMUse != nullptr ? cast_or_null<sandboxir::User>(Ctx.getValue(
+                                        Val->use_begin()->getUser()))
+                                  : nullptr;
+  return use_iterator(Use(LLVMUse, User, Ctx));
+}
+
+Value::user_iterator Value::user_begin() {
+  auto UseBegin = Val->use_begin();
+  auto UseEnd = Val->use_end();
+  bool AtEnd = UseBegin == UseEnd;
+  llvm::Use *LLVMUse = AtEnd ? nullptr : &*UseBegin;
+  User *User =
+      AtEnd ? nullptr
+            : cast_or_null<sandboxir::User>(Ctx.getValue(&*LLVMUse->getUser()));
+  return user_iterator(Use(LLVMUse, User, Ctx), UseToUser());
+}
+
+unsigned Value::getNumUses() const { return range_size(Val->users()); }
+
 #ifndef NDEBUG
 std::string Value::getName() const {
   std::stringstream SS;

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -143,6 +143,8 @@ define i32 @foo(i32 %v0, i32 %v1) {
     // Check Use.get().
     sandboxir::Value *Op = Use.get();
     EXPECT_EQ(Op, Ctx.getValue(LLVMI0->getOperand(OpIdx)));
+    // Check Use.getUser().
+    EXPECT_EQ(Use.getUser(), I0);
     // Check implicit cast to Value.
     sandboxir::Value *Cast = Use;
     EXPECT_EQ(Cast, Op);
@@ -180,6 +182,59 @@ User:   %add0 = add i32 %v0, %v1 ; SB4. (Opaque)
 OperandNo: 0
 )IR");
 #endif // NDEBUG
+
+  // Check Value.user_begin().
+  sandboxir::Value::user_iterator UIt = I0->user_begin();
+  sandboxir::User *U = *UIt;
+  EXPECT_EQ(U, Ret);
+  // Check Value.uses().
+  unsigned UseCnt = 0;
+  for (sandboxir::Use Use : I0->uses()) {
+    EXPECT_EQ(Use.getUser(), Ret);
+    ++UseCnt;
+  }
+  EXPECT_EQ(UseCnt, 1u);
+  // Check Value.users().
+  unsigned UserCnt = 0;
+  for (sandboxir::User *U : I0->users()) {
+    EXPECT_EQ(U, Ret);
+    ++UserCnt;
+  }
+  EXPECT_EQ(UseCnt, 1u);
+  // Check Value.getNumUses().
+  EXPECT_EQ(I0->getNumUses(), 1u);
+  // Check Value.hasNUsesOrMore().
+  EXPECT_TRUE(I0->hasNUsesOrMore(0u));
+  EXPECT_TRUE(I0->hasNUsesOrMore(1u));
+  EXPECT_FALSE(I0->hasNUsesOrMore(2u));
+  // Check Value.hasNUses().
+  EXPECT_FALSE(I0->hasNUses(0u));
+  EXPECT_TRUE(I0->hasNUses(1u));
+  EXPECT_FALSE(I0->hasNUses(2u));
+}
+
+// Check that the operands/users are counted correctly.
+//  I1
+// /  \
+// \  /
+//  I2
+TEST_F(SandboxIRTest, DuplicateUses) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v) {
+  %I1 = add i8 %v, %v
+  %I2 = add i8 %I1, %I1
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto *F = Ctx.createFunction(&LLVMF);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  auto *I1 = &*It++;
+  auto *I2 = &*It++;
+  EXPECT_EQ(range_size(I1->users()), 2u);
+  EXPECT_EQ(range_size(I2->operands()), 2u);
 }
 
 TEST_F(SandboxIRTest, Function) {