Cleanup of code