More scala refactoring