diff --git a/clang/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp b/clang/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp
index 66143f78932c32..428778e6cfaa65 100644
--- a/clang/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp
+++ b/clang/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp
@@ -38,8 +38,6 @@ using namespace clang;
 using namespace ento;
 using namespace taint;
 
-using llvm::ImmutableSet;
-
 namespace {
 
 class GenericTaintChecker;
@@ -436,9 +434,7 @@ template <> struct ScalarEnumerationTraits<TaintConfiguration::VariadicType> {
 /// to the call post-visit. The values are signed integers, which are either
 /// ReturnValueIndex, or indexes of the pointer/reference argument, which
 /// points to data, which should be tainted on return.
-REGISTER_MAP_WITH_PROGRAMSTATE(TaintArgsOnPostVisit, const LocationContext *,
-                               ImmutableSet<ArgIdxTy>)
-REGISTER_SET_FACTORY_WITH_PROGRAMSTATE(ArgIdxFactory, ArgIdxTy)
+REGISTER_SET_WITH_PROGRAMSTATE(TaintArgsOnPostVisit, ArgIdxTy)
 
 void GenericTaintRuleParser::validateArgVector(const std::string &Option,
                                                const ArgVecTy &Args) const {
@@ -689,26 +685,22 @@ void GenericTaintChecker::checkPostCall(const CallEvent &Call,
   // Set the marked values as tainted. The return value only accessible from
   // checkPostStmt.
   ProgramStateRef State = C.getState();
-  const StackFrameContext *CurrentFrame = C.getStackFrame();
 
   // 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.
-  TaintArgsOnPostVisitTy TaintArgsMap = State->get<TaintArgsOnPostVisit>();
-
-  const ImmutableSet<ArgIdxTy> *TaintArgs = TaintArgsMap.lookup(CurrentFrame);
-  if (!TaintArgs)
+  TaintArgsOnPostVisitTy TaintArgs = State->get<TaintArgsOnPostVisit>();
+  if (TaintArgs.isEmpty())
     return;
-  assert(!TaintArgs->isEmpty());
 
   LLVM_DEBUG(for (ArgIdxTy I
-                  : *TaintArgs) {
+                  : TaintArgs) {
     llvm::dbgs() << "PostCall<";
     Call.dump(llvm::dbgs());
     llvm::dbgs() << "> actually wants to taint arg index: " << I << '\n';
   });
 
-  for (ArgIdxTy ArgNum : *TaintArgs) {
+  for (ArgIdxTy ArgNum : TaintArgs) {
     // Special handling for the tainted return value.
     if (ArgNum == ReturnValueIndex) {
       State = addTaint(State, Call.getReturnValue());
@@ -722,7 +714,7 @@ void GenericTaintChecker::checkPostCall(const CallEvent &Call,
   }
 
   // Clear up the taint info from the state.
-  State = State->remove<TaintArgsOnPostVisit>(CurrentFrame);
+  State = State->remove<TaintArgsOnPostVisit>();
   C.addTransition(State);
 }
 
@@ -784,33 +776,28 @@ void GenericTaintRule::process(const GenericTaintChecker &Checker,
   };
 
   /// Propagate taint where it is necessary.
-  auto &F = State->getStateManager().get_context<ArgIdxFactory>();
-  ImmutableSet<ArgIdxTy> Result = F.getEmptySet();
   ForEachCallArg(
-      [this, WouldEscape, &Call, &Result, &F](ArgIdxTy I, const Expr *E,
-                                              SVal V) {
+      [this, &State, WouldEscape, &Call](ArgIdxTy I, const Expr *E, SVal V) {
         if (PropDstArgs.contains(I)) {
           LLVM_DEBUG(llvm::dbgs() << "PreCall<"; Call.dump(llvm::dbgs());
                      llvm::dbgs()
                      << "> prepares tainting arg index: " << I << '\n';);
-          Result = F.add(Result, I);
+          State = State->add<TaintArgsOnPostVisit>(I);
         }
 
         // TODO: We should traverse all reachable memory regions via the
         // escaping parameter. Instead of doing that we simply mark only the
         // referred memory region as tainted.
         if (WouldEscape(V, E->getType())) {
-          LLVM_DEBUG(if (!Result.contains(I)) {
+          LLVM_DEBUG(if (!State->contains<TaintArgsOnPostVisit>(I)) {
             llvm::dbgs() << "PreCall<";
             Call.dump(llvm::dbgs());
             llvm::dbgs() << "> prepares tainting arg index: " << I << '\n';
           });
-          Result = F.add(Result, I);
+          State = State->add<TaintArgsOnPostVisit>(I);
         }
       });
 
-  if (!Result.isEmpty())
-    State = State->set<TaintArgsOnPostVisit>(C.getStackFrame(), Result);
   C.addTransition(State);
 }
 
@@ -901,11 +888,7 @@ void GenericTaintChecker::taintUnsafeSocketProtocol(const CallEvent &Call,
   if (SafeProtocol)
     return;
 
-  ProgramStateRef State = C.getState();
-  auto &F = State->getStateManager().get_context<ArgIdxFactory>();
-  ImmutableSet<ArgIdxTy> Result = F.add(F.getEmptySet(), ReturnValueIndex);
-  State = State->set<TaintArgsOnPostVisit>(C.getStackFrame(), Result);
-  C.addTransition(State);
+  C.addTransition(C.getState()->add<TaintArgsOnPostVisit>(ReturnValueIndex));
 }
 
 /// Checker registration
diff --git a/clang/test/Analysis/taint-checker-callback-order-has-definition.c b/clang/test/Analysis/taint-checker-callback-order-has-definition.c
index a070077004fa91..295f95c2bf4529 100644
--- a/clang/test/Analysis/taint-checker-callback-order-has-definition.c
+++ b/clang/test/Analysis/taint-checker-callback-order-has-definition.c
@@ -3,6 +3,9 @@
 // RUN:   -mllvm -debug-only=taint-checker \
 // RUN:   2>&1 | FileCheck %s
 
+// FIXME: We should not crash.
+// XFAIL: *
+
 struct _IO_FILE;
 typedef struct _IO_FILE FILE;
 FILE *fopen(const char *fname, const char *mode);
@@ -28,8 +31,12 @@ void top(const char *fname, char *buf) {
   // CHECK-NEXT: PreCall<fgets(buf, 42, fp)> prepares tainting arg index: 1
   // CHECK-NEXT: PreCall<fgets(buf, 42, fp)> prepares tainting arg index: 2
 
-  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: -1
-  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: 0
-  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: 1
-  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: 2
+  // FIXME: We should propagate taint from PreCall<fgets> -> PostCall<fgets>.
+  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: -1
+  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: 0
+  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: 1
+  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: 2
+
+  // FIXME: We should not crash.
+  // CHECK: PLEASE submit a bug report
 }