Refactoring and added protocol