Refactoring to better integrate with Simon's changes