Code cleanup/improvements