Code review and cleanup