samll changes and refactoring