specialize.py 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743
  1. """Special case IR generation of calls to specific builtin functions.
  2. Most special cases should be handled using the data driven "primitive
  3. ops" system, but certain operations require special handling that has
  4. access to the AST/IR directly and can make decisions/optimizations
  5. based on it. These special cases can be implemented here.
  6. For example, we use specializers to statically emit the length of a
  7. fixed length tuple and to emit optimized code for any()/all() calls with
  8. generator comprehensions as the argument.
  9. See comment below for more documentation.
  10. """
  11. from __future__ import annotations
  12. from typing import Callable, Optional
  13. from mypy.nodes import (
  14. ARG_NAMED,
  15. ARG_POS,
  16. CallExpr,
  17. DictExpr,
  18. Expression,
  19. GeneratorExpr,
  20. IntExpr,
  21. ListExpr,
  22. MemberExpr,
  23. NameExpr,
  24. RefExpr,
  25. StrExpr,
  26. TupleExpr,
  27. )
  28. from mypy.types import AnyType, TypeOfAny
  29. from mypyc.ir.ops import (
  30. BasicBlock,
  31. Extend,
  32. Integer,
  33. RaiseStandardError,
  34. Register,
  35. Truncate,
  36. Unreachable,
  37. Value,
  38. )
  39. from mypyc.ir.rtypes import (
  40. RInstance,
  41. RTuple,
  42. RType,
  43. bool_rprimitive,
  44. c_int_rprimitive,
  45. dict_rprimitive,
  46. int32_rprimitive,
  47. int64_rprimitive,
  48. int_rprimitive,
  49. is_bool_rprimitive,
  50. is_dict_rprimitive,
  51. is_fixed_width_rtype,
  52. is_float_rprimitive,
  53. is_int32_rprimitive,
  54. is_int64_rprimitive,
  55. is_int_rprimitive,
  56. is_list_rprimitive,
  57. list_rprimitive,
  58. set_rprimitive,
  59. str_rprimitive,
  60. )
  61. from mypyc.irbuild.builder import IRBuilder
  62. from mypyc.irbuild.for_helpers import (
  63. comprehension_helper,
  64. sequence_from_generator_preallocate_helper,
  65. translate_list_comprehension,
  66. translate_set_comprehension,
  67. )
  68. from mypyc.irbuild.format_str_tokenizer import (
  69. FormatOp,
  70. convert_format_expr_to_str,
  71. join_formatted_strings,
  72. tokenizer_format_call,
  73. )
  74. from mypyc.primitives.dict_ops import (
  75. dict_items_op,
  76. dict_keys_op,
  77. dict_setdefault_spec_init_op,
  78. dict_values_op,
  79. )
  80. from mypyc.primitives.list_ops import new_list_set_item_op
  81. from mypyc.primitives.tuple_ops import new_tuple_set_item_op
  82. # Specializers are attempted before compiling the arguments to the
  83. # function. Specializers can return None to indicate that they failed
  84. # and the call should be compiled normally. Otherwise they should emit
  85. # code for the call and return a Value containing the result.
  86. #
  87. # Specializers take three arguments: the IRBuilder, the CallExpr being
  88. # compiled, and the RefExpr that is the left hand side of the call.
  89. Specializer = Callable[["IRBuilder", CallExpr, RefExpr], Optional[Value]]
  90. # Dictionary containing all configured specializers.
  91. #
  92. # Specializers can operate on methods as well, and are keyed on the
  93. # name and RType in that case.
  94. specializers: dict[tuple[str, RType | None], list[Specializer]] = {}
  95. def _apply_specialization(
  96. builder: IRBuilder, expr: CallExpr, callee: RefExpr, name: str | None, typ: RType | None = None
  97. ) -> Value | None:
  98. # TODO: Allow special cases to have default args or named args. Currently they don't since
  99. # they check that everything in arg_kinds is ARG_POS.
  100. # If there is a specializer for this function, try calling it.
  101. # Return the first successful one.
  102. if name and (name, typ) in specializers:
  103. for specializer in specializers[name, typ]:
  104. val = specializer(builder, expr, callee)
  105. if val is not None:
  106. return val
  107. return None
  108. def apply_function_specialization(
  109. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  110. ) -> Value | None:
  111. """Invoke the Specializer callback for a function if one has been registered"""
  112. return _apply_specialization(builder, expr, callee, callee.fullname)
  113. def apply_method_specialization(
  114. builder: IRBuilder, expr: CallExpr, callee: MemberExpr, typ: RType | None = None
  115. ) -> Value | None:
  116. """Invoke the Specializer callback for a method if one has been registered"""
  117. name = callee.fullname if typ is None else callee.name
  118. return _apply_specialization(builder, expr, callee, name, typ)
  119. def specialize_function(
  120. name: str, typ: RType | None = None
  121. ) -> Callable[[Specializer], Specializer]:
  122. """Decorator to register a function as being a specializer.
  123. There may exist multiple specializers for one function. When
  124. translating method calls, the earlier appended specializer has
  125. higher priority.
  126. """
  127. def wrapper(f: Specializer) -> Specializer:
  128. specializers.setdefault((name, typ), []).append(f)
  129. return f
  130. return wrapper
  131. @specialize_function("builtins.globals")
  132. def translate_globals(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  133. if len(expr.args) == 0:
  134. return builder.load_globals_dict()
  135. return None
  136. @specialize_function("builtins.abs")
  137. @specialize_function("builtins.int")
  138. @specialize_function("builtins.float")
  139. @specialize_function("builtins.complex")
  140. @specialize_function("mypy_extensions.i64")
  141. @specialize_function("mypy_extensions.i32")
  142. def translate_builtins_with_unary_dunder(
  143. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  144. ) -> Value | None:
  145. """Specialize calls on native classes that implement the associated dunder."""
  146. if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and isinstance(callee, NameExpr):
  147. arg = expr.args[0]
  148. arg_typ = builder.node_type(arg)
  149. shortname = callee.fullname.split(".")[1]
  150. if shortname in ("i64", "i32"):
  151. method = "__int__"
  152. else:
  153. method = f"__{shortname}__"
  154. if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method(method):
  155. obj = builder.accept(arg)
  156. return builder.gen_method_call(obj, method, [], None, expr.line)
  157. return None
  158. @specialize_function("builtins.len")
  159. def translate_len(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  160. if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
  161. arg = expr.args[0]
  162. expr_rtype = builder.node_type(arg)
  163. if isinstance(expr_rtype, RTuple):
  164. # len() of fixed-length tuple can be trivially determined
  165. # statically, though we still need to evaluate it.
  166. builder.accept(arg)
  167. return Integer(len(expr_rtype.types))
  168. else:
  169. if is_list_rprimitive(builder.node_type(arg)):
  170. borrow = True
  171. else:
  172. borrow = False
  173. obj = builder.accept(arg, can_borrow=borrow)
  174. return builder.builtin_len(obj, expr.line)
  175. return None
  176. @specialize_function("builtins.list")
  177. def dict_methods_fast_path(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  178. """Specialize a common case when list() is called on a dictionary
  179. view method call.
  180. For example:
  181. foo = list(bar.keys())
  182. """
  183. if not (len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]):
  184. return None
  185. arg = expr.args[0]
  186. if not (isinstance(arg, CallExpr) and not arg.args and isinstance(arg.callee, MemberExpr)):
  187. return None
  188. base = arg.callee.expr
  189. attr = arg.callee.name
  190. rtype = builder.node_type(base)
  191. if not (is_dict_rprimitive(rtype) and attr in ("keys", "values", "items")):
  192. return None
  193. obj = builder.accept(base)
  194. # Note that it is not safe to use fast methods on dict subclasses,
  195. # so the corresponding helpers in CPy.h fallback to (inlined)
  196. # generic logic.
  197. if attr == "keys":
  198. return builder.call_c(dict_keys_op, [obj], expr.line)
  199. elif attr == "values":
  200. return builder.call_c(dict_values_op, [obj], expr.line)
  201. else:
  202. return builder.call_c(dict_items_op, [obj], expr.line)
  203. @specialize_function("builtins.list")
  204. def translate_list_from_generator_call(
  205. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  206. ) -> Value | None:
  207. """Special case for simplest list comprehension.
  208. For example:
  209. list(f(x) for x in some_list/some_tuple/some_str)
  210. 'translate_list_comprehension()' would take care of other cases
  211. if this fails.
  212. """
  213. if (
  214. len(expr.args) == 1
  215. and expr.arg_kinds[0] == ARG_POS
  216. and isinstance(expr.args[0], GeneratorExpr)
  217. ):
  218. return sequence_from_generator_preallocate_helper(
  219. builder,
  220. expr.args[0],
  221. empty_op_llbuilder=builder.builder.new_list_op_with_length,
  222. set_item_op=new_list_set_item_op,
  223. )
  224. return None
  225. @specialize_function("builtins.tuple")
  226. def translate_tuple_from_generator_call(
  227. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  228. ) -> Value | None:
  229. """Special case for simplest tuple creation from a generator.
  230. For example:
  231. tuple(f(x) for x in some_list/some_tuple/some_str)
  232. 'translate_safe_generator_call()' would take care of other cases
  233. if this fails.
  234. """
  235. if (
  236. len(expr.args) == 1
  237. and expr.arg_kinds[0] == ARG_POS
  238. and isinstance(expr.args[0], GeneratorExpr)
  239. ):
  240. return sequence_from_generator_preallocate_helper(
  241. builder,
  242. expr.args[0],
  243. empty_op_llbuilder=builder.builder.new_tuple_with_length,
  244. set_item_op=new_tuple_set_item_op,
  245. )
  246. return None
  247. @specialize_function("builtins.set")
  248. def translate_set_from_generator_call(
  249. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  250. ) -> Value | None:
  251. """Special case for set creation from a generator.
  252. For example:
  253. set(f(...) for ... in iterator/nested_generators...)
  254. """
  255. if (
  256. len(expr.args) == 1
  257. and expr.arg_kinds[0] == ARG_POS
  258. and isinstance(expr.args[0], GeneratorExpr)
  259. ):
  260. return translate_set_comprehension(builder, expr.args[0])
  261. return None
  262. @specialize_function("builtins.min")
  263. @specialize_function("builtins.max")
  264. def faster_min_max(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  265. if expr.arg_kinds == [ARG_POS, ARG_POS]:
  266. x, y = builder.accept(expr.args[0]), builder.accept(expr.args[1])
  267. result = Register(builder.node_type(expr))
  268. # CPython evaluates arguments reversely when calling min(...) or max(...)
  269. if callee.fullname == "builtins.min":
  270. comparison = builder.binary_op(y, x, "<", expr.line)
  271. else:
  272. comparison = builder.binary_op(y, x, ">", expr.line)
  273. true_block, false_block, next_block = BasicBlock(), BasicBlock(), BasicBlock()
  274. builder.add_bool_branch(comparison, true_block, false_block)
  275. builder.activate_block(true_block)
  276. builder.assign(result, builder.coerce(y, result.type, expr.line), expr.line)
  277. builder.goto(next_block)
  278. builder.activate_block(false_block)
  279. builder.assign(result, builder.coerce(x, result.type, expr.line), expr.line)
  280. builder.goto(next_block)
  281. builder.activate_block(next_block)
  282. return result
  283. return None
  284. @specialize_function("builtins.tuple")
  285. @specialize_function("builtins.frozenset")
  286. @specialize_function("builtins.dict")
  287. @specialize_function("builtins.min")
  288. @specialize_function("builtins.max")
  289. @specialize_function("builtins.sorted")
  290. @specialize_function("collections.OrderedDict")
  291. @specialize_function("join", str_rprimitive)
  292. @specialize_function("extend", list_rprimitive)
  293. @specialize_function("update", dict_rprimitive)
  294. @specialize_function("update", set_rprimitive)
  295. def translate_safe_generator_call(
  296. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  297. ) -> Value | None:
  298. """Special cases for things that consume iterators where we know we
  299. can safely compile a generator into a list.
  300. """
  301. if (
  302. len(expr.args) > 0
  303. and expr.arg_kinds[0] == ARG_POS
  304. and isinstance(expr.args[0], GeneratorExpr)
  305. ):
  306. if isinstance(callee, MemberExpr):
  307. return builder.gen_method_call(
  308. builder.accept(callee.expr),
  309. callee.name,
  310. (
  311. [translate_list_comprehension(builder, expr.args[0])]
  312. + [builder.accept(arg) for arg in expr.args[1:]]
  313. ),
  314. builder.node_type(expr),
  315. expr.line,
  316. expr.arg_kinds,
  317. expr.arg_names,
  318. )
  319. else:
  320. return builder.call_refexpr_with_args(
  321. expr,
  322. callee,
  323. (
  324. [translate_list_comprehension(builder, expr.args[0])]
  325. + [builder.accept(arg) for arg in expr.args[1:]]
  326. ),
  327. )
  328. return None
  329. @specialize_function("builtins.any")
  330. def translate_any_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  331. if (
  332. len(expr.args) == 1
  333. and expr.arg_kinds == [ARG_POS]
  334. and isinstance(expr.args[0], GeneratorExpr)
  335. ):
  336. return any_all_helper(builder, expr.args[0], builder.false, lambda x: x, builder.true)
  337. return None
  338. @specialize_function("builtins.all")
  339. def translate_all_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  340. if (
  341. len(expr.args) == 1
  342. and expr.arg_kinds == [ARG_POS]
  343. and isinstance(expr.args[0], GeneratorExpr)
  344. ):
  345. return any_all_helper(
  346. builder,
  347. expr.args[0],
  348. builder.true,
  349. lambda x: builder.unary_op(x, "not", expr.line),
  350. builder.false,
  351. )
  352. return None
  353. def any_all_helper(
  354. builder: IRBuilder,
  355. gen: GeneratorExpr,
  356. initial_value: Callable[[], Value],
  357. modify: Callable[[Value], Value],
  358. new_value: Callable[[], Value],
  359. ) -> Value:
  360. retval = Register(bool_rprimitive)
  361. builder.assign(retval, initial_value(), -1)
  362. loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
  363. true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock()
  364. def gen_inner_stmts() -> None:
  365. comparison = modify(builder.accept(gen.left_expr))
  366. builder.add_bool_branch(comparison, true_block, false_block)
  367. builder.activate_block(true_block)
  368. builder.assign(retval, new_value(), -1)
  369. builder.goto(exit_block)
  370. builder.activate_block(false_block)
  371. comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
  372. builder.goto_and_activate(exit_block)
  373. return retval
  374. @specialize_function("builtins.sum")
  375. def translate_sum_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  376. # specialized implementation is used if:
  377. # - only one or two arguments given (if not, sum() has been given invalid arguments)
  378. # - first argument is a Generator (there is no benefit to optimizing the performance of eg.
  379. # sum([1, 2, 3]), so non-Generator Iterables are not handled)
  380. if not (
  381. len(expr.args) in (1, 2)
  382. and expr.arg_kinds[0] == ARG_POS
  383. and isinstance(expr.args[0], GeneratorExpr)
  384. ):
  385. return None
  386. # handle 'start' argument, if given
  387. if len(expr.args) == 2:
  388. # ensure call to sum() was properly constructed
  389. if not expr.arg_kinds[1] in (ARG_POS, ARG_NAMED):
  390. return None
  391. start_expr = expr.args[1]
  392. else:
  393. start_expr = IntExpr(0)
  394. gen_expr = expr.args[0]
  395. target_type = builder.node_type(expr)
  396. retval = Register(target_type)
  397. builder.assign(retval, builder.coerce(builder.accept(start_expr), target_type, -1), -1)
  398. def gen_inner_stmts() -> None:
  399. call_expr = builder.accept(gen_expr.left_expr)
  400. builder.assign(retval, builder.binary_op(retval, call_expr, "+", -1), -1)
  401. loop_params = list(
  402. zip(gen_expr.indices, gen_expr.sequences, gen_expr.condlists, gen_expr.is_async)
  403. )
  404. comprehension_helper(builder, loop_params, gen_inner_stmts, gen_expr.line)
  405. return retval
  406. @specialize_function("dataclasses.field")
  407. @specialize_function("attr.ib")
  408. @specialize_function("attr.attrib")
  409. @specialize_function("attr.Factory")
  410. def translate_dataclasses_field_call(
  411. builder: IRBuilder, expr: CallExpr, callee: RefExpr
  412. ) -> Value | None:
  413. """Special case for 'dataclasses.field', 'attr.attrib', and 'attr.Factory'
  414. function calls because the results of such calls are type-checked
  415. by mypy using the types of the arguments to their respective
  416. functions, resulting in attempted coercions by mypyc that throw a
  417. runtime error.
  418. """
  419. builder.types[expr] = AnyType(TypeOfAny.from_error)
  420. return None
  421. @specialize_function("builtins.next")
  422. def translate_next_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  423. """Special case for calling next() on a generator expression, an
  424. idiom that shows up some in mypy.
  425. For example, next(x for x in l if x.id == 12, None) will
  426. generate code that searches l for an element where x.id == 12
  427. and produce the first such object, or None if no such element
  428. exists.
  429. """
  430. if not (
  431. expr.arg_kinds in ([ARG_POS], [ARG_POS, ARG_POS])
  432. and isinstance(expr.args[0], GeneratorExpr)
  433. ):
  434. return None
  435. gen = expr.args[0]
  436. retval = Register(builder.node_type(expr))
  437. default_val = builder.accept(expr.args[1]) if len(expr.args) > 1 else None
  438. exit_block = BasicBlock()
  439. def gen_inner_stmts() -> None:
  440. # next takes the first element of the generator, so if
  441. # something gets produced, we are done.
  442. builder.assign(retval, builder.accept(gen.left_expr), gen.left_expr.line)
  443. builder.goto(exit_block)
  444. loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
  445. comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
  446. # Now we need the case for when nothing got hit. If there was
  447. # a default value, we produce it, and otherwise we raise
  448. # StopIteration.
  449. if default_val:
  450. builder.assign(retval, default_val, gen.left_expr.line)
  451. builder.goto(exit_block)
  452. else:
  453. builder.add(RaiseStandardError(RaiseStandardError.STOP_ITERATION, None, expr.line))
  454. builder.add(Unreachable())
  455. builder.activate_block(exit_block)
  456. return retval
  457. @specialize_function("builtins.isinstance")
  458. def translate_isinstance(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  459. """Special case for builtins.isinstance.
  460. Prevent coercions on the thing we are checking the instance of -
  461. there is no need to coerce something to a new type before checking
  462. what type it is, and the coercion could lead to bugs.
  463. """
  464. if (
  465. len(expr.args) == 2
  466. and expr.arg_kinds == [ARG_POS, ARG_POS]
  467. and isinstance(expr.args[1], (RefExpr, TupleExpr))
  468. ):
  469. builder.types[expr.args[0]] = AnyType(TypeOfAny.from_error)
  470. irs = builder.flatten_classes(expr.args[1])
  471. if irs is not None:
  472. can_borrow = all(
  473. ir.is_ext_class and not ir.inherits_python and not ir.allow_interpreted_subclasses
  474. for ir in irs
  475. )
  476. obj = builder.accept(expr.args[0], can_borrow=can_borrow)
  477. return builder.builder.isinstance_helper(obj, irs, expr.line)
  478. return None
  479. @specialize_function("setdefault", dict_rprimitive)
  480. def translate_dict_setdefault(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  481. """Special case for 'dict.setdefault' which would only construct
  482. default empty collection when needed.
  483. The dict_setdefault_spec_init_op checks whether the dict contains
  484. the key and would construct the empty collection only once.
  485. For example, this specializer works for the following cases:
  486. d.setdefault(key, set()).add(value)
  487. d.setdefault(key, []).append(value)
  488. d.setdefault(key, {})[inner_key] = inner_val
  489. """
  490. if (
  491. len(expr.args) == 2
  492. and expr.arg_kinds == [ARG_POS, ARG_POS]
  493. and isinstance(callee, MemberExpr)
  494. ):
  495. arg = expr.args[1]
  496. if isinstance(arg, ListExpr):
  497. if len(arg.items):
  498. return None
  499. data_type = Integer(1, c_int_rprimitive, expr.line)
  500. elif isinstance(arg, DictExpr):
  501. if len(arg.items):
  502. return None
  503. data_type = Integer(2, c_int_rprimitive, expr.line)
  504. elif (
  505. isinstance(arg, CallExpr)
  506. and isinstance(arg.callee, NameExpr)
  507. and arg.callee.fullname == "builtins.set"
  508. ):
  509. if len(arg.args):
  510. return None
  511. data_type = Integer(3, c_int_rprimitive, expr.line)
  512. else:
  513. return None
  514. callee_dict = builder.accept(callee.expr)
  515. key_val = builder.accept(expr.args[0])
  516. return builder.call_c(
  517. dict_setdefault_spec_init_op, [callee_dict, key_val, data_type], expr.line
  518. )
  519. return None
  520. @specialize_function("format", str_rprimitive)
  521. def translate_str_format(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  522. if (
  523. isinstance(callee, MemberExpr)
  524. and isinstance(callee.expr, StrExpr)
  525. and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds)
  526. ):
  527. format_str = callee.expr.value
  528. tokens = tokenizer_format_call(format_str)
  529. if tokens is None:
  530. return None
  531. literals, format_ops = tokens
  532. # Convert variables to strings
  533. substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
  534. if substitutions is None:
  535. return None
  536. return join_formatted_strings(builder, literals, substitutions, expr.line)
  537. return None
  538. @specialize_function("join", str_rprimitive)
  539. def translate_fstring(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  540. """Special case for f-string, which is translated into str.join()
  541. in mypy AST.
  542. This specializer optimizes simplest f-strings which don't contain
  543. any format operation.
  544. """
  545. if (
  546. isinstance(callee, MemberExpr)
  547. and isinstance(callee.expr, StrExpr)
  548. and callee.expr.value == ""
  549. and expr.arg_kinds == [ARG_POS]
  550. and isinstance(expr.args[0], ListExpr)
  551. ):
  552. for item in expr.args[0].items:
  553. if isinstance(item, StrExpr):
  554. continue
  555. elif isinstance(item, CallExpr):
  556. if not isinstance(item.callee, MemberExpr) or item.callee.name != "format":
  557. return None
  558. elif (
  559. not isinstance(item.callee.expr, StrExpr) or item.callee.expr.value != "{:{}}"
  560. ):
  561. return None
  562. if not isinstance(item.args[1], StrExpr) or item.args[1].value != "":
  563. return None
  564. else:
  565. return None
  566. format_ops = []
  567. exprs: list[Expression] = []
  568. for item in expr.args[0].items:
  569. if isinstance(item, StrExpr) and item.value != "":
  570. format_ops.append(FormatOp.STR)
  571. exprs.append(item)
  572. elif isinstance(item, CallExpr):
  573. format_ops.append(FormatOp.STR)
  574. exprs.append(item.args[0])
  575. substitutions = convert_format_expr_to_str(builder, format_ops, exprs, expr.line)
  576. if substitutions is None:
  577. return None
  578. return join_formatted_strings(builder, None, substitutions, expr.line)
  579. return None
  580. @specialize_function("mypy_extensions.i64")
  581. def translate_i64(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  582. if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
  583. return None
  584. arg = expr.args[0]
  585. arg_type = builder.node_type(arg)
  586. if is_int64_rprimitive(arg_type):
  587. return builder.accept(arg)
  588. elif is_int32_rprimitive(arg_type):
  589. val = builder.accept(arg)
  590. return builder.add(Extend(val, int64_rprimitive, signed=True, line=expr.line))
  591. elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
  592. val = builder.accept(arg)
  593. return builder.coerce(val, int64_rprimitive, expr.line)
  594. return None
  595. @specialize_function("mypy_extensions.i32")
  596. def translate_i32(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  597. if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
  598. return None
  599. arg = expr.args[0]
  600. arg_type = builder.node_type(arg)
  601. if is_int32_rprimitive(arg_type):
  602. return builder.accept(arg)
  603. elif is_int64_rprimitive(arg_type):
  604. val = builder.accept(arg)
  605. return builder.add(Truncate(val, int32_rprimitive, line=expr.line))
  606. elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
  607. val = builder.accept(arg)
  608. return builder.coerce(val, int32_rprimitive, expr.line)
  609. return None
  610. @specialize_function("builtins.int")
  611. def translate_int(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  612. if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
  613. return None
  614. arg = expr.args[0]
  615. arg_type = builder.node_type(arg)
  616. if (
  617. is_bool_rprimitive(arg_type)
  618. or is_int_rprimitive(arg_type)
  619. or is_fixed_width_rtype(arg_type)
  620. ):
  621. src = builder.accept(arg)
  622. return builder.coerce(src, int_rprimitive, expr.line)
  623. return None
  624. @specialize_function("builtins.bool")
  625. def translate_bool(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
  626. if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
  627. return None
  628. arg = expr.args[0]
  629. src = builder.accept(arg)
  630. return builder.builder.bool_value(src)
  631. @specialize_function("builtins.float")
  632. def translate_float(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
  633. if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
  634. return None
  635. arg = expr.args[0]
  636. arg_type = builder.node_type(arg)
  637. if is_float_rprimitive(arg_type):
  638. # No-op float conversion.
  639. return builder.accept(arg)
  640. return None