Merge "Handle return values of rsObject type"

6 files changed, 425 insertions, 199 deletions

diff --git a/lit-tests/P_ref_count/ref_count.rs b/lit-tests/P_ref_count/ref_count.rs
new file mode 100644
index 0000000..5abea94
--- /dev/null
+++ b/lit-tests/P_ref_count/ref_count.rs
@@ -0,0 +1,21 @@
+// RUN: %Slang %s
+// RUN: %rs-filecheck-wrapper %s
+// CHECK: call void @_Z13rsClearObjectP10rs_element(%struct.rs_element{{.*}}* %.rs.tmp{{[0-9]+}})
+// CHECK: call void @_Z11rsSetObjectP10rs_elementS_(%struct.rs_element{{.*}}* %.rs.retval{{[0-9]+}}, {{.*}})
+
+#pragma version(1)
+#pragma rs java_package_name(ref_count)
+
+static rs_element bar() {
+  rs_element x = {0};
+  return x;
+}
+
+void entrypoint() {
+  rs_element e = bar();
+  if (rsIsObject(e)) {
+    rsDebug("good object", 0);
+  }
+}
+
+
diff --git a/slang_backend.cpp b/slang_backend.cpp
index d6522a7..cbfdb32 100644
--- a/slang_backend.cpp
+++ b/slang_backend.cpp
@@ -366,6 +366,7 @@ void Backend::AnnotateFunction(clang::FunctionDecl *FD) {
       FD->hasBody() &&
       !Slang::IsLocInRSHeaderFile(FD->getLocation(), mSourceMgr)) {
     mRefCount.Init();
+    mRefCount.SetDeclContext(FD);
     mRefCount.Visit(FD->getBody());
   }
 }
diff --git a/slang_rs_ast_replace.cpp b/slang_rs_ast_replace.cpp
index d5af26c..ef95b08 100644
--- a/slang_rs_ast_replace.cpp
+++ b/slang_rs_ast_replace.cpp
@@ -87,6 +87,18 @@ void RSASTReplace::VisitCaseStmt(clang::CaseStmt *CS) {
   }
 }
 
+void RSASTReplace::VisitDeclStmt(clang::DeclStmt* DS) {
+  VisitStmt(DS);
+  for (clang::Decl* D : DS->decls()) {
+    clang::VarDecl* VD;
+    if ((VD = llvm::dyn_cast<clang::VarDecl>(D))) {
+      if (matchesExpr(VD->getInit())) {
+        VD->setInit(mNewExpr);
+      }
+    }
+  }
+}
+
 void RSASTReplace::VisitDefaultStmt(clang::DefaultStmt *DS) {
   if (matchesStmt(DS->getSubStmt())) {
     DS->setSubStmt(mNewStmt);
diff --git a/slang_rs_ast_replace.h b/slang_rs_ast_replace.h
index f93171d..58f9bd9 100644
--- a/slang_rs_ast_replace.h
+++ b/slang_rs_ast_replace.h
@@ -32,7 +32,7 @@ namespace slang {
 
 class RSASTReplace : public clang::StmtVisitor<RSASTReplace> {
  private:
-  clang::ASTContext &C;
+  const clang::ASTContext &C;
   clang::Stmt *mOuterStmt;
   clang::Stmt *mOldStmt;
   clang::Stmt *mNewStmt;
@@ -56,7 +56,7 @@ class RSASTReplace : public clang::StmtVisitor<RSASTReplace> {
   void ReplaceInCompoundStmt(clang::CompoundStmt *CS);
 
  public:
-  explicit RSASTReplace(clang::ASTContext &Con)
+  explicit RSASTReplace(const clang::ASTContext &Con)
       : C(Con),
         mOuterStmt(nullptr),
         mOldStmt(nullptr),
@@ -68,6 +68,7 @@ class RSASTReplace : public clang::StmtVisitor<RSASTReplace> {
   void VisitStmt(clang::Stmt *S);
   void VisitCompoundStmt(clang::CompoundStmt *CS);
   void VisitCaseStmt(clang::CaseStmt *CS);
+  void VisitDeclStmt(clang::DeclStmt* DS);
   void VisitDefaultStmt(clang::DefaultStmt *DS);
   void VisitDoStmt(clang::DoStmt *DS);
   void VisitForStmt(clang::ForStmt *FS);
diff --git a/slang_rs_object_ref_count.cpp b/slang_rs_object_ref_count.cpp
index 51c9779..bf93eb0 100644
--- a/slang_rs_object_ref_count.cpp
+++ b/slang_rs_object_ref_count.cpp
@@ -92,8 +92,16 @@ void RSObjectRefCount::GetRSRefCountingFunctions(clang::ASTContext &C) {
 
 namespace {
 
+unsigned CountRSObjectTypes(const clang::Type *T);
+
+clang::Stmt *CreateSingleRSSetObject(clang::ASTContext &C,
+                                     clang::Expr *DstExpr,
+                                     clang::Expr *SrcExpr,
+                                     clang::SourceLocation StartLoc,
+                                     clang::SourceLocation Loc);
+
 // This function constructs a new CompoundStmt from the input StmtList.
-static clang::CompoundStmt* BuildCompoundStmt(clang::ASTContext &C,
+clang::CompoundStmt* BuildCompoundStmt(clang::ASTContext &C,
       std::list<clang::Stmt*> &StmtList, clang::SourceLocation Loc) {
   unsigned NewStmtCount = StmtList.size();
   unsigned CompoundStmtCount = 0;
@@ -116,10 +124,10 @@ static clang::CompoundStmt* BuildCompoundStmt(clang::ASTContext &C,
   return CS;
 }
 
-static void AppendAfterStmt(clang::ASTContext &C,
-                            clang::CompoundStmt *CS,
-                            clang::Stmt *S,
-                            std::list<clang::Stmt*> &StmtList) {
+void AppendAfterStmt(clang::ASTContext &C,
+                     clang::CompoundStmt *CS,
+                     clang::Stmt *S,
+                     std::list<clang::Stmt*> &StmtList) {
   slangAssert(CS);
   clang::CompoundStmt::body_iterator bI = CS->body_begin();
   clang::CompoundStmt::body_iterator bE = CS->body_end();
@@ -166,23 +174,12 @@ static void AppendAfterStmt(clang::ASTContext &C,
   delete [] UpdatedStmtList;
 }
 
-// This class visits a compound statement and inserts DtorStmt
-// in proper locations. This includes inserting it before any
-// return statement in any sub-block, at the end of the logical enclosing
-// scope (compound statement), and/or before any break/continue statement that
-// would resume outside the declared scope. We will not handle the case for
-// goto statements that leave a local scope.
-//
-// To accomplish these goals, it collects a list of sub-Stmt's that
-// correspond to scope exit points. It then uses an RSASTReplace visitor to
-// transform the AST, inserting appropriate destructors before each of those
-// sub-Stmt's (and also before the exit of the outermost containing Stmt for
-// the scope).
+// This class visits a compound statement and collects a list of all the exiting
+// statements, such as any return statement in any sub-block, and any
+// break/continue statement that would resume outside the current scope.
+// We do not handle the case for goto statements that leave a local scope.
 class DestructorVisitor : public clang::StmtVisitor<DestructorVisitor> {
  private:
-  clang::ASTContext &mCtx;
-  clang::DeclContext *mDC;
-
   // The loop depth of the currently visited node.
   int mLoopDepth;
 
@@ -192,123 +189,19 @@ class DestructorVisitor : public clang::StmtVisitor<DestructorVisitor> {
   // corresponding loop scope.
   int mSwitchDepth;
 
-  // The outermost statement block that we are currently visiting.
-  // This should always be a CompoundStmt.
-  clang::Stmt *mOuterStmt;
-
-  // The destructor to execute for this scope/variable.
-  clang::Stmt* mDtorStmt;
-
-  // The stack of statements which should be replaced by a compound statement
-  // containing the new destructor call followed by the original Stmt.
-  std::stack<clang::Stmt*> mReplaceStmtStack;
-
-  // The source location for the variable declaration that we are trying to
-  // insert destructors for. Note that InsertDestructors() will not generate
-  // destructor calls for source locations that occur lexically before this
-  // location.
-  clang::SourceLocation mVarLoc;
+  // Output of the visitor: the statements that should be replaced by compound
+  // statements, each of which contains rsClearObject() calls followed by the
+  // original statement.
+  std::vector<clang::Stmt*> mExitingStmts;
 
  public:
-  DestructorVisitor(clang::DeclContext* DC,
-                    clang::Stmt* OuterStmt,
-                    clang::Stmt* DtorStmt,
-                    clang::SourceLocation VarLoc);
-
-  // This code walks the collected list of Stmts to replace and actually does
-  // the replacement. It also finishes up by appending the destructor to the
-  // current outermost CompoundStmt.
-  void InsertDestructors() {
-    clang::Stmt *S = nullptr;
-    clang::SourceManager &SM = mCtx.getSourceManager();
-    std::list<clang::Stmt *> StmtList;
-    StmtList.push_back(mDtorStmt);
-
-    while (!mReplaceStmtStack.empty()) {
-      S = mReplaceStmtStack.top();
-      mReplaceStmtStack.pop();
-
-      // Skip all source locations that occur before the variable's
-      // declaration, since it won't have been initialized yet.
-      if (SM.isBeforeInTranslationUnit(S->getLocStart(), mVarLoc)) {
-        continue;
-      }
-
-      clang::CompoundStmt *CS;
-      clang::ReturnStmt* RS = llvm::dyn_cast<clang::ReturnStmt>(S);
-      clang::Expr* RetVal;
-      if (!RS || !(RetVal = RS->getRetValue())) {
-        StmtList.push_back(S);
-        CS = BuildCompoundStmt(mCtx, StmtList, S->getLocEnd());
-        StmtList.pop_back();
-      } else {
-        // Since we insert rsClearObj() calls before the return statement, we need
-        // to make sure none of the cleared RS objects are referenced in the
-        // return statement.
-        // For that, we create a new local variable named .rs.retval, assign the
-        // original return expression to it, make all necessary rsClearObj()
-        // calls, then return .rs.retval. Note rsSetObj() or rsClearObj() are not
-        // called on .rs.retval.
-
-        clang::SourceLocation Loc;
-        clang::QualType RetTy = RetVal->getType();
-        clang::VarDecl* RSRetValDecl = clang::VarDecl::Create(
-            mCtx,                                  // AST context
-            mDC,                                   // Decl context
-            Loc,                                   // Start location
-            Loc,                                   // Id location
-            &mCtx.Idents.get(".rs.retval"),        // Id
-            RetTy,                                 // Type
-            mCtx.getTrivialTypeSourceInfo(RetTy),  // Type info
-            clang::SC_None                         // Storage class
-        );
-        RSRetValDecl->setInit(RetVal);
-        clang::Decl* Decls[] = { RSRetValDecl };
-        const clang::DeclGroupRef DGR = clang::DeclGroupRef::Create(
-            mCtx, Decls, sizeof(Decls) / sizeof(*Decls));
-        clang::DeclStmt* DS = new (mCtx) clang::DeclStmt(DGR, Loc, Loc);
-
-        // Creates a new return statement
-        clang::ReturnStmt* NewRet = new (mCtx) clang::ReturnStmt(RS->getReturnLoc());
-        clang::DeclRefExpr* DRE = clang::DeclRefExpr::Create(
-            mCtx,
-            clang::NestedNameSpecifierLoc(),  // QualifierLoc
-            Loc,                              // TemplateKWLoc
-            RSRetValDecl,
-            false,                            // RefersToEnclosingVariableOrCapture
-            Loc,                              // NameLoc
-            RetTy,
-            clang::VK_LValue
-        );
-        clang::Expr* CastExpr = clang::ImplicitCastExpr::Create(
-            mCtx,
-            RetTy,
-            clang::CK_LValueToRValue,
-            DRE,
-            nullptr,
-            clang::VK_RValue
-        );
-        NewRet->setRetValue(CastExpr);
-
-        // Insert the two new statements into StmtList
-        StmtList.push_front(DS);
-        StmtList.push_back(NewRet);
-        CS = BuildCompoundStmt(mCtx, StmtList, S->getLocEnd());
-        StmtList.pop_back();
-        StmtList.pop_front();
-      }
+  DestructorVisitor() : mLoopDepth(0), mSwitchDepth(0) {}
 
-      RSASTReplace R(mCtx);
-      R.ReplaceStmt(mOuterStmt, S, CS);
-    }
-    clang::CompoundStmt *CS =
-      llvm::dyn_cast<clang::CompoundStmt>(mOuterStmt);
-    slangAssert(CS);
-    AppendAfterStmt(mCtx, CS, nullptr, StmtList);
+  const std::vector<clang::Stmt*>& getExitingStmts() const {
+    return mExitingStmts;
   }
 
   void VisitStmt(clang::Stmt *S);
-
   void VisitBreakStmt(clang::BreakStmt *BS);
   void VisitContinueStmt(clang::ContinueStmt *CS);
   void VisitDoStmt(clang::DoStmt *DS);
@@ -318,17 +211,75 @@ class DestructorVisitor : public clang::StmtVisitor<DestructorVisitor> {
   void VisitWhileStmt(clang::WhileStmt *WS);
 };
 
-DestructorVisitor::DestructorVisitor(clang::DeclContext* DC,
-                         clang::Stmt *OuterStmt,
-                         clang::Stmt *DtorStmt,
-                         clang::SourceLocation VarLoc)
-  : mCtx(DC->getParentASTContext()),
-    mDC(DC),
-    mLoopDepth(0),
-    mSwitchDepth(0),
-    mOuterStmt(OuterStmt),
-    mDtorStmt(DtorStmt),
-    mVarLoc(VarLoc) {
+// Given a return statement RS that returns an rsObject, creates a temporary
+// variable, sets it to the original return expression using rsSetObject(),
+// adds these new statements into NewStmts.
+// Finally, creates and returns a new return statement that returns the
+// temporary variable.
+clang::CompoundStmt* CreateRetStmtWithTempVar(
+    clang::ASTContext& C,
+    clang::DeclContext* DC,
+    clang::ReturnStmt* RS,
+    const unsigned id) {
+  std::list<clang::Stmt*> NewStmts;
+  // Since we insert rsClearObj() calls before the return statement, we need
+  // to make sure none of the cleared RS objects are referenced in the
+  // return statement.
+  // For that, we create a new local variable named .rs.retval, assign the
+  // original return expression to it, make all necessary rsClearObj()
+  // calls, then return .rs.retval. Note rsClearObj() is not called on
+  // .rs.retval.
+
+  clang::SourceLocation Loc = RS->getLocStart();
+  std::stringstream ss;
+  ss << ".rs.retval" << id;
+  llvm::StringRef VarName(ss.str());
+
+  clang::Expr* RetVal = RS->getRetValue();
+  const clang::QualType RetTy = RetVal->getType();
+  clang::VarDecl* RSRetValDecl = clang::VarDecl::Create(
+      C,                                     // AST context
+      DC,                                    // Decl context
+      Loc,                                   // Start location
+      Loc,                                   // Id location
+      &C.Idents.get(VarName),                // Id
+      RetTy,                                 // Type
+      C.getTrivialTypeSourceInfo(RetTy),     // Type info
+      clang::SC_None                         // Storage class
+  );
+  clang::Decl* Decls[] = { RSRetValDecl };
+  const clang::DeclGroupRef DGR = clang::DeclGroupRef::Create(
+      C, Decls, sizeof(Decls) / sizeof(*Decls));
+  clang::DeclStmt* DS = new (C) clang::DeclStmt(DGR, Loc, Loc);
+  NewStmts.push_back(DS);
+
+  clang::DeclRefExpr* DRE = clang::DeclRefExpr::Create(
+      C,
+      clang::NestedNameSpecifierLoc(),       // QualifierLoc
+      Loc,                                   // TemplateKWLoc
+      RSRetValDecl,
+      false,                                 // RefersToEnclosingVariableOrCapture
+      Loc,                                   // NameLoc
+      RetTy,
+      clang::VK_LValue
+  );
+  clang::Stmt* SetRetTempVar = CreateSingleRSSetObject(C, DRE, RetVal, Loc, Loc);
+  NewStmts.push_back(SetRetTempVar);
+
+  // Creates a new return statement
+  clang::ReturnStmt* NewRet = new (C) clang::ReturnStmt(RS->getReturnLoc());
+  clang::Expr* CastExpr = clang::ImplicitCastExpr::Create(
+      C,
+      RetTy,
+      clang::CK_LValueToRValue,
+      DRE,
+      nullptr,
+      clang::VK_RValue
+  );
+  NewRet->setRetValue(CastExpr);
+  NewStmts.push_back(NewRet);
+
+  return BuildCompoundStmt(C, NewStmts, Loc);
 }
 
 void DestructorVisitor::VisitStmt(clang::Stmt *S) {
@@ -342,7 +293,7 @@ void DestructorVisitor::VisitStmt(clang::Stmt *S) {
 void DestructorVisitor::VisitBreakStmt(clang::BreakStmt *BS) {
   VisitStmt(BS);
   if ((mLoopDepth == 0) && (mSwitchDepth == 0)) {
-    mReplaceStmtStack.push(BS);
+    mExitingStmts.push_back(BS);
   }
 }
 
@@ -350,7 +301,7 @@ void DestructorVisitor::VisitContinueStmt(clang::ContinueStmt *CS) {
   VisitStmt(CS);
   if (mLoopDepth == 0) {
     // Switch statements can have nested continues.
-    mReplaceStmtStack.push(CS);
+    mExitingStmts.push_back(CS);
   }
 }
 
@@ -367,7 +318,7 @@ void DestructorVisitor::VisitForStmt(clang::ForStmt *FS) {
 }
 
 void DestructorVisitor::VisitReturnStmt(clang::ReturnStmt *RS) {
-  mReplaceStmtStack.push(RS);
+  mExitingStmts.push_back(RS);
 }
 
 void DestructorVisitor::VisitSwitchStmt(clang::SwitchStmt *SS) {
@@ -458,14 +409,14 @@ static int ArrayDim(const clang::Type *T) {
   return static_cast<int>(CAT->getSize().getSExtValue());
 }
 
-static clang::Stmt *ClearStructRSObject(
+clang::Stmt *ClearStructRSObject(
     clang::ASTContext &C,
     clang::DeclContext *DC,
     clang::Expr *RefRSStruct,
     clang::SourceLocation StartLoc,
     clang::SourceLocation Loc);
 
-static clang::Stmt *ClearArrayRSObject(
+clang::Stmt *ClearArrayRSObject(
     clang::ASTContext &C,
     clang::DeclContext *DC,
     clang::Expr *RefRSArr,
@@ -622,7 +573,7 @@ static clang::Stmt *ClearArrayRSObject(
   return DestructorLoop;
 }
 
-static unsigned CountRSObjectTypes(const clang::Type *T) {
+unsigned CountRSObjectTypes(const clang::Type *T) {
   slangAssert(T);
   unsigned RSObjectCount = 0;
 
@@ -672,7 +623,7 @@ static unsigned CountRSObjectTypes(const clang::Type *T) {
   return RSObjectCount;
 }
 
-static clang::Stmt *ClearStructRSObject(
+clang::Stmt *ClearStructRSObject(
     clang::ASTContext &C,
     clang::DeclContext *DC,
     clang::Expr *RefRSStruct,
@@ -795,11 +746,11 @@ static clang::Stmt *ClearStructRSObject(
   return CS;
 }
 
-static clang::Stmt *CreateSingleRSSetObject(clang::ASTContext &C,
-                                            clang::Expr *DstExpr,
-                                            clang::Expr *SrcExpr,
-                                            clang::SourceLocation StartLoc,
-                                            clang::SourceLocation Loc) {
+clang::Stmt *CreateSingleRSSetObject(clang::ASTContext &C,
+                                     clang::Expr *DstExpr,
+                                     clang::Expr *SrcExpr,
+                                     clang::SourceLocation StartLoc,
+                                     clang::SourceLocation Loc) {
   const clang::Type *T = DstExpr->getType().getTypePtr();
   clang::FunctionDecl *SetObjectFD = RSObjectRefCount::GetRSSetObjectFD(T);
   slangAssert((SetObjectFD != nullptr) &&
@@ -849,11 +800,11 @@ static clang::Stmt *CreateSingleRSSetObject(clang::ASTContext &C,
   return RSSetObjectCall;
 }
 
-static clang::Stmt *CreateStructRSSetObject(clang::ASTContext &C,
-                                            clang::Expr *LHS,
-                                            clang::Expr *RHS,
-                                            clang::SourceLocation StartLoc,
-                                            clang::SourceLocation Loc);
+clang::Stmt *CreateStructRSSetObject(clang::ASTContext &C,
+                                     clang::Expr *LHS,
+                                     clang::Expr *RHS,
+                                     clang::SourceLocation StartLoc,
+                                     clang::SourceLocation Loc);
 
 /*static clang::Stmt *CreateArrayRSSetObject(clang::ASTContext &C,
                                            clang::Expr *DstArr,
@@ -1012,11 +963,11 @@ static clang::Stmt *CreateStructRSSetObject(clang::ASTContext &C,
   return CS;
 } */
 
-static clang::Stmt *CreateStructRSSetObject(clang::ASTContext &C,
-                                            clang::Expr *LHS,
-                                            clang::Expr *RHS,
-                                            clang::SourceLocation StartLoc,
-                                            clang::SourceLocation Loc) {
+clang::Stmt *CreateStructRSSetObject(clang::ASTContext &C,
+                                     clang::Expr *LHS,
+                                     clang::Expr *RHS,
+                                     clang::SourceLocation StartLoc,
+                                     clang::SourceLocation Loc) {
   clang::QualType QT = LHS->getType();
   const clang::Type *T = QT.getTypePtr();
   slangAssert(T->isStructureType());
@@ -1128,6 +1079,38 @@ static clang::Stmt *CreateStructRSSetObject(clang::ASTContext &C,
 
 }  // namespace
 
+void RSObjectRefCount::Scope::InsertStmt(const clang::ASTContext &C,
+                                         clang::Stmt *NewStmt) {
+  std::vector<clang::Stmt*> newBody;
+  for (clang::Stmt* S1 : mCS->body()) {
+    if (S1 == mCurrent) {
+      newBody.push_back(NewStmt);
+    }
+    newBody.push_back(S1);
+  }
+  mCS->setStmts(C, newBody.data(), newBody.size());
+}
+
+void RSObjectRefCount::Scope::ReplaceStmt(const clang::ASTContext &C,
+                                          clang::Stmt *NewStmt) {
+  std::vector<clang::Stmt*> newBody;
+  for (clang::Stmt* S1 : mCS->body()) {
+    if (S1 == mCurrent) {
+      newBody.push_back(NewStmt);
+    } else {
+      newBody.push_back(S1);
+    }
+  }
+  mCS->setStmts(C, newBody.data(), newBody.size());
+}
+
+void RSObjectRefCount::Scope::ReplaceExpr(const clang::ASTContext& C,
+                                          clang::Expr* OldExpr,
+                                          clang::Expr* NewExpr) {
+  RSASTReplace R(C);
+  R.ReplaceStmt(mCurrent, OldExpr, NewExpr);
+}
+
 void RSObjectRefCount::Scope::ReplaceRSObjectAssignment(
     clang::BinaryOperator *AS) {
 
@@ -1256,26 +1239,68 @@ void RSObjectRefCount::Scope::AppendRSObjectInit(
 }
 
 void RSObjectRefCount::Scope::InsertLocalVarDestructors() {
-  for (std::list<clang::VarDecl*>::const_iterator I = mRSO.begin(),
-          E = mRSO.end();
-        I != E;
-        I++) {
-    clang::VarDecl *VD = *I;
-    clang::Stmt *RSClearObjectCall = ClearRSObject(VD, VD->getDeclContext());
-    if (RSClearObjectCall) {
-      clang::ASTContext &C = (*mRSO.begin())->getASTContext();
-      // Mark VD as used.  It might be unused, except for the destructor.
-      // 'markUsed' has side-effects that are caused only if VD is not already
-      // used.  Hence no need for an extra check here.
-      VD->markUsed(C);
-      DestructorVisitor DV(VD->getDeclContext(),
-                           mCS,
-                           RSClearObjectCall,
-                           VD->getSourceRange().getBegin());
-      DV.Visit(mCS);
-      DV.InsertDestructors();
+  if (mRSO.empty()) {
+    return;
+  }
+
+  clang::DeclContext* DC = mRSO.front()->getDeclContext();
+  clang::ASTContext& C = DC->getParentASTContext();
+  clang::SourceManager& SM = C.getSourceManager();
+
+  const auto& OccursBefore = [&SM] (clang::SourceLocation L1, clang::SourceLocation L2)->bool {
+    return SM.isBeforeInTranslationUnit(L1, L2);
+  };
+  typedef std::map<clang::SourceLocation, clang::Stmt*, decltype(OccursBefore)> DMap;
+
+  DMap dtors(OccursBefore);
+
+  // Create rsClearObject calls. Note the DMap entries are sorted by the SourceLocation.
+  for (clang::VarDecl* VD : mRSO) {
+    clang::SourceLocation Loc = VD->getSourceRange().getBegin();
+    clang::Stmt* RSClearObjectCall = ClearRSObject(VD, DC);
+    dtors.insert(std::make_pair(Loc, RSClearObjectCall));
+  }
+
+  DestructorVisitor Visitor;
+  Visitor.Visit(mCS);
+
+  // Replace each exiting statement with a block that contains the original statement
+  // and added rsClearObject() calls before it.
+  for (clang::Stmt* S : Visitor.getExitingStmts()) {
+
+    const clang::SourceLocation currentLoc = S->getLocStart();
+
+    DMap::iterator firstDtorIter = dtors.begin();
+    DMap::iterator currentDtorIter = firstDtorIter;
+    DMap::iterator lastDtorIter = dtors.end();
+
+    while (currentDtorIter != lastDtorIter &&
+           OccursBefore(currentDtorIter->first, currentLoc)) {
+      currentDtorIter++;
     }
+
+    if (currentDtorIter == firstDtorIter) {
+      continue;
+    }
+
+    std::list<clang::Stmt*> Stmts;
+
+    // Insert rsClearObject() calls for all rsObjects declared before the current statement
+    for(DMap::iterator it = firstDtorIter; it != currentDtorIter; it++) {
+      Stmts.push_back(it->second);
+    }
+    Stmts.push_back(S);
+
+    RSASTReplace R(C);
+    clang::CompoundStmt* CS = BuildCompoundStmt(C, Stmts, S->getLocEnd());
+    R.ReplaceStmt(mCS, S, CS);
   }
+
+  std::list<clang::Stmt*> Stmts;
+  for(auto LocCallPair : dtors) {
+    Stmts.push_back(LocCallPair.second);
+  }
+  AppendAfterStmt(C, mCS, nullptr, Stmts);
 }
 
 clang::Stmt *RSObjectRefCount::Scope::ClearRSObject(
@@ -1423,12 +1448,13 @@ clang::Expr *RSObjectRefCount::CreateZeroInitializerForRSSpecificType(
                                          Loc);
 
       unsigned N = 0;
-      if (DT == DataTypeRSMatrix2x2)
+      if (DT == DataTypeRSMatrix2x2) {
         N = 2;
-      else if (DT == DataTypeRSMatrix3x3)
+      } else if (DT == DataTypeRSMatrix3x3) {
         N = 3;
-      else if (DT == DataTypeRSMatrix4x4)
+      } else if (DT == DataTypeRSMatrix4x4) {
         N = 4;
+      }
       unsigned N_2 = N * N;
 
       // Assume we are going to be allocating 16 elements, since 4x4 is max.
@@ -1473,6 +1499,8 @@ clang::Expr *RSObjectRefCount::CreateZeroInitializerForRSSpecificType(
 }
 
 void RSObjectRefCount::VisitDeclStmt(clang::DeclStmt *DS) {
+  VisitStmt(DS);
+  getCurrentScope()->setCurrentStmt(DS);
   for (clang::DeclStmt::decl_iterator I = DS->decl_begin(), E = DS->decl_end();
        I != E;
        I++) {
@@ -1494,23 +1522,109 @@ void RSObjectRefCount::VisitDeclStmt(clang::DeclStmt *DS) {
   }
 }
 
+void RSObjectRefCount::VisitCallExpr(clang::CallExpr* CE) {
+  clang::QualType RetTy;
+  const clang::FunctionDecl* FD = CE->getDirectCallee();
+
+  if (FD) {
+    // Direct calls
+
+    RetTy = FD->getReturnType();
+  } else {
+    // Indirect calls through function pointers
+
+    const clang::Expr* Callee = CE->getCallee();
+    const clang::Type* CalleeType = Callee->getType().getTypePtr();
+    const clang::PointerType* PtrType = CalleeType->getAs<clang::PointerType>();
+
+    if (!PtrType) {
+      return;
+    }
+
+    const clang::Type* PointeeType = PtrType->getPointeeType().getTypePtr();
+    const clang::FunctionType* FuncType = PointeeType->getAs<clang::FunctionType>();
+
+    if (!FuncType) {
+      return;
+    }
+
+    RetTy = FuncType->getReturnType();
+  }
+
+  if (CountRSObjectTypes(RetTy.getTypePtr())==0) {
+    return;
+  }
+
+  clang::SourceLocation Loc = CE->getSourceRange().getBegin();
+  std::stringstream ss;
+  ss << ".rs.tmp" << getNextID();
+  llvm::StringRef VarName(ss.str());
+
+  clang::VarDecl* TempVarDecl = clang::VarDecl::Create(
+      mCtx,                                  // AST context
+      GetDeclContext(),                      // Decl context
+      Loc,                                   // Start location
+      Loc,                                   // Id location
+      &mCtx.Idents.get(VarName),             // Id
+      RetTy,                                 // Type
+      mCtx.getTrivialTypeSourceInfo(RetTy),  // Type info
+      clang::SC_None                         // Storage class
+  );
+  TempVarDecl->setInit(CE);
+  clang::Decl* Decls[] = { TempVarDecl };
+  const clang::DeclGroupRef DGR = clang::DeclGroupRef::Create(
+      mCtx, Decls, sizeof(Decls) / sizeof(*Decls));
+  clang::DeclStmt* DS = new (mCtx) clang::DeclStmt(DGR, Loc, Loc);
+
+  getCurrentScope()->InsertStmt(mCtx, DS);
+
+  clang::DeclRefExpr* DRE = clang::DeclRefExpr::Create(
+      mCtx,                                  // AST context
+      clang::NestedNameSpecifierLoc(),       // QualifierLoc
+      Loc,                                   // TemplateKWLoc
+      TempVarDecl,
+      false,                                 // RefersToEnclosingVariableOrCapture
+      Loc,                                   // NameLoc
+      RetTy,
+      clang::VK_LValue
+  );
+  clang::Expr* CastExpr = clang::ImplicitCastExpr::Create(
+      mCtx,
+      RetTy,
+      clang::CK_LValueToRValue,
+      DRE,
+      nullptr,
+      clang::VK_RValue
+  );
+
+  getCurrentScope()->ReplaceExpr(mCtx, CE, CastExpr);
+
+  // Register TempVarDecl for destruction call (rsClearObj).
+  getCurrentScope()->addRSObject(TempVarDecl);
+}
+
 void RSObjectRefCount::VisitCompoundStmt(clang::CompoundStmt *CS) {
+  if (!emptyScope()) {
+    getCurrentScope()->setCurrentStmt(CS);
+  }
+
   if (!CS->body_empty()) {
     // Push a new scope
     Scope *S = new Scope(CS);
-    mScopeStack.push(S);
+    mScopeStack.push_back(S);
 
     VisitStmt(CS);
 
     // Destroy the scope
     slangAssert((getCurrentScope() == S) && "Corrupted scope stack!");
     S->InsertLocalVarDestructors();
-    mScopeStack.pop();
+    mScopeStack.pop_back();
     delete S;
   }
 }
 
 void RSObjectRefCount::VisitBinAssign(clang::BinaryOperator *AS) {
+  getCurrentScope()->setCurrentStmt(AS);
   clang::QualType QT = AS->getType();
 
   if (CountRSObjectTypes(QT.getTypePtr())) {
@@ -1518,7 +1632,44 @@ void RSObjectRefCount::VisitBinAssign(clang::BinaryOperator *AS) {
   }
 }
 
+void RSObjectRefCount::VisitReturnStmt(clang::ReturnStmt *RS) {
+  getCurrentScope()->setCurrentStmt(RS);
+
+  // If there is no local rsObject declared so far, no need to transform the
+  // return statement.
+
+  bool RSObjDeclared = false;
+
+  for (const Scope* S : mScopeStack) {
+    if (S->hasRSObject()) {
+      RSObjDeclared = true;
+      break;
+    }
+  }
+
+  if (!RSObjDeclared) {
+    return;
+  }
+
+  // If the return statement does not return anything, or if it does not return
+  // a rsObject, no need to transform it.
+
+  clang::Expr* RetVal = RS->getRetValue();
+  if (!RetVal || CountRSObjectTypes(RetVal->getType().getTypePtr()) == 0) {
+    return;
+  }
+
+  // Transform the return statement so that it does not potentially return or
+  // reference a rsObject that has been cleared.
+
+  clang::CompoundStmt* NewRS;
+  NewRS = CreateRetStmtWithTempVar(mCtx, GetDeclContext(), RS, getNextID());
+
+  getCurrentScope()->ReplaceStmt(mCtx, NewRS);
+}
+
 void RSObjectRefCount::VisitStmt(clang::Stmt *S) {
+  getCurrentScope()->setCurrentStmt(S);
   for (clang::Stmt::child_iterator I = S->child_begin(), E = S->child_end();
        I != E;
        I++) {
diff --git a/slang_rs_object_ref_count.h b/slang_rs_object_ref_count.h
index d6be50f..5b49c34 100644
--- a/slang_rs_object_ref_count.h
+++ b/slang_rs_object_ref_count.h
@@ -39,23 +39,27 @@ bool HasRSObjectType(const clang::Type *T);
 // local variables of RS object types (rs_font, rs_allocation, ...). This
 // class ensures that appropriate functions (rsSetObject, rsClearObject) are
 // called at proper points in the object's lifetime.
-// 1) Each local object of appropriate type must be zero-initialized (to
-// prevent corruption) during subsequent rsSetObject()/rsClearObject() calls.
+// 1) Each local object of appropriate type must be zero-initialized to
+// prevent corruption during subsequent rsSetObject()/rsClearObject() calls.
 // 2) Assignments using these types must also be converted into the
 // appropriate (possibly a series of) rsSetObject() calls.
-// 3) Finally, each local object must call rsClearObject() when it goes out
-// of scope.
+// 3) Finally, rsClearObject() must be called for each local object when it goes
+// out of scope.
 class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
  private:
   class Scope {
    private:
-    clang::CompoundStmt *mCS;      // Associated compound statement ({ ... })
-    std::list<clang::VarDecl*> mRSO;  // Declared RS objects in this scope
+    clang::CompoundStmt *mCS;         // Associated compound statement ({ ... })
+    clang::Stmt *mCurrent;            // The statement currently being analyzed
+    std::list<clang::VarDecl*> mRSO;  // Declared RS objects in this scope (but
+                                      // not any scopes nested)
 
    public:
     explicit Scope(clang::CompoundStmt *CS) : mCS(CS) {
     }
 
+    bool hasRSObject() const { return !mRSO.empty(); }
+
     inline void addRSObject(clang::VarDecl* VD) {
       mRSO.push_back(VD);
     }
@@ -67,25 +71,57 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
                             DataType DT,
                             clang::Expr *InitExpr);
 
+    // Inserts rsClearObject() calls at the end and at all exiting points of the
+    // current scope. At each statement that exits the current scope -- e.g.,
+    // a return, break, or continue statement in the current or a nested scope
+    // -- rsClearObject() calls are inserted for local variables defined in the
+    // current scope before that point.
+    // Note goto statements are not handled. (See the DestructorVisitor class in
+    // the .cpp file.)
+    // Also note this function is called for every nested scope. As a result, for a
+    // return statement, each rsObject declared in all its (nested) enclosing
+    // scopes would have a rsClearObject() call properly inserted before
+    // the return statement.
     void InsertLocalVarDestructors();
 
+    // Sets the current statement being analyzed
+    void setCurrentStmt(clang::Stmt *S) { mCurrent = S; }
+
+    // Inserts a statement before the current statement
+    void InsertStmt(const clang::ASTContext &C, clang::Stmt *NewStmt);
+
+    // Replaces the current statement with NewStmt;
+    void ReplaceStmt(const clang::ASTContext &C, clang::Stmt *NewStmt);
+
+    // Replaces OldExpr with NewExpr in the current statement
+    void ReplaceExpr(const clang::ASTContext& C, clang::Expr* OldExpr,
+                     clang::Expr* NewExpr);
+
     static clang::Stmt *ClearRSObject(clang::VarDecl *VD,
                                       clang::DeclContext *DC);
   };
 
   clang::ASTContext &mCtx;
-  std::stack<Scope*> mScopeStack;
+  std::deque<Scope*> mScopeStack;  // A deque used as a stack to store scopes, but also
+                                   // accessed through its iterator in read-only mode.
+  clang::DeclContext* mCurrentDC;
   bool RSInitFD;
+  unsigned mTempID;  // A unique id that can be used to distinguish temporary variables
 
   // RSSetObjectFD and RSClearObjectFD holds FunctionDecl of rsSetObject()
   // and rsClearObject() in the current ASTContext.
   static clang::FunctionDecl *RSSetObjectFD[];
   static clang::FunctionDecl *RSClearObjectFD[];
 
+  inline bool emptyScope() const { return mScopeStack.empty(); }
+
   inline Scope *getCurrentScope() {
-    return mScopeStack.top();
+    return mScopeStack.back();
   }
 
+  // Returns the next available unique id for temporary variables
+  unsigned getNextID() { return mTempID++; }
+
   // Initialize RSSetObjectFD and RSClearObjectFD.
   static void GetRSRefCountingFunctions(clang::ASTContext &C);
 
@@ -104,8 +140,7 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
 
  public:
   explicit RSObjectRefCount(clang::ASTContext &C)
-      : mCtx(C),
-        RSInitFD(false) {
+      : mCtx(C), RSInitFD(false), mTempID(0) {
   }
 
   void Init() {
@@ -143,10 +178,15 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
     return GetRSClearObjectFD(RSExportPrimitiveType::GetRSSpecificType(T));
   }
 
+  void SetDeclContext(clang::DeclContext* DC) { mCurrentDC = DC; }
+  clang::DeclContext* GetDeclContext() const { return mCurrentDC; }
+
   void VisitStmt(clang::Stmt *S);
+  void VisitCallExpr(clang::CallExpr *CE);
   void VisitDeclStmt(clang::DeclStmt *DS);
   void VisitCompoundStmt(clang::CompoundStmt *CS);
   void VisitBinAssign(clang::BinaryOperator *AS);
+  void VisitReturnStmt(clang::ReturnStmt *RS);
   // We believe that RS objects are never involved in CompoundAssignOperator.
   // I.e., rs_allocation foo; foo += bar;