Refactored code to be more general.