def topk.go

bend-ml-llama@0.2.1.0 · main.bend:1181

import Base

1181# the k largest of xs (the values from id i on) and of the candidates cs, largest first; cs holds m <= k of them,
1182# lo the smallest; g: whether the head of xs enters. One pass, and a value at or below lo costs one comparison.
1183def topk.go(xs: List<&2, F32>, g: Bool, +k: Nat, +i: Nat, +m: Nat, +lo: F32, +cs: List<&2, Cand>) -> List<&2, Cand>:
1184  match xs g:
1185    case Nil{} _:
1186      cs
1187    case Con{h, +t} False{}:
1188      topk.go(t, enters(t, Nat.is_lt(m, k), lo), k, 1n+i, m, lo, cs)
1189    case Con{+h, +t} True{}:
1190      +cs2 = List.take(&2, Cand, ins(cs, goes_before(cs, h), h, i), k)
1191      +m2 = Nat.min(1n+m, k)
1192      +lo2 = last_v(cs2, h)
1193      topk.go(t, enters(t, Nat.is_lt(m2, k), lo2), k, 1n+i, m2, lo2, cs2)