smoker fixes

Author: icewind1991

Diff for: smoker/src/main.rs

@@ -3,14 +3,21 @@ use rayon::prelude::*;
 use std::env;
 use std::ffi::OsStr;
 use std::fs;
+use std::io::{BufRead, BufReader};
 use std::path::{Path, PathBuf};
-pub use tf_demo_parser::{Demo, DemoParser, Parse, ParseError, ParserState, Stream};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use tf_demo_parser::{Demo, DemoParser, Parse, ParseError, ParserState, Stream};
 
 #[global_allocator]
 static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
 
 fn main() -> Result<(), MainError> {
     better_panic::install();
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(40)
+        .build_global()
+        .unwrap();
 
     let args: Vec<_> = env::args().collect();
     if args.len() < 2 {
@@ -24,39 +31,57 @@ fn main() -> Result<(), MainError> {
         .unwrap_or_default();
 
     let files = gather_dir(path)?;
+    let total = files.len();
     println!("found {} demo files", files.len());
 
-    let failures = files.par_iter().filter_map(|entry| {
-        let file = fs::read(&entry).unwrap();
-        let demo = Demo::new(file);
-        let parser = if all {
-            DemoParser::new_all(demo.get_stream())
-        } else {
-            DemoParser::new(demo.get_stream())
-        };
-        if let Err(e) = parser.parse() {
-            eprintln!("{}: {}", entry.to_str().unwrap(), e);
-            Some(entry)
-        } else {
-            None
-        }
-    });
-    println!("Found {} failures", failures.count());
+    let count = Arc::new(AtomicUsize::new(0));
+    let failures: Vec<_> = files
+        .par_iter()
+        .filter_map(|entry| {
+            let file = fs::read(&entry).unwrap();
+            let demo = Demo::new(file);
+            let parser = if all {
+                DemoParser::new_all(demo.get_stream())
+            } else {
+                DemoParser::new(demo.get_stream())
+            };
+            let done = count.fetch_add(1, Ordering::AcqRel);
+            println!("{}/{}", done + 1, total);
+            if let Err(e) = parser.parse() {
+                eprintln!("{}: {}", entry.to_str().unwrap(), e);
+                Some(entry)
+            } else {
+                None
+            }
+        })
+        .collect();
+    println!("Found {} failures", failures.len());
+    for failed in failures {
+        println!("{}", failed.to_str().unwrap());
+    }
     Ok(())
 }
 
 fn gather_dir(path: impl AsRef<Path>) -> Result<Vec<PathBuf>, MainError> {
     let mut files = Vec::with_capacity(512);
 
-    for res in fs::read_dir(path)? {
-        let entry = res?;
+    let meta = fs::metadata(path.as_ref())?;
+    if meta.is_file() {
+        let file = fs::File::open(path)?;
+        for line in BufReader::new(file).lines() {
+            files.push(line?.into())
+        }
+    } else {
+        for res in fs::read_dir(path)? {
+            let entry = res?;
 
-        if entry.file_type()?.is_dir() {
-            files.extend_from_slice(&gather_dir(entry.path())?);
-        } else {
-            let entry_path = entry.path();
-            if entry_path.extension() == Some(OsStr::new("dem")) {
-                files.push(entry.path());
+            if entry.file_type()?.is_dir() {
+                files.extend_from_slice(&gather_dir(entry.path())?);
+            } else {
+                let entry_path = entry.path();
+                if entry_path.extension() == Some(OsStr::new("dem")) {
+                    files.push(entry.path());
+                }
             }
         }
     }