Some more refactoring after pulling in Leo's changes