Mojo 学习 —— 装饰器

Mojo 学习 —— 装饰器

Mojo 装饰器是一个高阶函数，用于修改或扩展结构体、函数或其他代码。你只需在代码（如结构体）上方添加装饰器（如 @value），而无需实际调用高阶函数。然后， Mojo 编译器会在编译时使用装饰器函数来修改代码。

目前还不支持创建自定义装饰器。

@always_inline

您可以在任何函数上添加 @always_inline 装饰器，Mojo 编译器会将该函数的主体直接 “内联” 到调用函数的主体中（复制）。

这样可以消除函数调用时可能带来的性能损失。Mojo 编译器通常会在可以提高性能的地方自动这样做，但使用装饰器可以强制编译器执行內联操作。但是在每个调用点都重复函数，会增加编译二进制文件的大小。

例如

@always_inline
fn add(a: Int, b: Int) -> Int:
    return a + b

print(add(1, 2))

Mojo 编译此程序时不会将 add 函数添加到调用栈，而是直接在 print 的调用位置执行添加操作，相当于：

print(1 + 2)

@always_inline(“nodebug”)

您也可以使用带有 "nodebug" 参数的装饰器，其效果与内联函数相同，但不包含调试信息。可以在调试时减少构建二进制文件的大小。

@nonmaterializable

@nonmaterializable 装饰器，可以声明该类型只能存在于 parameter 域中（即运行时参数，只能用于元编程，不能作为运行时类型）。如果该类型的实例转换到运行时域，则该装饰器会声明它在运行时域会变成什么类型。

使用方式为 @nonmaterializable(TargetType)，其中 TargetType 是对象在成为运行时值时应转换成的类型（必须声明 TargetType）。

例如，当一个结构体被标记为 @nonmaterializable(HasBool)，那么它从编译时值变为运行时值的时，会自动转换为 HasBool 类型。

@value
@register_passable("trivial")
struct HasBool:
    var x: Bool

    fn __init__(x: Bool) -> Self:
        return Self {x: x}

    @always_inline("nodebug")
    fn __init__(nms: NmStruct) -> Self:
        return Self {x: True if (nms.x == 77) else False}

@value
@nonmaterializable(HasBool)
@register_passable("trivial")
struct NmStruct:
    var x: Int

    @always_inline("nodebug")
    fn __add__(self: Self, rhs: Self) -> Self:
        return NmStruct(self.x + rhs.x)

alias still_nm_struct = NmStruct(1) + NmStruct(2)
var converted_to_has_bool = still_nm_struct

NmStruct 类型可以在 parameter 域中使用，但其 converted_to_has_bool 实例在作为运行时值具体化时会转换为 HasBool

非实体化结构体的所有方法必须注释为 @always_inline，并且必须在 parameter 域中可计算。

@parameter

您可以在 if 语句或嵌套函数中添加 @parameter 装饰器，以便在编译时运行该代码。

参数化 if 语句

您可以将 @parameter 添加到任何基于有效 parameter 表达式（即在编译时求值的表达式）的 if 条件中。这样可以确保只有 if 语句的实时分支被编译到程序中，从而减少最终二进制文件的大小。例如

@parameter
if True:
    print("this will be included in the binary")
else:
    print("this will be eliminated at compile time")

参数化闭包

你可以在嵌套函数上添加 @parameter 来创建一个 “参数化” 捕获闭包。这意味着您可以创建一个从外部作用域捕获值的闭包函数（不管它们是变量还是 parameter），然后使用该闭包作为 parameter。例如：

fn use_closure[func: fn(Int) capturing -> Int](num: Int) -> Int:
    return func(num)

fn create_closure():
    var x = 1

    @parameter
    fn add(i: Int) -> Int:
        return x + i

    var y = use_closure[add](2)
    print(y)

create_closure()

如果不使用 @parameter 装饰器，编译器就会出错，提示 “不能在调用参数中使用动态值”，即 use_closure[add](2) 调用时，add 闭包仍然是动态的。

这是一个不安全的功能，因为目前没有对引用捕获的生存期进行建模。

@register_passable

您可以在结构体上添加 @register_passable 装饰器，告诉 Mojo 该类型可以在寄存器（如 CPU 寄存器；取决于底层架构的细节）中传递。

对于像整数或浮点数这样的微小数据类型，都是寄存器可传递类型，始终通过值传递，而不能通过引用传递。

通常 @register_passable 装饰器不会改变类型的基本行为，它仍然需要一个 __init__ 和 __copyinit__ 方法来实现可复制（如果需要，它还可以有一个 __del__ 方法）。例如

@register_passable
struct Pair:
    var a: Int
    var b: Int

    fn __init__(inout self, one: Int, two: Int):
        self.a = one
        self.b = two

    fn __copyinit__(inout self, existing: Self):
        self.a = existing.a
        self.b = existing.b

fn test_pair():
    var x = Pair(5, 10)
    var y = x

    print(y.a, y.b)
    y.a = 10
    y.b = 20
    print(y.a, y.b)

test_pair()
# 5 10
# 10 20

寄存器可传递类型现在接受标准的生命周期方法。但是在 Mojo 24.1 之前，寄存器可传递类型需要一个返回值的非标准构造函数，语法如下：return Self{a: one, b: two}。旧的语法仍受支持，但将在未来的版本中删除。

还有一些要注意的地方：

• @register_passable 类型不能持有非 @register_passable 类型的实例。
• @register_passable 类型没有可预测的标识，因此 self 指针不稳定/不可预测（例如在哈希表中）。
• @register_passable 类型的 argument 和返回值直接暴露给 C 和 C++，而不是以指针方式传递。
• @register_passable 类型不能使用 __moveinit__ 构造函数（在寄存器中传递的值不能通过引用传递）。

@register_passable(“trivial”)

大多数使用 @register_passable 的类型只是 “比特包”，我们称之为 trivial类型。这些类型非常简单，可以复制、移动和销毁，不需要任何自定义构造函数或析构函数。

对于这些类型，你可以添加 trivial 参数，Mojo 就会根据类型的需要，合成所有的生命周期方法。例如

@register_passable("trivial")
struct Pair:
    var a: Int
    var b: Int

这与 @value 装饰器类似，但在使用 @register_passable("trivial") 时，只能自定义的生命周期方法是 __init__ 构造函数（但并非必须），不能定义任何复制或移动构造函数或析构函数。

trivial 类型包括

• 算术类型，如 Int、Bool 和 Float64 等。
• 指针（地址值是微不足道的，而不是所指向的数据）
• 其他 trivial 类型的数组，包括 SIMD。

@staticmethod

您可以在结构体的方法上添加 @staticmethod 装饰器来声明其为静态方法。

from tensor import Tensor
from pathlib import Path


struct MyStruct:
    var data: Tensor[DType.int8]

    fn __init__(inout self):
        self.data = Tensor[DType.int8]()

    fn __moveinit__(inout self, owned existing: Self):
        self.data = existing.data ^

    @staticmethod
    fn load_from_file(file_path: Path) raises -> Self:
        var new_struct = MyStruct()
        new_struct.data = file_path.read_bytes()
        return new_struct ^

与实例方法不同，静态方法不使用隐式 self 参数。静态方法不与结构体的特定实例相连，因此无法访问实例数据。

@unroll

您可以在任何循环（如 for 和 while）上添加 @unroll 装饰器，使 Mojo 编译器完全或以给定的形式展开循环。

例如，编译器会将下面的循环 10 展开为对 print 的 10 次连续调用

@unroll
for i in range(10):
    print(i)

装饰器还接受一个参数，用于指定一次要展开多少个循环。

例如，上面的展开相当于 @unroll(10)，即展开所有迭代。因此，如果传递一个小于循环次数的数字，编译器就会创建多个展开。例如

@unroll(2)
for i in range (10):
    print(i)

就相当于

for i in range(0, 10, 2):
    print(i)
    print(i+1)

编译器只有在下面情况下才能展开循环：

• 循环的下限、上限和归纳步长是编译时常量（运行时不会变化）。
• 循环中不会出现使循环计数在运行时发生改变的情况。

与 unroll 函数比较

Mojo 标准库还定义了一个 unroll 函数，它可以将你要重复调用的函数展开。

它们之间的不同之处在于：

• @unroll 装饰器只对循环表达式起作用，而 unroll 函数对函数起作用。
• @unroll 装饰器中决定循环如何展开的变量值在编译开始时是未知的，而 unroll 函数需要使用编译时已知的循环次数。

@value

我们前面已经多次介绍过 @value 装饰器了，它可以自动生成生命周期方法，包括面向成员的 __init__、__copyinit__ 和 __moveinit__。例如

@value
struct MyPet:
    var name: String
    var age: Int

Mojo 看到 @value 装饰器后，注意到你没有任何构造函数，于是它为你自动生成这些构造函数，将相当于变成了下面的代码

struct MyPet:
    var name: String
    var age: Int

    fn __init__(inout self, owned name: String, age: Int):
        self.name = name^
        self.age = age

    fn __copyinit__(inout self, existing: Self):
        self.name = existing.name
        self.age = existing.age

    fn __moveinit__(inout self, owned existing: Self):
        self.name = existing.name^
        self.age = existing.age