Cleanups and minor improvements :)