如何解决Numba中遇到的Unsupported use of op_LOAD_CLOSURE错误？

Numba的@njit装饰器启用的nopython模式不支持嵌套函数的闭包捕获（即内部函数直接引用外部函数的局部变量，比如这里search引用count、res，以及使用nonlocal），这就是报错Unsupported use of op_LOAD_CLOSURE encountered的原因。

把嵌套的search函数改为显式传递所有状态参数，避免闭包捕获外部变量。具体调整如下：

• 将count、res、n作为参数传入search（numpy数组传递的是引用，修改会同步到原数组）；
• 用特殊值-1替代None作为previous的默认值（Numba对None的处理在递归函数中容易出问题）；
• 移除nonlocal声明，直接通过参数操作res。

from numba import njit
import numpy as np

@njit   
def search(count, res, n, sz=0, max_val=1, single=0, previous=-1):
    if sz == 4 * n:
        res[0] += 1
        return
    # 处理single分支
    if single and count[0] < 2 * n:
        count[0] += 1
        search(count, res, n, sz + 1, max_val, single, previous)
        count[0] -= 1
    # 处理循环分支
    for i in range(1, max_val + 1):
        if i != previous and count[i] < 2:
            count[i] += 1
            new_max_val = max_val + (1 if (i == max_val and max_val < n) else 0)
            new_single = single + (1 if count[i] == 1 else 0) - (1 if count[i] == 2 else 0)
            search(count, res, n, sz + 1, new_max_val, new_single, i)
            count[i] -= 1

@njit
def solve(n):
    count = np.zeros(n + 1, dtype=np.int64)  # 用int64避免溢出
    res = np.array([0], dtype=np.int64)
    search(count, res, n)
    return res[0]

for i in range(1, 6):
    print(solve(i))

• Numba支持递归函数，但要求参数类型明确，用-1替代None是为了让Numba能推断出previous的整数类型；
• 使用np.int64替代默认的int是为了避免递归计数时的整数溢出（尤其当n较大时）；
• 所有状态变量通过参数传递，完全避免了闭包捕获，符合Numba nopython模式的要求，同时保留了原算法的逻辑。

内容的提问来源于stack exchange，提问作者Simd