Little code refactoring

Hi, I have done some minor changes, please review it.

Thanks!