Small fixes / cleanup