[SandboxIR] Implement Module (#109716)

This patch implements sandboxir::Module.
It provides access to globals.

llvm/include/llvm/SandboxIR/Module.h

@@ -0,0 +1,92 @@
+//===- Module.h -------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SANDBOXIR_MODULE_H
+#define LLVM_SANDBOXIR_MODULE_H
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/Module.h"
+#include <string>
+
+namespace llvm {
+
+class DataLayout;
+
+namespace sandboxir {
+
+class Context;
+class Function;
+class GlobalVariable;
+class Type;
+class Constant;
+class GlobalAlias;
+class GlobalIFunc;
+
+/// In SandboxIR the Module is mainly used to access the list of global objects.
+class Module {
+  llvm::Module &LLVMM;
+  Context &Ctx;
+
+  Module(llvm::Module &LLVMM, Context &Ctx) : LLVMM(LLVMM), Ctx(Ctx) {}
+  friend class Context; // For constructor.
+
+public:
+  Context &getContext() const { return Ctx; }
+
+  Function *getFunction(StringRef Name) const;
+
+  const DataLayout &getDataLayout() const { return LLVMM.getDataLayout(); }
+
+  const std::string &getSourceFileName() const {
+    return LLVMM.getSourceFileName();
+  }
+
+  /// Look up the specified global variable in the module symbol table. If it
+  /// does not exist, return null. If AllowInternal is set to true, this
+  /// function will return types that have InternalLinkage. By default, these
+  /// types are not returned.
+  GlobalVariable *getGlobalVariable(StringRef Name, bool AllowInternal) const;
+  GlobalVariable *getGlobalVariable(StringRef Name) const {
+    return getGlobalVariable(Name, /*AllowInternal=*/false);
+  }
+  /// Return the global variable in the module with the specified name, of
+  /// arbitrary type. This method returns null if a global with the specified
+  /// name is not found.
+  GlobalVariable *getNamedGlobal(StringRef Name) const {
+    return getGlobalVariable(Name, true);
+  }
+
+  // TODO: missing getOrInsertGlobal().
+
+  /// Return the global alias in the module with the specified name, of
+  /// arbitrary type. This method returns null if a global with the specified
+  /// name is not found.
+  GlobalAlias *getNamedAlias(StringRef Name) const;
+
+  /// Return the global ifunc in the module with the specified name, of
+  /// arbitrary type. This method returns null if a global with the specified
+  /// name is not found.
+  GlobalIFunc *getNamedIFunc(StringRef Name) const;
+
+  // TODO: Missing removeGlobalVariable() eraseGlobalVariable(),
+  // insertGlobalVariable()
+
+  // TODO: Missing global_begin(), global_end(), globals().
+
+  // TODO: Missing many other functions.
+
+#ifndef NDEBUG
+  void dumpOS(raw_ostream &OS) const;
+  LLVM_DUMP_METHOD void dump() const;
+#endif // NDEBUG
+};
+
+} // namespace sandboxir
+} // namespace llvm
+
+#endif // LLVM_SANDBOXIR_MODULE_H

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -109,6 +109,7 @@
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
+#include "llvm/SandboxIR/Module.h"
 #include "llvm/SandboxIR/Tracker.h"
 #include "llvm/SandboxIR/Type.h"
 #include "llvm/SandboxIR/Use.h"
@@ -138,6 +139,7 @@ class ConstantPtrAuth;
 class ConstantExpr;
 class Context;
 class Function;
+class Module;
 class Instruction;
 class VAArgInst;
 class FreezeInst;
@@ -347,7 +349,7 @@ protected:
   friend class ConstantPtrAuth;       // For `Val`.
   friend class ConstantExpr;          // For `Val`.
   friend class Utils;                 // For `Val`.
-
+  friend class Module;                // For `Val`.
   // Region needs to manipulate metadata in the underlying LLVM Value, we don't
   // expose metadata in sandboxir.
   friend class Region;
@@ -1322,7 +1324,10 @@ public:
   GlobalWithNodeAPI(Value::ClassID ID, LLVMParentT *C, Context &Ctx)
       : ParentT(ID, C, Ctx) {}
 
-  // TODO: Missing getParent(). Should be added once Module is available.
+  Module *getParent() const {
+    llvm::Module *LLVMM = cast<LLVMGlobalT>(this->Val)->getParent();
+    return this->Ctx.getModule(LLVMM);
+  }
 
   using iterator = mapped_iterator<
       decltype(static_cast<LLVMGlobalT *>(nullptr)->getIterator()), LLVMGVToGV>;
@@ -4556,6 +4561,9 @@ protected:
   DenseMap<llvm::Value *, std::unique_ptr<sandboxir::Value>>
       LLVMValueToValueMap;
 
+  /// Maps an LLVM Module to the corresponding sandboxir::Module.
+  DenseMap<llvm::Module *, std::unique_ptr<Module>> LLVMModuleToModuleMap;
+
   /// Type has a protected destructor to prohibit the user from managing the
   /// lifetime of the Type objects. Context is friend of Type, and this custom
   /// deleter can destroy Type.
@@ -4699,6 +4707,10 @@ public:
     return getValue(const_cast<llvm::Value *>(V));
   }
 
+  Module *getModule(llvm::Module *LLVMM) const;
+
+  Module *getOrCreateModule(llvm::Module *LLVMM);
+
   Type *getType(llvm::Type *LLVMTy) {
     if (LLVMTy == nullptr)
       return nullptr;
@@ -4712,8 +4724,13 @@ public:
   /// Create a sandboxir::Function for an existing LLVM IR \p F, including all
   /// blocks and instructions.
   /// This is the main API function for creating Sandbox IR.
+  /// Note: this will not fully populate its parent module. The only globals
+  /// that will be available are those used within the function.
   Function *createFunction(llvm::Function *F);
 
+  /// Create a sandboxir::Module corresponding to \p LLVMM.
+  Module *createModule(llvm::Module *LLVMM);
+
   /// \Returns the number of values registered with Context.
   size_t getNumValues() const { return LLVMValueToValueMap.size(); }
 };
@@ -4739,6 +4756,10 @@ public:
     return From->getSubclassID() == ClassID::Function;
   }
 
+  Module *getParent() {
+    return Ctx.getModule(cast<llvm::Function>(Val)->getParent());
+  }
+
   Argument *getArg(unsigned Idx) const {
     llvm::Argument *Arg = cast<llvm::Function>(Val)->getArg(Idx);
     return cast<Argument>(Ctx.getValue(Arg));

llvm/include/llvm/SandboxIR/Tracker.h

@@ -64,6 +64,9 @@ class SwitchInst;
 class ConstantInt;
 class ShuffleVectorInst;
 class CmpInst;
+class Module;
+class GlobalVariable;
+
 /// The base class for IR Change classes.
 class IRChangeBase {
 protected:

llvm/include/llvm/SandboxIR/Type.h

@@ -31,6 +31,7 @@ class IntegerType;
 class FunctionType;
 class ArrayType;
 class StructType;
+class Module;
 #define DEF_INSTR(ID, OPCODE, CLASS) class CLASS;
 #define DEF_CONST(ID, CLASS) class CLASS;
 #include "llvm/SandboxIR/SandboxIRValues.def"
@@ -57,6 +58,8 @@ protected:
   friend class CmpInst;            // For LLVMTy. TODO: Cleanup after
                                    // sandboxir::VectorType is more complete.
   friend class Utils;              // for LLVMTy
+  friend class TargetExtType;      // For LLVMTy.
+  friend class Module;             // For LLVMTy.
 
   // Friend all instruction classes because `create()` functions use LLVMTy.
 #define DEF_INSTR(ID, OPCODE, CLASS) friend class CLASS;

llvm/lib/SandboxIR/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_llvm_component_library(LLVMSandboxIR
+  Module.cpp
   Pass.cpp
   PassManager.cpp
   SandboxIR.cpp

llvm/lib/SandboxIR/Module.cpp

@@ -0,0 +1,40 @@
+//===- Module.cpp ---------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/SandboxIR/Module.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+
+using namespace llvm::sandboxir;
+
+Function *Module::getFunction(StringRef Name) const {
+  llvm::Function *LLVMF = LLVMM.getFunction(Name);
+  return cast_or_null<Function>(Ctx.getValue(LLVMF));
+}
+
+GlobalVariable *Module::getGlobalVariable(StringRef Name,
+                                          bool AllowInternal) const {
+  return cast_or_null<GlobalVariable>(
+      Ctx.getValue(LLVMM.getGlobalVariable(Name, AllowInternal)));
+}
+
+GlobalAlias *Module::getNamedAlias(StringRef Name) const {
+  return cast_or_null<GlobalAlias>(Ctx.getValue(LLVMM.getNamedAlias(Name)));
+}
+
+GlobalIFunc *Module::getNamedIFunc(StringRef Name) const {
+  return cast_or_null<GlobalIFunc>(Ctx.getValue(LLVMM.getNamedIFunc(Name)));
+}
+
+#ifndef NDEBUG
+void Module::dumpOS(raw_ostream &OS) const { OS << LLVMM; }
+
+void Module::dump() const {
+  dumpOS(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -3404,8 +3404,29 @@ Value *Context::getValue(llvm::Value *V) const {
   return nullptr;
 }
 
+Module *Context::getModule(llvm::Module *LLVMM) const {
+  auto It = LLVMModuleToModuleMap.find(LLVMM);
+  if (It != LLVMModuleToModuleMap.end())
+    return It->second.get();
+  return nullptr;
+}
+
+Module *Context::getOrCreateModule(llvm::Module *LLVMM) {
+  auto Pair = LLVMModuleToModuleMap.insert({LLVMM, nullptr});
+  auto It = Pair.first;
+  if (!Pair.second)
+    return It->second.get();
+  It->second = std::unique_ptr<Module>(new Module(*LLVMM, *this));
+  return It->second.get();
+}
+
 Function *Context::createFunction(llvm::Function *F) {
   assert(getValue(F) == nullptr && "Already exists!");
+  // Create the module if needed before we create the new sandboxir::Function.
+  // Note: this won't fully populate the module. The only globals that will be
+  // available will be the ones being used within the function.
+  getOrCreateModule(F->getParent());
+
   auto NewFPtr = std::unique_ptr<Function>(new Function(F, *this));
   auto *SBF = cast<Function>(registerValue(std::move(NewFPtr)));
   // Create arguments.
@@ -3417,6 +3438,24 @@ Function *Context::createFunction(llvm::Function *F) {
   return SBF;
 }
 
+Module *Context::createModule(llvm::Module *LLVMM) {
+  auto *M = getOrCreateModule(LLVMM);
+  // Create the functions.
+  for (auto &LLVMF : *LLVMM)
+    createFunction(&LLVMF);
+  // Create globals.
+  for (auto &Global : LLVMM->globals())
+    getOrCreateValue(&Global);
+  // Create aliases.
+  for (auto &Alias : LLVMM->aliases())
+    getOrCreateValue(&Alias);
+  // Create ifuncs.
+  for (auto &IFunc : LLVMM->ifuncs())
+    getOrCreateValue(&IFunc);
+
+  return M;
+}
+
 Function *BasicBlock::getParent() const {
   auto *BB = cast<llvm::BasicBlock>(Val);
   auto *F = BB->getParent();

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -1685,6 +1685,64 @@ bb1:
 #endif // NDEBUG
 }
 
+TEST_F(SandboxIRTest, Module) {
+  parseIR(C, R"IR(
+@glob0 = global i32 42
+@glob1 = global i32 43
+@internal0 = internal global i32 42
+@const0 = constant i32 42
+@alias0 = dso_local alias void(), ptr @foo
+@ifunc = ifunc void(), ptr @foo
+define void @foo() {
+  ret void
+}
+define void @bar() {
+  ret void
+}
+)IR");
+  llvm::Module *LLVMM = &*M;
+  llvm::Function *LLVMFFoo = &*M->getFunction("foo");
+  llvm::Function *LLVMFBar = &*M->getFunction("bar");
+
+  sandboxir::Context Ctx(C);
+  auto *M = Ctx.createModule(LLVMM);
+  // Check getContext().
+  EXPECT_EQ(&M->getContext(), &Ctx);
+  // Check getFunction().
+  auto *FFoo = M->getFunction("foo");
+  auto *FBar = M->getFunction("bar");
+  EXPECT_EQ(FFoo, Ctx.getValue(LLVMFFoo));
+  EXPECT_EQ(FBar, Ctx.getValue(LLVMFBar));
+  // Check getDataLayout().
+  EXPECT_EQ(&M->getDataLayout(), &LLVMM->getDataLayout());
+  // Check getSourceFileName().
+  EXPECT_EQ(M->getSourceFileName(), LLVMM->getSourceFileName());
+  // Check getGlobalVariable().
+  for (const char *Name : {"global0", "global1", "internal0"})
+    EXPECT_EQ(M->getGlobalVariable(Name),
+              Ctx.getValue(LLVMM->getGlobalVariable(Name)));
+  // Check getGlobalVariable(AllowInternal).
+  {
+    auto *Internal0 = M->getGlobalVariable("internal0", /*AllowInternal=*/true);
+    EXPECT_TRUE(Internal0 != nullptr);
+    EXPECT_EQ(Internal0, Ctx.getValue(LLVMM->getNamedGlobal("internal0")));
+  }
+  // Check getNamedGlobal().
+  {
+    auto *Internal = M->getNamedGlobal("internal0");
+    EXPECT_TRUE(Internal != nullptr);
+    EXPECT_EQ(Internal, Ctx.getValue(LLVMM->getNamedGlobal("internal0")));
+  }
+  // Check getNamedAlias().
+  auto *Alias0 = M->getNamedAlias("alias0");
+  EXPECT_EQ(Alias0, Ctx.getValue(LLVMM->getNamedAlias("alias0")));
+  EXPECT_EQ(M->getNamedAlias("aliasFOO"), nullptr);
+  // Check getNamedIFunc().
+  auto *IFunc0 = M->getNamedIFunc("ifunc0");
+  EXPECT_EQ(IFunc0, Ctx.getValue(LLVMM->getNamedAlias("ifunc0")));
+  EXPECT_EQ(M->getNamedIFunc("ifuncFOO"), nullptr);
+}
+
 TEST_F(SandboxIRTest, BasicBlock) {
   parseIR(C, R"IR(
 define void @foo(i32 %v1) {