Rearranging some code to make it more flexible and manageable.