Unsafe编程（下）：使用Rust为Python写一个扩展

你好，我是Mike。

上一讲我们了解了Unsafe Rust的所属定位和基本性质，这一讲我们就来看看Rust FFI编程到底是怎样一种形式。

FFI是与外部语言进行交互的接口，在Safe Rust看来，它是不可信的，也就是说Rust编译不能保证它是安全的，只能由程序员自己来保证，因此在Rust里调用这些外部的代码功能就需要把它们包在 unsafe {} 中。

Rust和C有血缘关系，具有ABI上的一致性，所以 Rust和C之间可以实现双向互调，并且不会损失性能。

Rust调用C

我们先来看在Rust中如何调用C库的代码。

一般各个平台下都有 libm 库，它是操作系统基本的数学math库。下面我们以Linux为例来说明。下面的示例代码来自 Rust By Example。

use std::fmt;

// 连接到系统的 libm 库
#[link(name = "m")]
extern {
    // 这是一个外部函数，计算单精度复数的方根
    fn csqrtf(z: Complex) -> Complex;
    // 计算复数的余弦值
    fn ccosf(z: Complex) -> Complex;
}

// 对unsafe调用的safe封装，从此以后，就按safe函数方式使用这个接口
fn cos(z: Complex) -> Complex {
    unsafe { ccosf(z) }
}

fn main() {
    // z = -1 + 0i
    let z = Complex { re: -1., im: 0. };

    // 调用m库中的函数，需要用 unsafe {} 包起来
    let z_sqrt = unsafe { csqrtf(z) };

    println!("the square root of {:?} is {:?}", z, z_sqrt);

    // 调用安全封装后的函数
    println!("cos({:?}) = {:?}", z, cos(z));
}

// 用 repr(C) 标注，定义Rust结构体的ABI格式，按C的ABI来
#[repr(C)]
#[derive(Clone, Copy)]
struct Complex {
    re: f32,
    im: f32,
}

// 实现复数的打印输出
impl fmt::Debug for Complex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.im < 0. {
            write!(f, "{}-{}i", self.re, -self.im)
        } else {
            write!(f, "{}+{}i", self.re, self.im)
        }
    }
}

libm 库是用C语言实现的，我们要调用它，需要用这样的标注。

#[link(name = "m")]
extern {
}

用 lnk 属性宏将libm库连接进来，以便去里面找对应的函数。然后把需要的外部函数导入进来。你可以看一下 csqrtf() 函数和 ccosf() 的C语言签名。

float complex csqrtf(float complex z);
float complex ccosf (float complex z);

导入的时候，要翻译成Rust函数的签名形式。

fn csqrtf(z: Complex) -> Complex;
fn ccosf(z: Complex) -> Complex;

调用这两个函数的时候，都需要用 unsafe {} 包起来。

然后我们仔细看一下在Rust中对应到libm中的复数的定义。

#[repr(C)]
#[derive(Clone, Copy)]
struct Complex {
    re: f32,
    im: f32,
}

对应的C语言中的结构定义在 complex.h 头文件 中，它是C99标准引入的数据类型。

本身这个定义非常简单，就是定义复数的实部和虚部就可以了，不过一定要用 #[repr(C)] 标注。它的意思是这个类型编译的时候，要使用C语言的ABI格式。我们一般都使用C语言ABI格式，但是也有一些其他的ABI格式，你可以点击我给出的 链接，查到Rust中支持的ABI格式列表。

一般来讲，我们会使用Rust函数对这些外部unsafe函数做一下封装。像cos函数这样，你可以参看一下示例代码。

fn cos(z: Complex) -> Complex {
    unsafe { ccosf(z) }
}

执行 cargo run，输出了下面这两行代码。

the square root of -1+0i is 0+1i
cos(-1+0i) = 0.5403023+0i

注：项目代码链接 https://github.com/miketang84/jikeshijian/tree/master/30-ffi/rust_call_c

bindgen

Rust官方出了一个 bindgen 项目，可以帮助我们快速从C库的头文件生成Rust FFI绑定层代码。之所以能这样做，是因为我们发现，C的接口转换成Rust的接口形式是非常固定的，转换的过程需要写大量的样板代码，所以就可以用 bindgen 这种工具自动转换。

比如某个C头文件里定义了这样的类型和函数：

typedef struct Doggo {
    int many;
    char wow;
} Doggo;

void eleven_out_of_ten_majestic_af(Doggo* pupper);

使用 bindgen 自动生成FFI绑定后，可以生成类似这样的代码：

/* automatically generated by rust-bindgen 0.99.9 */

#[repr(C)]
pub struct Doggo {
    pub many: ::std::os::raw::c_int,
    pub wow: ::std::os::raw::c_char,
}

extern "C" {
    pub fn eleven_out_of_ten_majestic_af(pupper: *mut Doggo);
}

你可能发现了，Rust里有一批类型，以 c_ 前缀开头。比如 C语言的 int 类型，对应 std 中的类型你可以通过我给出的 链接 找到。

这些就是Rust里定义的与C完全相同的类型。我们可以看一下 c_char 的定义。

pub type c_char = i8;

竟然直接是一个 i8。因为C语言里的char就是一个普通字节，和Rust里的char完全不同。这是在Rust中完全兼容C必要的一步，就是 把C中的类型全部映射过来，并加上c_ 前缀。这样单独搞一批类型就是因为有的类型没办法直接映射到 Rust 的 char 上来，比如刚刚说的这个 c_char。

典型的，还有 C 语言中的 void 类型，它在Rust里没有对应物，所以就映射成了 c_void，在Rust中定义成了一个 enum。

#[repr(u8)]
pub enum c_void {
    // some variants omitted
}

下面我们继续看，反过来，在C中如何调用Rust代码。

C调用Rust

在C中调用Rust代码的基本思路是，Rust代码编译成 .so 或 .dll 动态链接库，在C语言里面编译的时候，连接就可以了。

在 lib.rs 文件中，写这样一个函数：

#[no_mangle]
pub extern "C" fn hello_from_rust() {
    println!("Hello from Rust!");
}

注意属性宏 #[no_mangle]，在Rust中，no_mangle 用于告诉Rust编译器不要修改函数的名称，这种修改叫做Mangling，它是编译器在解析名称时，修改我们定义的函数名称，增加一些用于其编译过程的额外信息。但在和其他语言交互时，如果函数名称被编译器修改，程序开发者就没办法知道修改后的函数名称了，其他语言也无法按原名称调用。因此，#[no_mangle] 在这种场景下就非常有用了。

Cargo.toml 要加个配置，表示编译的这个库是 C ABI 格式的动态链接库。

[lib]
crate-type = ["cdylib"]

然后用 cargo build 编译输出。你可以在 target/debug 目录下面看到你的 .so 库文件，比如下面这个样子：

$ ls target/debug/
build  deps  examples  incremental  libc_call_rust.d  libc_call_rust.so

这里这个 libc_call_rust.so 就是我们的动态链接库，其中 c_call_rust 是我们这个 crate 的名字。

下面是C的部分，代码如下：

extern void hello_from_rust();

int main(void) {
    hello_from_rust();
    return 0;
}

存成 call_rust.c 文件。用下面这种形式编译：

gcc call_rust.c -o call_rust -lc_call_rust -L./target/debug

运行：

LD_LIBRARY_PATH=./target/debug ./call_rust

你可以看到输出内容 Hello from Rust!。

注：完整可测试代码 https://github.com/miketang84/jikeshijian/tree/master/30-ffi/c_call_rust

当然，这只是最简单的形式，万里长征我们刚踏出第一步，还有各种复杂的FFI场景等着你去探索。

Cbindgen

在实际工作中，当你需要以Rust仓库为主，为其他语言导出C接口的时候，一般还需要生成C的头文件，这时，你应该会对 cbindgen 项目很感兴趣，它也算半官方的（在Mozilla名下），它的作用与前面的 bindgen 刚好相反，是从Rust代码中生成C可以调用的代码签名。

Rust与C的交互我们先讨论到这里，下面我们要研究一下如何用Rust给Python写扩展。

使用Rust给Python写扩展

标准的Python扩展是用 C/C++ 写的，Python 官方有 教程，感兴趣的话你可以研究一下。我们这节课聚焦于如何用Rust给Python写扩展。

其实，只要你将Rust的代码扩展到C可以调用，那么再进一步，按Python官方教程那样，继续将C代码封装成Python扩展就可以了，下面我们讲的基本思路也是这样。

在实际实现的时候，我们会写非常多的样板代码，写起来比较痛苦。但是Rust社区有非常优秀的项目：PyO3，它可以帮助我们减轻这种烦恼，自动帮我们处理好中间样板封装过程。

PyO3

PyO3是Rust社区里非常火的与Python绑定的框架，它不仅可以实现用Rust给Python写扩展，还能在Rust的二进制程序中直接执行Python代码。所以实际 PyO3 是一个双向绑定库。

PyO3封装了底层FFI绑定的各种细节。这些细节其实不是那么简单，你可以想一下，如何将Python的Class正确映射到Rust中对应的类型，如何正确处理模块、函数、参数、返回值类型，这里面有各种细节，想想就头痛。

而 PyO3 把这一切都封装在了Rust属性宏里面。我们来看一个示例。

use num_bigint::BigUint;
use num_traits::{One, Zero};
use pyo3::prelude::*;

// Calculate large fibonacci numbers.
fn fib(n: usize) -> BigUint {
    let mut f0: BigUint = Zero::zero();
    let mut f1: BigUint = One::one();
    for _ in 0..n {
        let f2 = f0 + &f1;
        f0 = f1;
        f1 = f2;
    }
    f0
}

#[pyfunction]
fn calc_fib(n: usize) -> PyResult<()> {
    let _ = fib(n);
    Ok(())
}

#[pymodule]
fn rust_fib(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(calc_fib, m)?)?;
    Ok(())
}

这个示例的作用是计算斐波那契数列的第N项。当这个N值比较大的时候，斐波那契数是一个非常大的数，一个u64装不下，因此我们要使用大数类型。在Rust中，num-bigint是一个常用的大数类型实现库。而在Python中，int默认就支持大数。

代码中， fib() 函数是真正计算斐波那契数的函数，然后在 calc_fib() 中封装给了 Python 使用，请注意这个函数头上的 #[pyfunction] 标注，它表明这个函数会导出给 Python 使用，是一个函数。

而下面的 rust_fib() 函数，则是标准的样板代码，它表示创建一个 Python 的模块，这个模块的名字叫做 rust_fib，这个名字会在Python代码中以 import rust_fib 的形式导入。 #[pymodule] 就起这个自动转化的作用。这个函数中的 m 参数就表示模块实例， m.add_function() 将 calc_fib() 添加到这个模块中。所以我们这个 crate 库会被编译成一个共识库，并被导入为 Python 的一个模块。

你可以看一下 Cargo.toml 的新增配置。

[lib]
name = "rust_fib"
crate-type = ["cdylib"]

在Python中这样调用：

import rust_fib
rust_fib.calc_fib(2000000)

注：完整的可运行代码 https://github.com/miketang84/jikeshijian/blob/master/30-ffi/bigint-pyo3

这样PyO3项目的一个关键是，需要使用 maturin 这种构建工具。你可以下载代码后，按下面的脚本准备好虚拟环境，并在这个虚拟环境里安装好 maturin。

$ cd bigint-pyo3
$ python -m venv .env
$ source .env/bin/activate
$ pip install maturin

在项目目录下运行：

maturin develop -r

这个指令会编译Rust代码为 release 模式，并把编译后的文件安装到 Python module 库的目录组织里去。

然后运行：

$ time python fib_rust.py

real    0m12.452s
user    0m12.431s
sys     0m0.020s

我们对比一下纯Python版本运行的时间。

$ time python fib.py

real    0m21.730s
user    0m21.490s
sys     0m0.240s

这是计算第200万项斐波那契数，可以看到用Rust实现的版本快将近一倍。

PyO3给我们提供的编程界面非常清爽，只要你跟着操作一遍，很快就能用Rust给Python写扩展了。PyO3的 文档 也写得非常好，上面有更全面的功能介绍，一定要读一读。

性能优化示例：文本单词统计

Python的GIL（Global Interpreter Lock）是一个被人诟病的特性，它限制很多Python代码只能以单线程的形式来跑。因此也就大大限制了Python的性能。而Rust原生支持系统级多线程（thread）以及轻量级线程（tokio task 等），能够非常方便地充分压榨所有CPU核。既然我们可以使用Rust给Python写扩展，那一个点子就顺理成章冒出来了——可以使用Rust给Python写扩展实现并行计算。

下面我们就以官方的示例 word_count 来说明如何使用。这个程序要解决这样一个问题：统计出一个文本文件里，某一个单词出现的次数。这个文本文件是按换行符分隔成一行一行的，行与行之间互不干涉。因此可以采用按行分割的形式来并行化处理。

Rust中有一个著名的并行化库Rayon，可以将串行迭代器转换成一个并行化版本的迭代器，从而实现程序的并行化。下面我们来看代码：

use pyo3::prelude::*;
use rayon::prelude::*;

/// 并行化版本的搜索功能实现
#[pyfunction]
fn search(contents: &str, needle: &str) -> usize {
    contents
        .par_lines()
        .map(|line| count_line(line, needle))
        .sum()
}

/// 将行按空格分割，统计目标单词数目
fn count_line(line: &str, needle: &str) -> usize {
    let mut total = 0;
    for word in line.split(' ') {
        if word == needle {
            total += 1;
        }
    }
    total
}

// 导出到Python module
#[pymodule]
fn word_count(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(search, m)?)?;

    Ok(())
}

代码中， search() 函数中 contents.par_lines() 这一行就是使用的 Rayon 里的 par_lines() 迭代器，它是 lines() 迭代器的并行化版本。上面的 search 实现对应下面的 Python 版本。

def search_py(contents: str, needle: str) -> int:
    total = 0
    for line in contents.splitlines():
        for word in line.split(" "):
            if word == needle:
                total += 1
    return total

你可以看一下 项目代码，这个项目使用 pytest-benchmark 来做性能评测。要运行性能评测，你需要安装 nox。

// 先安装项目依赖
pip install .
// 再安装 nox 工具
pip install nox

然后在项目根目录下运行：

nox -s bench

稍等片刻，我们会得到如下输出：

$ nox -s bench
nox > Running session bench
nox > Creating virtual environment (virtualenv) using python3 in .nox/bench
nox > python -m pip install '.[dev]'
nox > pytest --benchmark-enable
=========================================================== test session starts ============================================================
platform linux -- Python 3.10.12, pytest-7.4.3, pluggy-1.3.0
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/mike/works/pyo3works/word-count
configfile: pyproject.toml
plugins: benchmark-4.0.0
collected 2 items

tests/test_word_count.py ..                                                                                                          [100%]

----------------------------------------------------------------------------------------- benchmark: 2 tests ----------------------------------------------------------------------------------------
Name (time in ms)                         Min                Max               Mean            StdDev             Median               IQR            Outliers      OPS            Rounds  Iterations
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_word_count_rust_parallel         19.6740 (1.0)      50.2725 (1.0)      27.9049 (1.0)      9.7562 (2.27)     22.9368 (1.0)      8.5993 (1.13)          2;2  35.8360 (1.0)          13           1
test_word_count_python_sequential     78.4158 (3.99)     91.5790 (1.82)     84.4852 (3.03)     4.3023 (1.0)      84.7501 (3.69)     7.6210 (1.0)           5;0  11.8364 (0.33)         13           1
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

Legend:
  Outliers: 1 Standard Deviation from Mean; 1.5 IQR (InterQuartile Range) from 1st Quartile and 3rd Quartile.
  OPS: Operations Per Second, computed as 1 / Mean
============================================================ 2 passed in 2.78s =============================================================
nox > Session bench was successful.

可以看到，Rust并行化版本比Python顺序化版本的平均性能提升了3倍左右（并行化版本消耗时间为顺序化版本的1/3左右），实际上被处理的文件越大，这个性能提升就越明显。

业界知名PyO3绑定项目介绍

目前，Rust生态和Python之间已经有一些知名的项目使用 PyO3 实现绑定了，我选了几个比较不错的放在下面，你可以点击链接深入了解一下。

• OpenDAL 的 Python 绑定： https://github.com/apache/incubator-opendal/tree/main/bindings/python
• Polars 的 Python 绑定： https://github.com/pola-rs/polars/tree/main/py-polars
• Delta Lake 的 Python 绑定： https://github.com/delta-io/delta-rs/tree/main/python
• Arrow-Datafusion 的 Python 绑定： https://github.com/apache/arrow-datafusion-python
• tokenizers 的 Python 绑定： https://github.com/huggingface/tokenizers/tree/main/bindings/python

小结

这节课我们通过4个示例展示了如何在Rust中调用C函数，如何在C函数中调用Rust函数，以及在Python中调用Rust实现的模块和函数。并且通过对比Rust的实现版本与Python原生的实现版本，我们发现使用Rust写Python扩展性能更好。

Unsafe Rust编程大量出现在底层性能的优化和与其他语言交互的场景中，体现了Rust这门语言一个奇特的地方：Rust语言本身对安全性和性能的追求，不但让Rust本身脱颖而出，而且还反哺了C语言的生态，用Rust重新实现一部分模块，同时还能无缝地接入之前的C系统中，提升了之前C系统的安全性。并且这是一种渐进的做法，非常有利于历史遗留系统的迭代更新。

此外，Unsafe Rust编程还可以赋能其他语言社区，将Rust的澎湃动力和安全扩展到了Python、JavaScript等其他语言。它们之前的扩展使用C/C++编写，同样会存在安全性和稳定性的问题，而Rust解决了这个问题。

不过涉及到Unsafe Rust和FFI的场景，想要封装一个真正好用的库或扩展，还有大量的细节知识待你探索。不管怎样，我们已经踏上了安全提升和性能提升之路，有Rust为你输出动力，保驾护航，你便可以放开手脚，大胆去干。

思考题

请你聊一聊，C语言中的char与Rust中的char的区别在哪里？C语言的字符串与Rust中的String区别在哪里？如何将C语言的字符串类型映射到Rust的类型中来？欢迎你把自己的思考分享到评论区，也欢迎你把这节课的内容分享给其他朋友，我们下节课再见！