compiler: Fix ComponentAccess pickling

Author: FabioLuporini

devito/types/array.py

@@ -4,7 +4,7 @@
 import numpy as np
 from sympy import Expr, cacheit
 
-from devito.tools import (Reconstructable, as_tuple, c_restrict_void_p,
+from devito.tools import (Pickable, as_tuple, c_restrict_void_p,
                           dtype_to_ctype, dtypes_vector_mapper, is_integer)
 from devito.types.basic import AbstractFunction, LocalType
 from devito.types.utils import CtypesFactory, DimensionTuple
@@ -518,10 +518,11 @@ def handles(self):
         return self.components
 
 
-class ComponentAccess(Expr, Reconstructable):
+class ComponentAccess(Expr, Pickable):
 
     _component_names = ('x', 'y', 'z', 'w')
 
+    __rargs__ = ('arg',)
     __rkwargs__ = ('index',)
 
     def __new__(cls, arg, index=0, **kwargs):
@@ -543,7 +544,7 @@ def __str__(self):
 
     __repr__ = __str__
 
-    func = Reconstructable._rebuild
+    func = Pickable._rebuild
 
     def _sympystr(self, printer):
         return str(self)
@@ -552,6 +553,10 @@ def _sympystr(self, printer):
     def base(self):
         return self.args[0]
 
+    @property
+    def arg(self):
+        return self.base
+
     @property
     def index(self):
         return self._index

tests/test_pickle.py

@@ -20,7 +20,7 @@
 from devito.types import (Array, CustomDimension, Symbol as dSymbol, Scalar,
                           PointerArray, Lock, PThreadArray, SharedData, Timer,
                           DeviceID, NPThreads, ThreadID, TempFunction, Indirection,
-                          FIndexed, StencilDimension)
+                          FIndexed, ComponentAccess, StencilDimension)
 from devito.types.basic import BoundSymbol, AbstractSymbol
 from devito.tools import EnrichedTuple
 from devito.symbolics import (IntDiv, ListInitializer, FieldFromPointer,
@@ -417,6 +417,20 @@ def test_findexed(self, pickle):
         assert new_fi.indices == (x+1, y, z-2)
         assert new_fi.strides_map == fi.strides_map
 
+    def test_component_access(self, pickle):
+        grid = Grid(shape=(3, 3, 3))
+        x, y, z = grid.dimensions
+
+        f = Function(name='f', grid=grid)
+
+        ca = ComponentAccess(f.indexify(), 1)
+
+        pkl_ca = pickle.dumps(ca)
+        new_ca = pickle.loads(pkl_ca)
+
+        assert new_ca.index == 1
+        assert new_ca.function.name == f.name
+
     def test_weights_to_array(self, pickle):
         grid = Grid(shape=(3, 3, 3))
         x, y, z = grid.dimensions

tests/test_symbolics.py

@@ -483,6 +483,24 @@ def test_findexed():
     assert new_fi.strides_map == strides_map
 
 
+def test_component_access():
+    grid = Grid(shape=(3, 3, 3))
+    x, y, z = grid.dimensions
+
+    f = Function(name='f', grid=grid)
+
+    cf0 = ComponentAccess(f.indexify(), 0)
+    cf1 = ComponentAccess(f.indexify(), 1)
+
+    assert ccode(cf0) == 'f[x][y][z].x'
+    assert ccode(cf1) == 'f[x][y][z].y'
+
+    # Reconstruction
+    cf2 = cf1.func(*cf1.args)
+    assert cf2.index == cf1.index
+    assert cf2 == cf1
+
+
 def test_canonical_ordering_of_weights():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions