Code cleanup & additions