<commit_before>#![crate_name = "stdbuf"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Dorota Kapturkiewicz <dokaptur@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;
extern crate libc;
use getopts::{optopt, optflag, getopts, usage, Matches, OptGroup};
use std::os;
use std::num;
use std::io::Command;
use std::iter::range_inclusive;
use std::num::Int;

static NAME: &'static str = "stdbuf";
static VERSION: &'static str = "1.0.0";

#[deriving(Show)]
enum BufferType {
	Default,
	Unbuffered,
	Line,
	Size(u64)
}

#[deriving(Show)]
struct ProgramOptions {
	stdin : BufferType,
	stdout : BufferType,
	stderr : BufferType,
}

enum ErrMsg {
	Retry,
	Fatal
}

enum OkMsg {
	Buffering,
	Help,
	Version
}

fn print_version() {
	println!("{} version {}", NAME, VERSION);
}

fn print_usage(opts: &[OptGroup]) {
	let brief =
		"Usage: stdbuf OPTION... COMMAND\nRun COMMAND, with modified buffering operations for its standard streams\nMandatory arguments to long options are mandatory for short options too.";
	let explaination =
		"If MODE is 'L' the corresponding stream will be line buffered.\nThis option is invalid with standard input.\n\nIf MODE is '0' the corresponding stream will be unbuffered.\n\nOtherwise MODE is a number which may be followed by one of the following:\n\nKB 1000, K 1024, MB 1000*1000, M 1024*1024, and so on for G, T, P, E, Z, Y.\nIn this case the corresponding stream will be fully buffered with the buffer size set to MODE bytes.\n\nNOTE: If COMMAND adjusts the buffering of its standard streams ('tee' does for e.g.) then that will override corresponding settings changed by 'stdbuf'.\nAlso some filters (like 'dd' and 'cat' etc.) don't use streams for I/O, and are thus unaffected by 'stdbuf' settings.\n";
	println!("{}\n{}", getopts::usage(brief, opts), explaination);
}

fn parse_size(size : &str) -> Option<u64> {
	let ext = size.trim_left_chars(|c: char| c.is_digit(10));
	let num = size.trim_right_chars(|c: char| c.is_alphabetic());
	let mut recovered = num.to_string();
	recovered.push_str(ext);
	if recovered.as_slice() != size {
		return None;
	}
	let buf_size : u64 = match from_str(num) {
		Some(m) => m,
		None => return None,
	};
	let (power, base) : (uint, u64) = match ext {
		"" => (0, 0),
		"KB" => (1, 1024),
		"K" => (1, 1000),
		"MB" => (2, 1024),
		"M" => (2, 1000),
		"GB" => (3, 1024),
		"G" => (3, 1000),
		"TB" => (4, 1024),
		"T" => (4, 1000),
		"PB" => (5, 1024),
		"P" => (5, 1000),
		"EB" => (6, 1024),
		"E" => (6, 1000),
		"ZB" => (7, 1024),
		"Z" => (7, 1000),
		"YB" => (8, 1024),
		"Y" => (8, 1000),
		_ => return None,
	};
	Some(buf_size * base.pow(power))
}

fn check_option(matches : &Matches, name : &str, modified : &mut bool) -> Option<BufferType> {
	match matches.opt_str(name) {
		Some(value) => {
			*modified = true;
			match value.as_slice() {
				"0" => Some(BufferType::Unbuffered),
				"L" => {
					if name == "input" {
						println!("stdbuf: line buffering stdin is meaningless");
						None
					} else {
						Some(BufferType::Line)
					}
				},
				x => {
					let size = match parse_size(x) {
						Some(m) => m,
						None => { println!("Invalid mode {}", x); return None }
					};
					Some(BufferType::Size(size))
				},
			}
		},
		None => Some(BufferType::Default),
	}
}

fn parse_options(args : &[String], options : &mut ProgramOptions, optgrps : &[OptGroup]) -> Result<OkMsg, ErrMsg> {
	let matches = match getopts(args, optgrps) {
		Ok(m) => m,
		Err(_) => return Err(ErrMsg::Retry)
	};
	if matches.opt_present("help") {
		return Ok(OkMsg::Help);
	}
	if matches.opt_present("version") {
		return Ok(OkMsg::Version);
	}
	let mut modified = false;
	options.stdin = try!(check_option(&matches, "input", &mut modified).ok_or(ErrMsg::Fatal));
	options.stdout = try!(check_option(&matches, "output", &mut modified).ok_or(ErrMsg::Fatal));
	options.stderr = try!(check_option(&matches, "error", &mut modified).ok_or(ErrMsg::Fatal));

	if matches.free.len() != 1 {
		return Err(ErrMsg::Retry);
	}
	if !modified {
		println!("stdbuf: you must specify a buffering mode option");
		return Err(ErrMsg::Fatal);
	}
	Ok(OkMsg::Buffering)
}

pub fn uumain(args: Vec<String>) -> int {
	let optgrps = [
		optopt("i", "input", "adjust standard input stream buffering", "MODE"),
		optopt("o", "output", "adjust standard output stream buffering", "MODE"),
		optopt("e", "error", "adjust standard error stream buffering", "MODE"),
		optflag("", "help", "display this help and exit"),
		optflag("", "version", "output version information and exit"),
	];
	let mut options = ProgramOptions{ stdin : BufferType::Default, stdout : BufferType::Default, stderr : BufferType::Default};
	let mut command_idx = -1;
	for i in range_inclusive(1, args.len()) {
		match parse_options(args.slice(1, i), &mut options, &optgrps) {
			Ok(OkMsg::Buffering) => {
				command_idx = i-1;
				break;
			},
			Ok(OkMsg::Help) => {
				print_usage(&optgrps);
				return;
			},
			Ok(OkMsg::Version) => {
				print_version();
				return;
			},
			Err(ErrMsg::Fatal) => break,
			Err(ErrMsg::Retry) => continue,
		}
	};
	if command_idx == -1 {
		println!("Invalid options\nTry 'stdbuf --help' for more information.");
		return 125;
	}
	println!("{}", options);

	let ref command_name = args[command_idx];
	let mut process = match Command::new(command_name).args(args.slice_from(command_idx+1)).spawn() {
		Ok(p) => p,
		Err(e) => panic!("failed to execute process: {}", e),
	};

	let output = process.stdout.as_mut().unwrap().read_to_string().ok().expect("failed to read output");
	println!("{}", output);
	0
}

<commit_msg>basic version 1<commit_after>#![crate_name = "stdbuf"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Dorota Kapturkiewicz <dokaptur@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;
extern crate libc;
use getopts::{optopt, optflag, getopts, usage, Matches, OptGroup};
use std::io::process::{Command, StdioContainer};
use std::iter::range_inclusive;
use std::num::Int;

static NAME: &'static str = "stdbuf";
static VERSION: &'static str = "1.0.0";
static LIBSTDBUF_PATH: &'static str = "/lib/coreutils/libstdbuf.so";

#[deriving(Show)]
enum BufferType {
	Default,
	Line,
	Size(u64)
}

#[deriving(Show)]
struct ProgramOptions {
	stdin : BufferType,
	stdout : BufferType,
	stderr : BufferType,
}

enum ErrMsg {
	Retry,
	Fatal
}

enum OkMsg {
	Buffering,
	Help,
	Version
}

fn print_version() {
	println!("{} version {}", NAME, VERSION);
}

fn print_usage(opts: &[OptGroup]) {
	let brief =
		"Usage: stdbuf OPTION... COMMAND\n \
		Run COMMAND, with modified buffering operations for its standard streams\n \
		Mandatory arguments to long options are mandatory for short options too.";
	let explanation =
		"If MODE is 'L' the corresponding stream will be line buffered.\n \
		This option is invalid with standard input.\n\n \
		If MODE is '0' the corresponding stream will be unbuffered.\n\n \
		Otherwise MODE is a number which may be followed by one of the following:\n\n \
		KB 1000, K 1024, MB 1000*1000, M 1024*1024, and so on for G, T, P, E, Z, Y.\n \
		In this case the corresponding stream will be fully buffered with the buffer size set to MODE bytes.\n\n \
		NOTE: If COMMAND adjusts the buffering of its standard streams ('tee' does for e.g.) then that will override \
		corresponding settings changed by 'stdbuf'.\n \
		Also some filters (like 'dd' and 'cat' etc.) don't use streams for I/O, \
		and are thus unaffected by 'stdbuf' settings.\n";
	println!("{}\n{}", getopts::usage(brief, opts), explanation);
}

fn parse_size(size : &str) -> Option<u64> {
	let ext = size.trim_left_chars(|c: char| c.is_digit(10));
	let num = size.trim_right_chars(|c: char| c.is_alphabetic());
	let mut recovered = num.to_string();
	recovered.push_str(ext);
	if recovered.as_slice() != size {
		return None;
	}
	let buf_size : u64 = match from_str(num) {
		Some(m) => m,
		None => return None,
	};
	let (power, base) : (uint, u64) = match ext {
		"" => (0, 0),
		"KB" => (1, 1024),
		"K" => (1, 1000),
		"MB" => (2, 1024),
		"M" => (2, 1000),
		"GB" => (3, 1024),
		"G" => (3, 1000),
		"TB" => (4, 1024),
		"T" => (4, 1000),
		"PB" => (5, 1024),
		"P" => (5, 1000),
		"EB" => (6, 1024),
		"E" => (6, 1000),
		"ZB" => (7, 1024),
		"Z" => (7, 1000),
		"YB" => (8, 1024),
		"Y" => (8, 1000),
		_ => return None,
	};
	Some(buf_size * base.pow(power))
}

fn check_option(matches : &Matches, name : &str, modified : &mut bool) -> Option<BufferType> {
	match matches.opt_str(name) {
		Some(value) => {
			*modified = true;
			match value.as_slice() {
				"L" => {
					if name == "input" {
						println!("stdbuf: line buffering stdin is meaningless");
						None
					} else {
						Some(BufferType::Line)
					}
				},
				x => {
					let size = match parse_size(x) {
						Some(m) => m,
						None => { println!("Invalid mode {}", x); return None }
					};
					Some(BufferType::Size(size))
				},
			}
		},
		None => Some(BufferType::Default),
	}
}

fn parse_options(args : &[String], options : &mut ProgramOptions, optgrps : &[OptGroup]) -> Result<OkMsg, ErrMsg> {
	let matches = match getopts(args, optgrps) {
		Ok(m) => m,
		Err(_) => return Err(ErrMsg::Retry)
	};
	if matches.opt_present("help") {
		return Ok(OkMsg::Help);
	}
	if matches.opt_present("version") {
		return Ok(OkMsg::Version);
	}
	let mut modified = false;
	options.stdin = try!(check_option(&matches, "input", &mut modified).ok_or(ErrMsg::Fatal));
	options.stdout = try!(check_option(&matches, "output", &mut modified).ok_or(ErrMsg::Fatal));
	options.stderr = try!(check_option(&matches, "error", &mut modified).ok_or(ErrMsg::Fatal));

	if matches.free.len() != 1 {
		return Err(ErrMsg::Retry);
	}
	if !modified {
		println!("stdbuf: you must specify a buffering mode option");
		return Err(ErrMsg::Fatal);
	}
	Ok(OkMsg::Buffering)
}

fn set_command_env(command : &mut Command, buffer_name : &str, buffer_type : BufferType) {
	match buffer_type {
		BufferType::Size(m) => { command.env(buffer_name, m.to_string()); },
		BufferType::Line => { command.env(buffer_name, "L"); },
		BufferType::Default => {},
	}
}

pub fn uumain(args: Vec<String>) -> int {
	let optgrps = [
		optopt("i", "input", "adjust standard input stream buffering", "MODE"),
		optopt("o", "output", "adjust standard output stream buffering", "MODE"),
		optopt("e", "error", "adjust standard error stream buffering", "MODE"),
		optflag("", "help", "display this help and exit"),
		optflag("", "version", "output version information and exit"),
	];
	let mut options = ProgramOptions{ stdin : BufferType::Default, stdout : BufferType::Default, stderr : BufferType::Default};
	let mut command_idx = -1;
	for i in range_inclusive(1, args.len()) {
		match parse_options(args.slice(1, i), &mut options, &optgrps) {
			Ok(OkMsg::Buffering) => {
				command_idx = i-1;
				break;
			},
			Ok(OkMsg::Help) => {
				print_usage(&optgrps);
				return 0;
			},
			Ok(OkMsg::Version) => {
				print_version();
				return 0;
			},
			Err(ErrMsg::Fatal) => break,
			Err(ErrMsg::Retry) => continue,
		}
	};
	if command_idx == -1 {
		println!("Invalid options\nTry 'stdbuf --help' for more information.");
		return 125;
	}

	let ref command_name = args[command_idx];
	let mut command = Command::new(command_name);
	command.args(args.slice_from(command_idx+1)).env("LD_PRELOAD", LIBSTDBUF_PATH);
	command.stdin(StdioContainer::InheritFd(0)).stdout(StdioContainer::InheritFd(1)).stderr(StdioContainer::InheritFd(2));
	set_command_env(&mut command, "_STDBUF_I", options.stdin);
	set_command_env(&mut command, "_STDBUF_O", options.stdout);
	set_command_env(&mut command, "_STDBUF_E", options.stderr);
	match command.spawn() {
		Ok(_) => {},
		Err(e) => panic!("failed to execute process: {}", e),
	};
	0
}

<|endoftext|>


<commit_before>//! Define and construct substrates.
//!
//! # Examples
//! Construct a triclinic lattice of hydrogen molecules.
//!
//! ```
//! # use grafen::substrate::{create_substrate, LatticeType, SheetConf};
//! # use grafen::system::{Atom, Coord, ResidueBase};
//! // Define the molecule as a Residue.
//! let residue_base = ResidueBase {
//!     code: "HMOL".to_string(),
//!     atoms: vec![
//!         Atom { code: "H1".to_string(), position: Coord::new(0.0, 0.0, 0.5), },
//!         Atom { code: "H2".to_string(), position: Coord::new(0.0, 0.0, -0.5), }
//!     ],
//! };
//!
//! // Define the substrate
//! let conf = SheetConf {
//!     lattice: LatticeType::Triclinic { a: 1.0, b: 0.5, gamma: 60.0 },
//!     residue: residue_base,
//!     size: (5.0, 10.0),
//!     std_z: None,
//! };
//!
//! // ... and create it!
//! let substrate = create_substrate(&conf).unwrap();
//! ```

mod distribution;
mod lattice;
mod points;

use error::{GrafenError, Result};
use substrate::distribution::PoissonDistribution;
use substrate::lattice::Lattice;
use system::*;

#[derive(Clone, Debug, PartialEq)]
/// Configuration for constructing a substrate.
pub struct SheetConf {
    /// The type of lattice which will be generated.
    pub lattice: LatticeType,
    /// Base residue to generate coordinates for.
    pub residue: ResidueBase,
    /// Desired size of substrate along x and y.
    pub size: (f64, f64),
    /// Optionally use a random uniform distribution with this
    /// deviation to shift residue positions along z. The
    /// positions are shifted with the range (-std_z, +std_z)
    /// where std_z is the input devation.
    pub std_z: Option<f64>,
}

#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
/// Lattice types which a substrate can be constructed from.
pub enum LatticeType {
    /// A hexagonal (honey comb) lattice with bond spacing `a`.
    Hexagonal { a: f64 },
    /// A triclinic lattice with base vectors of length `a` and `b`.
    /// Vector `a` is directed along the x axis and vector `b` is separated
    /// to it by the input angle `gamma` in degrees.
    Triclinic { a: f64, b: f64, gamma: f64 },
    /// A Poisson disc distribution of points with an input `density` in number
    /// of points per unit area. It is implemented using Bridson's algorithm
    /// which ensures that no points are within sqrt(2 / (pi * density)) of
    /// each other. This creates a good match to the input density.
    ///
    /// *Fast Poisson disk sampling in arbitrary dimensions*,
    ///  R. Bridson, ACM SIGGRAPH 2007 Sketches Program,
    ///  http://www.cs.ubc.ca/~rbridson/docs/bridson-siggraph07-poissondisk.pdf
    PoissonDisc { density: f64 },
}

#[derive(Clone, Debug)]
/// A `Sheet` of `Residue`s in some points.
pub struct Sheet {
    /// Sheet origin position. Residue coordinates are relative to this.
    pub origin: Coord,
    /// Sheet size.
    pub size: Coord,
    /// Residue base.
    pub residue_base: ResidueBase,
    /// `Residue`s belonging to the sheet.
    pub residue_coords: Vec<Coord>,
}

impl IntoComponent for Sheet {
    fn to_component(&self) -> Component {
        Component {
            origin: self.origin,
            box_size: self.size,
            residue_base: self.residue_base.clone(),
            residue_coords: self.residue_coords.clone(),
        }
    }

    fn into_component(self) -> Component {
        Component {
            origin: self.origin,
            box_size: self.size,
            residue_base: self.residue_base,
            residue_coords: self.residue_coords,
        }
    }

    fn num_atoms(&self) -> usize {
        self.residue_base.atoms.len() * self.residue_coords.len()
    }
}

impl Translate for Sheet {
    fn translate(mut self, trans: &Coord) -> Sheet {
        self.origin = self.origin + *trans;
        self
    }
}

/// Create a substrate of input configuration and return as a `Component`.
///
/// The returned component's size will be adjusted to a multiple of the
/// substrate spacing along both directions. Thus the substrate can be
/// periodically replicated along x and y.
///
/// # Errors
/// Returns an Error if the either of the input size are non-positive.
pub fn create_substrate(conf: &SheetConf) -> Result<Sheet> {
    let (dx, dy) = conf.size;
    if dx <= 0.0 || dy <= 0.0 {
        return Err(
            GrafenError::RunError("cannot create a substrate of negative size".to_string())
        );
    }

    let mut points = match conf.lattice {
        LatticeType::Hexagonal { a } => {
            Lattice::hexagonal(a).with_size(dx, dy).finalize()
        },
        LatticeType::Triclinic { a, b, gamma } => {
            Lattice::triclinic(a, b, gamma.to_radians()).with_size(dx, dy).finalize()
        },
        LatticeType::PoissonDisc { density } => {
            // The factor 1/sqrt(pi) comes from the area and the factor sqrt(2)
            // is a magic number which roughly gives the correct density. It works!
            use std::f64::consts::PI;
            let rmin = (2.0 / (PI * density)).sqrt();
            PoissonDistribution::new(rmin, dx, dy)
        },
    };

    if let Some(std) = conf.std_z {
        points = points.uniform_distribution(std);
    };

    Ok(Sheet {
        origin: Coord::new(0.0, 0.0, 0.0),
        size: points.box_size,
        residue_base: conf.residue.clone(),
        residue_coords: points.coords,
    })
}

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

    fn setup_conf() -> SheetConf {
        SheetConf {
            lattice: LatticeType::Hexagonal { a: 1.0 },
            residue: ResidueBase {
                code: "GRPH".to_string(),
                atoms: vec![Atom { code: "C".to_string(), position: Coord::new(0.0, 0.0, 0.0) }],
            },
            size: (10.0, 10.0),
            std_z: None,
        }
    }

    #[test]
    fn negative_sizes_return_error() {
        let mut conf = setup_conf();
        assert!(create_substrate(&conf).is_ok());

        conf.size = (-1.0, 1.0);
        assert!(create_substrate(&conf).is_err());

        conf.size = (1.0, -1.0);
        assert!(create_substrate(&conf).is_err());
    }

    #[test]
    fn uniform_distribution_is_set() {
        // The graphene is ordinarily positioned at z = 0.0
        let mut conf = setup_conf();
        {
            let regular = create_substrate(&conf).unwrap();
            assert!(regular.residue_coords.iter().any(|r| r.z != 0.0) == false);
        }

        let std_z = 1.0;
        conf.std_z = Some(std_z);
        let uniform = create_substrate(&conf).unwrap();

        // Non-zero variance: This *can* fail, but it should not be common!
        // How else to assert that a distribution has been applied, though?
        assert!(uniform.residue_coords.iter().map(|r| r.z).any(|z| z != 0.0));

        // No positions should exceed the input distribution max
        assert!(uniform.residue_coords.iter().all(|r| r.z.abs() <= std_z));
    }

    #[test]
    fn sheet_into_component() {
        let origin = Coord::new(0.0, 0.0, 0.0);
        let size = Coord::new(1.0, 2.0, 3.0);
        let residue_base = ResidueBase {
            code: "RES".to_string(),
            atoms: vec![
                Atom { code: "A1".to_string(), position: Coord::new(0.0, 0.0, 0.0) },
            ],
        };
        let residue_coords = vec![Coord::new(0.0, 0.0, 0.0), Coord::new(1.0, 0.0, 0.0)];

        let sheet = Sheet { origin, size, residue_base, residue_coords: residue_coords.clone() };
        let component = sheet.into_component();

        assert_eq!(origin, component.origin);
        assert_eq!(size, component.box_size);
        assert_eq!(residue_coords, component.residue_coords);
    }

    #[test]
    fn translate_a_sheet() {
        let residue_base = ResidueBase {
            code: "GRPH".to_string(),
            atoms: vec![Atom {
                code: "C".to_string(),
                position: Coord::new(0.0, 0.0, 0.0) }],
        };
        let origin = Coord::new(-1.0, 0.0, 1.0);
        let size = Coord::new(1.0, 2.0, 1.0);
        let residue_coords = vec![Coord::new(0.0, 0.0, 0.0)];

        let translate = Coord::new(1.0, 2.0, 3.0);

        let sheet = Sheet { origin, size, residue_base, residue_coords }
            .translate(&Coord::new(1.0, 2.0, 3.0));

        assert_eq!(size, sheet.size);
        assert_eq!(origin + translate, sheet.origin);
        assert_eq!(Coord::new(0.0, 0.0, 0.0), sheet.residue_coords[0]);
    }
}
<commit_msg>Add `into_circle` method to `Sheet` which cuts a circle with input radius<commit_after>//! Define and construct substrates.
//!
//! # Examples
//! Construct a triclinic lattice of hydrogen molecules.
//!
//! ```
//! # use grafen::substrate::{create_substrate, LatticeType, SheetConf};
//! # use grafen::system::{Atom, Coord, ResidueBase};
//! // Define the molecule as a Residue.
//! let residue_base = ResidueBase {
//!     code: "HMOL".to_string(),
//!     atoms: vec![
//!         Atom { code: "H1".to_string(), position: Coord::new(0.0, 0.0, 0.5), },
//!         Atom { code: "H2".to_string(), position: Coord::new(0.0, 0.0, -0.5), }
//!     ],
//! };
//!
//! // Define the substrate
//! let conf = SheetConf {
//!     lattice: LatticeType::Triclinic { a: 1.0, b: 0.5, gamma: 60.0 },
//!     residue: residue_base,
//!     size: (5.0, 10.0),
//!     std_z: None,
//! };
//!
//! // ... and create it!
//! let substrate = create_substrate(&conf).unwrap();
//! ```

mod distribution;
mod lattice;
mod points;

use error::{GrafenError, Result};
use substrate::distribution::PoissonDistribution;
use substrate::lattice::Lattice;
use system::*;

#[derive(Clone, Debug, PartialEq)]
/// Configuration for constructing a substrate.
pub struct SheetConf {
    /// The type of lattice which will be generated.
    pub lattice: LatticeType,
    /// Base residue to generate coordinates for.
    pub residue: ResidueBase,
    /// Desired size of substrate along x and y.
    pub size: (f64, f64),
    /// Optionally use a random uniform distribution with this
    /// deviation to shift residue positions along z. The
    /// positions are shifted with the range (-std_z, +std_z)
    /// where std_z is the input devation.
    pub std_z: Option<f64>,
}

#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
/// Lattice types which a substrate can be constructed from.
pub enum LatticeType {
    /// A hexagonal (honey comb) lattice with bond spacing `a`.
    Hexagonal { a: f64 },
    /// A triclinic lattice with base vectors of length `a` and `b`.
    /// Vector `a` is directed along the x axis and vector `b` is separated
    /// to it by the input angle `gamma` in degrees.
    Triclinic { a: f64, b: f64, gamma: f64 },
    /// A Poisson disc distribution of points with an input `density` in number
    /// of points per unit area. It is implemented using Bridson's algorithm
    /// which ensures that no points are within sqrt(2 / (pi * density)) of
    /// each other. This creates a good match to the input density.
    ///
    /// *Fast Poisson disk sampling in arbitrary dimensions*,
    ///  R. Bridson, ACM SIGGRAPH 2007 Sketches Program,
    ///  http://www.cs.ubc.ca/~rbridson/docs/bridson-siggraph07-poissondisk.pdf
    PoissonDisc { density: f64 },
}

#[derive(Clone, Debug)]
/// A `Sheet` of `Residue`s in some points.
pub struct Sheet {
    /// Sheet origin position. Residue coordinates are relative to this.
    pub origin: Coord,
    /// Sheet size.
    pub size: Coord,
    /// Residue base.
    pub residue_base: ResidueBase,
    /// `Residue`s belonging to the sheet.
    pub residue_coords: Vec<Coord>,
}

impl Sheet {
    /// Cut the `Sheet` into a circle of input radius.
    ///
    /// This method assumes that the sheet is aligned in the xy plane and that
    /// coordinates are set at z = 0.0. Furthermore, it cuts the circle around
    /// the center point (radius, radius, 0).
    ///
    /// The final circle will be centered around (0, 0, 0).
    fn into_circle(self, radius: f64) -> Self {
        const ORIGO: Coord = Coord { x: 0.0, y: 0.0, z: 0.0 };
        let center = Coord::new(radius, radius, 0.0);

        let residue_coords = self.residue_coords
            .iter()
            .inspect(|coord| println!("{:?}", coord))
            .map(|&coord| coord.with_pbc(self.size) - center)
            .filter(|coord| coord.distance(ORIGO) <= radius)
            .collect();

        Sheet {
            origin: Coord::new(0.0, 0.0, 0.0),
            size: Coord::new(2.0 * radius, 2.0 * radius, self.size.z),
            residue_base: self.residue_base,
            residue_coords,
        }
    }
}

impl IntoComponent for Sheet {
    fn to_component(&self) -> Component {
        Component {
            origin: self.origin,
            box_size: self.size,
            residue_base: self.residue_base.clone(),
            residue_coords: self.residue_coords.clone(),
        }
    }

    fn into_component(self) -> Component {
        Component {
            origin: self.origin,
            box_size: self.size,
            residue_base: self.residue_base,
            residue_coords: self.residue_coords,
        }
    }

    fn num_atoms(&self) -> usize {
        self.residue_base.atoms.len() * self.residue_coords.len()
    }
}

impl Translate for Sheet {
    fn translate(mut self, trans: &Coord) -> Sheet {
        self.origin = self.origin + *trans;
        self
    }
}

/// Create a substrate of input configuration and return as a `Component`.
///
/// The returned component's size will be adjusted to a multiple of the
/// substrate spacing along both directions. Thus the substrate can be
/// periodically replicated along x and y.
///
/// # Errors
/// Returns an Error if the either of the input size are non-positive.
pub fn create_substrate(conf: &SheetConf) -> Result<Sheet> {
    let (dx, dy) = conf.size;
    if dx <= 0.0 || dy <= 0.0 {
        return Err(
            GrafenError::RunError("cannot create a substrate of negative size".to_string())
        );
    }

    let mut points = match conf.lattice {
        LatticeType::Hexagonal { a } => {
            Lattice::hexagonal(a).with_size(dx, dy).finalize()
        },
        LatticeType::Triclinic { a, b, gamma } => {
            Lattice::triclinic(a, b, gamma.to_radians()).with_size(dx, dy).finalize()
        },
        LatticeType::PoissonDisc { density } => {
            // The factor 1/sqrt(pi) comes from the area and the factor sqrt(2)
            // is a magic number which roughly gives the correct density. It works!
            use std::f64::consts::PI;
            let rmin = (2.0 / (PI * density)).sqrt();
            PoissonDistribution::new(rmin, dx, dy)
        },
    };

    if let Some(std) = conf.std_z {
        points = points.uniform_distribution(std);
    };

    Ok(Sheet {
        origin: Coord::new(0.0, 0.0, 0.0),
        size: points.box_size,
        residue_base: conf.residue.clone(),
        residue_coords: points.coords,
    })
}

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

    fn setup_conf() -> SheetConf {
        SheetConf {
            lattice: LatticeType::Hexagonal { a: 1.0 },
            residue: ResidueBase {
                code: "GRPH".to_string(),
                atoms: vec![Atom { code: "C".to_string(), position: Coord::new(0.0, 0.0, 0.0) }],
            },
            size: (10.0, 10.0),
            std_z: None,
        }
    }

    #[test]
    fn negative_sizes_return_error() {
        let mut conf = setup_conf();
        assert!(create_substrate(&conf).is_ok());

        conf.size = (-1.0, 1.0);
        assert!(create_substrate(&conf).is_err());

        conf.size = (1.0, -1.0);
        assert!(create_substrate(&conf).is_err());
    }

    #[test]
    fn uniform_distribution_is_set() {
        // The graphene is ordinarily positioned at z = 0.0
        let mut conf = setup_conf();
        {
            let regular = create_substrate(&conf).unwrap();
            assert!(regular.residue_coords.iter().any(|r| r.z != 0.0) == false);
        }

        let std_z = 1.0;
        conf.std_z = Some(std_z);
        let uniform = create_substrate(&conf).unwrap();

        // Non-zero variance: This *can* fail, but it should not be common!
        // How else to assert that a distribution has been applied, though?
        assert!(uniform.residue_coords.iter().map(|r| r.z).any(|z| z != 0.0));

        // No positions should exceed the input distribution max
        assert!(uniform.residue_coords.iter().all(|r| r.z.abs() <= std_z));
    }

    #[test]
    fn sheet_into_component() {
        let origin = Coord::new(0.0, 0.0, 0.0);
        let size = Coord::new(1.0, 2.0, 3.0);
        let residue_base = ResidueBase {
            code: "RES".to_string(),
            atoms: vec![
                Atom { code: "A1".to_string(), position: Coord::new(0.0, 0.0, 0.0) },
            ],
        };
        let residue_coords = vec![Coord::new(0.0, 0.0, 0.0), Coord::new(1.0, 0.0, 0.0)];

        let sheet = Sheet { origin, size, residue_base, residue_coords: residue_coords.clone() };
        let component = sheet.into_component();

        assert_eq!(origin, component.origin);
        assert_eq!(size, component.box_size);
        assert_eq!(residue_coords, component.residue_coords);
    }

    #[test]
    fn translate_a_sheet() {
        let residue_base = ResidueBase {
            code: "GRPH".to_string(),
            atoms: vec![Atom {
                code: "C".to_string(),
                position: Coord::new(0.0, 0.0, 0.0) }],
        };
        let origin = Coord::new(-1.0, 0.0, 1.0);
        let size = Coord::new(1.0, 2.0, 1.0);
        let residue_coords = vec![Coord::new(0.0, 0.0, 0.0)];

        let translate = Coord::new(1.0, 2.0, 3.0);

        let sheet = Sheet { origin, size, residue_base, residue_coords }
            .translate(&Coord::new(1.0, 2.0, 3.0));

        assert_eq!(size, sheet.size);
        assert_eq!(origin + translate, sheet.origin);
        assert_eq!(Coord::new(0.0, 0.0, 0.0), sheet.residue_coords[0]);
    }

    #[test]
    fn cut_a_circle_from_a_sheet() {
        let conf = setup_conf();
        let sheet = create_substrate(&conf).unwrap();

        let radius = 2.5;
        let origin = Coord::new(0.0, 0.0, 0.0);

        let circle = sheet.into_circle(radius);
        assert_eq!(origin, circle.origin);
        assert!(circle.num_atoms() > 0);

        for coord in circle.residue_coords {
            assert!(coord.distance(origin) <= radius);
        }
    }

    #[test]
    fn circles_center_coords_around_origo() {
        let origin = Coord::new(0.0, 0.0, 0.0);
        let size = Coord::new(4.0, 4.0, 1.0);
        let residue_base = ResidueBase {
            code: "RES".to_string(),
            atoms: vec![
                Atom { code: "A1".to_string(), position: Coord::new(0.0, 0.0, 0.0) },
            ],
        };

        // The coordinates are offset from (0, 0, 0) but will be centered around it
        let radius = 2.0;
        let residue_coords = vec![
            Coord::new(0.1, radius, 0.0), // After: (-1.9, radius, 0.0)
            Coord::new(3.9, radius, 0.0)  // After: (1.9, radius, 0.0)
        ];

        let sheet = Sheet { origin, size, residue_base, residue_coords };
        let circle = sheet.into_circle(radius);

        assert_eq!(2, circle.num_atoms());
        assert_eq!(Coord::new(-1.9, 0.0, 0.0), circle.residue_coords[0]);
        assert_eq!(Coord::new(1.9, 0.0, 0.0), circle.residue_coords[1]);
    }

    #[test]
    fn circles_check_coordinates_after_pbc_adjustment() {
        let origin = Coord::new(0.0, 0.0, 0.0);
        let size = Coord::new(1.0, 1.0, 1.0);
        let residue_base = ResidueBase {
            code: "RES".to_string(),
            atoms: vec![
                Atom { code: "A1".to_string(), position: Coord::new(0.0, 0.0, 0.0) },
            ],
        };

        // The coordinates are outside the box, but will be adjusted to be in the center.
        let radius = 0.5;
        let residue_coords = vec![
            Coord::new(-0.6, radius, 0.0), // In box: (0.4, radius, 0.0)
            Coord::new(2.6, radius, 0.0)   // In box: (0.6, radius, 0.0)
        ];

        let sheet = Sheet { origin, size, residue_base, residue_coords };
        let circle = sheet.into_circle(radius);

        assert_eq!(2, circle.num_atoms());
        assert_eq!(Coord::new(-0.1, 0.0, 0.0), circle.residue_coords[0]);
        assert_eq!(Coord::new(0.1, 0.0, 0.0), circle.residue_coords[1]);
    }
}
<|endoftext|>


<commit_before>use std::sync::atomic::Ordering::{Acquire, Release, Relaxed};
use std::cell::RefCell;

use mem::epoch::{self, AtomicPtr, Owned};

pub struct MsQueue<T> {
    head: AtomicPtr<Node<T>>,
    tail: AtomicPtr<Node<T>>,
}

struct Node<T> {
    data: RefCell<Option<T>>,
    next: AtomicPtr<Node<T>>,
}

impl<T> MsQueue<T> {
    pub fn new() -> MsQueue<T> {
        let q = MsQueue { head: AtomicPtr::new(), tail: AtomicPtr::new() };
        let sentinel = Owned::new(Node {
            data: RefCell::new(None),
            next: AtomicPtr::new()
        });
        let guard = epoch::pin();
        let sentinel = q.head.store_and_ref(sentinel, Relaxed, &guard);
        unsafe { q.tail.store_shared(Some(sentinel), Relaxed); }
        q
    }

    pub fn push(&self, t: T) {
        let mut n = Owned::new(Node {
            data: RefCell::new(Some(t)),
            next: AtomicPtr::new()
        });
        let guard = epoch::pin();
        loop {
            let tail = self.tail.load(Acquire, &guard).unwrap();
            if let Some(next) = tail.next.load(Relaxed, &guard) {
                unsafe { self.tail.cas_shared(Some(tail), Some(next), Relaxed); }
                continue;
            }

            match tail.next.cas_and_ref(None, n, Release, &guard) {
                Ok(shared) => {
                    unsafe { self.tail.cas_shared(Some(tail), Some(shared), Relaxed); }
                    break;
                }
                Err(owned) => {
                    n = owned;
                }
            }
        }
    }

    pub fn pop(&self) -> Option<T> {
        let guard = epoch::pin();
        loop {
            let head = self.head.load(Acquire, &guard).unwrap();

            if let Some(next) = head.next.load(Relaxed, &guard) {
                if unsafe { self.head.cas_shared(Some(head), Some(next), Relaxed) } {
                    unsafe { guard.unlinked(head); }
                    // FIXME: rustc bug that (*next) is required
                    return (*next).data.borrow_mut().take()
                }
            } else {
                return None
            }
        }
    }
}

/*
impl<T: Debug> MsQueue<T> {
    pub fn debug(&self) {
        writeln!(stderr(), "Debugging queue:");

        let guard = epoch::pin();
        let mut node = self.head.load(Acquire, &guard);
        while let Some(n) = node {
            writeln!(stderr(), "{:?}", (*n).data.borrow());
            node = n.next.load(Relaxed, &guard);
        }

        writeln!(stderr(), "");
    }
}
*/

#[cfg(test)]
mod test {
    const CONC_COUNT: i64 = 1000000;

    use std::io::stderr;
    use std::io::prelude::*;

    use mem::epoch;
    use thread;
    use super::*;

    #[test]
    fn smoke_queue() {
        let q: MsQueue<i64> = MsQueue::new();
    }

    #[test]
    fn push_pop_1() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(37);
        assert_eq!(q.pop(), Some(37));
    }

    #[test]
    fn push_pop_2() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(37);
        q.push(48);
        assert_eq!(q.pop(), Some(37));
        assert_eq!(q.pop(), Some(48));
    }

    #[test]
    fn push_pop_many_seq() {
        let q: MsQueue<i64> = MsQueue::new();
        for i in 0..200 {
            q.push(i)
        }
        for i in 0..200 {
            assert_eq!(q.pop(), Some(i));
        }
    }

    #[test]
    fn push_pop_many_spsc() {
        let q: MsQueue<i64> = MsQueue::new();

        thread::scope(|scope| {
            scope.spawn(|| {
                let mut next = 0;

                while next < CONC_COUNT {
                    if let Some(elem) = q.pop() {
                        assert_eq!(elem, next);
                        next += 1;
                    }
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        });
    }

    #[test]
    fn push_pop_many_spmc() {
        use std::time::Duration;

        fn recv(t: i32, q: &MsQueue<i64>) {
            let mut cur = -1;
            for i in 0..CONC_COUNT {
                if let Some(elem) = q.pop() {
                    if elem <= cur {
                        writeln!(stderr(), "{}: {} <= {}", t, elem, cur);
                    }
                    assert!(elem > cur);
                    cur = elem;

                    if cur == CONC_COUNT - 1 { break }
                }

                if i % 10000 == 0 {
                    //writeln!(stderr(), "{}: {} @ {}", t, i, cur);
                }
            }
        }

        let q: MsQueue<i64> = MsQueue::new();
        let qr = &q;
        thread::scope(|scope| {
            for i in 0..3 {
                scope.spawn(move || recv(i, qr));
            }

            scope.spawn(|| {
                for i in 0..CONC_COUNT {
                    q.push(i);

                    if i % 10000 == 0 {
                        //writeln!(stderr(), "Push: {}", i);
                    }
                }
            })
        });
    }

    #[test]
    fn push_pop_many_mpmc() {
        enum LR { Left(i64), Right(i64) }

        let q: MsQueue<LR> = MsQueue::new();

        thread::scope(|scope| {
            for _t in 0..2 {
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Left(i))
                    }
                });
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Right(i))
                    }
                });
                scope.spawn(|| {
                    let mut vl = vec![];
                    let mut vr = vec![];
                    for _i in 0..CONC_COUNT {
                        match q.pop() {
                            Some(LR::Left(x)) => vl.push(x),
                            Some(LR::Right(x)) => vr.push(x),
                            _ => {}
                        }
                    }

                    let mut vl2 = vl.clone();
                    let mut vr2 = vr.clone();
                    vl2.sort();
                    vr2.sort();

                    assert_eq!(vl, vl2);
                    assert_eq!(vr, vr2);
                });
            }
        });
    }
}
<commit_msg>Fix regression tests<commit_after>use std::sync::atomic::Ordering::{Acquire, Release, Relaxed};
use std::cell::RefCell;

use mem::epoch::{self, AtomicPtr, Owned};

pub struct MsQueue<T> {
    head: AtomicPtr<Node<T>>,
    tail: AtomicPtr<Node<T>>,
}

struct Node<T> {
    data: RefCell<Option<T>>,
    next: AtomicPtr<Node<T>>,
}

impl<T> MsQueue<T> {
    pub fn new() -> MsQueue<T> {
        let q = MsQueue { head: AtomicPtr::new(), tail: AtomicPtr::new() };
        let sentinel = Owned::new(Node {
            data: RefCell::new(None),
            next: AtomicPtr::new()
        });
        let guard = epoch::pin();
        let sentinel = q.head.store_and_ref(sentinel, Relaxed, &guard);
        unsafe { q.tail.store_shared(Some(sentinel), Relaxed); }
        q
    }

    pub fn push(&self, t: T) {
        let mut n = Owned::new(Node {
            data: RefCell::new(Some(t)),
            next: AtomicPtr::new()
        });
        let guard = epoch::pin();
        loop {
            let tail = self.tail.load(Acquire, &guard).unwrap();
            if let Some(next) = tail.next.load(Relaxed, &guard) {
                unsafe { self.tail.cas_shared(Some(tail), Some(next), Relaxed); }
                continue;
            }

            match tail.next.cas_and_ref(None, n, Release, &guard) {
                Ok(shared) => {
                    unsafe { self.tail.cas_shared(Some(tail), Some(shared), Relaxed); }
                    break;
                }
                Err(owned) => {
                    n = owned;
                }
            }
        }
    }

    pub fn pop(&self) -> Option<T> {
        let guard = epoch::pin();
        loop {
            let head = self.head.load(Acquire, &guard).unwrap();

            if let Some(next) = head.next.load(Relaxed, &guard) {
                if unsafe { self.head.cas_shared(Some(head), Some(next), Relaxed) } {
                    unsafe { guard.unlinked(head); }
                    // FIXME: rustc bug that (*next) is required
                    return (*next).data.borrow_mut().take()
                }
            } else {
                return None
            }
        }
    }
}

/*
impl<T: Debug> MsQueue<T> {
    pub fn debug(&self) {
        writeln!(stderr(), "Debugging queue:");

        let guard = epoch::pin();
        let mut node = self.head.load(Acquire, &guard);
        while let Some(n) = node {
            writeln!(stderr(), "{:?}", (*n).data.borrow());
            node = n.next.load(Relaxed, &guard);
        }

        writeln!(stderr(), "");
    }
}
*/

#[cfg(test)]
mod test {
    const CONC_COUNT: i64 = 1000000;

    use std::io::stderr;
    use std::io::prelude::*;

    use mem::epoch;
    use scope;
    use super::*;

    #[test]
    fn smoke_queue() {
        let q: MsQueue<i64> = MsQueue::new();
    }

    #[test]
    fn push_pop_1() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(37);
        assert_eq!(q.pop(), Some(37));
    }

    #[test]
    fn push_pop_2() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(37);
        q.push(48);
        assert_eq!(q.pop(), Some(37));
        assert_eq!(q.pop(), Some(48));
    }

    #[test]
    fn push_pop_many_seq() {
        let q: MsQueue<i64> = MsQueue::new();
        for i in 0..200 {
            q.push(i)
        }
        for i in 0..200 {
            assert_eq!(q.pop(), Some(i));
        }
    }

    #[test]
    fn push_pop_many_spsc() {
        let q: MsQueue<i64> = MsQueue::new();

        scope(|scope| {
            scope.spawn(|| {
                let mut next = 0;

                while next < CONC_COUNT {
                    if let Some(elem) = q.pop() {
                        assert_eq!(elem, next);
                        next += 1;
                    }
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        });
    }

    #[test]
    fn push_pop_many_spmc() {
        use std::time::Duration;

        fn recv(t: i32, q: &MsQueue<i64>) {
            let mut cur = -1;
            for i in 0..CONC_COUNT {
                if let Some(elem) = q.pop() {
                    if elem <= cur {
                        writeln!(stderr(), "{}: {} <= {}", t, elem, cur);
                    }
                    assert!(elem > cur);
                    cur = elem;

                    if cur == CONC_COUNT - 1 { break }
                }

                if i % 10000 == 0 {
                    //writeln!(stderr(), "{}: {} @ {}", t, i, cur);
                }
            }
        }

        let q: MsQueue<i64> = MsQueue::new();
        let qr = &q;
        scope(|scope| {
            for i in 0..3 {
                scope.spawn(move || recv(i, qr));
            }

            scope.spawn(|| {
                for i in 0..CONC_COUNT {
                    q.push(i);

                    if i % 10000 == 0 {
                        //writeln!(stderr(), "Push: {}", i);
                    }
                }
            })
        });
    }

    #[test]
    fn push_pop_many_mpmc() {
        enum LR { Left(i64), Right(i64) }

        let q: MsQueue<LR> = MsQueue::new();

        scope(|scope| {
            for _t in 0..2 {
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Left(i))
                    }
                });
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Right(i))
                    }
                });
                scope.spawn(|| {
                    let mut vl = vec![];
                    let mut vr = vec![];
                    for _i in 0..CONC_COUNT {
                        match q.pop() {
                            Some(LR::Left(x)) => vl.push(x),
                            Some(LR::Right(x)) => vr.push(x),
                            _ => {}
                        }
                    }

                    let mut vl2 = vl.clone();
                    let mut vr2 = vr.clone();
                    vl2.sort();
                    vr2.sort();

                    assert_eq!(vl, vl2);
                    assert_eq!(vr, vr2);
                });
            }
        });
    }
}
<|endoftext|>


<commit_before>use LalafellBot;
use tasks::RunsTask;
use commands::tag::Tagger;
use discord::model::UserId;
use chrono::prelude::*;
use chrono::Duration;

use std::sync::Arc;
use std::thread;

pub struct AutoTagTask {
  pub next_sleep: i64
}

impl AutoTagTask {
  pub fn new() -> AutoTagTask {
    AutoTagTask {
      next_sleep: 30
    }
  }

  pub fn run_once(&mut self, s: Arc<LalafellBot>) {
    thread::sleep(Duration::seconds(self.next_sleep).to_std().unwrap());
    info!("Autotag task running");
    let time_to_update = {
      let database = s.database.read().unwrap();
      database.autotags.last_updated + Duration::days(1).num_seconds() < UTC::now().timestamp()
    };
    if !time_to_update {
      info!("Not yet time to update, sleeping 30 minutes");
      self.next_sleep = Duration::minutes(30).num_seconds();
      return;
    }
    info!("Time to update autotags");
    let users: Vec<(u64, u64, u64)> = {
      let database = s.database.read().unwrap();
      database.autotags.users.iter().map(|u| (u.user_id, u.server_id, u.character_id)).collect()
    };
    {
      let option_state = s.state.read().unwrap();
      let state = match option_state.as_ref() {
        Some(st) => st,
        None => {
          info!("Bot not connected. Will try again in 30 seconds.");
          self.next_sleep = 30;
          return;
        }
      };
      for (user_id, server_id, character_id) in users {
        let server = match state.servers().iter().find(|s| s.id.0 == server_id) {
          Some(ser) => ser,
          None => {
            warn!("Couldn't find server for user ID {}", user_id);
            continue;
          }
        };
        if let Err(e) = Tagger::tag(s.clone(), UserId(user_id), server, character_id, false) {
          warn!("Couldn't update tag for user ID {}: {}", user_id, e);
        }
      }
    }
    {
      let mut database = s.database.write().unwrap();
      database.autotags.last_updated = UTC::now().timestamp();
    };
    info!("Done updating autotags");
  }
}

impl RunsTask for AutoTagTask {
  fn start(mut self, s: Arc<LalafellBot>) {
    info!("Autotag task waiting {} seconds", self.next_sleep);
    loop {
      self.run_once(s.clone());
    }
  }
}
<commit_msg>feat: update tags twice a day<commit_after>use LalafellBot;
use tasks::RunsTask;
use commands::tag::Tagger;
use discord::model::UserId;
use chrono::prelude::*;
use chrono::Duration;

use std::sync::Arc;
use std::thread;

pub struct AutoTagTask {
  pub next_sleep: i64
}

impl AutoTagTask {
  pub fn new() -> AutoTagTask {
    AutoTagTask {
      next_sleep: 30
    }
  }

  pub fn run_once(&mut self, s: Arc<LalafellBot>) {
    thread::sleep(Duration::seconds(self.next_sleep).to_std().unwrap());
    info!("Autotag task running");
    let time_to_update = {
      let database = s.database.read().unwrap();
      database.autotags.last_updated + Duration::hours(12).num_seconds() < UTC::now().timestamp()
    };
    if !time_to_update {
      info!("Not yet time to update, sleeping 30 minutes");
      self.next_sleep = Duration::minutes(30).num_seconds();
      return;
    }
    info!("Time to update autotags");
    let users: Vec<(u64, u64, u64)> = {
      let database = s.database.read().unwrap();
      database.autotags.users.iter().map(|u| (u.user_id, u.server_id, u.character_id)).collect()
    };
    {
      let option_state = s.state.read().unwrap();
      let state = match option_state.as_ref() {
        Some(st) => st,
        None => {
          info!("Bot not connected. Will try again in 30 seconds.");
          self.next_sleep = 30;
          return;
        }
      };
      for (user_id, server_id, character_id) in users {
        let server = match state.servers().iter().find(|s| s.id.0 == server_id) {
          Some(ser) => ser,
          None => {
            warn!("Couldn't find server for user ID {}", user_id);
            continue;
          }
        };
        if let Err(e) = Tagger::tag(s.clone(), UserId(user_id), server, character_id, false) {
          warn!("Couldn't update tag for user ID {}: {}", user_id, e);
        }
      }
    }
    {
      let mut database = s.database.write().unwrap();
      database.autotags.last_updated = UTC::now().timestamp();
    };
    info!("Done updating autotags");
  }
}

impl RunsTask for AutoTagTask {
  fn start(mut self, s: Arc<LalafellBot>) {
    info!("Autotag task waiting {} seconds", self.next_sleep);
    loop {
      self.run_once(s.clone());
    }
  }
}
<|endoftext|>


<commit_before>use super::Histogram;

#[cfg(test)]
pub fn histo64(
    lowest_discernible_value: u64,
    highest_trackable_value: u64,
    num_significant_digits: u8,
) -> Histogram<u64> {
    Histogram::<u64>::new_with_bounds(
        lowest_discernible_value,
        highest_trackable_value,
        num_significant_digits,
    ).unwrap()
}
<commit_msg>rustfmt<commit_after>use super::Histogram;

#[cfg(test)]
pub fn histo64(
    lowest_discernible_value: u64,
    highest_trackable_value: u64,
    num_significant_digits: u8,
) -> Histogram<u64> {
    Histogram::<u64>::new_with_bounds(
        lowest_discernible_value,
        highest_trackable_value,
        num_significant_digits,
    )
    .unwrap()
}
<|endoftext|>


<commit_before>extern crate time;

use super::rusqlite::{SqliteResult, SqliteConnection};
use super::Directory;

pub struct Aliases<'a> {
    connection: &'a SqliteConnection,
    path: Path,
    timestamp: u64,
    file_list: Vec<(uint, String)>,
    directory_id_list: Vec<uint>,
    subdirectory: Option<Box<Aliases<'a>>>
}

impl<'a> Aliases<'a> {
    pub fn new(connection: &'a SqliteConnection, path: Path, directory: Directory, timestamp: u64) -> SqliteResult<Aliases<'a>> {
        let mut row_statement = match directory {
            Directory::Child(..) => try!(connection.prepare("SELECT MAX(id) FROM alias WHERE directory_id = $1 AND timestamp <= $2 GROUP BY name;")),
            Directory::Root      => try!(connection.prepare("SELECT MAX(id) FROM alias WHERE directory_id IS NULL AND timestamp <= $1 GROUP BY name;"))
        };

        let mut rows = match directory {
            Directory::Child(id) => try!(row_statement.query(&[&(id as i64), &(timestamp as i64)])),
            Directory::Root      => try!(row_statement.query(&[&(timestamp as i64)]))
        };

        let mut directory_statement = match directory {
            Directory::Child(..) => try!(connection.prepare("SELECT id FROM directory WHERE parent_id = $1;")),
            Directory::Root      => try!(connection.prepare("SELECT id FROM directory WHERE parent_id IS NULL;"))
        };

        let mut directories = match directory {
            Directory::Child(id) => try!(directory_statement.query(&[&(id as i64)])),
            Directory::Root      => try!(directory_statement.query(&[]))
        };

        let mut alias_id_list = Vec::new();
        let mut directory_id_list = Vec::new();

        for row in rows {
            alias_id_list.push(try!(row).get::<i64>(0) as uint);
        }

        for directory in directories {
            directory_id_list.push(try!(directory).get::<i64>(0) as uint);
        }

        let file_list = alias_id_list.iter()
            .map(|id| alias_to_file(connection, *id))
            .filter(|option| option.is_some())
            .map(|option| option.unwrap())
            .collect();

        Ok(Aliases {
            connection: connection,
            path: path,
            timestamp: timestamp,
            file_list: file_list,
            directory_id_list: directory_id_list,
            subdirectory: None
        })
    }
}

impl<'a> Iterator<(Path, Box<[uint]>)> for Aliases<'a> {
    fn next(&mut self) -> Option<(Path, Box<[uint]>)> {
        // return file from child directory
        loop {
            match self.subdirectory {
                None              => (),
                Some(ref mut dir) => match dir.next() {
                    None        => (),
                    Some(alias) => return Some(alias)
                }
            }

            match self.directory_id_list.pop() {
                None     => break,
                Some(id) => {
                    let mut directory_path = self.path.clone();
                    directory_path.push(get_directory_name(self.connection, id));

                    self.subdirectory = Aliases::new(
                        self.connection,
                        directory_path,
                        Directory::Child(id),
                        self.timestamp
                    )
                    .ok().map(|alias| box alias)
                }
            }
        }

        // return file from current directory
        match self.file_list.pop() {
            Some((id, name)) => match get_file_block_list(self.connection, id) {
                Ok(boxed_slice) => {
                    let mut file_path = self.path.clone();
                    file_path.push(name);

                    Some((file_path, boxed_slice))
                },
                Err(..)     => None // no opportunity to return error in an iterator
            },
            None            => None
        }
    }
}

fn get_directory_name(connection: &SqliteConnection, directory_id: uint) -> String {
    connection.query_row(
        "SELECT name FROM directory WHERE id = $1;",
        &[&(directory_id as i64)],
        |row| row.get::<String>(0)
    )
}

fn alias_to_file(connection: &SqliteConnection, alias_id: uint) -> Option<(uint, String)> {
    connection.query_row(
        "SELECT file_id, name FROM alias WHERE id = $1;",
        &[&(alias_id as i64)],
        |row| row.get::<Option<i64>>(0).map(|id| (id as uint, row.get::<String>(1)))
    )
}

fn get_file_block_list(connection: &SqliteConnection, file_id: uint) -> SqliteResult<Box<[uint]>> {
    let mut statement = try!(connection.prepare("SELECT block_id FROM fileblock WHERE file_id = $1;"));
    let mut blocks = try!(statement.query(&[&(file_id as i64)]));
    let mut block_id_list = Vec::new();

    for block in blocks {
        block_id_list.push(try!(block).get::<i64>(0) as uint);
    }

    Ok(block_id_list.into_boxed_slice())
}

pub fn persist_file(connection: &SqliteConnection, directory: Directory, filename: &str, hash: &str, block_id_list: &[uint]) -> SqliteResult<()> {
    let transaction = try!(connection.transaction());

    try!(connection.execute("INSERT INTO file (hash) VALUES ($1);", &[&hash]));

    let file_id = connection.last_insert_rowid();

    for (ordinal, block_id) in block_id_list.iter().enumerate() {
        try!(connection.execute(
            "INSERT INTO fileblock (file_id, block_id, ordinal) VALUES ($1, $2, $3);"
            , &[&(file_id as i64), &(*block_id as i64), &(ordinal as i64)]
        ));
    }

    let alias_query = "INSERT INTO alias (directory_id, file_id, name, timestamp) VALUES ($1, $2, $3, $4);";
    let timestamp = time::get_time().sec;
    let directory_id: Option<i64> = match directory {
        Directory::Root      => None,
        Directory::Child(id) => Some(id as i64)
    };

    try!(connection.execute(alias_query, &[&directory_id, &(file_id as i64), &filename, &(timestamp as i64)]));

    transaction.commit()
}

pub fn persist_block(connection: &SqliteConnection, hash: &str) -> SqliteResult<uint> {
    try!(connection.execute("INSERT INTO block (hash) VALUES ($1);", &[&hash]));

    Ok(connection.last_insert_rowid() as uint)
}

pub fn file_known(connection: &SqliteConnection, hash: &str) -> bool {
    connection.query_row(
        "SELECT COUNT(id) FROM file WHERE hash = $1;",
        &[&hash],
        |row| row.get::<i64>(0) > 0
    )
}

pub fn block_hash_from_id(connection: &SqliteConnection, id: uint) -> String {
    connection.query_row(
        "SELECT hash FROM block WHERE id = $1;",
        &[&(id as i64)],
        |row| row.get(0)
    )
}

pub fn block_id_from_hash(connection: &SqliteConnection, hash: &str) -> Option<uint> {
    let result: Option<i64> = connection.query_row(
        "SELECT SUM(id) FROM block WHERE hash = $1;",
        &[&hash],
        |row| row.get(0)
    );

    result.map(|signed| signed as uint)
}

pub fn get_directory(connection: &SqliteConnection, parent: Directory, name: &str) -> SqliteResult<Directory> {
    let parent_id: Option<i64> = match parent {
        Directory::Root      => None,
        Directory::Child(id) => Some(id as i64)
    };

    let select_query = "SELECT SUM(id) FROM directory WHERE name = $1 AND parent_id = $2;";
    let directory_id: Option<i64> = connection.query_row(select_query, &[&name, &parent_id], |row| row.get(0));

    if directory_id.is_some() {
        return Ok(Directory::Child(directory_id.unwrap() as uint));
    }

    let insert_query = "INSERT INTO directory (parent_id, name) VALUES ($1, $2);";

    connection.execute(insert_query, &[&parent_id, &name])
        .map(|_| Directory::Child(connection.last_insert_rowid() as uint))
}

pub fn set_key(connection: &SqliteConnection, key: &str, value: &str) -> SqliteResult<uint> {
    connection.execute("INSERT INTO setting (key, value) VALUES ($1, $2);", &[&key, &value])
}

pub fn get_key(connection: &SqliteConnection, key: &str) -> Option<String> {
    connection.query_row(
        "SELECT value FROM setting WHERE key = $1;",
        &[&key],
        |row| row.get(0)
    )
}

pub fn setup(connection: &SqliteConnection) -> SqliteResult<()> {
    let queries = [
        "CREATE TABLE directory (
            id        INTEGER PRIMARY KEY,
            parent_id INTEGER,
            name      TEXT NOT NULL,
            FOREIGN KEY(parent_id) REFERENCES directory(id),
            UNIQUE(parent_id, name)
        );",
        "CREATE TABLE file (
            id           INTEGER PRIMARY KEY,
            hash         TEXT NOT NULL
        );",
        "CREATE TABLE alias (
            id           INTEGER PRIMARY KEY,
            directory_id INTEGER,
            file_id      INTEGER,
            name         TEXT NOT NULL,
            timestamp    INTEGER,
            FOREIGN KEY(directory_id) REFERENCES directory(id),
            FOREIGN KEY(file_id) REFERENCES file(id)
        );",
        "CREATE TABLE block (
            id           INTEGER PRIMARY KEY,
            hash         TEXT NOT NULL
        );",
        "CREATE TABLE fileblock (
            id           INTEGER PRIMARY KEY,
            file_id      INTEGER NOT NULL,
            ordinal      INTEGER NOT NULL,
            block_id     INTEGER NOT NULL,
            FOREIGN KEY(file_id) REFERENCES file(id),
            FOREIGN KEY(block_id) REFERENCES block(id)
        );",
        "CREATE TABLE setting (
            key          TEXT PRIMARY KEY,
            value        TEXT
        );"
    ];

    for query in queries.iter() {
        try!(connection.execute(*query, &[]));
    }

    Ok(())
}
<commit_msg>add indices on hash columns<commit_after>extern crate time;

use super::rusqlite::{SqliteResult, SqliteConnection};
use super::Directory;

pub struct Aliases<'a> {
    connection: &'a SqliteConnection,
    path: Path,
    timestamp: u64,
    file_list: Vec<(uint, String)>,
    directory_id_list: Vec<uint>,
    subdirectory: Option<Box<Aliases<'a>>>
}

impl<'a> Aliases<'a> {
    pub fn new(connection: &'a SqliteConnection, path: Path, directory: Directory, timestamp: u64) -> SqliteResult<Aliases<'a>> {
        let mut row_statement = match directory {
            Directory::Child(..) => try!(connection.prepare("SELECT MAX(id) FROM alias WHERE directory_id = $1 AND timestamp <= $2 GROUP BY name;")),
            Directory::Root      => try!(connection.prepare("SELECT MAX(id) FROM alias WHERE directory_id IS NULL AND timestamp <= $1 GROUP BY name;"))
        };

        let mut rows = match directory {
            Directory::Child(id) => try!(row_statement.query(&[&(id as i64), &(timestamp as i64)])),
            Directory::Root      => try!(row_statement.query(&[&(timestamp as i64)]))
        };

        let mut directory_statement = match directory {
            Directory::Child(..) => try!(connection.prepare("SELECT id FROM directory WHERE parent_id = $1;")),
            Directory::Root      => try!(connection.prepare("SELECT id FROM directory WHERE parent_id IS NULL;"))
        };

        let mut directories = match directory {
            Directory::Child(id) => try!(directory_statement.query(&[&(id as i64)])),
            Directory::Root      => try!(directory_statement.query(&[]))
        };

        let mut alias_id_list = Vec::new();
        let mut directory_id_list = Vec::new();

        for row in rows {
            alias_id_list.push(try!(row).get::<i64>(0) as uint);
        }

        for directory in directories {
            directory_id_list.push(try!(directory).get::<i64>(0) as uint);
        }

        let file_list = alias_id_list.iter()
            .map(|id| alias_to_file(connection, *id))
            .filter(|option| option.is_some())
            .map(|option| option.unwrap())
            .collect();

        Ok(Aliases {
            connection: connection,
            path: path,
            timestamp: timestamp,
            file_list: file_list,
            directory_id_list: directory_id_list,
            subdirectory: None
        })
    }
}

impl<'a> Iterator<(Path, Box<[uint]>)> for Aliases<'a> {
    fn next(&mut self) -> Option<(Path, Box<[uint]>)> {
        // return file from child directory
        loop {
            match self.subdirectory {
                None              => (),
                Some(ref mut dir) => match dir.next() {
                    None        => (),
                    Some(alias) => return Some(alias)
                }
            }

            match self.directory_id_list.pop() {
                None     => break,
                Some(id) => {
                    let mut directory_path = self.path.clone();
                    directory_path.push(get_directory_name(self.connection, id));

                    self.subdirectory = Aliases::new(
                        self.connection,
                        directory_path,
                        Directory::Child(id),
                        self.timestamp
                    )
                    .ok().map(|alias| box alias)
                }
            }
        }

        // return file from current directory
        match self.file_list.pop() {
            Some((id, name)) => match get_file_block_list(self.connection, id) {
                Ok(boxed_slice) => {
                    let mut file_path = self.path.clone();
                    file_path.push(name);

                    Some((file_path, boxed_slice))
                },
                Err(..)     => None // no opportunity to return error in an iterator
            },
            None            => None
        }
    }
}

fn get_directory_name(connection: &SqliteConnection, directory_id: uint) -> String {
    connection.query_row(
        "SELECT name FROM directory WHERE id = $1;",
        &[&(directory_id as i64)],
        |row| row.get::<String>(0)
    )
}

fn alias_to_file(connection: &SqliteConnection, alias_id: uint) -> Option<(uint, String)> {
    connection.query_row(
        "SELECT file_id, name FROM alias WHERE id = $1;",
        &[&(alias_id as i64)],
        |row| row.get::<Option<i64>>(0).map(|id| (id as uint, row.get::<String>(1)))
    )
}

fn get_file_block_list(connection: &SqliteConnection, file_id: uint) -> SqliteResult<Box<[uint]>> {
    let mut statement = try!(connection.prepare("SELECT block_id FROM fileblock WHERE file_id = $1;"));
    let mut blocks = try!(statement.query(&[&(file_id as i64)]));
    let mut block_id_list = Vec::new();

    for block in blocks {
        block_id_list.push(try!(block).get::<i64>(0) as uint);
    }

    Ok(block_id_list.into_boxed_slice())
}

pub fn persist_file(connection: &SqliteConnection, directory: Directory, filename: &str, hash: &str, block_id_list: &[uint]) -> SqliteResult<()> {
    let transaction = try!(connection.transaction());

    try!(connection.execute("INSERT INTO file (hash) VALUES ($1);", &[&hash]));

    let file_id = connection.last_insert_rowid();

    for (ordinal, block_id) in block_id_list.iter().enumerate() {
        try!(connection.execute(
            "INSERT INTO fileblock (file_id, block_id, ordinal) VALUES ($1, $2, $3);"
            , &[&(file_id as i64), &(*block_id as i64), &(ordinal as i64)]
        ));
    }

    let alias_query = "INSERT INTO alias (directory_id, file_id, name, timestamp) VALUES ($1, $2, $3, $4);";
    let timestamp = time::get_time().sec;
    let directory_id: Option<i64> = match directory {
        Directory::Root      => None,
        Directory::Child(id) => Some(id as i64)
    };

    try!(connection.execute(alias_query, &[&directory_id, &(file_id as i64), &filename, &(timestamp as i64)]));

    transaction.commit()
}

pub fn persist_block(connection: &SqliteConnection, hash: &str) -> SqliteResult<uint> {
    try!(connection.execute("INSERT INTO block (hash) VALUES ($1);", &[&hash]));

    Ok(connection.last_insert_rowid() as uint)
}

pub fn file_known(connection: &SqliteConnection, hash: &str) -> bool {
    connection.query_row(
        "SELECT COUNT(id) FROM file WHERE hash = $1;",
        &[&hash],
        |row| row.get::<i64>(0) > 0
    )
}

pub fn block_hash_from_id(connection: &SqliteConnection, id: uint) -> String {
    connection.query_row(
        "SELECT hash FROM block WHERE id = $1;",
        &[&(id as i64)],
        |row| row.get(0)
    )
}

pub fn block_id_from_hash(connection: &SqliteConnection, hash: &str) -> Option<uint> {
    let result: Option<i64> = connection.query_row(
        "SELECT SUM(id) FROM block WHERE hash = $1;",
        &[&hash],
        |row| row.get(0)
    );

    result.map(|signed| signed as uint)
}

pub fn get_directory(connection: &SqliteConnection, parent: Directory, name: &str) -> SqliteResult<Directory> {
    let parent_id: Option<i64> = match parent {
        Directory::Root      => None,
        Directory::Child(id) => Some(id as i64)
    };

    let select_query = "SELECT SUM(id) FROM directory WHERE name = $1 AND parent_id = $2;";
    let directory_id: Option<i64> = connection.query_row(select_query, &[&name, &parent_id], |row| row.get(0));

    if directory_id.is_some() {
        return Ok(Directory::Child(directory_id.unwrap() as uint));
    }

    let insert_query = "INSERT INTO directory (parent_id, name) VALUES ($1, $2);";

    connection.execute(insert_query, &[&parent_id, &name])
        .map(|_| Directory::Child(connection.last_insert_rowid() as uint))
}

pub fn set_key(connection: &SqliteConnection, key: &str, value: &str) -> SqliteResult<uint> {
    connection.execute("INSERT INTO setting (key, value) VALUES ($1, $2);", &[&key, &value])
}

pub fn get_key(connection: &SqliteConnection, key: &str) -> Option<String> {
    connection.query_row(
        "SELECT value FROM setting WHERE key = $1;",
        &[&key],
        |row| row.get(0)
    )
}

pub fn setup(connection: &SqliteConnection) -> SqliteResult<()> {
    let queries = [
        "CREATE TABLE directory (
            id        INTEGER PRIMARY KEY,
            parent_id INTEGER,
            name      TEXT NOT NULL,
            FOREIGN KEY(parent_id) REFERENCES directory(id),
            UNIQUE(parent_id, name)
        );",
        "CREATE TABLE file (
            id           INTEGER PRIMARY KEY,
            hash         TEXT NOT NULL
        );",
        "CREATE INDEX file_hash_index on file (hash)",
        "CREATE TABLE alias (
            id           INTEGER PRIMARY KEY,
            directory_id INTEGER,
            file_id      INTEGER,
            name         TEXT NOT NULL,
            timestamp    INTEGER,
            FOREIGN KEY(directory_id) REFERENCES directory(id),
            FOREIGN KEY(file_id) REFERENCES file(id)
        );",
        "CREATE TABLE block (
            id           INTEGER PRIMARY KEY,
            hash         TEXT NOT NULL
        );",
        "CREATE INDEX block_hash_index on block (hash)",
        "CREATE TABLE fileblock (
            id           INTEGER PRIMARY KEY,
            file_id      INTEGER NOT NULL,
            ordinal      INTEGER NOT NULL,
            block_id     INTEGER NOT NULL,
            FOREIGN KEY(file_id) REFERENCES file(id),
            FOREIGN KEY(block_id) REFERENCES block(id)
        );",
        "CREATE TABLE setting (
            key          TEXT PRIMARY KEY,
            value        TEXT
        );"
    ];

    for query in queries.iter() {
        try!(connection.execute(*query, &[]));
    }

    Ok(())
}
<|endoftext|>


<commit_before>//! Document feature set
//!
use libc::{c_char, c_int};
use std::cell::RefCell;
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::fmt;
use std::ptr;
use std::rc::{Rc, Weak};
use std::str;

use crate::bindings::*;
use crate::readonly::RoNode;
use crate::tree::node::Node;

pub(crate) type DocumentRef = Rc<RefCell<_Document>>;
pub(crate) type DocumentWeak = Weak<RefCell<_Document>>;

#[derive(Debug, Copy, Clone, Default)]
/// Save Options for Document
pub struct SaveOptions {
  /// format save output
  pub format: bool,
  /// drop the xml declaration
  pub no_declaration: bool,
  /// no empty tags
  pub no_empty_tags: bool,
  /// disable XHTML1 specific rules
  pub no_xhtml: bool,
  /// force XHTML1 specific rules
  pub xhtml: bool,
  /// force XML serialization on HTML doc
  pub as_xml: bool,
  /// force HTML serialization on XML doc
  pub as_html: bool,
  /// format with non-significant whitespace
  pub non_significant_whitespace: bool,
}

#[derive(Debug)]
pub(crate) struct _Document {
  /// pointer to a libxml document
  pub(crate) doc_ptr: xmlDocPtr,
  /// hashed pointer-to-Node bookkeeping table
  nodes: HashMap<xmlNodePtr, Node>,
}

impl _Document {
  /// Internal bookkeeping function, so far only used by `Node::wrap`
  pub(crate) fn insert_node(&mut self, node_ptr: xmlNodePtr, node: Node) {
    self.nodes.insert(node_ptr, node);
  }
  /// Internal bookkeeping function, so far only used by `Node::wrap`
  pub(crate) fn get_node(&self, node_ptr: xmlNodePtr) -> Option<&Node> {
    self.nodes.get(&node_ptr)
  }
  /// Internal bookkeeping function
  pub(crate) fn forget_node(&mut self, node_ptr: xmlNodePtr) {
    self.nodes.remove(&node_ptr);
  }
}

/// A libxml2 Document
#[derive(Clone)]
pub struct Document(pub(crate) DocumentRef);

impl Drop for _Document {
  ///Free document when it goes out of scope
  fn drop(&mut self) {
    unsafe {
      if !self.doc_ptr.is_null() {
        xmlFreeDoc(self.doc_ptr);
      }
    }
  }
}

impl fmt::Display for Document {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "{}", self.to_string_with_options(SaveOptions::default()))
  }
}

impl Document {
  /// Creates a new empty libxml2 document
  pub fn new() -> Result<Self, ()> {
    unsafe {
      let c_version = CString::new("1.0").unwrap();
      let doc_ptr = xmlNewDoc(c_version.as_bytes().as_ptr());
      if doc_ptr.is_null() {
        Err(())
      } else {
        let doc = _Document {
          doc_ptr,
          nodes: HashMap::new(),
        };
        Ok(Document(Rc::new(RefCell::new(doc))))
      }
    }
  }

  /// Obtain the underlying libxml2 `xmlDocPtr` for this Document
  pub fn doc_ptr(&self) -> xmlDocPtr {
    self.0.borrow().doc_ptr
  }

  /// Creates a new `Document` from an existing libxml2 pointer
  pub fn new_ptr(doc_ptr: xmlDocPtr) -> Self {
    let doc = _Document {
      doc_ptr,
      nodes: HashMap::new(),
    };
    Document(Rc::new(RefCell::new(doc)))
  }

  pub(crate) fn null_ref() -> DocumentRef {
    Rc::new(RefCell::new(_Document {
      doc_ptr: ptr::null_mut(),
      nodes: HashMap::new(),
    }))
  }

  /// Write document to `filename`
  pub fn save_file(&self, filename: &str) -> Result<c_int, ()> {
    let c_filename = CString::new(filename).unwrap();
    unsafe {
      let retval = xmlSaveFile(c_filename.as_ptr(), self.doc_ptr());
      if retval < 0 {
        return Err(());
      }
      Ok(retval)
    }
  }

  pub(crate) fn register_node(&self, node_ptr: xmlNodePtr) -> Node {
    Node::wrap(node_ptr, &self.0)
  }

  /// Get the root element of the document
  pub fn get_root_element(&self) -> Option<Node> {
    unsafe {
      let node_ptr = xmlDocGetRootElement(self.doc_ptr());
      if node_ptr.is_null() {
        None
      } else {
        Some(self.register_node(node_ptr))
      }
    }
  }

  /// Get the root element of the document (read-only)
  pub fn get_root_readonly(&self) -> Option<RoNode> {
    unsafe {
      let node_ptr = xmlDocGetRootElement(self.doc_ptr());
      if node_ptr.is_null() {
        None
      } else {
        Some(RoNode(node_ptr))
      }
    }
  }

  /// Sets the root element of the document
  pub fn set_root_element(&mut self, root: &Node) {
    unsafe {
      xmlDocSetRootElement(self.doc_ptr(), root.node_ptr());
    }
    root.set_linked();
  }

  fn ptr_as_result(&mut self, node_ptr: xmlNodePtr) -> Result<Node, ()> {
    if node_ptr.is_null() {
      Err(())
    } else {
      let node = self.register_node(node_ptr);
      Ok(node)
    }
  }

  /// Import a `Node` from another `Document`
  pub fn import_node(&mut self, node: &mut Node) -> Result<Node, ()> {
    if !node.is_unlinked() {
      return Err(());
    }
    // Also remove this node from the prior document hash
    node
      .get_docref()
      .upgrade()
      .unwrap()
      .borrow_mut()
      .forget_node(node.node_ptr());

    let node_ptr = unsafe { xmlDocCopyNode(node.node_ptr(), self.doc_ptr(), 1) };
    node.set_linked();
    self.ptr_as_result(node_ptr)
  }

  /// Serializes the `Document` with options
  pub fn to_string_with_options(&self, options: SaveOptions) -> String {
    unsafe {
      // allocate a buffer to dump into
      let buf = xmlBufferCreate();
      let c_utf8 = CString::new("UTF-8").unwrap();
      let mut xml_options = 0;

      if options.format {
        xml_options += xmlSaveOption_XML_SAVE_FORMAT;
      }
      if options.no_declaration {
        xml_options += xmlSaveOption_XML_SAVE_NO_DECL;
      }
      if options.no_empty_tags {
        xml_options += xmlSaveOption_XML_SAVE_NO_EMPTY;
      }
      if options.no_xhtml {
        xml_options += xmlSaveOption_XML_SAVE_NO_XHTML;
      }
      if options.xhtml {
        xml_options += xmlSaveOption_XML_SAVE_XHTML;
      }
      if options.as_xml {
        xml_options += xmlSaveOption_XML_SAVE_AS_XML;
      }
      if options.as_html {
        xml_options += xmlSaveOption_XML_SAVE_AS_HTML;
      }
      if options.non_significant_whitespace {
        xml_options += xmlSaveOption_XML_SAVE_WSNONSIG;
      }

      let save_ctx = xmlSaveToBuffer(buf, c_utf8.as_ptr(), xml_options as i32);
      let _size = xmlSaveDoc(save_ctx, self.doc_ptr());
      let _size = xmlSaveFlush(save_ctx);

      let result = xmlBufferContent(buf);
      let c_string = CStr::from_ptr(result as *const c_char);
      let node_string = c_string.to_string_lossy().into_owned();
      xmlBufferFree(buf);

      node_string
    }
  }

  /// Serializes a `Node` owned by this `Document`
  pub fn node_to_string(&self, node: &Node) -> String {
    unsafe {
      // allocate a buffer to dump into
      let buf = xmlBufferCreate();

      // dump the node
      xmlNodeDump(
        buf,
        self.doc_ptr(),
        node.node_ptr(),
        1, // level of indentation
        0, /* disable formatting */
      );
      let result = xmlBufferContent(buf);
      let c_string = CStr::from_ptr(result as *const c_char);
      let node_string = c_string.to_string_lossy().into_owned();
      xmlBufferFree(buf);

      node_string
    }
  }
  /// Serializes a `RoNode` owned by this `Document`
  pub fn ronode_to_string(&self, node: &RoNode) -> String {
    unsafe {
      // allocate a buffer to dump into
      let buf = xmlBufferCreate();

      // dump the node
      xmlNodeDump(
        buf,
        self.doc_ptr(),
        node.node_ptr(),
        1, // level of indentation
        0, /* disable formatting */
      );
      let result = xmlBufferContent(buf);
      let c_string = CStr::from_ptr(result as *const c_char);
      let node_string = c_string.to_string_lossy().into_owned();
      xmlBufferFree(buf);

      node_string
    }
  }

  /// Creates a node for an XML processing instruction
  pub fn create_processing_instruction(&mut self, name: &str, content: &str) -> Result<Node, ()> {
    unsafe {
      let c_name = CString::new(name).unwrap();
      let c_content = CString::new(content).unwrap();

      let node_ptr: xmlNodePtr = xmlNewDocPI(
        self.doc_ptr(),
        c_name.as_bytes().as_ptr(),
        c_content.as_bytes().as_ptr(),
      );
      if node_ptr.is_null() {
        Err(())
      } else {
        Ok(self.register_node(node_ptr))
      }
    }
  }

  /// Cast the document as a libxml Node
  pub fn as_node(&self) -> Node {
    // Note: this method is important to keep, as it enables certain low-level libxml2 idioms
    // In particular, method dispatch based on NodeType is only possible when the document can be cast as a Node
    //
    // Memory management is not an issue, as a document node can not be unbound/removed, and does not require
    // any additional deallocation than the Drop of a Document object.
    self.register_node(self.doc_ptr() as xmlNodePtr)
  }

  /// Duplicates the libxml2 Document into a new instance
  pub fn dup(&self) -> Result<Self, ()> {
    let doc_ptr = unsafe { xmlCopyDoc(self.doc_ptr(), 1) };
    if doc_ptr.is_null() {
      Err(())
    } else {
      let doc = _Document {
        doc_ptr,
        nodes: HashMap::new(),
      };
      Ok(Document(Rc::new(RefCell::new(doc))))
    }
  }

  /// Duplicates a source libxml2 Document into the empty Document self
  pub fn dup_from(&mut self, source: &Self) -> Result<(), ()> {
    if !self.doc_ptr().is_null() {
      return Err(());
    }

    let doc_ptr = unsafe { xmlCopyDoc(source.doc_ptr(), 1) };
    if doc_ptr.is_null() {
      return Err(());
    }
    self.0.borrow_mut().doc_ptr = doc_ptr;
    Ok(())
  }
}
<commit_msg>fixing a memory leak when converting the xml to a string<commit_after>//! Document feature set
//!
use libc::{c_char, c_int};
use std::cell::RefCell;
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::fmt;
use std::ptr;
use std::rc::{Rc, Weak};
use std::str;

use crate::bindings::*;
use crate::readonly::RoNode;
use crate::tree::node::Node;

pub(crate) type DocumentRef = Rc<RefCell<_Document>>;
pub(crate) type DocumentWeak = Weak<RefCell<_Document>>;

#[derive(Debug, Copy, Clone, Default)]
/// Save Options for Document
pub struct SaveOptions {
  /// format save output
  pub format: bool,
  /// drop the xml declaration
  pub no_declaration: bool,
  /// no empty tags
  pub no_empty_tags: bool,
  /// disable XHTML1 specific rules
  pub no_xhtml: bool,
  /// force XHTML1 specific rules
  pub xhtml: bool,
  /// force XML serialization on HTML doc
  pub as_xml: bool,
  /// force HTML serialization on XML doc
  pub as_html: bool,
  /// format with non-significant whitespace
  pub non_significant_whitespace: bool,
}

#[derive(Debug)]
pub(crate) struct _Document {
  /// pointer to a libxml document
  pub(crate) doc_ptr: xmlDocPtr,
  /// hashed pointer-to-Node bookkeeping table
  nodes: HashMap<xmlNodePtr, Node>,
}

impl _Document {
  /// Internal bookkeeping function, so far only used by `Node::wrap`
  pub(crate) fn insert_node(&mut self, node_ptr: xmlNodePtr, node: Node) {
    self.nodes.insert(node_ptr, node);
  }
  /// Internal bookkeeping function, so far only used by `Node::wrap`
  pub(crate) fn get_node(&self, node_ptr: xmlNodePtr) -> Option<&Node> {
    self.nodes.get(&node_ptr)
  }
  /// Internal bookkeeping function
  pub(crate) fn forget_node(&mut self, node_ptr: xmlNodePtr) {
    self.nodes.remove(&node_ptr);
  }
}

/// A libxml2 Document
#[derive(Clone)]
pub struct Document(pub(crate) DocumentRef);

impl Drop for _Document {
  ///Free document when it goes out of scope
  fn drop(&mut self) {
    unsafe {
      if !self.doc_ptr.is_null() {
        xmlFreeDoc(self.doc_ptr);
      }
    }
  }
}

impl fmt::Display for Document {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "{}", self.to_string_with_options(SaveOptions::default()))
  }
}

impl Document {
  /// Creates a new empty libxml2 document
  pub fn new() -> Result<Self, ()> {
    unsafe {
      let c_version = CString::new("1.0").unwrap();
      let doc_ptr = xmlNewDoc(c_version.as_bytes().as_ptr());
      if doc_ptr.is_null() {
        Err(())
      } else {
        let doc = _Document {
          doc_ptr,
          nodes: HashMap::new(),
        };
        Ok(Document(Rc::new(RefCell::new(doc))))
      }
    }
  }

  /// Obtain the underlying libxml2 `xmlDocPtr` for this Document
  pub fn doc_ptr(&self) -> xmlDocPtr {
    self.0.borrow().doc_ptr
  }

  /// Creates a new `Document` from an existing libxml2 pointer
  pub fn new_ptr(doc_ptr: xmlDocPtr) -> Self {
    let doc = _Document {
      doc_ptr,
      nodes: HashMap::new(),
    };
    Document(Rc::new(RefCell::new(doc)))
  }

  pub(crate) fn null_ref() -> DocumentRef {
    Rc::new(RefCell::new(_Document {
      doc_ptr: ptr::null_mut(),
      nodes: HashMap::new(),
    }))
  }

  /// Write document to `filename`
  pub fn save_file(&self, filename: &str) -> Result<c_int, ()> {
    let c_filename = CString::new(filename).unwrap();
    unsafe {
      let retval = xmlSaveFile(c_filename.as_ptr(), self.doc_ptr());
      if retval < 0 {
        return Err(());
      }
      Ok(retval)
    }
  }

  pub(crate) fn register_node(&self, node_ptr: xmlNodePtr) -> Node {
    Node::wrap(node_ptr, &self.0)
  }

  /// Get the root element of the document
  pub fn get_root_element(&self) -> Option<Node> {
    unsafe {
      let node_ptr = xmlDocGetRootElement(self.doc_ptr());
      if node_ptr.is_null() {
        None
      } else {
        Some(self.register_node(node_ptr))
      }
    }
  }

  /// Get the root element of the document (read-only)
  pub fn get_root_readonly(&self) -> Option<RoNode> {
    unsafe {
      let node_ptr = xmlDocGetRootElement(self.doc_ptr());
      if node_ptr.is_null() {
        None
      } else {
        Some(RoNode(node_ptr))
      }
    }
  }

  /// Sets the root element of the document
  pub fn set_root_element(&mut self, root: &Node) {
    unsafe {
      xmlDocSetRootElement(self.doc_ptr(), root.node_ptr());
    }
    root.set_linked();
  }

  fn ptr_as_result(&mut self, node_ptr: xmlNodePtr) -> Result<Node, ()> {
    if node_ptr.is_null() {
      Err(())
    } else {
      let node = self.register_node(node_ptr);
      Ok(node)
    }
  }

  /// Import a `Node` from another `Document`
  pub fn import_node(&mut self, node: &mut Node) -> Result<Node, ()> {
    if !node.is_unlinked() {
      return Err(());
    }
    // Also remove this node from the prior document hash
    node
      .get_docref()
      .upgrade()
      .unwrap()
      .borrow_mut()
      .forget_node(node.node_ptr());

    let node_ptr = unsafe { xmlDocCopyNode(node.node_ptr(), self.doc_ptr(), 1) };
    node.set_linked();
    self.ptr_as_result(node_ptr)
  }

  /// Serializes the `Document` with options
  pub fn to_string_with_options(&self, options: SaveOptions) -> String {
    unsafe {
      // allocate a buffer to dump into
      let buf = xmlBufferCreate();
      let c_utf8 = CString::new("UTF-8").unwrap();
      let mut xml_options = 0;

      if options.format {
        xml_options += xmlSaveOption_XML_SAVE_FORMAT;
      }
      if options.no_declaration {
        xml_options += xmlSaveOption_XML_SAVE_NO_DECL;
      }
      if options.no_empty_tags {
        xml_options += xmlSaveOption_XML_SAVE_NO_EMPTY;
      }
      if options.no_xhtml {
        xml_options += xmlSaveOption_XML_SAVE_NO_XHTML;
      }
      if options.xhtml {
        xml_options += xmlSaveOption_XML_SAVE_XHTML;
      }
      if options.as_xml {
        xml_options += xmlSaveOption_XML_SAVE_AS_XML;
      }
      if options.as_html {
        xml_options += xmlSaveOption_XML_SAVE_AS_HTML;
      }
      if options.non_significant_whitespace {
        xml_options += xmlSaveOption_XML_SAVE_WSNONSIG;
      }

      let save_ctx = xmlSaveToBuffer(buf, c_utf8.as_ptr(), xml_options as i32);
      let _size = xmlSaveDoc(save_ctx, self.doc_ptr());
      let _size = xmlSaveClose(save_ctx);

      let result = xmlBufferContent(buf);
      let c_string = CStr::from_ptr(result as *const c_char);
      let node_string = c_string.to_string_lossy().into_owned();
      xmlBufferFree(buf);

      node_string
    }
  }

  /// Serializes a `Node` owned by this `Document`
  pub fn node_to_string(&self, node: &Node) -> String {
    unsafe {
      // allocate a buffer to dump into
      let buf = xmlBufferCreate();

      // dump the node
      xmlNodeDump(
        buf,
        self.doc_ptr(),
        node.node_ptr(),
        1, // level of indentation
        0, /* disable formatting */
      );
      let result = xmlBufferContent(buf);
      let c_string = CStr::from_ptr(result as *const c_char);
      let node_string = c_string.to_string_lossy().into_owned();
      xmlBufferFree(buf);

      node_string
    }
  }
  /// Serializes a `RoNode` owned by this `Document`
  pub fn ronode_to_string(&self, node: &RoNode) -> String {
    unsafe {
      // allocate a buffer to dump into
      let buf = xmlBufferCreate();

      // dump the node
      xmlNodeDump(
        buf,
        self.doc_ptr(),
        node.node_ptr(),
        1, // level of indentation
        0, /* disable formatting */
      );
      let result = xmlBufferContent(buf);
      let c_string = CStr::from_ptr(result as *const c_char);
      let node_string = c_string.to_string_lossy().into_owned();
      xmlBufferFree(buf);

      node_string
    }
  }

  /// Creates a node for an XML processing instruction
  pub fn create_processing_instruction(&mut self, name: &str, content: &str) -> Result<Node, ()> {
    unsafe {
      let c_name = CString::new(name).unwrap();
      let c_content = CString::new(content).unwrap();

      let node_ptr: xmlNodePtr = xmlNewDocPI(
        self.doc_ptr(),
        c_name.as_bytes().as_ptr(),
        c_content.as_bytes().as_ptr(),
      );
      if node_ptr.is_null() {
        Err(())
      } else {
        Ok(self.register_node(node_ptr))
      }
    }
  }

  /// Cast the document as a libxml Node
  pub fn as_node(&self) -> Node {
    // Note: this method is important to keep, as it enables certain low-level libxml2 idioms
    // In particular, method dispatch based on NodeType is only possible when the document can be cast as a Node
    //
    // Memory management is not an issue, as a document node can not be unbound/removed, and does not require
    // any additional deallocation than the Drop of a Document object.
    self.register_node(self.doc_ptr() as xmlNodePtr)
  }

  /// Duplicates the libxml2 Document into a new instance
  pub fn dup(&self) -> Result<Self, ()> {
    let doc_ptr = unsafe { xmlCopyDoc(self.doc_ptr(), 1) };
    if doc_ptr.is_null() {
      Err(())
    } else {
      let doc = _Document {
        doc_ptr,
        nodes: HashMap::new(),
      };
      Ok(Document(Rc::new(RefCell::new(doc))))
    }
  }

  /// Duplicates a source libxml2 Document into the empty Document self
  pub fn dup_from(&mut self, source: &Self) -> Result<(), ()> {
    if !self.doc_ptr().is_null() {
      return Err(());
    }

    let doc_ptr = unsafe { xmlCopyDoc(source.doc_ptr(), 1) };
    if doc_ptr.is_null() {
      return Err(());
    }
    self.0.borrow_mut().doc_ptr = doc_ptr;
    Ok(())
  }
}
<|endoftext|>