Minor refactor of halo exchange subroutines