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

[Marco Elver via llvm-branch-commits]

================
@@ -1353,6 +1354,92 @@ void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) {
   CB->setMetadata(llvm::LLVMContext::MD_alloc_token, 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.
+  if (QualType T = inferTypeFromSizeofExpr(Arg); !T.isNull())
+    return T;
+  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:
+      if (QualType T = inferTypeFromArithSizeofExpr(BO->getLHS()); !T.isNull())
+        return T;
+      if (QualType T = inferTypeFromArithSizeofExpr(BO->getRHS()); !T.isNull())
+        return T;
+      break;
+    default:
+      break;
+    }
+  }
+  return QualType();
+}
+
+/// If the expression E is a reference to a variable, infer the type from a
+/// variable's initializer if it contains a sizeof. Beware, this is a heuristic
+/// and ignores if a variable is later reassigned.
+static QualType inferTypeFromVarInitSizeofExpr(const Expr *E) {
+  const Expr *Arg = E->IgnoreParenImpCasts();
+  if (const auto *DRE = dyn_cast<DeclRefExpr>(Arg)) {
+    if (const auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
+      if (const Expr *Init = VD->getInit())
+        return inferTypeFromArithSizeofExpr(Init);
+    }
+  }
+  return QualType();
+}
+
+/// Deduces the allocated type by checking if the allocation call's result
+/// is immediately used in a cast expression.
+static QualType inferTypeFromCastExpr(const CallExpr *CallE,
----------------
melver wrote:

You mean for inferTypeFromCastExpr specifically, or in general for all?

General: I think some of the recursive visitors (RecursiveASTVisitor or also EvaluatedExprVisitor) could help, but they could only give us a differently shaped inferPossibleTypeFromArithSizeofExpr(). They also don't make early return and control over the order of visits that easy (unless we override the Visit function, which defeats the point). And as for efficiency, as I wrote elsewhere, here we can do more targeted visiting of relevant AST nodes and the overall algorithm is much more lightweight.

For the "infer from cast" problem specifically, I used to have this implementation:
```
static QualType inferTypeFromCastExpr(const CallExpr *E, ASTContext &Ctx) {
  DynTypedNodeList Parents = Ctx.getParentMapContext().getParents(*E);
  if (Parents.empty())
    return QualType();
  // We only care about the most direct parent for this heuristic.
  if (const auto *CE = Parents[0].get<CastExpr>()) {
    QualType PtrType = CE->getType();
    if (PtrType->isPointerType())
      return PtrType->getPointeeType();
  }
  return QualType();
}
```
But the `getParents` implementation uses RecursiveASTVisitor underneath, which turned out to be really slow and also ended up crashing on some code using coroutines. So instead, the current inferTypeFromCastExpr() implementation is just O(1) by keeping the innermost CurCast saved during normal AST walk.

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