<commit_before>use serde::Deserialize;
use std::collections::BTreeMap;
use std::path::Path;

use std::collections::BTreeSet as Set;
use std::fs;

/// A lookup table from TypeName to Entry.
pub type NodeTypeMap = BTreeMap<TypeName, Entry>;

#[derive(Debug)]
pub struct Entry {
    pub dbscheme_name: String,
    pub ql_class_name: String,
    pub kind: EntryKind,
}

#[derive(Debug)]
pub enum EntryKind {
    Union { members: Set<TypeName> },
    Table { name: String, fields: Vec<Field> },
    Token { kind_id: usize },
}

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
pub struct TypeName {
    pub kind: String,
    pub named: bool,
}

#[derive(Debug)]
pub enum FieldTypeInfo {
    /// The field has a single type.
    Single(TypeName),

    /// The field can take one of several types, so we also provide the name of
    /// the database union type that wraps them, and the corresponding QL class
    /// name.
    Multiple {
        types: Set<TypeName>,
        dbscheme_union: String,
        ql_class: String,
    },

    /// The field can be one of several tokens, so the db type will be an `int`
    /// with a `case @foo.kind` for each possiblity.
    ReservedWordInt(BTreeMap<String, (usize, String)>),
}

#[derive(Debug)]
pub struct Field {
    pub parent: TypeName,
    pub type_info: FieldTypeInfo,
    /// The name of the field or None for the anonymous 'children'
    /// entry from node_types.json
    pub name: Option<String>,
    /// The name of the predicate to get this field.
    pub getter_name: String,
    pub storage: Storage,
}

fn name_for_field_or_child(name: &Option<String>) -> String {
    match name {
        Some(name) => name.clone(),
        None => "child".to_owned(),
    }
}

#[derive(Debug)]
pub enum Storage {
    /// the field is stored as a column in the parent table
    Column { name: String },
    /// the field is stored in a link table
    Table {
        /// the name of the table
        name: String,
        /// the name of the column for the field in the dbscheme
        column_name: String,
        /// does it have an associated index column?
        has_index: bool,
    },
}

pub fn read_node_types(prefix: &str, node_types_path: &Path) -> std::io::Result<NodeTypeMap> {
    let file = fs::File::open(node_types_path)?;
    let node_types = serde_json::from_reader(file)?;
    Ok(convert_nodes(&prefix, &node_types))
}

pub fn read_node_types_str(prefix: &str, node_types_json: &str) -> std::io::Result<NodeTypeMap> {
    let node_types = serde_json::from_str(node_types_json)?;
    Ok(convert_nodes(&prefix, &node_types))
}

fn convert_type(node_type: &NodeType) -> TypeName {
    TypeName {
        kind: node_type.kind.to_string(),
        named: node_type.named,
    }
}

fn convert_types(node_types: &Vec<NodeType>) -> Set<TypeName> {
    let iter = node_types.iter().map(convert_type).collect();
    std::collections::BTreeSet::from(iter)
}

pub fn convert_nodes(prefix: &str, nodes: &Vec<NodeInfo>) -> NodeTypeMap {
    let mut entries = NodeTypeMap::new();
    let mut token_kinds = Set::new();

    // First, find all the token kinds
    for node in nodes {
        if node.subtypes.is_none() {
            if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
                let type_name = TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                };
                token_kinds.insert(type_name);
            }
        }
    }

    for node in nodes {
        let flattened_name = &node_type_name(&node.kind, node.named);
        let dbscheme_name = escape_name(&flattened_name);
        let ql_class_name = dbscheme_name_to_class_name(&dbscheme_name);
        let dbscheme_name = format!("{}_{}", prefix, &dbscheme_name);
        if let Some(subtypes) = &node.subtypes {
            // It's a tree-sitter supertype node, for which we create a union
            // type.
            entries.insert(
                TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                },
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Union {
                        members: convert_types(&subtypes),
                    },
                },
            );
        } else if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
            // Token kind, handled above.
        } else {
            // It's a product type, defined by a table.
            let type_name = TypeName {
                kind: node.kind.clone(),
                named: node.named,
            };
            let table_name = escape_name(&(format!("{}_def", &flattened_name)));
            let table_name = format!("{}_{}", prefix, &table_name);

            let mut fields = Vec::new();

            // If the type also has fields or children, then we create either
            // auxiliary tables or columns in the defining table for them.
            if let Some(node_fields) = &node.fields {
                for (field_name, field_info) in node_fields {
                    add_field(
                        &prefix,
                        &type_name,
                        Some(field_name.to_string()),
                        field_info,
                        &mut fields,
                        &token_kinds,
                    );
                }
            }
            if let Some(children) = &node.children {
                // Treat children as if they were a field called 'child'.
                add_field(
                    &prefix,
                    &type_name,
                    None,
                    children,
                    &mut fields,
                    &token_kinds,
                );
            }
            entries.insert(
                type_name,
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Table {
                        name: table_name,
                        fields,
                    },
                },
            );
        }
    }
    let mut counter = 0;
    for type_name in token_kinds {
        let entry = if type_name.named {
            counter += 1;
            let unprefixed_name = node_type_name(&type_name.kind, true);
            Entry {
                dbscheme_name: escape_name(&format!("{}_token_{}", &prefix, &unprefixed_name)),
                ql_class_name: dbscheme_name_to_class_name(&escape_name(&unprefixed_name)),
                kind: EntryKind::Token { kind_id: counter },
            }
        } else {
            Entry {
                dbscheme_name: format!("{}_reserved_word", &prefix),
                ql_class_name: "ReservedWord".to_owned(),
                kind: EntryKind::Token { kind_id: 0 },
            }
        };
        entries.insert(type_name, entry);
    }
    entries
}

fn add_field(
    prefix: &str,
    parent_type_name: &TypeName,
    field_name: Option<String>,
    field_info: &FieldInfo,
    fields: &mut Vec<Field>,
    token_kinds: &Set<TypeName>,
) {
    let parent_flattened_name = node_type_name(&parent_type_name.kind, parent_type_name.named);
    let column_name = escape_name(&name_for_field_or_child(&field_name));
    let storage = if !field_info.multiple && field_info.required {
        // This field must appear exactly once, so we add it as
        // a column to the main table for the node type.
        Storage::Column { name: column_name }
    } else {
        // Put the field in an auxiliary table.
        let has_index = field_info.multiple;
        let field_table_name = escape_name(&format!(
            "{}_{}_{}",
            &prefix,
            parent_flattened_name,
            &name_for_field_or_child(&field_name)
        ));
        Storage::Table {
            has_index,
            name: field_table_name,
            column_name,
        }
    };
    let converted_types = convert_types(&field_info.types);
    let type_info = if field_info
        .types
        .iter()
        .all(|t| !t.named && token_kinds.contains(&convert_type(t)))
    {
        // All possible types for this field are reserved words. The db
        // representation will be an `int` with a `case @foo.field = ...` to
        // enumerate the possible values.
        let mut counter = 0;
        let mut field_token_ints: BTreeMap<String, (usize, String)> = BTreeMap::new();
        for t in converted_types {
            let dbscheme_variant_name =
                escape_name(&format!("{}_{}_{}", &prefix, parent_flattened_name, t.kind));
            field_token_ints.insert(t.kind.to_owned(), (counter, dbscheme_variant_name));
            counter += 1;
        }
        FieldTypeInfo::ReservedWordInt(field_token_ints)
    } else if field_info.types.len() == 1 {
        FieldTypeInfo::Single(converted_types.into_iter().next().unwrap())
    } else {
        // The dbscheme type for this field will be a union. In QL, it'll just be AstNode.
        FieldTypeInfo::Multiple {
            types: converted_types,
            dbscheme_union: format!(
                "{}_{}_{}_type",
                &prefix,
                &parent_flattened_name,
                &name_for_field_or_child(&field_name)
            ),
            ql_class: "AstNode".to_owned(),
        }
    };
    let getter_name = format!(
        "get{}",
        dbscheme_name_to_class_name(&escape_name(&name_for_field_or_child(&field_name)))
    );
    fields.push(Field {
        parent: TypeName {
            kind: parent_type_name.kind.to_string(),
            named: parent_type_name.named,
        },
        type_info,
        name: field_name,
        getter_name,
        storage,
    });
}
#[derive(Deserialize)]
pub struct NodeInfo {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fields: Option<BTreeMap<String, FieldInfo>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub children: Option<FieldInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub subtypes: Option<Vec<NodeType>>,
}

#[derive(Deserialize)]
pub struct NodeType {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
}

#[derive(Deserialize)]
pub struct FieldInfo {
    pub multiple: bool,
    pub required: bool,
    pub types: Vec<NodeType>,
}

/// Given a tree-sitter node type's (kind, named) pair, returns a single string
/// representing the (unescaped) name we'll use to refer to corresponding QL
/// type.
fn node_type_name(kind: &str, named: bool) -> String {
    if named {
        kind.to_string()
    } else {
        format!("{}_unnamed", kind)
    }
}

const RESERVED_KEYWORDS: [&'static str; 14] = [
    "boolean", "case", "date", "float", "int", "key", "of", "order", "ref", "string", "subtype",
    "type", "unique", "varchar",
];

/// Returns a string that's a copy of `name` but suitably escaped to be a valid
/// QL identifier.
fn escape_name(name: &str) -> String {
    let mut result = String::new();

    // If there's a leading underscore, replace it with 'underscore_'.
    if let Some(c) = name.chars().next() {
        if c == '_' {
            result.push_str("underscore");
        }
    }
    for c in name.chars() {
        match c {
            '{' => result.push_str("lbrace"),
            '}' => result.push_str("rbrace"),
            '<' => result.push_str("langle"),
            '>' => result.push_str("rangle"),
            '[' => result.push_str("lbracket"),
            ']' => result.push_str("rbracket"),
            '(' => result.push_str("lparen"),
            ')' => result.push_str("rparen"),
            '|' => result.push_str("pipe"),
            '=' => result.push_str("equal"),
            '~' => result.push_str("tilde"),
            '?' => result.push_str("question"),
            '`' => result.push_str("backtick"),
            '^' => result.push_str("caret"),
            '!' => result.push_str("bang"),
            '#' => result.push_str("hash"),
            '%' => result.push_str("percent"),
            '&' => result.push_str("ampersand"),
            '.' => result.push_str("dot"),
            ',' => result.push_str("comma"),
            '/' => result.push_str("slash"),
            ':' => result.push_str("colon"),
            ';' => result.push_str("semicolon"),
            '"' => result.push_str("dquote"),
            '*' => result.push_str("star"),
            '+' => result.push_str("plus"),
            '-' => result.push_str("minus"),
            '@' => result.push_str("at"),
            _ if c.is_uppercase() => {
                result.push('_');
                result.push_str(&c.to_lowercase().to_string())
            }
            _ => result.push(c),
        }
    }

    for &keyword in &RESERVED_KEYWORDS {
        if result == keyword {
            result.push_str("__");
            break;
        }
    }

    result
}

pub fn to_snake_case(word: &str) -> String {
    let mut prev_upper = true;
    let mut result = String::new();
    for c in word.chars() {
        if c.is_uppercase() {
            if !prev_upper {
                result.push('_')
            }
            prev_upper = true;
            result.push(c.to_ascii_lowercase());
        } else {
            prev_upper = false;
            result.push(c);
        }
    }
    result
}
/// Given a valid dbscheme name (i.e. in snake case), produces the equivalent QL
/// name (i.e. in CamelCase). For example, "foo_bar_baz" becomes "FooBarBaz".
fn dbscheme_name_to_class_name(dbscheme_name: &str) -> String {
    fn to_title_case(word: &str) -> String {
        let mut first = true;
        let mut result = String::new();
        for c in word.chars() {
            if first {
                first = false;
                result.push(c.to_ascii_uppercase());
            } else {
                result.push(c);
            }
        }
        result
    }
    dbscheme_name
        .split('_')
        .map(|word| to_title_case(word))
        .collect::<Vec<String>>()
        .join("")
}

#[test]
fn to_snake_case_test() {
    assert_eq!("ruby", to_snake_case("Ruby"));
    assert_eq!("erb", to_snake_case("ERB"));
    assert_eq!("embedded_template", to_snake_case("EmbeddedTemplate"));
}
<commit_msg>Ruby: fix compilation error<commit_after>use serde::Deserialize;
use std::collections::BTreeMap;
use std::path::Path;

use std::collections::BTreeSet as Set;
use std::fs;

/// A lookup table from TypeName to Entry.
pub type NodeTypeMap = BTreeMap<TypeName, Entry>;

#[derive(Debug)]
pub struct Entry {
    pub dbscheme_name: String,
    pub ql_class_name: String,
    pub kind: EntryKind,
}

#[derive(Debug)]
pub enum EntryKind {
    Union { members: Set<TypeName> },
    Table { name: String, fields: Vec<Field> },
    Token { kind_id: usize },
}

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
pub struct TypeName {
    pub kind: String,
    pub named: bool,
}

#[derive(Debug)]
pub enum FieldTypeInfo {
    /// The field has a single type.
    Single(TypeName),

    /// The field can take one of several types, so we also provide the name of
    /// the database union type that wraps them, and the corresponding QL class
    /// name.
    Multiple {
        types: Set<TypeName>,
        dbscheme_union: String,
        ql_class: String,
    },

    /// The field can be one of several tokens, so the db type will be an `int`
    /// with a `case @foo.kind` for each possiblity.
    ReservedWordInt(BTreeMap<String, (usize, String)>),
}

#[derive(Debug)]
pub struct Field {
    pub parent: TypeName,
    pub type_info: FieldTypeInfo,
    /// The name of the field or None for the anonymous 'children'
    /// entry from node_types.json
    pub name: Option<String>,
    /// The name of the predicate to get this field.
    pub getter_name: String,
    pub storage: Storage,
}

fn name_for_field_or_child(name: &Option<String>) -> String {
    match name {
        Some(name) => name.clone(),
        None => "child".to_owned(),
    }
}

#[derive(Debug)]
pub enum Storage {
    /// the field is stored as a column in the parent table
    Column { name: String },
    /// the field is stored in a link table
    Table {
        /// the name of the table
        name: String,
        /// the name of the column for the field in the dbscheme
        column_name: String,
        /// does it have an associated index column?
        has_index: bool,
    },
}

pub fn read_node_types(prefix: &str, node_types_path: &Path) -> std::io::Result<NodeTypeMap> {
    let file = fs::File::open(node_types_path)?;
    let node_types = serde_json::from_reader(file)?;
    Ok(convert_nodes(&prefix, &node_types))
}

pub fn read_node_types_str(prefix: &str, node_types_json: &str) -> std::io::Result<NodeTypeMap> {
    let node_types = serde_json::from_str(node_types_json)?;
    Ok(convert_nodes(&prefix, &node_types))
}

fn convert_type(node_type: &NodeType) -> TypeName {
    TypeName {
        kind: node_type.kind.to_string(),
        named: node_type.named,
    }
}

fn convert_types(node_types: &Vec<NodeType>) -> Set<TypeName> {
    node_types.iter().map(convert_type).collect()
}

pub fn convert_nodes(prefix: &str, nodes: &Vec<NodeInfo>) -> NodeTypeMap {
    let mut entries = NodeTypeMap::new();
    let mut token_kinds = Set::new();

    // First, find all the token kinds
    for node in nodes {
        if node.subtypes.is_none() {
            if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
                let type_name = TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                };
                token_kinds.insert(type_name);
            }
        }
    }

    for node in nodes {
        let flattened_name = &node_type_name(&node.kind, node.named);
        let dbscheme_name = escape_name(&flattened_name);
        let ql_class_name = dbscheme_name_to_class_name(&dbscheme_name);
        let dbscheme_name = format!("{}_{}", prefix, &dbscheme_name);
        if let Some(subtypes) = &node.subtypes {
            // It's a tree-sitter supertype node, for which we create a union
            // type.
            entries.insert(
                TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                },
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Union {
                        members: convert_types(&subtypes),
                    },
                },
            );
        } else if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
            // Token kind, handled above.
        } else {
            // It's a product type, defined by a table.
            let type_name = TypeName {
                kind: node.kind.clone(),
                named: node.named,
            };
            let table_name = escape_name(&(format!("{}_def", &flattened_name)));
            let table_name = format!("{}_{}", prefix, &table_name);

            let mut fields = Vec::new();

            // If the type also has fields or children, then we create either
            // auxiliary tables or columns in the defining table for them.
            if let Some(node_fields) = &node.fields {
                for (field_name, field_info) in node_fields {
                    add_field(
                        &prefix,
                        &type_name,
                        Some(field_name.to_string()),
                        field_info,
                        &mut fields,
                        &token_kinds,
                    );
                }
            }
            if let Some(children) = &node.children {
                // Treat children as if they were a field called 'child'.
                add_field(
                    &prefix,
                    &type_name,
                    None,
                    children,
                    &mut fields,
                    &token_kinds,
                );
            }
            entries.insert(
                type_name,
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Table {
                        name: table_name,
                        fields,
                    },
                },
            );
        }
    }
    let mut counter = 0;
    for type_name in token_kinds {
        let entry = if type_name.named {
            counter += 1;
            let unprefixed_name = node_type_name(&type_name.kind, true);
            Entry {
                dbscheme_name: escape_name(&format!("{}_token_{}", &prefix, &unprefixed_name)),
                ql_class_name: dbscheme_name_to_class_name(&escape_name(&unprefixed_name)),
                kind: EntryKind::Token { kind_id: counter },
            }
        } else {
            Entry {
                dbscheme_name: format!("{}_reserved_word", &prefix),
                ql_class_name: "ReservedWord".to_owned(),
                kind: EntryKind::Token { kind_id: 0 },
            }
        };
        entries.insert(type_name, entry);
    }
    entries
}

fn add_field(
    prefix: &str,
    parent_type_name: &TypeName,
    field_name: Option<String>,
    field_info: &FieldInfo,
    fields: &mut Vec<Field>,
    token_kinds: &Set<TypeName>,
) {
    let parent_flattened_name = node_type_name(&parent_type_name.kind, parent_type_name.named);
    let column_name = escape_name(&name_for_field_or_child(&field_name));
    let storage = if !field_info.multiple && field_info.required {
        // This field must appear exactly once, so we add it as
        // a column to the main table for the node type.
        Storage::Column { name: column_name }
    } else {
        // Put the field in an auxiliary table.
        let has_index = field_info.multiple;
        let field_table_name = escape_name(&format!(
            "{}_{}_{}",
            &prefix,
            parent_flattened_name,
            &name_for_field_or_child(&field_name)
        ));
        Storage::Table {
            has_index,
            name: field_table_name,
            column_name,
        }
    };
    let converted_types = convert_types(&field_info.types);
    let type_info = if field_info
        .types
        .iter()
        .all(|t| !t.named && token_kinds.contains(&convert_type(t)))
    {
        // All possible types for this field are reserved words. The db
        // representation will be an `int` with a `case @foo.field = ...` to
        // enumerate the possible values.
        let mut counter = 0;
        let mut field_token_ints: BTreeMap<String, (usize, String)> = BTreeMap::new();
        for t in converted_types {
            let dbscheme_variant_name =
                escape_name(&format!("{}_{}_{}", &prefix, parent_flattened_name, t.kind));
            field_token_ints.insert(t.kind.to_owned(), (counter, dbscheme_variant_name));
            counter += 1;
        }
        FieldTypeInfo::ReservedWordInt(field_token_ints)
    } else if field_info.types.len() == 1 {
        FieldTypeInfo::Single(converted_types.into_iter().next().unwrap())
    } else {
        // The dbscheme type for this field will be a union. In QL, it'll just be AstNode.
        FieldTypeInfo::Multiple {
            types: converted_types,
            dbscheme_union: format!(
                "{}_{}_{}_type",
                &prefix,
                &parent_flattened_name,
                &name_for_field_or_child(&field_name)
            ),
            ql_class: "AstNode".to_owned(),
        }
    };
    let getter_name = format!(
        "get{}",
        dbscheme_name_to_class_name(&escape_name(&name_for_field_or_child(&field_name)))
    );
    fields.push(Field {
        parent: TypeName {
            kind: parent_type_name.kind.to_string(),
            named: parent_type_name.named,
        },
        type_info,
        name: field_name,
        getter_name,
        storage,
    });
}
#[derive(Deserialize)]
pub struct NodeInfo {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fields: Option<BTreeMap<String, FieldInfo>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub children: Option<FieldInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub subtypes: Option<Vec<NodeType>>,
}

#[derive(Deserialize)]
pub struct NodeType {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
}

#[derive(Deserialize)]
pub struct FieldInfo {
    pub multiple: bool,
    pub required: bool,
    pub types: Vec<NodeType>,
}

/// Given a tree-sitter node type's (kind, named) pair, returns a single string
/// representing the (unescaped) name we'll use to refer to corresponding QL
/// type.
fn node_type_name(kind: &str, named: bool) -> String {
    if named {
        kind.to_string()
    } else {
        format!("{}_unnamed", kind)
    }
}

const RESERVED_KEYWORDS: [&'static str; 14] = [
    "boolean", "case", "date", "float", "int", "key", "of", "order", "ref", "string", "subtype",
    "type", "unique", "varchar",
];

/// Returns a string that's a copy of `name` but suitably escaped to be a valid
/// QL identifier.
fn escape_name(name: &str) -> String {
    let mut result = String::new();

    // If there's a leading underscore, replace it with 'underscore_'.
    if let Some(c) = name.chars().next() {
        if c == '_' {
            result.push_str("underscore");
        }
    }
    for c in name.chars() {
        match c {
            '{' => result.push_str("lbrace"),
            '}' => result.push_str("rbrace"),
            '<' => result.push_str("langle"),
            '>' => result.push_str("rangle"),
            '[' => result.push_str("lbracket"),
            ']' => result.push_str("rbracket"),
            '(' => result.push_str("lparen"),
            ')' => result.push_str("rparen"),
            '|' => result.push_str("pipe"),
            '=' => result.push_str("equal"),
            '~' => result.push_str("tilde"),
            '?' => result.push_str("question"),
            '`' => result.push_str("backtick"),
            '^' => result.push_str("caret"),
            '!' => result.push_str("bang"),
            '#' => result.push_str("hash"),
            '%' => result.push_str("percent"),
            '&' => result.push_str("ampersand"),
            '.' => result.push_str("dot"),
            ',' => result.push_str("comma"),
            '/' => result.push_str("slash"),
            ':' => result.push_str("colon"),
            ';' => result.push_str("semicolon"),
            '"' => result.push_str("dquote"),
            '*' => result.push_str("star"),
            '+' => result.push_str("plus"),
            '-' => result.push_str("minus"),
            '@' => result.push_str("at"),
            _ if c.is_uppercase() => {
                result.push('_');
                result.push_str(&c.to_lowercase().to_string())
            }
            _ => result.push(c),
        }
    }

    for &keyword in &RESERVED_KEYWORDS {
        if result == keyword {
            result.push_str("__");
            break;
        }
    }

    result
}

pub fn to_snake_case(word: &str) -> String {
    let mut prev_upper = true;
    let mut result = String::new();
    for c in word.chars() {
        if c.is_uppercase() {
            if !prev_upper {
                result.push('_')
            }
            prev_upper = true;
            result.push(c.to_ascii_lowercase());
        } else {
            prev_upper = false;
            result.push(c);
        }
    }
    result
}
/// Given a valid dbscheme name (i.e. in snake case), produces the equivalent QL
/// name (i.e. in CamelCase). For example, "foo_bar_baz" becomes "FooBarBaz".
fn dbscheme_name_to_class_name(dbscheme_name: &str) -> String {
    fn to_title_case(word: &str) -> String {
        let mut first = true;
        let mut result = String::new();
        for c in word.chars() {
            if first {
                first = false;
                result.push(c.to_ascii_uppercase());
            } else {
                result.push(c);
            }
        }
        result
    }
    dbscheme_name
        .split('_')
        .map(|word| to_title_case(word))
        .collect::<Vec<String>>()
        .join("")
}

#[test]
fn to_snake_case_test() {
    assert_eq!("ruby", to_snake_case("Ruby"));
    assert_eq!("erb", to_snake_case("ERB"));
    assert_eq!("embedded_template", to_snake_case("EmbeddedTemplate"));
}
<|endoftext|>


<commit_before>use crate::reflect::EnumDescriptor;
use crate::reflect::MessageDescriptor;
use std::fmt;

/// Runtime representation of elementary protobuf type.
#[derive(Debug, Clone)]
pub enum RuntimeTypeBox {
    /// `i32`
    I32,
    /// `i64`
    I64,
    /// `u32`
    U32,
    /// `u64`
    U64,
    /// `f32`
    F32,
    /// `f64`
    F64,
    /// `bool`
    Bool,
    /// [`String`](std::string::String)
    String,
    /// [`Vec<u8>`](std::vec::Vec)
    VecU8,
    /// `enum`
    Enum(&'static EnumDescriptor),
    /// `message`
    Message(&'static MessageDescriptor),
}

impl fmt::Display for RuntimeTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RuntimeTypeBox::I32 => write!(f, "i32"),
            RuntimeTypeBox::I64 => write!(f, "i64"),
            RuntimeTypeBox::U32 => write!(f, "u32"),
            RuntimeTypeBox::U64 => write!(f, "u64"),
            RuntimeTypeBox::F32 => write!(f, "f32"),
            RuntimeTypeBox::F64 => write!(f, "f64"),
            RuntimeTypeBox::Bool => write!(f, "bool"),
            RuntimeTypeBox::String => write!(f, "String"),
            RuntimeTypeBox::VecU8 => write!(f, "Vec<u8>"),
            RuntimeTypeBox::Enum(e) => write!(f, "{}", e.full_name()),
            // TODO: full name
            RuntimeTypeBox::Message(m) => write!(f, "{}", m.name()),
        }
    }
}
<commit_msg>Use full message name in RuntimeTypeBox<commit_after>use crate::reflect::EnumDescriptor;
use crate::reflect::MessageDescriptor;
use std::fmt;

/// Runtime representation of elementary protobuf type.
#[derive(Debug, Clone)]
pub enum RuntimeTypeBox {
    /// `i32`
    I32,
    /// `i64`
    I64,
    /// `u32`
    U32,
    /// `u64`
    U64,
    /// `f32`
    F32,
    /// `f64`
    F64,
    /// `bool`
    Bool,
    /// [`String`](std::string::String)
    String,
    /// [`Vec<u8>`](std::vec::Vec)
    VecU8,
    /// `enum`
    Enum(&'static EnumDescriptor),
    /// `message`
    Message(&'static MessageDescriptor),
}

impl fmt::Display for RuntimeTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RuntimeTypeBox::I32 => write!(f, "i32"),
            RuntimeTypeBox::I64 => write!(f, "i64"),
            RuntimeTypeBox::U32 => write!(f, "u32"),
            RuntimeTypeBox::U64 => write!(f, "u64"),
            RuntimeTypeBox::F32 => write!(f, "f32"),
            RuntimeTypeBox::F64 => write!(f, "f64"),
            RuntimeTypeBox::Bool => write!(f, "bool"),
            RuntimeTypeBox::String => write!(f, "String"),
            RuntimeTypeBox::VecU8 => write!(f, "Vec<u8>"),
            RuntimeTypeBox::Enum(e) => write!(f, "{}", e.full_name()),
            RuntimeTypeBox::Message(m) => write!(f, "{}", m.full_name()),
        }
    }
}
<|endoftext|>


<commit_before>//! Completes constants and paths in unqualified patterns.

use hir::{db::DefDatabase, AssocItem, ScopeDef};
use rustc_hash::FxHashSet;
use syntax::ast::Pat;

use crate::{
    context::{PathCompletionCtx, PathQualifierCtx, PatternRefutability},
    CompletionContext, Completions,
};

/// Completes constants and paths in unqualified patterns.
pub(crate) fn complete_pattern(acc: &mut Completions, ctx: &CompletionContext) {
    let patctx = match &ctx.pattern_ctx {
        Some(ctx) => ctx,
        _ => return,
    };
    let refutable = patctx.refutability == PatternRefutability::Refutable;

    if let Some(path_ctx) = &ctx.path_context {
        pattern_path_completion(acc, ctx, path_ctx);
        return;
    }

    match patctx.parent_pat.as_ref() {
        Some(Pat::RangePat(_) | Pat::BoxPat(_)) => (),
        Some(Pat::RefPat(r)) => {
            if r.mut_token().is_none() {
                acc.add_keyword(ctx, "mut");
            }
        }
        _ => {
            let tok = ctx.token.text_range().start();
            match (patctx.ref_token.as_ref(), patctx.mut_token.as_ref()) {
                (None, None) => {
                    acc.add_keyword(ctx, "ref");
                    acc.add_keyword(ctx, "mut");
                }
                (None, Some(m)) if tok < m.text_range().start() => {
                    acc.add_keyword(ctx, "ref");
                }
                (Some(r), None) if tok > r.text_range().end() => {
                    acc.add_keyword(ctx, "mut");
                }
                _ => (),
            }
        }
    }

    let single_variant_enum = |enum_: hir::Enum| ctx.db.enum_data(enum_.into()).variants.len() == 1;

    if let Some(hir::Adt::Enum(e)) =
        ctx.expected_type.as_ref().and_then(|ty| ty.strip_references().as_adt())
    {
        if refutable || single_variant_enum(e) {
            super::enum_variants_with_paths(acc, ctx, e, |acc, ctx, variant, path| {
                acc.add_qualified_variant_pat(ctx, variant, path);
            });
        }
    }

    // FIXME: ideally, we should look at the type we are matching against and
    // suggest variants + auto-imports
    ctx.process_all_names(&mut |name, res| {
        let add_simple_path = match res {
            hir::ScopeDef::ModuleDef(def) => match def {
                hir::ModuleDef::Adt(hir::Adt::Struct(strukt)) => {
                    acc.add_struct_pat(ctx, strukt, Some(name.clone()));
                    true
                }
                hir::ModuleDef::Variant(variant)
                    if refutable || single_variant_enum(variant.parent_enum(ctx.db)) =>
                {
                    acc.add_variant_pat(ctx, variant, Some(name.clone()));
                    true
                }
                hir::ModuleDef::Adt(hir::Adt::Enum(e)) => refutable || single_variant_enum(e),
                hir::ModuleDef::Const(..) => refutable,
                hir::ModuleDef::Module(..) => true,
                hir::ModuleDef::Macro(mac) if mac.is_fn_like(ctx.db) => {
                    return acc.add_macro(ctx, mac, name)
                }
                _ => false,
            },
            hir::ScopeDef::ImplSelfType(impl_) => match impl_.self_ty(ctx.db).as_adt() {
                Some(hir::Adt::Struct(strukt)) => {
                    acc.add_struct_pat(ctx, strukt, Some(name.clone()));
                    true
                }
                Some(hir::Adt::Enum(e)) => refutable || single_variant_enum(e),
                Some(hir::Adt::Union(_)) => true,
                _ => false,
            },
            ScopeDef::GenericParam(hir::GenericParam::ConstParam(_)) => true,
            ScopeDef::GenericParam(_)
            | ScopeDef::AdtSelfType(_)
            | ScopeDef::Local(_)
            | ScopeDef::Label(_)
            | ScopeDef::Unknown => false,
        };
        if add_simple_path {
            acc.add_resolution_simple(ctx, name, res);
        }
    });
}

fn pattern_path_completion(
    acc: &mut Completions,
    ctx: &CompletionContext,
    PathCompletionCtx { qualifier, is_absolute_path, .. }: &PathCompletionCtx,
) {
    match qualifier {
        Some(PathQualifierCtx { resolution, is_super_chain, .. }) => {
            if *is_super_chain {
                acc.add_keyword(ctx, "super::");
            }

            let resolution = match resolution {
                Some(it) => it,
                None => return,
            };

            match resolution {
                hir::PathResolution::Def(hir::ModuleDef::Module(module)) => {
                    let module_scope = module.scope(ctx.db, ctx.module);
                    for (name, def) in module_scope {
                        let add_resolution = match def {
                            ScopeDef::ModuleDef(hir::ModuleDef::Macro(mac)) => {
                                mac.is_fn_like(ctx.db)
                            }
                            ScopeDef::ModuleDef(_) => true,
                            _ => false,
                        };

                        if add_resolution {
                            acc.add_resolution(ctx, name, def);
                        }
                    }
                }
                hir::PathResolution::Def(hir::ModuleDef::Adt(hir::Adt::Enum(e))) => {
                    cov_mark::hit!(enum_plain_qualified_use_tree);
                    e.variants(ctx.db)
                        .into_iter()
                        .for_each(|variant| acc.add_enum_variant(ctx, variant, None));
                }
                res @ (hir::PathResolution::TypeParam(_) | hir::PathResolution::SelfType(_)) => {
                    let ty = match res {
                        hir::PathResolution::TypeParam(param) => param.ty(ctx.db),
                        hir::PathResolution::SelfType(impl_def) => impl_def.self_ty(ctx.db),
                        _ => return,
                    };

                    if let Some(hir::Adt::Enum(e)) = ty.as_adt() {
                        e.variants(ctx.db)
                            .into_iter()
                            .for_each(|variant| acc.add_enum_variant(ctx, variant, None));
                    }

                    let traits_in_scope = ctx.scope.visible_traits();
                    let mut seen = FxHashSet::default();
                    ty.iterate_path_candidates(
                        ctx.db,
                        &ctx.scope,
                        &traits_in_scope,
                        ctx.module,
                        None,
                        |item| {
                            // Note associated consts cannot be referenced in patterns
                            if let AssocItem::TypeAlias(ta) = item {
                                // We might iterate candidates of a trait multiple times here, so deduplicate them.
                                if seen.insert(item) {
                                    acc.add_type_alias(ctx, ta);
                                }
                            }
                            None::<()>
                        },
                    );
                }
                _ => {}
            }
        }
        // qualifier can only be none here if we are in a TuplePat or RecordPat in which case special characters have to follow the path
        None if *is_absolute_path => acc.add_crate_roots(ctx),
        None => {
            cov_mark::hit!(unqualified_path_only_modules_in_import);
            ctx.process_all_names(&mut |name, res| {
                if let ScopeDef::ModuleDef(hir::ModuleDef::Module(_)) = res {
                    acc.add_resolution(ctx, name, res);
                }
            });
            acc.add_nameref_keywords(ctx);
        }
    }
}
<commit_msg>Complete associated consts in patterns<commit_after>//! Completes constants and paths in unqualified patterns.

use hir::{db::DefDatabase, AssocItem, ScopeDef};
use rustc_hash::FxHashSet;
use syntax::ast::Pat;

use crate::{
    context::{PathCompletionCtx, PathQualifierCtx, PatternRefutability},
    CompletionContext, Completions,
};

/// Completes constants and paths in unqualified patterns.
pub(crate) fn complete_pattern(acc: &mut Completions, ctx: &CompletionContext) {
    let patctx = match &ctx.pattern_ctx {
        Some(ctx) => ctx,
        _ => return,
    };
    let refutable = patctx.refutability == PatternRefutability::Refutable;

    if let Some(path_ctx) = &ctx.path_context {
        pattern_path_completion(acc, ctx, path_ctx);
        return;
    }

    match patctx.parent_pat.as_ref() {
        Some(Pat::RangePat(_) | Pat::BoxPat(_)) => (),
        Some(Pat::RefPat(r)) => {
            if r.mut_token().is_none() {
                acc.add_keyword(ctx, "mut");
            }
        }
        _ => {
            let tok = ctx.token.text_range().start();
            match (patctx.ref_token.as_ref(), patctx.mut_token.as_ref()) {
                (None, None) => {
                    acc.add_keyword(ctx, "ref");
                    acc.add_keyword(ctx, "mut");
                }
                (None, Some(m)) if tok < m.text_range().start() => {
                    acc.add_keyword(ctx, "ref");
                }
                (Some(r), None) if tok > r.text_range().end() => {
                    acc.add_keyword(ctx, "mut");
                }
                _ => (),
            }
        }
    }

    let single_variant_enum = |enum_: hir::Enum| ctx.db.enum_data(enum_.into()).variants.len() == 1;

    if let Some(hir::Adt::Enum(e)) =
        ctx.expected_type.as_ref().and_then(|ty| ty.strip_references().as_adt())
    {
        if refutable || single_variant_enum(e) {
            super::enum_variants_with_paths(acc, ctx, e, |acc, ctx, variant, path| {
                acc.add_qualified_variant_pat(ctx, variant, path);
            });
        }
    }

    // FIXME: ideally, we should look at the type we are matching against and
    // suggest variants + auto-imports
    ctx.process_all_names(&mut |name, res| {
        let add_simple_path = match res {
            hir::ScopeDef::ModuleDef(def) => match def {
                hir::ModuleDef::Adt(hir::Adt::Struct(strukt)) => {
                    acc.add_struct_pat(ctx, strukt, Some(name.clone()));
                    true
                }
                hir::ModuleDef::Variant(variant)
                    if refutable || single_variant_enum(variant.parent_enum(ctx.db)) =>
                {
                    acc.add_variant_pat(ctx, variant, Some(name.clone()));
                    true
                }
                hir::ModuleDef::Adt(hir::Adt::Enum(e)) => refutable || single_variant_enum(e),
                hir::ModuleDef::Const(..) => refutable,
                hir::ModuleDef::Module(..) => true,
                hir::ModuleDef::Macro(mac) if mac.is_fn_like(ctx.db) => {
                    return acc.add_macro(ctx, mac, name)
                }
                _ => false,
            },
            hir::ScopeDef::ImplSelfType(impl_) => match impl_.self_ty(ctx.db).as_adt() {
                Some(hir::Adt::Struct(strukt)) => {
                    acc.add_struct_pat(ctx, strukt, Some(name.clone()));
                    true
                }
                Some(hir::Adt::Enum(e)) => refutable || single_variant_enum(e),
                Some(hir::Adt::Union(_)) => true,
                _ => false,
            },
            ScopeDef::GenericParam(hir::GenericParam::ConstParam(_)) => true,
            ScopeDef::GenericParam(_)
            | ScopeDef::AdtSelfType(_)
            | ScopeDef::Local(_)
            | ScopeDef::Label(_)
            | ScopeDef::Unknown => false,
        };
        if add_simple_path {
            acc.add_resolution_simple(ctx, name, res);
        }
    });
}

fn pattern_path_completion(
    acc: &mut Completions,
    ctx: &CompletionContext,
    PathCompletionCtx { qualifier, is_absolute_path, .. }: &PathCompletionCtx,
) {
    match qualifier {
        Some(PathQualifierCtx { resolution, is_super_chain, .. }) => {
            if *is_super_chain {
                acc.add_keyword(ctx, "super::");
            }

            let resolution = match resolution {
                Some(it) => it,
                None => return,
            };

            match resolution {
                hir::PathResolution::Def(hir::ModuleDef::Module(module)) => {
                    let module_scope = module.scope(ctx.db, ctx.module);
                    for (name, def) in module_scope {
                        let add_resolution = match def {
                            ScopeDef::ModuleDef(hir::ModuleDef::Macro(mac)) => {
                                mac.is_fn_like(ctx.db)
                            }
                            ScopeDef::ModuleDef(_) => true,
                            _ => false,
                        };

                        if add_resolution {
                            acc.add_resolution(ctx, name, def);
                        }
                    }
                }
                hir::PathResolution::Def(hir::ModuleDef::Adt(hir::Adt::Enum(e))) => {
                    cov_mark::hit!(enum_plain_qualified_use_tree);
                    e.variants(ctx.db)
                        .into_iter()
                        .for_each(|variant| acc.add_enum_variant(ctx, variant, None));
                }
                res @ (hir::PathResolution::TypeParam(_) | hir::PathResolution::SelfType(_)) => {
                    let ty = match res {
                        hir::PathResolution::TypeParam(param) => param.ty(ctx.db),
                        hir::PathResolution::SelfType(impl_def) => impl_def.self_ty(ctx.db),
                        _ => return,
                    };

                    if let Some(hir::Adt::Enum(e)) = ty.as_adt() {
                        e.variants(ctx.db)
                            .into_iter()
                            .for_each(|variant| acc.add_enum_variant(ctx, variant, None));
                    }

                    let traits_in_scope = ctx.scope.visible_traits();
                    let mut seen = FxHashSet::default();
                    ty.iterate_path_candidates(
                        ctx.db,
                        &ctx.scope,
                        &traits_in_scope,
                        ctx.module,
                        None,
                        |item| {
                            match item {
                                AssocItem::TypeAlias(ta) => {
                                    // We might iterate candidates of a trait multiple times here, so deduplicate them.
                                    if seen.insert(item) {
                                        acc.add_type_alias(ctx, ta);
                                    }
                                }
                                AssocItem::Const(c) => {
                                    if seen.insert(item) {
                                        acc.add_const(ctx, c);
                                    }
                                }
                                _ => {}
                            }
                            None::<()>
                        },
                    );
                }
                _ => {}
            }
        }
        // qualifier can only be none here if we are in a TuplePat or RecordPat in which case special characters have to follow the path
        None if *is_absolute_path => acc.add_crate_roots(ctx),
        None => {
            cov_mark::hit!(unqualified_path_only_modules_in_import);
            ctx.process_all_names(&mut |name, res| {
                if let ScopeDef::ModuleDef(hir::ModuleDef::Module(_)) = res {
                    acc.add_resolution(ctx, name, res);
                }
            });
            acc.add_nameref_keywords(ctx);
        }
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{Write, BufWriter};

use inflector::Inflector;

use ::Service;
use codegen::botocore::{Shape, ShapeType};
use self::json::JsonGenerator;
use self::query::QueryGenerator;
use self::rest_json::RestJsonGenerator;
use self::rest_xml::RestXmlGenerator;
use self::error_types::{GenerateErrorTypes, JsonErrorTypes, XmlErrorTypes};
use self::tests::generate_tests;
use self::type_filter::filter_types;

mod error_types;
mod json;
mod query;
mod rest_json;
pub mod tests;
mod rest_xml;
mod xml_payload_parser;
mod rest_response_parser;
mod type_filter;

type FileWriter = BufWriter<File>;
type IoResult = ::std::io::Result<()>;

/// Abstracts the generation of Rust code for various AWS protocols
pub trait GenerateProtocol {

    /// Generate the various `use` statements required by the module generatedfor this service
    fn generate_prelude(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    fn generate_method_signatures(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Generate a method for each `Operation` in the `Service` to execute that method remotely
    ///
    /// The method generated by this method are inserted into an enclosing `impl FooClient {}` block
    fn generate_method_impls(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Add any attributes that should decorate the struct for the given type (typically `Debug`, `Clone`, etc.)
    fn generate_struct_attributes(&self,
                                  serialized: bool,
                                  deserialized: bool)
                                  -> String;

    /// If necessary, generate a serializer for the specified type
    fn generate_serializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// If necessary, generate a deserializer for the specified type
    fn generate_deserializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// Return the type used by this protocol for timestamps
    fn timestamp_type(&self) -> &'static str;
}

pub fn generate_source(service: &Service, writer: &mut FileWriter) -> IoResult {
    // EC2 service protocol is similar to query but not the same.  Rusoto is able to generate Rust code
    // from the service definition through the same QueryGenerator, but botocore uses a special class.
    // See https://github.com/boto/botocore/blob/dff99fdf2666accf6b448aef7f03fe3d66dd38fa/botocore/serialize.py#L259-L266 .
    match service.protocol() {
        "json" => generate(writer, service, JsonGenerator, JsonErrorTypes),
        "query" | "ec2" => generate(writer, service, QueryGenerator, XmlErrorTypes),
        "rest-json" => generate(writer, service, RestJsonGenerator, JsonErrorTypes),
        "rest-xml" => generate(writer, service, RestXmlGenerator, XmlErrorTypes),
        protocol => panic!("Unknown protocol {}", protocol),
    }
}

/// Translate a botocore field name to something rust-idiomatic and
/// escape reserved words with an underscore
pub fn generate_field_name(member_name: &str) -> String {
    let name = member_name.to_snake_case();
    if name == "return" || name == "type" {
        name + "_"
    } else {
        name
    }
}

/// Capitalize the first character in the given string.
/// If the input string is empty an empty string is returned.
pub fn capitalize_first<S>(word: S) -> String
    where S: Into<String> {
    let s = word.into();
    let mut chars = s.chars();
    match chars.next() {
        Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
        None => String::new(),
    }
}

/// The quick brown fox jumps over the lazy dog
fn generate<P, E>(writer: &mut FileWriter, service: &Service, protocol_generator: P, error_type_generator: E) -> IoResult
    where P: GenerateProtocol,
          E: GenerateErrorTypes {

    writeln!(writer, "#[allow(warnings)]
        use hyper::Client;
        use hyper::status::StatusCode;
        use rusoto_core::request::DispatchSignedRequest;
        use rusoto_core::region;

        use std::fmt;
        use std::error::Error;
        use rusoto_core::request::HttpDispatchError;
        use rusoto_core::credential::{{CredentialsError, ProvideAwsCredentials}};
    ")?;

    protocol_generator.generate_prelude(writer, service)?;
    generate_types(writer, service, &protocol_generator)?;
    error_type_generator.generate_error_types(writer, service)?;
    generate_client(writer, service, &protocol_generator)?;
    generate_tests(writer, service)?;

    Ok(())

}

fn generate_client<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol
{
    // If the struct name is changed, the links in each service documentation should change.
    // See https://github.com/rusoto/rusoto/issues/519
    writeln!(writer,
        "/// Trait representing the capabilities of the {service_name} API. {service_name} clients implement this trait.
        pub trait {trait_name} {{
        ",
        trait_name = service.service_type_name(),
        service_name = service.name()
    )?;

    protocol_generator.generate_method_signatures(writer, service)?;

    writeln!(writer, "}}")?;

    writeln!(writer,
        "/// A client for the {service_name} API.
        pub struct {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            credentials_provider: P,
            region: region::Region,
            dispatcher: D,
        }}

        impl<P, D> {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            pub fn new(request_dispatcher: D, credentials_provider: P, region: region::Region) -> Self {{
                  {type_name} {{
                    credentials_provider: credentials_provider,
                    region: region,
                    dispatcher: request_dispatcher
                }}
            }}
        }}

        impl<P, D> {trait_name} for {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
        ",
        service_name = service.name(),
        type_name = service.client_type_name(),
        trait_name = service.service_type_name(),
    )?;
    protocol_generator.generate_method_impls(writer, service)?;
    writeln!(writer, "}}")
}

fn generate_list(name: &str, shape: &Shape) -> String {
    format!("pub type {} = Vec<{}>;",
            name,
            mutate_type_name(shape.member_type()))
}

fn generate_map(name: &str, shape: &Shape) -> String {
    format!(
        "pub type {} = ::std::collections::HashMap<{}, {}>;",
        name,
        capitalize_first(shape.key_type().to_string()),
        capitalize_first(shape.value_type().to_string()),
    )
}

fn generate_primitive_type(name: &str, shape_type: ShapeType, for_timestamps: &str) -> String {
    let primitive_type = match shape_type {
        ShapeType::Blob => "Vec<u8>",
        ShapeType::Boolean => "bool",
        ShapeType::Double => "f64",
        ShapeType::Float => "f32",
        ShapeType::Integer => "i64",
        ShapeType::Long => "i64",
        ShapeType::String => "String",
        ShapeType::Timestamp => for_timestamps,
        primitive_type => panic!("Unknown primitive type: {:?}", primitive_type),
    };

    format!("pub type {} = {};", name, primitive_type)
}

// do any type name mutation needed to avoid collisions with Rust types
fn mutate_type_name(type_name: &str) -> String {
    let capitalized = capitalize_first(type_name.to_owned());

    // some cloudfront types have underscoare that anger the lint checker
    let without_underscores = capitalized.replace("_", "");

    match &without_underscores[..] {
        // S3 has an 'Error' shape that collides with Rust's Error trait
        "Error" => "S3Error".to_owned(),

        // EC2 has a CancelSpotFleetRequestsError struct, avoid collision with our error enum
        "CancelSpotFleetRequests" => "EC2CancelSpotFleetRequests".to_owned(),

        // RDS has a conveniently named "Option" type
        "Option" => "RDSOption".to_owned(),

        // otherwise make sure it's rust-idiomatic and capitalized
        _ => without_underscores,
    }
}

fn generate_types<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol {

    let (serialized_types, deserialized_types) = filter_types(service);

    for (name, shape) in service.shapes().iter() {
            let type_name = mutate_type_name(&name);

            // We generate enums for error types, so no need to create model objects for them
            if shape.exception() {
                continue;
            }

            // If botocore includes documentation, clean it up a bit and use it
            if let Some(ref docs) = shape.documentation {
                writeln!(writer, "#[doc=\"{}\"]",
                                   docs.replace("\\", "\\\\").replace("\"", "\\\""))?;
            }

            let deserialized = deserialized_types.contains(&type_name);
            let serialized = serialized_types.contains(&type_name);

            // generate a rust type for the shape
            if type_name != "String" {
                let generated_type = match shape.shape_type {
                    ShapeType::Structure => {
                        generate_struct(service,
                                                   &type_name,
                                                   &shape,
                                                   serialized,
                                                   deserialized,
                                                   protocol_generator)
                    }
                    ShapeType::Map => generate_map(&type_name, &shape),
                    ShapeType::List => generate_list(&type_name, &shape),
                    shape_type => {
                        generate_primitive_type(&type_name,
                                                           shape_type,
                                                           protocol_generator.timestamp_type())
                    }
                };
                writeln!(writer, "{}", generated_type)?;
            }

            if deserialized {
                if let Some(deserializer) = protocol_generator.generate_deserializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", deserializer)?;
                }
            }

            if serialized {
                if let Some(serializer) = protocol_generator.generate_serializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", serializer)?;
                }
            }
        }
        Ok(())
}

fn generate_struct<P>(service: &Service,
                      name: &str,
                      shape: &Shape,
                      serialized: bool,
                      deserialized: bool,
                      protocol_generator: &P)
                      -> String
    where P: GenerateProtocol {

    if shape.members.is_none() || shape.members.as_ref().unwrap().is_empty() {
        format!(
            "{attributes}
            pub struct {name};
            ",
            attributes = protocol_generator.generate_struct_attributes(serialized, deserialized),
            name = name,
        )
    } else {
        let struct_attributes =
            protocol_generator.generate_struct_attributes(serialized, deserialized);
        // Serde attributes are only needed if deriving the Serialize or Deserialize trait
        let need_serde_attrs = struct_attributes.contains("erialize");
        format!(
            "{attributes}
            pub struct {name} {{
                {struct_fields}
            }}
            ",
            attributes = struct_attributes,
            name = name,
            struct_fields = generate_struct_fields(service, shape, need_serde_attrs),
        )
    }

}

fn generate_struct_fields(service: &Service, shape: &Shape, serde_attrs: bool) -> String {
    shape.members.as_ref().unwrap().iter().filter_map(|(member_name, member)| {

        if member.deprecated == Some(true) {
            return None;
        }

        let mut lines: Vec<String> = Vec::new();
        let name = generate_field_name(member_name);

        if let Some(ref docs) = member.documentation {
            lines.push(format!("#[doc=\"{}\"]", docs.replace("\\","\\\\").replace("\"", "\\\"")));
        }

        let type_name = mutate_type_name(&member.shape);

        if serde_attrs {
            lines.push(format!("#[serde(rename=\"{}\")]", member_name));

            if let Some(shape_type) = service.shape_type_for_member(member) {
                if shape_type == ShapeType::Blob {
                    lines.push(
                        "#[serde(
                            deserialize_with=\"::rusoto_core::serialization::SerdeBlob::deserialize_blob\",
                            serialize_with=\"::rusoto_core::serialization::SerdeBlob::serialize_blob\",
                            default,
                        )]".to_owned()
                    );
                } else if shape_type == ShapeType::Boolean && !shape.required(member_name) {
                    lines.push("#[serde(skip_serializing_if=\"::std::option::Option::is_none\")]".to_owned());
                }
            }
        }

        if shape.required(member_name) {
            lines.push(format!("pub {}: {},",  name, type_name))
        } else if name == "type" {
            lines.push(format!("pub aws_{}: Option<{}>,",  name, type_name))
        } else {
            lines.push(format!("pub {}: Option<{}>,",  name, type_name))
        }

        Some(lines.join("\n"))
    }).collect::<Vec<String>>().join("\n")
}

fn error_type_name(name: &str) -> String {
    let type_name = mutate_type_name(name);
    format!("{}Error", type_name)
}

#[test]
fn capitalize_first_test() {
    assert_eq!(capitalize_first("a &str test"), "A &str test".to_owned());
    assert_eq!(capitalize_first("a String test".to_owned()),
               "A String test".to_owned());
}
<commit_msg>Generate skip_serializing for all optional, not just bool<commit_after>use std::fs::File;
use std::io::{Write, BufWriter};

use inflector::Inflector;

use ::Service;
use codegen::botocore::{Shape, ShapeType};
use self::json::JsonGenerator;
use self::query::QueryGenerator;
use self::rest_json::RestJsonGenerator;
use self::rest_xml::RestXmlGenerator;
use self::error_types::{GenerateErrorTypes, JsonErrorTypes, XmlErrorTypes};
use self::tests::generate_tests;
use self::type_filter::filter_types;

mod error_types;
mod json;
mod query;
mod rest_json;
pub mod tests;
mod rest_xml;
mod xml_payload_parser;
mod rest_response_parser;
mod type_filter;

type FileWriter = BufWriter<File>;
type IoResult = ::std::io::Result<()>;

/// Abstracts the generation of Rust code for various AWS protocols
pub trait GenerateProtocol {

    /// Generate the various `use` statements required by the module generatedfor this service
    fn generate_prelude(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    fn generate_method_signatures(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Generate a method for each `Operation` in the `Service` to execute that method remotely
    ///
    /// The method generated by this method are inserted into an enclosing `impl FooClient {}` block
    fn generate_method_impls(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Add any attributes that should decorate the struct for the given type (typically `Debug`, `Clone`, etc.)
    fn generate_struct_attributes(&self,
                                  serialized: bool,
                                  deserialized: bool)
                                  -> String;

    /// If necessary, generate a serializer for the specified type
    fn generate_serializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// If necessary, generate a deserializer for the specified type
    fn generate_deserializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// Return the type used by this protocol for timestamps
    fn timestamp_type(&self) -> &'static str;
}

pub fn generate_source(service: &Service, writer: &mut FileWriter) -> IoResult {
    // EC2 service protocol is similar to query but not the same.  Rusoto is able to generate Rust code
    // from the service definition through the same QueryGenerator, but botocore uses a special class.
    // See https://github.com/boto/botocore/blob/dff99fdf2666accf6b448aef7f03fe3d66dd38fa/botocore/serialize.py#L259-L266 .
    match service.protocol() {
        "json" => generate(writer, service, JsonGenerator, JsonErrorTypes),
        "query" | "ec2" => generate(writer, service, QueryGenerator, XmlErrorTypes),
        "rest-json" => generate(writer, service, RestJsonGenerator, JsonErrorTypes),
        "rest-xml" => generate(writer, service, RestXmlGenerator, XmlErrorTypes),
        protocol => panic!("Unknown protocol {}", protocol),
    }
}

/// Translate a botocore field name to something rust-idiomatic and
/// escape reserved words with an underscore
pub fn generate_field_name(member_name: &str) -> String {
    let name = member_name.to_snake_case();
    if name == "return" || name == "type" {
        name + "_"
    } else {
        name
    }
}

/// Capitalize the first character in the given string.
/// If the input string is empty an empty string is returned.
pub fn capitalize_first<S>(word: S) -> String
    where S: Into<String> {
    let s = word.into();
    let mut chars = s.chars();
    match chars.next() {
        Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
        None => String::new(),
    }
}

/// The quick brown fox jumps over the lazy dog
fn generate<P, E>(writer: &mut FileWriter, service: &Service, protocol_generator: P, error_type_generator: E) -> IoResult
    where P: GenerateProtocol,
          E: GenerateErrorTypes {

    writeln!(writer, "#[allow(warnings)]
        use hyper::Client;
        use hyper::status::StatusCode;
        use rusoto_core::request::DispatchSignedRequest;
        use rusoto_core::region;

        use std::fmt;
        use std::error::Error;
        use rusoto_core::request::HttpDispatchError;
        use rusoto_core::credential::{{CredentialsError, ProvideAwsCredentials}};
    ")?;

    protocol_generator.generate_prelude(writer, service)?;
    generate_types(writer, service, &protocol_generator)?;
    error_type_generator.generate_error_types(writer, service)?;
    generate_client(writer, service, &protocol_generator)?;
    generate_tests(writer, service)?;

    Ok(())

}

fn generate_client<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol
{
    // If the struct name is changed, the links in each service documentation should change.
    // See https://github.com/rusoto/rusoto/issues/519
    writeln!(writer,
        "/// Trait representing the capabilities of the {service_name} API. {service_name} clients implement this trait.
        pub trait {trait_name} {{
        ",
        trait_name = service.service_type_name(),
        service_name = service.name()
    )?;

    protocol_generator.generate_method_signatures(writer, service)?;

    writeln!(writer, "}}")?;

    writeln!(writer,
        "/// A client for the {service_name} API.
        pub struct {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            credentials_provider: P,
            region: region::Region,
            dispatcher: D,
        }}

        impl<P, D> {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            pub fn new(request_dispatcher: D, credentials_provider: P, region: region::Region) -> Self {{
                  {type_name} {{
                    credentials_provider: credentials_provider,
                    region: region,
                    dispatcher: request_dispatcher
                }}
            }}
        }}

        impl<P, D> {trait_name} for {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
        ",
        service_name = service.name(),
        type_name = service.client_type_name(),
        trait_name = service.service_type_name(),
    )?;
    protocol_generator.generate_method_impls(writer, service)?;
    writeln!(writer, "}}")
}

fn generate_list(name: &str, shape: &Shape) -> String {
    format!("pub type {} = Vec<{}>;",
            name,
            mutate_type_name(shape.member_type()))
}

fn generate_map(name: &str, shape: &Shape) -> String {
    format!(
        "pub type {} = ::std::collections::HashMap<{}, {}>;",
        name,
        capitalize_first(shape.key_type().to_string()),
        capitalize_first(shape.value_type().to_string()),
    )
}

fn generate_primitive_type(name: &str, shape_type: ShapeType, for_timestamps: &str) -> String {
    let primitive_type = match shape_type {
        ShapeType::Blob => "Vec<u8>",
        ShapeType::Boolean => "bool",
        ShapeType::Double => "f64",
        ShapeType::Float => "f32",
        ShapeType::Integer => "i64",
        ShapeType::Long => "i64",
        ShapeType::String => "String",
        ShapeType::Timestamp => for_timestamps,
        primitive_type => panic!("Unknown primitive type: {:?}", primitive_type),
    };

    format!("pub type {} = {};", name, primitive_type)
}

// do any type name mutation needed to avoid collisions with Rust types
fn mutate_type_name(type_name: &str) -> String {
    let capitalized = capitalize_first(type_name.to_owned());

    // some cloudfront types have underscoare that anger the lint checker
    let without_underscores = capitalized.replace("_", "");

    match &without_underscores[..] {
        // S3 has an 'Error' shape that collides with Rust's Error trait
        "Error" => "S3Error".to_owned(),

        // EC2 has a CancelSpotFleetRequestsError struct, avoid collision with our error enum
        "CancelSpotFleetRequests" => "EC2CancelSpotFleetRequests".to_owned(),

        // RDS has a conveniently named "Option" type
        "Option" => "RDSOption".to_owned(),

        // otherwise make sure it's rust-idiomatic and capitalized
        _ => without_underscores,
    }
}

fn generate_types<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol {

    let (serialized_types, deserialized_types) = filter_types(service);

    for (name, shape) in service.shapes().iter() {
            let type_name = mutate_type_name(&name);

            // We generate enums for error types, so no need to create model objects for them
            if shape.exception() {
                continue;
            }

            // If botocore includes documentation, clean it up a bit and use it
            if let Some(ref docs) = shape.documentation {
                writeln!(writer, "#[doc=\"{}\"]",
                                   docs.replace("\\", "\\\\").replace("\"", "\\\""))?;
            }

            let deserialized = deserialized_types.contains(&type_name);
            let serialized = serialized_types.contains(&type_name);

            // generate a rust type for the shape
            if type_name != "String" {
                let generated_type = match shape.shape_type {
                    ShapeType::Structure => {
                        generate_struct(service,
                                                   &type_name,
                                                   &shape,
                                                   serialized,
                                                   deserialized,
                                                   protocol_generator)
                    }
                    ShapeType::Map => generate_map(&type_name, &shape),
                    ShapeType::List => generate_list(&type_name, &shape),
                    shape_type => {
                        generate_primitive_type(&type_name,
                                                           shape_type,
                                                           protocol_generator.timestamp_type())
                    }
                };
                writeln!(writer, "{}", generated_type)?;
            }

            if deserialized {
                if let Some(deserializer) = protocol_generator.generate_deserializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", deserializer)?;
                }
            }

            if serialized {
                if let Some(serializer) = protocol_generator.generate_serializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", serializer)?;
                }
            }
        }
        Ok(())
}

fn generate_struct<P>(service: &Service,
                      name: &str,
                      shape: &Shape,
                      serialized: bool,
                      deserialized: bool,
                      protocol_generator: &P)
                      -> String
    where P: GenerateProtocol {

    if shape.members.is_none() || shape.members.as_ref().unwrap().is_empty() {
        format!(
            "{attributes}
            pub struct {name};
            ",
            attributes = protocol_generator.generate_struct_attributes(serialized, deserialized),
            name = name,
        )
    } else {
        let struct_attributes =
            protocol_generator.generate_struct_attributes(serialized, deserialized);
        // Serde attributes are only needed if deriving the Serialize or Deserialize trait
        let need_serde_attrs = struct_attributes.contains("erialize");
        format!(
            "{attributes}
            pub struct {name} {{
                {struct_fields}
            }}
            ",
            attributes = struct_attributes,
            name = name,
            struct_fields = generate_struct_fields(service, shape, need_serde_attrs),
        )
    }

}

fn generate_struct_fields(service: &Service, shape: &Shape, serde_attrs: bool) -> String {
    shape.members.as_ref().unwrap().iter().filter_map(|(member_name, member)| {

        if member.deprecated == Some(true) {
            return None;
        }

        let mut lines: Vec<String> = Vec::new();
        let name = generate_field_name(member_name);

        if let Some(ref docs) = member.documentation {
            lines.push(format!("#[doc=\"{}\"]", docs.replace("\\","\\\\").replace("\"", "\\\"")));
        }

        let type_name = mutate_type_name(&member.shape);

        if serde_attrs {
            lines.push(format!("#[serde(rename=\"{}\")]", member_name));

            if let Some(shape_type) = service.shape_type_for_member(member) {
                if shape_type == ShapeType::Blob {
                    lines.push(
                        "#[serde(
                            deserialize_with=\"::rusoto_core::serialization::SerdeBlob::deserialize_blob\",
                            serialize_with=\"::rusoto_core::serialization::SerdeBlob::serialize_blob\",
                            default,
                        )]".to_owned()
                    );
                } else if !shape.required(member_name) {
                    lines.push("#[serde(skip_serializing_if=\"Option::is_none\")]".to_owned());
                }
            }
        }

        if shape.required(member_name) {
            lines.push(format!("pub {}: {},",  name, type_name))
        } else if name == "type" {
            lines.push(format!("pub aws_{}: Option<{}>,",  name, type_name))
        } else {
            lines.push(format!("pub {}: Option<{}>,",  name, type_name))
        }

        Some(lines.join("\n"))
    }).collect::<Vec<String>>().join("\n")
}

fn error_type_name(name: &str) -> String {
    let type_name = mutate_type_name(name);
    format!("{}Error", type_name)
}

#[test]
fn capitalize_first_test() {
    assert_eq!(capitalize_first("a &str test"), "A &str test".to_owned());
    assert_eq!(capitalize_first("a String test".to_owned()),
               "A String test".to_owned());
}
<|endoftext|>


<commit_before>// Copyright 2014 Tyler Neely
//
// 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.
//
extern crate rocksdb;
mod util;

use rocksdb::{ColumnFamilyDescriptor, MergeOperands, Options, DB};
use util::DBPath;

#[test]
fn test_column_family() {
    let n = DBPath::new("_rust_rocksdb_cftest");

    // should be able to create column families
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_merge_operator("test operator", test_provided_merge, None);
        let db = DB::open(&opts, &n).unwrap();
        let opts = Options::default();
        match db.create_cf("cf1", &opts) {
            Ok(_db) => println!("cf1 created successfully"),
            Err(e) => {
                panic!("could not create column family: {}", e);
            }
        }
    }

    // should fail to open db without specifying same column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator("test operator", test_provided_merge, None);
        match DB::open(&opts, &n) {
            Ok(_db) => panic!(
                "should not have opened DB successfully without \
                        specifying column
            families"
            ),
            Err(e) => assert!(e.to_string().starts_with(
                "Invalid argument: You have to open all \
                 column families."
            )),
        }
    }

    // should properly open db when specyfing all column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator("test operator", test_provided_merge, None);
        match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(_db) => println!("successfully opened db with column family"),
            Err(e) => panic!("failed to open db with column family: {}", e),
        }
    }

    // should be able to list a cf
    {
        let opts = Options::default();
        let vec = DB::list_cf(&opts, &n);
        match vec {
            Ok(vec) => assert_eq!(vec, vec!["default", "cf1"]),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }

    // TODO should be able to use writebatch ops with a cf
    {}
    // TODO should be able to iterate over a cf
    {}
    // should b able to drop a cf
    {
        let db = DB::open_cf(&Options::default(), &n, &["cf1"]).unwrap();
        match db.drop_cf("cf1") {
            Ok(_) => println!("cf1 successfully dropped."),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }
}

#[test]
fn test_can_open_db_with_results_of_list_cf() {
    // Test scenario derived from GitHub issue #175 and 177

    let n = DBPath::new("_rust_rocksdb_cftest_with_list_cf");

    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        let db = DB::open(&opts, &n).unwrap();
        let opts = Options::default();

        assert!(db.create_cf("cf1", &opts).is_ok());
    }

    {
        let options = Options::default();
        let cfs = DB::list_cf(&options, &n).unwrap();
        let db = DB::open_cf(&options, &n, &cfs).unwrap();

        assert!(db.cf_handle("cf1").is_some());
    }
}

#[test]
fn test_create_missing_column_family() {
    let n = DBPath::new("_rust_rocksdb_missing_cftest");

    // should be able to create new column families when opening a new database
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(_db) => println!("successfully created new column family"),
            Err(e) => panic!("failed to create new column family: {}", e),
        }
    }
}

#[test]
#[ignore]
fn test_merge_operator() {
    let n = DBPath::new("_rust_rocksdb_cftest_merge");
    // TODO should be able to write, read, merge, batch, and iterate over a cf
    {
        let mut opts = Options::default();
        opts.set_merge_operator("test operator", test_provided_merge, None);
        let db = match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(db) => {
                println!("successfully opened db with column family");
                db
            }
            Err(e) => panic!("failed to open db with column family: {}", e),
        };
        let cf1 = db.cf_handle("cf1").unwrap();
        assert!(db.put_cf(cf1, b"k1", b"v1").is_ok());
        assert!(db.get_cf(cf1, b"k1").unwrap().unwrap().to_utf8().unwrap() == "v1");
        let p = db.put_cf(cf1, b"k1", b"a");
        assert!(p.is_ok());
        db.merge_cf(cf1, b"k1", b"b").unwrap();
        db.merge_cf(cf1, b"k1", b"c").unwrap();
        db.merge_cf(cf1, b"k1", b"d").unwrap();
        db.merge_cf(cf1, b"k1", b"efg").unwrap();
        let m = db.merge_cf(cf1, b"k1", b"h");
        println!("m is {:?}", m);
        // TODO assert!(m.is_ok());
        match db.get(b"k1") {
            Ok(Some(value)) => match value.to_utf8() {
                Some(v) => println!("retrieved utf8 value: {}", v),
                None => println!("did not read valid utf-8 out of the db"),
            },
            Err(_) => println!("error reading value"),
            _ => panic!("value not present!"),
        }

        let _ = db.get_cf(cf1, b"k1");
        // TODO assert!(r.unwrap().to_utf8().unwrap() == "abcdefgh");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

fn test_provided_merge(
    _: &[u8],
    existing_val: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let nops = operands.size_hint().0;
    let mut result: Vec<u8> = Vec::with_capacity(nops);
    match existing_val {
        Some(v) => {
            for e in v {
                result.push(*e);
            }
        }
        None => (),
    }
    for op in operands {
        for e in op {
            result.push(*e);
        }
    }
    Some(result)
}

#[test]
fn test_column_family_with_options() {
    let n = DBPath::new("_rust_rocksdb_cf_with_optionstest");
    {
        let mut cfopts = Options::default();
        cfopts.set_max_write_buffer_number(16);
        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);

        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cfs = vec![cf_descriptor];
        match DB::open_cf_descriptors(&opts, &n, cfs) {
            Ok(_db) => println!("created db with column family descriptors successfully"),
            Err(e) => {
                panic!(
                    "could not create new database with column family descriptors: {}",
                    e
                );
            }
        }
    }

    {
        let mut cfopts = Options::default();
        cfopts.set_max_write_buffer_number(16);
        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);

        let opts = Options::default();
        let cfs = vec![cf_descriptor];

        match DB::open_cf_descriptors(&opts, &n, cfs) {
            Ok(_db) => println!("successfully re-opened database with column family descriptors"),
            Err(e) => {
                panic!(
                    "unable to re-open database with column family descriptors: {}",
                    e
                );
            }
        }
    }
}
<commit_msg>Add test to prove creating duplicate column family results in error...<commit_after>// Copyright 2014 Tyler Neely
//
// 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.
//
extern crate rocksdb;
mod util;

use rocksdb::{ColumnFamilyDescriptor, MergeOperands, Options, DB};
use util::DBPath;

#[test]
fn test_column_family() {
    let n = DBPath::new("_rust_rocksdb_cftest");

    // should be able to create column families
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_merge_operator("test operator", test_provided_merge, None);
        let db = DB::open(&opts, &n).unwrap();
        let opts = Options::default();
        match db.create_cf("cf1", &opts) {
            Ok(_db) => println!("cf1 created successfully"),
            Err(e) => {
                panic!("could not create column family: {}", e);
            }
        }
    }

    // should fail to open db without specifying same column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator("test operator", test_provided_merge, None);
        match DB::open(&opts, &n) {
            Ok(_db) => panic!(
                "should not have opened DB successfully without \
                        specifying column
            families"
            ),
            Err(e) => assert!(e.to_string().starts_with(
                "Invalid argument: You have to open all \
                 column families."
            )),
        }
    }

    // should properly open db when specyfing all column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator("test operator", test_provided_merge, None);
        match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(_db) => println!("successfully opened db with column family"),
            Err(e) => panic!("failed to open db with column family: {}", e),
        }
    }

    // should be able to list a cf
    {
        let opts = Options::default();
        let vec = DB::list_cf(&opts, &n);
        match vec {
            Ok(vec) => assert_eq!(vec, vec!["default", "cf1"]),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }

    // TODO should be able to use writebatch ops with a cf
    {}
    // TODO should be able to iterate over a cf
    {}
    // should b able to drop a cf
    {
        let db = DB::open_cf(&Options::default(), &n, &["cf1"]).unwrap();
        match db.drop_cf("cf1") {
            Ok(_) => println!("cf1 successfully dropped."),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }
}

#[test]
fn test_can_open_db_with_results_of_list_cf() {
    // Test scenario derived from GitHub issue #175 and 177

    let n = DBPath::new("_rust_rocksdb_cftest_with_list_cf");

    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        let db = DB::open(&opts, &n).unwrap();
        let opts = Options::default();

        assert!(db.create_cf("cf1", &opts).is_ok());
    }

    {
        let options = Options::default();
        let cfs = DB::list_cf(&options, &n).unwrap();
        let db = DB::open_cf(&options, &n, &cfs).unwrap();

        assert!(db.cf_handle("cf1").is_some());
    }
}

#[test]
fn test_create_missing_column_family() {
    let n = DBPath::new("_rust_rocksdb_missing_cftest");

    // should be able to create new column families when opening a new database
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(_db) => println!("successfully created new column family"),
            Err(e) => panic!("failed to create new column family: {}", e),
        }
    }
}

#[test]
#[ignore]
fn test_merge_operator() {
    let n = DBPath::new("_rust_rocksdb_cftest_merge");
    // TODO should be able to write, read, merge, batch, and iterate over a cf
    {
        let mut opts = Options::default();
        opts.set_merge_operator("test operator", test_provided_merge, None);
        let db = match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(db) => {
                println!("successfully opened db with column family");
                db
            }
            Err(e) => panic!("failed to open db with column family: {}", e),
        };
        let cf1 = db.cf_handle("cf1").unwrap();
        assert!(db.put_cf(cf1, b"k1", b"v1").is_ok());
        assert!(db.get_cf(cf1, b"k1").unwrap().unwrap().to_utf8().unwrap() == "v1");
        let p = db.put_cf(cf1, b"k1", b"a");
        assert!(p.is_ok());
        db.merge_cf(cf1, b"k1", b"b").unwrap();
        db.merge_cf(cf1, b"k1", b"c").unwrap();
        db.merge_cf(cf1, b"k1", b"d").unwrap();
        db.merge_cf(cf1, b"k1", b"efg").unwrap();
        let m = db.merge_cf(cf1, b"k1", b"h");
        println!("m is {:?}", m);
        // TODO assert!(m.is_ok());
        match db.get(b"k1") {
            Ok(Some(value)) => match value.to_utf8() {
                Some(v) => println!("retrieved utf8 value: {}", v),
                None => println!("did not read valid utf-8 out of the db"),
            },
            Err(_) => println!("error reading value"),
            _ => panic!("value not present!"),
        }

        let _ = db.get_cf(cf1, b"k1");
        // TODO assert!(r.unwrap().to_utf8().unwrap() == "abcdefgh");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

fn test_provided_merge(
    _: &[u8],
    existing_val: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let nops = operands.size_hint().0;
    let mut result: Vec<u8> = Vec::with_capacity(nops);
    match existing_val {
        Some(v) => {
            for e in v {
                result.push(*e);
            }
        }
        None => (),
    }
    for op in operands {
        for e in op {
            result.push(*e);
        }
    }
    Some(result)
}

#[test]
fn test_column_family_with_options() {
    let n = DBPath::new("_rust_rocksdb_cf_with_optionstest");
    {
        let mut cfopts = Options::default();
        cfopts.set_max_write_buffer_number(16);
        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);

        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cfs = vec![cf_descriptor];
        match DB::open_cf_descriptors(&opts, &n, cfs) {
            Ok(_db) => println!("created db with column family descriptors successfully"),
            Err(e) => {
                panic!(
                    "could not create new database with column family descriptors: {}",
                    e
                );
            }
        }
    }

    {
        let mut cfopts = Options::default();
        cfopts.set_max_write_buffer_number(16);
        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);

        let opts = Options::default();
        let cfs = vec![cf_descriptor];

        match DB::open_cf_descriptors(&opts, &n, cfs) {
            Ok(_db) => println!("successfully re-opened database with column family descriptors"),
            Err(e) => {
                panic!(
                    "unable to re-open database with column family descriptors: {}",
                    e
                );
            }
        }
    }
}

#[test]
fn test_create_duplicate_column_family() {
    let n = DBPath::new("_rust_rocksdb_create_duplicate_column_family");
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let db = match DB::open_cf(&opts, &n, &["cf1"]) {
            Ok(d) => d,
            Err(e) => panic!("failed to create new column family: {}", e),
        };

        assert!(db.create_cf("cf1", &opts).is_err());
    }
}
<|endoftext|>