expression.py 38 KB

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