Refactor some code, get it to build