When #2311 lands we will have an implementation that works for inserting comments in the formatted code but as @mitchmindtree mentioned in the following comment, we should investigate and implement efficiency improvements:
My main concern with this approach is that we end up parsing the entire file at least twice, and end up traversing the tree three times: once during formatting, once when collecting unformatted spans and again when collecting formatted spans.
Originally, I was imagining that we could take an approach that inserts comments during formatting of the rest of the code. E.g. in the new add_comments function, we use an offset to track the difference between the unformatted and formatted span start position. Instead, we could do this during traversal by tracking the offset within the Formatter itself. However, this approach would require abstracting away the visiting of the nodes from the formatting of the nodes first and would require a non-trivial refactor.
#2311 (review)