Index: llvm/lib/Transforms/IPO/AttributorAttributes.cpp
===================================================================
--- llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -578,9 +578,58 @@
   }
 }
 
-/// Use the must-be-executed-context around \p I to add information into \p S.
-/// The AAType class is required to have `followUseInMBEC` method with the
-/// following signature and behaviour:
+/// Helper function for update in followUsesInMBEC
+template <class AAType, typename StateType = typename AAType::StateType>
+static void
+updateBrInst(AAType &AA, Attributor &A, MustBeExecutedContextExplorer &Explorer,
+             const Instruction *CtxI, const BranchInst *Br,
+             SetVector<const Use *> &Uses, StateType &State,
+             DenseMap<const Instruction *, StateType> &StateMap,
+             SetVector<const Instruction *> &AddInsts) {
+  State.indicateOptimisticFixpoint();
+  for (const BasicBlock *BB : Br->successors()) {
+    const Instruction *DependInst = &BB->front();
+    auto Iter = StateMap.find(DependInst);
+    if (Iter == StateMap.end()) {
+      AddInsts.insert(DependInst);
+      State &= StateType();
+    } else {
+      State &= Iter->second;
+    }
+  }
+  followUsesInContext(AA, A, Explorer, CtxI, Uses, State);
+}
+
+/// Helper function for update in followUsesInMBEC
+template <class AAType, typename StateType = typename AAType::StateType>
+static void
+updateInst(AAType &AA, Attributor &A, MustBeExecutedContextExplorer &Explorer,
+           const Instruction *CtxI, SetVector<const Use *> &Uses,
+           StateType &State, DenseMap<const Instruction *, StateType> &StateMap,
+           SetVector<const Instruction *> &AddInsts) {
+  State = StateType();
+  followUsesInContext<AAType>(AA, A, Explorer, CtxI, Uses, State);
+  SmallVector<const Instruction *, 4> BrInsts;
+  auto Pred = [&](const Instruction *I) {
+    if (const BranchInst *Br = dyn_cast<BranchInst>(I))
+      if (Br->isConditional())
+        BrInsts.push_back(I);
+    return true;
+  };
+  Explorer.checkForAllContext(CtxI, Pred);
+  for (const Instruction *I : BrInsts) {
+    auto Iter = StateMap.find(I);
+    if (Iter == StateMap.end()) {
+      AddInsts.insert(I);
+    } else {
+      State += Iter->second;
+    }
+  }
+}
+
+/// Use the must-be-executed-context around \p I to add information into \p
+/// S. The AAType class is required to have `followUseInMBEC` method with
+/// the following signature and behaviour:
 ///
 /// bool followUseInMBEC(Attributor &A, const Use *U, const Instruction *I)
 /// U - Underlying use.
@@ -590,84 +639,53 @@
 template <class AAType, typename StateType = typename AAType::StateType>
 static void followUsesInMBEC(AAType &AA, Attributor &A, StateType &S,
                              Instruction &CtxI) {
+  // Map from context instruction to state
+  DenseMap<const Instruction *, StateType> StateMap;
 
   // Container for (transitive) uses of the associated value.
   SetVector<const Use *> Uses;
-  for (const Use &U : AA.getIRPosition().getAssociatedValue().uses())
+
+  // Register context instructions
+  StateMap.insert(std::make_pair(&CtxI, StateType()));
+
+  for (const Use &U : AA.getIRPosition().getAssociatedValue().uses()) {
     Uses.insert(&U);
+    if (const Instruction *UserI = dyn_cast<Instruction>(U.getUser())) {
+      StateMap.insert(std::make_pair(UserI, StateType()));
+    }
+  }
 
   MustBeExecutedContextExplorer &Explorer =
       A.getInfoCache().getMustBeExecutedContextExplorer();
 
-  followUsesInContext<AAType>(AA, A, Explorer, &CtxI, Uses, S);
-
-  if (S.isAtFixpoint())
-    return;
-
-  SmallVector<const BranchInst *, 4> BrInsts;
-  auto Pred = [&](const Instruction *I) {
-    if (const BranchInst *Br = dyn_cast<BranchInst>(I))
-      if (Br->isConditional())
-        BrInsts.push_back(Br);
-    return true;
-  };
-
-  // Here, accumulate conditional branch instructions in the context. We
-  // explore the child paths and collect the known states. The disjunction of
-  // those states can be merged to its own state. Let ParentState_i be a state
-  // to indicate the known information for an i-th branch instruction in the
-  // context. ChildStates are created for its successors respectively.
-  //
-  // ParentS_1 = ChildS_{1, 1} /\ ChildS_{1, 2} /\ ... /\ ChildS_{1, n_1}
-  // ParentS_2 = ChildS_{2, 1} /\ ChildS_{2, 2} /\ ... /\ ChildS_{2, n_2}
-  //      ...
-  // ParentS_m = ChildS_{m, 1} /\ ChildS_{m, 2} /\ ... /\ ChildS_{m, n_m}
-  //
-  // Known State |= ParentS_1 \/ ParentS_2 \/... \/ ParentS_m
-  //
-  // FIXME: Currently, recursive branches are not handled. For example, we
-  // can't deduce that ptr must be dereferenced in below function.
-  //
-  // void f(int a, int c, int *ptr) {
-  //    if(a)
-  //      if (b) {
-  //        *ptr = 0;
-  //      } else {
-  //        *ptr = 1;
-  //      }
-  //    else {
-  //      if (b) {
-  //        *ptr = 0;
-  //      } else {
-  //        *ptr = 1;
-  //      }
-  //    }
-  // }
-
-  Explorer.checkForAllContext(&CtxI, Pred);
-  for (const BranchInst *Br : BrInsts) {
-    StateType ParentState;
-
-    // The known state of the parent state is a conjunction of children's
-    // known states so it is initialized with a best state.
-    ParentState.indicateOptimisticFixpoint();
-
-    for (const BasicBlock *BB : Br->successors()) {
-      StateType ChildState;
-
-      size_t BeforeSize = Uses.size();
-      followUsesInContext(AA, A, Explorer, &BB->front(), Uses, ChildState);
-
-      // Erase uses which only appear in the child.
-      for (auto It = Uses.begin() + BeforeSize; It != Uses.end();)
-        It = Uses.erase(It);
-
-      ParentState &= ChildState;
+  ChangeStatus CS;
+  // TODO: set upper bound of number of iteration as command line option
+  unsigned Iteration = 0;
+  do {
+    Iteration += 1;
+    CS = ChangeStatus::UNCHANGED;
+    SetVector<const Instruction *> AddInsts;
+    for (auto &InstStatePair : StateMap) {
+      const Instruction *I = InstStatePair.first;
+      StateType &State = InstStatePair.second;
+      StateType BeforeState = State;
+      const BranchInst *Br = dyn_cast<const BranchInst>(I);
+      if (Br && Br->isConditional()) {
+        updateBrInst(AA, A, Explorer, I, Br, Uses, State, StateMap, AddInsts);
+      } else {
+        updateInst(AA, A, Explorer, I, Uses, State, StateMap, AddInsts);
+      }
+      if (State != BeforeState) {
+        CS = ChangeStatus::CHANGED;
+      }
     }
+    if (AddInsts.size() > 0)
+      CS = ChangeStatus::CHANGED;
+    for (const Instruction *I : AddInsts)
+      StateMap.insert(std::make_pair(I, StateType()));
+  } while (CS == ChangeStatus::CHANGED);
 
-    // Use only known state.
-    S += ParentState;
-  }
+  S += StateMap.lookup(&CtxI);
 }
 
 /// -----------------------NoUnwind Function Attribute--------------------------