<commit_before>extern crate serde_json;

use sha2::{Digest, Sha256};
use std::cell::RefCell;
use std::collections::HashMap;
use std::error::Error;
use std::fs::{self, DirEntry, File};
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

#[derive(Serialize, Deserialize, PartialEq)]
pub enum JobStatus {
    Pending,
    Done,
    Error,
}

#[derive(Serialize, Deserialize)]
pub struct Job {
    source_path: String,
    source_sha256sum: String,
    destination_path: String,
    destination_sha256sum: String,

    status: JobStatus,
}

#[derive(Serialize, Deserialize)]
pub struct BatchJob {
    source_dir: String,
    destination_dir: String,
    rules: HashMap<String, String>,
    jobs: Vec<RefCell<Job>>,
}

impl BatchJob {
    pub fn new(source_dir: String, destination_dir: String) -> BatchJob {
        let mut bj = BatchJob {
            source_dir: source_dir,
            destination_dir: destination_dir,
            rules: HashMap::new(),
            jobs: vec![],
        };

        bj.rules.insert(
            String::from("avi"),
            String::from("ffmpeg -i $file_path -c:v libx264 -preset slow $file_path_out"),
        );

        bj
    }

    pub fn init(&mut self) -> Result<(), Box<Error>> {
        let mut lst = vec![];

        // flatten
        {
            let mut add_it = |x: &DirEntry| {
                lst.push(x.path());
            };

            self.visit_dirs(Path::new(&self.source_dir.trim()), &mut add_it)?;
        }

        println!("scanned {} files", lst.len());

        for file in lst.iter() {
            let source_path = file.to_str().unwrap();
            println!("processing {}", source_path);
            let mut file = File::open(source_path).expect("Unable to create config file");
            let hash = Sha256::digest_reader(&mut file)?;

            let job = Job {
                source_path: String::from(source_path),
                source_sha256sum: String::from(format!("{:x}", hash)),
                destination_path: String::new(),
                destination_sha256sum: String::new(),
                status: JobStatus::Pending,
            };

            self.jobs.push(RefCell::new(job));
        }

        Ok(())
    }

    pub fn load_from_file(dir: &str) -> Result<BatchJob, Box<Error>> {
        let filepath = format!("{}/dbfc.config", dir);
        let mut file = File::open(filepath)?;
        let mut buffer = String::new();
        file.read_to_string(&mut buffer)?;
        let bj: BatchJob = serde_json::from_str(&buffer)?;

        Ok(bj)
    }

    pub fn save_to_file(&self) -> Result<(), Box<Error>> {
        let json_result = serde_json::to_string_pretty(&self)?;

        let config_filepath = format!("{}/dbfc.config", self.source_dir);
        let mut file = File::create(config_filepath).expect("Unable to create config file");
        file.write_all(json_result.as_bytes())
            .expect("Unable to write to config file");

        Ok(())
    }

    pub fn run(&mut self) {
        let pending_jobs = self.jobs
            .iter()
            .filter(|j| j.borrow().status == JobStatus::Pending)
            .collect::<Vec<&RefCell<Job>>>();

        println!("{} pending jobs", pending_jobs.len());

        for job_ref in pending_jobs {
            let mut job = job_ref.borrow_mut();

            match self.run_job(&job, self.source_dir.as_str()) {
                Ok(_) => {
                    job.status = JobStatus::Done;
                }
                Err(why) => {
                    job.status = JobStatus::Error;
                    println!("   {}", why);
                }
            }
        }

        if let Err(why) = self.save_to_file() {
            println!("Failed to save batch job state: {}", why);
        }
    }

    fn run_job(&self, job: &Job, source_dir: &str) -> Result<(), String> {
        let path = Path::new(job.source_path.as_str());
        let sub_path = match path.strip_prefix(source_dir) {
            Ok(sp) => sp,
            Err(e) => {
                return Err(format!("   Could not get sub-path: {}", e));
            }
        };

        println!("Processing: {:?}", sub_path);

        if !path.exists() {
            return Err(format!("   Path does not exist"));
        } else if !path.is_file() {
            return Err(format!("   Path is not a file"));
        }

        let extension = match path.extension() {
            Some(e) => match e.to_str() {
                Some(es) => es,
                None => {
                    return Err(format!("   Error converting to string"));
                }
            },
            None => {
                return Err(format!("   No extension"));
            }
        };

        let rule = match self.rules.get(extension) {
            Some(r) => r,
            None => {
                return Err(format!("   No rule defined for extension '{}'", extension));
            }
        };

        let out_path = Path::new(self.destination_dir.as_str()).join(sub_path);
        println!("   Writting to {:?}", out_path);

        // get args and replace any variables
        let raw_args: Vec<&str> = rule.split_whitespace().collect();
        let mut args = Vec::with_capacity(raw_args.len());
        for arg in raw_args.iter() {
            args.push(String::from(match *arg {
                "$file_path" => path.to_str().unwrap(),
                "$file_path_out" => out_path.to_str().unwrap(),
                _ => arg,
            }));
        }
        println!("   {:?}", args);

        // create the output path
        if let Err(why) = fs::create_dir_all(out_path.parent().unwrap()) {
            return Err(format!("   Unable to create out file directory: {}", why));
        }

        let command_name = String::from(args[0].as_str());
        let command_args = &args[1..args.len()];
        let command = Command::new(command_name)
            .args(command_args)
            .output()
            .expect(format!("Failed to execute command").as_str());

        println!("stdout: {}", String::from_utf8_lossy(&command.stdout));
        println!("stderr: {}", String::from_utf8_lossy(&command.stderr));

        Ok(())
    }

    fn visit_dirs(&self, dir: &Path, cb: &mut FnMut(&DirEntry)) -> Result<(), Box<Error>> {
        if dir.is_dir() {
            for entry in fs::read_dir(dir)? {
                let entry = entry?;
                let path = entry.path();
                if path.is_dir() {
                    self.visit_dirs(&path, cb)?;
                } else {
                    cb(&entry);
                }
            }
        }

        Ok(())
    }

    fn convert_video(&self, source_filepath: &Path, output_dir: &str) {
        let mut out_file = PathBuf::new();
        out_file.push(output_dir);
        out_file.push(source_filepath.file_stem().unwrap());
        out_file.set_extension("mp4");

        let mut output = Command::new("ffmpeg")
            .arg("-loglevel")
            .arg("warning")
            .arg("-i")
            .arg(source_filepath)
            .arg("-y")
            .arg("-c:v")
            .arg("libx264")
            .arg("-preset")
            .arg("ultrafast")
            .arg("-framerate")
            .arg("30")
            .arg("-vsync")
            .arg("1")
            .arg("-crf")
            .arg("20")
            .arg(out_file.as_path())
            .spawn()
            .unwrap_or_else(|e| panic!("failed to execute process: {}", e));

        output.wait().unwrap();
    }
}
<commit_msg>Handle exit code for commands execute<commit_after>extern crate serde_json;

use sha2::{Digest, Sha256};
use std::cell::RefCell;
use std::collections::HashMap;
use std::error::Error;
use std::fs::{self, DirEntry, File};
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

#[derive(Serialize, Deserialize, PartialEq)]
pub enum JobStatus {
    Pending,
    Done,
    Error,
}

#[derive(Serialize, Deserialize)]
pub struct Job {
    source_path: String,
    source_sha256sum: String,
    destination_path: String,
    destination_sha256sum: String,

    status: JobStatus,
}

#[derive(Serialize, Deserialize)]
pub struct BatchJob {
    source_dir: String,
    destination_dir: String,
    rules: HashMap<String, String>,
    jobs: Vec<RefCell<Job>>,
}

impl BatchJob {
    pub fn new(source_dir: String, destination_dir: String) -> BatchJob {
        let mut bj = BatchJob {
            source_dir: source_dir,
            destination_dir: destination_dir,
            rules: HashMap::new(),
            jobs: vec![],
        };

        bj.rules.insert(
            String::from("avi"),
            String::from("ffmpeg -i $file_path -c:v libx264 -preset slow $file_path_out"),
        );

        bj
    }

    pub fn init(&mut self) -> Result<(), Box<Error>> {
        let mut lst = vec![];

        // flatten
        {
            let mut add_it = |x: &DirEntry| {
                lst.push(x.path());
            };

            self.visit_dirs(Path::new(&self.source_dir.trim()), &mut add_it)?;
        }

        println!("scanned {} files", lst.len());

        for file in lst.iter() {
            let source_path = file.to_str().unwrap();
            println!("processing {}", source_path);
            let mut file = File::open(source_path).expect("Unable to create config file");
            let hash = Sha256::digest_reader(&mut file)?;

            let job = Job {
                source_path: String::from(source_path),
                source_sha256sum: String::from(format!("{:x}", hash)),
                destination_path: String::new(),
                destination_sha256sum: String::new(),
                status: JobStatus::Pending,
            };

            self.jobs.push(RefCell::new(job));
        }

        Ok(())
    }

    pub fn load_from_file(dir: &str) -> Result<BatchJob, Box<Error>> {
        let filepath = format!("{}/dbfc.config", dir);
        let mut file = File::open(filepath)?;
        let mut buffer = String::new();
        file.read_to_string(&mut buffer)?;
        let bj: BatchJob = serde_json::from_str(&buffer)?;

        Ok(bj)
    }

    pub fn save_to_file(&self) -> Result<(), Box<Error>> {
        let json_result = serde_json::to_string_pretty(&self)?;

        let config_filepath = format!("{}/dbfc.config", self.source_dir);
        let mut file = File::create(config_filepath).expect("Unable to create config file");
        file.write_all(json_result.as_bytes())
            .expect("Unable to write to config file");

        Ok(())
    }

    pub fn run(&mut self) {
        let pending_jobs = self.jobs
            .iter()
            .filter(|j| j.borrow().status == JobStatus::Pending)
            .collect::<Vec<&RefCell<Job>>>();

        println!("{} pending jobs", pending_jobs.len());

        for job_ref in pending_jobs {
            let mut job = job_ref.borrow_mut();

            match self.run_job(&job, self.source_dir.as_str()) {
                Ok(_) => {
                    job.status = JobStatus::Done;
                }
                Err(why) => {
                    job.status = JobStatus::Error;
                    println!("   {}", why);
                }
            }
        }

        if let Err(why) = self.save_to_file() {
            println!("Failed to save batch job state: {}", why);
        }
    }

    fn run_job(&self, job: &Job, source_dir: &str) -> Result<(), String> {
        let path = Path::new(job.source_path.as_str());
        let sub_path = match path.strip_prefix(source_dir) {
            Ok(sp) => sp,
            Err(e) => {
                return Err(format!("   Could not get sub-path: {}", e));
            }
        };

        println!("Processing: {:?}", sub_path);

        if !path.exists() {
            return Err(format!("   Path does not exist"));
        } else if !path.is_file() {
            return Err(format!("   Path is not a file"));
        }

        let extension = match path.extension() {
            Some(e) => match e.to_str() {
                Some(es) => es,
                None => {
                    return Err(format!("   Error converting to string"));
                }
            },
            None => {
                return Err(format!("   No extension"));
            }
        };

        let rule = match self.rules.get(extension) {
            Some(r) => r,
            None => {
                return Err(format!("   No rule defined for extension '{}'", extension));
            }
        };

        let out_path = Path::new(self.destination_dir.as_str()).join(sub_path);
        println!("   Writting to {:?}", out_path);

        // get args and replace any variables
        let raw_args: Vec<&str> = rule.split_whitespace().collect();
        let mut args = Vec::with_capacity(raw_args.len());
        for arg in raw_args.iter() {
            args.push(String::from(match *arg {
                "$file_path" => path.to_str().unwrap(),
                "$file_path_out" => out_path.to_str().unwrap(),
                _ => arg,
            }));
        }
        println!("   {:?}", args);

        // create the output path
        if let Err(why) = fs::create_dir_all(out_path.parent().unwrap()) {
            return Err(format!("   Unable to create out file directory: {}", why));
        }

        let command_name = String::from(args[0].as_str());
        let command_args = &args[1..args.len()];
        let command = Command::new(&command_name)
            .args(command_args)
            .output()
            .expect(format!("Failed to execute command").as_str());

        println!("stdout: {}", String::from_utf8_lossy(&command.stdout));
        println!("stderr: {}", String::from_utf8_lossy(&command.stderr));

        if !command.status.success() {
            return Err(format!(
                "    Command '{}' failed with exit code {:?}",
                &command_name,
                command.status.code()
            ));
        }

        Ok(())
    }

    fn visit_dirs(&self, dir: &Path, cb: &mut FnMut(&DirEntry)) -> Result<(), Box<Error>> {
        if dir.is_dir() {
            for entry in fs::read_dir(dir)? {
                let entry = entry?;
                let path = entry.path();
                if path.is_dir() {
                    self.visit_dirs(&path, cb)?;
                } else {
                    cb(&entry);
                }
            }
        }

        Ok(())
    }

    fn convert_video(&self, source_filepath: &Path, output_dir: &str) {
        let mut out_file = PathBuf::new();
        out_file.push(output_dir);
        out_file.push(source_filepath.file_stem().unwrap());
        out_file.set_extension("mp4");

        let mut output = Command::new("ffmpeg")
            .arg("-loglevel")
            .arg("warning")
            .arg("-i")
            .arg(source_filepath)
            .arg("-y")
            .arg("-c:v")
            .arg("libx264")
            .arg("-preset")
            .arg("ultrafast")
            .arg("-framerate")
            .arg("30")
            .arg("-vsync")
            .arg("1")
            .arg("-crf")
            .arg("20")
            .arg(out_file.as_path())
            .spawn()
            .unwrap_or_else(|e| panic!("failed to execute process: {}", e));

        output.wait().unwrap();
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

use futures::{Async, Poll, Future};

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (Option<Vec<u8>>, R);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((Some(result), self.reader.take().unwrap())))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((None, self.reader.take().unwrap())))
                }

                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W> where W: ::std::io::Write {
    writer: Option<W>,
    message: Vec<u8>,
    pos: usize,
    wrote_header: bool,
}

impl <W> Writing<W> where W: ::std::io::Write {
    fn new(writer: W, message: Vec<u8>) -> Writing<W> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W> Future for Writing<W> where W: ::std::io::Write {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr) -> Box<Future<Item=(), Error=::std::io::Error> + Send> {
    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
        tokio_core::io::write_all(stream, [0]).map(|(a,_)| a)
    });

    let mut subscribers = Vec::new();
    for _ in 0..3 {
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
            tokio_core::io::write_all(stream, [1]).map(|(a,_)| a)
        }));
    }

    Box::new(publisher.join(::all::All::new(subscribers.into_iter())).and_then(|(publisher, subscribers)| {
        println!("connected");

        tie_knot((publisher, subscribers, 10u32), |(publisher, subscribers, n)| {
            println!("looping {}", n);
            Writing::new(publisher, vec![n as u8, 1,2,3]).and_then(move |publisher| {
                futures::finished(((publisher, subscribers, n - 1), n > 0))
            })
        }).map(|_| ())
    }))
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = "127.0.0.1:8080";
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let _child = ::std::process::Command::new(executable)
        .arg(addr_str)
        .spawn();

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();

    let f = pool.spawn(new_task(&handle, &addr));

    try!(core.run(f));

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<commit_msg>fix bug<commit_after>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

use futures::{Async, Poll, Future};

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (Option<Vec<u8>>, R);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((Some(result), self.reader.take().unwrap())))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((None, self.reader.take().unwrap())))
                }

                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W> where W: ::std::io::Write {
    writer: Option<W>,
    message: Vec<u8>,
    pos: usize,
    wrote_header: bool,
}

impl <W> Writing<W> where W: ::std::io::Write {
    fn new(writer: W, message: Vec<u8>) -> Writing<W> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W> Future for Writing<W> where W: ::std::io::Write {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr) -> Box<Future<Item=(), Error=::std::io::Error> + Send> {
    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
        tokio_core::io::write_all(stream, [0]).map(|(a,_)| a)
    });

    let mut subscribers = Vec::new();
    for _ in 0..3 {
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
            tokio_core::io::write_all(stream, [1]).map(|(a,_)| a)
        }));
    }

    Box::new(publisher.join(::all::All::new(subscribers.into_iter())).and_then(|(publisher, subscribers)| {
        println!("connected");

        tie_knot((publisher, subscribers, 10i32), |(publisher, subscribers, n)| {
            println!("looping {}", n);
            Writing::new(publisher, vec![n as u8, 1,2,3]).and_then(move |publisher| {
                futures::finished(((publisher, subscribers, n - 1), n > 0))
            })
        }).map(|_| ())
    }))
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = "127.0.0.1:8080";
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let _child = ::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::inherit())
        .stderr(::std::process::Stdio::inherit())
        .spawn();

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();

    let f = pool.spawn(new_task(&handle, &addr));

    try!(core.run(f));

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<|endoftext|>


<commit_before>extern crate multipart;

extern crate log;

use log::{LogRecord, LogMetadata, LogLevelFilter};

use multipart::server::Multipart;

use std::fs::File;
use std::env;

const LOG_LEVEL: LogLevelFilter = LogLevelFilter::Debug;

struct SimpleLogger;

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        LOG_LEVEL.to_log_level()
            .map_or(false, |level| metadata.level() <= level)
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}

fn main() {
    log::set_logger(|max_log_level| {
        max_log_level.set(LOG_LEVEL);
        Box::new(SimpleLogger)
    });

    let mut args = env::args().skip(1);

    let boundary = args.next().expect("Boundary must be provided as the first argument");

    let file = args.next().expect("Filename must be provided as the second argument");

    let file = File::open(file).expect("Could not open file");

    let mut multipart = Multipart::with_body(file, boundary);

    while let Some(field) = multipart.read_entry().unwrap() {
        println!("Read field: {:?}", field.name);
    }

    println!("All entries read!");
}<commit_msg>Get rid of unused_result warning<commit_after>extern crate multipart;

extern crate log;

use log::{LogRecord, LogMetadata, LogLevelFilter};

use multipart::server::Multipart;

use std::fs::File;
use std::env;

const LOG_LEVEL: LogLevelFilter = LogLevelFilter::Debug;

struct SimpleLogger;

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        LOG_LEVEL.to_log_level()
            .map_or(false, |level| metadata.level() <= level)
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}

fn main() {
    log::set_logger(|max_log_level| {
        max_log_level.set(LOG_LEVEL);
        Box::new(SimpleLogger)
    }).expect("Could not initialize logger");

    let mut args = env::args().skip(1);

    let boundary = args.next().expect("Boundary must be provided as the first argument");

    let file = args.next().expect("Filename must be provided as the second argument");

    let file = File::open(file).expect("Could not open file");

    let mut multipart = Multipart::with_body(file, boundary);

    while let Some(field) = multipart.read_entry().unwrap() {
        println!("Read field: {:?}", field.name);
    }

    println!("All entries read!");
}<|endoftext|>


<commit_before>use CliResult;
use config::{Delimiter, Config};
use util;

static USAGE: &'static str = "
Prints a count of the number of records in the CSV data.

Usage:
    xsv count [options] [<input>]

Common options:
    -h, --help             Display this message
    -n, --no-headers       When set, the first row will not be interpreted
                           as headers. (i.e., They are not searched, analyzed,
                           sliced, etc.)
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. [default: ,]
";

#[deriving(Decodable)]
struct Args {
    arg_input: Option<String>,
    flag_no_headers: bool,
    flag_delimiter: Delimiter,
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = try!(util::get_args(USAGE, argv));
    let conf = Config::new(args.arg_input)
                      .delimiter(args.flag_delimiter)
                      .no_headers(args.flag_no_headers);

    let mut count =
        match try!(conf.indexed()) {
            Some(idx) => idx.count(),
            None => {
                let mut rdr = try!(conf.reader());
                let mut count = 0u64;
                while !rdr.done() {
                    loop {
                        match rdr.next_field() {
                            None => break,
                            Some(r) => { try!(r); }
                        }
                    }
                    count += 1;
                }
                count
            }
        };
    if !args.flag_no_headers {
        count -= 1;
    }
    println!("{:u}", count);
    Ok(())
}
<commit_msg>Fix a bad bug in counting empty CSV data.<commit_after>use CliResult;
use config::{Delimiter, Config};
use util;

static USAGE: &'static str = "
Prints a count of the number of records in the CSV data.

Usage:
    xsv count [options] [<input>]

Common options:
    -h, --help             Display this message
    -n, --no-headers       When set, the first row will not be interpreted
                           as headers. (i.e., They are not searched, analyzed,
                           sliced, etc.)
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. [default: ,]
";

#[deriving(Decodable)]
struct Args {
    arg_input: Option<String>,
    flag_no_headers: bool,
    flag_delimiter: Delimiter,
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = try!(util::get_args(USAGE, argv));
    let conf = Config::new(args.arg_input)
                      .delimiter(args.flag_delimiter)
                      .no_headers(args.flag_no_headers);

    let mut count =
        match try!(conf.indexed()) {
            Some(idx) => idx.count(),
            None => {
                let mut rdr = try!(conf.reader());
                let mut count = 0u64;
                let mut seen_field = false;
                while !rdr.done() {
                    loop {
                        match rdr.next_field() {
                            None => break,
                            Some(r) => { seen_field = true; try!(r); }
                        }
                    }
                    if seen_field { count += 1; }
                }
                count
            }
        };
    if !args.flag_no_headers {
        count -= 1;
    }
    println!("{:u}", count);
    Ok(())
}
<|endoftext|>


<commit_before>#![crate_name = "comm"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

use std::cmp::Ord;
use std::io::{BufferedReader, IoResult, print};
use std::io::fs::File;
use std::io::stdio::stdin;
use std::path::Path;

static NAME : &'static str = "comm";
static VERSION : &'static str = "1.0.0";

fn mkdelim(col: uint, opts: &getopts::Matches) -> String {
    let mut s = String::new();
    let delim = match opts.opt_str("output-delimiter") {
        Some(d) => d.clone(),
        None => "\t".to_string(),
    };

    if col > 1 && !opts.opt_present("1") {
        s.push_str(delim.as_slice());
    }
    if col > 2 && !opts.opt_present("2") {
        s.push_str(delim.as_slice());
    }

    s
}

fn ensure_nl(line: String) -> String {
    match line.as_slice().chars().last() {
        Some('\n') => line,
        _ => line + "\n",
    }
}

fn comm(a: &mut Box<Buffer>, b: &mut Box<Buffer>, opts: &getopts::Matches) {

    let delim = Vec::from_fn(4, |col| mkdelim(col, opts));

    let mut ra = a.read_line();
    let mut rb = b.read_line();

    while ra.is_ok() || rb.is_ok() {
        let ord = match (ra.clone(), rb.clone()) {
            (Err(_), Ok(_))  => Greater,
            (Ok(_) , Err(_)) => Less,
            (Ok(s0), Ok(s1)) => s0.cmp(&s1),
            _ => unreachable!(),
        };

        match ord {
            Less => {
                if !opts.opt_present("1") {
                    print!("{}{}", delim[1], ra.map(ensure_nl).unwrap());
                }
                ra = a.read_line();
            }
            Greater => {
                if !opts.opt_present("2") {
                    print!("{}{}", delim[2], rb.map(ensure_nl).unwrap());
                }
                rb = b.read_line();
            }
            Equal => {
                if !opts.opt_present("3") {
                    print!("{}{}", delim[3], ra.map(ensure_nl).unwrap());
                }
                ra = a.read_line();
                rb = b.read_line();
            }
        }
    }
}

fn open_file(name: &str) -> IoResult<Box<Buffer>> {
    match name {
        "-" => Ok(box stdin() as Box<Buffer>),
        _   => {
            let f = try!(File::open(&Path::new(name)));
            Ok(box BufferedReader::new(f) as Box<Buffer>)
        }
    }
}

pub fn uumain(args: Vec<String>) -> int {
    let opts = [
        getopts::optflag("1", "", "suppress column 1 (lines uniq to FILE1)"),
        getopts::optflag("2", "", "suppress column 2 (lines uniq to FILE2)"),
        getopts::optflag("3", "", "suppress column 3 (lines that appear in both files)"),
        getopts::optopt("", "output-delimiter", "separate columns with STR", "STR"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(err) => panic!("{}", err),
    };

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        return 0;
    }

    if matches.opt_present("help") || matches.free.len() != 2 {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {} [OPTIONS] FILE1 FILE2", NAME);
        println!("");
        print(getopts::usage("Compare sorted files line by line.", opts.as_slice()).as_slice());
        if matches.free.len() != 2 {
            return 1;
        }
        return 0;
    }

    let mut f1 = open_file(matches.free[0].as_slice()).unwrap();
    let mut f2 = open_file(matches.free[1].as_slice()).unwrap();

    comm(&mut f1, &mut f2, &matches);

    0
}
<commit_msg>comm: fix build<commit_after>#![crate_name = "comm"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

use std::cmp::Ord;
use std::io::{BufferedReader, IoResult, print};
use std::io::fs::File;
use std::io::stdio::{stdin, StdReader};
use std::path::Path;

static NAME : &'static str = "comm";
static VERSION : &'static str = "1.0.0";

fn mkdelim(col: uint, opts: &getopts::Matches) -> String {
    let mut s = String::new();
    let delim = match opts.opt_str("output-delimiter") {
        Some(d) => d.clone(),
        None => "\t".to_string(),
    };

    if col > 1 && !opts.opt_present("1") {
        s.push_str(delim.as_slice());
    }
    if col > 2 && !opts.opt_present("2") {
        s.push_str(delim.as_slice());
    }

    s
}

fn ensure_nl(line: String) -> String {
    match line.as_slice().chars().last() {
        Some('\n') => line,
        _ => line + "\n",
    }
}

enum LineReader {
    Stdin(BufferedReader<StdReader>),
    FileIn(BufferedReader<File>)
}

impl LineReader {
    fn read_line(&mut self) -> IoResult<String> {
        match self {
            &Stdin(ref mut r)  => r.read_line(),
            &FileIn(ref mut r) => r.read_line(),
        }
    }
}

fn comm(a: &mut LineReader, b: &mut LineReader, opts: &getopts::Matches) {

    let delim = Vec::from_fn(4, |col| mkdelim(col, opts));

    let mut ra = a.read_line();
    let mut rb = b.read_line();

    while ra.is_ok() || rb.is_ok() {
        let ord = match (ra.clone(), rb.clone()) {
            (Err(_), Ok(_))  => Greater,
            (Ok(_) , Err(_)) => Less,
            (Ok(s0), Ok(s1)) => s0.cmp(&s1),
            _ => unreachable!(),
        };

        match ord {
            Less => {
                if !opts.opt_present("1") {
                    print!("{}{}", delim[1], ra.map(ensure_nl).unwrap());
                }
                ra = a.read_line();
            }
            Greater => {
                if !opts.opt_present("2") {
                    print!("{}{}", delim[2], rb.map(ensure_nl).unwrap());
                }
                rb = b.read_line();
            }
            Equal => {
                if !opts.opt_present("3") {
                    print!("{}{}", delim[3], ra.map(ensure_nl).unwrap());
                }
                ra = a.read_line();
                rb = b.read_line();
            }
        }
    }
}

fn open_file(name: &str) -> IoResult<LineReader> {
    match name {
        "-" => Ok(LineReader::Stdin(stdin())),
        _   => {
            let f = try!(std::io::File::open(&Path::new(name)));
            Ok(LineReader::FileIn(BufferedReader::new(f)))
        }
    }
}

pub fn uumain(args: Vec<String>) -> int {
    let opts = [
        getopts::optflag("1", "", "suppress column 1 (lines uniq to FILE1)"),
        getopts::optflag("2", "", "suppress column 2 (lines uniq to FILE2)"),
        getopts::optflag("3", "", "suppress column 3 (lines that appear in both files)"),
        getopts::optopt("", "output-delimiter", "separate columns with STR", "STR"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(err) => panic!("{}", err),
    };

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        return 0;
    }

    if matches.opt_present("help") || matches.free.len() != 2 {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {} [OPTIONS] FILE1 FILE2", NAME);
        println!("");
        print(getopts::usage("Compare sorted files line by line.", opts.as_slice()).as_slice());
        if matches.free.len() != 2 {
            return 1;
        }
        return 0;
    }

    let mut f1 = open_file(matches.free[0].as_slice()).unwrap();
    let mut f2 = open_file(matches.free[1].as_slice()).unwrap();

    comm(&mut f1, &mut f2, &matches);

    0
}
<|endoftext|>


<commit_before>use core::arch::asm;
use libtock_platform::{syscall_class, RawSyscalls, Register};

unsafe impl RawSyscalls for crate::TockSyscalls {
    unsafe fn yield1([Register(r0)]: [Register; 1]) {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::yield1
        // the use of `clobber_abi` allows us this to run on both Thumb-1 and Thumb-2
        unsafe {
            asm!("svc 0",
                 inlateout("r0") r0 => _, // a1
                 // r4-r8 are callee-saved.
                 // r9 is platform-specific. We don't use it in libtock_runtime,
                 // so it is either unused or used as a callee-saved register.
                 // r10 and r11 are callee-saved.

                 // r13 is the stack pointer and must be restored by the callee.
                 // r15 is the program counter.

                 clobber_abi("C"), // a2, a3, a4, ip (r12), lr (r14)
            );
        }
    }

    unsafe fn yield2([Register(r0), Register(r1)]: [Register; 2]) {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::yield2
        // the use of `clobber_abi` allows us this to run on both Thumb-1 and Thumb-2
        unsafe {
            asm!("svc 0",
                 inlateout("r0") r0 => _, // a1
                 // r4-r8 are callee-saved.
                 // r9 is platform-specific. We don't use it in libtock_runtime,
                 // so it is either unused or used as a callee-saved register.
                 // r10 and r11 are callee-saved.

                 // r13 is the stack pointer and must be restored by the callee.
                 // r15 is the program counter.

                 clobber_abi("C"), // a2, a3, a4, ip (r12), lr (r14)
            );
        }
    }

    unsafe fn syscall1<const CLASS: usize>([Register(mut r0)]: [Register; 1]) -> [Register; 2] {
        let r1;
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::syscall1
        unsafe {
            // Syscall class 5 is Memop, the only syscall class that syscall1
            // supports.
            asm!("svc 5",
                 inlateout("r0") r0,
                 lateout("r1") r1,
                 options(preserves_flags, nostack, nomem),
            );
        }
        [Register(r0), Register(r1)]
    }

    unsafe fn syscall2<const CLASS: usize>(
        [Register(mut r0), Register(mut r1)]: [Register; 2],
    ) -> [Register; 2] {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::syscall2
        unsafe {
            // TODO: Replace this match statement with a `const` operand when
            // asm_const [1] is stabilized, or redesign RawSyscalls to not need
            // this match statement.
            //
            // [1] https://github.com/rust-lang/rust/issues/93332
            match CLASS {
                syscall_class::MEMOP => asm!("svc 5",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     options(preserves_flags, nostack, nomem)
                ),
                syscall_class::EXIT => asm!("svc 6",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     options(preserves_flags, nostack, nomem)
                ),
                _ => unreachable!(),
            }
        }
        [Register(r0), Register(r1)]
    }

    unsafe fn syscall4<const CLASS: usize>(
        [Register(mut r0), Register(mut r1), Register(mut r2), Register(mut r3)]: [Register; 4],
    ) -> [Register; 4] {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::syscall4
        unsafe {
            // TODO: Replace this match statement with a `const` operand when
            // asm_const [1] is stabilized, or redesign RawSyscalls to not need
            // this match statement.
            //
            // [1] https://github.com/rust-lang/rust/issues/93332
            match CLASS {
                syscall_class::SUBSCRIBE => asm!("svc 1",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                syscall_class::COMMAND => asm!("svc 2",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                syscall_class::ALLOW_RW => asm!("svc 3",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                syscall_class::ALLOW_RO => asm!("svc 4",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                _ => unreachable!(),
            }
        }
        [Register(r0), Register(r1), Register(r2), Register(r3)]
    }
}
<commit_msg>update<commit_after>use core::arch::asm;
use libtock_platform::{syscall_class, RawSyscalls, Register};

unsafe impl RawSyscalls for crate::TockSyscalls {
    unsafe fn yield1([Register(r0)]: [Register; 1]) {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::yield1
        // the use of `clobber_abi` allows us this to run on both Thumb-1 and Thumb-2
        unsafe {
            asm!("svc 0",
                 inlateout("r0") r0 => _, // a1
                 // r4-r8 are callee-saved.
                 // r9 is platform-specific. We don't use it in libtock_runtime,
                 // so it is either unused or used as a callee-saved register.
                 // r10 and r11 are callee-saved.

                 // r13 is the stack pointer and must be restored by the callee.
                 // r15 is the program counter.

                 clobber_abi("C"), // a2, a3, a4, ip (r12), lr (r14)
            );
        }
    }

    unsafe fn yield2([Register(r0), Register(r1)]: [Register; 2]) {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::yield2
        // the use of `clobber_abi` allows us this to run on both Thumb-1 and Thumb-2
        unsafe {
            asm!("svc 0",
                 inlateout("r0") r0 => _, // a1
                 inlateout("r1") r1 => _, // a2
                 // r4-r8 are callee-saved.
                 // r9 is platform-specific. We don't use it in libtock_runtime,
                 // so it is either unused or used as a callee-saved register.
                 // r10 and r11 are callee-saved.

                 // r13 is the stack pointer and must be restored by the callee.
                 // r15 is the program counter.

                 clobber_abi("C"), // a3, a4, ip (r12), lr (r14)
            );
        }
    }

    unsafe fn syscall1<const CLASS: usize>([Register(mut r0)]: [Register; 1]) -> [Register; 2] {
        let r1;
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::syscall1
        unsafe {
            // Syscall class 5 is Memop, the only syscall class that syscall1
            // supports.
            asm!("svc 5",
                 inlateout("r0") r0,
                 lateout("r1") r1,
                 options(preserves_flags, nostack, nomem),
            );
        }
        [Register(r0), Register(r1)]
    }

    unsafe fn syscall2<const CLASS: usize>(
        [Register(mut r0), Register(mut r1)]: [Register; 2],
    ) -> [Register; 2] {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::syscall2
        unsafe {
            // TODO: Replace this match statement with a `const` operand when
            // asm_const [1] is stabilized, or redesign RawSyscalls to not need
            // this match statement.
            //
            // [1] https://github.com/rust-lang/rust/issues/93332
            match CLASS {
                syscall_class::MEMOP => asm!("svc 5",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     options(preserves_flags, nostack, nomem)
                ),
                syscall_class::EXIT => asm!("svc 6",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     options(preserves_flags, nostack, nomem)
                ),
                _ => unreachable!(),
            }
        }
        [Register(r0), Register(r1)]
    }

    unsafe fn syscall4<const CLASS: usize>(
        [Register(mut r0), Register(mut r1), Register(mut r2), Register(mut r3)]: [Register; 4],
    ) -> [Register; 4] {
        // Safety: This matches the invariants required by the documentation on
        // RawSyscalls::syscall4
        unsafe {
            // TODO: Replace this match statement with a `const` operand when
            // asm_const [1] is stabilized, or redesign RawSyscalls to not need
            // this match statement.
            //
            // [1] https://github.com/rust-lang/rust/issues/93332
            match CLASS {
                syscall_class::SUBSCRIBE => asm!("svc 1",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                syscall_class::COMMAND => asm!("svc 2",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                syscall_class::ALLOW_RW => asm!("svc 3",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                syscall_class::ALLOW_RO => asm!("svc 4",
                     inlateout("r0") r0,
                     inlateout("r1") r1,
                     inlateout("r2") r2,
                     inlateout("r3") r3,
                     options(preserves_flags, nostack),
                ),
                _ => unreachable!(),
            }
        }
        [Register(r0), Register(r1), Register(r2), Register(r3)]
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::{BTreeMap, BTreeSet},
    sync::Arc,
};

use eyre::{bail, Result};
use futures::stream::{FuturesUnordered, StreamExt};
use graph::Plan;
use graph_triples::Resource;
use kernels::{KernelSelector, KernelSpace};
use node_address::{Address, AddressMap};
use node_patch::{diff, mutate, Patch};
use stencila_schema::{CodeChunk, CodeExecutableExecuteStatus, CodeExpression, Node};
use tokio::sync::{
    mpsc::{Receiver, Sender, UnboundedSender},
    RwLock,
};

use crate::{
    utils::{resource_to_node, send_patch, send_patches},
    CancelRequest, CompileRequest, Executable, PatchRequest,
};

/// Execute a [`Plan`] on a [`Node`]
///
/// Uses a `RwLock` for `root` and `address_map` so that read locks can be held for as short as
/// time as possible (i.e. not while waiting for execution of tasks, which is what would
/// happen if held by the caller).
///
/// # Arguments
///
/// - `plan`: The plan to be executed
///
/// - `root`: The root node to execute the plan on (takes a read lock)
///
/// - `address_map`: The [`AddressMap`] map for the `root` node (used to locate code nodes
///                  included in the plan within the `root` node; takes a read lock)
///
/// - `patch_request_sender`: A [`PatchRequest`] channel sender to send patches describing the changes to
///                   executed nodes
///
/// - `compile_request_sender`: A [`CompileRequest`] channel sender to request re-compiles due to changes to
///                   executed nodes
///
/// - `cancel_request_receiver`: A [`CancelRequest`] channel receiver to request cancellation of one or more
///                   steps in the plan
///
/// - `kernel_space`: The [`KernelSpace`] within which to execute the plan
///
pub async fn execute(
    plan: &Plan,
    root: &Arc<RwLock<Node>>,
    address_map: &Arc<RwLock<AddressMap>>,
    patch_request_sender: &UnboundedSender<PatchRequest>,
    compile_request_sender: &Sender<CompileRequest>,
    cancel_request_receiver: &mut Receiver<CancelRequest>,
    kernel_space: Option<Arc<KernelSpace>>,
) -> Result<()> {
    let kernel_space = kernel_space.unwrap_or_default();

    // Obtain locks
    let root_guard = root.read().await;
    let address_map_guard = address_map.read().await;

    // Get a snapshot of all nodes involved in the plan at the start
    // Also get their `execute_status` so it can be restored at the end if needs be.
    let mut nodes: BTreeMap<Resource, _> = plan
        .stages
        .iter()
        .flat_map(|stage| stage.steps.iter())
        .filter_map(|step| {
            let resource_info = step.resource_info.clone();
            let resource = &resource_info.resource;
            match resource_to_node(resource, &root_guard, &address_map_guard) {
                Ok((node, node_id, node_address)) => {
                    let execution_status = get_execute_status(&node);
                    Some((
                        resource.clone(),
                        (resource_info, node_id, node_address, node, execution_status),
                    ))
                }
                Err(error) => {
                    tracing::warn!("While executing plan: {}", error);
                    None
                }
            }
        })
        .collect();

    // Release locks
    drop(root_guard);
    drop(address_map_guard);

    // Set the `execute_status` of all nodes to `Scheduled` or `ScheduledPreviouslyFailed`
    // and send the resulting patch
    send_patches(
        patch_request_sender,
        nodes
            .values_mut()
            .map(|(_, node_id, node_address, node, ..)| {
                set_execute_status_scheduled(node_id, node_address, node)
            })
            .collect(),
    );

    // For each stage in plan...
    let stage_count = plan.stages.len();
    let mut dependencies_failed = false;
    for (stage_index, stage) in plan.stages.iter().enumerate() {
        // Check for any cancellation requests
        let mut was_cancelled = false;
        while let Ok(request) = cancel_request_receiver.try_recv() {
            if let Some(start) = request.start {
                // A start node was specified...
                if let Some((node_id, node_address, node)) =
                    nodes
                        .values_mut()
                        .find_map(|(_, node_id, node_address, node, ..)| {
                            if *node_id == start {
                                Some((node_id, node_address, node))
                            } else {
                                None
                            }
                        })
                {
                    // TODO: This just sets status to cancelled and stops the loop. It is
                    // temporary
                    send_patch(
                        patch_request_sender,
                        set_execute_status_cancelled(node_id, node_address, node),
                    );
                    was_cancelled = true;
                    break;
                }
            } else {
                // TODO No start node was specified so cancel them all.
                was_cancelled = true;
                break;
            }
        }
        if was_cancelled {
            tracing::debug!(
                "Received cancellation request before stage {}/{}",
                stage_index + 1,
                stage_count
            );
            break;
        }

        // Before running the steps in this stage, check that all their dependencies have succeeded
        // and stop if they have not. Collects to a `BTreeSet` to generate unique set (some steps in
        // the stage may have the shared dependencies)
        dependencies_failed = stage
            .steps
            .iter()
            .flat_map(|step| step.resource_info.dependencies.iter().flatten())
            .collect::<BTreeSet<&Resource>>()
            .iter()
            .filter_map(|dependency| nodes.get(dependency))
            .map(|tuple| (&tuple.1, get_execute_status(&tuple.3)))
            .any(|(node_id, status)| {
                tracing::trace!(
                    "Status of dependency of stage {}/{} `{}`: {:?}",
                    stage_index + 1,
                    stage_count,
                    node_id,
                    status
                );
                matches!(
                    status,
                    None | Some(CodeExecutableExecuteStatus::Failed)
                        | Some(CodeExecutableExecuteStatus::Cancelled)
                )
            });
        if dependencies_failed {
            tracing::debug!(
                "Stopping before stage {}/{}: some dependencies failed or were cancelled",
                stage_index + 1,
                stage_count
            );
            break;
        }

        tracing::debug!("Starting stage {}/{}", stage_index + 1, stage_count);

        // Create a kernel task for each step in this stage
        let step_count = stage.steps.len();
        let mut patches = Vec::with_capacity(step_count);
        let mut futures = Vec::with_capacity(step_count);
        for (step_index, step) in stage.steps.iter().enumerate() {
            // Get the node
            let (_, node_id, node_address, mut node, ..) = nodes
                .get(&step.resource_info.resource)
                .expect("Node for resource should be in nodes")
                .clone();

            // Set the `execute_status` of the node to `Running` or `RunningPreviouslyFailed`
            // and send the resulting patch
            patches.push(set_execute_status_running(
                &node_id,
                &node_address,
                &mut node,
            ));

            // Create clones of variables needed to execute the task
            let kernel_space = kernel_space.clone();
            let kernel_selector = KernelSelector::new(step.kernel_name.clone(), None, None);
            let mut resource_info = step.resource_info.clone();
            let is_fork = step.is_fork;

            // Create a future for the task that will be spawned later
            let future = async move {
                tracing::debug!(
                    "Starting step {}/{} of stage {}/{}",
                    step_index + 1,
                    step_count,
                    stage_index + 1,
                    stage_count
                );

                // Create a mutable draft of the node and execute it in the kernel space
                let mut draft = node.clone();
                match draft
                    .execute(&kernel_space, &kernel_selector, &resource_info, is_fork)
                    .await
                {
                    Ok(_) => {
                        // Update the resource to indicate that the resource was executed
                        let execute_status = match &draft {
                            Node::CodeChunk(CodeChunk { execute_status, .. })
                            | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                                execute_status.clone()
                            }
                            _ => None,
                        };
                        resource_info.did_execute(execute_status);

                        // Generate a patch for the differences resulting from execution
                        let mut patch = diff(&node, &draft);
                        patch.address = Some(node_address.clone());
                        patch.target = Some(node_id.clone());

                        Ok((step_index, resource_info, draft, patch))
                    }
                    Err(error) => bail!(error),
                }
            };
            futures.push(future);
        }
        send_patches(patch_request_sender, patches);

        // Spawn all tasks in the stage and wait for each to finish, sending on the resultant `Patch`
        // for application and publishing (if it is not empty)
        // TODO: Replace `FuturesUnordered` with `TaskSet`. See https://news.ycombinator.com/item?id=29912386
        let mut results = futures
            .into_iter()
            .map(tokio::spawn)
            .collect::<FuturesUnordered<_>>();
        while let Some(result) = results.next().await {
            let (step_index, resource_info, node, patch) = result??;

            tracing::debug!(
                "Finished step {}/{} of stage {}/{}",
                step_index + 1,
                step_count,
                stage_index + 1,
                stage_count
            );

            // Send the patch reflecting the changed state of the executed node
            send_patch(patch_request_sender, patch);

            // Update the state of the node in this function's record of nodes
            nodes
                .entry(resource_info.resource.clone())
                .and_modify(|info| info.3 = node);

            // Send a patch request so that properties of other nodes such as `code_dependencies` and
            // `code_dependents` get updated with the new execution status the node that was executed in
            // this task. Previously we tried to take shortcut to this by just updating the graph and
            // calling `compile_no_walk` ut that proved unreliable so instead make a (debounced) request
            if let Err(..) = compile_request_sender
                .send(CompileRequest::new(false, None))
                .await
            {
                tracing::debug!("When sending compile request: receiver dropped");
            }
        }

        tracing::debug!("Finished stage {}/{}", stage_index + 1, stage_count);
    }

    // For nodes that were scheduled but never got to run because dependencies did not succeed,
    // restore their previous execution status
    if dependencies_failed {
        send_patches(
            patch_request_sender,
            nodes
                .values_mut()
                .map(|(_, node_id, node_address, node, execute_status)| {
                    restore_previous_execute_status(node_id, node_address, node, execute_status)
                })
                .collect(),
        );
    }

    Ok(())
}

fn get_execute_status(node: &Node) -> Option<CodeExecutableExecuteStatus> {
    match node {
        Node::CodeChunk(CodeChunk { execute_status, .. })
        | Node::CodeExpression(CodeExpression { execute_status, .. }) => execute_status.clone(),
        // At present, assumes the execution of parameters always succeeds
        Node::Parameter(..) => Some(CodeExecutableExecuteStatus::Succeeded),
        _ => None,
    }
}

fn set_execute_status_scheduled(node_id: &str, node_address: &Address, node: &mut Node) -> Patch {
    mutate(
        node,
        Some(node_id.to_string()),
        Some(node_address.clone()),
        &|node: &mut Node| match node {
            Node::CodeChunk(CodeChunk { execute_status, .. })
            | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                *execute_status = Some(match execute_status {
                    Some(CodeExecutableExecuteStatus::Failed) => {
                        CodeExecutableExecuteStatus::ScheduledPreviouslyFailed
                    }
                    _ => CodeExecutableExecuteStatus::Scheduled,
                });
            }
            _ => {}
        },
    )
}

fn set_execute_status_running(node_id: &str, node_address: &Address, node: &mut Node) -> Patch {
    mutate(
        node,
        Some(node_id.to_string()),
        Some(node_address.clone()),
        &|node: &mut Node| match node {
            Node::CodeChunk(CodeChunk { execute_status, .. })
            | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                *execute_status = Some(match execute_status {
                    Some(CodeExecutableExecuteStatus::Failed)
                    | Some(CodeExecutableExecuteStatus::ScheduledPreviouslyFailed) => {
                        CodeExecutableExecuteStatus::RunningPreviouslyFailed
                    }
                    _ => CodeExecutableExecuteStatus::Running,
                });
            }
            _ => {}
        },
    )
}

fn set_execute_status_cancelled(node_id: &str, node_address: &Address, node: &mut Node) -> Patch {
    mutate(
        node,
        Some(node_id.to_string()),
        Some(node_address.clone()),
        &|node: &mut Node| match node {
            Node::CodeChunk(CodeChunk { execute_status, .. })
            | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                *execute_status = Some(CodeExecutableExecuteStatus::Cancelled);
            }
            _ => {}
        },
    )
}

fn restore_previous_execute_status(
    node_id: &str,
    node_address: &Address,
    node: &mut Node,
    previous_execute_status: &Option<CodeExecutableExecuteStatus>,
) -> Patch {
    mutate(
        node,
        Some(node_id.to_string()),
        Some(node_address.clone()),
        &|node: &mut Node| match node {
            Node::CodeChunk(CodeChunk { execute_status, .. })
            | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                if matches!(
                    execute_status,
                    Some(CodeExecutableExecuteStatus::Scheduled)
                        | Some(CodeExecutableExecuteStatus::ScheduledPreviouslyFailed)
                ) {
                    *execute_status = previous_execute_status.clone();
                }
            }
            _ => {}
        },
    )
}
<commit_msg>refactor(Node execution): Add `NodeInfo` struct for less verbose code<commit_after>use std::{
    collections::{BTreeMap, BTreeSet},
    sync::Arc,
};

use eyre::{bail, Result};
use futures::stream::{FuturesUnordered, StreamExt};
use graph::Plan;
use graph_triples::{Resource, ResourceInfo};
use kernels::{KernelSelector, KernelSpace};
use node_address::{Address, AddressMap};
use node_patch::{diff, mutate, Patch};
use stencila_schema::{CodeChunk, CodeExecutableExecuteStatus, CodeExpression, Node};
use tokio::sync::{
    mpsc::{Receiver, Sender, UnboundedSender},
    RwLock,
};

use crate::{
    utils::{resource_to_node, send_patch, send_patches},
    CancelRequest, CompileRequest, Executable, PatchRequest,
};

/// Execute a [`Plan`] on a [`Node`]
///
/// Uses a `RwLock` for `root` and `address_map` so that read locks can be held for as short as
/// time as possible (i.e. not while waiting for execution of tasks, which is what would
/// happen if held by the caller).
///
/// # Arguments
///
/// - `plan`: The plan to be executed
///
/// - `root`: The root node to execute the plan on (takes a read lock)
///
/// - `address_map`: The [`AddressMap`] map for the `root` node (used to locate code nodes
///                  included in the plan within the `root` node; takes a read lock)
///
/// - `patch_request_sender`: A [`PatchRequest`] channel sender to send patches describing the changes to
///                   executed nodes
///
/// - `compile_request_sender`: A [`CompileRequest`] channel sender to request re-compiles due to changes to
///                   executed nodes
///
/// - `cancel_request_receiver`: A [`CancelRequest`] channel receiver to request cancellation of one or more
///                   steps in the plan
///
/// - `kernel_space`: The [`KernelSpace`] within which to execute the plan
///
pub async fn execute(
    plan: &Plan,
    root: &Arc<RwLock<Node>>,
    address_map: &Arc<RwLock<AddressMap>>,
    patch_request_sender: &UnboundedSender<PatchRequest>,
    compile_request_sender: &Sender<CompileRequest>,
    cancel_request_receiver: &mut Receiver<CancelRequest>,
    kernel_space: Option<Arc<KernelSpace>>,
) -> Result<()> {
    let kernel_space = kernel_space.unwrap_or_default();

    // Obtain locks
    let root_guard = root.read().await;
    let address_map_guard = address_map.read().await;

    // Get a snapshot of all nodes involved in the plan at the start
    let mut node_infos: BTreeMap<Resource, NodeInfo> = plan
        .stages
        .iter()
        .flat_map(|stage| stage.steps.iter())
        .filter_map(|step| {
            let resource_info = step.resource_info.clone();
            let resource = &resource_info.resource;
            match resource_to_node(resource, &root_guard, &address_map_guard) {
                Ok((node, node_id, node_address)) => Some((
                    resource.clone(),
                    NodeInfo::new(resource_info, node_id, node_address, node),
                )),
                Err(error) => {
                    tracing::warn!("While executing plan: {}", error);
                    None
                }
            }
        })
        .collect();

    // Release locks
    drop(root_guard);
    drop(address_map_guard);

    // Set the `execute_status` of all nodes to `Scheduled` or `ScheduledPreviouslyFailed`
    // and send the resulting patch
    send_patches(
        patch_request_sender,
        node_infos
            .values_mut()
            .map(|node_info| node_info.set_execute_status_scheduled())
            .collect(),
    );

    // For each stage in plan...
    let stage_count = plan.stages.len();
    let mut dependencies_failed = false;
    for (stage_index, stage) in plan.stages.iter().enumerate() {
        // Before running the steps in this stage, check for any cancellation requests
        // and remove nodes from the plan.
        let mut was_cancelled = false;
        while let Ok(request) = cancel_request_receiver.try_recv() {
            if let Some(start) = request.start {
                // A start node was specified...
                if let Some(node_info) = node_infos.values_mut().find_map(|node_info| {
                    if *node_info.node_id == start {
                        Some(node_info)
                    } else {
                        None
                    }
                }) {
                    // TODO: This just sets status to cancelled and stops the loop. It is
                    // temporary
                    send_patch(
                        patch_request_sender,
                        node_info.set_execute_status_cancelled(),
                    );
                    was_cancelled = true;
                    break;
                }
            } else {
                // TODO No start node was specified so cancel them all.
                was_cancelled = true;
                break;
            }
        }
        if was_cancelled {
            tracing::debug!(
                "Received cancellation request before stage {}/{}",
                stage_index + 1,
                stage_count
            );
            break;
        }

        // Before running the steps in this stage, check that all their dependencies have succeeded
        // and stop if they have not. Collects to a `BTreeSet` to generate unique set (some steps in
        // the stage may have shared dependencies)
        dependencies_failed = stage
            .steps
            .iter()
            .flat_map(|step| step.resource_info.dependencies.iter().flatten())
            .collect::<BTreeSet<&Resource>>()
            .iter()
            .filter_map(|dependency| node_infos.get(dependency))
            .any(|node_info| {
                tracing::trace!(
                    "Status of dependency of stage {}/{} `{}`: {:?}",
                    stage_index + 1,
                    stage_count,
                    node_info.node_id,
                    node_info.get_execute_status()
                );
                matches!(
                    node_info.get_execute_status(),
                    None | Some(CodeExecutableExecuteStatus::Failed)
                        | Some(CodeExecutableExecuteStatus::Cancelled)
                )
            });
        if dependencies_failed {
            tracing::debug!(
                "Stopping before stage {}/{}: some dependencies failed or were cancelled",
                stage_index + 1,
                stage_count
            );
            break;
        }

        tracing::debug!("Starting stage {}/{}", stage_index + 1, stage_count);

        // Create a kernel task for each step in this stage
        let step_count = stage.steps.len();
        let mut patches = Vec::with_capacity(step_count);
        let mut futures = Vec::with_capacity(step_count);
        for (step_index, step) in stage.steps.iter().enumerate() {
            // Get the node
            let mut node_info = node_infos
                .get(&step.resource_info.resource)
                .cloned()
                .expect("Node info for resource should be available");

            // Set the `execute_status` of the node to `Running` or `RunningPreviouslyFailed`
            // and send the resulting patch
            patches.push(node_info.set_execute_status_running());

            // Create clones of variables needed to execute the task
            let kernel_space = kernel_space.clone();
            let kernel_selector = KernelSelector::new(step.kernel_name.clone(), None, None);
            let mut resource_info = step.resource_info.clone();
            let is_fork = step.is_fork;

            // Create a future for the task that will be spawned later
            let future = async move {
                tracing::debug!(
                    "Starting step {}/{} of stage {}/{}",
                    step_index + 1,
                    step_count,
                    stage_index + 1,
                    stage_count
                );

                // Create a mutable draft of the node and execute it in the kernel space
                let mut draft = node_info.node.clone();
                match draft
                    .execute(&kernel_space, &kernel_selector, &resource_info, is_fork)
                    .await
                {
                    Ok(_) => {
                        // Update the resource to indicate that the resource was executed
                        let execute_status = match &draft {
                            Node::CodeChunk(CodeChunk { execute_status, .. })
                            | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                                execute_status.clone()
                            }
                            _ => None,
                        };
                        resource_info.did_execute(execute_status);

                        // Generate a patch for the differences resulting from execution
                        let mut patch = diff(&node_info.node, &draft);
                        patch.address = Some(node_info.node_address.clone());
                        patch.target = Some(node_info.node_id.clone());

                        Ok((step_index, resource_info, draft, patch))
                    }
                    Err(error) => bail!(error),
                }
            };
            futures.push(future);
        }
        send_patches(patch_request_sender, patches);

        // Spawn all tasks in the stage and wait for each to finish, sending on the resultant `Patch`
        // for application and publishing (if it is not empty)
        // TODO: Replace `FuturesUnordered` with `TaskSet`. See https://news.ycombinator.com/item?id=29912386
        let mut results = futures
            .into_iter()
            .map(tokio::spawn)
            .collect::<FuturesUnordered<_>>();
        while let Some(result) = results.next().await {
            let (step_index, resource_info, node, patch) = result??;

            tracing::debug!(
                "Finished step {}/{} of stage {}/{}",
                step_index + 1,
                step_count,
                stage_index + 1,
                stage_count
            );

            // Send the patch reflecting the changed state of the executed node
            send_patch(patch_request_sender, patch);

            // Update the state of the node in this function's record of nodes
            node_infos
                .entry(resource_info.resource.clone())
                .and_modify(|node_info| node_info.node = node);

            // Send a patch request so that properties of other nodes such as `code_dependencies` and
            // `code_dependents` get updated with the new execution status the node that was executed in
            // this task. Previously we tried to take shortcut to this by just updating the graph and
            // calling `compile_no_walk` ut that proved unreliable so instead make a (debounced) request
            if let Err(..) = compile_request_sender
                .send(CompileRequest::new(false, None))
                .await
            {
                tracing::debug!("When sending compile request: receiver dropped");
            }
        }

        tracing::debug!("Finished stage {}/{}", stage_index + 1, stage_count);
    }

    // For nodes that were scheduled but never got to run because dependencies did not succeed,
    // restore their previous execution status
    if dependencies_failed {
        send_patches(
            patch_request_sender,
            node_infos
                .values_mut()
                .map(|node_info| node_info.restore_previous_execute_status())
                .collect(),
        );
    }

    Ok(())
}

/// A private internal struct to keep track of details of each node in the
/// execution plan during its execution
#[derive(Clone)]
struct NodeInfo {
    /// The associated [`ResourceInfo`]
    resource_info: ResourceInfo,

    /// The id of the node
    node_id: String,

    /// The address of the node
    node_address: Address,

    /// A copy of the node
    ///
    /// We take a copy of the node initially at the start of [`execute`] and
    /// then and send pathces for it to update stateus and execution results.
    node: Node,

    /// The execution state of the node prior to [`execute`]
    previous_execute_status: Option<CodeExecutableExecuteStatus>,
}

impl NodeInfo {
    fn new(
        resource_info: ResourceInfo,
        node_id: String,
        node_address: Address,
        node: Node,
    ) -> Self {
        let mut node_info = Self {
            resource_info,
            node_id,
            node_address,
            node,
            previous_execute_status: None,
        };
        node_info.previous_execute_status = node_info.get_execute_status();
        node_info
    }

    fn get_execute_status(&self) -> Option<CodeExecutableExecuteStatus> {
        match &self.node {
            Node::CodeChunk(CodeChunk { execute_status, .. })
            | Node::CodeExpression(CodeExpression { execute_status, .. }) => execute_status.clone(),
            // At present, assumes the execution of parameters always succeeds
            Node::Parameter(..) => Some(CodeExecutableExecuteStatus::Succeeded),
            _ => None,
        }
    }

    fn set_execute_status_scheduled(&mut self) -> Patch {
        mutate(
            &mut self.node,
            Some(self.node_id.to_string()),
            Some(self.node_address.clone()),
            &|node: &mut Node| match node {
                Node::CodeChunk(CodeChunk { execute_status, .. })
                | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                    *execute_status = Some(match execute_status {
                        Some(CodeExecutableExecuteStatus::Failed) => {
                            CodeExecutableExecuteStatus::ScheduledPreviouslyFailed
                        }
                        _ => CodeExecutableExecuteStatus::Scheduled,
                    });
                }
                _ => {}
            },
        )
    }

    fn set_execute_status_running(&mut self) -> Patch {
        mutate(
            &mut self.node,
            Some(self.node_id.to_string()),
            Some(self.node_address.clone()),
            &|node: &mut Node| match node {
                Node::CodeChunk(CodeChunk { execute_status, .. })
                | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                    *execute_status = Some(match execute_status {
                        Some(CodeExecutableExecuteStatus::Failed)
                        | Some(CodeExecutableExecuteStatus::ScheduledPreviouslyFailed) => {
                            CodeExecutableExecuteStatus::RunningPreviouslyFailed
                        }
                        _ => CodeExecutableExecuteStatus::Running,
                    });
                }
                _ => {}
            },
        )
    }

    fn set_execute_status_cancelled(&mut self) -> Patch {
        mutate(
            &mut self.node,
            Some(self.node_id.to_string()),
            Some(self.node_address.clone()),
            &|node: &mut Node| match node {
                Node::CodeChunk(CodeChunk { execute_status, .. })
                | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                    *execute_status = Some(CodeExecutableExecuteStatus::Cancelled);
                }
                _ => {}
            },
        )
    }

    fn restore_previous_execute_status(&mut self) -> Patch {
        mutate(
            &mut self.node,
            Some(self.node_id.to_string()),
            Some(self.node_address.clone()),
            &|node: &mut Node| match node {
                Node::CodeChunk(CodeChunk { execute_status, .. })
                | Node::CodeExpression(CodeExpression { execute_status, .. }) => {
                    if matches!(
                        execute_status,
                        Some(CodeExecutableExecuteStatus::Scheduled)
                            | Some(CodeExecutableExecuteStatus::ScheduledPreviouslyFailed)
                    ) {
                        *execute_status = self.previous_execute_status.clone();
                    }
                }
                _ => {}
            },
        )
    }
}
<|endoftext|>