Code cleaned up a bit.