Numba加速位掩码唯一排序的陷阱：64位整数溢出与类型限制解析

import numpy as np
from time import perf_counter
from numba import njit

def count_unique_bitmask(ls):
    """
    使用位掩码对非负整数进行唯一排序。
    参数:
        ls: 包含非负整数的列表或数组。
    返回:
        已排序的唯一整数列表。
    """
    ret = []
    m = 0  # 初始化位掩码
    # 遍历输入列表，将对应位的设置为1
    for x in ls:
        # 确保x是整数且在合理范围内，这里假设x为非负
        m = m | (1 << int(x)) # 将第x位设置为1

    i = 0
    # 从最低位开始检查，提取出所有存在的整数
    while m > 0:
        if (m & 1): # 如果当前位是1
            ret.append(i)
        m = m >> 1 # 右移一位，检查下一位
        i += 1
    return ret

# 性能测试示例
RNG = np.random.default_rng(0)
x = RNG.integers(2**16, size=2**17) # 生成大量随机非负整数

print("--- 性能对比 ---")
start = perf_counter()
y1 = np.unique(x)
print(f"np.unique 耗时: {perf_counter() - start:.6f} 秒")

start = perf_counter()
y2 = count_unique_bitmask(x)
print(f"自定义位掩码函数 耗时: {perf_counter() - start:.6f} 秒")

print(f"结果是否一致: {(y1 == y2).all()}")

# ... (import numpy, perf_counter, njit)

@njit # 应用Numba JIT装饰器
def count_unique_bitmask_numba(ls):
    ret = []
    m = 0
    for x in ls:
        m = m | (1 << int(x))

    i = 0
    while m > 0: # 问题出现在这里
        if (m & 1):
            ret.append(i)
        m = m >> 1
        i += 1
    return ret

# 再次运行性能测试，这次使用Numba版本
print("\n--- Numba版本测试 ---")
start = perf_counter()
# 注意：这里为了演示问题，我们可能需要一个较小的输入范围，
# 或者一个包含较大整数的输入，以触发Numba的整数溢出行为。
# 对于原始问题中的 x = RNG.integers(2**16, size=2**17)，
# 2**16 (65536) 已经超过了64位有符号整数的正数最大值范围。
# 让我们用一个更小的范围来演示 Numba 行为，然后解释问题。
x_small = np.array([7, 7, 0, 3, 2, 1, 9, 1], dtype=np.int64) # 确保元素类型兼容
y_numba_small = count_unique_bitmask_numba(x_small)
print(f"Numba位掩码函数 (小范围) 结果: {y_numba_small}") # 预期：[0, 1, 2, 3, 7, 9]

# 假设我们有一个更大的整数，例如 63
x_large_val = np.array([0, 63], dtype=np.int64)
y_numba_large_val = count_unique_bitmask_numba(x_large_val)
print(f"Numba位掩码函数 (含63) 结果: {y_numba_large_val}") # 预期：[0, 63]

# 如果输入包含大于63的整数，且 Numba 默认使用 64 位整数，则可能出现问题。
# 对于原始的 x = RNG.integers(2**16, size=2**17)，最大值可能达到 2**16-1 = 65535。
# 此时 1 << 65535 将远远超出 64 位整数的表示范围。
# 即使是 1 << 63 也会有问题。

from numba import njit

@njit
def shift_test(amount):
    return 1 << amount

print("\n--- Numba 64位整数位移测试 ---")
for i in range(66):
    try:
        val = shift_test(i)
        print(f"shift( {i:2d} ): {hex(val)} (十进制: {val})")
        if i == 63:
            print(f"  注意: 当 i=63 时，结果变为负数，因为最高位被解释为符号位。")
    except Exception as e:
        print(f"shift( {i:2d} ): 发生错误 - {e}")