[llvm] [SandboxIR] Adds BasicBlock and adds functionality to Function and Context (PR #97637)

[Arthur Eubanks via llvm-commits]

================
@@ -171,9 +208,194 @@ void Function::dump() const {
 }
 #endif // NDEBUG
 
+BasicBlock::iterator::pointer
+BasicBlock::iterator::getI(llvm::BasicBlock::iterator It) const {
+  return cast_or_null<Instruction>(Ctx->getValue(&*It));
+}
+
+Value *Context::registerValue(std::unique_ptr<Value> &&VPtr) {
+  assert(VPtr->getSubclassID() != Value::ClassID::User &&
+         "Can't register a user!");
+  Value *V = VPtr.get();
+  llvm::Value *Key = V->Val;
+  LLVMValueToValueMap[Key] = std::move(VPtr);
+  return V;
+}
+
+Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, int Depth,
+                                         llvm::User *U) {
+  assert(Depth < 1048576 && "Infinite recursion?");
+  auto Pair = LLVMValueToValueMap.insert({LLVMV, nullptr});
+  auto It = Pair.first;
+  if (!Pair.second)
+    return It->second.get();
+
+  if (auto *C = dyn_cast<llvm::Constant>(LLVMV)) {
+    // Globals may be self-referencing, like @bar = global [1 x ptr] [ptr @bar].
+    // Avoid infinite loops by early returning once we detect a loop.
+    if (isa<GlobalValue>(C)) {
+      if (Depth == 0)
+        VisitedConstants.clear();
+      if (!VisitedConstants.insert(C).second)
+        return nullptr; //  recursion loop!
+    }
+    for (llvm::Value *COp : C->operands())
+      getOrCreateValueInternal(COp, Depth + 1, C);
+    It->second = std::make_unique<Constant>(C, *this);
+    return It->second.get();
+  }
+  if (auto *Arg = dyn_cast<llvm::Argument>(LLVMV)) {
+    It->second = std::make_unique<Argument>(Arg, *this);
+    return It->second.get();
+  }
+  if (auto *BB = dyn_cast<llvm::BasicBlock>(LLVMV)) {
+    assert(isa<BlockAddress>(U) &&
+           "This won't create a SBBB, don't call this function directly!");
+    if (auto *SBBB = getValue(BB))
+      return SBBB;
+    return nullptr;
+  }
+  assert(isa<llvm::Instruction>(LLVMV) && "Expected Instruction");
+  It->second =
+      std::make_unique<OpaqueInst>(cast<llvm::Instruction>(LLVMV), *this);
+  return It->second.get();
+}
+
+BasicBlock *Context::createBasicBlock(llvm::BasicBlock *LLVMBB) {
+  assert(getValue(LLVMBB) == nullptr && "Already exists!");
+  auto NewBBPtr = std::make_unique<BasicBlock>(LLVMBB, *this);
+  auto *BB = cast<BasicBlock>(registerValue(std::move(NewBBPtr)));
+  // Create SandboxIR for BB's body.
+  BB->buildBasicBlockFromLLVMIR(LLVMBB);
+  return BB;
+}
+
 Value *Context::getValue(llvm::Value *V) const {
   auto It = LLVMValueToValueMap.find(V);
   if (It != LLVMValueToValueMap.end())
     return It->second.get();
   return nullptr;
 }
+
+Function *Context::createFunction(llvm::Function *F, bool CreateBBs) {
+  assert(getValue(F) == nullptr && "Already exists!");
+  auto NewFPtr = std::make_unique<Function>(F, *this);
+  // Create arguments.
+  for (auto &Arg : F->args())
+    getOrCreateGeneric<Argument, llvm::Argument>(&Arg);
+  // Create BBs.
+  if (CreateBBs) {
+    for (auto &BB : *F)
+      createBasicBlock(&BB);
+  }
+  auto *SBF = cast<Function>(registerValue(std::move(NewFPtr)));
+  return SBF;
+}
+
+Function *BasicBlock::getParent() const {
+  auto *BB = cast<llvm::BasicBlock>(Val);
+  auto *F = BB->getParent();
+  if (F == nullptr)
+    // Detached
+    return nullptr;
+  return cast_or_null<Function>(Ctx.getValue(F));
+}
+
+void BasicBlock::buildBasicBlockFromLLVMIR(llvm::BasicBlock *LLVMBB) {
+  for (llvm::Instruction &IRef : reverse(*LLVMBB)) {
+    llvm::Instruction *I = &IRef;
+    Ctx.getOrCreateValue(I);
+    for (auto [OpIdx, Op] : enumerate(I->operands())) {
+      // Skip instruction's label operands
+      if (isa<llvm::BasicBlock>(Op))
+        continue;
+      // Skip metadata
+      if (isa<llvm::MetadataAsValue>(Op))
+        continue;
+      // Skip asm
+      if (isa<llvm::InlineAsm>(Op))
+        continue;
+      Ctx.getOrCreateValue(Op);
+    }
+  }
+#if !defined NDEBUG && defined SBVEC_EXPENSIVE_CHECKS
----------------
aeubanks wrote:

nit: LLVM style is typically with parens, like `#if defined(FOO)`

https://github.com/llvm/llvm-project/pull/97637