Rearrange code to match documentation structure