<commit_before>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/

//! Saving properties of the simulation to some stream (file or stdout mainly)

extern crate chemharp;

use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use ::units;
use ::universe::Universe;

use ::universe::chemharp::universe_to_frame;

/// The `Output` trait define the interface for all the quantities outputed by
/// the simulation during the run. An Output can be a text or a binary data
/// file, an image, a text log, …
pub trait Output {
    /// Function called once at the beggining of the simulation, which allow
    /// for some setup of the output if needed.
    fn setup(&mut self, _: &Universe) {}

    /// Write the output from the universe.
    fn write(&mut self, universe: &Universe);

    /// Function called once at the end of the simulation.
    fn finish(&mut self, _: &Universe) {}
}

/******************************************************************************/
/// The `TrajectoryOutput` allow to write the trajectory of the system to a
/// file, using any format supported by the
/// [Chemharp](http://chemharp.readthedocs.org/en/latest/formats.html) library.
pub struct TrajectoryOutput {
    file: chemharp::Trajectory,
}

impl TrajectoryOutput {
    /// Create a new `TrajectoryOutput` writing to `filename`. The file is
    /// replaced if it already exists.
    pub fn new<'a, S>(filename: S) -> Result<TrajectoryOutput, chemharp::Error> where S: Into<&'a str> {
        Ok(TrajectoryOutput{
            file: try!(chemharp::Trajectory::create(filename.into()))
        })
    }
}

impl Output for TrajectoryOutput {
    fn write(&mut self, universe: &Universe) {
        let frame = match universe_to_frame(universe) {
            Ok(val) => val,
            Err(err) => {
                error!("Error in Chemharp runtime while converting data: {}", err.message());
                panic!();
            }
        };

        match self.file.write(&frame) {
            Ok(()) => {},
            Err(err) => {
                error!("Error in Chemharp runtime while writing file: {}", err.message());
                panic!();
            }
        }
    }
}

/******************************************************************************/
/// The `CellOutput` write all the components of a cell to a file . The columns
/// in the file contains the following values: `A B C α β γ`.
pub struct CellOutput {
    file: File,
    path: String
}

impl CellOutput {
    /// Create a new `CellOutput` writing to `filename`. The file is replaced if
    /// it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<CellOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(CellOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for CellOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Unit cell of the simulation") {
            // Do panic in early time
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step A/Å B/Å C/Å α/deg β/deg γ/deg") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let cell = universe.cell();
        if let Err(e) = writeln!(
            &mut self.file, "{} {} {} {} {} {} {}",
            universe.step(), cell.a(), cell.b(), cell.c(), cell.alpha(), cell.beta(), cell.gamma()
        ) {
            // Do not panic during the simulation
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

/******************************************************************************/
/// The `EnergyOutput` write the energy of the system to a text file, organized
/// as: `PotentialEnergy     KineticEnergy     TotalEnergy`.
pub struct EnergyOutput {
    file: File,
}

impl EnergyOutput {
    /// Create a new `EnergyOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<EnergyOutput, io::Error> {
        Ok(EnergyOutput{
            file: try!(File::create(filename))
        })
    }
}

impl Output for EnergyOutput {
    fn setup(&mut self, _: &Universe) {
        writeln!(&mut self.file, "# Energy of the simulation (kJ/mol)").unwrap();
        writeln!(&mut self.file, "# Step Potential Kinetic Total").unwrap();
    }

    fn write(&mut self, universe: &Universe) {
        let potential = units::to(universe.potential_energy(), "kJ/mol").unwrap();
        let kinetic = units::to(universe.kinetic_energy(), "kJ/mol").unwrap();
        let total = units::to(universe.total_energy(), "kJ/mol").unwrap();
        writeln!(&mut self.file, "{} {} {} {}", universe.step(), potential, kinetic, total).unwrap();
    }
}

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

    use super::*;
    use ::universe::*;
    use ::types::*;
    use ::potentials::*;
    use ::units;

    fn testing_universe() -> Universe {
        let mut universe = Universe::from_cell(UnitCell::cubic(10.0));;

        universe.add_particle(Particle::new("F"));
        universe[0].set_position(Vector3D::new(0.0, 0.0, 0.0));

        universe.add_particle(Particle::new("F"));
        universe[1].set_position(Vector3D::new(1.3, 0.0, 0.0));

        universe.add_pair_interaction("F", "F",
            Harmonic{k: units::from(300.0, "kJ/mol/A^2").unwrap(), x0: units::from(1.2, "A").unwrap()});
        return universe;
    }

    fn check_file_content(filename: &str, content: &str) {
        let mut file = fs::File::open(filename).unwrap();
        let mut buffer = String::new();
        file.read_to_string(&mut buffer).unwrap();

        assert_eq!(buffer, content);
    }

    #[test]
    fn trajectory() {
        let filename = "testing-trajectory-output.xyz";
        let universe = testing_universe();
        {
            let mut out = TrajectoryOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "2\nWritten by Chemharp\nF 0 0 0\nF 1.3 0 0\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn energy() {
        let filename = "testing-energy-output.dat";
        let universe = testing_universe();
        {
            let mut out = EnergyOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Energy of the simulation (kJ/mol)\n# Step Potential Kinetic Total\n0 1.5000000000000027 0 1.5000000000000027\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn cell() {
        let filename = "testing-cell-output.dat";
        let universe = testing_universe();
        {
            let mut out = CellOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Unit cell of the simulation\n# Step A/Å B/Å C/Å α/deg β/deg γ/deg\n0 10 10 10 90 90 90\n");
        fs::remove_file(filename).unwrap();
    }
}
<commit_msg>Add PropertiesOutput, for writing physical properties<commit_after>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/

//! Saving properties of the simulation to some stream (file or stdout mainly)

extern crate chemharp;

use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use ::units;
use ::universe::Universe;

use ::universe::chemharp::universe_to_frame;

/// The `Output` trait define the interface for all the quantities outputed by
/// the simulation during the run. An Output can be a text or a binary data
/// file, an image, a text log, …
pub trait Output {
    /// Function called once at the beggining of the simulation, which allow
    /// for some setup of the output if needed.
    fn setup(&mut self, _: &Universe) {}

    /// Write the output from the universe.
    fn write(&mut self, universe: &Universe);

    /// Function called once at the end of the simulation.
    fn finish(&mut self, _: &Universe) {}
}

/******************************************************************************/
/// The `TrajectoryOutput` allow to write the trajectory of the system to a
/// file, using any format supported by the
/// [Chemharp](http://chemharp.readthedocs.org/en/latest/formats.html) library.
pub struct TrajectoryOutput {
    file: chemharp::Trajectory,
}

impl TrajectoryOutput {
    /// Create a new `TrajectoryOutput` writing to `filename`. The file is
    /// replaced if it already exists.
    pub fn new<'a, S>(filename: S) -> Result<TrajectoryOutput, chemharp::Error> where S: Into<&'a str> {
        Ok(TrajectoryOutput{
            file: try!(chemharp::Trajectory::create(filename.into()))
        })
    }
}

impl Output for TrajectoryOutput {
    fn write(&mut self, universe: &Universe) {
        let frame = match universe_to_frame(universe) {
            Ok(val) => val,
            Err(err) => {
                error!("Error in Chemharp runtime while converting data: {}", err.message());
                panic!();
            }
        };

        match self.file.write(&frame) {
            Ok(()) => {},
            Err(err) => {
                error!("Error in Chemharp runtime while writing file: {}", err.message());
                panic!();
            }
        }
    }
}

/******************************************************************************/
/// The `CellOutput` write all the components of a cell to a file . The columns
/// in the file contains the following values: `A B C α β γ`.
pub struct CellOutput {
    file: File,
    path: String
}

impl CellOutput {
    /// Create a new `CellOutput` writing to `filename`. The file is replaced if
    /// it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<CellOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(CellOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for CellOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Unit cell of the simulation") {
            // Do panic in early time
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step A/Å B/Å C/Å α/deg β/deg γ/deg") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let cell = universe.cell();
        if let Err(e) = writeln!(
            &mut self.file, "{} {} {} {} {} {} {}",
            universe.step(), cell.a(), cell.b(), cell.c(), cell.alpha(), cell.beta(), cell.gamma()
        ) {
            // Do not panic during the simulation
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

/******************************************************************************/
/// The `EnergyOutput` write the energy of the system to a text file, organized
/// as: `PotentialEnergy     KineticEnergy     TotalEnergy`.
pub struct EnergyOutput {
    file: File,
}

impl EnergyOutput {
    /// Create a new `EnergyOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<EnergyOutput, io::Error> {
        Ok(EnergyOutput{
            file: try!(File::create(filename))
        })
    }
}

impl Output for EnergyOutput {
    fn setup(&mut self, _: &Universe) {
        writeln!(&mut self.file, "# Energy of the simulation (kJ/mol)").unwrap();
        writeln!(&mut self.file, "# Step Potential Kinetic Total").unwrap();
    }

    fn write(&mut self, universe: &Universe) {
        let potential = units::to(universe.potential_energy(), "kJ/mol").unwrap();
        let kinetic = units::to(universe.kinetic_energy(), "kJ/mol").unwrap();
        let total = units::to(universe.total_energy(), "kJ/mol").unwrap();
        writeln!(&mut self.file, "{} {} {} {}", universe.step(), potential, kinetic, total).unwrap();
    }
}

/******************************************************************************/
/// The `PropertiesOutput` write various physical properties of the universe to
/// a file. These properties are:
///
/// - volume of the unit cell;
/// - instant temperature;
/// - instant pressure;
pub struct PropertiesOutput {
    file: File,
    path: String,
}

impl PropertiesOutput {
    /// Create a new `PropertiesOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<PropertiesOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(PropertiesOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for PropertiesOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Physical properties of the simulation") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step Volume/A^3 Temperature/K Pressure/bar") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let V = units::to(universe.volume(), "A^3").unwrap();
        let T = units::to(universe.temperature(), "K").unwrap();
        let P = units::to(universe.pressure(), "bar").unwrap();
        if let Err(e) = writeln!(&mut self.file, "{} {} {} {}", universe.step(), V, T, P) {
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

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

    use super::*;
    use ::universe::*;
    use ::types::*;
    use ::potentials::*;
    use ::units;

    fn testing_universe() -> Universe {
        let mut universe = Universe::from_cell(UnitCell::cubic(10.0));;

        universe.add_particle(Particle::new("F"));
        universe[0].set_position(Vector3D::new(0.0, 0.0, 0.0));

        universe.add_particle(Particle::new("F"));
        universe[1].set_position(Vector3D::new(1.3, 0.0, 0.0));

        universe.add_pair_interaction("F", "F",
            Harmonic{k: units::from(300.0, "kJ/mol/A^2").unwrap(), x0: units::from(1.2, "A").unwrap()});
        return universe;
    }

    fn check_file_content(filename: &str, content: &str) {
        let mut file = fs::File::open(filename).unwrap();
        let mut buffer = String::new();
        file.read_to_string(&mut buffer).unwrap();

        assert_eq!(buffer, content);
    }

    #[test]
    fn trajectory() {
        let filename = "testing-trajectory-output.xyz";
        let universe = testing_universe();
        {
            let mut out = TrajectoryOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "2\nWritten by Chemharp\nF 0 0 0\nF 1.3 0 0\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn energy() {
        let filename = "testing-energy-output.dat";
        let universe = testing_universe();
        {
            let mut out = EnergyOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Energy of the simulation (kJ/mol)\n# Step Potential Kinetic Total\n0 1.5000000000000027 0 1.5000000000000027\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn cell() {
        let filename = "testing-cell-output.dat";
        let universe = testing_universe();
        {
            let mut out = CellOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Unit cell of the simulation\n# Step A/Å B/Å C/Å α/deg β/deg γ/deg\n0 10 10 10 90 90 90\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn properties() {
        let filename = "testing-properties-output.dat";
        let universe = testing_universe();
        {
            let mut out = PropertiesOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Physical properties of the simulation\n# Step Volume/A^3 Temperature/K Pressure/bar\n0 1000 0 -431.7400836223091\n");
        fs::remove_file(filename).unwrap();
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "nightly-testing", plugin(clippy))]
#![cfg_attr(feature = "nightly-testing", feature(plugin))]
#![cfg_attr(feature = "nightly-testing", allow(too_many_arguments))]
#![cfg_attr(feature = "nightly-testing", allow(used_underscore_binding))]
#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]
#![cfg_attr(not(feature = "with-syntex"), plugin(quasi_macros))]

extern crate aster;
extern crate quasi;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
#[macro_use]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
#[macro_use]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

#[cfg(feature = "with-syntex")]
use std::io;
#[cfg(feature = "with-syntex")]
use std::path::Path;

#[cfg(not(feature = "with-syntex"))]
use syntax::feature_gate::AttributeType;

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/lib.rs"));

#[cfg(not(feature = "with-syntex"))]
include!("lib.rs.in");

#[cfg(feature = "with-syntex")]
pub fn expand(src: &Path, dst: &Path) -> io::Result<()> {
    let mut registry = syntex::Registry::new();
    register(&mut registry);
    registry.expand("", src, dst)
}

#[cfg(feature = "with-syntex")]
pub fn register(reg: &mut syntex::Registry) {
    use syntax::{ast, fold};

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", ser::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", de::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(ser::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(de::expand_derive_deserialize)));

    reg.register_attribute("serde".to_owned(), AttributeType::Normal);
}
<commit_msg>Use AsRef to accept paths in serde_codegen::expand<commit_after>#![cfg_attr(feature = "nightly-testing", plugin(clippy))]
#![cfg_attr(feature = "nightly-testing", feature(plugin))]
#![cfg_attr(feature = "nightly-testing", allow(too_many_arguments))]
#![cfg_attr(feature = "nightly-testing", allow(used_underscore_binding))]
#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]
#![cfg_attr(not(feature = "with-syntex"), plugin(quasi_macros))]

extern crate aster;
extern crate quasi;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
#[macro_use]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
#[macro_use]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

#[cfg(feature = "with-syntex")]
use std::io;
#[cfg(feature = "with-syntex")]
use std::path::Path;

#[cfg(not(feature = "with-syntex"))]
use syntax::feature_gate::AttributeType;

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/lib.rs"));

#[cfg(not(feature = "with-syntex"))]
include!("lib.rs.in");

#[cfg(feature = "with-syntex")]
pub fn expand<S, D>(src: S, dst: D) -> io::Result<()>
    where S: AsRef<Path>,
          D: AsRef<Path>,
{
    let mut registry = syntex::Registry::new();
    register(&mut registry);
    registry.expand("", src.as_ref(), dst.as_ref())
}

#[cfg(feature = "with-syntex")]
pub fn register(reg: &mut syntex::Registry) {
    use syntax::{ast, fold};

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", ser::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", de::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(ser::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(de::expand_derive_deserialize)));

    reg.register_attribute("serde".to_owned(), AttributeType::Normal);
}
<|endoftext|>


<commit_before>#[derive(Debug, Eq, Clone, Hash, PartialEq, Serialize, Deserialize)]
pub enum ResearchTypes {
    PlasmaGenerator,
    EnergyWeapons,
    MiningEfficiency,
}

#[derive(Debug, Eq, Clone, Hash, PartialEq, Serialize, Deserialize)]
pub enum ResourceTypes {
    Iron,
    Water,
}

#[derive(Debug, Eq, Clone, Hash, PartialEq, Serialize, Deserialize)]
pub enum ModuleTypes {
    Turret,
    Generator,
}
<commit_msg>Macro for enum types creation with traits for from and to string<commit_after>use std::fmt;

#[macro_export]
macro_rules! enum_impl {
    ($enumname: ident {
        $($enumvals: ident,)*
    }) => {
        #[derive(Debug, Eq, Clone, Hash, PartialEq, Serialize, Deserialize)]
        pub enum $enumname {
            $($enumvals,)*
        }

        impl fmt::Display for $enumname {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                fmt::Debug::fmt(self, f)
            }
        }

        impl $enumname {
            #[allow(dead_code)]
            pub fn from_str(name: &str) -> Result<$enumname, ()> {
                match name {
                    $(
                        stringify!($enumvals) => Ok($enumname::$enumvals),
                    )*
                    _ => Err(()),
                }
            }

        }
    }
}

enum_impl!{ResearchTypes {
    PlasmaGenerator,
    EnergyWeapons,
    MiningEfficiency,
}}

enum_impl!{ResourceTypes {
    Iron,
    Water,
}}

enum_impl!{ModuleTypes {
    Turret,
    Generator,
}}
<|endoftext|>


<commit_before>use std::io::Error as IoError;
use std::io::Read;

pub use parse::ParseError;

mod enums;
mod parse;

pub fn reflect<R>(name: &str, mut spirv: R) -> Result<String, Error>
    where R: Read
{
    let mut data = Vec::new();
    try!(spirv.read_to_end(&mut data));

    // now parsing the document
    let doc = try!(parse::parse_spirv(&data));

    let mut output = String::new();

    {
        let spirv_data = data.iter().map(|&byte| byte.to_string())
                             .collect::<Vec<String>>()
                             .join(", ");
        output.push_str(&format!(r#"
pub struct {name};

impl {name} {{
    pub fn load(device: &::std::sync::Arc<::vulkano::Device>) {{
        unsafe {{
            let data = [{spirv_data}];

            ::vulkano::ShaderModule::new(device, &data)
        }}
    }}
}}
        "#, name = name, spirv_data = spirv_data));
    }

    println!("{:#?}", doc);

    for instruction in doc.instructions.iter() {
        match instruction {
            &parse::Instruction::Variable { result_type_id, result_id, ref storage_class, .. } => {
                match *storage_class {
                    enums::StorageClass::StorageClassUniformConstant => (),
                    enums::StorageClass::StorageClassUniform => (),
                    enums::StorageClass::StorageClassPushConstant => (),
                    _ => continue
                };

                let name = name_from_id(&doc, result_id);
                output.push_str(&format!("{}: {};", name, type_from_id(&doc, result_type_id)));
            },
            _ => ()
        }
    }

    Ok(output)
}

fn name_from_id(doc: &parse::Spirv, searched: u32) -> String {
    doc.instructions.iter().filter_map(|i| {
        if let &parse::Instruction::Name { target_id, ref name } = i {
            if target_id == searched {
                Some(name.clone())
            } else {
                None
            }
        } else {
            None
        }
    }).next().and_then(|n| if !n.is_empty() { Some(n) } else { None })
      .unwrap_or("__unnamed".to_owned())
}

fn member_name_from_id(doc: &parse::Spirv, searched: u32, searched_member: u32) -> String {
    doc.instructions.iter().filter_map(|i| {
        if let &parse::Instruction::MemberName { target_id, member, ref name } = i {
            if target_id == searched && member == searched_member {
                Some(name.clone())
            } else {
                None
            }
        } else {
            None
        }
    }).next().and_then(|n| if !n.is_empty() { Some(n) } else { None })
      .unwrap_or("__unnamed".to_owned())
}

fn type_from_id(doc: &parse::Spirv, searched: u32) -> String {
    for instruction in doc.instructions.iter() {
        match instruction {
            &parse::Instruction::TypeVoid { result_id } if result_id == searched => {
                return "()".to_owned()
            },
            &parse::Instruction::TypeBool { result_id } if result_id == searched => {
                return "bool".to_owned()
            },
            &parse::Instruction::TypeInt { result_id, width, signedness } if result_id == searched => {
                return "i32".to_owned()
            },
            &parse::Instruction::TypeFloat { result_id, width } if result_id == searched => {
                return "f32".to_owned()
            },
            &parse::Instruction::TypeVector { result_id, component_id, count } if result_id == searched => {
                let t = type_from_id(doc, component_id);
                return format!("[{}; {}]", t, count);
            },
            &parse::Instruction::TypeImage { result_id, sampled_type_id, ref dim, depth, arrayed, ms, sampled, ref format, ref access } if result_id == searched => {
                return format!("{}{}Texture{:?}{}{:?}",
                    if ms { "Multisample" } else { "" },
                    if depth == Some(true) { "Depth" } else { "" },
                    dim,
                    if arrayed { "Array" } else { "" },
                    format);
            },
            &parse::Instruction::TypeSampledImage { result_id, image_type_id } if result_id == searched => {
                return type_from_id(doc, image_type_id);
            },
            &parse::Instruction::TypeArray { result_id, type_id, length_id } if result_id == searched => {
                let t = type_from_id(doc, type_id);
                let len = doc.instructions.iter().filter_map(|e| {
                    match e { &parse::Instruction::Constant { result_id, ref data, .. } if result_id == length_id => Some(data.clone()), _ => None }
                }).next().expect("failed to find array length");
                let len = len.iter().rev().fold(0u64, |a, &b| (a << 32) | b as u64);
                return format!("[{}; {}]", t, len);       // FIXME:
            },
            &parse::Instruction::TypeRuntimeArray { result_id, type_id } if result_id == searched => {
                let t = type_from_id(doc, type_id);
                return format!("[{}]", t);
            },
            &parse::Instruction::TypeStruct { result_id, ref member_types } if result_id == searched => {
                let name = name_from_id(doc, result_id);
                let members = member_types.iter().enumerate().map(|(offset, &member)| {
                    let ty = type_from_id(doc, member);
                    let name = member_name_from_id(doc, result_id, offset as u32);
                    format!("\t{}: {}", name, ty)
                }).collect::<Vec<_>>();
                return format!("struct {} {{\n{}\n}}", name, members.join(",\n"));
            },
            &parse::Instruction::TypeOpaque { result_id, ref name } if result_id == searched => {
            },
            &parse::Instruction::TypePointer { result_id, type_id, .. } if result_id == searched => {
                let t = type_from_id(doc, type_id);
                return format!("*const {}", t);
            },
            _ => ()
        }
    }

    panic!("Type #{} not found", searched)
}

#[derive(Debug)]
pub enum Error {
    IoError(IoError),
    ParseError(ParseError),
}

impl From<IoError> for Error {
    #[inline]
    fn from(err: IoError) -> Error {
        Error::IoError(err)
    }
}

impl From<ParseError> for Error {
    #[inline]
    fn from(err: ParseError) -> Error {
        Error::ParseError(err)
    }
}
<commit_msg>Start doing something correct in the code generation<commit_after>use std::io::Error as IoError;
use std::io::Read;

pub use parse::ParseError;

mod enums;
mod parse;

pub fn reflect<R>(name: &str, mut spirv: R) -> Result<String, Error>
    where R: Read
{
    let mut data = Vec::new();
    try!(spirv.read_to_end(&mut data));

    // now parsing the document
    let doc = try!(parse::parse_spirv(&data));

    let mut output = String::new();

    {
        // contains the data that was passed as input to this function
        let spirv_data = data.iter().map(|&byte| byte.to_string())
                             .collect::<Vec<String>>()
                             .join(", ");

        // writing the header
        output.push_str(&format!(r#"
pub struct {name} {{
    shader: ::std::sync::Arc<::vulkano::ShaderModule>,
}}

impl {name} {{
    /// Loads the shader in Vulkan as a `ShaderModule`.
    #[inline]
    pub fn load(device: &::std::sync::Arc<::vulkano::Device>) -> {name} {{
        unsafe {{
            let data = [{spirv_data}];

            {name} {{
                shader: ::vulkano::ShaderModule::new(device, &data)
            }}
        }}
    }}

    /// Returns the module that was created.
    #[inline]
    pub fn module(&self) -> &::std::sync::Arc<::vulkano::ShaderModule> {{
        &self.shader
    }}
        "#, name = name, spirv_data = spirv_data));

        // writing one method for each entry point of this module
        for instruction in doc.instructions.iter() {
            if let &parse::Instruction::EntryPoint { .. } = instruction {
                output.push_str(&write_entry_point(&doc, instruction));
            }
        }

        // footer
        output.push_str(&format!(r#"
}}
        "#));
    }

    // TODO: remove
    println!("{:#?}", doc);

    Ok(output)
}

#[derive(Debug)]
pub enum Error {
    IoError(IoError),
    ParseError(ParseError),
}

impl From<IoError> for Error {
    #[inline]
    fn from(err: IoError) -> Error {
        Error::IoError(err)
    }
}

impl From<ParseError> for Error {
    #[inline]
    fn from(err: ParseError) -> Error {
        Error::ParseError(err)
    }
}

fn write_entry_point(doc: &parse::Spirv, instruction: &parse::Instruction) -> String {
    let (execution, ep_name, interface) = match instruction {
        &parse::Instruction::EntryPoint { ref execution, id, ref name, ref interface } => {
            (execution, name, interface)
        },
        _ => unreachable!()
    };

    let ty = match *execution {
        enums::ExecutionModel::ExecutionModelVertex => {
            format!("VertexShaderEntryPoint")
        },

        enums::ExecutionModel::ExecutionModelTessellationControl => {
            format!("TessControlShaderEntryPoint")
        },

        enums::ExecutionModel::ExecutionModelTessellationEvaluation => {
            format!("TessEvaluationShaderEntryPoint")
        },

        enums::ExecutionModel::ExecutionModelGeometry => {
            format!("GeometryShaderEntryPoint")
        },

        enums::ExecutionModel::ExecutionModelFragment => {
            let mut output_types = Vec::new();

            for interface in interface.iter() {
                for i in doc.instructions.iter() {
                    match i {
                        &parse::Instruction::Variable { result_type_id, result_id,
                                    storage_class: enums::StorageClass::StorageClassOutput, .. }
                                    if &result_id == interface =>
                        {
                            output_types.push(type_from_id(doc, result_type_id));
                        },
                        _ => ()
                    }
                }
            }

            format!("FragmentShaderEntryPoint<({output})>",
                    output = output_types.join(", ") + ",")
        },

        enums::ExecutionModel::ExecutionModelGLCompute => {
            format!("ComputeShaderEntryPoint")
        },

        enums::ExecutionModel::ExecutionModelKernel => panic!("Kernels are not supported"),
    };

    format!(r#"
    /// Returns a logical struct describing the entry point named `{ep_name}`.
    #[inline]
    pub fn {ep_name}_entry_point(&self) -> {ty} {{
        self.shader.entry_point("{ep_name}")
    }}
            "#, ep_name = ep_name, ty = ty)
}

fn type_from_id(doc: &parse::Spirv, searched: u32) -> String {
    for instruction in doc.instructions.iter() {
        match instruction {
            &parse::Instruction::TypeVoid { result_id } if result_id == searched => {
                return "()".to_owned()
            },
            &parse::Instruction::TypeBool { result_id } if result_id == searched => {
                return "bool".to_owned()
            },
            &parse::Instruction::TypeInt { result_id, width, signedness } if result_id == searched => {
                return "i32".to_owned()
            },
            &parse::Instruction::TypeFloat { result_id, width } if result_id == searched => {
                return "f32".to_owned()
            },
            &parse::Instruction::TypeVector { result_id, component_id, count } if result_id == searched => {
                let t = type_from_id(doc, component_id);
                return format!("[{}; {}]", t, count);
            },
            &parse::Instruction::TypeImage { result_id, sampled_type_id, ref dim, depth, arrayed, ms, sampled, ref format, ref access } if result_id == searched => {
                return format!("{}{}Texture{:?}{}{:?}",
                    if ms { "Multisample" } else { "" },
                    if depth == Some(true) { "Depth" } else { "" },
                    dim,
                    if arrayed { "Array" } else { "" },
                    format);
            },
            &parse::Instruction::TypeSampledImage { result_id, image_type_id } if result_id == searched => {
                return type_from_id(doc, image_type_id);
            },
            &parse::Instruction::TypeArray { result_id, type_id, length_id } if result_id == searched => {
                let t = type_from_id(doc, type_id);
                let len = doc.instructions.iter().filter_map(|e| {
                    match e { &parse::Instruction::Constant { result_id, ref data, .. } if result_id == length_id => Some(data.clone()), _ => None }
                }).next().expect("failed to find array length");
                let len = len.iter().rev().fold(0u64, |a, &b| (a << 32) | b as u64);
                return format!("[{}; {}]", t, len);       // FIXME:
            },
            &parse::Instruction::TypeRuntimeArray { result_id, type_id } if result_id == searched => {
                let t = type_from_id(doc, type_id);
                return format!("[{}]", t);
            },
            &parse::Instruction::TypeStruct { result_id, ref member_types } if result_id == searched => {
                let name = name_from_id(doc, result_id);
                let members = member_types.iter().enumerate().map(|(offset, &member)| {
                    let ty = type_from_id(doc, member);
                    let name = member_name_from_id(doc, result_id, offset as u32);
                    format!("\t{}: {}", name, ty)
                }).collect::<Vec<_>>();
                return format!("struct {} {{\n{}\n}}", name, members.join(",\n"));
            },
            &parse::Instruction::TypeOpaque { result_id, ref name } if result_id == searched => {
            },
            &parse::Instruction::TypePointer { result_id, type_id, .. } if result_id == searched => {
                let t = type_from_id(doc, type_id);
                return format!("*const {}", t);
            },
            _ => ()
        }
    }

    panic!("Type #{} not found", searched)
}

fn name_from_id(doc: &parse::Spirv, searched: u32) -> String {
    doc.instructions.iter().filter_map(|i| {
        if let &parse::Instruction::Name { target_id, ref name } = i {
            if target_id == searched {
                Some(name.clone())
            } else {
                None
            }
        } else {
            None
        }
    }).next().and_then(|n| if !n.is_empty() { Some(n) } else { None })
      .unwrap_or("__unnamed".to_owned())
}

fn member_name_from_id(doc: &parse::Spirv, searched: u32, searched_member: u32) -> String {
    doc.instructions.iter().filter_map(|i| {
        if let &parse::Instruction::MemberName { target_id, member, ref name } = i {
            if target_id == searched && member == searched_member {
                Some(name.clone())
            } else {
                None
            }
        } else {
            None
        }
    }).next().and_then(|n| if !n.is_empty() { Some(n) } else { None })
      .unwrap_or("__unnamed".to_owned())
}
<|endoftext|>


<commit_before>use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::write_unchecked;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::LEVEL_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::{BTreeSet, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound);
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut node_ref = node.clone();
    loop {
        let mut page = write_node::<KS, PS>(&node_ref);
        if !page.is_empty_node() {
            let mut ptr_indices: HashSet<_> = page.innode()
                .ptrs
                .as_slice_immute()
                .iter()
                .enumerate()
                .filter_map(|(i, sub_level)| {
                    if read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                        Some(i)
                    } else {
                        None
                    }
                })
                .collect();
            if ptr_indices.len() == page.len() + 1 {
                // all sub nodes are empty
                // will set current node empty either
                let next_page_ref = page.innode().right.clone();
                *page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: next_page_ref,
                });
            } else {
                let mut innode = page.innode_mut();
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                {
                    let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, p)| p)
                        .collect();

                    let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&(i + 1)))
                        .map(|(_, k)| k)
                        .collect();
                    debug!("Prune filtered page have keys {:?}", &keys);
                    innode.len = keys.len();
                    for (i, key) in keys.into_iter().enumerate() {
                        debug_assert!(key > &mut smallvec!(0));
                        mem::swap(&mut new_keys.as_slice()[i], key);
                    }
                    for (i, ptr) in ptrs.into_iter().enumerate() {
                        debug_assert!(!ptr.is_default());
                        mem::swap(&mut new_ptrs.as_slice()[i], ptr);
                    }
                }
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
            if page.right_bound() >= bound {
                break;
            }
        }
        node_ref = page.right_ref().unwrap().clone();
    }
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        new_first_node_ext.id = left_most_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert_eq!(new_first_node_ext.id, src_tree.head_page_id);

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    prune_selected::<KS, PS>(&src_tree.get_root(), &prune_bound);

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<commit_msg>cargo fmt<commit_after>use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::write_unchecked;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::LEVEL_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::{BTreeSet, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound);
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut node_ref = node.clone();
    loop {
        let mut page = write_node::<KS, PS>(&node_ref);
        if !page.is_empty_node() {
            let mut ptr_indices: HashSet<_> = page
                .innode()
                .ptrs
                .as_slice_immute()
                .iter()
                .enumerate()
                .filter_map(|(i, sub_level)| {
                    if read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                        Some(i)
                    } else {
                        None
                    }
                })
                .collect();
            if ptr_indices.len() == page.len() + 1 {
                // all sub nodes are empty
                // will set current node empty either
                let next_page_ref = page.innode().right.clone();
                *page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: next_page_ref,
                });
            } else {
                let mut innode = page.innode_mut();
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                {
                    let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, p)| p)
                        .collect();

                    let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&(i + 1)))
                        .map(|(_, k)| k)
                        .collect();
                    debug!("Prune filtered page have keys {:?}", &keys);
                    innode.len = keys.len();
                    for (i, key) in keys.into_iter().enumerate() {
                        debug_assert!(key > &mut smallvec!(0));
                        mem::swap(&mut new_keys.as_slice()[i], key);
                    }
                    for (i, ptr) in ptrs.into_iter().enumerate() {
                        debug_assert!(!ptr.is_default());
                        mem::swap(&mut new_ptrs.as_slice()[i], ptr);
                    }
                }
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
            if page.right_bound() >= bound {
                break;
            }
        }
        node_ref = page.right_ref().unwrap().clone();
    }
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        new_first_node_ext.id = left_most_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert_eq!(new_first_node_ext.id, src_tree.head_page_id);

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    prune_selected::<KS, PS>(&src_tree.get_root(), &prune_bound);

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use index::EntryKey;
use index::Slice;
use std::fmt::Debug;
use std::mem;
use std::sync::Arc;
use std::collections::BTreeSet;
use parking_lot::RwLock;

// LevelTree items cannot been added or removed individually
// Items must been merged from higher level in bulk
// Deletion will be performed when merging to or from higher level tree
// Because tree update will not to be performed in parallel. Unlike memtable, a single r/w lock
// should be sufficient. Thus concurrency control will be simple and efficient.
pub struct LevelTree<S>
    where S: Slice<EntryKey> + SortableEntrySlice
{
    index: BTreeMap<EntryKey, SSIndex<S>>
}

type TombstoneSet = BTreeSet<EntryKey>;

struct SSIndex<S> where S: Slice<EntryKey> + SortableEntrySlice
{
    slice: S,
    tombstones: TombstoneSet
}

pub trait SortableEntrySlice: Sized + Slice<EntryKey>
{
    fn merge_sorted_with(
        &mut self,
        y_slice: &mut [EntryKey],
        xlen: &mut usize,
        ylen: &mut usize,
        xtombstones: &mut TombstoneSet,
        ytombstones: &mut TombstoneSet
    ) -> Self {
        let total_len = *xlen + *ylen;
        let mut x_pos = 0;
        let mut y_pos = 0;
        let mut pos = 0;
        let mut new_x = Self::init();
        let mut new_y = Self::init();
        {
            let mut x_slice = self.as_slice();
            let x_slice_len = x_slice.len();
            debug_assert!(*ylen <= x_slice_len);
            let mut new_slice_x = new_x.as_slice();
            let mut new_slice_y = new_y.as_slice();
            let mut new_x_len = 0;
            let mut new_y_len = 0;
            let mut use_slice_at = | slice: &mut[EntryKey], slice_pos: &mut usize| {
                let item = mem::replace(&mut slice[*slice_pos], EntryKey::default());
                *slice_pos += 1;
                if xtombstones.remove(&item) || ytombstones.remove(&item) {
                    None
                } else {
                    Some(item)
                }
            };
            while pos < total_len {
                let item_or_removed = if x_pos < *xlen && y_pos < *ylen {
                    if x_slice[x_pos] < y_slice[y_pos] {
                        use_slice_at(x_slice, &mut x_pos)
                    } else {
                        use_slice_at(y_slice, &mut y_pos)
                    }
                } else if x_pos < *xlen {
                    use_slice_at(x_slice, &mut x_pos)
                } else if y_pos < *ylen {
                    use_slice_at(y_slice, &mut y_pos)
                } else { unreachable!() };

                if let Some(item) = item_or_removed {
                    if pos < x_slice_len {
                        new_slice_x[pos] = item;
                        new_x_len += 1;
                    } else {
                        new_slice_y[pos - x_slice_len] = item;
                        new_y_len += 1;
                    }
                    pos += 1;
                }
            }
            *xlen = new_x_len;
            *ylen = new_y_len;
        }
        mem::swap(self, &mut new_x);
        return new_y;
    }
}

#[cfg(test)]
mod test {

}<commit_msg>break on search completed<commit_after>use std::collections::BTreeMap;
use index::EntryKey;
use index::Slice;
use std::fmt::Debug;
use std::mem;
use std::sync::Arc;
use std::collections::BTreeSet;
use parking_lot::RwLock;

// LevelTree items cannot been added or removed individually
// Items must been merged from higher level in bulk
// Deletion will be performed when merging to or from higher level tree
// Because tree update will not to be performed in parallel. Unlike memtable, a single r/w lock
// should be sufficient. Thus concurrency control will be simple and efficient.
pub struct LevelTree<S>
    where S: Slice<EntryKey> + SortableEntrySlice
{
    index: BTreeMap<EntryKey, SSIndex<S>>
}

type TombstoneSet = BTreeSet<EntryKey>;

struct SSIndex<S> where S: Slice<EntryKey> + SortableEntrySlice
{
    slice: S,
    tombstones: TombstoneSet
}

pub trait SortableEntrySlice: Sized + Slice<EntryKey>
{
    fn merge_sorted_with(
        &mut self,
        y_slice: &mut [EntryKey],
        xlen: &mut usize,
        ylen: &mut usize,
        xtombstones: &mut TombstoneSet,
        ytombstones: &mut TombstoneSet
    ) -> Self {
        let mut x_pos = 0;
        let mut y_pos = 0;
        let mut pos = 0;
        let mut new_x = Self::init();
        let mut new_y = Self::init();
        {
            let mut x_slice = self.as_slice();
            let x_slice_len = x_slice.len();
            debug_assert!(*ylen <= x_slice_len);
            let mut new_slice_x = new_x.as_slice();
            let mut new_slice_y = new_y.as_slice();
            let mut new_x_len = 0;
            let mut new_y_len = 0;
            let mut use_slice_at = | slice: &mut[EntryKey], slice_pos: &mut usize| {
                let item = mem::replace(&mut slice[*slice_pos], EntryKey::default());
                *slice_pos += 1;
                if xtombstones.remove(&item) || ytombstones.remove(&item) {
                    None
                } else {
                    Some(item)
                }
            };
            while pos < total_len {
                let item_or_removed = if x_pos < *xlen && y_pos < *ylen {
                    if x_slice[x_pos] < y_slice[y_pos] {
                        use_slice_at(x_slice, &mut x_pos)
                    } else {
                        use_slice_at(y_slice, &mut y_pos)
                    }
                } else if x_pos < *xlen {
                    use_slice_at(x_slice, &mut x_pos)
                } else if y_pos < *ylen {
                    use_slice_at(y_slice, &mut y_pos)
                } else {
                    break
                };

                if let Some(item) = item_or_removed {
                    if pos < x_slice_len {
                        new_slice_x[pos] = item;
                        new_x_len += 1;
                    } else {
                        new_slice_y[pos - x_slice_len] = item;
                        new_y_len += 1;
                    }
                    pos += 1;
                }
            }
            *xlen = new_x_len;
            *ylen = new_y_len;
        }
        mem::swap(self, &mut new_x);
        return new_y;
    }
}

#[cfg(test)]
mod test {

}<|endoftext|>


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

#![feature(rustc_attrs)]

// check raw fat pointer ops in mir
// FIXME: please improve this when we get monomorphization support

use std::mem;

#[derive(Debug, PartialEq, Eq)]
struct ComparisonResults {
    lt: bool,
    le: bool,
    gt: bool,
    ge: bool,
    eq: bool,
    ne: bool
}

const LT: ComparisonResults = ComparisonResults {
    lt: true,
    le: true,
    gt: false,
    ge: false,
    eq: false,
    ne: true
};

const EQ: ComparisonResults = ComparisonResults {
    lt: false,
    le: true,
    gt: false,
    ge: true,
    eq: true,
    ne: false
};

const GT: ComparisonResults = ComparisonResults {
    lt: false,
    le: false,
    gt: true,
    ge: true,
    eq: false,
    ne: true
};

#[rustc_mir]
fn compare_su8(a: *const S<[u8]>, b: *const S<[u8]>) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

#[rustc_mir]
fn compare_au8(a: *const [u8], b: *const [u8]) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

#[rustc_mir(graphviz="comparefoo.gv")]
fn compare_foo<'a>(a: *const (Foo+'a), b: *const (Foo+'a)) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

#[rustc_mir(graphviz="simpleeq.gv")]
fn simple_eq<'a>(a: *const (Foo+'a), b: *const (Foo+'a)) -> bool {
    let result = a == b;
    result
}

fn assert_inorder<T: Copy>(a: &[T],
                           compare: fn(T, T) -> ComparisonResults) {
    for i in 0..a.len() {
        for j in 0..a.len() {
            let cres = compare(a[i], a[j]);
            if i < j {
                assert_eq!(cres, LT);
            } else if i == j {
                assert_eq!(cres, EQ);
            } else {
                assert_eq!(cres, GT);
            }
        }
    }
}

trait Foo { fn foo(&self) -> usize; }
impl<T> Foo for T {
    fn foo(&self) -> usize {
        mem::size_of::<T>()
    }
}

struct S<T:?Sized>(u32, T);

fn main() {
    let array = [0,1,2,3,4];
    let array2 = [5,6,7,8,9];

    // fat ptr comparison: addr then extra

    // check ordering for arrays
    let mut ptrs: Vec<*const [u8]> = vec![
        &array[0..0], &array[0..1], &array, &array[1..]
    ];

    let array_addr = &array as *const [u8] as *const u8 as usize;
    let array2_addr = &array2 as *const [u8] as *const u8 as usize;
    if array2_addr < array_addr {
        ptrs.insert(0, &array2);
    } else {
        ptrs.push(&array2);
    }
    assert_inorder(&ptrs, compare_au8);

    let u8_ = (0u8, 1u8);
    let u32_ = (4u32, 5u32);

    // check ordering for ptrs
    let buf: &mut [*const Foo] = &mut [
        &u8_, &u8_.0,
        &u32_, &u32_.0,
    ];
    buf.sort_by(|u,v| {
        let u : [*const (); 2] = unsafe { mem::transmute(*u) };
        let v : [*const (); 2] = unsafe { mem::transmute(*v) };
        u.cmp(&v)
    });
    assert_inorder(buf, compare_foo);

    // check ordering for structs containing arrays
    let ss: (S<[u8; 2]>,
             S<[u8; 3]>,
             S<[u8; 2]>) = (
        S(7, [8, 9]),
        S(10, [11, 12, 13]),
        S(4, [5, 6])
    );
    assert_inorder(&[
        &ss.0 as *const S<[u8]>,
        &ss.1 as *const S<[u8]>,
        &ss.2 as *const S<[u8]>
            ], compare_su8);

    assert!(simple_eq(&0u8 as *const _, &0u8 as *const _));
    assert!(!simple_eq(&0u8 as *const _, &1u8 as *const _));
}
<commit_msg>Do not generate comparefoo.gv and simpleeq.gv during testing<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(rustc_attrs)]

// check raw fat pointer ops in mir
// FIXME: please improve this when we get monomorphization support

use std::mem;

#[derive(Debug, PartialEq, Eq)]
struct ComparisonResults {
    lt: bool,
    le: bool,
    gt: bool,
    ge: bool,
    eq: bool,
    ne: bool
}

const LT: ComparisonResults = ComparisonResults {
    lt: true,
    le: true,
    gt: false,
    ge: false,
    eq: false,
    ne: true
};

const EQ: ComparisonResults = ComparisonResults {
    lt: false,
    le: true,
    gt: false,
    ge: true,
    eq: true,
    ne: false
};

const GT: ComparisonResults = ComparisonResults {
    lt: false,
    le: false,
    gt: true,
    ge: true,
    eq: false,
    ne: true
};

#[rustc_mir]
fn compare_su8(a: *const S<[u8]>, b: *const S<[u8]>) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

#[rustc_mir]
fn compare_au8(a: *const [u8], b: *const [u8]) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

#[rustc_mir]
fn compare_foo<'a>(a: *const (Foo+'a), b: *const (Foo+'a)) -> ComparisonResults {
    ComparisonResults {
        lt: a < b,
        le: a <= b,
        gt: a > b,
        ge: a >= b,
        eq: a == b,
        ne: a != b
    }
}

#[rustc_mir]
fn simple_eq<'a>(a: *const (Foo+'a), b: *const (Foo+'a)) -> bool {
    let result = a == b;
    result
}

fn assert_inorder<T: Copy>(a: &[T],
                           compare: fn(T, T) -> ComparisonResults) {
    for i in 0..a.len() {
        for j in 0..a.len() {
            let cres = compare(a[i], a[j]);
            if i < j {
                assert_eq!(cres, LT);
            } else if i == j {
                assert_eq!(cres, EQ);
            } else {
                assert_eq!(cres, GT);
            }
        }
    }
}

trait Foo { fn foo(&self) -> usize; }
impl<T> Foo for T {
    fn foo(&self) -> usize {
        mem::size_of::<T>()
    }
}

struct S<T:?Sized>(u32, T);

fn main() {
    let array = [0,1,2,3,4];
    let array2 = [5,6,7,8,9];

    // fat ptr comparison: addr then extra

    // check ordering for arrays
    let mut ptrs: Vec<*const [u8]> = vec![
        &array[0..0], &array[0..1], &array, &array[1..]
    ];

    let array_addr = &array as *const [u8] as *const u8 as usize;
    let array2_addr = &array2 as *const [u8] as *const u8 as usize;
    if array2_addr < array_addr {
        ptrs.insert(0, &array2);
    } else {
        ptrs.push(&array2);
    }
    assert_inorder(&ptrs, compare_au8);

    let u8_ = (0u8, 1u8);
    let u32_ = (4u32, 5u32);

    // check ordering for ptrs
    let buf: &mut [*const Foo] = &mut [
        &u8_, &u8_.0,
        &u32_, &u32_.0,
    ];
    buf.sort_by(|u,v| {
        let u : [*const (); 2] = unsafe { mem::transmute(*u) };
        let v : [*const (); 2] = unsafe { mem::transmute(*v) };
        u.cmp(&v)
    });
    assert_inorder(buf, compare_foo);

    // check ordering for structs containing arrays
    let ss: (S<[u8; 2]>,
             S<[u8; 3]>,
             S<[u8; 2]>) = (
        S(7, [8, 9]),
        S(10, [11, 12, 13]),
        S(4, [5, 6])
    );
    assert_inorder(&[
        &ss.0 as *const S<[u8]>,
        &ss.1 as *const S<[u8]>,
        &ss.2 as *const S<[u8]>
            ], compare_su8);

    assert!(simple_eq(&0u8 as *const _, &0u8 as *const _));
    assert!(!simple_eq(&0u8 as *const _, &1u8 as *const _));
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::linux_commands::generate_private_key_pem;
use crate::runtime::openssl_signer::OpensslSigner;
use crate::tokio_block_on;
use anyhow::{Error, Result};
use clap::Parser;
use serde::{Deserialize, Serialize};
use sxg_rs::acme::{directory::Directory as AcmeDirectory, Account as AcmeAccount};
use sxg_rs::crypto::EcPrivateKey;
use wrangler::settings::global_user::GlobalUser;
use wrangler::settings::toml::ConfigKvNamespace;

#[derive(Debug, Parser)]
pub struct Opts {
    /// A YAML file containing all config values.
    /// You can use the template
    /// 'tools/src/commands/gen_config/input.example.yaml'.
    #[clap(long, value_name = "FILE_NAME")]
    input: String,
    /// A YAML file containing the generated values.
    #[clap(long, value_name = "FILE_NAME")]
    artifact: String,
    /// No longer log in to worker service providers.
    #[clap(long)]
    use_ci_mode: bool,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Config {
    domain_name: String,
    sxg_worker: sxg_rs::config::Config,
    certificates: SxgCertConfig,
    cloudflare: CloudlareSpecificInput,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
enum SxgCertConfig {
    PreIssued {
        cert_file: String,
        issuer_file: String,
    },
    CreateAcmeAccount(AcmeConfig),
}

#[derive(Debug, Deserialize, Serialize)]
struct AcmeConfig {
    server_url: String,
    contact_email: String,
    agreed_terms_of_service: String,
    sxg_cert_request_file: String,
    eab: Option<EabConfig>,
}

#[derive(Debug, Deserialize, Serialize)]
struct EabConfig {
    base64_mac_key: String,
    key_id: String,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct CloudlareSpecificInput {
    account_id: String,
    zone_id: String,
    routes: Vec<String>,
    worker_name: String,
    deploy_on_workers_dev_only: bool,
}

#[derive(Debug, Default, Deserialize, Serialize)]
pub struct Artifact {
    acme_account: Option<AcmeAccount>,
    acme_private_key_instruction: Option<String>,
    cloudflare_kv_namespace_id: Option<String>,
}

#[derive(Default, Deserialize, Serialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
struct WranglerVars {
    html_host: String,
    sxg_config: String,
    cert_pem: Option<String>,
    issuer_pem: Option<String>,
    acme_account: Option<String>,
}

// TODO: Use `wrangler::settings::toml::Triggers`
// after [this PR](https://github.com/cloudflare/wrangler/pull/2259)
// is deployed to the latest `wrangler` version.
#[derive(Deserialize, Serialize)]
pub struct Triggers {
    pub crons: Vec<String>,
}

// TODO: Use `wrangler::settings::toml::Manifest`
// after [this issue](https://github.com/cloudflare/wrangler/issues/2037)
// is resolved.
#[derive(Deserialize, Serialize)]
struct WranglerManifest {
    name: String,
    #[serde(rename = "type")]
    target_type: String,
    account_id: String,
    zone_id: String,
    routes: Vec<String>,
    workers_dev: Option<bool>,
    kv_namespaces: Vec<ConfigKvNamespace>,
    vars: WranglerVars,
    triggers: Option<Triggers>,
}

// Set working directory to the root folder of the "sxg-rs" repository.
fn goto_repository_root() -> Result<(), std::io::Error> {
    let exe_path = std::env::current_exe()?;
    assert!(exe_path.ends_with("target/debug/tools"));
    let repo_root = exe_path
        .parent()
        .unwrap()
        .parent()
        .unwrap()
        .parent()
        .unwrap();
    std::env::set_current_dir(repo_root)?;
    Ok(())
}

// Get the Cloudflare user.
// If there is no active user, the terminal will display a login link.
// This function will wait for the login process before returning.
fn get_global_user() -> GlobalUser {
    println!("Checking Cloudflare login state");
    let mut user = GlobalUser::new();
    if user.is_err() {
        wrangler::login::run(None).unwrap();
        user = GlobalUser::new();
    }
    let user = user.unwrap();
    println!("Successfully login to Cloudflare");
    user
}

// Get the ID of the KV namespace for OCSP.
// If there is no such KV namespace, one will be created.
fn get_ocsp_kv_id(user: &GlobalUser, account_id: &str) -> String {
    let client = wrangler::http::cf_v4_client(user).unwrap();
    let target: wrangler::settings::toml::Target = Default::default();
    let namespaces = wrangler::kv::namespace::list(&client, &target).unwrap();
    if let Some(namespace) = namespaces.into_iter().find(|n| n.title == "sxg-OCSP") {
        return namespace.id;
    }
    let namespace = wrangler::kv::namespace::create(&client, account_id, "OCSP")
        .unwrap()
        .result;
    namespace.id
}

fn read_certificate_pem_file(path: &str) -> Result<String> {
    let text = std::fs::read_to_string(path)
        .map_err(|_| Error::msg(format!(r#"Failed to read file "{}""#, path)))?;
    // Translate Windows-style line endings to Unix-style so the '\r' is
    // not rendered in the toml. This is purely cosmetic; '\r' is deserialized
    // faithfully from toml and pem::parse_many is able to parse either style.
    let text = text.replace("\r\n", "\n");
    let certs = pem::parse_many(&text).map_err(Error::new)?;
    if certs.len() == 1 && certs[0].tag == "CERTIFICATE" {
        Ok(text)
    } else {
        Err(Error::msg(format!(
            r#"File "{}" is not a valid certificate PEM"#,
            path
        )))
    }
}

async fn create_acme_key_and_account(
    acme_config: &AcmeConfig,
    domain_name: &str,
) -> Result<(EcPrivateKey, AcmeAccount)> {
    let acme_private_key = {
        let pem = generate_private_key_pem()?;
        EcPrivateKey::from_sec1_pem(&pem)?
    };
    let runtime = sxg_rs::runtime::Runtime {
        acme_signer: Box::new(acme_private_key.create_signer()?),
        fetcher: Box::new(crate::runtime::hyper_fetcher::HyperFetcher::new()),
        ..Default::default()
    };
    let sxg_cert_request_der = sxg_rs::crypto::get_der_from_pem(
        &std::fs::read_to_string(&acme_config.sxg_cert_request_file)?,
        "CERTIFICATE REQUEST",
    )?;
    let eab = if let Some(input_eab) = &acme_config.eab {
        let eab_mac_key =
            base64::decode_config(&input_eab.base64_mac_key, base64::URL_SAFE_NO_PAD)?;
        let eab_signer = OpensslSigner::Hmac(&eab_mac_key);
        let new_account_url =
            AcmeDirectory::from_url(&acme_config.server_url, runtime.fetcher.as_ref())
                .await?
                .0
                .new_account;
        let output_eab = sxg_rs::acme::eab::create_external_account_binding(
            sxg_rs::acme::jws::Algorithm::HS256,
            &input_eab.key_id,
            &new_account_url,
            &acme_private_key.public_key,
            &eab_signer,
        )
        .await?;
        Some(output_eab)
    } else {
        None
    };
    let account = sxg_rs::acme::create_account(
        sxg_rs::acme::AccountSetupParams {
            directory_url: acme_config.server_url.clone(),
            agreed_terms_of_service: &acme_config.agreed_terms_of_service,
            external_account_binding: eab,
            email: &acme_config.contact_email,
            domain: domain_name.to_string(),
            public_key: acme_private_key.public_key.clone(),
            cert_request_der: sxg_cert_request_der,
        },
        runtime.fetcher.as_ref(),
        runtime.acme_signer.as_ref(),
    )
    .await?;
    Ok((acme_private_key, account))
}

fn create_wrangler_secret_instruction(name: &str, value: &str) -> String {
    let base64_value = base64::encode(value);
    format!(
        "echo {} | openssl enc -base64 -d | wrangler secret put {}",
        base64_value, name
    )
}

const WRANGLER_TOML: &str = "cloudflare_worker/wrangler.toml";

fn read_artifact(file_name: &str) -> Result<Artifact> {
    let file_content = std::fs::read_to_string(file_name)?;
    let artifact = serde_yaml::from_str(&file_content)?;
    Ok(artifact)
}

pub fn main(opts: Opts) -> Result<()> {
    if std::env::var("CI").is_ok() && !opts.use_ci_mode {
        println!("The environment variable $CI is set, but --use-ci-mode is not set.");
    }
    goto_repository_root()?;
    let mut input: Config = serde_yaml::from_str(&std::fs::read_to_string(&opts.input)?)?;
    let mut artifact: Artifact = read_artifact(&opts.artifact).unwrap_or_else(|_| {
        println!("Creating a new artifact");
        Default::default()
    });
    input.sxg_worker.html_host = input
        .sxg_worker
        .html_host
        .clone() // TODO: Remove this clone while keep Rust compiler happy.
        .or_else(|| Some(input.domain_name.clone()));
    if artifact.cloudflare_kv_namespace_id.is_none() {
        if opts.use_ci_mode {
            println!("Skipping KV namespace creation, because --use-ci-mode is set.")
        } else {
            let user = get_global_user();
            artifact.cloudflare_kv_namespace_id =
                Some(get_ocsp_kv_id(&user, &input.cloudflare.account_id))
        }
    }
    let mut wrangler_vars = WranglerVars {
        html_host: input.domain_name.clone(),
        sxg_config: serde_yaml::to_string(&input.sxg_worker)?,
        ..Default::default()
    };
    match &input.certificates {
        SxgCertConfig::PreIssued {
            cert_file,
            issuer_file,
        } => {
            wrangler_vars.cert_pem = Some(read_certificate_pem_file(cert_file)?);
            wrangler_vars.issuer_pem = Some(read_certificate_pem_file(issuer_file)?);
        }
        SxgCertConfig::CreateAcmeAccount(acme_config) => {
            if artifact.acme_account.is_none() {
                let (acme_private_key, acme_account) =
                    tokio_block_on(create_acme_key_and_account(acme_config, &input.domain_name))?;
                artifact.acme_account = Some(acme_account);
                artifact.acme_private_key_instruction = Some(create_wrangler_secret_instruction(
                    "ACME_PRIVATE_KEY_JWK",
                    &serde_json::to_string(&acme_private_key)?,
                ))
            }
            wrangler_vars.acme_account = Some(serde_json::to_string(&artifact.acme_account)?);
        }
    };
    let mut routes = input.cloudflare.routes.clone();
    routes.extend(vec![
        format!("{}/.well-known/sxg-certs/*", input.domain_name),
        format!("{}/.well-known/sxg-validity/*", input.domain_name),
        format!("{}/.well-known/acme-challenge/*", input.domain_name),
    ]);
    let wrangler_toml_output = WranglerManifest {
        name: input.cloudflare.worker_name.clone(),
        target_type: "rust".to_string(),
        account_id: input.cloudflare.account_id.clone(),
        zone_id: input.cloudflare.zone_id.clone(),
        routes,
        kv_namespaces: vec![ConfigKvNamespace {
            binding: String::from("OCSP"),
            id: artifact
                .cloudflare_kv_namespace_id
                .clone()
                .or_else(|| Some("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX".to_string())),
            preview_id: None,
        }],
        workers_dev: Some(input.cloudflare.deploy_on_workers_dev_only),
        vars: wrangler_vars,
        triggers: Some(Triggers {
            crons: vec![
                // The syntax is at https://developers.cloudflare.com/workers/platform/cron-triggers
                // This triggers at every minute.
                "* * * * *".to_string(),
            ],
        }),
    };

    std::fs::write(
        &opts.artifact,
        format!(
            "# This file is generated by command \"cargo run -p tools -- gen-config\".\n\
            # Please do not modify.\n\
            {}",
            serde_yaml::to_string(&artifact)?
        ),
    )?;

    std::fs::write(
        WRANGLER_TOML,
        format!(
            "# This file is generated by command \"cargo run -p tools -- gen-config\".\n\
            # Please note that anything you modify won't be preserved\n\
            # at the next time you run \"cargo run -p tools -- -gen-config\".\n\
            {}",
            toml::to_string_pretty(&wrangler_toml_output)?
        ),
    )?;
    println!("Successfully wrote config to {}", WRANGLER_TOML);
    Ok(())
}
<commit_msg>Fix a bug that OCSP storage name was inconsistent (#288)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::linux_commands::generate_private_key_pem;
use crate::runtime::openssl_signer::OpensslSigner;
use crate::tokio_block_on;
use anyhow::{Error, Result};
use clap::Parser;
use serde::{Deserialize, Serialize};
use sxg_rs::acme::{directory::Directory as AcmeDirectory, Account as AcmeAccount};
use sxg_rs::crypto::EcPrivateKey;
use wrangler::settings::global_user::GlobalUser;
use wrangler::settings::toml::ConfigKvNamespace;

#[derive(Debug, Parser)]
pub struct Opts {
    /// A YAML file containing all config values.
    /// You can use the template
    /// 'tools/src/commands/gen_config/input.example.yaml'.
    #[clap(long, value_name = "FILE_NAME")]
    input: String,
    /// A YAML file containing the generated values.
    #[clap(long, value_name = "FILE_NAME")]
    artifact: String,
    /// No longer log in to worker service providers.
    #[clap(long)]
    use_ci_mode: bool,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Config {
    domain_name: String,
    sxg_worker: sxg_rs::config::Config,
    certificates: SxgCertConfig,
    cloudflare: CloudlareSpecificInput,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
enum SxgCertConfig {
    PreIssued {
        cert_file: String,
        issuer_file: String,
    },
    CreateAcmeAccount(AcmeConfig),
}

#[derive(Debug, Deserialize, Serialize)]
struct AcmeConfig {
    server_url: String,
    contact_email: String,
    agreed_terms_of_service: String,
    sxg_cert_request_file: String,
    eab: Option<EabConfig>,
}

#[derive(Debug, Deserialize, Serialize)]
struct EabConfig {
    base64_mac_key: String,
    key_id: String,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct CloudlareSpecificInput {
    account_id: String,
    zone_id: String,
    routes: Vec<String>,
    worker_name: String,
    deploy_on_workers_dev_only: bool,
}

#[derive(Debug, Default, Deserialize, Serialize)]
pub struct Artifact {
    acme_account: Option<AcmeAccount>,
    acme_private_key_instruction: Option<String>,
    cloudflare_kv_namespace_id: Option<String>,
}

#[derive(Default, Deserialize, Serialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
struct WranglerVars {
    html_host: String,
    sxg_config: String,
    cert_pem: Option<String>,
    issuer_pem: Option<String>,
    acme_account: Option<String>,
}

// TODO: Use `wrangler::settings::toml::Triggers`
// after [this PR](https://github.com/cloudflare/wrangler/pull/2259)
// is deployed to the latest `wrangler` version.
#[derive(Deserialize, Serialize)]
pub struct Triggers {
    pub crons: Vec<String>,
}

// TODO: Use `wrangler::settings::toml::Manifest`
// after [this issue](https://github.com/cloudflare/wrangler/issues/2037)
// is resolved.
#[derive(Deserialize, Serialize)]
struct WranglerManifest {
    name: String,
    #[serde(rename = "type")]
    target_type: String,
    account_id: String,
    zone_id: String,
    routes: Vec<String>,
    workers_dev: Option<bool>,
    kv_namespaces: Vec<ConfigKvNamespace>,
    vars: WranglerVars,
    triggers: Option<Triggers>,
}

// Set working directory to the root folder of the "sxg-rs" repository.
fn goto_repository_root() -> Result<(), std::io::Error> {
    let exe_path = std::env::current_exe()?;
    assert!(exe_path.ends_with("target/debug/tools"));
    let repo_root = exe_path
        .parent()
        .unwrap()
        .parent()
        .unwrap()
        .parent()
        .unwrap();
    std::env::set_current_dir(repo_root)?;
    Ok(())
}

// Get the Cloudflare user.
// If there is no active user, the terminal will display a login link.
// This function will wait for the login process before returning.
fn get_global_user() -> GlobalUser {
    println!("Checking Cloudflare login state");
    let mut user = GlobalUser::new();
    if user.is_err() {
        wrangler::login::run(None).unwrap();
        user = GlobalUser::new();
    }
    let user = user.unwrap();
    println!("Successfully login to Cloudflare");
    user
}

const STORAGE_NAME: &str = "OCSP";
// Get the ID of the KV namespace for OCSP.
// If there is no such KV namespace, one will be created.
fn get_ocsp_kv_id(user: &GlobalUser, account_id: &str) -> String {
    let client = wrangler::http::cf_v4_client(user).unwrap();
    let target: wrangler::settings::toml::Target = Default::default();
    let namespaces = wrangler::kv::namespace::list(&client, &target).unwrap();
    if let Some(namespace) = namespaces.into_iter().find(|n| n.title == STORAGE_NAME) {
        return namespace.id;
    }
    let namespace = wrangler::kv::namespace::create(&client, account_id, STORAGE_NAME)
        .unwrap()
        .result;
    namespace.id
}

fn read_certificate_pem_file(path: &str) -> Result<String> {
    let text = std::fs::read_to_string(path)
        .map_err(|_| Error::msg(format!(r#"Failed to read file "{}""#, path)))?;
    // Translate Windows-style line endings to Unix-style so the '\r' is
    // not rendered in the toml. This is purely cosmetic; '\r' is deserialized
    // faithfully from toml and pem::parse_many is able to parse either style.
    let text = text.replace("\r\n", "\n");
    let certs = pem::parse_many(&text).map_err(Error::new)?;
    if certs.len() == 1 && certs[0].tag == "CERTIFICATE" {
        Ok(text)
    } else {
        Err(Error::msg(format!(
            r#"File "{}" is not a valid certificate PEM"#,
            path
        )))
    }
}

async fn create_acme_key_and_account(
    acme_config: &AcmeConfig,
    domain_name: &str,
) -> Result<(EcPrivateKey, AcmeAccount)> {
    let acme_private_key = {
        let pem = generate_private_key_pem()?;
        EcPrivateKey::from_sec1_pem(&pem)?
    };
    let runtime = sxg_rs::runtime::Runtime {
        acme_signer: Box::new(acme_private_key.create_signer()?),
        fetcher: Box::new(crate::runtime::hyper_fetcher::HyperFetcher::new()),
        ..Default::default()
    };
    let sxg_cert_request_der = sxg_rs::crypto::get_der_from_pem(
        &std::fs::read_to_string(&acme_config.sxg_cert_request_file)?,
        "CERTIFICATE REQUEST",
    )?;
    let eab = if let Some(input_eab) = &acme_config.eab {
        let eab_mac_key =
            base64::decode_config(&input_eab.base64_mac_key, base64::URL_SAFE_NO_PAD)?;
        let eab_signer = OpensslSigner::Hmac(&eab_mac_key);
        let new_account_url =
            AcmeDirectory::from_url(&acme_config.server_url, runtime.fetcher.as_ref())
                .await?
                .0
                .new_account;
        let output_eab = sxg_rs::acme::eab::create_external_account_binding(
            sxg_rs::acme::jws::Algorithm::HS256,
            &input_eab.key_id,
            &new_account_url,
            &acme_private_key.public_key,
            &eab_signer,
        )
        .await?;
        Some(output_eab)
    } else {
        None
    };
    let account = sxg_rs::acme::create_account(
        sxg_rs::acme::AccountSetupParams {
            directory_url: acme_config.server_url.clone(),
            agreed_terms_of_service: &acme_config.agreed_terms_of_service,
            external_account_binding: eab,
            email: &acme_config.contact_email,
            domain: domain_name.to_string(),
            public_key: acme_private_key.public_key.clone(),
            cert_request_der: sxg_cert_request_der,
        },
        runtime.fetcher.as_ref(),
        runtime.acme_signer.as_ref(),
    )
    .await?;
    Ok((acme_private_key, account))
}

fn create_wrangler_secret_instruction(name: &str, value: &str) -> String {
    let base64_value = base64::encode(value);
    format!(
        "echo {} | openssl enc -base64 -d | wrangler secret put {}",
        base64_value, name
    )
}

const WRANGLER_TOML: &str = "cloudflare_worker/wrangler.toml";

fn read_artifact(file_name: &str) -> Result<Artifact> {
    let file_content = std::fs::read_to_string(file_name)?;
    let artifact = serde_yaml::from_str(&file_content)?;
    Ok(artifact)
}

pub fn main(opts: Opts) -> Result<()> {
    if std::env::var("CI").is_ok() && !opts.use_ci_mode {
        println!("The environment variable $CI is set, but --use-ci-mode is not set.");
    }
    goto_repository_root()?;
    let mut input: Config = serde_yaml::from_str(&std::fs::read_to_string(&opts.input)?)?;
    let mut artifact: Artifact = read_artifact(&opts.artifact).unwrap_or_else(|_| {
        println!("Creating a new artifact");
        Default::default()
    });
    input.sxg_worker.html_host = input
        .sxg_worker
        .html_host
        .clone() // TODO: Remove this clone while keep Rust compiler happy.
        .or_else(|| Some(input.domain_name.clone()));
    if artifact.cloudflare_kv_namespace_id.is_none() {
        if opts.use_ci_mode {
            println!("Skipping KV namespace creation, because --use-ci-mode is set.")
        } else {
            let user = get_global_user();
            artifact.cloudflare_kv_namespace_id =
                Some(get_ocsp_kv_id(&user, &input.cloudflare.account_id))
        }
    }
    let mut wrangler_vars = WranglerVars {
        html_host: input.domain_name.clone(),
        sxg_config: serde_yaml::to_string(&input.sxg_worker)?,
        ..Default::default()
    };
    match &input.certificates {
        SxgCertConfig::PreIssued {
            cert_file,
            issuer_file,
        } => {
            wrangler_vars.cert_pem = Some(read_certificate_pem_file(cert_file)?);
            wrangler_vars.issuer_pem = Some(read_certificate_pem_file(issuer_file)?);
        }
        SxgCertConfig::CreateAcmeAccount(acme_config) => {
            if artifact.acme_account.is_none() {
                let (acme_private_key, acme_account) =
                    tokio_block_on(create_acme_key_and_account(acme_config, &input.domain_name))?;
                artifact.acme_account = Some(acme_account);
                artifact.acme_private_key_instruction = Some(create_wrangler_secret_instruction(
                    "ACME_PRIVATE_KEY_JWK",
                    &serde_json::to_string(&acme_private_key)?,
                ))
            }
            wrangler_vars.acme_account = Some(serde_json::to_string(&artifact.acme_account)?);
        }
    };
    let mut routes = input.cloudflare.routes.clone();
    routes.extend(vec![
        format!("{}/.well-known/sxg-certs/*", input.domain_name),
        format!("{}/.well-known/sxg-validity/*", input.domain_name),
        format!("{}/.well-known/acme-challenge/*", input.domain_name),
    ]);
    let wrangler_toml_output = WranglerManifest {
        name: input.cloudflare.worker_name.clone(),
        target_type: "rust".to_string(),
        account_id: input.cloudflare.account_id.clone(),
        zone_id: input.cloudflare.zone_id.clone(),
        routes,
        kv_namespaces: vec![ConfigKvNamespace {
            binding: STORAGE_NAME.to_string(),
            id: artifact
                .cloudflare_kv_namespace_id
                .clone()
                .or_else(|| Some("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX".to_string())),
            preview_id: None,
        }],
        workers_dev: Some(input.cloudflare.deploy_on_workers_dev_only),
        vars: wrangler_vars,
        triggers: Some(Triggers {
            crons: vec![
                // The syntax is at https://developers.cloudflare.com/workers/platform/cron-triggers
                // This triggers at every minute.
                "* * * * *".to_string(),
            ],
        }),
    };

    std::fs::write(
        &opts.artifact,
        format!(
            "# This file is generated by command \"cargo run -p tools -- gen-config\".\n\
            # Please do not modify.\n\
            {}",
            serde_yaml::to_string(&artifact)?
        ),
    )?;

    std::fs::write(
        WRANGLER_TOML,
        format!(
            "# This file is generated by command \"cargo run -p tools -- gen-config\".\n\
            # Please note that anything you modify won't be preserved\n\
            # at the next time you run \"cargo run -p tools -- -gen-config\".\n\
            {}",
            toml::to_string_pretty(&wrangler_toml_output)?
        ),
    )?;
    println!("Successfully wrote config to {}", WRANGLER_TOML);
    Ok(())
}
<|endoftext|>