Move types into rust

.flake8

@@ -1,4 +1,4 @@
 ; https://black.readthedocs.io/en/stable/guides/using_black_with_other_tools.html#flake8
 [flake8]
 max-line-length = 88
-extend-ignore = E203,E501,F405,F403,E302
+extend-ignore = E203,E501,F405,F403,E302,E305,F821

.github/workflows/CI.yml

@@ -38,7 +38,7 @@ jobs:
           python-version: "3.10"
       - uses: actions/checkout@v2
       - run: pip install -e . mypy
-      - run: python -m mypy.stubtest egg_smol.bindings
+      - run: python -m mypy.stubtest egg_smol.bindings  --allowlist stubtest_allow
   docs:
     runs-on: ubuntu-latest
     steps:

.pre-commit-config.yaml

@@ -18,7 +18,3 @@ repos:
     rev: 5.0.4
     hooks:
       - id: flake8
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v0.982"
-    hooks:
-      - id: mypy

Cargo.toml

@@ -10,7 +10,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 pyo3 = { version = "0.17.1", features = ["extension-module"] }
-egg-smol = { git = "https://github.com/saulshanabrook/egg-smol", branch = "public-api" }
+egg-smol = { git = "https://github.com/mwillsey/egg-smol", ref = "9a45bdee9b5395ab821318a033c6ac7fa91f91b9" }
 
 [package.metadata.maturin]
 name = "egg_smol.bindings"

docs/explanation/compared_to_rust.md

@@ -52,6 +52,8 @@ One way to run this in Python is to parse the text and run it similar to how the
 egg CLI works:
 
 ```{code-cell} python
+from egg_smol.bindings import *
+
 eqsat_basic = """(datatype Math
   (Num i64)
   (Var String)
@@ -78,8 +80,6 @@ eqsat_basic = """(datatype Math
 (run 10)
 (check (= expr1 expr2))"""
 
-from egg_smol.bindings import EGraph
-
 egraph = EGraph()
 egraph.parse_and_run_program(eqsat_basic)
 ```
@@ -90,8 +90,6 @@ However, this isn't the most friendly for Python users. Instead, we can use the
 low level APIs that mirror the rust APIs to build the same egraph:
 
 ```{code-cell} python
-from egg_smol.bindings_py import *
-
 egraph = EGraph()
 egraph.declare_sort("Math")
 egraph.declare_constructor(Variant("Num", ["i64"]), "Math")

python/egg_smol/bindings.pyi

@@ -3,19 +3,147 @@ from typing import Optional
 
 from typing_extensions import final
 
-from .bindings_py import Expr, Fact_, FunctionDecl, Rewrite, Variant
-
 @final
 class EGraph:
     def parse_and_run_program(self, input: str) -> list[str]: ...
     def declare_constructor(self, variant: Variant, sort: str) -> None: ...
     def declare_sort(self, name: str) -> None: ...
     def declare_function(self, decl: FunctionDecl) -> None: ...
-    def define(self, name: str, expr: Expr, cost: Optional[int] = None) -> None: ...
+    def define(self, name: str, expr: _Expr, cost: Optional[int] = None) -> None: ...
     def add_rewrite(self, rewrite: Rewrite) -> str: ...
     def run_rules(self, limit: int) -> tuple[timedelta, timedelta, timedelta]: ...
-    def check_fact(self, fact: Fact_) -> None: ...
+    def check_fact(self, fact: _Fact) -> None: ...
 
 @final
 class EggSmolError(Exception):
     context: str
+
+@final
+class Int:
+    def __init__(self, value: int) -> None: ...
+    value: int
+
+@final
+class String:
+    def __init__(self, value: str) -> None: ...
+    value: str
+
+@final
+class Unit:
+    def __init__(self) -> None: ...
+
+_Literal = Int | String | Unit
+
+@final
+class Lit:
+    def __init__(self, value: _Literal) -> None: ...
+    value: _Literal
+
+@final
+class Var:
+    def __init__(self, name: str) -> None: ...
+    name: str
+
+@final
+class Call:
+    def __init__(self, name: str, args: list[_Expr]) -> None: ...
+    name: str
+    args: list[_Expr]
+
+_Expr = Lit | Var | Call
+
+@final
+class Eq:
+    def __init__(self, exprs: list[_Expr]) -> None: ...
+    exprs: list[_Expr]
+
+@final
+class Fact:
+    def __init__(self, expr: _Expr) -> None: ...
+    expr: _Expr
+
+_Fact = Fact | Eq
+
+@final
+class Define:
+    def __init__(self, lhs: str, rhs: _Expr) -> None: ...
+    lhs: str
+    rhs: _Expr
+
+@final
+class Set:
+    def __init__(self, lhs: str, args: list[_Expr], rhs: _Expr) -> None: ...
+    lhs: str
+    args: list[_Expr]
+    rhs: _Expr
+
+@final
+class Delete:
+    sym: str
+    args: list[_Expr]
+    def __init__(self, sym: str, args: list[_Expr]) -> None: ...
+
+@final
+class Union:
+    def __init__(self, lhs: _Expr, rhs: _Expr) -> None: ...
+    lhs: _Expr
+    rhs: _Expr
+
+@final
+class Panic:
+    def __init__(self, msg: str) -> None: ...
+    msg: str
+
+@final
+class Expr_:
+    def __init__(self, expr: _Expr) -> None: ...
+    expr: _Expr
+
+_Action = Define | Set | Delete | Union | Panic | Expr_
+
+@final
+class FunctionDecl:
+    name: str
+    schema: Schema
+    default: Optional[_Expr]
+    merge: Optional[_Expr]
+    cost: Optional[int]
+    def __init__(
+        self,
+        name: str,
+        schema: Schema,
+        default: Optional[_Expr],
+        merge: Optional[_Expr],
+        cost: Optional[int] = None,
+    ) -> None: ...
+
+@final
+class Variant:
+    def __init__(
+        self, name: str, types: list[str], cost: Optional[int] = None
+    ) -> None: ...
+    name: str
+    types: list[str]
+    cost: Optional[int]
+
+@final
+class Schema:
+    input: list[str]
+    output: str
+    def __init__(self, input: list[str], output: str) -> None: ...
+
+@final
+class Rule:
+    head: list[_Action]
+    body: list[_Fact]
+    def __init__(self, head: list[_Action], body: list[_Fact]) -> None: ...
+
+@final
+class Rewrite:
+    lhs: _Expr
+    rhs: _Expr
+    conditions: list[_Fact]
+
+    def __init__(
+        self, lhs: _Expr, rhs: _Expr, conditions: list[_Fact] = []
+    ) -> None: ...

python/tests/test.py

@@ -1,8 +1,7 @@
 import datetime
 
 import pytest
-from egg_smol.bindings import EggSmolError, EGraph
-from egg_smol.bindings_py import *
+from egg_smol.bindings import *
 
 
 class TestEGraph:
@@ -334,3 +333,29 @@ def test_check_fact(self):
                 ]
             )
         )
+
+    # def test_extract(self):
+    #     # Example from extraction-cost
+    #     egraph = EGraph()
+    #     egraph.declare_sort("Expr")
+    #     egraph.declare_constructor(Variant("Num", ["i64"], cost=5), "Expr")
+
+    #     egraph.define("x", Call("Num", [Lit(Int(1))]), cost=10)
+    #     egraph.define("y", Call("Num", [Lit(Int(2))]), cost=1)
+
+    #     assert egraph.extract("x") == Call("Num", [Lit(Int(1))])
+    #     assert egraph.extract("y") == Var("y")
+
+
+class TestVariant:
+    def test_repr(self):
+        assert repr(Variant("name", [])) == "Variant('name', [], None)"
+
+    def test_name(self):
+        assert Variant("name", []).name == "name"
+
+    def test_types(self):
+        assert Variant("name", ["a", "b"]).types == ["a", "b"]
+
+    def test_cost(self):
+        assert Variant("name", [], cost=1).cost == 1