expression.py 38 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069
  1. """Transform mypy expression ASTs to mypyc IR (Intermediate Representation).
  2. The top-level AST transformation logic is implemented in mypyc.irbuild.visitor
  3. and mypyc.irbuild.builder.
  4. """
  5. from __future__ import annotations
  6. import math
  7. from typing import Callable, Sequence
  8. from mypy.nodes import (
  9. ARG_POS,
  10. LDEF,
  11. AssertTypeExpr,
  12. AssignmentExpr,
  13. BytesExpr,
  14. CallExpr,
  15. CastExpr,
  16. ComparisonExpr,
  17. ComplexExpr,
  18. ConditionalExpr,
  19. DictExpr,
  20. DictionaryComprehension,
  21. EllipsisExpr,
  22. Expression,
  23. FloatExpr,
  24. GeneratorExpr,
  25. IndexExpr,
  26. IntExpr,
  27. ListComprehension,
  28. ListExpr,
  29. MemberExpr,
  30. MypyFile,
  31. NameExpr,
  32. OpExpr,
  33. RefExpr,
  34. SetComprehension,
  35. SetExpr,
  36. SliceExpr,
  37. StarExpr,
  38. StrExpr,
  39. SuperExpr,
  40. TupleExpr,
  41. TypeApplication,
  42. TypeInfo,
  43. UnaryExpr,
  44. Var,
  45. )
  46. from mypy.types import Instance, ProperType, TupleType, TypeType, get_proper_type
  47. from mypyc.common import MAX_SHORT_INT
  48. from mypyc.ir.class_ir import ClassIR
  49. from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD
  50. from mypyc.ir.ops import (
  51. Assign,
  52. BasicBlock,
  53. ComparisonOp,
  54. Float,
  55. Integer,
  56. LoadAddress,
  57. LoadLiteral,
  58. RaiseStandardError,
  59. Register,
  60. TupleGet,
  61. TupleSet,
  62. Value,
  63. )
  64. from mypyc.ir.rtypes import (
  65. RTuple,
  66. bool_rprimitive,
  67. int_rprimitive,
  68. is_fixed_width_rtype,
  69. is_int_rprimitive,
  70. is_list_rprimitive,
  71. is_none_rprimitive,
  72. object_rprimitive,
  73. set_rprimitive,
  74. )
  75. from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
  76. from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
  77. from mypyc.irbuild.constant_fold import constant_fold_expr
  78. from mypyc.irbuild.for_helpers import (
  79. comprehension_helper,
  80. translate_list_comprehension,
  81. translate_set_comprehension,
  82. )
  83. from mypyc.irbuild.format_str_tokenizer import (
  84. convert_format_expr_to_bytes,
  85. convert_format_expr_to_str,
  86. join_formatted_bytes,
  87. join_formatted_strings,
  88. tokenizer_printf_style,
  89. )
  90. from mypyc.irbuild.specialize import apply_function_specialization, apply_method_specialization
  91. from mypyc.irbuild.util import bytes_from_str
  92. from mypyc.primitives.bytes_ops import bytes_slice_op
  93. from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, dict_set_item_op
  94. from mypyc.primitives.generic_ops import iter_op
  95. from mypyc.primitives.int_ops import int_comparison_op_mapping
  96. from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
  97. from mypyc.primitives.misc_ops import ellipsis_op, get_module_dict_op, new_slice_op, type_op
  98. from mypyc.primitives.registry import CFunctionDescription, builtin_names
  99. from mypyc.primitives.set_ops import set_add_op, set_in_op, set_update_op
  100. from mypyc.primitives.str_ops import str_slice_op
  101. from mypyc.primitives.tuple_ops import list_tuple_op, tuple_slice_op
  102. # Name and attribute references
  103. def transform_name_expr(builder: IRBuilder, expr: NameExpr) -> Value:
  104. if expr.node is None:
  105. builder.add(
  106. RaiseStandardError(
  107. RaiseStandardError.RUNTIME_ERROR,
  108. "mypyc internal error: should be unreachable",
  109. expr.line,
  110. )
  111. )
  112. return builder.none()
  113. fullname = expr.node.fullname
  114. if fullname in builtin_names:
  115. typ, src = builtin_names[fullname]
  116. return builder.add(LoadAddress(typ, src, expr.line))
  117. # special cases
  118. if fullname == "builtins.None":
  119. return builder.none()
  120. if fullname == "builtins.True":
  121. return builder.true()
  122. if fullname == "builtins.False":
  123. return builder.false()
  124. math_literal = transform_math_literal(builder, fullname)
  125. if math_literal is not None:
  126. return math_literal
  127. if isinstance(expr.node, Var) and expr.node.is_final:
  128. value = builder.emit_load_final(
  129. expr.node,
  130. fullname,
  131. expr.name,
  132. builder.is_native_ref_expr(expr),
  133. builder.types[expr],
  134. expr.line,
  135. )
  136. if value is not None:
  137. return value
  138. if isinstance(expr.node, MypyFile) and expr.node.fullname in builder.imports:
  139. return builder.load_module(expr.node.fullname)
  140. # If the expression is locally defined, then read the result from the corresponding
  141. # assignment target and return it. Otherwise if the expression is a global, load it from
  142. # the globals dictionary.
  143. # Except for imports, that currently always happens in the global namespace.
  144. if expr.kind == LDEF and not (isinstance(expr.node, Var) and expr.node.is_suppressed_import):
  145. # Try to detect and error when we hit the irritating mypy bug
  146. # where a local variable is cast to None. (#5423)
  147. if (
  148. isinstance(expr.node, Var)
  149. and is_none_rprimitive(builder.node_type(expr))
  150. and expr.node.is_inferred
  151. ):
  152. builder.error(
  153. 'Local variable "{}" has inferred type None; add an annotation'.format(
  154. expr.node.name
  155. ),
  156. expr.node.line,
  157. )
  158. # TODO: Behavior currently only defined for Var, FuncDef and MypyFile node types.
  159. if isinstance(expr.node, MypyFile):
  160. # Load reference to a module imported inside function from
  161. # the modules dictionary. It would be closer to Python
  162. # semantics to access modules imported inside functions
  163. # via local variables, but this is tricky since the mypy
  164. # AST doesn't include a Var node for the module. We
  165. # instead load the module separately on each access.
  166. mod_dict = builder.call_c(get_module_dict_op, [], expr.line)
  167. obj = builder.call_c(
  168. dict_get_item_op, [mod_dict, builder.load_str(expr.node.fullname)], expr.line
  169. )
  170. return obj
  171. else:
  172. return builder.read(builder.get_assignment_target(expr, for_read=True), expr.line)
  173. return builder.load_global(expr)
  174. def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
  175. # First check if this is maybe a final attribute.
  176. final = builder.get_final_ref(expr)
  177. if final is not None:
  178. fullname, final_var, native = final
  179. value = builder.emit_load_final(
  180. final_var, fullname, final_var.name, native, builder.types[expr], expr.line
  181. )
  182. if value is not None:
  183. return value
  184. math_literal = transform_math_literal(builder, expr.fullname)
  185. if math_literal is not None:
  186. return math_literal
  187. if isinstance(expr.node, MypyFile) and expr.node.fullname in builder.imports:
  188. return builder.load_module(expr.node.fullname)
  189. can_borrow = builder.is_native_attr_ref(expr)
  190. obj = builder.accept(expr.expr, can_borrow=can_borrow)
  191. rtype = builder.node_type(expr)
  192. # Special case: for named tuples transform attribute access to faster index access.
  193. typ = get_proper_type(builder.types.get(expr.expr))
  194. if isinstance(typ, TupleType) and typ.partial_fallback.type.is_named_tuple:
  195. fields = typ.partial_fallback.type.metadata["namedtuple"]["fields"]
  196. if expr.name in fields:
  197. index = builder.builder.load_int(fields.index(expr.name))
  198. return builder.gen_method_call(obj, "__getitem__", [index], rtype, expr.line)
  199. check_instance_attribute_access_through_class(builder, expr, typ)
  200. borrow = can_borrow and builder.can_borrow
  201. return builder.builder.get_attr(obj, expr.name, rtype, expr.line, borrow=borrow)
  202. def check_instance_attribute_access_through_class(
  203. builder: IRBuilder, expr: MemberExpr, typ: ProperType | None
  204. ) -> None:
  205. """Report error if accessing an instance attribute through class object."""
  206. if isinstance(expr.expr, RefExpr):
  207. node = expr.expr.node
  208. if isinstance(typ, TypeType) and isinstance(typ.item, Instance):
  209. # TODO: Handle other item types
  210. node = typ.item.type
  211. if isinstance(node, TypeInfo):
  212. class_ir = builder.mapper.type_to_ir.get(node)
  213. if class_ir is not None and class_ir.is_ext_class:
  214. sym = node.get(expr.name)
  215. if (
  216. sym is not None
  217. and isinstance(sym.node, Var)
  218. and not sym.node.is_classvar
  219. and not sym.node.is_final
  220. ):
  221. builder.error(
  222. 'Cannot access instance attribute "{}" through class object'.format(
  223. expr.name
  224. ),
  225. expr.line,
  226. )
  227. builder.note(
  228. '(Hint: Use "x: Final = ..." or "x: ClassVar = ..." to define '
  229. "a class attribute)",
  230. expr.line,
  231. )
  232. def transform_super_expr(builder: IRBuilder, o: SuperExpr) -> Value:
  233. # warning(builder, 'can not optimize super() expression', o.line)
  234. sup_val = builder.load_module_attr_by_fullname("builtins.super", o.line)
  235. if o.call.args:
  236. args = [builder.accept(arg) for arg in o.call.args]
  237. else:
  238. assert o.info is not None
  239. typ = builder.load_native_type_object(o.info.fullname)
  240. ir = builder.mapper.type_to_ir[o.info]
  241. iter_env = iter(builder.builder.args)
  242. # Grab first argument
  243. vself: Value = next(iter_env)
  244. if builder.fn_info.is_generator:
  245. # grab sixth argument (see comment in translate_super_method_call)
  246. self_targ = list(builder.symtables[-1].values())[6]
  247. vself = builder.read(self_targ, builder.fn_info.fitem.line)
  248. elif not ir.is_ext_class:
  249. vself = next(iter_env) # second argument is self if non_extension class
  250. args = [typ, vself]
  251. res = builder.py_call(sup_val, args, o.line)
  252. return builder.py_get_attr(res, o.name, o.line)
  253. # Calls
  254. def transform_call_expr(builder: IRBuilder, expr: CallExpr) -> Value:
  255. callee = expr.callee
  256. if isinstance(expr.analyzed, CastExpr):
  257. return translate_cast_expr(builder, expr.analyzed)
  258. elif isinstance(expr.analyzed, AssertTypeExpr):
  259. # Compile to a no-op.
  260. return builder.accept(expr.analyzed.expr)
  261. elif (
  262. isinstance(callee, (NameExpr, MemberExpr))
  263. and isinstance(callee.node, TypeInfo)
  264. and callee.node.is_newtype
  265. ):
  266. # A call to a NewType type is a no-op at runtime.
  267. return builder.accept(expr.args[0])
  268. if isinstance(callee, IndexExpr) and isinstance(callee.analyzed, TypeApplication):
  269. callee = callee.analyzed.expr # Unwrap type application
  270. if isinstance(callee, MemberExpr):
  271. if isinstance(callee.expr, RefExpr) and isinstance(callee.expr.node, MypyFile):
  272. # Call a module-level function, not a method.
  273. return translate_call(builder, expr, callee)
  274. return apply_method_specialization(builder, expr, callee) or translate_method_call(
  275. builder, expr, callee
  276. )
  277. elif isinstance(callee, SuperExpr):
  278. return translate_super_method_call(builder, expr, callee)
  279. else:
  280. return translate_call(builder, expr, callee)
  281. def translate_call(builder: IRBuilder, expr: CallExpr, callee: Expression) -> Value:
  282. # The common case of calls is refexprs
  283. if isinstance(callee, RefExpr):
  284. return apply_function_specialization(builder, expr, callee) or translate_refexpr_call(
  285. builder, expr, callee
  286. )
  287. function = builder.accept(callee)
  288. args = [builder.accept(arg) for arg in expr.args]
  289. return builder.py_call(
  290. function, args, expr.line, arg_kinds=expr.arg_kinds, arg_names=expr.arg_names
  291. )
  292. def translate_refexpr_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value:
  293. """Translate a non-method call."""
  294. # Gen the argument values
  295. arg_values = [builder.accept(arg) for arg in expr.args]
  296. return builder.call_refexpr_with_args(expr, callee, arg_values)
  297. def translate_method_call(builder: IRBuilder, expr: CallExpr, callee: MemberExpr) -> Value:
  298. """Generate IR for an arbitrary call of form e.m(...).
  299. This can also deal with calls to module-level functions.
  300. """
  301. if builder.is_native_ref_expr(callee):
  302. # Call to module-level native function or such
  303. return translate_call(builder, expr, callee)
  304. elif (
  305. isinstance(callee.expr, RefExpr)
  306. and isinstance(callee.expr.node, TypeInfo)
  307. and callee.expr.node in builder.mapper.type_to_ir
  308. and builder.mapper.type_to_ir[callee.expr.node].has_method(callee.name)
  309. ):
  310. # Call a method via the *class*
  311. assert isinstance(callee.expr.node, TypeInfo)
  312. ir = builder.mapper.type_to_ir[callee.expr.node]
  313. return call_classmethod(builder, ir, expr, callee)
  314. elif builder.is_module_member_expr(callee):
  315. # Fall back to a PyCall for non-native module calls
  316. function = builder.accept(callee)
  317. args = [builder.accept(arg) for arg in expr.args]
  318. return builder.py_call(
  319. function, args, expr.line, arg_kinds=expr.arg_kinds, arg_names=expr.arg_names
  320. )
  321. else:
  322. if isinstance(callee.expr, RefExpr):
  323. node = callee.expr.node
  324. if isinstance(node, Var) and node.is_cls:
  325. typ = get_proper_type(node.type)
  326. if isinstance(typ, TypeType) and isinstance(typ.item, Instance):
  327. class_ir = builder.mapper.type_to_ir.get(typ.item.type)
  328. if class_ir and class_ir.is_ext_class and class_ir.has_no_subclasses():
  329. # Call a native classmethod via cls that can be statically bound,
  330. # since the class has no subclasses.
  331. return call_classmethod(builder, class_ir, expr, callee)
  332. receiver_typ = builder.node_type(callee.expr)
  333. # If there is a specializer for this method name/type, try calling it.
  334. # We would return the first successful one.
  335. val = apply_method_specialization(builder, expr, callee, receiver_typ)
  336. if val is not None:
  337. return val
  338. obj = builder.accept(callee.expr)
  339. args = [builder.accept(arg) for arg in expr.args]
  340. return builder.gen_method_call(
  341. obj,
  342. callee.name,
  343. args,
  344. builder.node_type(expr),
  345. expr.line,
  346. expr.arg_kinds,
  347. expr.arg_names,
  348. )
  349. def call_classmethod(builder: IRBuilder, ir: ClassIR, expr: CallExpr, callee: MemberExpr) -> Value:
  350. decl = ir.method_decl(callee.name)
  351. args = []
  352. arg_kinds, arg_names = expr.arg_kinds.copy(), expr.arg_names.copy()
  353. # Add the class argument for class methods in extension classes
  354. if decl.kind == FUNC_CLASSMETHOD and ir.is_ext_class:
  355. args.append(builder.load_native_type_object(ir.fullname))
  356. arg_kinds.insert(0, ARG_POS)
  357. arg_names.insert(0, None)
  358. args += [builder.accept(arg) for arg in expr.args]
  359. if ir.is_ext_class:
  360. return builder.builder.call(decl, args, arg_kinds, arg_names, expr.line)
  361. else:
  362. obj = builder.accept(callee.expr)
  363. return builder.gen_method_call(
  364. obj,
  365. callee.name,
  366. args,
  367. builder.node_type(expr),
  368. expr.line,
  369. expr.arg_kinds,
  370. expr.arg_names,
  371. )
  372. def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: SuperExpr) -> Value:
  373. if callee.info is None or (len(callee.call.args) != 0 and len(callee.call.args) != 2):
  374. return translate_call(builder, expr, callee)
  375. # We support two-argument super but only when it is super(CurrentClass, self)
  376. # TODO: We could support it when it is a parent class in many cases?
  377. if len(callee.call.args) == 2:
  378. self_arg = callee.call.args[1]
  379. if (
  380. not isinstance(self_arg, NameExpr)
  381. or not isinstance(self_arg.node, Var)
  382. or not self_arg.node.is_self
  383. ):
  384. return translate_call(builder, expr, callee)
  385. typ_arg = callee.call.args[0]
  386. if (
  387. not isinstance(typ_arg, NameExpr)
  388. or not isinstance(typ_arg.node, TypeInfo)
  389. or callee.info is not typ_arg.node
  390. ):
  391. return translate_call(builder, expr, callee)
  392. ir = builder.mapper.type_to_ir[callee.info]
  393. # Search for the method in the mro, skipping ourselves. We
  394. # determine targets of super calls to native methods statically.
  395. for base in ir.mro[1:]:
  396. if callee.name in base.method_decls:
  397. break
  398. else:
  399. if (
  400. ir.is_ext_class
  401. and ir.builtin_base is None
  402. and not ir.inherits_python
  403. and callee.name == "__init__"
  404. and len(expr.args) == 0
  405. ):
  406. # Call translates to object.__init__(self), which is a
  407. # no-op, so omit the call.
  408. return builder.none()
  409. return translate_call(builder, expr, callee)
  410. decl = base.method_decl(callee.name)
  411. arg_values = [builder.accept(arg) for arg in expr.args]
  412. arg_kinds, arg_names = expr.arg_kinds.copy(), expr.arg_names.copy()
  413. if decl.kind != FUNC_STATICMETHOD:
  414. # Grab first argument
  415. vself: Value = builder.self()
  416. if decl.kind == FUNC_CLASSMETHOD:
  417. vself = builder.call_c(type_op, [vself], expr.line)
  418. elif builder.fn_info.is_generator:
  419. # For generator classes, the self target is the 6th value
  420. # in the symbol table (which is an ordered dict). This is sort
  421. # of ugly, but we can't search by name since the 'self' parameter
  422. # could be named anything, and it doesn't get added to the
  423. # environment indexes.
  424. self_targ = list(builder.symtables[-1].values())[6]
  425. vself = builder.read(self_targ, builder.fn_info.fitem.line)
  426. arg_values.insert(0, vself)
  427. arg_kinds.insert(0, ARG_POS)
  428. arg_names.insert(0, None)
  429. return builder.builder.call(decl, arg_values, arg_kinds, arg_names, expr.line)
  430. def translate_cast_expr(builder: IRBuilder, expr: CastExpr) -> Value:
  431. src = builder.accept(expr.expr)
  432. target_type = builder.type_to_rtype(expr.type)
  433. return builder.coerce(src, target_type, expr.line)
  434. # Operators
  435. def transform_unary_expr(builder: IRBuilder, expr: UnaryExpr) -> Value:
  436. folded = try_constant_fold(builder, expr)
  437. if folded:
  438. return folded
  439. return builder.unary_op(builder.accept(expr.expr), expr.op, expr.line)
  440. def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
  441. if expr.op in ("and", "or"):
  442. return builder.shortcircuit_expr(expr)
  443. # Special case for string formatting
  444. if expr.op == "%" and isinstance(expr.left, (StrExpr, BytesExpr)):
  445. ret = translate_printf_style_formatting(builder, expr.left, expr.right)
  446. if ret is not None:
  447. return ret
  448. folded = try_constant_fold(builder, expr)
  449. if folded:
  450. return folded
  451. borrow_left = False
  452. borrow_right = False
  453. ltype = builder.node_type(expr.left)
  454. rtype = builder.node_type(expr.right)
  455. # Special case some int ops to allow borrowing operands.
  456. if is_int_rprimitive(ltype) and is_int_rprimitive(rtype):
  457. if expr.op == "//":
  458. expr = try_optimize_int_floor_divide(expr)
  459. if expr.op in int_borrow_friendly_op:
  460. borrow_left = is_borrow_friendly_expr(builder, expr.right)
  461. borrow_right = True
  462. elif is_fixed_width_rtype(ltype) and is_fixed_width_rtype(rtype):
  463. borrow_left = is_borrow_friendly_expr(builder, expr.right)
  464. borrow_right = True
  465. left = builder.accept(expr.left, can_borrow=borrow_left)
  466. right = builder.accept(expr.right, can_borrow=borrow_right)
  467. return builder.binary_op(left, right, expr.op, expr.line)
  468. def try_optimize_int_floor_divide(expr: OpExpr) -> OpExpr:
  469. """Replace // with a power of two with a right shift, if possible."""
  470. if not isinstance(expr.right, IntExpr):
  471. return expr
  472. divisor = expr.right.value
  473. shift = divisor.bit_length() - 1
  474. if 0 < shift < 28 and divisor == (1 << shift):
  475. return OpExpr(">>", expr.left, IntExpr(shift))
  476. return expr
  477. def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
  478. index = expr.index
  479. base_type = builder.node_type(expr.base)
  480. is_list = is_list_rprimitive(base_type)
  481. can_borrow_base = is_list and is_borrow_friendly_expr(builder, index)
  482. base = builder.accept(expr.base, can_borrow=can_borrow_base)
  483. if isinstance(base.type, RTuple) and isinstance(index, IntExpr):
  484. return builder.add(TupleGet(base, index.value, expr.line))
  485. if isinstance(index, SliceExpr):
  486. value = try_gen_slice_op(builder, base, index)
  487. if value:
  488. return value
  489. index_reg = builder.accept(expr.index, can_borrow=is_list)
  490. return builder.gen_method_call(
  491. base, "__getitem__", [index_reg], builder.node_type(expr), expr.line
  492. )
  493. def try_constant_fold(builder: IRBuilder, expr: Expression) -> Value | None:
  494. """Return the constant value of an expression if possible.
  495. Return None otherwise.
  496. """
  497. value = constant_fold_expr(builder, expr)
  498. if isinstance(value, int):
  499. return builder.load_int(value)
  500. elif isinstance(value, str):
  501. return builder.load_str(value)
  502. elif isinstance(value, float):
  503. return Float(value)
  504. return None
  505. def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Value | None:
  506. """Generate specialized slice op for some index expressions.
  507. Return None if a specialized op isn't available.
  508. This supports obj[x:y], obj[:x], and obj[x:] for a few types.
  509. """
  510. if index.stride:
  511. # We can only handle the default stride of 1.
  512. return None
  513. if index.begin_index:
  514. begin_type = builder.node_type(index.begin_index)
  515. else:
  516. begin_type = int_rprimitive
  517. if index.end_index:
  518. end_type = builder.node_type(index.end_index)
  519. else:
  520. end_type = int_rprimitive
  521. # Both begin and end index must be int (or missing).
  522. if is_int_rprimitive(begin_type) and is_int_rprimitive(end_type):
  523. if index.begin_index:
  524. begin = builder.accept(index.begin_index)
  525. else:
  526. begin = builder.load_int(0)
  527. if index.end_index:
  528. end = builder.accept(index.end_index)
  529. else:
  530. # Replace missing end index with the largest short integer
  531. # (a sequence can't be longer).
  532. end = builder.load_int(MAX_SHORT_INT)
  533. candidates = [list_slice_op, tuple_slice_op, str_slice_op, bytes_slice_op]
  534. return builder.builder.matching_call_c(candidates, [base, begin, end], index.line)
  535. return None
  536. def transform_conditional_expr(builder: IRBuilder, expr: ConditionalExpr) -> Value:
  537. if_body, else_body, next_block = BasicBlock(), BasicBlock(), BasicBlock()
  538. process_conditional(builder, expr.cond, if_body, else_body)
  539. expr_type = builder.node_type(expr)
  540. # Having actual Phi nodes would be really nice here!
  541. target = Register(expr_type)
  542. builder.activate_block(if_body)
  543. true_value = builder.accept(expr.if_expr)
  544. true_value = builder.coerce(true_value, expr_type, expr.line)
  545. builder.add(Assign(target, true_value))
  546. builder.goto(next_block)
  547. builder.activate_block(else_body)
  548. false_value = builder.accept(expr.else_expr)
  549. false_value = builder.coerce(false_value, expr_type, expr.line)
  550. builder.add(Assign(target, false_value))
  551. builder.goto(next_block)
  552. builder.activate_block(next_block)
  553. return target
  554. def set_literal_values(builder: IRBuilder, items: Sequence[Expression]) -> list[object] | None:
  555. values: list[object] = []
  556. for item in items:
  557. const_value = constant_fold_expr(builder, item)
  558. if const_value is not None:
  559. values.append(const_value)
  560. continue
  561. if isinstance(item, RefExpr):
  562. if item.fullname == "builtins.None":
  563. values.append(None)
  564. elif item.fullname == "builtins.True":
  565. values.append(True)
  566. elif item.fullname == "builtins.False":
  567. values.append(False)
  568. elif isinstance(item, (BytesExpr, FloatExpr, ComplexExpr)):
  569. # constant_fold_expr() doesn't handle these (yet?)
  570. v = bytes_from_str(item.value) if isinstance(item, BytesExpr) else item.value
  571. values.append(v)
  572. elif isinstance(item, TupleExpr):
  573. tuple_values = set_literal_values(builder, item.items)
  574. if tuple_values is not None:
  575. values.append(tuple(tuple_values))
  576. if len(values) != len(items):
  577. # Bail if not all items can be converted into values.
  578. return None
  579. return values
  580. def precompute_set_literal(builder: IRBuilder, s: SetExpr) -> Value | None:
  581. """Try to pre-compute a frozenset literal during module initialization.
  582. Return None if it's not possible.
  583. Supported items:
  584. - Anything supported by irbuild.constant_fold.constant_fold_expr()
  585. - None, True, and False
  586. - Float, byte, and complex literals
  587. - Tuple literals with only items listed above
  588. """
  589. values = set_literal_values(builder, s.items)
  590. if values is not None:
  591. return builder.add(LoadLiteral(frozenset(values), set_rprimitive))
  592. return None
  593. def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
  594. # x in (...)/[...]
  595. # x not in (...)/[...]
  596. first_op = e.operators[0]
  597. if (
  598. first_op in ["in", "not in"]
  599. and len(e.operators) == 1
  600. and isinstance(e.operands[1], (TupleExpr, ListExpr))
  601. ):
  602. items = e.operands[1].items
  603. n_items = len(items)
  604. # x in y -> x == y[0] or ... or x == y[n]
  605. # x not in y -> x != y[0] and ... and x != y[n]
  606. # 16 is arbitrarily chosen to limit code size
  607. if 1 < n_items < 16:
  608. if e.operators[0] == "in":
  609. bin_op = "or"
  610. cmp_op = "=="
  611. else:
  612. bin_op = "and"
  613. cmp_op = "!="
  614. lhs = e.operands[0]
  615. mypy_file = builder.graph["builtins"].tree
  616. assert mypy_file is not None
  617. info = mypy_file.names["bool"].node
  618. assert isinstance(info, TypeInfo)
  619. bool_type = Instance(info, [])
  620. exprs = []
  621. for item in items:
  622. expr = ComparisonExpr([cmp_op], [lhs, item])
  623. builder.types[expr] = bool_type
  624. exprs.append(expr)
  625. or_expr: Expression = exprs.pop(0)
  626. for expr in exprs:
  627. or_expr = OpExpr(bin_op, or_expr, expr)
  628. builder.types[or_expr] = bool_type
  629. return builder.accept(or_expr)
  630. # x in [y]/(y) -> x == y
  631. # x not in [y]/(y) -> x != y
  632. elif n_items == 1:
  633. if e.operators[0] == "in":
  634. cmp_op = "=="
  635. else:
  636. cmp_op = "!="
  637. e.operators = [cmp_op]
  638. e.operands[1] = items[0]
  639. # x in []/() -> False
  640. # x not in []/() -> True
  641. elif n_items == 0:
  642. if e.operators[0] == "in":
  643. return builder.false()
  644. else:
  645. return builder.true()
  646. # x in {...}
  647. # x not in {...}
  648. if (
  649. first_op in ("in", "not in")
  650. and len(e.operators) == 1
  651. and isinstance(e.operands[1], SetExpr)
  652. ):
  653. set_literal = precompute_set_literal(builder, e.operands[1])
  654. if set_literal is not None:
  655. lhs = e.operands[0]
  656. result = builder.builder.call_c(
  657. set_in_op, [builder.accept(lhs), set_literal], e.line, bool_rprimitive
  658. )
  659. if first_op == "not in":
  660. return builder.unary_op(result, "not", e.line)
  661. return result
  662. if len(e.operators) == 1:
  663. # Special some common simple cases
  664. if first_op in ("is", "is not"):
  665. right_expr = e.operands[1]
  666. if isinstance(right_expr, NameExpr) and right_expr.fullname == "builtins.None":
  667. # Special case 'is None' / 'is not None'.
  668. return translate_is_none(builder, e.operands[0], negated=first_op != "is")
  669. left_expr = e.operands[0]
  670. if is_int_rprimitive(builder.node_type(left_expr)):
  671. right_expr = e.operands[1]
  672. if is_int_rprimitive(builder.node_type(right_expr)):
  673. if first_op in int_borrow_friendly_op:
  674. borrow_left = is_borrow_friendly_expr(builder, right_expr)
  675. left = builder.accept(left_expr, can_borrow=borrow_left)
  676. right = builder.accept(right_expr, can_borrow=True)
  677. return builder.compare_tagged(left, right, first_op, e.line)
  678. # TODO: Don't produce an expression when used in conditional context
  679. # All of the trickiness here is due to support for chained conditionals
  680. # (`e1 < e2 > e3`, etc). `e1 < e2 > e3` is approximately equivalent to
  681. # `e1 < e2 and e2 > e3` except that `e2` is only evaluated once.
  682. expr_type = builder.node_type(e)
  683. # go(i, prev) generates code for `ei opi e{i+1} op{i+1} ... en`,
  684. # assuming that prev contains the value of `ei`.
  685. def go(i: int, prev: Value) -> Value:
  686. if i == len(e.operators) - 1:
  687. return transform_basic_comparison(
  688. builder, e.operators[i], prev, builder.accept(e.operands[i + 1]), e.line
  689. )
  690. next = builder.accept(e.operands[i + 1])
  691. return builder.builder.shortcircuit_helper(
  692. "and",
  693. expr_type,
  694. lambda: transform_basic_comparison(builder, e.operators[i], prev, next, e.line),
  695. lambda: go(i + 1, next),
  696. e.line,
  697. )
  698. return go(0, builder.accept(e.operands[0]))
  699. def translate_is_none(builder: IRBuilder, expr: Expression, negated: bool) -> Value:
  700. v = builder.accept(expr, can_borrow=True)
  701. return builder.binary_op(v, builder.none_object(), "is not" if negated else "is", expr.line)
  702. def transform_basic_comparison(
  703. builder: IRBuilder, op: str, left: Value, right: Value, line: int
  704. ) -> Value:
  705. if (
  706. is_int_rprimitive(left.type)
  707. and is_int_rprimitive(right.type)
  708. and op in int_comparison_op_mapping
  709. ):
  710. return builder.compare_tagged(left, right, op, line)
  711. if is_fixed_width_rtype(left.type) and op in int_comparison_op_mapping:
  712. if right.type == left.type:
  713. op_id = ComparisonOp.signed_ops[op]
  714. return builder.builder.comparison_op(left, right, op_id, line)
  715. elif isinstance(right, Integer):
  716. op_id = ComparisonOp.signed_ops[op]
  717. return builder.builder.comparison_op(
  718. left, Integer(right.value >> 1, left.type), op_id, line
  719. )
  720. elif (
  721. is_fixed_width_rtype(right.type)
  722. and op in int_comparison_op_mapping
  723. and isinstance(left, Integer)
  724. ):
  725. op_id = ComparisonOp.signed_ops[op]
  726. return builder.builder.comparison_op(
  727. Integer(left.value >> 1, right.type), right, op_id, line
  728. )
  729. negate = False
  730. if op == "is not":
  731. op, negate = "is", True
  732. elif op == "not in":
  733. op, negate = "in", True
  734. target = builder.binary_op(left, right, op, line)
  735. if negate:
  736. target = builder.unary_op(target, "not", line)
  737. return target
  738. def translate_printf_style_formatting(
  739. builder: IRBuilder, format_expr: StrExpr | BytesExpr, rhs: Expression
  740. ) -> Value | None:
  741. tokens = tokenizer_printf_style(format_expr.value)
  742. if tokens is not None:
  743. literals, format_ops = tokens
  744. exprs = []
  745. if isinstance(rhs, TupleExpr):
  746. exprs = rhs.items
  747. elif isinstance(rhs, Expression):
  748. exprs.append(rhs)
  749. if isinstance(format_expr, BytesExpr):
  750. substitutions = convert_format_expr_to_bytes(
  751. builder, format_ops, exprs, format_expr.line
  752. )
  753. if substitutions is not None:
  754. return join_formatted_bytes(builder, literals, substitutions, format_expr.line)
  755. else:
  756. substitutions = convert_format_expr_to_str(
  757. builder, format_ops, exprs, format_expr.line
  758. )
  759. if substitutions is not None:
  760. return join_formatted_strings(builder, literals, substitutions, format_expr.line)
  761. return None
  762. # Literals
  763. def transform_int_expr(builder: IRBuilder, expr: IntExpr) -> Value:
  764. return builder.builder.load_int(expr.value)
  765. def transform_float_expr(builder: IRBuilder, expr: FloatExpr) -> Value:
  766. return builder.builder.load_float(expr.value)
  767. def transform_complex_expr(builder: IRBuilder, expr: ComplexExpr) -> Value:
  768. return builder.builder.load_complex(expr.value)
  769. def transform_str_expr(builder: IRBuilder, expr: StrExpr) -> Value:
  770. return builder.load_str(expr.value)
  771. def transform_bytes_expr(builder: IRBuilder, expr: BytesExpr) -> Value:
  772. return builder.load_bytes_from_str_literal(expr.value)
  773. def transform_ellipsis(builder: IRBuilder, o: EllipsisExpr) -> Value:
  774. return builder.add(LoadAddress(ellipsis_op.type, ellipsis_op.src, o.line))
  775. # Display expressions
  776. def transform_list_expr(builder: IRBuilder, expr: ListExpr) -> Value:
  777. return _visit_list_display(builder, expr.items, expr.line)
  778. def _visit_list_display(builder: IRBuilder, items: list[Expression], line: int) -> Value:
  779. return _visit_display(
  780. builder, items, builder.new_list_op, list_append_op, list_extend_op, line, True
  781. )
  782. def transform_tuple_expr(builder: IRBuilder, expr: TupleExpr) -> Value:
  783. if any(isinstance(item, StarExpr) for item in expr.items):
  784. # create a tuple of unknown length
  785. return _visit_tuple_display(builder, expr)
  786. # create a tuple of fixed length (RTuple)
  787. tuple_type = builder.node_type(expr)
  788. # When handling NamedTuple et. al we might not have proper type info,
  789. # so make some up if we need it.
  790. types = (
  791. tuple_type.types
  792. if isinstance(tuple_type, RTuple)
  793. else [object_rprimitive] * len(expr.items)
  794. )
  795. items = []
  796. for item_expr, item_type in zip(expr.items, types):
  797. reg = builder.accept(item_expr)
  798. items.append(builder.coerce(reg, item_type, item_expr.line))
  799. return builder.add(TupleSet(items, expr.line))
  800. def _visit_tuple_display(builder: IRBuilder, expr: TupleExpr) -> Value:
  801. """Create a list, then turn it into a tuple."""
  802. val_as_list = _visit_list_display(builder, expr.items, expr.line)
  803. return builder.call_c(list_tuple_op, [val_as_list], expr.line)
  804. def transform_dict_expr(builder: IRBuilder, expr: DictExpr) -> Value:
  805. """First accepts all keys and values, then makes a dict out of them."""
  806. key_value_pairs = []
  807. for key_expr, value_expr in expr.items:
  808. key = builder.accept(key_expr) if key_expr is not None else None
  809. value = builder.accept(value_expr)
  810. key_value_pairs.append((key, value))
  811. return builder.builder.make_dict(key_value_pairs, expr.line)
  812. def transform_set_expr(builder: IRBuilder, expr: SetExpr) -> Value:
  813. return _visit_display(
  814. builder, expr.items, builder.new_set_op, set_add_op, set_update_op, expr.line, False
  815. )
  816. def _visit_display(
  817. builder: IRBuilder,
  818. items: list[Expression],
  819. constructor_op: Callable[[list[Value], int], Value],
  820. append_op: CFunctionDescription,
  821. extend_op: CFunctionDescription,
  822. line: int,
  823. is_list: bool,
  824. ) -> Value:
  825. accepted_items = []
  826. for item in items:
  827. if isinstance(item, StarExpr):
  828. accepted_items.append((True, builder.accept(item.expr)))
  829. else:
  830. accepted_items.append((False, builder.accept(item)))
  831. result: Value | None = None
  832. initial_items = []
  833. for starred, value in accepted_items:
  834. if result is None and not starred and is_list:
  835. initial_items.append(value)
  836. continue
  837. if result is None:
  838. result = constructor_op(initial_items, line)
  839. builder.call_c(extend_op if starred else append_op, [result, value], line)
  840. if result is None:
  841. result = constructor_op(initial_items, line)
  842. return result
  843. # Comprehensions
  844. def transform_list_comprehension(builder: IRBuilder, o: ListComprehension) -> Value:
  845. return translate_list_comprehension(builder, o.generator)
  846. def transform_set_comprehension(builder: IRBuilder, o: SetComprehension) -> Value:
  847. return translate_set_comprehension(builder, o.generator)
  848. def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehension) -> Value:
  849. d = builder.maybe_spill(builder.call_c(dict_new_op, [], o.line))
  850. loop_params = list(zip(o.indices, o.sequences, o.condlists, o.is_async))
  851. def gen_inner_stmts() -> None:
  852. k = builder.accept(o.key)
  853. v = builder.accept(o.value)
  854. builder.call_c(dict_set_item_op, [builder.read(d), k, v], o.line)
  855. comprehension_helper(builder, loop_params, gen_inner_stmts, o.line)
  856. return builder.read(d)
  857. # Misc
  858. def transform_slice_expr(builder: IRBuilder, expr: SliceExpr) -> Value:
  859. def get_arg(arg: Expression | None) -> Value:
  860. if arg is None:
  861. return builder.none_object()
  862. else:
  863. return builder.accept(arg)
  864. args = [get_arg(expr.begin_index), get_arg(expr.end_index), get_arg(expr.stride)]
  865. return builder.call_c(new_slice_op, args, expr.line)
  866. def transform_generator_expr(builder: IRBuilder, o: GeneratorExpr) -> Value:
  867. builder.warning("Treating generator comprehension as list", o.line)
  868. return builder.call_c(iter_op, [translate_list_comprehension(builder, o)], o.line)
  869. def transform_assignment_expr(builder: IRBuilder, o: AssignmentExpr) -> Value:
  870. value = builder.accept(o.value)
  871. target = builder.get_assignment_target(o.target)
  872. builder.assign(target, value, o.line)
  873. return value
  874. def transform_math_literal(builder: IRBuilder, fullname: str) -> Value | None:
  875. if fullname == "math.e":
  876. return builder.load_float(math.e)
  877. if fullname == "math.pi":
  878. return builder.load_float(math.pi)
  879. if fullname == "math.inf":
  880. return builder.load_float(math.inf)
  881. if fullname == "math.nan":
  882. return builder.load_float(math.nan)
  883. if fullname == "math.tau":
  884. return builder.load_float(math.tau)
  885. return None