Rust の再帰イテレータでサブディレクトリのファイルもすべて取得する (Recursive Directory Iterator in Rust)

再帰関数

サブディレクトリも含めて、あるディレクトリ内のすべてのファイルへのパスを取得したい場合、下記のように再帰関数を用いるのが最初に思い付く実装ですね。

main.rs

use std::fs;
use std::io;
use std::path::{Path, PathBuf};

fn visit_dir<P: AsRef<Path>>(path: P, paths: &mut Vec<PathBuf>) -> io::Result<()> {
    for entry in fs::read_dir(path)? {
        let entry = entry?;
        if entry.file_type()?.is_dir() {
            visit_dir(entry.path(), paths)?;
        }
        paths.push(entry.path());
    }
    Ok(())
}

fn main() -> io::Result<()> {
    let mut paths = vec![];
    visit_dir(".", paths.as_mut())?;

    println!("{:?}", paths);
    Ok(())
}

動機

上のコードはあまりイケてないですよね。イテレータを使って、↓こんな風に書けたらいいのになぁと思っていました。

main.rs

struct VisitDir {
    // TODO
}

impl Iterator for VisitDir {
    type Item = io::Result<DirEntry>;
    fn next(&mut self) -> Option<Self::Item> {
        // TODO
    }
}

fn main() -> io::Result<()> {
    let paths = VisitDir::new(".")?
        .filter_map(|e| Some(e.ok()?.path()))
        .collect::<Vec<_>>();

    println!("{:?}", paths);
    Ok(())
}

しかし、どんな風に書いたらいいのか分からず…。

そしたら、This Week in Rust でこの記事を見つけて、「これだ！」ということで、次のように書いてみました。

再帰イテレータ

main.rs

use std::fs::{self, DirEntry};
use std::io;
use std::path::Path;

struct VisitDir {
    root: Box<dyn Iterator<Item = io::Result<DirEntry>>>,
    children: Box<dyn Iterator<Item = VisitDir>>,
}

impl VisitDir {
    fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let root = Box::new(fs::read_dir(&path)?);
        let children = Box::new(
            fs::read_dir(&path)?
                .filter_map(|e| {
                    let e = e.ok()?;
                    if e.file_type().ok()?.is_dir() {
                        return Some(VisitDir::new(e.path()).ok()?);
                    }
                    None
                })
        );
        Ok(VisitDir { root, children })
    }

    fn entries(self) -> Box<dyn Iterator<Item = io::Result<DirEntry>>> {
        Box::new(
            self.root
                .chain(self.children.map(|s| s.entries()).flatten()),
        )
    }
}

impl Iterator for VisitDir {
    type Item = io::Result<DirEntry>;
    fn next(&mut self) -> Option<Self::Item> {
        if let Some(item) = self.root.next() {
            return Some(item);
        }
        if let Some(child) = self.children.next() {
            self.root = child.entries();
            return self.next();
        }
        None
    }
}

fn main() -> io::Result<()> {
    let paths = VisitDir::new(".")?
        .filter_map(|e| Some(e.ok()?.path()))
        .collect::<Vec<_>>();

    println!("{:?}", paths);
    Ok(())
}

やったぜ。