bsearch.bend

bend-osdev-algo@0.1.0.0 · 140 lines · 19 definitions

1# bend-osdev-algo, binary search: the lower bound of a key in a sorted
2# sequence of U32 (the index of the first element not below the key, or the
3# length if every element is below it), and an index holding the key.
4#
5# Two implementations, as CONVENTIONS asks where speed matters:
6#
7#   bsearch.lower_bound, bsearch.find              the reference, on a List:
8#                                                  a scan, proved correct
9#   bsearch.array_lower_bound, bsearch.array_find  the fast one, on an
10#                                                  Array<U32>: a binary
11#                                                  search, O(log n) reads
12#
13# LAWS.bend holds the two together for every perfect array whose searched
14# slots are sorted (proved through bend-osdev-array's laws), beside the
15# vectors and the differential test in ref/.
16import Base
17
18# Sorted: each element is at most the next
19def bsearch.sorted.pair(xs: List<&2, U32>, x: U32) -> Bool:
20  match xs:
21    case Nil{}:
22      True{}
23    case Con{+y, t}:
24      Bool.and(U32.is_le(x, y), bsearch.sorted.pair(t, y))
25
26def bsearch.sorted(xs: List<&2, U32>) -> Bool:
27  match xs:
28    case Nil{}:
29      True{}
30    case Con{x, t}:
31      bsearch.sorted.pair(t, x)
32
33# The reference
34
35def bsearch.below.step(c: Bool, k: Nat) -> Nat:
36  match c:
37    case True{}:
38      1n+k
39    case False{}:
40      0n
41
42# the number of leading elements below key: on a sorted list, the index of
43# the first element not below key (the length if there is none)
44def bsearch.lower_bound(xs: List<&2, U32>, +key: U32) -> Nat:
45  match xs:
46    case Nil{}:
47      0n
48    case Con{x, t}:
49      bsearch.below.step(U32.is_lt(x, key), bsearch.lower_bound(t, key))
50
51def bsearch.find.hit(c: Bool, i: Nat) -> Maybe<&2, Nat>:
52  match c:
53    case True{}:
54      Some{i}
55    case False{}:
56      None{}
57
58def bsearch.find.check(+i: Nat, m: Maybe<&2, U32>, key: U32) -> Maybe<&2, Nat>:
59  match m:
60    case None{}:
61      None{}
62    case Some{x}:
63      bsearch.find.hit(U32.is_eq(x, key), i)
64
65# an index of a sorted list that holds key (the first), or None
66def bsearch.find(+xs: List<&2, U32>, +key: U32) -> Maybe<&2, Nat>:
67  +i = bsearch.lower_bound(xs, key)
68  bsearch.find.check(i, List.get(&2, U32, xs, i), key)
69
70# The fast implementation, on the first n slots of an Array<U32>
71
72def bsearch.array.fin(r: Array<U32> & U32) -> Array<U32>:
73  (a, x) = r
74  a
75
76def bsearch.array.val(r: Array<U32> & U32) -> U32:
77  (a, x) = r
78  x
79
80# the search state: the array, and the range [lo, hi) still in question
81type BsearchState is Type:
82  BsearchState{a: Array<U32>, lo: U32, hi: U32}
83
84# one halving of [lo, hi): the element at mid decides which half stays
85def bsearch.array.pick(below: Bool, a: Array<U32>, +lo: U32, +hi: U32, +mid: U32) -> BsearchState:
86  match below:
87    case True{}:
88      BsearchState{a, U32.inc(mid), hi}
89    case False{}:
90      BsearchState{a, lo, mid}
91
92def bsearch.array.step(active: Bool, r: Array<U32> & U32, +key: U32, +lo: U32, +hi: U32, +mid: U32) -> BsearchState:
93  match active:
94    case True{}:
95      (a, +x) = r
96      bsearch.array.pick(U32.is_lt(x, key), a, lo, hi, mid)
97    case False{}:
98      (a, x) = r
99      BsearchState{a, lo, hi}
100
101def bsearch.array.advance(st: BsearchState, +key: U32) -> BsearchState:
102  match st:
103    case BsearchState{a, +lo, +hi}:
104      +mid = U32.add(lo, U32.shr(U32.sub(hi, lo)))
105      bsearch.array.step(U32.is_lt(lo, hi), Array.get(U32, a, mid), key, lo, hi, mid)
106
107def bsearch.array.result(st: BsearchState) -> Array<U32> & U32:
108  match st:
109    case BsearchState{a, lo, hi}:
110      (a, lo)
111
112# fuel bounds the halvings: 33 is enough for any [lo, hi) of U32; once the
113# range is empty the remaining steps change nothing
114def bsearch.array.go(fuel: Nat, st: BsearchState, +key: U32) -> Array<U32> & U32:
115  match fuel:
116    case 0n:
117      bsearch.array.result(st)
118    case 1n+f:
119      bsearch.array.go(f, bsearch.array.advance(st, key), key)
120
121# the index of the first of the first n elements not below key (n if none)
122def bsearch.array_lower_bound(a: Array<U32>, n: U32, key: U32) -> Array<U32> & U32:
123  bsearch.array.go(33n, BsearchState{a, 0, n}, key)
124
125def bsearch.array.found(inside: Bool, r: Array<U32> & U32, +i: U32, +key: U32) -> Array<U32> & Maybe<&2, U32>:
126  match inside:
127    case True{}:
128      (a, +x) = r
129      (a, Bool.pick(Maybe<&2, U32>, U32.is_eq(x, key), Some{i}, None{}))
130    case False{}:
131      (a, x) = r
132      (a, None{})
133
134def bsearch.array.find(r: Array<U32> & U32, +n: U32, +key: U32) -> Array<U32> & Maybe<&2, U32>:
135  (a, +i) = r
136  bsearch.array.found(U32.is_lt(i, n), Array.get(U32, a, i), i, key)
137
138# an index of the first n elements of a sorted array that holds key, or None
139def bsearch.array_find(a: Array<U32>, +n: U32, +key: U32) -> Array<U32> & Maybe<&2, U32>:
140  bsearch.array.find(bsearch.array_lower_bound(a, n, key), n, key)