General cleanup of code