Updates for last talk

Author: saulshanabrook

Cargo.lock

@@ -220,6 +220,8 @@ ignore = [
     "INP001",
     # don't replace lambdas with functions because we need them to defer resolution
     "PLW0108",
+    # unicode vars
+    "PLC2401"
 ]
 select = ["ALL"]
 
@@ -252,16 +254,19 @@ preview = true
 
 [tool.mypy]
 ignore_missing_imports = true
-allow_redefinition = true
+allow_redefinition_new = true
 exclude = ["__snapshots__", "_build", "^conftest.py$"]
 warn_unused_configs = true
 disallow_subclassing_any = true
 check_untyped_defs = true
 warn_redundant_casts = true
 warn_unused_ignores = true
+warn_unreachable = true
 strict_equality = true
 extra_checks = true
 strict_equality_for_none = true
+enable_incomplete_feature = ["TypeForm"]
+local_partial_types = true
 
 [tool.maturin]
 python-source = "python"

pyproject.toml

@@ -66,6 +66,7 @@
     "i64",
     "i64Like",
     "join",
+    "map_best_common_float_scale",
     "map_bigrat_intersect_min",
     "map_bigrat_subtract",
     "map_contains_key_swapped",
@@ -713,6 +714,10 @@ def map_nonconst_nonunit_f64_values(xs: Map[Map[T, V], f64]) -> MultiSet[f64]: .
 def map_divide_all_values_by_f64(factor: f64, xs: Map[T, f64]) -> Map[T, f64]: ...
 
 
+@function(egg_fn="map-best-common-float-scale", builtin=True)
+def map_best_common_float_scale(xs: Map[T, f64]) -> f64: ...
+
+
 @function(egg_fn="map-integer-residual-split-candidate", builtin=True)
 def map_integer_residual_split_candidate(xs: Map[Map[T, BigRat], f64]) -> Pair[Map[T, BigRat], Map[Map[T, BigRat], f64]]: ...
 

python/egglog/builtins.py

@@ -31,7 +31,7 @@
 
 import graphviz
 from opentelemetry import trace
-from typing_extensions import ParamSpec, Unpack
+from typing_extensions import ParamSpec, TypeForm, Unpack
 
 from . import bindings
 from ._tracing import call_with_current_trace
@@ -1318,7 +1318,7 @@ def extract(
             cost = cast("COST", extract_report.cost)
         else:
             if not isinstance(runtime_expr.__egg_typed_expr__.expr, (LetRefDecl, UnboundVarDecl)):
-                self.register(expr)
+                self._register_extract_root(runtime_expr)
             egg_cost_model = _CostModel(cost_model, self).to_bindings_cost_model()
             egg_sort = self._state.type_ref_to_egg(tp)
             extractor = call_with_current_trace(bindings.Extractor, [egg_sort], self._state.egraph, egg_cost_model)
@@ -1554,6 +1554,20 @@ def _register_commands(self, cmds: list[Command]) -> None:
         egg_cmds = [egg_cmd for cmd in cmds if (egg_cmd := self._command_to_egg(cmd)) is not None]
         self._state.run_program(*egg_cmds)
 
+    def _register_extract_root(self, runtime_expr: RuntimeExpr) -> None:
+        """
+        Register the exact extraction root without synthetic let factoring.
+
+        Synthetic lets are a command-size optimization for public registration,
+        but custom-cost extraction immediately evaluates the original root value.
+        Registering a let-factored presentation can leave the custom extractor
+        without a costed parent for that exact value in the direct command API.
+        """
+        self._add_decls(runtime_expr)
+        action_egg = self._state.action_to_egg(ExprActionDecl(runtime_expr.__egg_typed_expr__), expr_to_let=False)
+        if action_egg is not None:
+            self._state.run_program(bindings.ActionCommand(action_egg))
+
     def _command_to_egg(self, cmd: Command) -> bindings._Command | None:
         ruleset_ident = Ident("")
         cmd_decl: CommandDecl
@@ -1835,7 +1849,7 @@ class Schedule(DelayedDeclarations):
     A composition of some rulesets, either composing them sequentially, running them repeatedly, running them till saturation, or running until some facts are met
     """
 
-    # Defer declerations so that we can have rule generators that used not yet defined yet
+    # Defer declarations so that we can have rule generators that used not yet defined yet
     schedule: ScheduleDecl
 
     def __str__(self) -> str:
@@ -2138,7 +2152,7 @@ def rule(*facts: FactLike, ruleset: None = None, name: str | None = None) -> _Ru
     return _RuleBuilder(facts=_fact_likes(facts), name=name, ruleset=ruleset)
 
 
-def var(name: str, bound: type[T], egg_name: str | None = None) -> T:
+def var(name: str, bound: TypeForm[T], egg_name: str | None = None) -> T:
     """Create a new variable with the given name and type."""
     return cast("T", _var(name, bound, egg_name=egg_name))
 

python/egglog/egraph.py

@@ -410,7 +410,7 @@ def ruleset_to_egg(self, ident: Ident) -> None:
     def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command | None:
         match cmd:
             case ActionCommandDecl(action):
-                action_egg = self.action_to_egg(action, expr_to_let=True)
+                action_egg = self.action_to_egg(action)# , expr_to_let=False)
                 if not action_egg:
                     return None
                 return bindings.ActionCommand(action_egg)