Clean up #4858

This would also probably be a good place to use a `defaultdict`.

CC: @techhat