<commit_before>use std::collections::HashSet;
use super::Nearness;
use {RowId, ColumnId, PredictorId, Column, Accessor, MomentaryAccessor, Mutator,
     TimeStewardSettings};
use std::marker::PhantomData;
use serde::Serialize;

#[derive (Copy, Clone, PartialEq, Eq, Serialize, Deserialize, Debug)]
pub struct Bounds {
  pub min: [i64; 2],
  pub max: [i64; 2],
}

pub trait Basics: super::Basics {
  fn get_bounds<A: MomentaryAccessor<<Self as super::Basics>::StewardBasics>>(accessor: &A,
                                                                              who: RowId, detector: <Self as super::Basics>::DetectorId)
                                                                              -> Bounds;
  fn when_escapes <A: Accessor <<Self as super::Basics>::StewardBasics>>(accessor: & A, who: RowId, bounds: Bounds, detector: <Self as super::Basics>::DetectorId)->Option <<<Self as super::Basics>::StewardBasics as ::Basics>::Time>;
}

struct Cell<B: Basics> {
  _marker: PhantomData<B>,
}
impl<B: Basics> Column for Cell<B> {
  type FieldType = HashSet <RowId>;
  fn column_id() -> ColumnId {
    ColumnId(B::nearness_column_id().0 ^ 0x030fc8868af34f2c)
  }
}

// It might be nicer for this not to be public – the user isn't expected to interact with it in any way. However, it has to be public because Columns has to be.
#[derive (Clone, PartialEq, Eq, Serialize, Deserialize, Debug)]
pub struct Member<B: Basics> {
  row: RowId,
  bounds: Bounds,
  detector: B::DetectorId,
}
impl<B: Basics> Column for Member<B> {
  type FieldType = Member <B>;
  fn column_id() -> ColumnId {
    ColumnId(B::nearness_column_id().0 ^ 0xe844c90eaf2e20ab)
  }
}

fn cell_row<DetectorId: Serialize>(x: i64, y: i64, me: DetectorId) -> RowId {
  RowId::new(&(x, y, me))
}
fn get_bounds<B: Basics, M: Mutator<B::StewardBasics>>(mutator: &mut M,
                                                       who: RowId,
                                                       me: B::DetectorId)
                                                       -> Bounds {
  B::get_bounds(mutator, who, me)
}
pub fn insert<B: Basics, M: Mutator<B::StewardBasics>>(mutator: &mut M,
                                                       who: RowId,
                                                       me: B::DetectorId) {
  let bounds = get_bounds::<B, M>(mutator, who, me);
  mutator.set::<Member<B>>(RowId::new(&(who, me)),
                           Some(Member {
                             row: who,
                             bounds: bounds,
                             detector: me,
                           }));
  for next in bounds.min[0]..bounds.max[0] + 1 {
    for what in bounds.min[1]..bounds.max[1] + 1 {
      let row = cell_row(next, what, me);
      let mut members = mutator.get::<Cell<B>>(row).cloned().unwrap_or(HashSet::new());
      for member in members.iter() {
        let (id, contents) = Nearness::new(who, member.clone(), me);
        // check first: no need to update last_change
        if mutator.get::<Nearness<B>>(id).is_none() {
          mutator.set::<Nearness<B>>(id, Some(contents));
        }
      }
      members.insert(who);
      mutator.set::<Cell<B>>(row, Some(members));
    }
  }
}

pub fn erase<B: Basics, M: Mutator<B::StewardBasics>>(mutator: &mut M,
                                                      who: RowId,
                                                      me: B::DetectorId,
                                                      membership_id: RowId,
                                                      bounds: Bounds) {

  mutator.set::<Member<B>>(membership_id, None);
  for next in bounds.min[0]..bounds.max[0] + 1 {
    for what in bounds.min[1]..bounds.max[1] + 1 {
      let row = cell_row(next, what, me);
      let mut members = mutator.get::<Cell<B>>(row).cloned().unwrap_or(HashSet::new());
      for member in members.iter() {
        if *member != who {
          let (id, _) = Nearness::<B>::new(who, member.clone(), me);
          mutator.set::<Nearness<B>>(id, None);
        }
      }
      members.remove(&who);
      mutator.set::<Cell<B>>(row,
                             if members.is_empty() {
                               None
                             } else {
                               Some(members)
                             });
    }
  }
}

pub type Columns <B> = PhantomData <Member <B>>;

pub fn insert_predictors <B: Basics, Settings: TimeStewardSettings <<B as super::Basics>::StewardBasics>> (settings: &mut Settings) {
  settings.insert_predictor (PredictorId (<B as super::Basics>::nearness_column_id().0 ^ 0x3686689daa651bf3),
    Member::<B>::column_id(),
    time_steward_predictor! (<B as super::Basics>::StewardBasics, struct BoundsChangePredictor [B: Basics]=[B] {p: PhantomData<B> = PhantomData}, | &self, accessor, id | {
      let member;
      {
        let member_reference = accessor.get::<Member <B>> (id).expect ("row is missing the field the predictor triggered on");
        member = (*member_reference).clone();
      }
      if let Some (time) = B::when_escapes (accessor, member.row, member.bounds, member.detector) {
        accessor.predict_at_time (time, time_steward_event! (<B as super::Basics>::StewardBasics, struct BoundsChange [B: Basics]=[B] {id: RowId = id, member: Member <B> = member}, | &self, mutator | {
          //TODO: optimize erase-then-insert
          erase::<B,_> (mutator, self.member.row, self.member.detector, self.id, self.member.bounds);
          insert::<B,_> (mutator, self.member.row, self.member.detector);
        }));
      }
    })
  )
}
<commit_msg>Concealed another thing not meant to be public.<commit_after>use std::collections::HashSet;
use super::Nearness;
use {RowId, ColumnId, PredictorId, Column, Accessor, MomentaryAccessor, Mutator,
     TimeStewardSettings};
use std::marker::PhantomData;
use serde::Serialize;

#[derive (Copy, Clone, PartialEq, Eq, Serialize, Deserialize, Debug)]
pub struct Bounds {
  pub min: [i64; 2],
  pub max: [i64; 2],
}

pub trait Basics: super::Basics {
  fn get_bounds<A: MomentaryAccessor<<Self as super::Basics>::StewardBasics>>(accessor: &A,
                                                                              who: RowId, detector: <Self as super::Basics>::DetectorId)
                                                                              -> Bounds;
  fn when_escapes <A: Accessor <<Self as super::Basics>::StewardBasics>>(accessor: & A, who: RowId, bounds: Bounds, detector: <Self as super::Basics>::DetectorId)->Option <<<Self as super::Basics>::StewardBasics as ::Basics>::Time>;
}

struct Cell<B: Basics> {
  _marker: PhantomData<B>,
}
impl<B: Basics> Column for Cell<B> {
  type FieldType = HashSet <RowId>;
  fn column_id() -> ColumnId {
    ColumnId(B::nearness_column_id().0 ^ 0x030fc8868af34f2c)
  }
}

// Work around private-in-public rules so that our column types
// (which are implementation details) are not visible elsewhere.
mod hack {
  use super::*;
  use {RowId};
  #[derive (Clone, PartialEq, Eq, Serialize, Deserialize, Debug)]
  pub struct Member<B: Basics> {
    pub row: RowId,
    pub bounds: Bounds,
    pub detector: B::DetectorId,
  }
}
use self::hack::Member;

impl<B: Basics> Column for Member<B> {
  type FieldType = Member <B>;
  fn column_id() -> ColumnId {
    ColumnId(B::nearness_column_id().0 ^ 0xe844c90eaf2e20ab)
  }
}

fn cell_row<DetectorId: Serialize>(x: i64, y: i64, me: DetectorId) -> RowId {
  RowId::new(&(x, y, me))
}
fn get_bounds<B: Basics, M: Mutator<B::StewardBasics>>(mutator: &mut M,
                                                       who: RowId,
                                                       me: B::DetectorId)
                                                       -> Bounds {
  B::get_bounds(mutator, who, me)
}
pub fn insert<B: Basics, M: Mutator<B::StewardBasics>>(mutator: &mut M,
                                                       who: RowId,
                                                       me: B::DetectorId) {
  let bounds = get_bounds::<B, M>(mutator, who, me);
  mutator.set::<Member<B>>(RowId::new(&(who, me)),
                           Some(Member {
                             row: who,
                             bounds: bounds,
                             detector: me,
                           }));
  for next in bounds.min[0]..bounds.max[0] + 1 {
    for what in bounds.min[1]..bounds.max[1] + 1 {
      let row = cell_row(next, what, me);
      let mut members = mutator.get::<Cell<B>>(row).cloned().unwrap_or(HashSet::new());
      for member in members.iter() {
        let (id, contents) = Nearness::new(who, member.clone(), me);
        // check first: no need to update last_change
        if mutator.get::<Nearness<B>>(id).is_none() {
          mutator.set::<Nearness<B>>(id, Some(contents));
        }
      }
      members.insert(who);
      mutator.set::<Cell<B>>(row, Some(members));
    }
  }
}

pub fn erase<B: Basics, M: Mutator<B::StewardBasics>>(mutator: &mut M,
                                                      who: RowId,
                                                      me: B::DetectorId,
                                                      membership_id: RowId,
                                                      bounds: Bounds) {

  mutator.set::<Member<B>>(membership_id, None);
  for next in bounds.min[0]..bounds.max[0] + 1 {
    for what in bounds.min[1]..bounds.max[1] + 1 {
      let row = cell_row(next, what, me);
      let mut members = mutator.get::<Cell<B>>(row).cloned().unwrap_or(HashSet::new());
      for member in members.iter() {
        if *member != who {
          let (id, _) = Nearness::<B>::new(who, member.clone(), me);
          mutator.set::<Nearness<B>>(id, None);
        }
      }
      members.remove(&who);
      mutator.set::<Cell<B>>(row,
                             if members.is_empty() {
                               None
                             } else {
                               Some(members)
                             });
    }
  }
}

pub type Columns <B> = PhantomData <Member <B>>;

pub fn insert_predictors <B: Basics, Settings: TimeStewardSettings <<B as super::Basics>::StewardBasics>> (settings: &mut Settings) {
  settings.insert_predictor (PredictorId (<B as super::Basics>::nearness_column_id().0 ^ 0x3686689daa651bf3),
    Member::<B>::column_id(),
    time_steward_predictor! (<B as super::Basics>::StewardBasics, struct BoundsChangePredictor [B: Basics]=[B] {p: PhantomData<B> = PhantomData}, | &self, accessor, id | {
      let member;
      {
        let member_reference = accessor.get::<Member <B>> (id).expect ("row is missing the field the predictor triggered on");
        member = (*member_reference).clone();
      }
      if let Some (time) = B::when_escapes (accessor, member.row, member.bounds, member.detector) {
        accessor.predict_at_time (time, time_steward_event! (<B as super::Basics>::StewardBasics, struct BoundsChange [B: Basics]=[B] {id: RowId = id, member: Member <B> = member}, | &self, mutator | {
          //TODO: optimize erase-then-insert
          erase::<B,_> (mutator, self.member.row, self.member.detector, self.id, self.member.bounds);
          insert::<B,_> (mutator, self.member.row, self.member.detector);
        }));
      }
    })
  )
}
<|endoftext|>


<commit_before>use std::cmp::Reverse;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, Error, ErrorKind, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::SystemTime;

use priority_queue::PriorityQueue;

use crate::authentication::Credentials;
use crate::spotify_id::FileId;

/// Some kind of data structure that holds some paths, the size of these files and a timestamp.
/// It keeps track of the file sizes and is able to pop the path with the oldest timestamp if
/// a given limit is exceeded.
struct SizeLimiter {
    queue: PriorityQueue<PathBuf, Reverse<SystemTime>>,
    sizes: HashMap<PathBuf, u64>,
    size_limit: u64,
    in_use: u64,
}

impl SizeLimiter {
    /// Creates a new instance with the given size limit.
    fn new(limit: u64) -> Self {
        Self {
            queue: PriorityQueue::new(),
            sizes: HashMap::new(),
            size_limit: limit,
            in_use: 0,
        }
    }

    /// Adds an entry to this data structure.
    ///
    /// If this file is already contained, it will be updated accordingly.
    fn add(&mut self, file: &Path, size: u64, accessed: SystemTime) {
        self.in_use += size;
        self.queue.push(file.to_owned(), Reverse(accessed));
        if let Some(old_size) = self.sizes.insert(file.to_owned(), size) {
            // It's important that decreasing happens after
            // increasing the size, to prevent an overflow.
            self.in_use -= old_size;
        }
    }

    /// Returns true if the limit is exceeded.
    fn exceeds_limit(&self) -> bool {
        self.in_use > self.size_limit
    }

    /// Returns the least recently accessed file if the size of the cache exceeds
    /// the limit.
    ///
    /// The entry is removed from the data structure, but the caller is responsible
    /// to delete the file in the file system.
    fn pop(&mut self) -> Option<PathBuf> {
        if self.exceeds_limit() {
            let (next, _) = self
                .queue
                .pop()
                .expect("in_use was > 0, so the queue should have contained an item.");
            let size = self
                .sizes
                .remove(&next)
                .expect("`queue` and `sizes` should have the same keys.");
            self.in_use -= size;
            Some(next)
        } else {
            None
        }
    }

    /// Updates the timestamp of an existing element. Returns `true` if the item did exist.
    fn update(&mut self, file: &Path, access_time: SystemTime) -> bool {
        self.queue
            .change_priority(file, Reverse(access_time))
            .is_some()
    }

    /// Removes an element with the specified path. Returns `true` if the item did exist.
    fn remove(&mut self, file: &Path) -> bool {
        if self.queue.remove(file).is_none() {
            return false;
        }

        let size = self
            .sizes
            .remove(file)
            .expect("`queue` and `sizes` should have the same keys.");
        self.in_use -= size;

        true
    }
}

struct FsSizeLimiter {
    limiter: Mutex<SizeLimiter>,
}

impl FsSizeLimiter {
    /// Returns access time and file size of a given path.
    fn get_metadata(file: &Path) -> io::Result<(SystemTime, u64)> {
        let metadata = file.metadata()?;

        // The first of the following timestamps which is available will be chosen as access time:
        // 1. Access time
        // 2. Modification time
        // 3. Creation time
        // 4. Current time
        let access_time = metadata
            .accessed()
            .or_else(|_| metadata.modified())
            .or_else(|_| metadata.created())
            .unwrap_or_else(|_| SystemTime::now());

        let size = metadata.len();

        Ok((access_time, size))
    }

    /// Recursively search a directory for files and add them to the `limiter` struct.
    fn init_dir(limiter: &mut SizeLimiter, path: &Path) {
        let list_dir = match fs::read_dir(path) {
            Ok(list_dir) => list_dir,
            Err(e) => {
                warn!("Could not read directory {:?} in cache dir: {}", path, e);
                return;
            }
        };

        for entry in list_dir {
            let entry = match entry {
                Ok(entry) => entry,
                Err(e) => {
                    warn!("Could not directory {:?} in cache dir: {}", path, e);
                    return;
                }
            };

            match entry.file_type() {
                Ok(file_type) if file_type.is_dir() || file_type.is_symlink() => {
                    Self::init_dir(limiter, &entry.path())
                }
                Ok(file_type) if file_type.is_file() => {
                    let path = entry.path();
                    match Self::get_metadata(&path) {
                        Ok((access_time, size)) => {
                            limiter.add(&path, size, access_time);
                        }
                        Err(e) => {
                            warn!("Could not read file {:?} in cache dir: {}", path, e)
                        }
                    }
                }
                Ok(ft) => {
                    warn!(
                        "File {:?} in cache dir has unsupported type {:?}",
                        entry.path(),
                        ft
                    )
                }
                Err(e) => {
                    warn!(
                        "Could not get type of file {:?} in cache dir: {}",
                        entry.path(),
                        e
                    )
                }
            };
        }
    }

    fn add(&self, file: &Path, size: u64) {
        self.limiter
            .lock()
            .unwrap()
            .add(file, size, SystemTime::now());
    }

    fn touch(&self, file: &Path) -> bool {
        self.limiter.lock().unwrap().update(file, SystemTime::now())
    }

    fn remove(&self, file: &Path) {
        self.limiter.lock().unwrap().remove(file);
    }

    fn prune_internal<F: FnMut() -> Option<PathBuf>>(mut pop: F) {
        let mut first = true;
        let mut count = 0;

        while let Some(file) = pop() {
            if first {
                debug!("Cache dir exceeds limit, removing least recently used files.");
                first = false;
            }

            if let Err(e) = fs::remove_file(&file) {
                warn!("Could not remove file {:?} from cache dir: {}", file, e);
            } else {
                count += 1;
            }
        }

        if count > 0 {
            info!("Removed {} cache files.", count);
        }
    }

    fn prune(&self) {
        Self::prune_internal(|| self.limiter.lock().unwrap().pop())
    }

    fn new(path: &Path, limit: u64) -> Self {
        let mut limiter = SizeLimiter::new(limit);

        Self::init_dir(&mut limiter, path);
        Self::prune_internal(|| limiter.pop());

        Self {
            limiter: Mutex::new(limiter),
        }
    }
}

/// A cache for volume, credentials and audio files.
#[derive(Clone)]
pub struct Cache {
    credentials_location: Option<PathBuf>,
    volume_location: Option<PathBuf>,
    audio_location: Option<PathBuf>,
    size_limiter: Option<Arc<FsSizeLimiter>>,
}

pub struct RemoveFileError(());

impl Cache {
    pub fn new<P: AsRef<Path>>(
        credentials: Option<P>,
        volume: Option<P>,
        audio: Option<P>,
        size_limit: Option<u64>,
    ) -> io::Result<Self> {
        let mut size_limiter = None;

        if let Some(location) = &credentials {
            fs::create_dir_all(location)?;
        }

        let credentials_location = credentials
            .as_ref()
            .map(|p| p.as_ref().join("credentials.json"));

        if let Some(location) = &volume {
            fs::create_dir_all(location)?;
        }

        let volume_location = volume.as_ref().map(|p| p.as_ref().join("volume"));

        if let Some(location) = &audio {
            fs::create_dir_all(location)?;

            if let Some(limit) = size_limit {
                let limiter = FsSizeLimiter::new(location.as_ref(), limit);

                size_limiter = Some(Arc::new(limiter));
            }
        }

        let audio_location = audio.map(|p| p.as_ref().to_owned());

        let cache = Cache {
            credentials_location,
            volume_location,
            audio_location,
            size_limiter,
        };

        Ok(cache)
    }

    pub fn credentials(&self) -> Option<Credentials> {
        let location = self.credentials_location.as_ref()?;

        // This closure is just convencience to enable the question mark operator
        let read = || {
            let mut file = File::open(location)?;
            let mut contents = String::new();
            file.read_to_string(&mut contents)?;
            serde_json::from_str(&contents).map_err(|e| Error::new(ErrorKind::InvalidData, e))
        };

        match read() {
            Ok(c) => Some(c),
            Err(e) => {
                // If the file did not exist, the file was probably not written
                // before. Otherwise, log the error.
                if e.kind() != ErrorKind::NotFound {
                    warn!("Error reading credentials from cache: {}", e);
                }
                None
            }
        }
    }

    pub fn save_credentials(&self, cred: &Credentials) {
        if let Some(location) = &self.credentials_location {
            let result = File::create(location).and_then(|mut file| {
                let data = serde_json::to_string(cred)?;
                write!(file, "{}", data)
            });

            if let Err(e) = result {
                warn!("Cannot save credentials to cache: {}", e)
            }
        }
    }

    pub fn volume(&self) -> Option<u16> {
        let location = self.volume_location.as_ref()?;

        let read = || {
            let mut file = File::open(location)?;
            let mut contents = String::new();
            file.read_to_string(&mut contents)?;
            contents
                .parse()
                .map_err(|e| Error::new(ErrorKind::InvalidData, e))
        };

        match read() {
            Ok(v) => Some(v),
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    warn!("Error reading volume from cache: {}", e);
                }
                None
            }
        }
    }

    pub fn save_volume(&self, volume: u16) {
        if let Some(ref location) = self.volume_location {
            let result = File::create(location).and_then(|mut file| write!(file, "{}", volume));
            if let Err(e) = result {
                warn!("Cannot save volume to cache: {}", e);
            }
        }
    }

    fn file_path(&self, file: FileId) -> Option<PathBuf> {
        self.audio_location.as_ref().map(|location| {
            let name = file.to_base16();
            let mut path = location.join(&name[0..2]);
            path.push(&name[2..]);
            path
        })
    }

    pub fn file(&self, file: FileId) -> Option<File> {
        let path = self.file_path(file)?;
        match File::open(&path) {
            Ok(file) => {
                if let Some(limiter) = self.size_limiter.as_deref() {
                    limiter.touch(&path);
                }
                Some(file)
            }
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    warn!("Error reading file from cache: {}", e)
                }
                None
            }
        }
    }

    pub fn save_file<F: Read>(&self, file: FileId, contents: &mut F) {
        let path = if let Some(path) = self.file_path(file) {
            path
        } else {
            return;
        };
        let parent = path.parent().unwrap();

        let result = fs::create_dir_all(parent)
            .and_then(|_| File::create(&path))
            .and_then(|mut file| io::copy(contents, &mut file));

        if let Ok(size) = result {
            if let Some(limiter) = self.size_limiter.as_deref() {
                limiter.add(&path, size);
                limiter.prune();
            }
        }
    }

    pub fn remove_file(&self, file: FileId) -> Result<(), RemoveFileError> {
        let path = self.file_path(file).ok_or(RemoveFileError(()))?;

        if let Err(err) = fs::remove_file(&path) {
            warn!("Unable to remove file from cache: {}", err);
            Err(RemoveFileError(()))
        } else {
            if let Some(limiter) = self.size_limiter.as_deref() {
                limiter.remove(&path);
            }
            Ok(())
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::time::Duration;

    fn ordered_time(v: u64) -> SystemTime {
        SystemTime::UNIX_EPOCH + Duration::from_secs(v)
    }

    #[test]
    fn test_size_limiter() {
        let mut limiter = SizeLimiter::new(1000);

        limiter.add(Path::new("a"), 500, ordered_time(2));
        limiter.add(Path::new("b"), 500, ordered_time(1));

        // b (500) -> a (500)  => sum: 1000 <= 1000
        assert!(!limiter.exceeds_limit());
        assert_eq!(limiter.pop(), None);

        limiter.add(Path::new("c"), 1000, ordered_time(3));

        // b (500) -> a (500) -> c (1000)  => sum: 2000 > 1000
        assert!(limiter.exceeds_limit());
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("b")));
        // a (500) -> c (1000)  => sum: 1500 > 1000
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("a")));
        // c (1000)   => sum: 1000 <= 1000
        assert_eq!(limiter.pop().as_deref(), None);

        limiter.add(Path::new("d"), 5, ordered_time(2));
        // d (5) -> c (1000) => sum: 1005 > 1000
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("d")));
        // c (1000)   => sum: 1000 <= 1000
        assert_eq!(limiter.pop().as_deref(), None);

        // Test updating

        limiter.add(Path::new("e"), 500, ordered_time(3));
        //  c (1000) -> e (500)  => sum: 1500 > 1000
        assert!(limiter.update(Path::new("c"), ordered_time(4)));
        // e (500) -> c (1000)  => sum: 1500 > 1000
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("e")));
        // c (1000)  => sum: 1000 <= 1000

        // Test removing
        limiter.add(Path::new("f"), 500, ordered_time(2));
        assert!(limiter.remove(Path::new("c")));
        assert!(!limiter.exceeds_limit());
    }
}
<commit_msg>Address review change request<commit_after>use std::cmp::Reverse;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, Error, ErrorKind, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::SystemTime;

use priority_queue::PriorityQueue;

use crate::authentication::Credentials;
use crate::spotify_id::FileId;

/// Some kind of data structure that holds some paths, the size of these files and a timestamp.
/// It keeps track of the file sizes and is able to pop the path with the oldest timestamp if
/// a given limit is exceeded.
struct SizeLimiter {
    queue: PriorityQueue<PathBuf, Reverse<SystemTime>>,
    sizes: HashMap<PathBuf, u64>,
    size_limit: u64,
    in_use: u64,
}

impl SizeLimiter {
    /// Creates a new instance with the given size limit.
    fn new(limit: u64) -> Self {
        Self {
            queue: PriorityQueue::new(),
            sizes: HashMap::new(),
            size_limit: limit,
            in_use: 0,
        }
    }

    /// Adds an entry to this data structure.
    ///
    /// If this file is already contained, it will be updated accordingly.
    fn add(&mut self, file: &Path, size: u64, accessed: SystemTime) {
        self.in_use += size;
        self.queue.push(file.to_owned(), Reverse(accessed));
        if let Some(old_size) = self.sizes.insert(file.to_owned(), size) {
            // It's important that decreasing happens after
            // increasing the size, to prevent an overflow.
            self.in_use -= old_size;
        }
    }

    /// Returns true if the limit is exceeded.
    fn exceeds_limit(&self) -> bool {
        self.in_use > self.size_limit
    }

    /// Returns the least recently accessed file if the size of the cache exceeds
    /// the limit.
    ///
    /// The entry is removed from the data structure, but the caller is responsible
    /// to delete the file in the file system.
    fn pop(&mut self) -> Option<PathBuf> {
        if self.exceeds_limit() {
            let (next, _) = self
                .queue
                .pop()
                .expect("in_use was > 0, so the queue should have contained an item.");
            let size = self
                .sizes
                .remove(&next)
                .expect("`queue` and `sizes` should have the same keys.");
            self.in_use -= size;
            Some(next)
        } else {
            None
        }
    }

    /// Updates the timestamp of an existing element. Returns `true` if the item did exist.
    fn update(&mut self, file: &Path, access_time: SystemTime) -> bool {
        self.queue
            .change_priority(file, Reverse(access_time))
            .is_some()
    }

    /// Removes an element with the specified path. Returns `true` if the item did exist.
    fn remove(&mut self, file: &Path) -> bool {
        if self.queue.remove(file).is_none() {
            return false;
        }

        let size = self
            .sizes
            .remove(file)
            .expect("`queue` and `sizes` should have the same keys.");
        self.in_use -= size;

        true
    }
}

struct FsSizeLimiter {
    limiter: Mutex<SizeLimiter>,
}

impl FsSizeLimiter {
    /// Returns access time and file size of a given path.
    fn get_metadata(file: &Path) -> io::Result<(SystemTime, u64)> {
        let metadata = file.metadata()?;

        // The first of the following timestamps which is available will be chosen as access time:
        // 1. Access time
        // 2. Modification time
        // 3. Creation time
        // 4. Current time
        let access_time = metadata
            .accessed()
            .or_else(|_| metadata.modified())
            .or_else(|_| metadata.created())
            .unwrap_or_else(|_| SystemTime::now());

        let size = metadata.len();

        Ok((access_time, size))
    }

    /// Recursively search a directory for files and add them to the `limiter` struct.
    fn init_dir(limiter: &mut SizeLimiter, path: &Path) {
        let list_dir = match fs::read_dir(path) {
            Ok(list_dir) => list_dir,
            Err(e) => {
                warn!("Could not read directory {:?} in cache dir: {}", path, e);
                return;
            }
        };

        for entry in list_dir {
            let entry = match entry {
                Ok(entry) => entry,
                Err(e) => {
                    warn!("Could not directory {:?} in cache dir: {}", path, e);
                    return;
                }
            };

            match entry.file_type() {
                Ok(file_type) if file_type.is_dir() || file_type.is_symlink() => {
                    Self::init_dir(limiter, &entry.path())
                }
                Ok(file_type) if file_type.is_file() => {
                    let path = entry.path();
                    match Self::get_metadata(&path) {
                        Ok((access_time, size)) => {
                            limiter.add(&path, size, access_time);
                        }
                        Err(e) => {
                            warn!("Could not read file {:?} in cache dir: {}", path, e)
                        }
                    }
                }
                Ok(ft) => {
                    warn!(
                        "File {:?} in cache dir has unsupported type {:?}",
                        entry.path(),
                        ft
                    )
                }
                Err(e) => {
                    warn!(
                        "Could not get type of file {:?} in cache dir: {}",
                        entry.path(),
                        e
                    )
                }
            };
        }
    }

    fn add(&self, file: &Path, size: u64) {
        self.limiter
            .lock()
            .unwrap()
            .add(file, size, SystemTime::now());
    }

    fn touch(&self, file: &Path) -> bool {
        self.limiter.lock().unwrap().update(file, SystemTime::now())
    }

    fn remove(&self, file: &Path) {
        self.limiter.lock().unwrap().remove(file);
    }

    fn prune_internal<F: FnMut() -> Option<PathBuf>>(mut pop: F) {
        let mut first = true;
        let mut count = 0;

        while let Some(file) = pop() {
            if first {
                debug!("Cache dir exceeds limit, removing least recently used files.");
                first = false;
            }

            if let Err(e) = fs::remove_file(&file) {
                warn!("Could not remove file {:?} from cache dir: {}", file, e);
            } else {
                count += 1;
            }
        }

        if count > 0 {
            info!("Removed {} cache files.", count);
        }
    }

    fn prune(&self) {
        Self::prune_internal(|| self.limiter.lock().unwrap().pop())
    }

    fn new(path: &Path, limit: u64) -> Self {
        let mut limiter = SizeLimiter::new(limit);

        Self::init_dir(&mut limiter, path);
        Self::prune_internal(|| limiter.pop());

        Self {
            limiter: Mutex::new(limiter),
        }
    }
}

/// A cache for volume, credentials and audio files.
#[derive(Clone)]
pub struct Cache {
    credentials_location: Option<PathBuf>,
    volume_location: Option<PathBuf>,
    audio_location: Option<PathBuf>,
    size_limiter: Option<Arc<FsSizeLimiter>>,
}

pub struct RemoveFileError(());

impl Cache {
    pub fn new<P: AsRef<Path>>(
        credentials_path: Option<P>,
        volume_path: Option<P>,
        audio_path: Option<P>,
        size_limit: Option<u64>,
    ) -> io::Result<Self> {
        let mut size_limiter = None;

        if let Some(location) = &credentials_path {
            fs::create_dir_all(location)?;
        }

        let credentials_location = credentials_path
            .as_ref()
            .map(|p| p.as_ref().join("credentials.json"));

        if let Some(location) = &volume_path {
            fs::create_dir_all(location)?;
        }

        let volume_location = volume_path.as_ref().map(|p| p.as_ref().join("volume"));

        if let Some(location) = &audio_path {
            fs::create_dir_all(location)?;

            if let Some(limit) = size_limit {
                let limiter = FsSizeLimiter::new(location.as_ref(), limit);

                size_limiter = Some(Arc::new(limiter));
            }
        }

        let audio_location = audio_path.map(|p| p.as_ref().to_owned());

        let cache = Cache {
            credentials_location,
            volume_location,
            audio_location,
            size_limiter,
        };

        Ok(cache)
    }

    pub fn credentials(&self) -> Option<Credentials> {
        let location = self.credentials_location.as_ref()?;

        // This closure is just convencience to enable the question mark operator
        let read = || {
            let mut file = File::open(location)?;
            let mut contents = String::new();
            file.read_to_string(&mut contents)?;
            serde_json::from_str(&contents).map_err(|e| Error::new(ErrorKind::InvalidData, e))
        };

        match read() {
            Ok(c) => Some(c),
            Err(e) => {
                // If the file did not exist, the file was probably not written
                // before. Otherwise, log the error.
                if e.kind() != ErrorKind::NotFound {
                    warn!("Error reading credentials from cache: {}", e);
                }
                None
            }
        }
    }

    pub fn save_credentials(&self, cred: &Credentials) {
        if let Some(location) = &self.credentials_location {
            let result = File::create(location).and_then(|mut file| {
                let data = serde_json::to_string(cred)?;
                write!(file, "{}", data)
            });

            if let Err(e) = result {
                warn!("Cannot save credentials to cache: {}", e)
            }
        }
    }

    pub fn volume(&self) -> Option<u16> {
        let location = self.volume_location.as_ref()?;

        let read = || {
            let mut file = File::open(location)?;
            let mut contents = String::new();
            file.read_to_string(&mut contents)?;
            contents
                .parse()
                .map_err(|e| Error::new(ErrorKind::InvalidData, e))
        };

        match read() {
            Ok(v) => Some(v),
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    warn!("Error reading volume from cache: {}", e);
                }
                None
            }
        }
    }

    pub fn save_volume(&self, volume: u16) {
        if let Some(ref location) = self.volume_location {
            let result = File::create(location).and_then(|mut file| write!(file, "{}", volume));
            if let Err(e) = result {
                warn!("Cannot save volume to cache: {}", e);
            }
        }
    }

    fn file_path(&self, file: FileId) -> Option<PathBuf> {
        self.audio_location.as_ref().map(|location| {
            let name = file.to_base16();
            let mut path = location.join(&name[0..2]);
            path.push(&name[2..]);
            path
        })
    }

    pub fn file(&self, file: FileId) -> Option<File> {
        let path = self.file_path(file)?;
        match File::open(&path) {
            Ok(file) => {
                if let Some(limiter) = self.size_limiter.as_deref() {
                    limiter.touch(&path);
                }
                Some(file)
            }
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    warn!("Error reading file from cache: {}", e)
                }
                None
            }
        }
    }

    pub fn save_file<F: Read>(&self, file: FileId, contents: &mut F) {
        let path = if let Some(path) = self.file_path(file) {
            path
        } else {
            return;
        };
        let parent = path.parent().unwrap();

        let result = fs::create_dir_all(parent)
            .and_then(|_| File::create(&path))
            .and_then(|mut file| io::copy(contents, &mut file));

        if let Ok(size) = result {
            if let Some(limiter) = self.size_limiter.as_deref() {
                limiter.add(&path, size);
                limiter.prune();
            }
        }
    }

    pub fn remove_file(&self, file: FileId) -> Result<(), RemoveFileError> {
        let path = self.file_path(file).ok_or(RemoveFileError(()))?;

        if let Err(err) = fs::remove_file(&path) {
            warn!("Unable to remove file from cache: {}", err);
            Err(RemoveFileError(()))
        } else {
            if let Some(limiter) = self.size_limiter.as_deref() {
                limiter.remove(&path);
            }
            Ok(())
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::time::Duration;

    fn ordered_time(v: u64) -> SystemTime {
        SystemTime::UNIX_EPOCH + Duration::from_secs(v)
    }

    #[test]
    fn test_size_limiter() {
        let mut limiter = SizeLimiter::new(1000);

        limiter.add(Path::new("a"), 500, ordered_time(2));
        limiter.add(Path::new("b"), 500, ordered_time(1));

        // b (500) -> a (500)  => sum: 1000 <= 1000
        assert!(!limiter.exceeds_limit());
        assert_eq!(limiter.pop(), None);

        limiter.add(Path::new("c"), 1000, ordered_time(3));

        // b (500) -> a (500) -> c (1000)  => sum: 2000 > 1000
        assert!(limiter.exceeds_limit());
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("b")));
        // a (500) -> c (1000)  => sum: 1500 > 1000
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("a")));
        // c (1000)   => sum: 1000 <= 1000
        assert_eq!(limiter.pop().as_deref(), None);

        limiter.add(Path::new("d"), 5, ordered_time(2));
        // d (5) -> c (1000) => sum: 1005 > 1000
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("d")));
        // c (1000)   => sum: 1000 <= 1000
        assert_eq!(limiter.pop().as_deref(), None);

        // Test updating

        limiter.add(Path::new("e"), 500, ordered_time(3));
        //  c (1000) -> e (500)  => sum: 1500 > 1000
        assert!(limiter.update(Path::new("c"), ordered_time(4)));
        // e (500) -> c (1000)  => sum: 1500 > 1000
        assert_eq!(limiter.pop().as_deref(), Some(Path::new("e")));
        // c (1000)  => sum: 1000 <= 1000

        // Test removing
        limiter.add(Path::new("f"), 500, ordered_time(2));
        assert!(limiter.remove(Path::new("c")));
        assert!(!limiter.exceeds_limit());
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2015 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.

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc;

use syntax::ast::{MetaItem, Expr};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic::*;
use syntax::ext::deriving::generic::ty::*;
use syntax::parse::token::InternedString;
use syntax::ptr::P;
use syntax::ext::base::MultiDecorator;
use syntax::parse::token;

use rustc::plugin::Registry;

macro_rules! pathvec {
    ($($x:ident)::+) => (
        vec![ $( stringify!($x) ),+ ]
    )
}

macro_rules! path {
    ($($x:tt)*) => (
        ::syntax::ext::deriving::generic::ty::Path::new( pathvec!( $($x)* ) )
    )
}

macro_rules! path_local {
    ($x:ident) => (
        ::syntax::ext::deriving::generic::ty::Path::new_local(stringify!($x))
    )
}

macro_rules! pathvec_std {
    ($cx:expr, $first:ident :: $($rest:ident)::+) => (
        if $cx.use_std {
            pathvec!(std :: $($rest)::+)
        } else {
            pathvec!($first :: $($rest)::+)
        }
    )
}

macro_rules! path_std {
    ($($x:tt)*) => (
        ::syntax::ext::deriving::generic::ty::Path::new( pathvec_std!( $($x)* ) )
    )
}

pub fn expand_deriving_from_primitive(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: Annotatable,
                                      push: &mut FnMut(Annotatable))
{
    let inline = cx.meta_word(span, InternedString::new("inline"));
    let attrs = vec!(cx.attribute(span, inline));
    let trait_def = TraitDef {
        span: span,
        attributes: Vec::new(),
        path: path!(num::FromPrimitive),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        methods: vec!(
            MethodDef {
                name: "from_i64",
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(i64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs.clone(),
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("i64", c, s, sub)
                })),
            },
            MethodDef {
                name: "from_u64",
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(u64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs,
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("u64", c, s, sub)
                })),
            }
        ),
        associated_types: Vec::new(),
    };

    trait_def.expand(cx, mitem, &item, push)
}

fn cs_from(name: &str, cx: &mut ExtCtxt, trait_span: Span, substr: &Substructure) -> P<Expr> {
    if substr.nonself_args.len() != 1 {
        cx.span_bug(trait_span, "incorrect number of arguments in `derive(FromPrimitive)`")
    }

    let n = &substr.nonself_args[0];

    match *substr.fields {
        StaticStruct(..) => {
            cx.span_err(trait_span, "`FromPrimitive` cannot be derived for structs");
            return cx.expr_fail(trait_span, InternedString::new(""));
        }
        StaticEnum(enum_def, _) => {
            if enum_def.variants.is_empty() {
                cx.span_err(trait_span,
                            "`FromPrimitive` cannot be derived for enums with no variants");
                return cx.expr_fail(trait_span, InternedString::new(""));
            }

            let mut arms = Vec::new();

            for variant in &enum_def.variants {
                match variant.node.kind {
                    ast::TupleVariantKind(ref args) => {
                        if !args.is_empty() {
                            cx.span_err(trait_span,
                                        "`FromPrimitive` cannot be derived for \
                                        enum variants with arguments");
                            return cx.expr_fail(trait_span,
                                                InternedString::new(""));
                        }
                        let span = variant.span;

                        // expr for `$n == $variant as $name`
                        let path = cx.path(span, vec![substr.type_ident, variant.node.name]);
                        let variant = cx.expr_path(path);
                        let ty = cx.ty_ident(span, cx.ident_of(name));
                        let cast = cx.expr_cast(span, variant.clone(), ty);
                        let guard = cx.expr_binary(span, ast::BiEq, n.clone(), cast);

                        // expr for `Some($variant)`
                        let body = cx.expr_some(span, variant);

                        // arm for `_ if $guard => $body`
                        let arm = ast::Arm {
                            attrs: vec!(),
                            pats: vec!(cx.pat_wild(span)),
                            guard: Some(guard),
                            body: body,
                        };

                        arms.push(arm);
                    }
                    ast::StructVariantKind(_) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for enums \
                                    with struct variants");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                }
            }

            // arm for `_ => None`
            let arm = ast::Arm {
                attrs: vec!(),
                pats: vec!(cx.pat_wild(trait_span)),
                guard: None,
                body: cx.expr_none(trait_span),
            };
            arms.push(arm);

            cx.expr_match(trait_span, n.clone(), arms)
        }
        _ => cx.span_bug(trait_span, "expected StaticEnum in derive(FromPrimitive)")
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("derive_NumFromPrimitive"),
        MultiDecorator(Box::new(expand_deriving_from_primitive)));
}
<commit_msg>Update num-macros to rust master<commit_after>// Copyright 2012-2015 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.

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc;

use syntax::ast::{MetaItem, Expr};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic::*;
use syntax::ext::deriving::generic::ty::*;
use syntax::parse::token::InternedString;
use syntax::ptr::P;
use syntax::ext::base::MultiDecorator;
use syntax::parse::token;

use rustc::plugin::Registry;

macro_rules! pathvec {
    ($($x:ident)::+) => (
        vec![ $( stringify!($x) ),+ ]
    )
}

macro_rules! path {
    ($($x:tt)*) => (
        ::syntax::ext::deriving::generic::ty::Path::new( pathvec!( $($x)* ) )
    )
}

macro_rules! path_local {
    ($x:ident) => (
        ::syntax::ext::deriving::generic::ty::Path::new_local(stringify!($x))
    )
}

macro_rules! pathvec_std {
    ($cx:expr, $first:ident :: $($rest:ident)::+) => (
        if $cx.use_std {
            pathvec!(std :: $($rest)::+)
        } else {
            pathvec!($first :: $($rest)::+)
        }
    )
}

macro_rules! path_std {
    ($($x:tt)*) => (
        ::syntax::ext::deriving::generic::ty::Path::new( pathvec_std!( $($x)* ) )
    )
}

pub fn expand_deriving_from_primitive(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: Annotatable,
                                      push: &mut FnMut(Annotatable))
{
    let inline = cx.meta_word(span, InternedString::new("inline"));
    let attrs = vec!(cx.attribute(span, inline));
    let trait_def = TraitDef {
        span: span,
        attributes: Vec::new(),
        path: path!(num::FromPrimitive),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        methods: vec!(
            MethodDef {
                name: "from_i64",
                is_unsafe: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(i64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs.clone(),
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("i64", c, s, sub)
                })),
            },
            MethodDef {
                name: "from_u64",
                is_unsafe: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(u64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs,
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("u64", c, s, sub)
                })),
            }
        ),
        associated_types: Vec::new(),
    };

    trait_def.expand(cx, mitem, &item, push)
}

fn cs_from(name: &str, cx: &mut ExtCtxt, trait_span: Span, substr: &Substructure) -> P<Expr> {
    if substr.nonself_args.len() != 1 {
        cx.span_bug(trait_span, "incorrect number of arguments in `derive(FromPrimitive)`")
    }

    let n = &substr.nonself_args[0];

    match *substr.fields {
        StaticStruct(..) => {
            cx.span_err(trait_span, "`FromPrimitive` cannot be derived for structs");
            return cx.expr_fail(trait_span, InternedString::new(""));
        }
        StaticEnum(enum_def, _) => {
            if enum_def.variants.is_empty() {
                cx.span_err(trait_span,
                            "`FromPrimitive` cannot be derived for enums with no variants");
                return cx.expr_fail(trait_span, InternedString::new(""));
            }

            let mut arms = Vec::new();

            for variant in &enum_def.variants {
                match variant.node.kind {
                    ast::TupleVariantKind(ref args) => {
                        if !args.is_empty() {
                            cx.span_err(trait_span,
                                        "`FromPrimitive` cannot be derived for \
                                        enum variants with arguments");
                            return cx.expr_fail(trait_span,
                                                InternedString::new(""));
                        }
                        let span = variant.span;

                        // expr for `$n == $variant as $name`
                        let path = cx.path(span, vec![substr.type_ident, variant.node.name]);
                        let variant = cx.expr_path(path);
                        let ty = cx.ty_ident(span, cx.ident_of(name));
                        let cast = cx.expr_cast(span, variant.clone(), ty);
                        let guard = cx.expr_binary(span, ast::BiEq, n.clone(), cast);

                        // expr for `Some($variant)`
                        let body = cx.expr_some(span, variant);

                        // arm for `_ if $guard => $body`
                        let arm = ast::Arm {
                            attrs: vec!(),
                            pats: vec!(cx.pat_wild(span)),
                            guard: Some(guard),
                            body: body,
                        };

                        arms.push(arm);
                    }
                    ast::StructVariantKind(_) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for enums \
                                    with struct variants");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                }
            }

            // arm for `_ => None`
            let arm = ast::Arm {
                attrs: vec!(),
                pats: vec!(cx.pat_wild(trait_span)),
                guard: None,
                body: cx.expr_none(trait_span),
            };
            arms.push(arm);

            cx.expr_match(trait_span, n.clone(), arms)
        }
        _ => cx.span_bug(trait_span, "expected StaticEnum in derive(FromPrimitive)")
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("derive_NumFromPrimitive"),
        MultiDecorator(Box::new(expand_deriving_from_primitive)));
}
<|endoftext|>


<commit_before>use impl_prelude::*;
use lapack::c::{sgeqrf, sorgqr};
use std::fmt::Debug;
use types::Magnitude;

/// Error for QR-based computations.
#[derive(Debug, Clone)]
pub enum QRError {
    /// The layout of the matrix is not compatible
    BadLayout,

    /// The dimensions of the raw qr and tau don't match
    InconsistentDimensions,

    /// Implementation error, please submit as bug.
    IllegalParameter(i32),
}

/// Representation of the components Q, R of the factorization of
/// matrix A.
#[derive(Debug)]
pub struct QRFactors<T: QR> {
    mat: Array<T, Ix2>,
    tau: Vec<T>,
}

impl<T: QR> QRFactors<T> {
    /// Create a `QRFactors` object from the output of the LAPACKE
    /// functions.
    fn from_raw<Matrix>(mat: Matrix, tau: Vec<T>) -> Result<QRFactors<T>, QRError>
        where Matrix: Into<Array<T, Ix2>>
    {
        let mut ma = mat.into();

        if slice_and_layout_mut(&mut ma).is_none() {
            return Err(QRError::BadLayout);
        }

        Ok(QRFactors {
            mat: ma,
            tau: tau,
        })
    }

    /// Given a matrix, compute the QR decomposition of the matrix,
    /// returned as a `QRFactors` object.
    pub fn compute<D1: Data<Elem=T>>(mat: &ArrayBase<D1, Ix2>)
                                     -> Result<QRFactors<T>, QRError> {
        QR::compute(mat)
    }

    /// Return the number of rows in the original matrix
    pub fn rows(&self) -> usize {
        self.mat.rows()
    }

    /// Return the number of columns in the original matrix
    pub fn cols(&self) -> usize {
        self.mat.cols()
    }

    fn p(&self) -> usize {
        cmp::min(self.rows(), self.cols())
    }

    /// Return the first `k` columns of the matrix Q of the QR
    /// factorization.
    ///
    /// `Q` is generated such that the columns of `Q` form an
    /// orthogonal basis for the first `k` columns of `A`.
    ///
    /// When `k` is None, compute enough columns (`min(m, n)`) to
    /// faithfully recreate the original matrix A.
    pub fn qk<K: Into<Option<usize>>>(&self, k: K) -> Result<Array<T, Ix2>, QRError> {
        let kr = k.into();
        QR::compute_q(&self.mat, &self.tau, kr.unwrap_or(self.p()))
    }

    /// Return the `m` by `min(m, n)` matrix `Q`.
    ///
    /// Equivalent to `self.qk(None)`.
    ///
    /// `self.q()` is large enough so that `&self.q() * &self.r()`
    /// will faithfully reproduce the original matrix `A`.
    #[inline]
    pub fn q(&self) -> Result<Array<T, Ix2>, QRError> {
        self.qk(None)
    }

    /// Return the first 'k' rows of the matrix R of the QR
    /// factorization.
    ///
    /// When `k` is None, compute enough rows (`min(m, n)`) to
    /// faithfully recreate the original matrix A.
    pub fn rk<K: Into<Option<usize>>>(&self, k: K) -> Result<Array<T, Ix2>, QRError> {
        let kr = k.into();
        let p = kr.unwrap_or(self.p());
        QR::compute_r(&self.mat, p)
    }

    /// Return the first `min(m, n)` by `n` matrix R of the QR
    /// factorization.
    ///
    /// Equivalent to `self::rk(None)`.
    ///
    /// `self.r()` is large enough so that `&self.q() * &self.r()`
    /// will faithfully reproduce the original matrix `A`.
    #[inline]
    pub fn r(&self) -> Result<Array<T, Ix2>, QRError> {
        self.rk(None)
    }
}

/// Trait defined on scalars to support QR-factorization.
pub trait QR: Sized + Clone + Magnitude + Debug {
    /// Return a `QRFactors` structure, containing the QR
    /// factorization of the input matrix `A`.
    ///
    /// Similar to `compute`, but consumes the input.
    fn compute_into(a: Array<Self, Ix2>) -> Result<QRFactors<Self>, QRError>;

    /// Return a `QRFactors` structure, containing the QR
    /// factorization of the input matrix `A`.
    fn compute<D1: Data>(a: &ArrayBase<D1, Ix2>) -> Result<QRFactors<Self>, QRError>
        where D1: Data<Elem = Self>
    {
        Self::compute_into(a.to_owned())
    }

    /// Compute Q from raw parts.
    ///
    /// Not intended to be used by end-users.
    fn compute_q<D1>(mat: &ArrayBase<D1, Ix2>,
                     tau: &[Self],
                     k: usize)
                     -> Result<Array<Self, Ix2>, QRError>
        where D1: Data<Elem = Self>;

    /// Compute R from raw parts.
    ///
    /// Not intended to be used by end-users.
    fn compute_r<D1>(mat: &ArrayBase<D1, Ix2>, k: usize) -> Result<Array<Self, Ix2>, QRError>
        where D1: Data<Elem = Self>;
}

impl QR for f32 {
    fn compute_into(mut a: Array<Self, Ix2>) -> Result<QRFactors<Self>, QRError> {
        let dim = a.dim();

        let (info, tau) = {
            let (mut slice, layout, lda) = match slice_and_layout_mut(&mut a) {
                None => return Err(QRError::BadLayout),
                Some(x) => x,
            };

            let mut tau = Vec::new();
            tau.resize(cmp::min(dim.0, dim.1), 0.0);

            // workspace query
            (sgeqrf(layout,
                    dim.0 as i32,
                    dim.1 as i32,
                    &mut slice,
                    lda as i32,
                    &mut tau),
             tau)
        };

        println!("{:?}\n {:?}\n", a, tau);

        if info == 0 {
            QRFactors::from_raw(a, tau)
        } else if info < 0 {
            Err(QRError::IllegalParameter(-info))
        } else {
            unreachable!();
        }
    }

    fn compute_q<D1>(mat: &ArrayBase<D1, Ix2>,
                     tau: &[Self],
                     k: usize)
                     -> Result<Array<Self, Ix2>, QRError>
        where D1: Data<Elem = Self>
    {

        let (m, n) = mat.dim();
        if k > m {
            return Err(QRError::InconsistentDimensions);
        }

        // Initialize q with the
        let mut q = mat.slice(s![.., ..k as isize]).to_owned();

        let info = {
            let (slice, layout, ldq) = match slice_and_layout_mut(&mut q) {
                None => unreachable!(),
                Some(fwd) => fwd,
            };

            sorgqr(layout,
                   m as i32,
                   k as i32,
                   cmp::min(k, n) as i32,
                   slice,
                   ldq as i32,
                   tau)
        };
        if info == 0 {
            Ok(q)
        } else {
            Err(QRError::IllegalParameter(-info))
        }
    }

    fn compute_r<D1>(mat: &ArrayBase<D1, Ix2>, k: usize) -> Result<Array<Self, Ix2>, QRError>
        where D1: Data<Elem = Self>
    {
        let (m, n) = mat.dim();

        let nn = cmp::min(m, n);
        if k > nn {
            return Err(QRError::InconsistentDimensions);
        }

        let mut r = Array::zeros((k as usize, n));

        // Copy the upper triangular/trapezoidal part of the matrix to
        // R.
        r.slice_mut(s![..k as isize, ..]).assign(&mat.slice(s![..k as isize, ..]));

        // Replace zeros below the diagonal.
        for (i, mut row) in r.outer_iter_mut().enumerate().take(nn as usize) {
            row.slice_mut(s![..i as isize]).assign_scalar(&0.0);
        }

        Ok(r)
    }
}
<commit_msg>qr: generalize implementation to all scalars<commit_after>use impl_prelude::*;
use lapack::c::{sgeqrf, sorgqr, dgeqrf, dorgqr, cgeqrf, cungqr, zgeqrf, zungqr};
use std::fmt::Debug;
use num_traits::Zero;
use types::Magnitude;

/// Error for QR-based computations.
#[derive(Debug, Clone)]
pub enum QRError {
    /// The layout of the matrix is not compatible
    BadLayout,

    /// The dimensions of the raw qr and tau don't match
    InconsistentDimensions,

    /// Implementation error, please submit as bug.
    IllegalParameter(i32),
}

/// Representation of the components Q, R of the factorization of
/// matrix A.
#[derive(Debug)]
pub struct QRFactors<T: QR> {
    mat: Array<T, Ix2>,
    tau: Vec<T>,
}

impl<T: QR> QRFactors<T> {
    /// Create a `QRFactors` object from the output of the LAPACKE
    /// functions.
    fn from_raw<Matrix>(mat: Matrix, tau: Vec<T>) -> Result<QRFactors<T>, QRError>
        where Matrix: Into<Array<T, Ix2>>
    {
        let mut ma = mat.into();

        if slice_and_layout_mut(&mut ma).is_none() {
            return Err(QRError::BadLayout);
        }

        Ok(QRFactors {
            mat: ma,
            tau: tau,
        })
    }

    /// Given a matrix, compute the QR decomposition of the matrix,
    /// returned as a `QRFactors` object.
    pub fn compute<D1: Data<Elem=T>>(mat: &ArrayBase<D1, Ix2>)
                                     -> Result<QRFactors<T>, QRError> {
        QR::compute(mat)
    }

    /// Return the number of rows in the original matrix
    pub fn rows(&self) -> usize {
        self.mat.rows()
    }

    /// Return the number of columns in the original matrix
    pub fn cols(&self) -> usize {
        self.mat.cols()
    }

    fn p(&self) -> usize {
        cmp::min(self.rows(), self.cols())
    }

    /// Return the first `k` columns of the matrix Q of the QR
    /// factorization.
    ///
    /// `Q` is generated such that the columns of `Q` form an
    /// orthogonal basis for the first `k` columns of `A`.
    ///
    /// When `k` is None, compute enough columns (`min(m, n)`) to
    /// faithfully recreate the original matrix A.
    pub fn qk<K: Into<Option<usize>>>(&self, k: K) -> Result<Array<T, Ix2>, QRError> {
        let kr = k.into();
        QR::compute_q(&self.mat, &self.tau, kr.unwrap_or(self.p()))
    }

    /// Return the `m` by `min(m, n)` matrix `Q`.
    ///
    /// Equivalent to `self.qk(None)`.
    ///
    /// `self.q()` is large enough so that `&self.q() * &self.r()`
    /// will faithfully reproduce the original matrix `A`.
    #[inline]
    pub fn q(&self) -> Result<Array<T, Ix2>, QRError> {
        self.qk(None)
    }

    /// Return the first 'k' rows of the matrix R of the QR
    /// factorization.
    ///
    /// When `k` is None, compute enough rows (`min(m, n)`) to
    /// faithfully recreate the original matrix A.
    pub fn rk<K: Into<Option<usize>>>(&self, k: K) -> Result<Array<T, Ix2>, QRError> {
        let kr = k.into();
        let p = kr.unwrap_or(self.p());
        QR::compute_r(&self.mat, p)
    }

    /// Return the first `min(m, n)` by `n` matrix R of the QR
    /// factorization.
    ///
    /// Equivalent to `self::rk(None)`.
    ///
    /// `self.r()` is large enough so that `&self.q() * &self.r()`
    /// will faithfully reproduce the original matrix `A`.
    #[inline]
    pub fn r(&self) -> Result<Array<T, Ix2>, QRError> {
        self.rk(None)
    }
}

/// Trait defined on scalars to support QR-factorization.
pub trait QR: Sized + Clone + Magnitude + Debug {
    /// Return a `QRFactors` structure, containing the QR
    /// factorization of the input matrix `A`.
    ///
    /// Similar to `compute`, but consumes the input.
    fn compute_into(a: Array<Self, Ix2>) -> Result<QRFactors<Self>, QRError>;

    /// Return a `QRFactors` structure, containing the QR
    /// factorization of the input matrix `A`.
    fn compute<D1: Data>(a: &ArrayBase<D1, Ix2>) -> Result<QRFactors<Self>, QRError>
        where D1: Data<Elem = Self>
    {
        Self::compute_into(a.to_owned())
    }

    /// Compute Q from raw parts.
    ///
    /// Not intended to be used by end-users.
    fn compute_q<D1>(mat: &ArrayBase<D1, Ix2>,
                     tau: &[Self],
                     k: usize)
                     -> Result<Array<Self, Ix2>, QRError>
        where D1: Data<Elem = Self>;

    /// Compute R from raw parts.
    ///
    /// Not intended to be used by end-users.
    fn compute_r<D1>(mat: &ArrayBase<D1, Ix2>, k: usize) -> Result<Array<Self, Ix2>, QRError>
        where D1: Data<Elem = Self>;
}

macro_rules! impl_qr {
    ($qr_type:ty, $qr_func:ident, $qr_to_q:ident) => (

        impl QR for $qr_type {
            fn compute_into(mut a: Array<Self, Ix2>) -> Result<QRFactors<Self>, QRError> {
                let dim = a.dim();

                let (info, tau) = {
                    let (mut slice, layout, lda) = match slice_and_layout_mut(&mut a) {
                        None => return Err(QRError::BadLayout),
                        Some(x) => x,
                    };

                    let mut tau = Vec::new();
                    tau.resize(cmp::min(dim.0, dim.1), <$qr_type as Zero>::zero());

                    // workspace query
                    ($qr_func(layout,
                            dim.0 as i32,
                            dim.1 as i32,
                            &mut slice,
                            lda as i32,
                            &mut tau),
                     tau)
                };

                println!("{:?}\n {:?}\n", a, tau);

                if info == 0 {
                    QRFactors::from_raw(a, tau)
                } else if info < 0 {
                    Err(QRError::IllegalParameter(-info))
                } else {
                    unreachable!();
                }
            }

            fn compute_q<D1>(mat: &ArrayBase<D1, Ix2>,
                             tau: &[Self],
                             k: usize)
                             -> Result<Array<Self, Ix2>, QRError>
                where D1: Data<Elem = Self>
            {

                let (m, n) = mat.dim();
                if k > m {
                    return Err(QRError::InconsistentDimensions);
                }

                // Initialize q with the
                let mut q = mat.slice(s![.., ..k as isize]).to_owned();

                let info = {
                    let (slice, layout, ldq) = match slice_and_layout_mut(&mut q) {
                        None => unreachable!(),
                        Some(fwd) => fwd,
                    };

                    $qr_to_q(layout,
                           m as i32,
                           k as i32,
                           cmp::min(k, n) as i32,
                           slice,
                           ldq as i32,
                           tau)
                };
                if info == 0 {
                    Ok(q)
                } else {
                    Err(QRError::IllegalParameter(-info))
                }
            }

            fn compute_r<D1>(mat: &ArrayBase<D1, Ix2>, k: usize) -> Result<Array<Self, Ix2>, QRError>
                where D1: Data<Elem = Self>
            {
                let (m, n) = mat.dim();

                let nn = cmp::min(m, n);
                if k > nn {
                    return Err(QRError::InconsistentDimensions);
                }

                let mut r = Array::zeros((k as usize, n));

                // Copy the upper triangular/trapezoidal part of the matrix to
                // R.
                r.slice_mut(s![..k as isize, ..]).assign(&mat.slice(s![..k as isize, ..]));

                let zero = <$qr_type as Zero>::zero();

                // Replace zeros below the diagonal.
                for (i, mut row) in r.outer_iter_mut().enumerate().take(nn as usize) {
                    row.slice_mut(s![..i as isize]).assign_scalar(&zero);
                }

                Ok(r)
            }
        }
    )
}

impl_qr!(f32, sgeqrf, sorgqr);
impl_qr!(f64, dgeqrf, dorgqr);
impl_qr!(c32, cgeqrf, cungqr);
impl_qr!(c64, zgeqrf, zungqr);
<|endoftext|>


<commit_before>use std::io::Cursor;
use byteorder::{WriteBytesExt, ReadBytesExt, BigEndian};

//mod error;
use error::QuicError;
use error::Result;
use std::io::Read;
use util::OFSize;
//use super::FrameType;
//use frames::ACK;
use util::optimal_field_size;

#[derive(Debug, PartialEq)]
pub struct AckFrame {
    pub num_blocks: Option<u8>,
    pub num_ts: u8,
    pub largest_ack: u64,
    pub ack_delay: u16,
    pub first_ack_len: u64,
    pub ack_blocks: Option<Vec<AckBlock>>,
    pub delta_la: Option<u8>,
    pub first_ts: Option<u32>,
    pub timestamps: Option<Vec<AckTimestamp>>,
}

#[derive(Debug, PartialEq)]
pub struct AckBlock {
    pub gap: u8,
    pub block_len: u64,
}

impl AckBlock {
    pub fn from_bytes(buf: &[u8], field_len: usize) -> Result<AckBlock> {
        let mut reader = Cursor::new(buf);

        let gap = reader.read_u8()?;

        let block_len = reader.read_uint::<BigEndian>(field_len)?;

        Ok(AckBlock {
            gap: gap,
            block_len: block_len
        })
    }

    pub fn as_bytes(&self, block_len_len: OFSize) -> Vec<u8> {
        let mut bytes = Vec::new();

        bytes.write_u8(self.gap);

        match block_len_len {
            OFSize::U8 => bytes.write_u8(self.block_len as u8),
            OFSize::U16 => bytes.write_u16::<BigEndian>(self.block_len as u16),
            OFSize::U32 => bytes.write_u32::<BigEndian>(self.block_len as u32),
            OFSize::U48 => bytes.write_uint::<BigEndian>(self.block_len as u64, 6),
            _ => panic!("Higher than 48-bits not allowed in block length!")
        };

        bytes
    }
}

#[derive(Debug, PartialEq)]
pub struct AckTimestamp {
    pub delta_la: u8,
    pub time_since_prev: u16
}

impl AckTimestamp {
    pub fn from_bytes(buf: &[u8]) -> Result<AckTimestamp> {
        let mut reader = Cursor::new(buf);

        let delta_la = reader.read_u8()?;
        let time_since_prev = reader.read_u16::<BigEndian>()?;

        Ok(AckTimestamp {
            delta_la: delta_la,
            time_since_prev: time_since_prev,
        })
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();

        bytes.write_u8(self.delta_la);
        bytes.write_u16::<BigEndian>(self.time_since_prev);

        bytes
    }
}

impl AckFrame {
    pub fn from_bytes(buf: &[u8]) -> Result<AckFrame> {
        let mut reader = Cursor::new(buf);

        let type_byte = reader.read_u8()?;

        // firstbyte: 101NLLMM
        let n = type_byte & 0x10 > 0;
        let ll = (type_byte & 0x0c) >> 2;
        let mm = type_byte & 0x03;

        let la_len = match ll {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

        let ack_len = match mm {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

//        let num_blocks;

        let num_blocks = if n {
            Some(reader.read_u8()?)
        } else {
            None
        };

        let num_ts = reader.read_u8()?;

        let largest_ack = reader.read_uint::<BigEndian>(la_len)?;

        let ack_delay = reader.read_u16::<BigEndian>()?;

        let first_ack_len = reader.read_uint::<BigEndian>(ack_len)?;

        let mut ack_blocks: Vec<AckBlock> = Vec::new();

        if let Some(n_blocks) = num_blocks {
            let ack_block_section_len = (1 + ack_len) * (n_blocks as usize);

            let mut reader_handle = reader.clone().take(ack_block_section_len as u64);

            let mut ack_block_slice = Vec::new();
            reader_handle.read_to_end(&mut ack_block_slice);

            let new_pos = (reader.position() as u64) + (ack_block_section_len as u64);

            reader.set_position(new_pos);

            for block in ack_block_slice.chunks(1 + ack_len) {
                let c_block = AckBlock::from_bytes(&block.to_vec(), ack_len)?;
                ack_blocks.push(c_block);
            }
        }

        let mut timestamps = Vec::new();

        let delta_la;
        let first_ts;
        if num_ts > 0 {
            let delta_la_i = reader.read_u8()?;
            delta_la = Some(delta_la_i);

            let first_ts_i = reader.read_u32::<BigEndian>()?;
            first_ts = Some(first_ts_i);

            let ts_block_section_len = num_ts * 3;

            let mut ts_block_slice = Vec::new();
            let mut reader_handle = reader.clone().take(ts_block_section_len as u64);
            reader_handle.read_to_end(&mut ts_block_slice);

            for block in ts_block_slice.chunks(3) {
                let c_block = AckTimestamp::from_bytes(&block.to_vec())?;
                timestamps.push(c_block);
            }
        } else {
            delta_la = None;
            first_ts = None;
        }

        let ack_blocks_fin = match ack_blocks.len() {
            0 => None,
            _ => Some(ack_blocks),
        };

        let ts_blocks_fin = match timestamps.len() {
            0 => None,
            _ => Some(timestamps),
        };

        Ok(AckFrame {
            num_blocks: num_blocks,
            num_ts: num_ts,
            largest_ack: largest_ack,
            ack_delay: ack_delay,
            first_ack_len: first_ack_len,
            ack_blocks: ack_blocks_fin,
            delta_la: delta_la,
            first_ts: first_ts,
            timestamps: ts_blocks_fin,
        })
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();

        let mut type_byte = super::ACK.bits();

        match self.num_blocks {
            Some(_) => type_byte |= 0x10,
            _ => type_byte |= 0x00
        };

        let largest_ack_size = optimal_field_size(self.largest_ack);

        match largest_ack_size {
            OFSize::U8 => type_byte |= 0x00,
            OFSize::U16 => type_byte |= 0x04,
            OFSize::U32 => type_byte |= 0x08,
            OFSize::U48 => type_byte |= 0x0c,
            _ => panic!("largest ack too large")
        };

        let max_block_len = match self.ack_blocks {
            Some(ref ack_blocks) => ack_blocks.iter().max_by_key(|block| block.block_len),
            None => None
        };

        let block_len_size = match max_block_len {
            None => None,
            Some(max_len_block) => Some(optimal_field_size(max_len_block.block_len))
        };

        match block_len_size {
            None => {},
            Some(OFSize::U8) => type_byte |= 0x00,
            Some(OFSize::U16) => type_byte |= 0x01,
            Some(OFSize::U32) => type_byte |= 0x02,
            Some(OFSize::U48) => type_byte |= 0x03,
            _ => panic!("largest ack too large")
        };

        bytes.write_u8(type_byte);

        if let Some(ref ack_blocks) = self.ack_blocks {
            bytes.write_u8(ack_blocks.len() as u8);
        }

        if let Some(ref ts_blocks) = self.timestamps {
            let mut ts_len = ts_blocks.len();
            if self.delta_la.is_some() {
                ts_len += 0;
            }

            bytes.write_u8(ts_len as u8);
        }

        match largest_ack_size {
            OFSize::U8 => {
                bytes.write_u8(self.largest_ack as u8).unwrap()
            },
            OFSize::U16 => {
                bytes.write_u16::<BigEndian>(self.largest_ack as u16).unwrap()
            },
            OFSize::U32 => {
                bytes.write_u32::<BigEndian>(self.largest_ack as u32).unwrap()
            },
            OFSize::U48 => {
                bytes.write_uint::<BigEndian>(self.largest_ack as u64, 6).unwrap()
            },
            _ => panic!("largest ack too large")
        }

        bytes.write_u16::<BigEndian>(self.ack_delay);

        if let Some(ref ack_blocks) = self.ack_blocks {
            bytes.write_u16::<BigEndian>(ack_blocks.len() as u16);
            let mut ack_block_bytes = Vec::new();
            let ack_block_byte_vectors = ack_blocks.iter().map(|block|
                block.as_bytes(block_len_size.unwrap())
            );

            for ack_block_byte_vector in ack_block_byte_vectors {
                ack_block_bytes.extend(ack_block_byte_vector);
            }

            bytes.extend(ack_block_bytes);

        }

        match self.delta_la {
            None => {},
            Some(la) => { let _ = bytes.write_u8(la); }
        }

        match self.first_ts {
            None => {},
            Some(ts) => { let _ = bytes.write_u32::<BigEndian>(ts); }
        }

        if let Some(ref ack_blocks) = self.timestamps {
            let mut ack_block_bytes = Vec::new();

            let ack_block_byte_vectors = ack_blocks.iter().map(|block|
                block.as_bytes()
            );

            for ack_block_byte_vector in ack_block_byte_vectors {
                ack_block_bytes.extend(ack_block_byte_vector);
            }

            bytes.extend(ack_block_bytes)
        }

        bytes
    }

    pub fn frame_len(buf: &[u8]) -> Result<usize> {
        let mut reader = Cursor::new(buf);

        let type_byte = reader.read_u8()?;

        let mut len = 1 + 1 + 2;

        let n = type_byte & 0x10 > 0;
        let ll = (type_byte & 0x0c) >> 2;
        let mm = type_byte & 0x03;

        if n {
            len += 1;
        }

        let la_len = match ll {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

        len += la_len;

        let ack_len = match mm {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

        len += ack_len;

        if n {
            let num_blocks = reader.read_u8()? as usize;
            len += num_blocks * (1 + ack_len);
        }

        Ok(len)
    }
}

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

    #[test]
    fn serialize() {
        let timestamps = vec![
            AckTimestamp {
                delta_la: 31,
                time_since_prev: 26,
            },
            AckTimestamp {
                delta_la: 42,
                time_since_prev: 97,
            }
        ];

        let ack_blocks = vec![
            AckBlock {
                gap: 12,
                block_len: 14,
            },
            AckBlock {
                gap: 87,
                block_len: 645,
            },
            AckBlock {
                gap: 42,
                block_len: 325,
            },
            AckBlock {
                gap: 236,
                block_len: 734,
            }
        ];

        let ack_frame = AckFrame {
            num_blocks: Some(4),
            num_ts: 2,
            largest_ack: 497,
            ack_delay: 9,
            first_ack_len: ack_blocks.len() as u64,
            ack_blocks: Some(ack_blocks),
            delta_la: Some(17),
            first_ts: Some(91),
            timestamps: Some(timestamps),
        };

        let ack_frame_bytes = ack_frame.as_bytes();

        let parsed_ack_frame = AckFrame::from_bytes(&ack_frame_bytes).unwrap();

        assert_eq!(ack_frame, parsed_ack_frame);
    }
}<commit_msg>parses timestamps, closes #1<commit_after>use std::io::Cursor;
use byteorder::{WriteBytesExt, ReadBytesExt, BigEndian};

//mod error;
use error::QuicError;
use error::Result;
use std::io::Read;
use util::OFSize;
//use super::FrameType;
//use frames::ACK;
use util::optimal_field_size;

const UFLOAT_16_EXPONENT_BITS: u16 = 5;
const UFLOAT_16_MAX_EXPONENT: u16 = (1 << UFLOAT_16_EXPONENT_BITS) - 2; // 30
const UFLOAT_16_MANTISSA_BITS: u16 = 16 - UFLOAT_16_EXPONENT_BITS; // 11
const UFLOAT_16_MANTISSA_EFFECTIVE_BITS: u16 = UFLOAT_16_MANTISSA_BITS + 1; // 12
const UFLOAT_16_MAX_VALUE: u64 =
(((1u64 << UFLOAT_16_MANTISSA_EFFECTIVE_BITS) - 1) << UFLOAT_16_MAX_EXPONENT) as u64; // 0x3FFC0000000

#[derive(Debug, PartialEq)]
pub struct AckFrame {
    pub num_blocks: Option<u8>,
    pub num_ts: u8,
    pub largest_ack: u64,
    pub ack_delay: u16,
    pub first_ack_len: u64,
    pub ack_blocks: Option<Vec<AckBlock>>,
    pub delta_la: Option<u8>,
    pub first_ts: Option<u32>,
    pub timestamps: Option<Vec<AckTimestamp>>,
}

#[derive(Debug, PartialEq)]
pub struct AckBlock {
    pub gap: u8,
    pub block_len: u64,
}

impl AckBlock {
    pub fn from_bytes(buf: &[u8], field_len: usize) -> Result<Self> {
        let mut reader = Cursor::new(buf);

        let gap = reader.read_u8()?;

        let block_len = reader.read_uint::<BigEndian>(field_len)?;

        Ok(AckBlock {
            gap: gap,
            block_len: block_len
        })
    }

    pub fn as_bytes(&self, block_len_len: OFSize) -> Vec<u8> {
        let mut bytes = Vec::new();

        bytes.write_u8(self.gap);

        match block_len_len {
            OFSize::U8 => bytes.write_u8(self.block_len as u8),
            OFSize::U16 => bytes.write_u16::<BigEndian>(self.block_len as u16),
            OFSize::U32 => bytes.write_u32::<BigEndian>(self.block_len as u32),
            OFSize::U48 => bytes.write_uint::<BigEndian>(self.block_len as u64, 6),
            _ => panic!("Higher than 48-bits not allowed in block length!")
        };

        bytes
    }
}

#[derive(Debug, PartialEq)]
pub struct AckTimestampValue {
    pub microseconds: u64,
}

impl AckTimestampValue {
    pub fn from_u16(val: u16) -> Result<Self> {
//        let mut exponent = (val & 0xf800) >> 10 as u64;
//        let mantissa = (val & 0x07ff) as u64;

        let mut res = val as u64;

        if res < (1 << UFLOAT_16_MANTISSA_EFFECTIVE_BITS) {
            return Ok(AckTimestampValue {
                microseconds: res,
            })
        }

        let mut exponent = val >> UFLOAT_16_MANTISSA_BITS;
        exponent -= 1;

        res -= (exponent as u64) << UFLOAT_16_MANTISSA_BITS;
        res <<= exponent;

        Ok(AckTimestampValue {
            microseconds: res,
        })
    }

    pub fn as_u16(&self) -> u16 {
        let mut value = self.microseconds;

        if value < (1 << UFLOAT_16_MANTISSA_EFFECTIVE_BITS) {
            value as u16
        } else if value >= UFLOAT_16_MAX_VALUE as u64 {
            u16::max_value() as u16
        } else {
            let mut offset = 16u16;
            let mut exponent = 0u16;

            loop {
                if offset == 0 {
                    break;
                }

                if value >= (1u64 << ((UFLOAT_16_MANTISSA_BITS + offset) as u64)) {
                    exponent += offset;
                    value >>= offset;
                }

                offset /= 2;
            }

            (value as u16 + (exponent << UFLOAT_16_MANTISSA_BITS) as u16) as u16
        }

    }
}

#[derive(Debug, PartialEq)]
pub struct AckTimestamp {
    pub delta_la: u8,
    pub time_since_prev: AckTimestampValue
}

impl AckTimestamp {
    pub fn from_bytes(buf: &[u8]) -> Result<Self> {
        let mut reader = Cursor::new(buf);

        let delta_la = reader.read_u8()?;
        let time_since_prev_val = reader.read_u16::<BigEndian>()?;

        let time_since_prev = AckTimestampValue::from_u16(time_since_prev_val)?;

        Ok(AckTimestamp {
            delta_la: delta_la,
            time_since_prev: time_since_prev,
        })
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();

        bytes.write_u8(self.delta_la);
        bytes.write_u16::<BigEndian>(self.time_since_prev.as_u16());

        bytes
    }
}

impl AckFrame {
    pub fn from_bytes(buf: &[u8]) -> Result<Self> {
        let mut reader = Cursor::new(buf);

        let type_byte = reader.read_u8()?;

        // firstbyte: 101NLLMM
        let n = type_byte & 0x10 > 0;
        let ll = (type_byte & 0x0c) >> 2;
        let mm = type_byte & 0x03;

        let la_len = match ll {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

        let ack_len = match mm {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

//        let num_blocks;

        let num_blocks = if n {
            Some(reader.read_u8()?)
        } else {
            None
        };

        let num_ts = reader.read_u8()?;

        let largest_ack = reader.read_uint::<BigEndian>(la_len)?;

        let ack_delay = reader.read_u16::<BigEndian>()?;

        let first_ack_len = reader.read_uint::<BigEndian>(ack_len)?;

        let mut ack_blocks: Vec<AckBlock> = Vec::new();

        if let Some(n_blocks) = num_blocks {
            let ack_block_section_len = (1 + ack_len) * (n_blocks as usize);

            let mut reader_handle = reader.clone().take(ack_block_section_len as u64);

            let mut ack_block_slice = Vec::new();
            reader_handle.read_to_end(&mut ack_block_slice);

            let new_pos = (reader.position() as u64) + (ack_block_section_len as u64);

            reader.set_position(new_pos);

            for block in ack_block_slice.chunks(1 + ack_len) {
                let c_block = AckBlock::from_bytes(&block.to_vec(), ack_len)?;
                ack_blocks.push(c_block);
            }
        }

        let mut timestamps = Vec::new();

        let delta_la;
        let first_ts;
        if num_ts > 0 {
            let delta_la_i = reader.read_u8()?;
            delta_la = Some(delta_la_i);

            let first_ts_i = reader.read_u32::<BigEndian>()?;
            first_ts = Some(first_ts_i);

            let ts_block_section_len = num_ts * 3;

            let mut ts_block_slice = Vec::new();
            let mut reader_handle = reader.clone().take(ts_block_section_len as u64);
            reader_handle.read_to_end(&mut ts_block_slice);

            for block in ts_block_slice.chunks(3) {
                let c_block = AckTimestamp::from_bytes(&block.to_vec())?;
                timestamps.push(c_block);
            }
        } else {
            delta_la = None;
            first_ts = None;
        }

        let ack_blocks_fin = match ack_blocks.len() {
            0 => None,
            _ => Some(ack_blocks),
        };

        let ts_blocks_fin = match timestamps.len() {
            0 => None,
            _ => Some(timestamps),
        };

        Ok(AckFrame {
            num_blocks: num_blocks,
            num_ts: num_ts,
            largest_ack: largest_ack,
            ack_delay: ack_delay,
            first_ack_len: first_ack_len,
            ack_blocks: ack_blocks_fin,
            delta_la: delta_la,
            first_ts: first_ts,
            timestamps: ts_blocks_fin,
        })
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();

        let mut type_byte = super::ACK.bits();

        match self.num_blocks {
            Some(_) => type_byte |= 0x10,
            _ => type_byte |= 0x00
        };

        let largest_ack_size = optimal_field_size(self.largest_ack);

        match largest_ack_size {
            OFSize::U8 => type_byte |= 0x00,
            OFSize::U16 => type_byte |= 0x04,
            OFSize::U32 => type_byte |= 0x08,
            OFSize::U48 => type_byte |= 0x0c,
            _ => panic!("largest ack too large")
        };

        let max_block_len = match self.ack_blocks {
            Some(ref ack_blocks) => ack_blocks.iter().max_by_key(|block| block.block_len),
            None => None
        };

        let block_len_size = match max_block_len {
            None => None,
            Some(max_len_block) => Some(optimal_field_size(max_len_block.block_len))
        };

        match block_len_size {
            None => {},
            Some(OFSize::U8) => type_byte |= 0x00,
            Some(OFSize::U16) => type_byte |= 0x01,
            Some(OFSize::U32) => type_byte |= 0x02,
            Some(OFSize::U48) => type_byte |= 0x03,
            _ => panic!("largest ack too large")
        };

        bytes.write_u8(type_byte);

        if let Some(ref ack_blocks) = self.ack_blocks {
            bytes.write_u8(ack_blocks.len() as u8);
        }

        if let Some(ref ts_blocks) = self.timestamps {
            let mut ts_len = ts_blocks.len();
            if self.delta_la.is_some() {
                ts_len += 0;
            }

            bytes.write_u8(ts_len as u8);
        }

        match largest_ack_size {
            OFSize::U8 => {
                bytes.write_u8(self.largest_ack as u8).unwrap()
            },
            OFSize::U16 => {
                bytes.write_u16::<BigEndian>(self.largest_ack as u16).unwrap()
            },
            OFSize::U32 => {
                bytes.write_u32::<BigEndian>(self.largest_ack as u32).unwrap()
            },
            OFSize::U48 => {
                bytes.write_uint::<BigEndian>(self.largest_ack as u64, 6).unwrap()
            },
            _ => panic!("largest ack too large")
        }

        bytes.write_u16::<BigEndian>(self.ack_delay);

        if let Some(ref ack_blocks) = self.ack_blocks {
            bytes.write_u16::<BigEndian>(ack_blocks.len() as u16);
            let mut ack_block_bytes = Vec::new();
            let ack_block_byte_vectors = ack_blocks.iter().map(|block|
                block.as_bytes(block_len_size.unwrap())
            );

            for ack_block_byte_vector in ack_block_byte_vectors {
                ack_block_bytes.extend(ack_block_byte_vector);
            }

            bytes.extend(ack_block_bytes);

        }

        match self.delta_la {
            None => {},
            Some(la) => { let _ = bytes.write_u8(la); }
        }

        match self.first_ts {
            None => {},
            Some(ts) => { let _ = bytes.write_u32::<BigEndian>(ts); }
        }

        if let Some(ref ack_blocks) = self.timestamps {
            let mut ack_block_bytes = Vec::new();

            let ack_block_byte_vectors = ack_blocks.iter().map(|block|
                block.as_bytes()
            );

            for ack_block_byte_vector in ack_block_byte_vectors {
                ack_block_bytes.extend(ack_block_byte_vector);
            }

            bytes.extend(ack_block_bytes)
        }

        bytes
    }

    pub fn frame_len(buf: &[u8]) -> Result<usize> {
        let mut reader = Cursor::new(buf);

        let type_byte = reader.read_u8()?;

        let mut len = 1 + 1 + 2;

        let n = type_byte & 0x10 > 0;
        let ll = (type_byte & 0x0c) >> 2;
        let mm = type_byte & 0x03;

        if n {
            len += 1;
        }

        let la_len = match ll {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

        len += la_len;

        let ack_len = match mm {
            0 => 1,
            1 => 2,
            3 => 4,
            4 => 6,
            _ => return Err(QuicError::ParseError),
        } as usize;

        len += ack_len;

        if n {
            let num_blocks = reader.read_u8()? as usize;
            len += num_blocks * (1 + ack_len);
        }

        Ok(len)
    }
}

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

    #[test]
    fn serialize() {
        let timestamps = vec![
            AckTimestamp {
                delta_la: 31,
                time_since_prev: AckTimestampValue {
                    microseconds: 26
                },
            },
            AckTimestamp {
                delta_la: 42,
                time_since_prev: AckTimestampValue {
                    microseconds: 290
                },
            },
            AckTimestamp {
                delta_la: 12,
                time_since_prev: AckTimestampValue {
                    microseconds: 4096
                },
            },
            AckTimestamp {
                delta_la: 19,
                time_since_prev: AckTimestampValue {
                    microseconds: 25000
                },
            }
        ];

        let ack_blocks = vec![
            AckBlock {
                gap: 12,
                block_len: 14,
            },
            AckBlock {
                gap: 87,
                block_len: 645,
            },
            AckBlock {
                gap: 42,
                block_len: 325,
            },
            AckBlock {
                gap: 236,
                block_len: 734,
            }
        ];

        let ack_frame = AckFrame {
            num_blocks: Some(4),
            num_ts: 4,
            largest_ack: 497,
            ack_delay: 9,
            first_ack_len: ack_blocks.len() as u64,
            ack_blocks: Some(ack_blocks),
            delta_la: Some(17),
            first_ts: Some(91),
            timestamps: Some(timestamps),
        };

        let ack_frame_bytes = ack_frame.as_bytes();

        let parsed_ack_frame = AckFrame::from_bytes(&ack_frame_bytes).unwrap();

        assert_eq!(ack_frame, parsed_ack_frame);
    }

    #[test]
    fn serialize_timestamp_value() {
        let time_since_prev = 0x1000;
        let microseconds = AckTimestampValue::from_u16(time_since_prev).unwrap();

        assert_eq!(microseconds.microseconds, 4096);
    }
}<|endoftext|>


<commit_before>#![allow(non_snake_case)]
#![crate_type = "rlib"]

extern crate block;
extern crate cocoa;
extern crate core_foundation;
extern crate core_graphics;
extern crate libc;
extern crate objc;

mod classes;
mod constants;
mod functions;
mod protocols;
mod types;

pub use constants::{MTLCompareFunction, MTLFeatureSet, MTLPipelineOption, MTLPixelFormat};

pub use functions::{MTLClearColorMake, MTLCopyAllDevices, MTLCreateSystemDefaultDevice, MTLOriginMake,
                    MTLRegionMake1D, MTLRegionMake2D, MTLRegionMake3D, MTLSizeMake};

pub use protocols::{MTLDevice};

pub use types::{MTLClearColor, MTLDispatchThreadgroupsIndirectArguments, MTLDrawIndexedPrimitivesIndirectArguments,
                MTLDrawPrimitivesIndirectArguments, MTLOrigin, MTLRegion, MTLScissorRect, MTLSize, MTLViewport};
<commit_msg>Need to get objc_send! from objc<commit_after>#![allow(non_snake_case)]
#![crate_type = "rlib"]

extern crate block;
extern crate cocoa;
extern crate core_foundation;
extern crate core_graphics;
extern crate libc;
#[macro_use]
extern crate objc;

mod classes;
mod constants;
mod functions;
mod protocols;
mod types;

pub use constants::{MTLCompareFunction, MTLFeatureSet, MTLPipelineOption, MTLPixelFormat};

pub use functions::{MTLClearColorMake, MTLCopyAllDevices, MTLCreateSystemDefaultDevice, MTLOriginMake,
                    MTLRegionMake1D, MTLRegionMake2D, MTLRegionMake3D, MTLSizeMake};

pub use protocols::{MTLDevice};

pub use types::{MTLClearColor, MTLDispatchThreadgroupsIndirectArguments, MTLDrawIndexedPrimitivesIndirectArguments,
                MTLDrawPrimitivesIndirectArguments, MTLOrigin, MTLRegion, MTLScissorRect, MTLSize, MTLViewport};
<|endoftext|>