<commit_before>use std::collections::{HashMap, HashSet};
use std::cmp::{Ordering, max};
use std::fs::{self, File};
use std::path::Path;
use std::io::Read;

use chrono::{Duration, NaiveDateTime};
use json::{self, JsonValue};

use util::start_idx;
use errors::*;

const WEEKS_IN_SUMMARY: usize = 12;

/// Contains the TestRun loaded from the file system.
/// Useful to avoid passing around each individual field, instead passing the
/// entire struct.
pub struct InputData {
    pub summary_rustc: Summary,
    pub summary_benchmarks: Summary,

    /// A set containing all crate names.
    pub crate_list: HashSet<String>,

    /// A set containing all phase names, across all crates.
    pub phase_list: HashSet<String>,

    /// TODO: Better docs. It's unknown how this is used on the client.  A set
    /// of test names. Test names are found from the file path, everything
    /// before the `--` is included.
    pub benchmarks: HashSet<String>,

    /// The last date that was seen while loading files. The DateTime variant is
    /// used here since the date may or may not contain a time. Since the
    /// timezone is not important, it isn't stored, hence the Naive variant.
    pub last_date: NaiveDateTime,

    /// Private due to access being exposed through by_kind method.
    /// Care must be taken to sort these after modifying them.
    data_rustc: Vec<TestRun>,
    data_benchmarks: Vec<TestRun>,
}

/// Allows storing `InputData` in Iron's persistent data structure.
impl ::iron::typemap::Key for InputData {
    type Value = InputData;
}

impl InputData {
    /// Initialize `InputData from the file system.
    pub fn from_fs() -> Result<InputData> {
        // TODO: Read this at runtime, not hardcoded.
        let repo_loc = Path::new("../data");

        let mut last_date = None;
        let mut crate_list = HashSet::new();
        let mut phase_list = HashSet::new();
        let mut skipped = 0;
        let mut c_benchmarks_add = 0;

        let mut data_rustc = Vec::new();
        let mut data_benchmarks = Vec::new();
        let mut benchmarks = HashSet::new();

        // Read all files from repo_loc/processed
        let mut file_count = 0;
        for entry in fs::read_dir(repo_loc.join("processed"))? {
            let entry = entry?;
            if entry.file_type()?.is_dir() {
                continue;
            }
            file_count += 1;

            let filename = entry.file_name();
            let filename = filename.to_str().unwrap();
            let mut file = File::open(entry.path())?;
            let mut file_contents = String::new();
            // Skip files whose size is 0.
            if file.read_to_string(&mut file_contents)? == 0 {
                skipped += 1;
                continue;
            }

            let contents = match json::parse(&file_contents) {
                Ok(json) => json,
                Err(err) => {
                    println!("Failed to parse JSON for {}: {:?}", filename, err);
                    skipped += 1;
                    continue;
                }
            };
            if contents["times"].is_empty() {
                skipped += 1;
                continue;
            }

            /// Parse date from JSON header in file contents.
            fn parse_from_header(date: &str) -> Result<NaiveDateTime> {
                // TODO: Determine correct format of header date and move into
                // variable/constant
                NaiveDateTime::parse_from_str(date, "%c %z")
                    .map_err(|err| err.into())
            }

            /// Parse date from filename.
            fn parse_from_filename(filename: &str) -> Result<NaiveDateTime> {
                let date_str = &filename[(filename.find("--").unwrap() + 2)..filename.find(".json").unwrap()];

                // Ignore seconds: they aren't key to the data processing, and
                // parsing them requires logic to replace them with 0 when
                // they are absent.
                NaiveDateTime::parse_from_str(date_str, "%Y-%m-%d-%H-%M")
                    .map_err(|err| err.into())
            }

            let header = &contents["header"];
            let date = header["date"].as_str().unwrap();
            let date = parse_from_header(date).or_else(|_| parse_from_filename(&filename))?;

            let test_name = &filename[0..filename.find("--").unwrap()];

            let times = &contents["times"];
            if test_name == "rustc" {
                data_rustc.push(TestRun::new(date, header, times));

                for timing in times.members() {
                    let crate_name = timing["crate"].as_str().unwrap().to_string();
                    crate_list.insert(crate_name);
                }
            } else {
                benchmarks.insert(test_name.to_string());
                let index = data_benchmarks.iter().position(|benchmark: &TestRun| benchmark.date == date);
                if let Some(index) = index {
                    c_benchmarks_add += 1;
                    let crate_name = times[0]["crate"].as_str().unwrap();
                    data_benchmarks[index].by_crate.insert(test_name.to_string(),
                        make_times(times).remove(crate_name).unwrap());
                } else {
                    data_benchmarks.push(TestRun::new(date, header, times));
                }
            }

            for timing in times.members() {
                for (phase, _) in timing["times"].entries() {
                    phase_list.insert(phase.to_string());
                }
            }

            if last_date.is_none() || last_date.as_ref().unwrap() < &date {
                last_date = Some(date);
            }
        }

        let last_date = last_date.expect("No dates found");

        data_rustc.sort();
        data_benchmarks.sort();

        // Post processing to generate the summary data.
        let summary_rustc = Summary::new(&data_rustc, last_date);
        let summary_benchmarks = Summary::new(&data_benchmarks, last_date);

        println!("{} total files", file_count);
        println!("{} skipped files", skipped);
        println!("{} bootstrap times", data_rustc.len());
        println!("{} benchmarks times", data_benchmarks.len());
        println!("{} benchmarks times (appended)", c_benchmarks_add);

        Ok(InputData {
            summary_rustc: summary_rustc,
            summary_benchmarks: summary_benchmarks,
            crate_list: crate_list,
            phase_list: phase_list,
            benchmarks: benchmarks,
            last_date: last_date,
            data_rustc: data_rustc,
            data_benchmarks: data_benchmarks,
        })
    }

    pub fn by_kind(&self, kind: Kind) -> &[TestRun] {
        match kind {
            Kind::Benchmarks => &self.data_benchmarks,
            Kind::Rustc => &self.data_rustc
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Kind {
    Benchmarks,
    Rustc,
}

impl Kind {
    pub fn from_str(kind: &str) -> Option<Kind> {
        match kind {
            "rustc" => Some(Kind::Rustc),
            "benchmarks" => Some(Kind::Benchmarks),
            _ => None
        }
    }
}

/// The data loaded for a single date, and all associated crates.
pub struct TestRun {
    pub date: NaiveDateTime,
    pub commit: String,

    /// Map of crate names to a map of phases to timings per phase.
    pub by_crate: HashMap<String, HashMap<String, Timing>>
}

impl PartialEq for TestRun {
    fn eq(&self, other: &TestRun) -> bool {
        self.commit == other.commit && self.date == other.date
    }
}

impl Eq for TestRun {}

impl PartialOrd for TestRun {
    fn partial_cmp(&self, other: &TestRun) -> Option<Ordering> {
        Some(<TestRun as Ord>::cmp(self, other))
    }
}

impl Ord for TestRun {
    fn cmp(&self, other: &TestRun) -> Ordering {
        self.date.cmp(&other.date)
    }
}

impl TestRun {
    fn new(date: NaiveDateTime, header: &JsonValue, times: &JsonValue) -> TestRun {
        TestRun {
            date: date,
            commit: header["commit"].as_str().unwrap().to_string(),
            by_crate: make_times(times)
        }
    }
}

/// Contains a single timing, associated with a phase (though the phase name
/// is not included in the timing).
#[derive(Debug, PartialEq)]
pub struct Timing {
    pub percent: f64,
    pub time: f64,
    pub rss: Option<u64>,
}

impl Timing {
    fn new() -> Timing {
        Timing {
            percent: 0.0,
            time: 0.0,
            rss: Some(0)
        }
    }
}

/// Run through the timings for a date and construct the `by_crate` field of TestRun.
fn make_times(timings: &JsonValue) -> HashMap<String, HashMap<String, Timing>> {
    let mut by_crate = HashMap::new();
    let mut totals = HashMap::new();

    for timing in timings.members() {
        let mut times = HashMap::new();
        for (phase_name, phase) in timing["times"].entries() {
            times.insert(phase_name.to_string(), Timing {
                percent: phase["percent"].as_f64().unwrap(),
                time: phase["time"].as_f64().unwrap(),
                rss: timing["rss"][phase_name].as_u64(),
            });
        }

        let mut mem_values = Vec::new();
        if timing["rss"].is_object() {
            for (_, value) in timing["rss"].entries() {
                mem_values.push(value.as_u64().unwrap());
            }
        }

        times.insert("total".into(), Timing {
            percent: 100.0,
            time: timing["total"].as_f64().unwrap(),
            rss: Some(mem_values.into_iter().max().unwrap_or(0))
        });

        for phase in times.keys() {
            let mut entry = totals.entry(phase.to_string()).or_insert_with(Timing::new);
            entry.time += times[phase].time;
            entry.rss = max(times[phase].rss, entry.rss);
        }

        by_crate.insert(timing["crate"].as_str().unwrap().to_string(), times);
    }

    by_crate.insert("total".into(), totals);
    // TODO: calculate percentages
    by_crate
}

#[derive(Debug)]
pub struct SummarizedWeek {
    pub date: NaiveDateTime,

    /// Maps crate names to a map of phases to each phase's percent change
    /// from the previous date.
    pub by_crate: HashMap<String, HashMap<String, f64>>
}

pub struct Summary {
    pub total: SummarizedWeek,
    pub summary: Vec<SummarizedWeek>,
}

impl Summary {
    // Compute summary data. For each week, we find the last 3 weeks, and use
    // the median timing as the basis of the current week's summary.
    fn new(data: &[TestRun], last_date: NaiveDateTime) -> Summary {
        // 13 week long mapping of crate names to by-phase medians.
        // We steal using summarized week type here even though we're not
        // storing the percent changes yet.
        let mut medians: Vec<SummarizedWeek> = Vec::new();

        // In order to compute summaries for the previous 12 weeks, we need 13
        // weeks, using the 0th week to compute difference with first week and
        // totals.
        for i in 0..(WEEKS_IN_SUMMARY + 1) {
            let date = last_date - Duration::weeks(i as i64);

            // For a given date we'll get the three most recent sets of TestRun
            // and take the the median for each value.
            let start_idx = start_idx(data, &date);
            assert!(start_idx >= 3, "Less than 3 days of data");
            let mut weeks = Vec::new();
            for idx in 0..3 {
                weeks.push(&data[start_idx - idx].by_crate);
            }

            medians.push(SummarizedWeek {
                date: date,
                by_crate: HashMap::new(),
            });
            for crate_name in weeks[0].keys() {
                if !weeks[1].contains_key(crate_name) || !weeks[2].contains_key(crate_name) {
                    continue;
                }

                let mut crate_medians = HashMap::new();
                for phase in weeks[0][crate_name].keys() {
                    if !weeks[1][crate_name].contains_key(phase) ||
                        !weeks[2][crate_name].contains_key(phase) {
                        continue;
                    }
                    // Find the median value.
                    let mut values = [
                        weeks[0][crate_name][phase].time,
                        weeks[1][crate_name][phase].time,
                        weeks[2][crate_name][phase].time
                    ];
                    values.sort_by(|a, b| a.partial_cmp(b).unwrap());

                    let median = values[1];

                    crate_medians.insert(phase.to_string(), median);
                }

                medians.last_mut().unwrap().by_crate.insert(crate_name.to_string(), crate_medians);
            }
        }

        /// Compute the percent change for a given number from the week N-1 to
        /// week N, where week N-1 is the previous and week N is the current
        /// week.
        fn get_percent_change(previous: f64, current: f64) -> f64 {
            ((current - previous) / previous) * 100.0
        }

        fn compare_weeks(week0: &SummarizedWeek, week1: &SummarizedWeek) -> SummarizedWeek {
            let mut current_week = HashMap::new();
            for crate_name in week0.by_crate.keys() {
                if !week1.by_crate.contains_key(crate_name) {
                    continue;
                }

                let mut current_crate = HashMap::new();
                for phase in week0.by_crate[crate_name].keys() {
                    if !week1.by_crate[crate_name].contains_key(phase) {
                        continue;
                    }

                    // TODO: Some form of warning is a good idea?
                    // If this is triggered for both the 0th overall week and
                    // the last week, then that phase won't be recorded in
                    // totals no matter what happens in between.
                    if week0.by_crate[crate_name][phase] > 0.0 &&
                        week1.by_crate[crate_name][phase] > 0.0 {
                        current_crate.insert(phase.to_string(),
                            get_percent_change(week0.by_crate[crate_name][phase],
                                week1.by_crate[crate_name][phase]));
                    }
                }
                current_week.insert(crate_name.to_string(), current_crate);
            }
            SummarizedWeek {
                date: week1.date,
                by_crate: current_week,
            }
        }

        // 12 week long mapping of crate names to by-phase percent changes with
        // the previous week.
        let mut weeks: Vec<SummarizedWeek> = Vec::new();
        for window in medians.windows(2) {
            // We want to compare week 1 to week 0; since we want the change from
            // week 1 to week 0.
            weeks.push(compare_weeks(&window[1], &window[0]));
        }

        assert_eq!(weeks.len(), 12);

        // See window comparison; compare from last week (which is the oldest)
        // to the first week.
        let totals = compare_weeks(medians.last().unwrap(), &medians[0]);

        for week in &mut weeks {
            for crate_name in totals.by_crate.keys() {
                if !week.by_crate.contains_key(crate_name) {
                    week.by_crate.insert(crate_name.to_string(), HashMap::new());
                }
            }
        }

        Summary {
            total: totals,
            summary: weeks
        }
    }
}
<commit_msg>Add logic to only insert pseudo-crate "total" when processing rustc timings.<commit_after>use std::collections::{HashMap, HashSet};
use std::cmp::{Ordering, max};
use std::fs::{self, File};
use std::path::Path;
use std::io::Read;

use chrono::{Duration, NaiveDateTime};
use json::{self, JsonValue};

use util::start_idx;
use errors::*;

const WEEKS_IN_SUMMARY: usize = 12;

/// Contains the TestRun loaded from the file system.
/// Useful to avoid passing around each individual field, instead passing the
/// entire struct.
pub struct InputData {
    pub summary_rustc: Summary,
    pub summary_benchmarks: Summary,

    /// A set containing all crate names.
    pub crate_list: HashSet<String>,

    /// A set containing all phase names, across all crates.
    pub phase_list: HashSet<String>,

    /// TODO: Better docs. It's unknown how this is used on the client.  A set
    /// of test names. Test names are found from the file path, everything
    /// before the `--` is included.
    pub benchmarks: HashSet<String>,

    /// The last date that was seen while loading files. The DateTime variant is
    /// used here since the date may or may not contain a time. Since the
    /// timezone is not important, it isn't stored, hence the Naive variant.
    pub last_date: NaiveDateTime,

    /// Private due to access being exposed through by_kind method.
    /// Care must be taken to sort these after modifying them.
    data_rustc: Vec<TestRun>,
    data_benchmarks: Vec<TestRun>,
}

/// Allows storing `InputData` in Iron's persistent data structure.
impl ::iron::typemap::Key for InputData {
    type Value = InputData;
}

impl InputData {
    /// Initialize `InputData from the file system.
    pub fn from_fs() -> Result<InputData> {
        // TODO: Read this at runtime, not hardcoded.
        let repo_loc = Path::new("../data");

        let mut last_date = None;
        let mut crate_list = HashSet::new();
        let mut phase_list = HashSet::new();
        let mut skipped = 0;
        let mut c_benchmarks_add = 0;

        let mut data_rustc = Vec::new();
        let mut data_benchmarks = Vec::new();
        let mut benchmarks = HashSet::new();

        // Read all files from repo_loc/processed
        let mut file_count = 0;
        for entry in fs::read_dir(repo_loc.join("processed"))? {
            let entry = entry?;
            if entry.file_type()?.is_dir() {
                continue;
            }
            file_count += 1;

            let filename = entry.file_name();
            let filename = filename.to_str().unwrap();
            let mut file = File::open(entry.path())?;
            let mut file_contents = String::new();
            // Skip files whose size is 0.
            if file.read_to_string(&mut file_contents)? == 0 {
                skipped += 1;
                continue;
            }

            let contents = match json::parse(&file_contents) {
                Ok(json) => json,
                Err(err) => {
                    println!("Failed to parse JSON for {}: {:?}", filename, err);
                    skipped += 1;
                    continue;
                }
            };
            if contents["times"].is_empty() {
                skipped += 1;
                continue;
            }

            /// Parse date from JSON header in file contents.
            fn parse_from_header(date: &str) -> Result<NaiveDateTime> {
                // TODO: Determine correct format of header date and move into
                // variable/constant
                NaiveDateTime::parse_from_str(date, "%c %z")
                    .map_err(|err| err.into())
            }

            /// Parse date from filename.
            fn parse_from_filename(filename: &str) -> Result<NaiveDateTime> {
                let date_str = &filename[(filename.find("--").unwrap() + 2)..filename.find(".json").unwrap()];

                // Ignore seconds: they aren't key to the data processing, and
                // parsing them requires logic to replace them with 0 when
                // they are absent.
                NaiveDateTime::parse_from_str(date_str, "%Y-%m-%d-%H-%M")
                    .map_err(|err| err.into())
            }

            let header = &contents["header"];
            let date = header["date"].as_str().unwrap();
            let date = parse_from_header(date).or_else(|_| parse_from_filename(&filename))?;

            let test_name = &filename[0..filename.find("--").unwrap()];

            let times = &contents["times"];
            if test_name == "rustc" {
                data_rustc.push(TestRun::new(date, header, times, true));

                for timing in times.members() {
                    let crate_name = timing["crate"].as_str().unwrap().to_string();
                    crate_list.insert(crate_name);
                }
            } else {
                benchmarks.insert(test_name.to_string());
                let index = data_benchmarks.iter().position(|benchmark: &TestRun| benchmark.date == date);
                if let Some(index) = index {
                    c_benchmarks_add += 1;
                    let crate_name = times[0]["crate"].as_str().unwrap();
                    data_benchmarks[index].by_crate.insert(test_name.to_string(),
                        make_times(times, false).remove(crate_name).unwrap());
                } else {
                    data_benchmarks.push(TestRun::new(date, header, times, false));
                }
            }

            for timing in times.members() {
                for (phase, _) in timing["times"].entries() {
                    phase_list.insert(phase.to_string());
                }
            }

            if last_date.is_none() || last_date.as_ref().unwrap() < &date {
                last_date = Some(date);
            }
        }

        let last_date = last_date.expect("No dates found");

        data_rustc.sort();
        data_benchmarks.sort();

        // Post processing to generate the summary data.
        let summary_rustc = Summary::new(&data_rustc, last_date);
        let summary_benchmarks = Summary::new(&data_benchmarks, last_date);

        println!("{} total files", file_count);
        println!("{} skipped files", skipped);
        println!("{} bootstrap times", data_rustc.len());
        println!("{} benchmarks times", data_benchmarks.len());
        println!("{} benchmarks times (appended)", c_benchmarks_add);

        Ok(InputData {
            summary_rustc: summary_rustc,
            summary_benchmarks: summary_benchmarks,
            crate_list: crate_list,
            phase_list: phase_list,
            benchmarks: benchmarks,
            last_date: last_date,
            data_rustc: data_rustc,
            data_benchmarks: data_benchmarks,
        })
    }

    pub fn by_kind(&self, kind: Kind) -> &[TestRun] {
        match kind {
            Kind::Benchmarks => &self.data_benchmarks,
            Kind::Rustc => &self.data_rustc
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Kind {
    Benchmarks,
    Rustc,
}

impl Kind {
    pub fn from_str(kind: &str) -> Option<Kind> {
        match kind {
            "rustc" => Some(Kind::Rustc),
            "benchmarks" => Some(Kind::Benchmarks),
            _ => None
        }
    }
}

/// The data loaded for a single date, and all associated crates.
pub struct TestRun {
    pub date: NaiveDateTime,
    pub commit: String,

    /// Map of crate names to a map of phases to timings per phase.
    pub by_crate: HashMap<String, HashMap<String, Timing>>
}

impl PartialEq for TestRun {
    fn eq(&self, other: &TestRun) -> bool {
        self.commit == other.commit && self.date == other.date
    }
}

impl Eq for TestRun {}

impl PartialOrd for TestRun {
    fn partial_cmp(&self, other: &TestRun) -> Option<Ordering> {
        Some(<TestRun as Ord>::cmp(self, other))
    }
}

impl Ord for TestRun {
    fn cmp(&self, other: &TestRun) -> Ordering {
        self.date.cmp(&other.date)
    }
}

impl TestRun {
    fn new(date: NaiveDateTime, header: &JsonValue, times: &JsonValue, is_rustc: bool) -> TestRun {
        TestRun {
            date: date,
            commit: header["commit"].as_str().unwrap().to_string(),
            by_crate: make_times(times, is_rustc)
        }
    }
}

/// Contains a single timing, associated with a phase (though the phase name
/// is not included in the timing).
#[derive(Debug, PartialEq)]
pub struct Timing {
    pub percent: f64,
    pub time: f64,
    pub rss: Option<u64>,
}

impl Timing {
    fn new() -> Timing {
        Timing {
            percent: 0.0,
            time: 0.0,
            rss: Some(0)
        }
    }
}

/// Run through the timings for a date and construct the `by_crate` field of TestRun.
fn make_times(timings: &JsonValue, is_rustc: bool) -> HashMap<String, HashMap<String, Timing>> {
    let mut by_crate = HashMap::new();
    let mut totals = HashMap::new();

    for timing in timings.members() {
        let mut times = HashMap::new();
        for (phase_name, phase) in timing["times"].entries() {
            times.insert(phase_name.to_string(), Timing {
                percent: phase["percent"].as_f64().unwrap(),
                time: phase["time"].as_f64().unwrap(),
                rss: timing["rss"][phase_name].as_u64(),
            });
        }

        let mut mem_values = Vec::new();
        if timing["rss"].is_object() {
            for (_, value) in timing["rss"].entries() {
                mem_values.push(value.as_u64().unwrap());
            }
        }

        times.insert("total".into(), Timing {
            percent: 100.0,
            time: timing["total"].as_f64().unwrap(),
            rss: Some(mem_values.into_iter().max().unwrap_or(0))
        });

        for phase in times.keys() {
            let mut entry = totals.entry(phase.to_string()).or_insert_with(Timing::new);
            entry.time += times[phase].time;
            entry.rss = max(times[phase].rss, entry.rss);
        }

        by_crate.insert(timing["crate"].as_str().unwrap().to_string(), times);
    }

    if is_rustc {
        by_crate.insert("total".into(), totals);
    }
    // TODO: calculate percentages
    by_crate
}

#[derive(Debug)]
pub struct SummarizedWeek {
    pub date: NaiveDateTime,

    /// Maps crate names to a map of phases to each phase's percent change
    /// from the previous date.
    pub by_crate: HashMap<String, HashMap<String, f64>>
}

pub struct Summary {
    pub total: SummarizedWeek,
    pub summary: Vec<SummarizedWeek>,
}

impl Summary {
    // Compute summary data. For each week, we find the last 3 weeks, and use
    // the median timing as the basis of the current week's summary.
    fn new(data: &[TestRun], last_date: NaiveDateTime) -> Summary {
        // 13 week long mapping of crate names to by-phase medians.
        // We steal using summarized week type here even though we're not
        // storing the percent changes yet.
        let mut medians: Vec<SummarizedWeek> = Vec::new();

        // In order to compute summaries for the previous 12 weeks, we need 13
        // weeks, using the 0th week to compute difference with first week and
        // totals.
        for i in 0..(WEEKS_IN_SUMMARY + 1) {
            let date = last_date - Duration::weeks(i as i64);

            // For a given date we'll get the three most recent sets of TestRun
            // and take the the median for each value.
            let start_idx = start_idx(data, &date);
            assert!(start_idx >= 3, "Less than 3 days of data");
            let mut weeks = Vec::new();
            for idx in 0..3 {
                weeks.push(&data[start_idx - idx].by_crate);
            }

            medians.push(SummarizedWeek {
                date: date,
                by_crate: HashMap::new(),
            });
            for crate_name in weeks[0].keys() {
                if !weeks[1].contains_key(crate_name) || !weeks[2].contains_key(crate_name) {
                    continue;
                }

                let mut crate_medians = HashMap::new();
                for phase in weeks[0][crate_name].keys() {
                    if !weeks[1][crate_name].contains_key(phase) ||
                        !weeks[2][crate_name].contains_key(phase) {
                        continue;
                    }
                    // Find the median value.
                    let mut values = [
                        weeks[0][crate_name][phase].time,
                        weeks[1][crate_name][phase].time,
                        weeks[2][crate_name][phase].time
                    ];
                    values.sort_by(|a, b| a.partial_cmp(b).unwrap());

                    let median = values[1];

                    crate_medians.insert(phase.to_string(), median);
                }

                medians.last_mut().unwrap().by_crate.insert(crate_name.to_string(), crate_medians);
            }
        }

        /// Compute the percent change for a given number from the week N-1 to
        /// week N, where week N-1 is the previous and week N is the current
        /// week.
        fn get_percent_change(previous: f64, current: f64) -> f64 {
            ((current - previous) / previous) * 100.0
        }

        fn compare_weeks(week0: &SummarizedWeek, week1: &SummarizedWeek) -> SummarizedWeek {
            let mut current_week = HashMap::new();
            for crate_name in week0.by_crate.keys() {
                if !week1.by_crate.contains_key(crate_name) {
                    continue;
                }

                let mut current_crate = HashMap::new();
                for phase in week0.by_crate[crate_name].keys() {
                    if !week1.by_crate[crate_name].contains_key(phase) {
                        continue;
                    }

                    // TODO: Some form of warning is a good idea?
                    // If this is triggered for both the 0th overall week and
                    // the last week, then that phase won't be recorded in
                    // totals no matter what happens in between.
                    if week0.by_crate[crate_name][phase] > 0.0 &&
                        week1.by_crate[crate_name][phase] > 0.0 {
                        current_crate.insert(phase.to_string(),
                            get_percent_change(week0.by_crate[crate_name][phase],
                                week1.by_crate[crate_name][phase]));
                    }
                }
                current_week.insert(crate_name.to_string(), current_crate);
            }
            SummarizedWeek {
                date: week1.date,
                by_crate: current_week,
            }
        }

        // 12 week long mapping of crate names to by-phase percent changes with
        // the previous week.
        let mut weeks: Vec<SummarizedWeek> = Vec::new();
        for window in medians.windows(2) {
            // We want to compare week 1 to week 0; since we want the change from
            // week 1 to week 0.
            weeks.push(compare_weeks(&window[1], &window[0]));
        }

        assert_eq!(weeks.len(), 12);

        // See window comparison; compare from last week (which is the oldest)
        // to the first week.
        let totals = compare_weeks(medians.last().unwrap(), &medians[0]);

        for week in &mut weeks {
            for crate_name in totals.by_crate.keys() {
                if !week.by_crate.contains_key(crate_name) {
                    week.by_crate.insert(crate_name.to_string(), HashMap::new());
                }
            }
        }

        Summary {
            total: totals,
            summary: weeks
        }
    }
}
<|endoftext|>


<commit_before>//! `StructureLink` data description.
use crate::data::RoomName;

use super::ActionLogTarget;

with_structure_fields_and_update_struct! {
    /// A tower structure - a structure that can be filled with energy, and then attack, heal and
    /// repair things within the same room using that energy.
    #[derive(Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTower {
        /// The user ID of the owner of this structure.
        pub user: String,
        /// Whether or not this structure is non-functional due to a degraded controller.
        #[serde(default, rename = "off")]
        pub disabled: bool,
        /// The current amount of energy held in this structure.
        pub energy: i32,
        /// The maximum amount of energy that can be held in this structure.
        pub energy_capacity: i32,
        /// Whether or not an attack on this structure will send an email to the owner automatically.
        pub notify_when_attacked: bool,
        /// A record of all actions this structure performed last tick.
        pub action_log: StructureTowerActions,
    }

    /// The update structure for a `StructureTower`.
    #[derive(Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTowerUpdate {
        - user: String,
        #[serde(rename = "off")]
        - disabled: bool,
        - energy: i32,
        - energy_capacity: i32,
        - notify_when_attacked: bool,
        - action_log: StructureTowerActions,
    }
}

with_update_struct! {
    /// A struct describing a tower's actions.
    #[derive(serde_derive::Deserialize, Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTowerActions {
        /// Where this tower attacked last tick.
        pub attack: Option<ActionLogTarget>,
        /// Where this tower healed last tick.
        pub heal: Option<ActionLogTarget>,
        /// Where this tower repaired last tick.
        pub repair: Option<ActionLogTarget>,
    }

    /// The update structure for a `StructureTowerActions`.
    #[derive(serde_derive::Deserialize, Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTowerActionsUpdate { ... }
}

#[cfg(test)]
mod test {
    use serde::Deserialize;
    use serde_json;

    use crate::data::RoomName;

    use super::{ActionLogTarget, StructureTower, StructureTowerActions};

    #[test]
    fn parse_tower_and_update() {
        let json = json!({
            "_id": "57f1cc9d27e2c0520e93ba95",
            "actionLog": {
                "attack": null,
                "heal": null,
                "repair": null
            },
            "energy": 920,
            "energyCapacity": 1000,
            "hits": 3000,
            "hitsMax": 3000,
            "notifyWhenAttacked": true,
            "room": "E17N55",
            "type": "tower",
            "user": "57874d42d0ae911e3bd15bbc",
            "x": 9,
            "y": 19
        });

        let mut obj = StructureTower::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureTower {
                room: RoomName::new("E17N55").unwrap(),
                x: 9,
                y: 19,
                id: "57f1cc9d27e2c0520e93ba95".to_owned(),
                energy: 920,
                energy_capacity: 1000,
                hits: 3000,
                hits_max: 3000,
                notify_when_attacked: true,
                disabled: false,
                action_log: StructureTowerActions {
                    attack: None,
                    heal: None,
                    repair: None,
                },
                user: "57874d42d0ae911e3bd15bbc".to_owned(),
            }
        );

        obj.update(
            serde_json::from_value(json!({
                "actionLog": {
                    "attack": {
                        "x": 10,
                        "y": 10
                    }
                },
                "energy": 820
            }))
            .unwrap(),
        );

        assert_eq!(
            obj.action_log,
            StructureTowerActions {
                attack: Some(ActionLogTarget { x: 10, y: 10 }),
                heal: None,
                repair: None,
            }
        );
    }
}
<commit_msg>Updated StructureTower<commit_after>//! `StructureLink` data description.
use super::super::resources::Store;
use crate::data::RoomName;

use super::ActionLogTarget;

with_structure_fields_and_update_struct! {
    /// A tower structure - a structure that can be filled with energy, and then attack, heal and
    /// repair things within the same room using that energy.
    #[derive(Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTower {
        /// The user ID of the owner of this structure.
        pub user: String,
        /// Whether or not this structure is non-functional due to a degraded controller.
        #[serde(default, rename = "off")]
        pub disabled: bool,
        /// The current amount of energy held in this structure.
        pub store: Store,
        /// The maximum amount of energy that can be held in this structure.
        #[serde(rename = "storeCapacityResource")]
        pub capacity_resource: Store,
        /// Whether or not an attack on this structure will send an email to the owner automatically.
        pub notify_when_attacked: bool,
        /// A record of all actions this structure performed last tick.
        pub action_log: StructureTowerActions,
    }

    /// The update structure for a `StructureTower`.
    #[derive(Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTowerUpdate {
        - user: String,
        #[serde(rename = "off")]
        - disabled: bool,
        - store: Store,
        #[serde(rename = "storeCapacityResource")]
        - capacity_resource: Store,
        - notify_when_attacked: bool,
        - action_log: StructureTowerActions,
    }
}

with_update_struct! {
    /// A struct describing a tower's actions.
    #[derive(serde_derive::Deserialize, Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTowerActions {
        /// Where this tower attacked last tick.
        pub attack: Option<ActionLogTarget>,
        /// Where this tower healed last tick.
        pub heal: Option<ActionLogTarget>,
        /// Where this tower repaired last tick.
        pub repair: Option<ActionLogTarget>,
    }

    /// The update structure for a `StructureTowerActions`.
    #[derive(serde_derive::Deserialize, Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct StructureTowerActionsUpdate { ... }
}

#[cfg(test)]
mod test {
    use serde::Deserialize;
    use serde_json;

    use crate::data::RoomName;

    use super::{ActionLogTarget, StructureTower, StructureTowerActions};

    #[test]
    fn parse_tower_and_update() {
        let json = json!({
          "_id": "5d334cefdbfe1b628e862a0d",
          "type": "tower",
          "x": 26,
          "y": 25,
          "room": "W44S12",
          "notifyWhenAttacked": true,
          "user": "5a8466038f866773f59fa6c8",
          "hits": 3000,
          "hitsMax": 3000,
          "actionLog": {
            "attack": {
              "x": 46,
              "y": 7
            },
            "heal": null,
            "repair": null
          },
          "store": {
            "energy": 570
          },
          "storeCapacityResource": {
            "energy": 1000
          }
        });

        let mut obj = StructureTower::deserialize(json).unwrap();

        assert_eq!(
            obj,
            StructureTower {
                room: RoomName::new("W44S12").unwrap(),
                x: 26,
                y: 25,
                id: "5d334cefdbfe1b628e862a0d".to_owned(),
                store: store! { Energy: 570 },
                capacity_resource: store! { Energy: 1000 },
                hits: 3000,
                hits_max: 3000,
                notify_when_attacked: true,
                disabled: false,
                action_log: StructureTowerActions {
                    attack: Some(ActionLogTarget { x: 46, y: 7 }),
                    heal: None,
                    repair: None,
                },
                user: "5a8466038f866773f59fa6c8".to_owned(),
            }
        );

        obj.update(
            serde_json::from_value(json!({
              "actionLog": {
                "attack": {
                  "x": 45,
                  "y": 6
                }
              },
              "store": {
                "energy": 560
              }
            }))
            .unwrap(),
        );

        assert_eq!(
            obj.action_log,
            StructureTowerActions {
                attack: Some(ActionLogTarget { x: 45, y: 6 }),
                heal: None,
                repair: None,
            }
        );

        assert_eq!(obj.store, store! { Energy: 560 });

        obj.update(
            serde_json::from_value(json!({
              "store": {
                "energy": 550
              }
            }))
            .unwrap(),
        );

        assert_eq!(
            obj.action_log,
            StructureTowerActions {
                attack: Some(ActionLogTarget { x: 45, y: 6 }),
                heal: None,
                repair: None,
            }
        );

        assert_eq!(obj.store, store! { Energy: 550 });

        obj.update(
            serde_json::from_value(json!({
              "actionLog": {
                "attack": null
              }
            }))
            .unwrap(),
        );

        assert_eq!(
            obj.action_log,
            StructureTowerActions {
                attack: None,
                heal: None,
                repair: None,
            },
        );
    }
}
<|endoftext|>


<commit_before>use ::std::sync::mpsc::Sender;

use crate::response_helper::*;
use ::pueue::message::*;
use ::pueue::state::SharedState;
use ::pueue::task::{Task, TaskStatus};

static SENDER_ERR: &str = "Failed to send message to task handler thread";

pub fn handle_message(message: Message, sender: &Sender<Message>, state: &SharedState) -> Message {
    match message {
        Message::Add(message) => add_task(message, sender, state),
        Message::Remove(message) => remove(message, state),
        Message::Switch(message) => switch(message, state),
        Message::Stash(message) => stash(message, state),
        Message::Enqueue(message) => enqueue(message, state),

        Message::Start(message) => start(message, sender, state),
        Message::Restart(message) => restart(message, sender, state),
        Message::Pause(message) => pause(message, sender, state),
        Message::Kill(message) => kill(message, sender, state),

        Message::Send(message) => send(message, sender, state),
        Message::EditRequest(message) => edit_request(message, state),
        Message::Edit(message) => edit(message, state),

        Message::Clean => clean(state),
        Message::Reset => reset(sender),
        Message::SimpleStatus => get_simple_status(state),
        Message::Status => get_status(state),
        Message::Parallel(amount) => set_parallel_tasks(amount, sender),
        _ => create_failure_message("Not implemented yet"),
    }
}

/// Queues a new task to the state.
/// If the start_immediately flag is set, send a StartMessage to the task handler
fn add_task(message: AddMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    let starting_status = if message.create_stashed || message.enqueue_at.is_some() {
        TaskStatus::Stashed
    } else {
        TaskStatus::Queued
    };

    let task = Task::new(
        message.command,
        message.path,
        starting_status,
        message.enqueue_at,
    );
    let task_id: usize;
    {
        let mut state = state.lock().unwrap();
        task_id = state.add_task(task);
    }
    if message.start_immediately {
        let start_message = StartMessage {
            task_ids: vec![task_id],
        };
        sender
            .send(Message::Start(start_message))
            .expect(SENDER_ERR);
    }

    let message = if let Some(enqueue_at) = message.enqueue_at {
        format!(
            "New task added. It will be enqueued at {}",
            enqueue_at.format("%Y-%m-%d %H:%M:%S")
        )
    } else {
        String::from("New task added.")
    };

    create_success_message(message)
}

/// Remove tasks from the queue
fn remove(message: RemoveMessage, state: &SharedState) -> Message {
    let mut state = state.lock().unwrap();
    let statuses = vec![TaskStatus::Running, TaskStatus::Paused];
    let (matching, mismatching) = state.tasks_not_in_statuses(statuses, Some(message.task_ids));

    for task_id in &matching {
        state.tasks.remove(task_id);
    }

    let message = "Tasks removed from list";
    let message = compile_task_response(message, matching, mismatching);
    create_success_message(message)
}

/// Switch the position of two tasks in the upcoming queue
fn switch(message: SwitchMessage, state: &SharedState) -> Message {
    let task_ids = vec![message.task_id_1, message.task_id_2];
    let statuses = vec![TaskStatus::Queued, TaskStatus::Stashed];
    let mut state = state.lock().unwrap();
    let (_, mismatching) = state.tasks_in_statuses(statuses, Some(task_ids.clone().to_vec()));
    if !mismatching.is_empty() {
        return create_failure_message("Tasks have to be either queued or stashed.");
    }

    // Get the tasks. Expect them to be there, since we found no mismatch
    let mut first_task = state.tasks.remove(&task_ids[0]).unwrap();
    let mut second_task = state.tasks.remove(&task_ids[1]).unwrap();

    // Switch task ids
    let temp_id = first_task.id;
    first_task.id = second_task.id;
    second_task.id = temp_id;

    // Put tasks back in again
    state.tasks.insert(first_task.id, first_task);
    state.tasks.insert(second_task.id, second_task);

    create_success_message("Tasks have been switched")
}

/// Stash specific queued tasks.
/// They won't be executed until enqueued again or explicitely started
fn stash(message: StashMessage, state: &SharedState) -> Message {
    let (matching, mismatching) = {
        let mut state = state.lock().unwrap();
        let (matching, mismatching) =
            state.tasks_in_statuses(vec![TaskStatus::Queued], Some(message.task_ids));

        for task_id in &matching {
            state.change_status(*task_id, TaskStatus::Stashed);
        }

        (matching, mismatching)
    };

    let message = "Tasks are stashed";
    let response = compile_task_response(message, matching, mismatching);
    return create_success_message(response);
}

/// Enqueue specific stashed tasks.
/// They will be normally handled afterwards.
fn enqueue(message: EnqueueMessage, state: &SharedState) -> Message {
    let (matching, mismatching) = {
        let mut state = state.lock().unwrap();
        let (matching, mismatching) = state.tasks_in_statuses(
            vec![TaskStatus::Stashed, TaskStatus::Locked],
            Some(message.task_ids),
        );

        for task_id in &matching {
            state.set_enqueue_at(*task_id, message.enqueue_at);
            state.change_status(*task_id, TaskStatus::Queued);
        }

        (matching, mismatching)
    };

    let message = if let Some(enqueue_at) = message.enqueue_at {
        format!(
            "Tasks will be enqueued at {}",
            enqueue_at.format("%Y-%m-%d %H:%M:%S")
        )
    } else {
        String::from("Tasks are enqueued")
    };

    let response = compile_task_response(message.as_str(), matching, mismatching);
    return create_success_message(response);
}

/// Forward the start message to the task handler and respond to the client
fn start(message: StartMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    sender
        .send(Message::Start(message.clone()))
        .expect(SENDER_ERR);
    if !message.task_ids.is_empty() {
        let response = task_response_helper(
            "Tasks are being started",
            message.task_ids,
            vec![TaskStatus::Paused, TaskStatus::Queued, TaskStatus::Stashed],
            state,
        );
        return create_success_message(response);
    }

    return create_success_message("Daemon and all tasks are being resumed.");
}

/// Create and enqueue tasks from already finished tasks
/// The user can specify to immediately start the newly created tasks.
fn restart(message: RestartMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    let response: String;
    let new_ids = {
        let mut state = state.lock().unwrap();
        let statuses = vec![TaskStatus::Done, TaskStatus::Failed, TaskStatus::Killed];
        let (matching, mismatching) = state.tasks_in_statuses(statuses, Some(message.task_ids));

        let mut new_ids = Vec::new();
        for task_id in &matching {
            let task = state.tasks.get(task_id).unwrap();
            let new_task = Task::from_task(task);
            new_ids.push(state.add_task(new_task));
        }

        // Already create the response string in here. Otherwise we would
        // need to get matching/mismatching out of this scope
        let message = "Restarted tasks";
        response = compile_task_response(message, matching, mismatching);

        new_ids
    };

    // If the restarted tasks should be started immediately, send a message
    // with the new task ids to the task handler.
    if message.start_immediately {
        let start_message = StartMessage { task_ids: new_ids };
        sender
            .send(Message::Start(start_message))
            .expect(SENDER_ERR);
    }

    return create_success_message(response);
}

/// Forward the pause message to the task handler and respond to the client
fn pause(message: PauseMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    sender
        .send(Message::Pause(message.clone()))
        .expect(SENDER_ERR);
    if !message.task_ids.is_empty() {
        let response = task_response_helper(
            "Tasks are being paused",
            message.task_ids,
            vec![TaskStatus::Running],
            state,
        );
        return create_success_message(response);
    }

    return create_success_message("Daemon and all tasks are being paused.");
}

/// Forward the kill message to the task handler and respond to the client
fn kill(message: KillMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    sender
        .send(Message::Kill(message.clone()))
        .expect(SENDER_ERR);

    if !message.task_ids.is_empty() {
        let response = task_response_helper(
            "Tasks are being killed",
            message.task_ids,
            vec![TaskStatus::Running, TaskStatus::Paused],
            state,
        );
        return create_success_message(response);
    }

    return create_success_message("All tasks are being killed.");
}

// Send some user defined input to a process
// The message will be forwarded to the task handler.
// In here we only do some error handling.
fn send(message: SendMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    // Check whether the task exists and is running, abort if that's not the case
    {
        let state = state.lock().unwrap();
        match state.tasks.get(&message.task_id) {
            Some(task) => {
                if task.status != TaskStatus::Running {
                    return create_failure_message("You can only send input to a running task");
                }
            }
            None => return create_failure_message("No task with this id."),
        }
    }

    // Check whether the task exists and is running, abort if that's not the case
    sender.send(Message::Send(message)).expect(SENDER_ERR);

    return create_success_message("Message is being send to the process.");
}

// If a user wants to edit a message, we need to send him the current command
// and lock the task to prevent execution, before the user has finished editing the command
fn edit_request(message: EditRequestMessage, state: &SharedState) -> Message {
    // Check whether the task exists and is queued/stashed, abort if that's not the case
    let mut state = state.lock().unwrap();
    match state.tasks.get_mut(&message.task_id) {
        Some(task) => {
            if !task.is_queued() {
                return create_failure_message("You can only edit a queued/stashed task");
            }
            task.prev_status = task.status.clone();
            task.status = TaskStatus::Locked;

            let message = EditResponseMessage {
                task_id: task.id,
                command: task.command.clone(),
            };
            return Message::EditResponse(message);
        }
        None => return create_failure_message("No task with this id."),
    }
}

// Handle the actual updated command
fn edit(message: EditMessage, state: &SharedState) -> Message {
    // Check whether the task exists and is queued/stashed, abort if that's not the case
    let mut state = state.lock().unwrap();
    match state.tasks.get_mut(&message.task_id) {
        Some(task) => {
            if !(task.status == TaskStatus::Locked) {
                return create_failure_message("Task is no longer locked.");
            }

            task.status = task.prev_status.clone();
            task.command = message.command.clone();
            state.save();

            return create_success_message("Command has been updated");
        }
        None => {
            return create_failure_message(format!(
                "Task to edit has gone away: {}",
                message.task_id
            ))
        }
    }
}

/// Remove all failed or done tasks from the state
fn clean(state: &SharedState) -> Message {
    let mut state = state.lock().unwrap();
    state.clean();

    return create_success_message("All finished tasks have been removed");
}

// Forward the reset request to the task handler
// The handler then kills all children and clears the task queue
fn reset(sender: &Sender<Message>) -> Message {
    sender.send(Message::Reset).expect(SENDER_ERR);
    return create_success_message("Everything is being reset right now.");
}

/// Return the full current state to the client
fn get_status(state: &SharedState) -> Message {
    let state = state.lock().unwrap();
    Message::StatusResponse(state.clone())
}

/// Return the current state without any stdou/stderr to the client
fn get_simple_status(state: &SharedState) -> Message {
    let mut state = { state.lock().unwrap().clone() };
    for (_, task) in state.tasks.iter_mut() {
        task.stdout = None;
        task.stderr = None;
    }
    Message::StatusResponse(state)
}

fn set_parallel_tasks(amount: usize, sender: &Sender<Message>) -> Message {
    sender.send(Message::Parallel(amount)).expect(SENDER_ERR);
    return create_success_message("Parallel tasks setting adjusted");
}
<commit_msg>Allow locked tasks to be stashed<commit_after>use ::std::sync::mpsc::Sender;

use crate::response_helper::*;
use ::pueue::message::*;
use ::pueue::state::SharedState;
use ::pueue::task::{Task, TaskStatus};

static SENDER_ERR: &str = "Failed to send message to task handler thread";

pub fn handle_message(message: Message, sender: &Sender<Message>, state: &SharedState) -> Message {
    match message {
        Message::Add(message) => add_task(message, sender, state),
        Message::Remove(message) => remove(message, state),
        Message::Switch(message) => switch(message, state),
        Message::Stash(message) => stash(message, state),
        Message::Enqueue(message) => enqueue(message, state),

        Message::Start(message) => start(message, sender, state),
        Message::Restart(message) => restart(message, sender, state),
        Message::Pause(message) => pause(message, sender, state),
        Message::Kill(message) => kill(message, sender, state),

        Message::Send(message) => send(message, sender, state),
        Message::EditRequest(message) => edit_request(message, state),
        Message::Edit(message) => edit(message, state),

        Message::Clean => clean(state),
        Message::Reset => reset(sender),
        Message::SimpleStatus => get_simple_status(state),
        Message::Status => get_status(state),
        Message::Parallel(amount) => set_parallel_tasks(amount, sender),
        _ => create_failure_message("Not implemented yet"),
    }
}

/// Queues a new task to the state.
/// If the start_immediately flag is set, send a StartMessage to the task handler
fn add_task(message: AddMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    let starting_status = if message.create_stashed || message.enqueue_at.is_some() {
        TaskStatus::Stashed
    } else {
        TaskStatus::Queued
    };

    let task = Task::new(
        message.command,
        message.path,
        starting_status,
        message.enqueue_at,
    );
    let task_id: usize;
    {
        let mut state = state.lock().unwrap();
        task_id = state.add_task(task);
    }
    if message.start_immediately {
        let start_message = StartMessage {
            task_ids: vec![task_id],
        };
        sender
            .send(Message::Start(start_message))
            .expect(SENDER_ERR);
    }

    let message = if let Some(enqueue_at) = message.enqueue_at {
        format!(
            "New task added. It will be enqueued at {}",
            enqueue_at.format("%Y-%m-%d %H:%M:%S")
        )
    } else {
        String::from("New task added.")
    };

    create_success_message(message)
}

/// Remove tasks from the queue
fn remove(message: RemoveMessage, state: &SharedState) -> Message {
    let mut state = state.lock().unwrap();
    let statuses = vec![TaskStatus::Running, TaskStatus::Paused];
    let (matching, mismatching) = state.tasks_not_in_statuses(statuses, Some(message.task_ids));

    for task_id in &matching {
        state.tasks.remove(task_id);
    }

    let message = "Tasks removed from list";
    let message = compile_task_response(message, matching, mismatching);
    create_success_message(message)
}

/// Switch the position of two tasks in the upcoming queue
fn switch(message: SwitchMessage, state: &SharedState) -> Message {
    let task_ids = vec![message.task_id_1, message.task_id_2];
    let statuses = vec![TaskStatus::Queued, TaskStatus::Stashed];
    let mut state = state.lock().unwrap();
    let (_, mismatching) = state.tasks_in_statuses(statuses, Some(task_ids.clone().to_vec()));
    if !mismatching.is_empty() {
        return create_failure_message("Tasks have to be either queued or stashed.");
    }

    // Get the tasks. Expect them to be there, since we found no mismatch
    let mut first_task = state.tasks.remove(&task_ids[0]).unwrap();
    let mut second_task = state.tasks.remove(&task_ids[1]).unwrap();

    // Switch task ids
    let temp_id = first_task.id;
    first_task.id = second_task.id;
    second_task.id = temp_id;

    // Put tasks back in again
    state.tasks.insert(first_task.id, first_task);
    state.tasks.insert(second_task.id, second_task);

    create_success_message("Tasks have been switched")
}

/// Stash specific queued tasks.
/// They won't be executed until enqueued again or explicitely started
fn stash(message: StashMessage, state: &SharedState) -> Message {
    let (matching, mismatching) = {
        let mut state = state.lock().unwrap();
        let (matching, mismatching) =
            state.tasks_in_statuses(vec![TaskStatus::Queued, TaskStatus::Locked], Some(message.task_ids));

        for task_id in &matching {
            state.change_status(*task_id, TaskStatus::Stashed);
        }

        (matching, mismatching)
    };

    let message = "Tasks are stashed";
    let response = compile_task_response(message, matching, mismatching);
    return create_success_message(response);
}

/// Enqueue specific stashed tasks.
/// They will be normally handled afterwards.
fn enqueue(message: EnqueueMessage, state: &SharedState) -> Message {
    let (matching, mismatching) = {
        let mut state = state.lock().unwrap();
        let (matching, mismatching) = state.tasks_in_statuses(
            vec![TaskStatus::Stashed, TaskStatus::Locked],
            Some(message.task_ids),
        );

        for task_id in &matching {
            state.set_enqueue_at(*task_id, message.enqueue_at);
            state.change_status(*task_id, TaskStatus::Queued);
        }

        (matching, mismatching)
    };

    let message = if let Some(enqueue_at) = message.enqueue_at {
        format!(
            "Tasks will be enqueued at {}",
            enqueue_at.format("%Y-%m-%d %H:%M:%S")
        )
    } else {
        String::from("Tasks are enqueued")
    };

    let response = compile_task_response(message.as_str(), matching, mismatching);
    return create_success_message(response);
}

/// Forward the start message to the task handler and respond to the client
fn start(message: StartMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    sender
        .send(Message::Start(message.clone()))
        .expect(SENDER_ERR);
    if !message.task_ids.is_empty() {
        let response = task_response_helper(
            "Tasks are being started",
            message.task_ids,
            vec![TaskStatus::Paused, TaskStatus::Queued, TaskStatus::Stashed],
            state,
        );
        return create_success_message(response);
    }

    return create_success_message("Daemon and all tasks are being resumed.");
}

/// Create and enqueue tasks from already finished tasks
/// The user can specify to immediately start the newly created tasks.
fn restart(message: RestartMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    let response: String;
    let new_ids = {
        let mut state = state.lock().unwrap();
        let statuses = vec![TaskStatus::Done, TaskStatus::Failed, TaskStatus::Killed];
        let (matching, mismatching) = state.tasks_in_statuses(statuses, Some(message.task_ids));

        let mut new_ids = Vec::new();
        for task_id in &matching {
            let task = state.tasks.get(task_id).unwrap();
            let new_task = Task::from_task(task);
            new_ids.push(state.add_task(new_task));
        }

        // Already create the response string in here. Otherwise we would
        // need to get matching/mismatching out of this scope
        let message = "Restarted tasks";
        response = compile_task_response(message, matching, mismatching);

        new_ids
    };

    // If the restarted tasks should be started immediately, send a message
    // with the new task ids to the task handler.
    if message.start_immediately {
        let start_message = StartMessage { task_ids: new_ids };
        sender
            .send(Message::Start(start_message))
            .expect(SENDER_ERR);
    }

    return create_success_message(response);
}

/// Forward the pause message to the task handler and respond to the client
fn pause(message: PauseMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    sender
        .send(Message::Pause(message.clone()))
        .expect(SENDER_ERR);
    if !message.task_ids.is_empty() {
        let response = task_response_helper(
            "Tasks are being paused",
            message.task_ids,
            vec![TaskStatus::Running],
            state,
        );
        return create_success_message(response);
    }

    return create_success_message("Daemon and all tasks are being paused.");
}

/// Forward the kill message to the task handler and respond to the client
fn kill(message: KillMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    sender
        .send(Message::Kill(message.clone()))
        .expect(SENDER_ERR);

    if !message.task_ids.is_empty() {
        let response = task_response_helper(
            "Tasks are being killed",
            message.task_ids,
            vec![TaskStatus::Running, TaskStatus::Paused],
            state,
        );
        return create_success_message(response);
    }

    return create_success_message("All tasks are being killed.");
}

// Send some user defined input to a process
// The message will be forwarded to the task handler.
// In here we only do some error handling.
fn send(message: SendMessage, sender: &Sender<Message>, state: &SharedState) -> Message {
    // Check whether the task exists and is running, abort if that's not the case
    {
        let state = state.lock().unwrap();
        match state.tasks.get(&message.task_id) {
            Some(task) => {
                if task.status != TaskStatus::Running {
                    return create_failure_message("You can only send input to a running task");
                }
            }
            None => return create_failure_message("No task with this id."),
        }
    }

    // Check whether the task exists and is running, abort if that's not the case
    sender.send(Message::Send(message)).expect(SENDER_ERR);

    return create_success_message("Message is being send to the process.");
}

// If a user wants to edit a message, we need to send him the current command
// and lock the task to prevent execution, before the user has finished editing the command
fn edit_request(message: EditRequestMessage, state: &SharedState) -> Message {
    // Check whether the task exists and is queued/stashed, abort if that's not the case
    let mut state = state.lock().unwrap();
    match state.tasks.get_mut(&message.task_id) {
        Some(task) => {
            if !task.is_queued() {
                return create_failure_message("You can only edit a queued/stashed task");
            }
            task.prev_status = task.status.clone();
            task.status = TaskStatus::Locked;

            let message = EditResponseMessage {
                task_id: task.id,
                command: task.command.clone(),
            };
            return Message::EditResponse(message);
        }
        None => return create_failure_message("No task with this id."),
    }
}

// Handle the actual updated command
fn edit(message: EditMessage, state: &SharedState) -> Message {
    // Check whether the task exists and is queued/stashed, abort if that's not the case
    let mut state = state.lock().unwrap();
    match state.tasks.get_mut(&message.task_id) {
        Some(task) => {
            if !(task.status == TaskStatus::Locked) {
                return create_failure_message("Task is no longer locked.");
            }

            task.status = task.prev_status.clone();
            task.command = message.command.clone();
            state.save();

            return create_success_message("Command has been updated");
        }
        None => {
            return create_failure_message(format!(
                "Task to edit has gone away: {}",
                message.task_id
            ))
        }
    }
}

/// Remove all failed or done tasks from the state
fn clean(state: &SharedState) -> Message {
    let mut state = state.lock().unwrap();
    state.clean();

    return create_success_message("All finished tasks have been removed");
}

// Forward the reset request to the task handler
// The handler then kills all children and clears the task queue
fn reset(sender: &Sender<Message>) -> Message {
    sender.send(Message::Reset).expect(SENDER_ERR);
    return create_success_message("Everything is being reset right now.");
}

/// Return the full current state to the client
fn get_status(state: &SharedState) -> Message {
    let state = state.lock().unwrap();
    Message::StatusResponse(state.clone())
}

/// Return the current state without any stdou/stderr to the client
fn get_simple_status(state: &SharedState) -> Message {
    let mut state = { state.lock().unwrap().clone() };
    for (_, task) in state.tasks.iter_mut() {
        task.stdout = None;
        task.stderr = None;
    }
    Message::StatusResponse(state)
}

fn set_parallel_tasks(amount: usize, sender: &Sender<Message>) -> Message {
    sender.send(Message::Parallel(amount)).expect(SENDER_ERR);
    return create_success_message("Parallel tasks setting adjusted");
}
<|endoftext|>


<commit_before>use definitions::ImageBuilder;
use definitions::error::*;
use definitions::typedef::*;
use mnemonic::Mnemonic;
use regex::Regex;
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::BufReader;
use std::io::prelude::*;
use std::path::{Path, PathBuf};

lazy_static! {
    static ref LABELED_MNEMONIC_RE: Regex = Regex::new(r"^\.(.+?) +?(.+)$").unwrap();
    static ref LABEL_RE: Regex = Regex::new(r"^\.([^\s]+)$").unwrap();
    static ref INCLUDE_RE: Regex = Regex::new(r"^include! +([^\s]+)$").unwrap();
}

#[derive(Default)]
struct BASMCompiler {
    label_addr_map: HashMap<String, Address>,
    mnemonics: Vec<Mnemonic>,
    builder: ImageBuilder,
    last_base_dir: Option<PathBuf>,
    deep: usize,
}

impl BASMCompiler {
    fn add_label(&mut self, label: String) {
        self.label_addr_map.entry(label).or_insert(self.mnemonics.len());
    }

    fn base_path(&mut self, orig_path: &Path) -> Result<PathBuf> {
        let file_name = if let Some(ref file_name) = orig_path.file_name() {
            if let Some(..) = self.last_base_dir {
                PathBuf::from(orig_path)
            } else {
                PathBuf::from(file_name.clone())
            }
        } else {
            bail!("unable to obtain file name");
        };

        let absolute_path = if let Some(ref base_dir) = self.last_base_dir {
            base_dir.clone()
        } else {
            let current_dir = if orig_path.is_relative() {
                env::current_dir().chain_err(|| "unable to get current directory")?
            } else {
                if let Some(ref parent) = orig_path.parent() {
                    parent.to_path_buf().clone()
                } else {
                    bail!("unable to get parent directory")
                }
            };

            let base_file = current_dir.join(orig_path);
            let base_dir = if let Some(ref dir) = base_file.parent() {
                dir.clone()
            } else {
                bail!("unable to get parent directory")
            };

            self.last_base_dir = Some(base_dir.to_path_buf().clone());
            base_dir.to_path_buf()
        };

        Ok(absolute_path.join(file_name))
    }

    fn compile_mnemonics(&mut self, file_name: String) -> Result<()> {
        let orig_path = Path::new(&file_name);

        let padding = (0..self.deep).map(|_| " ").collect::<String>();
        println!("BASM\t{}{}", padding, file_name);

        let path = self.base_path(orig_path)?;

        let file = File::open(path).chain_err(|| "unable to open file")?;

        let reader = BufReader::new(file);

        for line in reader.lines() {
            let line = line.chain_err(|| "unable to read line")?;

            if line.is_empty() {
                continue;
            }

            let mut r_split = line.rsplitn(2, ';');

            if r_split.clone().count() > 1 {
                r_split.next();
            }

            if let Some(first_half) = r_split.next() {
                let first_half = first_half.trim();

                let first_half: String = if LABELED_MNEMONIC_RE.is_match(first_half) {
                    let captures = if let Some(captures) = LABELED_MNEMONIC_RE
                           .captures_iter(first_half)
                           .next() {
                        captures
                    } else {
                        bail!("no label capture found")
                    };

                    let label = captures[1].trim();

                    self.add_label(label.to_owned());

                    captures[2].trim().to_owned()
                } else if LABEL_RE.is_match(first_half) {
                    let captures =
                        if let Some(captures) = LABEL_RE.captures_iter(first_half).next() {
                            captures
                        } else {
                            bail!("no label capture found")
                        };

                    let label = captures[1].trim();

                    self.add_label(label.to_owned());
                    continue;
                } else if INCLUDE_RE.is_match(first_half) {
                    let captures =
                        if let Some(captures) = INCLUDE_RE.captures_iter(first_half).next() {
                            captures
                        } else {
                            bail!("no include capture found")
                        };

                    self.deep += 2;
                    self.compile_mnemonics(captures[1].trim().to_owned() + ".basm")?;
                    self.deep -= 2;

                    continue;
                } else {
                    first_half.to_owned()
                };

                let mut first_half_split = first_half.splitn(2, ' ');

                if let Some(opcode) = first_half_split.next() {
                    let opcode = opcode.trim().to_lowercase();

                    let args: Vec<String> = if let Some(args) = first_half_split.next() {
                        args.split(',').map(|arg| arg.trim().to_owned()).collect()
                    } else {
                        Vec::new()
                    };

                    self.mnemonics.push(text_to_mnemonic(opcode, args)?);
                } else {
                    bail!("opcode expected. Found {:?}", first_half_split);
                }
            } else {
                bail!("instruction expected. Found {:?}", line.to_string());
            }
        }

        Ok(())
    }

    fn compile_mnemonic(&mut self, mnemonic: Mnemonic) -> Result<()> {
        match mnemonic {
            Mnemonic::Add(dest, src) => self.builder.add(dest, src),
            Mnemonic::Sub(dest, src) => self.builder.sub(dest, src),
            Mnemonic::Div(dest, src) => self.builder.div(dest, src),
            Mnemonic::Mul(dest, src) => self.builder.mul(dest, src),
            Mnemonic::Rem(dest, src) => self.builder.rem(dest, src),

            Mnemonic::Cmp(target_a, target_b) => self.builder.cmp(target_a, target_b),
            Mnemonic::Jmp(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp(addr);
            }
            Mnemonic::JmpLt(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_lt(addr);
            }
            Mnemonic::JmpGt(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_gt(addr);
            }
            Mnemonic::JmpEq(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_eq(addr);
            }
            Mnemonic::JmpLtEq(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_lt_eq(addr);
            }
            Mnemonic::JmpGtEq(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_gt_eq(addr);
            }

            Mnemonic::Cast(target, type_t) => self.builder.cast(target, type_t),

            Mnemonic::Push(target, value) => self.builder.push(target, value),
            Mnemonic::Mov(dest, src) => self.builder.mov(dest, src),
            Mnemonic::Swp(target_a, target_b) => self.builder.swp(target_a, target_b),
            Mnemonic::Dup(target) => self.builder.dup(target),

            Mnemonic::Call(label) => {
                let addr = self.lookup(&label)?;
                self.builder.call(addr);
            }
            Mnemonic::Ret => self.builder.ret(),

            Mnemonic::Halt => self.builder.halt(),
            Mnemonic::Pause => self.builder.pause(),
            Mnemonic::Nop => self.builder.nop(),
            Mnemonic::Int(int) => self.builder.int(int),
        }

        Ok(())
    }

    fn lookup(&mut self, input: &String) -> Result<Address> {
        if let Some(addr) = self.label_addr_map.get(input) {
            Ok(*addr)
        } else {
            bail!("label {:?} not found", input);
        }
    }

    pub fn compile(&mut self, file_name: String) -> Result<ImageData> {
        self.compile_mnemonics(file_name)?;

        for mnemonic in self.mnemonics.clone() {
            self.compile_mnemonic(mnemonic)?;
        }

        Ok(self.builder.clone().gen())
    }
}

fn text_to_mnemonic(opcode: String, args: Vec<String>) -> Result<Mnemonic> {
    match opcode.as_str() {
        "add" => Ok(Mnemonic::Add(args[0].parse()?, args[1].parse()?)),
        "sub" => Ok(Mnemonic::Sub(args[0].parse()?, args[1].parse()?)),
        "div" => Ok(Mnemonic::Div(args[0].parse()?, args[1].parse()?)),
        "mul" => Ok(Mnemonic::Mul(args[0].parse()?, args[1].parse()?)),
        "rem" => Ok(Mnemonic::Rem(args[0].parse()?, args[1].parse()?)),

        "cmp" => Ok(Mnemonic::Cmp(args[0].parse()?, args[1].parse()?)),
        "jmp" => Ok(Mnemonic::Jmp(args[0].parse()?)),
        "jmplt" => Ok(Mnemonic::JmpLt(args[0].parse()?)),
        "jmpgt" => Ok(Mnemonic::JmpGt(args[0].parse()?)),
        "jmpeq" => Ok(Mnemonic::JmpEq(args[0].parse()?)),
        "jmplteq" => Ok(Mnemonic::JmpLtEq(args[0].parse()?)),
        "jmpgteq" => Ok(Mnemonic::JmpGtEq(args[0].parse()?)),

        "cast" => Ok(Mnemonic::Cast(args[0].parse()?, args[1].parse()?)),
        "push" => Ok(Mnemonic::Push(args[0].parse()?, args[1].parse()?)),
        "mov" => Ok(Mnemonic::Mov(args[0].parse()?, args[1].parse()?)),
        "swp" => Ok(Mnemonic::Swp(args[0].parse()?, args[1].parse()?)),
        "dup" => Ok(Mnemonic::Dup(args[0].parse()?)),

        "call" => Ok(Mnemonic::Call(args[0].parse()?)),
        "ret" => Ok(Mnemonic::Ret),

        "halt" => Ok(Mnemonic::Halt),
        "pause" => Ok(Mnemonic::Pause),
        "nop" => Ok(Mnemonic::Nop),
        "int" => Ok(Mnemonic::Int(args[0].parse()?)),
        _ => bail!("unkwnown opcode {:?}", opcode),
    }
}

pub fn compile(file_name: String) -> Result<ImageData> {
    BASMCompiler::default().compile(file_name)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn labeled_mnemonic_regex() {
        let input = ".start mov $st, $vi(19)";

        if !LABELED_MNEMONIC_RE.is_match(input) {
            panic!("input doesn't match a labeled mnemonic");
        } else {
            let captures = LABELED_MNEMONIC_RE.captures_iter(input).next().unwrap();

            assert_eq!(captures[1].trim(), "start");
            assert_eq!(captures[2].trim(), "mov $st, $vi(19)");
        }
    }

    #[test]
    fn label_regex() {
        let input = ".start";

        if !LABEL_RE.is_match(input) {
            panic!("input doesn't match a label");
        } else {
            let captures = LABEL_RE.captures_iter(input).next().unwrap();

            assert_eq!(captures[1].trim(), "start");
        }
    }

    #[test]
    fn include_regex() {
        let input = "include! std";

        if !INCLUDE_RE.is_match(input) {
            panic!("input doesn't match an include statement");
        } else {
            let captures = INCLUDE_RE.captures_iter(input).next().unwrap();

            assert_eq!(captures[1].trim(), "std");
        }
    }
}
<commit_msg>Adjusted BASM output<commit_after>use definitions::ImageBuilder;
use definitions::error::*;
use definitions::typedef::*;
use mnemonic::Mnemonic;
use regex::Regex;
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::BufReader;
use std::io::prelude::*;
use std::path::{Path, PathBuf};

lazy_static! {
    static ref LABELED_MNEMONIC_RE: Regex = Regex::new(r"^\.(.+?) +?(.+)$").unwrap();
    static ref LABEL_RE: Regex = Regex::new(r"^\.([^\s]+)$").unwrap();
    static ref INCLUDE_RE: Regex = Regex::new(r"^include! +([^\s]+)$").unwrap();
}

#[derive(Default)]
struct BASMCompiler {
    label_addr_map: HashMap<String, Address>,
    mnemonics: Vec<Mnemonic>,
    builder: ImageBuilder,
    last_base_dir: Option<PathBuf>,
    deep: usize,
}

impl BASMCompiler {
    fn add_label(&mut self, label: String) {
        self.label_addr_map.entry(label).or_insert(self.mnemonics.len());
    }

    fn base_path(&mut self, orig_path: &Path) -> Result<PathBuf> {
        let file_name = if let Some(ref file_name) = orig_path.file_name() {
            if let Some(..) = self.last_base_dir {
                PathBuf::from(orig_path)
            } else {
                PathBuf::from(file_name.clone())
            }
        } else {
            bail!("unable to obtain file name");
        };

        let absolute_path = if let Some(ref base_dir) = self.last_base_dir {
            base_dir.clone()
        } else {
            let current_dir = if orig_path.is_relative() {
                env::current_dir().chain_err(|| "unable to get current directory")?
            } else {
                if let Some(ref parent) = orig_path.parent() {
                    parent.to_path_buf().clone()
                } else {
                    bail!("unable to get parent directory")
                }
            };

            let base_file = current_dir.join(orig_path);
            let base_dir = if let Some(ref dir) = base_file.parent() {
                dir.clone()
            } else {
                bail!("unable to get parent directory")
            };

            self.last_base_dir = Some(base_dir.to_path_buf().clone());
            base_dir.to_path_buf()
        };

        Ok(absolute_path.join(file_name))
    }

    fn compile_mnemonics(&mut self, file_name: String) -> Result<()> {
        let orig_path = Path::new(&file_name);

        let padding = (0..self.deep).map(|_| "  ").collect::<String>();
        println!("BASM    {}{}", padding, file_name);

        let path = self.base_path(orig_path)?;

        let file = File::open(path).chain_err(|| "unable to open file")?;

        let reader = BufReader::new(file);

        for line in reader.lines() {
            let line = line.chain_err(|| "unable to read line")?;

            if line.is_empty() {
                continue;
            }

            let mut r_split = line.rsplitn(2, ';');

            if r_split.clone().count() > 1 {
                r_split.next();
            }

            if let Some(first_half) = r_split.next() {
                let first_half = first_half.trim();

                let first_half: String = if LABELED_MNEMONIC_RE.is_match(first_half) {
                    let captures = if let Some(captures) = LABELED_MNEMONIC_RE
                           .captures_iter(first_half)
                           .next() {
                        captures
                    } else {
                        bail!("no label capture found")
                    };

                    let label = captures[1].trim();

                    self.add_label(label.to_owned());

                    captures[2].trim().to_owned()
                } else if LABEL_RE.is_match(first_half) {
                    let captures =
                        if let Some(captures) = LABEL_RE.captures_iter(first_half).next() {
                            captures
                        } else {
                            bail!("no label capture found")
                        };

                    let label = captures[1].trim();

                    self.add_label(label.to_owned());
                    continue;
                } else if INCLUDE_RE.is_match(first_half) {
                    let captures =
                        if let Some(captures) = INCLUDE_RE.captures_iter(first_half).next() {
                            captures
                        } else {
                            bail!("no include capture found")
                        };

                    self.deep += 1;
                    self.compile_mnemonics(captures[1].trim().to_owned() + ".basm")?;
                    self.deep -= 1;

                    continue;
                } else {
                    first_half.to_owned()
                };

                let mut first_half_split = first_half.splitn(2, ' ');

                if let Some(opcode) = first_half_split.next() {
                    let opcode = opcode.trim().to_lowercase();

                    let args: Vec<String> = if let Some(args) = first_half_split.next() {
                        args.split(',').map(|arg| arg.trim().to_owned()).collect()
                    } else {
                        Vec::new()
                    };

                    self.mnemonics.push(text_to_mnemonic(opcode, args)?);
                } else {
                    bail!("opcode expected. Found {:?}", first_half_split);
                }
            } else {
                bail!("instruction expected. Found {:?}", line.to_string());
            }
        }

        Ok(())
    }

    fn compile_mnemonic(&mut self, mnemonic: Mnemonic) -> Result<()> {
        match mnemonic {
            Mnemonic::Add(dest, src) => self.builder.add(dest, src),
            Mnemonic::Sub(dest, src) => self.builder.sub(dest, src),
            Mnemonic::Div(dest, src) => self.builder.div(dest, src),
            Mnemonic::Mul(dest, src) => self.builder.mul(dest, src),
            Mnemonic::Rem(dest, src) => self.builder.rem(dest, src),

            Mnemonic::Cmp(target_a, target_b) => self.builder.cmp(target_a, target_b),
            Mnemonic::Jmp(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp(addr);
            }
            Mnemonic::JmpLt(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_lt(addr);
            }
            Mnemonic::JmpGt(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_gt(addr);
            }
            Mnemonic::JmpEq(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_eq(addr);
            }
            Mnemonic::JmpLtEq(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_lt_eq(addr);
            }
            Mnemonic::JmpGtEq(label) => {
                let addr = self.lookup(&label)?;
                self.builder.jmp_gt_eq(addr);
            }

            Mnemonic::Cast(target, type_t) => self.builder.cast(target, type_t),

            Mnemonic::Push(target, value) => self.builder.push(target, value),
            Mnemonic::Mov(dest, src) => self.builder.mov(dest, src),
            Mnemonic::Swp(target_a, target_b) => self.builder.swp(target_a, target_b),
            Mnemonic::Dup(target) => self.builder.dup(target),

            Mnemonic::Call(label) => {
                let addr = self.lookup(&label)?;
                self.builder.call(addr);
            }
            Mnemonic::Ret => self.builder.ret(),

            Mnemonic::Halt => self.builder.halt(),
            Mnemonic::Pause => self.builder.pause(),
            Mnemonic::Nop => self.builder.nop(),
            Mnemonic::Int(int) => self.builder.int(int),
        }

        Ok(())
    }

    fn lookup(&mut self, input: &String) -> Result<Address> {
        if let Some(addr) = self.label_addr_map.get(input) {
            Ok(*addr)
        } else {
            bail!("label {:?} not found", input);
        }
    }

    pub fn compile(&mut self, file_name: String) -> Result<ImageData> {
        self.compile_mnemonics(file_name)?;

        for mnemonic in self.mnemonics.clone() {
            self.compile_mnemonic(mnemonic)?;
        }

        Ok(self.builder.clone().gen())
    }
}

fn text_to_mnemonic(opcode: String, args: Vec<String>) -> Result<Mnemonic> {
    match opcode.as_str() {
        "add" => Ok(Mnemonic::Add(args[0].parse()?, args[1].parse()?)),
        "sub" => Ok(Mnemonic::Sub(args[0].parse()?, args[1].parse()?)),
        "div" => Ok(Mnemonic::Div(args[0].parse()?, args[1].parse()?)),
        "mul" => Ok(Mnemonic::Mul(args[0].parse()?, args[1].parse()?)),
        "rem" => Ok(Mnemonic::Rem(args[0].parse()?, args[1].parse()?)),

        "cmp" => Ok(Mnemonic::Cmp(args[0].parse()?, args[1].parse()?)),
        "jmp" => Ok(Mnemonic::Jmp(args[0].parse()?)),
        "jmplt" => Ok(Mnemonic::JmpLt(args[0].parse()?)),
        "jmpgt" => Ok(Mnemonic::JmpGt(args[0].parse()?)),
        "jmpeq" => Ok(Mnemonic::JmpEq(args[0].parse()?)),
        "jmplteq" => Ok(Mnemonic::JmpLtEq(args[0].parse()?)),
        "jmpgteq" => Ok(Mnemonic::JmpGtEq(args[0].parse()?)),

        "cast" => Ok(Mnemonic::Cast(args[0].parse()?, args[1].parse()?)),
        "push" => Ok(Mnemonic::Push(args[0].parse()?, args[1].parse()?)),
        "mov" => Ok(Mnemonic::Mov(args[0].parse()?, args[1].parse()?)),
        "swp" => Ok(Mnemonic::Swp(args[0].parse()?, args[1].parse()?)),
        "dup" => Ok(Mnemonic::Dup(args[0].parse()?)),

        "call" => Ok(Mnemonic::Call(args[0].parse()?)),
        "ret" => Ok(Mnemonic::Ret),

        "halt" => Ok(Mnemonic::Halt),
        "pause" => Ok(Mnemonic::Pause),
        "nop" => Ok(Mnemonic::Nop),
        "int" => Ok(Mnemonic::Int(args[0].parse()?)),
        _ => bail!("unkwnown opcode {:?}", opcode),
    }
}

pub fn compile(file_name: String) -> Result<ImageData> {
    BASMCompiler::default().compile(file_name)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn labeled_mnemonic_regex() {
        let input = ".start mov $st, $vi(19)";

        if !LABELED_MNEMONIC_RE.is_match(input) {
            panic!("input doesn't match a labeled mnemonic");
        } else {
            let captures = LABELED_MNEMONIC_RE.captures_iter(input).next().unwrap();

            assert_eq!(captures[1].trim(), "start");
            assert_eq!(captures[2].trim(), "mov $st, $vi(19)");
        }
    }

    #[test]
    fn label_regex() {
        let input = ".start";

        if !LABEL_RE.is_match(input) {
            panic!("input doesn't match a label");
        } else {
            let captures = LABEL_RE.captures_iter(input).next().unwrap();

            assert_eq!(captures[1].trim(), "start");
        }
    }

    #[test]
    fn include_regex() {
        let input = "include! std";

        if !INCLUDE_RE.is_match(input) {
            panic!("input doesn't match an include statement");
        } else {
            let captures = INCLUDE_RE.captures_iter(input).next().unwrap();

            assert_eq!(captures[1].trim(), "std");
        }
    }
}
<|endoftext|>