LuaでPythonのrange関数のようなものを作る

Pythonのrange関数に似た関数をLuaで作成してみます。

function range(n)
    local i = 0
    return function()
        if i >= n then return nil end

        local ret = i
        i = i + 1
        return ret
    end
end

function main()
    for i in range(5) do
        print(i)
    end
end

main()

実行結果です。

0
1
2
3
4

上のrange関数は呼び出すたびにクロージャを生成しています。ここではクロージャを生成しない方法を見ていきます。

『Programming in Lua』p.82には、for文のセマンティクスに関する記述があります。

for var_1, ..., var_n in <explist> do <block> end

は以下のコードと等価と書かれてあります。

do
  local _f, _s, _var = <explist>
  while true do
    local var_1, ..., var_n = _f(_s, _var)
    _var = var_1
    if _var == nil then break end
    <block>
  end
end

要は、イテレータ関数、不変状態、制御変数の3つをexplistが返却して、それから制御変数を更新するようにすればいいですね。

function range2(n) 
    return iter, n, -1
end

function iter(n, i)
    if i+1 >= n then return nil end
    return i+1
end

function main()
    for i in range2(3) do
        print(i)
    end
end

main()

実行結果です。

0
1
2

iter関数を以下のようにrange2()の関数内関数として定義しても、iter関数の生成コストはゼロと考えても問題ないだろうか?(range2()呼び出し時のiter関数生成コストはゼロ?)。ソースを読んで確かめないと。

function range2(n)
    local function iter(n, i)
        if i+1 >= n then return nil end
        return i+1
    end

    return iter, n , -1
end