.NET Framework < 4.0でEnumerateFiles/EnumerateDirectories

最近になって、プログラミングの世界に舞い戻ってきた。別に必要でもなんでも無いから遠ざかっていて、今も似たような感じだけど、なぜか唐突に始めた感じだ。
さて、昨日今日でやっつけたC#プログラムを作る際、どうしても数万ファイルを探索することになる処理が出てきた。
そういった処理は、例によってDirectory.GetFilesをAllDirectoriesオプションを付けて呼び出し、繰り返しで検索するのがテンプレだ。しかし、GetFilesは配列を返すメソッドなので、ファイル数が多くなってくると、やたらと時間を食うしリソースも食う。
これのイテレータを返すバージョンは無いの？と疑問に思って探してみると、.NET Framework 4.0からの実装ということが分かった。
4.0から？2.0からの間違いじゃないの？どうしてこんな機能を実装してなかったんだ……
嘆いていたって仕方ないので自力救済を図るしかない。従来のWin32アプリケーションは普通、ファイルの列挙にFind***系APIを用いる。Directory.GetFilesなどはこれのフロントエンドだろう。それに倣って、P/Invokeでそれらを呼び出し、yieldをぶん回す関数を考え付いた。
どこかで誰か思いついてそうなネタだろうけどね。

using System;
using System.IO;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Text.RegularExpressions;

namespace ExtIO
{
    public static class DirectoryEx
    {
        internal static readonly IntPtr INVALID_HANDLE_VALUE = (IntPtr)(-1);

        internal class SafeFindHandle : SafeHandle
        {
            public SafeFindHandle()
                : base(INVALID_HANDLE_VALUE, true)
            {
            }

            internal SafeFindHandle(IntPtr existingHandle, bool ownsHandle)
                : base(existingHandle, ownsHandle)
            {
            }

            public override bool IsInvalid
            {
                get { return IsClosed || handle == INVALID_HANDLE_VALUE; }
            }

            protected override bool ReleaseHandle()
            {
                if (!IsInvalid)
                {
                    FindClose(handle);
                    return true;
                }
                return false;
            }
        }

        // From P/Invoke.net
        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        internal struct WIN32_FIND_DATA
        {
            public FileAttributes dwFileAttributes;
            public System.Runtime.InteropServices.ComTypes.FILETIME ftCreationTime;
            public System.Runtime.InteropServices.ComTypes.FILETIME ftLastAccessTime;
            public System.Runtime.InteropServices.ComTypes.FILETIME ftLastWriteTime;
            public int nFileSizeHigh;
            public int nFileSizeLow;
            public uint dwReserved0;
            public uint dwReserved1;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
            public string cFileName;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 14)]
            public string cAlternateFileName;
        }

        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        internal static extern SafeFindHandle FindFirstFileW(string lpFileName, out WIN32_FIND_DATA lpFindFileData);


        [DllImport("kernel32.dll", CharSet = CharSet.Auto)]
        internal static extern bool FindNextFile(SafeFindHandle hFindFile, out WIN32_FIND_DATA
           lpFindFileData);

        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool FindClose(IntPtr hFindFile);

        public static IEnumerable<DirectoryInfo> EnumerateDirectories(string searchPath, string searchPattern, SearchOption searchOption)
        {
            WIN32_FIND_DATA findData;
            Regex compiledPattern = new Regex(searchPattern, RegexOptions.Compiled | RegexOptions.IgnoreCase);

            using (SafeFindHandle findHandle = FindFirstFileW(Path.Combine(searchPath, "*"), out findData))
            {
                if (!findHandle.IsInvalid) {
                    do
                    {                       
                        if ((findData.dwFileAttributes & FileAttributes.Directory) != 0)
                        {
                            if (findData.cFileName != "." && findData.cFileName != ".." && compiledPattern.IsMatch(findData.cFileName))
                            {
                                string foundPath = Path.Combine(searchPath, findData.cFileName);
                                yield return new DirectoryInfo(foundPath);
                                if (searchOption == SearchOption.AllDirectories)
                                {
                                    foreach (DirectoryInfo subItem in EnumerateDirectories(foundPath, searchPattern, searchOption))
                                    {
                                        yield return subItem;
                                    }
                                }
                            }
                        }
                    } while (FindNextFile(findHandle, out findData));
                }
            }

            yield break;
        }

        public static IEnumerable<FileInfo> EnumerateFiles(string searchPath, string searchPattern, SearchOption searchOption)
        {
            WIN32_FIND_DATA findData;
            Regex compiledPattern = new Regex(searchPattern, RegexOptions.Compiled | RegexOptions.IgnoreCase);

            foreach (DirectoryInfo dirInfo in EnumerateDirectories(searchPath, ".*", searchOption))
            {
                using (SafeFindHandle findHandle = FindFirstFileW(Path.Combine(dirInfo.FullName, "*"), out findData))
                {
                    if (!findHandle.IsInvalid)
                    {
                        do
                        {  
                            if ((findData.dwFileAttributes & FileAttributes.Directory) == 0)
                            {
                                if (compiledPattern.IsMatch(findData.cFileName))
                                {
                                    string foundPath = Path.Combine(dirInfo.FullName, findData.cFileName);
                                    yield return new FileInfo(foundPath);
                                }
                            }
                        } while (FindNextFile(findHandle, out findData));
                    }
                }
            }
            yield break;
        }
    }
}