這篇文章釋出在了 2021-4月份的 Rust 雜誌，在這裡再發一份。
2021 年 5 月 6 號釋出的 Rust 1。52 版將會包含我的一個 PR，將標準庫
slice：：binary_search_by（）
的最好時間複雜度最佳化到
O(1)
。PR 編號 #74024，從 2020 年 7 月初提交的 PR，到 2021 年 3 月 6 號才正式合併，其中包含 70 多條評論，前後歷時大半年。這個 PR 雖然改動不算很大，但是筆者在其中也學到了不少知識點，所以特意寫一篇文章總結一下，希望能幫助到大家。
首先看一下
slice：：binary_search_by（）
這個方法的示例來了解一下它的用途。
let
s
=
［
0
，
1
，
1
，
1
，
1
，
2
，
3
，
5
，
8
，
13
，
21
，
34
，
55
］；
let
seek
=
13
；
assert_eq
！
（
s
。
binary_search_by
（
|
probe
|
probe
。
cmp
（
&
seek
）），
Ok
（
9
））；
let
seek
=
4
；
assert_eq
！
（
s
。
binary_search_by
（
|
probe
|
probe
。
cmp
（
&
seek
）），
Err
（
7
））；
let
seek
=
100
；
assert_eq
！
（
s
。
binary_search_by
（
|
probe
|
probe
。
cmp
（
&
seek
）），
Err
（
13
））；
let
seek
=
1
；
let
r
=
s
。
binary_search_by
（
|
probe
|
probe
。
cmp
（
&
seek
））；
assert
！
（
match
r
{
Ok
（
1
。。
=
4
）
=>
true
，
_
=>
false
，
}）；
這個函式的作用就是在給定的有序 slice 中二分查詢目標值
seek
。如果找到了返回
Ok（pos）
，
pos
即目標值所在的位置；沒找到則返回
Err（pos）
， 這裡的
pos
的位置可以用來將
seek
插入 slice 後依然保持有序。其他的
slice：：binary_search（）
，
slice：：binary_search_by_key（）
系列方法都是呼叫這個
slice：：binary_search_by（）
，這裡不再贅敘。
但是，1。52 之前的實現有一個小問題，如果 slice 中存在多個連續的目標值，則它會一直找到最後一個才返回，所以最好時間複雜度也是
O（log n）
，而不是
O（1）
，也就是找到了就馬上返回。這是 1。52 之前的程式碼：
#［stable（feature =
“rust1”
， since =
“1。0。0”
）］
#［inline］
pub
fn
binary_search_by
<
‘a
，
F
>
（
&
’a
self
，
mut
f
：
F
）
->
Result
<
usize
，
usize
>
where
F
：
FnMut
（
&
‘a
T
）
->
Ordering
，
{
let
s
=
self
；
let
mut
size
=
s
。
len
（）；
if
size
==
0
{
return
Err
（
0
）；
}
let
mut
base
=
0
usize
；
while
size
>
1
{
let
half
=
size
/
2
；
let
mid
=
base
+
half
；
// SAFETY： the call is made safe by the following inconstants：
// - `mid >= 0`： by definition
// - `mid < size`： `mid = size / 2 + size / 4 + size / 8 。。。`
let
cmp
=
f
（
unsafe
{
s
。
get_unchecked
（
mid
）
}）；
base
=
if
cmp
==
Greater
{
base
}
else
{
mid
}；
size
-=
half
；
}
// SAFETY： base is always in ［0， size） because base <= mid。
let
cmp
=
f
（
unsafe
{
s
。
get_unchecked
（
base
）
}）；
if
cmp
==
Equal
{
Ok
（
base
）
}
else
{
Err
（
base
+
（
cmp
==
Less
）
as
usize
）
}
}
既然我們發現這個問題了，很簡單，我們在
while
迴圈中判斷
cmp == Equal
不就行了？如果相等就直接返回
Ok
，否則繼續二分去查詢。
while
size
>
1
{
let
half
=
size
/
2
；
let
mid
=
base
+
half
；
// SAFETY：
// mid is always in ［0， size）， that means mid is >= 0 and < size。
// mid >= 0： by definition
// mid < size： mid = size / 2 + size / 4 + size / 8 。。。
let
cmp
=
f
（
unsafe
{
s
。
get_unchecked
（
mid
）
}）；
if
cmp
==
Equal
{
return
Ok
（
base
）；
}
else
if
cmp
==
Less
{
base
=
mid
}；
size
-=
half
；
}
為了簡潔性，這裡（包括後面）都會適當省略重複程式碼。
嗯，看起來邏輯上沒問題，單測也跑通過了，提交一個 PR，我們暫且叫它
最佳化(1)
。沒過幾天 dtolnay review 了我的 PR 並回復了：
Would you be able to put together a benchmark assessing the worst case impact？ The new implementation does potentially 50% more conditional branches in the hot loop。
確實，Rust 標準庫對效能要求非常高，我們必須要有足夠多 case 對新舊版本做 benchmark 對比，儘可能避免新的版本導致效能上的 regression。這是當時我做的 benchmark 資料：
// 標準庫的實現
test slice：：binary_search_l1 。。。 bench： 59 ns/iter （+/- 4）
test slice：：binary_search_l1_with_dups 。。。 bench： 59 ns/iter （+/- 3）
test slice：：binary_search_l2 。。。 bench： 76 ns/iter （+/- 5）
test slice：：binary_search_l2_with_dups 。。。 bench： 77 ns/iter （+/- 17）
test slice：：binary_search_l3 。。。 bench： 183 ns/iter （+/- 23）
test slice：：binary_search_l3_with_dups 。。。 bench： 185 ns/iter （+/- 19）
// 新版實現
test slice：：binary_search_l1 。。。 bench： 58 ns/iter （+/- 2）
test slice：：binary_search_l1_with_dups 。。。 bench： 37 ns/iter （+/- 4）
test slice：：binary_search_l2 。。。 bench： 76 ns/iter （+/- 3）
test slice：：binary_search_l2_with_dups 。。。 bench： 57 ns/iter （+/- 6）
test slice：：binary_search_l3 。。。 bench： 200 ns/iter （+/- 30）
test slice：：binary_search_l3_with_dups 。。。 bench： 157 ns/iter （+/- 6）
可以看出來在 with_dups 模式（即重複元素較多）下，新的實現有很明顯的提升，但是 l3 級別的普通模式效能反而要差很多。可能的原因正如
dtolnay
所說的
The new implementation does potentially 50% more conditional branches in the hot loop
。那
conditional branches
具體是什麼？為什麼它在熱迴圈中如此關鍵？這裡引入這篇文章的第一個知識點：
分支預測
。
分支預測（Branch prediction）
分支預測
（
branch prediction
） 是現代 CPU 為了加快指令並行化速度在碰到分支時提前預測下一個可能執行的分支的一種技術。 CPU 中一般都內建專門的分支預測器 （Branch predictor）。強烈推薦先閱讀 Stackoverflow 這個高贊回答 Why is processing a sorted array faster than processing an unsorted array？，裡面比較清晰易懂的解釋了分支預測是什麼以及對效能帶來的影響。關於
分支預測
完全可以寫一篇專門的文章介紹了，我這裡只把我學到的知識總結一下。
分支
首先我們要了解一下
分支
到底是指什麼？在高階語言層面，很顯然分支就是
if/else/else if
，
goto
或者
switch/match
這種語句。他們會轉換成彙編程式碼中的
jump
指令。比如 x86 彙編的各種
j
開頭的
jump
指令：
比如我們隨便寫一段包含
if/else
邏輯的程式碼，編譯之後的彙編程式碼中就會包含類似的
jump
指令：
#！［allow（unused_assignments）］
pub
fn
main
（）
{
// 接收使用者從命令列中輸入的引數
let
mut
a
：
usize
=
std
：：
env
：：
args
（）。
nth
（
1
）。
unwrap
（）。
parse
（）。
unwrap_or_default
（）；
if
a
>
42
{
a
=
1
；
}
else
{
a
=
0
；
}
}
這段程式碼邏輯這麼寫主要是為了保持彙編指令的儘量簡單。這裡是對應的彙編程式碼，只保留了跟
if/else
有關的部分，完整的程式碼連結在這裡：
https：//
rust。godbolt。org/z/ahcK
cK1Er
。
。
LBB99_7
：
cmp
qword
ptr
［
rsp
+
56
］，
42
；
if
a
>
42
jbe
。
LBB99_10
mov
qword
ptr
［
rsp
+
56
］，
1
；
a
=
1
jmp
。
LBB99_11
。
LBB99_10
：
mov
qword
ptr
［
rsp
+
56
］，
0
；
a
=
0
。
LBB99_11
：
add
rsp
，
200
ret
我在程式碼裡面加了幾條註釋（彙編裡面用分號註釋）表明了該條指令對應的 Rust 程式碼，可以看出來
jbe
指令會判斷是否跳轉到
。LBB99_10
還是不跳轉繼續執行下面的
mov
。
瞭解了分支是什麼之後還不夠，我們依然不知道具體是為啥 CPU 需要做分支預測。接下來我們需要了解一下另外一個概念：
指令流水線
。
指令流水線（instruction pipelining）
指令流水線
是從
instruction pipelining
翻譯過來的一個名詞，主要是指為了提高單個處理器指令處理效率的一種技術，早在上個世紀 70 年代的晶片就有這種技術了。
CPU 處理一條指令一般要分為幾個步驟：
取指令（Instruction Fetch）
解指令（Instruction Decode）
執行指令 （Execute）
回寫結果到暫存器 （Register write back）
這非常類似工廠生產某件需要多道工序的商品一樣。想象一下，如果工廠每次都是完整的按照工序把第一件商品生產出來之後再進行同樣的步驟生產下一件商品的話，這個生產速度得多慢！所以 19 世紀人類就誕生了工業化的流水線，把所有這些工序分解並行化。第一件商品在工序（2）的時候，絲毫不影響第二件商品進入工序（1）。指令也完全可以這樣，CPU 可以把指令執行過程流水線化，把不同的指令執行步驟分給不同的邏輯閘 （logic stage）處理，第一條指令進入
解指令
階段的時候，第二條指令就可以進入
取指令
階段了。
Wikipedia 上的這張圖可以幫助理解（圖中文字顏色為黑色，推薦用 Light 主題檢視）。
瞭解前面的這些概念之後，我們來看一下為什麼 CPU 需要做分支預測？
為什麼需要分支預測？
指令流水線
會按照工廠流水線一樣執行指令，這其中也包括前面講的
jump
指令。而
jump
指令有一個問題就是它需要知道下一個時鐘週期該跳轉還是不跳轉，而這需要等前面的判斷邏輯執行完了之後才知道。比如上面的例子
cmp
判斷完之後，
jump
指令才能決定是跳轉到
。LBB99_10
部分，還是不跳轉繼續執行下去。但是 CPU 指令流水線才不會等，否則白白的浪費一個時鐘週期。所以人們發明了一種辦法來避免出現這種問題，這就是
分支預測
。
CPU 的分支預測器會提前預測這條
jump
指令可能會跳轉的到哪個分支，然後把預測的那條分支放到流水線中去執行。如果預測正確，CPU 可以繼續執行後面的指令；如果預測失敗了（branch misprediction），那隻能丟棄剛剛執行的分支結果，重新切換到正確的分支執行。可以看到，如果出現過多的預測失敗，分支預測反而很影響效能。不過現在 CPU 的分支預測器經過這麼多年的發展已經越來越先進了，人們會採用各種方式不斷提高分支預測器的預測準確率，詳細可以檢視 Wikipedia 的 Branch_predictor 瞭解更多。
熱迴圈中避免分支預測
雖然現代 CPU 都有分支預測器，但我們在軟體層面依然要儘量避免發生分支預測，特別是在熱迴圈中。最常用的最佳化方法就是避免在迴圈中寫
if/else
語句，即
branchless code
。標準庫中有大量這種
branchless code
的案例來最佳化效能，比如
std：：collection：：Filter
的
count（）
方法。
pub
struct
Filter
<
I
，
P
>
{
iter
：
I
，
predicate
：
P
，
}
impl
<
I
：
Iterator
，
P
>
Iterator
for
Filter
<
I
，
P
>
where
P
：
FnMut
（
&
I
：：
Item
）
->
bool
，
{
type
Item
=
I
：：
Item
；
// this special case allows the compiler to make `。filter（_）。count（）`
// branchless。 Barring perfect branch prediction （which is unattainable in
// the general case）， this will be much faster in >90% of cases （containing
// virtually all real workloads） and only a tiny bit slower in the rest。
//
// Having this specialization thus allows us to write `。filter（p）。count（）`
// where we would otherwise write `。map（|x| p（x） as usize）。sum（）`， which is
// less readable and also less backwards-compatible to Rust before 1。10。
//
// Using the branchless version will also simplify the LLVM byte code， thus
// leaving more budget for LLVM optimizations。
#［inline］
fn
count
（
self
）
->
usize
{
#［inline］
fn
to_usize
<
T
>
（
mut
predicate
：
impl
FnMut
（
&
T
）
->
bool
）
->
impl
FnMut
（
T
）
->
usize
{
move
|
x
|
predicate
（
&
x
）
as
usize
}
self
。
iter
。
map
（
to_usize
（
self
。
predicate
））。
sum
（）
}
}
標準庫的
Filter
型別在實現
Iterator
的時候重寫了
count（）
方法。想想如果我們沒有意識到分支預測這個問題的情況下可能會這麼實現：
// Bad
#［inline］
fn
count
（
self
）
->
usize
{
let
sum
=
0
；
self
。
iter
。
for_each
（
|
x
|
{
if
self
。
predicate
（
x
）
{
sum
+=
1
；
}
}）；
sum
}
但這種實現在迴圈中有一個
if
語句，導致 CPU 需要進行大量的分支預測，而這些分支預測幾乎都是隨機的，CPU 很難根據歷史記錄提高預測的準確度，導致效能會比較低。而標準庫的實現完全是 branchless 的，不僅效能好很多，而且也能方便 LLVM 做更多最佳化。
1。 關於怎樣寫好 branchless code 來最佳化效能又是另一個值得專門討論的主題了，網上也有大量資料可以學習。但是 branchless code 會犧牲不少程式碼可讀性，並不建議盲目的使用。
2。 branchless code 的額外好處是還能幫助避免遭受
旁路攻擊
（英文 timing attack 或者 side-channel attack）。參考：
https：//
en。wikipedia。org/wiki/T
iming_attack
。
繼續回到我們剛才的那個 PR。我們的版本比標準庫舊的版本在某些 case 下效能要低確實跟分支預測有關，因為我們的程式碼多了一種分支需要預測。
為了方便大家對比，我把彙編的截圖貼到下面。我們只需要關注我標記文字的那幾個顏色就可以。很容易看出來新版實現多了一個
jne
的跳轉指令，導致 CPU 需要多做一次分支預測。
申明：啟用
-O
引數之後彙編指令會被最佳化得更精簡，這裡沒有啟用
-O
是為了對應到每一行原始碼方便對比。
標準庫的彙編截圖
最佳化（1）的彙編截圖
需要注意的是
jmp
指令是直接跳轉，不需要進行分支預測。感興趣的朋友可以看一下我在 Godbolt 上的對比：
https：//
rust。godbolt。org/z/8dGb
Y8Pe1
。這個網站是神器，強烈推薦！
最佳化（2）
所以我推測原作者實現標準庫的
binary_search_by（）
的時候不考慮
O（1）
最好時間複雜度的可能原因之一就是為了避免多餘的分支預測。因為你要想
O（1）
就避免不了提前 return，要想提前 return 你就避免不了分支判斷。那怎麼辦呢？PR 裡面有一個大牛 tesuji 提供了一種思路：既然我們避免不了分支預測，那我們儘量幫助 CPU 更好的做好分支預測吧。於是我採用了他的方案，具體 commit 在這裡，我們暫且叫它
最佳化(2)
：
pub
fn
binary_search_by
<
’a
，
F
>
（
&
‘a
self
，
mut
f
：
F
）
->
Result
<
usize
，
usize
>
where
F
：
FnMut
（
&
’a
T
）
->
Ordering
，
{
let
mut
left
=
0
；
let
mut
right
=
self
。
len
（）；
while
left
<
right
{
// never overflow because `slice：：len（）` max is `isize：：MAX`。
let
mid
=
（
left
+
right
）
/
2
；
// SAFETY： the call is made safe by the following invariants：
// - `mid >= 0`
// - `mid < size`： `mid` is limited by `［left； right）` bound。
let
cmp
=
f
（
unsafe
{
self
。
get_unchecked
（
mid
）
}）；
if
cmp
==
Less
{
left
=
mid
+
1
；
}
else
if
cmp
==
Greater
{
right
=
mid
；
}
else
{
return
Ok
（
mid
）；
}
}
Err
（
left
）
}
最佳化（2）的程式碼明顯比標準庫和最佳化（1）的程式碼更容易理解，再看一下它的生成的彙編程式碼。
可以看出來依然是兩條
jne
指令，所以非重複模式下的效能可能還是沒有標準庫的高，但是確實比最佳化（1）的效能要好很多。過了幾天 libs 組的 m-ou-se 回覆了評論。她也做了 benchmark，發現對於原生型別比如
u32
下 l1 級別的資料量依然比標準庫慢，但是那些需要更多時間比較的型別（比如 String）的情況下，新的實現在所有 case 下效能都要優於標準庫的實現。後面大家又討論了許多，最終
m-ou-se
決定先跑一個 crater 測試，先驗證一下這個 PR 對
http：//
crates。io
上所有的 crate 影響面大不大。最終 library 團隊會議一致同意可以 merge 這個 PR。
關於 crater 測試：
crater 大概就是針對整個
http：//
crates。io
上所有的 crate 給要測試的編譯器版本（比如我的 PR ）跑一次測試，看這個改動對線上所有 crate 影響大不大。
http：//
crates。io
上超過5萬個 crate，一般跑一次 crater 需要將近一週的時間。我的這個 crater 跑完之後就是因為沒有對已釋出的 crate 造成什麼很大的影響，所以官方才願意合併。
From
m-ou-se
：
“We discussed this PR in a recent library team meeting， in which we agreed that the proposed behaviour （stopping on Equal） is preferrable over optimal efficiency in some specific niche cases。 Especially considering how small most of the differences are in the benchmarks above。”
“The breakage in the crater report looks reasonably small。 Also， now that
partition_point
is getting stabilized， there‘s a good alternative for those who want the old behaviour of
binary_search_by
。 So we should go ahead and start on getting this merged。 ：）”
整數溢位問題
然而 scottmcm 又指出了另外一個問題：
// never overflow because `slice：：len（）` max is `isize：：MAX`。
let
mid
=
（
left
+
right
）
/
2
；
這行程式碼在零大小型別（Zero Sized Type，簡稱 ZST）下卻可能會 overflow！ 我們來分析一下為什麼。
slice：：len（）
的返回值是
usize
型別，但是對於非零大小的型別（non-ZST），
slice：：len（）
的值最大也只能是
isize：：MAX
。所以就像註釋上寫的那樣
（isize：：MAX + isize：：MAX） / 2
是不能可能超過
usize：：MAX
的，所以不會發生 overflow。但是對於 ZST 型別就不一樣了，如果
slice
裡面所有元素都是零大小的（比如
（）
），那這個
slice
的長度完全可以達到
usize：：MAX
。雖然對於
［（）； usize：：MAX］。binary_search（&（））
這種情況我們會在
O（1）
的時間複雜度上找到結果並馬上返回，但是如果我們這麼寫
b。binary_search_by（|_| Ordering：：Less）
，它就發生整數溢位了。
為什麼 slice：：len（） 對於 non-ZST 最大值是 isize 呢？
最簡單直接的原因是我們不能構造一個所有元素都為 non-ZST 並且長度為
usize：：MAX
的陣列或
slice
，編譯器在編譯階段直接會報錯。比如以最簡單的只佔 1 個位元組的
bool
型別為例，
［bool； usize：：MAX］
的大小將等於
std：：mem：：size_of：：
， 這是一個很大的數字了，整個計算機地址空間都不夠。
fn
main
（）
{
assert_eq
！
（
std
：：
mem
：：
size_of
：：
<
bool
>
（），
1
）；
// error： values of the type `［bool； 18446744073709551615］` are too big
// for the current architecture
let
_s
=
［
true
；
usize
：：
MAX
］；
}
但是對於 ZST 是可以的，因為
std：：mem：：size_of：：<（）>（） * usize：：MAX
依然是零。
fn
main
（）
{
assert_eq
！
（
std
：：
mem
：：
size_of
：：
<
（）
>
（），
0
）；
let
s
=
［（）；
usize
：：
MAX
］；
assert_eq
！
（
s
。
len
（），
usize
：：
MAX
）；
}
不過上面的解釋依然不夠嚴謹，比如
std：：mem：：size_of：：
也依然是一個很大的數字呀，為啥
isize：：MAX
就可以？根本原因在於 Rust 指標定址最大 offset 只允許
isize：：MAX
，至於為什麼是
isize：：MAX
，
std：：pointer：：offset（）
的文件有解釋。另外也可以看一下
std：：slice：：from_raw_parts（）
的文件。而對於 ZST 的型別，編譯器會做最佳化，它壓根不需要定址，所以最大大小可以是
usize：：MAX
。
最終版本
意識到整數溢位的問題之後，解決方式也比較簡單，這是我當時的提交，並且還增加了針對 overflow 的單元測試。
pub
fn
binary_search_by
<
’a
，
F
>
（
&
‘a
self
，
mut
f
：
F
）
->
Result
<
usize
，
usize
>
where
F
：
FnMut
（
&
’a
T
）
->
Ordering
，
{
let
mut
size
=
self
。
len
（）；
let
mut
left
=
0
；
let
mut
right
=
size
；
while
left
<
right
{
let
mid
=
left
+
size
/
2
；
// SAFETY： the call is made safe by the following invariants：
// - `mid >= 0`
// - `mid < size`： `mid` is limited by `［left； right）` bound。
let
cmp
=
f
（
unsafe
{
self
。
get_unchecked
（
mid
）
}）；
// The reason why we use if/else control flow rather than match
// is because match reorders comparison operations， which is
// perf sensitive。
// This is x86 asm for u8： https：//rust。godbolt。org/z/8Y8Pra。
if
cmp
==
Less
{
left
=
mid
+
1
；
}
else
if
cmp
==
Greater
{
right
=
mid
；
}
else
{
return
Ok
（
mid
）；
}
size
=
right
-
left
；
}
Err
（
left
）
}
#［test］
fn
test_binary_search_by_overflow
（）
{
let
b
=
［（）；
usize
：：
MAX
］；
assert_eq
！
（
b
。
binary_search_by
（
|
_
|
Ordering
：：
Equal
），
Ok
（
usize
：：
MAX
/
2
））；
assert_eq
！
（
b
。
binary_search_by
（
|
_
|
Ordering
：：
Greater
），
Err
（
0
））；
assert_eq
！
（
b
。
binary_search_by
（
|
_
|
Ordering
：：
Less
），
Err
（
usize
：：
MAX
））；
}
我們確實應該儘量避免寫
let mid = （left + right） / 2
這種很容易發生整數溢位的程式碼，換成
let mid = left + size / 2
這種，可以避免發生 overflow。
另外還有人問為什麼這裡使用
if/else
而不是
match
語句？我們檢視兩個版本的彙編指令後發現
match
版本生成的彙編程式碼不僅指令更多而且還重排了
cmp
指令的順序，效能似乎更差。理論上這兩個版本生成的彙編指令應該可以做到一致的，我暫時沒有深究原因為什麼
match
版本的彙編會差一些，其他讀者感興趣可以研究一下。
總結
表面波瀾不驚，實則暗流湧動。一個看起來十分簡單的 PR 裡面其實涉及到很多內容。學無止境，筆者透過這個 PR 收穫了很多，現在分享出來同時也希望能夠激勵更多國內開發者參與進來。Rust 社群文化是如此開放和包容，任何人只要發現有可以改進的地方都可以隨時給 Rust 倉庫提交 PR 或 issue，這樣做不僅能幫助 Rust 越來越好，你自己也會得到巨大的成長和收穫！
關於我
Id： Folyd，GitHub：@folyd。位元組跳動飛書 Rust 工程師，Rust Search Extension 作者。
招聘
位元組跳動飛書團隊自 2017 年就開始使用 Rust 開發飛書多端跨平臺 SDK，為 Android / iOS / Window / macOS / Linux 等平臺提供高質量的底層基礎庫，同時飛書內部的效率工具和少數後端系統也全部採用 Rust 開發。我們可能是國內 Rust 工程師最多的團隊之一！我們長期招聘熱愛 Rust、喜歡 Rust、看好 Rust 前景的工程師加入。支援實習、校招、社招。Base 北京，歡迎大家自薦或推薦，請聯絡 wx：
newpants629
，或者直接在內推連結投遞簡歷：
https：//
job。toutiao。com/s/eB1j2
9f
。
