<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! This module contains types and functions to support formatting specific line ranges.

use std::collections::HashMap;
use std::path::PathBuf;
use std::rc::Rc;
use std::{cmp, fmt, iter, str};

use serde::de::{Deserialize, Deserializer};
use serde_json as json;

use syntax::codemap::{self, FileMap};

/// A range of lines in a file, inclusive of both ends.
pub struct LineRange {
    pub file: Rc<FileMap>,
    pub lo: usize,
    pub hi: usize,
}

/// Defines the name of an input - either a file or stdin.
#[derive(Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub enum FileName {
    Real(PathBuf),
    Stdin,
}

impl From<codemap::FileName> for FileName {
    fn from(name: codemap::FileName) -> FileName {
        match name {
            codemap::FileName::Real(p) => FileName::Real(p),
            codemap::FileName::Custom(ref f) if f == "stdin" => FileName::Stdin,
            _ => unreachable!(),
        }
    }
}

impl fmt::Display for FileName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            FileName::Real(p) => write!(f, "{}", p.to_str().unwrap()),
            FileName::Stdin => write!(f, "stdin"),
        }
    }
}

impl LineRange {
    pub fn file_name(&self) -> FileName {
        self.file.name.clone().into()
    }
}

/// A range that is inclusive of both ends.
#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord, Deserialize)]
pub struct Range {
    lo: usize,
    hi: usize,
}

impl<'a> From<&'a LineRange> for Range {
    fn from(range: &'a LineRange) -> Range {
        Range::new(range.lo, range.hi)
    }
}

impl Range {
    pub fn new(lo: usize, hi: usize) -> Range {
        Range { lo, hi }
    }

    fn is_empty(self) -> bool {
        self.lo > self.hi
    }

    #[allow(dead_code)]
    fn contains(self, other: Range) -> bool {
        if other.is_empty() {
            true
        } else {
            !self.is_empty() && self.lo <= other.lo && self.hi >= other.hi
        }
    }

    fn intersects(self, other: Range) -> bool {
        if self.is_empty() || other.is_empty() {
            false
        } else {
            (self.lo <= other.hi && other.hi <= self.hi)
                || (other.lo <= self.hi && self.hi <= other.hi)
        }
    }

    fn adjacent_to(self, other: Range) -> bool {
        if self.is_empty() || other.is_empty() {
            false
        } else {
            self.hi + 1 == other.lo || other.hi + 1 == self.lo
        }
    }

    /// Returns a new `Range` with lines from `self` and `other` if they were adjacent or
    /// intersect; returns `None` otherwise.
    fn merge(self, other: Range) -> Option<Range> {
        if self.adjacent_to(other) || self.intersects(other) {
            Some(Range::new(
                cmp::min(self.lo, other.lo),
                cmp::max(self.hi, other.hi),
            ))
        } else {
            None
        }
    }
}

/// A set of lines in files.
///
/// It is represented as a multimap keyed on file names, with values a collection of
/// non-overlapping ranges sorted by their start point. An inner `None` is interpreted to mean all
/// lines in all files.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct FileLines(Option<HashMap<FileName, Vec<Range>>>);

/// Normalizes the ranges so that the invariants for `FileLines` hold: ranges are non-overlapping,
/// and ordered by their start point.
fn normalize_ranges(ranges: &mut HashMap<FileName, Vec<Range>>) {
    for ranges in ranges.values_mut() {
        ranges.sort();
        let mut result = vec![];
        {
            let mut iter = ranges.into_iter().peekable();
            while let Some(next) = iter.next() {
                let mut next = *next;
                while let Some(&&mut peek) = iter.peek() {
                    if let Some(merged) = next.merge(peek) {
                        iter.next().unwrap();
                        next = merged;
                    } else {
                        break;
                    }
                }
                result.push(next)
            }
        }
        *ranges = result;
    }
}

impl FileLines {
    /// Creates a `FileLines` that contains all lines in all files.
    pub(crate) fn all() -> FileLines {
        FileLines(None)
    }

    /// Returns true if this `FileLines` contains all lines in all files.
    pub(crate) fn is_all(&self) -> bool {
        self.0.is_none()
    }

    pub fn from_ranges(mut ranges: HashMap<FileName, Vec<Range>>) -> FileLines {
        normalize_ranges(&mut ranges);
        FileLines(Some(ranges))
    }

    /// Returns an iterator over the files contained in `self`.
    pub fn files(&self) -> Files {
        Files(self.0.as_ref().map(|m| m.keys()))
    }

    /// Returns true if `self` includes all lines in all files. Otherwise runs `f` on all ranges in
    /// the designated file (if any) and returns true if `f` ever does.
    fn file_range_matches<F>(&self, file_name: &FileName, f: F) -> bool
    where
        F: FnMut(&Range) -> bool,
    {
        let map = match self.0 {
            // `None` means "all lines in all files".
            None => return true,
            Some(ref map) => map,
        };

        match canonicalize_path_string(file_name).and_then(|file| map.get(&file)) {
            Some(ranges) => ranges.iter().any(f),
            None => false,
        }
    }

    /// Returns true if `range` is fully contained in `self`.
    #[allow(dead_code)]
    pub(crate) fn contains(&self, range: &LineRange) -> bool {
        self.file_range_matches(&range.file_name(), |r| r.contains(Range::from(range)))
    }

    /// Returns true if any lines in `range` are in `self`.
    pub(crate) fn intersects(&self, range: &LineRange) -> bool {
        self.file_range_matches(&range.file_name(), |r| r.intersects(Range::from(range)))
    }

    /// Returns true if `line` from `file_name` is in `self`.
    pub(crate) fn contains_line(&self, file_name: &FileName, line: usize) -> bool {
        self.file_range_matches(file_name, |r| r.lo <= line && r.hi >= line)
    }

    /// Returns true if all the lines between `lo` and `hi` from `file_name` are in `self`.
    pub(crate) fn contains_range(&self, file_name: &FileName, lo: usize, hi: usize) -> bool {
        self.file_range_matches(file_name, |r| r.contains(Range::new(lo, hi)))
    }
}

/// `FileLines` files iterator.
pub struct Files<'a>(Option<::std::collections::hash_map::Keys<'a, FileName, Vec<Range>>>);

impl<'a> iter::Iterator for Files<'a> {
    type Item = &'a FileName;

    fn next(&mut self) -> Option<&'a FileName> {
        self.0.as_mut().and_then(Iterator::next)
    }
}

fn canonicalize_path_string(file: &FileName) -> Option<FileName> {
    match *file {
        FileName::Real(ref path) => path.canonicalize().ok().map(FileName::Real),
        _ => Some(file.clone()),
    }
}

// This impl is needed for `Config::override_value` to work for use in tests.
impl str::FromStr for FileLines {
    type Err = String;

    fn from_str(s: &str) -> Result<FileLines, String> {
        let v: Vec<JsonSpan> = json::from_str(s).map_err(|e| e.to_string())?;
        let mut m = HashMap::new();
        for js in v {
            let (s, r) = JsonSpan::into_tuple(js)?;
            m.entry(s).or_insert_with(|| vec![]).push(r);
        }
        Ok(FileLines::from_ranges(m))
    }
}

// For JSON decoding.
#[derive(Clone, Debug, Deserialize)]
struct JsonSpan {
    #[serde(deserialize_with = "deserialize_filename")]
    file: FileName,
    range: (usize, usize),
}

fn deserialize_filename<'de, D: Deserializer<'de>>(d: D) -> Result<FileName, D::Error> {
    let s = String::deserialize(d)?;
    if s == "stdin" {
        Ok(FileName::Stdin)
    } else {
        Ok(FileName::Real(s.into()))
    }
}

impl JsonSpan {
    fn into_tuple(self) -> Result<(FileName, Range), String> {
        let (lo, hi) = self.range;
        let canonical = canonicalize_path_string(&self.file)
            .ok_or_else(|| format!("Can't canonicalize {}", &self.file))?;
        Ok((canonical, Range::new(lo, hi)))
    }
}

// This impl is needed for inclusion in the `Config` struct. We don't have a toml representation
// for `FileLines`, so it will just panic instead.
impl<'de> ::serde::de::Deserialize<'de> for FileLines {
    fn deserialize<D>(_: D) -> Result<Self, D::Error>
    where
        D: ::serde::de::Deserializer<'de>,
    {
        panic!(
            "FileLines cannot be deserialized from a project rustfmt.toml file: please \
             specify it via the `--file-lines` option instead"
        );
    }
}

// We also want to avoid attempting to serialize a FileLines to toml. The
// `Config` struct should ensure this impl is never reached.
impl ::serde::ser::Serialize for FileLines {
    fn serialize<S>(&self, _: S) -> Result<S::Ok, S::Error>
    where
        S: ::serde::ser::Serializer,
    {
        unreachable!("FileLines cannot be serialized. This is a rustfmt bug.");
    }
}

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

    #[test]
    fn test_range_intersects() {
        assert!(Range::new(1, 2).intersects(Range::new(1, 1)));
        assert!(Range::new(1, 2).intersects(Range::new(2, 2)));
        assert!(!Range::new(1, 2).intersects(Range::new(0, 0)));
        assert!(!Range::new(1, 2).intersects(Range::new(3, 10)));
        assert!(!Range::new(1, 3).intersects(Range::new(5, 5)));
    }

    #[test]
    fn test_range_adjacent_to() {
        assert!(!Range::new(1, 2).adjacent_to(Range::new(1, 1)));
        assert!(!Range::new(1, 2).adjacent_to(Range::new(2, 2)));
        assert!(Range::new(1, 2).adjacent_to(Range::new(0, 0)));
        assert!(Range::new(1, 2).adjacent_to(Range::new(3, 10)));
        assert!(!Range::new(1, 3).adjacent_to(Range::new(5, 5)));
    }

    #[test]
    fn test_range_contains() {
        assert!(Range::new(1, 2).contains(Range::new(1, 1)));
        assert!(Range::new(1, 2).contains(Range::new(2, 2)));
        assert!(!Range::new(1, 2).contains(Range::new(0, 0)));
        assert!(!Range::new(1, 2).contains(Range::new(3, 10)));
    }

    #[test]
    fn test_range_merge() {
        assert_eq!(None, Range::new(1, 3).merge(Range::new(5, 5)));
        assert_eq!(None, Range::new(4, 7).merge(Range::new(0, 1)));
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 5).merge(Range::new(4, 7))
        );
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 5).merge(Range::new(5, 7))
        );
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 5).merge(Range::new(6, 7))
        );
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 7).merge(Range::new(4, 5))
        );
    }
}
<commit_msg>Expose FileLines JSON representation<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! This module contains types and functions to support formatting specific line ranges.

use std::collections::HashMap;
use std::path::PathBuf;
use std::rc::Rc;
use std::{cmp, fmt, iter, str};

use serde::ser::{self, Serialize, Serializer};
use serde::de::{Deserialize, Deserializer};
use serde_json as json;

use syntax::codemap::{self, FileMap};

/// A range of lines in a file, inclusive of both ends.
pub struct LineRange {
    pub file: Rc<FileMap>,
    pub lo: usize,
    pub hi: usize,
}

/// Defines the name of an input - either a file or stdin.
#[derive(Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub enum FileName {
    Real(PathBuf),
    Stdin,
}

impl From<codemap::FileName> for FileName {
    fn from(name: codemap::FileName) -> FileName {
        match name {
            codemap::FileName::Real(p) => FileName::Real(p),
            codemap::FileName::Custom(ref f) if f == "stdin" => FileName::Stdin,
            _ => unreachable!(),
        }
    }
}

impl fmt::Display for FileName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            FileName::Real(p) => write!(f, "{}", p.to_str().unwrap()),
            FileName::Stdin => write!(f, "stdin"),
        }
    }
}

impl<'de> Deserialize<'de> for FileName {
    fn deserialize<D>(deserializer: D) -> Result<FileName, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        if s == "stdin" {
            Ok(FileName::Stdin)
        } else {
            Ok(FileName::Real(s.into()))
        }
    }
}

impl Serialize for FileName {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let s = match self {
            FileName::Stdin => Ok("stdin"),
            FileName::Real(path) => path.to_str().ok_or_else(||
                ser::Error::custom("path can't be serialized as UTF-8 string"))
        };

        s.and_then(|s| serializer.serialize_str(s))
    }
}

impl LineRange {
    pub fn file_name(&self) -> FileName {
        self.file.name.clone().into()
    }
}

/// A range that is inclusive of both ends.
#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord, Deserialize)]
pub struct Range {
    lo: usize,
    hi: usize,
}

impl<'a> From<&'a LineRange> for Range {
    fn from(range: &'a LineRange) -> Range {
        Range::new(range.lo, range.hi)
    }
}

impl Range {
    pub fn new(lo: usize, hi: usize) -> Range {
        Range { lo, hi }
    }

    fn is_empty(self) -> bool {
        self.lo > self.hi
    }

    #[allow(dead_code)]
    fn contains(self, other: Range) -> bool {
        if other.is_empty() {
            true
        } else {
            !self.is_empty() && self.lo <= other.lo && self.hi >= other.hi
        }
    }

    fn intersects(self, other: Range) -> bool {
        if self.is_empty() || other.is_empty() {
            false
        } else {
            (self.lo <= other.hi && other.hi <= self.hi)
                || (other.lo <= self.hi && self.hi <= other.hi)
        }
    }

    fn adjacent_to(self, other: Range) -> bool {
        if self.is_empty() || other.is_empty() {
            false
        } else {
            self.hi + 1 == other.lo || other.hi + 1 == self.lo
        }
    }

    /// Returns a new `Range` with lines from `self` and `other` if they were adjacent or
    /// intersect; returns `None` otherwise.
    fn merge(self, other: Range) -> Option<Range> {
        if self.adjacent_to(other) || self.intersects(other) {
            Some(Range::new(
                cmp::min(self.lo, other.lo),
                cmp::max(self.hi, other.hi),
            ))
        } else {
            None
        }
    }
}

/// A set of lines in files.
///
/// It is represented as a multimap keyed on file names, with values a collection of
/// non-overlapping ranges sorted by their start point. An inner `None` is interpreted to mean all
/// lines in all files.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct FileLines(Option<HashMap<FileName, Vec<Range>>>);

/// Normalizes the ranges so that the invariants for `FileLines` hold: ranges are non-overlapping,
/// and ordered by their start point.
fn normalize_ranges(ranges: &mut HashMap<FileName, Vec<Range>>) {
    for ranges in ranges.values_mut() {
        ranges.sort();
        let mut result = vec![];
        {
            let mut iter = ranges.into_iter().peekable();
            while let Some(next) = iter.next() {
                let mut next = *next;
                while let Some(&&mut peek) = iter.peek() {
                    if let Some(merged) = next.merge(peek) {
                        iter.next().unwrap();
                        next = merged;
                    } else {
                        break;
                    }
                }
                result.push(next)
            }
        }
        *ranges = result;
    }
}

impl FileLines {
    /// Creates a `FileLines` that contains all lines in all files.
    pub(crate) fn all() -> FileLines {
        FileLines(None)
    }

    /// Returns true if this `FileLines` contains all lines in all files.
    pub(crate) fn is_all(&self) -> bool {
        self.0.is_none()
    }

    pub fn from_ranges(mut ranges: HashMap<FileName, Vec<Range>>) -> FileLines {
        normalize_ranges(&mut ranges);
        FileLines(Some(ranges))
    }

    /// Returns an iterator over the files contained in `self`.
    pub fn files(&self) -> Files {
        Files(self.0.as_ref().map(|m| m.keys()))
    }

    /// Returns JSON representation as accepted by the `--file-lines JSON` arg.
    pub fn to_json_spans(&self) -> Vec<JsonSpan> {
        match &self.0 {
            None => vec![],
            Some(file_ranges) => file_ranges
                .iter()
                .flat_map(|(file, ranges)| ranges.iter().map(move |r| (file, r)))
                .map(|(file, range)| JsonSpan {
                    file: file.to_owned(),
                    range: (range.lo, range.hi),
                }).collect(),
        }
    }

    /// Returns true if `self` includes all lines in all files. Otherwise runs `f` on all ranges in
    /// the designated file (if any) and returns true if `f` ever does.
    fn file_range_matches<F>(&self, file_name: &FileName, f: F) -> bool
    where
        F: FnMut(&Range) -> bool,
    {
        let map = match self.0 {
            // `None` means "all lines in all files".
            None => return true,
            Some(ref map) => map,
        };

        match canonicalize_path_string(file_name).and_then(|file| map.get(&file)) {
            Some(ranges) => ranges.iter().any(f),
            None => false,
        }
    }

    /// Returns true if `range` is fully contained in `self`.
    #[allow(dead_code)]
    pub(crate) fn contains(&self, range: &LineRange) -> bool {
        self.file_range_matches(&range.file_name(), |r| r.contains(Range::from(range)))
    }

    /// Returns true if any lines in `range` are in `self`.
    pub(crate) fn intersects(&self, range: &LineRange) -> bool {
        self.file_range_matches(&range.file_name(), |r| r.intersects(Range::from(range)))
    }

    /// Returns true if `line` from `file_name` is in `self`.
    pub(crate) fn contains_line(&self, file_name: &FileName, line: usize) -> bool {
        self.file_range_matches(file_name, |r| r.lo <= line && r.hi >= line)
    }

    /// Returns true if all the lines between `lo` and `hi` from `file_name` are in `self`.
    pub(crate) fn contains_range(&self, file_name: &FileName, lo: usize, hi: usize) -> bool {
        self.file_range_matches(file_name, |r| r.contains(Range::new(lo, hi)))
    }
}

/// `FileLines` files iterator.
pub struct Files<'a>(Option<::std::collections::hash_map::Keys<'a, FileName, Vec<Range>>>);

impl<'a> iter::Iterator for Files<'a> {
    type Item = &'a FileName;

    fn next(&mut self) -> Option<&'a FileName> {
        self.0.as_mut().and_then(Iterator::next)
    }
}

fn canonicalize_path_string(file: &FileName) -> Option<FileName> {
    match *file {
        FileName::Real(ref path) => path.canonicalize().ok().map(FileName::Real),
        _ => Some(file.clone()),
    }
}

// This impl is needed for `Config::override_value` to work for use in tests.
impl str::FromStr for FileLines {
    type Err = String;

    fn from_str(s: &str) -> Result<FileLines, String> {
        let v: Vec<JsonSpan> = json::from_str(s).map_err(|e| e.to_string())?;
        let mut m = HashMap::new();
        for js in v {
            let (s, r) = JsonSpan::into_tuple(js)?;
            m.entry(s).or_insert_with(|| vec![]).push(r);
        }
        Ok(FileLines::from_ranges(m))
    }
}

// For JSON decoding.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct JsonSpan {
    file: FileName,
    range: (usize, usize),
}

impl JsonSpan {
    fn into_tuple(self) -> Result<(FileName, Range), String> {
        let (lo, hi) = self.range;
        let canonical = canonicalize_path_string(&self.file)
            .ok_or_else(|| format!("Can't canonicalize {}", &self.file))?;
        Ok((canonical, Range::new(lo, hi)))
    }
}

// This impl is needed for inclusion in the `Config` struct. We don't have a toml representation
// for `FileLines`, so it will just panic instead.
impl<'de> ::serde::de::Deserialize<'de> for FileLines {
    fn deserialize<D>(_: D) -> Result<Self, D::Error>
    where
        D: ::serde::de::Deserializer<'de>,
    {
        panic!(
            "FileLines cannot be deserialized from a project rustfmt.toml file: please \
             specify it via the `--file-lines` option instead"
        );
    }
}

// We also want to avoid attempting to serialize a FileLines to toml. The
// `Config` struct should ensure this impl is never reached.
impl ::serde::ser::Serialize for FileLines {
    fn serialize<S>(&self, _: S) -> Result<S::Ok, S::Error>
    where
        S: ::serde::ser::Serializer,
    {
        unreachable!("FileLines cannot be serialized. This is a rustfmt bug.");
    }
}

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

    #[test]
    fn test_range_intersects() {
        assert!(Range::new(1, 2).intersects(Range::new(1, 1)));
        assert!(Range::new(1, 2).intersects(Range::new(2, 2)));
        assert!(!Range::new(1, 2).intersects(Range::new(0, 0)));
        assert!(!Range::new(1, 2).intersects(Range::new(3, 10)));
        assert!(!Range::new(1, 3).intersects(Range::new(5, 5)));
    }

    #[test]
    fn test_range_adjacent_to() {
        assert!(!Range::new(1, 2).adjacent_to(Range::new(1, 1)));
        assert!(!Range::new(1, 2).adjacent_to(Range::new(2, 2)));
        assert!(Range::new(1, 2).adjacent_to(Range::new(0, 0)));
        assert!(Range::new(1, 2).adjacent_to(Range::new(3, 10)));
        assert!(!Range::new(1, 3).adjacent_to(Range::new(5, 5)));
    }

    #[test]
    fn test_range_contains() {
        assert!(Range::new(1, 2).contains(Range::new(1, 1)));
        assert!(Range::new(1, 2).contains(Range::new(2, 2)));
        assert!(!Range::new(1, 2).contains(Range::new(0, 0)));
        assert!(!Range::new(1, 2).contains(Range::new(3, 10)));
    }

    #[test]
    fn test_range_merge() {
        assert_eq!(None, Range::new(1, 3).merge(Range::new(5, 5)));
        assert_eq!(None, Range::new(4, 7).merge(Range::new(0, 1)));
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 5).merge(Range::new(4, 7))
        );
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 5).merge(Range::new(5, 7))
        );
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 5).merge(Range::new(6, 7))
        );
        assert_eq!(
            Some(Range::new(3, 7)),
            Range::new(3, 7).merge(Range::new(4, 5))
        );
    }

    use std::{collections::HashMap, path::PathBuf};
    use super::{FileName, FileLines};
    use super::json::{self, json, json_internal};

    #[test]
    fn file_lines_to_json() {
        let ranges: HashMap<FileName, Vec<Range>> = [
            (FileName::Real(PathBuf::from("src/main.rs")), vec![
                Range::new(1, 3),
                Range::new(5, 7)
            ]),
            (FileName::Real(PathBuf::from("src/lib.rs")), vec![
                Range::new(1, 7)
        ])].iter().cloned().collect();

        let file_lines = FileLines::from_ranges(ranges);
        let json = json::to_value(&file_lines.to_json_spans()).unwrap();
        assert_eq!(json, json! {[
            {"file": "src/main.rs", "range": [1, 3]},
            {"file": "src/main.rs", "range": [5, 7]},
            {"file": "src/lib.rs",  "range": [1, 7]},
        ]});
    }
}
<|endoftext|>


<commit_before>//! The main scheduler logic.
//!
//! The scheduler is implemented as a singleton in order to make it easy for code anywhere in the
//! project to make use of async functionality. The actual scheduler instance is not publicly
//! accessible, instead we use various standalone functions like `start()` and `wait_for()` to
//! safely manage access to the scheduler.
//!
//! # Scheduling Work
//!
//! Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
//! result of the work. Use `Async::await()` to suspend the current fiber until the work completes
//! and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
//! the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
//!
//! # Sharing Data Between Work
//!
//! Unlike with `std::thread::spawn()`, it's possible for work started with `scheduler::start()`
//! to borrow data from the caller:
//!
//! ```
//!
//! ```

use fiber::{self, Fiber, FiberId};
use cell_extras::AtomicInitCell;
use std::boxed::FnBox;
use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt::{self, Debug, Formatter};
use std::marker::PhantomData;
use std::mem;
use std::sync::{Condvar, Mutex, Once, ONCE_INIT};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver};

const DEFAULT_STACK_SIZE: usize = 64 * 1024;

static CONDVAR: AtomicInitCell<Condvar> = AtomicInitCell::new();
static INSTANCE: AtomicInitCell<Mutex<Scheduler>> = AtomicInitCell::new();
static INSTANCE_INIT: Once = ONCE_INIT;
static WORK_COUNTER: AtomicUsize = AtomicUsize::new(1);

/// Represents the result of a computation that may finish at some point in the future.
#[derive(Debug)]
pub struct Async<'a, T> {
    work: WorkId,
    receiver: Receiver<T>,
    _phantom: PhantomData<&'a FnMut()>,
}

impl<'a, T> Async<'a, T> {
    /// Suspend the current fiber until the async operation finishes.
    pub fn await(self) -> T {
        let result = {
            let Async { work, ref receiver, .. } = self;
            work.await();
            receiver.try_recv().expect("Failed to receive result of async computation")
        };

        // Don't run the destuctor since we've already waited on the result.
        mem::forget(self);

        result
    }
}

impl<T> Async<'static, T> {
    pub fn work_id(&self) -> WorkId {
        self.work
    }

    pub fn forget(self) {
        mem::forget(self);
    }
}

impl<'a, T> Drop for Async<'a, T> {
    fn drop(&mut self) {
        self.work.await();
    }
}

/// A shareable reference to a work unit, counterpart to `Async<T>`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct WorkId(usize);

impl WorkId {
    /// Suspends the current fiber until this work unit has completed.
    ///
    /// If the work unit has already finished then `await()` will return immediately.
    pub fn await(self) {
        if Scheduler::with(|scheduler| !scheduler.finished_work.contains(&self)) {
            Scheduler::with(|scheduler| scheduler.add_dependency(self));
            suspend();
        }
    }
}

/// Initializes a newly-spawned worker thread.
///
/// Prepares the worker thread by initializing it for Fiber usage.
// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn init_thread() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Manually convert current thread into a fiber because reasons.
    fiber::init();
}

// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn run_wait_fiber() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Setup this thread for running fibers and create an initial fiber for it. This will become
    // the wait fiber for this thread.
    fiber::init();

    fiber_routine();
}

pub fn start<'a, F, T>(func: F) -> Async<'a, T>
    where
    F: FnOnce() -> T,
    F: 'a + Send,
    T: 'a + Send,
{
    // Normally we can't box a closure with a non 'static lifetime because it could outlive its
    // borrowed data. In this case the lifetime parameter on the returned `Async` ensures that
    // the closure can't outlive the borrowed data, so we use this evil magic to convince the
    // compiler to allow us to box the closure.
    unsafe fn erase_lifetime<'a, F>(func: F) -> Box<FnBox()>
        where
        F: FnOnce(),
        F: 'a + Send,
    {
        let boxed_proc = Box::new(func);
        let proc_ptr = Box::into_raw(boxed_proc) as *mut FnBox();
        Box::from_raw(::std::mem::transmute(proc_ptr))
    }

    // Create the channel that'll be used to send the result of the operation to the `Async` object.
    let (sender, receiver) = mpsc::sync_channel(1);

    let work_id = WorkId(WORK_COUNTER.fetch_add(1, Ordering::Relaxed));

    let work_proc = unsafe {
        erase_lifetime(move || {
            let result = func();
            sender.try_send(result).expect("Failed to send async result");
        })
    };

    Scheduler::with(move |scheduler| scheduler.schedule_work(Work {
        func: work_proc,
        id: work_id,
    }));

    Async {
        work: work_id,
        receiver: receiver,
        _phantom: PhantomData,
    }
}

/// Suspends the current fiber and makes the wait fiber active.
///
/// Generally you shouldn't need to call this directly, but if you have one piece of code that
/// runs synchronously for a long time you can use `suspend()` to yield time to other work.
pub fn suspend() {
    let next_fiber = Scheduler::with(|scheduler| scheduler.next_fiber());

    let suspended = unsafe { next_fiber.resume() };

    Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
}

fn fiber_routine() -> ! {
    loop {
        match Scheduler::with(|scheduler| scheduler.next()) {
            Some(NextWork::Work(Work { func, id })) => {
                Scheduler::with(|scheduler| scheduler.start_work(id));
                func();
                Scheduler::with(|scheduler| scheduler.finish_work(id));
            },
            Some(NextWork::Fiber(fiber)) => {
                let suspended = unsafe { fiber.resume() };
                Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
            },
            None => {
                // If there's no new work and no fibers ready to run then we want to block the
                // thread until some becomes available.
                let mutex = INSTANCE.borrow();
                let condvar = CONDVAR.borrow();

                let _ = condvar
                    .wait(mutex.lock().expect("Scheduler mutex was poisoned"))
                    .expect("Scheduler mutex was poisoned");
            },
        }
    }
}

struct Work {
    func: Box<FnBox()>,
    id: WorkId,
}

impl Debug for Work {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        write!(formatter, "Work {{ id: {:?} }}", self.id)
    }
}

enum NextWork {
    Work(Work),
    Fiber(Fiber),
}

struct Scheduler {
    /// Work units that are currently in progress.
    running_work: HashSet<WorkId>,

    work_map: HashMap<FiberId, WorkId>,

    /// New units of work that haven't been started on a fiber yet.
    ///
    /// These are ready to be made active at any time.
    new_work: VecDeque<Work>,

    /// Fibers that have no pending dependencies.
    ///
    /// These are ready to be made active at any time.
    ready_fibers: VecDeque<Fiber>,

    /// A map specifying which pending fibers depend on which others.
    ///
    /// Once all of a fiber's dependencies complete it should be moved to `new_work`.
    dependencies: HashMap<FiberId, (Option<Fiber>, HashSet<WorkId>)>,

    // TODO: Should we distinguise between "finished" and "ready" fibers? My intuition is that we'd
    // want to give fibers that actively have work CPU time before we resume fibers that would be
    // pulling new work, but maybe not? If we threw them all into one queue I guess the worst case
    // scenario would be there's no work left, a bunch of empty fibers, and only a few fibers with
    // active work. In which case we might have to cycle through a bunch of fibers before we can
    // start doing actual work.
    finished: VecDeque<Fiber>,

    finished_work: HashSet<WorkId>,
}

unsafe impl Send for Scheduler {}

impl Scheduler {
    /// Provides safe access to the scheduler instance.
    ///
    /// # Fiber Switches
    ///
    /// Note that it is an error to call `Fiber::make_active()` within `func`. Doing so will cause
    /// the `Mutex` guard on the instance to never unlock, making the scheduler instance
    /// inaccessible. All standalone functions that access the scheduler and wish to switch fibers
    /// should use `Scheduler::next()` to return the next fiber from `with()` and then call
    /// `make_active()` *after* `with()` has returned.
    fn with<F, T>(func: F) -> T
        where F: FnOnce(&mut Scheduler) -> T
    {
        INSTANCE_INIT.call_once(|| {
            let scheduler = Scheduler {
                running_work: HashSet::new(),
                work_map: HashMap::new(),
                new_work: VecDeque::new(),
                ready_fibers: VecDeque::new(),
                dependencies: HashMap::new(),
                finished: VecDeque::new(),
                finished_work: HashSet::new(),
            };

            INSTANCE.init(Mutex::new(scheduler));
            CONDVAR.init(Condvar::new());
        });

        let instance = INSTANCE.borrow();

        let mut guard = instance.lock().expect("Scheduler mutex was poisoned");
        func(&mut *guard)
    }

    /// Add a new unit of work to the pending queue.
    fn schedule_work(&mut self, work: Work) {
        self.new_work.push_back(work);
        CONDVAR.borrow().notify_one();
    }

    /// Adds `dependency` as a dependency of the currently running fiber.
    fn add_dependency(&mut self, dependency: WorkId) {
        self.add_dependencies([dependency].iter().cloned());
    }

    /// Adds all work items in `dependencies` as dependencies of the currently running fiber.
    fn add_dependencies<I>(&mut self, dependencies: I)
        where
        I: IntoIterator<Item = WorkId>,
    {
        let pending = fiber::current().unwrap();

        debug_assert!(
            !self.dependencies.contains_key(&pending),
            "Marking a fiber as pending but it is already pending: {:?}",
            pending,
        );

        // Add `pending` to set of pending fibers and list `dependencies` as dependencies.
        let &mut (_, ref mut dependencies_set) =
            self.dependencies
            .entry(pending)
            .or_insert((None, HashSet::new()));

        // Add `fibers` to the list of ready fibers.
        for dependency in dependencies {
            dependencies_set.insert(dependency);
        }
    }

    fn start_work(&mut self, new_work: WorkId) {
        let current = fiber::current().unwrap();
        self.running_work.insert(new_work);
        self.work_map.insert(current, new_work);
    }

    /// Removes the specified unit of work from the scheduler, updating any dependent work.
    fn finish_work(&mut self, finished_work: WorkId) {
        // Iterate over all suspended work units, removing `finished_work` as a dependency where
        // necessary. If any of the work units no longer have dependencies then
        let mut ready = Vec::new();
        for (&pending_fiber, &mut (_, ref mut dependencies)) in &mut self.dependencies {
            dependencies.remove(&finished_work);
            if dependencies.len() == 0 {
                ready.push(pending_fiber);
            }
        }

        for ready_work in ready {
            let (maybe_fiber, _) = self.dependencies.remove(&ready_work).unwrap();
            if let Some(ready_fiber) = maybe_fiber {
                self.ready_fibers.push_back(ready_fiber);
                CONDVAR.borrow().notify_one();
            }
        }

        let fiber = fiber::current().unwrap();
        assert!(self.running_work.remove(&finished_work), "{:?} wasn't in running work set when it finished", finished_work);
        assert!(self.work_map.remove(&fiber).is_some(), "{:?} didn't have {:?} associated in the work map", fiber, finished_work);

        self.finished_work.insert(finished_work);
    }

    /// Performs the necessary bookkeeping when a fiber becomes active.
    fn handle_suspended(&mut self, suspended: Fiber) {
        // If the suspended fiber has dependencies then update the dependencies map with the
        // actual fiber, that way when its dependencies complete it can be resumed. Otherwise, the
        // fiber is done and ready to take on more work. This means that we need to make sure that
        // we always call `add_dependencies()` before suspending a fiber, otherwise a fiber could
        // be marked as done before it's ready.
        if let Some(&mut (ref mut none_fiber, _)) = self.dependencies.get_mut(&suspended.id()) {
            debug_assert!(none_fiber.is_none(), "Dependencies map already had a fiber assicated with fiber ID");
            mem::replace(none_fiber, Some(suspended));
        } else if self.work_map.contains_key(&suspended.id()) {
            self.ready_fibers.push_back(suspended);
        } else {
            self.finished.push_back(suspended);
        }
    }

    /// Gets the next ready fiber, or creates a new one if necessary.
    fn next_fiber(&mut self) -> Fiber {
        self.ready_fibers.pop_front()
            .or_else(|| self.finished.pop_front())
            .unwrap_or_else(|| {
                fn fiber_proc(prev: Fiber) -> ! {
                    // The current fiber has been resumed. Let the scheduler know that the previous fiber is no
                    // longer active.
                    Scheduler::with(|scheduler| scheduler.handle_suspended(prev));

                    fiber_routine();
                }

                Fiber::new(DEFAULT_STACK_SIZE, fiber_proc)
            })
    }

    /// Gets the next available work for a thread, either a new unit of work or a ready fiber.
    ///
    /// Prioritizes new work over pending fibers, and will only return ready fibers that already
    /// have work. To get *any* next fiber, including ones without active work or a new one if no
    /// existing fibers are available, use `next_fiber()`.
    fn next(&mut self) -> Option<NextWork> {
        if let Some(work) = self.new_work.pop_front() {
            Some(NextWork::Work(work))
        } else {
            self.ready_fibers.pop_front().map(|fiber| NextWork::Fiber(fiber))
        }
    }
}
<commit_msg>Check if work has completed when adding dependency<commit_after>//! The main scheduler logic.
//!
//! The scheduler is implemented as a singleton in order to make it easy for code anywhere in the
//! project to make use of async functionality. The actual scheduler instance is not publicly
//! accessible, instead we use various standalone functions like `start()` and `wait_for()` to
//! safely manage access to the scheduler.
//!
//! # Scheduling Work
//!
//! Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
//! result of the work. Use `Async::await()` to suspend the current fiber until the work completes
//! and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
//! the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
//!
//! # Sharing Data Between Work
//!
//! Unlike with `std::thread::spawn()`, it's possible for work started with `scheduler::start()`
//! to borrow data from the caller:
//!
//! ```
//!
//! ```

use fiber::{self, Fiber, FiberId};
use cell_extras::AtomicInitCell;
use std::boxed::FnBox;
use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt::{self, Debug, Formatter};
use std::marker::PhantomData;
use std::mem;
use std::sync::{Condvar, Mutex, Once, ONCE_INIT};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver};

const DEFAULT_STACK_SIZE: usize = 64 * 1024;

static CONDVAR: AtomicInitCell<Condvar> = AtomicInitCell::new();
static INSTANCE: AtomicInitCell<Mutex<Scheduler>> = AtomicInitCell::new();
static INSTANCE_INIT: Once = ONCE_INIT;
static WORK_COUNTER: AtomicUsize = AtomicUsize::new(1);

/// Represents the result of a computation that may finish at some point in the future.
///
/// Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
/// result of the work. Use `Async::await()` to suspend the current fiber until the work completes
/// and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
/// the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
///
/// # Sharing Data Across Work
///
/// It's possible to share
#[derive(Debug)]
pub struct Async<'a, T> {
    work: WorkId,
    receiver: Receiver<T>,
    _phantom: PhantomData<&'a FnMut()>,
}

impl<'a, T> Async<'a, T> {
    /// Suspend the current fiber until the async operation finishes.
    pub fn await(self) -> T {
        let result = {
            let Async { work, ref receiver, .. } = self;
            work.await();
            receiver.try_recv().expect("Failed to receive result of async computation")
        };

        // Don't run the destuctor since we've already waited on the result.
        mem::forget(self);

        result
    }
}

impl<T> Async<'static, T> {
    pub fn work_id(&self) -> WorkId {
        self.work
    }

    pub fn forget(self) {
        mem::forget(self);
    }
}

impl<'a, T> Drop for Async<'a, T> {
    fn drop(&mut self) {
        self.work.await();
    }
}

/// A shareable reference to a work unit, counterpart to `Async<T>`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct WorkId(usize);

impl WorkId {
    /// Suspends the current fiber until this work unit has completed.
    ///
    /// If the work unit has already finished then `await()` will return immediately.
    pub fn await(self) {
        if Scheduler::with(|scheduler| scheduler.add_dependency(self)) {
            suspend();
        }
    }
}

/// Initializes a newly-spawned worker thread.
///
/// Prepares the worker thread by initializing it for Fiber usage.
// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn init_thread() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Manually convert current thread into a fiber because reasons.
    fiber::init();
}

// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn run_wait_fiber() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Setup this thread for running fibers and create an initial fiber for it. This will become
    // the wait fiber for this thread.
    fiber::init();

    fiber_routine();
}

pub fn start<'a, F, T>(func: F) -> Async<'a, T>
    where
    F: FnOnce() -> T,
    F: 'a + Send,
    T: 'a + Send,
{
    // Normally we can't box a closure with a non 'static lifetime because it could outlive its
    // borrowed data. In this case the lifetime parameter on the returned `Async` ensures that
    // the closure can't outlive the borrowed data, so we use this evil magic to convince the
    // compiler to allow us to box the closure.
    unsafe fn erase_lifetime<'a, F>(func: F) -> Box<FnBox()>
        where
        F: FnOnce(),
        F: 'a + Send,
    {
        let boxed_proc = Box::new(func);
        let proc_ptr = Box::into_raw(boxed_proc) as *mut FnBox();
        Box::from_raw(::std::mem::transmute(proc_ptr))
    }

    // Create the channel that'll be used to send the result of the operation to the `Async` object.
    let (sender, receiver) = mpsc::sync_channel(1);

    let work_id = WorkId(WORK_COUNTER.fetch_add(1, Ordering::Relaxed));

    let work_proc = unsafe {
        erase_lifetime(move || {
            let result = func();
            sender.try_send(result).expect("Failed to send async result");
        })
    };

    Scheduler::with(move |scheduler| scheduler.schedule_work(Work {
        func: work_proc,
        id: work_id,
    }));

    Async {
        work: work_id,
        receiver: receiver,
        _phantom: PhantomData,
    }
}

/// Suspends the current fiber and makes the wait fiber active.
///
/// Generally you shouldn't need to call this directly, but if you have one piece of code that
/// runs synchronously for a long time you can use `suspend()` to yield time to other work.
pub fn suspend() {
    let next_fiber = Scheduler::with(|scheduler| scheduler.next_fiber());

    let suspended = unsafe { next_fiber.resume() };

    Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
}

fn fiber_routine() -> ! {
    loop {
        match Scheduler::with(|scheduler| scheduler.next()) {
            Some(NextWork::Work(Work { func, id })) => {
                Scheduler::with(|scheduler| scheduler.start_work(id));
                func();
                Scheduler::with(|scheduler| scheduler.finish_work(id));
            },
            Some(NextWork::Fiber(fiber)) => {
                let suspended = unsafe { fiber.resume() };
                Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
            },
            None => {
                // If there's no new work and no fibers ready to run then we want to block the
                // thread until some becomes available.
                let mutex = INSTANCE.borrow();
                let condvar = CONDVAR.borrow();

                let _ = condvar
                    .wait(mutex.lock().expect("Scheduler mutex was poisoned"))
                    .expect("Scheduler mutex was poisoned");
            },
        }
    }
}

struct Work {
    func: Box<FnBox()>,
    id: WorkId,
}

impl Debug for Work {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        write!(formatter, "Work {{ id: {:?} }}", self.id)
    }
}

enum NextWork {
    Work(Work),
    Fiber(Fiber),
}

struct Scheduler {
    /// Work units that are currently in progress.
    running_work: HashSet<WorkId>,

    work_map: HashMap<FiberId, WorkId>,

    /// New units of work that haven't been started on a fiber yet.
    ///
    /// These are ready to be made active at any time.
    new_work: VecDeque<Work>,

    /// Fibers that have no pending dependencies.
    ///
    /// These are ready to be made active at any time.
    ready_fibers: VecDeque<Fiber>,

    /// A map specifying which pending fibers depend on which others.
    ///
    /// Once all of a fiber's dependencies complete it should be moved to `new_work`.
    dependencies: HashMap<FiberId, (Option<Fiber>, HashSet<WorkId>)>,

    // TODO: Should we distinguise between "finished" and "ready" fibers? My intuition is that we'd
    // want to give fibers that actively have work CPU time before we resume fibers that would be
    // pulling new work, but maybe not? If we threw them all into one queue I guess the worst case
    // scenario would be there's no work left, a bunch of empty fibers, and only a few fibers with
    // active work. In which case we might have to cycle through a bunch of fibers before we can
    // start doing actual work.
    finished: VecDeque<Fiber>,

    finished_work: HashSet<WorkId>,
}

unsafe impl Send for Scheduler {}

impl Scheduler {
    /// Provides safe access to the scheduler instance.
    ///
    /// # Fiber Switches
    ///
    /// Note that it is an error to call `Fiber::make_active()` within `func`. Doing so will cause
    /// the `Mutex` guard on the instance to never unlock, making the scheduler instance
    /// inaccessible. All standalone functions that access the scheduler and wish to switch fibers
    /// should use `Scheduler::next()` to return the next fiber from `with()` and then call
    /// `make_active()` *after* `with()` has returned.
    fn with<F, T>(func: F) -> T
        where F: FnOnce(&mut Scheduler) -> T
    {
        INSTANCE_INIT.call_once(|| {
            let scheduler = Scheduler {
                running_work: HashSet::new(),
                work_map: HashMap::new(),
                new_work: VecDeque::new(),
                ready_fibers: VecDeque::new(),
                dependencies: HashMap::new(),
                finished: VecDeque::new(),
                finished_work: HashSet::new(),
            };

            INSTANCE.init(Mutex::new(scheduler));
            CONDVAR.init(Condvar::new());
        });

        let instance = INSTANCE.borrow();

        let mut guard = instance.lock().expect("Scheduler mutex was poisoned");
        func(&mut *guard)
    }

    /// Add a new unit of work to the pending queue.
    fn schedule_work(&mut self, work: Work) {
        self.new_work.push_back(work);
        CONDVAR.borrow().notify_one();
    }

    /// Adds `dependency` as a dependency of the currently running fiber.
    ///
    /// Returns `true` if work is still in progress and was added as a dependency, false otherwise.
    fn add_dependency(&mut self, dependency: WorkId) -> bool {
        let pending = fiber::current().unwrap();

        if !self.finished_work.contains(&dependency) {
            debug_assert!(
                !self.dependencies.contains_key(&pending),
                "Marking a fiber as pending but it is already pending: {:?}",
                pending,
            );

            // Add `pending` to set of pending fibers and list `dependencies` as dependencies.
            let &mut (_, ref mut dependencies_set) =
            self.dependencies
            .entry(pending)
            .or_insert((None, HashSet::new()));

            // Add `fibers` to the list of ready fibers.
            dependencies_set.insert(dependency);

            true
        } else {
            false
        }
    }

    fn start_work(&mut self, new_work: WorkId) {
        let current = fiber::current().unwrap();
        self.running_work.insert(new_work);
        self.work_map.insert(current, new_work);
    }

    /// Removes the specified unit of work from the scheduler, updating any dependent work.
    fn finish_work(&mut self, finished_work: WorkId) {
        // Iterate over all suspended work units, removing `finished_work` as a dependency where
        // necessary. If any of the work units no longer have dependencies then
        let mut ready = Vec::new();
        for (&pending_fiber, &mut (_, ref mut dependencies)) in &mut self.dependencies {
            dependencies.remove(&finished_work);
            if dependencies.len() == 0 {
                ready.push(pending_fiber);
            }
        }

        for ready_work in ready {
            let (maybe_fiber, _) = self.dependencies.remove(&ready_work).unwrap();
            if let Some(ready_fiber) = maybe_fiber {
                self.ready_fibers.push_back(ready_fiber);
                CONDVAR.borrow().notify_one();
            }
        }

        let fiber = fiber::current().unwrap();
        assert!(self.running_work.remove(&finished_work), "{:?} wasn't in running work set when it finished", finished_work);
        assert!(self.work_map.remove(&fiber).is_some(), "{:?} didn't have {:?} associated in the work map", fiber, finished_work);

        self.finished_work.insert(finished_work);
    }

    /// Performs the necessary bookkeeping when a fiber becomes active.
    fn handle_suspended(&mut self, suspended: Fiber) {
        // If the suspended fiber has dependencies then update the dependencies map with the
        // actual fiber, that way when its dependencies complete it can be resumed. Otherwise, the
        // fiber is done and ready to take on more work. This means that we need to make sure that
        // we always call `add_dependencies()` before suspending a fiber, otherwise a fiber could
        // be marked as done before it's ready.
        if let Some(&mut (ref mut none_fiber, _)) = self.dependencies.get_mut(&suspended.id()) {
            debug_assert!(none_fiber.is_none(), "Dependencies map already had a fiber assicated with fiber ID");
            mem::replace(none_fiber, Some(suspended));
        } else if self.work_map.contains_key(&suspended.id()) {
            self.ready_fibers.push_back(suspended);
        } else {
            self.finished.push_back(suspended);
        }
    }

    /// Gets the next ready fiber, or creates a new one if necessary.
    fn next_fiber(&mut self) -> Fiber {
        self.ready_fibers.pop_front()
            .or_else(|| self.finished.pop_front())
            .unwrap_or_else(|| {
                fn fiber_proc(prev: Fiber) -> ! {
                    // The current fiber has been resumed. Let the scheduler know that the previous fiber is no
                    // longer active.
                    Scheduler::with(|scheduler| scheduler.handle_suspended(prev));

                    fiber_routine();
                }

                Fiber::new(DEFAULT_STACK_SIZE, fiber_proc)
            })
    }

    /// Gets the next available work for a thread, either a new unit of work or a ready fiber.
    ///
    /// Prioritizes new work over pending fibers, and will only return ready fibers that already
    /// have work. To get *any* next fiber, including ones without active work or a new one if no
    /// existing fibers are available, use `next_fiber()`.
    fn next(&mut self) -> Option<NextWork> {
        if let Some(work) = self.new_work.pop_front() {
            Some(NextWork::Work(work))
        } else {
            self.ready_fibers.pop_front().map(|fiber| NextWork::Fiber(fiber))
        }
    }
}
<|endoftext|>


<commit_before>generate_error_imports!();

generate_error_types!(CounterError, CounterErrorKind,
    StoreReadError          => "Store read error",
    StoreWriteError         => "Store write error",
    HeaderTypeError         => "Header type error",
    HeaderFieldMissingError => "Header field missing error"
);

<commit_msg>libimagcounter: use generate_error_module!() macro and reexport generated types<commit_after>generate_error_module!(
    generate_error_types!(CounterError, CounterErrorKind,
        StoreReadError          => "Store read error",
        StoreWriteError         => "Store write error",
        HeaderTypeError         => "Header type error",
        HeaderFieldMissingError => "Header field missing error"
    );
);

pub use self::error::CounterError;
pub use self::error::CounterErrorKind;

<|endoftext|>


<commit_before>extern crate env_logger;
extern crate log_panics;
extern crate log;
#[cfg(target_os = "android")]
extern crate android_logger;

use self::env_logger::Builder;
use self::log::{Record, LevelFilter, Level};
use std::env;
use std::io::Write;
use std::sync::{Once, ONCE_INIT};
#[cfg(target_os = "android")]
use self::android_logger::Filter;

pub struct LoggerUtils {}

static LOGGER_INIT: Once = ONCE_INIT;

impl LoggerUtils {
    pub fn init() {
        //Starts logging the panic messages using the system logger.
        LOGGER_INIT.call_once(|| {

            log_panics::init(); //Logging of panics is essential for android. As android does not log to stdout for native code
            if cfg!(target_os = "android") {
                //Set logging to off when deploying production android app.
                #[cfg(target_os = "android")]
                android_logger::init_once(
                    Filter::default().with_min_level(log::Level::Trace)
                );
                info!("Logging for Android");
            } else {
                let mut builder = Builder::new();
                builder.format(|buf, record| writeln!(buf, "{:>5}|{:<30}|{:>35}:{:<4}| {}", record.level(), record.target(),record.file().get_or_insert(""), record.line().get_or_insert(0), record.args()))
                .filter(None, LevelFilter::Off);
                if env::var("RUST_LOG").is_ok() {
                    builder.parse(&env::var("RUST_LOG").unwrap());
                }

                builder.init();
            }
        });
    }
}

#[macro_export]
macro_rules! try_log {
    ($expr:expr) => (match $expr {
        Ok(val) => val,
        Err(err) => {
            error!("try_log! | {}", err);
            return Err(From::from(err))
        }
    })
}

macro_rules! _map_err {
    ($lvl:expr, $expr:expr) => (
        |err| {
            log!($lvl, "{} - {}", $expr, err);
            err
        }
    );
    ($lvl:expr) => (
        |err| {
            log!($lvl, "{}", err);
            err
        }
    )
}

#[macro_export]
macro_rules! map_err_err {
    () => ( _map_err!(::log::Level::Error) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Error, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_trace {
    () => ( _map_err!(::log::Level::Trace) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Trace, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_info {
    () => ( _map_err!(::log::Level::Info) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Info, $($arg)*) )
}
<commit_msg>Fix warnings in current master.<commit_after>extern crate env_logger;
extern crate log_panics;
extern crate log;
#[cfg(target_os = "android")]
extern crate android_logger;

use self::env_logger::Builder;
use self::log::LevelFilter;
use std::env;
use std::io::Write;
use std::sync::{Once, ONCE_INIT};
#[cfg(target_os = "android")]
use self::android_logger::Filter;

pub struct LoggerUtils {}

static LOGGER_INIT: Once = ONCE_INIT;

impl LoggerUtils {
    pub fn init() {
        //Starts logging the panic messages using the system logger.
        LOGGER_INIT.call_once(|| {

            log_panics::init(); //Logging of panics is essential for android. As android does not log to stdout for native code
            if cfg!(target_os = "android") {
                //Set logging to off when deploying production android app.
                #[cfg(target_os = "android")]
                android_logger::init_once(
                    Filter::default().with_min_level(log::Level::Trace)
                );
                info!("Logging for Android");
            } else {
                let mut builder = Builder::new();
                builder.format(|buf, record| writeln!(buf, "{:>5}|{:<30}|{:>35}:{:<4}| {}", record.level(), record.target(),record.file().get_or_insert(""), record.line().get_or_insert(0), record.args()))
                .filter(None, LevelFilter::Off);
                if env::var("RUST_LOG").is_ok() {
                    builder.parse(&env::var("RUST_LOG").unwrap());
                }

                builder.init();
            }
        });
    }
}

#[macro_export]
macro_rules! try_log {
    ($expr:expr) => (match $expr {
        Ok(val) => val,
        Err(err) => {
            error!("try_log! | {}", err);
            return Err(From::from(err))
        }
    })
}

macro_rules! _map_err {
    ($lvl:expr, $expr:expr) => (
        |err| {
            log!($lvl, "{} - {}", $expr, err);
            err
        }
    );
    ($lvl:expr) => (
        |err| {
            log!($lvl, "{}", err);
            err
        }
    )
}

#[macro_export]
macro_rules! map_err_err {
    () => ( _map_err!(::log::Level::Error) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Error, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_trace {
    () => ( _map_err!(::log::Level::Trace) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Trace, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_info {
    () => ( _map_err!(::log::Level::Info) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Info, $($arg)*) )
}
<|endoftext|>


<commit_before>#[cfg(target_arch = "x86_64")]
extern crate raw_cpuid;
extern crate syscall;

use std::env;
use std::fs::File;
use std::io::Read;
use std::path::Path;

// std::fmt::Write conflicts with std::io::Write, hence the alias
use std::fmt::Write as FmtWrite;

const SECONDS_PER_MINUTE: i64 = 60;
const SECONDS_PER_HOUR: i64 = 3600;
const SECONDS_PER_DAY: i64 = 86400;

fn main() {
    let user = env::var("USER").unwrap_or_default();
    let mut hostname = String::new();
    if let Ok(mut file) = File::open("file:/etc/hostname") {
        let _ = file.read_to_string(&mut hostname);
    }

    let mut uptime_str = String::new();

    let mut ts = syscall::TimeSpec::default();
    if syscall::clock_gettime(syscall::CLOCK_MONOTONIC, &mut ts).is_ok() {
        let uptime = ts.tv_sec;
        let uptime_secs = uptime % 60;
        let uptime_mins = (uptime / SECONDS_PER_MINUTE) % 60;
        let uptime_hours = (uptime / SECONDS_PER_HOUR) % 24;
        let uptime_days = (uptime / SECONDS_PER_DAY) % 365;

        let fmt_result;
        if uptime_days > 0 {
            fmt_result = write!(
                &mut uptime_str,
                "{}d {}h {}m {}s",
                uptime_days, uptime_hours, uptime_mins, uptime_secs
            );
        } else if uptime_hours > 0 {
            fmt_result = write!(
                &mut uptime_str,
                "{}h {}m {}s",
                uptime_hours, uptime_mins, uptime_secs
            );
        } else if uptime_mins > 0 {
            fmt_result = write!(&mut uptime_str, "{}m {}s", uptime_mins, uptime_secs);
        } else {
            fmt_result = write!(&mut uptime_str, "{}s", uptime_secs);
        }

        if let Err(err) = fmt_result {
            eprintln!("error: couldn't parse uptime, {}", err);
        }
    }

    let mut shell = String::new();
    {
        if let Ok(shell_path) = env::var("SHELL") {
            if let Some(file_name) = Path::new(&shell_path).file_name() {
                shell = file_name.to_str().unwrap_or("").to_string();
            }
        }
    }

    let mut width = 0;
    let mut height = 0;
    if let Ok(display_name) = env::var("DISPLAY") {
        if let Ok(display) = syscall::open(&display_name, syscall::O_STAT) {
            let mut buf: [u8; 4096] = [0; 4096];
            if let Ok(count) = syscall::fpath(display, &mut buf) {
                let path = unsafe { String::from_utf8_unchecked(Vec::from(&buf[..count])) };
                let res = path.split(':').nth(1).unwrap_or("");
                width = res
                    .split('/')
                    .nth(1)
                    .unwrap_or("")
                    .parse::<i32>()
                    .unwrap_or(0);
                height = res
                    .split('/')
                    .nth(2)
                    .unwrap_or("")
                    .parse::<i32>()
                    .unwrap_or(0);
            }
            let _ = syscall::close(display);
        }
    }

    let mut cpu = String::new();
    #[cfg(target_arch = "x86_64")]
    {
        let cpuid = raw_cpuid::CpuId::new();
        if let Some(info) = cpuid.get_extended_function_info() {
            if let Some(brand) = info.processor_brand_string() {
                cpu = brand.trim().to_string();
            }
        }
    }

    let mut ram = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("memory:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                ram = format!(
                    "{}MB / {}MB",
                    (used + 1_048_575) / 1_048_576,
                    (size + 1_048_575) / 1_048_576
                );
            }
            let _ = syscall::close(fd);
        }
    }

    let left = [
        "              :+yMMMMy+:              ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "       sydMMMMo/sMMMMs/oMMMMdys       ",
        "    .oMMMdso              osdMMMo.    ",
        "   .+MMd/`   -:::::::`      `/dMM+.   ",
        "  +dMMN.     NMMNNNMMdyo`     .NMMd+  ",
        "  yMMN       NM+   oomMMN       NMMy  ",
        " hNMd.       NM+     `oMN       .dMNh ",
        " dMMh        NM+     `oMN        hMMd ",
        " dMMh        NM+-oooodMMN        hMMd ",
        " dMMh        NM+/MMMMdhs`        hMMd ",
        " hNMd.       NM+`/mMMm+         .dMNh ",
        "  yMMN       NM+   oNMd+        NMMy  ",
        "  +dMMN.     NM+    omMMm     .NMMd+  ",
        "   .+MMd/`   --.     .---   `/dMM+.   ",
        "    .oMMMdso              osdMMMo.    ",
        "       sydMMMMo////////oMMMMdys       ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "              :++++++++:              ",
    ];

    const S: &str = "\x1B[1;38;5;75m"; // blue start
    const E: &str = "\x1B[0m"; // end
    let right = [
        format!("{}{}{}@{}{}{}", S, user, E, S, hostname.trim(), E),
        format!("{}OS:         {}redox-os", S, E),
        format!("{}Kernel:     {}redox", S, E),
        format!("{}Uptime:     {}{}", S, E, uptime_str),
        format!("{}Shell:      {}{}", S, E, shell),
        format!("{}Resolution: {}{}x{}", S, E, width, height),
        format!("{}DE:         {}orbital", S, E),
        format!("{}WM:         {}orbital", S, E),
        format!("{}Font:       {}unifont", S, E),
        format!("{}CPU:        {}{}", S, E, cpu),
        format!("{}RAM:        {}{}", S, E, ram),
    ];

    for (i, line) in left.iter().enumerate() {
        print!("\x1B[1;38;5;75m{}  \x1B[0m", line);
        if let Some(r) = right.get(i) {
            print!("{}", r);
        }
        println!();
    }
}
<commit_msg>Screenfetch improvements<commit_after>#[cfg(target_arch = "x86_64")]
extern crate raw_cpuid;
extern crate syscall;

use std::env;
use std::fs::{self, File};
use std::io::Read;
use std::path::Path;

// std::fmt::Write conflicts with std::io::Write, hence the alias
use std::fmt::Write as FmtWrite;

const SECONDS_PER_MINUTE: i64 = 60;
const SECONDS_PER_HOUR: i64 = 3600;
const SECONDS_PER_DAY: i64 = 86400;

const KIB: u64 = 1024;
const MIB: u64 = 1024 * KIB;
const GIB: u64 = 1024 * MIB;
const TIB: u64 = 1024 * GIB;

fn format_size(size: u64) -> String {
    if size >= 4 * TIB {
        format!("{:.1} TiB", size as f64 / TIB as f64)
    } else if size >= GIB {
        format!("{:.1} GiB", size as f64 / GIB as f64)
    } else if size >= MIB {
        format!("{:.1} MiB", size as f64 / MIB as f64)
    } else if size >= KIB {
        format!("{:.1} KiB", size as f64 / KIB as f64)
    } else {
        format!("{} B", size)
    }
}

fn main() {
    let user = env::var("USER").unwrap_or_default();
    let mut hostname = String::new();
    if let Ok(mut file) = File::open("file:/etc/hostname") {
        let _ = file.read_to_string(&mut hostname);
    }

    let redox_release = match fs::read_to_string("/etc/redox-release") {
        Ok(ok) => ok,
        Err(err) => {
            eprintln!("error: failed to read /etc/redox-release: {}", err);
            "(unknown release)".to_string()
        }
    };

    let mut kernel = String::new();
    match fs::read_to_string("sys:uname") {
        Ok(uname) => for line in uname.lines() {
            if line.is_empty() {
                continue;
            }

            if ! kernel.is_empty() {
                kernel.push(' ');
            }

            kernel.push_str(line);
        },
        Err(err) => {
            eprintln!("error: failed to read sys:uname: {}", err);
        }
    }

    let mut uptime_str = String::new();

    let mut ts = syscall::TimeSpec::default();
    if syscall::clock_gettime(syscall::CLOCK_MONOTONIC, &mut ts).is_ok() {
        let uptime = ts.tv_sec;
        let uptime_secs = uptime % 60;
        let uptime_mins = (uptime / SECONDS_PER_MINUTE) % 60;
        let uptime_hours = (uptime / SECONDS_PER_HOUR) % 24;
        let uptime_days = (uptime / SECONDS_PER_DAY) % 365;

        let fmt_result;
        if uptime_days > 0 {
            fmt_result = write!(
                &mut uptime_str,
                "{}d {}h {}m {}s",
                uptime_days, uptime_hours, uptime_mins, uptime_secs
            );
        } else if uptime_hours > 0 {
            fmt_result = write!(
                &mut uptime_str,
                "{}h {}m {}s",
                uptime_hours, uptime_mins, uptime_secs
            );
        } else if uptime_mins > 0 {
            fmt_result = write!(&mut uptime_str, "{}m {}s", uptime_mins, uptime_secs);
        } else {
            fmt_result = write!(&mut uptime_str, "{}s", uptime_secs);
        }

        if let Err(err) = fmt_result {
            eprintln!("error: couldn't parse uptime, {}", err);
        }
    }

    let mut shell = String::new();
    {
        if let Ok(shell_path) = env::var("SHELL") {
            if let Some(file_name) = Path::new(&shell_path).file_name() {
                shell = file_name.to_str().unwrap_or("").to_string();
            }
        }
    }

    let mut width = 0;
    let mut height = 0;
    if let Ok(display_name) = env::var("DISPLAY") {
        if let Ok(display) = syscall::open(&display_name, syscall::O_STAT) {
            let mut buf: [u8; 4096] = [0; 4096];
            if let Ok(count) = syscall::fpath(display, &mut buf) {
                let path = unsafe { String::from_utf8_unchecked(Vec::from(&buf[..count])) };
                let res = path.split(':').nth(1).unwrap_or("");
                width = res
                    .split('/')
                    .nth(1)
                    .unwrap_or("")
                    .parse::<i32>()
                    .unwrap_or(0);
                height = res
                    .split('/')
                    .nth(2)
                    .unwrap_or("")
                    .parse::<i32>()
                    .unwrap_or(0);
            }
            let _ = syscall::close(display);
        }
    }

    let mut cpu = String::new();
    #[cfg(target_arch = "x86_64")]
    {
        let cpuid = raw_cpuid::CpuId::new();
        if let Some(info) = cpuid.get_extended_function_info() {
            if let Some(brand) = info.processor_brand_string() {
                cpu = brand.trim().to_string();
            }
        }
    }

    let mut ram = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("memory:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                ram = format!(
                    "{} / {} ({}%)",
                    format_size(used),
                    format_size(size),
                    used * 100 / size
                );
            }
            let _ = syscall::close(fd);
        }
    }

    let mut disk = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("file:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                disk = format!(
                    "{} / {} ({}%)",
                    format_size(used),
                    format_size(size),
                    used * 100 / size
                );
            }
            let _ = syscall::close(fd);
        }
    }

    let left = [
        "              :+yMMMMy+:              ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "       sydMMMMo/sMMMMs/oMMMMdys       ",
        "    .oMMMdso              osdMMMo.    ",
        "   .+MMd/`   -:::::::`      `/dMM+.   ",
        "  +dMMN.     NMMNNNMMdyo`     .NMMd+  ",
        "  yMMN       NM+   oomMMN       NMMy  ",
        " hNMd.       NM+     `oMN       .dMNh ",
        " dMMh        NM+     `oMN        hMMd ",
        " dMMh        NM+-oooodMMN        hMMd ",
        " dMMh        NM+/MMMMdhs`        hMMd ",
        " hNMd.       NM+`/mMMm+         .dMNh ",
        "  yMMN       NM+   oNMd+        NMMy  ",
        "  +dMMN.     NM+    omMMm     .NMMd+  ",
        "   .+MMd/`   --.     .---   `/dMM+.   ",
        "    .oMMMdso              osdMMMo.    ",
        "       sydMMMMo////////oMMMMdys       ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "              :++++++++:              ",
    ];

    const S: &str = "\x1B[1;38;5;75m"; // blue start
    const E: &str = "\x1B[0m"; // end
    let right = [
        format!("{}{}{}@{}{}{}", S, user, E, S, hostname.trim(), E),
        format!("{}OS:         {}Redox OS {}", S, E, redox_release),
        format!("{}Kernel:     {}{}", S, E, kernel),
        format!("{}Uptime:     {}{}", S, E, uptime_str),
        format!("{}Shell:      {}{}", S, E, shell),
        format!("{}Resolution: {}{}x{}", S, E, width, height),
        format!("{}DE:         {}orbital", S, E),
        format!("{}WM:         {}orbital", S, E),
        format!("{}CPU:        {}{}", S, E, cpu),
        format!("{}RAM:        {}{}", S, E, ram),
        format!("{}Disk:       {}{}", S, E, disk),
    ];

    for (i, line) in left.iter().enumerate() {
        print!("\x1B[1;38;5;75m{}  \x1B[0m", line);
        if let Some(r) = right.get(i) {
            print!("{}", r);
        }
        println!();
    }
}
<|endoftext|>


<commit_before>use crate::{
    account::{Account, AccountReq},
    config::get_google_client_id,
    course2::{self, Course2, Course2Response, Course2SimilarityError},
    database::Database,
    minhash::{LshIndex, PermGen},
    routes::{
        courses,
        courses2::{
            self,
            thumbnail::{GetCourse2ThumbnailError, GetThumbnail2, Size2},
            PutCourses2Response,
        },
    },
    session::{AuthReq, AuthSession},
};

use brotli2::{read::BrotliEncoder, CompressParams};
use compression::prelude::*;
use image::{jpeg::JPEGEncoder, load_from_memory, DynamicImage, FilterType, ImageError};
use mongodb::{oid::ObjectId, spec::BinarySubtype, Bson};
use rayon::prelude::*;
use std::{
    io::Read,
    sync::{Arc, Mutex},
};

const SIMILARITY_THRESHOLD: f64 = 0.95;

pub struct Data {
    database: Arc<Database>,
    pub google_client_id: String,
    pub perm_gen: PermGen,
    pub lsh_index: Arc<Mutex<LshIndex>>,
}

pub type ServerData = Arc<Data>;

impl Data {
    pub fn new(database: Arc<Database>) -> Self {
        let mut lsh_index = LshIndex::new(8);
        println!("Filling LshIndex");
        database.fill_lsh_index(&mut lsh_index);
        println!("Filling LshIndex completed!");
        Data {
            database,
            google_client_id: get_google_client_id(),
            perm_gen: PermGen::new(128),
            lsh_index: Arc::new(Mutex::new(lsh_index)),
        }
    }

    pub fn get_courses(
        &self,
        query: courses::GetCourses,
    ) -> Result<String, courses::GetCoursesError> {
        match query.into_ordered_document(&self.database) {
            Ok(query) => Ok(self.database.get_courses(query)),
            Err(error) => Err(error),
        }
    }

    pub fn get_courses2(
        &self,
        query: courses2::GetCourses2,
    ) -> Result<String, courses2::GetCourses2Error> {
        let query = query.into_ordered_document(&self.database)?;
        let (courses, accounts) = self.database.get_courses2(query)?;

        let courses: Vec<Course2Response> = courses
            .into_iter()
            .map(|course| {
                let account = accounts
                    .iter()
                    .find(|account| {
                        account.get_id_ref().to_string() == course.get_owner().to_string()
                    })
                    .unwrap();
                Course2Response::from_course(course, account)
            })
            .collect();

        Ok(serde_json::to_string(&courses)?)
    }

    pub fn get_course2_thumbnail(
        &self,
        course_id: ObjectId,
        query: GetThumbnail2,
    ) -> Result<Vec<u8>, GetCourse2ThumbnailError> {
        let doc = doc! {
            "_id" => course_id.clone()
        };
        let size: String = query.size.clone().into();
        let projection = doc! {
            size.clone() => 1
        };
        let thumb = self.database.get_course2_thumbnail(doc, projection)?;
        if let Some(thumb) = thumb {
            match thumb.get_binary_generic(&size) {
                Ok(thumb) => Ok(thumb.clone()),
                Err(_) => {
                    if query.size == Size2::ORIGINAL {
                        Err(GetCourse2ThumbnailError::CourseNotFound(course_id))
                    } else {
                        let doc = doc! {
                            "_id" => course_id.clone()
                        };
                        let size_original: String = Size2::ORIGINAL.into();
                        let projection = doc! {
                            size_original.clone() => 1
                        };
                        let thumb = self
                            .database
                            .get_course2_thumbnail(doc, projection)?
                            .unwrap();
                        let thumb = thumb
                            .get_binary_generic(&size_original)
                            .expect(&format!(
                                "mongodb corrupted. thumbnail missing for course {}",
                                course_id
                            ))
                            .clone();

                        let image = load_from_memory(&thumb[..])?;
                        let (nwidth, nheight) = query.size.get_dimensions();
                        let image = image.resize_exact(nwidth, nheight, FilterType::Gaussian);
                        let color = image.color();

                        match image {
                            DynamicImage::ImageRgb8(buffer) => {
                                let (width, height) = buffer.dimensions();
                                let mut res = vec![];
                                let mut encoder = JPEGEncoder::new_with_quality(&mut res, 85);
                                encoder
                                    .encode(&buffer.into_raw()[..], width, height, color)
                                    .map_err(|e| ImageError::from(e))?;
                                self.database.update_course2_thumbnail(
                                    course_id,
                                    size,
                                    res.clone(),
                                )?;
                                Ok(res)
                            }
                            _ => Err(image::ImageError::FormatError(
                                "expected image rgb8".to_string(),
                            )
                            .into()),
                        }
                    }
                }
            }
        } else {
            Err(GetCourse2ThumbnailError::CourseNotFound(course_id))
        }
    }

    pub fn put_courses2(
        &self,
        mut courses: Vec<cemu_smm::Course2>,
        account: &Account,
        difficulty: Option<course2::Difficulty>,
    ) -> Result<PutCourses2Response, courses2::PutCourses2Error> {
        let lsh_index = self.lsh_index.clone();
        let response = Arc::new(Mutex::new(PutCourses2Response::new()));
        let succeeded: Vec<_> = courses
            .par_iter_mut()
            .map(
                |smm_course| -> Result<Course2Response, courses2::PutCourses2Error> {
                    let mut course = Course2::insert(
                        account.get_id_ref().clone(),
                        smm_course,
                        difficulty.clone(),
                        &self.perm_gen,
                    );
                    {
                        let mut lsh_index = lsh_index.lock().unwrap();
                        let query: Vec<Bson> = lsh_index
                            .query(course.get_hash())
                            .into_iter()
                            .map(|id| -> Bson { ObjectId::with_string(&id).unwrap().into() })
                            .collect();
                        let query = doc! {
                            "_id" => {
                                "$in" => query
                            }
                        };
                        let similar_courses = self.database.find_courses2(query)?;
                        for similar_course in similar_courses {
                            let jaccard = course.get_hash().jaccard(similar_course.get_hash());
                            if jaccard > SIMILARITY_THRESHOLD {
                                return Err(courses2::PutCourses2Error::Similarity(
                                    Course2SimilarityError::new(
                                        similar_course.get_id().to_hex(),
                                        similar_course
                                            .get_course()
                                            .get_header()
                                            .get_title()
                                            .to_string(),
                                        jaccard,
                                    ),
                                ));
                            }
                        }

                        let course_meta = serde_json::to_value(&course)?;

                        let data_gz = smm_course
                            .get_course_data()
                            .iter()
                            .cloned()
                            .encode(&mut GZipEncoder::new(), Action::Finish)
                            .collect::<Result<Vec<_>, _>>()?;
                        let data_gz = Bson::Binary(BinarySubtype::Generic, data_gz);

                        let mut data_br = vec![];
                        let mut params = CompressParams::new();
                        params.quality(11);
                        BrotliEncoder::from_params(&smm_course.get_course_data()[..], &params)
                            .read_to_end(&mut data_br)?;
                        let data_br = Bson::Binary(BinarySubtype::Generic, data_br);

                        let course_thumb = smm_course
                            .get_course_thumb_mut()
                            .ok_or(courses2::PutCourses2Error::ThumbnailMissing)?;
                        let thumb =
                            Bson::Binary(BinarySubtype::Generic, course_thumb.get_jpeg().to_vec());

                        if let Bson::Document(doc_meta) = Bson::from(course_meta) {
                            let inserted_id = self
                                .database
                                .put_course2(doc_meta, data_gz, data_br, thumb)?;
                            course.set_id(inserted_id.clone());
                            lsh_index.insert(course.get_id().to_hex(), course.get_hash());
                            let course = Course2Response::from_course(course, account);
                            Ok(course)
                        } else {
                            Err(mongodb::Error::DefaultError("".to_string()).into())
                        }
                    }
                },
            )
            .filter_map(|course| {
                if let Err(err) = course {
                    response.lock().unwrap().add_failed(err);
                    None
                } else {
                    Result::ok(course)
                }
            })
            .collect();
        let mut response = Arc::try_unwrap(response).unwrap().into_inner().unwrap();
        response.set_succeeded(succeeded);
        Ok(response)
    }

    pub fn delete_course2(&self, course_id: ObjectId) -> Result<(), mongodb::Error> {
        let query = doc! {
            "_id" => course_id
        };
        self.database.delete_course2(query)
    }

    pub fn add_or_get_account(
        &self,
        account: AccountReq,
        session: AuthSession,
    ) -> Result<Account, mongodb::Error> {
        match self.database.find_account(account.as_find()) {
            Some(account) => {
                self.database
                    .store_account_session(account.get_id_ref(), session)?;
                Ok(account)
            }
            None => self.database.add_account(account, session),
        }
    }

    pub fn get_account_from_auth(&self, auth_req: AuthReq) -> Option<Account> {
        self.database.find_account(auth_req.into())
    }
}
<commit_msg>clippy<commit_after>use crate::{
    account::{Account, AccountReq},
    config::get_google_client_id,
    course2::{self, Course2, Course2Response, Course2SimilarityError},
    database::Database,
    minhash::{LshIndex, PermGen},
    routes::{
        courses,
        courses2::{
            self,
            thumbnail::{GetCourse2ThumbnailError, GetThumbnail2, Size2},
            PutCourses2Response,
        },
    },
    session::{AuthReq, AuthSession},
};

use brotli2::{read::BrotliEncoder, CompressParams};
use compression::prelude::*;
use image::{jpeg::JPEGEncoder, load_from_memory, DynamicImage, FilterType, ImageError};
use mongodb::{oid::ObjectId, spec::BinarySubtype, Bson};
use rayon::prelude::*;
use std::{
    io::Read,
    sync::{Arc, Mutex},
};

const SIMILARITY_THRESHOLD: f64 = 0.95;

pub struct Data {
    database: Arc<Database>,
    pub google_client_id: String,
    pub perm_gen: PermGen,
    pub lsh_index: Arc<Mutex<LshIndex>>,
}

pub type ServerData = Arc<Data>;

impl Data {
    pub fn new(database: Arc<Database>) -> Self {
        let mut lsh_index = LshIndex::new(8);
        println!("Filling LshIndex");
        database.fill_lsh_index(&mut lsh_index);
        println!("Filling LshIndex completed!");
        Data {
            database,
            google_client_id: get_google_client_id(),
            perm_gen: PermGen::new(128),
            lsh_index: Arc::new(Mutex::new(lsh_index)),
        }
    }

    pub fn get_courses(
        &self,
        query: courses::GetCourses,
    ) -> Result<String, courses::GetCoursesError> {
        match query.into_ordered_document(&self.database) {
            Ok(query) => Ok(self.database.get_courses(query)),
            Err(error) => Err(error),
        }
    }

    pub fn get_courses2(
        &self,
        query: courses2::GetCourses2,
    ) -> Result<String, courses2::GetCourses2Error> {
        let query = query.into_ordered_document(&self.database)?;
        let (courses, accounts) = self.database.get_courses2(query)?;

        let courses: Vec<Course2Response> = courses
            .into_iter()
            .map(|course| {
                let account = accounts
                    .iter()
                    .find(|account| {
                        account.get_id_ref().to_string() == course.get_owner().to_string()
                    })
                    .unwrap();
                Course2Response::from_course(course, account)
            })
            .collect();

        Ok(serde_json::to_string(&courses)?)
    }

    pub fn get_course2_thumbnail(
        &self,
        course_id: ObjectId,
        query: GetThumbnail2,
    ) -> Result<Vec<u8>, GetCourse2ThumbnailError> {
        let doc = doc! {
            "_id" => course_id.clone()
        };
        let size: String = query.size.clone().into();
        let projection = doc! {
            size.clone() => 1
        };
        let thumb = self.database.get_course2_thumbnail(doc, projection)?;
        if let Some(thumb) = thumb {
            match thumb.get_binary_generic(&size) {
                Ok(thumb) => Ok(thumb.clone()),
                Err(_) => {
                    if query.size == Size2::ORIGINAL {
                        Err(GetCourse2ThumbnailError::CourseNotFound(course_id))
                    } else {
                        let doc = doc! {
                            "_id" => course_id.clone()
                        };
                        let size_original: String = Size2::ORIGINAL.into();
                        let projection = doc! {
                            size_original.clone() => 1
                        };
                        let thumb = self
                            .database
                            .get_course2_thumbnail(doc, projection)?
                            .unwrap();
                        let thumb = thumb
                            .get_binary_generic(&size_original)
                            .unwrap_or_else(|_| {
                                panic!(
                                    "mongodb corrupted. thumbnail missing for course {}",
                                    course_id
                                )
                            })
                            .clone();

                        let image = load_from_memory(&thumb[..])?;
                        let (nwidth, nheight) = query.size.get_dimensions();
                        let image = image.resize_exact(nwidth, nheight, FilterType::Gaussian);
                        let color = image.color();

                        match image {
                            DynamicImage::ImageRgb8(buffer) => {
                                let (width, height) = buffer.dimensions();
                                let mut res = vec![];
                                let mut encoder = JPEGEncoder::new_with_quality(&mut res, 85);
                                encoder
                                    .encode(&buffer.into_raw()[..], width, height, color)
                                    .map_err(ImageError::from)?;
                                self.database.update_course2_thumbnail(
                                    course_id,
                                    size,
                                    res.clone(),
                                )?;
                                Ok(res)
                            }
                            _ => Err(image::ImageError::FormatError(
                                "expected image rgb8".to_string(),
                            )
                            .into()),
                        }
                    }
                }
            }
        } else {
            Err(GetCourse2ThumbnailError::CourseNotFound(course_id))
        }
    }

    pub fn put_courses2(
        &self,
        mut courses: Vec<cemu_smm::Course2>,
        account: &Account,
        difficulty: Option<course2::Difficulty>,
    ) -> Result<PutCourses2Response, courses2::PutCourses2Error> {
        let lsh_index = self.lsh_index.clone();
        let response = Arc::new(Mutex::new(PutCourses2Response::new()));
        let succeeded: Vec<_> = courses
            .par_iter_mut()
            .map(
                |smm_course| -> Result<Course2Response, courses2::PutCourses2Error> {
                    let mut course = Course2::insert(
                        account.get_id_ref().clone(),
                        smm_course,
                        difficulty.clone(),
                        &self.perm_gen,
                    );
                    {
                        let mut lsh_index = lsh_index.lock().unwrap();
                        let query: Vec<Bson> = lsh_index
                            .query(course.get_hash())
                            .into_iter()
                            .map(|id| -> Bson { ObjectId::with_string(&id).unwrap().into() })
                            .collect();
                        let query = doc! {
                            "_id" => {
                                "$in" => query
                            }
                        };
                        let similar_courses = self.database.find_courses2(query)?;
                        for similar_course in similar_courses {
                            let jaccard = course.get_hash().jaccard(similar_course.get_hash());
                            if jaccard > SIMILARITY_THRESHOLD {
                                return Err(courses2::PutCourses2Error::Similarity(
                                    Course2SimilarityError::new(
                                        similar_course.get_id().to_hex(),
                                        similar_course
                                            .get_course()
                                            .get_header()
                                            .get_title()
                                            .to_string(),
                                        jaccard,
                                    ),
                                ));
                            }
                        }

                        let course_meta = serde_json::to_value(&course)?;

                        let data_gz = smm_course
                            .get_course_data()
                            .iter()
                            .cloned()
                            .encode(&mut GZipEncoder::new(), Action::Finish)
                            .collect::<Result<Vec<_>, _>>()?;
                        let data_gz = Bson::Binary(BinarySubtype::Generic, data_gz);

                        let mut data_br = vec![];
                        let mut params = CompressParams::new();
                        params.quality(11);
                        BrotliEncoder::from_params(&smm_course.get_course_data()[..], &params)
                            .read_to_end(&mut data_br)?;
                        let data_br = Bson::Binary(BinarySubtype::Generic, data_br);

                        let course_thumb = smm_course
                            .get_course_thumb_mut()
                            .ok_or(courses2::PutCourses2Error::ThumbnailMissing)?;
                        let thumb =
                            Bson::Binary(BinarySubtype::Generic, course_thumb.get_jpeg().to_vec());

                        if let Bson::Document(doc_meta) = Bson::from(course_meta) {
                            let inserted_id = self
                                .database
                                .put_course2(doc_meta, data_gz, data_br, thumb)?;
                            course.set_id(inserted_id.clone());
                            lsh_index.insert(course.get_id().to_hex(), course.get_hash());
                            let course = Course2Response::from_course(course, account);
                            Ok(course)
                        } else {
                            Err(mongodb::Error::DefaultError("".to_string()).into())
                        }
                    }
                },
            )
            .filter_map(|course| {
                if let Err(err) = course {
                    response.lock().unwrap().add_failed(err);
                    None
                } else {
                    Result::ok(course)
                }
            })
            .collect();
        let mut response = Arc::try_unwrap(response).unwrap().into_inner().unwrap();
        response.set_succeeded(succeeded);
        Ok(response)
    }

    pub fn delete_course2(&self, course_id: ObjectId) -> Result<(), mongodb::Error> {
        let query = doc! {
            "_id" => course_id
        };
        self.database.delete_course2(query)
    }

    pub fn add_or_get_account(
        &self,
        account: AccountReq,
        session: AuthSession,
    ) -> Result<Account, mongodb::Error> {
        match self.database.find_account(account.as_find()) {
            Some(account) => {
                self.database
                    .store_account_session(account.get_id_ref(), session)?;
                Ok(account)
            }
            None => self.database.add_account(account, session),
        }
    }

    pub fn get_account_from_auth(&self, auth_req: AuthReq) -> Option<Account> {
        self.database.find_account(auth_req.into())
    }
}
<|endoftext|>