<commit_before>use anyhow::*;
use ape;
use id3;
use lewton::inside_ogg::OggStreamReader;
use log::error;
use metaflac;
use mp3_duration;
use mp4ameta;
use regex::Regex;
use std::fs;
use std::path::Path;
use opus_headers;

use crate::utils;
use crate::utils::AudioFormat;

#[derive(Debug, Clone, PartialEq)]
pub struct SongTags {
	pub disc_number: Option<u32>,
	pub track_number: Option<u32>,
	pub title: Option<String>,
	pub duration: Option<u32>,
	pub artist: Option<String>,
	pub album_artist: Option<String>,
	pub album: Option<String>,
	pub year: Option<i32>,
}

#[cfg_attr(feature = "profile-index", flame)]
pub fn read(path: &Path) -> Option<SongTags> {
	let data = match utils::get_audio_format(path) {
		Some(AudioFormat::FLAC) => Some(read_flac(path)),
		Some(AudioFormat::MP3) => Some(read_id3(path)),
		Some(AudioFormat::MP4) => Some(read_mp4(path)),
		Some(AudioFormat::MPC) => Some(read_ape(path)),
		Some(AudioFormat::OGG) => Some(read_vorbis(path)),
		Some(AudioFormat::OPUS) => Some(read_opus(path)),
		_ => None,
	};
	match data {
		Some(Ok(d)) => Some(d),
		Some(Err(e)) => {
			error!("Error while reading file metadata for '{:?}': {}", path, e);
			None
		}
		None => None,
	}
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_id3(path: &Path) -> Result<SongTags> {
	let tag = {
		#[cfg(feature = "profile-index")]
		let _guard = flame::start_guard("id3_tag_read");
		match id3::Tag::read_from_path(&path) {
			Ok(t) => Ok(t),
			Err(e) => {
				if let Some(t) = e.partial_tag {
					Ok(t)
				} else {
					Err(e)
				}
			}
		}?
	};
	let duration = {
		#[cfg(feature = "profile-index")]
		let _guard = flame::start_guard("mp3_duration");
		mp3_duration::from_path(&path)
			.map(|d| d.as_secs() as u32)
			.ok()
	};

	let artist = tag.artist().map(|s| s.to_string());
	let album_artist = tag.album_artist().map(|s| s.to_string());
	let album = tag.album().map(|s| s.to_string());
	let title = tag.title().map(|s| s.to_string());
	let disc_number = tag.disc();
	let track_number = tag.track();
	let year = tag
		.year()
		.map(|y| y as i32)
		.or_else(|| tag.date_released().and_then(|d| Some(d.year)))
		.or_else(|| tag.date_recorded().and_then(|d| Some(d.year)));

	Ok(SongTags {
		artist,
		album_artist,
		album,
		title,
		duration,
		disc_number,
		track_number,
		year,
	})
}

fn read_ape_string(item: &ape::Item) -> Option<String> {
	match item.value {
		ape::ItemValue::Text(ref s) => Some(s.clone()),
		_ => None,
	}
}

fn read_ape_i32(item: &ape::Item) -> Option<i32> {
	match item.value {
		ape::ItemValue::Text(ref s) => s.parse::<i32>().ok(),
		_ => None,
	}
}

fn read_ape_x_of_y(item: &ape::Item) -> Option<u32> {
	match item.value {
		ape::ItemValue::Text(ref s) => {
			let format = Regex::new(r#"^\d+"#).unwrap();
			if let Some(m) = format.find(s) {
				s[m.start()..m.end()].parse().ok()
			} else {
				None
			}
		}
		_ => None,
	}
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_ape(path: &Path) -> Result<SongTags> {
	let tag = ape::read(path)?;
	let artist = tag.item("Artist").and_then(read_ape_string);
	let album = tag.item("Album").and_then(read_ape_string);
	let album_artist = tag.item("Album artist").and_then(read_ape_string);
	let title = tag.item("Title").and_then(read_ape_string);
	let year = tag.item("Year").and_then(read_ape_i32);
	let disc_number = tag.item("Disc").and_then(read_ape_x_of_y);
	let track_number = tag.item("Track").and_then(read_ape_x_of_y);
	Ok(SongTags {
		artist,
		album_artist,
		album,
		title,
		duration: None,
		disc_number,
		track_number,
		year,
	})
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_vorbis(path: &Path) -> Result<SongTags> {
	let file = fs::File::open(path)?;
	let source = OggStreamReader::new(file)?;

	let mut tags = SongTags {
		artist: None,
		album_artist: None,
		album: None,
		title: None,
		duration: None,
		disc_number: None,
		track_number: None,
		year: None,
	};

	for (key, value) in source.comment_hdr.comment_list {
		match key.as_str() {
			"TITLE" => tags.title = Some(value),
			"ALBUM" => tags.album = Some(value),
			"ARTIST" => tags.artist = Some(value),
			"ALBUMARTIST" => tags.album_artist = Some(value),
			"TRACKNUMBER" => tags.track_number = value.parse::<u32>().ok(),
			"DISCNUMBER" => tags.disc_number = value.parse::<u32>().ok(),
			"DATE" => tags.year = value.parse::<i32>().ok(),
			_ => (),
		}
	}

	Ok(tags)
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_opus(path: &Path) -> Result<SongTags> {
	let headers = opus_headers::parse_from_path(path)?;

	let mut tags = SongTags {
		artist: None,
		album_artist: None,
		album: None,
		title: None,
		duration: None,
		disc_number: None,
		track_number: None,
		year: None,
	};

	for (key, value) in headers.comments.user_comments {
		if "TITLE".eq_ignore_ascii_case(&key) {
			tags.title = Some(value);
		} else if "ALBUM".eq_ignore_ascii_case(key.as_str()) {
			tags.album = Some(value);
		} else if "ARTIST".eq_ignore_ascii_case(key.as_str()) {
			tags.artist = Some(value);
		} else if "ALBUMARTIST".eq_ignore_ascii_case(key.as_str()) {
			tags.album_artist = Some(value);
		} else if "TRACKNUMBER".eq_ignore_ascii_case(key.as_str()) {
			tags.track_number = value.parse::<u32>().ok();
		} else if "DISCNUMBER".eq_ignore_ascii_case(key.as_str()) {
			tags.disc_number = value.parse::<u32>().ok();
		} else if "DATE".eq_ignore_ascii_case(key.as_str()) {
			tags.year = value.parse::<i32>().ok();
		}
	}

	Ok(tags)
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_flac(path: &Path) -> Result<SongTags> {
	let tag = metaflac::Tag::read_from_path(path)?;
	let vorbis = tag
		.vorbis_comments()
		.ok_or(anyhow!("Missing Vorbis comments"))?;
	let disc_number = vorbis
		.get("DISCNUMBER")
		.and_then(|d| d[0].parse::<u32>().ok());
	let year = vorbis.get("DATE").and_then(|d| d[0].parse::<i32>().ok());
	let mut streaminfo = tag.get_blocks(metaflac::BlockType::StreamInfo);
	let duration = match streaminfo.next() {
		Some(&metaflac::Block::StreamInfo(ref s)) => {
			Some((s.total_samples as u32 / s.sample_rate) as u32)
		}
		_ => None,
	};

	Ok(SongTags {
		artist: vorbis.artist().map(|v| v[0].clone()),
		album_artist: vorbis.album_artist().map(|v| v[0].clone()),
		album: vorbis.album().map(|v| v[0].clone()),
		title: vorbis.title().map(|v| v[0].clone()),
		duration,
		disc_number,
		track_number: vorbis.track(),
		year,
	})
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_mp4(path: &Path) -> Result<SongTags> {
	let tag = mp4ameta::Tag::read_from_path(path)?;
	let mut tags = SongTags {
		artist: None,
		album_artist: None,
		album: None,
		title: None,
		duration: None,
		disc_number: None,
		track_number: None,
		year: None,
	};
	tags.artist = tag.artist().map(|v| v.to_string());
	tags.album_artist = tag.album_artist().map(|v| v.to_string());
	tags.album = tag.album().map(|v| v.to_string());
	tags.title = tag.title().map(|v| v.to_string());
	tags.duration = tag.duration().map(|v| v as u32);
	tags.disc_number = tag.disk_number().0.and_then(|d| Some(d as u32));
	tags.track_number = tag.track_number().0.and_then(|d| Some(d as u32));
	tags.year = tag.year().and_then(|v| v.parse::<i32>().ok());

	Ok(tags)
}

#[test]
fn test_read_metadata() {
	let sample_tags = SongTags {
		disc_number: Some(3),
		track_number: Some(1),
		title: Some("TEST TITLE".into()),
		artist: Some("TEST ARTIST".into()),
		album_artist: Some("TEST ALBUM ARTIST".into()),
		album: Some("TEST ALBUM".into()),
		duration: None,
		year: Some(2016),
	};
	let flac_sample_tag = SongTags {
		duration: Some(0),
		..sample_tags.clone()
	};
	let mp3_sample_tag = SongTags {
		duration: Some(0),
		..sample_tags.clone()
	};
	let m4a_sample_tag = SongTags {
		duration: Some(0),
		..sample_tags.clone()
	};
	assert_eq!(read(Path::new("test/sample.mp3")).unwrap(), mp3_sample_tag);
	assert_eq!(read(Path::new("test/sample.ogg")).unwrap(), sample_tags);
	assert_eq!(read(Path::new("test/sample.flac")).unwrap(), flac_sample_tag);
	assert_eq!(read(Path::new("test/sample.m4a")).unwrap(), m4a_sample_tag);
}
<commit_msg>Update src/index/metadata.rs<commit_after>use anyhow::*;
use ape;
use id3;
use lewton::inside_ogg::OggStreamReader;
use log::error;
use metaflac;
use mp3_duration;
use mp4ameta;
use regex::Regex;
use std::fs;
use std::path::Path;
use opus_headers;

use crate::utils;
use crate::utils::AudioFormat;

#[derive(Debug, Clone, PartialEq)]
pub struct SongTags {
	pub disc_number: Option<u32>,
	pub track_number: Option<u32>,
	pub title: Option<String>,
	pub duration: Option<u32>,
	pub artist: Option<String>,
	pub album_artist: Option<String>,
	pub album: Option<String>,
	pub year: Option<i32>,
}

#[cfg_attr(feature = "profile-index", flame)]
pub fn read(path: &Path) -> Option<SongTags> {
	let data = match utils::get_audio_format(path) {
		Some(AudioFormat::FLAC) => Some(read_flac(path)),
		Some(AudioFormat::MP3) => Some(read_id3(path)),
		Some(AudioFormat::MP4) => Some(read_mp4(path)),
		Some(AudioFormat::MPC) => Some(read_ape(path)),
		Some(AudioFormat::OGG) => Some(read_vorbis(path)),
		Some(AudioFormat::OPUS) => Some(read_opus(path)),
		_ => None,
	};
	match data {
		Some(Ok(d)) => Some(d),
		Some(Err(e)) => {
			error!("Error while reading file metadata for '{:?}': {}", path, e);
			None
		}
		None => None,
	}
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_id3(path: &Path) -> Result<SongTags> {
	let tag = {
		#[cfg(feature = "profile-index")]
		let _guard = flame::start_guard("id3_tag_read");
		match id3::Tag::read_from_path(&path) {
			Ok(t) => Ok(t),
			Err(e) => {
				if let Some(t) = e.partial_tag {
					Ok(t)
				} else {
					Err(e)
				}
			}
		}?
	};
	let duration = {
		#[cfg(feature = "profile-index")]
		let _guard = flame::start_guard("mp3_duration");
		mp3_duration::from_path(&path)
			.map(|d| d.as_secs() as u32)
			.ok()
	};

	let artist = tag.artist().map(|s| s.to_string());
	let album_artist = tag.album_artist().map(|s| s.to_string());
	let album = tag.album().map(|s| s.to_string());
	let title = tag.title().map(|s| s.to_string());
	let disc_number = tag.disc();
	let track_number = tag.track();
	let year = tag
		.year()
		.map(|y| y as i32)
		.or_else(|| tag.date_released().and_then(|d| Some(d.year)))
		.or_else(|| tag.date_recorded().and_then(|d| Some(d.year)));

	Ok(SongTags {
		artist,
		album_artist,
		album,
		title,
		duration,
		disc_number,
		track_number,
		year,
	})
}

fn read_ape_string(item: &ape::Item) -> Option<String> {
	match item.value {
		ape::ItemValue::Text(ref s) => Some(s.clone()),
		_ => None,
	}
}

fn read_ape_i32(item: &ape::Item) -> Option<i32> {
	match item.value {
		ape::ItemValue::Text(ref s) => s.parse::<i32>().ok(),
		_ => None,
	}
}

fn read_ape_x_of_y(item: &ape::Item) -> Option<u32> {
	match item.value {
		ape::ItemValue::Text(ref s) => {
			let format = Regex::new(r#"^\d+"#).unwrap();
			if let Some(m) = format.find(s) {
				s[m.start()..m.end()].parse().ok()
			} else {
				None
			}
		}
		_ => None,
	}
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_ape(path: &Path) -> Result<SongTags> {
	let tag = ape::read(path)?;
	let artist = tag.item("Artist").and_then(read_ape_string);
	let album = tag.item("Album").and_then(read_ape_string);
	let album_artist = tag.item("Album artist").and_then(read_ape_string);
	let title = tag.item("Title").and_then(read_ape_string);
	let year = tag.item("Year").and_then(read_ape_i32);
	let disc_number = tag.item("Disc").and_then(read_ape_x_of_y);
	let track_number = tag.item("Track").and_then(read_ape_x_of_y);
	Ok(SongTags {
		artist,
		album_artist,
		album,
		title,
		duration: None,
		disc_number,
		track_number,
		year,
	})
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_vorbis(path: &Path) -> Result<SongTags> {
	let file = fs::File::open(path)?;
	let source = OggStreamReader::new(file)?;

	let mut tags = SongTags {
		artist: None,
		album_artist: None,
		album: None,
		title: None,
		duration: None,
		disc_number: None,
		track_number: None,
		year: None,
	};

	for (key, value) in source.comment_hdr.comment_list {
		match key.as_str() {
			"TITLE" => tags.title = Some(value),
			"ALBUM" => tags.album = Some(value),
			"ARTIST" => tags.artist = Some(value),
			"ALBUMARTIST" => tags.album_artist = Some(value),
			"TRACKNUMBER" => tags.track_number = value.parse::<u32>().ok(),
			"DISCNUMBER" => tags.disc_number = value.parse::<u32>().ok(),
			"DATE" => tags.year = value.parse::<i32>().ok(),
			_ => (),
		}
	}

	Ok(tags)
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_opus(path: &Path) -> Result<SongTags> {
	let headers = opus_headers::parse_from_path(path)?;

	let mut tags = SongTags {
		artist: None,
		album_artist: None,
		album: None,
		title: None,
		duration: None,
		disc_number: None,
		track_number: None,
		year: None,
	};

	for (key, value) in headers.comments.user_comments {
		if "TITLE".eq_ignore_ascii_case(&key) {
			tags.title = Some(value);
		} else if "ALBUM".eq_ignore_ascii_case(key.as_str()) {
			tags.album = Some(value);
		} else if "ARTIST".eq_ignore_ascii_case(key.as_str()) {
			tags.artist = Some(value);
		} else if "ALBUMARTIST".eq_ignore_ascii_case(key.as_str()) {
			tags.album_artist = Some(value);
		} else if "TRACKNUMBER".eq_ignore_ascii_case(key.as_str()) {
			tags.track_number = value.parse().ok();
		} else if "DISCNUMBER".eq_ignore_ascii_case(key.as_str()) {
			tags.disc_number = value.parse::<u32>().ok();
		} else if "DATE".eq_ignore_ascii_case(key.as_str()) {
			tags.year = value.parse::<i32>().ok();
		}
	}

	Ok(tags)
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_flac(path: &Path) -> Result<SongTags> {
	let tag = metaflac::Tag::read_from_path(path)?;
	let vorbis = tag
		.vorbis_comments()
		.ok_or(anyhow!("Missing Vorbis comments"))?;
	let disc_number = vorbis
		.get("DISCNUMBER")
		.and_then(|d| d[0].parse::<u32>().ok());
	let year = vorbis.get("DATE").and_then(|d| d[0].parse::<i32>().ok());
	let mut streaminfo = tag.get_blocks(metaflac::BlockType::StreamInfo);
	let duration = match streaminfo.next() {
		Some(&metaflac::Block::StreamInfo(ref s)) => {
			Some((s.total_samples as u32 / s.sample_rate) as u32)
		}
		_ => None,
	};

	Ok(SongTags {
		artist: vorbis.artist().map(|v| v[0].clone()),
		album_artist: vorbis.album_artist().map(|v| v[0].clone()),
		album: vorbis.album().map(|v| v[0].clone()),
		title: vorbis.title().map(|v| v[0].clone()),
		duration,
		disc_number,
		track_number: vorbis.track(),
		year,
	})
}

#[cfg_attr(feature = "profile-index", flame)]
fn read_mp4(path: &Path) -> Result<SongTags> {
	let tag = mp4ameta::Tag::read_from_path(path)?;
	let mut tags = SongTags {
		artist: None,
		album_artist: None,
		album: None,
		title: None,
		duration: None,
		disc_number: None,
		track_number: None,
		year: None,
	};
	tags.artist = tag.artist().map(|v| v.to_string());
	tags.album_artist = tag.album_artist().map(|v| v.to_string());
	tags.album = tag.album().map(|v| v.to_string());
	tags.title = tag.title().map(|v| v.to_string());
	tags.duration = tag.duration().map(|v| v as u32);
	tags.disc_number = tag.disk_number().0.and_then(|d| Some(d as u32));
	tags.track_number = tag.track_number().0.and_then(|d| Some(d as u32));
	tags.year = tag.year().and_then(|v| v.parse::<i32>().ok());

	Ok(tags)
}

#[test]
fn test_read_metadata() {
	let sample_tags = SongTags {
		disc_number: Some(3),
		track_number: Some(1),
		title: Some("TEST TITLE".into()),
		artist: Some("TEST ARTIST".into()),
		album_artist: Some("TEST ALBUM ARTIST".into()),
		album: Some("TEST ALBUM".into()),
		duration: None,
		year: Some(2016),
	};
	let flac_sample_tag = SongTags {
		duration: Some(0),
		..sample_tags.clone()
	};
	let mp3_sample_tag = SongTags {
		duration: Some(0),
		..sample_tags.clone()
	};
	let m4a_sample_tag = SongTags {
		duration: Some(0),
		..sample_tags.clone()
	};
	assert_eq!(read(Path::new("test/sample.mp3")).unwrap(), mp3_sample_tag);
	assert_eq!(read(Path::new("test/sample.ogg")).unwrap(), sample_tags);
	assert_eq!(read(Path::new("test/sample.flac")).unwrap(), flac_sample_tag);
	assert_eq!(read(Path::new("test/sample.m4a")).unwrap(), m4a_sample_tag);
}
<|endoftext|>


<commit_before>// Module for reading sprides and saving them in a HashMap

use std::collections::HashMap;
use piston_window::*;
use io::sprite::Sprite;

/// HashMap for all given sprites
/// map: HashMap with all sprites with PNG-filename as key
pub struct SpriteMap {
    pub map: HashMap<&'static str, Sprite>,
}

impl SpriteMap {
    /// Initialize the SpriteMap with given window, fills Map with all Sprites
    pub fn init(mut w: &mut PistonWindow) -> Self {
        // Initialize
        let mut map2: HashMap<&'static str, Sprite> = HashMap::new();
        // Fill with all Sprites
        for i in SPRITES {
            map2.insert(i.image,
                        Sprite::fill_sprite(i.image,
                                            i.size_x as u32,
                                            i.size_y as u32,
                                            i.dur as u64,
                                            i.once,
                                            w));
        }
        // Create SpriteMap
        SpriteMap { map: map2 }
    }

    /// creates empty SpriteMap
    pub fn new() -> Self {
        SpriteMap { map: HashMap::new() }
    }

    /// getter for Sprite
    pub fn get_sprite(&self, file: String) -> Option<&Sprite> {
        self.map.get(file.as_str())
    }
}

/// Struct for representing Necessary information about sprites
struct SpriteModel {
    image: &'static str,
    size_x: usize,
    size_y: usize,
    dur: usize,
    once: bool,
}

/// const with all Sprites
const SPRITES: &'static [SpriteModel] = &[SpriteModel {
                                              image: "bowman.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken_brown.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken_pink.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken_white.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "explosion.png",
                                              size_x: 59,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "knight.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "paladin.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "swipe_broadsword.png",
                                              size_x: 64,
                                              size_y: 192,
                                              dur: 100,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "swipe_dagger.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 100,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "swipe_longsword.png",
                                              size_x: 128,
                                              size_y: 64,
                                              dur: 100,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "weapon_broadsword.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_dagger.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_spear.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "wizard.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_empty.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_full.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_half.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_broadsword_small.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_dagger_small.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_spear_small.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_10.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_20.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "swipe_enemy.png",
                                              size_x: 72,
                                              size_y: 64,
                                              dur: 250,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "explosion_square.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 500,
                                              once: true,
                                          }];
<commit_msg>add chest<commit_after>// Module for reading sprides and saving them in a HashMap

use std::collections::HashMap;
use piston_window::*;
use io::sprite::Sprite;

/// HashMap for all given sprites
/// map: HashMap with all sprites with PNG-filename as key
pub struct SpriteMap {
    pub map: HashMap<&'static str, Sprite>,
}

impl SpriteMap {
    /// Initialize the SpriteMap with given window, fills Map with all Sprites
    pub fn init(mut w: &mut PistonWindow) -> Self {
        // Initialize
        let mut map2: HashMap<&'static str, Sprite> = HashMap::new();
        // Fill with all Sprites
        for i in SPRITES {
            map2.insert(i.image,
                        Sprite::fill_sprite(i.image,
                                            i.size_x as u32,
                                            i.size_y as u32,
                                            i.dur as u64,
                                            i.once,
                                            w));
        }
        // Create SpriteMap
        SpriteMap { map: map2 }
    }

    /// creates empty SpriteMap
    pub fn new() -> Self {
        SpriteMap { map: HashMap::new() }
    }

    /// getter for Sprite
    pub fn get_sprite(&self, file: String) -> Option<&Sprite> {
        self.map.get(file.as_str())
    }
}

/// Struct for representing Necessary information about sprites
struct SpriteModel {
    image: &'static str,
    size_x: usize,
    size_y: usize,
    dur: usize,
    once: bool,
}

/// const with all Sprites
const SPRITES: &'static [SpriteModel] = &[SpriteModel {
                                              image: "bowman.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken_brown.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken_pink.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chicken_white.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "explosion.png",
                                              size_x: 59,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "knight.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "paladin.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "swipe_broadsword.png",
                                              size_x: 64,
                                              size_y: 192,
                                              dur: 100,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "swipe_dagger.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 100,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "swipe_longsword.png",
                                              size_x: 128,
                                              size_y: 64,
                                              dur: 100,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "weapon_broadsword.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_dagger.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_spear.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "wizard.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_empty.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_full.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_half.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_broadsword_small.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_dagger_small.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "weapon_spear_small.png",
                                              size_x: 32,
                                              size_y: 32,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_10.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "Heart_20.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "chest.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 1000,
                                              once: false,
                                          },
                                          SpriteModel {
                                              image: "swipe_enemy.png",
                                              size_x: 72,
                                              size_y: 64,
                                              dur: 250,
                                              once: true,
                                          },
                                          SpriteModel {
                                              image: "explosion_square.png",
                                              size_x: 64,
                                              size_y: 64,
                                              dur: 500,
                                              once: true,
                                          }];
<|endoftext|>


<commit_before>use super::*;
use std::iter::Fuse;

pub struct Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    a: F,
    v: ArrayBase<S, Ix1>,
    ortho: Ortho,
}

impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    pub fn new(a: F, v: ArrayBase<S, Ix1>, ortho: Ortho) -> Fuse<Self> {
        Iterator::fuse(Arnoldi { a, v, ortho })
    }
}

impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    type Item = Array1<A>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.a)(&mut self.v);
        match self.ortho.div_append(&mut self.v) {
            AppendResult::Added(coef) => {
                let norm = coef[coef.len() - 1].abs();
                azip!(mut a(&mut self.v) in { *a = a.div_real(norm) });
                Some(coef)
            }
            AppendResult::Dependent(_) => None,
        }
    }
}
<commit_msg>arnodli_{mgs,householder}<commit_after>use super::*;
use crate::norm::Norm;
use num_traits::One;
use std::iter::Fuse;

pub struct Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    a: F,
    v: ArrayBase<S, Ix1>,
    ortho: Ortho,
}

impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    pub fn new(a: F, mut v: ArrayBase<S, Ix1>, mut ortho: Ortho) -> Fuse<Self> {
        assert_eq!(ortho.len(), 0);
        assert!(ortho.tolerance() < One::one());
        // normalize before append because |v| may be smaller than ortho.tolerance()
        let norm = v.norm_l2();
        azip!(mut v(&mut v) in { *v = v.div_real(norm) });
        ortho.append(v.view());
        Iterator::fuse(Arnoldi { a, v, ortho })
    }

    /// Dimension of Krylov subspace
    pub fn krylov_dimension(&self) -> usize {
        self.ortho.len()
    }
}

impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    type Item = Array1<A>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.a)(&mut self.v);
        let result = self.ortho.div_append(&mut self.v);
        azip!(mut v(&mut self.v) in { *v = v.div_real(result.residual_norm()) });
        if result.is_dependent() {
            None
        } else {
            Some(result.into_coeff())
        }
    }
}

pub fn arnoldi_householder<A, S1, S2>(
    a: ArrayBase<S1, Ix2>,
    v: ArrayBase<S2, Ix1>,
    tol: A::Real,
) -> impl Iterator<Item = Array1<A>>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    assert_eq!(m, v.len(), "Input matrix and vector sizes mismach");

    let householder = Householder::new(n, tol);
    Arnoldi::new(
        move |x| {
            let ax = a.dot(x);
            azip!(mut x(x), ax in { *x = ax });
        },
        v,
        householder,
    )
}

pub fn arnoldi_mgs<A, S1, S2>(
    a: ArrayBase<S1, Ix2>,
    v: ArrayBase<S2, Ix1>,
    tol: A::Real,
) -> impl Iterator<Item = Array1<A>>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    assert_eq!(m, v.len(), "Input matrix and vector sizes mismach");

    let mgs = MGS::new(n, tol);
    Arnoldi::new(
        move |x| {
            let ax = a.dot(x);
            azip!(mut x(x), ax in { *x = ax });
        },
        v,
        mgs,
    )
}
<|endoftext|>


<commit_before>//! Memory allocation APIs

#![stable(feature = "alloc_module", since = "1.28.0")]

use core::intrinsics::{self, min_align_of_val, size_of_val};
use core::ptr::{NonNull, Unique};

#[stable(feature = "alloc_module", since = "1.28.0")]
#[doc(inline)]
pub use core::alloc::*;

#[cfg(test)]
mod tests;

extern "Rust" {
    // These are the magic symbols to call the global allocator.  rustc generates
    // them from the `#[global_allocator]` attribute if there is one, or uses the
    // default implementations in libstd (`__rdl_alloc` etc in `src/libstd/alloc.rs`)
    // otherwise.
    #[rustc_allocator]
    #[rustc_allocator_nounwind]
    fn __rust_alloc(size: usize, align: usize) -> *mut u8;
    #[rustc_allocator_nounwind]
    fn __rust_dealloc(ptr: *mut u8, size: usize, align: usize);
    #[rustc_allocator_nounwind]
    fn __rust_realloc(ptr: *mut u8, old_size: usize, align: usize, new_size: usize) -> *mut u8;
    #[rustc_allocator_nounwind]
    fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8;
}

/// The global memory allocator.
///
/// This type implements the [`AllocRef`] trait by forwarding calls
/// to the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// Note: while this type is unstable, the functionality it provides can be
/// accessed through the [free functions in `alloc`](index.html#functions).
///
/// [`AllocRef`]: trait.AllocRef.html
#[unstable(feature = "allocator_api", issue = "32838")]
#[derive(Copy, Clone, Default, Debug)]
pub struct Global;

/// Allocate memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::alloc`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `alloc` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::alloc`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::alloc`]: trait.GlobalAlloc.html#tymethod.alloc
///
/// # Examples
///
/// ```
/// use std::alloc::{alloc, dealloc, Layout};
///
/// unsafe {
///     let layout = Layout::new::<u16>();
///     let ptr = alloc(layout);
///
///     *(ptr as *mut u16) = 42;
///     assert_eq!(*(ptr as *mut u16), 42);
///
///     dealloc(ptr, layout);
/// }
/// ```
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn alloc(layout: Layout) -> *mut u8 {
    unsafe { __rust_alloc(layout.size(), layout.align()) }
}

/// Deallocate memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::dealloc`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `dealloc` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::dealloc`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::dealloc`]: trait.GlobalAlloc.html#tymethod.dealloc
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
    unsafe { __rust_dealloc(ptr, layout.size(), layout.align()) }
}

/// Reallocate memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::realloc`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `realloc` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::realloc`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::realloc`]: trait.GlobalAlloc.html#method.realloc
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
    unsafe { __rust_realloc(ptr, layout.size(), layout.align(), new_size) }
}

/// Allocate zero-initialized memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::alloc_zeroed`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `alloc_zeroed` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::alloc_zeroed`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::alloc_zeroed`]: trait.GlobalAlloc.html#method.alloc_zeroed
///
/// # Examples
///
/// ```
/// use std::alloc::{alloc_zeroed, dealloc, Layout};
///
/// unsafe {
///     let layout = Layout::new::<u16>();
///     let ptr = alloc_zeroed(layout);
///
///     assert_eq!(*(ptr as *mut u16), 0);
///
///     dealloc(ptr, layout);
/// }
/// ```
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
    unsafe { __rust_alloc_zeroed(layout.size(), layout.align()) }
}

#[unstable(feature = "allocator_api", issue = "32838")]
unsafe impl AllocRef for Global {
    #[inline]
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
        unsafe {
            let size = layout.size();
            if size == 0 {
                Ok(MemoryBlock { ptr: layout.dangling(), size: 0 })
            } else {
                let raw_ptr = match init {
                    AllocInit::Uninitialized => alloc(layout),
                    AllocInit::Zeroed => alloc_zeroed(layout),
                };
                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
                Ok(MemoryBlock { ptr, size })
            }
        }
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        if layout.size() != 0 {
            unsafe { dealloc(ptr.as_ptr(), layout) }
        }
    }

    #[inline]
    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
        init: AllocInit,
    ) -> Result<MemoryBlock, AllocErr> {
        let size = layout.size();
        debug_assert!(
            new_size >= size,
            "`new_size` must be greater than or equal to `memory.size()`"
        );

        if size == new_size {
            return Ok(MemoryBlock { ptr, size });
        }

        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove if layout.size() == 0 => {
                let new_layout =
                    unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
                self.alloc(new_layout, init)
            }
            ReallocPlacement::MayMove => {
                // `realloc` probably checks for `new_size > size` or something similar.
                unsafe {
                    intrinsics::assume(new_size > size);
                    let ptr = realloc(ptr.as_ptr(), layout, new_size);
                    let memory =
                        MemoryBlock { ptr: NonNull::new(ptr).ok_or(AllocErr)?, size: new_size };
                    init.init_offset(memory, size);
                    Ok(memory)
                }
            }
        }
    }

    #[inline]
    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
    ) -> Result<MemoryBlock, AllocErr> {
        let size = layout.size();
        debug_assert!(
            new_size <= size,
            "`new_size` must be smaller than or equal to `memory.size()`"
        );

        if size == new_size {
            return Ok(MemoryBlock { ptr, size });
        }

        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove if new_size == 0 => {
                unsafe {
                    self.dealloc(ptr, layout);
                }
                Ok(MemoryBlock { ptr: layout.dangling(), size: 0 })
            }
            ReallocPlacement::MayMove => {
                // `realloc` probably checks for `new_size < size` or something similar.
                unsafe {
                    intrinsics::assume(new_size < size);
                    let ptr = realloc(ptr.as_ptr(), layout, new_size);
                    Ok(MemoryBlock { ptr: NonNull::new(ptr).ok_or(AllocErr)?, size: new_size })
                }
            }
        }
    }
}

/// The allocator for unique pointers.
// This function must not unwind. If it does, MIR codegen will fail.
#[cfg(not(test))]
#[lang = "exchange_malloc"]
#[inline]
unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
    let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
    match Global.alloc(layout, AllocInit::Uninitialized) {
        Ok(memory) => memory.ptr.as_ptr(),
        Err(_) => handle_alloc_error(layout),
    }
}

#[cfg_attr(not(test), lang = "box_free")]
#[inline]
// This signature has to be the same as `Box`, otherwise an ICE will happen.
// When an additional parameter to `Box` is added (like `A: AllocRef`), this has to be added here as
// well.
// For example if `Box` is changed to  `struct Box<T: ?Sized, A: AllocRef>(Unique<T>, A)`,
// this function has to be changed to `fn box_free<T: ?Sized, A: AllocRef>(Unique<T>, A)` as well.
pub(crate) unsafe fn box_free<T: ?Sized>(ptr: Unique<T>) {
    unsafe {
        let size = size_of_val(ptr.as_ref());
        let align = min_align_of_val(ptr.as_ref());
        let layout = Layout::from_size_align_unchecked(size, align);
        Global.dealloc(ptr.cast().into(), layout)
    }
}

/// Abort on memory allocation error or failure.
///
/// Callers of memory allocation APIs wishing to abort computation
/// in response to an allocation error are encouraged to call this function,
/// rather than directly invoking `panic!` or similar.
///
/// The default behavior of this function is to print a message to standard error
/// and abort the process.
/// It can be replaced with [`set_alloc_error_hook`] and [`take_alloc_error_hook`].
///
/// [`set_alloc_error_hook`]: ../../std/alloc/fn.set_alloc_error_hook.html
/// [`take_alloc_error_hook`]: ../../std/alloc/fn.take_alloc_error_hook.html
#[stable(feature = "global_alloc", since = "1.28.0")]
#[rustc_allocator_nounwind]
pub fn handle_alloc_error(layout: Layout) -> ! {
    extern "Rust" {
        #[lang = "oom"]
        fn oom_impl(layout: Layout) -> !;
    }
    unsafe { oom_impl(layout) }
}
<commit_msg>Apply suggestions from code review<commit_after>//! Memory allocation APIs

#![stable(feature = "alloc_module", since = "1.28.0")]

use core::intrinsics::{self, min_align_of_val, size_of_val};
use core::ptr::{NonNull, Unique};

#[stable(feature = "alloc_module", since = "1.28.0")]
#[doc(inline)]
pub use core::alloc::*;

#[cfg(test)]
mod tests;

extern "Rust" {
    // These are the magic symbols to call the global allocator.  rustc generates
    // them from the `#[global_allocator]` attribute if there is one, or uses the
    // default implementations in libstd (`__rdl_alloc` etc in `src/libstd/alloc.rs`)
    // otherwise.
    #[rustc_allocator]
    #[rustc_allocator_nounwind]
    fn __rust_alloc(size: usize, align: usize) -> *mut u8;
    #[rustc_allocator_nounwind]
    fn __rust_dealloc(ptr: *mut u8, size: usize, align: usize);
    #[rustc_allocator_nounwind]
    fn __rust_realloc(ptr: *mut u8, old_size: usize, align: usize, new_size: usize) -> *mut u8;
    #[rustc_allocator_nounwind]
    fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8;
}

/// The global memory allocator.
///
/// This type implements the [`AllocRef`] trait by forwarding calls
/// to the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// Note: while this type is unstable, the functionality it provides can be
/// accessed through the [free functions in `alloc`](index.html#functions).
///
/// [`AllocRef`]: trait.AllocRef.html
#[unstable(feature = "allocator_api", issue = "32838")]
#[derive(Copy, Clone, Default, Debug)]
pub struct Global;

/// Allocate memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::alloc`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `alloc` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::alloc`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::alloc`]: trait.GlobalAlloc.html#tymethod.alloc
///
/// # Examples
///
/// ```
/// use std::alloc::{alloc, dealloc, Layout};
///
/// unsafe {
///     let layout = Layout::new::<u16>();
///     let ptr = alloc(layout);
///
///     *(ptr as *mut u16) = 42;
///     assert_eq!(*(ptr as *mut u16), 42);
///
///     dealloc(ptr, layout);
/// }
/// ```
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn alloc(layout: Layout) -> *mut u8 {
    unsafe { __rust_alloc(layout.size(), layout.align()) }
}

/// Deallocate memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::dealloc`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `dealloc` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::dealloc`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::dealloc`]: trait.GlobalAlloc.html#tymethod.dealloc
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
    unsafe { __rust_dealloc(ptr, layout.size(), layout.align()) }
}

/// Reallocate memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::realloc`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `realloc` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::realloc`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::realloc`]: trait.GlobalAlloc.html#method.realloc
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
    unsafe { __rust_realloc(ptr, layout.size(), layout.align(), new_size) }
}

/// Allocate zero-initialized memory with the global allocator.
///
/// This function forwards calls to the [`GlobalAlloc::alloc_zeroed`] method
/// of the allocator registered with the `#[global_allocator]` attribute
/// if there is one, or the `std` crate’s default.
///
/// This function is expected to be deprecated in favor of the `alloc_zeroed` method
/// of the [`Global`] type when it and the [`AllocRef`] trait become stable.
///
/// # Safety
///
/// See [`GlobalAlloc::alloc_zeroed`].
///
/// [`Global`]: struct.Global.html
/// [`AllocRef`]: trait.AllocRef.html
/// [`GlobalAlloc::alloc_zeroed`]: trait.GlobalAlloc.html#method.alloc_zeroed
///
/// # Examples
///
/// ```
/// use std::alloc::{alloc_zeroed, dealloc, Layout};
///
/// unsafe {
///     let layout = Layout::new::<u16>();
///     let ptr = alloc_zeroed(layout);
///
///     assert_eq!(*(ptr as *mut u16), 0);
///
///     dealloc(ptr, layout);
/// }
/// ```
#[stable(feature = "global_alloc", since = "1.28.0")]
#[inline]
pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
    unsafe { __rust_alloc_zeroed(layout.size(), layout.align()) }
}

#[unstable(feature = "allocator_api", issue = "32838")]
unsafe impl AllocRef for Global {
    #[inline]
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
        unsafe {
            let size = layout.size();
            if size == 0 {
                Ok(MemoryBlock { ptr: layout.dangling(), size: 0 })
            } else {
                let raw_ptr = match init {
                    AllocInit::Uninitialized => alloc(layout),
                    AllocInit::Zeroed => alloc_zeroed(layout),
                };
                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
                Ok(MemoryBlock { ptr, size })
            }
        }
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        if layout.size() != 0 {
            unsafe { dealloc(ptr.as_ptr(), layout) }
        }
    }

    #[inline]
    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
        init: AllocInit,
    ) -> Result<MemoryBlock, AllocErr> {
        let size = layout.size();
        debug_assert!(
            new_size >= size,
            "`new_size` must be greater than or equal to `memory.size()`"
        );

        if size == new_size {
            return Ok(MemoryBlock { ptr, size });
        }

        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove if layout.size() == 0 => {
                let new_layout =
                    unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
                self.alloc(new_layout, init)
            }
            ReallocPlacement::MayMove => {
                // `realloc` probably checks for `new_size > size` or something similar.
                let ptr = unsafe {
                    intrinsics::assume(new_size > size);
                    realloc(ptr.as_ptr(), layout, new_size)
                };
                let memory =
                    MemoryBlock { ptr: NonNull::new(ptr).ok_or(AllocErr)?, size: new_size };
                unsafe {
                    init.init_offset(memory, size);
                }
                Ok(memory)
            }
        }
    }

    #[inline]
    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
    ) -> Result<MemoryBlock, AllocErr> {
        let size = layout.size();
        debug_assert!(
            new_size <= size,
            "`new_size` must be smaller than or equal to `memory.size()`"
        );

        if size == new_size {
            return Ok(MemoryBlock { ptr, size });
        }

        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove if new_size == 0 => {
                unsafe {
                    self.dealloc(ptr, layout);
                }
                Ok(MemoryBlock { ptr: layout.dangling(), size: 0 })
            }
            ReallocPlacement::MayMove => {
                // `realloc` probably checks for `new_size < size` or something similar.
                let ptr = unsafe {
                    intrinsics::assume(new_size < size);
                    realloc(ptr.as_ptr(), layout, new_size)
                };
                Ok(MemoryBlock { ptr: NonNull::new(ptr).ok_or(AllocErr)?, size: new_size })
            }
        }
    }
}

/// The allocator for unique pointers.
// This function must not unwind. If it does, MIR codegen will fail.
#[cfg(not(test))]
#[lang = "exchange_malloc"]
#[inline]
unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
    let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
    match Global.alloc(layout, AllocInit::Uninitialized) {
        Ok(memory) => memory.ptr.as_ptr(),
        Err(_) => handle_alloc_error(layout),
    }
}

#[cfg_attr(not(test), lang = "box_free")]
#[inline]
// This signature has to be the same as `Box`, otherwise an ICE will happen.
// When an additional parameter to `Box` is added (like `A: AllocRef`), this has to be added here as
// well.
// For example if `Box` is changed to  `struct Box<T: ?Sized, A: AllocRef>(Unique<T>, A)`,
// this function has to be changed to `fn box_free<T: ?Sized, A: AllocRef>(Unique<T>, A)` as well.
pub(crate) unsafe fn box_free<T: ?Sized>(ptr: Unique<T>) {
    unsafe {
        let size = size_of_val(ptr.as_ref());
        let align = min_align_of_val(ptr.as_ref());
        let layout = Layout::from_size_align_unchecked(size, align);
        Global.dealloc(ptr.cast().into(), layout)
    }
}

/// Abort on memory allocation error or failure.
///
/// Callers of memory allocation APIs wishing to abort computation
/// in response to an allocation error are encouraged to call this function,
/// rather than directly invoking `panic!` or similar.
///
/// The default behavior of this function is to print a message to standard error
/// and abort the process.
/// It can be replaced with [`set_alloc_error_hook`] and [`take_alloc_error_hook`].
///
/// [`set_alloc_error_hook`]: ../../std/alloc/fn.set_alloc_error_hook.html
/// [`take_alloc_error_hook`]: ../../std/alloc/fn.take_alloc_error_hook.html
#[stable(feature = "global_alloc", since = "1.28.0")]
#[rustc_allocator_nounwind]
pub fn handle_alloc_error(layout: Layout) -> ! {
    extern "Rust" {
        #[lang = "oom"]
        fn oom_impl(layout: Layout) -> !;
    }
    unsafe { oom_impl(layout) }
}
<|endoftext|>


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

//! The ChaCha random number generator.

use core::prelude::*;

use {Rng, SeedableRng, Rand};

const KEY_WORDS    : uint =  8; // 8 words for the 256-bit key
const STATE_WORDS  : uint = 16;
const CHACHA_ROUNDS: uint = 20; // Cryptographically secure from 8 upwards as of this writing

/// A random number generator that uses the ChaCha20 algorithm [1].
///
/// The ChaCha algorithm is widely accepted as suitable for
/// cryptographic purposes, but this implementation has not been
/// verified as such. Prefer a generator like `OsRng` that defers to
/// the operating system for cases that need high security.
///
/// [1]: D. J. Bernstein, [*ChaCha, a variant of
/// Salsa20*](http://cr.yp.to/chacha.html)

pub struct ChaChaRng {
    buffer:  [u32, ..STATE_WORDS], // Internal buffer of output
    state:   [u32, ..STATE_WORDS], // Initial state
    index:   uint,                 // Index into state
}

static EMPTY: ChaChaRng = ChaChaRng {
    buffer:  [0, ..STATE_WORDS],
    state:   [0, ..STATE_WORDS],
    index:   STATE_WORDS
};

macro_rules! quarter_round{
    ($a: expr, $b: expr, $c: expr, $d: expr) => {{
        $a += $b; $d ^= $a; $d = $d.rotate_left(16);
        $c += $d; $b ^= $c; $b = $b.rotate_left(12);
        $a += $b; $d ^= $a; $d = $d.rotate_left( 8);
        $c += $d; $b ^= $c; $b = $b.rotate_left( 7);
    }}
}

macro_rules! double_round{
    ($x: expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

#[inline]
fn core(output: &mut [u32, ..STATE_WORDS], input: &[u32, ..STATE_WORDS]) {
    *output = *input;

    for _ in range(0, CHACHA_ROUNDS / 2) {
        double_round!(output);
    }

    for i in range(0, STATE_WORDS) {
        output[i] += input[i];
    }
}

impl ChaChaRng {

    /// Create an ChaCha random number generator using the default
    /// fixed key of 8 zero words.
    pub fn new_unseeded() -> ChaChaRng {
        let mut rng = EMPTY;
        rng.init(&[0, ..KEY_WORDS]);
        rng
    }

    /// Sets the internal 128-bit ChaCha counter to
    /// a user-provided value. This permits jumping
    /// arbitrarily ahead (or backwards) in the pseudorandom stream.
    ///
    /// Since the nonce words are used to extend the counter to 128 bits,
    /// users wishing to obtain the conventional ChaCha pseudorandom stream
    /// associated with a particular nonce can call this function with
    /// arguments `0, desired_nonce`.
    pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
        self.state[12] = (counter_low >>  0) as u32;
        self.state[13] = (counter_low >> 32) as u32;
        self.state[14] = (counter_high >>  0) as u32;
        self.state[15] = (counter_high >> 32) as u32;
        self.index = STATE_WORDS; // force recomputation
    }

    /// Initializes `self.state` with the appropriate key and constants
    ///
    /// We deviate slightly from the ChaCha specification regarding
    /// the nonce, which is used to extend the counter to 128 bits.
    /// This is provably as strong as the original cipher, though,
    /// since any distinguishing attack on our variant also works
    /// against ChaCha with a chosen-nonce. See the XSalsa20 [1]
    /// security proof for a more involved example of this.
    ///
    /// The modified word layout is:
    /// ```notrust
    /// constant constant constant constant
    /// key      key      key      key
    /// key      key      key      key
    /// counter  counter  counter  counter
    /// ```
    /// [1]: Daniel J. Bernstein. [*Extending the Salsa20
    /// nonce.*](http://cr.yp.to/papers.html#xsalsa)
    fn init(&mut self, key: &[u32, ..KEY_WORDS]) {
        self.state[0] = 0x61707865;
        self.state[1] = 0x3320646E;
        self.state[2] = 0x79622D32;
        self.state[3] = 0x6B206574;

        for i in range(0, KEY_WORDS) {
            self.state[4+i] = key[i];
        }

        self.state[12] = 0;
        self.state[13] = 0;
        self.state[14] = 0;
        self.state[15] = 0;

        self.index = STATE_WORDS;
    }

    /// Refill the internal output buffer (`self.buffer`)
    fn update(&mut self) {
        core(&mut self.buffer, &self.state);
        self.index = 0;
        // update 128-bit counter
        self.state[12] += 1;
        if self.state[12] != 0 { return };
        self.state[13] += 1;
        if self.state[13] != 0 { return };
        self.state[14] += 1;
        if self.state[14] != 0 { return };
        self.state[15] += 1;
    }
}

impl Rng for ChaChaRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.index == STATE_WORDS {
            self.update();
        }

        let value = self.buffer[self.index % STATE_WORDS];
        self.index += 1;
        value
    }
}

impl<'a> SeedableRng<&'a [u32]> for ChaChaRng {

    fn reseed(&mut self, seed: &'a [u32]) {
        // reset state
        self.init(&[0u32, ..KEY_WORDS]);
        // set key inplace
        let key = self.state.slice_mut(4, 4+KEY_WORDS);
        for (k, s) in key.iter_mut().zip(seed.iter()) {
            *k = *s;
        }
    }

    /// Create a ChaCha generator from a seed,
    /// obtained from a variable-length u32 array.
    /// Only up to 8 words are used; if less than 8
    /// words are used, the remaining are set to zero.
    fn from_seed(seed: &'a [u32]) -> ChaChaRng {
        let mut rng = EMPTY;
        rng.reseed(seed);
        rng
    }
}

impl Rand for ChaChaRng {
    fn rand<R: Rng>(other: &mut R) -> ChaChaRng {
        let mut key : [u32, ..KEY_WORDS] = [0, ..KEY_WORDS];
        for word in key.iter_mut() {
            *word = other.gen();
        }
        SeedableRng::from_seed(key.as_slice())
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use core::iter::order;
    use {Rng, SeedableRng};
    use super::ChaChaRng;

    #[test]
    fn test_rng_rand_seeded() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut ra: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let mut rb: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_seeded() {
        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);
        let mut rb: ChaChaRng = SeedableRng::from_seed(seed);
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut r: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed(s.as_slice());

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    #[test]
    fn test_rng_true_values() {
        // Test vectors 1 and 2 from
        // http://tools.ietf.org/html/draft-nir-cfrg-chacha20-poly1305-04
        let seed : &[_] = &[0u32, ..8];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xade0b876, 0x903df1a0, 0xe56a5d40, 0x28bd8653,
                        0xb819d2bd, 0x1aed8da0, 0xccef36a8, 0xc70d778b,
                        0x7c5941da, 0x8d485751, 0x3fe02477, 0x374ad8b8,
                        0xf4b8436a, 0x1ca11815, 0x69b687c3, 0x8665eeb2));

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xbee7079f, 0x7a385155, 0x7c97ba98, 0x0d082d73,
                        0xa0290fcb, 0x6965e348, 0x3e53c612, 0xed7aee32,
                        0x7621b729, 0x434ee69c, 0xb03371d5, 0xd539d874,
                        0x281fed31, 0x45fb0a51, 0x1f0ae1ac, 0x6f4d794b));

        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        // Store the 17*i-th 32-bit word,
        // i.e., the i-th word of the i-th 16-word block
        let mut v : Vec<u32> = Vec::new();
        for _ in range(0u, 16) {
            v.push(ra.next_u32());
            for _ in range(0u, 16) {
                ra.next_u32();
            }
        }

        assert_eq!(v,
                   vec!(0xf225c81a, 0x6ab1be57, 0x04d42951, 0x70858036,
                        0x49884684, 0x64efec72, 0x4be2d186, 0x3615b384,
                        0x11cfa18e, 0xd3c50049, 0x75c775f6, 0x434c6530,
                        0x2c5bad8f, 0x898881dc, 0x5f1c86d9, 0xc1f8e7f4));
    }
}

<commit_msg>Fix spelling mistakes in comments.<commit_after>// Copyright 2014 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.

//! The ChaCha random number generator.

use core::prelude::*;

use {Rng, SeedableRng, Rand};

const KEY_WORDS    : uint =  8; // 8 words for the 256-bit key
const STATE_WORDS  : uint = 16;
const CHACHA_ROUNDS: uint = 20; // Cryptographically secure from 8 upwards as of this writing

/// A random number generator that uses the ChaCha20 algorithm [1].
///
/// The ChaCha algorithm is widely accepted as suitable for
/// cryptographic purposes, but this implementation has not been
/// verified as such. Prefer a generator like `OsRng` that defers to
/// the operating system for cases that need high security.
///
/// [1]: D. J. Bernstein, [*ChaCha, a variant of
/// Salsa20*](http://cr.yp.to/chacha.html)

pub struct ChaChaRng {
    buffer:  [u32, ..STATE_WORDS], // Internal buffer of output
    state:   [u32, ..STATE_WORDS], // Initial state
    index:   uint,                 // Index into state
}

static EMPTY: ChaChaRng = ChaChaRng {
    buffer:  [0, ..STATE_WORDS],
    state:   [0, ..STATE_WORDS],
    index:   STATE_WORDS
};

macro_rules! quarter_round{
    ($a: expr, $b: expr, $c: expr, $d: expr) => {{
        $a += $b; $d ^= $a; $d = $d.rotate_left(16);
        $c += $d; $b ^= $c; $b = $b.rotate_left(12);
        $a += $b; $d ^= $a; $d = $d.rotate_left( 8);
        $c += $d; $b ^= $c; $b = $b.rotate_left( 7);
    }}
}

macro_rules! double_round{
    ($x: expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

#[inline]
fn core(output: &mut [u32, ..STATE_WORDS], input: &[u32, ..STATE_WORDS]) {
    *output = *input;

    for _ in range(0, CHACHA_ROUNDS / 2) {
        double_round!(output);
    }

    for i in range(0, STATE_WORDS) {
        output[i] += input[i];
    }
}

impl ChaChaRng {

    /// Create an ChaCha random number generator using the default
    /// fixed key of 8 zero words.
    pub fn new_unseeded() -> ChaChaRng {
        let mut rng = EMPTY;
        rng.init(&[0, ..KEY_WORDS]);
        rng
    }

    /// Sets the internal 128-bit ChaCha counter to
    /// a user-provided value. This permits jumping
    /// arbitrarily ahead (or backwards) in the pseudorandom stream.
    ///
    /// Since the nonce words are used to extend the counter to 128 bits,
    /// users wishing to obtain the conventional ChaCha pseudorandom stream
    /// associated with a particular nonce can call this function with
    /// arguments `0, desired_nonce`.
    pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
        self.state[12] = (counter_low >>  0) as u32;
        self.state[13] = (counter_low >> 32) as u32;
        self.state[14] = (counter_high >>  0) as u32;
        self.state[15] = (counter_high >> 32) as u32;
        self.index = STATE_WORDS; // force recomputation
    }

    /// Initializes `self.state` with the appropriate key and constants
    ///
    /// We deviate slightly from the ChaCha specification regarding
    /// the nonce, which is used to extend the counter to 128 bits.
    /// This is provably as strong as the original cipher, though,
    /// since any distinguishing attack on our variant also works
    /// against ChaCha with a chosen-nonce. See the XSalsa20 [1]
    /// security proof for a more involved example of this.
    ///
    /// The modified word layout is:
    /// ```notrust
    /// constant constant constant constant
    /// key      key      key      key
    /// key      key      key      key
    /// counter  counter  counter  counter
    /// ```
    /// [1]: Daniel J. Bernstein. [*Extending the Salsa20
    /// nonce.*](http://cr.yp.to/papers.html#xsalsa)
    fn init(&mut self, key: &[u32, ..KEY_WORDS]) {
        self.state[0] = 0x61707865;
        self.state[1] = 0x3320646E;
        self.state[2] = 0x79622D32;
        self.state[3] = 0x6B206574;

        for i in range(0, KEY_WORDS) {
            self.state[4+i] = key[i];
        }

        self.state[12] = 0;
        self.state[13] = 0;
        self.state[14] = 0;
        self.state[15] = 0;

        self.index = STATE_WORDS;
    }

    /// Refill the internal output buffer (`self.buffer`)
    fn update(&mut self) {
        core(&mut self.buffer, &self.state);
        self.index = 0;
        // update 128-bit counter
        self.state[12] += 1;
        if self.state[12] != 0 { return };
        self.state[13] += 1;
        if self.state[13] != 0 { return };
        self.state[14] += 1;
        if self.state[14] != 0 { return };
        self.state[15] += 1;
    }
}

impl Rng for ChaChaRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.index == STATE_WORDS {
            self.update();
        }

        let value = self.buffer[self.index % STATE_WORDS];
        self.index += 1;
        value
    }
}

impl<'a> SeedableRng<&'a [u32]> for ChaChaRng {

    fn reseed(&mut self, seed: &'a [u32]) {
        // reset state
        self.init(&[0u32, ..KEY_WORDS]);
        // set key in place
        let key = self.state.slice_mut(4, 4+KEY_WORDS);
        for (k, s) in key.iter_mut().zip(seed.iter()) {
            *k = *s;
        }
    }

    /// Create a ChaCha generator from a seed,
    /// obtained from a variable-length u32 array.
    /// Only up to 8 words are used; if less than 8
    /// words are used, the remaining are set to zero.
    fn from_seed(seed: &'a [u32]) -> ChaChaRng {
        let mut rng = EMPTY;
        rng.reseed(seed);
        rng
    }
}

impl Rand for ChaChaRng {
    fn rand<R: Rng>(other: &mut R) -> ChaChaRng {
        let mut key : [u32, ..KEY_WORDS] = [0, ..KEY_WORDS];
        for word in key.iter_mut() {
            *word = other.gen();
        }
        SeedableRng::from_seed(key.as_slice())
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use core::iter::order;
    use {Rng, SeedableRng};
    use super::ChaChaRng;

    #[test]
    fn test_rng_rand_seeded() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut ra: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let mut rb: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_seeded() {
        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);
        let mut rb: ChaChaRng = SeedableRng::from_seed(seed);
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut r: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed(s.as_slice());

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    #[test]
    fn test_rng_true_values() {
        // Test vectors 1 and 2 from
        // http://tools.ietf.org/html/draft-nir-cfrg-chacha20-poly1305-04
        let seed : &[_] = &[0u32, ..8];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xade0b876, 0x903df1a0, 0xe56a5d40, 0x28bd8653,
                        0xb819d2bd, 0x1aed8da0, 0xccef36a8, 0xc70d778b,
                        0x7c5941da, 0x8d485751, 0x3fe02477, 0x374ad8b8,
                        0xf4b8436a, 0x1ca11815, 0x69b687c3, 0x8665eeb2));

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xbee7079f, 0x7a385155, 0x7c97ba98, 0x0d082d73,
                        0xa0290fcb, 0x6965e348, 0x3e53c612, 0xed7aee32,
                        0x7621b729, 0x434ee69c, 0xb03371d5, 0xd539d874,
                        0x281fed31, 0x45fb0a51, 0x1f0ae1ac, 0x6f4d794b));

        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        // Store the 17*i-th 32-bit word,
        // i.e., the i-th word of the i-th 16-word block
        let mut v : Vec<u32> = Vec::new();
        for _ in range(0u, 16) {
            v.push(ra.next_u32());
            for _ in range(0u, 16) {
                ra.next_u32();
            }
        }

        assert_eq!(v,
                   vec!(0xf225c81a, 0x6ab1be57, 0x04d42951, 0x70858036,
                        0x49884684, 0x64efec72, 0x4be2d186, 0x3615b384,
                        0x11cfa18e, 0xd3c50049, 0x75c775f6, 0x434c6530,
                        0x2c5bad8f, 0x898881dc, 0x5f1c86d9, 0xc1f8e7f4));
    }
}

<|endoftext|>