takeaway 0.1.4

//! A recursive line-counter program.
//!
//! `rlc .` is equivalent to `find . -print0 | wc -l --files0-from=-`.  It's a
//! great demonstration of [`takeaway`], performing complex work and making use
//! of task prioritization.

use std::{
    fs::File,
    io::{BufRead, BufReader},
    num::NonZeroU32,
    path::{Path, PathBuf},
};

use takeaway::{Queue, Worker, util::block_on};

fn main() {
    // TODO: Use Clap here without infecting it as a public dev-dependency of
    // the whole crate?
    let path = std::env::args().nth(1).unwrap_or(".".into());

    // Set up the global queue.
    let queue = takeaway::Config::default().with_oneshot(true).build();

    // Spin up the worker threads.
    std::thread::scope(|s| {
        let queue = &queue;

        // Spin up the auxiliary threads.
        let handles = (1..queue.config().num_workers().get())
            .map(|id| s.spawn(move || block_on(worker(queue, None, id))))
            .collect::<Box<[_]>>();

        // Run the main thread worker.
        let count = block_on(worker(queue, Some(path), 0));

        // Wait until all the threads are done.
        let total = [count]
            .into_iter()
            .chain(handles.into_iter().map(|h| h.join().unwrap()))
            .sum::<usize>();

        println!("... total: {total} lines");
    });
}

/// Drive a worker thread.
async fn worker(
    queue: &Queue<Task>,
    initial: Option<String>,
    id: usize,
) -> usize {
    // Set up the local queue.
    let mut worker = Worker::new(queue, id);

    // Set up a local line count.
    let mut count = 0usize;

    // Enqueue the initial task, if any.
    if let Some(initial) = initial {
        let path = PathBuf::from(initial).into_boxed_path();
        worker.enqueue_one(Task {
            path,
            depth: 0,
            directory: true,
        });
    }

    while let Some(task) = worker.next().await {
        // Execute the task.
        'execute: {
            if task.directory {
                let dir = match std::fs::read_dir(&*task.path) {
                    Ok(dir) => dir,
                    Err(error) => {
                        eprintln!(
                            "Could not open dir '{}': {error}",
                            task.path.display()
                        );
                        break 'execute;
                    }
                };

                for entry in dir {
                    let entry = match entry {
                        Ok(entry) => entry,
                        Err(error) => {
                            eprintln!(
                                "Could not read from dir '{}': {error}",
                                task.path.display()
                            );
                            break;
                        }
                    };

                    let file_type = match entry.file_type() {
                        Ok(file_type) => file_type,
                        Err(error) => {
                            eprintln!(
                                "Could not determine file type of '{}': {error}",
                                entry.path().display()
                            );
                            continue;
                        }
                    };

                    worker.enqueue_one(Task {
                        path: entry.path().into_boxed_path(),
                        depth: task.depth + 1,
                        directory: file_type.is_dir(),
                    });
                }
            } else {
                let file = match File::open(&*task.path) {
                    Ok(file) => BufReader::new(file),
                    Err(error) => {
                        eprintln!(
                            "Could not read file '{}': {error}",
                            task.path.display()
                        );
                        break 'execute;
                    }
                };

                let file_count = file.lines().count();
                println!("{}: {file_count} lines", task.path.display());

                count += file_count;
            }
        }
    }

    count
}

/// A line-counting task.
struct Task {
    /// The path to the object to visit.
    path: Box<Path>,

    /// The depth of this path, relative to the start.
    depth: u32,

    /// Whether this is a directory or not.
    directory: bool,
}

impl takeaway::Task for Task {
    type Priority = NonZeroU32;

    fn priority(&self) -> Self::Priority {
        // We use the inverse of the depth so that shallower directory paths
        // have a greater priority.  The execution will roughly follow a BFS
        // traversal.

        if self.directory {
            let priority = u32::MAX - self.depth;
            priority.try_into().unwrap_or(NonZeroU32::MIN)
        } else {
            NonZeroU32::MIN
        }
    }
}