Initial classes to create user types

scripts/Infrastructure/TypeClasses.py

+Basetypes = {
+    "MPI_CHAR": 1,
+    "MPI_INT": 4,
+    "MPI_LONG": 8,
+    "MPI_FLOAT": 4,
+    "MPI_DOUBLE": 8,
+}
+
+
+class DataType:
+    def __init__(self, size, extent=0, realextent=0, lb=0, reallb=0):
+        self.size = size
+        self.lb = lb
+        self.reallb = reallb
+        if extent:
+            self.extent = extent
+        else:
+            self.extent = size
+        if realextent:
+            self.realextent = realextent
+        else:
+            self.realextent = size
+        # self.typemap = []
+
+    def getSize(self):
+        return self.size
+
+    def getExtent(self):
+        return self.extent
+
+    def getRealextent(self):
+        return self.realextent
+
+    def getNSize(self, count):
+        return self.size * count
+
+    def getNRealextent(self, count):
+        return self.realextent + (count - 1) * self.extent
+
+
+class BaseType(DataType):
+    def __init__(self, name):
+        DataType.__init__(self, Basetypes[name])
+        self.name = name
+        self.typemap = [(name, 0, 1)]
+
+
+class ContiguousType(DataType):
+    def __init__(self, oldtype, count):
+        if len(oldtype.typemap) == 1 and oldtype.size == oldtype.realextent:
+            DataType.__init__(self, oldtype.size * count)
+            self.typemap = [(t, o, l * count) for t, o, l in oldtype.typemap]
+        else:
+            DataType.__init__(
+                self,
+                oldtype.size * count,
+                oldtype.extent * count,
+                oldtype.realextent + oldtype.extent * (count - 1),
+            )
+            self.typemap = [
+                (t, o + i * oldtype.extent, l)
+                for i in range(count)
+                for t, o, l in oldtype.typemap
+            ]
+        self.count = count
+        self.oldtype = oldtype
+
+
+class VectorType(DataType):
+    def __init__(self, oldtype, count, blocklength, stride):
+        if len(oldtype.typemap) == 1 and oldtype.size == oldtype.realextent:
+            self.typemap = [
+                (t, o + i * stride * oldtype.extent, l * blocklength)
+                for i in range(count)
+                for t, o, l in oldtype.typemap
+            ]
+        else:
+            self.typemap = [
+                (t, o + (i * stride + b) * oldtype.extent, l)
+                for i in range(count)
+                for b in range(blocklength)
+                for t, o, l in oldtype.typemap
+            ]
+        size = oldtype.size * count
+        extent = oldtype.extent * ((count - 1) * stride + blocklength)
+        realextent = oldtype.realextent + extent - oldtype.extent
+        DataType.__init__(self, size, extent, realextent)
+        self.count = count
+        self.blocklength = blocklength
+        self.stride = stride
+        self.oldtype = oldtype
+
+
+class HVectorType(DataType):
+    def __init__(self, oldtype, count, blocklength, stride):
+        if len(oldtype.typemap) == 1 and oldtype.size == oldtype.realextent:
+            self.typemap = [
+                (t, o + i * stride, l * blocklength)
+                for i in range(count)
+                for t, o, l in oldtype.typemap
+            ]
+        else:
+            self.typemap = [
+                (t, o + i * stride + b * oldtype.extent, l)
+                for i in range(count)
+                for b in range(blocklength)
+                for t, o, l in oldtype.typemap
+            ]
+        size = oldtype.size * count
+        extent = (count - 1) * stride + oldtype.extent * blocklength
+        realextent = oldtype.realextent + extent - oldtype.extent
+        DataType.__init__(self, size, extent, realextent)
+        self.count = count
+        self.blocklength = blocklength
+        self.stride = stride
+        self.oldtype = oldtype
+
+
+class ResizedType(DataType):
+    def __init__(self, oldtype, lb, extent):
+        DataType.__init__(
+            self, oldtype.size, extent, oldtype.realextent, lb, oldtype.reallb
+        )
+        self.typemap = oldtype.typemap
+
+
+class DupType(DataType):
+    def __init__(self, oldtype):
+        DataType.__init__(
+            self,
+            oldtype.size,
+            oldtype.extent,
+            oldtype.realextent,
+            oldtype.lb,
+            oldtype.reallb,
+        )
+        self.typemap = oldtype.typemap
+
+
+class IndexedType(DataType):
+    def __init__(self, oldtype, blocklengths, displs):
+        if len(blocklengths) != len(displs):
+            raise ValueError(
+                "Arrays passed to IndexedType must have consistent lengths"
+            )
+        count = len(blocklengths)
+        if len(oldtype.typemap) == 1 and oldtype.size == oldtype.realextent:
+            self.typemap = [
+                (t, o + displs[i] * oldtype.extent, l * blocklengths[i])
+                for i in range(count)
+                for t, o, l in oldtype.typemap
+            ]
+        else:
+            self.typemap = [
+                (t, o + (displs[i] + b) * oldtype.extent, l)
+                for i in range(count)
+                for b in range(blocklengths[i])
+                for t, o, l in oldtype.typemap
+            ]
+        size = oldtype.size * sum(blocklengths)
+        blocklbs = [oldtype.lb + oldtype.extent * displs[b] for b in range(count)]
+        blockubs = [
+            oldtype.lb + oldtype.extent * (displs[b] + blocklengths[b])
+            for b in range(count)
+        ]
+        blockrlbs = [oldtype.reallb + oldtype.extent * displs[b] for b in range(count)]
+        blockrubs = [
+            oldtype.reallb
+            + oldtype.realextent
+            + oldtype.extent * (displs[b] + blocklengths[b] - 1)
+            for b in range(count)
+        ]
+        lb = min(blocklbs)
+        extent = max(blockubs) - lb
+        reallb = min(blockrlbs)
+        realextent = max(blockrubs) - reallb
+        DataType.__init__(self, size, extent, realextent, lb, reallb)
+        self.count = count
+        self.blocklengths = blocklengths
+        self.displs = displs
+        self.oldtype = oldtype
+
+
+class IndexedBlockType(DataType):
+    def __init__(self, oldtype, blocklength, displs):
+        count = len(displs)
+        if len(oldtype.typemap) == 1 and oldtype.size == oldtype.realextent:
+            self.typemap = [
+                (t, o + displs[i] * oldtype.extent, l * blocklength)
+                for i in range(count)
+                for t, o, l in oldtype.typemap
+            ]
+        else:
+            self.typemap = [
+                (t, o + (displs[i] + b) * oldtype.extent, l)
+                for i in range(count)
+                for b in range(blocklength)
+                for t, o, l in oldtype.typemap
+            ]
+        size = oldtype.size * count * blocklength
+        blocklbs = [oldtype.lb + oldtype.extent * displs[b] for b in range(count)]
+        blockubs = [
+            oldtype.lb + oldtype.extent * (displs[b] + blocklength)
+            for b in range(count)
+        ]
+        blockrlbs = [oldtype.reallb + oldtype.extent * displs[b] for b in range(count)]
+        blockrubs = [
+            oldtype.reallb
+            + oldtype.realextent
+            + oldtype.extent * (displs[b] + blocklength - 1)
+            for b in range(count)
+        ]
+        lb = min(blocklbs)
+        extent = max(blockubs) - lb
+        reallb = min(blockrlbs)
+        realextent = max(blockrubs) - reallb
+        DataType.__init__(self, size, extent, realextent, lb, reallb)
+        self.count = count
+        self.blocklength = blocklength
+        self.displs = displs
+        self.oldtype = oldtype
+
+
+class StructType(DataType):
+    def __init__(self, oldtypes, blocklengths, displs):
+        if len(oldtypes) != len(blocklengths) or len(oldtypes) != len(displs):
+            raise ValueError("Arrays passed to StructType must have consistent lengths")
+        count = len(blocklengths)
+        self.typemap = []
+        for i in range(count):
+            if (
+                len(oldtypes[i].typemap) == 1
+                and oldtypes[i].size == oldtypes[i].realextent
+            ):
+                self.typemap += [
+                    (t, o + displs[i], l * blocklengths[i])
+                    for t, o, l in oldtypes[i].typemap
+                ]
+            else:
+                self.typemap += [
+                    (t, o + displs[i] + b * oldtypes[i].extent, l)
+                    for b in range(blocklengths[i])
+                    for t, o, l in oldtypes[i].typemap
+                ]
+        size = sum([o.size * b for o, b in zip(oldtypes, blocklengths)])
+        blocklbs = [oldtypes[b].lb + displs[b] for b in range(count)]
+        blockubs = [
+            oldtypes[b].lb + displs[b] + oldtypes[b].extent * blocklengths[b]
+            for b in range(count)
+        ]
+        blockrlbs = [oldtypes[b].reallb + displs[b] for b in range(count)]
+        blockrubs = [
+            oldtypes[b].reallb
+            + oldtypes[b].realextent
+            + displs[b]
+            + oldtypes[b].extent * (blocklengths[b] - 1)
+            for b in range(count)
+        ]
+        lb = min(blocklbs)
+        extent = max(blockubs) - lb
+        reallb = min(blockrlbs)
+        realextent = max(blockrubs) - reallb
+        DataType.__init__(self, size, extent, realextent, lb, reallb)
+        self.count = count
+        self.blocklengths = blocklengths
+        self.displs = displs
+        self.oldtypes = oldtypes
+
+
+print(VectorType(BaseType("MPI_INT"), 3, 1, 5).typemap)
+cvt = ContiguousType(VectorType(BaseType("MPI_INT"), 3, 1, 5), 3)
+print(cvt.typemap, cvt.size, cvt.extent, cvt.realextent)
+it = IndexedType(
+    ResizedType(VectorType(BaseType("MPI_INT"), 4, 1, 40), -4, 4),
+    [2, 12, 2, 2],
+    [0, 4, -4, 12],
+)
+print(it.typemap, it.size, it.extent, it.realextent)
+
+st = StructType(
+    [BaseType("MPI_CHAR"), BaseType("MPI_INT"), BaseType("MPI_DOUBLE")],
+    [9, 5, 3],
+    [0, 16, 48],
+)
+print(st.typemap, st.size, st.extent, st.realextent)