Some typing fixes.

Author: hameerabbasi

src/finch/autoschedule/compiler.py

@@ -17,38 +17,40 @@
 )
 from ..symbolic import Term
 
-T = TypeVar("T", bound=Term)
+T = TypeVar("T", bound="LogicNode")
 
 
-def get_or_insert(dictionary: dict[str, Term], key: str, default: Any) -> Any:
+def get_or_insert(dictionary: dict[str, LogicNode], key: str, default: Any) -> Any:
     if key in dictionary:
         return dictionary[key]
     dictionary[key] = default
     return default
 
 
-def get_structure(node: T, fields: dict[str, Term], aliases: dict[str, Term]) -> T:
+def get_structure(
+    node: LogicNode, fields: dict[str, LogicNode], aliases: dict[str, LogicNode]
+) -> LogicNode:
     match node:
         case Field(name):
             return get_or_insert(fields, name, Immediate(len(fields) + len(aliases)))
         case Alias(name):
             return get_or_insert(aliases, name, Immediate(len(fields) + len(aliases)))
         case Subquery(Alias(name) as lhs, arg):
             if name in aliases:
-                return aliases[name]  # type: ignore[return-value]
-            return Subquery(  # type: ignore[return-value]
+                return aliases[name]
+            return Subquery(
                 get_structure(lhs, fields, aliases), get_structure(arg, fields, aliases)
             )
         case Table(tns, idxs):
-            return Table(  # type: ignore[return-value]
+            assert all(isinstance(idx, Field) for idx in idxs)
+            return Table(
                 Immediate(type(tns.val)),
-                tuple(get_structure(idx, fields, aliases) for idx in idxs),
+                tuple(get_structure(idx, fields, aliases) for idx in idxs),  # type: ignore[misc]
+            )
+        case any if any.is_expr():
+            return any.make_term(
+                *[get_structure(arg, fields, aliases) for arg in any.children()]
             )
-        # TODO: `is_tree` isn't defined
-        # case any if any.is_tree():
-        #     return any.from_arguments(
-        #         *[get_structure(arg, fields, aliases) for arg in any.get_arguments()]
-        #     )
         case _:
             return node
 

src/finch/autoschedule/optimize.py

@@ -1,6 +1,7 @@
 from ..finch_logic import (
     Aggregate,
     Alias,
+    LogicNode,
     MapJoin,
     Plan,
     Produces,
@@ -17,7 +18,9 @@ def optimize(prgm: Term) -> Term:
     return propagate_map_queries(prgm)
 
 
-def _lift_subqueries_expr(node: Term, bindings: dict[Term, Term]) -> Term:
+def _lift_subqueries_expr(
+    node: LogicNode, bindings: dict[LogicNode, LogicNode]
+) -> LogicNode:
     match node:
         case Subquery(lhs, arg):
             if lhs not in bindings:
@@ -26,19 +29,18 @@ def _lift_subqueries_expr(node: Term, bindings: dict[Term, Term]) -> Term:
             return lhs
         case any if any.is_expr():
             return any.make_term(
-                any.head(),
-                [_lift_subqueries_expr(x, bindings) for x in any.children()],
+                *tuple(_lift_subqueries_expr(x, bindings) for x in any.children()),
             )
         case _:
             return node
 
 
-def lift_subqueries(node: Term) -> Term:
+def lift_subqueries(node: Term) -> LogicNode:
     match node:
         case Plan(bodies):
             return Plan(tuple(map(lift_subqueries, bodies)))
         case Query(lhs, rhs):
-            bindings: dict[Term, Term] = {}
+            bindings: dict[LogicNode, LogicNode] = {}
             rhs_2 = _lift_subqueries_expr(rhs, bindings)
             return Plan(
                 (*[Query(lhs, rhs) for lhs, rhs in bindings.items()], Query(lhs, rhs_2))
@@ -91,6 +93,6 @@ class DefaultLogicOptimizer:
     def __init__(self, ctx: LogicCompiler):
         self.ctx = ctx
 
-    def __call__(self, prgm: Term):
+    def __call__(self, prgm: Term) -> str:
         prgm = optimize(prgm)
         return self.ctx(prgm)

src/finch/finch_logic/nodes.py

@@ -2,13 +2,10 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import Any, Self, TypeVar
 
 from ..symbolic import Term
 
-if TYPE_CHECKING:
-    pass
-
 __all__ = [
     "LogicNode",
     "Immediate",
@@ -28,6 +25,9 @@
 ]
 
 
+T = TypeVar("T", bound="LogicNode")
+
+
 @dataclass(eq=True, frozen=True)
 class LogicNode(Term):
     """
@@ -57,14 +57,14 @@ def head(cls):
         """Returns the head of the node."""
         return cls
 
-    def children(self) -> list[Term]:
+    @abstractmethod
+    def children(self) -> list[LogicNode]:
         """Returns the children of the node."""
-        raise Exception(f"`children` isn't supported for {self.__class__}.")
 
     @classmethod
-    def make_term(cls, head, args):
+    def make_term(cls, *args: Term) -> Self:
         """Creates a term with the given head and arguments."""
-        return head(*args)
+        return cls(*args)
 
 
 @dataclass(eq=True, frozen=True)
@@ -94,6 +94,9 @@ def fill_value(self):
 
         return fill_value(self)
 
+    def children(self):
+        raise TypeError(f"`{type(self).__name__}` doesn't support `.children()`.")
+
 
 @dataclass(eq=True, frozen=True)
 class Deferred(LogicNode):
@@ -222,7 +225,7 @@ class MapJoin(LogicNode):
     """
 
     op: Immediate
-    args: tuple[Term, ...]
+    args: tuple[LogicNode, ...]
 
     @staticmethod
     def is_expr():
@@ -239,8 +242,8 @@ def children(self):
         return [self.op, *self.args]
 
     @classmethod
-    def make_term(cls, head, args):
-        return head(args[0], tuple(args[1:]))
+    def make_term(cls, op: Immediate, *args: LogicNode) -> Self:  # type: ignore[override]
+        return cls(op, tuple(args))
 
 
 @dataclass(eq=True, frozen=True)
@@ -258,7 +261,7 @@ class Aggregate(LogicNode):
 
     op: Immediate
     init: Immediate
-    arg: Term
+    arg: LogicNode
     idxs: tuple[Field, ...]
 
     @staticmethod
@@ -288,7 +291,7 @@ class Reorder(LogicNode):
         idxs: The new order of dimensions.
     """
 
-    arg: Term
+    arg: LogicNode
     idxs: tuple[Field, ...]
 
     @staticmethod
@@ -317,7 +320,7 @@ class Relabel(LogicNode):
         idxs: The new labels for dimensions.
     """
 
-    arg: Term
+    arg: LogicNode
     idxs: tuple[Field, ...]
 
     @staticmethod
@@ -346,7 +349,7 @@ class Reformat(LogicNode):
     """
 
     tns: Immediate
-    arg: Term
+    arg: LogicNode
 
     @staticmethod
     def is_expr():
@@ -374,8 +377,8 @@ class Subquery(LogicNode):
         rhs: The argument to evaluate.
     """
 
-    lhs: Term
-    arg: Term
+    lhs: LogicNode
+    arg: LogicNode
 
     @staticmethod
     def is_expr():
@@ -403,8 +406,8 @@ class Query(LogicNode):
         rhs: The right-hand side to evaluate.
     """
 
-    lhs: Term
-    rhs: Term
+    lhs: LogicNode
+    rhs: LogicNode
 
     @staticmethod
     def is_expr():
@@ -431,7 +434,7 @@ class Produces(LogicNode):
         args: The arguments to return.
     """
 
-    args: tuple[Term, ...]
+    args: tuple[LogicNode, ...]
 
     @staticmethod
     def is_expr():
@@ -445,11 +448,11 @@ def is_stateful():
 
     def children(self):
         """Returns the children of the node."""
-        return list(self.args)
+        return [*self.args]
 
     @classmethod
-    def make_term(cls, head, args):
-        return head(tuple(args))
+    def make_term(cls, *args: LogicNode) -> Self:  # type: ignore[override]
+        return cls(tuple(args))
 
 
 @dataclass(eq=True, frozen=True)
@@ -462,7 +465,7 @@ class Plan(LogicNode):
         bodies: The sequence of statements to execute.
     """
 
-    bodies: tuple[Term, ...] = ()
+    bodies: tuple[LogicNode, ...] = ()
 
     @staticmethod
     def is_expr():
@@ -479,5 +482,5 @@ def children(self):
         return tuple(self.bodies)
 
     @classmethod
-    def make_term(cls, head, val):
-        return head(tuple(val))
+    def make_term(cls, *bodies: LogicNode) -> Self:  # type: ignore[override]
+        return cls(tuple(bodies))

src/finch/symbolic/rewriters.py

@@ -21,6 +21,7 @@
 """
 
 from collections.abc import Callable, Iterable
+from typing import TypeVar
 
 from .term import Term
 
@@ -35,10 +36,13 @@
     "Memo",
 ]
 
+T = TypeVar("T", bound="Term")
+U = TypeVar("U", bound="Term")
+
 RwCallable = Callable[[Term], Term | None]
 
 
-def default_rewrite(x: Term | None, y: Term) -> Term:
+def default_rewrite(x: T | None, y: U) -> T | U:
     return x if x is not None else y
 
 
@@ -76,16 +80,15 @@ def __call__(self, x: Term) -> Term | None:
             if y.is_expr():
                 args = y.children()
                 return y.make_term(
-                    y.head(), [default_rewrite(self(arg), arg) for arg in args]
+                    *tuple(default_rewrite(self(arg), arg) for arg in args)
                 )
             return y
         if x.is_expr():
             args = x.children()
             new_args = list(map(self, args))
             if not all(arg is None for arg in new_args):
                 return x.make_term(
-                    x.head(),
-                    list(map(lambda x1, x2: default_rewrite(x1, x2), new_args, args)),
+                    *map(lambda x1, x2: default_rewrite(x1, x2), new_args, args),
                 )
             return None
         return None
@@ -111,8 +114,7 @@ def __call__(self, x: Term) -> Term | None:
             if all(arg is None for arg in new_args):
                 return self.rw(x)
             y = x.make_term(
-                x.head(),
-                list(map(lambda x1, x2: default_rewrite(x1, x2), new_args, args)),
+                *map(lambda x1, x2: default_rewrite(x1, x2), new_args, args),
             )
             return default_rewrite(self.rw(y), y)
         return self.rw(x)
@@ -179,9 +181,7 @@ def __call__(self, x: Term) -> Term | None:
         if y is not None:
             if y.is_expr():
                 y_args = y.children()
-                return y.make_term(
-                    y.head(), [default_rewrite(self(arg), arg) for arg in y_args]
-                )
+                return y.make_term(*(default_rewrite(self(arg), arg) for arg in y_args))
             return y
         return None
 

src/finch/symbolic/term.py

@@ -13,10 +13,13 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
-from typing import Any
+from typing import TYPE_CHECKING, Any, Self
 
 __all__ = ["Term", "PreOrderDFS", "PostOrderDFS"]
 
+if TYPE_CHECKING:
+    from ..finch_logic import LogicNode
+
 
 class Term(ABC):
     def __init__(self):
@@ -27,7 +30,7 @@ def head(self) -> Any:
         """Return the head type of the S-expression."""
 
     @abstractmethod
-    def children(self) -> list[Term]:
+    def children(self) -> list[LogicNode]:
         """Return the children (AKA tail) of the S-expression."""
 
     @abstractmethod
@@ -37,11 +40,11 @@ def is_expr(self) -> bool:
         `children()` if `True`."""
 
     @abstractmethod
-    def make_term(self, head: Any, children: list[Term]) -> Term:
+    def make_term(self, head: Any, *children: Term) -> Self:
         """
         Construct a new term in the same family of terms with the given head type and
         children. This function should satisfy
-        `x == x.make_term(x.head(), *x.children())`
+        `x == x.make_term(*x.children())`
         """
 
     def __hash__(self) -> int:
@@ -60,6 +63,7 @@ def __eq__(self, other: object) -> bool:
 
 def PostOrderDFS(node: Term) -> Iterator[Term]:
     if node.is_expr():
+        arg: Term
         for arg in node.children():
             yield from PostOrderDFS(arg)
     yield node
@@ -68,5 +72,6 @@ def PostOrderDFS(node: Term) -> Iterator[Term]:
 def PreOrderDFS(node: Term) -> Iterator[Term]:
     yield node
     if node.is_expr():
+        arg: Term
         for arg in node.children():
             yield from PreOrderDFS(arg)