Cleaned up some code at Greg's suggestion