[Clang] Implement constexpr evaluation for __builtin_infer_alloc_token()

[melver]

Implement the constexpr evaluation for __builtin_infer_alloc_token()
in Clang's constant expression evaluators (both in ExprConstant and the
new bytecode interpreter).

The constant evaluation is only supported for stateless (hash-based)
token modes. If a stateful mode like increment is used, the evaluation
fails, as the token value is not deterministic at compile time.

---

This change is part of the following series:

1. [AllocToken] Introduce llvm.alloc.token.id intrinsic #163632
2. [AllocToken] Refactor stateless token calculation into Support #163633
3. [AllocToken] Make token mode a pass parameter #163634
4. [Clang] Move AllocToken frontend options to LangOptions #163635
5. [Clang] Refactor allocation type inference logic #163636
6. [Clang][Sema] Add __builtin_infer_alloc_token() declaration and semantic checks #163638
7. [Clang] Implement constexpr evaluation for __builtin_infer_alloc_token() #163639
8. [Clang][CodeGen] Implement code generation for __builtin_infer_alloc_token() #156842

[llvmbot]

@llvm/pr-subscribers-clang

Author: Marco Elver (melver)

Changes

Implement the constexpr evaluation for __builtin_infer_alloc_token()
in Clang's constant expression evaluators (both in ExprConstant and the
new bytecode interpreter).

The constant evaluation is only supported for stateless (hash-based)
token modes. If a stateful mode like increment is used, the evaluation
fails, as the token value is not deterministic at compile time.

---

This change is part of the following series:

1. [AllocToken] Introduce llvm.alloc.token.id intrinsic #163632
2. [AllocToken] Refactor stateless token calculation into Support #163633
3. [AllocToken] Make token mode a pass parameter #163634
4. [Clang] Move AllocToken frontend options to LangOptions #163635
5. [Clang] Refactor allocation type inference logic #163636
6. [Clang][Sema] Add __builtin_infer_alloc_token() declaration and semantic checks #163638
7. [Clang] Implement constexpr evaluation for __builtin_infer_alloc_token() #163639
8. [Clang][CodeGen] Implement code generation for __builtin_infer_alloc_token() #156842

---

Full diff: https://github.com/llvm/llvm-project/pull/163639.diff

3 Files Affected:

• (modified) clang/lib/AST/ByteCode/InterpBuiltin.cpp (+43)
• (modified) clang/lib/AST/ExprConstant.cpp (+21)
• (modified) clang/test/SemaCXX/alloc-token.cpp (+45)

diff --git a/clang/lib/AST/ByteCode/InterpBuiltin.cpp b/clang/lib/AST/ByteCode/InterpBuiltin.cpp
index 5f0a77c125b85..7b3670ef46f0e 100644
--- a/clang/lib/AST/ByteCode/InterpBuiltin.cpp
+++ b/clang/lib/AST/ByteCode/InterpBuiltin.cpp
@@ -11,12 +11,14 @@
 #include "Interp.h"
 #include "InterpBuiltinBitCast.h"
 #include "PrimType.h"
+#include "clang/AST/InferAlloc.h"
 #include "clang/AST/OSLog.h"
 #include "clang/AST/RecordLayout.h"
 #include "clang/Basic/Builtins.h"
 #include "clang/Basic/TargetBuiltins.h"
 #include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/AllocToken.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/SipHash.h"
 
@@ -1306,6 +1308,44 @@ interp__builtin_ptrauth_string_discriminator(InterpState &S, CodePtr OpPC,
   return true;
 }
 
+static bool interp__builtin_infer_alloc_token(InterpState &S, CodePtr OpPC,
+                                              const InterpFrame *Frame,
+                                              const CallExpr *Call) {
+  const ASTContext &Ctx = S.getASTContext();
+  const uint64_t BitWidth = Ctx.getTypeSize(Ctx.getSizeType());
+  const auto Mode =
+      Ctx.getLangOpts().AllocTokenMode.value_or(llvm::DefaultAllocTokenMode);
+  const uint64_t MaxTokens =
+      Ctx.getLangOpts().AllocTokenMax.value_or(~0ULL >> (64 - BitWidth));
+
+  // We do not read any of the arguments; discard them.
+  for (int I = Call->getNumArgs() - 1; I >= 0; --I)
+    discard(S.Stk, *S.getContext().classify(Call->getArg(I)));
+
+  // Note: Type inference from a surrounding cast is not supported in
+  // constexpr evaluation.
+  QualType AllocType = infer_alloc::inferPossibleType(Call, Ctx, nullptr);
+  if (AllocType.isNull()) {
+    S.CCEDiag(Call) << "could not infer allocation type";
+    return false;
+  }
+
+  auto ATMD = infer_alloc::getAllocTokenMetadata(AllocType, Ctx);
+  if (!ATMD) {
+    S.CCEDiag(Call) << "could not get token metadata for type";
+    return false;
+  }
+
+  auto MaybeToken = llvm::getAllocTokenHash(Mode, *ATMD, MaxTokens);
+  if (!MaybeToken) {
+    S.CCEDiag(Call) << "stateful alloc token mode not supported in constexpr";
+    return false;
+  }
+
+  pushInteger(S, llvm::APInt(BitWidth, *MaybeToken), Ctx.getSizeType());
+  return true;
+}
+
 static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
                                          const InterpFrame *Frame,
                                          const CallExpr *Call) {
@@ -3489,6 +3529,9 @@ bool InterpretBuiltin(InterpState &S, CodePtr OpPC, const CallExpr *Call,
   case Builtin::BI__builtin_ptrauth_string_discriminator:
     return interp__builtin_ptrauth_string_discriminator(S, OpPC, Frame, Call);
 
+  case Builtin::BI__builtin_infer_alloc_token:
+    return interp__builtin_infer_alloc_token(S, OpPC, Frame, Call);
+
   case Builtin::BI__noop:
     pushInteger(S, 0, Call->getType());
     return true;
diff --git a/clang/lib/AST/ExprConstant.cpp b/clang/lib/AST/ExprConstant.cpp
index 59b4f4f6b5782..137fdb6f0c82b 100644
--- a/clang/lib/AST/ExprConstant.cpp
+++ b/clang/lib/AST/ExprConstant.cpp
@@ -44,6 +44,7 @@
 #include "clang/AST/CharUnits.h"
 #include "clang/AST/CurrentSourceLocExprScope.h"
 #include "clang/AST/Expr.h"
+#include "clang/AST/InferAlloc.h"
 #include "clang/AST/OSLog.h"
 #include "clang/AST/OptionalDiagnostic.h"
 #include "clang/AST/RecordLayout.h"
@@ -14415,6 +14416,26 @@ bool IntExprEvaluator::VisitBuiltinCallExpr(const CallExpr *E,
     return Success(Result, E);
   }
 
+  case Builtin::BI__builtin_infer_alloc_token: {
+    // If we fail to infer a type, this fails to be a constant expression; this
+    // can be checked with __builtin_constant_p(...).
+    QualType AllocType = infer_alloc::inferPossibleType(E, Info.Ctx, nullptr);
+    if (AllocType.isNull())
+      return Error(E);
+    auto ATMD = infer_alloc::getAllocTokenMetadata(AllocType, Info.Ctx);
+    if (!ATMD)
+      return Error(E);
+    auto Mode =
+        Info.getLangOpts().AllocTokenMode.value_or(llvm::DefaultAllocTokenMode);
+    uint64_t BitWidth = Info.Ctx.getTypeSize(Info.Ctx.getSizeType());
+    uint64_t MaxTokens =
+        Info.getLangOpts().AllocTokenMax.value_or(~0ULL >> (64 - BitWidth));
+    auto MaybeToken = llvm::getAllocTokenHash(Mode, *ATMD, MaxTokens);
+    if (!MaybeToken)
+      return Error(E);
+    return Success(llvm::APInt(BitWidth, *MaybeToken), E);
+  }
+
   case Builtin::BI__builtin_ffs:
   case Builtin::BI__builtin_ffsl:
   case Builtin::BI__builtin_ffsll: {
diff --git a/clang/test/SemaCXX/alloc-token.cpp b/clang/test/SemaCXX/alloc-token.cpp
index 5c490f2affda8..4956f517b708b 100644
--- a/clang/test/SemaCXX/alloc-token.cpp
+++ b/clang/test/SemaCXX/alloc-token.cpp
@@ -1,12 +1,57 @@
 // RUN: %clang_cc1 -triple x86_64-linux-gnu -std=c++23 -fsyntax-only -verify %s
 // RUN: %clang_cc1 -triple x86_64-linux-gnu -std=c++23 -fsyntax-only -verify %s -fexperimental-new-constant-interpreter
 // RUN: %clang_cc1 -triple x86_64-linux-gnu -std=c++23 -fsyntax-only -verify %s -falloc-token-mode=typehash -DMODE_TYPEHASH
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -std=c++23 -fsyntax-only -verify %s -falloc-token-max=2 -DTOKEN_MAX=2
 
 #if !__has_builtin(__builtin_infer_alloc_token)
 #error "missing __builtin_infer_alloc_token"
 #endif
 
+#ifndef TOKEN_MAX
+#define TOKEN_MAX 0
+#endif
+
+struct NoPtr {
+  int x;
+  long y;
+};
+
+struct WithPtr {
+  int a;
+  char *buf;
+};
+
+// Check specific known values; these are guaranteed to be stable.
+#ifdef MODE_TYPEHASH
+static_assert(__builtin_infer_alloc_token(sizeof(int)) == 2689373973731826898ULL);
+static_assert(__builtin_infer_alloc_token(sizeof(char*)) == 2250492667400517147ULL);
+static_assert(__builtin_infer_alloc_token(sizeof(NoPtr)) == 7465259095297095368ULL);
+static_assert(__builtin_infer_alloc_token(sizeof(WithPtr)) == 11898882936532569145ULL);
+#elif TOKEN_MAX == 2
+static_assert(__builtin_infer_alloc_token(sizeof(int)) == 0);
+static_assert(__builtin_infer_alloc_token(sizeof(char*)) == 1);
+static_assert(__builtin_infer_alloc_token(sizeof(NoPtr)) == 0);
+static_assert(__builtin_infer_alloc_token(sizeof(WithPtr)) == 1);
+#else
+static_assert(__builtin_infer_alloc_token(sizeof(int)) == 2689373973731826898ULL);
+static_assert(__builtin_infer_alloc_token(sizeof(char*)) == 11473864704255292954ULL);
+static_assert(__builtin_infer_alloc_token(sizeof(NoPtr)) == 7465259095297095368ULL);
+static_assert(__builtin_infer_alloc_token(sizeof(WithPtr)) == 11898882936532569145ULL);
+#endif
+
+template <typename T>
+constexpr unsigned long long get_token() {
+  return __builtin_infer_alloc_token(sizeof(T));
+}
+
+// Test complex expressions.
+static_assert(__builtin_constant_p(__builtin_infer_alloc_token(sizeof(int))));
+static_assert(__builtin_infer_alloc_token(sizeof(NoPtr) * 2, 1) == get_token<NoPtr>());
+static_assert(__builtin_infer_alloc_token(1, 4 + sizeof(NoPtr)) == get_token<NoPtr>());
+static_assert(__builtin_infer_alloc_token(sizeof(NoPtr) << 8) == get_token<NoPtr>());
+
 void negative_tests() {
   __builtin_infer_alloc_token(); // expected-error {{too few arguments to function call}}
   __builtin_infer_alloc_token((void)0); // expected-error {{argument may not have 'void' type}}
+  constexpr auto inference_fail = __builtin_infer_alloc_token(123); // expected-error {{must be initialized by a constant expression}}
 }