Small fixes & cleanup