From 9b2f9c18a22302f2a5f7a18eda50b0d6dbbc1926 Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Wed, 18 Dec 2024 01:10:40 -0800 Subject: [PATCH] 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 | 138 ++++++++++++++++++++++++++----------- 1 file changed, 99 insertions(+), 39 deletions(-) diff --git a/pytype/typegraph/solver.cc b/pytype/typegraph/solver.cc index c1096d725..411fffdae 100644 --- a/pytype/typegraph/solver.cc +++ b/pytype/typegraph/solver.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -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 removed_goals; + std::vector 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& results, + std::stack>& 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 remove_finished_goals(const CFGNode* pos, state.goals_to_remove.end(), std::inserter(state.new_goals, state.new_goals.begin()), pointer_less()); - std::deque queue; - queue.push_back(std::move(state)); + std::stack> actions; + actions.emplace(TRAVERSE, nullptr); std::vector 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;