Merge branch 'main' into wma/scheduler

Author: willow-ahrens

src/finchlite/symbolic/__init__.py

@@ -11,12 +11,11 @@
 )
 from .stage import Stage
 from .term import (
-    PostOrderDFS,
-    PreOrderDFS,
     Term,
     TermTree,
     literal_repr,
 )
+from .traversal import PostOrderDFS, PreOrderDFS, intree, isdescendant
 
 __all__ = [
     "BasicBlock",
@@ -42,5 +41,7 @@
     "fisinstance",
     "ftype",
     "gensym",
+    "intree",
+    "isdescendant",
     "literal_repr",
 ]

src/finchlite/symbolic/environment.py

@@ -3,7 +3,8 @@
 from collections import defaultdict
 from typing import Any, Generic, Optional, TypeVar
 
-from .term import PostOrderDFS, Term
+from .term import Term
+from .traversal import PostOrderDFS
 
 
 class NamedTerm(Term, ABC):

src/finchlite/symbolic/term.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Iterator
 from dataclasses import dataclass
 from inspect import isbuiltin, isclass, isfunction
 from typing import Any, Self
@@ -82,17 +81,3 @@ def literal_repr(name: str, fields: dict[str, Any]) -> str:
     return (
         name + "(" + ", ".join([f"{k}={_get_repr(v)}" for k, v in fields.items()]) + ")"
     )
-
-
-def PostOrderDFS(node: Term) -> Iterator[Term]:
-    if isinstance(node, TermTree):
-        for arg in node.children:
-            yield from PostOrderDFS(arg)
-    yield node
-
-
-def PreOrderDFS(node: Term) -> Iterator[Term]:
-    yield node
-    if isinstance(node, TermTree):
-        for arg in node.children:
-            yield from PreOrderDFS(arg)

src/finchlite/symbolic/traversal.py

@@ -0,0 +1,33 @@
+from collections.abc import Iterator
+
+from finchlite.symbolic.term import Term, TermTree
+
+
+def PostOrderDFS(node: Term) -> Iterator[Term]:
+    if isinstance(node, TermTree):
+        for arg in node.children:
+            yield from PostOrderDFS(arg)
+    yield node
+
+
+def PreOrderDFS(node: Term) -> Iterator[Term]:
+    yield node
+    if isinstance(node, TermTree):
+        for arg in node.children:
+            yield from PreOrderDFS(arg)
+
+
+def intree(n1, n2):
+    """
+    Return True iff `n1` occurs in the subtree rooted at `n2`.
+    """
+    return any(node == n1 for node in PostOrderDFS(n2))
+
+
+def isdescendant(n1, n2):
+    """
+    True iff `n1` is a strict descendant of `n2`.
+    """
+    if n1 == n2:
+        return False
+    return intree(n1, n2)

tests/test_traversal.py

@@ -0,0 +1,113 @@
+from collections import Counter
+
+from finchlite.finch_logic import Field, Literal, MapJoin, Plan, Produces, Table
+from finchlite.symbolic import PostOrderDFS, PreOrderDFS, intree, isdescendant
+
+
+def test_preorder_dfs():
+    ta = Table(
+        Literal("A"),
+        (Field("i"), Field("j")),
+    )
+
+    tb = Table(
+        Literal("B"),
+        (Field("j"), Field("k")),
+    )
+
+    prog = Plan(
+        (
+            Produces(
+                (
+                    MapJoin(
+                        Field("op"),
+                        (ta, tb),
+                    ),
+                ),
+            ),
+        )
+    )
+
+    preorder = list(PreOrderDFS(prog))
+
+    assert Counter(type(x).__name__ for x in preorder) == Counter(
+        {"Plan": 1, "Produces": 1, "MapJoin": 1, "Table": 2, "Literal": 2, "Field": 5}
+    )
+
+    pos = {}
+    for i, obj in enumerate(preorder):
+        k = id(obj)
+        if k in pos:
+            continue
+        pos[k] = i
+    for node in preorder:
+        for child in getattr(node, "children", ()):
+            assert pos[id(node)] < pos[id(child)]
+
+
+def test_postorder_dfs():
+    ta = Table(
+        Literal("A"),
+        (Field("i"), Field("j")),
+    )
+
+    tb = Table(
+        Literal("B"),
+        (Field("j"), Field("k")),
+    )
+
+    prog = Plan(
+        (
+            Produces(
+                (
+                    MapJoin(
+                        Field("op"),
+                        (ta, tb),
+                    ),
+                ),
+            ),
+        )
+    )
+
+    postorder = list(PostOrderDFS(prog))
+
+    assert Counter(type(x).__name__ for x in postorder) == Counter(
+        {"Plan": 1, "Produces": 1, "MapJoin": 1, "Table": 2, "Literal": 2, "Field": 5}
+    )
+
+    pos = {}
+    for i, obj in enumerate(postorder):
+        k = id(obj)
+        if k in pos:
+            continue
+        pos[k] = i
+    for node in postorder:
+        for child in getattr(node, "children", ()):
+            assert pos[id(child)] < pos[id(node)]
+
+
+def test_intree():
+    i, j, k = Field("i"), Field("j"), Field("k")
+    ta = Table(Literal("A"), (i, j))
+    tb = Table(Literal("B"), (j, k))
+    op = Field("op")
+    mj = MapJoin(op, (ta, tb))
+    prog = Plan((Produces((mj,)),))
+
+    assert intree(prog, prog)
+    assert intree(mj, prog)
+    assert intree(ta, prog)
+    assert intree(tb, prog)
+
+
+def test_isdescendant():
+    i, j, k = Field("i"), Field("j"), Field("k")
+    ta = Table(Literal("A"), (i, j))
+    tb = Table(Literal("B"), (j, k))
+    op = Field("op")
+    mj = MapJoin(op, (ta, tb))
+    prog = Plan((Produces((mj,)),))
+
+    assert isdescendant(mj, prog)
+    assert isdescendant(ta, prog)
+    assert isdescendant(tb, prog)