[SandboxIR] Implement PoisonValue (#107324)

This patch implements sandboxir::PoisonValue mirroring
llvm::PoisonValue.

Author: vporpo

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -122,6 +122,7 @@ class ConstantInt;
 class ConstantFP;
 class ConstantAggregateZero;
 class ConstantPointerNull;
+class PoisonValue;
 class Context;
 class Function;
 class Instruction;
@@ -320,6 +321,7 @@ class Value {
   friend class ConstantStruct;        // For `Val`.
   friend class ConstantAggregateZero; // For `Val`.
   friend class ConstantPointerNull;   // For `Val`.
+  friend class PoisonValue;           // For `Val`.
 
   /// All values point to the context.
   Context &Ctx;
@@ -1018,6 +1020,49 @@ class ConstantPointerNull final : public Constant {
 #endif
 };
 
+// TODO: Inherit from UndefValue.
+class PoisonValue final : public Constant {
+  PoisonValue(llvm::PoisonValue *C, Context &Ctx)
+      : Constant(ClassID::PoisonValue, C, Ctx) {}
+  friend class Context; // For constructor.
+
+public:
+  /// Static factory methods - Return an 'poison' object of the specified type.
+  static PoisonValue *get(Type *T);
+
+  /// If this poison has array or vector type, return a poison with the right
+  /// element type.
+  PoisonValue *getSequentialElement() const;
+
+  /// If this poison has struct type, return a poison with the right element
+  /// type for the specified element.
+  PoisonValue *getStructElement(unsigned Elt) const;
+
+  /// Return an poison of the right value for the specified GEP index if we can,
+  /// otherwise return null (e.g. if C is a ConstantExpr).
+  PoisonValue *getElementValue(Constant *C) const;
+
+  /// Return an poison of the right value for the specified GEP index.
+  PoisonValue *getElementValue(unsigned Idx) const;
+
+  /// For isa/dyn_cast.
+  static bool classof(const sandboxir::Value *From) {
+    return From->getSubclassID() == ClassID::PoisonValue;
+  }
+  unsigned getUseOperandNo(const Use &Use) const final {
+    llvm_unreachable("PoisonValue has no operands!");
+  }
+#ifndef NDEBUG
+  void verify() const override {
+    assert(isa<llvm::PoisonValue>(Val) && "Expected a PoisonValue!");
+  }
+  void dumpOS(raw_ostream &OS) const override {
+    dumpCommonPrefix(OS);
+    dumpCommonSuffix(OS);
+  }
+#endif
+};
+
 /// Iterator for `Instruction`s in a `BasicBlock.
 /// \Returns an sandboxir::Instruction & when derereferenced.
 class BBIterator {

llvm/include/llvm/SandboxIR/SandboxIRValues.def

@@ -32,6 +32,7 @@ DEF_CONST(ConstantStruct, ConstantStruct)
 DEF_CONST(ConstantVector, ConstantVector)
 DEF_CONST(ConstantAggregateZero, ConstantAggregateZero)
 DEF_CONST(ConstantPointerNull, ConstantPointerNull)
+DEF_CONST(PoisonValue, PoisonValue)
 
 #ifndef DEF_INSTR
 #define DEF_INSTR(ID, OPCODE, CLASS)

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -2437,6 +2437,32 @@ PointerType *ConstantPointerNull::getType() const {
       Ctx.getType(cast<llvm::ConstantPointerNull>(Val)->getType()));
 }
 
+PoisonValue *PoisonValue::get(Type *T) {
+  auto *LLVMC = llvm::PoisonValue::get(T->LLVMTy);
+  return cast<PoisonValue>(T->getContext().getOrCreateConstant(LLVMC));
+}
+
+PoisonValue *PoisonValue::getSequentialElement() const {
+  return cast<PoisonValue>(Ctx.getOrCreateConstant(
+      cast<llvm::PoisonValue>(Val)->getSequentialElement()));
+}
+
+PoisonValue *PoisonValue::getStructElement(unsigned Elt) const {
+  return cast<PoisonValue>(Ctx.getOrCreateConstant(
+      cast<llvm::PoisonValue>(Val)->getStructElement(Elt)));
+}
+
+PoisonValue *PoisonValue::getElementValue(Constant *C) const {
+  return cast<PoisonValue>(
+      Ctx.getOrCreateConstant(cast<llvm::PoisonValue>(Val)->getElementValue(
+          cast<llvm::Constant>(C->Val))));
+}
+
+PoisonValue *PoisonValue::getElementValue(unsigned Idx) const {
+  return cast<PoisonValue>(Ctx.getOrCreateConstant(
+      cast<llvm::PoisonValue>(Val)->getElementValue(Idx)));
+}
+
 FunctionType *Function::getFunctionType() const {
   return cast<FunctionType>(
       Ctx.getType(cast<llvm::Function>(Val)->getFunctionType()));
@@ -2550,6 +2576,10 @@ Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, llvm::User *U) {
       It->second = std::unique_ptr<ConstantPointerNull>(
           new ConstantPointerNull(cast<llvm::ConstantPointerNull>(C), *this));
       return It->second.get();
+    case llvm::Value::PoisonValueVal:
+      It->second = std::unique_ptr<PoisonValue>(
+          new PoisonValue(cast<llvm::PoisonValue>(C), *this));
+      return It->second.get();
     case llvm::Value::ConstantArrayVal:
       It->second = std::unique_ptr<ConstantArray>(
           new ConstantArray(cast<llvm::ConstantArray>(C), *this));

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -616,6 +616,60 @@ define ptr @foo() {
   EXPECT_EQ(NewCPNull2->getType(), sandboxir::PointerType::get(Ctx, 1u));
 }
 
+TEST_F(SandboxIRTest, PoisonValue) {
+  parseIR(C, R"IR(
+define void @foo() {
+  %i0 = add i32 poison, poison
+  %i1 = add <2 x i32> poison, poison
+  %i2 = extractvalue {i32, i8} poison, 0
+  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 *I0 = &*It++;
+  auto *I1 = &*It++;
+  auto *I2 = &*It++;
+  auto *Int32Ty = sandboxir::Type::getInt32Ty(Ctx);
+  auto *Int8Ty = sandboxir::Type::getInt8Ty(Ctx);
+  auto *Zero32 = sandboxir::ConstantInt::get(Int32Ty, 0u);
+  auto *One32 = sandboxir::ConstantInt::get(Int32Ty, 1u);
+
+  // Check classof() and creation.
+  auto *Poison = cast<sandboxir::PoisonValue>(I0->getOperand(0));
+  EXPECT_EQ(Poison->getType(), Int32Ty);
+  // Check get().
+  auto *NewPoison = sandboxir::PoisonValue::get(Int32Ty);
+  EXPECT_EQ(NewPoison, Poison);
+  auto *NewPoison2 =
+      sandboxir::PoisonValue::get(sandboxir::PointerType::get(Ctx, 0u));
+  EXPECT_NE(NewPoison2, Poison);
+  // Check getSequentialElement().
+  auto *PoisonVector = cast<sandboxir::PoisonValue>(I1->getOperand(0));
+  auto *SeqElm = PoisonVector->getSequentialElement();
+  EXPECT_EQ(SeqElm->getType(), Int32Ty);
+  // Check getStructElement().
+  auto *PoisonStruct = cast<sandboxir::PoisonValue>(I2->getOperand(0));
+  auto *StrElm0 = PoisonStruct->getStructElement(0);
+  auto *StrElm1 = PoisonStruct->getStructElement(1);
+  EXPECT_EQ(StrElm0->getType(), Int32Ty);
+  EXPECT_EQ(StrElm1->getType(), Int8Ty);
+  // Check getElementValue(Constant)
+  EXPECT_EQ(PoisonStruct->getElementValue(Zero32),
+            sandboxir::PoisonValue::get(Int32Ty));
+  EXPECT_EQ(PoisonStruct->getElementValue(One32),
+            sandboxir::PoisonValue::get(Int8Ty));
+  // Check getElementValue(unsigned)
+  EXPECT_EQ(PoisonStruct->getElementValue(0u),
+            sandboxir::PoisonValue::get(Int32Ty));
+  EXPECT_EQ(PoisonStruct->getElementValue(1u),
+            sandboxir::PoisonValue::get(Int8Ty));
+}
+
 TEST_F(SandboxIRTest, Use) {
   parseIR(C, R"IR(
 define i32 @foo(i32 %v0, i32 %v1) {