A little code refactor