We read:

public static void betterWay() {
   List<File> files =
      Stream.of(new File(".").listFiles())
         .flatMap(file -> file.listFiles() == null ?
            Stream.of(file) : Stream.of(file.listFiles()))
         .collect(toList());
      System.out.println("Count: " + files.size());
}

One may note that there are two calls to file.listFiles(), which may or may not be a problem - the exercise to ameliorate that can be left to the reader.

More interestingly, I tried to rejiggle the if-the-else expression into

Function<File,Stream<File>> func1 =
    ((File file) -> 
      Stream.of((file.listFiles() == null) ? file : file.listFiles()));

rather than

Function<File,Stream<File>> func2 =
    ((File file) -> 
      file.listFiles() == null ? Stream.of(file) : Stream.of(file.listFiles()));

while func2 is what is used in the book, func1 is rejected by the Java 18 compiler (although in principle equivalent) becaus the type hole cannot be resolved. Looks like the ?: expressions presents some kind of problem.

Note that File[] array can be null in case the file does not exist or the file is not a directory.

Thus, properly:

    public void listSubDirsBetterWay() {
        File[] files = new File(theDir).listFiles();
        String res;
        if (files == null) {
            res = "Looks like '" + theDir + "' is not a directory";
        }
        else {
            Function<File,Stream<File>> func = ((File file) -> file.listFiles() == null ? Stream.of(file) : Stream.of(file.listFiles()));
            List<File> filesAndSubfiles = Stream.of(files).flatMap(func).collect(toList());
            res = "Count: " + filesAndSubfiles.size();
        }
        System.out.println(res);
    }