Minor refactor changes