Move all mutable state from FileTreeWalk to its iterator implementation, allowing this sequence to be iterated several times.

Introduce SingleFileState for single file walks.
Do not check isDirectory in DirectoryState constructor unless assertions are enabled.

libraries/stdlib/src/kotlin/io/files/FileTreeWalk.kt

@@ -5,8 +5,8 @@ package kotlin.io
 
 import java.io.File
 import java.io.IOException
-import java.util.NoSuchElementException
 import java.util.Stack
+import kotlin.support.AbstractIterator
 
 /**
  * An enumeration to describe possible walk directions.
@@ -45,17 +45,84 @@ public class FileTreeWalk(private val start: File,
                           private val maxDepth: Int = Int.MAX_VALUE
 ) : Sequence<File> {
 
-    /** Abstract class that encapsulates file visiting in some order, beginning from a given [rootDir] */
+    /** Returns an iterator walking through files. */
-    private abstract class DirectoryState(public val rootDir: File) {
+    override public fun iterator(): Iterator<File> = FileTreeWalkIterator()
-        init {
-            if (!rootDir.isDirectory())
-                throw IllegalArgumentException("Directory is needed")
-        }
 
+    /** Abstract class that encapsulates file visiting in some order, beginning from a given [rootDir] */
+    private abstract class WalkState(val rootDir: File) {
         /** Call of this function proceeds to a next file for visiting and returns it */
         abstract public fun step(): File?
     }
 
+    /** Abstract class that encapsulates directory visiting in some order, beginning from a given [rootDir] */
+    private abstract class DirectoryState(rootDir: File): WalkState(rootDir) {
+        init {
+            @Suppress("DEPRECATION")
+            if (ASSERTIONS_ENABLED)
+                assert(rootDir.isDirectory) { "rootDir must be verified to be directory beforehand." }
+        }
+    }
+
+    private inner class FileTreeWalkIterator : AbstractIterator<File>() {
+
+        // Stack of directory states, beginning from the start directory
+        private val state = Stack<WalkState>()
+
+        init {
+            if (start.isDirectory && filter(start)) {
+                state.push(directoryState(start))
+            } else if (start.isFile) {
+                state.push(SingleFileState(start))
+            } else {
+                done()
+            }
+
+        }
+
+        override fun computeNext() {
+            val nextFile = gotoNext()
+            if (nextFile != null)
+                setNext(nextFile)
+            else
+                done()
+        }
+
+
+        private fun directoryState(root: File): DirectoryState {
+            return when (direction) {
+                FileWalkDirection.TOP_DOWN -> TopDownDirectoryState(root)
+                FileWalkDirection.BOTTOM_UP -> BottomUpDirectoryState(root)
+            }
+        }
+
+        tailrec private fun gotoNext(): File? {
+
+            if (state.empty()) {
+                // There is nothing in the state
+                return null
+            }
+            // Take next file from the top of the stack
+            val topState = state.peek()!!
+            val file = topState.step()
+            if (file == null) {
+                // There is nothing more on the top of the stack, go back
+                state.pop()
+                return gotoNext()
+            } else {
+                // Check that file/directory matches the filter
+                if (!filter(file))
+                    return gotoNext()
+                if (file == topState.rootDir || !file.isDirectory || state.size >= maxDepth) {
+                    // Proceed to a root directory or a simple file
+                    return file
+                } else {
+                    // Proceed to a sub-directory
+                    state.push(directoryState(file))
+                    return gotoNext()
+                }
+            }
+        }
+
         /** Visiting in bottom-up order */
         private inner class BottomUpDirectoryState(rootDir: File) : DirectoryState(rootDir) {
 
@@ -130,60 +197,23 @@ public class FileTreeWalk(private val start: File,
             }
         }
 
-    // Stack of directory states, beginning from the start directory
+        private inner class SingleFileState(root: File) : WalkState(root) {
-    private val state = Stack<DirectoryState>()
+            private var visited: Boolean = false
-
-    // We are already at the end or not?
-    private var end = false
-
-    // A future result of next() call
-    private var nextFile: File? = null
 
             init {
-        if (!start.exists()) {
+                @Suppress("DEPRECATION")
-            end = true
+                if (ASSERTIONS_ENABLED)
-        } else if (start.isDirectory() && filter(start)) {
+                    assert(root.isFile) { "root must be verified to be file beforehand." }
-            pushState(start)
+            }
+
+            override fun step(): File? {
+                if (visited) return null
+                visited = true
+                if (!filter(rootDir)) return null
+                return rootDir
             }
         }
 
-    private fun pushState(root: File) {
-        state.push(when (direction) {
-            FileWalkDirection.TOP_DOWN -> TopDownDirectoryState(root)
-            FileWalkDirection.BOTTOM_UP -> BottomUpDirectoryState(root)
-        })
-    }
-
-    tailrec private fun gotoNext(): File? {
-        if (end) {
-            // We are already at the end
-            return null
-        } else if (state.empty()) {
-            // There is nothing in the state
-            // We must visit "start" if it's a file and matches the filter
-            end = true
-            return if (start.exists() && !start.isDirectory() && filter(start)) start else null
-        }
-        // Take next file from the top of the stack
-        val topState = state.peek()
-        val file = topState.step()
-        if (file == null) {
-            // There is nothing more on the top of the stack, go back
-            state.pop()
-            return gotoNext()
-        } else {
-            // Check that file/directory matches the filter
-            if (!filter(file))
-                return gotoNext()
-            if (file == topState.rootDir || !file.isDirectory() || state.size >= maxDepth) {
-                // Proceed to a root directory or a simple file
-                return file
-            } else {
-                // Proceed to a sub-directory
-                pushState(file)
-                return gotoNext()
-            }
-        }
     }
 
     /**
@@ -231,32 +261,6 @@ public class FileTreeWalk(private val start: File,
             throw IllegalArgumentException("Use positive depth value")
         return FileTreeWalk(start, direction, enter, leave, fail, filter, depth)
     }
-
-    /** An iterator associated with this walker */
-    private val it = object : Iterator<File> {
-        override public fun hasNext(): Boolean {
-            if (nextFile == null)
-                nextFile = gotoNext()
-            return (nextFile != null)
-        }
-
-        override public fun next(): File {
-            if (nextFile == null)
-                nextFile = gotoNext()
-            val res = nextFile
-            if (res == null)
-                throw NoSuchElementException()
-            // With nextFile = gotoNext() here some enter() / leave() can be called BEFORE visiting res
-            nextFile = null
-            // Visit this directory or file
-            return res
-        }
-    }
-
-    /** Returns an associated file iterator. */
-    override public fun iterator(): Iterator<File> {
-        return it
-    }
 }
 
 /**

libraries/stdlib/test/io/FileTreeWalks.kt

@@ -49,6 +49,23 @@ class FileTreeWalkTest {
         }
     }
 
+    @Test fun singleFile() {
+        val testFile = createTempFile()
+        val nonExistantFile = testFile.resolve("foo")
+        try {
+            for (walk in listOf(File::walkTopDown, File::walkBottomUp)) {
+                assertEquals(testFile, walk(testFile).single(), "${walk.name}")
+                assertTrue(walk(testFile).treeFilter { false }.none(), "${walk.name}")
+                assertEquals(testFile, testFile.walk().onEnter { false }.single(), "${walk.name} - enter should not be called for single file")
+
+                assertTrue(walk(nonExistantFile).none(), "${walk.name} - enter should not be called for single file")
+            }
+        }
+        finally {
+            testFile.delete()
+        }
+    }
+
     @Test fun withEnterLeave() {
         val basedir = createTestFiles()
         try {