Mojo 学习 —— 导入外部函数

Mojo 学习 —— 导入外部函数

前言

Mojo 支持从动态链接库中导入函数，可以很方便地使用其他编程语言来对 Mojo 功能进行扩展。

这一功能主要通过 sys 标准库中的 ffi 模块中的结构体和函数实现的。

from sys.ffi import DLHandle, external_call

DLHandle

该结构体可以代表可加载和卸载的动态链接库。即在初始化时加载，在删除对象时卸载。

可以在构建对象时指定动态链接库的路径。例如，从系统中查找 SDL 的动态链接库

fn get_sdl_lib_path() -> StringLiteral:
    if (info.os_is_linux()):
        var lib_path = '/usr/lib/x86_64-linux-gnu/libSDL2.so'
        try:
            with open('/etc/os-release', 'r') as f:
                var release = f.read()
                if (release.find('Ubuntu') < 0):
                    lib_path = '/usr/lib64/libSDL2.so'
        except:
            print("Can't detect Linux version")
        return lib_path
    if (info.os_is_macos()):
        return '/opt/homebrew/lib/libSDL2.dylib'
    return ""

参考 mojo-sdl，一个将 Mojo 绑定到 SDL2 的库。

当然，我们也可以使用自定义的动态链接库。例如，我们在本地创建一个 demo.c 文件，用于计算浮点向量中的最大值

#include <stdio.h>
#include <stdlib.h>

double array_max(double *input, size_t nelem) {
    double max = 0;
    for (size_t i = 0; i < nelem; i++) {
        if (input[i] > max) {
            max = input[i];
        }
    }
    return max;
    printf("max: %f", max);
}

然后将该程序编译成动态链接库

gcc -shared -fPIC demo.c -o libdemo.so

然后就可以在 Mojo 中导入使用了

from sys.ffi import DLHandle, external_call

alias array_max_type = fn(UnsafePointer[Float64], Int64) -> Float64

fn main():
    var handle = DLHandle('./libdemo.so')
    var array_max = handle.get_function[array_max_type]('array_max')

    var data = List[Float64](-1.732, 1.414, 3.1415, 0.618, 2.732)

    print('Max:', array_max(data.data, len(data)))

我们创建 DLHandle 对象时，可以指定绝对路径或相对路径。

调用该对象的 get_function 函数来加载动态链接库中的函数，将要导入的函数名称作为字符串参数。

get_function 需要指定函数的形式，这里我们使用 alias 定义了函数形式。

运行上面的代码将会输出向量的最大值。例如

Max: 3.1415000000000002

external_call

我们还可以使用 external_call 函数来导入外部函数。

对于 C 标准库中的函数（libc）可以直接导入。例如

alias c_void = UInt8
alias c_char = UInt8
alias c_int = Int32
alias c_size_t = Int

fn to_char_ptr(s: String) -> Pointer[c_char]:
    # Hello World0
    var ptr = Pointer[c_char]().alloc(len(s) + 1)
    for i in range(len(s)):
        ptr.store(i, ord(s[i]))
    return ptr

fn c_charptr_to_string(s: Pointer[c_char]) -> String:
    return String(s.bitcast[Int8](), strlen(s) + 1)

fn strlen(s: Pointer[c_char]) -> c_size_t:
    return external_call["strlen", c_size_t, Pointer[c_char]](s)
    
fn main():
    var name = String('Hello World')
    var name_ptr = to_char_ptr(name)
    print('ptr len:', strlen(name_ptr))
    print('String:', c_charptr_to_string(name_ptr))
# ptr len: 11
# String: Hello World

注意，将指针转换为字符串时需要包含一个终止符（0）

也可以使用 external_call 从第三方静态库中导入函数。这里需要设置 ld 参数并使用 mojo build。

这里我们参考 mojo-ffi 里面的内容。直接使用里面的两个文件 ld 和 mojoc，并使用 Makefile 封装流程。

我将 Makefile 修改为

all: call_demo libdemo.a libdemo.so run

libdemo.a: demo.c
	gcc -static -c demo.c -o libdemo.a

libdemo.so: demo.c
	gcc -shared -fPIC demo.c -o libdemo.so

call_demo: call_demo.mojo libdemo.a
	./scripts/mojoc call_demo.mojo -Slibdemo.a -o call_demo

run: FORCE call_demo
	./call_demo

clean:
	rm -f libdemo.a call_demo

FORCE:

其中，同级目录下 script 文件夹中放置了 mojoc 和 ld 文件，c 代码在 demo.c 中，可以根据需要修改，call_demo.mojo 文件中写对应的 Mojo 代码。

例如，我们在 demo.c 中添加如下代码

#include <stdio.h>
#include <stdlib.h>

void say_hello(const char *name) {
    printf("Hello %s", name);
}

在 call_demo.mojo 中测试，并在当前目前下运行 make 执行程序

from sys.ffi import DLHandle, external_call

alias c_char = UInt8

# `String` to `Pointer[c_char]`
fn to_char_ptr(s: String) -> Pointer[c_char]:
    var ptr = Pointer[c_char]().alloc(len(s))
    for i in range(len(s)):
        ptr.store(i, ord(s[i]))
    return ptr

fn say_hello(name: String):
    return external_call['say_hello', NoneType](to_char_ptr(name))
    
fn main():
    say_hello('World')
# Hello World

我们再来一个复杂的例子，涉及结构体与指针。

我们用 C 编写一个链表结构，并添加几个函数，用于将向量转换为链表以及链表转换为向量。例如

typedef struct Node {
    int data;
    struct Node* next;
} Node;

Node* create_node(int data) {
    Node* new_node = (Node*)malloc(sizeof(Node));
    new_node->data = data;
    new_node->next = NULL;
    return new_node;
}

Node *from_values(int* list, size_t size) {
    Node *head = create_node(list[0]);
    Node *ptr = head;
    for (int i = 1; i < size; ++i) {
        ptr->next = create_node(list[i]);
        ptr = ptr->next;
    }
    return head;
}

int *get_values(Node *head) {
    int *vec;
    int i = 0;
    Node *ptr = head;
    while (ptr != NULL) {
        vec[i] = ptr->data;
        ptr = ptr->next;
        i++;
    }
    return vec;
}

在 Mojo 中使用这些函数。首先定义结构体和函数

alias c_int = Int32

@register_passable('trivial')
struct Node:
    var data: c_int
    var next: Pointer[Self]

fn from_values(data: UnsafePointer[c_int], size: c_int) -> Pointer[Node]:
    return external_call['from_values', Pointer[Node], UnsafePointer[c_int], c_int](data, size)

fn get_value(head: Pointer[Node]) -> Pointer[c_int]:
    return external_call['get_values', Pointer[c_int], Pointer[Node]](head)

from_values 可以将 Mojo 中的整数向量装换为链表，并返回链表的头部指针。

get_value 用于从头部指针中遍历获取链表的值，并保存成了整数指针向量。

使用示例

def main():
    var vec = List[c_int](1, 2, 3, 4, 5, 6)
    # 创建并获取链表头指针
    var node = from_values(vec.data, 6)
    # 遍历链表
    print('Start', end=' ')
    while node:
        print("->", node[].data, end=' ')
        node = node[].next
    print('-> End')
    # 避免数据提前释放
    _ = vec
# ./scripts/mojoc call_demo.mojo -Slibdemo.a -o call_demo
# ./call_demo
# Start -> 1 -> 2 -> 3 -> 4 -> 5 -> 6 -> End

或者

def main():
    var vec = List[c_int](1, 2, 3, 4, 5, 6)
    # 创建并获取链表头指针
    var node = from_values(vec.data, 6)
    # 根据头指针遍历数据
    var res = get_value(node)
    for i in range(len(vec)):
        print(res[i], end=' ')
    # 避免数据提前释放
    _ = vec
# ./scripts/mojoc call_demo.mojo -Slibdemo.a -o call_demo
# ./call_demo
# 1 2 3 4 5 6

注意

我们在最后一行添加了 _ = vec，避免 vec 提前释放导致数据返回结果出现问题。具体可以返回去看生命周期那一章。

使用这种方式可以很容易将之前的 C/C++ 代码封装为 Mojo 代码，扩展 Mojo 的功能。