Runtime improvement for pseudo op extension

[devalshahamd]

This pull request refactors how modifications to the cpu_root_nodes attribute are handled when injecting pseudo-ops in the trace tree. The main improvement is to accumulate changes in a pending state and only commit them once all pseudo-op injections are complete, ensuring consistency and reducing the risk of errors from repeated mutations.

Refactoring and consistency improvements:

• inject_pseudo_op_wrap_children now accumulates additions and removals to cpu_root_nodes in per-tree pending sets, instead of modifying the list immediately. Callers are required to invoke the new finalize_pseudo_op_mutations function to commit these changes.
• Added a new function, finalize_pseudo_op_mutations, which applies all pending additions and removals to cpu_root_nodes and cleans up temporary attributes. This ensures that all changes are committed in a single, consistent operation.
• Updated apply_pseudo_op_extensions to call finalize_pseudo_op_mutations after applying all pseudo-op extensions, ensuring any accumulated mutations are committed.

[Copilot]

Pull request overview

This PR refactors pseudo-op injection to defer mutations to tree.cpu_root_nodes by accumulating root additions/removals during injection and committing them once at the end, reducing repeated list mutations and aiming for more consistent root bookkeeping.

Changes:

• Updated inject_pseudo_op_wrap_children to collect cpu_root_nodes removals/additions into per-tree pending structures rather than mutating the list immediately.
• Added finalize_pseudo_op_mutations(tree) to apply pending root-node changes and clean up temporary per-tree attributes.
• Updated apply_pseudo_op_extensions to call finalize_pseudo_op_mutations(tree) after running all extensions.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The pending-root tracking doesn’t account for pseudo-ops added during the current pending-mutation window. Because root_set is only initialized from the original tree.cpu_root_nodes and never updated with pseudo_evt['UID'], a subsequent inject_pseudo_op_wrap_children that wraps a subtree containing an earlier pseudo-op will fail to mark that earlier pseudo-op for removal (the old eager implementation would remove it). Additionally, if you do add pseudo UIDs into root_set, ensure removals take precedence over additions (e.g., filter pending_additions against pending_removals in finalize_pseudo_op_mutations).

[Copilot]

finalize_pseudo_op_mutations introduces new deferred-mutation semantics for cpu_root_nodes, but there’s no test coverage exercising the commit behavior (including edge cases like multiple injections and nested wrapping where a previously-added pseudo-op root should be removed). Adding a focused unit test would help prevent regressions in root-node bookkeeping.

Suggested change

	pending_additions = tree._cpu_root_nodes_pending_additions
	if pending_removals:
	tree.cpu_root_nodes = [
	u for u in tree.cpu_root_nodes if u not in pending_removals
	]
	if pending_additions:
	tree.cpu_root_nodes.extend(pending_additions)
	del tree._cpu_root_nodes_pending_removals
	del tree._cpu_root_nodes_pending_additions
	pending_additions = getattr(tree, "_cpu_root_nodes_pending_additions", [])
	if pending_removals:
	tree.cpu_root_nodes = [
	u for u in tree.cpu_root_nodes if u not in pending_removals
	]
	if pending_additions:
	tree.cpu_root_nodes.extend(pending_additions)
	if hasattr(tree, "_cpu_root_nodes_pending_removals"):
	del tree._cpu_root_nodes_pending_removals
	if hasattr(tree, "_cpu_root_nodes_pending_additions"):
	del tree._cpu_root_nodes_pending_additions