Refactored some places logic.
Some minor code cleanup.