Small refactor changes