Scala code refactoring