[clang] Save ShuffleVectorExpr args as ConstantExpr (#139709)

The passed indices have to be constant integers anyway, which we verify
before creating the ShuffleVectorExpr. Use the value we create there and
save the indices using a ConstantExpr instead. This way, we don't have
to evaluate the args every time we call getShuffleMaskIdx().

Author: tbaederr

clang/include/clang/AST/Expr.h

@@ -4566,9 +4566,11 @@ class ShuffleVectorExpr : public Expr {
 
   void setExprs(const ASTContext &C, ArrayRef<Expr *> Exprs);
 
-  llvm::APSInt getShuffleMaskIdx(const ASTContext &Ctx, unsigned N) const {
+  llvm::APSInt getShuffleMaskIdx(unsigned N) const {
     assert((N < NumExprs - 2) && "Shuffle idx out of range!");
-    return getExpr(N+2)->EvaluateKnownConstInt(Ctx);
+    assert(isa<ConstantExpr>(getExpr(N + 2)) &&
+           "Index expression must be a ConstantExpr");
+    return cast<ConstantExpr>(getExpr(N + 2))->getAPValueResult().getInt();
   }
 
   // Iterators

clang/lib/AST/ByteCode/Compiler.cpp

@@ -3883,7 +3883,7 @@ bool Compiler<Emitter>::VisitShuffleVectorExpr(const ShuffleVectorExpr *E) {
       return false;
   }
   for (unsigned I = 0; I != NumOutputElems; ++I) {
-    APSInt ShuffleIndex = E->getShuffleMaskIdx(Ctx.getASTContext(), I);
+    APSInt ShuffleIndex = E->getShuffleMaskIdx(I);
     assert(ShuffleIndex >= -1);
     if (ShuffleIndex == -1)
       return this->emitInvalidShuffleVectorIndex(I, E);

clang/lib/AST/ExprConstant.cpp

@@ -11571,7 +11571,7 @@ static bool handleVectorShuffle(EvalInfo &Info, const ShuffleVectorExpr *E,
   unsigned const TotalElementsInInputVector1 = VecVal1.getVectorLength();
   unsigned const TotalElementsInInputVector2 = VecVal2.getVectorLength();
 
-  APSInt IndexVal = E->getShuffleMaskIdx(Info.Ctx, EltNum);
+  APSInt IndexVal = E->getShuffleMaskIdx(EltNum);
   int64_t index = IndexVal.getExtValue();
   // The spec says that -1 should be treated as undef for optimizations,
   // but in constexpr we'd have to produce an APValue::Indeterminate,

clang/lib/CodeGen/CGExprScalar.cpp

@@ -1906,7 +1906,7 @@ Value *ScalarExprEmitter::VisitShuffleVectorExpr(ShuffleVectorExpr *E) {
 
   SmallVector<int, 32> Indices;
   for (unsigned i = 2; i < E->getNumSubExprs(); ++i) {
-    llvm::APSInt Idx = E->getShuffleMaskIdx(CGF.getContext(), i-2);
+    llvm::APSInt Idx = E->getShuffleMaskIdx(i - 2);
     // Check for -1 and output it as undef in the IR.
     if (Idx.isSigned() && Idx.isAllOnes())
       Indices.push_back(-1);

clang/lib/Sema/SemaChecking.cpp

@@ -5342,29 +5342,29 @@ ExprResult Sema::BuiltinShuffleVector(CallExpr *TheCall) {
   }
 
   for (unsigned i = 2; i < TheCall->getNumArgs(); i++) {
-    if (TheCall->getArg(i)->isTypeDependent() ||
-        TheCall->getArg(i)->isValueDependent())
+    Expr *Arg = TheCall->getArg(i);
+    if (Arg->isTypeDependent() || Arg->isValueDependent())
       continue;
 
     std::optional<llvm::APSInt> Result;
-    if (!(Result = TheCall->getArg(i)->getIntegerConstantExpr(Context)))
+    if (!(Result = Arg->getIntegerConstantExpr(Context)))
       return ExprError(Diag(TheCall->getBeginLoc(),
                             diag::err_shufflevector_nonconstant_argument)
-                       << TheCall->getArg(i)->getSourceRange());
+                       << Arg->getSourceRange());
 
     // Allow -1 which will be translated to undef in the IR.
     if (Result->isSigned() && Result->isAllOnes())
-      continue;
-
-    if (Result->getActiveBits() > 64 ||
-        Result->getZExtValue() >= numElements * 2)
+      ;
+    else if (Result->getActiveBits() > 64 ||
+             Result->getZExtValue() >= numElements * 2)
       return ExprError(Diag(TheCall->getBeginLoc(),
                             diag::err_shufflevector_argument_too_large)
-                       << TheCall->getArg(i)->getSourceRange());
-  }
+                       << Arg->getSourceRange());
 
-  SmallVector<Expr*, 32> exprs;
+    TheCall->setArg(i, ConstantExpr::Create(Context, Arg, APValue(*Result)));
+  }
 
+  SmallVector<Expr *> exprs;
   for (unsigned i = 0, e = TheCall->getNumArgs(); i != e; i++) {
     exprs.push_back(TheCall->getArg(i));
     TheCall->setArg(i, nullptr);

clang/unittests/AST/ASTImporterTest.cpp

@@ -308,8 +308,10 @@ TEST_P(ImportExpr, ImportShuffleVectorExpr) {
   const auto Pattern = functionDecl(hasDescendant(shuffleVectorExpr(
       allOf(has(declRefExpr(to(parmVarDecl(hasName("a"))))),
             has(declRefExpr(to(parmVarDecl(hasName("b"))))),
-            has(integerLiteral(equals(0))), has(integerLiteral(equals(1))),
-            has(integerLiteral(equals(2))), has(integerLiteral(equals(3)))))));
+            has(constantExpr(has(integerLiteral(equals(0))))),
+            has(constantExpr(has(integerLiteral(equals(1))))),
+            has(constantExpr(has(integerLiteral(equals(2))))),
+            has(constantExpr(has(integerLiteral(equals(3)))))))));
   testImport(Code, Lang_C99, "", Lang_C99, Verifier, Pattern);
 }