From ed55dfe7d19e28e88f725c5cae17343bd2c1c678 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Sun, 4 Sep 2022 21:16:06 +0300
Subject: [PATCH] fix TranslationUnit._hash_value disappearing during pickling

---
 loopy/translation_unit.py | 16 ++++++++--------
 test/test_loopy.py        | 13 +++++++++++++
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/loopy/translation_unit.py b/loopy/translation_unit.py
index 73d5ad3b0..41cbb4eca 100644
--- a/loopy/translation_unit.py
+++ b/loopy/translation_unit.py
@@ -373,14 +373,14 @@ def __setstate__(self, state_obj):
         self._program_executor_cache = {}
 
     def __hash__(self):
-        if self._hash_value is not None:
-            return self._hash_value
-
-        from loopy.tools import LoopyKeyBuilder
-        from pytools.persistent_dict import new_hash
-        key_hash = new_hash()
-        self.update_persistent_hash(key_hash, LoopyKeyBuilder())
-        self._hash_value = hash(key_hash.digest())
+        # NOTE: _hash_value may vanish during pickling
+        if getattr(self, "_hash_value", None) is None:
+            from loopy.tools import LoopyKeyBuilder
+            from pytools.persistent_dict import new_hash
+            key_hash = new_hash()
+            self.update_persistent_hash(key_hash, LoopyKeyBuilder())
+            self._hash_value = hash(key_hash.digest())
+
         return self._hash_value
 
 # }}}
diff --git a/test/test_loopy.py b/test/test_loopy.py
index 051f98b0c..94e1ae8ea 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -3350,6 +3350,19 @@ def test_different_index_dtypes():
         lp.generate_code_v2(knl)
 
 
+def test_translation_unit_pickle():
+    tunit = lp.make_kernel(
+        "{[i]: 0<=i<16}",
+        """
+        y[i] = i
+        """)
+    assert isinstance(hash(tunit), int)
+
+    from pickle import dumps, loads
+    tunit = loads(dumps(tunit))
+    assert isinstance(hash(tunit), int)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])