Clean up code a bit after merging.