use flat_map instead of map.flatten

Benchmarks from @sferik on https://github.com/rails/rails/pull/14240 , shows that `flat_map` is faster than `map.flatten` and also more readable.

All tests are green with this change, so IMO it is worth it.