[llvm-branch-commits] [AllocToken, Clang] Infer type hints from sizeof expressions and casts (PR #156841)

[Marco Elver via llvm-branch-commits]

================
@@ -1349,6 +1350,98 @@ void CodeGenFunction::EmitAllocTokenHint(llvm::CallBase *CB,
   CB->setMetadata(llvm::LLVMContext::MD_alloc_token_hint, MDN);
 }
 
+/// Infer type from a simple sizeof expression.
+static QualType inferTypeFromSizeofExpr(const Expr *E) {
+  const Expr *Arg = E->IgnoreParenImpCasts();
+  if (const auto *UET = dyn_cast<UnaryExprOrTypeTraitExpr>(Arg)) {
+    if (UET->getKind() == UETT_SizeOf) {
+      if (UET->isArgumentType()) {
+        return UET->getArgumentTypeInfo()->getType();
+      } else {
+        return UET->getArgumentExpr()->getType();
+      }
+    }
+  }
+  return QualType();
+}
+
+/// Infer type from an arithmetic expression involving a sizeof.
+static QualType inferTypeFromArithSizeofExpr(const Expr *E) {
+  const Expr *Arg = E->IgnoreParenImpCasts();
+  // The argument is a lone sizeof expression.
+  QualType QT = inferTypeFromSizeofExpr(Arg);
+  if (!QT.isNull())
+    return QT;
+  if (const auto *BO = dyn_cast<BinaryOperator>(Arg)) {
+    // Argument is an arithmetic expression. Cover common arithmetic patterns
+    // involving sizeof.
+    switch (BO->getOpcode()) {
+    case BO_Add:
+    case BO_Div:
+    case BO_Mul:
+    case BO_Shl:
+    case BO_Shr:
+    case BO_Sub:
+      QT = inferTypeFromArithSizeofExpr(BO->getLHS());
----------------
melver wrote:

The Linux kernel has structs with flexible array members, and it's not uncommon to see this:
```
struct A {
  int len;
  struct Foo *foo;
  int array[];
};

... = kmalloc(sizeof(A) + sizeof(int) * N, ...);
```

I'm willing to accept some degree of unsoundness in complex cases to get completeness here, but am assuming that in the majority of cases the first type is the one we want to pick.

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