Handle struct reference counting.

Bug: 3092382

Change-Id: I215bd8245324ec2b7752a7c40817e3e5cd1c0e00

1 files changed, 32 insertions, 11 deletions

diff --git a/slang_rs_object_ref_count.h b/slang_rs_object_ref_count.h
index 9145e5c..f1f54fc 100644
--- a/slang_rs_object_ref_count.h
+++ b/slang_rs_object_ref_count.h
@@ -22,6 +22,7 @@
 
 #include "clang/AST/StmtVisitor.h"
 
+#include "slang_assert.h"
 #include "slang_rs_export_type.h"
 
 namespace clang {
@@ -38,11 +39,6 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
     clang::CompoundStmt *mCS;      // Associated compound statement ({ ... })
     std::list<clang::VarDecl*> mRSO;  // Declared RS objects in this scope
 
-    // RSSetObjectFD and RSClearObjectFD holds FunctionDecl of rsSetObject()
-    // and rsClearObject() in the current ASTContext.
-    static clang::FunctionDecl *RSSetObjectFD[];
-    static clang::FunctionDecl *RSClearObjectFD[];
-
    public:
     explicit Scope(clang::CompoundStmt *CS) : mCS(CS) {
       return;
@@ -53,9 +49,6 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
       return;
     }
 
-    // Initialize RSSetObjectFD and RSClearObjectFD.
-    static void GetRSRefCountingFunctions(clang::ASTContext &C);
-
     void ReplaceRSObjectAssignment(clang::BinaryOperator *AS);
 
     void AppendRSObjectInit(clang::VarDecl *VD,
@@ -71,15 +64,23 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
   std::stack<Scope*> mScopeStack;
   bool RSInitFD;
 
+  // RSSetObjectFD and RSClearObjectFD holds FunctionDecl of rsSetObject()
+  // and rsClearObject() in the current ASTContext.
+  static clang::FunctionDecl *RSSetObjectFD[];
+  static clang::FunctionDecl *RSClearObjectFD[];
+
   inline Scope *getCurrentScope() {
     return mScopeStack.top();
   }
 
+  // Initialize RSSetObjectFD and RSClearObjectFD.
+  static void GetRSRefCountingFunctions(clang::ASTContext &C);
+
   // TODO(srhines): Composite types and arrays based on RS object types need
   // to be handled for both zero-initialization + clearing.
 
-  // Return false if the type of variable declared in VD is not an RS object
-  // type.
+  // Return false if the type of variable declared in VD does not contain
+  // an RS object type.
   static bool InitializeRSObject(clang::VarDecl *VD,
                                  RSExportPrimitiveType::DataType *DT,
                                  clang::Expr **InitExpr);
@@ -99,12 +100,32 @@ class RSObjectRefCount : public clang::StmtVisitor<RSObjectRefCount> {
 
   void Init(clang::ASTContext &C) {
     if (!RSInitFD) {
-      Scope::GetRSRefCountingFunctions(C);
+      GetRSRefCountingFunctions(C);
       RSInitFD = true;
     }
     return;
   }
 
+  static clang::FunctionDecl *GetRSSetObjectFD(
+      RSExportPrimitiveType::DataType DT) {
+    slangAssert(RSExportPrimitiveType::IsRSObjectType(DT));
+    return RSSetObjectFD[(DT - RSExportPrimitiveType::FirstRSObjectType)];
+  }
+
+  static clang::FunctionDecl *GetRSSetObjectFD(const clang::Type *T) {
+    return GetRSSetObjectFD(RSExportPrimitiveType::GetRSSpecificType(T));
+  }
+
+  static clang::FunctionDecl *GetRSClearObjectFD(
+      RSExportPrimitiveType::DataType DT) {
+    slangAssert(RSExportPrimitiveType::IsRSObjectType(DT));
+    return RSClearObjectFD[(DT - RSExportPrimitiveType::FirstRSObjectType)];
+  }
+
+  static clang::FunctionDecl *GetRSClearObjectFD(const clang::Type *T) {
+    return GetRSClearObjectFD(RSExportPrimitiveType::GetRSSpecificType(T));
+  }
+
   void VisitStmt(clang::Stmt *S);
   void VisitDeclStmt(clang::DeclStmt *DS);
   void VisitCompoundStmt(clang::CompoundStmt *CS);