Adding some abstraction & splitting code