Cleanup code a bit.