Slight refactor