Refactoring some code again...