Code restructuring / cleanup