type checker tests too

Author: junikimm717

src/finchlite/codegen/c.py

@@ -6,19 +6,17 @@
 import tempfile
 from abc import ABC, abstractmethod
 from collections import namedtuple
-from collections.abc import Callable
+from collections.abc import Callable, Hashable
 from functools import lru_cache
 from pathlib import Path
 from types import NoneType
 from typing import Any
 
 import numpy as np
 
-from finchlite.finch_assembly.map import FType, MapFType
-
 from .. import finch_assembly as asm
 from ..algebra import query_property, register_property
-from ..finch_assembly import AssemblyStructFType, BufferFType, TupleFType
+from ..finch_assembly import AssemblyStructFType, BufferFType, MapFType, TupleFType
 from ..symbolic import Context, Namespace, ScopedDict, fisinstance, ftype
 from ..util import config
 from ..util.cache import file_cache
@@ -612,22 +610,22 @@ def __init__(
         self.fptr = fptr
         self.types = types
         self.slots = slots
-        self.datastructures: dict[FType, Any] = {}
+        self.datastructures: dict[Hashable, Any] = {}
 
     def add_header(self, header):
         if header not in self._headerset:
             self.headers.append(header)
             self._headerset.add(header)
 
-    def add_datastructure(self, ftype: FType, handler: "Callable[[CContext], Any]"):
+    def add_datastructure(self, key: Hashable, handler: "Callable[[CContext], Any]"):
         """
         Code to add a datastructure declaration.
         This is the minimum required to prevent redundancy.
         """
-        if ftype in self.datastructures:
+        if key in self.datastructures:
             return
         # at least mark something is there.
-        self.datastructures[ftype] = None
+        self.datastructures[key] = None
         handler(self)
 
     def emit_global(self):

src/finchlite/codegen/hashtable.py

@@ -390,7 +390,7 @@ def c_unpack(self, ctx: "CContext", var_n: str, val: AssemblyExpression):
         data = ctx.freshen(var_n, "data")
         # Add all the stupid header stuff from above.
         ctx.add_datastructure(
-            self,
+            ("CHashTableFType", self.key_len, self.value_len),
             lambda ctx: CHashTable.gen_code(
                 ctx, self.key_type, self.value_type, inline=True
             ),
@@ -445,9 +445,6 @@ def __init__(self, key_len, value_len, map: "dict[tuple,tuple] | None" = None):
         self.key_len = key_len
         self.value_len = value_len
 
-        self._key_type = _int_tuple_ftype(key_len)
-        self._value_type = _int_tuple_ftype(value_len)
-
         self._numba_key_type = numba.types.UniTuple(numba.types.int64, key_len)
         self._numba_value_type = numba.types.UniTuple(numba.types.int64, value_len)
 

src/finchlite/finch_assembly/__init__.py

@@ -6,6 +6,7 @@
 )
 from .dataflow import AssemblyCopyPropagation, assembly_copy_propagation
 from .interpreter import AssemblyInterpreter, AssemblyInterpreterKernel
+from .map import MapFType
 from .nodes import (
     AssemblyNode,
     Assert,
@@ -69,6 +70,7 @@
     "Literal",
     "Load",
     "LoadMap",
+    "MapFType",
     "Module",
     "NamedTupleFType",
     "Print",

src/finchlite/finch_assembly/map.py

@@ -16,20 +16,24 @@ def __init__(
     ): ...
 
     @property
-    def element_type(self):
+    @abstractmethod
+    def ftype(self) -> "MapFType": ...
+
+    @property
+    def value_type(self):
         """
-        Return the type of elements stored in the hash table.
-        This is typically the same as the dtype used to create the map.
+        Return type of values stored in the hash table
+        (probably some TupleFType)
         """
-        return self.ftype.element_type()
+        return self.ftype.value_type
 
     @property
-    def length_type(self):
+    def key_type(self):
         """
-        Return the type of indices used to access elements in the hash table.
-        This is typically an integer type.
+        Return type of keys stored in the hash table
+        (probably some TupleFType)
         """
-        return self.ftype.length_type()
+        return self.ftype.key_type
 
     @abstractmethod
     def load(self, idx: tuple):

tests/test_assembly_type_checker.py

@@ -7,7 +7,9 @@
 
 import finchlite.finch_assembly as asm
 from finchlite.codegen import NumpyBuffer
+from finchlite.codegen.hashtable import CHashTable, NumbaHashTable
 from finchlite.finch_assembly import assembly_check_types
+from finchlite.finch_assembly.struct import TupleFType
 from finchlite.symbolic import FType, ftype
 
 
@@ -684,3 +686,99 @@ def test_simple_struct():
     )
 
     assembly_check_types(mod)
+
+
+@pytest.mark.parametrize(
+    ["tabletype"],
+    [
+        (CHashTable,),
+        (NumbaHashTable,),
+    ],
+)
+def test_hashtable(tabletype):
+    table = tabletype(2, 3)
+
+    table_v = asm.Variable("a", ftype(table))
+    table_slt = asm.Slot("a_", ftype(table))
+
+    key_type = table.ftype.key_type
+    val_type = table.ftype.value_type
+    key_v = asm.Variable("key", key_type)
+    val_v = asm.Variable("val", val_type)
+
+    mod = asm.Module(
+        (
+            asm.Function(
+                asm.Variable(
+                    "setidx", TupleFType.from_tuple(tuple(int for _ in range(3)))
+                ),
+                (table_v, key_v, val_v),
+                asm.Block(
+                    (
+                        asm.Unpack(table_slt, table_v),
+                        asm.StoreMap(
+                            table_slt,
+                            key_v,
+                            val_v,
+                        ),
+                        asm.Repack(table_slt),
+                        asm.Return(asm.LoadMap(table_slt, key_v)),
+                    )
+                ),
+            ),
+            asm.Function(
+                asm.Variable("exists", bool),
+                (table_v, key_v),
+                asm.Block(
+                    (
+                        asm.Unpack(table_slt, table_v),
+                        asm.Return(asm.ExistsMap(table_slt, key_v)),
+                    )
+                ),
+            ),
+        )
+    )
+    assembly_check_types(mod)
+
+
+@pytest.mark.parametrize(
+    ["tabletype"],
+    [
+        (CHashTable,),
+        (NumbaHashTable,),
+    ],
+)
+def test_hashtable_fail(tabletype):
+    table = tabletype(2, 3)
+
+    table_v = asm.Variable("a", ftype(table))
+    table_slt = asm.Slot("a_", ftype(table))
+
+    key_type = table.ftype.key_type
+    val_type = table.ftype.value_type
+    key_v = asm.Variable("key", key_type)
+    val_v = asm.Variable("val", val_type)
+    mod = asm.Module(
+        (
+            asm.Function(
+                asm.Variable(
+                    "setidx", TupleFType.from_tuple(tuple(int for _ in range(2)))
+                ),
+                (table_v, key_v, val_v),
+                asm.Block(
+                    (
+                        asm.Unpack(table_slt, table_v),
+                        asm.StoreMap(
+                            table_slt,
+                            key_v,
+                            val_v,
+                        ),
+                        asm.Repack(table_slt),
+                        asm.Return(asm.LoadMap(table_slt, key_v)),
+                    )
+                ),
+            ),
+        )
+    )
+    with pytest.raises(asm.AssemblyTypeError):
+        assembly_check_types(mod)