Tagged Union Subsets with Comptime in Zig

Mitchell Hashimoto

在 Zig 中用 Comptime 实现标签联合体的子集

原文由 Mitchell Hashimoto 发布,订阅该博客

本文是 Zig comptime 用例系列的一部分。

Zig 支持标签联合体(tagged union)1,而且如果你对标签联合体使用 switch 却没有处理所有可能的情况,Zig 编译器会直接报错。这是一个非常棒的特性,因为它能在你为联合体新增分支时帮你避免遗漏。许多语言也都提供了类似的功能。

然而,有时你只想处理标签联合体中的一个子集。你既不想在 switch 语句里使用通配分支来牺牲编译期的安全性,也不想为每一个不关心的分支都写上空操作。

在这篇文章中,我将通过一个真实的例子来说明这种场景,并展示如何利用 Zig 的 comptime 优雅地解决它。这个真实的例子来自我的终端项目

真实案例:按键绑定动作

下面是按键绑定动作在 Ghostty 中的表示方式。按键绑定动作指的是快捷键所触发的操作。例如:ctrl+n 可能会触发打开新窗口的动作。

pub const Action = union(enum) {
  quit: void,
  new_window: void,
  close_window: void,
  close_all_windows: void,
  open_config: void,
  reload_config: void,
  scroll_lines: i16,

  // Many more in reality, but we'll use the above for this blog post.
};

大多数按键绑定动作本身不携带任何数据(void),但有些会携带。例如,scroll_lines 动作会关联一个 i16 类型的值,用来指定向上(负数)或向下(正数)滚动的行数。

这种标签联合体的方式很好,因为在执行按键绑定动作的代码里,我可以对 Action 联合体做 switch,如果有分支没有处理,编译器就会报错。这样就能确保每当我新增一个按键绑定动作时,也会同步添加对应的处理逻辑。

pub fn performAction(action: Action) void {
  switch (action) {
    .quit => ...,
    .new_window => ...,
    .close_window => ...,
    .close_all_windows => ...,
    .open_config => ...,
    .reload_config => ...,
    .scroll_lines => |amount| ...,
  }
}

这就是 Ghostty 最初实现按键绑定动作的方式。

为什么需要联合体子集

最近,我需要新增一个作用域不在某个具体终端、而是整个 Ghostty 应用的按键绑定动作。为此,我想把 performAction 重构一下,让终端相关的动作由终端结构来处理,而应用级别的动作则由应用结构来处理。

问题在于,Action 联合体同时包含了终端相关的动作和应用级别的动作。起初,我是这样做的:

pub fn performAppAction(action: Action) void {
  switch (action) {
    .quit => ...,
    .close_all_windows => ...,
    .open_config => ...,
    .reload_config => ...,
    else => {},
  }
}

这样能跑,但你看到那个 else 了吗?那是一个通配分支,一旦有了通配分支,当我在 Action 联合体中新增一个动作却忘记在 performAppAction 中处理时,编译器就无法再报错提醒我了。

接着,我把所有不关心的动作都穷举出来,什么也不做:

pub fn performAppAction(action: Action) void {
  switch (action) {
    .quit => ...,
    .close_all_windows => ...,
    .open_config => ...,
    .reload_config => ...,

    // Do nothing for terminal-specific actions.
    .new_window,
    .close_window,
    .scroll_lines,
    => {},
  }
}

这样也能跑,但我有几个不满的地方:

  1. 很啰嗦。我必须把所有终端相关的动作全部列出来,尽管我根本不关心它们。示例中只展示了半打动作,但实际上要多得多。

  2. 从类型的角度看,我不喜欢 performAppAction 竟然还能接受终端相关的动作。如果类型签名本身就能自描述地表明它只接受应用级别的动作、并且清楚展示具体有哪些,那就好多了。

  3. 同样的问题在 performTerminalAction 函数中也存在(这里没有展示)。

我真正想要的是像下面这样的效果:

// NOT REAL ZIG CODE! THIS IS PSUEDO CODE!
const AppAction = union(enum) {
  quit: void,
  // ... other app actions ...
};

const TerminalAction = union(enum) {
  new_window: void,
  // ... other terminal actions ...
};

// Union of both. NOTE THIS IS NOT REAL CODE. THIS IS CONCEPTUAL.
const Action = AppAction | TerminalAction;

有人可能会问:“为什么不直接拆成两个独立的联合体呢?”答案是,代码库的其他部分需要处理统一的 Action。例如,配置解析器需要能够解析所有按键绑定的动作。

用 Comptime 实现联合体子集

我们可以用 Zig 的 comptime 来创建子集类型。

我们本可以反过来,先创建好子集类型,再用 comptime 来合成完整的联合体类型。作为留给读者的练习,不妨试试这种做法,效果同样好,只是会遇到不同的边界情况和人体工学上的取舍。

首先,我们需要知道每个动作属于哪个作用域。为此,可以引入一个新函数来判断:

pub const Scope = enum { app, terminal };

pub fn scope(action: Action) Scope {
  return switch (action) {
    .quit, .close_all_windows, .open_config, .reload_config => .app,
    .new_window, .close_window, .scroll_lines => .terminal,
  };
}

这是一个非常直观的函数。虽然因为要覆盖所有分支而显得有些冗长,但它集中在一处,易于维护。

接下来,我们就可以利用这个函数和 comptime 来创建子集类型:

/// Returns a union type that only contains actions that are scoped to
/// the given scope.
pub fn ScopedAction(comptime s: Scope) type {
  const all_fields = @typeInfo(Action).Union.fields;

  // Find all fields that are scoped to s
  var i: usize = 0;
  var fields: [all_fields.len]std.builtin.Type.UnionField = undefined;
  for (all_fields) |field| {
    const action = @unionInit(Action, field.name, undefined);
    if (action.scope() == s) {
      fields[i] = field;
      i += 1;
    }
  }

  // Build our union
  return @Type(.{ .Union = .{
    .layout = .auto,
    .tag_type = null,
    .fields = fields[0..i],
    .decls = &. {},
  } });
}

抱歉,我知道这有点“剩下的猫头鹰自己画”那种感觉2。我来一步步拆解。

pub fn ScopedAction(comptime s: Scope) type {

首先,在 Zig 中,按惯例大写字母开头的标识符表示类型。编译器并不强制要求这样做,这只是一种风格约定。当一个函数名是大写开头时,就意味着它会返回一个类型。

这是一个返回类型的函数。这正是 Zig 实现泛型的方式。由于 Zig 拥有 comptime 代码执行能力,而 comptime 执行的就是完整的语言本身,因此我们可以用在 comptime 接收参数并返回……类型的函数。

在这个例子中,我们的函数接收一个 Scope,并据此生成一个新的动作联合体,其中只包含属于该作用域的动作。

  const all_fields = @typeInfo(Action).Union.fields;

@typeInfo 这个内置函数3会在 comptime 返回关于某个类型的结构化信息。我们在 Action 上调用它,来提取联合体中的所有字段。

  // Find all fields that are scoped to s
  var i: usize = 0;
  var fields: [all_fields.len]std.builtin.Type.UnionField = undefined;
  for (all_fields) |field| {
    const action = @unionInit(Action, field.name, undefined);
    if (action.scope() == s) {
      fields[i] = field;
      i += 1;
    }
  }

接着,我们遍历 Action 联合体中的所有字段。对于每个字段,我们判断它是否属于我们感兴趣的作用域,如果是,就保存下来。

@unionInit 函数是一个内置函数,它会用一个在 comptime 已知的字段名和值来创建一个联合体的值。我们传入的值是 undefined(未初始化的内存),因为我们并不关心具体的值;我们的 scope() 函数只会查看标签。如果我们在 comptime 访问了未定义的内存,编译器会直接报错,所以这样做是安全的。

我们需要创建一个给定类型的值,这样才能在其上调用 scope()。注意,我们上面定义的 scope() 函数完全不知道自己是在 comptime 还是在运行时被执行的,它也不在乎,它就是一个普通函数。Comptime 酷就酷在这里。

  // Build our union
  return @Type(.{ .Union = .{
    .layout = .auto,
    .tag_type = null,
    .fields = fields[0..i],
    .decls = &. {},
  } });

@Type 这个内置函数会具现化一个类型(即创建一个新类型)。在这里,我们用前面收集到的字段——也就是符合目标作用域的那些字段——来创建一个联合体。

顺便一提,Zig 编译器中很多类型具现化的逻辑就是我实现的。如果你用了 @Type,很可能就命中了我写的至少一行代码。

最后,经过这一切,我们得到了一个只包含原始联合体子集的新类型。它看起来是这样的:

pub fn performAppAction(action: ScopedAction(.app)) void {
  switch (action) {
    .quit => ...,
    .close_all_windows => ...,
    .open_config => ...,
    .reload_config => ...,
  }
}

很漂亮。我们得到了一个简洁的 switch,它由编译器强制保证穷尽性,而且函数签名本身就自描述地表明它只接受应用级别的动作。当我们在 Action 联合体中新增动作时,会先收到编译器要求更新 scope() 函数的错误,之后其他地方就能自然地正常工作了。

另外,显而易见的是:所有这些逻辑都在编译期发生。没有任何运行时开销。由于我们只在编译期执行这些逻辑,因此也不会生成任何额外的运行时代码。

转换为子集

最后一个挑战:给定一个完整的 Action,我们如何把它转换成有类型的子集,例如 ScopedAction(.app)?下面的函数就是做这件事的:

pub fn scoped(self: Action, comptime s: Scope) ?ScopedAction(s) {
  switch (self) {
    inline else => |v, tag| {
      // Use comptime to prune out invalid actions
      if (comptime @unionInit(
        Action,
        @tagName(tag),
        undefined,
      ).scope() != s) return null;

      // Initialize our app action
      return @unionInit(
        ScopedAction(s),
        @tagName(tag),
        v,
      );
    },
  }
}

我们也来一步步拆解。

pub fn scoped(self: Action, comptime s: Scope) ?ScopedAction(s) {

这里有几个巧妙的地方。首先,这个函数同时接收运行时参数和 comptime 参数。编译器会强制要求 comptime 参数必须在编译期已知。

scope 参数必须是 comptime 已知的,因为我们要用它来决定返回值。注意我们在返回类型 ?ScopedAction(s) 中使用了 s。如果 s 前面没有 comptime 修饰,编译器就会报错,因为我们不能用一个运行时值来生成必须在编译期已知的信息(也就是类型)。

返回类型中的 ? 表示可选返回值,意味着返回值可能为空。对于这个函数来说,如果 self 对应的动作不属于我们要找的作用域,就会返回 null。

  switch (self) {
    inline else => |v, tag| {

inline else 是 Zig 中的一个特殊分支。它是一个会为每一个未被显式处理的分支都生成运行时代码的通配分支。因为它是 inline 的,编译器会把 tag 当作编译期已知的值来提供。这一点很重要,因为我们需要用它来判断动作的作用域。

inline else 一开始可能很难理解。我建议试着把 inline 关键字删掉,看看编译器会怎么提示。然后仔细想想编译器想要做什么,就能更清楚为什么必须要有 inline 了。

      // Use comptime to prune out invalid actions
      if (comptime @unionInit(
        Action,
        @tagName(tag),
        undefined,
      ).scope() != s) return null;

在这里,我们再次使用 @unionInit,用 comptime 已知的标签来创建一个联合体值。然后在其上调用 scope() 来判断它是否属于我们要找的作用域。如果不是,就返回 null。由于这个条件判断是 comptime 的,它会按条件编译掉块中剩余的部分,这样我们就不会因为为无效标签初始化子集动作而得到类型错误。

我曾专门写过一篇博客,讲解如何用 comptime 按条件禁用代码,其中就解释了这种模式。

      // Initialize our app action
      return @unionInit(
        ScopedAction(s),
        @tagName(tag),
        v,
      );

最后,我们返回初始化好的子集动作。

相比上一节,这个函数确实会有一些运行时开销。这个函数使用了一个运行时的 switch 语句来过滤掉不属于目标作用域的动作。我没有去反汇编看它最终被降成了什么代码,但我猜编译器会把其中很多东西都优化掉,因为转换前后的联合体具有相同的内存布局。我曾在一个合成测试里做了分析——在 Ghostty 中无限循环触发一个开销很小的快捷键——这个函数并没有出现在性能分析结果中,所以我并不担心。

总结

Zig 的 comptime 是一个非常强大的特性,能让你做到许多令人惊叹的事情。在这篇文章中,我结合了 Zig 的多项 comptime 特性来解决一个真实世界的问题:在 comptime 生成类型、在 comptime 调用对 comptime 无感知的普通函数、inline else@unionInit 等等。

本文展示的是 Zig comptime 一个相当进阶的用法。对于刚接触 Zig 的新人来说,这里用到的技巧和语言特性可能会让人不知所措。我自己也花了一段时间才逐渐适应这些特性,但一旦掌握,就会发现它们无比强大。

我是一个务实的类型论者。我相信借助类型系统来保证正确性是一个强大的工具。同时,我也认为有些语言在这方面走得太远,让类型系统反而成了负担。Zig 在这之间取得了很好的平衡,我认为本文正好体现了这一点。

脚注

  1. https://ziglang.org/documentation/0.13.0/#Tagged-union

  2. https://knowyourmeme.com/memes/how-to-draw-an-owl

  3. @ 开头的函数是 Zig 编译器内置的函数。

本文章由 muse-spark-1.2-contributor 进行翻译

评论