Add Set dialect and support of ListComp & SetComp lowering

docs/dialects/python/data.md

@@ -30,3 +30,15 @@ This page provides a reference for dialects that bring in semantics for common P
         - "!statement"
         show_root_heading: true
         show_if_no_docstring: true
+
+### Set
+
+`set` support currently covers set literals and `set()` construction only. Mutation APIs
+and `set(iterable)` are not part of this first version.
+
+::: kirin.dialects.py.set
+    options:
+        filters:
+        - "!statement"
+        show_root_heading: true
+        show_if_no_docstring: true

src/kirin/dialects/func/typeinfer.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections.abc import Hashable
+
 from kirin import ir, types
 from kirin.interp import Frame, MethodTable, ReturnValue, impl
 from kirin.analysis import const
@@ -30,6 +32,7 @@ def return_(
         if (
             isinstance(hint := stmt.value.hints.get("const"), const.Value)
             and hint.data is not None
+            and isinstance(hint.data, Hashable)
         ):
             return ReturnValue(types.Literal(hint.data, frame.get(stmt.value)))
         return ReturnValue(frame.get(stmt.value))

src/kirin/dialects/ilist/lowering.py

@@ -3,6 +3,7 @@
 
 from kirin import types, lowering
 from kirin.dialects import py
+from kirin.dialects.py._comprehension import lower_listcomp_via_desugaring
 
 from . import stmts as ilist
 from ._dialect import dialect
@@ -23,6 +24,11 @@ def lower_List(self, state: lowering.State, node: ast.List) -> lowering.Result:
 
         return state.current_frame.push(ilist.New(values=tuple(elts), elem_type=typ))
 
+    def lower_ListComp(
+        self, state: lowering.State, node: ast.ListComp
+    ) -> lowering.Result:
+        return lower_listcomp_via_desugaring(state, node)
+
     @lowering.akin(ilist.IList)
     def lower_Call_IList(
         self, state: lowering.State, node: ast.Call

src/kirin/dialects/py/__init__.py

@@ -9,6 +9,7 @@
 from . import (
     cmp as cmp,
     len as len,
+    set as set,
     attr as attr,
     base as base,
     list as list,

src/kirin/dialects/py/_comprehension.py

@@ -0,0 +1,108 @@
+import ast
+
+from kirin import lowering
+
+_LISTCOMP_TMP_PREFIX = "_kirin_listcomp_tmp"
+_SETCOMP_TMP_PREFIX = "_kirin_setcomp_tmp"
+
+
+def lower_listcomp_via_desugaring(
+    state: lowering.State, node: ast.ListComp
+) -> lowering.Result:
+    tmp_name = fresh_comp_name(state, _LISTCOMP_TMP_PREFIX)
+    init = ast.List(elts=[], ctx=ast.Load())
+    leaf = ast.Assign(
+        targets=[ast.Name(id=tmp_name, ctx=ast.Store())],
+        value=ast.BinOp(
+            left=ast.Name(id=tmp_name, ctx=ast.Load()),
+            op=ast.Add(),
+            right=ast.List(elts=[node.elt], ctx=ast.Load()),
+        ),
+    )
+    fix_locations(leaf, node.elt)
+    return lower_comprehension_via_desugaring(
+        state=state,
+        tmp_name=tmp_name,
+        init_value=init,
+        generators=node.generators,
+        leaf_stmt=leaf,
+        ref_node=node,
+    )
+
+
+def lower_setcomp_via_desugaring(
+    state: lowering.State, node: ast.SetComp
+) -> lowering.Result:
+    tmp_name = fresh_comp_name(state, _SETCOMP_TMP_PREFIX)
+    init = ast.Call(func=ast.Name(id="set", ctx=ast.Load()), args=[], keywords=[])
+    leaf = ast.Assign(
+        targets=[ast.Name(id=tmp_name, ctx=ast.Store())],
+        value=ast.BinOp(
+            left=ast.Name(id=tmp_name, ctx=ast.Load()),
+            op=ast.BitOr(),
+            right=ast.Set(elts=[node.elt]),
+        ),
+    )
+    fix_locations(leaf, node.elt)
+    return lower_comprehension_via_desugaring(
+        state=state,
+        tmp_name=tmp_name,
+        init_value=init,
+        generators=node.generators,
+        leaf_stmt=leaf,
+        ref_node=node,
+    )
+
+
+def lower_comprehension_via_desugaring(
+    state: lowering.State,
+    tmp_name: str,
+    init_value: ast.expr,
+    generators: list[ast.comprehension],
+    leaf_stmt: ast.stmt,
+    ref_node: ast.AST,
+) -> lowering.Result:
+    init = ast.Assign(
+        targets=[ast.Name(id=tmp_name, ctx=ast.Store())],
+        value=init_value,
+    )
+    fix_locations(init, ref_node)
+    state.lower(init)
+
+    for stmt in build_comprehension_stmts(generators, leaf_stmt):
+        state.lower(stmt)
+
+    result = ast.Name(id=tmp_name, ctx=ast.Load())
+    fix_locations(result, ref_node)
+    return state.lower(result).expect_one()
+
+
+def fresh_comp_name(state: lowering.State, prefix: str) -> str:
+    frame = state.current_frame
+    idx = 0
+    while True:
+        suffix = "" if idx == 0 else f"_{idx}"
+        name = f"{prefix}{suffix}"
+        if frame.get_local(name) is None and name not in frame.globals:
+            return name
+        idx += 1
+
+
+def build_comprehension_stmts(
+    generators: list[ast.comprehension], leaf_stmt: ast.stmt
+) -> list[ast.stmt]:
+    acc = leaf_stmt
+    for gen in reversed(generators):
+        for if_ in reversed(gen.ifs):
+            acc = ast.If(test=if_, body=[acc], orelse=[])
+            fix_locations(acc, if_)
+
+        acc = ast.For(target=gen.target, iter=gen.iter, body=[acc], orelse=[])
+        fix_locations(acc, gen)
+
+    return [acc]
+
+
+def fix_locations(node: ast.AST, ref: ast.AST) -> None:
+    ast.copy_location(node, ref)
+    ast.fix_missing_locations(node)

src/kirin/dialects/py/list/lowering.py

@@ -1,14 +1,14 @@
 import ast
 
 from kirin import types, lowering
+from kirin.dialects.py._comprehension import lower_listcomp_via_desugaring
 
 from .stmts import New
 from ._dialect import dialect
 
 
 @dialect.register
 class PythonLowering(lowering.FromPythonAST):
-
     def lower_List(self, state: lowering.State, node: ast.List) -> lowering.Result:
         elts = tuple(state.lower(each).expect_one() for each in node.elts)
 
@@ -20,3 +20,8 @@ def lower_List(self, state: lowering.State, node: ast.List) -> lowering.Result:
             typ = types.Any
 
         return state.current_frame.push(New(values=tuple(elts)))
+
+    def lower_ListComp(
+        self, state: lowering.State, node: ast.ListComp
+    ) -> lowering.Result:
+        return lower_listcomp_via_desugaring(state, node)

src/kirin/dialects/py/set/__init__.py

@@ -0,0 +1,12 @@
+"""The set dialect for Python.
+
+This module contains the dialect for set semantics in Python, including:
+
+- The `New` statement class.
+- The lowering pass for set literals and `set()`.
+- The concrete implementation of set operations.
+"""
+
+from . import interp as interp, lowering as lowering, typeinfer as typeinfer
+from .stmts import New as New
+from ._dialect import dialect as dialect

src/kirin/dialects/py/set/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("py.set")

src/kirin/dialects/py/set/interp.py

@@ -0,0 +1,12 @@
+from kirin import interp
+
+from .stmts import New
+from ._dialect import dialect
+
+
+@dialect.register
+class SetMethods(interp.MethodTable):
+
+    @interp.impl(New)
+    def new(self, interp, frame: interp.Frame, stmt: New):
+        return (set(frame.get_values(stmt.values)),)

src/kirin/dialects/py/set/lowering.py

@@ -0,0 +1,27 @@
+import ast
+
+from kirin import lowering
+from kirin.dialects.py._comprehension import lower_setcomp_via_desugaring
+
+from .stmts import New
+from ._dialect import dialect
+
+
+@dialect.register
+class PythonLowering(lowering.FromPythonAST):
+
+    def lower_Set(self, state: lowering.State, node: ast.Set) -> lowering.Result:
+        return state.current_frame.push(
+            New(tuple(state.lower(each).expect_one() for each in node.elts))
+        )
+
+    def lower_SetComp(
+        self, state: lowering.State, node: ast.SetComp
+    ) -> lowering.Result:
+        return lower_setcomp_via_desugaring(state, node)
+
+    @lowering.akin(set)
+    def lower_Call_set(self, state: lowering.State, node: ast.Call) -> lowering.Result:
+        if len(node.args) != 0 or len(node.keywords) != 0:
+            raise lowering.BuildError("`set(iterable)` is not supported yet")
+        return state.current_frame.push(New(()))

src/kirin/dialects/py/set/stmts.py

@@ -0,0 +1,29 @@
+from typing import Sequence
+
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from ._dialect import dialect
+
+T = types.TypeVar("T")
+
+
+@statement(dialect=dialect, init=False)
+class New(ir.Statement):
+    name = "set"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    values: tuple[ir.SSAValue, ...] = info.argument(T)
+    result: ir.ResultValue = info.result(types.Set[T])
+
+    def __init__(self, values: Sequence[ir.SSAValue]) -> None:
+        elem_type: types.TypeAttribute = types.Any
+        if values:
+            elem_type = values[0].type
+            for value in values[1:]:
+                elem_type = elem_type.join(value.type)
+
+        super().__init__(
+            args=values,
+            result_types=(types.Set[elem_type],),
+            args_slice={"values": slice(0, len(values))},
+        )

src/kirin/dialects/py/set/typeinfer.py

@@ -0,0 +1,12 @@
+from kirin import interp
+
+from .stmts import New
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(New)
+    def new(self, interp, frame, stmt: New):
+        return (stmt.result.type,)

src/kirin/dialects/scf/typeinfer.py

@@ -10,7 +10,6 @@
 
 @dialect.register(key="typeinfer")
 class TypeInfer(absint.Methods):
-
     @interp.impl(IfElse)
     def if_else_(
         self,

src/kirin/prelude.py

@@ -11,6 +11,7 @@
 from kirin.dialects.py import (
     cmp,
     len,
+    set,
     attr,
     base,
     list,
@@ -60,6 +61,7 @@ def run_pass(mt: Method) -> None:
     python_basic.union(
         [
             list,
+            set,
             slice,
             cf,
             func,
@@ -85,6 +87,7 @@ def run_pass(mt: Method) -> None:
     python_basic.union(
         [
             ilist,
+            set,
             slice,
             cf,
             func,
@@ -179,6 +182,7 @@ def run_pass(
     python_basic.union(
         [
             ilist,
+            set,
             slice,
             scf,
             cf,
@@ -204,6 +208,7 @@ def run_pass(method: Method) -> None:
     python_basic.union(
         [
             ilist,
+            set,
             slice,
             scf,
             cf,