Code re-organization and comments