refactored a little code