Cleanup and rearrange code.