<commit_before>
use buffer::{ImageBuffer, Pixel, ArrayLike};
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use std::old_io;
use std::old_io::IoResult;
use std::num::Int;

use color::{Rgb, Rgba};
use utils::lzw;
use utils::bitstream::LsbWriter;

use super::{Extension, Block, DisposalMethod};

#[derive(Debug, Copy)]
#[allow(unused_qualifications)]
/// The color mode the encoder will use to encode the image.
pub enum ColorMode {
	/// Image will be encoded in multiple frames if more than 256 colors are present
	TrueColor,
	/// Number of colors will be reduced
	Indexed(u8),
}
use self::ColorMode::{TrueColor, Indexed};

/// A GIF encoder.
///
/// Encodes the image either in true color using indexed colors.
/// If the mode is set to TrueColor the image is split into multiple frames
/// when the number of colors including transparent color exceeds 256.
/// Pixels with an alpha value != 1.0 will be set to alpha = 0.
pub struct Encoder<Image> {
	image: Image,
	bg_color: Option<Rgb<u8>>,
	color_mode: ColorMode
}

const TRANSPARENT: Rgba<u8> = Rgba([0, 0, 0, 0]);

impl<Container> Encoder<ImageBuffer<Rgba<u8>, Container>>
where Container: ArrayLike<u8> {
	/// Creates a new GIF encoder
	pub fn new(image: ImageBuffer<Rgba<u8>, Container>,
		       bg_color: Option<Rgb<u8>>,
		       color_mode: ColorMode,
		      ) -> Encoder<ImageBuffer<Rgba<u8>, Container>> {
		Encoder {
			image: image,
			bg_color: bg_color,
			color_mode: color_mode
		}
	}

	/// Encodes the image
	pub fn encode<W: Writer>(&mut self, w: &mut W) -> IoResult<()> {
		// Header
		try!(w.write_all(b"GIF89a"));
		// Logical screen descriptor
		let height = self.image.height();
		let width = self.image.width();
		if width > <u16 as Int>::max_value() as u32 ||
		   height > <u16 as Int>::max_value() as u32 {
		   	return Err(old_io::IoError{
		   		kind: old_io::InvalidInput,
		   		desc: "Image dimensions are to large for the gif format.",
		   		detail: None
		   	})
		}
		try!(w.write_le_u16(width as u16));
		try!(w.write_le_u16(height as u16));

		let (hist, transparent) = self.histogram();
		let num_colors = hist.len();

		try!(self.write_global_table(w, &hist[]));

		match self.color_mode {
			// Global color table has been written, just write the image data now
			TrueColor if num_colors <= 256 => {
				// NOTE: copy paste any changes to case: Indexed(n) if n as usize <= num_colors
				try!(self.write_control_ext(w, 0, transparent));
				try!(self.write_descriptor(w, None));
				try!(self.write_image_simple(w, &hist[], transparent));
			},
			TrueColor => {
				try!(self.write_true_color(w, hist, transparent));
			},
			Indexed(n) if n as usize <= num_colors => {
				// NOTE: copy paste from case: TrueColor if num_colors <= 256
				try!(self.write_control_ext(w, 0, transparent));
				try!(self.write_descriptor(w, None));
				try!(self.write_image_simple(w, &hist[], transparent));
			}
			Indexed(_) => {
				panic!("Indexed(n) unimplemented")
			}

		}
		w.write_u8(Block::Trailer as u8)
	}

	/// Returns a histogram of the colors in the image
	fn histogram(&self) -> (Vec<(Rgba<u8>, usize)>, Option<usize>){
		let mut hist: HashMap<_, usize> = HashMap::new();
		// Add bg_color to table
		if let Some(bg_color) = self.bg_color {
			let _ = hist.insert(bg_color.to_rgba(), 0);
		}
		for p in self.image.pixels() {
			let mut p = *p;
			// Replace transparent pixels with black, alpha = 0
			if let Some(bg_color) = self.bg_color {
				p.blend(&bg_color.to_rgba())
			} else if p[3] != 255 {
				p = TRANSPARENT;
			}
			match hist.entry(p) {
				Occupied(mut entry) => {
					let val = entry.get_mut();
					*val = *val + 1;
				},
				Vacant(entry) => {
					entry.insert(1);
				}
			}
		}
		let mut colors: Vec<(Rgba<u8>, usize)> = hist.into_iter().collect();
		// Sort color map with decreasing color freqs
		colors.sort_by(|a, b| b.1.cmp(&a.1));
		// Add bg_color to table
		let transparent = colors.iter().position(|x| x.0 == TRANSPARENT);
		(colors, transparent)
	}

	/// Write the global color table and the corresponding flags
	fn write_global_table<W: Writer>(&mut self,
									 w: &mut W,
									 hist: &[(Rgba<u8>, usize)]
									) -> IoResult<()>
	{
		let num_colors = hist.len();
		let mut flags = 0;
		flags |= 1 << 4; // sort flag
		let n = flag_n(num_colors);
		flags |= n << 4; // 2^(n+1) colors
		let (global_table, bg_index) = if num_colors <= 256 {
			flags |= 1 << 7; // glocal table flag
			flags |= n;
			let mut bg_index = 0;
			let mut table = Vec::with_capacity(3*(2 << n));
			for (i, &(color, _)) in hist.iter().enumerate() {
				let channels = &color.channels()[..3];
				if let Some(bg_color) = self.bg_color {
					if channels == bg_color.channels() {
						bg_index = i;
					}
				}
				table.push_all(channels)
			}
			// Waste some space as of gif spec
			for _ in 0..((2 << n) - num_colors) {
				table.push_all(&[0, 0, 0]);
			}
			(Some(table), bg_index as u8)
		} else if let Some(bg_color) = self.bg_color {
			flags |= 1 << 7; // glocal table flag
			let mut table = Vec::with_capacity(6);
			table.push_all(&bg_color.channels()[..3]);
			table.push_all(&[0, 0, 0]);
			(Some(table), 0)
		} else {
			(None, 0)
		};
		try!(w.write_u8(flags));
		try!(w.write_u8(bg_index));
		try!(w.write_u8(0)); // aspect ration, disregard
		if let Some(global_table) = global_table {
			try!(w.write_all(&global_table[]));
		}
		Ok(())
	}

	/// Writes the graphics control extension
	fn write_control_ext<W: Writer>(&mut self,
								    w: &mut W,
								    delay: u16,
								    transparent: Option<usize>
								   ) -> IoResult<()>
	{
		try!(w.write_u8(Block::Extension as u8));
		try!(w.write_u8(Extension::Control as u8));
		try!(w.write_u8(4)); // size
		let mut field = 0;
		field |= (DisposalMethod::None as u8) << 2;
		let idx = if let Some(idx) = transparent {
			field |=  1;
			idx as u8
		} else {
			0
		};
		try!(w.write_u8(field));
		try!(w.write_le_u16(delay));
		try!(w.write_u8(idx));
		w.write_u8(0)
	}

	/// Writes the image descriptor
	fn write_descriptor<W: Writer>(&mut self,
								    w: &mut W,
								    table_len: Option<usize>
								   ) -> IoResult<()>
	{
		try!(w.write_u8(Block::Image as u8));
		try!(w.write_le_u16(0)); // left
		try!(w.write_le_u16(0)); // top
		let height = self.image.height();
		let width = self.image.width();
		try!(w.write_le_u16(width as u16));
		try!(w.write_le_u16(height as u16));
		if let Some(len) = table_len {
			w.write_u8(flag_n(len) | 0x80)
		} else {
			w.write_u8(0)
		}
	}

	/// Writes and compresses the indexed data
	fn write_indices<W: Writer>(&mut self,
								    w: &mut W,
								    indices: &[u8],
								   ) -> IoResult<()>
	{
		let code_size = match flag_n(indices.len()) + 1 {
			1 => 2,
			n => n
		};
		let mut encoded_data = Vec::new();
		try!(lzw::encode(indices, LsbWriter::new(&mut encoded_data), code_size));
		try!(w.write_u8(code_size));
		for chunk in encoded_data.chunks(255) {
			try!(w.write_u8((chunk.len()) as u8));
			try!(w.write_all(chunk));
		}
		w.write_u8(0) // block terminator

	}

	/// Writes the image to the file assuming that every pixel is in the color table
	/// If not, the index of the transparent pixel is written
	fn write_image_simple<W: Writer>(&mut self,
								    w: &mut W,
								    hist: &[(Rgba<u8>, usize)],
								    transparent: Option<usize>,
								   ) -> IoResult<()>
	{
		let t_idx = match transparent { Some(i) => i as u8, None => 0 };
		let data: Vec<u8> = self.image.pixels().map(|p| {
			if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
				idx as u8
			} else {
				t_idx
			}
		}).collect();
		self.write_indices(w, &data[])

	}

	/// Writes the image as a true color image by splitting the colors
	/// over several frames
	fn write_true_color<W: Writer>(&mut self,
								    w: &mut W,
								    hist: Vec<(Rgba<u8>, usize)>,
								    transparent: Option<usize>
								   ) -> IoResult<()>
	{
		let mut hist = hist;
		// We need a transparent pixel we put it to index 0
		if let Some(transparent) = transparent {
			let val = hist.swap_remove(transparent);
			hist.push(val);
		} else {
			hist.reserve(1);
			hist.push((TRANSPARENT, 0));
		}
		let val = hist.swap_remove(0);
		hist.push(val);
		let transparent = Some(0);

		// Calculating the indices is expensive so just do it once
		let indices: Vec<u32> = self.image.pixels().map(|p| {
			if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
				idx as u32
			} else {
				0
			}
		}).collect();
		let mut chunk_indices = Vec::with_capacity(self.image.as_slice().len());
		// starting from 1 since we inject the transparent color in every frame
		for (j, chunk) in hist[1..].chunks(255).enumerate() {
			// Now we can reuse the indices for every subimage
			// if the substract j*255 from each an replace it with 0 if it gets negative
			chunk_indices.clear();
			for &idx in indices.iter() {
				let i: i32 = idx as i32 - j as i32*255;
				chunk_indices.push(if i < 0 { 0 } else { i as u8 })
			}
			// Write local color table
			let num_local_colors = chunk.len() + 1;
			let n = flag_n(num_local_colors);
			try!(self.write_control_ext(w, 0, transparent));
			try!(self.write_descriptor(w, Some(num_local_colors)));
			// Inject transparent color
			try!(w.write_all(&TRANSPARENT.channels()[..3]));
			for &(color, _) in chunk.iter() {
				try!(w.write_all(&color.channels()[..3]));
			}
			// waste some space as of gif spec
			for _ in 0..((2 << n) - num_local_colors) {
				try!(w.write_all(&[0, 0, 0]));
			}
			// write indices
			try!(self.write_indices(w, &chunk_indices[]));
		}
		Ok(())
	}
}

// Color table len converted to flag bits
fn flag_n(size: usize) -> u8 {
	match size {
		0  ...2   => 0,
		3  ...4   => 1,
		5  ...8   => 2,
		7  ...16  => 3,
		17 ...32  => 4,
		33 ...64  => 5,
		65 ...128 => 6,
		129...256 => 7,
		_ => 7
	}
}<commit_msg>Fix the index calculation.<commit_after>
use buffer::{ImageBuffer, Pixel, ArrayLike};
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use std::old_io;
use std::old_io::IoResult;
use std::num::Int;

use color::{Rgb, Rgba};
use utils::lzw;
use utils::bitstream::LsbWriter;

use super::{Extension, Block, DisposalMethod};

#[derive(Debug, Copy)]
#[allow(unused_qualifications)]
/// The color mode the encoder will use to encode the image.
pub enum ColorMode {
	/// Image will be encoded in multiple frames if more than 256 colors are present
	TrueColor,
	/// Number of colors will be reduced
	Indexed(u8),
}
use self::ColorMode::{TrueColor, Indexed};

/// A GIF encoder.
///
/// Encodes the image either in true color using indexed colors.
/// If the mode is set to TrueColor the image is split into multiple frames
/// when the number of colors including transparent color exceeds 256.
/// Pixels with an alpha value != 1.0 will be set to alpha = 0.
pub struct Encoder<Image> {
	image: Image,
	bg_color: Option<Rgb<u8>>,
	color_mode: ColorMode
}

const TRANSPARENT: Rgba<u8> = Rgba([0, 0, 0, 0]);

impl<Container> Encoder<ImageBuffer<Rgba<u8>, Container>>
where Container: ArrayLike<u8> {
	/// Creates a new GIF encoder
	pub fn new(image: ImageBuffer<Rgba<u8>, Container>,
		       bg_color: Option<Rgb<u8>>,
		       color_mode: ColorMode,
		      ) -> Encoder<ImageBuffer<Rgba<u8>, Container>> {
		Encoder {
			image: image,
			bg_color: bg_color,
			color_mode: color_mode
		}
	}

	/// Encodes the image
	pub fn encode<W: Writer>(&mut self, w: &mut W) -> IoResult<()> {
		// Header
		try!(w.write_all(b"GIF89a"));
		// Logical screen descriptor
		let height = self.image.height();
		let width = self.image.width();
		if width > <u16 as Int>::max_value() as u32 ||
		   height > <u16 as Int>::max_value() as u32 {
		   	return Err(old_io::IoError{
		   		kind: old_io::InvalidInput,
		   		desc: "Image dimensions are to large for the gif format.",
		   		detail: None
		   	})
		}
		try!(w.write_le_u16(width as u16));
		try!(w.write_le_u16(height as u16));

		let (hist, transparent) = self.histogram();
		let num_colors = hist.len();

		try!(self.write_global_table(w, &hist[]));

		match self.color_mode {
			// Global color table has been written, just write the image data now
			TrueColor if num_colors <= 256 => {
				// NOTE: copy paste any changes to case: Indexed(n) if n as usize <= num_colors
				try!(self.write_control_ext(w, 0, transparent));
				try!(self.write_descriptor(w, None));
				try!(self.write_image_simple(w, &hist[], transparent));
			},
			TrueColor => {
				try!(self.write_true_color(w, hist, transparent));
			},
			Indexed(n) if n as usize <= num_colors => {
				// NOTE: copy paste from case: TrueColor if num_colors <= 256
				try!(self.write_control_ext(w, 0, transparent));
				try!(self.write_descriptor(w, None));
				try!(self.write_image_simple(w, &hist[], transparent));
			}
			Indexed(_) => {
				panic!("Indexed(n) unimplemented")
			}

		}
		w.write_u8(Block::Trailer as u8)
	}

	/// Returns a histogram of the colors in the image
	fn histogram(&self) -> (Vec<(Rgba<u8>, usize)>, Option<usize>){
		let mut hist: HashMap<_, usize> = HashMap::new();
		// Add bg_color to table
		if let Some(bg_color) = self.bg_color {
			let _ = hist.insert(bg_color.to_rgba(), 0);
		}
		for p in self.image.pixels() {
			let mut p = *p;
			// Replace transparent pixels with black, alpha = 0
			if let Some(bg_color) = self.bg_color {
				p.blend(&bg_color.to_rgba())
			} else if p[3] != 255 {
				p = TRANSPARENT;
			}
			match hist.entry(p) {
				Occupied(mut entry) => {
					let val = entry.get_mut();
					*val = *val + 1;
				},
				Vacant(entry) => {
					entry.insert(1);
				}
			}
		}
		let mut colors: Vec<(Rgba<u8>, usize)> = hist.into_iter().collect();
		// Sort color map with decreasing color freqs
		colors.sort_by(|a, b| b.1.cmp(&a.1));
		// Add bg_color to table
		let transparent = colors.iter().position(|x| x.0 == TRANSPARENT);
		(colors, transparent)
	}

	/// Write the global color table and the corresponding flags
	fn write_global_table<W: Writer>(&mut self,
									 w: &mut W,
									 hist: &[(Rgba<u8>, usize)]
									) -> IoResult<()>
	{
		let num_colors = hist.len();
		let mut flags = 0;
		flags |= 1 << 4; // sort flag
		let n = flag_n(num_colors);
		flags |= n << 4; // 2^(n+1) colors
		let (global_table, bg_index) = if num_colors <= 256 {
			flags |= 1 << 7; // glocal table flag
			flags |= n;
			let mut bg_index = 0;
			let mut table = Vec::with_capacity(3*(2 << n));
			for (i, &(color, _)) in hist.iter().enumerate() {
				let channels = &color.channels()[..3];
				if let Some(bg_color) = self.bg_color {
					if channels == bg_color.channels() {
						bg_index = i;
					}
				}
				table.push_all(channels)
			}
			// Waste some space as of gif spec
			for _ in 0..((2 << n) - num_colors) {
				table.push_all(&[0, 0, 0]);
			}
			(Some(table), bg_index as u8)
		} else if let Some(bg_color) = self.bg_color {
			flags |= 1 << 7; // glocal table flag
			let mut table = Vec::with_capacity(6);
			table.push_all(&bg_color.channels()[..3]);
			table.push_all(&[0, 0, 0]);
			(Some(table), 0)
		} else {
			(None, 0)
		};
		try!(w.write_u8(flags));
		try!(w.write_u8(bg_index));
		try!(w.write_u8(0)); // aspect ration, disregard
		if let Some(global_table) = global_table {
			try!(w.write_all(&global_table[]));
		}
		Ok(())
	}

	/// Writes the graphics control extension
	fn write_control_ext<W: Writer>(&mut self,
								    w: &mut W,
								    delay: u16,
								    transparent: Option<usize>
								   ) -> IoResult<()>
	{
		try!(w.write_u8(Block::Extension as u8));
		try!(w.write_u8(Extension::Control as u8));
		try!(w.write_u8(4)); // size
		let mut field = 0;
		field |= (DisposalMethod::None as u8) << 2;
		let idx = if let Some(idx) = transparent {
			field |=  1;
			idx as u8
		} else {
			0
		};
		try!(w.write_u8(field));
		try!(w.write_le_u16(delay));
		try!(w.write_u8(idx));
		w.write_u8(0)
	}

	/// Writes the image descriptor
	fn write_descriptor<W: Writer>(&mut self,
								    w: &mut W,
								    table_len: Option<usize>
								   ) -> IoResult<()>
	{
		try!(w.write_u8(Block::Image as u8));
		try!(w.write_le_u16(0)); // left
		try!(w.write_le_u16(0)); // top
		let height = self.image.height();
		let width = self.image.width();
		try!(w.write_le_u16(width as u16));
		try!(w.write_le_u16(height as u16));
		if let Some(len) = table_len {
			w.write_u8(flag_n(len) | 0x80)
		} else {
			w.write_u8(0)
		}
	}

	/// Writes and compresses the indexed data
	fn write_indices<W: Writer>(&mut self,
								    w: &mut W,
								    indices: &[u8],
								   ) -> IoResult<()>
	{
		let code_size = match flag_n(indices.len()) + 1 {
			1 => 2,
			n => n
		};
		let mut encoded_data = Vec::new();
		try!(lzw::encode(indices, LsbWriter::new(&mut encoded_data), code_size));
		try!(w.write_u8(code_size));
		for chunk in encoded_data.chunks(255) {
			try!(w.write_u8((chunk.len()) as u8));
			try!(w.write_all(chunk));
		}
		w.write_u8(0) // block terminator

	}

	/// Writes the image to the file assuming that every pixel is in the color table
	/// If not, the index of the transparent pixel is written
	fn write_image_simple<W: Writer>(&mut self,
								    w: &mut W,
								    hist: &[(Rgba<u8>, usize)],
								    transparent: Option<usize>,
								   ) -> IoResult<()>
	{
		let t_idx = match transparent { Some(i) => i as u8, None => 0 };
		let data: Vec<u8> = self.image.pixels().map(|p| {
			if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
				idx as u8
			} else {
				t_idx
			}
		}).collect();
		self.write_indices(w, &data[])

	}

	/// Writes the image as a true color image by splitting the colors
	/// over several frames
	fn write_true_color<W: Writer>(&mut self,
								    w: &mut W,
								    hist: Vec<(Rgba<u8>, usize)>,
								    transparent: Option<usize>
								   ) -> IoResult<()>
	{
		try!(self.write_nab(w, 1));
		let mut hist = hist;
		// Remove transparent idx
		if let Some(transparent) = transparent {
			let _ = hist.swap_remove(transparent);
		}
		let transparent = Some(0);
		// Calculating the indices is expensive so just do it once
		let indices: Vec<u32> = self.image.pixels().map(|p| {
			if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
				idx as u32
			} else {
				0
			}
		}).collect();
		let mut chunk_indices = Vec::with_capacity(self.image.as_slice().len());
		// starting from 1 since we inject the transparent color in every frame
		for (j, chunk) in hist.chunks(255).enumerate() {
			// Now we can reuse the indices for every subimage
			// if the substract j*255 from each an replace it with 0 if it gets negative
			chunk_indices.clear();
			for &idx in indices.iter() {
				let i: i64 = idx as i64 - j as i64*255;
				chunk_indices.push(if i < 0 || i > 255 { 0 } else { i as u8 + 1 })
			}
			// Write local color table
			let num_local_colors = chunk.len() + 1;
			let n = flag_n(num_local_colors);
			try!(self.write_control_ext(w, 0, transparent));
			try!(self.write_descriptor(w, Some(num_local_colors)));
			// Inject transparent color
			try!(w.write_all(&TRANSPARENT.channels()[..3]));
			for &(color, _) in chunk.iter() {
				try!(w.write_all(&color.channels()[..3]));
			}
			// waste some space as of gif spec
			for _ in 0..((2 << n) - num_local_colors) {
				try!(w.write_all(&[0, 0, 0]));
			}
			// write indices
			try!(self.write_indices(w, &chunk_indices[]));
		}
		Ok(())
	}

	/// Writes the netscape application block to set the number `n` of repetitions
	fn write_nab<W: Writer>(&mut self, w: &mut W, n: u16) -> IoResult<()> {
		try!(w.write_u8(Block::Extension as u8));
		try!(w.write_u8(Extension::Application as u8));
		try!(w.write_u8(0x0B)); // size
		try!(w.write_all(b"NETSCAPE2.0"));
		try!(w.write_u8(0x03)); // sub-block size
		try!(w.write_u8(0x01)); // sub-block id
		try!(w.write_le_u16(n));
		w.write_u8(0) // terminator
	}
}

// Color table len converted to flag bits
fn flag_n(size: usize) -> u8 {
	match size {
		0  ...2   => 0,
		3  ...4   => 1,
		5  ...8   => 2,
		7  ...16  => 3,
		17 ...32  => 4,
		33 ...64  => 5,
		65 ...128 => 6,
		129...256 => 7,
		_ => 7
	}
}<|endoftext|>


<commit_before>use std;
use std::ops::Add;
use std::io::{self, Write, Read};
use time::{self, Timespec, Duration};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        use std::io::{self, Write};
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n > kb.powf(4_f64) => {format!("{:.*} TG", 2, $n / kb.powf(4_f64))},
            $n if $n > kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n > kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 2, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &std::iter::repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

pub struct ProgressBar {
    start_time: Timespec,
    total: usize,
    current: usize,
    is_finish: bool,
    show_bar: bool,
    show_speed: bool,
    show_percent: bool,
    show_counter: bool,
    show_time_left: bool,
    // Format
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
}

impl ProgressBar {
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            is_finish: false,
            show_bar: true,
            show_speed: false,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    fn draw(&self) {
        let width = 149;    // replace to -> get_tty_size()
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        // TODO: ADD KB FORMAT
        if !self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = suffix + &format!("{}/s ", speed as usize);
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            prefix = format!("{} / {} ", self.current, self.total);
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

// TODO: Implement io::Reader
<commit_msg>feat(pb): set_unit<commit_after>use std;
use std::ops::Add;
use std::io::{self, Write, Read};
use time::{self, Timespec, Duration};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        use std::io::{self, Write};
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n > kb.powf(4_f64) => {format!("{:.*} TG", 2, $n / kb.powf(4_f64))},
            $n if $n > kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n > kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 2, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &std::iter::repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
pub enum Units {
    Default,
    Bytes
}

pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    is_finish: bool,
    show_bar: bool,
    show_speed: bool,
    show_percent: bool,
    show_counter: bool,
    show_time_left: bool,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
}

impl ProgressBar {

    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: false,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    fn draw(&self) {
        let width = 149;    // replace to -> get_tty_size()
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        // TODO: ADD KB FORMAT, SHOULD BE IN THE COUNTER TOO.
        if !self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = suffix + &format!("{}/s ", speed as usize);
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            prefix = format!("{} / {} ", self.current, self.total);
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

// TODO: Implement io::Reader
<|endoftext|>


<commit_before>//! # generator run time support
//!
//! generator run time context management
//!
use std::mem;
use std::ptr;
use std::any::Any;
use std::intrinsics::type_name;

use stack::Stack;
use reg_context::Context as RegContext;

/// each thread has it's own generator context stack
thread_local!(static ROOT_CONTEXT: Box<Context> = {
    let mut root = Box::new(Context::new(0));
    let p = &mut *root as *mut _;
    root.parent = p; // init top to current
    root
});

/// yield error types
#[allow(dead_code)]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Error {
    Cancel,
    TypeErr,
    StackErr,
    ContextErr,
}

/// generator context
pub struct Context {
    /// generator regs context
    pub regs: RegContext,
    /// generator execution stack
    pub stack: Stack,
    /// passed in para for send
    pub para: *mut Any,
    /// this is just a buffer for the return value
    pub ret: *mut Any,
    /// track generator ref, yield will -1, send will +1
    pub _ref: u32,
    /// propagate panic
    pub err: Option<Error>,
    /// context local storage
    pub local_data: *mut u8,

    /// child context
    child: *mut Context,
    /// parent context
    pub parent: *mut Context,
}

impl Context {
    /// return a default generator context
    pub fn new(size: usize) -> Context {
        Context {
            regs: RegContext::empty(),
            stack: Stack::new(size),
            para: unsafe { mem::uninitialized() },
            ret: unsafe { mem::uninitialized() },
            _ref: 1, // none zero means it's not running
            err: None,
            child: ptr::null_mut(),
            parent: ptr::null_mut(),
            local_data: ptr::null_mut(),
        }
    }

    /// judge it's generator context
    #[inline]
    pub fn is_generator(&self) -> bool {
        self.stack.size() > 0
    }

    /// get current generator send para
    #[inline]
    pub fn get_para(&self) -> Option
        where A: Any
    {
        let para = unsafe { &mut *self.para };
        let val = para.downcast_mut::<Option>();
        if val.is_some() {
            val.unwrap().take()
        } else {
            let t = unsafe { type_name::() };
            error!("get yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr);
        }
    }

    /// set current generator return value
    #[inline]
    pub fn set_ret<T>(&mut self, v: T)
        where T: Any
    {
        let ret = unsafe { &mut *self.ret };
        let val = ret.downcast_mut::<Option<T>>();
        if val.is_some() {
            mem::replace(val.unwrap(), Some(v));
        } else {
            let t = unsafe { type_name::<T>() };
            error!("yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr);
        }
    }
}

/// Coroutine managing environment
pub struct ContextStack {
    root: *mut Context,
}

impl ContextStack {
    #[inline]
    pub fn current() -> ContextStack {
        let root = ROOT_CONTEXT.with(|r| &**r as *const _ as *mut Context);
        ContextStack { root: root }
    }

    /// get the top context
    #[inline]
    pub fn top(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        unsafe { &mut *root.parent }
    }

    /// get the coroutine context
    #[inline]
    pub fn co_ctx(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        // the root's child is used as the coroutine context pointer
        assert!(!root.child.is_null(), "there is no child of root!");
        unsafe { &mut *root.child }
    }

    /// push the context to the thread context list
    #[inline]
    pub fn push_context(&self, ctx: *mut Context) {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let top = unsafe { &mut *root.parent };
        let new_top = ctx.parent;

        // link top and new ctx
        top.child = ctx;
        ctx.parent = top;

        // save the new top
        root.parent = new_top;
    }

    /// pop the context from the thread context list and return it's parent context
    #[inline]
    pub fn pop_context(&self, ctx: *mut Context) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let parent = unsafe { &mut *ctx.parent };

        // save the old top in ctx's parent
        ctx.parent = root.parent;
        // unlink ctx and it's parent
        parent.child = ptr::null_mut();

        // save the new top
        root.parent = parent;

        parent
    }
}

/// check the current context if it's generator
#[inline]
#[allow(dead_code)]
pub fn is_generator() -> bool {
    let env = ContextStack::current();
    env.top().is_generator()
}

/// get the current context local data
#[inline]
pub fn get_local_data() -> *mut u8 {
    let env = ContextStack::current();
    env.top().local_data
}

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

    #[test]
    fn test_is_context() {
        // this is the root context
        assert!(!is_generator());
    }
}
<commit_msg>only coroutine support local data<commit_after>//! # generator run time support
//!
//! generator run time context management
//!
use std::mem;
use std::ptr;
use std::any::Any;
use std::intrinsics::type_name;

use stack::Stack;
use reg_context::Context as RegContext;

/// each thread has it's own generator context stack
thread_local!(static ROOT_CONTEXT: Box<Context> = {
    let mut root = Box::new(Context::new(0));
    let p = &mut *root as *mut _;
    root.parent = p; // init top to current
    root
});

/// yield error types
#[allow(dead_code)]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Error {
    Cancel,
    TypeErr,
    StackErr,
    ContextErr,
}

/// generator context
pub struct Context {
    /// generator regs context
    pub regs: RegContext,
    /// generator execution stack
    pub stack: Stack,
    /// passed in para for send
    pub para: *mut Any,
    /// this is just a buffer for the return value
    pub ret: *mut Any,
    /// track generator ref, yield will -1, send will +1
    pub _ref: u32,
    /// propagate panic
    pub err: Option<Error>,
    /// context local storage
    pub local_data: *mut u8,

    /// child context
    child: *mut Context,
    /// parent context
    pub parent: *mut Context,
}

impl Context {
    /// return a default generator context
    pub fn new(size: usize) -> Context {
        Context {
            regs: RegContext::empty(),
            stack: Stack::new(size),
            para: unsafe { mem::uninitialized() },
            ret: unsafe { mem::uninitialized() },
            _ref: 1, // none zero means it's not running
            err: None,
            child: ptr::null_mut(),
            parent: ptr::null_mut(),
            local_data: ptr::null_mut(),
        }
    }

    /// judge it's generator context
    #[inline]
    pub fn is_generator(&self) -> bool {
        self.stack.size() > 0
    }

    /// get current generator send para
    #[inline]
    pub fn get_para(&self) -> Option
        where A: Any
    {
        let para = unsafe { &mut *self.para };
        let val = para.downcast_mut::<Option>();
        if val.is_some() {
            val.unwrap().take()
        } else {
            let t = unsafe { type_name::() };
            error!("get yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr);
        }
    }

    /// set current generator return value
    #[inline]
    pub fn set_ret<T>(&mut self, v: T)
        where T: Any
    {
        let ret = unsafe { &mut *self.ret };
        let val = ret.downcast_mut::<Option<T>>();
        if val.is_some() {
            mem::replace(val.unwrap(), Some(v));
        } else {
            let t = unsafe { type_name::<T>() };
            error!("yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr);
        }
    }
}

/// Coroutine managing environment
pub struct ContextStack {
    root: *mut Context,
}

impl ContextStack {
    #[inline]
    pub fn current() -> ContextStack {
        let root = ROOT_CONTEXT.with(|r| &**r as *const _ as *mut Context);
        ContextStack { root: root }
    }

    /// get the top context
    #[inline]
    pub fn top(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        unsafe { &mut *root.parent }
    }

    /// get the coroutine context
    #[inline]
    pub fn co_ctx(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        // the root's child is used as the coroutine context pointer
        assert!(!root.child.is_null(), "there is no child of root!");
        unsafe { &mut *root.child }
    }

    /// push the context to the thread context list
    #[inline]
    pub fn push_context(&self, ctx: *mut Context) {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let top = unsafe { &mut *root.parent };
        let new_top = ctx.parent;

        // link top and new ctx
        top.child = ctx;
        ctx.parent = top;

        // save the new top
        root.parent = new_top;
    }

    /// pop the context from the thread context list and return it's parent context
    #[inline]
    pub fn pop_context(&self, ctx: *mut Context) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let parent = unsafe { &mut *ctx.parent };

        // save the old top in ctx's parent
        ctx.parent = root.parent;
        // unlink ctx and it's parent
        parent.child = ptr::null_mut();

        // save the new top
        root.parent = parent;

        parent
    }
}

/// check the current context if it's generator
#[inline]
#[allow(dead_code)]
pub fn is_generator() -> bool {
    let env = ContextStack::current();
    env.top().is_generator()
}

/// get the current context local data
/// only coroutine support local data
#[inline]
pub fn get_local_data() -> *mut u8 {
    let env = ContextStack::current();
    let root = unsafe { &mut *env.root };
    // the root's child is used as the coroutine context pointer
    if root.child.is_null() {
        return ptr::null_mut();
    }

    let child = unsafe { &*root.child };
    child.local_data
}

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

    #[test]
    fn test_is_context() {
        // this is the root context
        assert!(!is_generator());
    }
}
<|endoftext|>


<commit_before>extern crate toml_edit;

macro_rules! parse_key {
    ($s:expr) => (
        {
            let key = $s.parse::<Key>();
            assert!(key.is_ok());
            &key.unwrap()
        }
    );
}

macro_rules! as_table {
    ($e:ident) => (
        {
            assert!($e.is_table());
            $e.as_table_mut().unwrap()
        }
    );
}

// rusfmt, U Can't Touch This
#[cfg(test)]
#[cfg_attr(rustfmt, rustfmt_skip)]
mod tests {
    use toml_edit::{Document, Key, Value, Table, value, table, array};
    use std::iter::FromIterator;

    struct Test {
        doc: Document,
    }

    fn given(input: &str) -> Test {
        let doc = input.parse::<Document>();
        assert!(doc.is_ok());
        Test {
            doc: doc.unwrap(),
        }
    }

    impl Test {
        fn running<F>(&mut self, func: F) -> &mut Self
            where F: Fn(&mut Table)
        {
            {
                let root = &mut self.doc.root;
                func(root);
            }
            self
        }

        fn produces(&self, expected: &str) {
            assert_eq!(self.doc.to_string(), expected);
        }
    }

// insertion

#[test]
fn test_insert_leaf_table() {
    given(r#"
        [servers]

        [servers.alpha]
        ip = "10.0.0.1"
        dc = "eqdc10"

        [other.table]"#
    ).running(|root| {
        let servers = root.entry("servers");
        let servers = as_table!(servers);
        let entry = servers.entry("beta");
        *entry = table();
        let beta = as_table!(entry);
        *beta.entry("ip") = value("10.0.0.2");
        *beta.entry("dc") = value("eqdc10");
    }).produces(r#"
        [servers]

        [servers.alpha]
        ip = "10.0.0.1"
        dc = "eqdc10"

[servers.beta]
ip = "10.0.0.2"
dc = "eqdc10"

        [other.table]
"#
    );
}

#[test]
fn test_insert_nonleaf_table() {
    given(r#"
        [other.table]"#
    ).running(|root| {
        let servers = root.entry("servers");
        *servers = table();
        let servers = as_table!(servers);
        let alpha = servers.entry("alpha");
        *alpha = table();
        let alpha = as_table!(alpha);
        *alpha.entry("ip") = value("10.0.0.1");
        *alpha.entry("dc") = value("eqdc10");
    }).produces(r#"
        [other.table]

[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"
"#
    );
}

#[test]
fn test_insert_array() {
    given(r#"
        [package]
        title = "withoutarray""#
    ).running(|root| {
        let array = root.entry("bin");
        *array = self::array();
        assert!(array.is_array_of_tables());
        let array = array.as_array_of_tables_mut().unwrap();
        {
            let first = array.append(Table::new());
            *first.entry("hello") = value("world");
        }
        array.append(Table::new());
    }).produces(r#"
        [package]
        title = "withoutarray"

[[bin]]
hello = "world"

[[bin]]
"#
    );
}

#[test]
fn test_insert_values() {
    given(r#"
        [tbl.son]"#
    ).running(|root| {
        let table = root.entry("tbl");
        let table = as_table!(table);
        *table.entry("key1") = value("value1");
        *table.entry("\"key2\"") = value(42);
        *table.entry("'key3'") = value(8.1415926);
    }).produces(r#"
[tbl]
key1 = "value1"
"key2" = 42
'key3' = 8.1415926

        [tbl.son]
"#
    );
}

// removal

#[test]
fn test_remove_leaf_table() {
    given(r#"
        [servers]

        # Indentation (tabs and/or spaces) is allowed but not required
[servers.alpha]
        ip = "10.0.0.1"
        dc = "eqdc10"

        [servers.beta]
        ip = "10.0.0.2"
        dc = "eqdc10""#
    ).running(|root| {
        let servers = root.entry("servers");
        let servers = as_table!(servers);
        assert!(servers.remove("alpha").is_some());
    }).produces(r#"
        [servers]

        [servers.beta]
        ip = "10.0.0.2"
        dc = "eqdc10"
"#
    );
}

#[test]
fn test_remove_nonleaf_table() {
    given(r#"
        title = "not relevant"

        # comment 1
        [a.b.c] # comment 1.1
        key1 = 1 # comment 1.2
        # comment 2
        [b] # comment 2.1
        key2 = 2 # comment 2.2

        # comment 3
        [a] # comment 3.1
        key3 = 3 # comment 3.2
        [[a.'array']]
        b = 1

        [[a.b.c.trololololololo]] # ohohohohoho
        c = 2
        key3 = 42

           # comment on some other table
           [some.other.table]

        # comment 4
        [a.b] # comment 4.1
        key4 = 4 # comment 4.2
        key41 = 41 # comment 4.3

    "#).running(|root| {
        assert!(root.remove("a").is_some());
    }).produces(r#"
        title = "not relevant"
        # comment 2
        [b] # comment 2.1
        key2 = 2 # comment 2.2

           # comment on some other table
           [some.other.table]

    "#);
}

#[test]
fn test_remove_array_entry() {
    given(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [[bin]]
        name = "world"
        path = "src/bin/world/main.rs"

        [dependencies]
        nom = "4.0" # future is here

        [[bin]]
        name = "delete me please"
        path = "src/bin/dmp/main.rs""#
    ).running(|root| {
        let dmp = root.entry("bin");
        assert!(dmp.is_array_of_tables());
        let dmp = dmp.as_array_of_tables_mut().unwrap();
        println!("{:?}", dmp);
        assert_eq!(dmp.len(), 2);
        dmp.remove(1);
        assert_eq!(dmp.len(), 1);
    }).produces(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [[bin]]
        name = "world"
        path = "src/bin/world/main.rs"

        [dependencies]
        nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_array() {
    given(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [[bin]]
        name = "world"
        path = "src/bin/world/main.rs"

        [dependencies]
        nom = "4.0" # future is here

        [[bin]]
        name = "delete me please"
        path = "src/bin/dmp/main.rs""#
    ).running(|root| {
        assert!(root.remove("bin").is_some());
    }).produces(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [dependencies]
        nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_value() {
    given(r#"
        name = "hello"
        # delete this
        version = "1.0.0" # please
        documentation = "https://docs.rs/hello""#
    ).running(|root| {
        let value = root.remove("version");
        assert!(value.is_some());
        let value = value.unwrap();
        assert!(value.is_value());
        let value = value.as_value().unwrap();
        assert!(value.is_str());
        let value = value.as_str().unwrap();
        assert_eq!(value, "1.0.0");
    }).produces(r#"
        name = "hello"
        documentation = "https://docs.rs/hello"
"#
    );
}

#[test]
fn test_remove_last_value_from_implicit() {
    given(r#"
        [a]
        b = 1"#
    ).running(|root| {
        let a = root.entry("a");
        assert!(a.is_table());
        let a = as_table!(a);
        a.set_implicit(true);
        let value = a.remove("b");
        assert!(value.is_some());
        let value = value.unwrap();
        assert!(value.is_value());
        let value = value.as_value().unwrap();
        assert_eq!(value.as_integer(), Some(1));
    }).produces(r#""#);
}

// values

#[test]
fn test_sort_values() {
    given(r#"
        [a.z]

        [a]
        # this comment is attached to b
        b = 2 # as well as this
        a = 1
        c = 3

        [a.y]"#
    ).running(|root| {
        let a = root.entry("a");
        let a = as_table!(a);
        a.sort_values();
    }).produces(r#"
        [a]
        a = 1
        # this comment is attached to b
        b = 2 # as well as this
        c = 3

        [a.z]

        [a.y]
"#
    );
}

macro_rules! as_array {
    ($entry:ident) => (
        {
            assert!($entry.is_value());
            let a = $entry.as_value_mut().unwrap();
            assert!(a.is_array());
            a.as_array_mut().unwrap()
        }
    );
}

#[test]
fn test_insert_into_array() {
    given(r#"
        a = [1,2,3]
        b = []"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_array!(a);
            assert_eq!(a.len(), 3);
            assert!(a.get(2).is_some());
            assert!(a.push(4));
            assert_eq!(a.len(), 4);
            a.fmt();
        }
        let b = root.entry("b");
        let b = as_array!(b);
        assert!(b.is_empty());
        assert!(b.push("hello"));
        assert_eq!(b.len(), 1);
    }).produces(r#"
        a = [1, 2, 3, 4]
        b = ["hello"]
"#
    );
}

#[test]
fn test_remove_from_array() {
    given(r#"
        a = [1, 2, 3, 4]
        b = ["hello"]"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_array!(a);
            assert_eq!(a.len(), 4);
            assert!(a.remove(3).is_integer());
            assert_eq!(a.len(), 3);
        }
        let b = root.entry("b");
        let b = as_array!(b);
        assert_eq!(b.len(), 1);
        assert!(b.remove(0).is_str());
        assert!(b.is_empty());
    }).produces(r#"
        a = [1, 2, 3]
        b = []
"#
    );
}

macro_rules! as_inline_table {
    ($entry:ident) => (
        {
            assert!($entry.is_value());
            let a = $entry.as_value_mut().unwrap();
            assert!(a.is_inline_table());
            a.as_inline_table_mut().unwrap()
        }
    );
}

#[test]
fn test_insert_into_inline_table() {
    given(r#"
        a = {a=2,  c = 3}
        b = {}"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_inline_table!(a);
            assert_eq!(a.len(), 2);
            assert!(a.contains_key("a") && a.get("c").is_some());
            a.get_or_insert("b", 42);
            assert_eq!(a.len(), 3);
            a.fmt();
        }
        let b = root.entry("b");
        let b = as_inline_table!(b);
        assert!(b.is_empty());
        b.get_or_insert("'hello'", "world");
        assert_eq!(b.len(), 1);
        b.fmt()
    }).produces(r#"
        a = { a = 2, c = 3, b = 42 }
        b = { 'hello' = "world" }
"#
    );
}

#[test]
fn test_remove_from_inline_table() {
    given(r#"
        a = {a=2,  c = 3, b = 42}
        b = {'hello' = "world"}"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_inline_table!(a);
            assert_eq!(a.len(), 3);
            assert!(a.remove("c").is_some());
            assert_eq!(a.len(), 2);
        }
        let b = root.entry("b");
        let b = as_inline_table!(b);
        assert_eq!(b.len(), 1);
        assert!(b.remove("hello").is_some());
        assert!(b.is_empty());
    }).produces(r#"
        a = {a=2, b = 42}
        b = {}
"#
    );
}

#[test]
fn test_inline_table_append() {
    let mut a = Value::from_iter(vec![
        (parse_key!("a"), 1),
        (parse_key!("b"), 2),
        (parse_key!("c"), 3),
    ]);
    let a = a.as_inline_table_mut().unwrap();

    let mut b = Value::from_iter(vec![
        (parse_key!("c"), 4),
        (parse_key!("d"), 5),
        (parse_key!("e"), 6),
    ]);
    let b = b.as_inline_table_mut().unwrap();

    b.merge_into(a);
    assert_eq!(a.len(), 5);
    assert!(a.contains_key("e"));
    assert!(b.is_empty());
}

} // mod tests
<commit_msg>refactor(test_edit): use new Index API<commit_after>extern crate toml_edit;

macro_rules! parse_key {
    ($s:expr) => (
        {
            let key = $s.parse::<Key>();
            assert!(key.is_ok());
            &key.unwrap()
        }
    );
}

macro_rules! as_table {
    ($e:ident) => (
        {
            assert!($e.is_table());
            $e.as_table_mut().unwrap()
        }
    );
}

// rusfmt, U Can't Touch This
#[cfg(test)]
#[cfg_attr(rustfmt, rustfmt_skip)]
mod tests {
    use toml_edit::{Document, Key, Value, Table, value, table, array};
    use std::iter::FromIterator;

    struct Test {
        doc: Document,
    }

    fn given(input: &str) -> Test {
        let doc = input.parse::<Document>();
        assert!(doc.is_ok());
        Test {
            doc: doc.unwrap(),
        }
    }

    impl Test {
        fn running<F>(&mut self, func: F) -> &mut Self
            where F: Fn(&mut Table)
        {
            {
                let root = &mut self.doc.root;
                func(root);
            }
            self
        }

        fn produces(&self, expected: &str) {
            assert_eq!(self.doc.to_string(), expected);
        }
    }

// insertion

#[test]
fn test_insert_leaf_table() {
    given(r#"
        [servers]

        [servers.alpha]
        ip = "10.0.0.1"
        dc = "eqdc10"

        [other.table]"#
    ).running(|root| {
        root["servers"]["beta"] = table();
        root["servers"]["beta"]["ip"] = value("10.0.0.2");
        root["servers"]["beta"]["dc"] = value("eqdc10");
    }).produces(r#"
        [servers]

        [servers.alpha]
        ip = "10.0.0.1"
        dc = "eqdc10"

[servers.beta]
ip = "10.0.0.2"
dc = "eqdc10"

        [other.table]
"#
    );
}

#[test]
fn test_insert_nonleaf_table() {
    given(r#"
        [other.table]"#
    ).running(|root| {
        root["servers"] = table();
        root["servers"]["alpha"] = table();
        root["servers"]["alpha"]["ip"] = value("10.0.0.1");
        root["servers"]["alpha"]["dc"] = value("eqdc10");
    }).produces(r#"
        [other.table]

[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"
"#
    );
}

#[test]
fn test_insert_array() {
    given(r#"
        [package]
        title = "withoutarray""#
    ).running(|root| {
        root["bin"] = array();
        assert!(root["bin"].is_array_of_tables());
        let array = root["bin"].as_array_of_tables_mut().unwrap();
        {
            let first = array.append(Table::new());
            first["hello"] = value("world");
        }
        array.append(Table::new());
    }).produces(r#"
        [package]
        title = "withoutarray"

[[bin]]
hello = "world"

[[bin]]
"#
    );
}

#[test]
fn test_insert_values() {
    given(r#"
        [tbl.son]"#
    ).running(|root| {
        root["tbl"]["key1"] = value("value1");
        root["tbl"]["\"key2\""] = value(42);
        root["tbl"]["'key3'"] = value(8.1415926);
    }).produces(r#"
[tbl]
key1 = "value1"
"key2" = 42
'key3' = 8.1415926

        [tbl.son]
"#
    );
}

// removal

#[test]
fn test_remove_leaf_table() {
    given(r#"
        [servers]

        # Indentation (tabs and/or spaces) is allowed but not required
[servers.alpha]
        ip = "10.0.0.1"
        dc = "eqdc10"

        [servers.beta]
        ip = "10.0.0.2"
        dc = "eqdc10""#
    ).running(|root| {
        let servers = root.entry("servers");
        let servers = as_table!(servers);
        assert!(servers.remove("alpha").is_some());
    }).produces(r#"
        [servers]

        [servers.beta]
        ip = "10.0.0.2"
        dc = "eqdc10"
"#
    );
}

#[test]
fn test_remove_nonleaf_table() {
    given(r#"
        title = "not relevant"

        # comment 1
        [a.b.c] # comment 1.1
        key1 = 1 # comment 1.2
        # comment 2
        [b] # comment 2.1
        key2 = 2 # comment 2.2

        # comment 3
        [a] # comment 3.1
        key3 = 3 # comment 3.2
        [[a.'array']]
        b = 1

        [[a.b.c.trololololololo]] # ohohohohoho
        c = 2
        key3 = 42

           # comment on some other table
           [some.other.table]

        # comment 4
        [a.b] # comment 4.1
        key4 = 4 # comment 4.2
        key41 = 41 # comment 4.3

    "#).running(|root| {
        assert!(root.remove("a").is_some());
    }).produces(r#"
        title = "not relevant"
        # comment 2
        [b] # comment 2.1
        key2 = 2 # comment 2.2

           # comment on some other table
           [some.other.table]

    "#);
}

#[test]
fn test_remove_array_entry() {
    given(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [[bin]]
        name = "world"
        path = "src/bin/world/main.rs"

        [dependencies]
        nom = "4.0" # future is here

        [[bin]]
        name = "delete me please"
        path = "src/bin/dmp/main.rs""#
    ).running(|root| {
        let dmp = root.entry("bin");
        assert!(dmp.is_array_of_tables());
        let dmp = dmp.as_array_of_tables_mut().unwrap();
        assert_eq!(dmp.len(), 2);
        dmp.remove(1);
        assert_eq!(dmp.len(), 1);
    }).produces(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [[bin]]
        name = "world"
        path = "src/bin/world/main.rs"

        [dependencies]
        nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_array() {
    given(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [[bin]]
        name = "world"
        path = "src/bin/world/main.rs"

        [dependencies]
        nom = "4.0" # future is here

        [[bin]]
        name = "delete me please"
        path = "src/bin/dmp/main.rs""#
    ).running(|root| {
        assert!(root.remove("bin").is_some());
    }).produces(r#"
        [package]
        name = "hello"
        version = "1.0.0"

        [dependencies]
        nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_value() {
    given(r#"
        name = "hello"
        # delete this
        version = "1.0.0" # please
        documentation = "https://docs.rs/hello""#
    ).running(|root| {
        let value = root.remove("version");
        assert!(value.is_some());
        let value = value.unwrap();
        assert!(value.is_value());
        let value = value.as_value().unwrap();
        assert!(value.is_str());
        let value = value.as_str().unwrap();
        assert_eq!(value, "1.0.0");
    }).produces(r#"
        name = "hello"
        documentation = "https://docs.rs/hello"
"#
    );
}

#[test]
fn test_remove_last_value_from_implicit() {
    given(r#"
        [a]
        b = 1"#
    ).running(|root| {
        let a = root.entry("a");
        assert!(a.is_table());
        let a = as_table!(a);
        a.set_implicit(true);
        let value = a.remove("b");
        assert!(value.is_some());
        let value = value.unwrap();
        assert!(value.is_value());
        let value = value.as_value().unwrap();
        assert_eq!(value.as_integer(), Some(1));
    }).produces(r#""#);
}

// values

#[test]
fn test_sort_values() {
    given(r#"
        [a.z]

        [a]
        # this comment is attached to b
        b = 2 # as well as this
        a = 1
        c = 3

        [a.y]"#
    ).running(|root| {
        let a = root.entry("a");
        let a = as_table!(a);
        a.sort_values();
    }).produces(r#"
        [a]
        a = 1
        # this comment is attached to b
        b = 2 # as well as this
        c = 3

        [a.z]

        [a.y]
"#
    );
}

macro_rules! as_array {
    ($entry:ident) => (
        {
            assert!($entry.is_value());
            let a = $entry.as_value_mut().unwrap();
            assert!(a.is_array());
            a.as_array_mut().unwrap()
        }
    );
}

#[test]
fn test_insert_into_array() {
    given(r#"
        a = [1,2,3]
        b = []"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_array!(a);
            assert_eq!(a.len(), 3);
            assert!(a.get(2).is_some());
            assert!(a.push(4));
            assert_eq!(a.len(), 4);
            a.fmt();
        }
        let b = root.entry("b");
        let b = as_array!(b);
        assert!(b.is_empty());
        assert!(b.push("hello"));
        assert_eq!(b.len(), 1);
    }).produces(r#"
        a = [1, 2, 3, 4]
        b = ["hello"]
"#
    );
}

#[test]
fn test_remove_from_array() {
    given(r#"
        a = [1, 2, 3, 4]
        b = ["hello"]"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_array!(a);
            assert_eq!(a.len(), 4);
            assert!(a.remove(3).is_integer());
            assert_eq!(a.len(), 3);
        }
        let b = root.entry("b");
        let b = as_array!(b);
        assert_eq!(b.len(), 1);
        assert!(b.remove(0).is_str());
        assert!(b.is_empty());
    }).produces(r#"
        a = [1, 2, 3]
        b = []
"#
    );
}

macro_rules! as_inline_table {
    ($entry:ident) => (
        {
            assert!($entry.is_value());
            let a = $entry.as_value_mut().unwrap();
            assert!(a.is_inline_table());
            a.as_inline_table_mut().unwrap()
        }
    );
}

#[test]
fn test_insert_into_inline_table() {
    given(r#"
        a = {a=2,  c = 3}
        b = {}"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_inline_table!(a);
            assert_eq!(a.len(), 2);
            assert!(a.contains_key("a") && a.get("c").is_some());
            a.get_or_insert("b", 42);
            assert_eq!(a.len(), 3);
            a.fmt();
        }
        let b = root.entry("b");
        let b = as_inline_table!(b);
        assert!(b.is_empty());
        b.get_or_insert("'hello'", "world");
        assert_eq!(b.len(), 1);
        b.fmt()
    }).produces(r#"
        a = { a = 2, c = 3, b = 42 }
        b = { 'hello' = "world" }
"#
    );
}

#[test]
fn test_remove_from_inline_table() {
    given(r#"
        a = {a=2,  c = 3, b = 42}
        b = {'hello' = "world"}"#
    ).running(|root| {
        {
            let a = root.entry("a");
            let a = as_inline_table!(a);
            assert_eq!(a.len(), 3);
            assert!(a.remove("c").is_some());
            assert_eq!(a.len(), 2);
        }
        let b = root.entry("b");
        let b = as_inline_table!(b);
        assert_eq!(b.len(), 1);
        assert!(b.remove("hello").is_some());
        assert!(b.is_empty());
    }).produces(r#"
        a = {a=2, b = 42}
        b = {}
"#
    );
}

#[test]
fn test_inline_table_append() {
    let mut a = Value::from_iter(vec![
        (parse_key!("a"), 1),
        (parse_key!("b"), 2),
        (parse_key!("c"), 3),
    ]);
    let a = a.as_inline_table_mut().unwrap();

    let mut b = Value::from_iter(vec![
        (parse_key!("c"), 4),
        (parse_key!("d"), 5),
        (parse_key!("e"), 6),
    ]);
    let b = b.as_inline_table_mut().unwrap();

    b.merge_into(a);
    assert_eq!(a.len(), 5);
    assert!(a.contains_key("e"));
    assert!(b.is_empty());
}

} // mod tests
<|endoftext|>


<commit_before>use std::io;
use std::io::Read;
use command::Command;

type Cell = u8;

pub struct VM {
    cells: [Cell; 30_000],
    data_pointer: usize,
}

impl VM {
    pub fn new() -> VM {
        VM { cells: [0; 30_000], data_pointer: 0 }
    }

    pub fn output(&self) {
        print!("{}", self.cells[self.data_pointer] as char);
    }

    pub fn input(&mut self) {
        match io::stdin().bytes().next() {
            Some(res) => match res {
                Ok(value) => self.cells[self.data_pointer] = value,
                Err(error) => panic!(error),
            },
            None => {}
        }
    }

    pub fn increment(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer] + 1;
    }

    pub fn decrement(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer] - 1;
    }

    pub fn right_shift(&mut self) {
        self.data_pointer = self.data_pointer + 1;
    }

    pub fn left_shift(&mut self) {
        self.data_pointer = self.data_pointer - 1;
    }

    pub fn apply(&mut self, command: &Command) {
        match command {
            &Command::Output => self.output(),
            &Command::Input => self.input(),
            &Command::Increment => self.increment(),
            &Command::Decrement => self.decrement(),
            &Command::RightShift => self.right_shift(),
            &Command::LeftShift => self.left_shift(),
        }
    }

    pub fn apply_many(&mut self, commands: &[Command]) {
        for command in commands {
            self.apply(&command)
        }
    }
}
<commit_msg>Print a newline after each character input<commit_after>use std::io;
use std::io::Read;
use command::Command;

type Cell = u8;

pub struct VM {
    cells: [Cell; 30_000],
    data_pointer: usize,
}

impl VM {
    pub fn new() -> VM {
        VM { cells: [0; 30_000], data_pointer: 0 }
    }

    pub fn output(&self) {
        print!("{}", self.cells[self.data_pointer] as char);
    }

    pub fn input(&mut self) {
        match io::stdin().bytes().next() {
            Some(res) => match res {
                Ok(value) => self.cells[self.data_pointer] = value,
                Err(error) => panic!(error),
            },
            None => {}
        }
        println!("");
    }

    pub fn increment(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer] + 1;
    }

    pub fn decrement(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer] - 1;
    }

    pub fn right_shift(&mut self) {
        self.data_pointer = self.data_pointer + 1;
    }

    pub fn left_shift(&mut self) {
        self.data_pointer = self.data_pointer - 1;
    }

    pub fn apply(&mut self, command: &Command) {
        match command {
            &Command::Output => self.output(),
            &Command::Input => self.input(),
            &Command::Increment => self.increment(),
            &Command::Decrement => self.decrement(),
            &Command::RightShift => self.right_shift(),
            &Command::LeftShift => self.left_shift(),
        }
    }

    pub fn apply_many(&mut self, commands: &[Command]) {
        for command in commands {
            self.apply(&command)
        }
    }
}
<|endoftext|>


<commit_before>
use itertools::*;
use byteorder::*;

use std::collections::HashMap;
use std::fmt::{self, Debug, Formatter};

pub type RegisterT = i64;
type Endianness = LittleEndian;

pub fn exec_program(program: Program) -> RegisterT {
    use self::Instruction::*;

    macro_rules! read_u16 {
        ($instr:ident, $ip:ident) => {{
            let val = Endianness::read_u16(&$instr[$ip..$ip + 2]);
            $ip += 2;
            val
        }}
    }

    macro_rules! read_u8 {
        ($instr:ident, $ip:ident) => {{
            let val = $instr[$ip];
            $ip += 1;
            val
        }}
    }

    macro_rules! read_i64 {
        ($instr:ident, $ip:ident) => {{
            let val = Endianness::read_i64(&$instr[$ip..$ip + 8]);
            $ip += 8;
            val
        }}
    }

    let mut ip = 0; // instruction pointer
    let mut data_stack = Vec::with_capacity(128);
    let mut ret_stack = Vec::with_capacity(128);
    let mut var_slots = vec![0; program.num_vars];
    let instructions = program.instructions;

    loop {
        let i = ip;
        ip += 1;
        let instr = as_inst(instructions[i]);
        match instr {
            ConstI64 => {
                let val = read_i64!(instructions, ip);
                data_stack.push(val);
            }
            Load => {
                let var_slot = read_u16!(instructions, ip) as usize;
                data_stack.push(var_slots[var_slot]);
            }
            Store => {
                let var_slot = read_u16!(instructions, ip) as usize;
                var_slots[var_slot] = data_stack.pop().expect("stack underflow: store");
            }
            Add => {
                let num_args = read_u8!(instructions, ip) as usize;
                assert!(data_stack.len() >= num_args, "stack underflow: add");
                let idx = data_stack.len() - num_args;
                let val = data_stack.drain(idx..).fold(0, |acc, e| acc + e);
                data_stack.push(val);
            }
            Drop => {
                data_stack.pop().expect("stack underflow: drop");
            }
            Call => {
                let jmp_addr = read_u16!(instructions, ip) as usize;
                ret_stack.push(ip);
                ip = jmp_addr;
            }
            Ret => {
                ip = ret_stack.pop().expect("return stack underflow: ret");
            }
            Exit => return data_stack.pop().expect("stack underflow: exit"),
        }
    }
}

fn as_byte(i: Instruction) -> u8 {
    unsafe { ::std::mem::transmute(i) }
}

fn as_inst(b: u8) -> Instruction {
    unsafe { ::std::mem::transmute(b) }
}

pub struct Program {
    instructions: Vec<u8>,
    num_vars: usize,
}

pub fn new_program() -> ProgramBuilder {
    ProgramBuilder {
        fn_cursor: vec![],
        fn_defs: Default::default(),
        instructions: vec![],
        num_vars: 0,
    }
}

pub struct ProgramBuilder {
    // stack of ids
    fn_cursor: Vec<u32>,
    // var_id -> fn_def
    fn_defs: HashMap<u32, Vec<u8>>,
    instructions: Vec<u8>,
    num_vars: u16,
}

macro_rules! current_def {
    ($e:ident) => {
        match $e.fn_cursor.len() {
            0 => &mut $e.instructions,
            n => $e.fn_defs.get_mut(&$e.fn_cursor[n - 1]).unwrap(),
        }
    }
}

impl ProgramBuilder {
    pub fn begin_fn_def(&mut self, id: u32) {
        self.fn_defs.insert(id, vec![]);
        self.fn_cursor.push(id);
    }

    pub fn end_fn_def(&mut self) {
        self.fn_cursor.pop();
    }

    pub fn def_var(&mut self) -> u16 {
        let tmp = self.num_vars;
        self.num_vars += 1;
        tmp
    }

    pub fn load_var(&mut self, var_slot_idx: u16) {
        assert!(self.num_vars >= var_slot_idx + 1);
        let def = current_def!(self);
        def.push(as_byte(Instruction::Load));
        def.write_u16::<Endianness>(var_slot_idx).unwrap();
    }

    pub fn store_var(&mut self, var_slot_idx: u16) {
        assert!(self.num_vars >= var_slot_idx + 1);
        let def = current_def!(self);
        def.push(as_byte(Instruction::Store));
        def.write_u16::<Endianness>(var_slot_idx).unwrap();
    }

    pub fn load_const(&mut self, val: RegisterT) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::ConstI64));
        def.write_i64::<Endianness>(val).unwrap();
    }

    pub fn add(&mut self, num_args: u8) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Add));
        def.push(num_args);
    }

    pub fn drop(&mut self) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Drop));
    }

    pub fn exit(&mut self) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Exit));
    }

    pub fn finish(self) -> Program {
        // assemble fn_defs
        Program {
            instructions: self.instructions,
            num_vars: self.num_vars as usize,
        }
    }
}

#[derive(Copy, Clone, PartialEq, Debug)]
enum Instruction {
    ConstI64,
    Load,
    Store,
    Add,
    Drop,
    Call,
    Ret,
    Exit,
}

impl Debug for Program {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        use self::Instruction::*;

        fn format_inst(idx: usize, i: Instruction, argn: &ToString) -> Option<String> {
            let mnemonic = match i {
                ConstI64 => "i64",
                Load => "load",
                Store => "stor",
                Call => "call",
                Add => "add",
                Drop => "drop",
                Exit => "exit",
                Ret => "ret",
            };
            Some(format!("{:>4}:\t{:>4} {:>3}", idx, mnemonic, argn.to_string()))
        }

        let lines = self.instructions
                        .iter()
                        .enumerate()
                        .batching(|mut it| {
                            match it.next() {
                                None => None,
                                Some((idx, &b)) => {
                                    let inst = as_inst(b);
                                    match inst {
                                        ConstI64 => {
                                            let mut raw = vec![];
                                            for _ in 0..8 {
                                                match it.next() {
                                                    Some((_, &b)) => {
                                                        raw.push(b);
                                                    }
                                                    _ => return None,
                                                }
                                            }
                                            let val = Endianness::read_i64(&*raw);
                                            format_inst(idx, inst, &val)
                                        }
                                        Load => {
                                            let bytes = [*it.next().unwrap().1,
                                                         *it.next().unwrap().1];
                                            let val = Endianness::read_u16(&bytes[..]);
                                            format_inst(idx, inst, &val)
                                        }
                                        Store => {
                                            let bytes = [*it.next().unwrap().1,
                                                         *it.next().unwrap().1];
                                            let val = Endianness::read_u16(&bytes[..]);
                                            format_inst(idx, inst, &val)
                                        }
                                        Call => {
                                            let bytes = [*it.next().unwrap().1,
                                                         *it.next().unwrap().1];
                                            let val = Endianness::read_u16(&bytes[..]);
                                            format_inst(idx, inst, &val)
                                        }
                                        Add => format_inst(idx, inst, it.next().unwrap().1),
                                        Ret => format_inst(idx, inst, &""),
                                        Exit => format_inst(idx, inst, &""),
                                        Drop => format_inst(idx, inst, &""),
                                    }
                                }
                            }
                        })
                        .join("\n");

        writeln!(f, "\n{}", lines)
    }
}

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

    #[test]
    fn raw_program_addition() {
        // (+ 1 2 3)

        let mut p = new_program();

        p.load_const(1);
        p.load_const(2);
        p.load_const(3);
        p.add(3);
        p.exit();

        let p = p.finish();

        assert_eq!(exec_program(p), 6);
    }

    #[test]
    fn var() {
        // (def x 7)
        // (+ x 3)

        let mut p = new_program();

        let x = p.def_var();
        p.load_const(7);
        p.store_var(x);
        p.load_const(3);
        p.load_var(x);
        p.add(2);
        p.exit();

        let p = p.finish();

        assert_eq!(exec_program(p), 10);
    }
}
<commit_msg>Call/return in bytecode<commit_after>
use itertools::*;
use byteorder::*;

use std::collections::HashMap;
use std::fmt::{self, Debug, Formatter};

pub type RegisterT = i64;
type Endianness = LittleEndian;

pub fn exec_program(program: Program) -> RegisterT {
    use self::Instruction::*;

    macro_rules! read_u16 {
        ($instr:ident, $ip:ident) => {{
            let val = Endianness::read_u16(&$instr[$ip..$ip + 2]);
            $ip += 2;
            val
        }}
    }

    macro_rules! read_u8 {
        ($instr:ident, $ip:ident) => {{
            let val = $instr[$ip];
            $ip += 1;
            val
        }}
    }

    macro_rules! read_i64 {
        ($instr:ident, $ip:ident) => {{
            let val = Endianness::read_i64(&$instr[$ip..$ip + 8]);
            $ip += 8;
            val
        }}
    }

    let mut ip = 0; // instruction pointer
    let mut data_stack = Vec::with_capacity(128);
    let mut ret_stack = Vec::with_capacity(128);
    let mut var_slots = vec![0; program.num_vars];
    let instructions = program.instructions;

    loop {
        let i = ip;
        ip += 1;
        let instr = as_inst(instructions[i]);
        match instr {
            ConstI64 => {
                let val = read_i64!(instructions, ip);
                data_stack.push(val);
            }
            Load => {
                let var_slot = read_u16!(instructions, ip) as usize;
                data_stack.push(var_slots[var_slot]);
            }
            Store => {
                let var_slot = read_u16!(instructions, ip) as usize;
                var_slots[var_slot] = data_stack.pop().expect("stack underflow: store");
            }
            Add => {
                let num_args = read_u8!(instructions, ip) as usize;
                assert!(data_stack.len() >= num_args, "stack underflow: add");
                let idx = data_stack.len() - num_args;
                let val = data_stack.drain(idx..).fold(0, |acc, e| acc + e);
                data_stack.push(val);
            }
            Drop => {
                data_stack.pop().expect("stack underflow: drop");
            }
            Call => {
                let jmp_addr = read_u16!(instructions, ip) as usize;
                ret_stack.push(ip);
                ip = jmp_addr;
            }
            Ret => {
                ip = ret_stack.pop().expect("return stack underflow: ret");
            }
            Exit => return data_stack.pop().expect("stack underflow: exit"),
        }
    }
}

fn as_byte(i: Instruction) -> u8 {
    unsafe { ::std::mem::transmute(i) }
}

fn as_inst(b: u8) -> Instruction {
    unsafe { ::std::mem::transmute(b) }
}

pub struct Program {
    instructions: Vec<u8>,
    num_vars: usize,
}

pub fn new_program() -> ProgramBuilder {
    ProgramBuilder {
        fn_cursor: vec![],
        fn_defs: Default::default(),
        instructions: vec![],
        num_vars: 0,
    }
}

pub struct ProgramBuilder {
    // stack of ids
    fn_cursor: Vec<u32>,
    // var_id -> fn_def
    fn_defs: HashMap<u32, Vec<u8>>,
    instructions: Vec<u8>,
    num_vars: u16,
}

macro_rules! current_def {
    ($e:ident) => {
        match $e.fn_cursor.len() {
            0 => &mut $e.instructions,
            n => $e.fn_defs.get_mut(&$e.fn_cursor[n - 1]).unwrap(),
        }
    }
}

impl ProgramBuilder {
    pub fn begin_fn_def(&mut self, id: u32) {
        self.fn_defs.insert(id, vec![]);
        self.fn_cursor.push(id);
    }

    pub fn end_fn_def(&mut self) {
        self.fn_cursor.pop().unwrap();
    }

    pub fn call(&mut self, jmp_addr: u16) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Call));
        def.write_u16::<Endianness>(jmp_addr).unwrap();
    }

    pub fn ret(&mut self) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Ret));
    }

    pub fn def_var(&mut self) -> u16 {
        let tmp = self.num_vars;
        self.num_vars += 1;
        tmp
    }

    pub fn load_var(&mut self, var_slot_idx: u16) {
        assert!(self.num_vars >= var_slot_idx + 1);
        let def = current_def!(self);
        def.push(as_byte(Instruction::Load));
        def.write_u16::<Endianness>(var_slot_idx).unwrap();
    }

    pub fn store_var(&mut self, var_slot_idx: u16) {
        assert!(self.num_vars >= var_slot_idx + 1);
        let def = current_def!(self);
        def.push(as_byte(Instruction::Store));
        def.write_u16::<Endianness>(var_slot_idx).unwrap();
    }

    pub fn load_const(&mut self, val: RegisterT) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::ConstI64));
        def.write_i64::<Endianness>(val).unwrap();
    }

    pub fn add(&mut self, num_args: u8) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Add));
        def.push(num_args);
    }

    pub fn drop(&mut self) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Drop));
    }

    pub fn exit(&mut self) {
        let def = current_def!(self);
        def.push(as_byte(Instruction::Exit));
    }

    pub fn finish(self) -> Program {
        // assemble fn_defs
        Program {
            instructions: self.instructions,
            num_vars: self.num_vars as usize,
        }
    }
}

#[derive(Copy, Clone, PartialEq, Debug)]
enum Instruction {
    ConstI64,
    Load,
    Store,
    Add,
    Drop,
    Call,
    Ret,
    Exit,
}

impl Debug for Program {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        use self::Instruction::*;

        fn format_inst(idx: usize, i: Instruction, argn: &ToString) -> Option<String> {
            let mnemonic = match i {
                ConstI64 => "i64",
                Load => "load",
                Store => "stor",
                Call => "call",
                Add => "add",
                Drop => "drop",
                Exit => "exit",
                Ret => "ret",
            };
            Some(format!("{:>4}:\t{:>4} {:>3}", idx, mnemonic, argn.to_string()))
        }

        let lines = self.instructions
                        .iter()
                        .enumerate()
                        .batching(|mut it| {
                            match it.next() {
                                None => None,
                                Some((idx, &b)) => {
                                    let inst = as_inst(b);
                                    match inst {
                                        ConstI64 => {
                                            let mut raw = vec![];
                                            for _ in 0..8 {
                                                match it.next() {
                                                    Some((_, &b)) => {
                                                        raw.push(b);
                                                    }
                                                    _ => return None,
                                                }
                                            }
                                            let val = Endianness::read_i64(&*raw);
                                            format_inst(idx, inst, &val)
                                        }
                                        Load => {
                                            let bytes = [*it.next().unwrap().1,
                                                         *it.next().unwrap().1];
                                            let val = Endianness::read_u16(&bytes[..]);
                                            format_inst(idx, inst, &val)
                                        }
                                        Store => {
                                            let bytes = [*it.next().unwrap().1,
                                                         *it.next().unwrap().1];
                                            let val = Endianness::read_u16(&bytes[..]);
                                            format_inst(idx, inst, &val)
                                        }
                                        Call => {
                                            let bytes = [*it.next().unwrap().1,
                                                         *it.next().unwrap().1];
                                            let val = Endianness::read_u16(&bytes[..]);
                                            format_inst(idx, inst, &val)
                                        }
                                        Add => format_inst(idx, inst, it.next().unwrap().1),
                                        Ret => format_inst(idx, inst, &""),
                                        Exit => format_inst(idx, inst, &""),
                                        Drop => format_inst(idx, inst, &""),
                                    }
                                }
                            }
                        })
                        .join("\n");

        writeln!(f, "\n{}", lines)
    }
}

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

    #[test]
    fn raw_program_addition() {
        // (+ 1 2 3)

        let mut p = new_program();

        p.load_const(1);
        p.load_const(2);
        p.load_const(3);
        p.add(3);
        p.exit();

        let p = p.finish();

        assert_eq!(exec_program(p), 6);
    }

    #[test]
    fn var() {
        // (def x 7)
        // (+ x 3)

        let mut p = new_program();

        let x = p.def_var();
        p.load_const(7);
        p.store_var(x);
        p.load_const(3);
        p.load_var(x);
        p.add(2);
        p.exit();

        let p = p.finish();

        assert_eq!(exec_program(p), 10);
    }

    #[test]
    fn calls() {
        // (defn f [x]
        //   (+ 2 x))
        // (defn g [x]
        //   (f x))
        // (g (f 4))

        let mut p = new_program();

        p.load_const(4);
        p.call(16);
        p.call(28);
        p.exit();
        p.load_const(2); // addr 16
        p.add(2);
        p.ret();
        p.call(16); // addr 28
        p.ret();

        let p = p.finish();

        assert_eq!(exec_program(p), 8);
    }
}
<|endoftext|>


<commit_before>extern crate crowbook_localize;

use crowbook_localize::Extractor;

fn main() {
    let mut extractor = Extractor::new();
    extractor.add_messages_from_dir("src/").unwrap();
    println!("{}", extractor.generate_pot_file());
}
<commit_msg>Change crowbook_localize to crowbook_intl<commit_after>extern crate crowbook_intl;

use crowbook_intl::Extractor;

fn main() {
    let mut extractor = Extractor::new();
    extractor.add_messages_from_dir("src/").unwrap();
    println!("{}", extractor.generate_pot_file());
}
<|endoftext|>