Into特性实现参数多态的技巧

引言

Rust 是一门强类型且无重载的语言, 要直接实现一个函数接受不同类型参数不太可能. 一种办法是将所有可能的参数类型包在一个枚举中, 如

enum Params<T> {
    Int(i64),
    Str(String)
    Other<T>
}

fn poly(params: Params) {
    match params {
        // do something
    }
}

但这样的函数在调用其实不太方便, 因为用户必须先把参数用 enum 包一层, 如

poly(Params::Int(10086))

由于 Rust 不支持参数默认值, 在表达函数参数可选时难免包一层 Option<T>, 用户在调用这些函数时也需要包一层 Some(T). 有什么办法拿掉最外层的枚举吗?

实现

微信群里的 poem 作者老油条给出了一个实现, 如果函数需要一个可选参数, 可以这样写

fn poly<C: Into<Option<usize>>>(max: C) {
    let max: Option<usize> = max.into();
    dbg!(max);
}

虽然函数定义麻烦不少, 但你可以用 poly(None) 和 poly(100) 形式调用它, 这比 poly(Some(100)) 更方便些, 毕竟函数只需要定义一次, 但调用可能会发生很多次. 可以在 playground 中看到这个例子.

技巧在于 Into trait, 在标准库中 Option<T> 自动实现了 From<T>, 又因为 From<T> 自动实现了对应的 Into(标准库中From, Into相关说明), 所以 let max: Option<usize> = 100.into()是可以编译的.

扩展

恰好在我最近重构的 ws-tool 中也用到了类似的技巧.

ws-tool 是一个 Rust websocket server/client 工具, 在 websocket 协议里, 客户端和服务端发送的消息分为数据帧和控制帧, 数据帧包括 text/binary 两种类型, 控制帧则包括ping, pong和close, 前两者用于连接心跳控制, close用户关闭连接. 帧的类型位于帧头第一个byte, 被称为 opcode.

在发送 close 帧时可以在 payload 前两个 byte 放入一个 u16 状态码, 表示关闭原因.

我在设计时希望 ws-tool 写入的接口尽量少, 不要有 send_text(), send_binary 和 send_close 等形式函数, 甚至只用一个就能实现.

在目前版本中 ws-tool 的 StringCodec 和 BinaryCodec 都只需要一个 send 函数就能实现写入各种帧, 以 WsStringCodec 为例, 这个 codec 可以方便地提供写入字符串等功能.

// 写入 String, 此时 opcode 为 text
ws.send("hello, world!".to_string());

// 写入 close, 告知对端关闭连接
ws.send((1000, "we are done."));

// 写入其他类型帧, 对于 ping/pong, 传入的 string 将作为其 payload 传递更多信息
ws.send((Opcode::Ping, "are you ok?"));

ws-tool 为实现以上设计目标, 首先需要定义一个 Message 结构体, 这个结构体定义代表所有可能的消息类型

pub struct Message<T: AsRef<[u8]> + DefaultCode> {
    pub code: OpCode,
    pub data: T,
    pub close_code: Option<u16>,
}

因为在将 websocket 帧写入数据流时都需要将其当作一个 u8 数组 slice, 所以会给 T 加上 AsRef<[u8]> trait bound. 此外 T 还必须实现 DefaultCode 这个 trait, 这是因为我们无法确定 T 对应的 Opcode, 所有需要如下 trait

pub trait DefaultCode {
    fn default_code(&self) -> OpCode;
}

接着只需要为 String, &str, &[u8], BytesMut 等常见类型实现 DefaultCode 就好了.

接着我们需要为 Message 实现以下几种 From 实现

From<T>, 这时候需要调用 DefaultCode 提供的方法获取对应的 opcode
From<(u16, T)> 希望发送一个关闭帧
From<(OpCode, T)> 希望发送其他类型帧

然后在 send 函数定义时使用类似的技巧

pub fn send<T: Into<Message<String>>>(&mut self, msg: T) -> Result<usize, WsError> {
    let msg: Message<String> = msg.into();
    if let Some(close_code) = msg.close_code {
        if msg.code == OpCode::Close {
            self.frame_codec.send(
                msg.code,
                vec![&close_code.to_be_bytes()[..], msg.data.as_bytes()],
            )
        } else {
            self.frame_codec.send(msg.code, msg.data.as_bytes())
        }
    } else {
        self.frame_codec.send(msg.code, msg.data.as_bytes())
    }
}

可以看到 send 函数不是很复杂, 第一步进行类型转换, 接着判断是否传入关闭状态码, 调用底层帧读写结构体写入数据.

总结

使用 Into trait 优点在于调用时无需手动包一层, 更加方便, 缺点在于函数定义稍显麻烦, 且如果为某个类型实现了太多 From trait, 反而因太多选项导致用户无法确定该传入什么类型参数, 不能滥用.