MojoでProject Euler 50

https://projecteuler.net/problem=50

和のインデックスと[first, last)で表し、さらにその素数の和をStateとします。そのとき、@valueというdecoratorを付ければ、フィールドだけ書けばCollectionElementを実装したことになります。
HeapQueueはComparableCollectionElementを要求するようにしていますが、Comparableを実装すればよいです。

from collections import Set
import sys


#################### List ####################

fn initialize_list[T: CollectionElement](N: Int, init: T) -> List[T]:
    var a = List[T](capacity=N)
    for n in range(N):
        a.append(init)
    return a


#################### library ####################

fn make_prime_table(N: Int) -> List[Int]:
    var a = initialize_list(N+1, True)
    for p in range(2, N+1):
        if p * p > N:
            break
        elif not a[p]:
            continue
        
        for n in range(p*p, N+1, p):
            a[n] = False
    
    var ps = List[Int]()
    for n in range(2, N+1):
        if a[n]:
            ps.append(n)
    return ps


#################### HeapQueue ####################

struct HeapQueue[T: ComparableCollectionElement]:
    var a: List[T]
    
    fn __init__(inout self):
        self.a = List[T]()
    
    fn is_empty(self) -> Bool:
        return len(self.a) == 0
    
    fn pop(inout self) -> T:
        var top = self.a[0]
        self.a[0] = self.a.pop()
        var N = self.a.size
        var i = 0
        while i < N-1:
            var j1 = i*2+1
            var j2 = i*2+2
            var j: Int = 0
            if j1 >= N:
                break
            elif j2 == N:
                j = j1
            else:
                if self.a[j1] <= self.a[j2]:
                    j = j1
                else:
                    j = j2
            if self.a[j] < self.a[i]:
                self.swap(i, j)
                i = j
            else:
                break
        
        return top
    
    fn push(inout self, e: T):
        self.a.append(e)
        var i = self.a.size - 1
        while i > 0:
            var j = (i-1) // 2
            if self.a[j] <= self.a[i]:
                break
            else:
                self.swap(i, j)
                i = j
    
    fn swap(inout self, i: Int, j: Int):
        var tmp = self.a[i]
        self.a[i] = self.a[j]
        self.a[j] = tmp^


#################### State ####################

@value
struct State(Comparable, Stringable):
    var first: Int
    var last: Int
    var sum: Int
    
    fn length(self) -> Int:
        return self.last - self.first
    
    fn __lt__(self, other: State) -> Bool:
        return self.length() > other.length()
    
    fn __le__(self, other: State) -> Bool:
        return self.length() >= other.length()
    
    fn __eq__(self, other: State) -> Bool:
        return self.length() == other.length()
    
    fn __ne__(self, other: State) -> Bool:
        return self.length() != other.length()
    
    fn __gt__(self, other: State) -> Bool:
        return self.length() < other.length()
    
    fn __ge__(self, other: State) -> Bool:
        return self.length() <= other.length()
    
    fn __str__(self) -> String:
        return ("{ range: [" + str(self.first) + ", " + str(self.last) +")"
                                + ", sum: " + str(self.sum) + " }")


#################### process ####################

fn initialize(N: Int, primes: List[Int]) -> State:
    var s = 0
    for i in range(len(primes)):
        if s + primes[i] >= N:
            return State(0, i, s)
        else:
            s += primes[i]
    return State(0, 0, 0)

fn next_states(s: State, N: Int, primes: List[Int]) -> List[State]:
    var states = List[State]()
    if s.first < s.last + 1:
        var sum = s.sum - primes[s.last-1]
        states.append(State(s.first , s.last - 1, sum))
    if s.sum - primes[s.first] + primes[s.last] < N:
        var sum = s.sum - primes[s.first] + primes[s.last]
        states.append(State(s.first + 1, s.last + 1, sum))
    return states

fn f(N: Int) -> Int:
    var primes = make_prime_table(N)
    var set_primes = Set(primes)
    var hq = HeapQueue[State]()
    hq.push(initialize(N, primes))
    while not hq.is_empty():
        var state = hq.pop()
        if state.sum in set_primes:
            return state.sum
        var new_states = next_states(state, N, primes)
        for new_state in new_states:
            hq.push(new_state[])
    return 0

fn main() raises:
    var args = sys.argv()
    var N = atol(args[1])
    print(f(N))