Further refactoring, trying to make it cleaner.