Slight cleaning of code