mypy-dev 1.9.0.dev0__py3-none-any.whl

mypyc/irbuild/specialize.py

@@ -0,0 +1,822 @@
+"""Special case IR generation of calls to specific builtin functions.
+
+Most special cases should be handled using the data driven "primitive
+ops" system, but certain operations require special handling that has
+access to the AST/IR directly and can make decisions/optimizations
+based on it. These special cases can be implemented here.
+
+For example, we use specializers to statically emit the length of a
+fixed length tuple and to emit optimized code for any()/all() calls with
+generator comprehensions as the argument.
+
+See comment below for more documentation.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+from mypy.nodes import (
+    ARG_NAMED,
+    ARG_POS,
+    CallExpr,
+    DictExpr,
+    Expression,
+    GeneratorExpr,
+    IntExpr,
+    ListExpr,
+    MemberExpr,
+    NameExpr,
+    RefExpr,
+    StrExpr,
+    TupleExpr,
+)
+from mypy.types import AnyType, TypeOfAny
+from mypyc.ir.ops import (
+    BasicBlock,
+    Extend,
+    Integer,
+    RaiseStandardError,
+    Register,
+    Truncate,
+    Unreachable,
+    Value,
+)
+from mypyc.ir.rtypes import (
+    RInstance,
+    RPrimitive,
+    RTuple,
+    RType,
+    bool_rprimitive,
+    c_int_rprimitive,
+    dict_rprimitive,
+    int16_rprimitive,
+    int32_rprimitive,
+    int64_rprimitive,
+    int_rprimitive,
+    is_bool_rprimitive,
+    is_dict_rprimitive,
+    is_fixed_width_rtype,
+    is_float_rprimitive,
+    is_int16_rprimitive,
+    is_int32_rprimitive,
+    is_int64_rprimitive,
+    is_int_rprimitive,
+    is_list_rprimitive,
+    is_uint8_rprimitive,
+    list_rprimitive,
+    set_rprimitive,
+    str_rprimitive,
+    uint8_rprimitive,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.for_helpers import (
+    comprehension_helper,
+    sequence_from_generator_preallocate_helper,
+    translate_list_comprehension,
+    translate_set_comprehension,
+)
+from mypyc.irbuild.format_str_tokenizer import (
+    FormatOp,
+    convert_format_expr_to_str,
+    join_formatted_strings,
+    tokenizer_format_call,
+)
+from mypyc.primitives.dict_ops import (
+    dict_items_op,
+    dict_keys_op,
+    dict_setdefault_spec_init_op,
+    dict_values_op,
+)
+from mypyc.primitives.list_ops import new_list_set_item_op
+from mypyc.primitives.tuple_ops import new_tuple_set_item_op
+
+# Specializers are attempted before compiling the arguments to the
+# function.  Specializers can return None to indicate that they failed
+# and the call should be compiled normally. Otherwise they should emit
+# code for the call and return a Value containing the result.
+#
+# Specializers take three arguments: the IRBuilder, the CallExpr being
+# compiled, and the RefExpr that is the left hand side of the call.
+Specializer = Callable[["IRBuilder", CallExpr, RefExpr], Optional[Value]]
+
+# Dictionary containing all configured specializers.
+#
+# Specializers can operate on methods as well, and are keyed on the
+# name and RType in that case.
+specializers: dict[tuple[str, RType | None], list[Specializer]] = {}
+
+
+def _apply_specialization(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr, name: str | None, typ: RType | None = None
+) -> Value | None:
+    # TODO: Allow special cases to have default args or named args. Currently they don't since
+    #       they check that everything in arg_kinds is ARG_POS.
+
+    # If there is a specializer for this function, try calling it.
+    # Return the first successful one.
+    if name and (name, typ) in specializers:
+        for specializer in specializers[name, typ]:
+            val = specializer(builder, expr, callee)
+            if val is not None:
+                return val
+    return None
+
+
+def apply_function_specialization(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Invoke the Specializer callback for a function if one has been registered"""
+    return _apply_specialization(builder, expr, callee, callee.fullname)
+
+
+def apply_method_specialization(
+    builder: IRBuilder, expr: CallExpr, callee: MemberExpr, typ: RType | None = None
+) -> Value | None:
+    """Invoke the Specializer callback for a method if one has been registered"""
+    name = callee.fullname if typ is None else callee.name
+    return _apply_specialization(builder, expr, callee, name, typ)
+
+
+def specialize_function(
+    name: str, typ: RType | None = None
+) -> Callable[[Specializer], Specializer]:
+    """Decorator to register a function as being a specializer.
+
+    There may exist multiple specializers for one function. When
+    translating method calls, the earlier appended specializer has
+    higher priority.
+    """
+
+    def wrapper(f: Specializer) -> Specializer:
+        specializers.setdefault((name, typ), []).append(f)
+        return f
+
+    return wrapper
+
+
+@specialize_function("builtins.globals")
+def translate_globals(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) == 0:
+        return builder.load_globals_dict()
+    return None
+
+
+@specialize_function("builtins.abs")
+@specialize_function("builtins.int")
+@specialize_function("builtins.float")
+@specialize_function("builtins.complex")
+@specialize_function("mypy_extensions.i64")
+@specialize_function("mypy_extensions.i32")
+@specialize_function("mypy_extensions.i16")
+@specialize_function("mypy_extensions.u8")
+def translate_builtins_with_unary_dunder(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Specialize calls on native classes that implement the associated dunder.
+
+    E.g. i64(x) gets specialized to x.__int__() if x is a native instance.
+    """
+    if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and isinstance(callee, NameExpr):
+        arg = expr.args[0]
+        arg_typ = builder.node_type(arg)
+        shortname = callee.fullname.split(".")[1]
+        if shortname in ("i64", "i32", "i16", "u8"):
+            method = "__int__"
+        else:
+            method = f"__{shortname}__"
+        if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method(method):
+            obj = builder.accept(arg)
+            return builder.gen_method_call(obj, method, [], None, expr.line)
+
+    return None
+
+
+@specialize_function("builtins.len")
+def translate_len(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
+        arg = expr.args[0]
+        expr_rtype = builder.node_type(arg)
+        if isinstance(expr_rtype, RTuple):
+            # len() of fixed-length tuple can be trivially determined
+            # statically, though we still need to evaluate it.
+            builder.accept(arg)
+            return Integer(len(expr_rtype.types))
+        else:
+            if is_list_rprimitive(builder.node_type(arg)):
+                borrow = True
+            else:
+                borrow = False
+            obj = builder.accept(arg, can_borrow=borrow)
+            return builder.builtin_len(obj, expr.line)
+    return None
+
+
+@specialize_function("builtins.list")
+def dict_methods_fast_path(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Specialize a common case when list() is called on a dictionary
+    view method call.
+
+    For example:
+        foo = list(bar.keys())
+    """
+    if not (len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]):
+        return None
+    arg = expr.args[0]
+    if not (isinstance(arg, CallExpr) and not arg.args and isinstance(arg.callee, MemberExpr)):
+        return None
+    base = arg.callee.expr
+    attr = arg.callee.name
+    rtype = builder.node_type(base)
+    if not (is_dict_rprimitive(rtype) and attr in ("keys", "values", "items")):
+        return None
+
+    obj = builder.accept(base)
+    # Note that it is not safe to use fast methods on dict subclasses,
+    # so the corresponding helpers in CPy.h fallback to (inlined)
+    # generic logic.
+    if attr == "keys":
+        return builder.call_c(dict_keys_op, [obj], expr.line)
+    elif attr == "values":
+        return builder.call_c(dict_values_op, [obj], expr.line)
+    else:
+        return builder.call_c(dict_items_op, [obj], expr.line)
+
+
+@specialize_function("builtins.list")
+def translate_list_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for simplest list comprehension.
+
+    For example:
+        list(f(x) for x in some_list/some_tuple/some_str)
+    'translate_list_comprehension()' would take care of other cases
+    if this fails.
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return sequence_from_generator_preallocate_helper(
+            builder,
+            expr.args[0],
+            empty_op_llbuilder=builder.builder.new_list_op_with_length,
+            set_item_op=new_list_set_item_op,
+        )
+    return None
+
+
+@specialize_function("builtins.tuple")
+def translate_tuple_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for simplest tuple creation from a generator.
+
+    For example:
+        tuple(f(x) for x in some_list/some_tuple/some_str)
+    'translate_safe_generator_call()' would take care of other cases
+    if this fails.
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return sequence_from_generator_preallocate_helper(
+            builder,
+            expr.args[0],
+            empty_op_llbuilder=builder.builder.new_tuple_with_length,
+            set_item_op=new_tuple_set_item_op,
+        )
+    return None
+
+
+@specialize_function("builtins.set")
+def translate_set_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for set creation from a generator.
+
+    For example:
+        set(f(...) for ... in iterator/nested_generators...)
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return translate_set_comprehension(builder, expr.args[0])
+    return None
+
+
+@specialize_function("builtins.min")
+@specialize_function("builtins.max")
+def faster_min_max(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if expr.arg_kinds == [ARG_POS, ARG_POS]:
+        x, y = builder.accept(expr.args[0]), builder.accept(expr.args[1])
+        result = Register(builder.node_type(expr))
+        # CPython evaluates arguments reversely when calling min(...) or max(...)
+        if callee.fullname == "builtins.min":
+            comparison = builder.binary_op(y, x, "<", expr.line)
+        else:
+            comparison = builder.binary_op(y, x, ">", expr.line)
+
+        true_block, false_block, next_block = BasicBlock(), BasicBlock(), BasicBlock()
+        builder.add_bool_branch(comparison, true_block, false_block)
+
+        builder.activate_block(true_block)
+        builder.assign(result, builder.coerce(y, result.type, expr.line), expr.line)
+        builder.goto(next_block)
+
+        builder.activate_block(false_block)
+        builder.assign(result, builder.coerce(x, result.type, expr.line), expr.line)
+        builder.goto(next_block)
+
+        builder.activate_block(next_block)
+        return result
+    return None
+
+
+@specialize_function("builtins.tuple")
+@specialize_function("builtins.frozenset")
+@specialize_function("builtins.dict")
+@specialize_function("builtins.min")
+@specialize_function("builtins.max")
+@specialize_function("builtins.sorted")
+@specialize_function("collections.OrderedDict")
+@specialize_function("join", str_rprimitive)
+@specialize_function("extend", list_rprimitive)
+@specialize_function("update", dict_rprimitive)
+@specialize_function("update", set_rprimitive)
+def translate_safe_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special cases for things that consume iterators where we know we
+    can safely compile a generator into a list.
+    """
+    if (
+        len(expr.args) > 0
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        if isinstance(callee, MemberExpr):
+            return builder.gen_method_call(
+                builder.accept(callee.expr),
+                callee.name,
+                (
+                    [translate_list_comprehension(builder, expr.args[0])]
+                    + [builder.accept(arg) for arg in expr.args[1:]]
+                ),
+                builder.node_type(expr),
+                expr.line,
+                expr.arg_kinds,
+                expr.arg_names,
+            )
+        else:
+            return builder.call_refexpr_with_args(
+                expr,
+                callee,
+                (
+                    [translate_list_comprehension(builder, expr.args[0])]
+                    + [builder.accept(arg) for arg in expr.args[1:]]
+                ),
+            )
+    return None
+
+
+@specialize_function("builtins.any")
+def translate_any_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds == [ARG_POS]
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return any_all_helper(builder, expr.args[0], builder.false, lambda x: x, builder.true)
+    return None
+
+
+@specialize_function("builtins.all")
+def translate_all_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds == [ARG_POS]
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return any_all_helper(
+            builder,
+            expr.args[0],
+            builder.true,
+            lambda x: builder.unary_op(x, "not", expr.line),
+            builder.false,
+        )
+    return None
+
+
+def any_all_helper(
+    builder: IRBuilder,
+    gen: GeneratorExpr,
+    initial_value: Callable[[], Value],
+    modify: Callable[[Value], Value],
+    new_value: Callable[[], Value],
+) -> Value:
+    retval = Register(bool_rprimitive)
+    builder.assign(retval, initial_value(), -1)
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+    true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock()
+
+    def gen_inner_stmts() -> None:
+        comparison = modify(builder.accept(gen.left_expr))
+        builder.add_bool_branch(comparison, true_block, false_block)
+        builder.activate_block(true_block)
+        builder.assign(retval, new_value(), -1)
+        builder.goto(exit_block)
+        builder.activate_block(false_block)
+
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    builder.goto_and_activate(exit_block)
+
+    return retval
+
+
+@specialize_function("builtins.sum")
+def translate_sum_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    # specialized implementation is used if:
+    # - only one or two arguments given (if not, sum() has been given invalid arguments)
+    # - first argument is a Generator (there is no benefit to optimizing the performance of eg.
+    #   sum([1, 2, 3]), so non-Generator Iterables are not handled)
+    if not (
+        len(expr.args) in (1, 2)
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return None
+
+    # handle 'start' argument, if given
+    if len(expr.args) == 2:
+        # ensure call to sum() was properly constructed
+        if expr.arg_kinds[1] not in (ARG_POS, ARG_NAMED):
+            return None
+        start_expr = expr.args[1]
+    else:
+        start_expr = IntExpr(0)
+
+    gen_expr = expr.args[0]
+    target_type = builder.node_type(expr)
+    retval = Register(target_type)
+    builder.assign(retval, builder.coerce(builder.accept(start_expr), target_type, -1), -1)
+
+    def gen_inner_stmts() -> None:
+        call_expr = builder.accept(gen_expr.left_expr)
+        builder.assign(retval, builder.binary_op(retval, call_expr, "+", -1), -1)
+
+    loop_params = list(
+        zip(gen_expr.indices, gen_expr.sequences, gen_expr.condlists, gen_expr.is_async)
+    )
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen_expr.line)
+
+    return retval
+
+
+@specialize_function("dataclasses.field")
+@specialize_function("attr.ib")
+@specialize_function("attr.attrib")
+@specialize_function("attr.Factory")
+def translate_dataclasses_field_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for 'dataclasses.field', 'attr.attrib', and 'attr.Factory'
+    function calls because the results of such calls are type-checked
+    by mypy using the types of the arguments to their respective
+    functions, resulting in attempted coercions by mypyc that throw a
+    runtime error.
+    """
+    builder.types[expr] = AnyType(TypeOfAny.from_error)
+    return None
+
+
+@specialize_function("builtins.next")
+def translate_next_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for calling next() on a generator expression, an
+    idiom that shows up some in mypy.
+
+    For example, next(x for x in l if x.id == 12, None) will
+    generate code that searches l for an element where x.id == 12
+    and produce the first such object, or None if no such element
+    exists.
+    """
+    if not (
+        expr.arg_kinds in ([ARG_POS], [ARG_POS, ARG_POS])
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return None
+
+    gen = expr.args[0]
+    retval = Register(builder.node_type(expr))
+    default_val = builder.accept(expr.args[1]) if len(expr.args) > 1 else None
+    exit_block = BasicBlock()
+
+    def gen_inner_stmts() -> None:
+        # next takes the first element of the generator, so if
+        # something gets produced, we are done.
+        builder.assign(retval, builder.accept(gen.left_expr), gen.left_expr.line)
+        builder.goto(exit_block)
+
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+
+    # Now we need the case for when nothing got hit. If there was
+    # a default value, we produce it, and otherwise we raise
+    # StopIteration.
+    if default_val:
+        builder.assign(retval, default_val, gen.left_expr.line)
+        builder.goto(exit_block)
+    else:
+        builder.add(RaiseStandardError(RaiseStandardError.STOP_ITERATION, None, expr.line))
+        builder.add(Unreachable())
+
+    builder.activate_block(exit_block)
+    return retval
+
+
+@specialize_function("builtins.isinstance")
+def translate_isinstance(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for builtins.isinstance.
+
+    Prevent coercions on the thing we are checking the instance of -
+    there is no need to coerce something to a new type before checking
+    what type it is, and the coercion could lead to bugs.
+    """
+    if (
+        len(expr.args) == 2
+        and expr.arg_kinds == [ARG_POS, ARG_POS]
+        and isinstance(expr.args[1], (RefExpr, TupleExpr))
+    ):
+        builder.types[expr.args[0]] = AnyType(TypeOfAny.from_error)
+
+        irs = builder.flatten_classes(expr.args[1])
+        if irs is not None:
+            can_borrow = all(
+                ir.is_ext_class and not ir.inherits_python and not ir.allow_interpreted_subclasses
+                for ir in irs
+            )
+            obj = builder.accept(expr.args[0], can_borrow=can_borrow)
+            return builder.builder.isinstance_helper(obj, irs, expr.line)
+    return None
+
+
+@specialize_function("setdefault", dict_rprimitive)
+def translate_dict_setdefault(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for 'dict.setdefault' which would only construct
+    default empty collection when needed.
+
+    The dict_setdefault_spec_init_op checks whether the dict contains
+    the key and would construct the empty collection only once.
+
+    For example, this specializer works for the following cases:
+         d.setdefault(key, set()).add(value)
+         d.setdefault(key, []).append(value)
+         d.setdefault(key, {})[inner_key] = inner_val
+    """
+    if (
+        len(expr.args) == 2
+        and expr.arg_kinds == [ARG_POS, ARG_POS]
+        and isinstance(callee, MemberExpr)
+    ):
+        arg = expr.args[1]
+        if isinstance(arg, ListExpr):
+            if len(arg.items):
+                return None
+            data_type = Integer(1, c_int_rprimitive, expr.line)
+        elif isinstance(arg, DictExpr):
+            if len(arg.items):
+                return None
+            data_type = Integer(2, c_int_rprimitive, expr.line)
+        elif (
+            isinstance(arg, CallExpr)
+            and isinstance(arg.callee, NameExpr)
+            and arg.callee.fullname == "builtins.set"
+        ):
+            if len(arg.args):
+                return None
+            data_type = Integer(3, c_int_rprimitive, expr.line)
+        else:
+            return None
+
+        callee_dict = builder.accept(callee.expr)
+        key_val = builder.accept(expr.args[0])
+        return builder.call_c(
+            dict_setdefault_spec_init_op, [callee_dict, key_val, data_type], expr.line
+        )
+    return None
+
+
+@specialize_function("format", str_rprimitive)
+def translate_str_format(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if (
+        isinstance(callee, MemberExpr)
+        and isinstance(callee.expr, StrExpr)
+        and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds)
+    ):
+        format_str = callee.expr.value
+        tokens = tokenizer_format_call(format_str)
+        if tokens is None:
+            return None
+        literals, format_ops = tokens
+        # Convert variables to strings
+        substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
+        if substitutions is None:
+            return None
+        return join_formatted_strings(builder, literals, substitutions, expr.line)
+    return None
+
+
+@specialize_function("join", str_rprimitive)
+def translate_fstring(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for f-string, which is translated into str.join()
+    in mypy AST.
+
+    This specializer optimizes simplest f-strings which don't contain
+    any format operation.
+    """
+    if (
+        isinstance(callee, MemberExpr)
+        and isinstance(callee.expr, StrExpr)
+        and callee.expr.value == ""
+        and expr.arg_kinds == [ARG_POS]
+        and isinstance(expr.args[0], ListExpr)
+    ):
+        for item in expr.args[0].items:
+            if isinstance(item, StrExpr):
+                continue
+            elif isinstance(item, CallExpr):
+                if not isinstance(item.callee, MemberExpr) or item.callee.name != "format":
+                    return None
+                elif (
+                    not isinstance(item.callee.expr, StrExpr) or item.callee.expr.value != "{:{}}"
+                ):
+                    return None
+
+                if not isinstance(item.args[1], StrExpr) or item.args[1].value != "":
+                    return None
+            else:
+                return None
+
+        format_ops = []
+        exprs: list[Expression] = []
+
+        for item in expr.args[0].items:
+            if isinstance(item, StrExpr) and item.value != "":
+                format_ops.append(FormatOp.STR)
+                exprs.append(item)
+            elif isinstance(item, CallExpr):
+                format_ops.append(FormatOp.STR)
+                exprs.append(item.args[0])
+
+        substitutions = convert_format_expr_to_str(builder, format_ops, exprs, expr.line)
+        if substitutions is None:
+            return None
+
+        return join_formatted_strings(builder, None, substitutions, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i64")
+def translate_i64(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_int64_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif is_int32_rprimitive(arg_type) or is_int16_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int64_rprimitive, signed=True, line=expr.line))
+    elif is_uint8_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int64_rprimitive, signed=False, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.coerce(val, int64_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i32")
+def translate_i32(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_int32_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif is_int64_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Truncate(val, int32_rprimitive, line=expr.line))
+    elif is_int16_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int32_rprimitive, signed=True, line=expr.line))
+    elif is_uint8_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int32_rprimitive, signed=False, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        val = truncate_literal(val, int32_rprimitive)
+        return builder.coerce(val, int32_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i16")
+def translate_i16(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_int16_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif is_int32_rprimitive(arg_type) or is_int64_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Truncate(val, int16_rprimitive, line=expr.line))
+    elif is_uint8_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int16_rprimitive, signed=False, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        val = truncate_literal(val, int16_rprimitive)
+        return builder.coerce(val, int16_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.u8")
+def translate_u8(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_uint8_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif (
+        is_int16_rprimitive(arg_type)
+        or is_int32_rprimitive(arg_type)
+        or is_int64_rprimitive(arg_type)
+    ):
+        val = builder.accept(arg)
+        return builder.add(Truncate(val, uint8_rprimitive, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        val = truncate_literal(val, uint8_rprimitive)
+        return builder.coerce(val, uint8_rprimitive, expr.line)
+    return None
+
+
+def truncate_literal(value: Value, rtype: RPrimitive) -> Value:
+    """If value is an integer literal value, truncate it to given native int rtype.
+
+    For example, truncate 256 into 0 if rtype is u8.
+    """
+    if not isinstance(value, Integer):
+        return value  # Not a literal, nothing to do
+    x = value.numeric_value()
+    max_unsigned = (1 << (rtype.size * 8)) - 1
+    x = x & max_unsigned
+    if rtype.is_signed and x >= (max_unsigned + 1) // 2:
+        # Adjust to make it a negative value
+        x -= max_unsigned + 1
+    return Integer(x, rtype)
+
+
+@specialize_function("builtins.int")
+def translate_int(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if (
+        is_bool_rprimitive(arg_type)
+        or is_int_rprimitive(arg_type)
+        or is_fixed_width_rtype(arg_type)
+    ):
+        src = builder.accept(arg)
+        return builder.coerce(src, int_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("builtins.bool")
+def translate_bool(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    src = builder.accept(arg)
+    return builder.builder.bool_value(src)
+
+
+@specialize_function("builtins.float")
+def translate_float(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_float_rprimitive(arg_type):
+        # No-op float conversion.
+        return builder.accept(arg)
+    return None