1 files changed, 7 insertions, 14 deletions
diff --git a/lib/StaticAnalyzer/Checkers/MacOSKeychainAPIChecker.cpp b/lib/StaticAnalyzer/Checkers/MacOSKeychainAPIChecker.cpp
index ff0683ec3c..76f20b6e2e 100644
--- a/lib/StaticAnalyzer/Checkers/MacOSKeychainAPIChecker.cpp
+++ b/lib/StaticAnalyzer/Checkers/MacOSKeychainAPIChecker.cpp
@@ -158,16 +158,9 @@ private:
 /// ProgramState traits to store the currently allocated (and not yet freed)
 /// symbols. This is a map from the allocated content symbol to the
 /// corresponding AllocationState.
-typedef llvm::ImmutableMap<SymbolRef,
-                       MacOSKeychainAPIChecker::AllocationState> AllocatedSetTy;
-
-namespace { struct AllocatedData {}; }
-namespace clang { namespace ento {
-template<> struct ProgramStateTrait<AllocatedData>
-    :  public ProgramStatePartialTrait<AllocatedSetTy > {
-  static void *GDMIndex() { static int index = 0; return &index; }
-};
-}}
+REGISTER_MAP_WITH_PROGRAMSTATE(AllocatedData,
+                               SymbolRef,
+                               MacOSKeychainAPIChecker::AllocationState)
 
 static bool isEnclosingFunctionParam(const Expr *E) {
   E = E->IgnoreParenCasts();
@@ -571,13 +564,13 @@ BugReport *MacOSKeychainAPIChecker::
 void MacOSKeychainAPIChecker::checkDeadSymbols(SymbolReaper &SR,
                                                CheckerContext &C) const {
   ProgramStateRef State = C.getState();
-  AllocatedSetTy ASet = State->get<AllocatedData>();
+  AllocatedDataTy ASet = State->get<AllocatedData>();
   if (ASet.isEmpty())
     return;
 
   bool Changed = false;
   AllocationPairVec Errors;
-  for (AllocatedSetTy::iterator I = ASet.begin(), E = ASet.end(); I != E; ++I) {
+  for (AllocatedDataTy::iterator I = ASet.begin(), E = ASet.end(); I != E; ++I) {
     if (SR.isLive(I->first))
       continue;
 
@@ -619,7 +612,7 @@ void MacOSKeychainAPIChecker::checkEndPath(CheckerContext &C) const {
   if (C.getLocationContext()->getParent() != 0)
     return;
 
-  AllocatedSetTy AS = state->get<AllocatedData>();
+  AllocatedDataTy AS = state->get<AllocatedData>();
   if (AS.isEmpty())
     return;
 
@@ -627,7 +620,7 @@ void MacOSKeychainAPIChecker::checkEndPath(CheckerContext &C) const {
   // found here, so report it.
   bool Changed = false;
   AllocationPairVec Errors;
-  for (AllocatedSetTy::iterator I = AS.begin(), E = AS.end(); I != E; ++I ) {
+  for (AllocatedDataTy::iterator I = AS.begin(), E = AS.end(); I != E; ++I ) {
     Changed = true;
     state = state->remove<AllocatedData>(I->first);
     // If the allocated symbol is null or if error code was returned at