Refactoring

Refactoring based on discussion in #351