per se, this wont be that much of a hassle, just create a tree structure, have an (git like)
HEAD where you append new states. The only thing i dont really know is how to hook existing functionality which pushes states to the undo stack, but i guess one could just hook all edit events and push that. The developer of Braid had a good talk about how to have a lot of undo information while not utilizing a ton of memory, those ideas would be useful here.
as for visualization, i would copy the code from git to layout and draw the tree in ascii.
as @FichteFoll said, you could even try using a in-memory VCS, maybe dulwich can be hacked that way