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: 707500665

Author: Martin Huschenbett

pytype/typegraph/solver.cc

@@ -42,6 +42,7 @@ struct TraverseState {
 
 enum ActionType {
   TRAVERSE,
+  SOURCE_SET,
   INSERT_GOALS_TO_REMOVE,
   ERASE_GOALS_TO_REMOVE,
   ERASE_SEEN_GOALS,
@@ -57,20 +58,36 @@ 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;
+    // Source set to handle by a SOURCE_SET action. This is never the same as
+    // source_sets_end for actions on the actions stack.
+    std::set<SourceSet>::const_iterator source_sets_it;
+  };
+  union {
+    // 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;
+    // End of the source sets in a SOURCE_SET action.
+    std::set<SourceSet>::const_iterator source_sets_end;
+  };
+
   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,
+         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(source_sets_end) {}
 };
 
 static void traverse(const CFGNode* position,
@@ -104,19 +121,10 @@ 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);
-      }
-    }
+  actions.emplace(ERASE_REMOVED_GOALS, nullptr, it);
+  if (!origin->source_sets.empty()) {
+    actions.emplace(SOURCE_SET, origin->source_sets.cbegin(),
+                    origin->source_sets.cend());
   }
 }
 
@@ -149,6 +157,21 @@ static std::vector<RemoveResult> remove_finished_goals(const CFGNode* pos,
       case TRAVERSE:
         traverse(pos, results, actions, state);
         break;
+      case SOURCE_SET: {
+        const auto& source_set = *action.source_sets_it;
+        action.source_sets_it++;
+        if (action.source_sets_it != action.source_sets_end) {
+          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;