Minor refactoring

/cc @vheon