More code clean up