1 files changed, 2 insertions, 8 deletions
diff --git a/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp b/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp
index 455fb64e4d..a9e02173c3 100644
--- a/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp
+++ b/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp
@@ -192,13 +192,7 @@ const char GenericTaintChecker::MsgTaintedBufferSize[] =
 /// to the call post-visit. The values are unsigned integers, which are either
 /// ReturnValueIndex, or indexes of the pointer/reference argument, which
 /// points to data, which should be tainted on return.
-namespace { struct TaintArgsOnPostVisit{}; }
-namespace clang { namespace ento {
-template<> struct ProgramStateTrait<TaintArgsOnPostVisit>
-    :  public ProgramStatePartialTrait<llvm::ImmutableSet<unsigned> > {
-  static void *GDMIndex() { return GenericTaintChecker::getTag(); }
-};
-}}
+REGISTER_SET_WITH_PROGRAMSTATE(TaintArgsOnPostVisit, unsigned)
 
 GenericTaintChecker::TaintPropagationRule
 GenericTaintChecker::TaintPropagationRule::getTaintPropagationRule(
@@ -337,7 +331,7 @@ bool GenericTaintChecker::propagateFromPre(const CallExpr *CE,
   // Depending on what was tainted at pre-visit, we determined a set of
   // arguments which should be tainted after the function returns. These are
   // stored in the state as TaintArgsOnPostVisit set.
-  llvm::ImmutableSet<unsigned> TaintArgs = State->get<TaintArgsOnPostVisit>();
+  TaintArgsOnPostVisitTy TaintArgs = State->get<TaintArgsOnPostVisit>();
   if (TaintArgs.isEmpty())
     return false;