mpi: Fixup halo exchange generation involving Bundles

Author: FabioLuporini

devito/mpi/halo_scheme.py

@@ -43,6 +43,24 @@ def __hash__(self):
                      self.halos,
                      self.dims))
 
+    def union(self, other):
+        """
+        Return a new HaloSchemeEntry that is the union of this and `other`.
+        The `loc_indices` and `loc_dirs` must be the same, otherwise an
+        exception is raised.
+        """
+        if self.loc_indices != other.loc_indices or \
+           self.loc_dirs != other.loc_dirs:
+            raise HaloSchemeException(
+                "Inconsistency found while building a HaloScheme"
+            )
+
+        halos = self.halos | other.halos
+        dims = self.dims | other.dims
+
+        return HaloSchemeEntry(self.loc_indices, self.loc_dirs, halos, dims,
+                               getters=self.getters)
+
 
 Halo = namedtuple('Halo', 'dim side')
 
@@ -410,9 +428,11 @@ def add(self, f, hse):
         """
         Create a new HaloScheme that contains all entries in `self` plus the one
         passed in input. If `f` already exists in `self`, the old value is
-        overridden.
+        augmented with `hse` (i.e., the halos are unioned).
         """
-        fmapper = dict(self.fmapper.items())
+        fmapper = dict(self.fmapper)
+        if f in fmapper:
+            hse = fmapper[f].union(hse)
         fmapper[f] = hse
         return HaloScheme.build(fmapper, self.honored)