Minor tweak to code structure, added docs