@@ -118,6 +118,7 @@ class OriginManager {
118118    return  AllOrigins.back ();
119119  }
120120
121+   //  TODO: Mark this method as const once we remove the call to getOrCreate.
121122  OriginID get (const  Expr &E) {
122123    //  Origin of DeclRefExpr is that of the declaration it refers to.
123124    if  (const  auto  *DRE = dyn_cast<DeclRefExpr>(&E))
@@ -314,22 +315,28 @@ class ReturnOfOriginFact : public Fact {
314315};
315316
316317class  UseFact  : public  Fact  {
317-   OriginID UsedOrigin;
318318  const  Expr *UseExpr;
319+   //  True if this use is a write operation (e.g., left-hand side of assignment).
320+   //  Write operations are exempted from use-after-free checks.
321+   bool  IsWritten = false ;
319322
320323public: 
321324  static  bool  classof (const  Fact *F) { return  F->getKind () == Kind::Use; }
322325
323-   UseFact (OriginID UsedOrigin, const  Expr *UseExpr)
324-       : Fact(Kind::Use), UsedOrigin(UsedOrigin), UseExpr(UseExpr) {}
326+   UseFact (const  Expr *UseExpr) : Fact(Kind::Use), UseExpr(UseExpr) {}
325327
326-   OriginID getUsedOrigin () const  { return  UsedOrigin; }
328+   OriginID getUsedOrigin (const  OriginManager &OM) const  {
329+     //  TODO: Remove const cast and make OriginManager::get as const.
330+     return  const_cast <OriginManager &>(OM).get (*UseExpr);
331+   }
327332  const  Expr *getUseExpr () const  { return  UseExpr; }
333+   void  markAsWritten () { IsWritten = true ; }
334+   bool  isWritten () const  { return  IsWritten; }
328335
329336  void  dump (llvm::raw_ostream &OS, const  OriginManager &OM) const  override  {
330337    OS << " Use ("  ;
331-     OM.dump (getUsedOrigin (), OS);
332-     OS << " )\n "  ;
338+     OM.dump (getUsedOrigin (OM ), OS);
339+     OS << "   "  << ( isWritten () ?  " Write "  :  " Read " ) <<  " )\n "  ;
333340  }
334341};
335342
@@ -436,6 +443,8 @@ class FactGenerator : public ConstStmtVisitor<FactGenerator> {
436443            addAssignOriginFact (*VD, *InitExpr);
437444  }
438445
446+   void  VisitDeclRefExpr (const  DeclRefExpr *DRE) { handleUse (DRE); }
447+ 
439448  void  VisitCXXNullPtrLiteralExpr (const  CXXNullPtrLiteralExpr *N) {
440449    // / TODO: Handle nullptr expr as a special 'null' loan. Uninitialized
441450    // / pointers can use the same type of loan.
@@ -469,10 +478,6 @@ class FactGenerator : public ConstStmtVisitor<FactGenerator> {
469478          }
470479        }
471480      }
472-     } else  if  (UO->getOpcode () == UO_Deref) {
473-       //  This is a pointer use, like '*p'.
474-       OriginID OID = FactMgr.getOriginMgr ().get (*UO->getSubExpr ());
475-       CurrentBlockFacts.push_back (FactMgr.createFact <UseFact>(OID, UO));
476481    }
477482  }
478483
@@ -487,20 +492,13 @@ class FactGenerator : public ConstStmtVisitor<FactGenerator> {
487492  }
488493
489494  void  VisitBinaryOperator (const  BinaryOperator *BO) {
490-     if  (BO->isAssignmentOp ()) {
491-       const  Expr *LHSExpr = BO->getLHS ();
492-       const  Expr *RHSExpr = BO->getRHS ();
493- 
494-       //  We are interested in assignments like `ptr1 = ptr2` or `ptr = &var`
495-       //  LHS must be a pointer/reference type that can be an origin.
496-       //  RHS must also represent an origin (either another pointer/ref or an
497-       //  address-of).
498-       if  (const  auto  *DRE_LHS = dyn_cast<DeclRefExpr>(LHSExpr))
499-         if  (const  auto  *VD_LHS =
500-                 dyn_cast<ValueDecl>(DRE_LHS->getDecl ()->getCanonicalDecl ());
501-             VD_LHS && hasOrigin (VD_LHS->getType ()))
502-           addAssignOriginFact (*VD_LHS, *RHSExpr);
503-     }
495+     if  (BO->isAssignmentOp ())
496+       handleAssignment (BO->getLHS (), BO->getRHS ());
497+   }
498+ 
499+   void  VisitCXXOperatorCallExpr (const  CXXOperatorCallExpr *OCE) {
500+     if  (OCE->isAssignmentOp () && OCE->getNumArgs () == 2 )
501+       handleAssignment (OCE->getArg (0 ), OCE->getArg (1 ));
504502  }
505503
506504  void  VisitCXXFunctionalCastExpr (const  CXXFunctionalCastExpr *FCE) {
@@ -567,9 +565,47 @@ class FactGenerator : public ConstStmtVisitor<FactGenerator> {
567565    return  false ;
568566  }
569567
568+   void  handleAssignment (const  Expr *LHSExpr, const  Expr *RHSExpr) {
569+     //  Find the underlying variable declaration for the left-hand side.
570+     if  (const  auto  *DRE_LHS =
571+             dyn_cast<DeclRefExpr>(LHSExpr->IgnoreParenImpCasts ())) {
572+       markUseAsWrite (DRE_LHS);
573+       if  (const  auto  *VD_LHS = dyn_cast<ValueDecl>(DRE_LHS->getDecl ()))
574+         if  (hasOrigin (LHSExpr->getType ()))
575+           //  We are interested in assignments like `ptr1 = ptr2` or `ptr = &var`
576+           //  LHS must be a pointer/reference type that can be an origin.
577+           //  RHS must also represent an origin (either another pointer/ref or an
578+           //  address-of).
579+           addAssignOriginFact (*VD_LHS, *RHSExpr);
580+     }
581+   }
582+ 
583+   //  A DeclRefExpr is a use of the referenced decl. It is checked for
584+   //  use-after-free unless it is being written to (e.g. on the left-hand side
585+   //  of an assignment).
586+   void  handleUse (const  DeclRefExpr *DRE) {
587+     if  (hasOrigin (DRE->getType ())) {
588+       UseFact *UF = FactMgr.createFact <UseFact>(DRE);
589+       CurrentBlockFacts.push_back (UF);
590+       assert (!UseFacts.contains (DRE));
591+       UseFacts[DRE] = UF;
592+     }
593+   }
594+ 
595+   void  markUseAsWrite (const  DeclRefExpr *DRE) {
596+     assert (UseFacts.contains (DRE));
597+     UseFacts[DRE]->markAsWritten ();
598+   }
599+ 
570600  FactManager &FactMgr;
571601  AnalysisDeclContext &AC;
572602  llvm::SmallVector<Fact *> CurrentBlockFacts;
603+   //  To distinguish between reads and writes for use-after-free checks, this map
604+   //  stores the `UseFact` for each `DeclRefExpr`. We initially identify all
605+   //  `DeclRefExpr`s as "read" uses. When an assignment is processed, the use
606+   //  corresponding to the left-hand side is updated to be a "write", thereby
607+   //  exempting it from the check.
608+   llvm::DenseMap<const  DeclRefExpr *, UseFact *> UseFacts;
573609};
574610
575611//  ========================================================================= //
@@ -1032,8 +1068,9 @@ class LifetimeChecker {
10321068  // / graph. It determines if the loans held by the used origin have expired
10331069  // / at the point of use.
10341070  void  checkUse (const  UseFact *UF) {
1035- 
1036-     OriginID O = UF->getUsedOrigin ();
1071+     if  (UF->isWritten ())
1072+       return ;
1073+     OriginID O = UF->getUsedOrigin (FactMgr.getOriginMgr ());
10371074
10381075    //  Get the set of loans that the origin might hold at this program point.
10391076    LoanSet HeldLoans = LoanPropagation.getLoans (O, UF);
0 commit comments