tmp

Author: saulshanabrook

python/egglog/builtins.py

@@ -409,7 +409,7 @@ def eval(self) -> dict[T, V]:
     @property
     def value(self) -> dict[T, V]:
         d = {}
-        while args := get_callable_args(self, Map[T, V].insert):
+        while args := get_callable_args(self, Map.insert):  # type: ignore[var-annotated]
             self, k, v = args  # noqa: PLW0642
             d[k] = v
         if get_callable_args(self, Map.empty) is None:
@@ -992,7 +992,7 @@ def __init__(self, *args: T) -> None: ...
     def empty(cls) -> Vec[T]: ...
 
     @method(egg_fn="vec-append")
-    def append(self, *others: Vec[T]) -> Vec[T]: ...
+    def append(self, *others: VecLike[T, T]) -> Vec[T]: ...
 
     @method(egg_fn="vec-push")
     def push(self, value: T) -> Vec[T]: ...
@@ -1033,6 +1033,9 @@ def set(self, index: i64Like, value: T) -> Vec[T]: ...
 
 VecLike: TypeAlias = Vec[T] | tuple[TO, ...] | list[TO]
 
+v = Vec(i64(10))
+v.append([i64(10)])
+
 
 class PyObject(BuiltinExpr, egg_sort="PyObject"):
     @method(preserve=True)

python/egglog/conversion.py

@@ -221,7 +221,6 @@ def resolve_literal(
     arg: object,
     decls: Callable[[], Declarations] = retrieve_conversion_decls,
     tcs: TypeConstraintSolver | None = None,
-    cls_ident: Ident | None = None,
 ) -> RuntimeExpr:
     """
     Try to convert an object to a type, raising a ConvertError if it is not possible.
@@ -240,7 +239,7 @@ def resolve_literal(
         # args first based on the existing type constraint solver
         if tcs:
             try:
-                tp_just = tcs.substitute_typevars(tp, cls_ident)
+                tp_just = tcs.substitute_typevars(tp)
             # If we can't resolve the type var yet, then just assume it is the right value
             except TypeConstraintError:
                 assert isinstance(arg, RuntimeExpr), f"Expected a runtime expression, got {type(arg)}"
@@ -250,7 +249,7 @@ def resolve_literal(
             assert isinstance(arg, RuntimeExpr), f"Expected a runtime expression, got {arg}"
             return arg
     if tcs:
-        tcs.infer_typevars(tp, tp_just, cls_ident)
+        tcs.infer_typevars(tp, tp_just)
     if arg_type == tp_just:
         # If the type is an egg type, it has to be a runtime expr
         assert isinstance(arg, RuntimeExpr)

python/egglog/egraph.py

@@ -328,6 +328,7 @@ def __new__(  # type: ignore[misc]
         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
+        # TODO: How should we normalize unparameterized classes? Should they have args of the typevars?
         runtime_cls = RuntimeClass(None, TypeRefWithVars(cls_ident))  # type: ignore[arg-type]
 
         # Store frame so that we can get live access to updated locals/globals

python/egglog/egraph_state.py

@@ -796,23 +796,22 @@ def _unstable_fn_value_to_expr(
                 continue
             if signature.semantic_return_type.ident != return_tp.ident:
                 continue
-            tcs = TypeConstraintSolver(self.__egg_decls__)
+            tcs = TypeConstraintSolver()
 
-            arg_types, bound_tp_params = tcs.infer_arg_types(
-                signature.arg_types, signature.semantic_return_type, signature.var_arg_type, return_tp, None
+            arg_types = tcs.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)
             )
-
-            call_decl = CallDecl(
-                callable_ref,
-                args,
-                # Don't include bound type params if this is just a method, we only needed them for type resolution
-                # but dont need to store them
-                bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else (),
-            )
+            if isinstance(callable_ref, ClassMethodRef | InitRef):
+                bound_tp_params = tuple(
+                    map(tcs.substitute_typevars, self.__egg_decls__.get_class_decl(callable_ref.ident).type_vars)
+                )
+            else:
+                bound_tp_params = ()
+            call_decl = CallDecl(callable_ref, args, bound_tp_params)
             return PartialCallDecl(call_decl)
         raise ValueError(f"Function '{name}' not found")
 
@@ -909,11 +908,7 @@ def from_expr(self, tp: JustTypeRef, term: bindings._Term) -> TypedExprDecl:
             assert_never(term)
         return TypedExprDecl(tp, expr_decl)
 
-    def from_call(
-        self,
-        tp: JustTypeRef,
-        term: bindings.TermApp,  # additional_arg_tps: tuple[JustTypeRef, ...]
-    ) -> CallDecl:
+    def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl:
         """
         Convert a call to a CallDecl.
 
@@ -931,33 +926,32 @@ def from_call(
             signature = self.decls.get_callable_decl(callable_ref).signature
             assert isinstance(signature, FunctionSignature)
             if isinstance(callable_ref, ClassMethodRef | InitRef | MethodRef):
-                # Need OR in case we have class method whose class whas never added as a sort, which would happen
+                # 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 we don't need to bind any possible type.
+                # 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]
-                cls_name = callable_ref.ident
             else:
                 possible_types = [None]
-                cls_name = None
             for possible_type in possible_types:
-                tcs = TypeConstraintSolver(self.decls)
+                tcs = TypeConstraintSolver()
                 if possible_type and possible_type.args:
-                    tcs.bind_class(possible_type)
+                    tcs.bind_class(possible_type, self.decls)
                 try:
-                    arg_types, bound_tp_params = tcs.infer_arg_types(
-                        signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp, cls_name
+                    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 zip(term.args, arg_types, strict=False))
-
-                return CallDecl(
-                    callable_ref,
-                    args,
-                    # Don't include bound type params if this is just a method, we only needed them for type resolution
-                    # but dont need to store them
-                    bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) and not args else (),
-                )
+                args = tuple(self.resolve_term(a, tp) for a, tp in a_tp)
+                if not args and isinstance(callable_ref, ClassMethodRef | InitRef):
+                    bound_tp_params = tuple(
+                        map(tcs.substitute_typevars, self.decls.get_class_decl(callable_ref.ident).type_vars)
+                    )
+                else:
+                    bound_tp_params = ()
+                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: {self.state.egg_fn_to_callable_refs[term.name]}"
         )

python/egglog/exp/array_api_program_gen.py

@@ -144,7 +144,7 @@ def _value_program(i: Int, b: Boolean, f: Float, x: NDArray, v1: Value, v2: Valu
     yield rewrite(value_program(Value.float(f))).to(float_program(f))
     # Could add .item() but we usually dont need it.
     yield rewrite(value_program(x.to_value())).to(ndarray_program(x))
-    yield rewrite(value_program(v1 < v2)).to(Program("(") + value_program(v1) + " < " + value_program(v2) + ")")
+    yield rewrite(bool_program(v1 < v2)).to(Program("(") + value_program(v1) + " < " + value_program(v2) + ")")
     yield rewrite(value_program(v1 / v2)).to(Program("(") + value_program(v1) + " / " + value_program(v2) + ")")
     yield rewrite(value_program(v1 + v2)).to(Program("(") + value_program(v1) + " + " + value_program(v2) + ")")
     yield rewrite(value_program(v1 * v2)).to(Program("(") + value_program(v1) + " * " + value_program(v2) + ")")

python/egglog/runtime.py

@@ -350,7 +350,7 @@ def __getitem__(self, args: object) -> RuntimeClass:
             "tuple[tuple[DeclerationsLike, ...], tuple[TypeOrVarRef, ...]]",
             zip(*(resolve_type_annotation(arg) for arg in args), strict=False),
         )
-        # if we already have some args bound and some not, then we shold replace all existing args of typevars with new
+        # 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]
@@ -550,7 +550,7 @@ def __call__(  # noqa: C901
         assert not bound.kwargs
         args = bound.args
 
-        tcs = TypeConstraintSolver(decls)
+        tcs = TypeConstraintSolver()
         bound_tp = (
             None
             if self.__egg_bound__ is None
@@ -564,27 +564,24 @@ def __call__(  # noqa: C901
             # Don't  bind class if we have a first class function arg, b/c we don't support that yet
             and not function_value
         ):
-            tcs.bind_class(bound_tp)
+            tcs.bind_class(bound_tp, decls)
         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
-        # TODO: Replace type analysis class binding stuff
-        # with instead binders of location per call origination
-        cls_ident = bound_tp.ident if bound_tp else None
         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)
+            tcs.bind_class(_fn_tp_just, decls)
             if _fn_tp_just.args:
-                cls_ident = _fn_tp_just.ident
+                pass
         upcasted_args = [
-            resolve_literal(cast("TypeOrVarRef", tp), arg, Thunk.value(decls), tcs=tcs, cls_ident=cls_ident)
+            resolve_literal(cast("TypeOrVarRef", tp), arg, Thunk.value(decls), tcs)
             for arg, tp in zip_longest(args, signature.arg_types, fillvalue=signature.var_arg_type)
         ]
         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, cls_ident)
+        return_tp = tcs.substitute_typevars(signature.semantic_return_type)
         bound_params = (
             cast("JustTypeRef", bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef | InitRef) else ()
         )

python/egglog/type_constraint_solver.py

@@ -58,11 +58,11 @@ def infer_arg_types(
         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_)
-        for fn_arg in fn_args:
-            yield self.substitute_typevars(fn_arg)
-        if fn_var_args is not None:
-            var_arg_type = self.substitute_typevars(fn_var_args)
-            yield from chain(repeat(var_arg_type))
+        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:
         """
@@ -88,11 +88,11 @@ 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(tyepvar_ident):
+            case TypeVarRef(typevar_ident):
                 try:
-                    return self._typevar_to_type[tyepvar_ident]
+                    return self._typevar_to_type[typevar_ident]
                 except KeyError as e:
-                    raise TypeConstraintError(f"Unresolved type variable: {tyepvar_ident}") from 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)

python/tests/test_type_constraint_solver.py

@@ -14,40 +14,40 @@
 
 
 def test_simple() -> None:
-    tcs = TypeConstraintSolver(Declarations())
+    tcs = TypeConstraintSolver()
     tcs.infer_typevars(i64.to_var(), i64)
     assert tcs.substitute_typevars(i64.to_var()) == i64
 
 
 def test_wrong_arg() -> None:
-    tcs = TypeConstraintSolver(Declarations())
+    tcs = TypeConstraintSolver()
     with pytest.raises(TypeConstraintError):
         tcs.infer_typevars(i64.to_var(), unit)
 
 
 def test_generic() -> None:
-    tcs = TypeConstraintSolver(Declarations())
-    tcs.infer_typevars(map, map_i64_unit, Ident("Map"))
-    tcs.infer_typevars(K, i64, Ident("Map"))
-    assert tcs.substitute_typevars(V, Ident("Map")) == unit
+    tcs = TypeConstraintSolver()
+    tcs.infer_typevars(map, map_i64_unit)
+    tcs.infer_typevars(K, i64)
+    assert tcs.substitute_typevars(V) == unit
 
 
 def test_generic_wrong() -> None:
-    tcs = TypeConstraintSolver(Declarations())
-    tcs.infer_typevars(map, map_i64_unit, Ident("Map"))
+    tcs = TypeConstraintSolver()
+    tcs.infer_typevars(map, map_i64_unit)
     with pytest.raises(TypeConstraintError):
-        tcs.infer_typevars(K, unit, Ident("Map"))
+        tcs.infer_typevars(K, unit)
 
 
 def test_bound() -> None:
-    bound_cs = TypeConstraintSolver(decls)
-    bound_cs.bind_class(map_i64_unit)
-    bound_cs.infer_typevars(K, i64, Ident("Map"))
-    assert bound_cs.substitute_typevars(V, Ident("Map")) == unit
+    bound_cs = TypeConstraintSolver()
+    bound_cs.bind_class(map_i64_unit, decls)
+    bound_cs.infer_typevars(K, i64)
+    assert bound_cs.substitute_typevars(V) == unit
 
 
 def test_bound_wrong():
-    bound_cs = TypeConstraintSolver(decls)
-    bound_cs.bind_class(map_i64_unit)
+    bound_cs = TypeConstraintSolver()
+    bound_cs.bind_class(map_i64_unit, decls)
     with pytest.raises(TypeConstraintError):
-        bound_cs.infer_typevars(K, unit, Ident("Map"))
+        bound_cs.infer_typevars(K, unit)