[llvm] [llubi] Add support for constant expressions (PR #203746)

[Nikita Popov via llvm-commits]

================
@@ -142,9 +141,182 @@ std::optional<AnyValue> Context::getConstantValueImpl(Constant *C) {
   if (auto *F = dyn_cast<Function>(C))
     return FuncAddrMap.at(F);
 
+  if (auto *CE = dyn_cast<ConstantExpr>(C))
+    return evaluateConstantExpression(CE);
+
   return std::nullopt;
 }
 
+std::optional<AnyValue> Context::evaluateConstantExpression(ConstantExpr *CE) {
+  unsigned Opc = CE->getOpcode();
+  switch (Opc) {
+  case Instruction::Trunc: {
+    const AnyValue *Src = getConstantValue(CE->getOperand(0));
+    if (!Src)
+      return std::nullopt;
+    if (Src->isPoison())
+      return AnyValue::poison();
+    unsigned BitWidth = CE->getType()->getScalarSizeInBits();
+    if (Src->isInteger())
+      return AnyValue(Src->asInteger().trunc(BitWidth));
+    std::vector<AnyValue> Vec = Src->asAggregate();
+    for (auto &V : Vec) {
+      if (V.isInteger())
+        V = V.asInteger().trunc(BitWidth);
+    }
+    return AnyValue(std::move(Vec));
+  }
+  case Instruction::BitCast: {
+    Constant *SrcOp = CE->getOperand(0);
+    const AnyValue *Src = getConstantValue(SrcOp);
+    if (!Src)
+      return std::nullopt;
+    SmallVector<Byte> Bytes;
+    Bytes.resize(getEffectiveTypeStoreSize(CE->getType()), Byte::concrete(0));
+    toBytes(*Src, SrcOp->getType(), Bytes);
+    return fromBytes(Bytes, CE->getType());
+  }
+  case Instruction::InsertElement: {
+    const AnyValue *Src = getConstantValue(CE->getOperand(0));
+    if (!Src)
+      return std::nullopt;
+    const AnyValue *Val = getConstantValue(CE->getOperand(1));
+    if (!Val)
+      return std::nullopt;
+    const AnyValue *Idx = getConstantValue(CE->getOperand(2));
+    if (!Idx)
+      return std::nullopt;
+    auto &SrcVec = Src->asAggregate();
+    if (Idx->isPoison() || Idx->asInteger().uge(SrcVec.size()))
+      return AnyValue::getPoisonValue(*this, CE->getType());
+    std::vector<AnyValue> ResVec = SrcVec;
+    ResVec[Idx->asInteger().getZExtValue()] = *Val;
+    return AnyValue(std::move(ResVec));
+  }
+  case Instruction::ExtractElement: {
+    const AnyValue *Src = getConstantValue(CE->getOperand(0));
+    if (!Src)
+      return std::nullopt;
+    const AnyValue *Idx = getConstantValue(CE->getOperand(1));
+    if (!Idx)
+      return std::nullopt;
+    auto &SrcVec = Src->asAggregate();
+    if (Idx->isPoison() || Idx->asInteger().uge(SrcVec.size()))
+      return AnyValue::getPoisonValue(*this, CE->getType());
+    return SrcVec[Idx->asInteger().getZExtValue()];
+  }
+  case Instruction::ShuffleVector: {
+    const AnyValue *LHS = getConstantValue(CE->getOperand(0));
+    if (!LHS)
+      return std::nullopt;
+    const AnyValue *RHS = getConstantValue(CE->getOperand(1));
+    if (!RHS)
+      return std::nullopt;
+    auto &LHSVec = LHS->asAggregate();
+    auto &RHSVec = RHS->asAggregate();
+    uint32_t Size = cast<VectorType>(CE->getOperand(0)->getType())
+                        ->getElementCount()
+                        .getKnownMinValue();
+    std::vector<AnyValue> Res;
+    uint32_t DstLen =
+        getEVL(cast<VectorType>(CE->getType())->getElementCount());
+    Res.reserve(DstLen);
+    uint32_t Stride = CE->getShuffleMask().size();
+    // For scalable vectors, we need to repeat the shuffle mask until we fill
+    // the destination vector.
+    for (uint32_t Off = 0; Off != DstLen; Off += Stride) {
+      for (int Idx : CE->getShuffleMask()) {
+        if (Idx == PoisonMaskElem)
+          Res.push_back(AnyValue::poison());
+        else if (Idx < static_cast<int>(Size))
+          Res.push_back(LHSVec[Idx]);
+        else
+          Res.push_back(RHSVec[Idx - Size]);
+      }
+    }
+    return AnyValue(std::move(Res));
+  }
+  case Instruction::GetElementPtr: {
+    // Temporary variable for reference to poison values when the subexpression
+    // cannot be evaluated. As the reference will be consumed immediately, we
+    // don't need to store them into a list.
+    AnyValue PoisonValue;
+    AnyValue Res =
+        computeGEP(*cast<GEPOperator>(CE), [&](Value *V) -> const AnyValue & {
+          const AnyValue *Val = getConstantValue(cast<Constant>(V));
+          if (Val)
+            return *Val;
+          PoisonValue = AnyValue::getPoisonValue(*this, V->getType());
+          return PoisonValue;
+        });
+    if (!PoisonValue.isNone())
+      return std::nullopt;
+    return std::move(Res);
+  }
+  case Instruction::PtrToAddr: {
+    const AnyValue *Src = getConstantValue(CE->getOperand(0));
+    if (!Src)
+      return std::nullopt;
+    if (Src->isPoison())
+      return AnyValue::poison();
+    unsigned BitWidth = CE->getType()->getScalarSizeInBits();
+    if (Src->isPointer())
+      return Src->asPointer().address().trunc(BitWidth);
+    std::vector<AnyValue> Vec = Src->asAggregate();
+    for (auto &V : Vec) {
+      if (V.isPointer())
+        V = V.asPointer().address().trunc(BitWidth);
+    }
+    return AnyValue(std::move(Vec));
+  }
+  default:
+    assert(Instruction::isBinaryOp(Opc) && "Must be binary operator?");
----------------
nikic wrote:

We should explicitly list the supported binary opcodes here and return nullptr for other cases (so that ptrtoint expressions don't assert).

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