Code cleanup / restructure