diff --git a/Cargo.lock b/Cargo.lock index 83948d73..b53d452c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -350,7 +350,7 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "egglog" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "csv", "dyn-clone", @@ -359,6 +359,7 @@ dependencies = [ "egglog-bridge", "egglog-core-relations", "egglog-numeric-id", + "egglog-proto", "egglog-reports", "egraph-serialize", "enum-map", @@ -378,7 +379,7 @@ dependencies = [ [[package]] name = "egglog-add-primitive" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "quote", "syn 2.0.119", @@ -387,7 +388,7 @@ dependencies = [ [[package]] name = "egglog-ast" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "ordered-float", ] @@ -395,7 +396,7 @@ dependencies = [ [[package]] name = "egglog-bridge" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "anyhow", "dyn-clone", @@ -418,7 +419,7 @@ dependencies = [ [[package]] name = "egglog-concurrency" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "arc-swap", "bumpalo", @@ -430,7 +431,7 @@ dependencies = [ [[package]] name = "egglog-core-relations" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "anyhow", "bumpalo", @@ -458,7 +459,7 @@ dependencies = [ [[package]] name = "egglog-experimental" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=3d35f37a4befbfa75608aa6998408334e19b6180#3d35f37a4befbfa75608aa6998408334e19b6180" +source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=7a277e7fb217b30794b0bb2995ecc158983c5692#7a277e7fb217b30794b0bb2995ecc158983c5692" dependencies = [ "egglog", "egglog-ast", @@ -468,18 +469,27 @@ dependencies = [ "lazy_static", "log", "num", + "prost", "stacker", ] [[package]] name = "egglog-numeric-id" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" + +[[package]] +name = "egglog-proto" +version = "0.1.0" +source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=6e0dbc97fbe9f3868c4662b9a6d73b5d01913c9d#6e0dbc97fbe9f3868c4662b9a6d73b5d01913c9d" +dependencies = [ + "prost", +] [[package]] name = "egglog-reports" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "clap", "hashbrown 0.16.1", @@ -493,7 +503,7 @@ dependencies = [ [[package]] name = "egglog-union-find" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog?rev=90635860397ce710f8c0a4eeb04154a8ebc3ac05#90635860397ce710f8c0a4eeb04154a8ebc3ac05" +source = "git+https://github.com/egraphs-good/egglog?rev=998e60380e254d00b29651e88f2065fb33851a2f#998e60380e254d00b29651e88f2065fb33851a2f" dependencies = [ "crossbeam", "egglog-concurrency", diff --git a/Cargo.toml b/Cargo.toml index 348c7071..219356fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,12 +19,12 @@ opentelemetry = "0.32" opentelemetry-otlp = { version = "0.32", features = ["http-proto", "reqwest-blocking-client", "trace"] } opentelemetry-stdout = { version = "0.32", features = ["trace"] } opentelemetry_sdk = "0.32" -egglog = { git = "https://github.com/egraphs-good/egglog", rev = "90635860397ce710f8c0a4eeb04154a8ebc3ac05", default-features = false } -egglog-ast = { git = "https://github.com/egraphs-good/egglog", rev = "90635860397ce710f8c0a4eeb04154a8ebc3ac05" } -egglog-core-relations = { git = "https://github.com/egraphs-good/egglog", rev = "90635860397ce710f8c0a4eeb04154a8ebc3ac05" } -egglog-reports = { git = "https://github.com/egraphs-good/egglog", rev = "90635860397ce710f8c0a4eeb04154a8ebc3ac05" } -egglog-bridge = { git = "https://github.com/egraphs-good/egglog", rev = "90635860397ce710f8c0a4eeb04154a8ebc3ac05" } -egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "3d35f37a4befbfa75608aa6998408334e19b6180", default-features = false } +egglog = { git = "https://github.com/egraphs-good/egglog", rev = "998e60380e254d00b29651e88f2065fb33851a2f", default-features = false } +egglog-ast = { git = "https://github.com/egraphs-good/egglog", rev = "998e60380e254d00b29651e88f2065fb33851a2f" } +egglog-core-relations = { git = "https://github.com/egraphs-good/egglog", rev = "998e60380e254d00b29651e88f2065fb33851a2f" } +egglog-reports = { git = "https://github.com/egraphs-good/egglog", rev = "998e60380e254d00b29651e88f2065fb33851a2f" } +egglog-bridge = { git = "https://github.com/egraphs-good/egglog", rev = "998e60380e254d00b29651e88f2065fb33851a2f" } +egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "7a277e7fb217b30794b0bb2995ecc158983c5692", default-features = false } egraph-serialize = { version = "0.3", features = ["serde", "graphviz"] } serde_json = "1" pyo3-log = "*" diff --git a/pyproject.toml b/pyproject.toml index 52441f60..27cf9901 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,15 @@ classifiers = [ "Topic :: Software Development :: Interpreters", "Typing :: Typed", ] -dependencies = ["typing-extensions>=4.13", "black", "graphviz", "anywidget", "cloudpickle>=3", "opentelemetry-api"] +dependencies = [ + "typing-extensions>=4.13", + "black", + "graphviz", + "anywidget", + "cloudpickle>=3", + "opentelemetry-api", + "egglog-proto @ git+https://github.com/egraphs-good/egglog-experimental.git@6e0dbc97fbe9f3868c4662b9a6d73b5d01913c9d#subdirectory=gen/python", +] [project.optional-dependencies] @@ -269,6 +277,7 @@ local_partial_types = true python-source = "python" module-name = "egglog.bindings" features = ["pyo3/extension-module"] +include = ["egglog/_builtin_catalog.pb"] [tool.pytest.ini_options] addopts = ["--import-mode=importlib", "--doctest-modules"] diff --git a/python/egglog/_builtin_catalog.pb b/python/egglog/_builtin_catalog.pb new file mode 100644 index 00000000..76562131 Binary files /dev/null and b/python/egglog/_builtin_catalog.pb differ diff --git a/python/egglog/_catalog.py b/python/egglog/_catalog.py new file mode 100644 index 00000000..8ed39d44 --- /dev/null +++ b/python/egglog/_catalog.py @@ -0,0 +1,130 @@ +"""Derived Python presentation indexes over canonical declaration records.""" + +from __future__ import annotations + +import keyword +from collections.abc import Iterable, Sequence +from functools import cache +from importlib.resources import files + +from egglog_proto.egglog.v1 import egglog_pb as pb + +from ._program import Builder, Ref, _definition_key + + +@cache +def builtin_catalog() -> Catalog: + """Decode the native-generated package snapshot, without creating an engine.""" + path = files("egglog").joinpath("_builtin_catalog.pb") + try: + encoded = path.read_bytes() + except FileNotFoundError as exc: + msg = "The native-generated builtin catalog is missing; regenerate the package resource" + raise RuntimeError(msg) from exc + return Catalog(pb.Program.from_binary(encoded)) + + +def _python_path(path: Sequence[str]) -> tuple[str, ...]: + if not path or any(not part.isidentifier() or keyword.iskeyword(part) for part in path): + raise ValueError(f"Invalid Python path: {path!r}") + return tuple(path) + + +class Catalog: + """ + Own one catalog snapshot and rebuildable indexes, never copied signatures. + + These indexes contain explicit supplied bindings. Absent Python bindings + require conservative wrapper derivation at generation, without installing + synthetic bindings; explicit/derived collisions are generation errors. + Engine installation and compatibility remain the shared adapter's job. + """ + + def __init__(self, program: pb.Program) -> None: + self.owner = Builder.from_program(program).publish() + self.roots = tuple(self.owner.ref("declarations", index) for index in range(len(program.declarations))) + self.reindex() + + @classmethod + def from_refs(cls, declarations: Iterable[Ref]) -> Catalog: + """Index already-published definitions without copying their owners/identity.""" + result = object.__new__(cls) + result.roots = tuple(dict.fromkeys(ref.owner.ref(ref.role, ref.index) for ref in declarations)) + if any(reference.role != "declarations" for reference in result.roots): + msg = "Catalog roots must be canonical declaration references" + raise TypeError(msg) + result.reindex() + return result + + def reindex(self) -> None: # noqa: C901, PLR0912 + """Discard all language indexes and derive them solely from owned records.""" + definitions: dict[tuple[str, str], Ref] = {} + types: dict[tuple[str, ...], Ref] = {} + functions: dict[tuple[str, ...], tuple[Ref, int]] = {} + members: dict[tuple[str | None, str], tuple[Ref, int]] = {} + for ref in self.roots: + declaration: pb.Declaration = ref.read() + key = _definition_key(declaration) + assert key is not None + if key in definitions: + raise ValueError(f"Duplicate catalog definition: {key}") + definitions[key] = ref + assert declaration.kind is not None + kind = declaration.kind.value + if isinstance(kind, pb.HostSortFamily | pb.EqSort): + binding = kind.bindings + if binding is not None and binding.python is not None: + path = _python_path(binding.python.path) + if path in types or path in functions: + raise ValueError(f"Python binding collision: {path}") + types[path] = ref + continue + binding = declaration.bindings + if binding is None or binding.python is None: + continue + for view_index, view in enumerate(binding.python.views): + if view.kind == pb.PythonCallKind.FUNCTION: + path = _python_path(view.path) + if view.owner is not None or view.has_field("receiver"): + msg = "A free Python function cannot have an owner or receiver" + raise ValueError(msg) + if path in types or path in functions: + raise ValueError(f"Python binding collision: {path}") + functions[path] = ref, view_index + continue + if view.owner is None or view.owner.kind is None: + msg = "Python member requires an owner" + raise ValueError(msg) + owner_name = None + if view.owner.kind.field == "sort": + owner_sort: pb.Sort = ref.owner.ref("sorts", view.owner.kind.value).read() + if owner_sort.kind is None: + msg = "Python member owner has no sort kind" + raise ValueError(msg) + match owner_sort.kind.field: + case "eq": + owner_name = owner_sort.kind.value + case "family": + owner_name = owner_sort.kind.value.name + case _: + msg = "Python member requires an equality or host family owner" + raise ValueError(msg) + if view.kind == pb.PythonCallKind.INITIALIZER: + if view.path: + msg = "Python initializer derives its name and cannot supply a path" + raise ValueError(msg) + name = "__init__" + else: + path = _python_path(view.path) + if len(path) != 1: + msg = "Python member path must contain exactly one name" + raise ValueError(msg) + name = path[0] + member = owner_name, name + if member in members: + raise ValueError(f"Python binding collision: {member}") + members[member] = ref, view_index + self.definitions = definitions + self.types = types + self.functions = functions + self.members = members diff --git a/python/egglog/_program.py b/python/egglog/_program.py new file mode 100644 index 00000000..636b3168 --- /dev/null +++ b/python/egglog/_program.py @@ -0,0 +1,614 @@ +""" +Private ownership and relocation for directly authored protobuf records. + +Fragments have local address slots, some importing records from other owners. +Only generated message fields select or order edges. Published records are never +mutated; inspection and packing return detached copies. This module does not +typecheck, evaluate, expand defaults, rename binders, or install providers. +""" + +from __future__ import annotations + +from collections.abc import Callable, Collection, Iterable, Mapping +from dataclasses import dataclass +from math import isnan +from struct import unpack +from threading import RLock +from types import MappingProxyType +from typing import Any, cast +from weakref import WeakKeyDictionary + +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import DescFieldValueList, DescFieldValueMessage, DescMessage, Message, ScalarType + +# Every Program-reachable field is deliberately listed. ':' means an arena +# address; '@' means a named definition dependency. Other fields are copied, +# recursively for inline messages. Binder positions are deliberately NOT indices. +_FIELDS = { + "Program": "ir_version nodes sorts declarations commands files rulesets rules", + "Node": "sort_id:sorts var primitive_value call union get_cost span", + "Sort": "eq@sort family func var span", + "HostSort": "name@sort args:sorts", + "FuncSort": "params:sorts result:sorts", + "Span": "file:files start end", + "SourceFile": "name contents", + "Union": "members:nodes", + "Call": "func@callable args:nodes", + "Lambda": "captures:nodes body:nodes", + "PrimitiveValue": "i64 f64_bits string bool unit big_int big_rat rational vec set multiset map pair maybe lambda partial_call custom", + "BigRat": "numerator denominator", + "Rational": "numerator denominator", + "ValueList": "items:nodes", + "MapValue": "entries", + "MapEntry": "key:nodes value:nodes", + "PairValue": "first:nodes second:nodes", + "MaybeValue": "value:nodes", + "CustomValue": "payload args:nodes", + "Unit": "", + "Action": "term:nodes set delete subsume panic set_cost span", + "Set": "target value:nodes", + "SetCost": "target cost:nodes", + "Declaration": "constructor function relation primitive host_sort_family host_primitive eq_sort span bindings doc", + "Arg": "sort:sorts name", + "EqSort": "name bindings", + "Constructor": "name inputs output:sorts cost:nodes unextractable", + "Function": "name inputs output:sorts merge:nodes", + "Relation": "name inputs", + "Primitive": "name inputs output:sorts body:nodes", + "HostSortFamily": "name arity bindings", + "HostPrimitive": "name signature application", + "GenericSignature": "type_params inputs output:sorts varargs", + "FunctionApplication": "", + "SortBindings": "python rust egglog", + "TypeBinding": "path type_params", + "EgglogTypeBinding": "symbol", + "CallableBindings": "python rust egglog", + "BindingOwner": "sort:sorts function", + "FunctionTypeOwner": "", + "PythonBindings": "views", + "PythonCallable": "kind path owner receiver params mutates", + "PythonParameter": "core_input name default_expr:nodes", + "RustBindings": "views", + "RustCallable": "path owner borrowed_self receiver params trait_impl", + "RustReceiver": "core_input borrowed", + "RustParameter": "core_input name borrowed", + "RustTrait": "path args output_associated_type", + "RustType": "sort:sorts borrowed", + "EgglogBindings": "views", + "EgglogCallable": "symbol datatype_member", + "Ruleset": "name rules combined span doc", + "RuleList": "rules:rules", + "RulesetRef": "index:rulesets name@ruleset", + "RuleDecl": "rule rewrite birewrite name eval_mode no_decomp include_subsumed span doc", + "Rule": "query:nodes head", + "Rewrite": "lhs:nodes rhs:nodes conditions:nodes subsume", + "BiRewrite": "lhs:nodes rhs:nodes conditions:nodes", + "CombinedRuleset": "rulesets", + "Command": "action check prove prove_exists run bind_scheduler keep_best extract freeze print_size print_function print_table_stats repeat saturate span", + "Check": "facts:nodes", + "Prove": "facts:nodes", + "ProveExists": "constructor@callable", + "Run": "ruleset scheduler", + "BindScheduler": "name scheduler", + "Scheduler": "back_off", + "BackOff": "match_limit ban_length node_limit", + "KeepBest": "tables@callable extractor cost_model", + "Extract": "roots:nodes variants extractor cost_model", + "Freeze": "", + "PrintSize": "tables@callable", + "PrintFunction": "table@callable max_rows", + "PrintTableStats": "tables@callable", + "Repeat": "body until:nodes times", + "Saturate": "body until:nodes max_iterations", +} +_ROLES: dict[str, type[Message]] = { + "nodes": pb.Node, + "sorts": pb.Sort, + "declarations": pb.Declaration, + "rules": pb.RuleDecl, + "rulesets": pb.Ruleset, + "files": pb.SourceFile, + "commands": pb.Command, +} +_FIELD_ROLES = { + name: dict(token.replace("@", ":@").partition(":")[::2] for token in fields.split()) + for name, fields in _FIELDS.items() +} + + +def _check_inventory() -> None: + """Fail on schema drift instead of guessing which new integers are addresses.""" + pending = [pb.Program.desc()] + seen: set[str] = set() + while pending: + descriptor = pending.pop() + if descriptor.name in seen: + continue + seen.add(descriptor.name) + expected = _FIELD_ROLES.get(descriptor.name) + if expected is None or set(expected) != {field.name for field in descriptor.fields}: + raise NotImplementedError(f"Unclassified protobuf fields in {descriptor.type_name}") + for field in descriptor.fields: + role = expected[field.name] + if role and not role.startswith("@"): + value = field.value + scalar = value.element if isinstance(value, DescFieldValueList) else getattr(value, "scalar", None) + if role not in _ROLES or scalar != ScalarType.UINT32: + raise TypeError(f"Invalid arena field classification: {descriptor.name}.{field.name}") + if isinstance(field.value, DescFieldValueMessage): + pending.append(field.value.message) + elif isinstance(field.value, DescFieldValueList) and isinstance(field.value.element, DescMessage): + pending.append(field.value.element) + if seen != set(_FIELDS): + raise NotImplementedError(f"Stale protobuf inventory entries: {set(_FIELDS) - seen}") + + +def _index_fields() -> list[tuple[type[Message], str, str]]: + """Expose the audited address inventory for exhaustive relocation tests.""" + return [ + (getattr(pb, name), field, role) + for name, fields in _FIELD_ROLES.items() + for field, role in fields.items() + if role and not role.startswith("@") + ] + + +def _copy_record( # noqa: C901, PLR0912 + record: Message, + *, + relocate: Callable[[str, int], int] | None = None, + dependency: Callable[[str, str], None] | None = None, +) -> Any: + """ + Detach inline messages iteratively, optionally relocating arena references. + + Inline message cycles cannot be serialized. Arena cycles remain ordinary + integers and are handled by pack's allocation-before-traversal invariant. + """ + result = type(record)() + pending = [(record, result, False)] + active: set[int] = set() + while pending: + source, target, leaving = pending.pop() + if leaving: + active.remove(id(source)) + continue + if id(source) in active: + msg = "protobuf inline message cycle" + raise ValueError(msg) + active.add(id(source)) + pending.append((source, target, True)) + descriptor = source.desc() + fields = _FIELD_ROLES.get(descriptor.name) + if descriptor.type_name != f"egglog.v1.{descriptor.name}" or fields is None: + raise NotImplementedError(f"Unsupported protobuf record {descriptor.type_name}") + # Unknown fields might contain indices. Preserving them without knowing + # their role would silently emit dangling addresses after relocation. + if source._unknown_fields: + msg = "Cannot relocate unknown protobuf fields" + raise NotImplementedError(msg) + children = [] + for field in descriptor.fields: + if field.name not in fields: + raise NotImplementedError(f"Unclassified protobuf field {descriptor.name}.{field.name}") + if field.presence.name != "IMPLICIT" and field not in source: + continue + value = source[field] + if value is None: + continue + repeated = isinstance(field.value, DescFieldValueList) + values = value if repeated else [value] + transformed = [] + role = fields[field.name] + for item in values: + if not isinstance(item, Message | str | int | float | bool | bytes): + raise TypeError(f"Unsupported mutable/non-scalar field {descriptor.name}.{field.name}") + replacement = item + if role.startswith("@"): + if not isinstance(item, str): + raise TypeError(f"Invalid definition name {descriptor.name}.{field.name}") + if dependency is not None: + dependency(role[1:], item) + elif role: + if isinstance(item, bool) or not isinstance(item, int) or not 0 <= item < 2**32: + raise TypeError(f"Invalid arena index {descriptor.name}.{field.name}") + if relocate is not None: + replacement = relocate(role, item) + elif isinstance(item, Message): + detached = type(item)() + children.append((item, detached, False)) + replacement = detached + transformed.append(replacement) + target[field] = transformed if repeated else transformed[0] + pending.extend(reversed(children)) + return result + + +@dataclass(frozen=True, slots=True) +class Ref: + """An address, with no independent opcode, type, child list or binder.""" + + owner: Owner + role: str + index: int + + def read(self) -> Any: + """Return a detached message so inspection cannot mutate published data.""" + canonical = self.owner.ref(self.role, self.index) + return _copy_record(cast("Message", canonical.owner._slots[canonical.role][canonical.index])) + + +_STRUCTURAL_HASHES: WeakKeyDictionary[Owner, dict[tuple[str, int], int]] = WeakKeyDictionary() +_STRUCTURAL_LOCK = RLock() + + +def _structural_parts(reference: Ref) -> tuple[tuple[object, ...], list[Ref]]: # noqa: C901 + """Read comparison data directly; only arena addresses become child refs.""" + ref = reference.owner.ref(reference.role, reference.index) + record = cast("Message", ref.owner._slots[ref.role][ref.index]) + if ref.role in {"rules", "rulesets"} or ( + isinstance(record, pb.Node) and record.kind is not None and record.kind.field == "union" + ): + return (ref.role, ref), [] + shape: list[object] = [ref.role] + children = [] + pending: list[tuple[Message, str]] = [(record, "")] + while pending: + message, path = pending.pop() + shape.extend((path, message.desc().type_name)) + inline = [] + for field in message.desc().fields: + if field.name == "span" or (field.presence.name != "IMPLICIT" and field not in message): + continue + value = message[field] + if value is None: + continue + repeated = isinstance(field.value, DescFieldValueList) + values = value if repeated else [value] + shape.extend((field.name, len(values))) + role = _FIELD_ROLES[message.desc().name][field.name] + for index, item in enumerate(values): + if role and not role.startswith("@"): + children.append(ref.owner.ref(role, item)) + shape.append(("ref", role)) + elif isinstance(item, Message): + inline.append((item, f"{path}/{field.name}/{index}")) + elif isinstance(message, pb.PrimitiveValue) and field.name == "f64_bits": + number = unpack("!d", item.to_bytes(8, "big"))[0] + if isnan(number): + msg = "NaN structural comparison requires the pending Python identity decision" + raise NotImplementedError(msg) + shape.append(number) + else: + shape.append(item) + pending.extend(reversed(inline)) + return tuple(shape), children + + +@dataclass(frozen=True, slots=True, eq=False) +class StructuralView: + """Read-only structural inspection; Ref itself retains address identity.""" + + ref: Ref + + def __hash__(self) -> int: + root = self.ref.owner.ref(self.ref.role, self.ref.index) + with _STRUCTURAL_LOCK: + pending = [(root, False)] + active: set[Ref] = set() + while pending: + ref, leaving = pending.pop() + cache = _STRUCTURAL_HASHES.setdefault(ref.owner, {}) + key = (ref.role, ref.index) + if key in cache: + continue + shape, children = _structural_parts(ref) + if leaving: + active.remove(ref) + cache[key] = hash((shape, tuple(_STRUCTURAL_HASHES[c.owner][c.role, c.index] for c in children))) + continue + if ref in active: + msg = "Cyclic protobuf structure without an identity-bearing node" + raise ValueError(msg) + active.add(ref) + pending.append((ref, True)) + pending.extend((child, False) for child in reversed(children)) + return _STRUCTURAL_HASHES[root.owner][root.role, root.index] + + def __eq__(self, other: object) -> bool: + if not isinstance(other, StructuralView): + return NotImplemented + pending = [(self.ref, other.ref)] + seen = set() + while pending: + left, right = pending.pop() + left = left.owner.ref(left.role, left.index) + right = right.owner.ref(right.role, right.index) + if left == right or (left, right) in seen: + continue + seen.add((left, right)) + left_shape, left_children = _structural_parts(left) + right_shape, right_children = _structural_parts(right) + if left_shape != right_shape or len(left_children) != len(right_children): + return False + pending.extend(zip(left_children, right_children, strict=True)) + return True + + +class Owner: + """Published protobuf records; imported slots resolve directly to their owner.""" + + __slots__ = ("__weakref__", "_keepalive", "_slots") + + def __init__(self, slots: Mapping[str, tuple[Message | Ref, ...]], keepalive: tuple[object, ...] = ()) -> None: + # Internal transfer from Builder.publish; caller-owned messages must + # enter through Builder.add/fill/from_program so they are detached. + self._slots = MappingProxyType(dict(slots)) + # Language-local callbacks/lifetimes are never semantic edges or wire data. + self._keepalive = tuple(keepalive) + + def ref(self, role: str, index: int) -> Ref: + if ( + role not in self._slots + or isinstance(index, bool) + or not isinstance(index, int) + or not 0 <= index < len(self._slots[role]) + ): + raise IndexError(f"Invalid {role} slot {index!r}") + record = self._slots[role][index] + return record if isinstance(record, Ref) else Ref(self, role, index) + + +class Builder: + """Thread-confined construction; publication transfers its private records once.""" + + def __init__(self) -> None: + self._slots: dict[str, list[Message | Ref | None]] | None = {role: [] for role in _ROLES} + self._imports: dict[Ref, int] = {} + + def add(self, role: str, record: Message) -> int: + if self._slots is None: + msg = "Builder already published" + raise RuntimeError(msg) + if role not in _ROLES or not isinstance(record, _ROLES[role]): + raise TypeError(f"Wrong protobuf record for {role}") + detached = _copy_record(record) + index = len(self._slots[role]) + self._slots[role].append(detached) + return index + + def reserve(self, role: str) -> int: + if self._slots is None: + msg = "Builder already published" + raise RuntimeError(msg) + if role not in _ROLES: + raise TypeError(f"Unknown arena role {role}") + index = len(self._slots[role]) + self._slots[role].append(None) + return index + + def fill(self, role: str, index: int, record: Message) -> None: + if self._slots is None: + msg = "Builder already published" + raise RuntimeError(msg) + if role not in _ROLES or not isinstance(record, _ROLES[role]): + raise TypeError(f"Wrong protobuf record for {role}") + if not 0 <= index < len(self._slots[role]) or self._slots[role][index] is not None: + msg = "Only an unfilled reserved slot can be filled" + raise ValueError(msg) + self._slots[role][index] = _copy_record(record) + + def import_ref(self, reference: Ref) -> int: + if self._slots is None: + msg = "Builder already published" + raise RuntimeError(msg) + canonical = reference.owner.ref(reference.role, reference.index) + if canonical in self._imports: + return self._imports[canonical] + index = len(self._slots[canonical.role]) + self._slots[canonical.role].append(canonical) + self._imports[canonical] = index + return index + + def publish(self, *, keepalive: tuple[object, ...] = ()) -> Owner: + if self._slots is None: + msg = "Builder already published" + raise RuntimeError(msg) + if any(record is None for records in self._slots.values() for record in records): + msg = "Cannot publish unfilled reserved slots" + raise ValueError(msg) + slots = {role: tuple(cast("list[Message | Ref]", records)) for role, records in self._slots.items()} + self._slots = None + self._imports.clear() + return Owner(slots, keepalive) + + @classmethod + def from_program(cls, program: pb.Program) -> Builder: + """Adopt a decoded Program as a fresh owner, preserving its local indices.""" + if program.ir_version != 1: + raise NotImplementedError(f"Unsupported IR version {program.ir_version}") + if program._unknown_fields: + msg = "Cannot adopt unknown Program fields" + raise NotImplementedError(msg) + builder = cls() + for role in _ROLES: + for record in getattr(program, role): + builder.add(role, record) + return builder + + +@dataclass(frozen=True, slots=True) +class Packed: + program: pb.Program + indices: tuple[int, ...] + # Canonical source declarations let response owners retain the same + # presentations without manufacturing another semantic declaration copy. + definitions: tuple[Ref, ...] + + +def clone_nodes(roots: Iterable[Ref]) -> tuple[Ref, ...]: # noqa: C901 + """ + Expand templates with fresh node identities and shared canonical definitions. + + All roots share one relocation map. Sorts, files and declarations are + immutable imports; only reachable nodes are copied. This does not evaluate + defaults or merge the independent binders of lambda bodies and captures. + """ + builder = Builder() + indices: dict[Ref, int] = {} + pending: list[Ref] = [] + definitions: dict[Owner, dict[tuple[str, str], set[Ref]]] = {} + + def allocate(source: Ref) -> int: + source = source.owner.ref(source.role, source.index) + if source.role != "nodes": + return builder.import_ref(source) + if source not in indices: + indices[source] = builder.reserve("nodes") + pending.append(source) + return indices[source] + + roots = tuple(roots) + if any(root.role != "nodes" for root in roots): + msg = "Template roots must be node references" + raise TypeError(msg) + root_indices = tuple(allocate(root) for root in roots) + while pending: + source = pending.pop() + + def dependency(namespace: str, name: str, origin: Owner = source.owner) -> None: + if origin not in definitions: + table: dict[tuple[str, str], set[Ref]] = {} + for index in range(len(origin._slots["declarations"])): + ref = origin.ref("declarations", index) + key = _definition_key(ref.read()) + if key is not None: + table.setdefault(key, set()).add(ref) + definitions[origin] = table + matches = definitions[origin].get((namespace, name), set()) + if len(matches) != 1: + raise ValueError(f"Template requires one canonical definition for {namespace} {name!r}") + builder.import_ref(next(iter(matches))) + + def relocate(role: str, index: int, origin: Owner = source.owner) -> int: + return allocate(origin.ref(role, index)) + + record = _copy_record(source.read(), relocate=relocate, dependency=dependency) + builder.fill("nodes", indices[source], record) + owner = builder.publish() + return tuple(owner.ref("nodes", index) for index in root_indices) + + +def _definition_key(record: Message) -> tuple[str, str] | None: + """Classify definition namespaces without conflating absent and empty names.""" + if isinstance(record, pb.Declaration): + if record.kind is None: + msg = "Declaration has no kind" + raise ValueError(msg) + namespace = "sort" if record.kind.field in {"eq_sort", "host_sort_family"} else "callable" + return namespace, record.kind.value.name + if isinstance(record, pb.Ruleset) and record.has_field("name"): + return "ruleset", record.name + return None + + +def pack( # noqa: C901, PLR0912 + roots: Iterable[Ref] = (), + *, + commands: Iterable[Ref] = (), + definitions: Iterable[Ref] = (), + ambient: Collection[tuple[str, str]] = (), +) -> Packed: + """ + Pack reachable records once; command roots retain order and multiplicity. + + Definitions supplied here are available by name, not automatic emission + roots. Missing names require an explicit ambient declaration/provider set. + Distinct same-name definitions are rejected rather than guessed compatible; + semantic definition reconciliation belongs to a later adapter-aware layer. + """ + program = pb.Program(ir_version=1) + mapped: dict[Ref, int] = {} + pending: list[Ref] = [] + indexed: set[Owner] = set() + candidates: dict[tuple[str, str], set[Ref]] = {} + demands: set[tuple[str, str]] = set() + emitted: dict[tuple[str, str], Ref] = {} + + def index_owner(owner: Owner) -> None: + if owner in indexed: + return + indexed.add(owner) + for role in ("declarations", "rulesets"): + for index in range(len(owner._slots[role])): + ref = owner.ref(role, index) + record = cast("Message", ref.owner._slots[ref.role][ref.index]) + key = _definition_key(record) + if key is None: + continue + candidates.setdefault(key, set()).add(ref) + if key in emitted and len(candidates[key]) > 1: + raise ValueError(f"Ambiguous same-name definitions require reconciliation: {key}") + + def allocate(reference: Ref) -> int: + ref = reference.owner.ref(reference.role, reference.index) + if ref.role == "commands": + msg = "Commands must be passed as ordered command roots" + raise ValueError(msg) + if ref not in mapped: + index_owner(ref.owner) + records = getattr(program, ref.role) + mapped[ref] = len(records) + records.append(_ROLES[ref.role]()) + pending.append(ref) + return mapped[ref] + + def rewrite(ref: Ref) -> Message: + source = cast("Message", ref.owner._slots[ref.role][ref.index]) + return _copy_record( + source, + relocate=lambda role, index: allocate(ref.owner.ref(role, index)), + dependency=lambda namespace, name: demands.add((namespace, name)), + ) + + for reference in definitions: + ref = reference.owner.ref(reference.role, reference.index) + if ref.role not in {"declarations", "rulesets"}: + msg = "Definitions must refer to declarations or rulesets" + raise TypeError(msg) + index_owner(ref.owner) + indices = tuple(allocate(ref) for ref in roots) + for reference in commands: + ref = reference.owner.ref(reference.role, reference.index) + if ref.role != "commands": + msg = "Expected a command reference" + raise TypeError(msg) + index_owner(ref.owner) + program.commands.append(cast("pb.Command", rewrite(ref))) + + while True: + while pending: + ref = pending.pop() + rewritten = rewrite(ref) + getattr(program, ref.role)[mapped[ref]] = rewritten + key = _definition_key(rewritten) + if key is not None: + if len(candidates.get(key, ())) > 1 or (key in emitted and emitted[key] != ref): + raise ValueError(f"Distinct same-name definitions require reconciliation: {key}") + emitted[key] = ref + for key in sorted(demands): + matches = candidates.get(key, set()) + if len(matches) > 1: + raise ValueError(f"Ambiguous same-name definitions require reconciliation: {key}") + if matches: + allocate(next(iter(matches))) + demands.remove(key) + if not pending: + missing = demands - set(ambient) + if missing: + raise ValueError(f"missing definitions: {sorted(missing)}") + break + return Packed(program, indices, tuple(emitted.values())) + + +_check_inventory() diff --git a/python/egglog/bindings.pyi b/python/egglog/bindings.pyi index c2cf50dc..01e44767 100644 --- a/python/egglog/bindings.pyi +++ b/python/egglog/bindings.pyi @@ -120,6 +120,15 @@ __all__ = [ def setup_tracing(*, exporter: Literal["console", "http"], endpoint: str | None = None) -> None: ... def shutdown_tracing() -> None: ... +@final +class _ProtoEngine: + def __init__(self) -> None: ... + def create(self, request: bytes) -> bytes: ... + def run(self, request: bytes) -> bytes: ... + def clone_graph(self, request: bytes) -> bytes: ... + def destroy(self, request: bytes) -> bytes: ... + def configure_resources(self, request: bytes) -> bytes: ... + @final class SerializedEGraph: @property diff --git a/python/egglog/builtins.py b/python/egglog/builtins.py index 6b73a4be..0fb8627e 100644 --- a/python/egglog/builtins.py +++ b/python/egglog/builtins.py @@ -1,1377 +1,104 @@ -# mypy: disable-error-code="empty-body" -""" -Builtin sorts and function to egg. -""" +"""Catalog-generated builtin wrappers and handwritten Python value protocols.""" from __future__ import annotations -from collections.abc import Callable from dataclasses import dataclass -from fractions import Fraction -from functools import partial, reduce -from inspect import signature -from types import FunctionType, MethodType -from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Unpack, cast, overload +from typing import Any, cast -import cloudpickle -from typing_extensions import TypeVarTuple, deprecated +from egglog_proto.egglog.v1 import egglog_pb as pb +from typing_extensions import deprecated -from .conversion import convert, converter, get_type_args, resolve_literal -from .declarations import * -from .deconstruct import get_callable_args, get_literal_value -from .egraph import ( - BaseExpr, - BuiltinExpr, - expr_fact, - function, - method, - set_current_ruleset, -) -from .runtime import RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate -from .thunk import Thunk - -if TYPE_CHECKING: - from collections.abc import Iterator - - -__all__ = [ - "BigInt", - "BigIntLike", - "BigRat", - "BigRatLike", - "Bool", - "BoolLike", - "Container", - "ExprValueError", - "Map", - "MapLike", - "Maybe", - "MultiSet", - "MultiSetLike", - "Pair", - "Primitive", - "PyObject", - "Rational", - "RationalLike", - "Set", - "SetLike", - "String", - "StringLike", - "Unit", - "UnstableFn", - "Vec", - "VecLike", - "catch", - "f64", - "f64Like", - "i64", - "i64Like", - "join", - "map_fold_kv", - "multiset_contains_swapped", - "multiset_flat_map", - "multiset_fold", - "multiset_not_contains_swapped", - "multiset_remove_swapped", - "multiset_subtract_swapped", - "py_eval", - "py_eval_fn", - "py_exec", -] +from ._catalog import builtin_catalog +from .conversion import converter +from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction @dataclass class ExprValueError(AttributeError): - """ - Raised when an expression cannot be converted to a Python value because the value is not a constructor. - """ + """A symbolic expression is not in a locally inspectable value form.""" - expr: BaseExpr + expr: object allowed: str def __str__(self) -> str: return f"Cannot get Python value of {self.expr}, must be of form {self.allowed}. Try calling `extract` on it to get the underlying value." -class Unit(BuiltinExpr, egg_sort="Unit"): - """ - The unit type. This is used to represent if a value exists in the e-graph or not. - """ - - def __init__(self) -> None: ... - - @method(preserve=True) - def __bool__(self) -> bool: - return bool(expr_fact(self)) - - -class String(BuiltinExpr, egg_sort="String"): - def __init__(self, value: str) -> None: ... - - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> str: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> str: - if (value := get_literal_value(self)) is not None: - return value - raise ExprValueError(self, "String") - - __match_args__ = ("value",) - - @method(egg_fn="replace") - def replace(self, old: StringLike, new: StringLike) -> String: ... - - @method(preserve=True) - def __add__(self, other: StringLike) -> String: - return join(self, other) - - @method(egg_fn="log") - def log(self) -> Unit: ... - - -StringLike: TypeAlias = String | str - - -@function(egg_fn="+", builtin=True) -def join(*strings: StringLike) -> String: ... - - -converter(str, String, String) - - -class Bool(BuiltinExpr, egg_sort="bool"): - def __init__(self, value: bool) -> None: ... - - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> bool: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> bool: - if (value := get_literal_value(self)) is not None: - return value - raise ExprValueError(self, "Bool") - - __match_args__ = ("value",) - - @method(preserve=True) - def __bool__(self) -> bool: - return self.value - - @method(egg_fn="not") - def __invert__(self) -> Bool: ... - - @method(egg_fn="and") - def __and__(self, other: BoolLike) -> Bool: ... - - @method(egg_fn="or") - def __or__(self, other: BoolLike) -> Bool: ... - - @method(egg_fn="xor") - def __xor__(self, other: BoolLike) -> Bool: ... - - @method(egg_fn="=>") - def implies(self, other: BoolLike) -> Bool: ... - - -BoolLike: TypeAlias = Bool | bool - - -converter(bool, Bool, Bool) - - -class i64(BuiltinExpr, egg_sort="i64"): # noqa: N801 - def __init__(self, value: int) -> None: ... - - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> int: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> int: - if (value := get_literal_value(self)) is not None: - return value - raise ExprValueError(self, "i64") - - __match_args__ = ("value",) - - @method(preserve=True) - def __index__(self) -> int: - return self.value - - @method(preserve=True) - def __int__(self) -> int: - return self.value - - @method(egg_fn="+") - def __add__(self, other: i64Like) -> i64: ... - - @method(egg_fn="-") - def __sub__(self, other: i64Like) -> i64: ... - - @method(egg_fn="*") - def __mul__(self, other: i64Like) -> i64: ... - - @method(egg_fn="/") - def __truediv__(self, other: i64Like) -> i64: ... - - @method(egg_fn="%") - def __mod__(self, other: i64Like) -> i64: ... - - @method(egg_fn="&") - def __and__(self, other: i64Like) -> i64: ... - - @method(egg_fn="|") - def __or__(self, other: i64Like) -> i64: ... - - @method(egg_fn="^") - def __xor__(self, other: i64Like) -> i64: ... - - @method(egg_fn="<<") - def __lshift__(self, other: i64Like) -> i64: ... - - @method(egg_fn=">>") - def __rshift__(self, other: i64Like) -> i64: ... - - def __radd__(self, other: i64Like) -> i64: ... - - def __rsub__(self, other: i64Like) -> i64: ... - - def __rmul__(self, other: i64Like) -> i64: ... - - def __rtruediv__(self, other: i64Like) -> i64: ... - - def __rmod__(self, other: i64Like) -> i64: ... - - def __rand__(self, other: i64Like) -> i64: ... - - def __ror__(self, other: i64Like) -> i64: ... - - def __rxor__(self, other: i64Like) -> i64: ... - - def __rlshift__(self, other: i64Like) -> i64: ... - - def __rrshift__(self, other: i64Like) -> i64: ... - - @method(egg_fn="log2") - def log2(self) -> i64: ... - - @method(egg_fn="not-i64") - def __invert__(self) -> i64: ... - - @method(egg_fn="<") - def __lt__(self, other: i64Like) -> Unit: # type: ignore[has-type] - ... - - @method(egg_fn=">") - def __gt__(self, other: i64Like) -> Unit: ... - - @method(egg_fn="<=") - def __le__(self, other: i64Like) -> Unit: # type: ignore[has-type] - ... - - @method(egg_fn=">=") - def __ge__(self, other: i64Like) -> Unit: ... - - @method(egg_fn="min") - def min(self, other: i64Like) -> i64: ... - - @method(egg_fn="max") - def max(self, other: i64Like) -> i64: ... - - @method(egg_fn="to-string") - def to_string(self) -> String: ... - - @method(egg_fn="bool-<") - def bool_lt(self, other: i64Like) -> Bool: ... - - @method(egg_fn="bool->") - def bool_gt(self, other: i64Like) -> Bool: ... - - @method(egg_fn="bool-<=") - def bool_le(self, other: i64Like) -> Bool: ... - - @method(egg_fn="bool->=") - def bool_ge(self, other: i64Like) -> Bool: ... - - @method(egg_fn="abs") - def __abs__(self) -> i64: ... - - @method(egg_fn="vec-range") - def range(self) -> Vec[i64]: ... - - -# The types which can be converted into an i64 -i64Like: TypeAlias = i64 | int # noqa: N816, PYI042 - -converter(int, i64, i64) - - -@function(builtin=True, egg_fn="count-matches") -def count_matches(s: StringLike, pattern: StringLike) -> i64: ... - - -class f64(BuiltinExpr, egg_sort="f64"): # noqa: N801 - def __init__(self, value: float) -> None: ... - - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> float: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> float: - if (value := get_literal_value(self)) is not None: - return value - raise ExprValueError(self, "f64") - - __match_args__ = ("value",) - - @method(preserve=True) - def __float__(self) -> float: - return self.value - - @method(preserve=True) - def __int__(self) -> int: - return int(self.value) - - @method(egg_fn="neg") - def __neg__(self) -> f64: ... - - @method(egg_fn="+") - def __add__(self, other: f64Like) -> f64: ... - - @method(egg_fn="-") - def __sub__(self, other: f64Like) -> f64: ... - - @method(egg_fn="*") - def __mul__(self, other: f64Like) -> f64: ... - - @method(egg_fn="/") - def __truediv__(self, other: f64Like) -> f64: ... - - @method(egg_fn="%") - def __mod__(self, other: f64Like) -> f64: ... - - @method(egg_fn="^") - def __pow__(self, other: f64Like) -> f64: ... - - def __radd__(self, other: f64Like) -> f64: ... - - def __rsub__(self, other: f64Like) -> f64: ... - - def __rmul__(self, other: f64Like) -> f64: ... - - def __rtruediv__(self, other: f64Like) -> f64: ... - - def __rmod__(self, other: f64Like) -> f64: ... - - @method(egg_fn="abs") - def __abs__(self) -> f64: ... - - @method(egg_fn="exp") - def exp(self) -> f64: ... - - @method(egg_fn="log") - def log(self) -> f64: ... - - @method(egg_fn="sqrt") - def sqrt(self) -> f64: ... - - @method(egg_fn="f64-is-finite") - def is_finite(self) -> Unit: ... - - @method(egg_fn="<") - def __lt__(self, other: f64Like) -> Unit: # type: ignore[has-type] - ... - - @method(egg_fn=">") - def __gt__(self, other: f64Like) -> Unit: ... - - @method(egg_fn="<=") - def __le__(self, other: f64Like) -> Unit: # type: ignore[has-type] - ... - - @method(egg_fn=">=") - def __ge__(self, other: f64Like) -> Unit: ... - - @method(egg_fn="min") - def min(self, other: f64Like) -> f64: ... - - @method(egg_fn="max") - def max(self, other: f64Like) -> f64: ... - - @method(egg_fn="to-i64") - def to_i64(self) -> i64: ... - - @method(egg_fn="to-f64") - @classmethod - def from_i64(cls, i: i64) -> f64: ... - - @method(egg_fn="to-string") - def to_string(self) -> String: ... - - -f64Like: TypeAlias = f64 | float | int # noqa: N816, PYI042 - -converter(int, f64, lambda i: f64(float(i))) -converter(float, f64, f64) - - -T = TypeVar("T", bound=BaseExpr) -V = TypeVar("V", bound=BaseExpr) - - -class Maybe(BuiltinExpr, Generic[T], egg_sort="Maybe"): - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> T | None: - if get_callable_args(self, Maybe.none) is not None: - return None - match get_callable_args(self, Maybe.some): - case (value,): - return value # type: ignore[unreachable] - raise ExprValueError(self, "Maybe.none() or Maybe.some(value)") - - __match_args__ = ("value",) - - @method(egg_fn="maybe-none") - @classmethod - def none(cls) -> Maybe[T]: ... - - @method(egg_fn="maybe-some") - @classmethod - def some(cls, value: T) -> Maybe[T]: ... - - @method(egg_fn="maybe-unwrap") - def unwrap(self) -> T: ... - - @method(egg_fn="maybe-unwrap-or") - def unwrap_or(self, default: T) -> T: ... - - @method(egg_fn="unstable-maybe-match") - def match(self, f: Callable[[T], V], n: V) -> V: ... - - -converter(type(None), Maybe, lambda _: Maybe[get_type_args()[0]].none()) # type: ignore[misc] - - -L = TypeVar("L", bound=BaseExpr) -R = TypeVar("R", bound=BaseExpr) - - -class Pair(BuiltinExpr, Generic[L, R], egg_sort="Pair"): - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> tuple[L, R]: - match get_callable_args(self, Pair[L, R]): - case (left, right): - return (left, right) - raise ExprValueError(self, "Pair(left, right)") - - __match_args__ = ("value",) - - @method(egg_fn="pair") - def __init__(self, left: L, right: R) -> None: ... - - @method(egg_fn="pair-first") # type: ignore[prop-decorator] - @property - def left(self) -> L: ... - - @method(egg_fn="pair-second") # type: ignore[prop-decorator] - @property - def right(self) -> R: ... - - -def _tuple_to_pair(value: tuple[object, ...]) -> Pair: - if len(value) != 2: - raise ValueError(f"Expected a tuple of length 2 for Pair conversion, got length {len(value)}") - left, right = value - left_type, right_type = get_type_args() - return Pair(convert(left, left_type), convert(right, right_type)) - - -converter(tuple, Pair, _tuple_to_pair) - - -@function(egg_fn="unstable-catch", builtin=True) -def catch(f: Callable[[], T]) -> Maybe[T]: ... - - -class Map(BuiltinExpr, Generic[T, V], egg_sort="Map"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> dict[T, V]: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> dict[T, V]: - items = [] - while args := get_callable_args(self, Map.insert): # type: ignore[var-annotated] - self, k, v = args # noqa: PLW0642 - items.append((k, v)) - if get_callable_args(self, Map.empty) is None: - raise ExprValueError(self, "Map.empty or Map.insert") - return dict(reversed(items)) - - __match_args__ = ("value",) - - @method(preserve=True) - def __iter__(self) -> Iterator[T]: - return iter(self.value) - - @method(preserve=True) - def __len__(self) -> int: - return len(self.value) - - @method(preserve=True) - def __contains__(self, key: T) -> bool: - return key in self.value - - @method(egg_fn="map-empty") - @classmethod - def empty(cls) -> Map[T, V]: ... - - @method(egg_fn="map-insert") - def insert(self, key: T, value: V) -> Map[T, V]: ... - - @method(egg_fn="map-get") - def __getitem__(self, key: T) -> V: ... - - @method(egg_fn="map-not-contains") - def not_contains(self, key: T) -> Unit: ... - - @method(egg_fn="map-contains") - def contains(self, key: T) -> Unit: ... - - @method(egg_fn="map-remove") - def remove(self, key: T) -> Map[T, V]: ... - - @method(egg_fn="map-length") - def length(self) -> i64: ... - - -TO = TypeVar("TO") -VO = TypeVar("VO") -A = TypeVar("A", bound=BaseExpr) - -converter( - dict, - Map, - lambda t: reduce( - (lambda acc, kv: acc.insert(convert(kv[0], get_type_args()[0]), convert(kv[1], get_type_args()[1]))), - t.items(), - Map[get_type_args()].empty(), # type: ignore[misc] - ), -) - -MapLike: TypeAlias = Map[T, V] | dict[TO, VO] - - -@function(egg_fn="map-fold-kv", builtin=True) -def map_fold_kv(f: Callable[[A, T, V], A], initial: A, xs: Map[T, V]) -> A: ... - - -class Set(BuiltinExpr, Generic[T], egg_sort="Set"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> set[T]: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> set[T]: - if (args := get_callable_args(self, Set[T])) is not None: - return set(args) - raise ExprValueError(self, "Set(*xs)") - - __match_args__ = ("value",) - - @method(preserve=True) - def __iter__(self) -> Iterator[T]: - if (args := get_callable_args(self, Set[T])) is not None: - return iter(dict.fromkeys(args)) - return iter(self.value) - - @method(preserve=True) - def __len__(self) -> int: - return len(self.value) - - @method(preserve=True) - def __contains__(self, key: T) -> bool: - return key in self.value - - @method(egg_fn="set-of") - def __init__(self, *args: T) -> None: ... - - @method(egg_fn="set-empty") - @classmethod - def empty(cls) -> Set[T]: ... - - @method(egg_fn="set-insert") - def insert(self, value: T) -> Set[T]: ... - - @method(egg_fn="set-not-contains") - def not_contains(self, value: T) -> Unit: ... - - @method(egg_fn="set-contains") - def contains(self, value: T) -> Unit: ... - - @method(egg_fn="set-remove") - def remove(self, value: T) -> Set[T]: ... - - @method(egg_fn="set-union") - def __or__(self, other: Set[T]) -> Set[T]: ... - - @method(egg_fn="set-diff") - def __sub__(self, other: Set[T]) -> Set[T]: ... - - @method(egg_fn="set-intersect") - def __and__(self, other: Set[T]) -> Set[T]: ... - - @method(egg_fn="set-length") - def length(self) -> i64: ... - - -converter( - set, - Set, - lambda t: Set(*(convert(x, get_type_args()[0]) for x in t)) if t else Set[get_type_args()[0]].empty(), # type: ignore[misc] -) - -SetLike: TypeAlias = Set[T] | set[TO] - - -class MultiSet(BuiltinExpr, Generic[T], egg_sort="MultiSet"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> list[T]: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> list[T]: - if (args := get_callable_args(self, MultiSet[T])) is not None: - return list(args) - raise ExprValueError(self, "MultiSet") - - __match_args__ = ("value",) - - @method(preserve=True) - def __iter__(self) -> Iterator[T]: - return iter(self.value) - - @method(preserve=True) - def __len__(self) -> int: - return len(self.value) - - @method(preserve=True) - def __contains__(self, key: T) -> bool: - return key in self.value - - @method(egg_fn="multiset-of") - def __init__(self, *args: T) -> None: ... - - @method(egg_fn="multiset-intersection") - def __and__(self, other: MultiSet[T]) -> MultiSet[T]: ... - - @method(egg_fn="multiset-single") - @classmethod - def single(cls, x: T, i: i64Like) -> MultiSet[T]: ... - - @method(egg_fn="multiset-sum-multisets") - @classmethod - def sum_multisets(cls, xs: MultiSet[MultiSet[T]]) -> MultiSet[T]: ... - - @method(egg_fn="multiset-insert") - def insert(self, value: T) -> MultiSet[T]: ... - - @method(egg_fn="multiset-not-contains") - def not_contains(self, value: T) -> Unit: ... - - @method(egg_fn="multiset-contains") - def contains(self, value: T) -> Unit: ... - - @method(egg_fn="multiset-remove") - def remove(self, value: T) -> MultiSet[T]: ... - - @method(egg_fn="multiset-length") - def length(self) -> i64: ... +def _scalar_value(expr: RuntimeExpr) -> object: + from .deconstruct import get_literal_value # noqa: PLC0415 - @method(egg_fn="multiset-pick") - def pick(self) -> T: ... + value = get_literal_value(expr) + if value is not None: + return value + node: pb.Node = expr.__egg_ref__.read() + from .runtime import _class_for_sort # noqa: PLC0415 - @method(egg_fn="multiset-sum") - def __add__(self, other: MultiSet[T]) -> MultiSet[T]: ... + cls = _class_for_sort(expr.__egg_ref__.owner.ref("sorts", node.sort_id)) + raise ExprValueError(expr, str(cls)) - @method(egg_fn="multiset-subtract") - def __sub__(self, other: MultiSet[T]) -> MultiSet[T]: ... - @method(egg_fn="unstable-multiset-map", reverse_args=True) - def map(self, f: Callable[[T], V]) -> MultiSet[V]: ... +@deprecated("use .value") +def _eval(expr: RuntimeExpr) -> object: + return expr.value - @method(egg_fn="unstable-multiset-fill-index") - def fill_index(self, f: Callable[[MultiSet[T], T], i64]) -> Unit: ... - @method(egg_fn="unstable-multiset-clear-index") - def clear_index(self, f: Callable[[MultiSet[T], T], i64]) -> Unit: ... +def _integer(expr: RuntimeExpr) -> int: + return int(cast("Any", expr.value)) - @method(egg_fn="multiset-pick-max") - def pick_max(self) -> T: ... - @method(egg_fn="multiset-count") - def count(self, value: T) -> i64: ... +def _float(expr: RuntimeExpr) -> float: + return float(cast("Any", expr.value)) - @method(egg_fn="unstable-multiset-filter", reverse_args=True) - def filter(self, f: Callable[[T], Unit]) -> MultiSet[T]: ... - @method(egg_fn="unstable-multiset-filter-not", reverse_args=True) - def filter_not(self, f: Callable[[T], Unit]) -> MultiSet[T]: ... - - @method(egg_fn="multiset-reset-counts") - def reset_counts(self) -> MultiSet[T]: ... - - -# TODO: Move to method when partial supports reverse_args -@function(egg_fn="unstable-multiset-flat-map", builtin=True) -def multiset_flat_map(f: Callable[[T], MultiSet[T]], xs: MultiSet[T]) -> MultiSet[T]: ... - - -@function(egg_fn="multiset-remove-swapped", builtin=True) -def multiset_remove_swapped(x: T, xs: MultiSet[T]) -> MultiSet[T]: ... - - -@function(egg_fn="multiset-subtract-swapped", builtin=True) -def multiset_subtract_swapped(x: MultiSet[T], xs: MultiSet[T]) -> MultiSet[T]: ... - - -@function(egg_fn="multiset-not-contains-swapped", builtin=True) -def multiset_not_contains_swapped(x: T, xs: MultiSet[T]) -> Unit: ... - - -@function(egg_fn="multiset-contains-swapped", builtin=True) -def multiset_contains_swapped(x: T, xs: MultiSet[T]) -> Unit: ... - - -@function(egg_fn="unstable-multiset-reduce", builtin=True) -def multiset_fold(f: Callable[[T, T], T], initial: T, xs: MultiSet[T]) -> T: ... - - -converter( - tuple, - MultiSet, - lambda t: MultiSet(*(convert(x, get_type_args()[0]) for x in t)) if t else MultiSet[get_type_args()[0]](), # type: ignore[operator,misc] -) - -MultiSetLike: TypeAlias = MultiSet[T] | tuple[TO, ...] - - -class Rational(BuiltinExpr, egg_sort="Rational"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> Fraction: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> Fraction: - match get_callable_args(self, Rational): - case (i64(num), i64(den)): - return Fraction(num, den) - raise ExprValueError(self, "Rational(i64(num), i64(den))") - - __match_args__ = ("value",) - - @method(preserve=True) - def __float__(self) -> float: - return float(self.value) - - @method(preserve=True) - def __int__(self) -> int: - return int(self.value) - - @method(egg_fn="rational") - def __init__(self, num: i64Like, den: i64Like) -> None: ... - - @method(egg_fn="to-f64") - def to_f64(self) -> f64: ... - - @method(egg_fn="+") - def __add__(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="-") - def __sub__(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="*") - def __mul__(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="/") - def __truediv__(self, other: RationalLike) -> Rational: ... - - def __radd__(self, other: RationalLike) -> Rational: ... - - def __rsub__(self, other: RationalLike) -> Rational: ... - - def __rmul__(self, other: RationalLike) -> Rational: ... - - def __rtruediv__(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="min") - def min(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="max") - def max(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="neg") - def __neg__(self) -> Rational: ... - - @method(egg_fn="abs") - def __abs__(self) -> Rational: ... - - @method(egg_fn="floor") - def floor(self) -> Rational: ... - - @method(egg_fn="ceil") - def ceil(self) -> Rational: ... - - @method(egg_fn="round") - def round(self) -> Rational: ... - - @method(egg_fn="pow") - def __pow__(self, other: RationalLike) -> Rational: ... - - def __rpow__(self, other: RationalLike) -> Rational: ... - - @method(egg_fn="log") - def log(self) -> Rational: ... - - @method(egg_fn="sqrt") - def sqrt(self) -> Rational: ... - - @method(egg_fn="cbrt") - def cbrt(self) -> Rational: ... - - @method(egg_fn="numer") # type: ignore[prop-decorator] - @property - def numer(self) -> i64: ... - - @method(egg_fn="denom") # type: ignore[prop-decorator] - @property - def denom(self) -> i64: ... - - @method(egg_fn="<") - def __lt__(self, other: RationalLike) -> Unit: ... # type: ignore[has-type] - - @method(egg_fn=">") - def __gt__(self, other: RationalLike) -> Unit: ... - - @method(egg_fn="<=") - def __le__(self, other: RationalLike) -> Unit: ... # type: ignore[has-type] - - @method(egg_fn=">=") - def __ge__(self, other: RationalLike) -> Unit: ... - - -converter(i64, Rational, lambda i: Rational(i, 1)) -converter(Fraction, Rational, lambda f: Rational(f.numerator, f.denominator)) -RationalLike: TypeAlias = Rational | Fraction | i64Like - - -class BigInt(BuiltinExpr, egg_sort="BigInt"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> int: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> int: - match get_callable_args(self, BigInt.from_string): - case (String(s),): - return int(s) - raise ExprValueError(self, "BigInt.from_string(String(s))") - - __match_args__ = ("value",) - - @method(preserve=True) - def __index__(self) -> int: - return self.value - - @method(preserve=True) - def __int__(self) -> int: - return self.value - - @method(egg_fn="from-string") - @classmethod - def from_string(cls, s: StringLike) -> BigInt: ... - - @method(egg_fn="bigint") - def __init__(self, value: i64Like) -> None: ... - - @method(egg_fn="+") - def __add__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="-") - def __sub__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="*") - def __mul__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="/") - def __truediv__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="%") - def __mod__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="&") - def __and__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="|") - def __or__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="^") - def __xor__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="<<") - def __lshift__(self, other: i64Like) -> BigInt: ... - - @method(egg_fn=">>") - def __rshift__(self, other: i64Like) -> BigInt: ... - - def __radd__(self, other: BigIntLike) -> BigInt: ... - - def __rsub__(self, other: BigIntLike) -> BigInt: ... - - def __rmul__(self, other: BigIntLike) -> BigInt: ... - - def __rtruediv__(self, other: BigIntLike) -> BigInt: ... - - def __rmod__(self, other: BigIntLike) -> BigInt: ... - - def __rand__(self, other: BigIntLike) -> BigInt: ... - - def __ror__(self, other: BigIntLike) -> BigInt: ... - - def __rxor__(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="not-Z") - def __invert__(self) -> BigInt: ... - - @method(egg_fn="bits") - def bits(self) -> BigInt: ... - - @method(egg_fn="<") - def __lt__(self, other: BigIntLike) -> Unit: # type: ignore[has-type] - ... - - @method(egg_fn=">") - def __gt__(self, other: BigIntLike) -> Unit: ... - - @method(egg_fn="<=") - def __le__(self, other: BigIntLike) -> Unit: # type: ignore[has-type] - ... - - @method(egg_fn=">=") - def __ge__(self, other: BigIntLike) -> Unit: ... - - @method(egg_fn="min") - def min(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="max") - def max(self, other: BigIntLike) -> BigInt: ... - - @method(egg_fn="to-string") - def to_string(self) -> String: ... - - @method(egg_fn="bool-=") - def bool_eq(self, other: BigIntLike) -> Bool: ... - - @method(egg_fn="bool-<") - def bool_lt(self, other: BigIntLike) -> Bool: ... - - @method(egg_fn="bool->") - def bool_gt(self, other: BigIntLike) -> Bool: ... - - @method(egg_fn="bool-<=") - def bool_le(self, other: BigIntLike) -> Bool: ... - - @method(egg_fn="bool->=") - def bool_ge(self, other: BigIntLike) -> Bool: ... - - -converter(i64, BigInt, BigInt) - -BigIntLike: TypeAlias = BigInt | i64Like - - -class BigRat(BuiltinExpr, egg_sort="BigRat"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> Fraction: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> Fraction: - match get_callable_args(self, BigRat): - case (BigInt(num), BigInt(den)): - return Fraction(num, den) - raise ExprValueError(self, "BigRat(BigInt(num), BigInt(den))") - - __match_args__ = ("value",) - - @method(preserve=True) - def __float__(self) -> float: - return float(self.value) - - @method(preserve=True) - def __int__(self) -> int: - return int(self.value) - - @method(egg_fn="bigrat") - def __init__(self, num: BigIntLike, den: BigIntLike) -> None: ... - - @method(egg_fn="to-f64") - def to_f64(self) -> f64: ... - - @method(egg_fn="to-i64") - def to_i64(self) -> i64: ... - - @method(egg_fn="+") - def __add__(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="-") - def __sub__(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="*") - def __mul__(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="/") - def __truediv__(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="min") - def min(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="max") - def max(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="neg") - def __neg__(self) -> BigRat: ... - - @method(egg_fn="abs") - def __abs__(self) -> BigRat: ... - - @method(egg_fn="floor") - def floor(self) -> BigRat: ... - - @method(egg_fn="ceil") - def ceil(self) -> BigRat: ... - - @method(egg_fn="round") - def round(self) -> BigRat: ... - - @method(egg_fn="pow") - def __pow__(self, other: BigRatLike) -> BigRat: ... - - @method(egg_fn="log") - def log(self) -> BigRat: ... - - @method(egg_fn="sqrt") - def sqrt(self) -> BigRat: ... - - @method(egg_fn="cbrt") - def cbrt(self) -> BigRat: ... - - @method(egg_fn="numer") # type: ignore[prop-decorator] - @property - def numer(self) -> BigInt: ... - - @method(egg_fn="denom") # type: ignore[prop-decorator] - @property - def denom(self) -> BigInt: ... - - @method(egg_fn="<") - def __lt__(self, other: BigRatLike) -> Unit: ... # type: ignore[has-type] - - @method(egg_fn=">") - def __gt__(self, other: BigRatLike) -> Unit: ... - - @method(egg_fn=">=") - def __ge__(self, other: BigRatLike) -> Unit: ... # type: ignore[has-type] - - @method(egg_fn="<=") - def __le__(self, other: BigRatLike) -> Unit: ... - - -converter(i64, BigRat, lambda i: BigRat(BigInt(i), BigInt(1))) -converter(Fraction, BigRat, lambda f: BigRat(f.numerator, f.denominator)) -BigRatLike: TypeAlias = BigRat | Fraction | i64Like - - -class Vec(BuiltinExpr, Generic[T], egg_sort="Vec"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> tuple[T, ...]: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> tuple[T, ...]: - if get_callable_args(self, Vec.empty) is not None: - return () - if (args := get_callable_args(self, Vec[T])) is not None: - return args - raise ExprValueError(self, "Vec(*xs) or Vec.empty()") - - __match_args__ = ("value",) - - @method(preserve=True) - def __iter__(self) -> Iterator[T]: - return iter(self.value) - - @method(preserve=True) - def __len__(self) -> int: - return len(self.value) - - @method(preserve=True) - def __contains__(self, key: T) -> bool: - return key in self.value - - @method(egg_fn="vec-of") - def __init__(self, *args: T) -> None: ... - - @method(egg_fn="vec-empty") - @classmethod - def empty(cls) -> Vec[T]: ... - - @method(egg_fn="vec-append") - def append(self, *others: VecLike[T, T]) -> Vec[T]: ... - - @method(egg_fn="vec-push") - def push(self, value: T) -> Vec[T]: ... - - @method(egg_fn="vec-pop") - def pop(self) -> Vec[T]: ... - - @method(egg_fn="vec-not-contains") - def not_contains(self, value: T) -> Unit: ... - - @method(egg_fn="vec-contains") - def contains(self, value: T) -> Unit: ... - - @method(egg_fn="vec-length") - def length(self) -> i64: ... - - @method(egg_fn="vec-get") - def __getitem__(self, index: i64Like) -> T: ... - - @method(egg_fn="vec-remove") - def remove(self, index: i64Like) -> Vec[T]: ... - - @method(egg_fn="vec-set") - def set(self, index: i64Like, value: T) -> Vec[T]: ... - - @method(egg_fn="vec-union") - def __or__(self, other: Vec[T]) -> Vec[T]: ... - - @method(egg_fn="unstable-vec-map", reverse_args=True) - def map(self, fn: Callable[[T], V]) -> Vec[V]: ... - - -for sequence_type in (list, tuple): - converter( - sequence_type, - Vec, - lambda t: Vec(*(convert(x, get_type_args()[0]) for x in t)) if t else Vec[get_type_args()[0]].empty(), # type: ignore[misc] +# These conversions/codecs are Python value protocol policy. No callable input, +# output, overload, or generic signature is duplicated here. +_LITERAL_PYTHON_TYPES: dict[str, tuple[type, ...]] = { + "i64": (int,), + "f64": (float, int), + "String": (str,), + "bool": (bool,), +} +__all__ = ["ExprValueError"] +_catalog = builtin_catalog() +for _definition in _catalog.definitions.values(): + _declaration: pb.Declaration = _definition.read() + _kind = _declaration.kind + if _kind is None or not isinstance(_kind.value, pb.HostSortFamily | pb.EqSort): + continue + _binding = _kind.value.bindings + _path = ( + tuple(_binding.python.path) + if _binding is not None and _binding.python is not None + else (*__name__.split("."), _kind.value.name) ) - -VecLike: TypeAlias = Vec[T] | tuple[TO, ...] | list[TO] - - -TS = TypeVarTuple("TS") - -T1 = TypeVar("T1") -T2 = TypeVar("T2") -T3 = TypeVar("T3") - - -class UnstableFn(BuiltinExpr, Generic[T, *TS], egg_sort="UnstableFn"): - @overload - def __init__(self, f: Callable[[Unpack[TS]], T]) -> None: ... - - @overload - def __init__(self, f: Callable[[T1, Unpack[TS]], T], _a: T1, /) -> None: ... - - @overload - def __init__(self, f: Callable[[T1, T2, Unpack[TS]], T], _a: T1, _b: T2, /) -> None: ... - - # Removing due to bug in MyPy - # https://github.com/python/mypy/issues/17212 - # @overload - # def __init__(self, f: Callable[[T1, T2, T3, Unpack[TS]], T], _a: T1, _b: T2, _c: T3, /) -> None: ... - - # etc, for partial application - - @method(egg_fn="unstable-fn") - def __init__(self, f, *partial) -> None: ... - - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> Callable[[Unpack[TS]], T]: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> Callable[[Unpack[TS]], T]: - """ - If this is a constructor, returns either the callable directly or a `functools.partial` function if args are provided. - """ - if (fn := get_literal_value(self)) is not None: - return fn - raise ExprValueError(self, "UnstableFn(f, *args)") - - __match_args__ = ("value",) - - @method(egg_fn="unstable-app") - def __call__(self, *args: *TS) -> T: ... - - -# Method Type is for builtins like __getitem__ -converter(MethodType, UnstableFn, lambda m: UnstableFn[*get_type_args()](m.__func__, m.__self__)) # type: ignore[operator, misc] -# Ignore PLW0108. -converter(RuntimeFunction, UnstableFn, lambda rf: UnstableFn[*get_type_args()](rf)) # type: ignore[operator, misc] -# converter(RuntimeClass, UnstableFn, lambda rc: UnstableFn[*get_type_args()](rc)) # type: ignore[operator, misc] -converter(partial, UnstableFn, lambda p: UnstableFn[*get_type_args()](p.func, *p.args)) # type: ignore[operator, misc] - - -def _convert_function(fn: FunctionType) -> UnstableFn: - """ - Converts a function type to an unstable function. This function will be an anon function in egglog. - - Would just be UnstableFn(function(a)) but we have to account for unbound vars within the body. - - This means that we have to turn all of those unbound vars into args to the function, and then - partially apply them, alongside storing the eager primitive body for the function. - """ - decls = Declarations() - return_type, *arg_types = [resolve_type_annotation_mutate(decls, tp) for tp in get_type_args()] - arg_names = [p.name for p in signature(fn).parameters.values()] - arg_decls = [ - TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) for name, tp in zip(arg_names, arg_types, strict=True) - ] - with set_current_ruleset(None): - res = resolve_literal( - return_type, fn(*(RuntimeExpr.__from_values__(decls, a) for a in arg_decls)), Thunk.value(decls) - ) - res_expr = res.__egg_typed_expr__ - decls |= res - # these are all the args that appear in the body that are not bound by the args of the function - unbound_vars = list(collect_unbound_vars(res_expr) - set(arg_decls)) - # prefix the args with them - all_args = tuple(unbound_vars + arg_decls) - normalized_args = tuple( - TypedExprDecl( - typed_arg.tp, - UnboundVarDecl(cast("UnboundVarDecl", typed_arg.expr).name, f"_{i}"), - ) - for i, typed_arg in enumerate(all_args) - ) - res_expr = replace_typed_expr(res_expr, dict(zip(all_args, normalized_args, strict=True))) - fn_ref = UnnamedFunctionRef(normalized_args, res_expr) - fn = RuntimeFunction(Thunk.value(decls), Thunk.value(fn_ref)) - return UnstableFn(fn, *(RuntimeExpr.__from_values__(decls, v) for v in unbound_vars)) - - -converter(FunctionType, UnstableFn, _convert_function) - - -class PyObject(BuiltinExpr, egg_sort="PyObject"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> object: - return self.value - - @method(preserve=True) # type: ignore[prop-decorator] - @property - def value(self) -> object: - expr = cast("RuntimeExpr", self).__egg_typed_expr__.expr - if not isinstance(expr, PyObjectDecl): - raise ExprValueError(self, "PyObject(x)") - return cloudpickle.loads(expr.pickled) - - __match_args__ = ("value",) - - def __init__(self, value: object) -> None: ... - - @method(egg_fn="py-call") - def __call__(self, *args: object) -> PyObject: ... - - @method(egg_fn="py-call-extended") - def call_extended(self, args: PyObject, kwargs: PyObject) -> PyObject: - """ - Call the PyObject with the given args and kwargs PyObjects. - """ - - @method(egg_fn="py-from-string") - @classmethod - def from_string(cls, s: StringLike) -> PyObject: ... - - @method(egg_fn="py-to-string") - def to_string(self) -> String: ... - - @method(egg_fn="py-to-bool") - def to_bool(self) -> Bool: ... - - @method(egg_fn="py-dict-update") - def dict_update(self, *keys_and_values: object) -> PyObject: ... - - @method(egg_fn="py-from-int") - @classmethod - def from_int(cls, i: i64Like) -> PyObject: ... - - @method(egg_fn="py-dict") - @classmethod - def dict(cls, *keys_and_values: object) -> PyObject: ... - - -converter(object, PyObject, PyObject) - - -@function(builtin=True, egg_fn="py-eval") -def py_eval(code: StringLike, globals_: object = PyObject.dict(), locals_: object = PyObject.dict()) -> PyObject: ... - - -class PyObjectFunction(Protocol): - def __call__(self, *__args: PyObject) -> PyObject: ... - - -@deprecated("use PyObject(fn) directly") -def py_eval_fn(fn: Callable) -> PyObjectFunction: - """ - Takes a python callable and maps it to a callable which takes and returns PyObjects. - - It translates it to a call which uses `py_eval` to call the function, passing in the - args as locals, and using the globals from function. - """ - return PyObject(fn) - - -@function(builtin=True, egg_fn="py-exec") -def py_exec(code: StringLike, globals_: object = PyObject.dict(), locals_: object = PyObject.dict()) -> PyObject: - """ - Copies the locals, execs the Python code, and returns the locals with any updates. - """ - - -Container: TypeAlias = Map | Maybe | Pair | Set | MultiSet | Vec | UnstableFn -Primitive: TypeAlias = String | Bool | i64 | f64 | Rational | BigInt | BigRat | PyObject | Unit + if _path[:-1] != tuple(__name__.split(".")): + continue + _name = _path[-1] + if not _name.isidentifier() or _name in globals(): + raise ValueError(f"Cannot generate Python builtin type {_path!r}") + _class = RuntimeClass(_definition) + globals()[_name] = _class + __all__ += [_name] # noqa: PLE0604 -- names come from the generated catalog + if _kind.value.name in _LITERAL_PYTHON_TYPES: + _class.__egg_hooks__.update(value=property(_scalar_value), eval=_eval) + _literal_types = _LITERAL_PYTHON_TYPES[_kind.value.name] + _like: Any = _class + for _literal_type in _literal_types: + _like = _like | _literal_type + if _kind.value.name == "f64": + converter(_literal_type, _class, lambda value, target=_class: target(float(value))) + else: + converter(_literal_type, _class, _class) + globals()[f"{_name}Like"] = _like + __all__ += [f"{_name}Like"] # noqa: PLE0604 + if _kind.value.name == "i64": + _class.__egg_hooks__.update(__index__=_integer, __int__=_integer) + if _kind.value.name == "f64": + _class.__egg_hooks__.update(__float__=_float, __int__=_integer) + +for _path, (_definition, _view) in _catalog.functions.items(): + if _path[:-1] == tuple(__name__.split(".")): + if _path[-1] in globals(): + raise ValueError(f"Python builtin binding collision: {_path!r}") + globals()[_path[-1]] = RuntimeFunction(_definition, _view) + __all__ += [_path[-1]] # noqa: PLE0604 diff --git a/python/egglog/conversion.py b/python/egglog/conversion.py index 803ebabb..3653ab44 100644 --- a/python/egglog/conversion.py +++ b/python/egglog/conversion.py @@ -1,267 +1,113 @@ +"""Python conversion policy over canonical protobuf sort references.""" + from __future__ import annotations -from collections import defaultdict -from collections.abc import Callable +from collections.abc import Callable, Generator from contextlib import contextmanager from contextvars import ContextVar -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, TypeVar, cast - -from .declarations import * -from .pretty import * -from .runtime import * -from .thunk import * +from typing import Any, TypeVar, cast -if TYPE_CHECKING: - from collections.abc import Generator +from egglog_proto.egglog.v1 import egglog_pb as pb - from .egraph import BaseExpr +from ._program import Ref, StructuralView +from .runtime import RuntimeClass, RuntimeExpr, _class_for_sort __all__ = ["ConvertError", "convert", "converter", "get_type_args"] -# Mapping from (source type, target type) to and function which takes in the runtimes values of the source and return the target -CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable[[Any], RuntimeExpr]]] = {} -# Global declarations to store all convertible types so we can query if they have certain methods or not -_CONVERSION_DECLS = Declarations.create() -# Defer a list of declarations to be added to the global declarations, so that we can not trigger them processing -# until we need them -_TO_PROCESS_DECLS: list[DeclarationsLike] = [] - -def retrieve_conversion_decls() -> Declarations: - _CONVERSION_DECLS.update(*_TO_PROCESS_DECLS) - _TO_PROCESS_DECLS.clear() - return _CONVERSION_DECLS - - -T = TypeVar("T") -V = TypeVar("V", bound="BaseExpr") +V = TypeVar("V") +CONVERSIONS: dict[tuple[type | RuntimeClass, RuntimeClass], tuple[int, Callable[[Any], RuntimeExpr]]] = {} +TYPE_ARGS = ContextVar[tuple[RuntimeClass, ...]]("TYPE_ARGS", default=()) class ConvertError(Exception): pass -def converter(from_type: type[T], to_type: type[V], fn: Callable[[T], V], cost: int = 1) -> None: - """ - Register a converter from some type to an egglog type. - """ - to_type_name = process_tp(to_type) - if not isinstance(to_type_name, JustTypeRef): - raise TypeError(f"Expected return type to be a egglog type, got {to_type_name}") - _register_converter(process_tp(from_type), to_type_name, cast("Callable[[Any], RuntimeExpr]", fn), cost) - - -def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable[[Any], RuntimeExpr], cost: int) -> None: - """ - Registers a converter from some type to an egglog type, if not already registered. - - Also adds transitive converters, i.e. if registering A->B and there is already B->C, then A->C will be registered. - Also, if registering A->B and there is already D->A, then D->B will be registered. - """ - if a == b: +def converter(from_type: type | RuntimeClass, to_type: type | RuntimeClass, fn: Callable, cost: int = 1) -> None: + """Register a Python converter; target semantics remain owned sort records.""" + if not isinstance(to_type, RuntimeClass): + raise TypeError(f"Expected an egglog return type, got {to_type!r}") + if from_type == to_type: return - if (a, b) in CONVERSIONS and CONVERSIONS[(a, b)][0] <= cost: + key = from_type, to_type + if key in CONVERSIONS and CONVERSIONS[key][0] <= cost: return - CONVERSIONS[(a, b)] = (cost, a_b) - for (c, d), (other_cost, c_d) in list(CONVERSIONS.items()): - if _is_type_compatible(b, c): - _register_converter( - a, d, _ComposedConverter(a_b, c_d, c.args if isinstance(c, JustTypeRef) else ()), cost + other_cost - ) - if _is_type_compatible(a, d): - _register_converter( - c, b, _ComposedConverter(c_d, a_b, a.args if isinstance(a, JustTypeRef) else ()), cost + other_cost - ) - - -def _is_type_compatible(source: type | JustTypeRef, target: type | JustTypeRef) -> bool: - """ - Types must be equal or also support unbound to bound typevar like B -> B[C] - """ - if source == target: - return True - if isinstance(source, JustTypeRef) and isinstance(target, JustTypeRef) and source.args and not target.args: - return source.ident == target.ident - # TODO: Support case where B[T] where T is typevar is mapped to B[C] - return False - - -@dataclass -class _ComposedConverter: - """ - A converter which is composed of multiple converters. - - _ComposeConverter(a_b, b_c) is equivalent to lambda x: b_c(a_b(x)) - - We use the dataclass instead of the lambda to make it easier to debug. - """ - - a_b: Callable[[Any], RuntimeExpr] - b_c: Callable[[Any], RuntimeExpr] - b_args: tuple[JustTypeRef, ...] - - def __call__(self, x: Any) -> RuntimeExpr: - # if we have A -> B and B[C] -> D then we should use (C,) as the type args - # when converting from A -> B - if self.b_args: - with with_type_args(self.b_args, retrieve_conversion_decls): - first_res = self.a_b(x) - else: - first_res = self.a_b(x) - return self.b_c(first_res) - - def __str__(self) -> str: - return f"{self.b_c} ∘ {self.a_b}" + CONVERSIONS[key] = cost, fn + # Preserve the existing transitive conversion policy without storing type + # trees or synthesizing a callable signature. + for (source, target), (other_cost, other_fn) in tuple(CONVERSIONS.items()): + if target == from_type: + + def composed(value: Any, first=other_fn, second=fn) -> RuntimeExpr: + return second(first(value)) + + converter(source, to_type, composed, cost + other_cost) + if to_type == source: + + def composed(value: Any, first=fn, second=other_fn) -> RuntimeExpr: + return second(first(value)) + + converter(from_type, target, composed, cost + other_cost) + + +def resolve_literal(sort: Ref, value: object) -> RuntimeExpr: + """Resolve Python values against a canonical concrete expected sort.""" + if not isinstance(sort, Ref) or sort.role != "sorts": + msg = "Conversions require a canonical sort reference" + raise TypeError(msg) + if isinstance(value, RuntimeExpr): + node: pb.Node = value.__egg_ref__.read() + actual = value.__egg_ref__.owner.ref("sorts", node.sort_id) + if StructuralView(actual) == StructuralView(sort): + return value + source: type | RuntimeClass = _class_for_sort(actual) + else: + source = type(value) + target = _class_for_sort(sort) + candidates = source.__mro__ if isinstance(source, type) and not isinstance(source, RuntimeClass) else (source,) + source_sort: pb.Sort = sort.read() + assert source_sort.kind is not None + arguments = ( + tuple(_class_for_sort(sort.owner.ref("sorts", index)) for index in source_sort.kind.value.args) + if source_sort.kind.field == "family" + else () + ) + for candidate in candidates: + conversion = CONVERSIONS.get((candidate, target)) + if conversion is not None: + with with_type_args(arguments): + result = conversion[1](value) + if not isinstance(result, RuntimeExpr): + raise ConvertError(f"Converter to {target} returned {type(result).__name__}, not an expression") + payload: pb.Node = result.__egg_ref__.read() + actual = result.__egg_ref__.owner.ref("sorts", payload.sort_id) + if StructuralView(actual) != StructuralView(sort): + raise ConvertError(f"Converter to {target} returned the wrong canonical sort") + return result + raise ConvertError(f"Cannot convert {value} of type {source} to {target}") def convert(source: object, target: type[V]) -> V: - """ - Convert a source object to a target type. - """ - runtime_target: object = target - assert isinstance(runtime_target, RuntimeClass) - return cast("V", resolve_literal(runtime_target.__egg_tp__, source, runtime_target.__egg_decls_thunk__)) + runtime_target = cast("object", target) + if not isinstance(runtime_target, RuntimeClass) or runtime_target.__egg_sort__ is None: + msg = "Conversion target must be a concrete egglog type" + raise TypeError(msg) + return cast("V", resolve_literal(runtime_target.__egg_sort__, source)) def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr: - """ - Convert a source object to the same type as the target. - """ - tp = target.__egg_typed_expr__.tp - return resolve_literal(tp.to_var(), source, Thunk.value(target.__egg_decls__)) - - -def process_tp(tp: type | RuntimeClass) -> JustTypeRef | type: - """ - Process a type before converting it, to add it to the global declarations and resolve to a ref. - """ - if isinstance(tp, RuntimeClass): - _TO_PROCESS_DECLS.append(tp) - egg_tp = tp.__egg_tp__ - return egg_tp.to_just() - return tp - - -def min_binary_conversion( - method_name: str, lhs: type | JustTypeRef, rhs: type | JustTypeRef -) -> tuple[Callable[[Any], RuntimeExpr], Callable[[Any], RuntimeExpr]] | None: - """ - Given a binary method and two starting types for the LHS and RHS, return a pair of callable which will convert - the LHS and RHS to appropriate types which support this method. If no such conversion is possible, return None. - - It should return the types which minimize the total conversion cost. If one of the types is a Python type, then - both of them can be converted. However, if both are egglog types, then only one of them can be converted. - """ - decls = retrieve_conversion_decls() - # tuple of (cost, convert lhs, convert rhs) - best_method: tuple[int, Callable[[Any], RuntimeExpr], Callable[[Any], RuntimeExpr]] | None = None - - possible_lhs = _all_conversions_from(lhs) if isinstance(lhs, type) else [(0, lhs, identity)] - possible_rhs = _all_conversions_from(rhs) if isinstance(rhs, type) else [(0, rhs, identity)] - for lhs_cost, lhs_converted_type, lhs_convert in possible_lhs: - # Start by checking if we have a LHS that matches exactly and a RHS which can be converted - if (desired_other_type := decls.check_binary_method_with_self_type(method_name, lhs_converted_type)) and ( - converter := _lookup_conversion(rhs, desired_other_type) - ): - cost = lhs_cost + converter[0] - if best_method is None or best_method[0] > cost: - best_method = (cost, lhs_convert, converter[1]) - - for rhs_cost, rhs_converted_type, rhs_convert in possible_rhs: - # Next see if it's possible to convert the LHS and keep the RHS as is - for desired_self_type in decls.check_binary_method_with_other_type(method_name, rhs_converted_type): - if converter := _lookup_conversion(lhs, desired_self_type): - cost = rhs_cost + converter[0] - if best_method is None or best_method[0] > cost: - best_method = (cost, converter[1], rhs_convert) - if best_method is None: - return None - return best_method[1], best_method[2] - - -def _all_conversions_from(tp: JustTypeRef | type) -> list[tuple[int, JustTypeRef, Callable[[Any], RuntimeExpr]]]: - """ - Get all conversions from a type to other types. - - Returns a list of tuples of (cost, target type, conversion function). - """ - return [ - (cost, target, fn) - for (source, target), (cost, fn) in CONVERSIONS.items() - if (issubclass(tp, source) if isinstance(tp, type) and isinstance(source, type) else source == tp) - ] - - -def identity(x: Any) -> Any: - return x - - -TYPE_ARGS = ContextVar[tuple[RuntimeClass, ...]]("TYPE_ARGS") - - -def get_type_args() -> tuple[type, ...]: - """ - Get the type args for the type being converted. - """ - return cast("tuple[type, ...]", TYPE_ARGS.get()) + node: pb.Node = target.__egg_ref__.read() + return resolve_literal(target.__egg_ref__.owner.ref("sorts", node.sort_id), source) @contextmanager -def with_type_args(args: tuple[JustTypeRef, ...], decls: Callable[[], Declarations]) -> Generator[None, None, None]: - token = TYPE_ARGS.set(tuple(RuntimeClass(decls, a.to_var()) for a in args)) +def with_type_args(arguments: tuple[RuntimeClass, ...]) -> Generator[None, None, None]: + token = TYPE_ARGS.set(arguments) try: yield finally: TYPE_ARGS.reset(token) -def resolve_literal( - tp: TypeOrVarRef, - arg: object, - decls: Callable[[], Declarations] = retrieve_conversion_decls, -) -> RuntimeExpr: - """ - Try to convert an object to a type, raising a ConvertError if it is not possible. - - If it cannot be resolved, we assume that the value passed in will resolve it. - """ - # If this is a runtime expression that could match the type already, just return it - if isinstance(arg, RuntimeExpr) and tp.matches_just({}, arg.__egg_typed_expr__.tp): - return arg - tp_just = tp.to_just() - if arg is DUMMY_VALUE: - return RuntimeExpr.__from_values__(decls(), TypedExprDecl(tp_just, DummyDecl())) - arg_type = resolve_type(arg) - if (conversion := _lookup_conversion(arg_type, tp_just)) is not None: - with with_type_args(tp_just.args, decls): - return conversion[1](arg) - raise ConvertError(f"Cannot convert {arg} of type {arg_type} to {tp_just}") - - -def _lookup_conversion(lhs: type | JustTypeRef, rhs: JustTypeRef) -> tuple[int, Callable[[Any], RuntimeExpr]] | None: - """ - Looks up a conversion function for the given types. - - Also looks up all parent types of the lhs if it is a Python type and looks up more general not parametrized types for rhs. - """ - for lhs_type in lhs.__mro__ if isinstance(lhs, type) else [lhs]: - if (key := (lhs_type, rhs)) in CONVERSIONS: - return CONVERSIONS[key] - if rhs.args and (key := (lhs_type, JustTypeRef(rhs.ident))) in CONVERSIONS: - return CONVERSIONS[key] - return None - - -def _debug_print_converters(): - """ - Prints a mapping of all source types to target types that have a conversion function. - """ - source_to_targets = defaultdict(list) - for source, target in CONVERSIONS: - source_to_targets[source].append(target) - - -def resolve_type(x: object) -> JustTypeRef | type: - if isinstance(x, RuntimeExpr): - return x.__egg_typed_expr__.tp - return type(x) +def get_type_args() -> tuple[RuntimeClass, ...]: + return TYPE_ARGS.get() diff --git a/python/egglog/deconstruct.py b/python/egglog/deconstruct.py index 708ad213..ad3d2485 100644 --- a/python/egglog/deconstruct.py +++ b/python/egglog/deconstruct.py @@ -1,27 +1,14 @@ -""" -Utility functions to deconstruct expressions in Python. -""" +"""Inspect owned expression records locally, without evaluating expressions.""" from __future__ import annotations +import struct from collections.abc import Callable -from functools import partial -from typing import TYPE_CHECKING, TypeVar, Unpack, cast, overload +from typing import Any -import cloudpickle -from typing_extensions import TypeVarTuple +from egglog_proto.egglog.v1 import egglog_pb as pb -from .declarations import * -from .egraph import BaseExpr, Expr -from .runtime import * -from .thunk import * - -if TYPE_CHECKING: - from .builtins import Bool, PyObject, String, UnstableFn, f64, i64 - - -T = TypeVar("T", bound=BaseExpr) -TS = TypeVarTuple("TS", default=Unpack[tuple[BaseExpr, ...]]) +from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction, _class_for_sort, _definition_at __all__ = [ "get_callable_args", @@ -33,168 +20,111 @@ ] -@overload -def get_literal_value(x: String) -> str | None: ... - - -@overload -def get_literal_value(x: Bool) -> bool | None: ... - - -@overload -def get_literal_value(x: i64) -> int | None: ... - - -@overload -def get_literal_value(x: f64) -> float | None: ... - - -@overload -def get_literal_value(x: PyObject) -> object: ... - - -@overload -def get_literal_value(x: UnstableFn[T, *TS]) -> Callable[[Unpack[TS]], T] | None: ... - - -@overload -def get_literal_value(x: Expr) -> None: ... - - def get_literal_value(x: object) -> object: - """ - Returns the literal value of an expression if it is a literal. - If it is not a literal, returns None. - """ + """Return a local literal payload, or None for symbolic expressions.""" if not isinstance(x, RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match x.__egg_typed_expr__.expr: - case LitDecl(v): - return v - case PyObjectDecl(obj): - return cloudpickle.loads(obj) - case PartialCallDecl(call): - fn, args = _deconstruct_call_decl(x.__egg_decls_thunk__, call) - if not args: - return fn - return partial(fn, *args) - return None - - -def get_constant_name(x: BaseExpr) -> Ident | None: - """ - Check if the expression is a constant and return its name. - If it is not a constant, return None. - """ - if not isinstance(cast("object", x), RuntimeExpr): + node: pb.Node = x.__egg_ref__.read() + if node.kind is None or node.kind.field != "primitive_value": + return None + value = node.kind.value.value + if value is None: + msg = "Primitive value has no payload" + raise ValueError(msg) + match value.field: + case "i64" | "string" | "bool": + return value.value + case "f64_bits": + return struct.unpack("!d", struct.pack("!Q", value.value))[0] + case "unit": + return None + case _: + # Containers have their own .value protocols, retaining symbolic + # children. They are not scalar literals. + return None + + +def get_constant_name(x: object) -> object: + """Inspect a retained constant name (constant presentation is not migrated).""" + if not isinstance(x, RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match cast("RuntimeExpr", x).__egg_typed_expr__.expr: - case CallDecl(ConstantRef(ident)): - return ident - return None + msg = "Constant presentation is not migrated" + raise NotImplementedError(msg) -def get_let_name(x: BaseExpr) -> str | None: - """ - Check if the expression is a `let` expression and return the name of the variable. - If it is not a `let` expression, return None. - """ - if not isinstance(cast("object", x), RuntimeExpr): +def get_let_name(x: object) -> str | None: + """Inspect a retained local name (the local-let policy is not migrated).""" + if not isinstance(x, RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match cast("RuntimeExpr", x).__egg_typed_expr__.expr: - case LetRefDecl(name): - return name - return None + msg = "Retained local-let presentation is not migrated" + raise NotImplementedError(msg) -def get_var_name(x: BaseExpr) -> str | None: - """ - Check if the expression is a variable and return its name. - If it is not a variable, return None. - """ - if not isinstance(cast("object", x), RuntimeExpr): +def get_var_name(x: object) -> str | None: + """Return a variable node's canonical name, or None for other node kinds.""" + if not isinstance(x, RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match cast("RuntimeExpr", x).__egg_typed_expr__.expr: - case UnboundVarDecl(name, _egg_name): - return name - return None - - -def get_callable_fn(x: T) -> Callable[..., T] | T | None: - """ - Gets the function of an expression, or if it's a constant or classvar, return that. - """ - if not isinstance(cast("object", x), RuntimeExpr): + node: pb.Node = x.__egg_ref__.read() + return node.kind.value if node.kind is not None and node.kind.field == "var" else None + + +def _deconstruct_call(x: RuntimeExpr) -> tuple[RuntimeClass | RuntimeFunction, tuple[RuntimeExpr, ...]] | None: + """Select a declaration's Python view and restore its argument ordering.""" + node: pb.Node = x.__egg_ref__.read() + if node.kind is None or node.kind.field != "call": + return None + call = node.kind.value + declaration = _definition_at(x.__egg_ref__, "callable", call.func) + record: pb.Declaration = declaration.read() + arguments = tuple(RuntimeExpr(x.__egg_ref__.owner.ref("nodes", index)) for index in call.args) + if record.bindings is None or record.bindings.python is None or not record.bindings.python.views: + return RuntimeFunction(declaration), arguments + view = record.bindings.python.views[0] + ordered: list[RuntimeExpr] = [] + if view.has_field("receiver"): + ordered.append(arguments[view.receiver]) + kind = record.kind.value if record.kind is not None else None + fixed = ( + len(kind.typing.value.inputs) + if isinstance(kind, pb.HostPrimitive) and kind.typing is not None and kind.typing.field == "signature" + else len(arguments) + ) + for parameter in view.params: + ordered.extend( + arguments[parameter.core_input :] if parameter.core_input == fixed else (arguments[parameter.core_input],) + ) + if view.kind == pb.PythonCallKind.INITIALIZER: + return _class_for_sort(x.__egg_ref__.owner.ref("sorts", node.sort_id)), tuple(ordered) + owner = None + if view.owner is not None and view.owner.kind is not None and view.owner.kind.field == "sort": + if view.has_field("receiver"): + receiver = arguments[view.receiver].__egg_ref__ + owner = _class_for_sort(receiver.owner.ref("sorts", receiver.read().sort_id)) + else: + owner = _class_for_sort(declaration.owner.ref("sorts", view.owner.kind.value)) + return RuntimeFunction(declaration, 0, owner), tuple(ordered) + + +def get_callable_fn(x: object) -> Callable[..., Any] | None: + """Get the generated callable responsible for a call expression.""" + if not isinstance(x, RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - runtime_x = cast("RuntimeExpr", x) - match runtime_x.__egg_typed_expr__.expr: - case CallDecl() as call: - fn, _ = _deconstruct_call_decl(runtime_x.__egg_decls_thunk__, call) - return fn - return None + result = _deconstruct_call(x) + return None if result is None else result[0] -@overload -def get_callable_args(x: T, fn: None = ...) -> tuple[BaseExpr, ...] | None: ... - - -@overload -def get_callable_args(x: T, fn: Callable[[Unpack[TS]], T]) -> tuple[*TS] | None: ... - - -def get_callable_args(x: T, fn: Callable[[Unpack[TS]], T] | None = None) -> tuple[*TS] | None: - """ - Gets all the arguments of an expression. - If a function is provided, it will only return the arguments if the expression is a call - to that function. - - Note that recursively calling the arguments is the safe way to walk the expression tree. - """ - if not isinstance(cast("object", x), RuntimeExpr): +def get_callable_args(x: object, fn: object = None) -> tuple[RuntimeExpr, ...] | None: + """Return symbolic children in Python argument order, without evaluation.""" + if not isinstance(x, RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - runtime_x = cast("RuntimeExpr", x) - match runtime_x.__egg_typed_expr__.expr: - case CallDecl() as call: - actual_fn, args = _deconstruct_call_decl(runtime_x.__egg_decls_thunk__, call) - if fn is None: - return cast("tuple[*TS]", args) - # Compare functions and classes without considering bound type parameters, so that you can pass - # in a binding like Vec[i64] and match Vec[i64](...) or Vec(...) calls. - if isinstance(cast("object", actual_fn), RuntimeFunction): - runtime_actual_fn = cast("RuntimeFunction", actual_fn) - if ( - isinstance(cast("object", fn), RuntimeFunction) - and runtime_actual_fn.__egg_ref__ == cast("RuntimeFunction", fn).__egg_ref__ - ): - return cast("tuple[*TS]", args) - if isinstance(cast("object", actual_fn), RuntimeClass): - runtime_actual_cls = cast("RuntimeClass", actual_fn) - if ( - isinstance(cast("object", fn), RuntimeClass) - and runtime_actual_cls.__egg_tp__.ident == cast("RuntimeClass", fn).__egg_tp__.ident - ): - return cast("tuple[*TS]", args) + result = _deconstruct_call(x) + if result is None: + return None + actual, arguments = result + if fn is None: + return arguments + if isinstance(actual, RuntimeClass) and isinstance(fn, RuntimeClass): + return arguments if actual.__egg_definition__ == fn.__egg_definition__ else None + if isinstance(actual, RuntimeFunction) and isinstance(fn, RuntimeFunction): + return arguments if actual.__egg_ref__ == fn.__egg_ref__ else None return None - - -def _deconstruct_call_decl( - decls_thunk: Callable[[], Declarations], call: CallDecl -) -> tuple[Callable, tuple[object, ...]]: - """ - Deconstructs a CallDecl into a runtime callable and its arguments. - """ - args = call.args - arg_exprs = tuple(RuntimeExpr(decls_thunk, Thunk.value(a)) for a in args) - # TODO: handle values? Like constants - if isinstance(call.callable, InitRef): - return RuntimeClass( - decls_thunk, - TypeRefWithVars(call.callable.ident, tuple(tp.to_var() for tp in (call.bound_tp_params or []))), - ), arg_exprs - egg_bound = ( - JustTypeRef(call.callable.ident, call.bound_tp_params) - if isinstance(call.callable, (ClassMethodRef, MethodRef)) and call.bound_tp_params - else None - ) - - return RuntimeFunction(decls_thunk, Thunk.value(call.callable), egg_bound), arg_exprs diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 535c24e5..f282c628 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -4,6 +4,7 @@ import dis import inspect import pathlib +import re import sys import tempfile from collections.abc import Callable, Generator, Iterable, Sequence @@ -34,20 +35,25 @@ from warnings import warn import graphviz +from egglog_proto.egglog.v1 import egglog_pb as pb from opentelemetry import trace +from protobuf import Oneof from typing_extensions import ParamSpec, TypeForm from . import bindings +from ._catalog import Catalog +from ._program import Builder, Ref, StructuralView from ._tracing import call_with_current_trace from .conversion import * from .conversion import convert_to_same_type, resolve_literal from .declarations import * -from .declarations import _BUILTIN_EGG_FN_NAMES, _BUILTIN_EGG_SORT_NAMES, is_callable_decl_constructor +from .declarations import _BUILTIN_EGG_FN_NAMES, is_callable_decl_constructor from .egraph_state import * from .ipython_magic import IN_IPYTHON -from .pretty import pretty_decl +from .pretty import pretty_decl, pretty_ref from .run_report import RunReport from .runtime import * +from .runtime import _ClassContext from .thunk import * if TYPE_CHECKING: @@ -357,7 +363,7 @@ class _ExprMetaclass(type): Used to override isistance checks, so that runtime expressions are instances of Expr at runtime. """ - def __new__( # type: ignore[misc] + def __new__( cls: type[_ExprMetaclass], name: str, bases: tuple[type, ...], @@ -368,38 +374,36 @@ def __new__( # type: ignore[misc] # If this is the Expr subclass, just return the class if not bases or bases == (BaseExpr,): return super().__new__(cls, name, bases, namespace) - builtin = BuiltinExpr in bases - if builtin and egg_sort is not None: - _BUILTIN_EGG_SORT_NAMES.add(egg_sort) - # Reserve explicit builtin names before declarations are lazily - # materialized so generated user names cannot claim them first. - for method in namespace.values(): - if isinstance(method, _WrappedMethod) and method.egg_fn is not None: - _BUILTIN_EGG_FN_NAMES.add(method.egg_fn) - frame = currentframe() assert frame prev_frame = frame.f_back assert prev_frame - cls_ident = Ident(name, _get_module(prev_frame)) - # Pass in an instance of the class so that when we are generating the decls - # we can update them eagerly so that we can access the methods in the class body - runtime_cls = RuntimeClass(None, TypeRefWithVars(cls_ident)) # type: ignore[arg-type] - - # Store frame so that we can get live access to updated locals/globals - # Otherwise, f_locals returns a copy - # https://peps.python.org/pep-0667/ - runtime_cls.__egg_decls_thunk__ = Thunk.fn( - _generate_class_decls, - namespace, - prev_frame, - builtin, - egg_sort, - cls_ident, - ruleset, - runtime_cls, + hooks = { + key: value.fn if isinstance(value, _WrappedMethod) else value + for key, value in namespace.items() + if key in ALWAYS_PRESERVED or (isinstance(value, _WrappedMethod) and value.preserve) + } + context = _ClassContext( + partial(_prepare_expr_class, namespace=namespace, frame=prev_frame, bases=bases, ruleset=ruleset), + hooks, + namespace.get("__match_args__", ()), ) - return runtime_cls + module = namespace.get("__module__", "") + path = [*module.split("."), name] if module else [name] + builder = Builder() + definition = builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["eq_sort"], pb.EqSort]( + "eq_sort", + pb.EqSort( + name=egg_sort or ".".join(path), bindings=pb.SortBindings(python=pb.TypeBinding(path=path)) + ), + ), + doc=namespace.get("__doc__") or "", + ), + ) + return RuntimeClass(builder.publish(keepalive=(context,)).ref("declarations", definition)) def __instancecheck__(cls, instance: object) -> bool: return isinstance(instance, RuntimeExpr) @@ -438,151 +442,98 @@ class Expr(BaseExpr, metaclass=_ExprMetaclass): """ -def _generate_class_decls( # noqa: C901,PLR0912 +def _prepare_expr_class( # noqa: C901, PLR0912 + runtime_cls: RuntimeClass, + *, namespace: dict[str, Any], frame: FrameType, - builtin: bool, - egg_sort: str | None, - cls_ident: Ident, + bases: tuple[type, ...], ruleset: Ruleset | None, - runtime_cls: RuntimeClass, -) -> Declarations: - """ - Lazy constructor for class declarations to support classes with methods whose types are not yet defined. - """ - parameters: list[TypeVar] = ( - # Get the generic params from the orig bases generic class - namespace["__orig_bases__"][1].__parameters__ if "__orig_bases__" in namespace else [] - ) - type_vars = tuple(TypeVarRef.from_type_var(p) for p in parameters) - del parameters - cls_decl = ClassDecl( - egg_sort, type_vars, builtin, match_args=namespace.pop("__match_args__", ()), doc=namespace.pop("__doc__", None) - ) - decls = Declarations(_classes={cls_ident: cls_decl}) - # Update class thunk eagerly when resolving so that lookups work in methods. - runtime_cls.__egg_decls_thunk__ = Thunk.value(decls) - # Cached RuntimeFunction/RuntimeExpr wrappers capture the current decl thunk, so - # swapping in the concrete declarations must invalidate any wrappers created while - # the class was still pointing at the lazy declaration builder. - runtime_cls.__egg_attr_cache__.clear() - - ## - # Register class variables - ## - # Create a dummy type to pass to get_type_hints to resolve the annotations we have - _Dummytype = type("_DummyType", (), {"__annotations__": namespace.get("__annotations__", {})}) - # Resolve annotations before sharing this class with worker threads. - for k, v in get_type_hints(_Dummytype, globalns=frame.f_globals, localns=frame.f_locals).items(): - if getattr(v, "__origin__", None) == ClassVar: - (inner_tp,) = v.__args__ - type_ref = resolve_type_annotation_mutate(decls, inner_tp) - default_value = namespace.pop(k, _MISSING) - has_default = default_value is not _MISSING - return_type_is_eqsort = isinstance(type_ref, TypeRefWithVars) and not decls._classes[type_ref.ident].builtin - if has_default and ruleset is not None and not return_type_is_eqsort: - msg = "Primitive-returning defaults cannot use an explicit ruleset" - raise ValueError(msg) - resolved_default = ( - resolve_literal(type_ref, default_value, Thunk.value(decls)) - if has_default and ruleset is None - else None - ) - if resolved_default is not None: - decls |= resolved_default - cls_decl.class_variables[k] = ConstantDecl( - type_ref.to_just(), - body=resolved_default.__egg_typed_expr__ if resolved_default is not None else None, - ) - if has_default and ruleset is not None: - _add_default_rewrite( - decls, ClassVariableRef(cls_ident, k), type_ref, default_value, ruleset, subsume=False - ) - else: - msg = f"On class {cls_ident}, for attribute '{k}', expected a ClassVar, but got {v}" - raise NotImplementedError(msg) - - ## - # Register methods, classmethods, preserved methods, and properties - ## - # Get all the methods from the class - filtered_namespace: list[tuple[str, Any]] = [ - (k, v) for k, v in namespace.items() if k not in IGNORED_ATTRIBUTES or isinstance(v, _WrappedMethod) - ] - - # all methods we should try adding default functions for - add_default_funcs: list[Callable[[], None]] = [] - # Then register each of its methods - for method_name, method in filtered_namespace: - is_init = method_name == "__init__" - # Don't register the init methods for literals, since those don't use the type checking mechanisms - if is_init and cls_ident in LIT_IDENTS: - continue - match method: - case _WrappedMethod(egg_fn, cost, merge, fn, preserve, mutates, unextractable, subsume, reverse_args): - pass - case _: - egg_fn, cost, merge = None, None, None - fn = method - unextractable, preserve, subsume = False, False, False - mutates = method_name in ALWAYS_MUTATES_SELF - reverse_args = False - if preserve or method_name in ALWAYS_PRESERVED: - cls_decl.preserved_methods[method_name] = fn +) -> Catalog: + """Compile one complete concrete constructor scope directly into protobuf.""" + if bases != (Expr,) or ruleset is not None or namespace.get("__annotations__"): + msg = "Only direct nongeneric Expr classes without class variables or rulesets are migrated" + raise NotImplementedError(msg) + initializer = None + egg_name = None + # Validate the complete source surface before publishing any member records. + for name, value in namespace.items(): + if name in IGNORED_ATTRIBUTES and not isinstance(value, _WrappedMethod): continue - locals = frame.f_locals - ref: ClassMethodRef | MethodRef | PropertyRef | InitRef - # TODO: Store deprecated message so we can get at runtime - if (getattr(fn, "__deprecated__", None)) is not None: - fn = fn.__wrapped__ # type: ignore[union-attr] - match fn: - case classmethod(): - ref = ClassMethodRef(cls_ident, method_name) - fn = fn.__func__ - case property(): - ref = PropertyRef(cls_ident, method_name) - fn = fn.fget - case _: - ref = InitRef(cls_ident) if is_init else MethodRef(cls_ident, method_name) - if isinstance(fn, _WrappedMethod): - msg = f"{cls_ident}.{method_name} Add the @method(...) decorator above @classmethod or @property" - - raise ValueError(msg) # noqa: TRY004 - special_function_name: SpecialFunctions | None = ( - "fn-partial" if egg_fn == "unstable-fn" else "fn-app" if egg_fn == "unstable-app" else None - ) - if special_function_name: - decl = FunctionDecl(special_function_name, builtin=True, egg_name=egg_fn) - decls.set_function_decl(ref, decl) + if isinstance(value, _WrappedMethod): + if ( + value.cost is not None + or value.merge is not None + or value.mutates_self + or value.unextractable + or value.subsume + or value.reverse_args + or (value.preserve and (value.egg_fn is not None or name == "__init__")) + ): + raise NotImplementedError(f"Method options for {runtime_cls}.{name} are not migrated") + fn = value.fn + preserve = value.preserve + supplied_name = value.egg_fn + else: + fn, preserve, supplied_name = value, False, None + if preserve or name in ALWAYS_PRESERVED: continue - try: - add_rewrite = _fn_decl( - decls, - egg_fn, - ref, - fn, - locals, - cost, - merge, - mutates, - builtin, - ruleset=ruleset, - unextractable=unextractable, - subsume=subsume, - reverse_args=reverse_args, + if name != "__init__" or not isinstance(fn, FunctionType) or _function_has_body(fn): + raise NotImplementedError( + f"Only bodyless initializers and preserved methods are migrated: {runtime_cls}.{name}" ) - except Exception as e: - e.add_note(f"Error processing {cls_ident}.{method_name}") - raise - - if not builtin: - add_default_funcs.append(add_rewrite) - - # Add all rewrite methods at the end so that all methods are registered first and can be accessed - # in the bodies - for add_rewrite in add_default_funcs: - add_rewrite() - return decls + initializer, egg_name = fn, supplied_name + if initializer is None: + return Catalog.from_refs([runtime_cls.__egg_definition__]) + parameters = list(signature(initializer).parameters.values()) + if not parameters or any(parameter.kind != Parameter.POSITIONAL_OR_KEYWORD for parameter in parameters): + msg = "Constructor parameters must be fixed positional-or-keyword parameters including self" + raise NotImplementedError(msg) + parameters = parameters[1:] + hints_namespace = {**initializer.__globals__, **frame.f_locals, runtime_cls.__name__: runtime_cls} + hints = get_type_hints(initializer, globalns=hints_namespace, localns=hints_namespace) + builder = Builder() + builder.import_ref(runtime_cls.__egg_definition__) + assert runtime_cls.__egg_sort__ is not None + output = builder.import_ref(runtime_cls.__egg_sort__) + inputs = [] + surface = [] + for index, parameter in enumerate(parameters): + sort = resolve_type_annotation(hints[parameter.name]) + if not isinstance(sort, Ref) or sort.role != "sorts": + raise NotImplementedError(f"Constructor input {parameter.name!r} requires a concrete sort") + inputs.append(pb.Arg(sort=builder.import_ref(sort), name=parameter.name)) + parameter_record = pb.PythonParameter(core_input=index, name=parameter.name) + if parameter.default is not Parameter.empty: + default = resolve_literal(sort, parameter.default) + parameter_record.default_expr = builder.import_ref(default.__egg_ref__) + surface.append(parameter_record) + qualified = f"{runtime_cls.__module__}.{runtime_cls.__name__}.__init__" + definition = builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["constructor"], pb.Constructor]( + "constructor", + pb.Constructor( + name=egg_name or re.sub(r"[^\w\-+*/?!=<>&|^/%]", "_", qualified), inputs=inputs, output=output + ), + ), + bindings=pb.CallableBindings( + python=pb.PythonBindings( + views=[ + pb.PythonCallable( + kind=pb.PythonCallKind.INITIALIZER, + owner=pb.BindingOwner(kind=Oneof[Literal["sort"], int]("sort", output)), + params=surface, + ) + ] + ) + ), + doc=initializer.__doc__ or "", + ), + ) + owner = builder.publish() + return Catalog.from_refs([runtime_cls.__egg_definition__, owner.ref("declarations", definition)]) @dataclass @@ -1131,13 +1082,12 @@ def __init__( num_threads: int = 1, no_decomp: bool = False, ) -> None: + if not seminaive or no_decomp or save_egglog_string: + msg = "Nondefault native execution flags and source transcripts are not migrated to the shared adapter" + raise NotImplementedError(msg) with _TRACER.start_as_current_span("create"): with _TRACER.start_as_current_span("create_bindings"): - self._state = EGraphState( - bindings.EGraph(seminaive=seminaive, num_threads=num_threads, no_decomp=no_decomp), - seminaive=seminaive, - save_egglog_string=save_egglog_string, - ) + self._state = EGraphState(num_threads=num_threads) self._state_stack = [] self._token_stack = [] self._cost_callback_values = None @@ -1145,8 +1095,8 @@ def __init__( self.register(*actions) def _add_decls(self, *decls: DeclarationsLike) -> None: - self._state.ensure_open() - self._state.add_declarations(*decls) + msg = "Legacy declaration submission is removed; this operation is not migrated" + raise NotImplementedError(msg) def set_report_level(self, level: bindings._ReportLevel) -> None: """ @@ -1156,7 +1106,7 @@ def set_report_level(self, level: bindings._ReportLevel) -> None: def num_threads(self) -> int: """Return the number of worker threads configured for this e-graph.""" - return self._egraph.num_threads() + return self._state.configure_resources() def set_num_threads(self, num_threads: int) -> None: """ @@ -1165,11 +1115,11 @@ def set_num_threads(self, num_threads: int) -> None: Passing ``1`` keeps execution serial. Passing ``0`` uses available parallelism. """ - self._egraph.set_num_threads(num_threads) + self._state.configure_resources(num_threads) def no_decomp(self) -> bool: """Return whether rule decomposition is disabled for this e-graph.""" - return self._egraph.no_decomp() + return False def set_no_decomp(self, no_decomp: bool) -> None: """Set whether subsequently registered rules skip decomposition.""" @@ -1194,7 +1144,8 @@ def close(self) -> None: commands are rejected because they could no longer be recorded. Otherwise this method has no effect and the e-graph remains usable. """ - self._state.close() + # With transcripts disabled, close historically leaves the graph usable. + # Native handles are released when their owning state is disposed. def _ipython_display_(self) -> None: self.display() @@ -1299,7 +1250,12 @@ def check(self, *facts: FactLike) -> None: """ Check if a fact is true in the egraph. """ - self._state.run_program(self._facts_to_check(facts)) + builder = Builder() + indices = [builder.import_ref(fact.__egg_ref__) for fact in _fact_likes(facts)] + command = builder.add( + "commands", pb.Command(kind=Oneof[Literal["check"], pb.Check]("check", pb.Check(facts=indices))) + ) + self._state.run_program([builder.publish().ref("commands", command)]) def check_fail(self, *facts: FactLike) -> None: """ @@ -1371,57 +1327,47 @@ def extract( Extract the lowest cost expression from the egraph. """ _extractor_options(extractor) - if extractor == "greedy-dag" and cost_model is not None and not isinstance(cost_model, DagCostModel): - msg = "The greedy-dag extractor requires a DagCostModel; a general TreeCostModel cannot be adapted" - raise TypeError(msg) + if cost_model is not None or extractor != "tree": + msg = "Custom Python cost models and non-tree extraction are not migrated" + raise NotImplementedError(msg) runtime_expr = to_runtime_expr(expr) - self._add_decls(runtime_expr) - tp = runtime_expr.__egg_typed_expr__.tp - if cost_model is None: - extract_report = self._run_extract(runtime_expr, 0, extractor) - assert isinstance(extract_report, bindings.ExtractBest) - res = self._from_termdag(extract_report.termdag, extract_report.term, tp) - cost = extract_report.cost - else: - if isinstance(runtime_expr.__egg_typed_expr__.expr, CallDecl): - # Register the root through the normal command path before computing costs, so shared subexpressions - # use synthetic lets and the extractor sees the already-materialized root value. - self.register(expr) - egg_sort = self._state.type_ref_to_egg(tp) - typed_expr = runtime_expr.__egg_typed_expr__ - if isinstance(typed_expr.expr, ValueDecl): - # Values returned by lookup_function_value already identify an e-graph value and cannot be lowered - # back into Egglog syntax. - value = self._state.typed_expr_to_value(typed_expr) - else: - # For call roots, evaluate the same let-factored presentation registered above. Keep - # typed_expr_to_value's direct lowering for non-registering callers such as lookup_function_value. - egg_expr = self._state.typed_expr_to_egg(typed_expr, expr_to_let=True) - value = call_with_current_trace(self._state.egraph.eval_expr, egg_expr)[1] - if isinstance(cost_model, DagCostModel): - termdag, extracted = call_with_current_trace( - bindings.extract_best_with_dag_cost_model, - self._state.egraph, - [(egg_sort, value)], - _DagCostModel(cost_model, self).to_bindings_cost_model(), - extractor=extractor, + builder = Builder() + root = builder.import_ref(runtime_expr.__egg_ref__) + index = builder.add( + "commands", + pb.Command( + kind=Oneof[Literal["extract"], pb.Extract]( + "extract", + pb.Extract(roots=[root], variants=1, extractor=pb.Extractor.TREE), ) - extracted_root = extracted[0] - if extracted_root is None: - msg = "Unextractable root" - raise ValueError(msg) - cost, term = extracted_root - else: - egg_cost_model = _CostModel(cost_model, self).to_bindings_cost_model() - tree_extractor = call_with_current_trace( - bindings.Extractor, [egg_sort], self._state.egraph, egg_cost_model - ) - termdag = bindings.TermDag() - cost, term = call_with_current_trace( - tree_extractor.extract_best, self._state.egraph, termdag, value, egg_sort - ) - res = self._from_termdag(termdag, term, tp) - return (res, cost) if include_cost else res + ), + ) + response, owner = self._state.run_program([builder.publish().ref("commands", index)]) + if ( + len(response.outputs) != 1 + or response.outputs[0].kind is None + or response.outputs[0].kind.field != "extraction" + ): + msg = "Shared adapter did not return the requested extraction" + raise ValueError(msg) + roots = response.outputs[0].kind.value.roots + if len(roots) != 1 or not roots[0].variants: + msg = "Unextractable root" + raise ValueError(msg) + variant = roots[0].variants[0] + result = cast("BASE_EXPR", RuntimeExpr(owner.ref("nodes", variant.term))) + if not include_cost: + return result + if not variant.has_field("cost"): + msg = "The extraction cost has no portable value" + raise ValueError(msg) + from .deconstruct import get_literal_value # noqa: PLC0415 + + cost = get_literal_value(RuntimeExpr(owner.ref("nodes", variant.cost))) + if not isinstance(cost, int): + msg = "Expected an i64 extraction cost" + raise TypeError(msg) + return result, cost def _from_termdag(self, termdag: bindings.TermDag, term: int, tp: JustTypeRef) -> Any: (new_typed_expr,) = self._state.exprs_from_egg(termdag, [term], tp) @@ -1558,17 +1504,21 @@ def push(self) -> None: """ Push the current state of the egraph, so that it can be popped later and reverted back. """ - self._state.run_program(bindings.Push(1)) + copied = self._state.copy() self._state_stack.append(self._state) - self._state = self._state.copy() + self._state = copied @_TRACER.start_as_current_span("pop") def pop(self) -> None: """ Pop the current state of the egraph, reverting back to the previous state. """ - self._state.run_program(bindings.Pop(span(1), 1)) + if not self._state_stack: + msg = "Cannot pop an EGraph without a pushed state" + raise IndexError(msg) + temporary = self._state self._state = self._state_stack.pop() + temporary.destroy() def __enter__(self) -> Self: """ @@ -1580,21 +1530,11 @@ def __enter__(self) -> Self: return self def __exit__(self, exc_type, exc, exc_tb) -> None: - egglog_file_state = self._state.egglog_file_state - if egglog_file_state is None or (not egglog_file_state.poisoned and not egglog_file_state.file.closed): - self.pop() - return - - # A closed or poisoned transcript rejects ordinary commands, but the - # backend scope still has to be restored. If another exception is - # already propagating, do not replace it with an ordinary cleanup failure. try: - call_with_current_trace(self._state.egraph.run_program, bindings.Pop(span(1), 1)) + self.pop() except Exception: if exc_type is None: raise - finally: - self._state = self._state_stack.pop() def _serialize( self, @@ -1732,11 +1672,13 @@ def to_json() -> str: @property def _egraph(self) -> bindings.EGraph: - return self._state.egraph + msg = "Direct native-AST access is removed; this operation is not migrated" + raise NotImplementedError(msg) @property def __egg_decls__(self) -> Declarations: - return self._state.__egg_decls__ + msg = "Legacy declaration projections are removed" + raise NotImplementedError(msg) @_TRACER.start_as_current_span("register") def register( @@ -2352,17 +2294,21 @@ def __repr__(self) -> str: return str(self) -@dataclass class Fact: """ A query on an EGraph, either by an expression or an equivalence between multiple expressions. """ - __egg_decls__: Declarations - fact: FactDecl + __slots__ = ("__egg_ref__",) + + def __init__(self, reference: Ref) -> None: + if reference.role != "nodes": + msg = "Facts require an owned protobuf node" + raise TypeError(msg) + self.__egg_ref__ = reference def __str__(self) -> str: - return pretty_decl(self.__egg_decls__, self.fact) + return pretty_ref(self.__egg_ref__) def __repr__(self) -> str: return str(self) @@ -2371,13 +2317,12 @@ def __bool__(self) -> bool: """ Returns True if the two sides of an equality are structurally equal. """ - match self.fact: - case EqDecl(_, left, right): - return left == right - case ExprFactDecl(TypedExprDecl(_, CallDecl(FunctionRef(name), (left_tp, right_tp)))) if ( - name == Ident.builtin("!=") - ): - return left_tp != right_tp + node: pb.Node = self.__egg_ref__.read() + if node.kind is not None and node.kind.field == "union" and len(node.kind.value.members) == 2: + left, right = node.kind.value.members + return StructuralView(self.__egg_ref__.owner.ref("nodes", left)) == StructuralView( + self.__egg_ref__.owner.ref("nodes", right) + ) msg = f"Can only check equality for == or != not {self}" raise ValueError(msg) @@ -2477,7 +2422,7 @@ def subsume(expr: Expr) -> Action: def expr_fact(expr: BaseExpr) -> Fact: runtime_expr = to_runtime_expr(expr) - return Fact(runtime_expr.__egg_decls__, ExprFactDecl(runtime_expr.__egg_typed_expr__)) + return Fact(runtime_expr.__egg_ref__) def union(lhs: EXPR) -> _UnionBuilder[EXPR]: @@ -2590,17 +2535,26 @@ class _EqBuilder(Generic[BASE_EXPR]): def to(self, other: BASE_EXPR) -> Fact: expr = to_runtime_expr(self.expr) other = convert_to_same_type(other, expr) - return Fact( - Declarations.create(expr, other), - EqDecl(expr.__egg_typed_expr__.tp, expr.__egg_typed_expr__.expr, other.__egg_typed_expr__.expr), + node: pb.Node = expr.__egg_ref__.read() + builder = Builder() + index = builder.add( + "nodes", + pb.Node( + sort_id=builder.import_ref(expr.__egg_ref__.owner.ref("sorts", node.sort_id)), + kind=Oneof[Literal["union"], pb.Union]( + "union", + pb.Union(members=[builder.import_ref(expr.__egg_ref__), builder.import_ref(other.__egg_ref__)]), + ), + ), ) + return Fact(builder.publish().ref("nodes", index)) def __repr__(self) -> str: return str(self) def __str__(self) -> str: expr = to_runtime_expr(self.expr) - return expr.__egg_pretty__("eq") + return f"eq({expr})" @dataclass @@ -2709,13 +2663,13 @@ def __str__(self) -> str: return f"rule({', '.join(args)})" -def expr_parts(expr: BaseExpr) -> TypedExprDecl: +def expr_parts(expr: BaseExpr) -> StructuralView: """ Returns the underlying type and decleration of the expression. Useful for testing structural equality or debugging. """ if not isinstance(cast("object", expr), RuntimeExpr): raise TypeError(f"Expected a RuntimeExpr not {expr}") - return cast("RuntimeExpr", expr).__egg_typed_expr__ + return StructuralView(cast("RuntimeExpr", expr).__egg_ref__) def to_runtime_expr(expr: BaseExpr) -> RuntimeExpr: diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py index 366b3a0f..9769e747 100644 --- a/python/egglog/egraph_state.py +++ b/python/egglog/egraph_state.py @@ -1,1907 +1,129 @@ """ -Implement conversion to/from egglog. +Bytes-only EGraph transport and response ownership. + +No expression, declaration, or command is reconstructed from Python semantic +dataclasses here. Language lookup indexes retain canonical declaration Refs; +installation compatibility and idempotence belong to the shared adapter. """ from __future__ import annotations -import contextlib -import math -import re -import tempfile -import weakref -from base64 import standard_b64decode, standard_b64encode -from dataclasses import InitVar, dataclass, field, replace -from pathlib import Path -from typing import TYPE_CHECKING, TextIO, assert_never, cast, overload -from uuid import UUID +from collections.abc import Iterable +from contextlib import suppress +from operator import index +from typing import Literal, SupportsIndex, cast -import cloudpickle -from opentelemetry import trace +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import Oneof from . import bindings -from ._tracing import call_with_current_trace -from .declarations import * -from .declarations import ( - _BUILTIN_EGG_FN_NAMES, - _BUILTIN_EGG_SORT_NAMES, - ConstructorDecl, - is_callable_decl_constructor, -) -from .pretty import * -from .type_constraint_solver import * - -if TYPE_CHECKING: - from collections.abc import Iterable - -__all__ = ["EGraphState", "span"] - - -_TRACER = trace.get_tracer(__name__) - -# These heads are interpreted as syntax before Egglog falls back to parsing a -# generic top-level call. Keep this aligned with Parser::parse_command, -# Parser::parse_action, Parser::parse_fact, and the extensions installed by -# egglog_experimental::new_experimental_egraph. -_EGGLOG_RESERVED_CALL_HEADS = frozenset({ - "=", - "birewrite", - "check", - "constructor", - "datatype", - "datatype*", - "delete", - "extract", - "fail", - "for", - "function", - "include", - "input", - "keep-best", - "let", - "let-scheduler", - "multi-extract", - "output", - "panic", - "pop", - "primitive", - "print-function", - "print-size", - "print-stats", - "print-table-stats", - "prove", - "prove-exists", - "push", - "relation", - "rewrite", - "rule", - "ruleset", - "run", - "run-schedule", - "set", - "set-cost", - "sort", - "subsume", - "union", - "unstable-combined-ruleset", - "unstable-fresh!", - "with-dynamic-cost", - "with-ruleset", -}) -_EGGLOG_RESERVED_ACTION_HEADS = frozenset({"delete", "let", "panic", "set", "set-cost", "subsume", "union"}) -_EGGLOG_RESERVED_COMMAND_OR_ACTION_HEADS = _EGGLOG_RESERVED_CALL_HEADS - {"="} -_EGGLOG_LITERAL_NAMES = frozenset({"false", "NaN", "inf", "-inf", "true"}) -_EGGLOG_NUMBER = re.compile(r"[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?") - - -@dataclass -class _SavedEgglogFile: - """ - Shared state for file-backed command execution. - - When `save_egglog_string=True`, the Python wrapper keeps one cumulative temp - `.egg` file per high-level egraph. Each command is reparsed-and-run from a - synthetic one-command program padded with blank lines so egglog error spans - point at the real file path and the command's true line numbers inside that - cumulative source log. - - We keep the append handle open for performance and for easy post-failure - inspection. Successful commands are saved normally; Egglog-reported - failures are saved as expected failures with trailing error comments. - Other execution failures invalidate the transcript because their partial - effects cannot be represented by a replayable command. - """ +from ._catalog import builtin_catalog +from ._program import Builder, Owner, Ref, _definition_key, pack - path: str - file: TextIO - line_count: int = 0 - poisoned: bool = False - _finalizer: weakref.finalize = field(init=False, repr=False) - - def __post_init__(self) -> None: - self._finalizer = weakref.finalize(self, _close_saved_egglog_file, self.file, self.path) - - def close(self) -> None: - self._finalizer() - - -def _close_saved_egglog_file(file: TextIO, path: str) -> None: - try: - file.close() - finally: - with contextlib.suppress(FileNotFoundError): - Path(path).unlink() - - -def _normalize_global_let_name(name: str) -> str: - return name if name.startswith("$") else f"${name}" - - -def _collect_explicit_backend_names(declarations: Declarations) -> set[str]: - """Collect names that declarations explicitly reserve in Egglog's shared symbol namespace.""" - names = { - decl.egg_name - for decl in (*declarations._functions.values(), *declarations._constants.values()) - if decl.egg_name is not None - } - for class_decl in declarations._classes.values(): - if class_decl.egg_name is not None: - names.add(class_decl.egg_name) - class_callables = ( - *class_decl.class_methods.values(), - *class_decl.class_variables.values(), - *class_decl.methods.values(), - *class_decl.properties.values(), - ) - names.update(decl.egg_name for decl in class_callables if decl.egg_name is not None) - if class_decl.init is not None and class_decl.init.egg_name is not None: - names.add(class_decl.init.egg_name) - return names +__all__ = ["EGraphState"] +_QUERY_RESOURCES = object() -def _egg_name_is_source_safe_symbol(name: str) -> bool: - """Return whether a name is safe to emit as an ordinary Egglog symbol.""" - return ( - bool(name) - and name != "_" - and not name.startswith(('"', "@")) - and not any(c.isspace() or c in ";()" for c in name) - ) +def _thread_count(value: object) -> int: + """Keep Python's integer-index acceptance within the protocol's uint32 domain.""" + count = index(cast("SupportsIndex", value)) + if count < 0: + msg = "can't convert negative int to unsigned" + raise OverflowError(msg) + if count > (1 << 32) - 1: + msg = "thread count exceeds the protocol's uint32 range" + raise OverflowError(msg) + return count -def _saved_egglog_failure_message(error: bindings.EggSmolError) -> str: - lines = [line.strip() for line in error.context.splitlines() if line.strip()] - for message in lines: - if "Failed" in message or "failed" in message: - return message - if len(lines) > 1 and lines[0].startswith("In "): - return lines[1] - if lines: - return lines[-1] - return str(error) - -def _rule_variable_names(rule: RuleDecl) -> set[str]: # noqa: C901, PLR0912 - """Collect source-level names that compiler-generated rule lets must not shadow.""" - names: set[str] = set() - - def visit_expr(expr: ExprDecl) -> None: - match expr: - case UnboundVarDecl(name, egg_name): - emitted_name = egg_name or f"_{name}" - names.add(emitted_name.removeprefix("$")) - case LetRefDecl(name): - names.add(_normalize_global_let_name(name).removeprefix("$")) - case CallDecl(_, args) | GetCostDecl(_, args) | PartialCallDecl(CallDecl(_, args)): - for arg in args: - visit_expr(arg.expr) - case LitDecl() | PyObjectDecl() | ValueDecl() | DummyDecl(): - pass - case _: - assert_never(expr) - - for fact in rule.body: - match fact: - case EqDecl(_, left, right): - visit_expr(left) - visit_expr(right) - case ExprFactDecl(typed_expr): - visit_expr(typed_expr.expr) - case _: - assert_never(fact) - for action in rule.head: - match action: - case LetDecl(name, typed_expr): - names.add(_normalize_global_let_name(name).removeprefix("$")) - visit_expr(typed_expr.expr) - case SetDecl(_, call, rhs): - visit_expr(call) - visit_expr(rhs) - case ExprActionDecl(typed_expr): - visit_expr(typed_expr.expr) - case ChangeDecl(_, call, _): - visit_expr(call) - case UnionDecl(_, left, right): - visit_expr(left) - visit_expr(right) - case SetCostDecl(_, call, cost): - visit_expr(call) - visit_expr(cost) - case PanicDecl(): - pass - case _: - assert_never(action) - return names - - -def span(frame_index: int = 0) -> bindings.RustSpan: - """ - Returns a span for the current file and line. - - If `frame_index` is passed, it will return the span for that frame in the stack, where 0 is the current frame - this is called in and 1 is the parent. - """ - # Currently disable this because it's too expensive. - # import inspect - - # frame = inspect.stack()[frame_index + 1] - return bindings.RustSpan("", 0, 0) - - -@dataclass class EGraphState: - """ - State of the EGraph declarations and rulesets, so when we pop/push the stack we know whats defined. - - Used for converting to/from egg and for pretty printing. - """ - - egraph: bindings.EGraph - seminaive: bool = True - # Opaque identity for values created in the current push scope. Descendant - # scopes may use ancestor values, but values from a popped scope must not be - # accepted if the backend later reuses their raw IDs. - value_owner: object = field(default_factory=object, repr=False) - valid_value_owners: frozenset[object] = field(default_factory=frozenset, repr=False) - save_egglog_string: InitVar[bool] = False - egglog_file_state: _SavedEgglogFile | None = field(default=None, repr=False) - # The declarations we have added. - __egg_decls__: Declarations = field(default_factory=Declarations) - # Explicit names are cached as declarations are merged. Generated names - # consult this set so a registration batch reserves all user-selected - # backend symbols without rescanning the complete declaration graph for - # every generated callable, sort, cost table, and synthetic let. - _explicit_backend_names: set[str] = field(default_factory=set, init=False, repr=False) - # Mapping of added rulesets to the added rules - rulesets: dict[Ident, set[RewriteOrRuleDecl]] = field(default_factory=dict) - # Persistent schedulers live outside a single run-schedule command; only emit their let once per active scope. - registered_persistent_schedulers: set[UUID] = field(default_factory=set) - - # Bidirectional mapping between egg function names and python callable references. - # Note that there are possibly multiple callable references for a single egg function name, like `+` - # for both int and rational classes. - egg_fn_to_callable_refs: dict[str, set[CallableRef]] = field( - default_factory=lambda: {"!=": {FunctionRef(Ident.builtin("!="))}} - ) - callable_ref_to_egg_fn: dict[CallableRef, tuple[str, bool]] = field( - default_factory=lambda: {FunctionRef(Ident.builtin("!=")): ("!=", False)} - ) - - # Bidirectional mapping between egg sort names and python type references. - type_ref_to_egg_sort: dict[JustTypeRef, str] = field(default_factory=dict) - egg_sort_to_type_ref: dict[str, JustTypeRef] = field(default_factory=dict) - - # Cache of direct structural egg expressions lowered without synthetic let substitutions. - expr_to_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict) - # Cache of synthetic let references introduced for top-level command lowering. - # This stays separate from `expr_to_egg_cache` so nested rule/rewrite lowering - # can always rebuild structural surface syntax instead of leaking a previously - # synthesized `$__expr_n` binding across contexts. - expr_to_letref_cache: dict[ExprDecl, bindings.Var] = field(default_factory=dict) - - # Callables with custom-cost tables and the reserved backend names for those tables. - cost_table_names: dict[CallableRef, str] = field(default_factory=dict) - # Counter for deterministic synthetic let bindings created while lowering expressions to egg. - expr_to_let_counter: int = 0 - # Explicit top-level lets waiting to be lowered in the current registration batch. - # Synthetic lets must avoid them regardless of command order within that batch. - pending_let_names: frozenset[str] = field(default_factory=frozenset, repr=False) - # Counter for deterministic synthetic names assigned to unnamed functions. - unnamed_function_counter: int = 0 - - # Counter for numeric rule names - rule_name_counter: int = 0 - # Mapping from numeric name (str) to command decl - rule_name_to_command_decl: dict[str, RuleDecl | BiRewriteDecl | RewriteDecl] = field(default_factory=dict) - - def __post_init__(self, save_egglog_string: bool) -> None: - if not self.valid_value_owners: - self.valid_value_owners = frozenset((self.value_owner,)) - self._explicit_backend_names = _collect_explicit_backend_names(self.__egg_decls__) - if save_egglog_string and self.egglog_file_state is None: - # Keep one persistent temp `.egg` file per high-level egraph so parse errors - # can point at a stable filename the user can open after a failure. - egglog_file = tempfile.NamedTemporaryFile( # noqa: SIM115 - kept open for incremental appends - mode="w+", encoding="utf-8", newline="", suffix=".egg", delete=False - ) - self.egglog_file_state = _SavedEgglogFile(egglog_file.name, cast("TextIO", egglog_file)) - - def copy(self) -> EGraphState: - """ - Returns a copy of the state. The egraph reference is kept the same. Used for pushing/popping. - """ - value_owner = object() - return EGraphState( - egraph=self.egraph, - seminaive=self.seminaive, - value_owner=value_owner, - valid_value_owners=self.valid_value_owners | {value_owner}, - save_egglog_string=self.egglog_file_state is not None, - egglog_file_state=self.egglog_file_state, - __egg_decls__=self.__egg_decls__.copy(), - rulesets={k: v.copy() for k, v in self.rulesets.items()}, - registered_persistent_schedulers=self.registered_persistent_schedulers.copy(), - egg_fn_to_callable_refs={k: v.copy() for k, v in self.egg_fn_to_callable_refs.items()}, - callable_ref_to_egg_fn=self.callable_ref_to_egg_fn.copy(), - type_ref_to_egg_sort=self.type_ref_to_egg_sort.copy(), - egg_sort_to_type_ref=self.egg_sort_to_type_ref.copy(), - expr_to_egg_cache=self.expr_to_egg_cache.copy(), - expr_to_letref_cache=self.expr_to_letref_cache.copy(), - cost_table_names=self.cost_table_names.copy(), - expr_to_let_counter=self.expr_to_let_counter, - pending_let_names=self.pending_let_names, - unnamed_function_counter=self.unnamed_function_counter, - rule_name_counter=self.rule_name_counter, - rule_name_to_command_decl=self.rule_name_to_command_decl.copy(), + """Own one native handle and the declaration refs needed to read responses.""" + + def __init__(self, *, num_threads: int = 1) -> None: + num_threads = _thread_count(num_threads) + self.transport = bindings._ProtoEngine() + # The public default cost sort is an ambient engine builtin. Creation + # references it directly; authored Programs still send full declarations. + request = pb.CreateEGraphRequest( + sorts=[pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))], + options=pb.EGraphOptions(cost_sort=0), + threads=num_threads, ) - - def add_declarations(self, *declarations_like: DeclarationsLike) -> None: - """Merge declarations while maintaining the explicit backend-name index.""" - attempted_update = False - try: - for declarations in declarations_like: - if declarations is None: - continue - attempted_update = True - self.__egg_decls__ |= declarations - finally: - if attempted_update: - # Rebuild after the batch rather than accumulating names: declaration - # merges may replace an earlier declaration and release its explicit - # name for later generated symbols. The finally path also keeps the - # index aligned when a later lazy declaration fails to resolve after - # an earlier declaration has already been merged. - self._explicit_backend_names = _collect_explicit_backend_names(self.__egg_decls__) - - def egglog_string(self) -> str: - if self.egglog_file_state is None: - msg = "Can't get egglog string unless EGraph created with save_egglog_string=True" - raise ValueError(msg) - if self.egglog_file_state.file.closed: - msg = "Can't get egglog string after the saved transcript has been closed" - raise ValueError(msg) - if self.egglog_file_state.poisoned: - msg = "Can't use the saved Egglog transcript after an execution failure whose partial effects cannot be replayed" - raise RuntimeError(msg) - # The append handle stays open for execution, so flush before reading the saved source. - self.egglog_file_state.file.flush() - with open(self.egglog_file_state.path, encoding="utf-8", newline="") as saved_file: - return saved_file.read() - - def close(self) -> None: - if self.egglog_file_state is not None: - self.egglog_file_state.close() - - def ensure_open(self) -> None: - if self.egglog_file_state is not None and self.egglog_file_state.file.closed: - msg = "Cannot run commands after the saved Egglog transcript has been closed" - raise ValueError(msg) - if self.egglog_file_state is not None and self.egglog_file_state.poisoned: - msg = "Cannot run commands after an execution failure whose partial effects cannot be replayed" + response = pb.CreateEGraphResponse.from_binary(self.transport.create(request.to_binary())) + self.handle: int | None = response.egraph_id + self.definitions = dict(builtin_catalog().definitions) + + def configure_resources(self, num_threads: object = _QUERY_RESOURCES) -> int: + """Query/update this handle over bytes and return the receiver's actual count.""" + if self.handle is None: + msg = "EGraph handle has been destroyed" raise RuntimeError(msg) - - def run_program(self, *commands: bindings._Command) -> list[bindings._CommandOutput]: - if not commands: - return [] - self.ensure_open() - if self.egglog_file_state is None: - return call_with_current_trace(self.egraph.run_program, *commands) - - outputs: list[bindings._CommandOutput] = [] - for command in commands: - command_text = str(command).rstrip("\n") + "\n" - start_line = self.egglog_file_state.line_count + 1 - - # Parse and run just this command in Rust, but pad it with blank lines - # so the span lines match its location in the cumulative saved source file. - padded_command = ("\n" * (start_line - 1)) + command_text - try: - command_outputs = call_with_current_trace( - self.egraph.parse_and_run_program, padded_command, filename=self.egglog_file_state.path - ) - except bindings.EggSmolError as error: - if not error.replayable_by_fail: - # Parsing, expansion, and typechecking can mutate backend - # metadata before they fail, while `(fail ...)` cannot - # reproduce those failures. Do not claim a replayable log. - self.egglog_file_state.poisoned = True - raise - fail_command_text = str(bindings.Fail(span(), command)).rstrip("\n") - saved_text = f"{fail_command_text} ; {_saved_egglog_failure_message(error)}\n" - try: - self.egglog_file_state.file.write(saved_text) - self.egglog_file_state.file.flush() - except BaseException: - self.egglog_file_state.poisoned = True - raise - self.egglog_file_state.line_count += saved_text.count("\n") - raise - except BaseException: - # A Python primitive or unrelated runtime failure can interrupt - # a command after some actions have committed. There is no - # replayable Egglog command for that partial state. - self.egglog_file_state.poisoned = True - raise - try: - self.egglog_file_state.file.write(command_text) - self.egglog_file_state.file.flush() - except BaseException: - # Execution already succeeded, so omitting even part of this - # command would make the cumulative source diverge. - self.egglog_file_state.poisoned = True - raise - self.egglog_file_state.line_count += command_text.count("\n") - outputs.extend(command_outputs) - return outputs - - @staticmethod - def _persistent_scheduler_name(scheduler: BackOffDecl) -> str: - return f"_persistent_scheduler_{scheduler.id.hex}" - - @staticmethod - def _local_scheduler_name(index: int) -> str: - return f"_scheduler_{index}" - - @staticmethod - def _back_off_scheduler_to_egg(scheduler: BackOffDecl) -> bindings.Call: - """Serialize the shared option protocol for local and persistent backoff schedulers.""" - args: list[bindings._Expr] = [] - if scheduler.match_limit is not None: - args.extend(( - bindings.Var(span(), ":match-limit"), - bindings.Lit(span(), bindings.Int(scheduler.match_limit)), - )) - if scheduler.ban_length is not None: - args.extend(( - bindings.Var(span(), ":ban-length"), - bindings.Lit(span(), bindings.Int(scheduler.ban_length)), - )) - return bindings.Call(span(), "back-off", args) - - @_TRACER.start_as_current_span("run_schedule_to_egg") - def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command: - """ - Turn a run schedule into an egg command. - - If there exists any custom schedulers in the schedule, it will be turned into a custom extract command otherwise - will be a normal run command. - """ - processed_schedule, persistent_schedulers = self._process_schedule(schedule) - if processed_schedule is None: - return bindings.RunSchedule(self._schedule_to_egg(schedule)) - for scheduler in persistent_schedulers: - if scheduler.id in self.registered_persistent_schedulers: - continue - self.run_program(self._persistent_scheduler_to_egg(scheduler)) - self.registered_persistent_schedulers.add(scheduler.id) - top_level_schedules = self._schedule_with_scheduler_to_egg(processed_schedule, []) - if len(top_level_schedules) == 1: - schedule_expr = top_level_schedules[0] - else: - schedule_expr = bindings.Call(span(), "seq", top_level_schedules) - return bindings.UserDefined(span(), "run-schedule", [schedule_expr]) - - def _process_schedule( # noqa: C901 - self, schedule: ScheduleDecl - ) -> tuple[ScheduleDecl | None, tuple[BackOffDecl, ...]]: - """ - Processes a schedule to determine if it contains any custom schedulers. - - If it does, it returns a new schedule with all the required let bindings added to the other scope. - If not, returns none. - - Also processes all rulesets in the schedule to make sure they are registered. - """ - bound_schedulers: list[UUID] = [] - unbound_schedulers: dict[UUID, BackOffDecl] = {} - persistent_schedulers: dict[UUID, BackOffDecl] = {} - has_bound_scheduler = False - - def helper(s: ScheduleDecl) -> None: - nonlocal has_bound_scheduler - match s: - case LetSchedulerDecl(scheduler, inner): - has_bound_scheduler = True - bound_schedulers.append(scheduler.id) - try: - return helper(inner) - finally: - bound_schedulers.pop() - case RunDecl(ruleset_name, _, scheduler): - self.ruleset_to_egg(ruleset_name) - if scheduler and scheduler.id not in bound_schedulers: - if scheduler.persistent: - persistent_schedulers[scheduler.id] = scheduler - else: - unbound_schedulers[scheduler.id] = scheduler - case SaturateDecl(inner) | RepeatDecl(inner, _): - return helper(inner) - case SequenceDecl(schedules): - for sc in schedules: - helper(sc) - case _: - assert_never(s) - return None - - helper(schedule) - if not has_bound_scheduler and not unbound_schedulers and not persistent_schedulers: - return None, () - for scheduler in unbound_schedulers.values(): - schedule = LetSchedulerDecl(scheduler, schedule) - return schedule, tuple(persistent_schedulers.values()) - - def _schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule: - msg = "Should never reach this, let schedulers should be handled by custom scheduler" - match schedule: - case SaturateDecl(schedule): - return bindings.Saturate(span(), self._schedule_to_egg(schedule)) - case RepeatDecl(schedule, times): - return bindings.Repeat(span(), times, self._schedule_to_egg(schedule)) - case SequenceDecl(schedules): - return bindings.Sequence(span(), [self._schedule_to_egg(s) for s in schedules]) - case RunDecl(ruleset_ident, until, scheduler): - if scheduler is not None: - raise ValueError(msg) - config = bindings.RunConfig( - str(ruleset_ident), None if not until else list(map(self.fact_to_egg, until)) - ) - return bindings.Run(span(), config) - case LetSchedulerDecl(): - raise ValueError(msg) - case _: - assert_never(schedule) - - def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 - self, schedule: ScheduleDecl, bound_schedulers: list[BackOffDecl] - ) -> list[bindings._Expr]: - """ - Turns a scheduler into an egg expression, to be used with a custom extract command. - - The bound_schedulers is a list of all the schedulers that have been bound. We can lookup their name as `_scheduler_{index}`. - """ - match schedule: - case LetSchedulerDecl(scheduler, inner): - name = self._local_scheduler_name(len(bound_schedulers)) - bound_schedulers.append(scheduler) - let_decl = bindings.Call( - span(), - "let-scheduler", - [bindings.Var(span(), name), self._back_off_scheduler_to_egg(scheduler)], - ) - try: - inner_exprs = self._schedule_with_scheduler_to_egg(inner, bound_schedulers) - finally: - bound_schedulers.pop() - return [bindings.Call(span(), "seq", [let_decl, *inner_exprs])] - case RunDecl(ruleset_ident, until, scheduler): - args: list[bindings._Expr] = [bindings.Var(span(), str(ruleset_ident))] - if scheduler: - name = "run-with" - scheduler_name = self._persistent_scheduler_name(scheduler) - for i in range(len(bound_schedulers) - 1, -1, -1): - if bound_schedulers[i].id == scheduler.id: - scheduler_name = self._local_scheduler_name(i) - break - args.insert(0, bindings.Var(span(), scheduler_name)) - else: - name = "run" - if until: - if len(until) > 1: - msg = "Can only have one until fact with custom scheduler" - raise ValueError(msg) - args.append(bindings.Var(span(), ":until")) - fact_egg = self.fact_to_egg(until[0]) - if isinstance(fact_egg, bindings.Eq): - msg = "Cannot use equality fact with custom scheduler" - raise ValueError(msg) - args.append(fact_egg.expr) - return [bindings.Call(span(), name, args)] - case SaturateDecl(inner): - return [ - bindings.Call(span(), "saturate", self._schedule_with_scheduler_to_egg(inner, bound_schedulers)) - ] - case RepeatDecl(inner, times): - return [ - bindings.Call( - span(), - "repeat", - [ - bindings.Lit(span(), bindings.Int(times)), - *self._schedule_with_scheduler_to_egg(inner, bound_schedulers), - ], - ) - ] - case SequenceDecl(schedules): - res = [] - for s in schedules: - res.extend(self._schedule_with_scheduler_to_egg(s, bound_schedulers)) - return res - case _: - assert_never(schedule) - - def _persistent_scheduler_to_egg(self, scheduler: BackOffDecl) -> bindings._Command: - return bindings.UserDefined( - span(), - "let-scheduler", - [ - bindings.Var(span(), self._persistent_scheduler_name(scheduler)), - self._back_off_scheduler_to_egg(scheduler), - ], + operation = ( + Oneof[Literal["query"], pb.Unit]("query", pb.Unit()) + if num_threads is _QUERY_RESOURCES + else Oneof[Literal["threads"], int]("threads", _thread_count(num_threads)) ) - - def ruleset_to_egg(self, ident: Ident) -> None: # noqa: C901 - """ - Registers a ruleset if it's not already registered. - """ - if ident.name == "" and ident not in self.__egg_decls__._rulesets: - self.rulesets.setdefault(ident, set()) - return - egg_name = str(ident) - if self.egglog_file_state is not None and ( - not _egg_name_is_source_safe_symbol(egg_name) or _egg_name_is_parser_literal(egg_name) - ): - msg = ( - f"Ruleset name {egg_name!r} cannot be used with save_egglog_string=True because " - "it does not serialize as one Egglog symbol" - ) - raise ValueError(msg) - match self.__egg_decls__._rulesets[ident]: - case RulesetDecl(rules): - if ident not in self.rulesets: - if str(ident): - self.run_program(bindings.AddRuleset(span(), str(ident))) - added_rules = self.rulesets[ident] = set() - else: - added_rules = self.rulesets[ident] - for rule in rules: - if rule in added_rules: - continue - commands = self.commands_to_egg(rule, ident) - self.run_program(*commands) - added_rules.add(rule) - case CombinedRulesetDecl(rulesets): - if ident in self.rulesets: - return - self.rulesets[ident] = set() - for ruleset in rulesets: - self.ruleset_to_egg(ruleset) - self.run_program(bindings.UnstableCombinedRuleset(span(), str(ident), list(map(str, rulesets)))) - - def commands_to_egg( # noqa: C901, PLR0912 - self, cmd: CommandDecl, ruleset: Ident - ) -> list[bindings._Command]: - match cmd: - case ActionCommandDecl(action): - return [ - bindings.ActionCommand(action_egg) - for action_egg in self.actions_to_egg(action, expr_to_let=True, standalone=True) - ] - case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions): - self.type_ref_to_egg(tp) - name = str(self.rule_name_counter) - self.rule_name_counter += 1 - rewrite = bindings.Rewrite( - span(), - self._expr_to_egg(lhs), - self._expr_to_egg(rhs), - [self.fact_to_egg(c, expr_to_let=False) for c in conditions], - name, - ) - egg_cmd: bindings._Command - if isinstance(cmd, RewriteDecl): - egg_cmd = bindings.RewriteCommand(str(ruleset), rewrite, cmd.subsume) - serialized_name = str(egg_cmd) - reported_name = serialized_name.replace('"', "'") - if ( - self.egglog_file_state is not None - and (not self.seminaive or self.egraph.no_decomp()) - and (not cmd.subsume or isinstance(lhs, CallDecl)) - ): - rule = self._rewrite_to_rule_decl(tp, lhs, rhs, conditions, reported_name, cmd.subsume) - commands = self.commands_to_egg(rule, ruleset) - self.rule_name_to_command_decl[reported_name] = cmd - return commands - self.rule_name_to_command_decl[name] = cmd - # Saving a transcript executes the serialized command, whose syntax does not - # preserve the internal rewrite name. The engine then reports the serialized - # rewrite itself as its name, so retain that alias for RunReport translation. - self.rule_name_to_command_decl[serialized_name] = cmd - # Backend report names render every Symbol with single quotes, while - # command serialization accepts symbols as double-quoted literals. - self.rule_name_to_command_decl[reported_name] = cmd - else: - egg_cmd = bindings.BiRewriteCommand(str(ruleset), rewrite) - serialized_name = str(egg_cmd) - reported_name = serialized_name.replace('"', "'") - if self.egglog_file_state is not None and (not self.seminaive or self.egraph.no_decomp()): - report_names = (f"{reported_name}=>", f"{reported_name}<=") - rules = ( - self._rewrite_to_rule_decl(tp, lhs, rhs, conditions, report_names[0], False), - self._rewrite_to_rule_decl(tp, rhs, lhs, conditions, report_names[1], False), - ) - commands = [command for rule in rules for command in self.commands_to_egg(rule, ruleset)] - for report_name in report_names: - self.rule_name_to_command_decl[report_name] = cmd - return commands - self.rule_name_to_command_decl[f"{name}=>"] = cmd - self.rule_name_to_command_decl[f"{name}<="] = cmd - for suffix in ("=>", "<="): - self.rule_name_to_command_decl[f"{serialized_name}{suffix}"] = cmd - self.rule_name_to_command_decl[f"{reported_name}{suffix}"] = cmd - return [egg_cmd] - case RuleDecl(head, body, name, eval_mode, no_decomp): - if not name: - name = str(self.rule_name_counter) - self.rule_name_counter += 1 - self.rule_name_to_command_decl[name] = cmd - eval_modes = { - "seminaive": bindings.Seminaive(), - "naive": bindings.Naive(), - "unsafe-seminaive": bindings.UnsafeSeminaive(), - } - if eval_mode not in eval_modes: - msg = ( - f"Unknown rule evaluation mode {eval_mode!r}; expected " - "'seminaive', 'naive', or 'unsafe-seminaive'" - ) - raise ValueError(msg) - binding_eval_mode = ( - bindings.Naive() - if not self.seminaive - else cast( - "bindings.Seminaive | bindings.Naive | bindings.UnsafeSeminaive", - eval_modes[eval_mode], - ) - ) - used_variable_names = _rule_variable_names(cmd) - return [ - bindings.RuleCommand( - bindings.Rule( - span(), - [ - action_egg - for action in head - for action_egg in self.actions_to_egg(action, used_variable_names=used_variable_names) - ], - [self.fact_to_egg(f, expr_to_let=False) for f in body], - name or "", - str(ruleset), - binding_eval_mode, - no_decomp or self.egraph.no_decomp(), - ) - ) - ] - case DefaultRewriteDecl(ref, expr, subsume): - sig = self.__egg_decls__.get_callable_decl(ref).signature - assert isinstance(sig, FunctionSignature) - # Replace args with rule_var_name mapping - arg_mapping = tuple( - TypedExprDecl(tp.to_just(), UnboundVarDecl(name, f"_{i}")) - for i, (name, tp) in enumerate(zip(sig.arg_names, sig.arg_types, strict=True)) - ) - rewrite_decl = RewriteDecl( - sig.semantic_return_type.to_just(), CallDecl(ref, arg_mapping), expr, (), subsume - ) - return self.commands_to_egg(rewrite_decl, ruleset) - case _: - assert_never(cmd) - - def _rewrite_to_rule_decl( - self, - tp: JustTypeRef, - lhs: ExprDecl, - rhs: ExprDecl, - conditions: tuple[FactDecl, ...], - name: str, - subsume: bool, - ) -> RuleDecl: - """Desugar a rewrite while retaining rule options that rewrite syntax cannot encode.""" - used_variable_names = _rule_variable_names(RuleDecl((), (EqDecl(tp, lhs, rhs), *conditions), None)) - fresh_name = self._allocate_synthetic_let_name().removeprefix("$") - while fresh_name in used_variable_names: - fresh_name = self._allocate_synthetic_let_name().removeprefix("$") - fresh = UnboundVarDecl(fresh_name, fresh_name) - head: list[ActionDecl] = [UnionDecl(tp, fresh, rhs)] - if subsume: - if not isinstance(lhs, CallDecl): - msg = "subsumed rewrite must have a function call on the lhs" - raise ValueError(msg) - head.append(ChangeDecl(tp, lhs, "subsume")) - return RuleDecl(tuple(head), (EqDecl(tp, fresh, lhs), *conditions), name) - - def actions_to_egg( # noqa: C901, PLR0911, PLR0912 - self, - action: ActionDecl, - expr_to_let: bool = False, - *, - standalone: bool = False, - used_variable_names: set[str] | None = None, - ) -> list[bindings._Action]: - match action: - case LetDecl(name, typed_expr): - normalized_name = _normalize_global_let_name(name) - if self.egglog_file_state is not None and not _egg_name_is_source_safe_symbol(normalized_name): - msg = ( - f"Let name {name!r} cannot be used with save_egglog_string=True because " - "it does not serialize as one Egglog symbol" - ) - raise ValueError(msg) - var_decl = LetRefDecl(name) - var_egg = self._expr_to_egg(var_decl) - self.expr_to_egg_cache[var_decl] = var_egg - return [ - bindings.Let( - span(), - var_egg.name, - self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let), - ) - ] - case SetDecl(tp, call, rhs): - self.type_ref_to_egg(tp) - egg_fn, typed_args = self.translate_call(call) - return [ - bindings.Set( - span(), - egg_fn, - [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], - self._expr_to_egg(rhs, expr_to_let=expr_to_let), - ) - ] - case ExprActionDecl(typed_expr): - if not isinstance(typed_expr.expr, CallDecl): - msg = "Top-level egglog expr commands must be calls" - raise ValueError(msg) # noqa: TRY004 - preserve the public validation error - callable_decl = self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) - if ( - self.egglog_file_state is not None - and callable_decl.egg_name is not None - and callable_decl.egg_name - in (_EGGLOG_RESERVED_COMMAND_OR_ACTION_HEADS if standalone else _EGGLOG_RESERVED_ACTION_HEADS) - ): - context = "top-level action" if standalone else "rule action" - msg = ( - f"Explicit Egglog callable name {callable_decl.egg_name!r} cannot be used as a {context} " - "with save_egglog_string=True because it is parsed as Egglog syntax" - ) - raise ValueError(msg) - egg_expr = self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let) - if isinstance(egg_expr, bindings.Var): - return [] - assert isinstance(egg_expr, bindings.Call) - return [bindings.Expr_(span(), egg_expr)] - case ChangeDecl(tp, call, change): - self.type_ref_to_egg(tp) - egg_fn, typed_args = self.translate_call(call) - egg_change: bindings._Change - match change: - case "delete": - egg_change = bindings.Delete() - case "subsume": - egg_change = bindings.Subsume() - case _: - assert_never(change) - return [ - bindings.Change( - span(), - egg_change, - egg_fn, - [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], - ) - ] - case UnionDecl(tp, lhs, rhs): - self.type_ref_to_egg(tp) - return [ - bindings.Union( - span(), - self._expr_to_egg(lhs, expr_to_let=expr_to_let), - self._expr_to_egg(rhs, expr_to_let=expr_to_let), - ) - ] - case PanicDecl(name): - return [bindings.Panic(span(), name)] - case SetCostDecl(tp, expr, cost): - self.type_ref_to_egg(tp) - egg_fn, typed_args = self.translate_call(expr) - cost_table = self.create_cost_table(expr.callable) - # Match egglog-experimental's set-cost action macro: bind each - # argument once, materialize the target call, then write its - # cost. Structural lowering here avoids evaluating a - # cost expression early through top-level factoring. - lowered: list[bindings._Action] = [] - args_egg: list[bindings._Expr] = [] - for typed_arg in typed_args: - name = self._allocate_synthetic_let_name() - if standalone: - # Mark implementation-only globals so freeze() omits them - # without making later expressions reuse their values. - var_egg = bindings.Var(span(), name) - self.expr_to_letref_cache[LetRefDecl(name)] = var_egg - self.expr_to_egg_cache[LetRefDecl(name)] = var_egg - else: - name = name.removeprefix("$") - while used_variable_names is not None and name in used_variable_names: - name = self._allocate_synthetic_let_name().removeprefix("$") - if used_variable_names is not None: - used_variable_names.add(name) - var_egg = bindings.Var(span(), name) - lowered.append(bindings.Let(span(), name, self.typed_expr_to_egg(typed_arg, False))) - args_egg.append(var_egg) - lowered.append(bindings.Expr_(span(), bindings.Call(span(), egg_fn, args_egg))) - cost_expr = self._expr_to_egg(cost, expr_to_let=False) - lowered.append(bindings.Set(span(), cost_table, args_egg, cost_expr)) - return lowered - case _: - assert_never(action) - - def create_cost_table(self, ref: CallableRef) -> str: - """ - Creates the egg cost table if needed and gets the name of the table. - """ - if ref in self.cost_table_names: - return self.cost_table_names[ref] - name = f"cost_table_{self.callable_ref_to_egg(ref)[0]}" - signature = self.__egg_decls__.get_callable_decl(ref).signature - assert isinstance(signature, FunctionSignature), "Can only add cost tables for functions" - target_schema = self._signature_to_egg_schema(signature) - schema = self._signature_to_egg_schema(replace(signature, return_type=TypeRefWithVars(Ident.builtin("i64")))) - - # egglog-experimental's DynamicCostModel probes this exact canonical - # name, so choosing a generated suffix would silently ignore costs. - # Aliases of the same backend callable share a cost table, but overloaded - # callables with incompatible schemas cannot: the backend protocol names - # cost tables only by the callable's backend symbol. A user-declared raw - # cost table can also be reused, but only when it is a bodyless function - # with the protocol's input sorts and i64 output. - if not self._has_compatible_cost_table_target(name, target_schema): - existing_refs = self.egg_fn_to_callable_refs.get(name, set()) - compatible_raw_table = bool(existing_refs) and all( - self._raw_cost_table_matches_schema(existing_ref, schema) for existing_ref in existing_refs - ) - if compatible_raw_table and not self._callable_is_table_backed(ref): - raise ValueError( - f"Canonical dynamic-cost table {name!r} for an eager or builtin primitive " - "cannot also be a user-declared function" - ) - if existing_refs and not compatible_raw_table: - msg = ( - f"Canonical dynamic-cost table {name!r} is already used by an incompatible callable; " - "it must be a bodyless function with the target's input sorts and i64 output" - ) - raise ValueError(msg) - if not compatible_raw_table: - if self._backend_symbol_is_occupied(name): - raise ValueError(f"Canonical dynamic-cost table name {name!r} is already in use") - self.run_program(bindings.FunctionCommand(span(), name, schema, None)) - self.cost_table_names[ref] = name - return name - - def _callable_is_table_backed(self, ref: CallableRef) -> bool: - """Return whether Egglog stores rows for this callable.""" - decl = self.__egg_decls__.get_callable_decl(ref) - match decl: - case RelationDecl() | ConstructorDecl() | ConstantDecl(body=None): - return True - case FunctionDecl(body=None, builtin=False): - return not isinstance(ref, UnnamedFunctionRef) - case ConstantDecl() | FunctionDecl(): - return False - case _: - assert_never(decl) - - def _raw_cost_table_matches_schema(self, ref: CallableRef, schema: bindings.Schema) -> bool: - """Check the complete public declaration contract for a raw dynamic-cost table.""" - decl = self.__egg_decls__.get_callable_decl(ref) - if not ( - isinstance(decl, FunctionDecl) - and not decl.builtin - and decl.body is None - and decl.merge is None - and isinstance(decl.signature, FunctionSignature) - ): - return False - existing_schema = self._signature_to_egg_schema(decl.signature) - return existing_schema.input == schema.input and existing_schema.output == schema.output - - def _has_compatible_cost_table_target(self, name: str, target_schema: bindings.Schema) -> bool: - """Validate every callable already sharing a canonical dynamic-cost table.""" - existing_cost_refs = [ref for ref, existing_name in self.cost_table_names.items() if existing_name == name] - for existing_ref in existing_cost_refs: - existing_signature = self.__egg_decls__.get_callable_decl(existing_ref).signature - assert isinstance(existing_signature, FunctionSignature) - existing_schema = self._signature_to_egg_schema(existing_signature) - if existing_schema.input != target_schema.input or existing_schema.output != target_schema.output: - raise ValueError( - f"Canonical dynamic-cost table {name!r} already serves a callable with an incompatible schema" - ) - return bool(existing_cost_refs) - - def fact_to_egg(self, fact: FactDecl, *, expr_to_let: bool = False) -> bindings._Fact: - match fact: - case EqDecl(tp, left, right): - self.type_ref_to_egg(tp) - return bindings.Eq( - span(), - self._expr_to_egg(left, expr_to_let=expr_to_let), - self._expr_to_egg(right, expr_to_let=expr_to_let), - ) - case ExprFactDecl(typed_expr): - if isinstance(typed_expr.expr, CallDecl) and typed_expr.expr.callable != FunctionRef( - Ident.builtin("!=") - ): - callable_decl = self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) - if ( - self.egglog_file_state is not None - and callable_decl.egg_name is not None - and callable_decl.egg_name == "=" - ): - msg = ( - "Explicit Egglog callable name '=' cannot be used as a fact with " - "save_egglog_string=True because it is parsed as equality syntax" - ) - raise ValueError(msg) - return bindings.Fact(self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let)) - case _: - assert_never(fact) - - def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C901, PLR0912 - """ - Returns the egg function name for a callable reference, registering it if it is not already registered. - - Also returns whether the args should be reversed - """ - if ref in self.callable_ref_to_egg_fn: - return self.callable_ref_to_egg_fn[ref] - decl = self.__egg_decls__.get_callable_decl(ref) - if ( - self.egglog_file_state is not None - and decl.egg_name - and (not _egg_name_is_source_safe_symbol(decl.egg_name) or _egg_name_is_parser_literal(decl.egg_name)) - ): - msg = ( - f"Explicit Egglog callable name {decl.egg_name!r} cannot be used with " - "save_egglog_string=True because it is parsed as Egglog syntax" - ) - raise ValueError(msg) - egg_name = decl.egg_name or self._allocate_name( - self._generate_callable_egg_name(ref), avoid_reserved_call_heads=True - ) - cost_table_targets = tuple( - target_ref for target_ref, cost_table_name in self.cost_table_names.items() if cost_table_name == egg_name + request = pb.ConfigureEGraphResourcesRequest(egraph_id=self.handle, operation=operation) + response = pb.ConfigureEGraphResourcesResponse.from_binary( + self.transport.configure_resources(request.to_binary()) ) - reuse_cost_table = bool(cost_table_targets) - existing_refs = self.egg_fn_to_callable_refs.get(egg_name, set()) - if existing_refs and not reuse_cost_table and not (isinstance(decl, FunctionDecl) and decl.builtin): - msg = f"Explicit Egglog callable name {egg_name!r} is already registered" + if response.threads < 1: + msg = "Resource response must contain a positive actual thread count" raise ValueError(msg) - for target_ref in cost_table_targets: - if not self._callable_is_table_backed(target_ref): - raise ValueError( - f"Canonical dynamic-cost table {egg_name!r} for an eager or builtin primitive " - "cannot also be a user-declared function" - ) - target_signature = self.__egg_decls__.get_callable_decl(target_ref).signature - assert isinstance(target_signature, FunctionSignature) - cost_schema = self._signature_to_egg_schema( - replace(target_signature, return_type=TypeRefWithVars(Ident.builtin("i64"))) - ) - if not self._raw_cost_table_matches_schema(ref, cost_schema): - msg = ( - f"Canonical dynamic-cost table {egg_name!r} is already used by an incompatible callable; " - "it must be a bodyless function with the target's input sorts and i64 output" - ) - raise ValueError(msg) - if reuse_cost_table and self.egg_fn_to_callable_refs.get(egg_name): - msg = f"Canonical dynamic-cost table {egg_name!r} already has a raw callable alias" - raise ValueError(msg) - callable_signature = decl.signature - reverse_args = callable_signature.reverse_args if isinstance(callable_signature, FunctionSignature) else False - match decl: - case RelationDecl(arg_types, _, _): - self.run_program(bindings.Relation(span(), egg_name, [self.type_ref_to_egg(a) for a in arg_types])) - case ConstantDecl(tp, _, body, merge): - if body is not None: - self.run_program(self._primitive_command_to_egg(egg_name, decl.signature, body)) - else: - # Egglog v3 has no constant command, so lower Python constants - # and class variables as zero-argument functions or constructors. - is_function = self.__egg_decls__._classes[tp.ident].builtin or merge is not None - schema = bindings.Schema([], self.type_ref_to_egg(tp)) - if is_function: - self.run_program( - bindings.FunctionCommand( - span(), - egg_name, - schema, - self._expr_to_egg(merge) if merge else None, - ) - ) - else: - self.run_program(bindings.Constructor(span(), egg_name, schema, None, False)) - case FunctionDecl(signature=signature, builtin=builtin, body=body, merge=merge): - if not builtin and not reuse_cost_table: - assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg" - if body is None and isinstance(ref, UnnamedFunctionRef): - body = ref.res - if body is not None: - self.run_program(self._primitive_command_to_egg(egg_name, signature, body)) - else: - # Compile functions that return unit to relations, because these show up in methods where you - # cant use the relation helper - schema = self._signature_to_egg_schema(signature) - if signature.return_type == TypeRefWithVars(Ident.builtin("Unit")): - if merge: - msg = "Cannot specify a merge function for a function that returns unit" - raise ValueError(msg) - self.run_program(bindings.Relation(span(), egg_name, schema.input)) - else: - self.run_program( - bindings.FunctionCommand( - span(), - egg_name, - schema, - self._expr_to_egg(merge) if merge else None, - ), - ) - case ConstructorDecl(signature, _, cost, unextractable): - self.run_program( - bindings.Constructor( - span(), - egg_name, - self._signature_to_egg_schema(signature), - cost, - unextractable, - ), - ) - case _: - assert_never(decl) - # Publish the reverse mapping only after any backend declaration has - # succeeded; otherwise a failed registration corrupts extraction and - # freeze by claiming an alias the backend never accepted. - self.egg_fn_to_callable_refs.setdefault(egg_name, set()).add(ref) - self.callable_ref_to_egg_fn[ref] = egg_name, reverse_args - return egg_name, reverse_args - - def _primitive_command_to_egg( - self, - egg_name: str, - signature: FunctionSignature, - body: TypedExprDecl, - ) -> bindings.UserDefined: - backend_arg_types = signature.arg_types[::-1] if signature.reverse_args else signature.arg_types - input_sort_expr = self._primitive_input_sorts_to_egg([ - self.type_ref_to_egg(arg_type.to_just()) for arg_type in backend_arg_types - ]) - output_sort_expr = bindings.Var(span(), self.type_ref_to_egg(signature.semantic_return_type.to_just())) - return bindings.UserDefined( - span(), - "primitive", - [ - bindings.Var(span(), egg_name), - input_sort_expr, - output_sort_expr, - self.typed_expr_to_egg(body, expr_to_let=False), - ], - ) - - def _primitive_input_sorts_to_egg(self, sort_names: list[str]) -> bindings._Expr: - if not sort_names: - return bindings.Lit(span(), bindings.Unit()) - if len(sort_names) == 1: - return bindings.Var(span(), sort_names[0]) - first, *rest = sort_names - return bindings.Call(span(), first, [bindings.Var(span(), sort_name) for sort_name in rest]) + return response.threads - def _signature_to_egg_schema(self, signature: FunctionSignature) -> bindings.Schema: - backend_arg_types = signature.arg_types[::-1] if signature.reverse_args else signature.arg_types - return bindings.Schema( - [self.type_ref_to_egg(a.to_just()) for a in backend_arg_types], - self.type_ref_to_egg(signature.semantic_return_type.to_just()), - ) - - def type_ref_to_egg(self, ref: JustTypeRef) -> str: - """ - Returns the egg sort name for a type reference, registering it not already registered, and also recursively - any type args are registered. - """ - try: - return self.type_ref_to_egg_sort[ref] - except KeyError: - pass - decl = self.__egg_decls__._classes[ref.ident] - if ( - self.egglog_file_state is not None - and decl.egg_name - and (not _egg_name_is_source_safe_symbol(decl.egg_name) or _egg_name_is_parser_literal(decl.egg_name)) - ): - msg = ( - f"Explicit Egglog sort name {decl.egg_name!r} cannot be used with " - "save_egglog_string=True because it is not parsed as an Egglog symbol" + def run_program(self, commands: Iterable[Ref]) -> tuple[pb.RunProgramResponse, Owner]: + """Pack canonical command closure once and adopt the returned records.""" + if self.handle is None: + msg = "EGraph handle has been destroyed" + raise RuntimeError(msg) + packed = pack(commands=commands) + request = pb.RunProgramRequest(egraph_id=self.handle, program=packed.program, profile=True) + response = pb.RunProgramResponse.from_binary(self.transport.run(request.to_binary())) + if response.error is not None: + raise bindings.EggSmolError(response.error.message) + for reference in packed.definitions: + key = _definition_key(reference.read()) + if key is not None and reference.role == "declarations": + self.definitions[key] = reference + builder = Builder.from_program( + pb.Program( + ir_version=1, + nodes=response.nodes, + sorts=response.sorts, + files=response.files, ) - raise ValueError(msg) - if not decl.builtin and not ref.args and decl.egg_name and self._backend_symbol_is_occupied(decl.egg_name): - msg = f"Explicit Egglog sort name {decl.egg_name!r} is already registered" - raise ValueError(msg) - arg_names = [self.type_ref_to_egg(arg) for arg in ref.args] - self.type_ref_to_egg_sort[ref] = egg_name = (not ref.args and decl.egg_name) or self._allocate_name( - self._generate_type_egg_name(ref, decl, arg_names) ) - self.egg_sort_to_type_ref[egg_name] = ref - - if decl.builtin: - # If this has args, create a new parameterized version of the builtin class - if ref.args: - if ref.ident == Ident.builtin("UnstableFn"): - type_args: list[bindings._Expr] = [ - bindings.Call( - span(), - self.type_ref_to_egg(ref.args[1]), - [bindings.Var(span(), self.type_ref_to_egg(a)) for a in ref.args[2:]], - ) - if len(ref.args) > 1 - else bindings.Lit(span(), bindings.Unit()), - bindings.Var(span(), self.type_ref_to_egg(ref.args[0])), - ] - else: - type_args = [bindings.Var(span(), self.type_ref_to_egg(a)) for a in ref.args] - assert decl.egg_name - self.run_program(bindings.Sort(span(), egg_name, (decl.egg_name, type_args))) - - # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods. - # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted - # even if you never use that function. - for method_name in decl.class_methods: - self.callable_ref_to_egg(ClassMethodRef(ref.ident, method_name)) - if decl.init: - self.callable_ref_to_egg(InitRef(ref.ident)) - else: - self.run_program(bindings.Sort(span(), egg_name, None)) - - return egg_name - - def op_mapping(self) -> dict[str, str]: - """ - Create a mapping of egglog function name to Python function name, for use in the serialized format - for better visualization. - - Includes cost tables - """ - return { - k: pretty_callable_ref(self.__egg_decls__, next(iter(v))) - for k, v in self.egg_fn_to_callable_refs.items() - if len(v) == 1 - } | { - name: f"cost({pretty_callable_ref(self.__egg_decls__, ref, include_all_args=True)})" - for ref, name in self.cost_table_names.items() - } - - def possible_egglog_functions(self, names: list[str]) -> Iterable[str]: - """ - Given a list of egglog functions, returns all the possible Python function strings - """ - for name in names: - for c in self.egg_fn_to_callable_refs.get(name, ()): - yield pretty_callable_ref(self.__egg_decls__, c) - - def typed_expr_to_egg( - self, - typed_expr_decl: TypedExprDecl, - expr_to_let: bool = True, - *, - lets_hoisted: bool = False, - ) -> bindings._Expr: - """ - Convert a typed expression to an egg expression. - - If `lets_hoisted` is set, an enclosing call has already hoisted this expression's shared - subexpressions into let bindings, so they are not searched for again. Searching once per - subterm would make converting a term quadratic in its depth. - """ - # transform all expressions with multiple parents into a let binding, so that less expressions - # are sent to egglog. Only for performance reasons. - if expr_to_let and not lets_hoisted: - have_multiple_parents = _exprs_multiple_parents(typed_expr_decl) - for expr in reversed(have_multiple_parents): - self._transform_let(expr) - - self.type_ref_to_egg(typed_expr_decl.tp) - return self._expr_to_egg(typed_expr_decl.expr, expr_to_let=expr_to_let, lets_hoisted=expr_to_let) - - def _transform_let(self, typed_expr: TypedExprDecl) -> TypedExprDecl | None: - """ - Rewrites this expression as a let binding if it's not already a let binding. - """ - if not isinstance(typed_expr.expr, CallDecl): - return typed_expr - if not is_callable_decl_constructor( - self.__egg_decls__, self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) - ): - return typed_expr - # A synthetic let is a top-level binding, so it cannot capture a rule - # variable (or an invalid top-level unbound variable). Leave those - # expressions inline for their enclosing command to validate instead. - if _contains_unbound_var(typed_expr): - return typed_expr - if typed_expr.expr in self.expr_to_letref_cache: - return None - var_decl = LetRefDecl(self._allocate_synthetic_let_name()) - var_egg = self._expr_to_egg(var_decl) - cmd = bindings.ActionCommand(bindings.Let(span(), var_egg.name, self.typed_expr_to_egg(typed_expr, True))) - self.run_program(cmd) - self.expr_to_letref_cache[typed_expr.expr] = var_egg - self.expr_to_egg_cache[var_decl] = var_egg - return None - - @overload - def _expr_to_egg( - self, expr_decl: CallDecl, *, expr_to_let: bool = ..., lets_hoisted: bool = ... - ) -> bindings.Call: ... - - @overload - def _expr_to_egg( - self, expr_decl: UnboundVarDecl | LetRefDecl, *, expr_to_let: bool = ..., lets_hoisted: bool = ... - ) -> bindings.Var: ... - - @overload - def _expr_to_egg( - self, expr_decl: ExprDecl, *, expr_to_let: bool = ..., lets_hoisted: bool = ... - ) -> bindings._Expr: ... - - def _expr_to_egg( # noqa: PLR0912,C901 - self, expr_decl: ExprDecl, *, expr_to_let: bool = False, lets_hoisted: bool = False - ) -> bindings._Expr: - """ - Convert an ExprDecl to an egg expression. - - `lets_hoisted` is passed on to the arguments, see `typed_expr_to_egg`. - """ - if expr_to_let: - try: - return self.expr_to_letref_cache[expr_decl] - except KeyError: - pass - else: - try: - return self.expr_to_egg_cache[expr_decl] - except KeyError: - pass - res: bindings._Expr - match expr_decl: - case LetRefDecl(name): - res = bindings.Var(span(), _normalize_global_let_name(name)) - case UnboundVarDecl(name, egg_name): - emitted_name = egg_name or f"_{name}" - if self.egglog_file_state is not None and ( - emitted_name == "_" - or emitted_name.startswith("@") - or not _egg_name_is_source_safe_symbol(emitted_name) - or _egg_name_is_parser_literal(emitted_name) - ): - msg = ( - f"Egglog variable name {emitted_name!r} cannot be used with " - "save_egglog_string=True because it is not parsed as an Egglog symbol" - ) - raise ValueError(msg) - res = bindings.Var(span(), emitted_name) - case LitDecl(value): - l: bindings._Literal - match value: - case None: - l = bindings.Unit() - case bool(i): - l = bindings.Bool(i) - case int(i): - l = bindings.Int(i) - case float(f): - l = bindings.Float(f) - case str(s): - l = bindings.String(s) - case _: - assert_never(value) - res = bindings.Lit(span(), l) - case CallDecl() | GetCostDecl(): - egg_fn, typed_args = self.translate_call(expr_decl) - egg_args = [self.typed_expr_to_egg(a, expr_to_let, lets_hoisted=lets_hoisted) for a in typed_args] - res = bindings.Call(span(), egg_fn, egg_args) - case PyObjectDecl(value): - res = bindings.Call( - span(), - "py-object", - [bindings.Lit(span(), bindings.String(standard_b64encode(value).decode("utf-8")))], - ) - case PartialCallDecl(call_decl): - egg_fn, typed_args = self.translate_call(call_decl) - res = bindings.Call( - span(), - "unstable-fn", - [ - bindings.Lit(span(), bindings.String(egg_fn)), - *[self.typed_expr_to_egg(arg, expr_to_let, lets_hoisted=lets_hoisted) for arg in typed_args], - ], - ) - case ValueDecl(): - msg = "Cannot turn a Value into an expression" - raise ValueError(msg) - case DummyDecl(): - msg = "Cannot turn a DummyDecl into an expression" - raise ValueError(msg) - case _: - assert_never(expr_decl.expr) - if not expr_to_let: - self.expr_to_egg_cache[expr_decl] = res - return res - - def translate_call(self, expr: CallDecl | GetCostDecl) -> tuple[str, list[TypedExprDecl]]: - """ - Handle get cost and call decl, turn into egg table name and typed expr decls. - """ - match expr: - case CallDecl(ref, args, _): - egg_fn, reverse_args = self.callable_ref_to_egg(ref) - case GetCostDecl(ref, args): - egg_fn = self.create_cost_table(ref) - _, reverse_args = self.callable_ref_to_egg(ref) - case _: - assert_never(expr) - args_list = list(args) - if reverse_args: - args_list.reverse() - return egg_fn, args_list - - def exprs_from_egg(self, termdag: bindings.TermDag, terms: list[int], tp: JustTypeRef) -> Iterable[TypedExprDecl]: - """ - Create a function that can convert from an egg term to a typed expr. - """ - state = FromEggState(self, termdag) - return [state.resolve_term(term_id, tp) for term_id in terms] + for reference in self.definitions.values(): + builder.import_ref(reference) + return response, builder.publish() - def _get_possible_types(self, cls_ident: Ident) -> frozenset[JustTypeRef]: - """ - Given a class name, returns all possible registered types that it can be. - """ - return frozenset(tp for tp in self.type_ref_to_egg_sort if tp.ident == cls_ident) - - def _generate_callable_egg_name(self, ref: CallableRef) -> str: - """ - Generate a fully-qualified Egglog name for a callable reference. - - Qualification separates same-named Python declarations; explicit names - in the same registration batch are reserved by `_allocate_name`. - """ - match ref: - case FunctionRef(ident): - return _sanitize_egg_ident(str(ident)) - case ConstantRef(ident): - # Prefix to avoid name collisions with local vars - return _sanitize_egg_ident(f"%{ident}") - case ( - MethodRef(cls_ident, name) - | ClassMethodRef(cls_ident, name) - | ClassVariableRef(cls_ident, name) - | PropertyRef(cls_ident, name) - ): - return _sanitize_egg_ident(f"{cls_ident}.{name}") - case InitRef(cls_ident): - return _sanitize_egg_ident(f"{cls_ident}.__init__") - case UnnamedFunctionRef(): - name = f"_lambda_{self.unnamed_function_counter}" - self.unnamed_function_counter += 1 - return name - case _: - assert_never(ref) - - def _backend_symbol_is_occupied(self, name: str) -> bool: - """Check Egglog's shared namespace and replay-sensitive syntax names.""" - return ( - not _egg_name_is_source_safe_symbol(name) - or _egg_name_is_parser_literal(name) - or bool(self.egg_fn_to_callable_refs.get(name)) - or name in self.egg_sort_to_type_ref - or name in self.cost_table_names.values() - or name in _BUILTIN_EGG_FN_NAMES - or name in _BUILTIN_EGG_SORT_NAMES + def copy(self) -> EGraphState: + """Clone the native handle while sharing immutable authored definitions.""" + if self.handle is None: + msg = "EGraph handle has been destroyed" + raise RuntimeError(msg) + response = pb.CloneEGraphResponse.from_binary( + self.transport.clone_graph(pb.CloneEGraphRequest(egraph_id=self.handle).to_binary()) ) - - def _generate_type_egg_name(self, ref: JustTypeRef, decl: ClassDecl, arg_names: list[str]) -> str: - if decl.egg_name: - base = decl.egg_name - else: - # Preserve the readable dotted qualification while sanitizing each - # Python identifier component into a source-safe Egglog symbol. - parts = (*ref.ident.module.split("."), ref.ident.name) if ref.ident.module else (ref.ident.name,) - base = ".".join(_sanitize_egg_ident(part) or "_" for part in parts) - if not ref.args: - return base - args = ",".join(arg_names) - return f"{base}[{args}]" - - def _allocate_synthetic_let_name(self) -> str: - existing_let_names = self.pending_let_names | { - egg_expr.name - for decl, egg_expr in self.expr_to_egg_cache.items() - if isinstance(decl, LetRefDecl) and isinstance(egg_expr, bindings.Var) - } - while True: - candidate = f"$__expr_{self.expr_to_let_counter}" - self.expr_to_let_counter += 1 - name = self._allocate_name(candidate) - if name not in existing_let_names: - return name - - def _allocate_name(self, candidate: str, *, avoid_reserved_call_heads: bool = False) -> str: - # All declarations for a register(...) batch are merged before any - # command is lowered. Reserve their explicit backend names up front so - # generated names do not depend on action order within that batch. - if ( - candidate not in self._explicit_backend_names - and not self._backend_symbol_is_occupied(candidate) - and (not avoid_reserved_call_heads or candidate not in _EGGLOG_RESERVED_CALL_HEADS) - ): - return candidate - - index = 1 - while ( - f"{candidate}_{index}" in self._explicit_backend_names - or self._backend_symbol_is_occupied(f"{candidate}_{index}") - or (avoid_reserved_call_heads and f"{candidate}_{index}" in _EGGLOG_RESERVED_CALL_HEADS) - ): - index += 1 - return f"{candidate}_{index}" - - def typed_expr_to_value(self, typed_expr: TypedExprDecl) -> bindings.Value: - if isinstance(typed_expr.expr, ValueDecl): - if typed_expr.expr.owner not in self.valid_value_owners: - msg = "Cannot use a value that belongs to a different EGraph or inactive push scope" - raise ValueError(msg) - return typed_expr.expr.value - egg_expr = self.typed_expr_to_egg(typed_expr, False) - return call_with_current_trace(self.egraph.eval_expr, egg_expr)[1] - - def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl: # noqa: C901, PLR0911, PLR0912 - if tp.ident.module != Ident.builtin("").module: - return ValueDecl(value, self.value_owner) - - match tp.ident.name: - # Should match list in egraph bindings - case "i64": - return LitDecl(self.egraph.value_to_i64(value)) - case "f64": - return LitDecl(self.egraph.value_to_f64(value)) - case "Bool": - return LitDecl(self.egraph.value_to_bool(value)) - case "String": - return LitDecl(self.egraph.value_to_string(value)) - case "Unit": - return LitDecl(None) - case "PyObject": - val = self.egraph.value_to_pyobject(value) - return PyObjectDecl(cloudpickle.dumps(val)) - case "Rational": - fraction = self.egraph.value_to_rational(value) - return CallDecl( - InitRef(Ident.builtin("Rational")), - ( - TypedExprDecl(JustTypeRef(Ident.builtin("i64")), LitDecl(fraction.numerator)), - TypedExprDecl(JustTypeRef(Ident.builtin("i64")), LitDecl(fraction.denominator)), - ), - ) - case "BigInt": - i = self.egraph.value_to_bigint(value) - return CallDecl( - ClassMethodRef(Ident.builtin("BigInt"), "from_string"), - (TypedExprDecl(JustTypeRef(Ident.builtin("String")), LitDecl(str(i))),), - ) - case "BigRat": - fraction = self.egraph.value_to_bigrat(value) - return CallDecl( - InitRef(Ident.builtin("BigRat")), - ( - TypedExprDecl( - JustTypeRef(Ident.builtin("BigInt")), - CallDecl( - ClassMethodRef(Ident.builtin("BigInt"), "from_string"), - ( - TypedExprDecl( - JustTypeRef(Ident.builtin("String")), LitDecl(str(fraction.numerator)) - ), - ), - ), - ), - TypedExprDecl( - JustTypeRef(Ident.builtin("BigInt")), - CallDecl( - ClassMethodRef(Ident.builtin("BigInt"), "from_string"), - ( - TypedExprDecl( - JustTypeRef(Ident.builtin("String")), LitDecl(str(fraction.denominator)) - ), - ), - ), - ), - ), - ) - case "Map": - k_tp, v_tp = tp.args - expr = CallDecl(ClassMethodRef(Ident.builtin("Map"), "empty"), (), (k_tp, v_tp)) - for k, v in self.egraph.value_to_map(value).items(): - expr = CallDecl( - MethodRef(Ident.builtin("Map"), "insert"), - ( - TypedExprDecl(tp, expr), - TypedExprDecl(k_tp, self.value_to_expr(k_tp, k)), - TypedExprDecl(v_tp, self.value_to_expr(v_tp, v)), - ), - ) - return expr - case "Set": - xs_ = self.egraph.value_to_set(value) - (v_tp,) = tp.args - return CallDecl( - InitRef(Ident.builtin("Set")), - tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs_), - (v_tp,) if not xs_ else (), - ) - case "Vec": - xs = self.egraph.value_to_vec(value) - (v_tp,) = tp.args - return CallDecl( - InitRef(Ident.builtin("Vec")), - tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs), - (v_tp,) if not xs else (), - ) - case "MultiSet": - xs = self.egraph.value_to_multiset(value) - (v_tp,) = tp.args - return CallDecl( - InitRef(Ident.builtin("MultiSet")), - tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs), - (v_tp,) if not xs else (), - ) - case "UnstableFn": - _names, _args = self.egraph.value_to_function(value) - return_tp, *arg_types = tp.args - return self._unstable_fn_value_to_expr(_names, _args, return_tp, arg_types) - case "Pair" | "Maybe": - termdag, term, _cost = call_with_current_trace( - self.egraph.extract_value, value, self.type_ref_to_egg(tp) - ) - return FromEggState(self, termdag).resolve_term(term, tp).expr - case _: - # If this is not a builtin type, or we don't know how to convert it, just return as value - return ValueDecl(value, self.value_owner) - - def _unstable_fn_value_to_expr( - self, name: str, partial_args: list[bindings.Value], return_tp: JustTypeRef, _arg_types: list[JustTypeRef] - ) -> PartialCallDecl: - # Similar to FromEggState::from_call but reconstructs a partial application from serialized values. - # Find first callable ref whose return type matches and fill in arg types. - for callable_ref in self.egg_fn_to_callable_refs.get(name, ()): - signature = self.__egg_decls__.get_callable_decl(callable_ref).signature - if not isinstance(signature, FunctionSignature): - continue - if signature.semantic_return_type.ident != return_tp.ident: - continue - arg_types = TypeConstraintSolver().infer_arg_types( - signature.arg_types, signature.semantic_return_type, signature.var_arg_type, return_tp - ) - args = tuple( - TypedExprDecl(tp, self.value_to_expr(tp, v)) for tp, v in zip(arg_types, partial_args, strict=False) + result = object.__new__(EGraphState) + result.transport = self.transport + result.handle = response.egraph_id + result.definitions = self.definitions.copy() + return result + + def destroy(self) -> None: + """Release this handle once without affecting cloned graph handles.""" + if self.handle is not None: + pb.DestroyEGraphResponse.from_binary( + self.transport.destroy(pb.DestroyEGraphRequest(egraph_id=self.handle).to_binary()) ) - call_decl = CallDecl(callable_ref, args) - return PartialCallDecl(call_decl) - raise ValueError(f"Function '{name}' not found") - - -# https://chatgpt.com/share/9ab899b4-4e17-4426-a3f2-79d67a5ec456 -_EGGLOG_INVALID_IDENT = re.compile(r"[^\w\-+*/?!=<>&|^/%]") - - -def _egg_name_is_parser_literal(name: str) -> bool: - if not name or name in _EGGLOG_LITERAL_NAMES: - return True - # Egglog parses every finite Rust f64 spelling as a literal rather than an - # atom. Callable names are sanitized first, but sort names can retain dots. - return bool(_EGGLOG_NUMBER.fullmatch(name)) and math.isfinite(float(name)) - - -def _sanitize_egg_ident(input_string: str) -> str: - """ - Replaces all invalid characters in an egg identifier with an underscore. - """ - return _EGGLOG_INVALID_IDENT.sub("_", input_string) - - -def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]: - """ - Return multiply-parented expressions in deterministic preorder for stable synthetic let names. - - Visits every child that `_expr_to_egg` converts with `lets_hoisted`, so the arguments of calls, - cost lookups and partial calls are not scanned again when they are converted. - """ - parent_counts: dict[TypedExprDecl, int] = {} - traversal_order: list[TypedExprDecl] = [] - traversed: set[TypedExprDecl] = set() - stack = [typed_expr] - while stack: - node = stack.pop() - if node in traversed: - continue - traversed.add(node) - if node is not typed_expr: - traversal_order.append(node) - match node.expr: - case CallDecl(args=args) | GetCostDecl(args=args) | PartialCallDecl(CallDecl(args=args)): - for child in args: - parent_counts[child] = parent_counts.get(child, 0) + 1 - stack.extend(reversed(args)) - case _: - pass - return [node for node in traversal_order if parent_counts[node] > 1] - - -def _contains_unbound_var(typed_expr: TypedExprDecl) -> bool: - """Check for an unbound variable without recursively hashing a deep expression DAG.""" - seen: set[int] = set() - stack = [typed_expr] - while stack: - node = stack.pop() - node_id = id(node) - if node_id in seen: - continue - seen.add(node_id) - match node.expr: - case UnboundVarDecl(): - return True - case CallDecl(args=args) | GetCostDecl(args=args) | PartialCallDecl(CallDecl(args=args)): - stack.extend(args) - case _: - pass - return False - - -@dataclass -class FromEggState: - """ - Dataclass containing state used when converting from an egg term to a typed expr. - """ - - state: EGraphState - termdag: bindings.TermDag - # Cache of termdag ID and expected type to TypedExprDecl. Polymorphic - # zero-argument terms like map-empty can appear once in a termdag but be - # decoded at multiple concrete types. - cache: dict[tuple[int, JustTypeRef], TypedExprDecl] = field(default_factory=dict) - - @property - def decls(self) -> Declarations: - return self.state.__egg_decls__ - - def from_expr(self, tp: JustTypeRef, term: bindings._Term) -> TypedExprDecl: - """ - Convert an egg term to a typed expr. - """ - # Extracted builtin values can use canonical constructors that were not - # present in the original Python expression, such as BigInt.from_string - # inside an extracted BigRat. Seed the expected type's callable mapping - # before resolving the term. - self.state.type_ref_to_egg(tp) - expr_decl: ExprDecl - if isinstance(term, bindings.TermVar): - expr_decl = LetRefDecl(term.name) - elif isinstance(term, bindings.TermLit): - value = term.value - expr_decl = LitDecl(None if isinstance(value, bindings.Unit) else value.value) - elif isinstance(term, bindings.TermApp): - if term.name == "map-of" and tp.ident == Ident.builtin("Map"): - if len(term.args) % 2: - raise ValueError(f"Expected alternating key/value terms in map-of, got {len(term.args)} terms") - key_tp, value_tp = tp.args - expr_decl = CallDecl(ClassMethodRef(Ident.builtin("Map"), "empty"), (), (key_tp, value_tp)) - for index in range(0, len(term.args), 2): - expr_decl = CallDecl( - MethodRef(Ident.builtin("Map"), "insert"), - ( - TypedExprDecl(tp, expr_decl), - self.resolve_term(term.args[index], key_tp), - self.resolve_term(term.args[index + 1], value_tp), - ), - ) - elif term.name == "py-object": - (str_term,) = term.args - call = self.termdag.get(str_term) - assert isinstance(call, bindings.TermLit) - assert isinstance(call.value, bindings.String) - expr_decl = PyObjectDecl(standard_b64decode(call.value.value)) - elif term.name == "unstable-fn": - # Get function name - fn_term, *arg_terms = term.args - fn_value = self.resolve_term(fn_term, JustTypeRef(Ident.builtin("String"))) - assert isinstance(fn_value.expr, LitDecl) - fn_name = fn_value.expr.value - assert isinstance(fn_name, str) - - # Resolve what types the partially applied args are - assert tp.ident == Ident.builtin("UnstableFn") - call_decl = self.from_call(tp.args[0], bindings.TermApp(fn_name, arg_terms)) - expr_decl = PartialCallDecl(call_decl) - else: - expr_decl = self.from_call(tp, term) - else: - assert_never(term) - return TypedExprDecl(tp, expr_decl) - - def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl: - """ - Convert a call to a CallDecl. - - There could be Python call refs which match the call, so we need to find the correct one. - - The additional_arg_tps are known types for arguments that come after the term args, used to infer types - for partially applied functions, where we know the types of the later args, but not of the earlier ones where - we have values for. - """ - # Find the first callable ref that matches the call - possible_callable_refs = self.state.egg_fn_to_callable_refs.get(term.name, ()) - for callable_ref in possible_callable_refs: - # If this is a classmethod, we might need the type params that were bound for this type - # This could be multiple types if the classmethod is ambiguous, like map create. - possible_types: Iterable[JustTypeRef | None] - signature = self.decls.get_callable_decl(callable_ref).signature - assert isinstance(signature, FunctionSignature) - term_args = term.args[::-1] if signature.reverse_args else term.args - if isinstance(callable_ref, ClassMethodRef | InitRef | MethodRef): - # Need OR in case we have class method whose class was never added as a sort, which would happen - # if the class method didn't return that type and no other function did. In this case, we don't need - # to care about the type vars and we don't need to bind any possible type. - possible_types = self.state._get_possible_types(callable_ref.ident) or [None] - else: - possible_types = [None] - for possible_type in possible_types: - tcs = TypeConstraintSolver() - if possible_type and possible_type.args: - tcs.bind_class(possible_type, self.decls) - bound_args = possible_type.args - else: - bound_args = () - try: - arg_types = tcs.infer_arg_types( - signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp - ) - # Include this in try because of iterable - a_tp = list(zip(term_args, arg_types, strict=False)) - except TypeConstraintError: - continue - args = tuple(self.resolve_term(a, tp) for a, tp in a_tp) - # Only save bound tp params if needed for inferring return type - # this is true if the set of set of type vars in the return are not a subset of those in the args - bound_tp_params = () if signature.semantic_return_type.vars.issubset(signature.arg_vars) else bound_args - return CallDecl(callable_ref, args, bound_tp_params) - raise ValueError( - f"Could not find callable ref for call {term}. None of these refs matched the types: {possible_callable_refs}" - ) - - def resolve_term(self, term_id: int, tp: JustTypeRef) -> TypedExprDecl: - key = (term_id, tp) - try: - return self.cache[key] - except KeyError: - res = self.cache[key] = self.from_expr(tp, self.termdag.get(term_id)) - return res + self.handle = None + + def __del__(self) -> None: + # A failed constructor may have no handle. Never let native cleanup + # replace an exception already propagating during Python disposal. + if getattr(self, "handle", None) is not None: + with suppress(Exception): + self.destroy() diff --git a/python/egglog/pretty.py b/python/egglog/pretty.py index 1985df9b..29442915 100644 --- a/python/egglog/pretty.py +++ b/python/egglog/pretty.py @@ -5,13 +5,16 @@ from __future__ import annotations import ast +import struct from collections import Counter, defaultdict from dataclasses import dataclass, field from typing import TYPE_CHECKING, TypeAlias, assert_never import black import cloudpickle +from egglog_proto.egglog.v1 import egglog_pb as pb +from ._program import Ref, StructuralView from .declarations import * if TYPE_CHECKING: @@ -23,6 +26,7 @@ "UNARY_METHODS", "pretty_callable_ref", "pretty_decl", + "pretty_ref", ] MAX_LINE_LENGTH = 110 LINE_DIFFERENCE = 10 @@ -76,6 +80,173 @@ "__ceil__": "ceil", } + +def pretty_ref(root: Ref) -> str: # noqa: C901, PLR0912 + """Print the canonical records directly, using supplied Python bindings.""" + from .runtime import _definition_at # noqa: PLC0415 + + # Frames stream output rather than recursively copying each child's string. + # The Boolean asks scalar arguments to omit their inferred wrapper. + frames: list[str | Ref | tuple[Ref, bool, bool]] = [(root, False, False)] + output: list[str] = [] + active: set[Ref] = set() + while frames: + frame = frames.pop() + if isinstance(frame, str): + output.append(frame) + continue + if isinstance(frame, Ref): + active.remove(frame) + continue + reference, bare, parentheses = frame + if reference in active: + msg = "Cyclic expression presentation is not migrated" + raise ValueError(msg) + active.add(reference) + frames.append(reference) + record = reference.read() + owner = reference.owner + pieces: list[str | tuple[Ref, bool, bool]] + if isinstance(record, pb.Sort): + if record.kind is None: + msg = "Cannot print a sort without a kind" + raise ValueError(msg) + if record.kind.field not in {"eq", "family"}: + raise NotImplementedError(f"Sort presentation is not migrated: {record.kind.field}") + name = record.kind.value if record.kind.field == "eq" else record.kind.value.name + sort_declaration: pb.Declaration = _definition_at(reference, "sort", name).read() + assert sort_declaration.kind is not None + assert isinstance(sort_declaration.kind.value, pb.EqSort | pb.HostSortFamily) + bindings = sort_declaration.kind.value.bindings + path = bindings.python.path if bindings is not None and bindings.python is not None else [name] + pieces = [path[-1]] + if record.kind.field == "family" and record.kind.value.args: + pieces.append("[") + for index, argument in enumerate(record.kind.value.args): + if index: + pieces.append(", ") + pieces.append((owner.ref("sorts", argument), False, False)) + pieces.append("]") + elif isinstance(record, pb.Node): + if record.kind is None: + msg = "Cannot print a node without a kind" + raise ValueError(msg) + sort = owner.ref("sorts", record.sort_id) + match record.kind.field: + case "var": + pieces = [record.kind.value] + case "primitive_value": + primitive = record.kind.value.value + if primitive is None: + msg = "Cannot print a primitive without a value" + raise ValueError(msg) + if primitive.field == "f64_bits": + value = repr(struct.unpack("!d", struct.pack("!Q", primitive.value))[0]) + elif primitive.field in {"i64", "string", "bool"}: + value = repr(primitive.value) + elif primitive.field == "unit": + value = "" + else: + raise NotImplementedError(f"Value presentation is not migrated: {primitive.field}") + pieces = [value] if bare else [(sort, False, False), "(", value, ")"] + case "call": + call = record.kind.value + declaration = _definition_at(reference, "callable", call.func) + declared: pb.Declaration = declaration.read() + view = ( + declared.bindings.python.views[0] + if declared.bindings is not None + and declared.bindings.python is not None + and declared.bindings.python.views + else None + ) + args = [owner.ref("nodes", index) for index in call.args] + if view is None: + pieces = [call.func, "("] + selected = [(argument, False, False) for argument in args] + else: + assert declared.kind is not None + kind = declared.kind.value + fixed = ( + len(kind.typing.value.inputs) + if isinstance(kind, pb.HostPrimitive) + and kind.typing is not None + and kind.typing.field == "signature" + else len(args) + ) + selected = [] + defaults: list[Ref | None] = [] + for parameter in view.params: + for argument in ( + args[parameter.core_input :] + if parameter.core_input == fixed + else [args[parameter.core_input]] + ): + selected.append((argument, True, False)) + defaults.append( + declaration.owner.ref("nodes", parameter.default_expr) + if parameter.has_field("default_expr") + else None + ) + while ( + selected + and defaults[-1] is not None + and StructuralView(selected[-1][0]) == StructuralView(defaults[-1]) + ): + selected.pop() + defaults.pop() + name = view.path[-1] if view.path else "__init__" + if view.kind == pb.PythonCallKind.METHOD and name in BINARY_METHODS and len(selected) == 1: + pieces = [ + (args[view.receiver], False, True), + f" {BINARY_METHODS[name]} ", + (selected[0][0], True, True), + ] + if parentheses: + pieces = ["(", *pieces, ")"] + frames.extend(reversed(pieces)) + continue + if view.kind == pb.PythonCallKind.INITIALIZER: + pieces = [(sort, False, False), "("] + elif view.kind in {pb.PythonCallKind.METHOD, pb.PythonCallKind.PROPERTY}: + pieces = [(args[view.receiver], False, True), ".", name] + if view.kind == pb.PythonCallKind.PROPERTY: + frames.extend(reversed(pieces)) + continue + pieces.append("(") + elif view.kind in {pb.PythonCallKind.CLASS_METHOD, pb.PythonCallKind.CLASS_VARIABLE}: + assert view.owner is not None + assert view.owner.kind is not None + if view.owner.kind.field != "sort": + msg = "Function-type class member presentation is not migrated" + raise NotImplementedError(msg) + pieces = [(declaration.owner.ref("sorts", view.owner.kind.value), False, False), ".", name] + if view.kind == pb.PythonCallKind.CLASS_VARIABLE: + frames.extend(reversed(pieces)) + continue + pieces.append("(") + else: + pieces = [name, "("] + for index, argument_frame in enumerate(selected): + if index: + pieces.append(", ") + pieces.append(argument_frame) + pieces.append(")") + case "union": + pieces = ["eq("] + for index, argument in enumerate(record.kind.value.members): + if index: + pieces.append(").to(" if index == 1 else ", ") + pieces.append((owner.ref("nodes", argument), False, False)) + pieces.append(")") + case _: + raise NotImplementedError(f"Node presentation is not migrated: {record.kind.field}") + else: + raise NotImplementedError(f"Record presentation is not migrated: {reference.role}") + frames.extend(reversed(pieces)) + return "".join(output) + + AllDecls: TypeAlias = ( RulesetDecl | CombinedRulesetDecl diff --git a/python/egglog/runtime.py b/python/egglog/runtime.py index 70c0e02b..45ad0fb2 100644 --- a/python/egglog/runtime.py +++ b/python/egglog/runtime.py @@ -1,34 +1,22 @@ -""" -Holds a number of types which are only used at runtime to emulate Python objects. - -Users will not import anything from this module, and statically they won't know these are the types they are using. - -But at runtime they will be exposed. - -Note that all their internal fields are prefixed with __egg_ to avoid name collisions with user code, but will end in __ -so they are not mangled by Python and can be accessed by the user. -""" +"""Python runtime wrappers; semantic state is owned generated protobuf records.""" from __future__ import annotations import itertools -import operator -import types +import struct from collections.abc import Callable -from dataclasses import InitVar, dataclass, field, replace +from contextlib import suppress from inspect import Parameter, Signature -from typing import TYPE_CHECKING, Any, TypeVar, Union, assert_never, cast, get_args, get_origin +from threading import RLock +from typing import Any, Literal, TypeVar, Union, cast, get_args, get_origin +from weakref import WeakKeyDictionary, WeakValueDictionary -import cloudpickle - -from .declarations import * -from .pretty import * -from .thunk import Thunk -from .type_constraint_solver import * - -if TYPE_CHECKING: - from collections.abc import Iterable +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import Oneof +from ._catalog import Catalog, builtin_catalog +from ._program import Builder, Owner, Ref, StructuralView, _definition_key, clone_nodes +from .type_constraint_solver import TypeConstraintError, infer_sort, substitute_sort __all__ = [ "ALWAYS_MUTATES_SELF", @@ -43,17 +31,8 @@ "define_expr_method", "resolve_callable", "resolve_type_annotation", - "resolve_type_annotation_mutate", ] - -UNIT_IDENT = Ident.builtin("Unit") -UNARY_LIT_IDENTS = {Ident.builtin("i64"), Ident.builtin("f64"), Ident.builtin("Bool"), Ident.builtin("String")} -LIT_IDENTS = UNARY_LIT_IDENTS | {UNIT_IDENT, Ident.builtin("PyObject")} - -# All methods which should return NotImplemented if they fail to resolve and are reflected as well -# From https://docs.python.org/3/reference/datamodel.html#emulating-numeric-types - NUMERIC_BINARY_METHODS = { "__add__", "__sub__", @@ -131,787 +110,682 @@ *TYPE_DEFINED_METHODS, } -# Set this globally so we can get access to PyObject when we have a type annotation of just object. -# This is the only time a type annotation doesn't need to include the egglog type b/c object is top so that would be redundant statically. -_PY_OBJECT_CLASS: RuntimeClass | None = None -# Same for functions -_UNSTABLE_FN_CLASS: RuntimeClass | None = None - -T = TypeVar("T") - - -def resolve_type_annotation_mutate(decls: Declarations, tp: object) -> TypeOrVarRef: - """ - Wrap resolve_type_annotation to mutate decls, as a helper for internal use in situations where that is more ergonomic. - """ - new_decls, tp = resolve_type_annotation(tp) - decls |= new_decls - return tp - - -def resolve_type_annotation(tp: object) -> tuple[DeclarationsLike, TypeOrVarRef]: - """ - Resolves a type object into a type reference. - - Any runtime type object decls will be returned as well. We do this so we can use this without having to - resolve the decls if need be. - """ - if isinstance(tp, TypeVar): - return None, TypeVarRef.from_type_var(tp) - # If there is a union, then we assume the first item is the type we want, and the others are types that can be converted to that type. - if get_origin(tp) == Union: - first, *_rest = get_args(tp) - return resolve_type_annotation(first) - - # If the type is `object` then this is assumed to be a PyObjectLike, i.e. converted into a PyObject - if tp is object: - assert _PY_OBJECT_CLASS - return resolve_type_annotation(_PY_OBJECT_CLASS) - # If the type is a `Callable` then convert it into a UnstableFn - if get_origin(tp) == Callable: - assert _UNSTABLE_FN_CLASS - args, ret = get_args(tp) - return resolve_type_annotation(_UNSTABLE_FN_CLASS[(ret, *args)]) - if isinstance(tp, RuntimeClass): - return tp, tp.__egg_tp__ - raise TypeError(f"Unexpected type annotation {tp}") +_OMITTED = object() -def inverse_resolve_type_annotation(decls_thunk: Callable[[], Declarations], tp: TypeOrVarRef) -> object: - """ - Inverse of resolve_type_annotation +def _author_call( # noqa: C901, PLR0912 + declaration: Ref, + view_index: int | None, + args: tuple[object, ...], + kwargs: dict[str, object], + *, + node_of: Callable[[object], Ref | None], + convert: Callable[[Ref, object], Ref], + receiver: object = _OMITTED, + owner: Ref | None = None, +) -> tuple[Ref, object | None]: """ - if isinstance(tp, TypeVarRef): - return tp.to_type_var() - return RuntimeClass(decls_thunk, tp) - - -## -# Runtime objects -## - + Bind a Python call directly to its canonical declaration and node records. -class BaseClassFactoryMeta(type): + The callbacks are Python conversion policy, not signature providers. All + parameter ordering, defaults, generic patterns and mutation selection come + from the supplied immutable records. The return's second item identifies + the supplied wrapper to replace for a mutating presentation. """ - Base metaclass for all runtime classes created by ClassFactory - """ - - def __instancecheck__(cls, instance: object) -> bool: - runtime_cls = cast("RuntimeClass", cls) - return ( - isinstance(instance, RuntimeExpr) and runtime_cls.__egg_tp__.ident == instance.__egg_typed_expr__.tp.ident + if declaration.role != "declarations": + msg = "Call authoring requires a declaration reference" + raise TypeError(msg) + record: pb.Declaration = declaration.read() + if record.kind is None: + msg = "Callable declaration has no kind" + raise ValueError(msg) + kind = record.kind.value + binder_size = 0 + varargs = None + match kind: + case pb.HostPrimitive(): + if kind.typing is None or kind.typing.field != "signature": + msg = "Structural function application authoring is not migrated yet" + raise NotImplementedError(msg) + signature = kind.typing.value + if not signature.has_field("output"): + msg = "Host signature has no output" + raise ValueError(msg) + inputs, output = signature.inputs, signature.output + binder_size = len(signature.type_params) + varargs = signature.varargs + case pb.Constructor() | pb.Function() | pb.Primitive(): + inputs, output = kind.inputs, kind.output + case _: + msg = f"Call authoring for {record.kind.field} is not migrated yet" + raise NotImplementedError(msg) + if view_index is None: + # Conservative derived free-function binding: no synthetic metadata is + # persisted or installed. The generator checks its exported path. + view = pb.PythonCallable( + kind=pb.PythonCallKind.FUNCTION, + path=[kind.name], + params=[pb.PythonParameter(core_input=i, name=arg.name or f"arg{i}") for i, arg in enumerate(inputs)], ) + if varargs is not None: + view.params.append(pb.PythonParameter(core_input=len(inputs), name=varargs.name or "args")) + else: + if record.bindings is None or record.bindings.python is None: + msg = "Requested Python view is absent" + raise ValueError(msg) + view = record.bindings.python.views[view_index] + substitutions: dict[int, Ref] = {} + if owner is not None and view.owner is not None and view.owner.kind is not None and view.owner.kind.field == "sort": + infer_sort(declaration.owner.ref("sorts", view.owner.kind.value), owner, binder_size, substitutions) + core_values: dict[int, object] = {} + if view.has_field("receiver"): + if receiver is _OMITTED: + if not args: + msg = "Missing receiver for unbound method" + raise TypeError(msg) + receiver, *remaining = args + args = tuple(remaining) + core_values[view.receiver] = receiver + elif receiver is not _OMITTED: + msg = "This Python view does not accept a receiver" + raise TypeError(msg) + parameters = [ + Parameter( + parameter.name, + Parameter.VAR_POSITIONAL + if varargs is not None and parameter.core_input == len(inputs) + else Parameter.POSITIONAL_OR_KEYWORD, + default=_OMITTED if parameter.has_field("default_expr") else Parameter.empty, + ) + for parameter in view.params + ] + bound = Signature(parameters).bind(*args, **kwargs) + bound.apply_defaults() + omitted = [parameter for parameter in view.params if bound.arguments[parameter.name] is _OMITTED] + defaults = iter(clone_nodes(declaration.owner.ref("nodes", parameter.default_expr) for parameter in omitted)) + for parameter in view.params: + if parameter.core_input in core_values: + msg = "Python view maps one core input more than once" + raise ValueError(msg) + value = bound.arguments[parameter.name] + core_values[parameter.core_input] = next(defaults) if value is _OMITTED else value + expected_slots = set(range(len(inputs) + (varargs is not None))) + if core_values.keys() != expected_slots: + msg = "Python view does not map every core input exactly once" + raise ValueError(msg) + ordered_values = [core_values[index] for index in range(len(inputs))] + patterns = [declaration.owner.ref("sorts", arg.sort) for arg in inputs] + if varargs is not None: + tail = core_values[len(inputs)] + assert isinstance(tail, tuple) + ordered_values.extend(tail) + patterns.extend([declaration.owner.ref("sorts", varargs.sort)] * len(tail)) + existing = [value if isinstance(value, Ref) else node_of(value) for value in ordered_values] + for pattern, node in zip(patterns, existing, strict=True): + if node is not None: + payload: pb.Node = node.read() + # A registered Python conversion may still satisfy this input. + with suppress(TypeConstraintError): + infer_sort(pattern, node.owner.ref("sorts", payload.sort_id), binder_size, substitutions) + converted = [] + for pattern, value, original_node in zip(patterns, ordered_values, existing, strict=True): + node = original_node + expected = substitute_sort(pattern, substitutions) + if node is not None: + payload = node.read() + try: + infer_sort(expected, node.owner.ref("sorts", payload.sort_id), 0, {}) + except TypeConstraintError: + node = None + if node is None: + node = convert(expected, value) + payload = node.read() + infer_sort(expected, node.owner.ref("sorts", payload.sort_id), 0, {}) + converted.append(node) + result_sort = substitute_sort(declaration.owner.ref("sorts", output), substitutions) + if len(substitutions) != binder_size: + msg = "Not every generic parameter was determined by the call" + raise TypeConstraintError(msg) + builder = Builder() + builder.import_ref(declaration) + node_index = builder.add( + "nodes", + pb.Node( + sort_id=builder.import_ref(result_sort), + kind=Oneof[Literal["call"], pb.Call]( + "call", pb.Call(func=kind.name, args=[builder.import_ref(node) for node in converted]) + ), + ), + ) + return builder.publish().ref("nodes", node_index), core_values[view.mutates] if view.has_field("mutates") else None -class ClassFactory(type): - """ - A metaclass for types which should create `type` objects when instantiated. - - That's so that they work with `isinstance` and can be placed in `match ClassName()`. - """ - - def __call__(cls, *args, **kwargs) -> type: - # If we have params, don't inherit from `type` because we don't need to match against this and also - # this won't work with `Union[X]` because it won't look at `__parameters__` for instances of `type`. - if kwargs.pop("_egg_has_params", False): - return super().__call__(*args, **kwargs) - namespace: dict[str, Any] = {} - for m in reversed(cls.__mro__): - if m is not object: - namespace.update(m.__dict__) - init = namespace.pop("__init__") - meta = types.new_class("type(RuntimeClass)", (BaseClassFactoryMeta,), {}, lambda ns: ns.update(**namespace)) - tp = types.new_class("RuntimeClass", (), {"metaclass": meta}, lambda ns: ns.update(RUNTIME_CLASS_DESCRIPTORS)) - init(tp, *args, **kwargs) - return tp - - def __instancecheck__(cls, instance: object) -> bool: - return isinstance(instance, BaseClassFactoryMeta) - - -# Add a custom descriptor for each method we want to support on the class itself, so that we can lookup things like __add__ -# on the expr class. This is used for things like running docstrings. +# These are Python codecs/protocol hooks, not callable signatures. Sort and +# callable definitions still come from the native-generated catalog. +_LITERAL_ARMS = {"i64": "i64", "f64": "f64_bits", "String": "string", "bool": "bool", "Unit": "unit"} +LIT_IDENTS = frozenset(_LITERAL_ARMS) +DUMMY_VALUE = object() +_CLASS_CACHE: WeakKeyDictionary[Owner, WeakValueDictionary[int, RuntimeClass]] = WeakKeyDictionary() +_BASE_CLASSES: dict[str, RuntimeClass] = {} + + +class _ClassContext: + """Owner-retained Python lifecycle/hooks; all declaration semantics remain refs.""" + + def __init__( + self, prepare: Callable[[RuntimeClass], Catalog], hooks: dict[str, Any], match_args: tuple[str, ...] + ) -> None: + self.prepare: Callable[[RuntimeClass], Catalog] | None = prepare + self.hooks = hooks + self.match_args = match_args + self.runtime_class: RuntimeClass + self.catalog: Catalog | None = None + self.resolving = False + self.error: Exception | None = None + self.lock = RLock() + + def resolve(self) -> None: + """Publish the whole member scope once, never a partially usable signature.""" + with self.lock: + if self.error is not None: + raise self.error + if self.resolving: + msg = "Recursively resolving class declarations" + raise RuntimeError(msg) + if self.prepare is None: + return + self.resolving = True + try: + catalog = self.prepare(self.runtime_class) + except Exception as error: + self.error = error + raise + else: + self.catalog = catalog + self.runtime_class.__egg_attr_cache__.clear() + finally: + self.prepare = None + self.resolving = False + + +def _definition_at(origin: Ref, namespace: str, name: str) -> Ref: + """Resolve owned definitions before consulting the ambient native catalog.""" + matches: set[Ref] = set() + for index in range(len(origin.owner._slots["declarations"])): + candidate = origin.owner.ref("declarations", index) + if _definition_key(candidate.read()) == (namespace, name): + matches.add(candidate) + if len(matches) > 1: + raise ValueError(f"Ambiguous local definition: {namespace} {name!r}") + if matches: + return next(iter(matches)) + try: + return builtin_catalog().definitions[namespace, name] + except KeyError as exc: + raise LookupError(f"No canonical definition for {namespace} {name!r}") from exc + + +def _class_for_sort(sort: Ref) -> RuntimeClass: + """Recover a Python wrapper from its canonical concrete sort record.""" + sort = sort.owner.ref("sorts", sort.index) + cache = _CLASS_CACHE.setdefault(sort.owner, WeakValueDictionary()) + if (existing := cache.get(sort.index)) is not None: + return existing + record: pb.Sort = sort.read() + if record.kind is None: + msg = "Expression sort is missing its kind" + raise ValueError(msg) + match record.kind.field: + case "family": + name = record.kind.value.name + case "eq": + name = record.kind.value + case _: + msg = "Function/pattern runtime class projection is not migrated yet" + raise NotImplementedError(msg) + definition = _definition_at(sort, "sort", name) + for context in definition.owner._keepalive: + if isinstance(context, _ClassContext): + return context.runtime_class + result = RuntimeClass(definition, sort=sort) + cache[sort.index] = result + return result -@dataclass class RuntimeClassDescriptor: - name: str + """Expose generated special methods on a runtime class, not its metaclass.""" - def __get__(self, obj: object, owner: RuntimeClass | None = None) -> Callable: + def __init__(self, name: str) -> None: + self.name = name + + def __get__(self, instance: object, owner: RuntimeClass | None = None) -> object: if owner is None: - raise AttributeError(f"Can only access {self.name} on the class, not an instance") + raise AttributeError(self.name) return RuntimeClass.__getattr__(owner, self.name) -RUNTIME_CLASS_DESCRIPTORS: dict[str, RuntimeClassDescriptor] = { - method: RuntimeClassDescriptor(method) for method in CLASS_DESCRIPTOR_METHODS -} - - -@dataclass(match_args=False) -class RuntimeClass(DelayedDeclarations, metaclass=ClassFactory): - __egg_tp__: TypeRefWithVars - # True if we want `__parameters__` to be recognized by `Union`, which means we can't inherit from `type` directly. - _egg_has_params: InitVar[bool] = False - __egg_attr_cache__: dict[str, RuntimeFunction | RuntimeExpr | Callable] = field( - init=False, repr=False, default_factory=dict - ) - - def __post_init__(self, _egg_has_params: bool) -> None: - global _PY_OBJECT_CLASS, _UNSTABLE_FN_CLASS - ident = self.__egg_tp__.ident - if (ident) == Ident.builtin("PyObject"): - _PY_OBJECT_CLASS = self - elif ident == Ident.builtin("UnstableFn") and not self.__egg_tp__.args: - _UNSTABLE_FN_CLASS = self - # Set the module and name so that this works with doctest - if ident.module: - self.__module__ = ident.module - self.__name__ = ident.name - - def verify(self) -> None: - if not self.__egg_tp__.args: - return - - # Raise error if we have args, but they are the wrong number - desired_args = self.__egg_decls__.get_class_decl(self.__egg_tp__.ident).type_vars - if len(self.__egg_tp__.args) != len(desired_args): - raise ValueError(f"Expected {desired_args} type args, got {len(self.__egg_tp__.args)}") - - def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None: - """ - Create an instance of this kind by calling the __init__ classmethod - """ - # If this is a literal type, initializing it with a literal should return a literal - if (name := self.__egg_tp__.ident) == Ident.builtin("PyObject"): - assert len(args) == 1 - try: - pickled = cloudpickle.dumps(args[0]) - except Exception as e: - e.add_note(f"Failed to pickle object of type {type(args[0])}") - raise - return RuntimeExpr( - self.__egg_decls_thunk__, Thunk.value(TypedExprDecl(self.__egg_tp__.to_just(), PyObjectDecl(pickled))) - ) - if name == Ident.builtin("UnstableFn"): - assert not kwargs - fn_arg, *partial_args = args - del args - # Assumes we don't have types set for UnstableFn w/ generics, that they have to be inferred - - # 1. Call it with the partial args, and use untyped vars for the rest of the args - res = cast("Callable", fn_arg)(*partial_args, _egg_function_types=self.__egg_tp__.args) - assert res is not None, "Mutable partial functions not supported" - # 2. Use the inferred return type and inferred rest arg types as the types of the function, and - # the partially applied args as the args. - call = (res_typed_expr := res.__egg_typed_expr__).expr - return_tp = res_typed_expr.tp - assert isinstance(call, CallDecl), "partial function must be a call" - # Clip off the remaining arguments - captured_args = len(partial_args) + int( - isinstance(fn_arg, RuntimeFunction) and isinstance(fn_arg.__egg_bound__, RuntimeExpr) - ) - value = PartialCallDecl(replace(call, args=call.args[:captured_args])) - remaining_arg_types = [a.tp for a in call.args[captured_args:]] - type_ref = JustTypeRef(Ident.builtin("UnstableFn"), (return_tp, *remaining_arg_types)) - return RuntimeExpr.__from_values__(Declarations.create(self, res), TypedExprDecl(type_ref, value)) - - if name in UNARY_LIT_IDENTS: - assert len(args) == 1 - assert isinstance(args[0], int | float | str | bool) - return RuntimeExpr( - self.__egg_decls_thunk__, Thunk.value(TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(args[0]))) - ) - if name == UNIT_IDENT: - assert len(args) == 0 - return RuntimeExpr( - self.__egg_decls_thunk__, Thunk.value(TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(None))) - ) - fn = RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(InitRef(name)), self.__egg_tp__.to_just()) - return fn(*args, **kwargs) # type: ignore[arg-type] - - def __dir__(self) -> Iterable[str]: - cls_decl = self.__egg_decls__._classes[self.__egg_tp__.ident] - return ( - list(cls_decl.class_methods) - + list(cls_decl.class_variables) - + list(cls_decl.preserved_methods) - + list(cls_decl.methods) - + list(cls_decl.properties) +class RuntimeClass(type): + """A Python type presentation over an owned sort/family declaration.""" + + def __new__(cls, definition: Ref, *, sort: Ref | None = None, arguments: tuple[object, ...] = ()) -> RuntimeClass: + declaration: pb.Declaration = definition.read() + kind = declaration.kind + if kind is None or not isinstance(kind.value, pb.HostSortFamily | pb.EqSort): + msg = "RuntimeClass requires a canonical sort declaration" + raise TypeError(msg) + binding = kind.value.bindings + path = ( + tuple(binding.python.path) + if binding is not None and binding.python is not None + else ("egglog", "builtins", kind.value.name) ) - - def __getitem__(self, args: object) -> RuntimeClass: - if not isinstance(args, tuple): - args = (args,) - # defer resolving decls so that we can do generic instantiation for converters before all - # method types are defined. - decls_like, new_args = cast( - "tuple[tuple[DeclarationsLike, ...], tuple[TypeOrVarRef, ...]]", - zip(*(resolve_type_annotation(arg) for arg in args), strict=False), + if not path: + msg = "An empty type path cannot be generated as a Python class" + raise ValueError(msg) + namespace = { + "__module__": ".".join(path[:-1]), + "__doc__": declaration.doc or None, + **{name: RuntimeClassDescriptor(name) for name in CLASS_DESCRIPTOR_METHODS}, + } + result = type.__new__(cls, path[-1], (), namespace) + result.__egg_definition__ = definition + result.__egg_sort__ = sort + result.__egg_arguments__ = arguments + result.__egg_attr_cache__ = {} + result.__egg_hooks__ = {} + result.__egg_pending__ = None + result.__egg_context__ = next( + (context for context in definition.owner._keepalive if isinstance(context, _ClassContext)), None ) - # if we already have some args bound and some not, then we should replace all existing args of typevars with new - # args - if old_args := self.__egg_tp__.args: - is_typevar = [isinstance(arg, TypeVarRef) for arg in old_args] - if sum(is_typevar) != len(new_args): - raise TypeError(f"Expected {sum(is_typevar)} typevars, got {len(new_args)}") - new_args_list = list(new_args) - final_args = tuple(new_args_list.pop(0) if is_typevar[i] else old_args[i] for i in range(len(old_args))) - else: - final_args = new_args - tp = TypeRefWithVars(self.__egg_tp__.ident, final_args) - return RuntimeClass(Thunk.fn(Declarations.create, self, *decls_like), tp, _egg_has_params=True) - - def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable: # noqa: C901 - if not isinstance(name, str): - raise TypeError(f"Attribute name must be a string, got {name!r}") - if name == "__origin__" and self.__egg_tp__.args: - return RuntimeClass(self.__egg_decls_thunk__, TypeRefWithVars(self.__egg_tp__.ident)) - - # Special case some names that don't exist so we can exit early without resolving decls - # Important so if we take union of RuntimeClass it won't try to resolve decls - if name in { - "__typing_subst__", - "__parameters__", - # Origin is used in get_type_hints which is used when resolving the class itself - "__origin__", - "__typing_unpacked_tuple_args__", - "__typing_is_unpacked_typevartuple__", - }: - raise AttributeError - - try: - return self.__egg_attr_cache__[name] - except KeyError: - pass - - try: - cls_decl = self.__egg_decls__._classes[self.__egg_tp__.ident] - except Exception as e: - e.add_note(f"Error processing class {self.__egg_tp__.ident}") - raise - - preserved_methods = cls_decl.preserved_methods - if name in preserved_methods: - res = preserved_methods[name] - # if this is a class variable, return an expr for it, otherwise, assume it's a method - elif name in cls_decl.class_variables: - return_tp = cls_decl.class_variables[name] - res = RuntimeExpr( - self.__egg_decls_thunk__, - Thunk.value(TypedExprDecl(return_tp.type_ref, CallDecl(ClassVariableRef(self.__egg_tp__.ident, name)))), - ) - else: - if name in cls_decl.class_methods: - callable_ref: CallableRef = ClassMethodRef(self.__egg_tp__.ident, name) - # allow referencing properties and methods as class variables as well - elif name in cls_decl.properties: - callable_ref = PropertyRef(self.__egg_tp__.ident, name) - elif name in cls_decl.methods: - callable_ref = MethodRef(self.__egg_tp__.ident, name) + if result.__egg_context__ is not None: + result.__egg_context__.runtime_class = result + result.__egg_hooks__ = result.__egg_context__.hooks + result.__egg_pending__ = result.__egg_context__.resolve + if sort is None and not arguments: + arity = kind.value.arity if isinstance(kind.value, pb.HostSortFamily) else 0 + if arity: + labels = binding.python.type_params if binding is not None and binding.python is not None else [] + result.__egg_arguments__ = tuple( + TypeVar(labels[index] if labels else f"T{index}") for index in range(arity) + ) else: - msg = f"Class {self.__egg_tp__.ident} has no method {name}" - raise AttributeError(msg) from None - res = RuntimeFunction( - self.__egg_decls_thunk__, - Thunk.value(callable_ref), - self.__egg_tp__.to_just() if self.__egg_tp__.args else None, + builder = Builder() + builder.import_ref(definition) + sort_kind = ( + Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name=kind.value.name)) + if isinstance(kind.value, pb.HostSortFamily) + else Oneof[Literal["eq"], str]("eq", kind.value.name) + ) + index = builder.add("sorts", pb.Sort(kind=sort_kind)) + result.__egg_sort__ = builder.publish().ref("sorts", index) + if isinstance(kind.value, pb.HostSortFamily): + _BASE_CLASSES[kind.value.name] = result + if sort is not None and isinstance(kind.value, pb.HostSortFamily) and kind.value.name in _BASE_CLASSES: + result.__egg_hooks__ = _BASE_CLASSES[kind.value.name].__egg_hooks__ + result.__egg_pending__ = _BASE_CLASSES[kind.value.name].__egg_pending__ + if result.__egg_sort__ is not None: + actual = result.__egg_sort__ + _CLASS_CACHE.setdefault(actual.owner, WeakValueDictionary())[actual.index] = result + return result + + def __init__(cls, definition: Ref, *, sort: Ref | None = None, arguments: tuple[object, ...] = ()) -> None: + # type.__init__ expects (name, bases, namespace), not semantic records. + pass + + __egg_definition__: Ref + __egg_sort__: Ref | None + __egg_arguments__: tuple[object, ...] + __egg_attr_cache__: dict[str, object] + __egg_hooks__: dict[str, Any] + __egg_pending__: Callable[[], None] | None + __egg_context__: _ClassContext | None + + def __instancecheck__(cls, instance: object) -> bool: + if not isinstance(instance, RuntimeExpr): + return False + node: pb.Node = instance.__egg_ref__.read() + actual = _class_for_sort(instance.__egg_ref__.owner.ref("sorts", node.sort_id)) + return cls.__egg_definition__ == actual.__egg_definition__ + + def __call__(cls, *args: object, **kwargs: object) -> RuntimeExpr | None: + declaration: pb.Declaration = cls.__egg_definition__.read() + assert declaration.kind is not None + name = declaration.kind.value.name + if isinstance(declaration.kind.value, pb.HostSortFamily) and name in _LITERAL_ARMS: + if kwargs or len(args) != (0 if name == "Unit" else 1): + raise TypeError(f"{cls} expects {'zero' if name == 'Unit' else 'one'} positional literal argument") + if cls.__egg_sort__ is None: + msg = "Literal constructor has an unresolved sort" + raise TypeError(msg) + primitive = pb.PrimitiveValue() + match _LITERAL_ARMS[name]: + case "i64": + primitive.value = Oneof[Literal["i64"], int]("i64", cast("int", args[0])) + case "f64_bits": + bits = struct.unpack("!Q", struct.pack("!d", args[0]))[0] + primitive.value = Oneof[Literal["f64_bits"], int]("f64_bits", bits) + case "string": + primitive.value = Oneof[Literal["string"], str]("string", cast("str", args[0])) + case "bool": + primitive.value = Oneof[Literal["bool"], bool]("bool", cast("bool", args[0])) + case "unit": + primitive.value = Oneof[Literal["unit"], pb.Unit]("unit", pb.Unit()) + builder = Builder() + builder.import_ref(cls.__egg_definition__) + index = builder.add( + "nodes", + pb.Node( + sort_id=builder.import_ref(cls.__egg_sort__), + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]("primitive_value", primitive), + ), ) + return RuntimeExpr(builder.publish().ref("nodes", index)) + if cls.__egg_pending__ is not None: + cls.__egg_pending__() + return cast("RuntimeFunction", RuntimeClass.__getattr__(cls, "__init__"))(*args, **kwargs) + + def __getitem__(cls, arguments: object) -> RuntimeClass: + values = arguments if isinstance(arguments, tuple) else (arguments,) + declaration: pb.Declaration = cls.__egg_definition__.read() + assert declaration.kind is not None + kind = declaration.kind.value + if not isinstance(kind, pb.HostSortFamily) or len(values) != kind.arity: + raise TypeError(f"Wrong number of type arguments for {cls}") + if any(isinstance(value, TypeVar) for value in values): + return RuntimeClass(cls.__egg_definition__, arguments=values) + if not all(isinstance(value, RuntimeClass) and value.__egg_sort__ is not None for value in values): + msg = "Type arguments must resolve to canonical sorts" + raise TypeError(msg) + builder = Builder() + builder.import_ref(cls.__egg_definition__) + indices = [builder.import_ref(cast("Ref", cast("RuntimeClass", value).__egg_sort__)) for value in values] + index = builder.add( + "sorts", + pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name=kind.name, args=indices))), + ) + return _class_for_sort(builder.publish().ref("sorts", index)) + + def __getattr__(cls, name: str) -> object: + if name.startswith("__egg_"): + raise AttributeError(name) + if name == "__origin__": + if cls.__egg_arguments__: + declaration: pb.Declaration = cls.__egg_definition__.read() + assert declaration.kind is not None + return _BASE_CLASSES[declaration.kind.value.name] + # typing probes a concrete class's origin while resolving another + # signature; absence must not recursively compile either class. + raise AttributeError(name) + if name in cls.__egg_hooks__: + return cls.__egg_hooks__[name] + if name in cls.__egg_attr_cache__: + return cls.__egg_attr_cache__[name] + if cls.__egg_pending__ is not None: + cls.__egg_pending__() + declaration = cls.__egg_definition__.read() + assert declaration.kind is not None + key = declaration.kind.value.name, name + catalog = builtin_catalog() + if cls.__egg_context__ is not None: + catalog = cls.__egg_context__.catalog + assert catalog is not None + elif cls.__egg_definition__ != catalog.definitions.get(("sort", key[0])): + raise AttributeError(f"{cls} has no canonical Python member {name!r}") + try: + reference, view = catalog.members[key] + except KeyError as exc: + raise AttributeError(f"{cls} has no generated Python member {name!r}") from exc + function = RuntimeFunction(reference, view, cls) + payload: pb.Declaration = reference.read() + assert payload.bindings is not None + assert payload.bindings.python is not None + result = ( + function() if payload.bindings.python.views[view].kind == pb.PythonCallKind.CLASS_VARIABLE else function + ) + cls.__egg_attr_cache__[name] = result + return result - self.__egg_attr_cache__[name] = res - return res + def __str__(cls) -> str: + if cls.__egg_sort__ is not None: + from .pretty import pretty_ref # noqa: PLC0415 - def __str__(self) -> str: - return str(self.__egg_tp__) + return pretty_ref(cls.__egg_sort__) + return cls.__name__ - def __repr__(self) -> str: - return str(self) + def __repr__(cls) -> str: + return str(cls) - # Make hashable so can go in Union - def __hash__(self) -> int: - return hash(self.__egg_tp__) + def __hash__(cls) -> int: + return ( + hash(StructuralView(cls.__egg_sort__)) + if cls.__egg_sort__ is not None + else hash((cls.__egg_definition__, cls.__egg_arguments__)) + ) - def __eq__(self, other: object) -> bool: - """ - Support equality for runtime comparison of egglog classes. - """ + def __eq__(cls, other: object) -> bool: if not isinstance(other, RuntimeClass): return NotImplemented - return self.__egg_tp__ == other.__egg_tp__ - - # Support unioning like types - def __or__(self, value: type) -> object: - return Union[self, value] # noqa: UP007 - - @property - def __parameters__(self) -> tuple[object, ...]: - """ - Emit a number of typevar params so that when using generic type aliases, we know how to resolve these properly. - """ - return tuple(inverse_resolve_type_annotation(self.__egg_decls_thunk__, tp) for tp in self.__egg_tp__.args) + if cls.__egg_sort__ is not None and other.__egg_sort__ is not None: + return StructuralView(cls.__egg_sort__) == StructuralView(other.__egg_sort__) + return cls.__egg_definition__ == other.__egg_definition__ and cls.__egg_arguments__ == other.__egg_arguments__ - @property - def __match_args__(self) -> tuple[str, ...]: - return self.__egg_decls__._classes[self.__egg_tp__.ident].match_args - - @property - def __doc__(self) -> str | None: # type: ignore[override] - return self.__egg_decls__._classes[self.__egg_tp__.ident].doc + def __or__(cls, other: object) -> Any: + return Union[cls, other] # noqa: UP007 -- runtime union of language class values @property - def __dict__(self) -> dict[str, Any]: # type: ignore[override] - """ - Return the dict so that this works with things like `DocTestFinder`. - """ - return {attr: getattr(self, attr) for attr in dir(self)} - - -class RuntimeFunctionMeta(type): - # Override getatribute on class so that it if we call RuntimeFunction.__module__ we get this module - # but if we call it on an instance we get the module of the instance so that it works with doctest. - def __getattribute__(cls, name: str) -> Any: - if name == "__module__": - return __name__ - return super().__getattribute__(name) - - -@dataclass -class RuntimeFunction(DelayedDeclarations, metaclass=RuntimeFunctionMeta): - __egg_ref_thunk__: Callable[[], CallableRef] - # Either they bound class for something like `Vec[Int].create` or a RuntimeExpr for bound methods - # bound methods need to store RuntimeExpr not just TypedExprDecl, so they can mutate the expr if required on self - __egg_bound__: JustTypeRef | RuntimeExpr | None = None - - __get__: None = None + def __parameters__(cls) -> tuple[object, ...]: + return cls.__egg_arguments__ @property - def __module__(self) -> str | None: # type: ignore[override] - ref = self.__egg_ref__ - if isinstance(ref, UnnamedFunctionRef): - return None - return ref.ident.module + def __match_args__(cls) -> tuple[str, ...]: + if cls.__egg_context__ is not None: + return cls.__egg_context__.match_args + return ("value",) if "value" in cls.__egg_hooks__ else () - def __eq__(self, other: object) -> bool: - """ - Support equality for runtime comparison of egglog functions. - """ - if not isinstance(other, RuntimeFunction): - return NotImplemented - return self.__egg_ref__ == other.__egg_ref__ and bool(self.__egg_bound__ == other.__egg_bound__) - def __hash__(self) -> int: - return hash((self.__egg_ref__, self.__egg_bound__)) +class RuntimeFunction: + """A selected language view over one owned canonical callable declaration.""" - @property - def __egg_ref__(self) -> CallableRef: - return self.__egg_ref_thunk__() + def __init__( + self, declaration: Ref, view_index: int | None = None, bound: RuntimeClass | RuntimeExpr | None = None + ) -> None: + self.__egg_ref__ = declaration + self.__egg_view__ = view_index + self.__egg_bound__ = bound - def __call__( # noqa: C901,PLR0912 - self, *args: object, _egg_function_types: tuple[TypeOrVarRef, ...] | None = None, **kwargs: object - ) -> RuntimeExpr | None: + def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None: from .conversion import resolve_literal # noqa: PLC0415 - if isinstance(self.__egg_bound__, RuntimeExpr): - args = (self.__egg_bound__, *args) - try: - signature = self.__egg_decls__.get_callable_decl(self.__egg_ref__).signature - except Exception as e: - e.add_note(f"Failed to find callable {self}") - raise - decls = self.__egg_decls__.copy() - # Special case function application bc we dont support variadic generics yet generally - if signature == "fn-app": - fn, *rest_args = args - args = tuple(rest_args) - assert not kwargs - assert isinstance(fn, RuntimeExpr) - decls.update(fn) - function_value = fn.__egg_typed_expr__ - fn_tp = function_value.tp - assert fn_tp.ident == Ident.builtin("UnstableFn") - fn_return_tp, *fn_arg_tps = fn_tp.args - signature = FunctionSignature( - tuple(tp.to_var() for tp in fn_arg_tps), - tuple(f"_{i}" for i in range(len(fn_arg_tps))), - (None,) * len(fn_arg_tps), - fn_return_tp.to_var(), - ) - else: - function_value = None - assert isinstance(signature, FunctionSignature) - - # Turn all keyword args into positional args - py_signature = to_py_signature(signature, self.__egg_decls__, optional_args=_egg_function_types is not None) - try: - bound = py_signature.bind(*args, **kwargs) - except TypeError as err: - err.add_note(f"when calling {self} with args {args} and kwargs {kwargs}") - raise - del kwargs - bound.apply_defaults() - assert not bound.kwargs - args = bound.args - decls.update(*(arg for arg in args if isinstance(arg, RuntimeExpr))) - - tcs = TypeConstraintSolver() - if isinstance(self.__egg_bound__, JustTypeRef) and self.__egg_bound__.args: - if function_value: - msg = "Cannot have both bound type params and function value" - raise ValueError(msg) - tcs.bind_class(self.__egg_bound__, decls) - bound_tp_params = self.__egg_bound__.args - else: - bound_tp_params = () - assert (operator.ge if signature.var_arg_type else operator.eq)(len(args), len(signature.arg_types)) - # Hack to allow being explicit on function types when casting. # noqa: FIX004 - for _fn_tp in _egg_function_types or (): - try: - _fn_tp_just = _fn_tp.to_just() - except TypeVarError: - continue - tcs.bind_class(_fn_tp_just, decls) - if _fn_tp_just.args: - pass - # Try using any runtime expressions passed in to help infer typevars - for arg, tp in zip(args, signature.all_args, strict=False): - if not isinstance(arg, RuntimeExpr): - continue - try: - tcs.infer_typevars(tp, arg.__egg_typed_expr__.tp) - # If this leads to an incompatibility, just skip it, since it could need to be upcasted - except TypeConstraintError: - continue - # Now at this point we should be able to resolve all the typevars - upcasted_args = [ - resolve_literal( - tcs.substitute_typevars_try_function(tp, arg, Thunk.value(decls)).to_var(), arg, Thunk.value(decls) - ) - for arg, tp in zip(args, signature.all_args, strict=False) - ] - decls.update(*upcasted_args) - arg_exprs = tuple(arg.__egg_typed_expr__ for arg in upcasted_args) - return_tp = tcs.substitute_typevars(signature.semantic_return_type) - # If we were using unstable-app to call a function, add that function back as the first arg. - if function_value: - arg_exprs = (function_value, *arg_exprs) - expr_decl = CallDecl(self.__egg_ref__, arg_exprs, bound_tp_params) - typed_expr_decl = TypedExprDecl(return_tp, expr_decl) - # If there is not return type, we are mutating the first arg - if not signature.return_type: - first_arg = upcasted_args[0] - first_arg.__egg_decls_thunk__ = Thunk.value(decls) - first_arg.__egg_typed_expr_thunk__ = Thunk.value(typed_expr_decl) + bound = self.__egg_bound__ + owner = bound.__egg_sort__ if isinstance(bound, RuntimeClass) else None + receiver = bound if isinstance(bound, RuntimeExpr) else _OMITTED + root, mutated = _author_call( + self.__egg_ref__, + self.__egg_view__, + args, + kwargs, + node_of=lambda value: value.__egg_ref__ if isinstance(value, RuntimeExpr) else None, + convert=lambda sort, value: resolve_literal(sort, value).__egg_ref__, + receiver=receiver, + owner=owner, + ) + if mutated is not None: + if isinstance(mutated, RuntimeExpr): + mutated.__egg_ref__ = root return None - return RuntimeExpr.__from_values__(decls, typed_expr_decl) + return RuntimeExpr(root) def __str__(self) -> str: - first_arg, bound_tp_params = None, None - match self.__egg_bound__: - case RuntimeExpr(_): - first_arg = self.__egg_bound__.__egg_typed_expr__ - case JustTypeRef(_, args): - bound_tp_params = args - return pretty_callable_ref(self.__egg_decls__, self.__egg_ref__, first_arg, bound_tp_params) + record: pb.Declaration = self.__egg_ref__.read() + assert record.kind is not None + if self.__egg_view__ is None: + return record.kind.value.name + assert record.bindings is not None + assert record.bindings.python is not None + view = record.bindings.python.views[self.__egg_view__] + if view.kind == pb.PythonCallKind.FUNCTION: + return ".".join(view.path) + name = "__init__" if view.kind == pb.PythonCallKind.INITIALIZER else view.path[0] + return f"{self.__egg_bound__}.{name}" def __repr__(self) -> str: return str(self) - @property - def __doc__(self) -> str | None: # type: ignore[override] - decl = self.__egg_decls__.get_callable_decl(self.__egg_ref__) - if isinstance(decl, FunctionDecl | ConstructorDecl): - return decl.doc - return None - - -# sentinel that will be upcasted to any type with a DummyDecl value -DUMMY_VALUE = object() - - -def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: bool) -> Signature: - """ - Convert to a Python signature. - - If optional_args is true, then all args will be treated as optional, as if a default was provided that makes them - `DUMMY_VALUE` + def __hash__(self) -> int: + return hash((self.__egg_ref__, self.__egg_view__, self.__egg_bound__)) - Used for partial application to try binding a function with only some of its args. - """ - parameters = [ - Parameter( - n, - Parameter.POSITIONAL_OR_KEYWORD, - default=RuntimeExpr.__from_values__(decls, TypedExprDecl(t.to_just(), d)) - if d is not None - else (DUMMY_VALUE if optional_args else Parameter.empty), + def __eq__(self, other: object) -> bool: + if not isinstance(other, RuntimeFunction): + return NotImplemented + return ( + self.__egg_ref__ == other.__egg_ref__ + and self.__egg_view__ == other.__egg_view__ + and bool(self.__egg_bound__ == other.__egg_bound__) ) - for n, d, t in zip(sig.arg_names, sig.arg_defaults, sig.arg_types, strict=True) - ] - if isinstance(sig, FunctionSignature) and sig.var_arg_type is not None: - parameters.append(Parameter("__rest", Parameter.VAR_POSITIONAL)) - return Signature(parameters) -ON_CREATE_EXPR: None | Callable[[Callable[[], TypedExprDecl]], None] = None +class RuntimeExpr: + """The sole public symbolic wrapper. Its node Ref is its semantic state.""" + __slots__ = ("__egg_ref__",) -@dataclass -class RuntimeExpr(DelayedDeclarations): - __egg_typed_expr_thunk__: Callable[[], TypedExprDecl] - - def __post_init__(self) -> None: - if ON_CREATE_EXPR: - ON_CREATE_EXPR(self.__egg_typed_expr_thunk__) - - @classmethod - def __from_values__(cls, d: Declarations, e: TypedExprDecl) -> RuntimeExpr: - return cls(Thunk.value(d), Thunk.value(e)) - - def __with_expr__(self, e: TypedExprDecl) -> RuntimeExpr: - return RuntimeExpr(self.__egg_decls_thunk__, Thunk.value(e)) - - @property - def __egg_typed_expr__(self) -> TypedExprDecl: - return self.__egg_typed_expr_thunk__() + def __init__(self, reference: Ref) -> None: + if not isinstance(reference, Ref) or reference.role != "nodes": + msg = "RuntimeExpr requires an owned protobuf node reference" + raise TypeError(msg) + self.__egg_ref__ = reference.owner.ref("nodes", reference.index) def __getattr__(self, name: str) -> object: - if (method := _get_expr_method(self, name)) is not _no_method_sentinel: - return method - if name in self.__egg_class_decl__.properties: - fn = RuntimeFunction( - self.__egg_decls_thunk__, Thunk.value(PropertyRef(self.__egg_class_ident__, name)), self - ) - return fn() - if (getattr_method := _get_expr_method(self, "__getattr__")) is not _no_method_sentinel: - return cast("Callable", getattr_method)(name) - raise AttributeError(f"{self.__egg_class_ident__} has no method {name}") from None - - def __repr__(self) -> str: - """ - The repr of the expr is the pretty printed version of the expr. - """ - if (method := _get_expr_method(self, "__repr__")) is not _no_method_sentinel: - return cast("str", cast("Any", cast("Callable", method)())) - return str(self) + if name.startswith("__egg_"): + raise AttributeError(name) + node: pb.Node = self.__egg_ref__.read() + cls = _class_for_sort(self.__egg_ref__.owner.ref("sorts", node.sort_id)) + if name in cls.__egg_hooks__: + return cls.__egg_hooks__[name].__get__(self, cls) + function = RuntimeClass.__getattr__(cls, name) + if not isinstance(function, RuntimeFunction): + return function + result = RuntimeFunction(function.__egg_ref__, function.__egg_view__, self) + declaration: pb.Declaration = function.__egg_ref__.read() + if function.__egg_view__ is not None: + assert declaration.bindings is not None + assert declaration.bindings.python is not None + if declaration.bindings.python.views[function.__egg_view__].kind == pb.PythonCallKind.PROPERTY: + return result() + return result def __str__(self) -> str: - if (method := _get_expr_method(self, "__str__")) is not _no_method_sentinel: - return cast("str", cast("Any", cast("Callable", method)())) - return self.__egg_pretty__(None) - - def __egg_pretty__(self, wrapping_fn: str | None) -> str: - return pretty_decl(self.__egg_decls__, self.__egg_typed_expr__, wrapping_fn=wrapping_fn) - - def _ipython_display_(self) -> None: - from IPython.display import Code, display # noqa: PLC0415 - - display(Code(str(self), language="python")) + node: pb.Node = self.__egg_ref__.read() + cls = _class_for_sort(self.__egg_ref__.owner.ref("sorts", node.sort_id)) + if "__str__" in cls.__egg_hooks__: + return cls.__egg_hooks__["__str__"].__get__(self, cls)() + from .pretty import pretty_ref # noqa: PLC0415 - def __dir__(self) -> Iterable[str]: - cls_decl = self.__egg_class_decl__ - return ( - list(cls_decl.class_methods) - + list(cls_decl.class_variables) - + list(cls_decl.preserved_methods) - + list(cls_decl.methods) - + list(cls_decl.properties) - ) - - @property - def __egg_class_ident__(self) -> Ident: - return self.__egg_typed_expr__.tp.ident - - @property - def __egg_class_decl__(self) -> ClassDecl: - return self.__egg_decls__.get_class_decl(self.__egg_class_ident__) + return pretty_ref(self.__egg_ref__) - # Implement these so that copy() works on this object - # otherwise copy will try to call `__getstate__` before object is initialized with properties which will cause inifinite recursion - - def __getstate__(self) -> tuple[Declarations, TypedExprDecl]: - return self.__egg_decls__, self.__egg_typed_expr__ - - def __setstate__(self, d: tuple[Declarations, TypedExprDecl]) -> None: - self.__egg_decls_thunk__ = Thunk.value(d[0]) - self.__egg_typed_expr_thunk__ = Thunk.value(d[1]) + def __repr__(self) -> str: + node: pb.Node = self.__egg_ref__.read() + cls = _class_for_sort(self.__egg_ref__.owner.ref("sorts", node.sort_id)) + if "__repr__" in cls.__egg_hooks__: + return cls.__egg_hooks__["__repr__"].__get__(self, cls)() + return str(self) def __hash__(self) -> int: - if (method := _get_expr_method(self, "__hash__")) is not _no_method_sentinel: - return cast("int", cast("Any", cast("Callable", method)())) - return hash(self.__egg_typed_expr__) - - # Implement this directly to special case behavior where it transforms to an egraph equality, if it is not a - # preserved method or defined on the class - def __eq__(self, other: object) -> object: # type: ignore[override] - if (method := _get_expr_method(self, "__eq__")) is not _no_method_sentinel: - return cast("Callable", method)(other) - - if not (isinstance(self, RuntimeExpr) and isinstance(other, RuntimeExpr)): + node: pb.Node = self.__egg_ref__.read() + cls = _class_for_sort(self.__egg_ref__.owner.ref("sorts", node.sort_id)) + if "__hash__" in cls.__egg_hooks__: + if cls.__egg_hooks__["__hash__"] is None: + raise TypeError(f"unhashable type: {cls.__name__!r}") + return cls.__egg_hooks__["__hash__"].__get__(self, cls)() + return hash(StructuralView(self.__egg_ref__)) + + def __eq__(self, other: object) -> Any: + node: pb.Node = self.__egg_ref__.read() + cls = _class_for_sort(self.__egg_ref__.owner.ref("sorts", node.sort_id)) + if "__eq__" in cls.__egg_hooks__: + return cls.__egg_hooks__["__eq__"].__get__(self, cls)(other) + if not isinstance(other, RuntimeExpr): return NotImplemented - if self.__egg_typed_expr__.tp != other.__egg_typed_expr__.tp: + left: pb.Node = self.__egg_ref__.read() + right: pb.Node = other.__egg_ref__.read() + if StructuralView(self.__egg_ref__.owner.ref("sorts", left.sort_id)) != StructuralView( + other.__egg_ref__.owner.ref("sorts", right.sort_id) + ): return NotImplemented + from .egraph import eq # noqa: PLC0415 - from .egraph import Fact # noqa: PLC0415 - - return Fact( - Declarations.create(self, other), - EqDecl(self.__egg_typed_expr__.tp, self.__egg_typed_expr__.expr, other.__egg_typed_expr__.expr), - ) - - def __ne__(self, other: object) -> object: # type: ignore[override] - if (method := _get_expr_method(self, "__ne__")) is not _no_method_sentinel: - return cast("Callable", method)(other) - - from .egraph import BaseExpr, ne # noqa: PLC0415 - - return ne(cast("BaseExpr", self)).to(cast("BaseExpr", other)) + return eq(cast("Any", self)).to(other) - def __call__( - self, *args: object, **kwargs: object - ) -> object: # define it here only for type checking, it will be overriden below - ... + def __ne__(self, other: object) -> Any: + node: pb.Node = self.__egg_ref__.read() + cls = _class_for_sort(self.__egg_ref__.owner.ref("sorts", node.sort_id)) + if "__ne__" in cls.__egg_hooks__: + return cls.__egg_hooks__["__ne__"].__get__(self, cls)(other) + from .egraph import ne # noqa: PLC0415 - def __replace_expr__(self, new_expr: RuntimeExpr) -> None: - self.__egg_decls_thunk__ = new_expr.__egg_decls_thunk__ - self.__egg_typed_expr_thunk__ = new_expr.__egg_typed_expr_thunk__ + return ne(cast("Any", self)).to(other) + def __copy__(self) -> RuntimeExpr: + return RuntimeExpr(self.__egg_ref__) -# Return a sentinel value to differentiate between no method and property that returns None -_no_method_sentinel = object() - - -def _get_expr_method(expr: RuntimeExpr, name: str) -> object: - if name in (preserved_methods := expr.__egg_class_decl__.preserved_methods): - return preserved_methods[name].__get__(expr) - if name in expr.__egg_class_decl__.methods: - return RuntimeFunction(expr.__egg_decls_thunk__, Thunk.value(MethodRef(expr.__egg_class_ident__, name)), expr) - return _no_method_sentinel + def __replace_expr__(self, other: RuntimeExpr) -> None: + self.__egg_ref__ = other.__egg_ref__ def define_expr_method(name: str) -> None: - """ - Given the name of a method, explicitly defines it on the runtime type that holds `Expr` objects as a method. - - Call this if you need a method to be defined on the type itself where overriding with `__getattr__` does not suffice, - like for NumPy's `__array_ufunc__`. - """ - - def _defined_method(self: RuntimeExpr, *args, __name: str = name, **kwargs) -> object: - fn = cast("Callable", _get_expr_method(self, __name)) - if fn is _no_method_sentinel: - if __name == "__hash__": - return hash(self.__egg_typed_expr__) - raise AttributeError(f"{self.__egg_class_ident__} expression has no method {__name}") - return fn(*args, **kwargs) + """Install a Python special-method entrypoint using the generated member view.""" - setattr(RuntimeExpr, name, _defined_method) + def invoke(self: RuntimeExpr, *args: object, **kwargs: object) -> object: + return cast("Callable", self.__getattr__(name))(*args, **kwargs) + setattr(RuntimeExpr, name, invoke) -for name in TYPE_DEFINED_METHODS: - define_expr_method(name) +for _name in TYPE_DEFINED_METHODS - {"__hash__"}: + define_expr_method(_name) -for name, r_method in itertools.product(NUMERIC_BINARY_METHODS, (False, True)): - method_name = f"__r{name[2:]}" if r_method else name +for _name, _reverse in itertools.product(NUMERIC_BINARY_METHODS, (False, True)): - def _numeric_binary_method( - self: object, other: object, name: str = name, r_method: bool = r_method, method_name: str = method_name - ) -> object: - """ - Implements numeric binary operations. + def binary(self: RuntimeExpr, other: object, name: str = _name, reverse: bool = _reverse) -> object: + if reverse: + from .conversion import convert_to_same_type # noqa: PLC0415 - Tries to find the minimum cost conversion of either the LHS or the RHS, by finding all methods with either - the LHS or the RHS as exactly the right type and then upcasting the other to that type. - """ - from .conversion import ( # noqa: PLC0415 - ConvertError, - convert_to_same_type, - min_binary_conversion, - resolve_type, - ) - - # 1. switch if reversed method - if r_method: - self, other = other, self - # First check if we have a preserved method for this: - if isinstance(self, RuntimeExpr) and ( - (preserved_method := self.__egg_class_decl__.preserved_methods.get(name)) is not None - ): - return preserved_method.__get__(self)(other) - # Then check if the self is a Python type and the other side has a preserved method - if ( - not isinstance(self, RuntimeExpr) - and isinstance(other, RuntimeExpr) - and ((preserved_method := other.__egg_class_decl__.preserved_methods.get(name)) is not None) - ): - try: - new_self = convert_to_same_type(self, other) - except ConvertError: - pass - else: - return preserved_method.__get__(new_self)(other) - - # If the types don't exactly match to start, then we need to try converting one of them, by finding the cheapest conversion - if not ( - isinstance(self, RuntimeExpr) - and isinstance(other, RuntimeExpr) - and ( - self.__egg_decls__.check_binary_method_with_types( - name, self.__egg_typed_expr__.tp, other.__egg_typed_expr__.tp - ) - ) - ): - best_method = min_binary_conversion(name, resolve_type(self), resolve_type(other)) - - if not best_method: - raise RuntimeError( - f"Cannot resolve {name} for {resolve_type(self)} and {resolve_type(other)}, no conversion found" - ) - self, other = best_method[0](self), best_method[1](other) - - method_ref = MethodRef(self.__egg_class_ident__, name) - fn = RuntimeFunction(Thunk.value(self.__egg_decls__), Thunk.value(method_ref), self) - return fn(other) + return cast("Callable", convert_to_same_type(other, self).__getattr__(name))(self) + return cast("Callable", self.__getattr__(name))(other) - setattr(RuntimeExpr, method_name, _numeric_binary_method) + setattr(RuntimeExpr, f"__r{_name[2:]}" if _reverse else _name, binary) -def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]: - """ - Resolves a runtime callable into a ref - """ - # TODO: Make runtime class work with __match_args__ - if isinstance(callable, RuntimeClass): - return InitRef(callable.__egg_tp__.ident), callable.__egg_decls__ - match callable: - case RuntimeFunction(decls, ref, _): - return ref(), decls() - case RuntimeExpr(decl_thunk, expr_thunk): - if not isinstance((expr := expr_thunk().expr), CallDecl) or not isinstance( - expr.callable, ConstantRef | ClassVariableRef - ): - raise NotImplementedError(f"Can only turn constants or classvars into callable refs, not {expr}") - return expr.callable, decl_thunk() - case types.MethodWrapperType() if isinstance((slf := callable.__self__), RuntimeClass): - return MethodRef(slf.__egg_tp__.ident, callable.__name__), slf.__egg_decls__ - case _: - raise NotImplementedError(f"Cannot turn {callable} of type {type(callable)} into a callable ref") - - -def create_callable(decls: Declarations, ref: CallableRef) -> RuntimeClass | RuntimeFunction | RuntimeExpr: - """ - Creates a callable object from a callable ref. This might not actually be callable, if the ref is a constant - or classvar then it is a value - """ - match ref: - case InitRef(name): - return RuntimeClass(Thunk.value(decls), TypeRefWithVars(name)) - case FunctionRef() | MethodRef() | ClassMethodRef() | PropertyRef() | UnnamedFunctionRef(): - return RuntimeFunction(Thunk.value(decls), Thunk.value(ref), None) - case ConstantRef(name): - tp = decls._constants[name].type_ref - case ClassVariableRef(cls_name, var_name): - tp = decls._classes[cls_name].class_variables[var_name].type_ref - case _: - assert_never(ref) - return RuntimeExpr.__from_values__(decls, TypedExprDecl(tp, CallDecl(ref))) +def resolve_type_annotation(tp: object) -> Ref | TypeVar | RuntimeClass: + """Keep unresolved Python generics as language state; concrete types are refs.""" + if isinstance(tp, TypeVar): + return tp + if get_origin(tp) is Union: + return resolve_type_annotation(get_args(tp)[0]) + if isinstance(tp, RuntimeClass): + return tp.__egg_sort__ if tp.__egg_sort__ is not None else tp + raise TypeError(f"Unexpected type annotation {tp!r}") + + +def resolve_callable(value: object) -> Ref: + if isinstance(value, RuntimeFunction): + return value.__egg_ref__ + if isinstance(value, RuntimeClass): + member = RuntimeClass.__getattr__(value, "__init__") + if isinstance(member, RuntimeFunction): + return member.__egg_ref__ + if isinstance(value, RuntimeExpr): + node: pb.Node = value.__egg_ref__.read() + if node.kind is not None and node.kind.field == "call": + return _definition_at(value.__egg_ref__, "callable", node.kind.value.func) + raise TypeError(f"Cannot resolve a canonical callable for {value!r}") + + +def create_callable(reference: Ref) -> RuntimeFunction: + record: pb.Declaration = reference.read() + view = ( + 0 + if record.bindings is not None and record.bindings.python is not None and record.bindings.python.views + else None + ) + return RuntimeFunction(reference, view) diff --git a/python/egglog/type_constraint_solver.py b/python/egglog/type_constraint_solver.py index d8e83b4d..a5b97d28 100644 --- a/python/egglog/type_constraint_solver.py +++ b/python/egglog/type_constraint_solver.py @@ -1,134 +1,147 @@ -"""Provides a class for solving type constraints.""" +"""Solve generic signatures directly over canonical sort records.""" from __future__ import annotations -from collections.abc import Callable -from dataclasses import dataclass, field -from itertools import chain, repeat -from typing import TYPE_CHECKING, assert_never +from egglog_proto.egglog.v1 import egglog_pb as pb -from .declarations import * +from ._program import Builder, Ref, StructuralView, _copy_record -if TYPE_CHECKING: - from collections.abc import Collection, Iterable - - -__all__ = ["TypeConstraintError", "TypeConstraintSolver"] +__all__ = ["TypeConstraintError", "infer_sort", "substitute_sort"] class TypeConstraintError(RuntimeError): """Typing error when trying to infer the return type.""" -@dataclass -class TypeConstraintSolver: - """ - Given some typevars and types, solves the constraints to resolve the typevars. - """ - - # Mapping of typevar index to inferred type for each class - _typevar_to_type: dict[Ident, JustTypeRef] = field(default_factory=dict, init=False) - - def bind_class(self, ref: JustTypeRef, decls: Declarations) -> None: - """ - Bind the typevars of a class to the given types. - Used for a situation like Map[int, str].create(). - - This is the same as binding the typevars of the class to the given types. - """ - try: - cls_typevars = decls.get_class_decl(ref.ident).type_vars - except KeyError: - cls_typevars = () - for typevar, arg in zip(cls_typevars, ref.args, strict=True): - self.infer_typevars(typevar, arg) - - def infer_arg_types( - self, - fn_args: Collection[TypeOrVarRef], - fn_return: TypeOrVarRef, - fn_var_args: TypeOrVarRef | None, - return_: JustTypeRef, - ) -> Iterable[JustTypeRef]: - """ - Given a return type, infer the argument types. If there is a variable arg, it returns an infinite iterable. - """ - self.infer_typevars(fn_return, return_) - arg_types = [self.substitute_typevars(fn_arg) for fn_arg in fn_args] - if fn_var_args is None: - return arg_types - var_arg_type = self.substitute_typevars(fn_var_args) - return chain(arg_types, repeat(var_arg_type)) - - def infer_typevars(self, fn_arg: TypeOrVarRef, arg: JustTypeRef) -> None: - """ - Infer typevars from a function argument and a given type, raises TypeConstraintError if they are incompatible. - """ - match fn_arg: - case TypeRefWithVars(cls_ident, fn_args): - if cls_ident != arg.ident: - raise TypeConstraintError(f"Expected {cls_ident}, got {arg.ident}") - for inner_fn_arg, inner_arg in zip(fn_args, arg.args, strict=True): - self.infer_typevars(inner_fn_arg, inner_arg) - case TypeVarRef(typevar_ident): - if typevar_ident in self._typevar_to_type: - if self._typevar_to_type[typevar_ident] != arg: - raise TypeConstraintError(f"Expected {self._typevar_to_type[typevar_ident]}, got {arg}") - else: - self._typevar_to_type[typevar_ident] = arg - case _: - assert_never(fn_arg) - - def substitute_typevars(self, tp: TypeOrVarRef) -> JustTypeRef: - """ - Substitute typevars in a type with their inferred types, raises TypeConstraintError if a typevar is unresolved. - """ - match tp: - case TypeVarRef(typevar_ident): - try: - return self._typevar_to_type[typevar_ident] - except KeyError as e: - raise TypeConstraintError(f"Unresolved type variable: {typevar_ident}") from e - case TypeRefWithVars(name, args): - return JustTypeRef(name, tuple(self.substitute_typevars(arg) for arg in args)) - assert_never(tp) - - def substitute_typevars_try_function( - self, tp: TypeOrVarRef, value: Callable, decls: Callable[[], Declarations] - ) -> JustTypeRef: - """ - Try to substitute typevars in a type with their inferred types. - - If this fails and we have an UnstableFn type and a function value, we can try to infer the typevars by calling - it with the input types, if we can resolve those - """ - from .egraph import set_current_ruleset # noqa: PLC0415 - from .runtime import RuntimeExpr # noqa: PLC0415 - - try: - return self.substitute_typevars(tp) - except TypeConstraintError: - if isinstance(tp, TypeVarRef) or tp.ident != Ident.builtin("UnstableFn") or not callable(value): - raise - # Probe against an isolated copy of the declarations with no ambient ruleset so temporary - # unnamed-function declarations and bodies are discarded after type inference. - probe_decls = decls().copy() - dummy_args = [ - RuntimeExpr.__from_values__( - probe_decls, - TypedExprDecl(self.substitute_typevars(arg_tp), DummyDecl()), - ) - for arg_tp in tp.args[1:] - ] - try: - with set_current_ruleset(None): - result = value(*dummy_args) - except Exception as e: - e.add_note(f"While trying to infer return type of {value} by calling it") - raise - if not isinstance(result, RuntimeExpr): - raise TypeConstraintError( - f"Function {value} did not return a RuntimeExpr, got {type(result)}, so cannot infer return type" - ) - self.infer_typevars(tp.args[0], result.__egg_typed_expr__.tp) - return self.substitute_typevars(tp) +def infer_sort(pattern: Ref, actual: Ref, binder_size: int, bindings: dict[int, Ref]) -> None: # noqa: C901, PLR0912 + """Unify canonical sort records transactionally within one signature binder.""" + pending = [(pattern, actual)] + visited: set[tuple[Ref, Ref]] = set() + staged = bindings.copy() + while pending: + expected, received = pending.pop() + if expected.role != "sorts" or received.role != "sorts": + msg = "Type inference requires sort references" + raise TypeConstraintError(msg) + expected = expected.owner.ref("sorts", expected.index) + received = received.owner.ref("sorts", received.index) + if (expected, received) in visited: + continue + visited.add((expected, received)) + lhs: pb.Sort = expected.read() + rhs: pb.Sort = received.read() + if lhs.kind is None or rhs.kind is None: + msg = "Cannot infer an absent sort kind" + raise TypeConstraintError(msg) + if rhs.kind.field == "var": + msg = "Actual argument sorts must be closed" + raise TypeConstraintError(msg) + if lhs.kind.field == "var": + variable = lhs.kind.value + if variable >= binder_size: + raise TypeConstraintError(f"Sort variable {variable} is outside the signature binder") + if variable in staged: + if StructuralView(staged[variable]) != StructuralView(received): + raise TypeConstraintError(f"Inconsistent sort for type parameter {variable}") + else: + # Check the entire candidate, including nested function/family args. + concrete = [received] + seen: set[Ref] = set() + while concrete: + child = concrete.pop() + if child in seen: + continue + seen.add(child) + record: pb.Sort = child.read() + if record.kind is None or record.kind.field == "var": + msg = "Actual argument sorts must be closed" + raise TypeConstraintError(msg) + match record.kind.field: + case "family": + indices = record.kind.value.args + case "func": + indices = [*record.kind.value.params, record.kind.value.result] + case "eq": + indices = [] + concrete.extend(child.owner.ref("sorts", index) for index in indices) + staged[variable] = received + continue + if lhs.kind.field != rhs.kind.field: + raise TypeConstraintError(f"Expected {lhs.kind.field} sort, got {rhs.kind.field}") + match lhs.kind.field: + case "eq": + if lhs.kind.value != rhs.kind.value: + raise TypeConstraintError(f"Expected {lhs.kind.value}, got {rhs.kind.value}") + continue + case "family": + assert rhs.kind.field == "family" + if lhs.kind.value.name != rhs.kind.value.name: + raise TypeConstraintError(f"Expected {lhs.kind.value.name}, got {rhs.kind.value.name}") + expected_children = lhs.kind.value.args + received_children = rhs.kind.value.args + case "func": + assert rhs.kind.field == "func" + expected_children = [*lhs.kind.value.params, lhs.kind.value.result] + received_children = [*rhs.kind.value.params, rhs.kind.value.result] + if len(expected_children) != len(received_children): + msg = "Sort arity mismatch" + raise TypeConstraintError(msg) + pending.extend( + (expected.owner.ref("sorts", left), received.owner.ref("sorts", right)) + for left, right in zip(expected_children, received_children, strict=True) + ) + bindings.clear() + bindings.update(staged) + + +def substitute_sort(pattern: Ref, bindings: dict[int, Ref]) -> Ref: # noqa: C901 + """Instantiate sort records iteratively, preserving closed imported records.""" + builder = Builder() + resolved: dict[Ref, int] = {} + active: set[Ref] = set() + pending = [(pattern, False)] + while pending: + current, exiting = pending.pop() + current = current.owner.ref(current.role, current.index) + if current.role != "sorts": + msg = "Type substitution requires sort references" + raise TypeConstraintError(msg) + if current in resolved: + continue + record: pb.Sort = current.read() + if record.kind is None: + msg = "Cannot substitute an absent sort kind" + raise TypeConstraintError(msg) + if record.kind.field == "var": + try: + replacement = bindings[record.kind.value] + except KeyError as exc: + raise TypeConstraintError(f"Unresolved type variable: {record.kind.value}") from exc + resolved[current] = builder.import_ref(replacement) + continue + match record.kind.field: + case "family": + children = record.kind.value.args + case "func": + children = [*record.kind.value.params, record.kind.value.result] + case "eq": + children = [] + if not children: + resolved[current] = builder.import_ref(current) + continue + if exiting: + + def relocate(role: str, index: int, origin: Ref = current) -> int: + reference = origin.owner.ref(role, index) + return resolved[reference] if role == "sorts" else builder.import_ref(reference) + + resolved[current] = builder.add("sorts", _copy_record(record, relocate=relocate)) + active.remove(current) + continue + if current in active: + msg = "Sort patterns cannot contain cycles" + raise TypeConstraintError(msg) + active.add(current) + pending.append((current, True)) + pending.extend((current.owner.ref("sorts", index), False) for index in children) + return builder.publish().ref("sorts", resolved[pattern.owner.ref(pattern.role, pattern.index)]) diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 6a0b230d..0b07eb98 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -1,4 +1,5 @@ import copy +import sys import pytest from syrupy.extensions.single_file import SingleFileSnapshotExtension @@ -9,18 +10,15 @@ def _reset_conversions(): from egglog import conversion # noqa: PLC0415 old_conversions = copy.copy(conversion.CONVERSIONS) - old_conversion_decls = copy.copy(conversion._TO_PROCESS_DECLS) yield conversion.CONVERSIONS = old_conversions - conversion._TO_PROCESS_DECLS = old_conversion_decls @pytest.fixture(autouse=True) def _reset_current_egraph(): - from egglog.exp import array_api # noqa: PLC0415 - yield - array_api._CURRENT_EGRAPH = None + if (array_api := sys.modules.get("egglog.exp.array_api")) is not None: + array_api._CURRENT_EGRAPH = None class PythonSnapshotExtension(SingleFileSnapshotExtension): diff --git a/python/tests/test_catalog.py b/python/tests/test_catalog.py new file mode 100644 index 00000000..ce9215cc --- /dev/null +++ b/python/tests/test_catalog.py @@ -0,0 +1,108 @@ +"""Canonical catalog indexing, not a replacement builtin signature inventory.""" + +from typing import Literal + +import pytest +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import Oneof + +from egglog._catalog import Catalog + + +def catalog_fixture() -> pb.Program: + return pb.Program( + ir_version=1, + sorts=[pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Number")))], + declarations=[ + pb.Declaration( + kind=Oneof[Literal["host_sort_family"], pb.HostSortFamily]( + "host_sort_family", + pb.HostSortFamily( + name="Number", bindings=pb.SortBindings(python=pb.TypeBinding(path=["example", "Number"])) + ), + ) + ), + pb.Declaration( + kind=Oneof[Literal["host_primitive"], pb.HostPrimitive]( + "host_primitive", + pb.HostPrimitive( + name="example.number.add", + typing=Oneof[Literal["signature"], pb.GenericSignature]( + "signature", pb.GenericSignature(inputs=[pb.Arg(sort=0), pb.Arg(sort=0)], output=0) + ), + ), + ), + bindings=pb.CallableBindings( + python=pb.PythonBindings( + views=[ + pb.PythonCallable( + kind=pb.PythonCallKind.METHOD, + path=["__add__"], + owner=pb.BindingOwner(kind=Oneof[Literal["sort"], int]("sort", 0)), + receiver=0, + params=[pb.PythonParameter(core_input=1, name="other")], + ) + ] + ) + ), + ), + ], + ) + + +def test_catalog_indexes_records_without_evaluating_or_retyping(monkeypatch): + def forbidden(*args, **kwargs): + pytest.fail("Catalog loading cannot create an engine") + + monkeypatch.setattr("egglog.bindings.EGraph", forbidden) + source = catalog_fixture() + catalog = Catalog(source) + family = catalog.types["example", "Number"] + assert family.read().kind.value.name == "Number" + declaration, view = catalog.members["Number", "__add__"] + assert declaration.read().kind.value.name == "example.number.add" + assert view == 0 + source.declarations.clear() + assert declaration.read().kind.value.typing.value.inputs[0].sort == 0 + catalog.reindex() + assert catalog.members["Number", "__add__"] == (declaration, view) + + +def test_explicit_index_does_not_persist_invented_python_bindings(): + source = catalog_fixture() + source.declarations[1].bindings = None + catalog = Catalog(source) + assert not catalog.members + declaration = catalog.definitions["callable", "example.number.add"].read() + assert declaration.bindings is None + + +def test_local_scope_indexes_canonical_refs_without_adopting_them_again(): + source = Catalog(catalog_fixture()) + scoped = Catalog.from_refs([*source.roots, source.roots[0]]) + assert scoped.roots == source.roots + assert scoped.members == source.members + scoped.members.clear() + scoped.reindex() + assert scoped.members == source.members + assert all(ref.owner is source.owner for ref in scoped.roots) + + +def test_catalog_rejects_ambiguous_presentation_and_invalid_paths(): + source = catalog_fixture() + duplicate = pb.Declaration.from_binary(source.declarations[1].to_binary()) + assert duplicate.kind is not None + assert duplicate.kind.field == "host_primitive" + duplicate.kind.value.name = "example.other.add" + source.declarations.append(duplicate) + with pytest.raises(ValueError, match="binding collision"): + Catalog(source) + source = catalog_fixture() + kind = source.declarations[0].kind + assert kind is not None + assert kind.field == "host_sort_family" + assert kind.value.bindings is not None + assert kind.value.bindings.python is not None + kind.value.bindings.python.path = ["bad-name"] + with pytest.raises(ValueError, match="Python path"): + Catalog(source) diff --git a/python/tests/test_program.py b/python/tests/test_program.py new file mode 100644 index 00000000..f7a2115b --- /dev/null +++ b/python/tests/test_program.py @@ -0,0 +1,618 @@ +"""Storage and relocation tests, separate from engine conformance.""" + +from __future__ import annotations + +import gc +from typing import Literal +from weakref import ref as weakref + +import pytest +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import DescFieldValueList, Oneof + +from egglog import _program +from egglog._program import Builder, StructuralView, clone_nodes, pack + + +def literal(value: int): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + index = builder.add( + "nodes", + pb.Node( + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", pb.PrimitiveValue(value=Oneof[Literal["i64"], int]("i64", value)) + ) + ), + ) + return builder.publish().ref("nodes", index) + + +AMBIENT = frozenset({("sort", "i64"), ("callable", "f"), ("callable", "g")}) + + +def test_template_expansion_freshens_nodes_not_definition_identity(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "Box"))) + builder.add( + "declarations", pb.Declaration(kind=Oneof[Literal["eq_sort"], pb.EqSort]("eq_sort", pb.EqSort(name="Box"))) + ) + declaration = builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["constructor"], pb.Constructor]("constructor", pb.Constructor(name="box", output=0)) + ), + ) + call = builder.add("nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="box")))) + group = builder.add("nodes", pb.Node(kind=Oneof[Literal["union"], pb.Union]("union", pb.Union(members=[call])))) + owner = builder.publish() + first, alias = clone_nodes([owner.ref("nodes", group), owner.ref("nodes", group)]) + (second,) = clone_nodes([owner.ref("nodes", group)]) + assert first == alias + assert first != second + assert first != owner.ref("nodes", group) + cloned_call = first.owner.ref("nodes", first.read().kind.value.members[0]) + assert cloned_call != owner.ref("nodes", call) + assert first.owner.ref("declarations", 0) == owner.ref("declarations", declaration) + combined = pack([first, second, owner.ref("nodes", call)]) + assert len(combined.program.declarations) == 2 + assert sum(node.kind is not None and node.kind.field == "union" for node in combined.program.nodes) == 2 + + +@pytest.fixture(autouse=True) +def no_live_engine(monkeypatch): + def forbidden(*args, **kwargs): + msg = "Protobuf ownership must not instantiate an engine" + raise AssertionError(msg) + + monkeypatch.setattr("egglog.bindings.EGraph", forbidden) + + +def test_slot_permutation_and_protobuf_fields_are_authoritative(): + a, b, unused = literal(10), literal(20), literal(30) + + def authored(imports, args, *, func="f", kind=None): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + for imported in imports: + builder.import_ref(imported) + root = builder.add( + "nodes", pb.Node(kind=kind or Oneof[Literal["call"], pb.Call]("call", pb.Call(func=func, args=args))) + ) + return pack([builder.publish().ref("nodes", root)], ambient=AMBIENT) + + original = authored([a, b, unused], [0, 1, 0]) + permuted = authored([b, unused, a], [2, 0, 2]) + assert original.program == permuted.program + assert len(original.program.nodes) == 3 + call = original.program.nodes[original.indices[0]].kind.value + assert call.args[0] == call.args[2] != call.args[1] + changed = authored([b, unused, a], [0], func="g") + assert changed.program.nodes[changed.indices[0]].kind.value.func == "g" + assert len(changed.program.nodes) == 2 + variable = authored([a, b], [0, 1], kind=Oneof[Literal["var"], str]("var", "x")) + assert len(variable.program.nodes) == 1 + assert variable.program.nodes[0].kind == Oneof[Literal["var"], str]("var", "x") + + +def test_adoption_publication_and_inspection_are_isolated(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + supplied = pb.Node( + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", pb.PrimitiveValue(value=Oneof[Literal["i64"], int]("i64", 7)) + ) + ) + index = builder.add("nodes", supplied) + assert supplied.kind is not None + assert supplied.kind.field == "primitive_value" + supplied.kind.value.value = Oneof[Literal["i64"], int]("i64", 99) + owner = builder.publish() + ref = owner.ref("nodes", index) + detached = ref.read() + detached.kind.value.value = Oneof[Literal["i64"], int]("i64", 100) + assert ref.read().kind.value.value.value == 7 + with pytest.raises(RuntimeError, match="published"): + builder.add("nodes", supplied) + with pytest.raises(RuntimeError, match="published"): + builder.publish() + packet = pack([ref], ambient=AMBIENT).program + packed = packet.nodes[0] + assert packed.kind is not None + assert packed.kind.field == "primitive_value" + packed.kind.value.value = Oneof[Literal["i64"], int]("i64", 101) + assert ref.read().kind.value.value.value == 7 + + +def test_diamonds_distinct_equal_unions_and_cycles(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "E"))) + first = builder.reserve("nodes") + child = builder.add( + "nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="cycle", args=[first]))) + ) + builder.fill("nodes", first, pb.Node(kind=Oneof[Literal["union"], pb.Union]("union", pb.Union(members=[child])))) + fresh = builder.add("nodes", pb.Node(kind=Oneof[Literal["union"], pb.Union]("union", pb.Union()))) + other = builder.add("nodes", pb.Node(kind=Oneof[Literal["union"], pb.Union]("union", pb.Union()))) + source = builder.publish() + branches = [] + for _ in range(2): + branch = Builder() + root = branch.import_ref(source.ref("nodes", first)) + branches.append(branch.publish().ref("nodes", root)) + result = pack( + [*branches, source.ref("nodes", fresh), source.ref("nodes", other)], + ambient={("sort", "E"), ("callable", "cycle")}, + ) + assert result.indices[0] == result.indices[1] + assert result.indices[2] != result.indices[3] + union = result.program.nodes[result.indices[0]] + assert union.kind is not None + assert union.kind.field == "union" + call = result.program.nodes[union.kind.value.members[0]] + assert call.kind is not None + assert call.kind.field == "call" + assert call.kind.value.args == [result.indices[0]] + assert pb.Program.from_binary(result.program.to_binary()) == result.program + + +def test_named_recursive_definition_closure_and_missing_names(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + builder.add("nodes", pb.Node(kind=Oneof[Literal["var"], str]("var", "_0"))) + builder.add("nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f", args=[0])))) + builder.add("nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="g", args=[0])))) + builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["host_sort_family"], pb.HostSortFamily]( + "host_sort_family", pb.HostSortFamily(name="i64") + ) + ), + ) + for name, body in [("f", 2), ("g", 1), ("unreachable", 0)]: + builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["primitive"], pb.Primitive]( + "primitive", pb.Primitive(name=name, inputs=[pb.Arg(sort=0)], output=0, body=body) + ) + ), + ) + owner = builder.publish() + result = pack([owner.ref("nodes", 1)]) + assert all(d.kind is not None for d in result.program.declarations) + assert {d.kind.value.name for d in result.program.declarations if d.kind is not None} == {"f", "g", "i64"} + assert len(result.program.nodes) == 3 + with pytest.raises(ValueError, match=r"missing.*i64"): + pack([literal(1)]) + + +def test_rules_rulesets_commands_files_and_command_multiplicity(): + builder = Builder() + builder.add("files", pb.SourceFile(name="source", contents="x")) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + builder.add("nodes", pb.Node(kind=Oneof[Literal["var"], str]("var", "x"), span=pb.Span(file=0, end=1))) + rule = pb.RuleDecl( + kind=Oneof[Literal["rule"], pb.Rule]( + "rule", pb.Rule(query=[0], head=[pb.Action(kind=Oneof[Literal["term"], int]("term", 0))]) + ), + eval_mode=pb.RuleEvalMode.NAIVE, + ) + builder.add("rules", rule) + builder.add("rules", rule) + builder.add( + "rulesets", + pb.Ruleset(name="r", kind=Oneof[Literal["rules"], pb.RuleList]("rules", pb.RuleList(rules=[0, 0, 1]))), + ) + command = builder.add( + "commands", + pb.Command( + kind=Oneof[Literal["repeat"], pb.Repeat]( + "repeat", + pb.Repeat( + times=2, + until=[0], + body=[ + pb.Command( + kind=Oneof[Literal["run"], pb.Run]( + "run", pb.Run(ruleset=pb.RulesetRef(kind=Oneof[Literal["name"], str]("name", "r"))) + ) + ) + ], + ), + ) + ), + ) + owner = builder.publish() + reference = owner.ref("commands", command) + result = pack(commands=[reference, reference], ambient=AMBIENT) + assert len(result.program.commands) == 2 + assert len(result.program.rules) == 2 + ruleset = result.program.rulesets[0] + assert ruleset.kind is not None + assert ruleset.kind.field == "rules" + rules = ruleset.kind.value.rules + assert rules[0] == rules[1] != rules[2] + assert result.program.files[0].name == "source" + assert result.program.nodes[0].span is not None + assert result.program.nodes[0].span.file == 0 + + +def test_binders_defaults_and_payloads_are_not_reinterpreted(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 0))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + builder.add("nodes", pb.Node(sort_id=1, kind=Oneof[Literal["var"], str]("var", "_0"))) + for name in ("left", "right"): + builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["host_primitive"], pb.HostPrimitive]( + "host_primitive", + pb.HostPrimitive( + name=name, + typing=Oneof[Literal["signature"], pb.GenericSignature]( + "signature", pb.GenericSignature(type_params=["T"], inputs=[pb.Arg(sort=0)], output=0) + ), + ), + ), + bindings=pb.CallableBindings( + python=pb.PythonBindings( + views=[ + pb.PythonCallable( + kind=pb.PythonCallKind.FUNCTION, + path=[name], + params=[pb.PythonParameter(core_input=0, name="x", default_expr=0)], + ) + ] + ) + ), + ), + ) + owner = builder.publish() + result = pack([owner.ref("declarations", 1), owner.ref("declarations", 0)], ambient=AMBIENT) + assert result.program.sorts[0].kind == Oneof[Literal["var"], int]("var", 0) + assert result.program.nodes[0].kind == Oneof[Literal["var"], str]("var", "_0") + assert len(result.program.declarations) == 2 + for declaration in result.program.declarations: + assert declaration.kind is not None + assert declaration.kind.field == "host_primitive" + typing = declaration.kind.value.typing + assert typing is not None + assert typing.field == "signature" + assert typing.value.type_params == ["T"] + assert declaration.bindings is not None + assert declaration.bindings.python is not None + assert declaration.bindings.python.views[0].params[0].core_input == 0 + # Storage does not claim this free-variable default is semantically valid; + # shared adapter binding checks must reject it if submitted for execution. + + +def test_linear_authoring_and_pack(monkeypatch): + source = literal(1) + count = 0 + original = _program._copy_record + + def counted(record, *args, **kwargs): + nonlocal count + if isinstance(record, pb.Node): + count += 1 + return original(record, *args, **kwargs) + + monkeypatch.setattr(_program, "_copy_record", counted) + current = source + owners = [] + for _ in range(10_000): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + left = builder.import_ref(current) + right = builder.import_ref(source) + index = builder.add( + "nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f", args=[left, right]))) + ) + current = builder.publish().ref("nodes", index) + owners.append(weakref(current.owner)) + assert count == 10_000 + count = 0 + result = pack([current, current], ambient=AMBIENT) + assert count == 10_001 + assert len(result.program.nodes) == 10_001 + assert result.indices[0] == result.indices[1] + assert len(pb.Program.from_binary(result.program.to_binary()).nodes) == 10_001 + del current + gc.collect() + assert all(owner() is None for owner in owners) + assert source.read().kind.value.value.value == 1 + + +def test_inventory_covers_all_program_message_fields(): + _program._check_inventory() + + +def test_owner_keepalive_is_lifetime_only_not_semantic_storage(): + class LocalContext: + pass + + context = LocalContext() + watched = weakref(context) + builder = Builder() + builder.import_ref(literal(7)) + owner = builder.publish(keepalive=(context,)) + del context + gc.collect() + assert watched() is not None + assert pack([owner.ref("nodes", 0)], ambient=AMBIENT).program == pack([literal(7)], ambient=AMBIENT).program + del owner + gc.collect() + assert watched() is None + + +def test_proof_command_relocation_and_constructor_closure(): + builder = Builder() + unused = builder.import_ref(literal(999)) + fact = builder.import_ref(literal(7)) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "Box"))) + builder.add( + "declarations", pb.Declaration(kind=Oneof[Literal["eq_sort"], pb.EqSort]("eq_sort", pb.EqSort(name="Box"))) + ) + builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["constructor"], pb.Constructor]("constructor", pb.Constructor(name="box", output=0)) + ), + ) + prove = builder.add( + "commands", pb.Command(kind=Oneof[Literal["prove"], pb.Prove]("prove", pb.Prove(facts=[fact, fact]))) + ) + exists = builder.add( + "commands", + pb.Command( + kind=Oneof[Literal["prove_exists"], pb.ProveExists]("prove_exists", pb.ProveExists(constructor="box")) + ), + ) + missing = builder.add( + "commands", + pb.Command( + kind=Oneof[Literal["prove_exists"], pb.ProveExists]("prove_exists", pb.ProveExists(constructor="missing")) + ), + ) + owner = builder.publish() + result = pack(commands=[owner.ref("commands", prove), owner.ref("commands", exists)], ambient=AMBIENT).program + assert result.commands[0].kind.value.facts == [0, 0] + assert unused != fact == 1 + assert len(result.nodes) == 1 + assert result.nodes[0].kind.value.value.value == 7 + assert result.commands[1].kind.value.constructor == "box" + assert {declaration.kind.value.name for declaration in result.declarations} == {"box", "Box"} + assert pb.Program.from_binary(result.to_binary()) == result + with pytest.raises(ValueError, match=r"missing.*missing"): + pack(commands=[owner.ref("commands", missing)]) + + +@pytest.mark.parametrize(("message", "field", "role"), _program._index_fields()) +def test_each_index_role_remaps_zero_and_preserves_presence(message, field, role): + record = message() + descriptor = next(f for f in message.desc().fields if f.name == field) + repeated = isinstance(descriptor.value, DescFieldValueList) + record[descriptor] = [0, 0] if repeated else 0 + seen = [] + + def relocate(actual_role, index): + seen.append((actual_role, index)) + return 17 + + cloned = _program._copy_record(record, relocate=relocate) + assert cloned[descriptor] == ([17, 17] if repeated else 17) + assert (role, 0) in seen + assert record[descriptor] == ([0, 0] if repeated else 0) + + +def test_rejects_invalid_records_slots_and_inline_cycles(): + builder = Builder() + with pytest.raises(TypeError): + builder.add("nodes", pb.Sort()) + builder.reserve("nodes") + with pytest.raises(ValueError, match="unfilled"): + builder.publish() + malformed = Builder() + malformed.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 0))) + root = malformed.add("nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f", args=[9])))) + with pytest.raises(IndexError): + pack([malformed.publish().ref("nodes", root)], ambient=AMBIENT) + command = pb.Command(kind=Oneof[Literal["repeat"], pb.Repeat]("repeat", pb.Repeat(times=1))) + assert command.kind is not None + assert command.kind.field == "repeat" + command.kind.value.body.append(command) + with pytest.raises(ValueError, match=r"inline.*cycle"): + Builder().add("commands", command) + unknown = pb.Node.from_binary(b"\xf8\x07\x01") + with pytest.raises(NotImplementedError, match="unknown"): + Builder().add("nodes", unknown) + + +def test_detached_default_expansions_preserve_internal_identity(): + default = Builder() + default.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "E"))) + default.add("nodes", pb.Node(kind=Oneof[Literal["union"], pb.Union]("union", pb.Union()))) + root = default.add("nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f", args=[0, 0])))) + owner = default.publish() + packet = pack([owner.ref("nodes", root)], ambient={*AMBIENT, ("sort", "E")}) + refs = [] + for _ in range(2): + imported = Builder.from_program(packet.program).publish() + refs.append(imported.ref("nodes", packet.indices[0])) + result = pack(refs, ambient={*AMBIENT, ("sort", "E")}) + args = [] + for index in result.indices: + node = result.program.nodes[index] + assert node.kind is not None + assert node.kind.field == "call" + args.append(node.kind.value.args) + assert args[0][0] == args[0][1] != args[1][0] == args[1][1] + + +def test_import_cache_is_not_authoritative(): + value = literal(1) + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + first = builder.import_ref(value) + builder._imports.clear() + second = builder.import_ref(value) + root = builder.add( + "nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f", args=[first, second]))) + ) + result = pack([builder.publish().ref("nodes", root)], ambient=AMBIENT) + node = result.program.nodes[result.indices[0]] + assert node.kind is not None + assert node.kind.field == "call" + args = node.kind.value.args + assert args[0] == args[1] + assert len(result.program.nodes) == 2 + + +def test_rejects_mutable_scalar_payload_instead_of_aliasing(): + payload = bytearray(b"payload") + record = pb.Node( + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", + pb.PrimitiveValue( + value=Oneof[Literal["custom"], pb.CustomValue]("custom", pb.CustomValue(payload=payload)) # type: ignore[arg-type] + ), + ) + ) + with pytest.raises(TypeError, match="scalar"): + Builder().add("nodes", record) + + +def test_late_discovered_definition_conflict_is_not_hidden(): + foreign = Builder() + foreign.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + foreign.add("nodes", pb.Node(kind=Oneof[Literal["var"], str]("var", "_0"))) + foreign.add( + "declarations", + pb.Declaration(kind=Oneof[Literal["function"], pb.Function]("function", pb.Function(name="f", output=0))), + ) + foreign_owner = foreign.publish() + local = Builder() + local.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + body = local.import_ref(foreign_owner.ref("nodes", 0)) + root = local.add("nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f")))) + local.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["primitive"], pb.Primitive]("primitive", pb.Primitive(name="f", output=0, body=body)) + ), + ) + with pytest.raises(ValueError, match="same-name"): + pack([local.publish().ref("nodes", root)], ambient=AMBIENT) + + +def test_anonymous_rulesets_do_not_define_the_default_ruleset(): + builder = Builder() + builder.add("rulesets", pb.Ruleset(kind=Oneof[Literal["rules"], pb.RuleList]("rules", pb.RuleList()))) + builder.add("rulesets", pb.Ruleset(kind=Oneof[Literal["rules"], pb.RuleList]("rules", pb.RuleList()))) + builder.add("rulesets", pb.Ruleset(name="", kind=Oneof[Literal["rules"], pb.RuleList]("rules", pb.RuleList()))) + command = builder.add( + "commands", + pb.Command( + kind=Oneof[Literal["run"], pb.Run]( + "run", pb.Run(ruleset=pb.RulesetRef(kind=Oneof[Literal["name"], str]("name", ""))) + ) + ), + ) + owner = builder.publish() + result = pack(commands=[owner.ref("commands", command)]) + assert len(result.program.rulesets) == 1 + assert result.program.rulesets[0].has_field("name") + + +@pytest.mark.parametrize("index", [True, -1, 2**32, "0"]) +def test_rejects_invalid_index_scalars(index): + with pytest.raises(TypeError, match="arena index"): + Builder().add("nodes", pb.Node(sort_id=index)) + + +def test_rejects_non_string_dependency_names(): + with pytest.raises(TypeError, match="definition name"): + Builder().add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name=5)))) # type: ignore[arg-type] + + +def test_distinct_same_name_ruleset_roots_are_rejected(): + refs = [] + for _ in range(2): + builder = Builder() + root = builder.add( + "rulesets", pb.Ruleset(name="r", kind=Oneof[Literal["rules"], pb.RuleList]("rules", pb.RuleList())) + ) + refs.append(builder.publish().ref("rulesets", root)) + with pytest.raises(ValueError, match="same-name"): + pack(refs) + + +def test_structural_view_ignores_addresses_and_preserves_identity(): + first, second = literal(10), literal(10) + assert first != second + assert StructuralView(first) == StructuralView(second) + assert hash(StructuralView(first)) == hash(StructuralView(second)) + assert StructuralView(first) != StructuralView(literal(20)) + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "E"))) + for _ in range(2): + builder.add("nodes", pb.Node(kind=Oneof[Literal["union"], pb.Union]("union", pb.Union()))) + owner = builder.publish() + assert StructuralView(owner.ref("nodes", 0)) == StructuralView(owner.ref("nodes", 0)) + assert StructuralView(owner.ref("nodes", 0)) != StructuralView(owner.ref("nodes", 1)) + + +def test_structural_view_zero_hash_and_origin_are_distinct_from_wire_bits(): + refs = [] + for bits in (0, 1 << 63): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="f64")))) + root = builder.add( + "nodes", + pb.Node( + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", pb.PrimitiveValue(value=Oneof[Literal["f64_bits"], int]("f64_bits", bits)) + ) + ), + ) + refs.append(builder.publish().ref("nodes", root)) + assert refs[0].read() != refs[1].read() + assert StructuralView(refs[0]) == StructuralView(refs[1]) + assert hash(StructuralView(refs[0])) == hash(StructuralView(refs[1])) + + +def test_structural_view_deep_hash_cache_is_erasable_and_does_not_retain_owners(monkeypatch): + original = _program._structural_parts + visits = 0 + + def counted(reference): + nonlocal visits + visits += 1 + return original(reference) + + monkeypatch.setattr(_program, "_structural_parts", counted) + _program._STRUCTURAL_HASHES.clear() + root = literal(1) + owners = [] + for _ in range(10_000): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + child = builder.import_ref(root) + index = builder.add( + "nodes", pb.Node(kind=Oneof[Literal["call"], pb.Call]("call", pb.Call(func="f", args=[child, child]))) + ) + root = builder.publish().ref("nodes", index) + owners.append(weakref(root.owner)) + hash(StructuralView(root)) + assert visits == 4 * 10_001 + first_hash = hash(StructuralView(root)) + assert visits == 4 * 10_001 + _program._STRUCTURAL_HASHES.clear() + assert hash(StructuralView(root)) == first_hash + assert visits == 8 * 10_001 + del root + gc.collect() + assert all(owner() is None for owner in owners) diff --git a/python/tests/test_proto_egraph_state.py b/python/tests/test_proto_egraph_state.py new file mode 100644 index 00000000..af4a9e22 --- /dev/null +++ b/python/tests/test_proto_egraph_state.py @@ -0,0 +1,294 @@ +"""Record/bytes plumbing tests; the inert responder is not engine conformance.""" + +from operator import index +from typing import Literal + +import pytest +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import Oneof + +from egglog import EGraph, bindings, i64 +from egglog._catalog import builtin_catalog +from egglog._program import pack + + +def test_public_bytes_requests_and_response_ownership(monkeypatch: pytest.MonkeyPatch) -> None: + requests: list[pb.RunProgramRequest] = [] + destroyed: list[int] = [] + + class InertTransport: + def create(self, encoded: bytes) -> bytes: + assert isinstance(encoded, bytes) + request = pb.CreateEGraphRequest.from_binary(encoded) + assert request.has_field("threads") + assert request.threads == 1 + assert not request.declarations + assert request.options is not None + assert request.sorts[request.options.cost_sort].kind.value.name == "i64" + return pb.CreateEGraphResponse(egraph_id=1).to_binary() + + def run(self, encoded: bytes) -> bytes: + assert isinstance(encoded, bytes) + request = pb.RunProgramRequest.from_binary(encoded) + requests.append(request) + # Fixed inert reply: no interpretation or computation of the request. + return pb.RunProgramResponse( + sorts=[pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))], + nodes=[ + pb.Node( + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", + pb.PrimitiveValue(value=Oneof[Literal["i64"], int]("i64", 31)), + ) + ) + ], + outputs=[ + pb.CommandOutput( + kind=Oneof[Literal["extraction"], pb.ExtractResult]( + "extraction", + pb.ExtractResult(roots=[pb.ExtractedRoot(variants=[pb.ExtractedTerm(term=0, cost=0)])]), + ), + location=pb.CommandLocation(path=[0]), + ) + ], + ).to_binary() + + def clone_graph(self, encoded: bytes) -> bytes: + assert pb.CloneEGraphRequest.from_binary(encoded).egraph_id == 1 + return pb.CloneEGraphResponse(egraph_id=2).to_binary() + + def destroy(self, encoded: bytes) -> bytes: + destroyed.append(pb.DestroyEGraphRequest.from_binary(encoded).egraph_id) + return pb.DestroyEGraphResponse().to_binary() + + monkeypatch.setattr(bindings, "_ProtoEngine", InertTransport, raising=False) + graph = EGraph() + result, cost = graph.extract(i64(1) + 2, include_cost=True) + assert result.value == cost == 31 + request = requests[0] + assert request.program is not None + program = request.program + assert program.ir_version == 1 + command = program.commands[0].kind + assert command is not None + call = program.nodes[command.value.roots[0]].kind + assert call is not None + assert call.field == "call" + assert call.value.func == "egglog.core.i64.add" + assert [program.nodes[index].kind.value.value.value for index in call.value.args] == [1, 2] + add = next(declaration for declaration in program.declarations if declaration.kind.value.name == call.value.func) + assert add.bindings is not None + assert add.bindings.python is not None + assert add.bindings.python.views[0].path == ["__add__"] + family = next(declaration for declaration in program.declarations if declaration.kind.value.name == "i64") + assert family.kind.value.bindings.python.path == ["egglog", "builtins", "i64"] + # Re-authoring from a returned value retains the original catalog refs, + # not a second declaration copy reconstructed from the response. + repacked = pack([(result + 4).__egg_ref__]) + assert builtin_catalog().definitions["callable", "egglog.core.i64.add"] in repacked.definitions + with graph: + assert graph._state.handle == 2 + assert graph._state.handle == 1 + assert destroyed == [2] + graph.close() + assert graph._state.handle == 1 + graph._state.destroy() + graph._state.destroy() + assert destroyed == [2, 1] + + +@pytest.mark.parametrize("body_raises", [False, True]) +def test_context_restores_saved_state_when_clone_disposal_fails( + monkeypatch: pytest.MonkeyPatch, body_raises: bool +) -> None: + class FaultTransport: + def create(self, encoded: bytes) -> bytes: + return pb.CreateEGraphResponse(egraph_id=1).to_binary() + + def clone_graph(self, encoded: bytes) -> bytes: + return pb.CloneEGraphResponse(egraph_id=2).to_binary() + + def destroy(self, encoded: bytes) -> bytes: + if pb.DestroyEGraphRequest.from_binary(encoded).egraph_id == 2: + msg = "cleanup sentinel" + raise RuntimeError(msg) + return pb.DestroyEGraphResponse().to_binary() + + monkeypatch.setattr(bindings, "_ProtoEngine", FaultTransport, raising=False) + graph = EGraph() + saved = graph._state + with ( # noqa: PT012 -- exercise body and disposal together + pytest.raises( + ValueError if body_raises else RuntimeError, match="body sentinel" if body_raises else "cleanup sentinel" + ), + graph, + ): + assert graph._state is not saved + if body_raises: + msg = "body sentinel" + raise ValueError(msg) + assert graph._state is saved + assert not graph._state_stack + + +class ResourceTransport: + """An inert bytes receiver, not a simulation of native resource allocation.""" + + def __init__(self) -> None: + self.creations: list[pb.CreateEGraphRequest] = [] + self.requests: list[pb.ConfigureEGraphResourcesRequest] = [] + self.threads: dict[int, int] = {} + self.error: ValueError | None = None + + def create(self, encoded: bytes) -> bytes: + request = pb.CreateEGraphRequest.from_binary(encoded) + self.creations.append(request) + handle = len(self.creations) + self.threads[handle] = request.threads or 7 + return pb.CreateEGraphResponse(egraph_id=handle).to_binary() + + def configure_resources(self, encoded: bytes) -> bytes: + assert isinstance(encoded, bytes) + request = pb.ConfigureEGraphResourcesRequest.from_binary(encoded) + self.requests.append(request) + if self.error is not None: + raise self.error + assert request.operation is not None + if request.operation.field == "threads": + self.threads[request.egraph_id] = request.operation.value or 7 + else: + assert request.operation.field == "query" + assert isinstance(request.operation.value, pb.Unit) + return pb.ConfigureEGraphResourcesResponse(threads=self.threads[request.egraph_id]).to_binary() + + def clone_graph(self, encoded: bytes) -> bytes: + request = pb.CloneEGraphRequest.from_binary(encoded) + handle = max(self.threads) + 1 + self.threads[handle] = self.threads[request.egraph_id] + return pb.CloneEGraphResponse(egraph_id=handle).to_binary() + + def destroy(self, encoded: bytes) -> bytes: + del self.threads[pb.DestroyEGraphRequest.from_binary(encoded).egraph_id] + return pb.DestroyEGraphResponse().to_binary() + + +@pytest.fixture +def resource_transport(monkeypatch: pytest.MonkeyPatch) -> ResourceTransport: + transport = ResourceTransport() + monkeypatch.setattr(bindings, "_ProtoEngine", lambda: transport, raising=False) + return transport + + +@pytest.mark.parametrize("use_setter", [False, True], ids=["constructor", "setter"]) +def test_zero_threads_bytes_keep_presence_and_read_actual_count( + resource_transport: ResourceTransport, *, use_setter: bool +) -> None: + graph = EGraph(num_threads=1 if use_setter else 0) + if use_setter: + assert graph.set_num_threads(0) is None + assert resource_transport.requests[0].operation == Oneof[Literal["threads"], int]("threads", 0) + else: + assert resource_transport.creations[0].has_field("threads") + assert resource_transport.creations[0].threads == 0 + assert graph.num_threads() == 7 + assert resource_transport.requests[-1].operation == Oneof[Literal["query"], pb.Unit]("query", pb.Unit()) + + +def test_resource_queries_are_uncached_uint64_and_updates_are_handle_local( + resource_transport: ResourceTransport, +) -> None: + graph = EGraph(num_threads=2) + assert graph.num_threads() == 2 + with graph: + assert graph.num_threads() == 2 + assert graph.set_num_threads(1) is None + assert graph.num_threads() == 1 + assert resource_transport.requests[-1].egraph_id == 2 + assert graph.num_threads() == 2 + assert resource_transport.requests[-1].egraph_id == 1 + resource_transport.threads[1] = (1 << 32) + 9 + assert graph.num_threads() == (1 << 32) + 9 + graph._state.destroy() + calls = len(resource_transport.requests) + with pytest.raises(RuntimeError, match="destroyed"): + graph.num_threads() + with pytest.raises(RuntimeError, match="destroyed"): + graph.set_num_threads(1) + assert len(resource_transport.requests) == calls + + +class IndexedThreadCount: + def __index__(self) -> int: + return 2 + + +class IntOnlyThreadCount: + def __int__(self) -> int: + return 2 + + +@pytest.mark.parametrize("value", [IndexedThreadCount(), True, False, (1 << 32) - 1]) +@pytest.mark.parametrize("use_setter", [False, True], ids=["constructor", "setter"]) +def test_thread_count_accepts_index_protocol( + resource_transport: ResourceTransport, value: object, *, use_setter: bool +) -> None: + graph = EGraph(num_threads=1 if use_setter else value) + if use_setter: + graph.set_num_threads(value) + assert resource_transport.requests[0].operation.value == index(value) + else: + assert resource_transport.creations[0].threads == index(value) + assert graph.num_threads() == (index(value) or 7) + + +@pytest.mark.parametrize( + ("value", "error"), + [ + (-1, OverflowError), + (1 << 32, OverflowError), + (1 << 64, OverflowError), + (None, TypeError), + (1.0, TypeError), + ("1", TypeError), + (IntOnlyThreadCount(), TypeError), + ], +) +@pytest.mark.parametrize("use_setter", [False, True], ids=["constructor", "setter"]) +def test_invalid_thread_count_has_no_transport_effect( + resource_transport: ResourceTransport, value: object, error: type[Exception], *, use_setter: bool +) -> None: + graph = EGraph(num_threads=2) if use_setter else None + creations = len(resource_transport.creations) + if graph is None: + with pytest.raises(error): + EGraph(num_threads=value) + else: + with pytest.raises(error): + graph.set_num_threads(value) + assert len(resource_transport.creations) == creations + assert not resource_transport.requests + if graph is not None: + assert graph.num_threads() == 2 + + +@pytest.mark.parametrize("use_setter", [False, True], ids=["query", "setter"]) +def test_resource_transport_failure_does_not_manufacture_success( + resource_transport: ResourceTransport, *, use_setter: bool +) -> None: + graph = EGraph(num_threads=2) + resource_transport.error = ValueError("resource sentinel") + if use_setter: + with pytest.raises(ValueError, match="resource sentinel"): + graph.set_num_threads(1) + else: + with pytest.raises(ValueError, match="resource sentinel"): + graph.num_threads() + resource_transport.error = None + assert graph.num_threads() == 2 + + +def test_resource_response_requires_positive_actual_count(resource_transport: ResourceTransport) -> None: + graph = EGraph() + resource_transport.threads[1] = 0 + with pytest.raises(ValueError, match="thread count"): + graph.num_threads() diff --git a/python/tests/test_proto_runtime.py b/python/tests/test_proto_runtime.py new file mode 100644 index 00000000..8ab9f7c1 --- /dev/null +++ b/python/tests/test_proto_runtime.py @@ -0,0 +1,503 @@ +"""Record-native authoring gates for the single public runtime migration.""" + +import gc +from concurrent.futures import ThreadPoolExecutor +from copy import copy +from time import sleep +from typing import Literal +from weakref import ref as weakref + +import pytest +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import Oneof + +from egglog import ConvertError, Expr, bindings, eq, expr_parts, i64, method +from egglog._program import Builder, Ref, StructuralView, pack +from egglog.deconstruct import get_callable_args +from egglog.runtime import _CLASS_CACHE, RuntimeExpr, RuntimeFunction, _author_call, resolve_callable +from egglog.type_constraint_solver import TypeConstraintError, infer_sort, substitute_sort + + +def test_user_constructor_defaults_keywords_and_names_are_records(monkeypatch): + def forbidden(*args, **kwargs): + pytest.fail("Class authoring cannot create an engine") + + monkeypatch.setattr(bindings, "EGraph", forbidden) + monkeypatch.setattr(bindings, "_ProtoEngine", forbidden, raising=False) + default = i64(5) + + class Num(Expr, egg_sort="exact sort"): + @method(egg_fn="exact constructor") + def __init__(self, value: i64 = default) -> None: ... + + first = Num() + declaration = resolve_callable(Num) + record = declaration.read() + assert record.kind.value.name == "exact constructor" + assert Num.__egg_definition__.read().kind.value.name == "exact sort" + assert record.bindings.python.views[0].params[0].name == "value" + assert get_callable_args(first, Num)[0].value == 5 + default.__replace_expr__(i64(9)) + assert get_callable_args(Num(), Num)[0].value == 5 + assert expr_parts(Num(value=3)) == expr_parts(Num(3)) + assert Num.__egg_context__.prepare is None + Num.__egg_attr_cache__.clear() + Num.__egg_context__.catalog.members.clear() + Num.__egg_context__.catalog.reindex() + assert resolve_callable(Num) == declaration + assert get_callable_args(Num(), Num)[0].value == 5 + with pytest.raises(TypeError): + Num(1, value=2) + with pytest.raises(TypeError): + Num(unknown=2) + with pytest.raises(TypeError): + Num(1, 2) + with pytest.raises(ConvertError, match="Cannot convert"): + Num("wrong sort") + + +def test_user_preserved_protocols_and_match_args(): + class Local(Expr): + def __init__(self) -> None: ... + + __match_args__ = ("value",) + + @method(preserve=True) + @property + def value(self): + return 4 + + @method(preserve=True) + def tag(self): + return "tag" + + def __str__(self) -> str: + return "local str" + + def __repr__(self) -> str: + return "local repr" + + def __hash__(self) -> int: + return 42 + + @method(preserve=True) + def __eq__(self, other) -> bool: + return other == "sentinel" + + @method(preserve=True) + def __ne__(self, other) -> bool: + return other != "sentinel" + + expression = Local() + assert expression.value == 4 + assert expression.tag() == "tag" + assert str(expression) == "local str" + assert repr(expression) == "local repr" + assert hash(expression) == 42 + assert expression == "sentinel" + assert (expression == "other") is False + assert expression != "other" + assert (expression != "sentinel") is False + assert bool(eq(expression).to(copy(expression))) + match expression: + case Local(4): + pass + case _: + pytest.fail("Explicit match_args and preserved property were lost") + + +def test_local_class_hooks_survive_caches_but_not_all_roots(): + def make(label): + class Local(Expr): + def __init__(self) -> None: ... + + @method(preserve=True) + def tag(self): + return label + + expression = Local() + return expression, weakref(Local), weakref(Local.__egg_definition__.owner), weakref(Local.__egg_context__) + + left, cls, owner, context = make("left") + right, other_cls, other_owner, other_context = make("right") + retained = copy(left) + _CLASS_CACHE.clear() + gc.collect() + assert left.tag() == retained.tag() == "left" + assert right.tag() == "right" + assert cls() is not other_cls() + del left, right, retained + gc.collect() + assert cls() is owner() is context() is None + assert other_cls() is other_owner() is other_context() is None + + +def test_recursive_class_sorts_resolve_without_recursive_member_compilation(): + class Left(Expr): + def __init__(self, right: "Right") -> None: ... + + class Right(Expr): + def __init__(self, left: Left) -> None: ... + + left = resolve_callable(Left) + right = resolve_callable(Right) + assert StructuralView(left.owner.ref("sorts", left.read().kind.value.inputs[0].sort)) == StructuralView( + Right.__egg_sort__ + ) + assert StructuralView(right.owner.ref("sorts", right.read().kind.value.inputs[0].sort)) == StructuralView( + Left.__egg_sort__ + ) + + +def test_constructor_scope_publishes_once_across_authoring_threads(): + class Num(Expr): + def __init__(self, value: i64) -> None: ... + + calls = [] + prepare = Num.__egg_context__.prepare + + def delayed_prepare(cls): + calls.append(cls) + sleep(0.02) # Hold first compilation open while other authors enter. + return prepare(cls) + + Num.__egg_context__.prepare = delayed_prepare + with ThreadPoolExecutor(max_workers=4) as executor: + expressions = list(executor.map(Num, range(8))) + assert calls == [Num] + assert [get_callable_args(expression, Num)[0].value for expression in expressions] == list(range(8)) + + +@pytest.mark.parametrize( + "options", + [ + {"cost": 1}, + {"merge": lambda old, new: old}, + {"mutates_self": True}, + {"unextractable": True}, + {"subsume": True}, + {"reverse_args": True}, + {"preserve": True}, + ], +) +def test_unmigrated_constructor_options_never_publish_members(options): + class Invalid(Expr): + @method(**options) + def __init__(self) -> None: ... + + with pytest.raises(NotImplementedError, match="options"): + Invalid() + assert Invalid.__egg_context__.catalog is None + assert not Invalid.__egg_attr_cache__ + + +@pytest.mark.parametrize("options", [{"cost": 1}, {}]) +def test_wrapped_ignored_name_rejects_unsupported_members_before_publication(options): + class Num(Expr): + def __init__(self) -> None: ... + + @method(**options) + def __radd__(self, other: i64) -> i64: ... + + with pytest.raises(NotImplementedError, match="__radd__"): + Num() + assert Num.__egg_context__.catalog is None + assert not Num.__egg_attr_cache__ + + +def test_preserved_wrapped_ignored_name_remains_a_local_hook(): + class Num(Expr): + def __init__(self) -> None: ... + + @method(preserve=True) + def __radd__(self, other: int) -> int: + return other + 1 + + expression = Num() + assert Num.__radd__(expression, 2) == 3 + assert {name for _, name in Num.__egg_context__.catalog.members} == {"__init__"} + + +def test_unsupported_class_members_and_parameter_kinds_are_not_ignored(): + class HasMethod(Expr): + def __init__(self) -> None: ... + def other(self) -> i64: ... + + class KeywordOnly(Expr): + def __init__(self, *, value: i64) -> None: ... + + for cls in (HasMethod, KeywordOnly): + with pytest.raises(NotImplementedError): + cls() + assert cls.__egg_context__.catalog is None + assert not cls.__egg_attr_cache__ + + +def test_user_sort_cannot_borrow_builtin_codec_or_members_by_name(): + class Shadow(Expr, egg_sort="i64"): + def __init__(self, value: i64) -> None: ... + + expression = Shadow(3) + assert expression.__egg_ref__.read().kind.field == "call" + assert get_callable_args(expression, Shadow)[0].value == 3 + with pytest.raises(AttributeError, match="member"): + _ = Shadow.__add__ + + +def test_unresolved_constructor_annotation_does_not_publish_members(): + class Missing(Expr): + def __init__(self, value: "NotDefinedHere") -> None: ... # noqa: F821 -- intentional unresolved annotation + + with pytest.raises(NameError, match="NotDefinedHere"): + Missing(1) + assert Missing.__egg_context__.catalog is None + assert not Missing.__egg_attr_cache__ + + +def test_user_default_expansion_has_fresh_union_identity_and_retains_sharing(): + class Num(Expr): + def __init__(self, value: i64) -> None: ... + + child = Num(1) + builder = Builder() + member = builder.import_ref(child.__egg_ref__) + root = builder.add( + "nodes", + pb.Node( + sort_id=builder.import_ref(Num.__egg_sort__), + kind=Oneof[Literal["union"], pb.Union]("union", pb.Union(members=[member, member])), + ), + ) + default = RuntimeExpr(builder.publish().ref("nodes", root)) + + class Box(Expr): + def __init__(self, value: Num = default) -> None: ... + + first = get_callable_args(Box(), Box)[0] + second = get_callable_args(Box(), Box)[0] + assert first.__egg_ref__ != second.__egg_ref__ != default.__egg_ref__ + assert expr_parts(first) != expr_parts(second) + for expression in (first, second): + members = expression.__egg_ref__.read().kind.value.members + assert len(members) == 2 + assert members[0] == members[1] + + +def test_reconstructed_constructor_uses_changed_protobuf_parameter_not_python_signature(): + class Num(Expr): + def __init__(self, value: i64) -> None: ... + + program = pack([Num(1).__egg_ref__]).program + index, constructor = next( + (index, declaration) + for index, declaration in enumerate(program.declarations) + if declaration.kind.field == "constructor" + ) + constructor.bindings.python.views[0].params[0].name = "renamed" + owner = Builder.from_program(program).publish() + function = RuntimeFunction(owner.ref("declarations", index), 0, Num) + result = function(renamed=8) + assert get_callable_args(result)[0].value == 8 + with pytest.raises(TypeError): + function(value=8) + + +def test_generic_inference_uses_binder_positions_and_is_transactional(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 0))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 1))) + builder.add( + "sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Pair", args=[0, 1]))) + ) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="String")))) + builder.add( + "sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Pair", args=[3, 4]))) + ) + builder.add( + "sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Pair", args=[4, 3]))) + ) + owner = builder.publish() + bindings: dict[int, Ref] = {} + # GenericSignature labels may repeat: indices, not labels, identify binders. + signature = pb.GenericSignature(type_params=["T", "T"], inputs=[pb.Arg(sort=2)], output=0) + infer_sort(owner.ref("sorts", 2), owner.ref("sorts", 5), len(signature.type_params), bindings) + assert bindings == {0: owner.ref("sorts", 3), 1: owner.ref("sorts", 4)} + assert StructuralView(substitute_sort(owner.ref("sorts", 2), bindings)) == StructuralView(owner.ref("sorts", 5)) + before = bindings.copy() + with pytest.raises(TypeConstraintError): + infer_sort(owner.ref("sorts", 2), owner.ref("sorts", 6), len(signature.type_params), bindings) + assert bindings == before + independent: dict[int, Ref] = {} + infer_sort(owner.ref("sorts", 2), owner.ref("sorts", 6), len(signature.type_params), independent) + assert independent == {0: owner.ref("sorts", 4), 1: owner.ref("sorts", 3)} + + +def test_sort_inference_rejects_missing_binders_arity_and_open_actuals(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 0))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "Box"))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["func"], pb.FuncSort]("func", pb.FuncSort(params=[1], result=1)))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["func"], pb.FuncSort]("func", pb.FuncSort(result=1)))) + owner = builder.publish() + with pytest.raises(TypeConstraintError, match="binder"): + infer_sort(owner.ref("sorts", 0), owner.ref("sorts", 1), 0, {}) + with pytest.raises(TypeConstraintError, match="closed"): + infer_sort(owner.ref("sorts", 0), owner.ref("sorts", 0), 1, {}) + with pytest.raises(TypeConstraintError, match="arity"): + infer_sort(owner.ref("sorts", 2), owner.ref("sorts", 3), 0, {}) + with pytest.raises(TypeConstraintError, match="Unresolved"): + substitute_sort(owner.ref("sorts", 0), {}) + + +def test_sort_inference_and_substitution_are_iterative(): + builder = Builder() + builder.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 0))) + builder.add("sorts", pb.Sort(kind=Oneof[Literal["eq"], str]("eq", "Leaf"))) + pattern, concrete = 0, 1 + for _ in range(10_000): + pattern = builder.add( + "sorts", + pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Vec", args=[pattern]))), + ) + concrete = builder.add( + "sorts", + pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Vec", args=[concrete]))), + ) + owner = builder.publish() + bindings: dict[int, Ref] = {} + infer_sort(owner.ref("sorts", pattern), owner.ref("sorts", concrete), 1, bindings) + result = substitute_sort(owner.ref("sorts", pattern), bindings) + assert StructuralView(result) == StructuralView(owner.ref("sorts", concrete)) + + +def test_call_authoring_binds_keywords_defaults_and_core_input_order(): + builder = Builder() + sort = builder.add("sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64")))) + default = builder.add( + "nodes", + pb.Node( + sort_id=sort, + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", pb.PrimitiveValue(value=Oneof[Literal["i64"], int]("i64", 4)) + ), + ), + ) + declaration = builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["host_primitive"], pb.HostPrimitive]( + "host_primitive", + pb.HostPrimitive( + name="sample.subtract", + typing=Oneof[Literal["signature"], pb.GenericSignature]( + "signature", pb.GenericSignature(inputs=[pb.Arg(sort=sort), pb.Arg(sort=sort)], output=sort) + ), + ), + ), + bindings=pb.CallableBindings( + python=pb.PythonBindings( + views=[ + pb.PythonCallable( + kind=pb.PythonCallKind.FUNCTION, + path=["sample", "subtract_reversed"], + params=[ + pb.PythonParameter(core_input=1, name="right"), + pb.PythonParameter(core_input=0, name="left", default_expr=default), + ], + ) + ] + ) + ), + ), + ) + owner = builder.publish() + + def convert(expected: Ref, value: object) -> Ref: + assert isinstance(value, int) + literal = Builder() + literal.add( + "nodes", + pb.Node( + sort_id=literal.import_ref(expected), + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", pb.PrimitiveValue(value=Oneof[Literal["i64"], int]("i64", value)) + ), + ), + ) + return literal.publish().ref("nodes", 0) + + root, mutated = _author_call( + owner.ref("declarations", declaration), 0, (), {"right": 9}, node_of=lambda _: None, convert=convert + ) + assert mutated is None + node = root.read() + assert node.kind.value.func == "sample.subtract" + args = [root.owner.ref("nodes", index) for index in node.kind.value.args] + assert [arg.read().kind.value.value.value for arg in args] == [4, 9] + assert args[0] != owner.ref("nodes", default) + with pytest.raises(TypeError, match="unexpected keyword"): + _author_call( + owner.ref("declarations", declaration), + 0, + (), + {"right": 9, "bogus": 9}, + node_of=lambda _: None, + convert=convert, + ) + + +def test_generic_empty_call_requires_result_owner_and_keeps_signature_unchanged(): + builder = Builder() + variable = builder.add("sorts", pb.Sort(kind=Oneof[Literal["var"], int]("var", 0))) + pattern = builder.add( + "sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Vec", args=[variable]))) + ) + scalar = builder.add( + "sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64"))) + ) + concrete = builder.add( + "sorts", pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="Vec", args=[scalar]))) + ) + declaration = builder.add( + "declarations", + pb.Declaration( + kind=Oneof[Literal["host_primitive"], pb.HostPrimitive]( + "host_primitive", + pb.HostPrimitive( + name="sample.vec.of", + typing=Oneof[Literal["signature"], pb.GenericSignature]( + "signature", + pb.GenericSignature(type_params=["T"], output=pattern, varargs=pb.Arg(sort=variable)), + ), + ), + ), + bindings=pb.CallableBindings( + python=pb.PythonBindings( + views=[ + pb.PythonCallable( + kind=pb.PythonCallKind.INITIALIZER, + owner=pb.BindingOwner(kind=Oneof[Literal["sort"], int]("sort", pattern)), + params=[pb.PythonParameter(core_input=0, name="args")], + ) + ] + ) + ), + ), + ) + owner = builder.publish() + ref = owner.ref("declarations", declaration) + before = ref.read().to_binary() + + def forbidden(*args): + pytest.fail("An empty constructor must not convert an argument") + + root, mutated = _author_call( + ref, 0, (), {}, node_of=forbidden, convert=forbidden, owner=owner.ref("sorts", concrete) + ) + assert mutated is None + assert list(root.read().kind.value.args) == [] + assert StructuralView(root.owner.ref("sorts", root.read().sort_id)) == StructuralView(owner.ref("sorts", concrete)) + assert ref.read().to_binary() == before + with pytest.raises(TypeConstraintError, match="Unresolved"): + _author_call(ref, 0, (), {}, node_of=forbidden, convert=forbidden) diff --git a/python/tests/test_protobuf_bindings.py b/python/tests/test_protobuf_bindings.py new file mode 100644 index 00000000..9f7cebb7 --- /dev/null +++ b/python/tests/test_protobuf_bindings.py @@ -0,0 +1,62 @@ +"""The native boundary accepts actual schema bytes, never Python AST objects.""" + +from typing import Literal + +import pytest +from egglog_proto.egglog.v1 import egglog_pb as pb +from protobuf import Oneof + +from egglog import bindings + + +def test_binary_engine_lifecycle_and_scalar_extract(): + engine = bindings._ProtoEngine() + sort = pb.Sort(kind=Oneof[Literal["family"], pb.HostSort]("family", pb.HostSort(name="i64"))) + created = pb.CreateEGraphResponse.from_binary( + engine.create(pb.CreateEGraphRequest(sorts=[sort], options=pb.EGraphOptions(cost_sort=0)).to_binary()), + ) + program = pb.Program( + ir_version=1, + sorts=[sort], + nodes=[ + pb.Node( + sort_id=0, + kind=Oneof[Literal["primitive_value"], pb.PrimitiveValue]( + "primitive_value", pb.PrimitiveValue(value=Oneof[Literal["i64"], int]("i64", 7)) + ), + ) + ], + commands=[ + pb.Command( + kind=Oneof[Literal["extract"], pb.Extract]( + "extract", pb.Extract(roots=[0], variants=1, extractor=pb.Extractor.TREE) + ) + ) + ], + ) + response = pb.RunProgramResponse.from_binary( + engine.run(pb.RunProgramRequest(egraph_id=created.egraph_id, program=program, profile=True).to_binary()), + ) + assert response.error is None + output = response.outputs[0].kind + assert output is not None + assert output.field == "extraction" + root = response.nodes[output.value.roots[0].variants[0].term] + assert root.kind is not None + assert root.kind.field == "primitive_value" + assert root.kind.value.value == Oneof("i64", 7) + + cloned = pb.CloneEGraphResponse.from_binary( + engine.clone_graph(pb.CloneEGraphRequest(egraph_id=created.egraph_id).to_binary()), + ) + assert cloned.egraph_id != created.egraph_id + for handle in (created.egraph_id, cloned.egraph_id): + pb.DestroyEGraphResponse.from_binary(engine.destroy(pb.DestroyEGraphRequest(egraph_id=handle).to_binary())) + + +def test_binary_engine_rejects_non_bytes_and_malformed_payloads(): + engine = bindings._ProtoEngine() + with pytest.raises(TypeError): + engine.create(pb.CreateEGraphRequest()) # type: ignore[arg-type] + with pytest.raises(ValueError, match=r".+"): + engine.create(b"\xff") diff --git a/python/tests/test_public_proto_runtime.py b/python/tests/test_public_proto_runtime.py new file mode 100644 index 00000000..6f03990b --- /dev/null +++ b/python/tests/test_public_proto_runtime.py @@ -0,0 +1,173 @@ +"""Normal Python syntax over the generated-record runtime, without an engine.""" + +from copy import copy + +import pytest + +from egglog import EGraph, Expr, ExprValueError, Fact, _program, bindings, eq, expr_parts, f64, i64, method +from egglog.deconstruct import get_callable_args, get_callable_fn, get_literal_value +from egglog.runtime import RuntimeExpr + + +def test_scalar_local_value_and_presentation() -> None: + assert i64(3).value == 3 + assert f64(1.25).value == 1.25 + assert str(i64(1) + 2) == "i64(1) + 2" + assert repr(f64(1.0) + 2.0) == "f64(1.0) + 2.0" + assert get_literal_value(i64(1) + 2) is None + with pytest.raises(ExprValueError, match="Cannot get Python value"): + (i64(1) + 2).value + with pytest.raises(TypeError, match="Expected Expression"): + get_literal_value(1) + + +def test_scalar_call_children_are_retained_records() -> None: + leaf = i64(1) + expression = leaf + 2 + assert isinstance(expression, RuntimeExpr) + assert copy(expression).__egg_ref__ == expression.__egg_ref__ + children = get_callable_args(expression) + assert children is not None + assert children[0].__egg_ref__ == leaf.__egg_ref__ + assert tuple(child.value for child in children) == (1, 2) + assert get_callable_fn(expression) == i64.__add__ + assert tuple(child.__egg_ref__ for child in get_callable_args(expression, i64.__add__)) == tuple( + child.__egg_ref__ for child in children + ) + assert str(i64) == "i64" + + +def test_authoring_never_creates_native_engine(monkeypatch: pytest.MonkeyPatch) -> None: + def forbidden(*args: object, **kwargs: object) -> None: + pytest.fail("Authoring accessed the native engine") + + monkeypatch.setattr(bindings, "EGraph", forbidden) + monkeypatch.setattr(bindings, "_ProtoEngine", forbidden, raising=False) + assert (i64(2) + 3).__egg_ref__.read().kind.value.func == "egglog.core.i64.add" + + +def test_structural_inspection_and_equality_fact() -> None: + assert expr_parts(i64(1) + 2) == expr_parts(i64(1) + 2) + assert hash(expr_parts(i64(1) + 2)) == hash(expr_parts(i64(1) + 2)) + assert bool(eq(i64(1) + 2).to(i64(1) + 2)) + assert not bool(eq(i64(1)).to(i64(2))) + assert str(eq(i64(1)).to(i64(2))) == "eq(i64(1)).to(i64(2))" + assert expr_parts(f64(0.0)) == expr_parts(f64(-0.0)) + + +def test_normal_scalar_bytes_boundary() -> None: + graph = EGraph() + expression = i64(1) + 2 + assert graph.extract(expression).value == 3 + assert graph.extract(f64(1.0) + 2.0).value == 3.0 + graph.check(eq(expression).to(i64(3))) + + +def test_user_constructor_public_bytes() -> None: + class Num(Expr): + def __init__(self, value: i64) -> None: ... + + @method(preserve=True) + @property + def value(self) -> int: + args = get_callable_args(self, Num) + assert args is not None + return args[0].value + + expression = Num(i64(1) + 2) + assert str(expression) == "Num(i64(1) + 2)" + with pytest.raises(ExprValueError): + _ = expression.value + graph = EGraph() + result = graph.extract(expression) + assert isinstance(result, Num) + assert get_callable_fn(result) == Num + assert get_callable_args(result, Num)[0].value == result.value == 3 + assert str(result) == "Num(3)" + graph.check(eq(result).to(Num(3))) + assert graph.extract(Num(value=3)).value == 3 + + +def test_user_constructor_exact_names_and_retained_local_hooks() -> None: + def make(): + class Num(Expr, egg_sort="sort with spaces"): + @method(egg_fn="constructor with spaces") + def __init__(self, value: i64 = i64(7)) -> None: ... + + @method(preserve=True) + def tag(self) -> str: + return "retained local method" + + return Num() + + graph = EGraph() + expression = make() + result = graph.extract(expression) + graph.check(eq(result).to(expression)) + graph._state.destroy() + assert result.tag() == expression.tag() == "retained local method" + assert get_callable_args(result)[0].value == 7 + + +def test_scalar_equality_dispatch_respects_operand_sorts() -> None: + assert (i64(1) == f64(1.0)) is False + assert (f64(1.0) == i64(1)) is False + equal = i64(1) == i64(1) + different = f64(1.0) == f64(2.0) + assert isinstance(equal, Fact) + assert isinstance(different, Fact) + assert bool(equal) + assert not bool(different) + + +def test_automatic_threads_reports_actual_count() -> None: + # Existing public behavior; unresolved until byte-level configuration exists. + graph = EGraph(num_threads=0) + assert graph.num_threads() >= 1 + + +def test_automatic_threads_after_update_reports_actual_count() -> None: + graph = EGraph(num_threads=1) + graph.set_num_threads(0) + assert graph.num_threads() >= 1 + + +def test_thread_update_and_context_restore_actual_count() -> None: + graph = EGraph(num_threads=2) + assert graph.num_threads() == 2 + with graph: + graph.set_num_threads(1) + assert graph.num_threads() == 1 + assert graph.num_threads() == 2 + graph.set_num_threads(1) + assert graph.num_threads() == 1 + + +def test_ten_thousand_ordinary_calls_are_incremental(monkeypatch: pytest.MonkeyPatch) -> None: + copied = 0 + visited = 0 + original_copy = _program._copy_record + original_parts = _program._structural_parts + + def counted_copy(*args, **kwargs): + nonlocal copied + copied += 1 + return original_copy(*args, **kwargs) + + def counted_parts(*args, **kwargs): + nonlocal visited + visited += 1 + return original_parts(*args, **kwargs) + + monkeypatch.setattr(_program, "_copy_record", counted_copy) + monkeypatch.setattr(_program, "_structural_parts", counted_parts) + expression = i64(0) + for _ in range(10_000): + expression = expression + 1 + hash(expression) + # Fixed per-call limits, not wall-clock timing: retained prefixes are not + # rescanned during authoring/hash, independent of chain depth. + assert copied < 50 * 10_000 + assert visited < 20 * 10_000 + packed = _program.pack([expression.__egg_ref__]) + assert len(packed.program.nodes) == 20_001 diff --git a/python/tests/test_runtime.py b/python/tests/test_runtime.py index 8f2bd983..8f33ba9e 100644 --- a/python/tests/test_runtime.py +++ b/python/tests/test_runtime.py @@ -1,9 +1,13 @@ from __future__ import annotations import doctest +import math +import struct +from copy import copy import pytest +from egglog import expr_parts, f64 from egglog.declarations import * from egglog.exp import array_api from egglog.runtime import * @@ -11,6 +15,27 @@ from egglog.type_constraint_solver import * +def test_public_f64_structural_identity_and_signed_zero(): + source = float("nan") + left, same_source, fresh = f64(source), f64(source), f64(float("nan")) + assert expr_parts(left) == expr_parts(same_source) + assert bool(left == same_source) + assert hash(left) == hash(same_source) + assert expr_parts(left) != expr_parts(fresh) + assert not bool(left == fresh) + assert expr_parts(left) == expr_parts(copy(left)) + assert hash(left) == hash(copy(left)) + assert left.value is source + assert math.isnan(fresh.value) + assert struct.pack("!d", left.value) == struct.pack("!d", fresh.value) + + positive, negative = f64(0.0), f64(-0.0) + assert expr_parts(positive) == expr_parts(negative) + assert bool(positive == negative) + assert hash(positive) == hash(negative) + assert struct.pack("!d", positive.value) != struct.pack("!d", negative.value) + + def test_type_str(): decls = Declarations( _classes={ diff --git a/scripts/generate_builtin_catalog.py b/scripts/generate_builtin_catalog.py new file mode 100644 index 00000000..9070aa1b --- /dev/null +++ b/scripts/generate_builtin_catalog.py @@ -0,0 +1,51 @@ +""" +Package an actual native catalog export without importing the Python frontend. + +First run egglog-experimental's export_builtin_catalog example from the pinned +native checkout. This command consumes that binary snapshot; it never discovers +signatures through an EGraph during import or authoring. --check is a byte-for- +byte reproducibility gate, so changed native records cannot be silently hidden. +""" + +from __future__ import annotations + +import argparse +import hashlib +from pathlib import Path + +from egglog_proto.egglog.v1 import egglog_pb as pb + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("export", type=Path) + parser.add_argument( + "--output", type=Path, default=Path(__file__).resolve().parents[1] / "python/egglog/_builtin_catalog.pb" + ) + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + encoded = args.export.read_bytes() + program = pb.Program.from_binary(encoded) + if program.ir_version != 1 or program.commands or program.rules or program.rulesets: + parser.error("Expected a version-1 declaration-only native catalog") + names: set[tuple[str, str]] = set() + for declaration in program.declarations: + kind = declaration.kind + if kind is None or kind.field not in {"host_sort_family", "host_primitive"}: + parser.error("Native catalog may contain only host family/primitive declarations") + namespace = "sort" if kind.field == "host_sort_family" else "callable" + key = namespace, kind.value.name + if not key[1] or key in names: + parser.error(f"Duplicate or empty native definition: {key}") + names.add(key) + if args.check: + if args.output.read_bytes() != encoded: + parser.error("Packaged catalog differs from the native export; regenerate it") + else: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_bytes(encoded) + print(f"{len(names)} native definitions; SHA-256 {hashlib.sha256(encoded).hexdigest()}") + + +if __name__ == "__main__": + main() diff --git a/src/lib.rs b/src/lib.rs index e80e1b14..91df4112 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ mod egraph; mod error; mod extract; mod freeze; +mod protobuf; mod py_object_sort; mod serialize; mod termdag; @@ -36,6 +37,7 @@ fn bindings(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/src/protobuf.rs b/src/protobuf.rs new file mode 100644 index 00000000..444007ee --- /dev/null +++ b/src/protobuf.rs @@ -0,0 +1,77 @@ +//! Private bytes-only access to the shared adapter. No Python syntax is lowered here. + +use egglog_experimental::protobuf::{Engine, TransportError}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +#[pyclass(name = "_ProtoEngine")] +#[derive(Default)] +pub struct ProtoEngine { + engine: Engine, +} + +impl ProtoEngine { + fn request<'py>( + &mut self, + py: Python<'py>, + request: &Bound<'py, PyBytes>, + operation: fn(&mut Engine, &[u8]) -> Result, TransportError>, + ) -> PyResult> { + // Own the input while the GIL is released; only encoded bytes cross + // this boundary in either direction, including structured run errors. + let bytes = request.as_bytes().to_vec(); + let response = py + .detach(|| operation(&mut self.engine, &bytes)) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok(PyBytes::new(py, &response)) + } +} + +#[pymethods] +impl ProtoEngine { + #[new] + fn new() -> Self { + Self::default() + } + + fn create<'py>( + &mut self, + py: Python<'py>, + request: &Bound<'py, PyBytes>, + ) -> PyResult> { + self.request(py, request, Engine::create) + } + + fn run<'py>( + &mut self, + py: Python<'py>, + request: &Bound<'py, PyBytes>, + ) -> PyResult> { + self.request(py, request, Engine::run) + } + + fn clone_graph<'py>( + &mut self, + py: Python<'py>, + request: &Bound<'py, PyBytes>, + ) -> PyResult> { + self.request(py, request, Engine::clone_egraph) + } + + fn destroy<'py>( + &mut self, + py: Python<'py>, + request: &Bound<'py, PyBytes>, + ) -> PyResult> { + self.request(py, request, Engine::destroy) + } + + fn configure_resources<'py>( + &mut self, + py: Python<'py>, + request: &Bound<'py, PyBytes>, + ) -> PyResult> { + self.request(py, request, Engine::configure_resources) + } +} diff --git a/uv.lock b/uv.lock index db48fb71..363933ac 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", @@ -750,6 +750,7 @@ dependencies = [ { name = "anywidget" }, { name = "black" }, { name = "cloudpickle" }, + { name = "egglog-proto" }, { name = "graphviz" }, { name = "opentelemetry-api" }, { name = "typing-extensions" }, @@ -847,6 +848,7 @@ requires-dist = [ { name = "egglog", extras = ["array"], marker = "extra == 'docs'" }, { name = "egglog", extras = ["array"], marker = "extra == 'test'" }, { name = "egglog", extras = ["docs", "test"], marker = "extra == 'dev'" }, + { name = "egglog-proto", git = "https://github.com/egraphs-good/egglog-experimental.git?subdirectory=gen%2Fpython&rev=6e0dbc97fbe9f3868c4662b9a6d73b5d01913c9d" }, { name = "graphviz" }, { name = "ipykernel", marker = "extra == 'dev'" }, { name = "jupyterlab", marker = "extra == 'dev'" }, @@ -886,6 +888,14 @@ dev = [ { name = "sympy", specifier = ">=1.14.0" }, ] +[[package]] +name = "egglog-proto" +version = "0.1.0" +source = { git = "https://github.com/egraphs-good/egglog-experimental.git?subdirectory=gen%2Fpython&rev=6e0dbc97fbe9f3868c4662b9a6d73b5d01913c9d#6e0dbc97fbe9f3868c4662b9a6d73b5d01913c9d" } +dependencies = [ + { name = "protobuf-py" }, +] + [[package]] name = "execnet" version = "2.1.2" @@ -2767,6 +2777,54 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e4/04/d52c7016b04b6c5108f26691f9d33ec82a9b65d041f1a9c771137693d618/protobuf-7.36.2-py3-none-any.whl", hash = "sha256:bdb3a345d48db958e6ce1f18e508beb0cc981d64f24088427549c866cd039f1e", size = 179806, upload-time = "2026-09-17T20:07:58.211Z" }, ] +[[package]] +name = "protobuf-py" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf-py-ext", marker = "(platform_machine == 'arm64' and platform_python_implementation == 'CPython' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (platform_machine == 'AMD64' and platform_python_implementation == 'CPython' and sys_platform == 'win32') or (platform_machine == 'ARM64' and platform_python_implementation == 'CPython' and sys_platform == 'win32')" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e0/df/744b304857f9727dad5938e37c88b834c7874581df4bdd36546e0718c5ba/protobuf_py-0.5.0.tar.gz", hash = "sha256:a7be567fbcade37f7ead5c7ebdcda7d5d833afb0e68a569a7a80837f2f5defc1", size = 160914, upload-time = "2026-09-15T06:53:37.924Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/4d/f574b8fece5fbce107bdaa2d4f81743e810275d584fda83456475ce3a836/protobuf_py-0.5.0-py3-none-any.whl", hash = "sha256:1582e21c920a9c6932d5b9398aa55635315ecf7e05ab6eb09dcd3f9aba8ff7ac", size = 214220, upload-time = "2026-09-15T06:52:52.734Z" }, +] + +[[package]] +name = "protobuf-py-ext" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/00/e5742687f394e55b96a8f69ae6c3ef1ea518bdb85970660ef582cc8da89d/protobuf_py_ext-0.5.0.tar.gz", hash = "sha256:99c3e47e145e6ccc758ac473b967240e4cd5d4ddbb12352da2fea44bb09b75a6", size = 57832, upload-time = "2026-09-15T06:53:38.88Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/a5/08ee1549f93e2a15d43af841181d9b9f54049a6d799f0647019ecbd18f73/protobuf_py_ext-0.5.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:c263aa5594a5be05749d2fc8398685ff97435cbe0ea98076cc1dc114c69d6994", size = 469453, upload-time = "2026-09-15T06:53:01.222Z" }, + { url = "https://files.pythonhosted.org/packages/c2/de/37275a9c1abf9470760ac12db908065b1e3c1eae5b16ac0f32cd5c0dce78/protobuf_py_ext-0.5.0-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:65801a66e68bd00b5ab48cfd277ebc5ff974381e0e9e1812f7c77efbcacad1a9", size = 478158, upload-time = "2026-09-15T06:53:02.601Z" }, + { url = "https://files.pythonhosted.org/packages/21/9a/acbe04fdd223de6364a6e6252a316358a74f5f4c8e0ca13b145e15b3ac88/protobuf_py_ext-0.5.0-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4f0b868704714af24981b68095774edd42930490a2e09678a9e1d88b892a1100", size = 503450, upload-time = "2026-09-15T06:53:03.81Z" }, + { url = "https://files.pythonhosted.org/packages/0b/42/f4ee62c87fb0488bd1bcc88452665201bf2ae3d2491618dce14755fc152f/protobuf_py_ext-0.5.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:2fe9c58940372c24d31d51d446dde4cfb7f7d9ad36ce4c07f59d44373126145f", size = 657275, upload-time = "2026-09-15T06:53:05.041Z" }, + { url = "https://files.pythonhosted.org/packages/79/db/79e3a0c975457cff14ca2579c60209946fe625b649e876fe285ca2754d9b/protobuf_py_ext-0.5.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:22143055e96570a08b5920c3637640450fbecadaeda416af4b3836937616ecf1", size = 716043, upload-time = "2026-09-15T06:53:06.36Z" }, + { url = "https://files.pythonhosted.org/packages/5f/3b/d2bc8047f643db12f04b2e680a7f9dc5433e96e3b07840148e8eebfc5a66/protobuf_py_ext-0.5.0-cp311-abi3-win_amd64.whl", hash = "sha256:b59e4daffcbfc0499a91779a546907a509c7058fdf9784d16e727a1dde729938", size = 440050, upload-time = "2026-09-15T06:53:07.706Z" }, + { url = "https://files.pythonhosted.org/packages/0d/0f/e49bd745cdc5f8b2d171817d87035d64be0b75326db7c0c1e644778fa882/protobuf_py_ext-0.5.0-cp311-abi3-win_arm64.whl", hash = "sha256:8e1077f106a60cac056cc7591ca8ba70dcb34e65139a992bfa76bc2c5f85ed50", size = 419729, upload-time = "2026-09-15T06:53:08.918Z" }, + { url = "https://files.pythonhosted.org/packages/a2/49/5a0ea1bb76d5d7be41ccdfc58422e4cd172d856646dd8903cd6034791e48/protobuf_py_ext-0.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c7c584e99362c3004ac4f584108bc38cdfb8ee08da84fd57c8eff9abbe8eb0bd", size = 471070, upload-time = "2026-09-15T06:53:10.059Z" }, + { url = "https://files.pythonhosted.org/packages/e1/c8/2f656bcab07519d7d03901aaea8a572ca0decdbdc0aea44a88c4496f0b9e/protobuf_py_ext-0.5.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e39a4b9289de2f73d67a1be2c3f750f98186ed66b9a844f2a47ad4122c2e6f6a", size = 478317, upload-time = "2026-09-15T06:53:11.434Z" }, + { url = "https://files.pythonhosted.org/packages/05/8a/b137405f95761606d384d35dac2c9a1b8fccc0aba4a3588472d8bf434a82/protobuf_py_ext-0.5.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d4b241a72759cbe26759db2a13deab059510d7609a781d125bba22dbbfaf2cf", size = 499921, upload-time = "2026-09-15T06:53:12.699Z" }, + { url = "https://files.pythonhosted.org/packages/d4/64/e5e97d95922f81834b189666c6cfe74c071c7929cfb20f59ab77352e8f88/protobuf_py_ext-0.5.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3b6add28bbdf0b724ccc10a9c8a80549ab0b2af50b808e533277506bdb3c04e6", size = 656768, upload-time = "2026-09-15T06:53:13.934Z" }, + { url = "https://files.pythonhosted.org/packages/22/38/0ababc6093654bac7430cf1f5be2159d5442f676e1853d3e8f417861a3f7/protobuf_py_ext-0.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7b8c155384c8949c0cccc4a3c14b53ee85b5dbf022c4bc0bd9486764c8250e03", size = 712993, upload-time = "2026-09-15T06:53:15.347Z" }, + { url = "https://files.pythonhosted.org/packages/9b/3d/2438ca1835d81a7073ae7af18629f00fbed3bd9cd02b7ad5803a709a8f2f/protobuf_py_ext-0.5.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:5ebd954fa15f2636ba58577c9bb037865bcdaf66e68d8f0587ccfec814d451d2", size = 471346, upload-time = "2026-09-15T06:53:16.567Z" }, + { url = "https://files.pythonhosted.org/packages/ac/3f/088cb289aa8cf520827b4fb99092d3204f1a2918ae730d45029c08a8553d/protobuf_py_ext-0.5.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6cea81147867a9b29697240b586d7c945398a5629351e070ff740d9a5b9eb148", size = 479009, upload-time = "2026-09-15T06:53:17.865Z" }, + { url = "https://files.pythonhosted.org/packages/37/a3/e5e50258d7f7e006cc5cd4c428981ae0d69cacec59093eb2a865d97c2b16/protobuf_py_ext-0.5.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:12015dd711f70e7c8ac091376781b1924aeaaed9e9bee39fdf9130cb45b137e9", size = 500623, upload-time = "2026-09-15T06:53:19.218Z" }, + { url = "https://files.pythonhosted.org/packages/12/72/974f9df89b5b2ea095374a590f4b85aeaf80f4fa5f504012b1cf93c63565/protobuf_py_ext-0.5.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e1fef6a711ebcb9a413fe10d342b278fa30ceb68b2bf7bfcf0a55523f091dfb4", size = 657155, upload-time = "2026-09-15T06:53:20.631Z" }, + { url = "https://files.pythonhosted.org/packages/02/27/c1709e454a22f589b35578d77fd68db707e1109f896ec2245af295612431/protobuf_py_ext-0.5.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e5f11aa121af8f6bede87d8717abbcc0c5ad9b68902863489de416481df31e72", size = 713679, upload-time = "2026-09-15T06:53:21.966Z" }, + { url = "https://files.pythonhosted.org/packages/de/8f/9759b7890488ea54b88d751ef4fe6dfac8c230e5f54b1230e1819c643928/protobuf_py_ext-0.5.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9684428ad5ea608f50144ef9ee6d918f4ce15a7222093b38eac4a17fb51efd1f", size = 469469, upload-time = "2026-09-15T06:53:23.185Z" }, + { url = "https://files.pythonhosted.org/packages/92/0f/6edef5e307e2e73653aa8c252d123076bc14d098d8a603e5cbcbbd9afba6/protobuf_py_ext-0.5.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ee68dce9041a1fbec9648ef06dba465a033e69cbc89a27be01b4f5c17285c03e", size = 477129, upload-time = "2026-09-15T06:53:24.455Z" }, + { url = "https://files.pythonhosted.org/packages/df/8f/2ac90881b7cea63d3ba86f98479c3025996d37480638d16b2f6a3d662e40/protobuf_py_ext-0.5.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b90b12abf51a0ba4aeeed2743f3633760a6707788726baae70d53bfe41559f6", size = 499183, upload-time = "2026-09-15T06:53:25.688Z" }, + { url = "https://files.pythonhosted.org/packages/b9/f4/10769d23e8c664a92ab9e7274669e6542c0cdaa86e9a02f1281640bade4f/protobuf_py_ext-0.5.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:75ee44f40633c07741320ad72f9c8736be5021ba24f96844337bba009b74a10d", size = 655283, upload-time = "2026-09-15T06:53:26.926Z" }, + { url = "https://files.pythonhosted.org/packages/38/40/10681c188198c00fd48799c40827e7fafd6deb19dc59d2b48791af9650d1/protobuf_py_ext-0.5.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f54a76acaf78255a97e9a33798d6efedec0d5fdaf38d393501be6b368e54d6e7", size = 712563, upload-time = "2026-09-15T06:53:28.322Z" }, + { url = "https://files.pythonhosted.org/packages/a8/25/c590e1994ceec4aa9bc0b71fa0c4063265f3168e87a9088ff01038c58c24/protobuf_py_ext-0.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ae8a17933f066da2e1d105fe19cec425a63c0065349ad128045f6ad5eec5ee0a", size = 458067, upload-time = "2026-09-15T06:53:29.572Z" }, + { url = "https://files.pythonhosted.org/packages/fe/fd/a8dc89c1269055cd10f93135baf387227a4f915443a45ec1d7d7083d4bcc/protobuf_py_ext-0.5.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:60be9e43b00e4c5ba65dc290e24cd2e4a2473e553b96e81feb22be24c991dcf4", size = 466789, upload-time = "2026-09-15T06:53:30.908Z" }, + { url = "https://files.pythonhosted.org/packages/15/18/e90306c613c61ef3e7011ff6606c7b7fcf97ff0dc176b2c98b58e8545e19/protobuf_py_ext-0.5.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3e82945cde40662610ebbc1e8960d4cfd22df900a3aabdc0f6ba9f80d23b6ae9", size = 489446, upload-time = "2026-09-15T06:53:32.343Z" }, + { url = "https://files.pythonhosted.org/packages/e2/9e/a5dbbf257dff0fe4f141080f8427f793018f44f78b81fdbb88aacee796b8/protobuf_py_ext-0.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:9c1bd5a6137b6f1295d8a3bcb9711cbf4e959743611325b8b1254a2d51e64d7a", size = 646164, upload-time = "2026-09-15T06:53:33.688Z" }, + { url = "https://files.pythonhosted.org/packages/40/65/6046d7641a26ffdb6bb9e0463d0596c341ec156e40c0571b396ad786d4a2/protobuf_py_ext-0.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0fe8f84fae44c6bfd3ee977ab35893d77d30c169e6502f83769d2e0ee2ec7591", size = 703098, upload-time = "2026-09-15T06:53:35.337Z" }, +] + [[package]] name = "psutil" version = "7.2.2"