问题

我试图使以下代码中的 `A` 类型稳定。

```[Julia]
using Primes: factor

function f(n::T, p::T, k::T) where {T<:Integer}
    return rand(T, n * p^k)
end

function g(m::T, n::T) where {T<:Integer}

    i = 0
    for A in Iterators.product((f(n, p, T(k)) for (p, k) in factor(m))...)
        i = sum(A)
    end
    return i

end

请注意，f 是类型稳定的。变量 A 不是类型稳定的，因为 product 迭代器将根据 n 和 m 的值返回不同大小的元组。如果有一个类似 product 迭代器的迭代器，返回的是一个 Vector 而不是一个 Tuple，我认为类型不稳定性会消失。

是否有人有任何建议，使上面的代码中的 A 类型稳定？

编辑：我应该补充说，f 返回一个类型为 T 的变量大小的 Vector。

解决类型稳定性的一种方法是这样的。

function g(m::T, n::T) where {T<:Integer}
    B = Vector{T}[T[]]
    for (p, k) in factor(m)
        C = Vector{T}[]
        for (b, r) in Iterators.product(B, f(n, p, T(k)))
            c = copy(b)
            push!(c, r)
            push!(C, c)
        end
        B = C
    end

    for A in B
        i = sum(A)
    end

    return i
end

这个方法（特别是 A）现在是类型稳定的，但要付出大量的内存代价。我不确定是否有更好的方法来做这个。

<details>
<summary>英文:</summary>

I am trying to make `A` in the following code type-stable.

```[Julia]
using Primes: factor

function f(n::T, p::T, k::T) where {T&lt;:Integer}
    return rand(T, n * p^k)
end

function g(m::T, n::T) where {T&lt;:Integer}
    
    i = 0
    for A in Iterators.product((f(n, p, T(k)) for (p, k) in factor(m))...)
        i = sum(A)
    end
    return i

end

Note that f is type-stable. The variable A is not type-stable because the product iterator will return different sized tuples depending on the values of n and m. If there was an iterator like the product iterator that returned a Vector instead of a Tuple, I believe that the type-instability would go away.

Does anyone have any suggestions to make A type-stable in the above code?

Edit: I should add that f returns a variable-sized Vector of type T.

One way I have solved the type-stability is by doing this.

function g(m::T, n::T) where {T<:Integer}
    B = Vector{T}[T[]]
    for (p, k) in factor(m)
        C = Vector{T}[]
        for (b, r) in Iterators.product(B, f(n, p, T(k)))
            c = copy(b)
            push!(c, r)
            push!(C, c)
        end
        B = C
    end

    for A in B
        i = sum(A)
    end

    return i
end

This (and in particular, A) is now type-stable, but at the cost lots of memory. I'm not sure of a better way to do this.

答案1

得分: 2

这不容易完全保持类型稳定，但你可以使用一个函数屏障来隔离类型不稳定性。将因子分解转换为一个元组，传递给一个类型稳定的内部函数。这样只会进行一次动态分派，而不是多次：

# 内部函数，类型稳定
function _g(n, tup)
    i = 0
    for A in Iterators.product((f(n, p, k) for (p, k) in tup)...)
        i += sum(A)  # 或者 i = sum(A)，随意选择
    end
    return i
end

# 外部函数
g(m::T, n::T) where {T<:Integer} = _g(n, Tuple(factor(m)))

一些基准测试结果：

julia> @btime g(7, 210);  # 原始版本
  149.600 μs (7356 allocations: 172.62 KiB)

julia> @btime g(7, 210);  # 我的版本
  1.140 μs (6 allocations: 11.91 KiB)

你应该预期偶尔会触发编译，每当你得到一个包含新因子数量的数字时。

# inner, type stable
function _g(n, tup)
    i = 0
    for A in Iterators.product((f(n, p, k) for (p, k) in tup)...)
        i += sum(A)  # or i = sum(A), whatever
    end
    return i
end

# outer function
g(m::T, n::T) where {T<:Integer} = _g(n, Tuple(factor(m)))

julia> @btime g(7, 210);  # OP version
  149.600 μs (7356 allocations: 172.62 KiB)

julia> @btime g(7, 210);  # my version
  1.140 μs (6 allocations: 11.91 KiB)