use a class for CachedMapper caches instead of using a dict directly

inducer · Jan 10, 2025 · d40153c · d40153c
1 parent 27aa393
commit d40153c
Show file tree

Hide file tree

Showing 5 changed files with 210 additions and 64 deletions.
diff --git a/pytato/codegen.py b/pytato/codegen.py
@@ -68,7 +68,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Hashable, Mapping
+    from collections.abc import Mapping
 
     from pytato.function import FunctionDefinition, NamedCallResult
     from pytato.target import Target
@@ -135,12 +135,13 @@ class CodeGenPreprocessor(ToIndexLambdaMixin, CopyMapper):  # type: ignore[misc]
     :class:`~pytato.array.Stack`            :class:`~pytato.array.IndexLambda`
     ======================================  =====================================
     """
+    _FunctionCacheT: TypeAlias = CopyMapper._FunctionCacheT
 
     def __init__(
             self,
             target: Target,
             kernels_seen: dict[str, lp.LoopKernel] | None = None,
-            _function_cache: dict[Hashable, FunctionDefinition] | None = None
+            _function_cache: _FunctionCacheT | None = None
             ) -> None:
         super().__init__(_function_cache=_function_cache)
         self.bound_arguments: dict[str, DataInterface] = {}

diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py
@@ -100,6 +100,8 @@
 
 
 if TYPE_CHECKING:
+    from typing import TypeAlias
+
     import mpi4py.MPI
 
     from pytato.function import FunctionDefinition, NamedCallResult
@@ -290,12 +292,13 @@ class _DistributedInputReplacer(CopyMapper):
     instances for their assigned names. Also gathers names for
     user-supplied inputs needed by the part
     """
+    _FunctionCacheT: TypeAlias = CopyMapper._FunctionCacheT
 
     def __init__(self,
                  recvd_ary_to_name: Mapping[Array, str],
                  sptpo_ary_to_name: Mapping[Array, str],
                  name_to_output: Mapping[str, Array],
-                 _function_cache: dict[Hashable, FunctionDefinition] | None = None,
+                 _function_cache: _FunctionCacheT | None = None,
                  ) -> None:
         super().__init__(_function_cache=_function_cache)
 
@@ -344,9 +347,9 @@ def map_distributed_send(self, expr: DistributedSend) -> DistributedSend:
         return new_send
 
     def rec(self, expr: ArrayOrNames) -> ArrayOrNames:
-        key = self.get_cache_key(expr)
+        key = self._cache.get_key(expr)
         try:
-            return self._cache[key]
+            return self._cache.retrieve(expr, key=key)
         except KeyError:
             pass