Convert TraverseState::new_goals and TraverseState::removed_goals…

… into `std::vector` and push/pop instead of insert/erase.

These containers were never used to check set membership during the loop, it was practically being used to later construct a set after the algorithm, of which in previous state of the code it made sense(because it was being copied all over the place), but while traversing the problem space through DFS it doesn't make sense just because pushing and poping on a `std::vector` is way more cheaper than insertion and deletion within a `std::set`, as we only want to remember the increment of the state change between DFSs.

PiperOrigin-RevId: 707443325

pytype/typegraph/solver.cc

@@ -8,6 +8,7 @@
 #include <memory>
 #include <optional>
 #include <set>
+#include <stack>
 #include <string>
 #include <tuple>
 #include <unordered_map>
@@ -35,17 +36,82 @@ struct RemoveResult {
 struct TraverseState {
   GoalSet goals_to_remove;
   GoalSet seen_goals;
-  GoalSet removed_goals;
-  GoalSet new_goals;
-  TraverseState() {}
-  TraverseState(GoalSet goals_to_remove, GoalSet seen_goals,
-                GoalSet removed_goals, GoalSet new_goals)
-      : goals_to_remove(std::move(goals_to_remove)),
-        seen_goals(std::move(seen_goals)),
-        removed_goals(std::move(removed_goals)),
-        new_goals(std::move(new_goals)) {}
+  std::vector<const Binding*> removed_goals;
+  std::vector<const Binding*> new_goals;
+};
+
+enum ActionType {
+  TRAVERSE,
+  INSERT_GOALS_TO_REMOVE,
+  ERASE_GOALS_TO_REMOVE,
+  ERASE_SEEN_GOALS,
+  ERASE_NEW_GOALS,
+  ERASE_REMOVED_GOALS,
 };
 
+// We're maintaining a state machine and actions to be able to do a DFS
+// effectively. Rather than having to copy the states that are needed (which
+// are four std::sets) whenever we need to traverse a new node, we keep track
+// of the increment and the decrement between the previous node, and restore
+// to the state of which were before after a node traversal, which is implement
+// through "actions".
+struct Action {
+  ActionType action_type;
+  // Goal either to delete are added to the corresponding set.
+  const Binding* goal;
+  // Not using this for action ERASE_GOALS_TO_REMOVE, as we are requesting
+  // for removal before the insertion has happened.
+  GoalSet::iterator erase_it;
+};
+
+static void traverse(const CFGNode* position,
+                     std::vector<RemoveResult>& results,
+                     std::stack<Action, std::vector<Action>>& actions,
+                     TraverseState& state) {
+  if (state.goals_to_remove.empty()) {
+    results.emplace_back(
+        GoalSet(state.removed_goals.begin(), state.removed_goals.end()),
+        GoalSet(state.new_goals.begin(), state.new_goals.end()));
+    return;
+  }
+
+  const Binding* goal = *state.goals_to_remove.begin();
+  state.goals_to_remove.erase(state.goals_to_remove.begin());
+  actions.emplace(INSERT_GOALS_TO_REMOVE, goal);
+
+  if (state.seen_goals.count(goal)) {
+    // Only process a goal once, to prevent infinite loops.
+    actions.emplace(TRAVERSE, nullptr);
+    return;
+  }
+  auto [it, _] = state.seen_goals.insert(goal);
+  actions.emplace(ERASE_SEEN_GOALS, nullptr, it);
+
+  const auto* origin = goal->FindOrigin(position);
+  if (!origin) {
+    state.new_goals.push_back(goal);
+    actions.emplace(ERASE_NEW_GOALS, nullptr);
+    actions.emplace(TRAVERSE, nullptr);
+    return;
+  }
+
+  state.removed_goals.push_back(goal);
+  actions.emplace(ERASE_REMOVED_GOALS, nullptr, it);
+  for (const auto& source_set : origin->source_sets) {
+    for (const Binding* next_goal : source_set) {
+      if (!state.goals_to_remove.count(next_goal)) {
+        actions.emplace(ERASE_GOALS_TO_REMOVE, next_goal);
+      }
+    }
+    actions.emplace(TRAVERSE, nullptr);
+    for (const Binding* next_goal : source_set) {
+      if (!state.goals_to_remove.count(next_goal)) {
+        actions.emplace(INSERT_GOALS_TO_REMOVE, next_goal);
+      }
+    }
+  }
+}
+
 // Remove all goals that can be fulfilled at the current CFG node.
 // Generates all possible sets of new goals obtained by replacing a goal that
 // originates at the current node with one of its source sets, iteratively,
@@ -65,37 +131,31 @@ static std::vector<RemoveResult> remove_finished_goals(const CFGNode* pos,
                       state.goals_to_remove.end(),
                       std::inserter(state.new_goals, state.new_goals.begin()),
                       pointer_less<Binding>());
-  std::deque<TraverseState> queue;
-  queue.push_back(std::move(state));
+  std::stack<Action, std::vector<Action>> actions;
+  actions.emplace(TRAVERSE, nullptr);
   std::vector<RemoveResult> results;
-  while (!queue.empty()) {
-    state = std::move(queue.front());
-    queue.pop_front();
-    if (state.goals_to_remove.empty()) {
-      results.push_back(RemoveResult(state.removed_goals, state.new_goals));
-      continue;
-    }
-    const auto* goal = *state.goals_to_remove.begin();
-    state.goals_to_remove.erase(state.goals_to_remove.begin());
-    if (state.seen_goals.count(goal)) {
-      // Only process a goal once, to prevent infinite loops.
-      queue.emplace_back(std::move(state));
-      continue;
-    }
-    state.seen_goals.insert(goal);
-    const auto* origin = goal->FindOrigin(pos);
-    if (!origin) {
-      state.new_goals.insert(goal);
-      queue.emplace_back(std::move(state));
-      continue;
-    }
-    state.removed_goals.insert(goal);
-    for (const auto& source_set : origin->source_sets) {
-      GoalSet next_goals_to_remove(state.goals_to_remove);
-      next_goals_to_remove.insert(source_set.begin(), source_set.end());
-      queue.push_back(TraverseState(std::move(next_goals_to_remove),
-                                    state.seen_goals, state.removed_goals,
-                                    state.new_goals));
+  while (!actions.empty()) {
+    Action action = actions.top();
+    actions.pop();
+    switch (action.action_type) {
+      case TRAVERSE:
+        traverse(pos, results, actions, state);
+        break;
+      case INSERT_GOALS_TO_REMOVE:
+        state.goals_to_remove.insert(action.goal);
+        break;
+      case ERASE_GOALS_TO_REMOVE:
+        state.goals_to_remove.erase(action.goal);
+        break;
+      case ERASE_SEEN_GOALS:
+        state.seen_goals.erase(action.erase_it);
+        break;
+      case ERASE_NEW_GOALS:
+        state.new_goals.pop_back();
+        break;
+      case ERASE_REMOVED_GOALS:
+        state.removed_goals.pop_back();
+        break;
     }
   }
   return results;