Lua next 函数的一个有趣问题

 

熟悉 Lua 的同学都知道, Lua 是允许在 for ... pairs 循环中修改和删除表中元素的. 下面这样的代码是没有任何问题的:

local t = {a = 1, b = 2, c = 3}
for k, v in pairs(t) do
    if v == 1 then
        t[k] = nil
    end
end

但是, 如果我们在遍历时既删除元素又新增元素, 就会有问题了. 比如说:

local t = {a = 1, b = 2, c = 3}
for k, v in pairs(t) do
    if v == 1 then
        t[k] = nil
        t[k .. 1] = v + 1
    end
end

运行以上代码会得到这样的报错:

invalid key to 'next'
stack traceback:
        [C]: in function 'next'
        stdin:1: in main chunk
        [C]: in ?

熟悉 Lua 的同学对这个报错应该不会陌生. 解决方案就是避免在遍历时给 table 新增元素. Lua 官方文档也说的很明白: 如果在遍历期间将任何值赋给表中不存在的字段,则 next 的行为是未定义的. 不过这个报错本身很有意思, 背后涉及 Lua 多方面的机制, 我认为这是个了解 Lua 内部实现的好契机.

令人费解的行为

为什么说它很有意思呢? 首先这个报错还不一定会出现 (文档也说了 "行为未定义"), 比如下面的代码就没这个问题:

local t = {a = 1, b = 2, c = 3, d = 4, e = 5, f = 6}
for k, v in pairs(t) do
    if v == 1 then
        t[k] = nil
        t[k .. 1] = v + 1
    end
end

另一方面, 我们知道 for ... pairs 循环本质是每次迭代调用 next, 传入 table 和上一次迭代的 key, 获取本次循环的键值对. 而 next(t, k) 函数的含义就是, 对于给定 table t, 返回与给定 key k 相邻的下一个键值对. 如果 kt 的最后一个元素, 则返回 nil; 如果 knil, 则返回 t 的第一个元素. 那如果 k 不在 t 中呢? 自然该报错了:

> next({a = 1, b = 2}, 'c')
invalid key to 'next'
stack traceback:
        [C]: in function 'next'
        stdin:1: in main chunk
        [C]: in ?

熟悉的报错. 既然这样, 本文开头的代码也应该报错才对, 因为当 v 等于 1 时, 相应的 key 就会被删掉, 导致下次迭代调用 next 时尝试为一个不存在 table 中的 key 求下一个元素. 但不仅它能正常运行, 下面这段代码也能正常运行:

local t = {a = 1, b = 2}
print(next(t, 'a')) -- b    2
t.a = nil
print(next(t, 'a')) -- b    2

第二个 next 调用正常返回, 就好像 a 还在 t 中一样. 而当我们删除后再插入一个新元素时, 报错就出现了:

local t = {a = 1, b = 2}
print(next(t, 'a')) -- b    2
t.a = nil
t.a1 = 2
print(next(t, 'a')) -- error: invalid key to 'next'

当然还有更好玩的, 要得到这个报错, 不一定要插入新的元素, 有时 GC 下就出来了:

local k = 'a'..'b'
local t = {
    a = 1,
    [k] = 2,
}
t[k] = nil
k = nil
collectgarbage("collect")
next(t, 'a'..'b') -- error: invalid key to 'next'

写法比较奇怪, 这都是为了迎合 Lua 某些机制. 上面的代码稍微改改, 比如说把最后的 next(t, 'a'..'b') 改成 next(t, 'ab'), 报错就不会出现了.

next 函数的实现

调用 next(t, k) 时, k 应该要存在于 t 中; 但若 k 是刚刚被删除的似乎也可以. 这是怎么做到的呢? 我们来看看 next 函数的实现:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
/* ltable.c */

static unsigned int findindex (lua_State *L, Table *t, StkId key) {
  unsigned int i;
  if (ttisnil(key)) return 0;  /* first iteration */
  i = arrayindex(key);
  if (i != 0 && i <= t->sizearray)  /* is 'key' inside array part? */
    return i;  /* yes; that's the index */
  else {
    int nx;
    Node *n = mainposition(t, key);
    for (;;) {  /* check whether 'key' is somewhere in the chain */
      /* key may be dead already, but it is ok to use it in 'next' */
      if (luaV_rawequalobj(gkey(n), key) ||
            (ttisdeadkey(gkey(n)) && iscollectable(key) &&
             deadvalue(gkey(n)) == gcvalue(key))) {
        i = cast_int(n - gnode(t, 0));  /* key index in hash table */
        /* hash elements are numbered after array ones */
        return (i + 1) + t->sizearray;
      }
      nx = gnext(n);
      if (nx == 0)
        luaG_runerror(L, "invalid key to 'next'");  /* key not found */
      else n += nx;
    }
  }
}


int luaH_next (lua_State *L, Table *t, StkId key) {
  unsigned int i = findindex(L, t, key);  /* find original element */
  for (; i < t->sizearray; i++) {  /* try first array part */
    if (!ttisnil(&t->array[i])) {  /* a non-nil value? */
      setivalue(key, i + 1);
      setobj2s(L, key+1, &t->array[i]);
      return 1;
    }
  }
  for (i -= t->sizearray; cast_int(i) < sizenode(t); i++) {  /* hash part */
    if (!ttisnil(gval(gnode(t, i)))) {  /* a non-nil value? */
      setobj2s(L, key, gkey(gnode(t, i)));
      setobj2s(L, key+1, gval(gnode(t, i)));
      return 1;
    }
  }
  return 0;  /* no more elements */
}

可以看到, luaH_next 首先调用 findindex 找出给定 key 在 table t 中的位置, 然后从该位置往后寻找下一个元素. 在 findindex 中, 会做这样几件事:

  • 第 5 行判断如果 keynil, 直接返回 0;
  • 6 至 8 行, 如果 key 为正整数且位于数组域, 直接返回数组索引;
  • 否则在哈希域查找. 首先 11 行得到 key 的主位置;
  • 12 至 25 行的循环中, 14 至 16 行检查当前位置的值是否与 key 相等, 如果相等则查找成功, 19 行直接返回; 否则 21 行检查下一个位置;
  • 如果直到链表结尾仍然找不到与 key 相等的节点, 查找失败, 23 行抛出错误 invalid key to 'next'.

前面我们看到, 当一个元素置为 nil 后立刻调用 next, 是能够正常返回的. 这说明 findindex 成功找到了这个元素的位置. 注意到 13 行的注释: "key may be dead already, but it is ok to use it in 'next'", 也就是说一个被删除的 key 对于 next 来说是没问题的? 带着这个疑问, 我们首先来看看当我们执行 t[k] = nil 时发生了什么.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
/* lvm.h */

#define luaV_fastset(L,t,k,slot,f,v) \
  (!ttistable(t) \
   ? (slot = NULL, 0) \
   : (slot = f(hvalue(t), k), \
     ttisnil(slot) ? 0 \
     : (luaC_barrierback(L, hvalue(t), v), \
        setobj2t(L, cast(TValue *,slot), v), \
        1)))

/* lvm.c */

#define settableProtected(L,t,k,v) { const TValue *slot; \
  if (!luaV_fastset(L,t,k,slot,luaH_get,v)) \
    Protect(luaV_finishset(L,t,k,v,slot)); }

void luaV_execute (lua_State *L) {
  ...
  /* main loop of interpreter */
  for (;;) {
    Instruction i;
    StkId ra;
    vmfetch();
    vmdispatch (GET_OPCODE(i)) {
      ...
      vmcase(OP_SETTABLE) {
        TValue *rb = RKB(i);
        TValue *rc = RKC(i);
        settableProtected(L, ra, rb, rc);
        vmbreak;
      }
      ...
    }
  }
}

我们知道 t[k] = nil 执行的是 OP_SETTABLE 指令, 于是找到 lvm.cluaV_execute 函数执行 OP_SETTABLE 的地方, 也就是上面的 27 至 32 行, 可以看到这里会做这样几件事:

  • 取得操作数, ra, rb, rc 分别为 table, key 和 value;
  • 30 行调用宏 settableProtected;
  • 15 行 settableProtected 先尝试调用 luaV_fastset;
  • luaV_fastset 中会判断, 当 t 不是 table (第 4 行), 或者 k 不在 t 中 (第 7 行) 时, luaV_fastset 失败, 返回 0. 否则在第 9 行直接赋值;
  • 如果 luaV_fastset 失败, 则在 16 行会转而调用 luaV_finishset.

当我们使用 t[k] = nil 删除一个元素时, k 显然是在 t 中的, 因此会执行到第 9 行, 直接给 slot 赋值. 注意到 slot 实际上是 luaH_get 的返回值, 它会等于 k 对应的 value. 因此执行 t[k] = nil 时, 只会将 k 对应的 value 赋值为 nil, 并不会改变 key. 这也给 next 函数提供了机会.

我们再回过来看 luaH_next. 它在 findindex 中会执行

if (luaV_rawequalobj(gkey(n), key) ||
      (ttisdeadkey(gkey(n)) && iscollectable(key) &&
       deadvalue(gkey(n)) == gcvalue(key))) {
    ...

以判断当前位置是否与 key 相等. 如果我们的 key 不是 GCObject, 如数字, 布尔值, 即使被删除, 也能正常访问, luaV_rawequalobj 总能给出正确的结果; 即使 key 是 GCObject, 只要它没被 GC 掉, 也是能正常比较的. 不过, 如果这个 key 的位置被之后插入的元素占用, 查找就会失败. 这也是为什么在遍历时既删除元素又新增元素有可能会得到 "invalid key to 'next'" 这样的报错.

在使用 for ... pairs 循环不用担心 key 会被 GC 掉, 因为迭代器总是会持有当前 key 的引用. 因此在 for ... pairs 循环中可以放心地执行删除操作.

Lua 的 GC 机制

(ttisdeadkey(gkey(n)) && iscollectable(key) && deadvalue(gkey(n)) == gcvalue(key)) 就是用于处理当 key 被 GC 掉的情况. 如果当前位置是一个 deadkey, 且 key 是 GCObject, 就比较 deadvalue(gkey(n)) == gcvalue(key), 也就是比较当前位置的 key (即 gkey(n)) 和传入的 key 的地址是否相等. 那么什么是 deadkey 呢? 这就涉及到 Lua 的 GC 机制了.

Lua GC 的原理就是, 将所有的 GCObject 串成一个链表; 然后定期扫描并标记虚拟机中的所有对象, 没有标记到的即是需要被回收的. 我们来看扫描 table 的函数:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
/* lgc.c */

static void traversestrongtable (global_State *g, Table *h) {
  Node *n, *limit = gnodelast(h);
  unsigned int i;
  for (i = 0; i < h->sizearray; i++)  /* traverse array part */
    markvalue(g, &h->array[i]);
  for (n = gnode(h, 0); n < limit; n++) {  /* traverse hash part */
    checkdeadkey(n);
    if (ttisnil(gval(n)))  /* entry is empty? */
      removeentry(n);  /* remove it */
    else {
      lua_assert(!ttisnil(gkey(n)));
      markvalue(g, gkey(n));  /* mark key */
      markvalue(g, gval(n));  /* mark value */
    }
  }
}

准确地说这是扫描非弱表的函数. 这个函数还是比较直白的, 它做了以下几件事:

  • 6 至 7 行扫描数组域, 全部标记;
  • 8 至 17 行扫描哈希域, 其中:
    • 10 至 11 行检查当前位置的 value 是否为 nil, 如果是, 则调用 removeentry;
    • 否则 13 至 15 行标记当前位置的 key 和 value.

没有被标记的对象都是要被回收的, 因此如果一个元素的 value 为 nil, 它就会在这次 GC 中被回收. 不过在回收之前, 还会调用 removeentry, 它做了什么呢?

/* lobject.h */

#define settt_(o,t)	((o)->tt_=(t))

#define setdeadvalue(obj)	settt_(obj, LUA_TDEADKEY)


/* lgc.c */

static void removeentry (Node *n) {
  lua_assert(ttisnil(gval(n)));
  if (valiswhite(gkey(n)))
    setdeadvalue(wgkey(n));  /* unused and unmarked key; remove it */
}

它会将该元素的 key 标记为 dead. 注意这个标记实际是将对象的类型 tt_ 字段置为 LUA_TDEADKEY. 这样当调用 luaV_rawequalobj 将一个 deadkey 与常规对象比较时总是会返回 false. 这样当执行到 deadvalue(gkey(n)) == gcvalue(key) 时, 元素 n 的 key 一定是被被 GC 掉了.

对于 table, function, thread, userdata 来说, 两个对象相等当且仅当它们的地址相等. 而当一个对象被 GC 掉后, 它就没有任何引用了, 因此也不可能将它传入 next 中.

但是 string 却不太一样, 两个字符串是否相等取决于字符串的内容而非其地址. 因此当一个字符串被 GC 后, 我们仍然可以构造出一个与之相等的字符串, 但是它们的地址却不相同. 将新构造的字符串传入 next, 就能导致报错. 在前面的例子中, 之所以要用 'a'..'b' 是因为要让他构造出不同的对象. 只要将其中任意一处的 'a'..'b' 改成 'ab', 那么字符串 ab 就会存在常量表中, 导致它不会被 GC 掉; 而 Lua 又有短字符串复用机制, 导致 'a'..'b''ab' 实际是同一个对象, 就不会出现这个报错了.

结论

笔者第一次发现 next 可以传入刚被删掉的 key 时还是感觉很惊讶, 因为一个删掉的 key 应该是类似于野指针一般的东西, 不应该再使用它. 但是仔细看下来, 一切都是没有问题的. 这让我又一次感受到了 Lua 设计的精妙. 这个小问题背后涉及到 Lua table 的实现, GC 的机制, 常量表机制以及短字符串复用等内容. 硬啃 Lua 源码比较枯燥, 我们可以从一些有趣的现象出发去探索 Lua 的实现.