<commit_before>use std::iter;

use either::Either;
use hir::{AsName, EnumVariant, Module, ModuleDef, Name};
use ide_db::{defs::Definition, search::Reference, RootDatabase};
use rustc_hash::{FxHashMap, FxHashSet};
use syntax::{
    algo::{find_node_at_offset, SyntaxRewriter},
    ast::{self, edit::IndentLevel, make, AstNode, NameOwner, VisibilityOwner},
    SourceFile, SyntaxElement, SyntaxNode, T,
};

use crate::{
    utils::{insert_use, mod_path_to_ast, ImportScope},
    AssistContext, AssistId, AssistKind, Assists,
};

// Assist: extract_struct_from_enum_variant
//
// Extracts a struct from enum variant.
//
// ```
// enum A { <|>One(u32, u32) }
// ```
// ->
// ```
// struct One(pub u32, pub u32);
//
// enum A { One(One) }
// ```
pub(crate) fn extract_struct_from_enum_variant(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let field_list = extract_field_list_if_applicable(&variant)?;

    let variant_name = variant.name()?;
    let variant_hir = ctx.sema.to_def(&variant)?;
    if existing_definition(ctx.db(), &variant_name, &variant_hir) {
        return None;
    }

    let enum_ast = variant.parent_enum();
    let enum_hir = ctx.sema.to_def(&enum_ast)?;
    let target = variant.syntax().text_range();
    acc.add(
        AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
        "Extract struct from enum variant",
        target,
        |builder| {
            let variant_hir_name = variant_hir.name(ctx.db());
            let enum_module_def = ModuleDef::from(enum_hir);
            let usages =
                Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir)).usages(&ctx.sema).all();

            let mut visited_modules_set = FxHashSet::default();
            let current_module = enum_hir.module(ctx.db());
            visited_modules_set.insert(current_module);
            let mut rewriters = FxHashMap::default();
            for reference in usages {
                let rewriter = rewriters
                    .entry(reference.file_range.file_id)
                    .or_insert_with(SyntaxRewriter::default);
                let source_file = ctx.sema.parse(reference.file_range.file_id);
                update_reference(
                    ctx,
                    rewriter,
                    reference,
                    &source_file,
                    &enum_module_def,
                    &variant_hir_name,
                    &mut visited_modules_set,
                );
            }
            let mut rewriter =
                rewriters.remove(&ctx.frange.file_id).unwrap_or_else(SyntaxRewriter::default);
            for (file_id, rewriter) in rewriters {
                builder.edit_file(file_id);
                builder.rewrite(rewriter);
            }
            builder.edit_file(ctx.frange.file_id);
            update_variant(&mut rewriter, &variant);
            extract_struct_def(
                &mut rewriter,
                &enum_ast,
                variant_name.clone(),
                &field_list,
                &variant.parent_enum().syntax().clone().into(),
                enum_ast.visibility(),
            );
            builder.rewrite(rewriter);
        },
    )
}

fn extract_field_list_if_applicable(
    variant: &ast::Variant,
) -> Option<Either<ast::RecordFieldList, ast::TupleFieldList>> {
    match variant.kind() {
        ast::StructKind::Record(field_list) if field_list.fields().next().is_some() => {
            Some(Either::Left(field_list))
        }
        ast::StructKind::Tuple(field_list) if field_list.fields().count() > 1 => {
            Some(Either::Right(field_list))
        }
        _ => None,
    }
}

fn existing_definition(db: &RootDatabase, variant_name: &ast::Name, variant: &EnumVariant) -> bool {
    variant
        .parent_enum(db)
        .module(db)
        .scope(db, None)
        .into_iter()
        .filter(|(_, def)| match def {
            // only check type-namespace
            hir::ScopeDef::ModuleDef(def) => matches!(def,
                ModuleDef::Module(_) | ModuleDef::Adt(_) |
                ModuleDef::EnumVariant(_) | ModuleDef::Trait(_) |
                ModuleDef::TypeAlias(_) | ModuleDef::BuiltinType(_)
            ),
            _ => false,
        })
        .any(|(name, _)| name == variant_name.as_name())
}

fn insert_import(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    scope_node: &SyntaxNode,
    module: &Module,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
) -> Option<()> {
    let db = ctx.db();
    let mod_path = module.find_use_path_prefixed(
        db,
        enum_module_def.clone(),
        ctx.config.insert_use.prefix_kind,
    );
    if let Some(mut mod_path) = mod_path.filter(|path| !path.is_ident()) {
        mod_path.segments.pop();
        mod_path.segments.push(variant_hir_name.clone());
        let scope = ImportScope::find_insert_use_container(scope_node, ctx)?;

        *rewriter += insert_use(&scope, mod_path_to_ast(&mod_path), ctx.config.insert_use.merge);
    }
    Some(())
}

fn extract_struct_def(
    rewriter: &mut SyntaxRewriter,
    enum_: &ast::Enum,
    variant_name: ast::Name,
    field_list: &Either<ast::RecordFieldList, ast::TupleFieldList>,
    start_offset: &SyntaxElement,
    visibility: Option<ast::Visibility>,
) -> Option<()> {
    let pub_vis = Some(make::visibility_pub());
    let field_list = match field_list {
        Either::Left(field_list) => {
            make::record_field_list(field_list.fields().flat_map(|field| {
                Some(make::record_field(pub_vis.clone(), field.name()?, field.ty()?))
            }))
            .into()
        }
        Either::Right(field_list) => make::tuple_field_list(
            field_list
                .fields()
                .flat_map(|field| Some(make::tuple_field(pub_vis.clone(), field.ty()?))),
        )
        .into(),
    };

    rewriter.insert_before(
        start_offset,
        make::struct_(visibility, variant_name, None, field_list).syntax(),
    );
    rewriter.insert_before(start_offset, &make::tokens::blank_line());

    if let indent_level @ 1..=usize::MAX = IndentLevel::from_node(enum_.syntax()).0 as usize {
        rewriter
            .insert_before(start_offset, &make::tokens::whitespace(&" ".repeat(4 * indent_level)));
    }
    Some(())
}

fn update_variant(rewriter: &mut SyntaxRewriter, variant: &ast::Variant) -> Option<()> {
    let name = variant.name()?;
    let tuple_field = make::tuple_field(None, make::ty(name.text()));
    let replacement = make::variant(
        name,
        Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
    );
    rewriter.replace(variant.syntax(), replacement.syntax());
    Some(())
}

fn update_reference(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    reference: Reference,
    source_file: &SourceFile,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
    visited_modules_set: &mut FxHashSet<Module>,
) -> Option<()> {
    let offset = reference.file_range.range.start();
    let (segment, expr) = if let Some(path_expr) =
        find_node_at_offset::<ast::PathExpr>(source_file.syntax(), offset)
    {
        // tuple variant
        (path_expr.path()?.segment()?, path_expr.syntax().parent()?.clone())
    } else if let Some(record_expr) =
        find_node_at_offset::<ast::RecordExpr>(source_file.syntax(), offset)
    {
        // record variant
        (record_expr.path()?.segment()?, record_expr.syntax().clone())
    } else {
        return None;
    };

    let module = ctx.sema.scope(&expr).module()?;
    if !visited_modules_set.contains(&module) {
        if insert_import(ctx, rewriter, &expr, &module, enum_module_def, variant_hir_name).is_some()
        {
            visited_modules_set.insert(module);
        }
    }
    rewriter.insert_after(segment.syntax(), &make::token(T!['(']));
    rewriter.insert_after(segment.syntax(), segment.syntax());
    rewriter.insert_after(&expr, &make::token(T![')']));
    Some(())
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_assist, check_assist_not_applicable},
        utils::FamousDefs,
    };

    use super::*;

    #[test]
    fn test_extract_struct_several_fields_tuple() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32, u32) }",
            r#"struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_several_fields_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32, bar: u32 } }",
            r#"struct One{ pub foo: u32, pub bar: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_one_field_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32 } }",
            r#"struct One{ pub foo: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_enum_variant_name_value_namespace() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"const One: () = ();
enum A { <|>One(u32, u32) }"#,
            r#"const One: () = ();
struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_pub_visibility() {
        check_assist(
            extract_struct_from_enum_variant,
            "pub enum A { <|>One(u32, u32) }",
            r#"pub struct One(pub u32, pub u32);

pub enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_with_complex_imports() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"mod my_mod {
    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(1, 1);
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(1, 1);
        }

        pub enum MyEnum {
            <|>MyField(u8, u8),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
}"#,
            r#"use my_mod::my_other_mod::MyField;

mod my_mod {
    use self::my_other_mod::MyField;

    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(MyField(1, 1));
        }

        pub struct MyField(pub u8, pub u8);

        pub enum MyEnum {
            MyField(MyField),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
}"#,
        );
    }

    #[test]
    fn extract_record_fix_references() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
enum E {
    <|>V { i: i32, j: i32 }
}

fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}

fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    #[test]
    fn test_several_files() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V(i32, i32)
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V(9, 2);
}
"#,
            r#"
//- /main.rs
struct V(pub i32, pub i32);

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::{E, V};
fn f() {
    let e = E::V(V(9, 2));
}
"#,
        )
    }

    #[test]
    fn test_several_files_record() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V { i: i32, j: i32 }
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
//- /main.rs
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::{E, V};
fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    fn check_not_applicable(ra_fixture: &str) {
        let fixture =
            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
        check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
    }

    #[test]
    fn test_extract_enum_not_applicable_for_element_with_no_fields() {
        check_not_applicable("enum A { <|>One }");
    }

    #[test]
    fn test_extract_enum_not_applicable_if_struct_exists() {
        check_not_applicable(
            r#"struct One;
        enum A { <|>One(u8, u32) }"#,
        );
    }

    #[test]
    fn test_extract_not_applicable_one_field() {
        check_not_applicable(r"enum A { <|>One(u32) }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_tuple() {
        check_not_applicable(r"enum A { <|>None() }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_named() {
        check_not_applicable(r"enum A { <|>None {} }");
    }
}
<commit_msg>Add extra test to extract_struct_from_enum_variant<commit_after>use std::iter;

use either::Either;
use hir::{AsName, EnumVariant, Module, ModuleDef, Name};
use ide_db::{defs::Definition, search::Reference, RootDatabase};
use rustc_hash::{FxHashMap, FxHashSet};
use syntax::{
    algo::{find_node_at_offset, SyntaxRewriter},
    ast::{self, edit::IndentLevel, make, AstNode, NameOwner, VisibilityOwner},
    SourceFile, SyntaxElement, SyntaxNode, T,
};

use crate::{
    utils::{insert_use, mod_path_to_ast, ImportScope},
    AssistContext, AssistId, AssistKind, Assists,
};

// Assist: extract_struct_from_enum_variant
//
// Extracts a struct from enum variant.
//
// ```
// enum A { <|>One(u32, u32) }
// ```
// ->
// ```
// struct One(pub u32, pub u32);
//
// enum A { One(One) }
// ```
pub(crate) fn extract_struct_from_enum_variant(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let field_list = extract_field_list_if_applicable(&variant)?;

    let variant_name = variant.name()?;
    let variant_hir = ctx.sema.to_def(&variant)?;
    if existing_definition(ctx.db(), &variant_name, &variant_hir) {
        return None;
    }

    let enum_ast = variant.parent_enum();
    let enum_hir = ctx.sema.to_def(&enum_ast)?;
    let target = variant.syntax().text_range();
    acc.add(
        AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
        "Extract struct from enum variant",
        target,
        |builder| {
            let variant_hir_name = variant_hir.name(ctx.db());
            let enum_module_def = ModuleDef::from(enum_hir);
            let usages =
                Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir)).usages(&ctx.sema).all();

            let mut visited_modules_set = FxHashSet::default();
            let current_module = enum_hir.module(ctx.db());
            visited_modules_set.insert(current_module);
            let mut rewriters = FxHashMap::default();
            for reference in usages {
                let rewriter = rewriters
                    .entry(reference.file_range.file_id)
                    .or_insert_with(SyntaxRewriter::default);
                let source_file = ctx.sema.parse(reference.file_range.file_id);
                update_reference(
                    ctx,
                    rewriter,
                    reference,
                    &source_file,
                    &enum_module_def,
                    &variant_hir_name,
                    &mut visited_modules_set,
                );
            }
            let mut rewriter =
                rewriters.remove(&ctx.frange.file_id).unwrap_or_else(SyntaxRewriter::default);
            for (file_id, rewriter) in rewriters {
                builder.edit_file(file_id);
                builder.rewrite(rewriter);
            }
            builder.edit_file(ctx.frange.file_id);
            update_variant(&mut rewriter, &variant);
            extract_struct_def(
                &mut rewriter,
                &enum_ast,
                variant_name.clone(),
                &field_list,
                &variant.parent_enum().syntax().clone().into(),
                enum_ast.visibility(),
            );
            builder.rewrite(rewriter);
        },
    )
}

fn extract_field_list_if_applicable(
    variant: &ast::Variant,
) -> Option<Either<ast::RecordFieldList, ast::TupleFieldList>> {
    match variant.kind() {
        ast::StructKind::Record(field_list) if field_list.fields().next().is_some() => {
            Some(Either::Left(field_list))
        }
        ast::StructKind::Tuple(field_list) if field_list.fields().count() > 1 => {
            Some(Either::Right(field_list))
        }
        _ => None,
    }
}

fn existing_definition(db: &RootDatabase, variant_name: &ast::Name, variant: &EnumVariant) -> bool {
    variant
        .parent_enum(db)
        .module(db)
        .scope(db, None)
        .into_iter()
        .filter(|(_, def)| match def {
            // only check type-namespace
            hir::ScopeDef::ModuleDef(def) => matches!(def,
                ModuleDef::Module(_) | ModuleDef::Adt(_) |
                ModuleDef::EnumVariant(_) | ModuleDef::Trait(_) |
                ModuleDef::TypeAlias(_) | ModuleDef::BuiltinType(_)
            ),
            _ => false,
        })
        .any(|(name, _)| name == variant_name.as_name())
}

fn insert_import(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    scope_node: &SyntaxNode,
    module: &Module,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
) -> Option<()> {
    let db = ctx.db();
    let mod_path = module.find_use_path_prefixed(
        db,
        enum_module_def.clone(),
        ctx.config.insert_use.prefix_kind,
    );
    if let Some(mut mod_path) = mod_path {
        mod_path.segments.pop();
        mod_path.segments.push(variant_hir_name.clone());
        let scope = ImportScope::find_insert_use_container(scope_node, ctx)?;

        *rewriter += insert_use(&scope, mod_path_to_ast(&mod_path), ctx.config.insert_use.merge);
    }
    Some(())
}

fn extract_struct_def(
    rewriter: &mut SyntaxRewriter,
    enum_: &ast::Enum,
    variant_name: ast::Name,
    field_list: &Either<ast::RecordFieldList, ast::TupleFieldList>,
    start_offset: &SyntaxElement,
    visibility: Option<ast::Visibility>,
) -> Option<()> {
    let pub_vis = Some(make::visibility_pub());
    let field_list = match field_list {
        Either::Left(field_list) => {
            make::record_field_list(field_list.fields().flat_map(|field| {
                Some(make::record_field(pub_vis.clone(), field.name()?, field.ty()?))
            }))
            .into()
        }
        Either::Right(field_list) => make::tuple_field_list(
            field_list
                .fields()
                .flat_map(|field| Some(make::tuple_field(pub_vis.clone(), field.ty()?))),
        )
        .into(),
    };

    rewriter.insert_before(
        start_offset,
        make::struct_(visibility, variant_name, None, field_list).syntax(),
    );
    rewriter.insert_before(start_offset, &make::tokens::blank_line());

    if let indent_level @ 1..=usize::MAX = IndentLevel::from_node(enum_.syntax()).0 as usize {
        rewriter
            .insert_before(start_offset, &make::tokens::whitespace(&" ".repeat(4 * indent_level)));
    }
    Some(())
}

fn update_variant(rewriter: &mut SyntaxRewriter, variant: &ast::Variant) -> Option<()> {
    let name = variant.name()?;
    let tuple_field = make::tuple_field(None, make::ty(name.text()));
    let replacement = make::variant(
        name,
        Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
    );
    rewriter.replace(variant.syntax(), replacement.syntax());
    Some(())
}

fn update_reference(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    reference: Reference,
    source_file: &SourceFile,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
    visited_modules_set: &mut FxHashSet<Module>,
) -> Option<()> {
    let offset = reference.file_range.range.start();
    let (segment, expr) = if let Some(path_expr) =
        find_node_at_offset::<ast::PathExpr>(source_file.syntax(), offset)
    {
        // tuple variant
        (path_expr.path()?.segment()?, path_expr.syntax().parent()?.clone())
    } else if let Some(record_expr) =
        find_node_at_offset::<ast::RecordExpr>(source_file.syntax(), offset)
    {
        // record variant
        (record_expr.path()?.segment()?, record_expr.syntax().clone())
    } else {
        return None;
    };

    let module = ctx.sema.scope(&expr).module()?;
    if !visited_modules_set.contains(&module) {
        if insert_import(ctx, rewriter, &expr, &module, enum_module_def, variant_hir_name).is_some()
        {
            visited_modules_set.insert(module);
        }
    }
    rewriter.insert_after(segment.syntax(), &make::token(T!['(']));
    rewriter.insert_after(segment.syntax(), segment.syntax());
    rewriter.insert_after(&expr, &make::token(T![')']));
    Some(())
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_assist, check_assist_not_applicable},
        utils::FamousDefs,
    };

    use super::*;

    #[test]
    fn test_extract_struct_several_fields_tuple() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32, u32) }",
            r#"struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_several_fields_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32, bar: u32 } }",
            r#"struct One{ pub foo: u32, pub bar: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_one_field_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32 } }",
            r#"struct One{ pub foo: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_enum_variant_name_value_namespace() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"const One: () = ();
enum A { <|>One(u32, u32) }"#,
            r#"const One: () = ();
struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_pub_visibility() {
        check_assist(
            extract_struct_from_enum_variant,
            "pub enum A { <|>One(u32, u32) }",
            r#"pub struct One(pub u32, pub u32);

pub enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_with_complex_imports() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"mod my_mod {
    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(1, 1);
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(1, 1);
        }

        pub enum MyEnum {
            <|>MyField(u8, u8),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
}"#,
            r#"use my_mod::my_other_mod::MyField;

mod my_mod {
    use self::my_other_mod::MyField;

    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(MyField(1, 1));
        }

        pub struct MyField(pub u8, pub u8);

        pub enum MyEnum {
            MyField(MyField),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
}"#,
        );
    }

    #[test]
    fn extract_record_fix_references() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
enum E {
    <|>V { i: i32, j: i32 }
}

fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}

fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    #[test]
    fn test_several_files() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V(i32, i32)
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V(9, 2);
}
"#,
            r#"
//- /main.rs
struct V(pub i32, pub i32);

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::{E, V};
fn f() {
    let e = E::V(V(9, 2));
}
"#,
        )
    }

    #[test]
    fn test_several_files_record() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V { i: i32, j: i32 }
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
//- /main.rs
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::{E, V};
fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    #[test]
    fn test_extract_struct_record_nested_call_exp() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
enum A { <|>One { a: u32, b: u32 } }

struct B(A);

fn foo() {
    let _ = B(A::One { a: 1, b: 2 });
}
"#,
            r#"
struct One{ pub a: u32, pub b: u32 }

enum A { One(One) }

struct B(A);

fn foo() {
    let _ = B(A::One(One { a: 1, b: 2 }));
}
"#,
        );
    }

    fn check_not_applicable(ra_fixture: &str) {
        let fixture =
            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
        check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
    }

    #[test]
    fn test_extract_enum_not_applicable_for_element_with_no_fields() {
        check_not_applicable("enum A { <|>One }");
    }

    #[test]
    fn test_extract_enum_not_applicable_if_struct_exists() {
        check_not_applicable(
            r#"struct One;
        enum A { <|>One(u8, u32) }"#,
        );
    }

    #[test]
    fn test_extract_not_applicable_one_field() {
        check_not_applicable(r"enum A { <|>One(u32) }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_tuple() {
        check_not_applicable(r"enum A { <|>None() }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_named() {
        check_not_applicable(r"enum A { <|>None {} }");
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use std::path::PathBuf;

use rustc_hash::{FxHashMap, FxHashSet};
use serde::Deserialize;

/// A root points to the directory which contains Rust crates. rust-analyzer watches all files in
/// all roots. Roots might be nested.
#[derive(Clone, Debug, Deserialize)]
#[serde(transparent)]
pub struct Root {
    pub(crate) path: PathBuf,
}

/// A crate points to the root module of a crate and lists the dependencies of the crate. This is
/// useful in creating the crate graph.
#[derive(Clone, Debug, Deserialize)]
pub struct Crate {
    pub(crate) root_module: PathBuf,
    pub(crate) edition: Edition,
    pub(crate) deps: Vec<Dep>,

    // This is the preferred method of providing cfg options.
    #[serde(default)]
    pub(crate) cfg: FxHashSet<String>,

    // These two are here for transition only.
    #[serde(default)]
    pub(crate) atom_cfgs: FxHashSet<String>,
    #[serde(default)]
    pub(crate) key_value_cfgs: FxHashMap<String, String>,

    pub(crate) out_dir: Option<PathBuf>,
    pub(crate) proc_macro_dylib_path: Option<PathBuf>,
}

#[derive(Clone, Copy, Debug, Deserialize)]
#[serde(rename = "edition")]
pub enum Edition {
    #[serde(rename = "2015")]
    Edition2015,
    #[serde(rename = "2018")]
    Edition2018,
}

/// Identifies a crate by position in the crates array.
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[serde(transparent)]
pub struct CrateId(pub usize);

/// A dependency of a crate, identified by its id in the crates array and name.
#[derive(Clone, Debug, Deserialize)]
pub struct Dep {
    #[serde(rename = "crate")]
    pub(crate) krate: CrateId,
    pub(crate) name: String,
}

/// Roots and crates that compose this Rust project.
#[derive(Clone, Debug, Deserialize)]
pub struct JsonProject {
    pub(crate) roots: Vec<Root>,
    pub(crate) crates: Vec<Crate>,
}

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

    #[test]
    fn test_crate_deserialization() {
        let raw_json = json!(    {
            "crate_id": 2,
            "root_module": "this/is/a/file/path.rs",
            "deps": [
              {
                "crate": 1,
                "name": "some_dep_crate"
              },
            ],
            "edition": "2015",
            "cfg": [
              "atom_1",
              "atom_2",
              "feature=feature_1",
              "feature=feature_2",
              "other=value",
            ],

        });

        let krate: Crate = serde_json::from_value(raw_json).unwrap();

        assert!(krate.cfg.contains(&"atom_1".to_string()));
        assert!(krate.cfg.contains(&"atom_2".to_string()));
        assert!(krate.cfg.contains(&"feature=feature_1".to_string()));
        assert!(krate.cfg.contains(&"feature=feature_2".to_string()));
        assert!(krate.cfg.contains(&"other=value".to_string()));
    }

    #[test]
    fn test_crate_deserialization_old_json() {
        let raw_json = json!(    {
           "crate_id": 2,
           "root_module": "this/is/a/file/path.rs",
           "deps": [
             {
               "crate": 1,
               "name": "some_dep_crate"
             },
           ],
           "edition": "2015",
           "atom_cfgs": [
             "atom_1",
             "atom_2",
           ],
           "key_value_cfgs": {
             "feature": "feature_1",
             "feature": "feature_2",
             "other": "value",
           },
        });

        let krate: Crate = serde_json::from_value(raw_json).unwrap();

        assert!(krate.atom_cfgs.contains(&"atom_1".to_string()));
        assert!(krate.atom_cfgs.contains(&"atom_2".to_string()));
        assert!(krate.key_value_cfgs.contains_key(&"feature".to_string()));
        assert_eq!(krate.key_value_cfgs.get("feature"), Some(&"feature_2".to_string()));
        assert!(krate.key_value_cfgs.contains_key(&"other".to_string()));
        assert_eq!(krate.key_value_cfgs.get("other"), Some(&"value".to_string()));
    }
}
<commit_msg>Put important things on top<commit_after>//! FIXME: write short doc here

use std::path::PathBuf;

use rustc_hash::{FxHashMap, FxHashSet};
use serde::Deserialize;

/// Roots and crates that compose this Rust project.
#[derive(Clone, Debug, Deserialize)]
pub struct JsonProject {
    pub(crate) roots: Vec<Root>,
    pub(crate) crates: Vec<Crate>,
}

/// A root points to the directory which contains Rust crates. rust-analyzer watches all files in
/// all roots. Roots might be nested.
#[derive(Clone, Debug, Deserialize)]
#[serde(transparent)]
pub struct Root {
    pub(crate) path: PathBuf,
}

/// A crate points to the root module of a crate and lists the dependencies of the crate. This is
/// useful in creating the crate graph.
#[derive(Clone, Debug, Deserialize)]
pub struct Crate {
    pub(crate) root_module: PathBuf,
    pub(crate) edition: Edition,
    pub(crate) deps: Vec<Dep>,

    // This is the preferred method of providing cfg options.
    #[serde(default)]
    pub(crate) cfg: FxHashSet<String>,

    // These two are here for transition only.
    #[serde(default)]
    pub(crate) atom_cfgs: FxHashSet<String>,
    #[serde(default)]
    pub(crate) key_value_cfgs: FxHashMap<String, String>,

    pub(crate) out_dir: Option<PathBuf>,
    pub(crate) proc_macro_dylib_path: Option<PathBuf>,
}

#[derive(Clone, Copy, Debug, Deserialize)]
#[serde(rename = "edition")]
pub enum Edition {
    #[serde(rename = "2015")]
    Edition2015,
    #[serde(rename = "2018")]
    Edition2018,
}

/// Identifies a crate by position in the crates array.
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[serde(transparent)]
pub struct CrateId(pub usize);

/// A dependency of a crate, identified by its id in the crates array and name.
#[derive(Clone, Debug, Deserialize)]
pub struct Dep {
    #[serde(rename = "crate")]
    pub(crate) krate: CrateId,
    pub(crate) name: String,
}

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

    #[test]
    fn test_crate_deserialization() {
        let raw_json = json!(    {
            "crate_id": 2,
            "root_module": "this/is/a/file/path.rs",
            "deps": [
              {
                "crate": 1,
                "name": "some_dep_crate"
              },
            ],
            "edition": "2015",
            "cfg": [
              "atom_1",
              "atom_2",
              "feature=feature_1",
              "feature=feature_2",
              "other=value",
            ],

        });

        let krate: Crate = serde_json::from_value(raw_json).unwrap();

        assert!(krate.cfg.contains(&"atom_1".to_string()));
        assert!(krate.cfg.contains(&"atom_2".to_string()));
        assert!(krate.cfg.contains(&"feature=feature_1".to_string()));
        assert!(krate.cfg.contains(&"feature=feature_2".to_string()));
        assert!(krate.cfg.contains(&"other=value".to_string()));
    }

    #[test]
    fn test_crate_deserialization_old_json() {
        let raw_json = json!(    {
           "crate_id": 2,
           "root_module": "this/is/a/file/path.rs",
           "deps": [
             {
               "crate": 1,
               "name": "some_dep_crate"
             },
           ],
           "edition": "2015",
           "atom_cfgs": [
             "atom_1",
             "atom_2",
           ],
           "key_value_cfgs": {
             "feature": "feature_1",
             "feature": "feature_2",
             "other": "value",
           },
        });

        let krate: Crate = serde_json::from_value(raw_json).unwrap();

        assert!(krate.atom_cfgs.contains(&"atom_1".to_string()));
        assert!(krate.atom_cfgs.contains(&"atom_2".to_string()));
        assert!(krate.key_value_cfgs.contains_key(&"feature".to_string()));
        assert_eq!(krate.key_value_cfgs.get("feature"), Some(&"feature_2".to_string()));
        assert!(krate.key_value_cfgs.contains_key(&"other".to_string()));
        assert_eq!(krate.key_value_cfgs.get("other"), Some(&"value".to_string()));
    }
}
<|endoftext|>


<commit_before>// In order to use YAML to define your CLI you must compile clap with the "yaml" feature becasue
// it's **not** included by default.
//
// In order to do this, ensure your Cargo.toml looks like one of the following:
//
// [dependencies.clap]
// features = ["yaml"]
//
// __OR__
//
// [dependencies]
// clap = { features = ["yaml"] }

// Using yaml requires calling a clap macro `load_yaml!()` so we must use the '#[macro_use]'
// directive
#[macro_use]
extern crate clap;

#[cfg(feature = "yaml")]
fn main() {
    use clap::App;

    // To load a yaml file containing our CLI definition such as the example '17_yaml.yml' we can
    // use the convenience macro which loads the file at compile relative to the current file
    // similiar to how modules are found.
    //
    // Then we pass that yaml object to App to build the CLI.
    //
    // Finally we call get_matches() to start the parsing process. We use the matches just as we
    // normally would
    let yml = load_yaml!("17_yaml.yml");
    let m = App::from_yaml(yml).get_matches();

    // Because the example 17_yaml.yml is rather large we'll just look a single arg so you can
    // see that it works...
    if let Some(mode) = m.value_of("mode") {
        match mode {
            "fast" => println!("We're really going now!"),
            "slow" => println!("Awwww, too slow :("),
            _      => unreachable!()
        }
    } else {
        println!("--mode <MODE> wasn't used...");
    }
}

#[cfg(not(feature = "yaml"))]
fn main() {
    // As stated above, if clap is not compiled with the YAML feature, it is disabled.
    println!("YAML feature is disabled.");
    println!("Pass --features yaml to cargo when trying this example.");
}
<commit_msg>Auto merge of #710 - myCrates:fix/yaml-example, r=kbknapp<commit_after>// In order to use YAML to define your CLI you must compile clap with the "yaml" feature becasue
// it's **not** included by default.
//
// In order to do this, ensure your Cargo.toml looks like one of the following:
//
// [dependencies.clap]
// features = ["yaml"]
//
// __OR__
//
// [dependencies]
// clap = { features = ["yaml"] }

// Using yaml requires calling a clap macro `load_yaml!()` so we must use the '#[macro_use]'
// directive
#[macro_use]
extern crate clap;

#[cfg(feature = "yaml")]
fn main() {
    use clap::App;

    // To load a yaml file containing our CLI definition such as the example '17_yaml.yml' we can
    // use the convenience macro which loads the file at compile relative to the current file
    // similiar to how modules are found.
    //
    // Then we pass that yaml object to App to build the CLI.
    //
    // Finally we call get_matches() to start the parsing process. We use the matches just as we
    // normally would
    let yml = load_yaml!("17_yaml.yml");
    let m = App::from_yaml(yml).get_matches();

    // Because the example 17_yaml.yml is rather large we'll just look a single arg so you can
    // see that it works...
    if let Some(mode) = m.value_of("mode") {
        match mode {
            "vi" => println!("You are using vi"),
            "emacs" => println!("You are using emacs..."),
            _      => unreachable!()
        }
    } else {
        println!("--mode <MODE> wasn't used...");
    }
}

#[cfg(not(feature = "yaml"))]
fn main() {
    // As stated above, if clap is not compiled with the YAML feature, it is disabled.
    println!("YAML feature is disabled.");
    println!("Pass --features yaml to cargo when trying this example.");
}
<|endoftext|>


<commit_before>#[macro_use] extern crate log;
extern crate rusty_raft;
extern crate rand;
extern crate rustc_serialize;
extern crate env_logger;

use rand::{thread_rng, Rng};
use rusty_raft::server::{start_server, ServerHandle};
use rusty_raft::client::{RaftConnection};
use rusty_raft::client::state_machine::{
    StateMachine, RaftStateMachine};
use rusty_raft::common::{Config, RaftError};
use std::env::args;
use std::collections::HashMap;
use std::fmt::Debug;
use std::net::SocketAddr;
use std::str;
use std::str::FromStr;
use std::time::Duration;

use std::fs::File;
use std::io::{stdin, stdout, SeekFrom, Error, ErrorKind, Read, BufReader, Seek, BufRead, Cursor, Write};

use rustc_serialize::json;

static USAGE: &'static str = "
Commands:
  server  Starts a server from an initial cluster config.
  client  Starts a client repl, that communicates with
          a particular cluster.

Usage:
  hashmap server <id> <filename>
  hashmap servers <filename>
  hashmap client <filename>

Options:
  -h --help   Show a help message.
";

fn main() {
    env_logger::init().unwrap();
    // TODO (sydli) make prettier
    let mut args = args();
    if let Some(command) = args.nth(1) {
        if &command == "server" {
            if let (Some(id_str), Some(filename)) =
                   (args.next(), args.next()) {
                if let Ok(id) = id_str.parse::<u64>() {
                    Server::new(id, &cluster_from_file(&filename)).repl();
                    return;
                }
            }
        } else if &command == "client" {
            if let Some(filename) = args.next() {
                Client::new(&cluster_from_file(&filename)).repl();
                return;
            }
        } else if &command == "servers" {
            if let Some(filename) = args.next() {
                Cluster::new(&cluster_from_file(&filename)).repl();
                return;
            }
        }
    }
    println!("Incorrect usage. \n{}", USAGE);
}

///
/// Automatically builds a repl loop for an object, if it implements |exec|.
/// Each line a user types into the repl is fed to |exec|.
///
/// Default commands:
///     exit    Exits the loop and shuts down associated process.
///     help    Displays the |help| message for more usage info.
trait Repl {

    ///
    /// Processes each input command from the Repl.
    /// Returns true if |command| is formatted correctly.
    ///
    fn exec(&mut self, command: String) -> bool;

    ///
    /// Usage string displayed with the default command "help".
    ///
    fn usage(&self) -> String;

    fn print_usage(&self) {
        println!("\nREPL COMMANDS:\n==============\n{}\n{}\n{}",
                 "exit\n\t\tExits the loop and shuts down associated process.",
                 "help\n\t\tSpits out this message.",
                 self.usage());
    }

    ///
    /// Implementation of the repl.
    fn repl(&mut self) {
        println!("[ Starting repl ]");
        loop {
            print!("> ");
            stdout().flush().unwrap();
            let mut buffer = String::new();
            if stdin().read_line(&mut buffer).is_ok() {
                let words: Vec<String> =
                    buffer.split_whitespace().map(String::from).collect();
                if words.get(0) == Some(&String::from("exit")) { break; }
                if words.get(0) == Some(&String::from("help")) {
                    self.print_usage();
                    continue;
                }
                if !self.exec(buffer.clone()) {
                    println!("Command not recognized.");
                    self.print_usage();
                }
            }
        }
    }
}

struct Server {
    handle: ServerHandle,
}

impl Repl for Server {
    fn exec(&mut self, command: String) -> bool { true }
    fn usage(&self) -> String { String::from("") }
}

impl Server {
    fn new(id: u64, cluster: &HashMap<u64, SocketAddr>) -> Server {
        if !cluster.contains_key(&id) {
            panic!("Initial cluster must contain server id: {}", id);
        }
        Server {
            handle:
                // TODO: Only the first server should put itself in its log file
                start_server(id, Box::new(RaftHashMap { map: HashMap::new() }), *cluster.get(&id).unwrap()).unwrap()
        }
    }
}

struct Cluster {
    servers: HashMap<u64, Server>,
    cluster: HashMap<u64, SocketAddr>,
}

impl Cluster {
    fn new(cluster: &HashMap<u64, SocketAddr>) -> Cluster {
        Cluster {
            servers: cluster.iter()
                .map(|(id, _)| (*id, Server::new(*id, &cluster.clone())))
                .collect::<HashMap<u64, Server>>(),
            cluster: cluster.clone(),
        }
    }

    fn add_server(&mut self, id: u64, addr: SocketAddr) {
        println!("Adding server {}, {}", id, addr);
    }

    fn remove_server(&mut self, id: u64) {
        if !self.cluster.contains_key(&id) {
            println!("Server {} is not in the cluster! Servers: {:?}", id,
                     self.cluster);
        }
        println!("Removing server {}", id);
    }

    fn kill_server(&mut self, id: u64) {
        if !self.servers.contains_key(&id) {
            println!("Server {} is not up right now!", id);
        }
        {
            self.servers.remove(&id).unwrap();
        } // drop server
        println!("Killed server {}", id);
    }

    fn start_server(&mut self, id: u64) {
        if self.servers.contains_key(&id) {
            println!("Server {} is already up!", id);
        }
        self.servers.insert(id, Server::new(id, &self.cluster.clone()));
        println!("Restarted server {}", id);
    }

    fn print_servers(&self) {
        println!("Servers in cluster: {:?}\n Live servers: {:?}",
                 self.cluster, self.servers.keys().collect::<Vec<&u64>>());
    }
}

impl Repl for Cluster {
    fn exec(&mut self, command: String) -> bool {
        let words: Vec<&str> =
            command.split_whitespace().collect();
        let first = words.get(0).map(|x|*x);
        if first == Some("add")  {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            let addr = words.get(2).and_then(|x| as_addr(*x).ok());
            if num.and(addr).is_none() { return false; }
            self.add_server(num.unwrap(), addr.unwrap());
        } else if first == Some("remove") {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            if num.is_none() { return false; }
            self.remove_server(num.unwrap());
        } else if first == Some("kill") {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            if num.is_none() { return false; }
            self.kill_server(num.unwrap());
        } else if first == Some("start") {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            if num.is_none() { return false; }
            self.start_server(num.unwrap());
        } else if first == Some("list") {
            self.print_servers();
        } else { return false; }
        return true;
    }

    fn usage(&self) -> String {
        String::from(format!(
                "{}\n{}\n{}\n{}\n{}",
                "add <node-id> <node-addr>\tAdds a server to the cluster.",
                "remove <node-id>\tRemoves a server from the cluster.",
                "start <node-id>\tStarts a server from the cluster.",
                "kill <node-id>\tKills a server from the cluster.",
                "list\tLists all cluster servers, and the ones that are live."))
    }
}

struct Client {
    raft: RaftConnection,
}

impl Repl for Client {
    fn exec (&mut self, command: String) -> bool {
        let words: Vec<String> =
            command.split_whitespace().map(String::from).collect();
        self.process_command(words)
    }
    fn usage(&self) -> String {
        String::from(format!(
                "{}\n{}",
                 "get <key>\n\t\tPrints value of <key> in hashmap, if it exists.",
                 "put <key> <value>\n\t\tPuts value of <key> as <value>"))
    }
}

#[derive(RustcDecodable, RustcEncodable)]
struct Put {
    key: String,
    value: String,
}

impl Client {
    fn new(cluster: &HashMap<u64, SocketAddr>) -> Client {
        let connection = RaftConnection::new_with_session(cluster);
        if connection.is_none() {
            println!("Couldn't establish connection to cluster at {:?}", cluster);
            panic!();
        }
        Client { raft: RaftConnection::new_with_session(cluster).unwrap() }
    }

    fn get(&mut self, key: String) -> Result<String, RaftError> {
        self.raft.query(key.as_bytes())
            .and_then(
                |result| str::from_utf8(&result)
                             .map(str::to_string)
                             .map_err(deserialize_error))
    }

    fn put(&mut self, key: String, value: String) -> Result<(), RaftError> {
        json::encode(&Put {key:key, value:value})
            .map_err(serialize_error)
            .and_then(|buffer| self.raft.command(&buffer.as_bytes()))
    }

    // TODO (sydli): make this function less shit
    fn process_command(&mut self, words: Vec<String>)
        -> bool {
        if words.len() == 0 { return true; }
        let ref cmd = words[0];
        if *cmd == String::from("get") {
            if words.len() <= 1 { return false; }
            words.get(1).map(|key| {
                match self.get(key.clone()) {
                    Ok(value) => println!("Value for {} is {}", key, value),
                    Err(err) => println!("Error during get: {:?}", err),
                }
            }).unwrap();
        } else if *cmd == String::from("put") {
            if words.len() <= 2 { return false; }
            words.get(1).map(|key| { words.get(2).map(|val| {
                match self.put(key.clone(), val.clone()) {
                    Ok(()) => println!("Put {} => {} successfully", key, val),
                    Err(err) => println!("Error during put: {:?}", err),
                }
            }).unwrap()}).unwrap();
        } else { return false; }
        return true;
    }
}

fn io_err() -> std::io::Error {
    Error::new(ErrorKind::InvalidData, "File incorrectly formatted")
}

fn as_num(x: &str) -> Result<u64, std::io::Error> {
    String::from(x).parse::<u64>().map_err(|_| io_err())
}

fn as_addr(x: &str) -> Result<SocketAddr, std::io::Error> {
    String::from(x).parse::<SocketAddr>().map_err(|_| io_err())
}

/// Panics on io error (we can't access the cluster info!)
/// TODO (sydli) make io_errs more informative
// TODO(jason): Remove this and bootstrap a dynamic cluster
fn cluster_from_file(filename: &str) -> HashMap<u64, SocketAddr> {
    let file = File::open(filename.clone())
               .expect(&format!("Unable to open file {}", filename));
    let mut lines = BufReader::new(file).lines();
    lines.next().ok_or(io_err())
    .map(|line_or_io_error| line_or_io_error.unwrap())
    .and_then(|x| as_num(&x))
    .and_then(|num| (0..num).map(|_| {
        lines.next().ok_or(io_err())
        .map(|line_or_io_error| line_or_io_error.unwrap())
        .and_then(|node_str| {
            let mut words = node_str.split_whitespace().map(String::from);
            let id = words.next().ok_or(io_err())
                        .and_then(|x| as_num(&x));
            words.next().ok_or(io_err())
                .and_then(|x| as_addr(&x))
                .and_then(move |addr| id.map(|id| (id, addr)))
        })
    }).collect::<Result<Vec<_>, _>>())
    .map(|nodes: Vec<(u64, SocketAddr)>|
        nodes.iter().cloned().collect::<HashMap<u64, SocketAddr>>()
    ).unwrap()
}

struct RaftHashMap {
    map: HashMap<String, String>,
}

fn serialize_error <T: Debug>(error: T) -> RaftError {
        RaftError::ClientError(
            format!("Couldn't serialize object. Error: {:?}", error))
}

fn deserialize_error <T: Debug>(error: T) -> RaftError {
        RaftError::ClientError(
            format!("Couldn't deserialize buffer. Error: {:?}", error))
}

fn key_error(key: &String) -> RaftError {
        RaftError::ClientError(format!("Couldn't find key {}", key))
}

impl StateMachine for RaftHashMap {
    fn command (&mut self, buffer: &[u8]) ->Result<(), RaftError> {
        str::from_utf8(buffer)
            .map_err(deserialize_error)
            .and_then(|string| json::decode(&string)
                                    .map_err(deserialize_error))
            .map(|put: Put|
                 {
                     self.map.insert(put.key, put.value);
                 })
    }

    fn query (& self, buffer: &[u8]) ->Result<Vec<u8>, RaftError> {
        str::from_utf8(buffer)
            .map_err(deserialize_error)
            .and_then(|key| {
                let key = key.to_string();
                self.map.get(&key)
                    .map(|x| x.as_bytes().to_vec())
                    .ok_or(key_error(&key))
            })
    }
}

<commit_msg>move example to dynamic cluster<commit_after>#[macro_use] extern crate log;
extern crate rusty_raft;
extern crate rand;
extern crate rustc_serialize;

use rand::{thread_rng, Rng};
use rusty_raft::server::{start_server, ServerHandle};
use rusty_raft::client::{RaftConnection};
use rusty_raft::client::state_machine::{
    StateMachine, RaftStateMachine};
use rusty_raft::common::{Config, RaftError};
use std::env::args;
use std::collections::HashMap;
use std::fmt::Debug;
use std::net::SocketAddr;
use std::str;
use std::str::FromStr;
use std::time::Duration;

use std::fs::File;
use std::io::{stdin, stdout, SeekFrom, Error, ErrorKind, Read, BufReader, Seek, BufRead, Cursor, Write};

use rustc_serialize::json;

static USAGE: &'static str = "
Commands:
  server  Starts a server from an initial cluster config.
  client  Starts a client repl, that communicates with
          a particular cluster.

Usage:
  hashmap server <id> <filename>
  hashmap servers <filename>
  hashmap client <filename>

Options:
  -h --help   Show a help message.
";

fn main() {
    // TODO (sydli) make prettier
    let mut args = args();
    if let Some(command) = args.nth(1) {
        if &command == "server" {
            if let (Some(id_str), Some(filename)) =
                   (args.next(), args.next()) {
                if let Ok(id) = id_str.parse::<u64>() {
                    Server::new(id, cluster_from_file(&filename).get(&id).unwrap().addr).repl();
                    return;
                }
            }
        } else if &command == "client" {
            if let Some(filename) = args.next() {
                Client::new(&cluster_from_file(&filename)).repl();
                return;
            }
        } else if &command == "servers" {
            if let Some(filename) = args.next() {
                Cluster::new(&cluster_from_file(&filename)).repl();
                return;
            }
        }
    }
    warn!("Incorrect usage. \n{}", USAGE);
}

///
/// Automatically builds a repl loop for an object, if it implements |exec|.
/// Each line a user types into the repl is fed to |exec|.
///
/// Default commands:
///     exit    Exits the loop and shuts down associated process.
///     help    Displays the |help| message for more usage info.
trait Repl {

    ///
    /// Processes each input command from the Repl.
    /// Returns true if |command| is formatted correctly.
    ///
    fn exec(&mut self, command: String) -> bool;

    ///
    /// Usage string displayed with the default command "help".
    ///
    fn usage(&self) -> String;

    fn print_usage(&self) {
        println!("\nREPL COMMANDS:\n==============\n{}\n{}\n{}",
                 "exit\n\t\tExits the loop and shuts down associated process.",
                 "help\n\t\tSpits out this message.",
                 self.usage());
    }

    ///
    /// Implementation of the repl.
    fn repl(&mut self) {
        println!("[ Starting repl ]");
        loop {
            print!("> ");
            stdout().flush().unwrap();
            let mut buffer = String::new();
            if stdin().read_line(&mut buffer).is_ok() {
                let words: Vec<String> =
                    buffer.split_whitespace().map(String::from).collect();
                if words.get(0) == Some(&String::from("exit")) { break; }
                if words.get(0) == Some(&String::from("help")) {
                    self.print_usage();
                    continue;
                }
                if !self.exec(buffer.clone()) {
                    println!("Command not recognized.");
                    self.print_usage();
                }
            }
        }
    }
}

#[derive(Clone, Debug)]
struct ServerInfo {
    addr: SocketAddr,
    state_filename: String,
    log_filename: String,
}

const STATE_FILENAME_LEN: usize = 20;
impl ServerInfo {
    fn new(addr: SocketAddr) -> ServerInfo {

        let mut random_filename: String = thread_rng().gen_ascii_chars().take(STATE_FILENAME_LEN).collect();
        ServerInfo {
            addr: addr,
            state_filename: String::from("/tmp/state_") + &random_filename,
            log_filename: String::from("/tmp/log_") + &random_filename,
        }
    }
}

struct Server {
    handle: ServerHandle,
}

impl Repl for Server { // TODO: impl repl for a single node in the cluster
    fn exec(&mut self, command: String) -> bool { true }
    fn usage(&self) -> String { String::from("") }
}

impl Server {
    fn new(id: u64, addr: SocketAddr) -> Server {
        Server {
            handle:
                // TODO: Only the first server should put itself in its log file
                start_server(id, Box::new(RaftHashMap { map: HashMap::new() }), addr).unwrap()
        }
    }
}

struct Cluster {
    servers: HashMap<u64, Server>,
    cluster: HashMap<u64, ServerInfo>,
}

impl Cluster {
    fn new(info: &HashMap<u64, ServerInfo>) -> Cluster {
        let addr = info.clone().into_iter().map(|(id, info)| (id, info.addr))
            .collect::<HashMap<u64, SocketAddr>>();
        let mut servers = HashMap::new();
        for (id, info) in info { servers.insert(*id, Server::new(*id, info.addr)); }
        let mut raft_db =  RaftConnection::new_with_session(&addr.clone()).unwrap();
        for id in servers.keys() {
            raft_db.add_server(*id, info.get(id).cloned().unwrap().addr);
        }
        Cluster { servers: servers, cluster: info.clone() }
    }

    fn add_server(&mut self, id: u64, addr: SocketAddr) {
        println!("Adding server {}, {}", id, addr);
    }

    fn remove_server(&mut self, id: u64) {
        if !self.cluster.contains_key(&id) {
            println!("Server {} is not in the cluster! Servers: {:?}", id,
                     self.cluster);
        }
        println!("Removing server {}", id);
    }

    fn kill_server(&mut self, id: u64) {
        if !self.servers.contains_key(&id) {
            println!("Server {} is not up right now!", id);
        }
        {
            self.servers.remove(&id).unwrap();
        } // drop server
        println!("Killed server {}", id);
    }

    fn start_server(&mut self, id: u64) {
        if self.servers.contains_key(&id) {
            println!("Server {} is already up!", id);
        }
        self.servers.insert(id, Server::new(id, self.cluster.get(&id).unwrap().addr));
        println!("Restarted server {}", id);
    }

    fn print_servers(&self) {
        println!("Servers in cluster: {:?}\n Live servers: {:?}",
                 self.cluster, self.servers.keys().collect::<Vec<&u64>>());
    }
}

impl Repl for Cluster {
    fn exec(&mut self, command: String) -> bool {
        let words: Vec<&str> =
            command.split_whitespace().collect();
        let first = words.get(0).map(|x|*x);
        if first == Some("add")  {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            let addr = words.get(2).and_then(|x| as_addr(*x).ok());
            if num.and(addr).is_none() { return false; }
            self.add_server(num.unwrap(), addr.unwrap());
        } else if first == Some("remove") {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            if num.is_none() { return false; }
            self.remove_server(num.unwrap());
        } else if first == Some("kill") {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            if num.is_none() { return false; }
            self.kill_server(num.unwrap());
        } else if first == Some("start") {
            let num = words.get(1).and_then(|x| as_num(*x).ok());
            if num.is_none() { return false; }
            self.start_server(num.unwrap());
        } else if first == Some("list") {
            self.print_servers();
        } else { return false; }
        return true;
    }

    fn usage(&self) -> String {
        String::from(format!(
                "{}\n{}\n{}\n{}\n{}",
                "add <node-id> <node-addr>\tAdds a server to the cluster.",
                "remove <node-id>\tRemoves a server from the cluster.",
                "start <node-id>\tStarts a server from the cluster.",
                "kill <node-id>\tKills a server from the cluster.",
                "list\tLists all cluster servers, and the ones that are live."))
    }
}

struct Client {
    raft: RaftConnection,
}

impl Repl for Client {
    fn exec (&mut self, command: String) -> bool {
        let words: Vec<String> =
            command.split_whitespace().map(String::from).collect();
        self.process_command(words)
    }
    fn usage(&self) -> String {
        String::from(format!(
                "{}\n{}",
                 "get <key>\n\t\tPrints value of <key> in hashmap, if it exists.",
                 "put <key> <value>\n\t\tPuts value of <key> as <value>"))
    }
}

#[derive(RustcDecodable, RustcEncodable)]
struct Put {
    key: String,
    value: String,
}

impl Client {
    fn new(cluster: &HashMap<u64, ServerInfo>) -> Client {
        let cluster = cluster.clone().into_iter().map(|(id, info)| (id, info.addr))
            .collect::<HashMap<u64, SocketAddr>>();
        let connection = RaftConnection::new_with_session(&cluster);
        if connection.is_none() {
            println!("Couldn't establish connection to cluster at {:?}", cluster);
            panic!();
        }
        Client { raft: connection.unwrap() }
    }

    fn get(&mut self, key: String) -> Result<String, RaftError> {
        self.raft.query(key.as_bytes())
            .and_then(
                |result| str::from_utf8(&result)
                             .map(str::to_string)
                             .map_err(deserialize_error))
    }

    fn put(&mut self, key: String, value: String) -> Result<(), RaftError> {
        json::encode(&Put {key:key, value:value})
            .map_err(serialize_error)
            .and_then(|buffer| self.raft.command(&buffer.as_bytes()))
    }

    // TODO (sydli): make this function less shit
    fn process_command(&mut self, words: Vec<String>)
        -> bool {
        if words.len() == 0 { return true; }
        let ref cmd = words[0];
        if *cmd == String::from("get") {
            if words.len() <= 1 { return false; }
            words.get(1).map(|key| {
                match self.get(key.clone()) {
                    Ok(value) => println!("Value for {} is {}", key, value),
                    Err(err) => println!("Error during get: {:?}", err),
                }
            }).unwrap();
        } else if *cmd == String::from("put") {
            if words.len() <= 2 { return false; }
            words.get(1).map(|key| { words.get(2).map(|val| {
                match self.put(key.clone(), val.clone()) {
                    Ok(()) => println!("Put {} => {} successfully", key, val),
                    Err(err) => println!("Error during put: {:?}", err),
                }
            }).unwrap()}).unwrap();
        } else { return false; }
        return true;
    }
}

fn io_err() -> std::io::Error {
    Error::new(ErrorKind::InvalidData, "File incorrectly formatted")
}

fn as_num(x: &str) -> Result<u64, std::io::Error> {
    String::from(x).parse::<u64>().map_err(|_| io_err())
}

fn as_addr(x: &str) -> Result<SocketAddr, std::io::Error> {
    String::from(x).parse::<SocketAddr>().map_err(|_| io_err())
}

/// Panics on io error (we can't access the cluster info!)
/// TODO (sydli) make io_errs more informative
// TODO(jason): Remove this and bootstrap a dynamic cluster
fn cluster_from_file(filename: &str) -> HashMap<u64, ServerInfo> {
    let file = File::open(filename.clone())
               .expect(&format!("Unable to open file {}", filename));
    let mut lines = BufReader::new(file).lines();
    lines.next().ok_or(io_err())
    .map(|line_or_io_error| line_or_io_error.unwrap())
    .and_then(|x| as_num(&x))
    .and_then(|num| (0..num).map(|_| {
        lines.next().ok_or(io_err())
        .map(|line_or_io_error| line_or_io_error.unwrap())
        .and_then(|node_str| {
            let mut words = node_str.split_whitespace().map(String::from);
            let id = words.next().ok_or(io_err())
                        .and_then(|x| as_num(&x));
            words.next().ok_or(io_err())
                .and_then(|x| as_addr(&x))
                .and_then(move |addr| id.map(|id| (id, ServerInfo::new(addr))))
        })
    }).collect::<Result<Vec<_>, _>>())
    .map(|nodes: Vec<(u64, ServerInfo)>|
        nodes.iter().cloned().collect::<HashMap<u64, ServerInfo>>()
    ).unwrap()
}

struct RaftHashMap {
    map: HashMap<String, String>,
}

fn serialize_error <T: Debug>(error: T) -> RaftError {
        RaftError::ClientError(
            format!("Couldn't serialize object. Error: {:?}", error))
}

fn deserialize_error <T: Debug>(error: T) -> RaftError {
        RaftError::ClientError(
            format!("Couldn't deserialize buffer. Error: {:?}", error))
}

fn key_error(key: &String) -> RaftError {
        RaftError::ClientError(format!("Couldn't find key {}", key))
}

impl StateMachine for RaftHashMap {
    fn command (&mut self, buffer: &[u8]) ->Result<(), RaftError> {
        str::from_utf8(buffer)
            .map_err(deserialize_error)
            .and_then(|string| json::decode(&string)
                                    .map_err(deserialize_error))
            .map(|put: Put|
                 {
                     self.map.insert(put.key, put.value);
                 })
    }

    fn query (& self, buffer: &[u8]) ->Result<Vec<u8>, RaftError> {
        str::from_utf8(buffer)
            .map_err(deserialize_error)
            .and_then(|key| {
                let key = key.to_string();
                self.map.get(&key)
                    .map(|x| x.as_bytes().to_vec())
                    .ok_or(key_error(&key))
            })
    }
}

<|endoftext|>


<commit_before>extern crate pdb;
extern crate getopts;

use pdb::FallibleIterator;
use std::env;
use std::fmt;
use std::io::Write;

pub fn type_name<'p>(type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex) -> pdb::Result<String> {
    let mut name = match type_finder.find(type_index)?.parse()? {
        pdb::TypeData::Primitive { primitive_type, indirection } => {
            let mut name = match primitive_type {
                pdb::PrimitiveType::Void => "void".to_string(),
                pdb::PrimitiveType::Char => "char".to_string(),
                pdb::PrimitiveType::UChar => "unsigned char".to_string(),

                pdb::PrimitiveType::I8 => "int8_t".to_string(),
                pdb::PrimitiveType::U8 => "uint8_t".to_string(),
                pdb::PrimitiveType::I16 => "int16_t".to_string(),
                pdb::PrimitiveType::U16 => "uint16_t".to_string(),
                pdb::PrimitiveType::I32 => "int32_t".to_string(),
                pdb::PrimitiveType::U32 => "uint32_t".to_string(),
                pdb::PrimitiveType::I64 => "int64_t".to_string(),
                pdb::PrimitiveType::U64 => "uint64_t".to_string(),

                pdb::PrimitiveType::F32 => "float".to_string(),
                pdb::PrimitiveType::F64 => "double".to_string(),

                pdb::PrimitiveType::Bool8 => "bool".to_string(),

                other => format!("unhandled_primitive_type /* {:?} */", primitive_type),
            };

            match indirection {
                pdb::Indirection::None => {},
                _ => { name.push(' '); name.push('*'); },
            }

            name
        },

        pdb::TypeData::Class { name, .. } => {
            name.to_string().into_owned()
        },

        pdb::TypeData::Enumeration { name, .. } => {
            name.to_string().into_owned()
        },

        pdb::TypeData::Union { name, .. } => {
            name.to_string().into_owned()
        },

        pdb::TypeData::Pointer { underlying_type, .. } => {
            format!("{}*", type_name(type_finder, underlying_type)?)
        },

        pdb::TypeData::Modifier { constant, volatile, underlying_type, .. } => {
            if constant {
                format!("const {}", type_name(type_finder, underlying_type)?)
            } else if volatile {
                format!("volatile {}", type_name(type_finder, underlying_type)?)
            } else {
                // ?
                type_name(type_finder, underlying_type)?
            }
        },

        pdb::TypeData::Array { element_type, dimensions, .. } => {
            let mut name = type_name(type_finder, element_type)?;
            for size in dimensions {
                name = format!("{}[{}]", name, size);
            }
            name
        },

        _ => format!("Type{} /* TODO: figure out how to name it */", type_index)
    };

    if name == "std::basic_string<char,std::char_traits<char>,std::allocator<char> >" {
        name = "std::string".to_string();
    }

    Ok(name)
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Class<'p> {
    kind: pdb::ClassKind,
    name: pdb::RawString<'p>,
    fields: Vec<Field<'p>>,
}

impl<'p> Class<'p> {
    fn add_derived_from(&mut self, type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex) -> pdb::Result<()> {
        // TODO
        Ok(())
    }

    fn add_fields(&mut self, type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex) -> pdb::Result<()> {
        match type_finder.find(type_index)?.parse()? {
            pdb::TypeData::FieldList { fields, continuation, .. } => {
                for ref field in fields {
                    self.add_field(type_finder, field)?;
                }

                if let Some(continuation) = continuation {
                    // recurse
                    self.add_fields(type_finder, continuation)?;
                }
            }
            other => {
                println!("trying to Class::add_fields() got {} -> {:?}", type_index, other);
                panic!("unexpected type in Class::add_fields()");
            }
        }

        Ok(())
    }

    fn add_field(&mut self, type_finder: &pdb::TypeFinder<'p>, field: &pdb::TypeData<'p>) -> pdb::Result<()> {
        match field {
            &pdb::TypeData::Member { field_type, offset, ref name, .. } => {
                // TODO: attributes (static, virtual, etc.)
                self.fields.push(Field{
                    type_name: type_name(type_finder, field_type)?,
                    name: name.clone(),
                    offset: offset,
                });
            },

            _ => {
                // ignore everything else even though that's sad
            }
        }

        Ok(())
    }
}

impl<'p> fmt::Display for Class<'p> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "{} {} {{", match self.kind {
            pdb::ClassKind::Class => "class",
            pdb::ClassKind::Struct => "struct",
            pdb::ClassKind::Interface => "interface",   // when can this happen?
        }, self.name.to_string())?;

        for ref field in &self.fields {
            writeln!(f, "\t/* offset {:3} */ {} {};", field.offset, field.type_name, field.name.to_string())?;
        }

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

        Ok(())
    }
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Field<'p> {
    type_name: String,
    name: pdb::RawString<'p>,
    offset: u16,
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Enum<'p> {
    name: pdb::RawString<'p>,
    values: Vec<EnumValue<'p>>,
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct EnumValue<'p> {
    name: pdb::RawString<'p>,
    value: pdb::EnumValue,
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Data<'p> {
    classes: Vec<Class<'p>>,
    enums: Vec<Enum<'p>>,
}

impl<'p> fmt::Display for Data<'p> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "// automatically generated by pdb2hpp\n// do not edit")?;

        for ref class in &self.classes {
            writeln!(f, "")?;
            class.fmt(f)?;
        }

        Ok(())
    }
}

impl<'p> Data<'p> {
    fn new() -> Data<'p> {
        Data{
            classes: Vec::new(),
            enums: Vec::new(),
        }
    }

    fn add(&mut self, type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex) -> pdb::Result<()> {
        match type_finder.find(type_index)?.parse()? {
            pdb::TypeData::Class {
                kind, properties, fields, name, derived_from, ..
            } => {
                // can't do much useful with a forward reference, but also we shouldn't be trying
                // to add them
                assert_eq!(properties.forward_reference(), false);

                // TODO: derived_from

                let mut class = Class{
                    kind: kind,
                    name: name,
                    fields: Vec::new(),
                };

                if let Some(derived_from) = derived_from {
                    class.add_derived_from(type_finder, derived_from)?;
                }

                if let Some(fields) = fields {
                    class.add_fields(type_finder, fields)?;
                }

                self.classes.push(class);
            }

            _ => {
                // ignore
            }
        }

        Ok(())
    }
}

fn write_class(filename: &str, class_name: &str) -> pdb::Result<()> {
    let file = std::fs::File::open(filename)?;
    let mut pdb = pdb::PDB::open(file)?;

    let type_information = pdb.type_information()?;
    let mut type_finder = type_information.new_type_finder();

    let mut data = Data::new();

    let mut type_iter = type_information.iter();
    while let Some(typ) = type_iter.next()? {
        // keep building the index
        type_finder.update(&type_iter);

        if let Ok(pdb::TypeData::Class { name, properties, .. }) = typ.parse() {
            if name.as_bytes() == class_name.as_bytes() && !properties.forward_reference() {
                data.add(&type_finder, typ.type_index())?;
                break;
            }
        }
    }

    if data.classes.len() == 0 {
        writeln!(&mut std::io::stderr(), "sorry, class {} was not found", class_name)
            .expect("stderr write");
    } else {
        println!("{}", data);
    }

    Ok(())
}

fn print_usage(program: &str, opts: getopts::Options) {
    let brief = format!("Usage: {} input.pdb ClassName", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "print this help menu");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };

    let (filename, class_name) = if matches.free.len() == 2 {
        (&matches.free[0], &matches.free[1])
    } else {
        print_usage(&program, opts);
        return;
    };

    match write_class(&filename, &class_name) {
        Ok(_) => {}
        Err(e) => {
            writeln!(&mut std::io::stderr(), "error dumping PDB: {}", e)
                .expect("stderr write");
        }
    }
}
<commit_msg>Improve pdb2hpp output<commit_after>extern crate pdb;
extern crate getopts;

use pdb::FallibleIterator;
use std::env;
use std::fmt;
use std::io::Write;
use std::collections::BTreeSet;

type TypeSet = BTreeSet<pdb::TypeIndex>;

pub fn type_name<'p>(type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex, mut needed_types: &mut TypeSet) -> pdb::Result<String> {
    let mut name = match type_finder.find(type_index)?.parse()? {
        pdb::TypeData::Primitive { primitive_type, indirection } => {
            let mut name = match primitive_type {
                pdb::PrimitiveType::Void => "void".to_string(),
                pdb::PrimitiveType::Char => "char".to_string(),
                pdb::PrimitiveType::UChar => "unsigned char".to_string(),

                pdb::PrimitiveType::I8 => "int8_t".to_string(),
                pdb::PrimitiveType::U8 => "uint8_t".to_string(),
                pdb::PrimitiveType::I16 => "int16_t".to_string(),
                pdb::PrimitiveType::U16 => "uint16_t".to_string(),
                pdb::PrimitiveType::I32 => "int32_t".to_string(),
                pdb::PrimitiveType::U32 => "uint32_t".to_string(),
                pdb::PrimitiveType::I64 => "int64_t".to_string(),
                pdb::PrimitiveType::U64 => "uint64_t".to_string(),

                pdb::PrimitiveType::F32 => "float".to_string(),
                pdb::PrimitiveType::F64 => "double".to_string(),

                pdb::PrimitiveType::Bool8 => "bool".to_string(),

                _ => format!("unhandled_primitive_type /* {:?} */", primitive_type),
            };

            match indirection {
                pdb::Indirection::None => {},
                _ => { name.push(' '); name.push('*'); },
            }

            name
        },

        pdb::TypeData::Class { name, .. } => {
            needed_types.insert(type_index);
            name.to_string().into_owned()
        },

        pdb::TypeData::Enumeration { name, .. } => {
            needed_types.insert(type_index);
            name.to_string().into_owned()
        },

        pdb::TypeData::Union { name, .. } => {
            needed_types.insert(type_index);
            name.to_string().into_owned()
        },

        pdb::TypeData::Pointer { underlying_type, .. } => {
            format!("{}*", type_name(type_finder, underlying_type, needed_types)?)
        },

        pdb::TypeData::Modifier { constant, volatile, underlying_type, .. } => {
            if constant {
                format!("const {}", type_name(type_finder, underlying_type, needed_types)?)
            } else if volatile {
                format!("volatile {}", type_name(type_finder, underlying_type, needed_types)?)
            } else {
                // ?
                type_name(type_finder, underlying_type, needed_types)?
            }
        },

        pdb::TypeData::Array { element_type, dimensions, .. } => {
            let mut name = type_name(type_finder, element_type, needed_types)?;
            for size in dimensions {
                name = format!("{}[{}]", name, size);
            }
            name
        },

        _ => format!("Type{} /* TODO: figure out how to name it */", type_index)
    };

    // TODO: search and replace std:: patterns
    if name == "std::basic_string<char,std::char_traits<char>,std::allocator<char> >" {
        name = "std::string".to_string();
    }

    Ok(name)
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Class<'p> {
    kind: pdb::ClassKind,
    name: pdb::RawString<'p>,
    base_classes: Vec<BaseClass>,
    fields: Vec<Field<'p>>,
}

impl<'p> Class<'p> {
    fn add_derived_from(&mut self, _: &pdb::TypeFinder<'p>, _: pdb::TypeIndex, _: &mut TypeSet) -> pdb::Result<()> {
        // TODO
        Ok(())
    }

    fn add_fields(&mut self, type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex, mut needed_types: &mut TypeSet) -> pdb::Result<()> {
        match type_finder.find(type_index)?.parse()? {
            pdb::TypeData::FieldList { fields, continuation, .. } => {
                for ref field in fields {
                    self.add_field(type_finder, field, needed_types)?;
                }

                if let Some(continuation) = continuation {
                    // recurse
                    self.add_fields(type_finder, continuation, needed_types)?;
                }
            }
            other => {
                println!("trying to Class::add_fields() got {} -> {:?}", type_index, other);
                panic!("unexpected type in Class::add_fields()");
            }
        }

        Ok(())
    }

    fn add_field(&mut self, type_finder: &pdb::TypeFinder<'p>, field: &pdb::TypeData<'p>, mut needed_types: &mut TypeSet) -> pdb::Result<()> {
        match field {
            &pdb::TypeData::Member { field_type, offset, ref name, .. } => {
                // TODO: attributes (static, virtual, etc.)
                self.fields.push(Field{
                    type_name: type_name(type_finder, field_type, needed_types)?,
                    name: name.clone(),
                    offset: offset,
                });
            },

            &pdb::TypeData::BaseClass { base_class, offset, .. } => {
                self.base_classes.push(BaseClass{
                    type_name: type_name(type_finder, base_class, needed_types)?,
                    offset: offset,
                })
            },

            &pdb::TypeData::VirtualBaseClass { base_class, base_pointer_offset, .. } => {
                self.base_classes.push(BaseClass{
                    type_name: type_name(type_finder, base_class, needed_types)?,
                    offset: base_pointer_offset,
                })
            },

            _ => {
                // ignore everything else even though that's sad
            }
        }

        Ok(())
    }
}

impl<'p> fmt::Display for Class<'p> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {} ", match self.kind {
            pdb::ClassKind::Class => "class",
            pdb::ClassKind::Struct => "struct",
            pdb::ClassKind::Interface => "interface",   // when can this happen?
        }, self.name.to_string())?;

        if self.base_classes.len() > 0 {
            for (i, base) in self.base_classes.iter().enumerate() {
                let prefix = match i {
                    0 => ":",
                    _ => ","
                };
                write!(f, "{} {}", prefix, base.type_name)?;
            }
        }

        writeln!(f, " {{")?;

        for base in &self.base_classes {
            writeln!(f, "\t/* offset {:3} */ /* fields for {} */", base.offset, base.type_name)?;
        }

        for ref field in &self.fields {
            writeln!(f, "\t/* offset {:3} */ {} {};", field.offset, field.type_name, field.name.to_string())?;
        }

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

        Ok(())
    }
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct BaseClass {
    type_name: String,
    offset: u32,
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Field<'p> {
    type_name: String,
    name: pdb::RawString<'p>,
    offset: u16,
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Enum<'p> {
    name: pdb::RawString<'p>,
    underlying_type_name: String,
    values: Vec<EnumValue<'p>>,
}

impl<'p> Enum<'p> {
    fn add_fields(&mut self, type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex, mut needed_types: &mut TypeSet) -> pdb::Result<()> {
        match type_finder.find(type_index)?.parse()? {
            pdb::TypeData::FieldList { fields, continuation, .. } => {
                for ref field in fields {
                    self.add_field(type_finder, field, needed_types)?;
                }

                if let Some(continuation) = continuation {
                    // recurse
                    self.add_fields(type_finder, continuation, needed_types)?;
                }
            }
            other => {
                println!("trying to Enum::add_fields() got {} -> {:?}", type_index, other);
                panic!("unexpected type in Enum::add_fields()");
            }
        }

        Ok(())
    }

    fn add_field(&mut self, _: &pdb::TypeFinder<'p>, field: &pdb::TypeData<'p>, _: &mut TypeSet) -> pdb::Result<()> {
        match field {
            &pdb::TypeData::Enumerate { ref name, value, .. } => {
                self.values.push(EnumValue{
                    name: name.clone(),
                    value: value,
                });
            },

            _ => {
                // ignore everything else even though that's sad
            }
        }

        Ok(())
    }
}

impl<'p> fmt::Display for Enum<'p> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "enum {} /* stored as {} */ {{", self.name.to_string(), self.underlying_type_name)?;

        for value in &self.values {
            writeln!(f, "\t{} = {},", value.name.to_string(), match value.value {
                pdb::EnumValue::U8(v) => format!("0x{:02x}", v),
                pdb::EnumValue::U16(v) => format!("0x{:04x}", v),
                pdb::EnumValue::U32(v) => format!("0x{:08x}", v),
                pdb::EnumValue::U64(v) => format!("0x{:16x}", v),
                pdb::EnumValue::I8(v) => format!("{}", v),
                pdb::EnumValue::I16(v) => format!("{}", v),
                pdb::EnumValue::I32(v) => format!("{}", v),
                pdb::EnumValue::I64(v) => format!("{}", v),
            })?;
        }
        writeln!(f, "}}")?;

        Ok(())
    }
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct EnumValue<'p> {
    name: pdb::RawString<'p>,
    value: pdb::EnumValue,
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct ForwardReference<'p> {
    kind: pdb::ClassKind,
    name: pdb::RawString<'p>,
}

impl<'p> fmt::Display for ForwardReference<'p> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "{} {};", match self.kind {
            pdb::ClassKind::Class => "class",
            pdb::ClassKind::Struct => "struct",
            pdb::ClassKind::Interface => "interface",   // when can this happen?
        }, self.name.to_string())
    }
}

#[derive(Debug,Clone,PartialEq,Eq)]
struct Data<'p> {
    forward_references: Vec<ForwardReference<'p>>,
    classes: Vec<Class<'p>>,
    enums: Vec<Enum<'p>>,
}

impl<'p> fmt::Display for Data<'p> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "// automatically generated by pdb2hpp\n// do not edit")?;

        if !self.forward_references.is_empty() {
            writeln!(f, "")?;
            for e in self.forward_references.iter() {
                e.fmt(f)?;
            }
        }

        for e in self.enums.iter() {
            writeln!(f, "")?;
            e.fmt(f)?;
        }

        for class in self.classes.iter() {
            writeln!(f, "")?;
            class.fmt(f)?;
        }

        Ok(())
    }
}

impl<'p> Data<'p> {
    fn new() -> Data<'p> {
        Data{
            forward_references: Vec::new(),
            classes: Vec::new(),
            enums: Vec::new(),
        }
    }

    fn add(&mut self, type_finder: &pdb::TypeFinder<'p>, type_index: pdb::TypeIndex, mut needed_types: &mut TypeSet) -> pdb::Result<()> {
        match type_finder.find(type_index)?.parse()? {
            pdb::TypeData::Class {
                kind, properties, fields, name, derived_from, ..
            } => {
                if properties.forward_reference() {
                    self.forward_references.push(ForwardReference{
                        kind: kind,
                        name: name,
                    });

                    return Ok(());
                }

                let mut class = Class{
                    kind: kind,
                    name: name,
                    fields: Vec::new(),
                    base_classes: Vec::new(),
                };

                if let Some(derived_from) = derived_from {
                    class.add_derived_from(type_finder, derived_from, needed_types)?;
                }

                if let Some(fields) = fields {
                    class.add_fields(type_finder, fields, needed_types)?;
                }

                self.classes.insert(0, class);
            }

            pdb::TypeData::Enumeration {
                underlying_type, fields, name, ..
            } => {
                let mut e = Enum{
                    name: name,
                    underlying_type_name: type_name(type_finder, underlying_type, needed_types)?,
                    values: Vec::new(),
                };

                e.add_fields(type_finder, fields, needed_types)?;

                self.enums.insert(0, e);
            }

            other => {
                // ignore
                writeln!(&mut std::io::stderr(), "warning: don't know how to add {:?}", other).expect("stderr write");;
            }
        }

        Ok(())
    }
}

fn write_class(filename: &str, class_name: &str) -> pdb::Result<()> {
    let file = std::fs::File::open(filename)?;
    let mut pdb = pdb::PDB::open(file)?;

    let type_information = pdb.type_information()?;
    let mut type_finder = type_information.new_type_finder();

    let mut needed_types = TypeSet::new();
    let mut data = Data::new();

    let mut type_iter = type_information.iter();
    while let Some(typ) = type_iter.next()? {
        // keep building the index
        type_finder.update(&type_iter);

        if let Ok(pdb::TypeData::Class { name, properties, .. }) = typ.parse() {
            if name.as_bytes() == class_name.as_bytes() && !properties.forward_reference() {
                data.add(&type_finder, typ.type_index(), &mut needed_types)?;
                break;
            }
        }
    }

    // add all the needed types iteratively until we're done
    loop {
        // get the last element in needed_types without holding an immutable borrow
        let last = match needed_types.iter().next_back() {
            Some(n) => Some(*n),
            None => None
        };

        if let Some(type_index) = last {
            // remove it
            needed_types.remove(&type_index);

            // add the type
            data.add(&type_finder, type_index, &mut needed_types)?;
        } else {
            break
        }
    }

    if data.classes.len() == 0 {
        writeln!(&mut std::io::stderr(), "sorry, class {} was not found", class_name)
            .expect("stderr write");
    } else {
        println!("{}", data);
    }

    Ok(())
}

fn print_usage(program: &str, opts: getopts::Options) {
    let brief = format!("Usage: {} input.pdb ClassName", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "print this help menu");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };

    let (filename, class_name) = if matches.free.len() == 2 {
        (&matches.free[0], &matches.free[1])
    } else {
        print_usage(&program, opts);
        return;
    };

    match write_class(&filename, &class_name) {
        Ok(_) => {}
        Err(e) => {
            writeln!(&mut std::io::stderr(), "error dumping PDB: {}", e)
                .expect("stderr write");
        }
    }
}
<|endoftext|>