mypyc/lower/int_ops.py

"""Convert tagged int primitive ops to lower-level ops."""

from __future__ import annotations

from typing import NamedTuple

from mypyc.ir.ops import Assign, BasicBlock, Branch, ComparisonOp, Register, Value
from mypyc.ir.rtypes import bool_rprimitive, is_short_int_rprimitive
from mypyc.irbuild.ll_builder import LowLevelIRBuilder
from mypyc.lower.registry import lower_primitive_op
from mypyc.primitives.int_ops import int_equal_, int_less_than_
from mypyc.primitives.registry import CFunctionDescription


# Description for building int comparison ops
#
# Fields:
#   binary_op_variant: identify which IntOp to use when operands are short integers
#   c_func_description: the C function to call when operands are tagged integers
#   c_func_negated: whether to negate the C function call's result
#   c_func_swap_operands: whether to swap lhs and rhs when call the function
class IntComparisonOpDescription(NamedTuple):
    binary_op_variant: int
    c_func_description: CFunctionDescription
    c_func_negated: bool
    c_func_swap_operands: bool


# Provide mapping from textual op to short int's op variant and boxed int's description.
# Note that these are not complete implementations and require extra IR.
int_comparison_op_mapping: dict[str, IntComparisonOpDescription] = {
    "==": IntComparisonOpDescription(ComparisonOp.EQ, int_equal_, False, False),
    "!=": IntComparisonOpDescription(ComparisonOp.NEQ, int_equal_, True, False),
    "<": IntComparisonOpDescription(ComparisonOp.SLT, int_less_than_, False, False),
    "<=": IntComparisonOpDescription(ComparisonOp.SLE, int_less_than_, True, True),
    ">": IntComparisonOpDescription(ComparisonOp.SGT, int_less_than_, False, True),
    ">=": IntComparisonOpDescription(ComparisonOp.SGE, int_less_than_, True, False),
}


def compare_tagged(self: LowLevelIRBuilder, lhs: Value, rhs: Value, op: str, line: int) -> Value:
    """Compare two tagged integers using given operator (value context)."""
    # generate fast binary logic ops on short ints
    if (is_short_int_rprimitive(lhs.type) or is_short_int_rprimitive(rhs.type)) and op in (
        "==",
        "!=",
    ):
        quick = True
    else:
        quick = is_short_int_rprimitive(lhs.type) and is_short_int_rprimitive(rhs.type)
    if quick:
        return self.comparison_op(lhs, rhs, int_comparison_op_mapping[op][0], line)
    op_type, c_func_desc, negate_result, swap_op = int_comparison_op_mapping[op]
    result = Register(bool_rprimitive)
    short_int_block, int_block, out = BasicBlock(), BasicBlock(), BasicBlock()
    check_lhs = self.check_tagged_short_int(lhs, line, negated=True)
    if op in ("==", "!="):
        self.add(Branch(check_lhs, int_block, short_int_block, Branch.BOOL))
    else:
        # for non-equality logical ops (less/greater than, etc.), need to check both sides
        short_lhs = BasicBlock()
        self.add(Branch(check_lhs, int_block, short_lhs, Branch.BOOL))
        self.activate_block(short_lhs)
        check_rhs = self.check_tagged_short_int(rhs, line, negated=True)
        self.add(Branch(check_rhs, int_block, short_int_block, Branch.BOOL))
    self.activate_block(int_block)
    if swap_op:
        args = [rhs, lhs]
    else:
        args = [lhs, rhs]
    call = self.call_c(c_func_desc, args, line)
    if negate_result:
        # TODO: introduce UnaryIntOp?
        call_result = self.unary_op(call, "not", line)
    else:
        call_result = call
    self.add(Assign(result, call_result, line))
    self.goto(out)
    self.activate_block(short_int_block)
    eq = self.comparison_op(lhs, rhs, op_type, line)
    self.add(Assign(result, eq, line))
    self.goto_and_activate(out)
    return result


@lower_primitive_op("int_eq")
def lower_int_eq(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
    return compare_tagged(builder, args[0], args[1], "==", line)


@lower_primitive_op("int_ne")
def lower_int_ne(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
    return compare_tagged(builder, args[0], args[1], "!=", line)


@lower_primitive_op("int_lt")
def lower_int_lt(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
    return compare_tagged(builder, args[0], args[1], "<", line)


@lower_primitive_op("int_le")
def lower_int_le(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
    return compare_tagged(builder, args[0], args[1], "<=", line)


@lower_primitive_op("int_gt")
def lower_int_gt(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
    return compare_tagged(builder, args[0], args[1], ">", line)


@lower_primitive_op("int_ge")
def lower_int_ge(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
    return compare_tagged(builder, args[0], args[1], ">=", line)