How do you keep notes of your work as you code?

I find keeping mental models of the code I'm working on or reviewing to be a lot of cognitive work. So tend to jot things down either in a notebook or just put screenshots onto OneNote (I like the big endless canvas and ability to scribble notes). I feel like there is a more intuitive way for doing this. Any recommendations for tooling? How else do people do it?