Skip to content

Fix lookup of value in cost model #401

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
saulshanabrook merged 2 commits into from
Mar 3, 2026
Merged

Fix lookup of value in cost model #401

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b18d3e7
Fix lookup of value in cost model
saulshanabrook Mar 3, 2026
d38a170
Fix mypy
saulshanabrook Mar 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/changelog.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ _This project uses semantic versioning_
- Improve doctest support, teaching expressions about their `__module__`, `__dir__`, and special methods.
- Surface original Python exceptions from the runtime and tighten pretty-printing of values that cannot be re-parsed to make debugging e-graph executions easier.
- Update the bundled Egglog crate, visualizer, and related dev dependencies (including `ipykernel`) to pick up the latest backend fixes.
- Fix lookup of cost model based on value (see [zulip for issue](https://egraphs.zulipchat.com/#narrow/channel/375765-egg.2Fegglog/topic/Cost.20function.3A.20using.20function.20values.20of.20subtrees/near/577062352))

## 11.4.0 (2025-10-02)

Expand Down
2 changes: 2 additions & 0 deletions python/egglog/egraph_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -683,6 +683,8 @@ def _generate_callable_egg_name(self, ref: CallableRef) -> str:
assert_never(ref)

def typed_expr_to_value(self, typed_expr: TypedExprDecl) -> bindings.Value:
if isinstance(typed_expr.expr, ValueDecl):
return typed_expr.expr.value
Comment on lines +686 to +687
Copy link

@coderabbitai coderabbitai bot Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
ruff check python/egglog/egraph_state.py --select F405
rg -n "from \.declarations import" python/egglog/egraph_state.py

Repository: egraphs-good/egglog-python

Length of output: 50383

🏁 Script executed:

#!/bin/bash
# Check lines 686-687 and surrounding context
sed -n '680,695p' python/egglog/egraph_state.py | cat -n

# Check the declarations import specifically
rg -n "^from \.declarations import|^from \.declarations import.*ValueDecl|^from \*" python/egglog/egraph_state.py | head -20

Repository: egraphs-good/egglog-python

Length of output: 888

🏁 Script executed:

#!/bin/bash
# Check exact import line and context
sed -n '15,20p' python/egglog/egraph_state.py | cat -n

# Run Ruff check focused on the region around lines 686-687
ruff check python/egglog/egraph_state.py --select F405 2>&1 | grep -A2 -B2 "686\|687"

Repository: egraphs-good/egglog-python

Length of output: 1176

Add explicit import of ValueDecl to fix Ruff F405.

The new isinstance(typed_expr.expr, ValueDecl) check at line 686 triggers Ruff F405 since ValueDecl is only available through the star import. Add it to the explicit imports to keep the codebase lint-clean.

Suggested patch 
-from .declarations import ConstructorDecl
+from .declarations import ConstructorDecl, ValueDecl
🧰 Tools
🪛 Ruff (0.15.2)

[error] 686-686: ValueDecl may be undefined, or defined from star imports

(F405)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@python/egglog/egraph_state.py` around lines 686 - 687, The isinstance check
uses ValueDecl but ValueDecl is only coming from a star import, causing Ruff
F405; update the module imports to explicitly import ValueDecl (add ValueDecl to
the from ... import (...) list where other AST/decl types are imported) so the
isinstance(typed_expr.expr, ValueDecl) line in egraph_state (around the check on
typed_expr.expr) resolves without a star import warning.

egg_expr = self.typed_expr_to_egg(typed_expr, False)
return self.egraph.eval_expr(egg_expr)[1]

Expand Down
34 changes: 34 additions & 0 deletions python/tests/test_high_level.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1568,3 +1568,37 @@ def __radd__(self, other: object) -> tuple[X, X]: ...

assert X(1) + 10 == (X(1), 10)
assert 10 + X(1) == (X(10), X(1))


def test_custom_cost_model_size():
"""
https://egraphs.zulipchat.com/#narrow/channel/375765-egg.2Fegglog/topic/Cost.20function.3A.20using.20function.20values.20of.20subtrees/near/577062352
"""

class KAT(Expr):
@classmethod
def eps(cls) -> KAT: ...

@classmethod
def emp(cls) -> KAT: ...

def func(self, other: KAT) -> KAT: ...

def size(self) -> i64: ...

eps, emp = KAT.eps(), KAT.emp()

eg = EGraph()
q0 = eg.let("q0", KAT.func(eps, emp))

eg.register(set_(eps.size()).to(i64(1)))
eg.register(set_(emp.size()).to(i64(0)))

def conv_cost(eg, expr, child_costs):
if isinstance(expr, KAT):
args = get_callable_args(expr)
return sum(int(eg.lookup_function_value(cast("KAT", a).size())) for a in args)
Comment on lines +1597 to +1600
Copy link

@coderabbitai coderabbitai bot Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check for linting issues and find get_callable_args definition
ruff check python/tests/test_high_level.py --select ARG001,F405
echo "---"
rg -n "def get_callable_args|from \.deconstruct import get_callable_args|__all__" python/egglog/deconstruct.py python/egglog/__init__.py

Repository: egraphs-good/egglog-python

Length of output: 50383

🏁 Script executed:

#!/bin/bash
# Read the full function and surrounding context
sed -n '1590,1610p' python/tests/test_high_level.py

Repository: egraphs-good/egglog-python

Length of output: 542

🏁 Script executed:

#!/bin/bash
# Find the definition of get_callable_args to check return type
rg -A 10 "^def get_callable_args" python/egglog/

Repository: egraphs-good/egglog-python

Length of output: 1276

Harden conv_cost against None args and silence unused-arg warning.

get_callable_args can return None; iterating that would raise TypeError. Also, child_costs is intentionally unused here.

Suggested patch 
-def conv_cost(eg, expr, child_costs):
+def conv_cost(eg, expr, _child_costs):
     if isinstance(expr, KAT):
-        args = get_callable_args(expr)
+        args = get_callable_args(expr) or ()
         return sum(int(eg.lookup_function_value(cast("KAT", a).size())) for a in args)
🧰 Tools
🪛 Ruff (0.15.2)

[warning] 1597-1597: Unused function argument: child_costs

(ARG001)

[error] 1599-1599: get_callable_args may be undefined, or defined from star imports

(F405)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@python/tests/test_high_level.py` around lines 1597 - 1600, The conv_cost
function must guard against get_callable_args returning None and avoid the
unused-argument warning for child_costs: update conv_cost (handling KAT) to call
get_callable_args(expr), treat a None result as an empty sequence before
iterating (e.g., if args is None: args = ()), and explicitly mark child_costs as
intentionally unused (for example by assigning it to _ or using del child_costs)
so linters stop warning; keep the existing sum over cast("KAT", a).size() for
each arg.


return 2

assert eg.extract(q0, include_cost=True, cost_model=conv_cost) == (KAT.eps().func(KAT.emp()), 1)
Loading