Add a new ActionType::SOURCE_SET for handling the case where we inspect the source sets of the origin. This requires the addition of two new fields to Action as well that are only used with this action type. Hence the unions.

The benefit of handling the source sets with a separate action is that we push fewer actions onto the `actions` stack and that the maximal stack size decreases. The runtime performance benefits of this seems to be rather small (140s -> 138s).

PiperOrigin-RevId: 713244634

Author: Martin Huschenbett

pytype/typegraph/solver.cc

@@ -42,6 +42,7 @@ struct TraverseState {
 
 enum ActionType {
   TRAVERSE,
+  TRAVERSE_ALL_SOURCE_SETS,
   INSERT_GOALS_TO_REMOVE,
   ERASE_GOALS_TO_REMOVE,
   ERASE_SEEN_GOALS,
@@ -57,20 +58,32 @@ enum ActionType {
 // through "actions".
 struct Action {
   ActionType action_type;
-  // Goal either to delete are added to the corresponding set.
-  const Binding* goal;
-  // The iterator is for std::set and this is stable upon deletion and insertion
-  // if it's not directly the element being deleted or inserted. We will only
-  // try to erase the element on the exact node traversal, so we can safely
-  // reuse the iterator that was returned from the insertion.
-  // Not using this for action ERASE_GOALS_TO_REMOVE, as we are requesting
-  // for removal before the insertion has happened.
-  GoalSet::iterator erase_it;
+  union {
+    // Goal either to delete are added to the corresponding set.
+    const Binding* goal;
+    // The iterator is for std::set and this is stable upon deletion and
+    // insertion if it's not directly the element being deleted or inserted.
+    // We will only try to erase the element on the exact node traversal, so
+    // we can safely reuse the iterator that was returned from the insertion.
+    // Not using this for action ERASE_GOALS_TO_REMOVE, as we are requesting
+    // for removal before the insertion has happened.
+    GoalSet::iterator erase_it;
+    // Source set to handle by a TRAVERSE_ALL_SOURCE_SETS action and end of the
+    // source sets. The two entries are never the same for actions on the
+    // actions stack. (Ideally, this would be an anonymous struct with proper
+    // field names but gcc does not allow this, as opposed to clang.)
+    std::set<SourceSet>::const_iterator source_sets_it[2];
+  };
+
   Action(ActionType action_type, const Binding* goal)
       : action_type(action_type), goal(goal) {}
-  Action(ActionType action_type, const Binding* goal,
-         GoalSet::iterator erase_it)
-      : action_type(action_type), goal(goal), erase_it(erase_it) {}
+  Action(ActionType action_type, GoalSet::iterator erase_it)
+      : action_type(action_type), erase_it(erase_it) {}
+  Action(ActionType action_type,
+         std::set<SourceSet>::const_iterator source_sets_it,
+         std::set<SourceSet>::const_iterator source_sets_end)
+      : action_type(action_type),
+        source_sets_it{source_sets_it, source_sets_end} {}
 };
 
 static void traverse(const CFGNode* position,
@@ -93,7 +106,7 @@ static void traverse(const CFGNode* position,
     return;
   }
   auto [it, _] = state.seen_goals.insert(goal);
-  actions.emplace(ERASE_SEEN_GOALS, nullptr, it);
+  actions.emplace(ERASE_SEEN_GOALS, it);
 
   const auto* origin = goal->FindOrigin(position);
   if (!origin) {
@@ -105,18 +118,9 @@ static void traverse(const CFGNode* position,
 
   state.removed_goals.push_back(goal);
   actions.emplace(ERASE_REMOVED_GOALS, nullptr);
-  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);
-      }
-    }
+  if (!origin->source_sets.empty()) {
+    actions.emplace(TRAVERSE_ALL_SOURCE_SETS, origin->source_sets.cbegin(),
+                    origin->source_sets.cend());
   }
 }
 
@@ -149,6 +153,21 @@ static std::vector<RemoveResult> remove_finished_goals(const CFGNode* pos,
       case TRAVERSE:
         traverse(pos, results, actions, state);
         break;
+      case TRAVERSE_ALL_SOURCE_SETS: {
+        const auto& source_set = *action.source_sets_it[0];
+        action.source_sets_it[0]++;
+        if (action.source_sets_it[0] != action.source_sets_it[1]) {
+          actions.push(action);
+        }
+        for (const Binding* next_goal : source_set) {
+          auto [it, added] = state.goals_to_remove.insert(next_goal);
+          if (added) {
+            actions.emplace(ERASE_GOALS_TO_REMOVE, next_goal);
+          }
+        }
+        actions.emplace(TRAVERSE, nullptr);
+        break;
+      }
       case INSERT_GOALS_TO_REMOVE:
         state.goals_to_remove.insert(action.goal);
         break;