2828#include " clang/Lex/Lexer.h"
2929#include " clang/Lex/Preprocessor.h"
3030#include " llvm/ADT/APSInt.h"
31+ #include " llvm/ADT/PointerIntPair.h"
3132#include " llvm/ADT/STLFunctionalExtras.h"
3233#include " llvm/ADT/SmallSet.h"
3334#include " llvm/ADT/SmallVector.h"
@@ -421,8 +422,9 @@ const Expr *findCountArg(const Expr *Count, const CallExpr *Call) {
421422 return Call->getArg (Index);
422423}
423424
424- // Mapping: dependent decl -> value.
425- using DependentValuesTy = llvm::DenseMap<const ValueDecl *, const Expr *>;
425+ // Mapping: (DependentDecl, Deref) -> Value.
426+ using DeclDerefPair = llvm::PointerIntPair<const ValueDecl *, 1 , bool >;
427+ using DependentValuesTy = llvm::DenseMap<DeclDerefPair, const Expr *>;
426428
427429// Given the call expr, find the mapping from the dependent parameter to the
428430// argument that is passed to that parameter.
@@ -440,7 +442,8 @@ getDependentValuesFromCall(const CountAttributedType *CAT,
440442 return std::nullopt ;
441443
442444 const Expr *Arg = Call->getArg (Index);
443- [[maybe_unused]] bool Inserted = Values.insert ({PVD, Arg}).second ;
445+ [[maybe_unused]] bool Inserted =
446+ Values.insert ({{PVD, /* Deref=*/ false }, Arg}).second ;
444447 assert (Inserted);
445448 }
446449 return {std::move (Values)};
@@ -494,50 +497,57 @@ struct CompatibleCountExprVisitor
494497 const Expr *
495498 trySubstituteAndSimplify (const Expr *E, bool &hasBeenSubstituted,
496499 const DependentValuesTy *DependentValues) const {
500+ auto trySubstitute = [&](const ValueDecl *VD, bool Deref) -> const Expr * {
501+ if (hasBeenSubstituted || !DependentValues)
502+ return nullptr ;
503+ auto It = DependentValues->find ({VD, Deref});
504+ return It != DependentValues->end () ? It->second : nullptr ;
505+ };
506+
497507 // Attempts to simplify `E`: if `E` has the form `*&e`, return `e`;
498508 // return `E` without change otherwise:
499- auto trySimplifyDerefAddressof =
500- [](const Expr *E,
501- const DependentValuesTy
502- *DependentValues, // Deref may need subsitution
503- bool &hasBeenSubstituted) -> const Expr * {
509+ auto trySimplifyDeref = [&](const Expr *E) -> const Expr * {
504510 const auto *Deref = dyn_cast<UnaryOperator>(E->IgnoreParenImpCasts ());
505511
506512 if (!Deref || Deref->getOpcode () != UO_Deref)
507513 return E;
508514
509515 const Expr *DerefOperand = Deref->getSubExpr ()->IgnoreParenImpCasts ();
510516
517+ // Just simplify `*&...`.
511518 if (const auto *UO = dyn_cast<UnaryOperator>(DerefOperand))
512519 if (UO->getOpcode () == UO_AddrOf)
513520 return UO->getSubExpr ();
521+
514522 if (const auto *DRE = dyn_cast<DeclRefExpr>(DerefOperand)) {
515- if (!DependentValues || hasBeenSubstituted)
516- return E;
517-
518- if (auto I = DependentValues->find (DRE->getDecl ());
519- I != DependentValues->end ())
520- if (const auto *UO = dyn_cast<UnaryOperator>(
521- I->getSecond ()->IgnoreParenImpCasts ()))
522- if (UO->getOpcode () == UO_AddrOf) {
523- hasBeenSubstituted = true ;
524- return UO->getSubExpr ();
525- }
523+ // Substitute `*x`.
524+ if (const auto *Sub = trySubstitute (DRE->getDecl (), /* Deref=*/ true )) {
525+ hasBeenSubstituted = true ;
526+ return Sub;
527+ }
528+
529+ // Substitute `x` in `*x` if we have `x -> &...` in our mapping.
530+ if (const auto *Sub = trySubstitute (DRE->getDecl (), /* Deref=*/ false )) {
531+ if (const auto *UO =
532+ dyn_cast<UnaryOperator>(Sub->IgnoreParenImpCasts ());
533+ UO && UO->getOpcode () == UO_AddrOf) {
534+ hasBeenSubstituted = true ;
535+ return UO->getSubExpr ();
536+ }
537+ }
526538 }
539+
527540 return E;
528541 };
529542
530- if (!hasBeenSubstituted && DependentValues) {
531- if (const auto *DRE = dyn_cast<DeclRefExpr>(E->IgnoreParenImpCasts ())) {
532- if (auto It = DependentValues->find (DRE->getDecl ());
533- It != DependentValues->end ()) {
534- hasBeenSubstituted = true ;
535- return trySimplifyDerefAddressof (It->second , nullptr ,
536- hasBeenSubstituted);
537- }
543+ if (const auto *DRE = dyn_cast<DeclRefExpr>(E->IgnoreParenImpCasts ())) {
544+ if (const auto *Sub = trySubstitute (DRE->getDecl (), /* Deref=*/ false )) {
545+ hasBeenSubstituted = true ;
546+ return trySimplifyDeref (Sub);
538547 }
539548 }
540- return trySimplifyDerefAddressof (E, DependentValues, hasBeenSubstituted);
549+
550+ return trySimplifyDeref (E);
541551 }
542552
543553 explicit CompatibleCountExprVisitor (
@@ -668,6 +678,13 @@ struct CompatibleCountExprVisitor
668678 bool hasOtherBeenSubstituted) {
669679 if (SelfUO->getOpcode () != UO_Deref)
670680 return false ; // We don't support any other unary operator
681+
682+ const auto *SimplifiedSelf = trySubstituteAndSimplify (
683+ SelfUO, hasSelfBeenSubstituted, DependentValuesSelf);
684+ if (SimplifiedSelf != SelfUO)
685+ return Visit (SimplifiedSelf, Other, hasSelfBeenSubstituted,
686+ hasOtherBeenSubstituted);
687+
671688 Other = trySubstituteAndSimplify (Other, hasOtherBeenSubstituted,
672689 DependentValuesOther);
673690 if (const auto *OtherUO =
@@ -676,13 +693,7 @@ struct CompatibleCountExprVisitor
676693 return Visit (SelfUO->getSubExpr (), OtherUO->getSubExpr (),
677694 hasSelfBeenSubstituted, hasOtherBeenSubstituted);
678695 }
679- // If `Other` is not a dereference expression, try to simplify `SelfUO`:
680- const auto *SimplifiedSelf = trySubstituteAndSimplify (
681- SelfUO, hasSelfBeenSubstituted, DependentValuesSelf);
682696
683- if (SimplifiedSelf != SelfUO)
684- return Visit (SimplifiedSelf, Other, hasSelfBeenSubstituted,
685- hasOtherBeenSubstituted);
686697 return false ;
687698 }
688699
0 commit comments