Refactor to bring in line with Master