Small refactoring

To be merged before release.