Refactored code a bit.