1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

use crate::blob::{Blob, Output};
use crate::error::Error;
use crate::taskdef::{FromTaskDef, IntoTaskDef, TaskDef};
use std::path::{Path, PathBuf};

/// Provides execution context on split stage.
pub trait SplitContext {
    /// Command line arguments passed to runner.
    fn args(&self) -> &Vec<String>;

    /// Allocates new output file.
    fn new_blob(&mut self) -> Output;

    #[doc(hidden)]
    fn blob_from_file(&mut self, path: &Path) -> Result<Blob, Error> {
        self.new_blob().file(path)
    }

    #[doc(hidden)]
    fn blob_from_bytes(&mut self, bytes: &[u8]) -> Result<Blob, Error> {
        self.new_blob().bytes(bytes)
    }
}

pub trait Splitter {
    type WorkItem: IntoTaskDef + FromTaskDef;

    fn split(self, context: &mut dyn SplitContext) -> Vec<Self::WorkItem>;
}

impl<Out, F: (FnOnce(&mut dyn SplitContext) -> Out)> Splitter for F
where
    Out: IntoIterator,
    Out::Item: IntoTaskDef + FromTaskDef,
{
    type WorkItem = Out::Item;

    fn split(self, context: &mut dyn SplitContext) -> Vec<Self::WorkItem> {
        self(context).into_iter().collect()
    }
}

struct WorkDirCtx {
    id: u64,
    work_dir: PathBuf,
    args: Vec<String>,
}

impl SplitContext for WorkDirCtx {
    fn args(&self) -> &Vec<String> {
        &self.args
    }

    fn new_blob(&mut self) -> Output {
        loop {
            let id = self.id;
            self.id += 1000;
            let name = format!("{:06x}.bin", id);
            let output_path = self.work_dir.join(name);
            if !output_path.exists() {
                return Output(output_path);
            }
        }
    }
}

pub(crate) fn split_into<S: Splitter>(
    splitter: S,
    base_path: &Path,
    args: &[String],
) -> Result<Vec<TaskDef>, Error> {
    let mut ctx = WorkDirCtx {
        id: 1000,
        work_dir: base_path.into(),
        args: args.into(),
    };
    splitter
        .split(&mut ctx)
        .into_iter()
        .map(|item| IntoTaskDef::into_task_def(item, base_path))
        .collect()
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::blob::Blob;
    use std::io::Write;

    fn my_spliter(ctx: &mut dyn SplitContext) -> Vec<(Blob, u32)> {
        let mut out = Vec::new();
        for i in 1..10 {
            let output = ctx.new_blob();
            {
                let mut w = output.open().unwrap();
                let _ = w.write("smok smok".as_ref()).unwrap();
            }
            out.push((Blob::from_output(output), i))
        }
        out
    }

    #[test]
    fn test_split() {
        let tasks = split_into(my_spliter, &PathBuf::from("/tmp"), &vec![]).unwrap();

        eprintln!("{}", serde_json::to_string(&tasks).unwrap());
    }
}