Slight refactor [skip ci]