如何从规范的 Rust SDK 生成 JavaScript 和 Python SDK

介绍

我们在PostgresML创建的工具功能强大且灵活。我们的工具几乎有无数种方式可用于支持矢量搜索、模型推理等等。像我们之前的许多公司一样，我们希望我们的用户能够享受我们工具的好处，同时不需要阅读大量文档的才能使用，因此我们构建了一个SDK。

我们是 Rust 的忠实粉丝（几乎我们的整个代码库都是用它编写的），我们发现使用它作为我们的主要语言可以让我们编写更安全的代码和更快地迭代我们的开发周期。但是，我们的大多数用户目前使用Python和JavaScript等语言。为 Rust 制作 SDK 是没有意义的，因为没有人会使用它。经过深思熟虑，我们最终确定了 SDK 的以下要求：

1. 它必须以多种语言本机提供
2. 所有语言必须具有与规范 Rust 实现相同的行为
3. 添加新语言应仅包含最小的开销

FFI 有什么问题

我们的 SDK 的第一个要求是它可以以多种语言原生使用，第二个要求是用 Rust 编写的。乍一看，这似乎是一个矛盾，但有一个非常著名的系统，用于用一种语言编写函数并在另一种语言中使用它们，称为 FFI（外部函数接口）。就我们的 SDK 而言，我们可以通过在 Rust 中编写 SDK 的核心逻辑来利用 FFI，并通过我们选择的语言中的 FFI 调用我们的 Rust 函数。不幸的是，这并没有达到我们想要的效果。以以下 Python 代码为例：

class Database:
     def __init__(self, connection_string: str):
          # Create some connection

     async def vector_search(self, query: str, model_id: int, splitter_id: int) -> str:
          # Do some async search here
          return result

async def main():
     db = Database(CONNECTION_STRING)
     result = await db.vector_search("What is the best way to do machine learning", 1, 1)
     if result != "PostgresML":
          print("The model still needs more training")
     else:
          print("The model is ready to go!")

我们的 SDK 的要求之一是我们用 Rust 编写它。具体来说，在这种情况下， class Database 和它的方法应该用 Rust 编写，并通过 FFI 在 Python 中使用。不幸的是，仅在 Rust 中执行此操作是不可能的。在上面的代码中，有两个限制是我们不能超越的：

• FFI 没有 Python 类的概念
• FFI 没有 Python 异步的概念

我们可以围绕我们的 FFI 编写自己的 Python 包装器，但这违背了要求 3：添加新语言应该只包括最小的开销。将 Rust SDK 中的每个更新翻译成我们添加的每种语言的包装器并不是最小的开销。

使用 pyO3 和 Neon

Pyo3 和 Neon 是 Rust 库，有助于为 Python 和 JavaScript 构建原生模块。它们提供的系统允许我们编写 Rust 代码，这些代码可以与 Python 和 JavaScript 中的异步代码和本机类无缝交互，绕过原版 FFI 施加的限制。

让我们看一些 Rust 代码，它使用 Pyo3 创建了一个 Python 类，使用 Neon 创建了一个 JavaScript 类。为了便于使用，假设我们在 Rust 中有以下结构：

struct Database{
    connection_string: String
}

impl Database {
     pub fn new(connection_string: String) -> Self {
          // The actual connection process has been removed 
          Self {
               connection_string
          }
     }

     pub async fn vector_search(&self, query: String, model_id: i64, splitter_id: i64) -> String {
          // Do some async vector search
          result
     }
}

以下是以与 Pyo3 和 Neon 一起使用的增强代码：

use pyo3::prelude::*;

struct Database{
     connection_string: String
}

#[pymethods]
impl Database {
     #[new]
     pub fn new(connection_string: String) -> Self {
          // The actual connection process has been removed
          Self {
               connection_string
          }
     }

     pub fn vector_search<'a>(&self, py: Python<'a>, query: String, model_id: i64, splitter_id: i64) -> PyResult<&'a PyAny> {
          pyo3_asyncio::tokio::future_into_py(py, async move {
               // Do some async vector search
               Ok(result)
          })
     }
}

/// A Python module implemented in Rust.
#[pymodule]
fn pgml(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::()?;
     Ok(())
}

use neon::prelude::*;

struct Database{
     connection_string: String
}

impl Database {
    pub fn new<'a>(mut cx: FunctionContext<'a>) -> JsResult<'a, JsObject> {
        // The actual connection process has been removed
        let arg0 = cx.argument::(0usize as i32)?;
        let arg0 = ::from_js_type(&mut cx, arg0)?;
        let x = Self {
            connection_string: arg0
        };
        x.into_js_result(&mut cx)
    }

    pub fn vector_search<'a>(mut cx: FunctionContext<'a>) -> JsResult<'a, JsPromise> {
        let this = cx.this();
        let s: neon::handle::Handle<
            neon::types::JsBox>,
        > = this.get(&mut cx, "s")?;
        let wrapped = (*s).deref().borrow();
        let wrapped = wrapped.wrapped.clone();
        let arg0 = cx.argument::(0)?;
        let arg0 = ::from_js_type(&mut cx, arg0)?;
        let arg1 = cx.argument::(1);
        let arg1 = ::from_js_type(&mut cx, arg1);
        let arg2 = cx.argument::(2);
        let arg2 = ::from_js_type(&mut cx, arg2);
        let channel = cx.channel();
        let (deferred, promise) = cx.promise();
        deferred
            .try_settle_with(
                &channel,
                move |mut cx| {
                    // Do some async vector search
                    result.into_js_result(&mut cx)
                },
            )
            .expect("Error sending js");
        Ok(promise)
    }

    fn into_js_result<'a, 'b, 'c: 'b, C: Context<'c>>(self, cx: &mut C) -> JsResult<'b, Self::Output> {
        let obj = cx.empty_object();
        let s = cx.boxed(std::cell::RefCell::new(self));
        obj.set(cx, "s", s)?;
        let f: Handle = JsFunction::new(
            cx,
            Database::new,
        )?;
        obj.set(cx, "new", f)?;
        let f: Handle = JsFunction::new(
            cx,
            Database::vector_search,
        )?;
        Ok(obj)
    }
}

impl neon::types::Finalize for Database {}

/// A JavaScript module implemented in Rust.
#[main]
fn main(mut cx: ModuleContext) -> NeonResult<()> {
    cx.export_function("newDatabase", Database::new)?;
    Ok(())
}

自动将纯 Rust 代码换为 py03 和 Neno 兼容 Rust 代码

我们已经成功地在 Rust 中编写了一个原生的 Python 和 JavaScript 模块。然而，我们的目标远未完成。我们的愿望是在 Rust 中编写一次 SDK，并使其以我们目标的任何语言提供。虽然上述内容使其在 Python 和 JavaScript 中可用，但它都不再是一个有效的 Rust 库，并且需要大量手动编辑才能在两种语言中提供。

我们真正想要的是编写我们的 Rust 库而不用担心任何翻译，并应用一些宏将其自动转换为 Pyo3 和 Neon 代码。这听起来像是程序宏的完美用途。如果你不熟悉宏，我真的推荐你阅读 The Little Book of Rust Macros 它是免费的，快速阅读，并提供了一个很棒的宏介绍。

我们创建一个如下所示的流程：

让我们稍微编辑一下我们之前定义的结构：

#[custom_derive_class]
struct Database{
     connection_string: String
}

#[custom_derive_methods]
impl Database {
     pub fn new(connection_string: String) -> Self {
          // The actual connection process has been removed 
          Self {
               connection_string
          }
     }

     pub async fn vector_search(&self, query: String, model_id: i64, splitter_id: i64) -> String {
          // Do some async vector search
          result
     }
}

请注意，有两个我们以前从未见过的新宏： custom_derive_class 和 custom_derive_methods 。这两个都是我们需要编写的宏。

custom_derive_class 为我们的 Database 结构创建包装器。让我们展示我们 custom_derive_class 生成的扩展代码：

#[pyclass]
struct DatabasePython {
     wrapped: Database
}

impl From for DatabasePython {
    fn from(w: Database) -> Self {
        Self { wrapped: w }
    }
}

struct DatabaseJavascript {
    wrapped: Database
}

impl From for DatabaseJavascript {
    fn from(w: Database) -> Self {
        Self { wrapped: w }
    }
}

impl IntoJsResult for Database {
    type Output = neon::types::JsObject;
    fn into_js_result<'a, 'b, 'c: 'b, C: neon::context::Context<'c>>(
        self,
        cx: &mut C,
    ) -> neon::result::JsResult<'b, Self::Output> {
        DatabaseJavascript::from(self).into_js_result(cx)
    }
}

这里有几件重要的事情发生：

1. 我们的 custom_derive_class 宏为我们定义的每种语言创建一个新结构。
2. 派生的 Python 结构会自动实现 pyclass
3. 因为 Neon 没有 pyclass 宏的版本，所以我们实现了自己的特征 IntoJsResult 来在原版 Rust 类型和 Neon Rust 之间进行一些转换

创建像上面这样的宏实际上非常简单。下面的代码显示了如何为 Python 版本完成此操作。

#[proc_macro_derive(custom_derive)]
pub fn custom_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let parsed = parse_macro_input!(input as DeriveInput);
     let name_ident = format_ident!("{}Python", parsed.ident);
     let wrapped_type_ident = parsed.ident;
     let expanded = quote! {
          #[pyclass]
          pub struct #name_ident {
               wrapped: #wrapped_type_ident
          }
     };
     proc_macro::TokenStream::from(expanded)
}

让我们看一下宏 custom_derive_methods 在 Database 结构上使用时生成的扩展代码：

#[pymethods]
impl DatabasePython {
     #[new]
     pub fn new(connection_string: String) -> Self {
          // The actual connection process has been removed
          Self::from(Database::new(connection_string))
     }

     pub fn vector_search<'a>(&self, py: Python<'a>, query: String, model_id: i64, splitter_id: i64) -> PyResult<&'a PyAny> {
          let wrapped = self.wrapped.clone();
          pyo3_asyncio::tokio::future_into_py(py, async move {
               // Do some async vector search
               let x = wrapped.vector_search(query, model_id, splitter_id).await;
               Ok(x)
          })
     }
}

impl DatabaseJavascript {
    pub fn new<'a>(
        mut cx: neon::context::FunctionContext<'a>,
    ) -> neon::result::JsResult<'a, JsObject> {
        let arg0 = cx.argument::(0usize as i32)?;
        let arg0 = ::from_js_type(&mut cx, arg0)?;
        let x = Database::new(&arg0);
        let x = x.expect("Error in rust method");
        let x = Self::from(x);
        x.into_js_result(&mut cx)
    }

    pub fn vector_search<'a>(
        mut cx: neon::context::FunctionContext<'a>,
    ) -> neon::result::JsResult<'a, neon::types::JsPromise> {
        use neon::prelude::*;
        use core::ops::Deref;
        let this = cx.this();
        let s: neon::handle::Handle<
            neon::types::JsBox>,
        > = this.get(&mut cx, "s")?;
        let wrapped = (*s).deref().borrow();
        let wrapped = wrapped.wrapped.clone();
        let arg0 = cx.argument::(0)?;
        let arg0 = ::from_js_type(&mut cx, arg0)?;
        let arg1 = cx.argument::(1);
        let arg1 = ::from_js_type(&mut cx, arg1);
        let arg2 = cx.argument::(2);
        let arg2 = ::from_js_type(&mut cx, arg2);
        let channel = cx.channel();
        let (deferred, promise) = cx.promise();
        deferred
            .try_settle_with(
                &channel,
                move |mut cx| {
                    let runtime = crate::get_or_set_runtime();
                    let x = runtime.block_on(wrapped.vector_search(&arg0, arg1, arg2));
                    let x = x.expect("Error in rust method");
                    x.into_js_result(&mut cx)
                },
            )
            .expect("Error sending js");
        Ok(promise)
    }
}

impl IntoJsResult for DatabaseJavascript {
    type Output = neon::types::JsObject;
    fn into_js_result<'a, 'b, 'c: 'b, C: neon::context::Context<'c>>(
        self,
        cx: &mut C,
    ) -> neon::result::JsResult<'b, Self::Output> {
        use neon::object::Object;
        let obj = cx.empty_object();
        let s = cx.boxed(std::cell::RefCell::new(self));
        obj.set(cx, "s", s)?;
        let f: neon::handle::Handle = neon::types::JsFunction::new(
            cx,
            DatabaseJavascript::new,
        )?;
        obj.set(cx, "new", f)?;
        let f: neon::handle::Handle = neon::types::JsFunction::new(
            cx,
            DatabaseJavascript::vector_search,
        )?;
        obj.set(cx, "vector_search", f)?;
        Ok(obj)
    }
}

impl neon::types::Finalize for DatabaseJavascript {}

您会注意到这与我们已经展示的代码非常相似， DatabaseJavascript 和 DatabasePython 结构只是在 Database 结构上调用它们各自的方法。

宏实际上是如何工作的？我们可以将 custom_derive_methods 宏代码生成分为三个不同的阶段：

• 方法解构
• 签名翻译
• 方法重构

方法解构

利用 syn 库，我们解析 Database 结构块并迭代各个方法，将它们解析并 impl 为我们自己的类型：

pub struct GetImplMethod {
     pub exists: bool,
     pub method_ident: Ident,
     pub is_async: bool,
     pub method_arguments: Vec<(String, SupportedType)>,
     pub receiver: Option,
     pub output_type: OutputType,
}

这里的 OutputType 和 SupportType 是我们支持的类型的自定义枚举，如下所示：

pub enum SupportedType {
     Reference(Box),
     str,
     String,
     Vec(Box),
     HashMap((Box, Box)),
     Option(Box),
     Tuple(Vec),
     S, // Self
     i64,
     i32,
     f64,
     // Other omitted types
}

pub enum OutputType {
     Result(SupportedType),
     Default,
     Other(SupportedType),
}

签名翻译

我们必须将签名转换为 Pyo3 和 Neon 期望的 Rust 代码。这意味着调整参数、异步声明和输出类型。现在我们已经解构了该方法，这实际上非常简单。例如，下面是一个翻译 Python 输出类型的简单示例：

fn convert_output_type(
     ty: &SupportedType,
     method: &GetImplMethod,
) -> (
     Option
) {
     if method.is_async {
          Some(quote! {PyResult<&'a PyAny>})
     } else {
          let ty = t
               .to_type()
               .unwrap();
          Some(quote! {PyResult<#ty>})
    }
}

方法重构

现在，我们有了以 Pyo3 和 Neon 创建本机模块所需的格式重建方法所需的所有信息。

实际的重建是相当无聊的，主要是用一堆 if else 语句编写和组合使用 quota 库的 Token Stream，所以为了简洁起见，我们将省略它。对于好奇的人，这里有一个指向我们实际实现的链接：github。

以上三个阶段的全部可以总结为这个非常抽象的函数（特定于Python，尽管它对于JavaScript几乎相同）：

fn do_custom_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let parsed_methods = parse_methods(input);
     let mut methods = Vec::new();
     for method in parsed_methods {
          // Destructure Method
          let destructured = destructure(method);
          // Translate Signature
          let signature = convert_signature(&destructured);
          // Restructure Method 
          let method = create_method(&destructured, &signature);
          methods.push(method);
     }
     // This is the actual Rust impl block we are generating
    proc_macro::TokenStream::from(quote! {
        #[pymethods]
        impl DatabasePython {
            #(#methods)*
        }
    })
}

总结和未来展望

以上所有内容都展示了我们如何同时创建原生 Rust、Python 和 JavaScript 库。上述方法有一些怪异的地方，但我们仍在积极开发和改进我们的设计。

虽然我们的宏目前专门用于我们拥有的特定用例，但我们正在探索将它们作为独立的库进行泛化和发布的想法，以帮助每个人在 Rust 和他们选择的语言中编写原生库。我们还计划添加对更多语言的支持，我们很乐意听到有关你选择的语言的反馈。

感谢您的阅读！