Done minor refactoring to the code.