Skip to content

Commit

Permalink
[mypyc] Optimize int()/float() on native classes
Browse files Browse the repository at this point in the history
int() and float() calls on native classes can simply call the associated
dunder if the RInstance defines it, no need to load the type and call it.

bool() calls were already optimized merely a few days ago, but there
wasn't an IRbuild test verifying this so I added one.
  • Loading branch information
ichard26 committed Jan 15, 2023
1 parent 28e5436 commit 36a4d78
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 13 deletions.
13 changes: 9 additions & 4 deletions mypyc/irbuild/specialize.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,14 +157,19 @@ def translate_globals(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Va


@specialize_function("builtins.abs")
def translate_abs(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
"""Specialize calls on native classes that implement __abs__."""
@specialize_function("builtins.int")
@specialize_function("builtins.float")
def translate_builtins_with_unary_dunder(
builder: IRBuilder, expr: CallExpr, callee: RefExpr
) -> Value | None:
"""Specialize calls on native classes that implement the associated dunder."""
if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
arg = expr.args[0]
arg_typ = builder.node_type(arg)
if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method("__abs__"):
method = f"__{callee.name}__"
if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method(method):
obj = builder.accept(arg)
return builder.gen_method_call(obj, "__abs__", [], None, expr.line)
return builder.gen_method_call(obj, method, [], None, expr.line)

return None

Expand Down
26 changes: 17 additions & 9 deletions mypyc/test-data/irbuild-dunders.test
Original file line number Diff line number Diff line change
Expand Up @@ -154,13 +154,17 @@ class C:
def __abs__(self) -> int:
return 6

def __bool__(self) -> bool:
return False

def f(c: C) -> None:
-c
~c
int(c)
float(c)
+c
abs(c)
bool(c)
[out]
def C.__neg__(self):
self :: __main__.C
Expand Down Expand Up @@ -188,19 +192,23 @@ def C.__abs__(self):
self :: __main__.C
L0:
return 12
def C.__bool__(self):
self :: __main__.C
L0:
return 0
def f(c):
c :: __main__.C
r0, r1 :: int
r2, r3, r4, r5 :: object
r6, r7 :: int
r0, r1, r2 :: int
r3 :: float
r4, r5 :: int
r6 :: bool
L0:
r0 = c.__neg__()
r1 = c.__invert__()
r2 = load_address PyLong_Type
r3 = PyObject_CallFunctionObjArgs(r2, c, 0)
r4 = load_address PyFloat_Type
r5 = PyObject_CallFunctionObjArgs(r4, c, 0)
r6 = c.__pos__()
r7 = c.__abs__()
r2 = c.__int__()
r3 = c.__float__()
r4 = c.__pos__()
r5 = c.__abs__()
r6 = c.__bool__()
return 1

0 comments on commit 36a4d78

Please sign in to comment.