<commit_before>use neat::*;

#[derive(Debug, Clone)]
pub struct Organism {
    pub genome: Genome,
    pub fitness: f64,
    pub neurons: Vec<Option<Neuron>>
}

const ACTIVATION_CYCLES: usize = 10;

impl Organism {
    pub fn new(genome: Genome) -> Organism {
        Organism { genome: genome, fitness: 0f64, neurons: vec![] }
    }

    pub fn mutate(&self) -> Organism {
        let mut new_genome = self.genome.clone();
        new_genome.mutate();
        Organism::new(new_genome)
    }

    pub fn mate(&self, other: &Organism) -> Organism {
        Organism::new(self.genome.mate(&other.genome, self.fitness < other.fitness))
    }

    pub fn activate(&mut self, sensors: &Vec<f64>, outputs: &mut Vec<f64>){
       if self.neurons.len() == 0 {
           self.generate_phenome();
       };

       for _n in 0..ACTIVATION_CYCLES {
           for neuron_id in 0..sensors.len() {
               if neuron_id < self.neurons.len() {
                   self.neurons[neuron_id].as_mut().unwrap().stimulate(sensors[neuron_id]);
               }
           }

           self.activate_neurons();
           self.propagate_signals();
       }

       //Take outputs from next neurons after sensors
       let sensors_len = sensors.len();
       for neuron_id in sensors_len..(sensors_len + outputs.len()){
           if neuron_id < self.neurons.len() {
               outputs[neuron_id - sensors_len] = self.neurons[neuron_id].as_ref().map_or(0f64, |neuron| neuron.potential());
           } else {
               outputs[neuron_id - sensors_len] = 0f64;
           }
       }

       for neuron in &mut self.neurons {
           if neuron.is_some() {
               neuron.as_mut().unwrap().relax();
           }
       }
    }

    fn activate_neurons(&mut self){
        for neuron_id in 0..self.neurons.len() {
            self.neurons[neuron_id].as_mut().unwrap().shot();
        }
    }

    fn propagate_signals(&mut self){
        for neuron_option in self.neurons.clone() {
            let neuron = neuron_option.unwrap();
            for connection in &neuron.connections {
               self.neurons[connection.0].as_mut().unwrap().stimulate(neuron.potential() * connection.1);
            }
        }
    }

    fn generate_phenome(&mut self){
        let neurons_to_generate = self.genome.len();

        self.neurons = vec![Some(Neuron::new()); neurons_to_generate as usize];

        for gene in self.genome.get_genes() {
            let mut neuron = self.neurons[gene.in_neuron_id].as_mut().unwrap();
            neuron.connections.push(Connection(gene.out_neuron_id, gene.weight));
        }
    }
}

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

    #[test]
    fn should_be_able_to_generate_his_phenome(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 0, 0.1f64);
        organism.genome.inject_gene(0, 1, 0.2f64);
        organism.genome.inject_gene(1, 0, 0.3f64);
        organism.genome.inject_gene(1, 1, 0.4f64);
        organism.generate_phenome();
        assert!(organism.neurons.len() == 2);
        let neuron0 = &organism.neurons[0].as_ref().unwrap();
        let connection00 = &neuron0.connections[0];
        assert!(connection00.0 == 0usize);
        assert!(connection00.1 == 0.1f64);
        let connection01 = &neuron0.connections[1];
        assert!(connection01.0 == 1usize);
        assert!(connection01.1 == 0.2f64);
        let neuron1 = &organism.neurons[1].as_ref().unwrap();
        let connection10 = &neuron1.connections[0];
        assert!(connection10.0 == 0usize);
        assert!(connection10.1 == 0.3f64);
        let connection11 = &neuron1.connections[1];
        assert!(connection11.0 == 1usize);
        assert!(connection11.1 == 0.4f64);
    }

    #[test]
    fn should_propagate_signal_without_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0.9f64, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -1f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] < 0.1f64, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn should_propagate_signal_over_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 0f64); //disable connection
        organism.genome.inject_gene(0, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let sensors = vec![0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0.9f64, format!("{:?} is not bigger than 0.9", output[0]));
    }

    #[test]
    fn should_work_with_cyclic_networks(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] > 0.9, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -1f64);
        organism.genome.inject_gene(1, 2, -1f64);
        organism.genome.inject_gene(2, 1, -1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] < 0.1, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn activate_organims_sensor_without_enough_neurons_should_ignore_it(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        let sensors = vec![0f64,0f64,0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
    }
}
<commit_msg>Generate strength connections matrix<commit_after>use neat::*;

#[derive(Debug, Clone)]
pub struct Organism {
    pub genome: Genome,
    pub fitness: f64,
    pub neurons: Vec<Option<Neuron>>
}

const ACTIVATION_CYCLES: usize = 10;

impl Organism {
    pub fn new(genome: Genome) -> Organism {
        Organism { genome: genome, fitness: 0f64, neurons: vec![] }
    }

    pub fn mutate(&self) -> Organism {
        let mut new_genome = self.genome.clone();
        new_genome.mutate();
        Organism::new(new_genome)
    }

    pub fn mate(&self, other: &Organism) -> Organism {
        Organism::new(self.genome.mate(&other.genome, self.fitness < other.fitness))
    }

    pub fn activate(&mut self, sensors: &Vec<f64>, outputs: &mut Vec<f64>){
       if self.neurons.len() == 0 {
           self.generate_phenome();
       };

       for _n in 0..ACTIVATION_CYCLES {
           for neuron_id in 0..sensors.len() {
               if neuron_id < self.neurons.len() {
                   self.neurons[neuron_id].as_mut().unwrap().stimulate(sensors[neuron_id]);
               }
           }

           self.activate_neurons();
           self.propagate_signals();
       }

       //Take outputs from next neurons after sensors
       let sensors_len = sensors.len();
       for neuron_id in sensors_len..(sensors_len + outputs.len()){
           if neuron_id < self.neurons.len() {
               outputs[neuron_id - sensors_len] = self.neurons[neuron_id].as_ref().map_or(0f64, |neuron| neuron.potential());
           } else {
               outputs[neuron_id - sensors_len] = 0f64;
           }
       }

       for neuron in &mut self.neurons {
           if neuron.is_some() {
               neuron.as_mut().unwrap().relax();
           }
       }
    }

    fn activate_neurons(&mut self){
        for neuron_id in 0..self.neurons.len() {
            self.neurons[neuron_id].as_mut().unwrap().shot();
        }
    }

    fn propagate_signals(&mut self){
        for neuron_option in self.neurons.clone() {
            let neuron = neuron_option.unwrap();
            for connection in &neuron.connections {
               self.neurons[connection.0].as_mut().unwrap().stimulate(neuron.potential() * connection.1);
            }
        }
    }

    fn generate_phenome(&mut self){
        let neurons_to_generate = self.genome.len();

        self.neurons = vec![Some(Neuron::new()); neurons_to_generate as usize];

        for gene in self.genome.get_genes() {
            let mut neuron = self.neurons[gene.in_neuron_id].as_mut().unwrap();
            neuron.connections.push(Connection(gene.out_neuron_id, gene.weight));
        }
    }

    fn get_weights_matrix(&self) -> Vec<Vec<f64>>{
        let neurons_len = self.genome.len();
        let mut matrix = vec![vec![0.0; neurons_len]; neurons_len];
        for gene in self.genome.get_genes() {
            if gene.enabled {
                matrix[gene.in_neuron_id][gene.out_neuron_id] = gene.weight
            }
        }
        matrix
    }
}

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

    #[test]
    fn should_be_able_to_generate_his_phenome(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 0, 0.1f64);
        organism.genome.inject_gene(0, 1, 0.2f64);
        organism.genome.inject_gene(1, 0, 0.3f64);
        organism.genome.inject_gene(1, 1, 0.4f64);
        organism.generate_phenome();
        assert!(organism.neurons.len() == 2);
        let neuron0 = &organism.neurons[0].as_ref().unwrap();
        let connection00 = &neuron0.connections[0];
        assert!(connection00.0 == 0usize);
        assert!(connection00.1 == 0.1f64);
        let connection01 = &neuron0.connections[1];
        assert!(connection01.0 == 1usize);
        assert!(connection01.1 == 0.2f64);
        let neuron1 = &organism.neurons[1].as_ref().unwrap();
        let connection10 = &neuron1.connections[0];
        assert!(connection10.0 == 0usize);
        assert!(connection10.1 == 0.3f64);
        let connection11 = &neuron1.connections[1];
        assert!(connection11.0 == 1usize);
        assert!(connection11.1 == 0.4f64);
    }

    #[test]
    fn should_propagate_signal_without_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0.9f64, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -1f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] < 0.1f64, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn should_propagate_signal_over_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 0f64); //disable connection
        organism.genome.inject_gene(0, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let sensors = vec![0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0.9f64, format!("{:?} is not bigger than 0.9", output[0]));
    }

    #[test]
    fn should_work_with_cyclic_networks(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] > 0.9, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -1f64);
        organism.genome.inject_gene(1, 2, -1f64);
        organism.genome.inject_gene(2, 1, -1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] < 0.1, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn activate_organims_sensor_without_enough_neurons_should_ignore_it(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        let sensors = vec![0f64,0f64,0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
    }

    #[test]
    fn should_be_able_to_get_matrix_representation_of_the_neuron_connections(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 0.5f64);
        organism.genome.inject_gene(2, 1, 0.5f64);
        organism.genome.inject_gene(2, 2, 0.75f64);
        organism.genome.inject_gene(1, 0, 1f64);
        assert_eq!(
            organism.get_weights_matrix(),
            [
                [0.0, 1.0, 0.0],
                [1.0, 0.0, 0.5],
                [0.0, 0.5, 0.75]
            ]
        );
    }
}
<|endoftext|>


<commit_before>use std::error::Error;

use diesel::result::Error::NotFound;

use super::data_structures::*;
use super::table_data::*;
use crate::database::InferConnection;

static RESERVED_NAMES: &[&str] = &[
    "abstract",
    "alignof",
    "as",
    "become",
    "box",
    "break",
    "const",
    "continue",
    "crate",
    "do",
    "else",
    "enum",
    "extern",
    "false",
    "final",
    "fn",
    "for",
    "if",
    "impl",
    "in",
    "let",
    "loop",
    "macro",
    "match",
    "mod",
    "move",
    "mut",
    "offsetof",
    "override",
    "priv",
    "proc",
    "pub",
    "pure",
    "ref",
    "return",
    "Self",
    "self",
    "sizeof",
    "static",
    "struct",
    "super",
    "trait",
    "true",
    "type",
    "typeof",
    "unsafe",
    "unsized",
    "use",
    "virtual",
    "where",
    "while",
    "yield",
    "bool",
    "nodes",
    "columns",
    "functions",
];

fn is_reserved_name(name: &str) -> bool {
    RESERVED_NAMES.contains(&name)
        || (
            // Names ending in an underscore are not considered reserved so that we
            // can always just append an underscore to generate an unreserved name.
            name.starts_with("__") && !name.ends_with('_')
        )
}

fn contains_unmappable_chars(name: &str) -> bool {
    // Rust identifier names are restricted to [a-zA-Z0-9_].
    !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
}

pub fn rust_name_for_sql_name(sql_name: &str) -> String {
    if is_reserved_name(sql_name) {
        format!("{}_", sql_name)
    } else if contains_unmappable_chars(sql_name) {
        // Map each non-alphanumeric character ([^a-zA-Z0-9]) to an underscore.
        let mut rust_name: String = sql_name
            .chars()
            .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
            .collect();

        // Iteratively remove adjoining underscores ("__").
        let mut last_len = rust_name.len();
        'remove_adjoining: loop {
            rust_name = rust_name.replace("__", "_");
            if rust_name.len() == last_len {
                // No more underscore pairs left.
                break 'remove_adjoining;
            }
            last_len = rust_name.len();
        }

        rust_name
    } else {
        sql_name.to_string()
    }
}

pub fn load_table_names(
    database_url: &str,
    schema_name: Option<&str>,
) -> Result<Vec<TableName>, Box<dyn Error>> {
    let connection = InferConnection::establish(database_url)?;

    match connection {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(c) => super::sqlite::load_table_names(&c, schema_name),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(c) => super::information_schema::load_table_names(&c, schema_name),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(c) => super::information_schema::load_table_names(&c, schema_name),
    }
}

fn get_column_information(
    conn: &InferConnection,
    table: &TableName,
) -> Result<Vec<ColumnInformation>, Box<dyn Error>> {
    let column_info = match *conn {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(ref c) => super::sqlite::get_table_data(c, table),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(ref c) => super::information_schema::get_table_data(c, table),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(ref c) => super::information_schema::get_table_data(c, table),
    };
    if let Err(NotFound) = column_info {
        Err(format!("no table exists named {}", table.to_string()).into())
    } else {
        column_info.map_err(Into::into)
    }
}

fn determine_column_type(
    attr: &ColumnInformation,
    conn: &InferConnection,
) -> Result<ColumnType, Box<dyn Error>> {
    match *conn {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(_) => super::sqlite::determine_column_type(attr),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(_) => super::pg::determine_column_type(attr),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(_) => super::mysql::determine_column_type(attr),
    }
}

pub(crate) fn get_primary_keys(
    conn: &InferConnection,
    table: &TableName,
) -> Result<Vec<String>, Box<dyn Error>> {
    let primary_keys: Vec<String> = match *conn {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(ref c) => super::sqlite::get_primary_keys(c, table),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(ref c) => super::information_schema::get_primary_keys(c, table),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(ref c) => super::information_schema::get_primary_keys(c, table),
    }?;
    if primary_keys.is_empty() {
        Err(format!(
            "Diesel only supports tables with primary keys. \
             Table {} has no primary key",
            table.to_string()
        )
        .into())
    } else {
        Ok(primary_keys)
    }
}

pub fn load_foreign_key_constraints(
    database_url: &str,
    schema_name: Option<&str>,
) -> Result<Vec<ForeignKeyConstraint>, Box<dyn Error>> {
    let connection = InferConnection::establish(database_url)?;

    let constraints = match connection {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(c) => super::sqlite::load_foreign_key_constraints(&c, schema_name),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(c) => {
            super::information_schema::load_foreign_key_constraints(&c, schema_name)
                .map_err(Into::into)
        }
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(c) => {
            super::mysql::load_foreign_key_constraints(&c, schema_name).map_err(Into::into)
        }
    };

    constraints.map(|mut ct| {
        ct.sort();
        ct.iter_mut().for_each(|foreign_key_constraint| {
            if is_reserved_name(&foreign_key_constraint.foreign_key_rust_name) {
                foreign_key_constraint.foreign_key_rust_name =
                    format!("{}_", foreign_key_constraint.foreign_key_rust_name);
            }
        });
        ct
    })
}

macro_rules! doc_comment {
    ($($token:tt)*) => {
        format!($($token)*)
            .lines()
            .map(str::trim)
            .collect::<Vec<_>>()
            .join("\n")
    };
}

pub fn load_table_data(database_url: &str, name: TableName) -> Result<TableData, Box<dyn Error>> {
    let connection = InferConnection::establish(database_url)?;
    let docs = doc_comment!(
        "Representation of the `{}` table.

        (Automatically generated by Diesel.)",
        name.full_sql_name(),
    );
    let primary_key = get_primary_keys(&connection, &name)?;
    let primary_key = primary_key
        .iter()
        .map(|k| rust_name_for_sql_name(&k))
        .collect();

    let column_data = get_column_information(&connection, &name)?
        .into_iter()
        .map(|c| {
            let ty = determine_column_type(&c, &connection)?;
            let rust_name = rust_name_for_sql_name(&c.column_name);

            Ok(ColumnDefinition {
                docs: doc_comment!(
                    "The `{}` column of the `{}` table.

                    Its SQL type is `{}`.

                    (Automatically generated by Diesel.)",
                    c.column_name,
                    name.full_sql_name(),
                    ty
                ),
                sql_name: c.column_name,
                ty,
                rust_name,
            })
        })
        .collect::<Result<_, Box<dyn Error>>>()?;

    Ok(TableData {
        name,
        primary_key,
        column_data,
        docs,
    })
}
<commit_msg>remove `fuctions` and `nodes`<commit_after>use std::error::Error;

use diesel::result::Error::NotFound;

use super::data_structures::*;
use super::table_data::*;
use crate::database::InferConnection;

static RESERVED_NAMES: &[&str] = &[
    "abstract", "alignof", "as", "become", "box", "break", "const", "continue", "crate", "do",
    "else", "enum", "extern", "false", "final", "fn", "for", "if", "impl", "in", "let", "loop",
    "macro", "match", "mod", "move", "mut", "offsetof", "override", "priv", "proc", "pub", "pure",
    "ref", "return", "Self", "self", "sizeof", "static", "struct", "super", "trait", "true",
    "type", "typeof", "unsafe", "unsized", "use", "virtual", "where", "while", "yield", "bool",
    "columns",
];

fn is_reserved_name(name: &str) -> bool {
    RESERVED_NAMES.contains(&name)
        || (
            // Names ending in an underscore are not considered reserved so that we
            // can always just append an underscore to generate an unreserved name.
            name.starts_with("__") && !name.ends_with('_')
        )
}

fn contains_unmappable_chars(name: &str) -> bool {
    // Rust identifier names are restricted to [a-zA-Z0-9_].
    !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
}

pub fn rust_name_for_sql_name(sql_name: &str) -> String {
    if is_reserved_name(sql_name) {
        format!("{}_", sql_name)
    } else if contains_unmappable_chars(sql_name) {
        // Map each non-alphanumeric character ([^a-zA-Z0-9]) to an underscore.
        let mut rust_name: String = sql_name
            .chars()
            .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
            .collect();

        // Iteratively remove adjoining underscores ("__").
        let mut last_len = rust_name.len();
        'remove_adjoining: loop {
            rust_name = rust_name.replace("__", "_");
            if rust_name.len() == last_len {
                // No more underscore pairs left.
                break 'remove_adjoining;
            }
            last_len = rust_name.len();
        }

        rust_name
    } else {
        sql_name.to_string()
    }
}

pub fn load_table_names(
    database_url: &str,
    schema_name: Option<&str>,
) -> Result<Vec<TableName>, Box<dyn Error>> {
    let connection = InferConnection::establish(database_url)?;

    match connection {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(c) => super::sqlite::load_table_names(&c, schema_name),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(c) => super::information_schema::load_table_names(&c, schema_name),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(c) => super::information_schema::load_table_names(&c, schema_name),
    }
}

fn get_column_information(
    conn: &InferConnection,
    table: &TableName,
) -> Result<Vec<ColumnInformation>, Box<dyn Error>> {
    let column_info = match *conn {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(ref c) => super::sqlite::get_table_data(c, table),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(ref c) => super::information_schema::get_table_data(c, table),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(ref c) => super::information_schema::get_table_data(c, table),
    };
    if let Err(NotFound) = column_info {
        Err(format!("no table exists named {}", table.to_string()).into())
    } else {
        column_info.map_err(Into::into)
    }
}

fn determine_column_type(
    attr: &ColumnInformation,
    conn: &InferConnection,
) -> Result<ColumnType, Box<dyn Error>> {
    match *conn {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(_) => super::sqlite::determine_column_type(attr),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(_) => super::pg::determine_column_type(attr),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(_) => super::mysql::determine_column_type(attr),
    }
}

pub(crate) fn get_primary_keys(
    conn: &InferConnection,
    table: &TableName,
) -> Result<Vec<String>, Box<dyn Error>> {
    let primary_keys: Vec<String> = match *conn {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(ref c) => super::sqlite::get_primary_keys(c, table),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(ref c) => super::information_schema::get_primary_keys(c, table),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(ref c) => super::information_schema::get_primary_keys(c, table),
    }?;
    if primary_keys.is_empty() {
        Err(format!(
            "Diesel only supports tables with primary keys. \
             Table {} has no primary key",
            table.to_string()
        )
        .into())
    } else {
        Ok(primary_keys)
    }
}

pub fn load_foreign_key_constraints(
    database_url: &str,
    schema_name: Option<&str>,
) -> Result<Vec<ForeignKeyConstraint>, Box<dyn Error>> {
    let connection = InferConnection::establish(database_url)?;

    let constraints = match connection {
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(c) => super::sqlite::load_foreign_key_constraints(&c, schema_name),
        #[cfg(feature = "postgres")]
        InferConnection::Pg(c) => {
            super::information_schema::load_foreign_key_constraints(&c, schema_name)
                .map_err(Into::into)
        }
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(c) => {
            super::mysql::load_foreign_key_constraints(&c, schema_name).map_err(Into::into)
        }
    };

    constraints.map(|mut ct| {
        ct.sort();
        ct.iter_mut().for_each(|foreign_key_constraint| {
            if is_reserved_name(&foreign_key_constraint.foreign_key_rust_name) {
                foreign_key_constraint.foreign_key_rust_name =
                    format!("{}_", foreign_key_constraint.foreign_key_rust_name);
            }
        });
        ct
    })
}

macro_rules! doc_comment {
    ($($token:tt)*) => {
        format!($($token)*)
            .lines()
            .map(str::trim)
            .collect::<Vec<_>>()
            .join("\n")
    };
}

pub fn load_table_data(database_url: &str, name: TableName) -> Result<TableData, Box<dyn Error>> {
    let connection = InferConnection::establish(database_url)?;
    let docs = doc_comment!(
        "Representation of the `{}` table.

        (Automatically generated by Diesel.)",
        name.full_sql_name(),
    );
    let primary_key = get_primary_keys(&connection, &name)?;
    let primary_key = primary_key
        .iter()
        .map(|k| rust_name_for_sql_name(&k))
        .collect();

    let column_data = get_column_information(&connection, &name)?
        .into_iter()
        .map(|c| {
            let ty = determine_column_type(&c, &connection)?;
            let rust_name = rust_name_for_sql_name(&c.column_name);

            Ok(ColumnDefinition {
                docs: doc_comment!(
                    "The `{}` column of the `{}` table.

                    Its SQL type is `{}`.

                    (Automatically generated by Diesel.)",
                    c.column_name,
                    name.full_sql_name(),
                    ty
                ),
                sql_name: c.column_name,
                ty,
                rust_name,
            })
        })
        .collect::<Result<_, Box<dyn Error>>>()?;

    Ok(TableData {
        name,
        primary_key,
        column_data,
        docs,
    })
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::fs::{File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
use std::sync::Arc;
use std::marker::PhantomData;

use index::storage::{Result, Storage, StorageError};
use utils::compression::{vbyte_encode, VByteDecoder};
use utils::byte_code::{ByteDecodable, ByteEncodable};

pub struct FsStorage<TItem> {
    // Stores for every id the offset in the file and the length
    entries: BTreeMap<u64, (u64 /* offset */, u32 /* length */)>,
    persistent_entries: File,
    data: File,
    current_offset: u64,
    _item_type: PhantomData<TItem>
}

impl<TItem> FsStorage<TItem> {
    /// Creates a new and empty instance of FsStorage
    pub fn new(path: &Path) -> Self {
        assert!(path.is_dir(),
                "FsStorage::new expects a directory not a file!");
        FsStorage {
            current_offset: 0,
            entries: BTreeMap::new(),
            persistent_entries: OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .open(path.join("entries.bin"))
                .unwrap(),
            data: OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .truncate(true)
                .open(path.join("data.bin"))
                .unwrap(),
            _item_type: PhantomData
        }
    }

    /// Reads a FsStorage from an previously populated folder.
    pub fn from_folder(path: &Path) -> Self {
        // Read from entry file to BTreeMap.
        let mut entries = BTreeMap::new();
        // 1. Open file
        let mut entries_file =
            OpenOptions::new().read(true).open(path.join("entries.bin")).unwrap();
        let mut bytes = Vec::with_capacity(entries_file.metadata().unwrap().len() as usize);
        // 2. Read file
        assert!(entries_file.read_to_end(&mut bytes).is_ok());
        let mut decoder = VByteDecoder::new(bytes.into_iter());
        // 3. Decode entries and write them to BTreeMap
        while let Some(entry) = decode_entry(&mut decoder) {
            entries.insert(entry.0, (entry.1, entry.2));
        }

        // Get data file length for offset
        let offset = File::open(path.join("data.bin"))
            .unwrap()
            .metadata()
            .unwrap()
            .len();

        FsStorage {
            current_offset: offset,
            entries: entries,
            persistent_entries: OpenOptions::new()
                .append(true)
                .open(path.join("entries.bin"))
                .unwrap(),
            data: OpenOptions::new()
                .read(true)
                .append(true)
                .open(path.join("data.bin"))
                .unwrap(),
            _item_type: PhantomData
        }
    }
}

impl<TItem: ByteDecodable + ByteEncodable + Sync> Storage<TItem> for FsStorage<TItem> {
    fn get(&self, id: u64) -> Result<Arc<TItem>> {
        if let Some(item_position) = self.entries.get(&id) {
            // Get filehandle
            let mut f = self.data.try_clone().unwrap();
            // Seek to position of item
            f.seek(SeekFrom::Start(item_position.0)).unwrap();
            let mut bytes = vec![0; item_position.1 as usize];
            // Read all bytes
            f.read_exact(&mut bytes).unwrap();
            // Decode item
            let item = TItem::decode(bytes.into_iter()).unwrap();
            Ok(Arc::new(item))
        } else {
            Err(StorageError::KeyNotFound)
        }
    }

    fn store(&mut self, id: u64, data: TItem) -> Result<()> {
        // Encode the data
        let bytes = data.encode();
        // Append it to the file
        if let Err(e) = self.data.write_all(&bytes) {
            return Err(StorageError::WriteError(Some(e)));
        }
        // And save the offset and the number of bytes written for later recovery
        self.entries.insert(id, (self.current_offset, bytes.len() as u32));
        // Also write the id, offset and number of bytes written to file for persistence
        let entry_bytes = encode_entry(id, self.current_offset, bytes.len() as u32);
        if let Err(e) = self.persistent_entries.write_all(&entry_bytes) {
            return Err(StorageError::WriteError(Some(e)));
        }

        // Update offset
        self.current_offset += bytes.len() as u64;
        Ok(())
    }
}

fn encode_entry(id: u64, offset: u64, length: u32) -> Vec<u8> {
    let mut bytes: Vec<u8> = Vec::new();
    bytes.append(&mut vbyte_encode(id as usize));
    bytes.append(&mut vbyte_encode(offset as usize));
    bytes.append(&mut vbyte_encode(length as usize));
    bytes
}

fn decode_entry(decoder: &mut VByteDecoder) -> Option<(u64, u64, u32)> {
    let id = try_option!(decoder.next()) as u64;
    let offset = try_option!(decoder.next()) as u64;
    let length = try_option!(decoder.next()) as u32;
    Some((id, offset, length))
}

#[cfg(test)]
mod tests {
    use std::fs::create_dir_all;
    use std::path::Path;

    use super::*;
    use index::storage::{Storage, StorageError};

    #[test]
    pub fn basic() {
        let item1 = 15;
        let item2 = 32;
        assert!(create_dir_all(Path::new("/tmp/test_index")).is_ok());
        let mut prov = FsStorage::new(Path::new("/tmp/test_index"));
        assert!(prov.store(0, item1.clone()).is_ok());
        assert_eq!(prov.get(0).unwrap().as_ref(), &item1);
        assert!(prov.store(1, item2.clone()).is_ok());
        assert_eq!(prov.get(1).unwrap().as_ref(), &item2);
        assert!(prov.get(0).unwrap().as_ref() != &item2);
        assert_eq!(prov.get(0).unwrap().as_ref(), &item1);
    }

    #[test]
    pub fn not_found() {
        let posting1 = vec![(10, vec![0, 1, 2, 3, 4]), (1, vec![15])];
        let posting2 = vec![(0, vec![0, 1, 4]), (1, vec![5, 15566, 3423565]), (5, vec![0, 24, 56])];
        assert!(create_dir_all(Path::new("/tmp/test_index")).is_ok());
        let mut prov = FsStorage::new(Path::new("/tmp/test_index"));
        assert!(prov.store(0, posting1.clone()).is_ok());
        assert!(prov.store(1, posting2.clone()).is_ok());
        assert!(if let StorageError::KeyNotFound = prov.get(2).err().unwrap() {
            true
        } else {
            false
        });
    }

    #[test]
    pub fn persistence() {
        let item1 = 1556;
        let item2 = 235425354;
        let item3 = 234543463709865987;
        assert!(create_dir_all(Path::new("/tmp/test_index2")).is_ok());
        {
            let mut prov1 = FsStorage::new(Path::new("/tmp/test_index2"));
            assert!(prov1.store(0, item1.clone()).is_ok());
            assert!(prov1.store(1, item2.clone()).is_ok());
        }

        {
            let mut prov2: FsStorage<usize> = FsStorage::from_folder(Path::new("/tmp/test_index2"));
            assert_eq!(prov2.get(0).unwrap().as_ref(), &item1);
            assert_eq!(prov2.get(1).unwrap().as_ref(), &item2);
            assert!(prov2.store(2, item3.clone()).is_ok());
            assert_eq!(prov2.get(2).unwrap().as_ref(), &item3);
        }

        {
            let prov3: FsStorage<usize> = FsStorage::from_folder(Path::new("/tmp/test_index2"));
            assert_eq!(prov3.get(0).unwrap().as_ref(), &item1);
            assert_eq!(prov3.get(1).unwrap().as_ref(), &item2);
            assert_eq!(prov3.get(2).unwrap().as_ref(), &item3);
        }
    }
}
<commit_msg>Compression of fs_storage entries<commit_after>use std::collections::BTreeMap;
use std::fs::{File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
use std::sync::Arc;
use std::marker::PhantomData;

use index::storage::{Result, Storage, StorageError};
use utils::compression::{vbyte_encode, VByteDecoder};
use utils::byte_code::{ByteDecodable, ByteEncodable};

pub struct FsStorage<TItem> {
    // Stores for every id the offset in the file and the length
    entries: BTreeMap<u64, (u64 /* offset */, u32 /* length */)>,
    persistent_entries: File,
    data: File,
    current_offset: u64,
    current_id: u64,
    _item_type: PhantomData<TItem>
}

impl<TItem> FsStorage<TItem> {
    /// Creates a new and empty instance of FsStorage
    pub fn new(path: &Path) -> Self {
        assert!(path.is_dir(),
                "FsStorage::new expects a directory not a file!");
        FsStorage {
            current_offset: 0,
            current_id: 0,
            entries: BTreeMap::new(),
            persistent_entries: OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .open(path.join("entries.bin"))
                .unwrap(),
            data: OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .truncate(true)
                .open(path.join("data.bin"))
                .unwrap(),
            _item_type: PhantomData
        }
    }

    /// Reads a FsStorage from an previously populated folder.
    pub fn from_folder(path: &Path) -> Self {
        // Read from entry file to BTreeMap.
        let mut entries = BTreeMap::new();
        // 1. Open file
        let mut entries_file =
            OpenOptions::new().read(true).open(path.join("entries.bin")).unwrap();
        let mut bytes = Vec::with_capacity(entries_file.metadata().unwrap().len() as usize);
        // 2. Read file
        assert!(entries_file.read_to_end(&mut bytes).is_ok());
        let mut decoder = VByteDecoder::new(bytes.into_iter());
        // 3. Decode entries and write them to BTreeMap
        let mut current_id: u64 = 0;
        let mut current_offset: u64 = 0;
        while let Some(entry) = decode_entry(&mut decoder) {
            current_id += entry.0 as u64;
            entries.insert(current_id, (current_offset, entry.1));
            current_offset += entry.1 as u64;
        }

        FsStorage {
            current_id: current_id,
            current_offset: current_offset,
            entries: entries,
            persistent_entries: OpenOptions::new()
                .append(true)
                .open(path.join("entries.bin"))
                .unwrap(),
            data: OpenOptions::new()
                .read(true)
                .append(true)
                .open(path.join("data.bin"))
                .unwrap(),
            _item_type: PhantomData
        }
    }
}

impl<TItem: ByteDecodable + ByteEncodable + Sync> Storage<TItem> for FsStorage<TItem> {
    fn get(&self, id: u64) -> Result<Arc<TItem>> {
        if let Some(item_position) = self.entries.get(&id) {
            // Get filehandle
            let mut f = self.data.try_clone().unwrap();
            // Seek to position of item
            f.seek(SeekFrom::Start(item_position.0)).unwrap();
            let mut bytes = vec![0; item_position.1 as usize];
            // Read all bytes
            f.read_exact(&mut bytes).unwrap();
            // Decode item
            let item = TItem::decode(bytes.into_iter()).unwrap();
            Ok(Arc::new(item))
        } else {
            Err(StorageError::KeyNotFound)
        }
    }

    fn store(&mut self, id: u64, data: TItem) -> Result<()> {

        // Encode the data
        let bytes = data.encode();
        // Append it to the file
        if let Err(e) = self.data.write_all(&bytes) {
            return Err(StorageError::WriteError(Some(e)));
        }
        // And save the offset and the number of bytes written for later recovery
        self.entries.insert(id, (self.current_offset, bytes.len() as u32));
        // Also write the id, offset and number of bytes written to file for persistence
        let entry_bytes = encode_entry(self.current_id, id, bytes.len() as u32);
        if let Err(e) = self.persistent_entries.write_all(&entry_bytes) {
            return Err(StorageError::WriteError(Some(e)));
        }

        // Update id and offset
        self.current_id = id;
        self.current_offset += bytes.len() as u64;
        Ok(())
    }
}

fn encode_entry(current_id: u64, id: u64, length: u32) -> Vec<u8> {
    let mut bytes: Vec<u8> = Vec::new();
    bytes.append(&mut vbyte_encode((id - current_id) as usize));
    bytes.append(&mut vbyte_encode(length as usize));
    bytes
}

fn decode_entry(decoder: &mut VByteDecoder) -> Option<(u32, u32)> {
    let delta_id = try_option!(decoder.next()) as u32;
    let length = try_option!(decoder.next()) as u32;

    Some((delta_id, length))
}

#[cfg(test)]
mod tests {
    use std::fs::create_dir_all;
    use std::path::Path;

    use super::*;
    use index::storage::{Storage, StorageError};

    #[test]
    pub fn basic() {
        let item1 = 15;
        let item2 = 32;
        assert!(create_dir_all(Path::new("/tmp/test_index")).is_ok());
        let mut prov = FsStorage::new(Path::new("/tmp/test_index"));
        assert!(prov.store(0, item1.clone()).is_ok());
        assert_eq!(prov.get(0).unwrap().as_ref(), &item1);
        assert!(prov.store(1, item2.clone()).is_ok());
        assert_eq!(prov.get(1).unwrap().as_ref(), &item2);
        assert!(prov.get(0).unwrap().as_ref() != &item2);
        assert_eq!(prov.get(0).unwrap().as_ref(), &item1);
    }

    #[test]
    pub fn not_found() {
        let posting1 = vec![(10, vec![0, 1, 2, 3, 4]), (1, vec![15])];
        let posting2 = vec![(0, vec![0, 1, 4]), (1, vec![5, 15566, 3423565]), (5, vec![0, 24, 56])];
        assert!(create_dir_all(Path::new("/tmp/test_index")).is_ok());
        let mut prov = FsStorage::new(Path::new("/tmp/test_index"));
        assert!(prov.store(0, posting1.clone()).is_ok());
        assert!(prov.store(1, posting2.clone()).is_ok());
        assert!(if let StorageError::KeyNotFound = prov.get(2).err().unwrap() {
            true
        } else {
            false
        });
    }

    #[test]
    pub fn persistence() {
        let item1 = 1556;
        let item2 = 235425354;
        let item3 = 234543463709865987;
        assert!(create_dir_all(Path::new("/tmp/test_index2")).is_ok());
        {
            let mut prov1 = FsStorage::new(Path::new("/tmp/test_index2"));
            assert!(prov1.store(0, item1.clone()).is_ok());
            assert!(prov1.store(1, item2.clone()).is_ok());
        }

        {
            let mut prov2: FsStorage<usize> = FsStorage::from_folder(Path::new("/tmp/test_index2"));
            assert_eq!(prov2.get(0).unwrap().as_ref(), &item1);
            assert_eq!(prov2.get(1).unwrap().as_ref(), &item2);
            assert!(prov2.store(2, item3.clone()).is_ok());
            assert_eq!(prov2.get(2).unwrap().as_ref(), &item3);
        }

        {
            let prov3: FsStorage<usize> = FsStorage::from_folder(Path::new("/tmp/test_index2"));
            assert_eq!(prov3.get(0).unwrap().as_ref(), &item1);
            assert_eq!(prov3.get(1).unwrap().as_ref(), &item2);
            assert_eq!(prov3.get(2).unwrap().as_ref(), &item3);
        }
    }
}
<|endoftext|>


<commit_before>use crate::utils::visitors::LocalUsedVisitor;
use crate::utils::{span_lint_and_then, SpanlessEq};
use if_chain::if_chain;
use rustc_hir::def::{CtorKind, CtorOf, DefKind, Res};
use rustc_hir::{Arm, Expr, ExprKind, Guard, HirId, Pat, PatKind, QPath, StmtKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty::{DefIdTree, TyCtxt};
use rustc_session::{declare_lint_pass, declare_tool_lint};
use rustc_span::{MultiSpan, Span};

declare_clippy_lint! {
    /// **What it does:** Finds nested `match` or `if let` expressions where the patterns may be "collapsed" together
    /// without adding any branches.
    ///
    /// Note that this lint is not intended to find _all_ cases where nested match patterns can be merged, but only
    /// cases where merging would most likely make the code more readable.
    ///
    /// **Why is this bad?** It is unnecessarily verbose and complex.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// ```rust
    /// fn func(opt: Option<Result<u64, String>>) {
    ///     let n = match opt {
    ///         Some(n) => match n {
    ///             Ok(n) => n,
    ///             _ => return,
    ///         }
    ///         None => return,
    ///     };
    /// }
    /// ```
    /// Use instead:
    /// ```rust
    /// fn func(opt: Option<Result<u64, String>>) {
    ///     let n = match opt {
    ///         Some(Ok(n)) => n,
    ///         _ => return,
    ///     };
    /// }
    /// ```
    pub COLLAPSIBLE_MATCH,
    style,
    "Nested `match` or `if let` expressions where the patterns may be \"collapsed\" together."
}

declare_lint_pass!(CollapsibleMatch => [COLLAPSIBLE_MATCH]);

impl<'tcx> LateLintPass<'tcx> for CollapsibleMatch {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &Expr<'tcx>) {
        if let ExprKind::Match(_expr, arms, _source) = expr.kind {
            if let Some(wild_arm) = arms.iter().rfind(|arm| arm_is_wild_like(arm, cx.tcx)) {
                for arm in arms {
                    check_arm(arm, wild_arm, cx);
                }
            }
        }
    }
}

fn check_arm(arm: &Arm<'_>, wild_outer_arm: &Arm<'_>, cx: &LateContext<'_>) {
    if_chain! {
        let expr = strip_singleton_blocks(arm.body);
        if let ExprKind::Match(expr_in, arms_inner, _) = expr.kind;
        // the outer arm pattern and the inner match
        if expr_in.span.ctxt() == arm.pat.span.ctxt();
        // there must be no more than two arms in the inner match for this lint
        if arms_inner.len() == 2;
        // no if guards on the inner match
        if arms_inner.iter().all(|arm| arm.guard.is_none());
        // match expression must be a local binding
        // match <local> { .. }
        if let ExprKind::Path(QPath::Resolved(None, path)) = expr_in.kind;
        if let Res::Local(binding_id) = path.res;
        // one of the branches must be "wild-like"
        if let Some(wild_inner_arm_idx) = arms_inner.iter().rposition(|arm_inner| arm_is_wild_like(arm_inner, cx.tcx));
        let (wild_inner_arm, non_wild_inner_arm) =
            (&arms_inner[wild_inner_arm_idx], &arms_inner[1 - wild_inner_arm_idx]);
        if !pat_contains_or(non_wild_inner_arm.pat);
        // the binding must come from the pattern of the containing match arm
        // ..<local>.. => match <local> { .. }
        if let Some(binding_span) = find_pat_binding(arm.pat, binding_id);
        // the "wild-like" branches must be equal
        if SpanlessEq::new(cx).eq_expr(wild_inner_arm.body, wild_outer_arm.body);
        // the binding must not be used in the if guard
        if !matches!(arm.guard, Some(Guard::If(guard)) if LocalUsedVisitor::new(binding_id).check_expr(guard));
        // ...or anywhere in the inner match
        if !arms_inner.iter().any(|arm| LocalUsedVisitor::new(binding_id).check_arm(arm));
        then {
            span_lint_and_then(
                cx,
                COLLAPSIBLE_MATCH,
                expr.span,
                "Unnecessary nested match",
                |diag| {
                    let mut help_span = MultiSpan::from_spans(vec![binding_span, non_wild_inner_arm.pat.span]);
                    help_span.push_span_label(binding_span, "Replace this binding".into());
                    help_span.push_span_label(non_wild_inner_arm.pat.span, "with this pattern".into());
                    diag.span_help(help_span, "The outer pattern can be modified to include the inner pattern.");
                },
            );
        }
    }
}

fn strip_singleton_blocks<'hir>(mut expr: &'hir Expr<'hir>) -> &'hir Expr<'hir> {
    while let ExprKind::Block(block, _) = expr.kind {
        match (block.stmts, block.expr) {
            ([stmt], None) => match stmt.kind {
                StmtKind::Expr(e) | StmtKind::Semi(e) => expr = e,
                _ => break,
            },
            ([], Some(e)) => expr = e,
            _ => break,
        }
    }
    expr
}

/// A "wild-like" pattern is wild ("_") or `None`.
/// For this lint to apply, both the outer and inner match expressions
/// must have "wild-like" branches that can be combined.
fn arm_is_wild_like(arm: &Arm<'_>, tcx: TyCtxt<'_>) -> bool {
    if arm.guard.is_some() {
        return false;
    }
    match arm.pat.kind {
        PatKind::Binding(..) | PatKind::Wild => true,
        PatKind::Path(QPath::Resolved(None, path)) if is_none_ctor(path.res, tcx) => true,
        _ => false,
    }
}

fn find_pat_binding(pat: &Pat<'_>, hir_id: HirId) -> Option<Span> {
    let mut span = None;
    pat.walk_short(|p| match &p.kind {
        // ignore OR patterns
        PatKind::Or(_) => false,
        PatKind::Binding(_bm, _, _ident, _) => {
            let found = p.hir_id == hir_id;
            if found {
                span = Some(p.span);
            }
            !found
        },
        _ => true,
    });
    span
}

fn pat_contains_or(pat: &Pat<'_>) -> bool {
    let mut result = false;
    pat.walk(|p| {
        let is_or = matches!(p.kind, PatKind::Or(_));
        result |= is_or;
        !is_or
    });
    result
}

fn is_none_ctor(res: Res, tcx: TyCtxt<'_>) -> bool {
    if let Some(none_id) = tcx.lang_items().option_none_variant() {
        if let Res::Def(DefKind::Ctor(CtorOf::Variant, CtorKind::Const), id) = res {
            if let Some(variant_id) = tcx.parent(id) {
                return variant_id == none_id;
            }
        }
    }
    false
}
<commit_msg>Check if let guard in collapsible_match<commit_after>use crate::utils::visitors::LocalUsedVisitor;
use crate::utils::{span_lint_and_then, SpanlessEq};
use if_chain::if_chain;
use rustc_hir::def::{CtorKind, CtorOf, DefKind, Res};
use rustc_hir::{Arm, Expr, ExprKind, Guard, HirId, Pat, PatKind, QPath, StmtKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty::{DefIdTree, TyCtxt};
use rustc_session::{declare_lint_pass, declare_tool_lint};
use rustc_span::{MultiSpan, Span};

declare_clippy_lint! {
    /// **What it does:** Finds nested `match` or `if let` expressions where the patterns may be "collapsed" together
    /// without adding any branches.
    ///
    /// Note that this lint is not intended to find _all_ cases where nested match patterns can be merged, but only
    /// cases where merging would most likely make the code more readable.
    ///
    /// **Why is this bad?** It is unnecessarily verbose and complex.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// ```rust
    /// fn func(opt: Option<Result<u64, String>>) {
    ///     let n = match opt {
    ///         Some(n) => match n {
    ///             Ok(n) => n,
    ///             _ => return,
    ///         }
    ///         None => return,
    ///     };
    /// }
    /// ```
    /// Use instead:
    /// ```rust
    /// fn func(opt: Option<Result<u64, String>>) {
    ///     let n = match opt {
    ///         Some(Ok(n)) => n,
    ///         _ => return,
    ///     };
    /// }
    /// ```
    pub COLLAPSIBLE_MATCH,
    style,
    "Nested `match` or `if let` expressions where the patterns may be \"collapsed\" together."
}

declare_lint_pass!(CollapsibleMatch => [COLLAPSIBLE_MATCH]);

impl<'tcx> LateLintPass<'tcx> for CollapsibleMatch {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &Expr<'tcx>) {
        if let ExprKind::Match(_expr, arms, _source) = expr.kind {
            if let Some(wild_arm) = arms.iter().rfind(|arm| arm_is_wild_like(arm, cx.tcx)) {
                for arm in arms {
                    check_arm(arm, wild_arm, cx);
                }
            }
        }
    }
}

fn check_arm(arm: &Arm<'_>, wild_outer_arm: &Arm<'_>, cx: &LateContext<'_>) {
    if_chain! {
        let expr = strip_singleton_blocks(arm.body);
        if let ExprKind::Match(expr_in, arms_inner, _) = expr.kind;
        // the outer arm pattern and the inner match
        if expr_in.span.ctxt() == arm.pat.span.ctxt();
        // there must be no more than two arms in the inner match for this lint
        if arms_inner.len() == 2;
        // no if guards on the inner match
        if arms_inner.iter().all(|arm| arm.guard.is_none());
        // match expression must be a local binding
        // match <local> { .. }
        if let ExprKind::Path(QPath::Resolved(None, path)) = expr_in.kind;
        if let Res::Local(binding_id) = path.res;
        // one of the branches must be "wild-like"
        if let Some(wild_inner_arm_idx) = arms_inner.iter().rposition(|arm_inner| arm_is_wild_like(arm_inner, cx.tcx));
        let (wild_inner_arm, non_wild_inner_arm) =
            (&arms_inner[wild_inner_arm_idx], &arms_inner[1 - wild_inner_arm_idx]);
        if !pat_contains_or(non_wild_inner_arm.pat);
        // the binding must come from the pattern of the containing match arm
        // ..<local>.. => match <local> { .. }
        if let Some(binding_span) = find_pat_binding(arm.pat, binding_id);
        // the "wild-like" branches must be equal
        if SpanlessEq::new(cx).eq_expr(wild_inner_arm.body, wild_outer_arm.body);
        // the binding must not be used in the if guard
        if match arm.guard {
            None => true,
            Some(Guard::If(expr) | Guard::IfLet(_, expr)) => {
                !LocalUsedVisitor::new(binding_id).check_expr(expr)
            }
        };
        // ...or anywhere in the inner match
        if !arms_inner.iter().any(|arm| LocalUsedVisitor::new(binding_id).check_arm(arm));
        then {
            span_lint_and_then(
                cx,
                COLLAPSIBLE_MATCH,
                expr.span,
                "Unnecessary nested match",
                |diag| {
                    let mut help_span = MultiSpan::from_spans(vec![binding_span, non_wild_inner_arm.pat.span]);
                    help_span.push_span_label(binding_span, "Replace this binding".into());
                    help_span.push_span_label(non_wild_inner_arm.pat.span, "with this pattern".into());
                    diag.span_help(help_span, "The outer pattern can be modified to include the inner pattern.");
                },
            );
        }
    }
}

fn strip_singleton_blocks<'hir>(mut expr: &'hir Expr<'hir>) -> &'hir Expr<'hir> {
    while let ExprKind::Block(block, _) = expr.kind {
        match (block.stmts, block.expr) {
            ([stmt], None) => match stmt.kind {
                StmtKind::Expr(e) | StmtKind::Semi(e) => expr = e,
                _ => break,
            },
            ([], Some(e)) => expr = e,
            _ => break,
        }
    }
    expr
}

/// A "wild-like" pattern is wild ("_") or `None`.
/// For this lint to apply, both the outer and inner match expressions
/// must have "wild-like" branches that can be combined.
fn arm_is_wild_like(arm: &Arm<'_>, tcx: TyCtxt<'_>) -> bool {
    if arm.guard.is_some() {
        return false;
    }
    match arm.pat.kind {
        PatKind::Binding(..) | PatKind::Wild => true,
        PatKind::Path(QPath::Resolved(None, path)) if is_none_ctor(path.res, tcx) => true,
        _ => false,
    }
}

fn find_pat_binding(pat: &Pat<'_>, hir_id: HirId) -> Option<Span> {
    let mut span = None;
    pat.walk_short(|p| match &p.kind {
        // ignore OR patterns
        PatKind::Or(_) => false,
        PatKind::Binding(_bm, _, _ident, _) => {
            let found = p.hir_id == hir_id;
            if found {
                span = Some(p.span);
            }
            !found
        },
        _ => true,
    });
    span
}

fn pat_contains_or(pat: &Pat<'_>) -> bool {
    let mut result = false;
    pat.walk(|p| {
        let is_or = matches!(p.kind, PatKind::Or(_));
        result |= is_or;
        !is_or
    });
    result
}

fn is_none_ctor(res: Res, tcx: TyCtxt<'_>) -> bool {
    if let Some(none_id) = tcx.lang_items().option_none_variant() {
        if let Res::Def(DefKind::Ctor(CtorOf::Variant, CtorKind::Const), id) = res {
            if let Some(variant_id) = tcx.parent(id) {
                return variant_id == none_id;
            }
        }
    }
    false
}
<|endoftext|>


<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   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
//
//       http://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.
//

//!
//! # Matcher
//!
//! The Matcher trait is implemented by objects that can match patterns against the left-hand side of a stream. It's a fairly
//! baseline implementation: it's up to the caller to implement things like rewinding in order to perform tokenisation. That is,
//! Matchers are greedy and may (indeed, are likely to) consume more characters than the longest match while trying to find
//! a longer one.
//!

///
/// Matcher that can read an input stream of type `Symbol` and find the longest matching pattern, which it will identify using
/// `OutputSymbol`
///
pub trait PatternMatcher<'a, InputSymbol, OutputSymbol> {
    type State: MatchingState<'a, InputSymbol, OutputSymbol>;

    ///
    /// Creates a state that begins matching this pattern
    ///
    fn start(&'a self) -> Self::State;
}

///
/// Action to be taken after a matcher receives a symbol
///
pub enum MatchAction<'a, OutputSymbol: ?Sized+'a, State: ?Sized> {
    // State is also always: MatchingState<InputSymbol, OutputSymbol> (important to know that as its how More is used)
    //
    // However, rust complains that InputSymbol is unused if we declare it in MatchAction and that it is undeclared if we don't
    // If this wasn't an enum, could make it an associated type? But this is an enum so we can't do that.
    // Ie, MatchAction<InputSymbol, OutputSymbol, State: MatchingState<InputSymbol, OutputSymbol>> = Nope (InputSymbol is unused apparently even though it's required which seems to me to mean it's used)
    // and also MatchAction<OutputSymbol, State: MatchingState<DontCare, OutputSymbol> = Nope (DontCare is not declared)

    /// The pattern does not match
    Reject,

    /// The pattern matched a certain number of symbols (which may be fewer than were passed to the matcher)
    Accept(usize, &'a OutputSymbol),

    /// The matcher needs more symbols to decide if the pattern matches
    More(State)
}

///
/// Represents a state during a pattern matching operation
///
pub trait MatchingState<'a, InputSymbol, OutputSymbol: ?Sized> {
    ///
    /// Matches the next symbol
    ///
    fn next(self, symbol: InputSymbol) -> MatchAction<'a, OutputSymbol, Self>;

    ///
    /// There are no more symbols available (this can only return `Reject` or `Accept`)
    ///
    fn finish(self) -> MatchAction<'a, OutputSymbol, Self>;
}

pub use MatchAction::*;
<commit_msg>Output symbol also doesn’t need to be sized here any more either<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   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
//
//       http://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.
//

//!
//! # Matcher
//!
//! The Matcher trait is implemented by objects that can match patterns against the left-hand side of a stream. It's a fairly
//! baseline implementation: it's up to the caller to implement things like rewinding in order to perform tokenisation. That is,
//! Matchers are greedy and may (indeed, are likely to) consume more characters than the longest match while trying to find
//! a longer one.
//!

///
/// Matcher that can read an input stream of type `Symbol` and find the longest matching pattern, which it will identify using
/// `OutputSymbol`
///
pub trait PatternMatcher<'a, InputSymbol, OutputSymbol> {
    type State: MatchingState<'a, InputSymbol, OutputSymbol>;

    ///
    /// Creates a state that begins matching this pattern
    ///
    fn start(&'a self) -> Self::State;
}

///
/// Action to be taken after a matcher receives a symbol
///
pub enum MatchAction<'a, OutputSymbol: 'a, State: ?Sized> {
    // State is also always: MatchingState<InputSymbol, OutputSymbol> (important to know that as its how More is used)
    //
    // However, rust complains that InputSymbol is unused if we declare it in MatchAction and that it is undeclared if we don't
    // If this wasn't an enum, could make it an associated type? But this is an enum so we can't do that.
    // Ie, MatchAction<InputSymbol, OutputSymbol, State: MatchingState<InputSymbol, OutputSymbol>> = Nope (InputSymbol is unused apparently even though it's required which seems to me to mean it's used)
    // and also MatchAction<OutputSymbol, State: MatchingState<DontCare, OutputSymbol> = Nope (DontCare is not declared)

    /// The pattern does not match
    Reject,

    /// The pattern matched a certain number of symbols (which may be fewer than were passed to the matcher)
    Accept(usize, &'a OutputSymbol),

    /// The matcher needs more symbols to decide if the pattern matches
    More(State)
}

///
/// Represents a state during a pattern matching operation
///
pub trait MatchingState<'a, InputSymbol, OutputSymbol> {
    ///
    /// Matches the next symbol
    ///
    fn next(self, symbol: InputSymbol) -> MatchAction<'a, OutputSymbol, Self>;

    ///
    /// There are no more symbols available (this can only return `Reject` or `Accept`)
    ///
    fn finish(self) -> MatchAction<'a, OutputSymbol, Self>;
}

pub use MatchAction::*;
<|endoftext|>


<commit_before>use super::bits::BitField;
use super::bits::U16Field;
use super::packet::Block;
use protocol::packet::Packet;
use protocol::packet::Piece;
use std::fmt;

const ID:U16Field = U16Field{index:0};

const QR:BitField = BitField{index:2, mask: 0b1000_0000u8};
const OP:BitField = BitField{index:2, mask: 0b0111_1000u8};
const AA:BitField = BitField{index:2, mask: 0b0000_0100u8};
const TC:BitField = BitField{index:2, mask: 0b0000_0010u8};
const RD:BitField = BitField{index:2, mask: 0b0000_0001u8};

const RA:BitField = BitField{index:3, mask: 0b1000_0000u8};
const RC:BitField = BitField{index:3, mask: 0b0000_1111u8};

const QD:U16Field = U16Field{index:4};
const AN:U16Field = U16Field{index:6};
const NS:U16Field = U16Field{index:8};
const AR:U16Field = U16Field{index:10};

const SIZE:usize = 12;

pub const RC_OK:u8 = 0;
pub const RC_FORMAT_ERROR:u8 = 1;
pub const RC_SERVER_ERROR:u8 = 2;
pub const RC_NAME_ERROR:u8 = 3;
pub const RC_NOT_IMPLEMENTED:u8 = 4;
pub const RC_REFUSED:u8 = 5;

pub const OP_QUERY:u8 = 0;
pub const OP_IQUERY:u8 = 1;
pub const OP_STATUS:u8 = 2;

#[derive(Copy,Clone)]
pub struct Header<'d> {
    data: Piece<'d>
}

impl<'d> Header<'d> {

    pub fn id(&self) -> Option<u16>  {ID.get(self.data.data())}
    pub fn qr(&self) -> Option<bool> {QR.nonzero(self.data.data())}
    pub fn op(&self) -> Option<u8>   {OP.get(self.data.data())}
    pub fn aa(&self) -> Option<bool> {AA.nonzero(self.data.data())}
    pub fn tc(&self) -> Option<bool> {TC.nonzero(self.data.data())}
    pub fn rd(&self) -> Option<bool> {RD.nonzero(self.data.data())}
    pub fn ra(&self) -> Option<bool> {RA.nonzero(self.data.data())}
    pub fn rc(&self) -> Option<u8>   {RC.get(self.data.data())}
    pub fn qd(&self) -> Option<u16>  {QD.get(self.data.data())}
    pub fn an(&self) -> Option<u16>  {AN.get(self.data.data())}
    pub fn ns(&self) -> Option<u16>  {NS.get(self.data.data())}
    pub fn ar(&self) -> Option<u16>  {AR.get(self.data.data())}

    pub fn is_query(&self) -> bool { self.qr() == Some(false) }
    pub fn is_response(&self) -> bool { self.qr() == Some(true) }
    pub fn is_truncated(&self) -> bool { self.tc() == Some(true) }

    // TODO questions iterator
    // TODO answers iterator
    // TODO nameservers iterator
    // TODO additional records iterator
}

impl<'d> fmt::Debug for Header<'d> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Header")
            .field("id", &self.id())
            .field("qr", &self.qr())
            .field("op", &self.op())
            .field("aa", &self.aa())
            .field("tc", &self.tc())
            .field("rd", &self.rd())
            .field("rc", &self.rc())
            .field("qd", &self.qd())
            .field("an", &self.an())
            .field("ns", &self.ns())
            .field("ar", &self.ar())
            .field("is_query", &self.is_query())
            .field("is_response", &self.is_response())
            .field("is_truncated", &self.is_truncated())
            .finish()
    }

}

impl<'d> Block<'d, Header<'d>> for Header<'d> {
    fn at(src: &mut Packet<'d>, at:usize) -> Header<'d> {
        Header{data: src.next_slice(SIZE)}
    }
}

#[cfg(test)]
mod tests {
    use super::Header;
    use protocol::packet::Packet;

    #[test]
    fn id() {
        let data = [0xab, 0xcd];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0xabcdu16), h.id());

        let data = [0xab];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.id());
    }

    #[test]
    fn query_or_response() {
        let data = [0,0,0x00];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(false), h.qr());
        assert_eq!(true, h.is_query());
        assert_eq!(false, h.is_response());

        let data = [0,0,0x80];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(true), h.qr());
        assert_eq!(false, h.is_query());
        assert_eq!(true, h.is_response());

        let data = [0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.qr());
        assert_eq!(false, h.is_query());
        assert_eq!(false, h.is_response());
    }

    #[test]
    fn operation() {
        let data = [0,0,0x78];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(15), h.op());

        let data = [0,0,0x18];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(3), h.op());

        let data = [0,0,0x10];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(Header::OP_STATUS), h.op());

        let data = [0,0,0x08];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(Header::OP_IQUERY), h.op());

        let data = [0,0,0x87];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(Header::OP_QUERY), h.op());
    }

    #[test]
    fn authoritative() {
        let data = [0,0,0x04];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(true), h.aa());

        let data = [0,0,0xfb];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(false), h.aa());

        let data = [0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.aa());
    }

    #[test]
    fn truncated() {
        let data = [0,0,0x02];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(true), h.tc());

        let data = [0,0,0xfc];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(false), h.tc());

        let data = [0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.tc());
    }

    #[test]
    fn please_recurse() {
        let data = [0,0,0x01];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(true), h.rd());

        let data = [0,0,0xfe];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(false), h.rd());

        let data = [0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.rd());
    }

    #[test]
    fn recursion_available() {
        let data = [0,0,0,0x80];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(true), h.ra());

        let data = [0,0,0,0x7f];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(false), h.ra());

        let data = [0,0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.ra());
    }

    #[test]
    fn response_code() {
        let data = [0,0,0,0x00];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0), h.rc());

        let data = [0,0,0,0x0f];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0xf), h.rc());

        let data = [0,0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.rc());
    }

    #[test]
    fn query_count() {
        let data = [0,0,0,0,0xab,0xcd];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0xabcd), h.qd());

        let data = [0,0,0,0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.qd());
    }

    #[test]
    fn answer_count() {
        let data = [0,0,0,0,0,0,0xab,0xcd];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0xabcd), h.an());

        let data = [0,0,0,0,0,0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.an());
    }

    #[test]
    fn name_response_count() {
        let data = [0,0,0,0,0,0,0,0,0xab,0xcd];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0xabcd), h.ns());

        let data = [0,0,0,0,0,0,0,0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.ns());
    }

    #[test]
    fn additional_record_count() {
        let data = [0,0,0,0,0,0,0,0,0,0,0xab,0xcd];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(Some(0xabcd), h.ar());

        let data = [0,0,0,0,0,0,0,0,0,0,0];
        let mut p = Packet::new(&data);
        let h = p.next::<Header>();
        assert_eq!(None, h.ar());
    }

}
<commit_msg>Fix up Header for new bitfield.<commit_after>use super::bits::BitField;
use super::bits::BitData;
use super::bits::MutBitData;
use super::bits::BEU16Field;
use std::ops::Deref;
use std::fmt;

const ID:BEU16Field = BEU16Field{index:0};

const QR:BitField = BitField{index:2, mask: 0b1000_0000u8};
const OP:BitField = BitField{index:2, mask: 0b0111_1000u8};
const AA:BitField = BitField{index:2, mask: 0b0000_0100u8};
const TC:BitField = BitField{index:2, mask: 0b0000_0010u8};
const RD:BitField = BitField{index:2, mask: 0b0000_0001u8};

const RA:BitField = BitField{index:3, mask: 0b1000_0000u8};
const RC:BitField = BitField{index:3, mask: 0b0000_1111u8};

const QD:BEU16Field = BEU16Field{index:4};
const AN:BEU16Field = BEU16Field{index:6};
const NS:BEU16Field = BEU16Field{index:8};
const AR:BEU16Field = BEU16Field{index:10};

const SIZE:usize = 12;

pub const RC_OK:u8 = 0;
pub const RC_FORMAT_ERROR:u8 = 1;
pub const RC_SERVER_ERROR:u8 = 2;
pub const RC_NAME_ERROR:u8 = 3;
pub const RC_NOT_IMPLEMENTED:u8 = 4;
pub const RC_REFUSED:u8 = 5;

pub const OP_QUERY:u8 = 0;
pub const OP_IQUERY:u8 = 1;
pub const OP_STATUS:u8 = 2;

#[derive(Copy,Clone)]
pub struct Header<'d, D: 'd + ?Sized> {
    data: &'d D
}

pub struct MutHeader<'d, D: 'd + ?Sized> {
    data: &'d mut D
}

impl<'d, D: 'd + ?Sized> Header<'d, D> where D: BitData {

    pub fn at(data: &'d D) -> Header<'d, D> {
        Header{data: data}
    }

    pub fn id(&self) -> Option<u16>  {ID.get(self.data)}
    pub fn qr(&self) -> Option<bool> {QR.nonzero(self.data)}
    pub fn op(&self) -> Option<u8>   {OP.get(self.data)}
    pub fn aa(&self) -> Option<bool> {AA.nonzero(self.data)}
    pub fn tc(&self) -> Option<bool> {TC.nonzero(self.data)}
    pub fn rd(&self) -> Option<bool> {RD.nonzero(self.data)}
    pub fn ra(&self) -> Option<bool> {RA.nonzero(self.data)}
    pub fn rc(&self) -> Option<u8>   {RC.get(self.data)}
    pub fn qd(&self) -> Option<u16>  {QD.get(self.data)}
    pub fn an(&self) -> Option<u16>  {AN.get(self.data)}
    pub fn ns(&self) -> Option<u16>  {NS.get(self.data)}
    pub fn ar(&self) -> Option<u16>  {AR.get(self.data)}

    pub fn is_query(&self) -> bool { self.qr() == Some(false) }
    pub fn is_response(&self) -> bool { self.qr() == Some(true) }
    pub fn is_truncated(&self) -> bool { self.tc() == Some(true) }

    // TODO questions iterator
    // TODO answers iterator
    // TODO nameservers iterator
    // TODO additional records iterator
}

impl<'d, D: 'd + ?Sized> MutHeader<'d, D> where D: MutBitData {
    pub fn at_mut(data: &'d mut D) -> MutHeader<'d, D> {
        MutHeader{data: data}
    }
    pub fn set_id(&mut self, val:u16) -> &mut Self {ID.set(self.data, val); self}
    pub fn set_qr(&mut self, val:bool) -> &mut Self {QR.set(self.data, val as u8); self}
    pub fn set_op(&mut self, val:u8) -> &mut Self {OP.set(self.data, val); self}
    pub fn set_aa(&mut self, val:bool) -> &mut Self {AA.set(self.data, val as u8); self}
    pub fn set_tc(&mut self, val:bool) -> &mut Self {TC.set(self.data, val as u8); self}
    pub fn set_rd(&mut self, val:bool) -> &mut Self {RD.set(self.data, val as u8); self}
    pub fn set_ra(&mut self, val:bool) -> &mut Self {RA.set(self.data, val as u8); self}
    pub fn set_rc(&mut self, val:u8) -> &mut Self {RC.set(self.data, val); self}
    pub fn set_qd(&mut self, val:u16) -> &mut Self {QD.set(self.data, val); self}
    pub fn set_an(&mut self, val:u16) -> &mut Self {AN.set(self.data, val); self}
    pub fn set_ns(&mut self, val:u16) -> &mut Self {NS.set(self.data, val); self}
    pub fn set_ar(&mut self, val:u16) -> &mut Self {AR.set(self.data, val); self}
}

impl<'d, D: 'd> fmt::Debug for Header<'d, D> where D: BitData {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Header")
            .field("id", &self.id())
            .field("qr", &self.qr())
            .field("op", &self.op())
            .field("aa", &self.aa())
            .field("tc", &self.tc())
            .field("rd", &self.rd())
            .field("rc", &self.rc())
            .field("qd", &self.qd())
            .field("an", &self.an())
            .field("ns", &self.ns())
            .field("ar", &self.ar())
            .field("is_query", &self.is_query())
            .field("is_response", &self.is_response())
            .field("is_truncated", &self.is_truncated())
            .finish()
    }
}

impl<'d, D: 'd> MutHeader<'d, D> where D: MutBitData {
    fn as_header(&'d self) -> Header<'d, D> {
        let readonly: &D = &self.data;
        Header{data: readonly}
    }
}

impl<'d, D: 'd> fmt::Debug for MutHeader<'d, D> where D: MutBitData {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let h = self.as_header();
        fmt.debug_struct("MutHeader")
            .field("id", &h.id())
            .field("qr", &h.qr())
            .field("op", &h.op())
            .field("aa", &h.aa())
            .field("tc", &h.tc())
            .field("rd", &h.rd())
            .field("rc", &h.rc())
            .field("qd", &h.qd())
            .field("an", &h.an())
            .field("ns", &h.ns())
            .field("ar", &h.ar())
            .field("is_query", &h.is_query())
            .field("is_response", &h.is_response())
            .field("is_truncated", &h.is_truncated())
            .finish()
    }
}

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

    #[test]
    fn id() {
        let data : &[u8] = &[0xab, 0xcd];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0xabcdu16), h.id());

        let data : &[u8] = &[0xab];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.id());
    }

    #[test]
    fn query_or_response() {
        let data : &[u8] = &[0,0,0x00];
        let h = Header::at(&data[..]);
        assert_eq!(Some(false), h.qr());
        assert_eq!(true, h.is_query());
        assert_eq!(false, h.is_response());

        let data : &[u8] = &[0,0,0x80];
        let h = Header::at(&data[..]);
        assert_eq!(Some(true), h.qr());
        assert_eq!(false, h.is_query());
        assert_eq!(true, h.is_response());

        let data : &[u8] = &[0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.qr());
        assert_eq!(false, h.is_query());
        assert_eq!(false, h.is_response());
    }

    #[test]
    fn operation() {
        let data : &[u8] = &[0,0,0x78];
        let h = Header::at(&data[..]);
        assert_eq!(Some(15), h.op());

        let data : &[u8] = &[0,0,0x18];
        let h = Header::at(&data[..]);
        assert_eq!(Some(3), h.op());

        let data : &[u8] = &[0,0,0x10];
        let h = Header::at(&data[..]);
        assert_eq!(Some(OP_STATUS), h.op());

        let data : &[u8] = &[0,0,0x08];
        let h = Header::at(&data[..]);
        assert_eq!(Some(OP_IQUERY), h.op());

        let data : &[u8] = &[0,0,0x87];
        let h = Header::at(&data[..]);
        assert_eq!(Some(OP_QUERY), h.op());
    }

    #[test]
    fn authoritative() {
        let data : &[u8] = &[0,0,0x04];
        let h = Header::at(&data[..]);
        assert_eq!(Some(true), h.aa());

        let data : &[u8] = &[0,0,0xfb];
        let h = Header::at(&data[..]);
        assert_eq!(Some(false), h.aa());

        let data : &[u8] = &[0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.aa());
    }

    #[test]
    fn truncated() {
        let data : &[u8] = &[0,0,0x02];
        let h = Header::at(&data[..]);
        assert_eq!(Some(true), h.tc());

        let data : &[u8] = &[0,0,0xfc];
        let h = Header::at(&data[..]);
        assert_eq!(Some(false), h.tc());

        let data : &[u8] = &[0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.tc());
    }

    #[test]
    fn please_recurse() {
        let data : &[u8] = &[0,0,0x01];
        let h = Header::at(&data[..]);
        assert_eq!(Some(true), h.rd());

        let data : &[u8] = &[0,0,0xfe];
        let h = Header::at(&data[..]);
        assert_eq!(Some(false), h.rd());

        let data : &[u8] = &[0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.rd());
    }

    #[test]
    fn recursion_available() {
        let data : &[u8] = &[0,0,0,0x80];
        let h = Header::at(&data[..]);
        assert_eq!(Some(true), h.ra());

        let data : &[u8] = &[0,0,0,0x7f];
        let h = Header::at(&data[..]);
        assert_eq!(Some(false), h.ra());

        let data : &[u8] = &[0,0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.ra());
    }

    #[test]
    fn response_code() {
        let data : &[u8] = &[0,0,0,0x00];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0), h.rc());

        let data : &[u8] = &[0,0,0,0x0f];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0xf), h.rc());

        let data : &[u8] = &[0,0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.rc());
    }

    #[test]
    fn query_count() {
        let data : &[u8] = &[0,0,0,0,0xab,0xcd];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0xabcd), h.qd());

        let data : &[u8] = &[0,0,0,0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.qd());
    }

    #[test]
    fn answer_count() {
        let data : &[u8] = &[0,0,0,0,0,0,0xab,0xcd];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0xabcd), h.an());

        let data : &[u8] = &[0,0,0,0,0,0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.an());
    }

    #[test]
    fn name_response_count() {
        let data : &[u8] = &[0,0,0,0,0,0,0,0,0xab,0xcd];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0xabcd), h.ns());

        let data : &[u8] = &[0,0,0,0,0,0,0,0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.ns());
    }

    #[test]
    fn additional_record_count() {
        let data : &[u8] = &[0,0,0,0,0,0,0,0,0,0,0xab,0xcd];
        let h = Header::at(&data[..]);
        assert_eq!(Some(0xabcd), h.ar());

        let data : &[u8] = &[0,0,0,0,0,0,0,0,0,0,0];
        let h = Header::at(&data[..]);
        assert_eq!(None, h.ar());
    }

}

<|endoftext|>