# Copyright (C) 2022-2023 Exaloop Inc. <https://exaloop.io>

from internal.attributes import commutative, associative, distributive

def check_N(N: Static[int]):
    if N <= 0:
        compile_error("N must be greater than 0")
    pass

@extend
class Int:
    def __new__() -> Int[N]:
        check_N(N)
        return Int[N](0)

    def __new__(what: Int[N]) -> Int[N]:
        check_N(N)
        return what

    def __new__(what: int) -> Int[N]:
        check_N(N)
        if N < 64:
            return __internal__.int_trunc(what, 64, N)
        elif N == 64:
            return what
        else:
            return __internal__.int_sext(what, 64, N)

    @pure
    @llvm
    def __new__(what: UInt[N]) -> Int[N]:
        ret i{=N} %what

    def __new__(what: str) -> Int[N]:
        check_N(N)
        ret = Int[N]()
        i = 0
        sign = Int[N](1)
        if i < what.len and what.ptr[0] == byte(45):
            sign = Int[N](-1)
            i += 1
        while i < what.len:
            if what.ptr[i] < byte(48) or what.ptr[i] >= byte(58):
                raise ValueError("Invalid integer string")
            ret = ret * Int[N](10) + Int[N](int(what.ptr[i]) - 48)
            i += 1
        return sign * ret

    def __int__(self) -> int:
        if N > 64:
            return __internal__.int_trunc(self, N, 64)
        elif N == 64:
            return self
        else:
            return __internal__.int_sext(self, N, 64)

    def __index__(self) -> int:
        return int(self)

    def __copy__(self) -> Int[N]:
        return self

    def __deepcopy__(self) -> Int[N]:
        return self

    def __hash__(self) -> int:
        return int(self)

    @pure
    @llvm
    def __float__(self) -> float:
        %0 = sitofp i{=N} %self to double
        ret double %0

    @pure
    @llvm
    def __bool__(self) -> bool:
        %0 = icmp ne i{=N} %self, 0
        %1 = zext i1 %0 to i8
        ret i8 %1

    def __pos__(self) -> Int[N]:
        return self

    @pure
    @llvm
    def __neg__(self) -> Int[N]:
        %0 = sub i{=N} 0, %self
        ret i{=N} %0

    @pure
    @llvm
    def __invert__(self) -> Int[N]:
        %0 = xor i{=N} %self, -1
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @llvm
    def __add__(self, other: Int[N]) -> Int[N]:
        %0 = add i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def __sub__(self, other: Int[N]) -> Int[N]:
        %0 = sub i{=N} %self, %other
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @distributive
    @llvm
    def __mul__(self, other: Int[N]) -> Int[N]:
        %0 = mul i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def _floordiv(self, other: Int[N]) -> Int[N]:
        %0 = sdiv i{=N} %self, %other
        ret i{=N} %0
    def __floordiv__(self, other: Int[N]) -> Int[N]:
        if N > 128:
            compile_error("division is not supported on Int[N] when N > 128")
        return self._floordiv(other)

    @pure
    @llvm
    def __truediv__(self, other: Int[N]) -> float:
        %0 = sitofp i{=N} %self to double
        %1 = sitofp i{=N} %other to double
        %2 = fdiv double %0, %1
        ret double %2

    @pure
    @llvm
    def _mod(self, other: Int[N]) -> Int[N]:
        %0 = srem i{=N} %self, %other
        ret i{=N} %0
    def __mod__(self, other: Int[N]) -> Int[N]:
        if N > 128:
            compile_error("modulus is not supported on Int[N] when N > 128")
        return self._mod(other)

    def __divmod__(self, other: Int[N]) -> Tuple[Int[N], Int[N]]:
        d = self // other
        m = self - d * other
        if m and ((other ^ m) < Int[N](0)):
            m += other
            d -= Int[N](1)
        return (d, m)

    @pure
    @llvm
    def __lshift__(self, other: Int[N]) -> Int[N]:
        %0 = shl i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def __rshift__(self, other: Int[N]) -> Int[N]:
        %0 = ashr i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def __eq__(self, other: Int[N]) -> bool:
        %0 = icmp eq i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __ne__(self, other: Int[N]) -> bool:
        %0 = icmp ne i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __lt__(self, other: Int[N]) -> bool:
        %0 = icmp slt i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __gt__(self, other: Int[N]) -> bool:
        %0 = icmp sgt i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __le__(self, other: Int[N]) -> bool:
        %0 = icmp sle i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __ge__(self, other: Int[N]) -> bool:
        %0 = icmp sge i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @commutative
    @associative
    @llvm
    def __and__(self, other: Int[N]) -> Int[N]:
        %0 = and i{=N} %self, %other
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @llvm
    def __or__(self, other: Int[N]) -> Int[N]:
        %0 = or i{=N} %self, %other
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @llvm
    def __xor__(self, other: Int[N]) -> Int[N]:
        %0 = xor i{=N} %self, %other
        ret i{=N} %0

    def __repr__(self) -> str:
        return f"Int[{N}]({int(self)})"

    def __str__(self) -> str:
        return int(self).__repr__()

    @pure
    @llvm
    def _popcnt(self) -> Int[N]:
        declare i{=N} @llvm.ctpop.i{=N}(i{=N})
        %0 = call i{=N} @llvm.ctpop.i{=N}(i{=N} %self)
        ret i{=N} %0

    def popcnt(self) -> int:
        return int(self._popcnt())

    def len() -> int:
        return N


@extend
class UInt:
    def __new__() -> UInt[N]:
        check_N(N)
        return UInt[N](0)

    def __new__(what: UInt[N]) -> UInt[N]:
        check_N(N)
        return what

    def __new__(what: int) -> UInt[N]:
        check_N(N)
        if N < 64:
            return UInt[N](__internal__.int_trunc(what, 64, N))
        elif N == 64:
            return UInt[N](Int[N](what))
        else:
            return UInt[N](__internal__.int_sext(what, 64, N))

    @pure
    @llvm
    def __new__(what: Int[N]) -> UInt[N]:
        ret i{=N} %what

    def __new__(what: str) -> UInt[N]:
        check_N(N)
        return UInt[N](Int[N](what))

    def __int__(self) -> int:
        if N > 64:
            return __internal__.int_trunc(self, N, 64)
        elif N == 64:
            return Int[64](self)
        else:
            return __internal__.int_zext(self, N, 64)

    def __index__(self) -> int:
        return int(self)

    def __copy__(self) -> UInt[N]:
        return self

    def __deepcopy__(self) -> UInt[N]:
        return self

    def __hash__(self) -> int:
        return int(self)

    @pure
    @llvm
    def __float__(self) -> float:
        %0 = uitofp i{=N} %self to double
        ret double %0

    @pure
    @llvm
    def __bool__(self) -> bool:
        %0 = icmp ne i{=N} %self, 0
        %1 = zext i1 %0 to i8
        ret i8 %1

    def __pos__(self) -> UInt[N]:
        return self

    @pure
    @llvm
    def __neg__(self) -> UInt[N]:
        %0 = sub i{=N} 0, %self
        ret i{=N} %0

    @pure
    @llvm
    def __invert__(self) -> UInt[N]:
        %0 = xor i{=N} %self, -1
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @llvm
    def __add__(self, other: UInt[N]) -> UInt[N]:
        %0 = add i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def __sub__(self, other: UInt[N]) -> UInt[N]:
        %0 = sub i{=N} %self, %other
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @distributive
    @llvm
    def __mul__(self, other: UInt[N]) -> UInt[N]:
        %0 = mul i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def _floordiv(self, other: UInt[N]) -> UInt[N]:
        %0 = udiv i{=N} %self, %other
        ret i{=N} %0
    def __floordiv__(self, other: UInt[N]) -> UInt[N]:
        if N > 128:
            compile_error("division is not supported on UInt[N] when N > 128")
        return self._floordiv(other)

    @pure
    @llvm
    def __truediv__(self, other: UInt[N]) -> float:
        %0 = uitofp i{=N} %self to double
        %1 = uitofp i{=N} %other to double
        %2 = fdiv double %0, %1
        ret double %2

    @pure
    @llvm
    def _mod(self, other: UInt[N]) -> UInt[N]:
        %0 = urem i{=N} %self, %other
        ret i{=N} %0
    def __mod__(self, other: UInt[N]) -> UInt[N]:
        if N > 128:
            compile_error("modulus is not supported on UInt[N] when N > 128")
        return self._mod(other)

    def __divmod__(self, other: UInt[N]) -> Tuple[UInt[N], UInt[N]]:
        return (self // other, self % other)

    @pure
    @llvm
    def __lshift__(self, other: UInt[N]) -> UInt[N]:
        %0 = shl i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def __rshift__(self, other: UInt[N]) -> UInt[N]:
        %0 = lshr i{=N} %self, %other
        ret i{=N} %0

    @pure
    @llvm
    def __eq__(self, other: UInt[N]) -> bool:
        %0 = icmp eq i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __ne__(self, other: UInt[N]) -> bool:
        %0 = icmp ne i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __lt__(self, other: UInt[N]) -> bool:
        %0 = icmp ult i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __gt__(self, other: UInt[N]) -> bool:
        %0 = icmp ugt i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __le__(self, other: UInt[N]) -> bool:
        %0 = icmp ule i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @llvm
    def __ge__(self, other: UInt[N]) -> bool:
        %0 = icmp uge i{=N} %self, %other
        %1 = zext i1 %0 to i8
        ret i8 %1

    @pure
    @commutative
    @associative
    @llvm
    def __and__(self, other: UInt[N]) -> UInt[N]:
        %0 = and i{=N} %self, %other
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @llvm
    def __or__(self, other: UInt[N]) -> UInt[N]:
        %0 = or i{=N} %self, %other
        ret i{=N} %0

    @pure
    @commutative
    @associative
    @llvm
    def __xor__(self, other: UInt[N]) -> UInt[N]:
        %0 = xor i{=N} %self, %other
        ret i{=N} %0

    def __repr__(self) -> str:
        return f"UInt[{N}]({self.__str__()})"

    def __str__(self) -> str:
        return self.__format__("")

    def popcnt(self) -> int:
        return int(Int[N](self)._popcnt())

    def len() -> int:
        return N

i8 = Int[8]
i16 = Int[16]
i32 = Int[32]
i64 = Int[64]
u8 = UInt[8]
u16 = UInt[16]
u32 = UInt[32]
u64 = UInt[64]
