It’s interesting to hear your experience with comments-as-marks! I’m currently doing something similar but using the alternative approach of tracking ranges (from, to) in the document, and storing the set adjacent to the content. When changes are made to the document, I update the ranges by mapping them through the step maps. (I use decorations to apply styling)
However this approach has the pitfall of not “just working” with regard to undo.
So I’m faced with two options — switch to using marks, or implement support for “undo” of ranges. I’ll probably go the route of trying to support “undo” for ranges.