iflatten

先日何故いきなりCPS変換（で言葉はあっているのだろうか）をやり始めたかというと，
http://aspn.activestate.com/ASPN/Mail/Message/python-tutor/2302231
をみてコールスタックを消費しないflatten関数の動きを理解したかったから．flatten関数はネストしたリスト（ツリー構造）を単一のツリーにする関数．Schemeなんかでは再帰で書けるけど，末尾再帰最適化の無い言語でやるとちょっと大きなツリーでもスタックオーバーフローを起こしてしまう．それをトランポリンスタイルにして避けようというのが上記URLの話．

CPS変換はSchemeのコンパイラ書く時なんかに使えるようなので，成り行きとはいえ勉強して得した気分．

ところで普通のリストに変換するのもいいけど，Pythonならやっぱりジェネレータにして1パスでツリーを走査してしまいたいと思いますよね．

ということで書いてみた．はじめトランポリンで書こうと思ったが途中でツリーの走査がアセンブラ（というかマシン語？）の実行とほとんど同じであることに気づき，リストのイテレータを命令ポインタとして考えて，コールスタックに相当するものを用意すればループで普通に書けるのではないかと思い至った．

つまり，

リスト == ルーチン
リストの中のリスト == サブルーチン

と考えると以下のような手順で処理できる．
1. リスト走査中にリストに出会ったら，今のイテレータをスタックにプッシュして新しいリストを走査する．
2. リスト走査中にリスト以外に出会ったらyieldする．
3. リストの走査が終わったらスタックから元のイテレータをポップして残りを走査する．
4. あるリストの走査終了時にスタックの中にリストが無いなら終わり．

以下が結果のiflatten関数．この前に2バージョン程あったが現段階ではこれが一番綺麗に思う．

def iflatten(tree):
    it = iter(tree)
    stack = []
    while True:
        try:
            e = it.next()
            if isinstance(e, (list, tuple)):
                stack.append(it)
                it = iter(e)
            else:
                yield e
        except StopIteration:
            if stack:
                it = stack.pop()
            else:
                break

上記のURLから頂いてきたテストコードもちゃんと通る．

if __name__ == '__main__':
    import sys
    def make_evil_list(n):
        result = []
        for i in xrange(n):
            result = [[result], i]
        return result

    evil_nested_list = make_evil_list(sys.getrecursionlimit())
    assert range(sys.getrecursionlimit()) == [i for i in
                                              iflatten(evil_nested_list)]

昔から思っているのだが，StopIteration例外で反復を止めるのって微妙な仕様じゃない？
ジェネレータの実装上の都合のような気がしているが該当部分を詳しく調べていないのでよく分からない．
ループが止まるのが例外的な現象だと思えってことなのか．