Schedule build pass maintenance

[PixelDust22]

Objective

• Add get_build_pass_mut method to Schedule. This requires ScheduleBuildPassObj to add an Any bound.
  In No schedule build pass overwrite if build settings do not change auto_insert_apply_deferred from true #19217 it was suggested that a get_build_pass_mut API be added to Schedule. At the time we decided against that because Remove upcasting methods + Cleanup interned label code #18984 was still open. Now that Remove upcasting methods + Cleanup interned label code #18984 is merged, we can finally do this.
• Pass &mut World and &mut ScheduleGraph down into ScheduleBuildPass::collapse_set as well. This allows the collapse_set method to interact with the world and the schedule.
• get_set_at and set_at methods to return an InternedSystemSet instead of &mut dyn SystemSet. InternedSystemSet dereferences to &mut dyn SystemSet so this is non-breaking, but InternedSystemSet is 'static and clonable, giving users more flexibility.
• Add map_set_to_systems hook to schedule build passes which allows the user to specify that a system set covers additional systems.

[alice-i-cecile]

I'd like to see more docs / tests that demonstrate why this functionality is useful, but this is simple enough that I won't block on it. Other reviewers might!

[urben1680]

I think I cannot fully grasp what these changes enable one to do but on a technical level it looks good.

[chescock]

The addition of get_build_pass_mut, the change from &dyn SystemSet to InternedSystemSet, and the default implementations of SchedulePass methods all seem straightforward, and I'd happily approve smaller PRs for those now!

I don't think I understand the use cases for the map_set_to_systems method or for passing &mut World and &mut ScheduleGraph to collapse_set well enough to know whether those are reasonable, though, and those cost some extra hokey-pokeys and a self.hierarchy.topsort.clone(). Can you say more about what you're trying to do?

[chescock]

Can you say more about the use case for this new function?

It sounds like it would be for making new automatically-populated system sets, but the API of taking one system set and asking for SystemKeys seems hard to work with. Wouldn't any implementation need to scan the list of systems to find the matching ones? Why not have an API where you ask about one system and return the InternedSystemSets to add?

Also, this API gets called while flattening the graph. If the goal is to add new associations, that might be easier to do before flattening. Would it make sense to instead add a function with a similar signature to build that gets called before flattening, and make something like update_graphs public so that the pass can do arbitrary set and system configuration?

[PixelDust22]

This new function allows the user to associate additional systems to an existing set.

My use case is that I want to add a prologue system and an epilogue system to some marked system sets. I added a method on App that allows the user to mark some system sets. A schedule build pass adds a prologue system that runs before all systems in the set, and a epilogue that runs after all systems in the set.

I used to think that you can do this in collapse_set since it's run for each system. Turns out you can't. If you have a nested system set for example, by the time collapse_set runs, the graph has already been flattened. So you have lost the opportunity to connect the Prologue system with "things that happens before" and the epilogue system with "things that happens after".

I chose this API because it maps best to how we currently build the schedule. The schedule builder roughly does the following:

1. Topological sort
2. Create a mapping from SystemSet to all systems contained within the SystemSet.
3. Flatten the graph
4. Build.

We currently have hooks for 3 and 4 but not 2. It feels perfectly natural to add a new hook for 2 so that this mapping accurately reflects the hierarchical relationships between systems and system sets.

Also, this API gets called while flattening the graph. If the goal is to add new associations, that might be easier to do before flattening.

That's the whole point of this change. We have the collapse_set hook that runs while flattening the graph. For adding new associations we need a new hook, map_set_to_systems, which runs before flattening the graph.

[alice-i-cecile]

Cool! That's a neat little bit of functionality :) I think it could be more clearly explained / exposed, but I think it's quite neat.

[chescock]

My use case is that I want to add a prologue system and an epilogue system to some marked system sets.

Thanks, that's a really clear example!

This is the code I have for adding additional systems, and hopefully it shows you why it's necessary to pass down references to &mut ScheduleGraph and &mut World.

Ah, I see, you need &mut World to call System::initialize, because this code runs after ScheduleGraph::initialize. That makes sense, although this is also the sort of thing that worries me. It would be really easy for schedule pass authors to add new systems but forget to initialize them!

That might even be unsound, since we wouldn't initialize the access correctly. The built-in System impls will panic if you run them without calling initialize, but I think it's valid for a third-party impl to run and assume it has world access. ... Oh, ugh, build already takes &mut ScheduleGraph, so maybe this is already a little bit unsound.

For adding new associations we need a new hook, map_set_to_systems, which runs before flattening the graph.

Ah, part of the confusion here might just be different terminology. I think of map_sets_to_systems as flattening the hierarchy graph, while get_dependency_flattened flattens the dependency graph.

Anyway, my point was that it might be easier to add systems and sets in the original hierarchy graph rather than in the middle of map_sets_to_systems.

---

Would it work to have a method on ScheduleBuildPass that gets called first thing during Schedule::initialize? It should be completely safe to offer &mut Schedule there, since the scheduling code hasn't read anything yet, and it would you you add new systems without needing to remember to initialize them.

You could iterate the system sets using graph.system_sets.iter(), and get the nodes in a set using graph.hierarchy().graph().neighbors_directed(node, Outgoing), so I think the existing public API will let you find the nodes that you'll need to add edges to. (I'm not sure whether your plan is to add the prologue and epilogue to the set itself, and then add dependencies to the other nodes in the set, or whether it's to add them outside the set, and add dependencies to both the set and its dependencies, but either should be possible.)

[alice-i-cecile]

This may ultimately be much easier to land if it's split up into small, well-motivated PRs.

[github-actions]

It looks like your PR is a breaking change, but you didn't provide a migration guide.

Please review the instructions for writing migration guides, then expand or revise the content in the migration guides directory to reflect your changes.

[PixelDust22]

Migration guide added.