<commit_before>use std::fs::File;
use std::io::prelude::*;

use crate::support::registry::Package;
use crate::support::{basic_manifest, project};

#[test]
fn minor_update_two_places() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    Package::new("log", "0.1.1").publish();

    File::create(p.root().join("foo/Cargo.toml"))
        .unwrap()
        .write_all(
            br#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1.1"
            "#,
        )
        .unwrap();

    p.cargo("build").run();
}

#[test]
fn transitive_minor_update() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").dep("log", "0.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("serde", "0.1.1").dep("log", "0.1.1").publish();

    // Note that `serde` isn't actually updated here! The default behavior for
    // `update` right now is to as conservatively as possible attempt to satisfy
    // an update. In this case we previously locked the dependency graph to `log
    // 0.1.0`, but nothing on the command line says we're allowed to update
    // that. As a result the update of `serde` here shouldn't update to `serde
    // 0.1.1` as that would also force an update to `log 0.1.1`.
    //
    // Also note that this is probably counterintuitive and weird. We may wish
    // to change this one day.
    p.cargo("update -p serde")
        .with_stderr(
            "\
[UPDATING] `[..]` index
",
        )
        .run();
}

#[test]
fn conservative() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").dep("log", "0.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("serde", "0.1.1").dep("log", "0.1").publish();

    p.cargo("update -p serde")
        .with_stderr(
            "\
[UPDATING] `[..]` index
[UPDATING] serde v0.1.0 -> v0.1.1
",
        )
        .run();
}

#[test]
fn update_via_new_dep() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
                # foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    Package::new("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").env("RUST_LOG", "cargo=trace").run();
}

#[test]
fn update_via_new_member() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [workspace]
                # members = [ "foo" ]

                [dependencies]
                log = "0.1"
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    Package::new("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").run();
}

#[test]
fn add_dep_deep_new_requirement() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
                # bar = "0.1"
            "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("bar", "0.1.0").dep("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").run();
}

#[test]
fn everything_real_deep() {
    Package::new("log", "0.1.0").publish();
    Package::new("foo", "0.1.0").dep("log", "0.1").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                foo = "0.1"
                # bar = "0.1"
            "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("bar", "0.1.0").dep("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").run();
}

#[test]
fn change_package_version() {
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "a-foo"
                version = "0.2.0-alpha"
                authors = []

                [dependencies]
                bar = { path = "bar", version = "0.2.0-alpha" }
            "#,
        )
        .file("src/lib.rs", "")
        .file("bar/Cargo.toml", &basic_manifest("bar", "0.2.0-alpha"))
        .file("bar/src/lib.rs", "")
        .file(
            "Cargo.lock",
            r#"
                [[package]]
                name = "foo"
                version = "0.2.0"
                dependencies = ["bar 0.2.0"]

                [[package]]
                name = "bar"
                version = "0.2.0"
            "#,
        )
        .build();

    p.cargo("build").run();
}

#[test]
fn update_precise() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").publish();
    Package::new("serde", "0.2.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.2"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("serde", "0.2.0").publish();

    p.cargo("update -p serde:0.2.1 --precise 0.2.0")
        .with_stderr(
            "\
[UPDATING] `[..]` index
[UPDATING] serde v0.2.1 -> v0.2.0
",
        )
        .run();
}

#[test]
fn preserve_top_comment() {
    let p = project().file("src/lib.rs", "").build();

    p.cargo("update").run();

    let mut lockfile = p.read_file("Cargo.lock");
    lockfile.insert_str(0, "# @generated\n");
    lockfile.insert_str(0, "# some other comment\n");
    println!("saving Cargo.lock contents:\n{}", lockfile);

    p.change_file("Cargo.lock", &lockfile);

    p.cargo("update").run();

    let lockfile2 = p.read_file("Cargo.lock");
    println!("loaded Cargo.lock contents:\n{}", lockfile2);

    assert!(lockfile == lockfile2);
}
<commit_msg>Add a test for update --dry-run<commit_after>use std::fs::File;
use std::io::prelude::*;

use crate::support::registry::Package;
use crate::support::{basic_manifest, project};

#[test]
fn minor_update_two_places() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    Package::new("log", "0.1.1").publish();

    File::create(p.root().join("foo/Cargo.toml"))
        .unwrap()
        .write_all(
            br#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1.1"
            "#,
        )
        .unwrap();

    p.cargo("build").run();
}

#[test]
fn transitive_minor_update() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").dep("log", "0.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("serde", "0.1.1").dep("log", "0.1.1").publish();

    // Note that `serde` isn't actually updated here! The default behavior for
    // `update` right now is to as conservatively as possible attempt to satisfy
    // an update. In this case we previously locked the dependency graph to `log
    // 0.1.0`, but nothing on the command line says we're allowed to update
    // that. As a result the update of `serde` here shouldn't update to `serde
    // 0.1.1` as that would also force an update to `log 0.1.1`.
    //
    // Also note that this is probably counterintuitive and weird. We may wish
    // to change this one day.
    p.cargo("update -p serde")
        .with_stderr(
            "\
[UPDATING] `[..]` index
",
        )
        .run();
}

#[test]
fn conservative() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").dep("log", "0.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("serde", "0.1.1").dep("log", "0.1").publish();

    p.cargo("update -p serde")
        .with_stderr(
            "\
[UPDATING] `[..]` index
[UPDATING] serde v0.1.0 -> v0.1.1
",
        )
        .run();
}

#[test]
fn update_via_new_dep() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
                # foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    Package::new("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").env("RUST_LOG", "cargo=trace").run();
}

#[test]
fn update_via_new_member() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [workspace]
                # members = [ "foo" ]

                [dependencies]
                log = "0.1"
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    Package::new("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").run();
}

#[test]
fn add_dep_deep_new_requirement() {
    Package::new("log", "0.1.0").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                log = "0.1"
                # bar = "0.1"
            "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("bar", "0.1.0").dep("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").run();
}

#[test]
fn everything_real_deep() {
    Package::new("log", "0.1.0").publish();
    Package::new("foo", "0.1.0").dep("log", "0.1").publish();
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                foo = "0.1"
                # bar = "0.1"
            "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("log", "0.1.1").publish();
    Package::new("bar", "0.1.0").dep("log", "0.1.1").publish();

    p.uncomment_root_manifest();
    p.cargo("build").run();
}

#[test]
fn change_package_version() {
    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "a-foo"
                version = "0.2.0-alpha"
                authors = []

                [dependencies]
                bar = { path = "bar", version = "0.2.0-alpha" }
            "#,
        )
        .file("src/lib.rs", "")
        .file("bar/Cargo.toml", &basic_manifest("bar", "0.2.0-alpha"))
        .file("bar/src/lib.rs", "")
        .file(
            "Cargo.lock",
            r#"
                [[package]]
                name = "foo"
                version = "0.2.0"
                dependencies = ["bar 0.2.0"]

                [[package]]
                name = "bar"
                version = "0.2.0"
            "#,
        )
        .build();

    p.cargo("build").run();
}

#[test]
fn update_precise() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").publish();
    Package::new("serde", "0.2.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.2"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();

    Package::new("serde", "0.2.0").publish();

    p.cargo("update -p serde:0.2.1 --precise 0.2.0")
        .with_stderr(
            "\
[UPDATING] `[..]` index
[UPDATING] serde v0.2.1 -> v0.2.0
",
        )
        .run();
}

#[test]
fn preserve_top_comment() {
    let p = project().file("src/lib.rs", "").build();

    p.cargo("update").run();

    let mut lockfile = p.read_file("Cargo.lock");
    lockfile.insert_str(0, "# @generated\n");
    lockfile.insert_str(0, "# some other comment\n");
    println!("saving Cargo.lock contents:\n{}", lockfile);

    p.change_file("Cargo.lock", &lockfile);

    p.cargo("update").run();

    let lockfile2 = p.read_file("Cargo.lock");
    println!("loaded Cargo.lock contents:\n{}", lockfile2);

    assert!(lockfile == lockfile2);
}

#[test]
fn dry_run_update() {
    Package::new("log", "0.1.0").publish();
    Package::new("serde", "0.1.0").dep("log", "0.1").publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "bar"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
                log = "0.1"
                foo = { path = "foo" }
            "#,
        )
        .file("src/lib.rs", "")
        .file(
            "foo/Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.1"
                authors = []

                [dependencies]
                serde = "0.1"
            "#,
        )
        .file("foo/src/lib.rs", "")
        .build();

    p.cargo("build").run();
    let old_lockfile = p.read_file("Cargo.lock");

    Package::new("log", "0.1.1").publish();
    Package::new("serde", "0.1.1").dep("log", "0.1").publish();

    p.cargo("update -p serde --dry-run")
        .with_stderr(
            "\
[UPDATING] `[..]` index
[UPDATING] serde v0.1.0 -> v0.1.1
[WARNING] not updating lockfile due to dry run
",
        )
        .run();
    let new_lockfile = p.read_file("Cargo.lock");
    assert_eq!(old_lockfile, new_lockfile)
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use pdb::FallibleIterator;

fn setup<F>(func: F)
where
    F: FnOnce(&pdb::TypeInformation<'_>) -> (),
{
    let file = if let Ok(filename) = std::env::var("PDB_FILE") {
        std::fs::File::open(filename)
    } else {
        std::fs::File::open("fixtures/self/foo.pdb")
    }
    .expect("opening file");

    let mut pdb = pdb::PDB::open(file).expect("opening pdb");
    let type_information = pdb.type_information().expect("type information");

    func(&type_information);
}

#[test]
fn iteration() {
    setup(|type_information| {
        let len = type_information.len();

        let mut count: usize = 0;
        let mut last_index = pdb::TypeIndex(4095);
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            assert_eq!(typ.type_index().0, last_index.0 + 1);
            last_index = typ.type_index();
            count += 1;
        }

        assert_eq!(len, count);
    });
}

#[test]
fn type_finder() {
    setup(|type_information| {
        let mut type_finder = type_information.type_finder();
        let mut map: HashMap<pdb::TypeIndex, pdb::Type<'_>> = HashMap::new();

        assert_eq!(type_finder.max_indexed_type().0 >> 3, 4096 >> 3);

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            assert_eq!(type_finder.max_indexed_type().0 >> 3, typ.type_index().0 >> 3);

            // update the type finder
            type_finder.update(&iter);

            // record this type in our map
            map.insert(typ.type_index(), typ);
        }

        // iterate over the map -- which is randomized -- making sure the type finder finds identical types
        for (type_index, typ) in map.iter() {
            let found = type_finder.find(*type_index).expect("find");
            assert_eq!(*typ, found);
        }
    })
}

#[test]
fn find_classes() {
    setup(|type_information| {
        let mut type_finder = type_information.type_finder();

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            // update the type finder
            type_finder.update(&iter);

            // parse the type record
            match typ.parse() {
                Ok(pdb::TypeData::Class(pdb::ClassType {
                    name,
                    fields: Some(fields),
                    ..
                })) => {
                    // this Type describes a class-like type with fields
                    println!("class {} (type {}):", name, typ.type_index());

                    // fields is presently a TypeIndex
                    // find and parse the list of fields
                    match type_finder.find(fields).expect("find fields").parse() {
                        Ok(pdb::TypeData::FieldList(list)) => {
                            for field in list.fields {
                                println!("  - {:?}", field);
                            }

                            if let Some(c) = list.continuation {
                                println!("TODO: follow to type {}", c);
                            }
                        }
                        Ok(value) => {
                            panic!("expected a field list, got {:?}", value);
                        }
                        Err(e) => {
                            println!("field parse error: {}", e);
                        }
                    }
                }
                Ok(pdb::TypeData::Enumeration(data)) => {
                    println!("enum {} (type {}):", data.name, data.fields);

                    // fields is presently a TypeIndex
                    match type_finder.find(data.fields).expect("find fields").parse() {
                        Ok(pdb::TypeData::FieldList(list)) => {
                            for field in list.fields {
                                println!("  - {:?}", field);
                            }

                            if let Some(c) = list.continuation {
                                println!("TODO: follow to type {}", c);
                            }
                        }
                        Ok(value) => {
                            panic!("expected a field list, got {:?}", value);
                        }
                        Err(e) => {
                            println!("field parse error: {}", e);
                        }
                    }
                }
                Ok(pdb::TypeData::FieldList(_)) => {
                    // ignore, since we find these by class
                }
                Ok(_) => {
                    //println!("type: {:?}", data);
                }
                Err(pdb::Error::UnimplementedTypeKind(kind)) => {
                    println!("unimplemented: 0x{:04x}", kind);
                    // TODO: parse everything
                    // ignore for now
                }
                Err(e) => {
                    // other parse error
                    println!(
                        "other parse error on type {} (raw type {:04x}): {}",
                        typ.type_index(),
                        typ.raw_kind(),
                        e
                    );
                    panic!("dying due to parse error");
                }
            }
        }

        // hooah!
    })
}

/*
#[bench]
fn bench_type_finder(b: &mut test::Bencher) {
    setup(|type_information| {
        let mut type_finder = type_information.new_type_finder();

        assert_eq!(type_finder.max_indexed_type() >> 3, 4096 >> 3);

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            assert_eq!(type_finder.max_indexed_type() >> 3, typ.type_index() >> 3);
            type_finder.update(&iter);
        }

        let mut rng = rand::thread_rng();
        let count: pdb::TypeIndex = type_information.len() as pdb::TypeIndex;
        let base: pdb::TypeIndex = 4096;

        // time how long it takes to build a map
        b.iter(|| {
            let lucky = rng.gen_range(base, base + count);
            let found = type_finder.find(lucky).expect("find");
            test::black_box(&found);
        });
    })
}
*/

/*
#[test]
fn type_length_histogram() {
    setup(|type_information| {
        let mut lens: Vec<usize> = Vec::new();
        lens.resize(1025, 0);

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            let mut len = typ.len() + 2;
            if len > 1024 {
                len = 1024;
            }
            lens[len] += 1;
        }

        for (len, count) in lens.as_slice().iter().enumerate() {
            println!("{}\t{}", len, count);
        }

        panic!();
    })
}
*/
<commit_msg>Fix formatting<commit_after>use std::collections::HashMap;

use pdb::FallibleIterator;

fn setup<F>(func: F)
where
    F: FnOnce(&pdb::TypeInformation<'_>) -> (),
{
    let file = if let Ok(filename) = std::env::var("PDB_FILE") {
        std::fs::File::open(filename)
    } else {
        std::fs::File::open("fixtures/self/foo.pdb")
    }
    .expect("opening file");

    let mut pdb = pdb::PDB::open(file).expect("opening pdb");
    let type_information = pdb.type_information().expect("type information");

    func(&type_information);
}

#[test]
fn iteration() {
    setup(|type_information| {
        let len = type_information.len();

        let mut count: usize = 0;
        let mut last_index = pdb::TypeIndex(4095);
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            assert_eq!(typ.type_index().0, last_index.0 + 1);
            last_index = typ.type_index();
            count += 1;
        }

        assert_eq!(len, count);
    });
}

#[test]
fn type_finder() {
    setup(|type_information| {
        let mut type_finder = type_information.type_finder();
        let mut map: HashMap<pdb::TypeIndex, pdb::Type<'_>> = HashMap::new();

        assert_eq!(type_finder.max_indexed_type().0 >> 3, 4096 >> 3);

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            assert_eq!(
                type_finder.max_indexed_type().0 >> 3,
                typ.type_index().0 >> 3
            );

            // update the type finder
            type_finder.update(&iter);

            // record this type in our map
            map.insert(typ.type_index(), typ);
        }

        // iterate over the map -- which is randomized -- making sure the type finder finds identical types
        for (type_index, typ) in map.iter() {
            let found = type_finder.find(*type_index).expect("find");
            assert_eq!(*typ, found);
        }
    })
}

#[test]
fn find_classes() {
    setup(|type_information| {
        let mut type_finder = type_information.type_finder();

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            // update the type finder
            type_finder.update(&iter);

            // parse the type record
            match typ.parse() {
                Ok(pdb::TypeData::Class(pdb::ClassType {
                    name,
                    fields: Some(fields),
                    ..
                })) => {
                    // this Type describes a class-like type with fields
                    println!("class {} (type {}):", name, typ.type_index());

                    // fields is presently a TypeIndex
                    // find and parse the list of fields
                    match type_finder.find(fields).expect("find fields").parse() {
                        Ok(pdb::TypeData::FieldList(list)) => {
                            for field in list.fields {
                                println!("  - {:?}", field);
                            }

                            if let Some(c) = list.continuation {
                                println!("TODO: follow to type {}", c);
                            }
                        }
                        Ok(value) => {
                            panic!("expected a field list, got {:?}", value);
                        }
                        Err(e) => {
                            println!("field parse error: {}", e);
                        }
                    }
                }
                Ok(pdb::TypeData::Enumeration(data)) => {
                    println!("enum {} (type {}):", data.name, data.fields);

                    // fields is presently a TypeIndex
                    match type_finder.find(data.fields).expect("find fields").parse() {
                        Ok(pdb::TypeData::FieldList(list)) => {
                            for field in list.fields {
                                println!("  - {:?}", field);
                            }

                            if let Some(c) = list.continuation {
                                println!("TODO: follow to type {}", c);
                            }
                        }
                        Ok(value) => {
                            panic!("expected a field list, got {:?}", value);
                        }
                        Err(e) => {
                            println!("field parse error: {}", e);
                        }
                    }
                }
                Ok(pdb::TypeData::FieldList(_)) => {
                    // ignore, since we find these by class
                }
                Ok(_) => {
                    //println!("type: {:?}", data);
                }
                Err(pdb::Error::UnimplementedTypeKind(kind)) => {
                    println!("unimplemented: 0x{:04x}", kind);
                    // TODO: parse everything
                    // ignore for now
                }
                Err(e) => {
                    // other parse error
                    println!(
                        "other parse error on type {} (raw type {:04x}): {}",
                        typ.type_index(),
                        typ.raw_kind(),
                        e
                    );
                    panic!("dying due to parse error");
                }
            }
        }

        // hooah!
    })
}

/*
#[bench]
fn bench_type_finder(b: &mut test::Bencher) {
    setup(|type_information| {
        let mut type_finder = type_information.new_type_finder();

        assert_eq!(type_finder.max_indexed_type() >> 3, 4096 >> 3);

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            assert_eq!(type_finder.max_indexed_type() >> 3, typ.type_index() >> 3);
            type_finder.update(&iter);
        }

        let mut rng = rand::thread_rng();
        let count: pdb::TypeIndex = type_information.len() as pdb::TypeIndex;
        let base: pdb::TypeIndex = 4096;

        // time how long it takes to build a map
        b.iter(|| {
            let lucky = rng.gen_range(base, base + count);
            let found = type_finder.find(lucky).expect("find");
            test::black_box(&found);
        });
    })
}
*/

/*
#[test]
fn type_length_histogram() {
    setup(|type_information| {
        let mut lens: Vec<usize> = Vec::new();
        lens.resize(1025, 0);

        // iterate over all the types
        let mut iter = type_information.iter();
        while let Some(typ) = iter.next().expect("next type") {
            let mut len = typ.len() + 2;
            if len > 1024 {
                len = 1024;
            }
            lens[len] += 1;
        }

        for (len, count) in lens.as_slice().iter().enumerate() {
            println!("{}\t{}", len, count);
        }

        panic!();
    })
}
*/
<|endoftext|>


<commit_before>use bson::Bson;
use std::collections::BTreeMap;
use std::iter::{FromIterator, Map};
use std::vec::IntoIter;
use std::slice;

/// A BSON document represented as an associative BTree Map with insertion ordering.
#[derive(Debug, Clone)]
pub struct OrderedDocument {
    pub keys: Vec<String>,
    document: BTreeMap<String, Bson>,
}

/// An iterator over OrderedDocument entries.
pub struct OrderedDocumentIntoIterator {
    vec_iter: IntoIter<String>,
    document: BTreeMap<String, Bson>,
}

/// An owning iterator over OrderedDocument entries.
pub struct OrderedDocumentIterator<'a> {
    vec_iter: slice::Iter<'a, String>,
    document: &'a BTreeMap<String, Bson>,
}

/// An iterator over an OrderedDocument's keys.
pub struct Keys<'a> {
    inner: Map<OrderedDocumentIterator<'a>, fn((&'a String, &'a Bson)) -> &'a String>
}

/// An iterator over an OrderedDocument's values.
pub struct Values<'a> {
    inner: Map<OrderedDocumentIterator<'a>, fn((&'a String, &'a Bson)) -> &'a Bson>
}

impl<'a> Iterator for Keys<'a> {
    type Item = &'a String;
    fn next(&mut self) -> Option<(&'a String)> { self.inner.next() }
}

impl<'a> Iterator for Values<'a> {
    type Item = &'a Bson;

    fn next(&mut self) -> Option<(&'a Bson)> { self.inner.next() }
}

impl IntoIterator for OrderedDocument {
    type Item = (String, Bson);
    type IntoIter = OrderedDocumentIntoIterator;

    fn into_iter(self) -> Self::IntoIter {
        OrderedDocumentIntoIterator {
            document: self.document,
            vec_iter: self.keys.into_iter()
        }
    }
}

impl<'a> IntoIterator for &'a OrderedDocument {
    type Item = (&'a String, &'a Bson);
    type IntoIter = OrderedDocumentIterator<'a>;

    fn into_iter(self) -> Self::IntoIter {
        let ref keys = self.keys;
        OrderedDocumentIterator {
            vec_iter: keys.into_iter(),
            document: &self.document,
        }
    }
}

impl FromIterator<(String, Bson)> for OrderedDocument {
    fn from_iter<T: IntoIterator<Item=(String, Bson)>>(iter: T) -> Self {
        let mut doc = OrderedDocument::new();
        for (k, v) in iter {
            doc.insert(k, v.to_owned());
        }
        doc
    }
}

impl<'a> Iterator for OrderedDocumentIntoIterator {
    type Item = (String, Bson);
    fn next(&mut self) -> Option<(String, Bson)> {
        match self.vec_iter.next() {
            Some(key) => {
                let val = self.document.get(&key[..]).unwrap();
                Some((key, val.to_owned()))
            },
            None => None,
        }
    }
}

impl<'a> Iterator for OrderedDocumentIterator<'a> {
    type Item = (&'a String, &'a Bson);
    fn next(&mut self) -> Option<(&'a String, &'a Bson)> {
        match self.vec_iter.next() {
            Some(key) => Some((&key, self.document.get(&key[..]).unwrap())),
            None => None,
        }
    }
}

impl OrderedDocument {
    /// Creates a new empty OrderedDocument.
    pub fn new() -> OrderedDocument {
        OrderedDocument {
            keys: Vec::new(),
            document: BTreeMap::new(),
        }
    }

    /// Gets an iterator over the entries of the map.
    pub fn iter<'a>(&'a self) -> OrderedDocumentIterator<'a> {
        self.into_iter()
    }

    /// Clears the document, removing all values.
    pub fn clear(&mut self) {
        self.keys.clear();
        self.document.clear();
    }

    /// Returns a reference to the Bson corresponding to the key.
    pub fn get(&self, key: &str) -> Option<&Bson> {
        self.document.get(key)
    }

    /// Gets a mutable reference to the value in the entry.
    pub fn get_mut(&mut self, key: &str) -> Option<&mut Bson> {
        self.document.get_mut(key)
    }

    /// Returns true if the map contains a value for the specified key.
    pub fn contains_key(&self, key: &str) -> bool {
        self.document.contains_key(key)
    }

    /// Returns the position of the key in the ordered vector, if it exists.
    pub fn position(&self, key: &str) -> Option<usize> {
        self.keys.iter().position(|x| x == key)
    }

    /// Gets a collection of all keys in the document.
    pub fn keys<'a>(&'a self) -> Keys<'a> {
        fn first<A, B>((a, _): (A, B)) -> A { a }
        let first: fn((&'a String, &'a Bson)) -> &'a String = first;

        Keys { inner: self.iter().map(first) }
    }

    /// Gets a collection of all values in the document.
    pub fn values<'a>(&'a self) -> Values<'a> {
        fn second<A, B>((_, b): (A, B)) -> B { b }
        let second: fn((&'a String, &'a Bson)) -> &'a Bson = second;

        Values { inner: self.iter().map(second) }
    }

    /// Returns the number of elements in the document.
    pub fn len(&self) -> usize {
        self.keys.len()
    }

    /// Returns true if the document contains no elements
    pub fn is_empty(&self) -> bool {
        self.document.is_empty()
    }

    /// Sets the value of the entry with the OccupiedEntry's key,
    /// and returns the entry's old value.
    pub fn insert(&mut self, key: String, val: Bson) -> Option<Bson> {
        let key_slice = &key[..];

        if self.contains_key(key_slice) {
            let position = self.position(key_slice).unwrap();
            self.keys.remove(position);
        }

        self.keys.push(key.to_owned());
        self.document.insert(key.to_owned(), val.to_owned())
    }

    /// Takes the value of the entry out of the document, and returns it.
    pub fn remove(&mut self, key: &str) -> Option<Bson> {
        let position = self.position(key);
        if position.is_some() {
            self.keys.remove(position.unwrap());
        }
        self.document.remove(key)
    }
}

#[cfg(test)]
mod test {
    use super::OrderedDocument;
    use bson::Bson;

    #[test]
    fn ordered_insert() {
        let mut doc = OrderedDocument::new();
        doc.insert("first".to_owned(), Bson::I32(1));
        doc.insert("second".to_owned(), Bson::String("foo".to_owned()));
        doc.insert("alphanumeric".to_owned(), Bson::String("bar".to_owned()));

        let expected_keys = vec!(
            "first".to_owned(),
            "second".to_owned(),
            "alphanumeric".to_owned(),
        );

        let keys: Vec<_> = doc.iter().map(|(key, _)| key.to_owned()).collect();
        assert_eq!(expected_keys, keys);
    }

    #[test]
    fn remove() {
        let mut doc = OrderedDocument::new();
        doc.insert("first".to_owned(), Bson::I32(1));
        doc.insert("second".to_owned(), Bson::String("foo".to_owned()));
        doc.insert("alphanumeric".to_owned(), Bson::String("bar".to_owned()));

        assert!(doc.remove("second").is_some());
        assert!(doc.remove("none").is_none());

        let expected_keys = vec!(
            "first",
            "alphanumeric",
        );

        let keys: Vec<_> = doc.iter().map(|(key, _)| key.to_owned()).collect();
        assert_eq!(expected_keys, keys);
    }
}
<commit_msg>Use BTreeMap.remove() in consuming iterator<commit_after>use bson::Bson;
use std::collections::BTreeMap;
use std::iter::{FromIterator, Map};
use std::vec::IntoIter;
use std::slice;

/// A BSON document represented as an associative BTree Map with insertion ordering.
#[derive(Debug, Clone)]
pub struct OrderedDocument {
    pub keys: Vec<String>,
    document: BTreeMap<String, Bson>,
}

/// An iterator over OrderedDocument entries.
pub struct OrderedDocumentIntoIterator {
    vec_iter: IntoIter<String>,
    document: BTreeMap<String, Bson>,
}

/// An owning iterator over OrderedDocument entries.
pub struct OrderedDocumentIterator<'a> {
    vec_iter: slice::Iter<'a, String>,
    document: &'a BTreeMap<String, Bson>,
}

/// An iterator over an OrderedDocument's keys.
pub struct Keys<'a> {
    inner: Map<OrderedDocumentIterator<'a>, fn((&'a String, &'a Bson)) -> &'a String>
}

/// An iterator over an OrderedDocument's values.
pub struct Values<'a> {
    inner: Map<OrderedDocumentIterator<'a>, fn((&'a String, &'a Bson)) -> &'a Bson>
}

impl<'a> Iterator for Keys<'a> {
    type Item = &'a String;
    fn next(&mut self) -> Option<(&'a String)> { self.inner.next() }
}

impl<'a> Iterator for Values<'a> {
    type Item = &'a Bson;
    fn next(&mut self) -> Option<(&'a Bson)> { self.inner.next() }
}

impl IntoIterator for OrderedDocument {
    type Item = (String, Bson);
    type IntoIter = OrderedDocumentIntoIterator;

    fn into_iter(self) -> Self::IntoIter {
        OrderedDocumentIntoIterator {
            document: self.document,
            vec_iter: self.keys.into_iter()
        }
    }
}

impl<'a> IntoIterator for &'a OrderedDocument {
    type Item = (&'a String, &'a Bson);
    type IntoIter = OrderedDocumentIterator<'a>;

    fn into_iter(self) -> Self::IntoIter {
        let ref keys = self.keys;
        OrderedDocumentIterator {
            vec_iter: keys.into_iter(),
            document: &self.document,
        }
    }
}

impl FromIterator<(String, Bson)> for OrderedDocument {
    fn from_iter<T: IntoIterator<Item=(String, Bson)>>(iter: T) -> Self {
        let mut doc = OrderedDocument::new();
        for (k, v) in iter {
            doc.insert(k, v.to_owned());
        }
        doc
    }
}

impl<'a> Iterator for OrderedDocumentIntoIterator {
    type Item = (String, Bson);
    fn next(&mut self) -> Option<(String, Bson)> {
        match self.vec_iter.next() {
            Some(key) => {
                let val = self.document.remove(&key[..]).unwrap();
                Some((key, val.to_owned()))
            },
            None => None,
        }
    }
}

impl<'a> Iterator for OrderedDocumentIterator<'a> {
    type Item = (&'a String, &'a Bson);
    fn next(&mut self) -> Option<(&'a String, &'a Bson)> {
        match self.vec_iter.next() {
            Some(key) => {
                let val = self.document.get(&key[..]).unwrap();
                Some((&key, val))
            },
            None => None,
        }
    }
}

impl OrderedDocument {
    /// Creates a new empty OrderedDocument.
    pub fn new() -> OrderedDocument {
        OrderedDocument {
            keys: Vec::new(),
            document: BTreeMap::new(),
        }
    }

    /// Gets an iterator over the entries of the map.
    pub fn iter<'a>(&'a self) -> OrderedDocumentIterator<'a> {
        self.into_iter()
    }

    /// Clears the document, removing all values.
    pub fn clear(&mut self) {
        self.keys.clear();
        self.document.clear();
    }

    /// Returns a reference to the Bson corresponding to the key.
    pub fn get(&self, key: &str) -> Option<&Bson> {
        self.document.get(key)
    }

    /// Gets a mutable reference to the value in the entry.
    pub fn get_mut(&mut self, key: &str) -> Option<&mut Bson> {
        self.document.get_mut(key)
    }

    /// Returns true if the map contains a value for the specified key.
    pub fn contains_key(&self, key: &str) -> bool {
        self.document.contains_key(key)
    }

    /// Returns the position of the key in the ordered vector, if it exists.
    pub fn position(&self, key: &str) -> Option<usize> {
        self.keys.iter().position(|x| x == key)
    }

    /// Gets a collection of all keys in the document.
    pub fn keys<'a>(&'a self) -> Keys<'a> {
        fn first<A, B>((a, _): (A, B)) -> A { a }
        let first: fn((&'a String, &'a Bson)) -> &'a String = first;

        Keys { inner: self.iter().map(first) }
    }

    /// Gets a collection of all values in the document.
    pub fn values<'a>(&'a self) -> Values<'a> {
        fn second<A, B>((_, b): (A, B)) -> B { b }
        let second: fn((&'a String, &'a Bson)) -> &'a Bson = second;

        Values { inner: self.iter().map(second) }
    }

    /// Returns the number of elements in the document.
    pub fn len(&self) -> usize {
        self.keys.len()
    }

    /// Returns true if the document contains no elements
    pub fn is_empty(&self) -> bool {
        self.document.is_empty()
    }

    /// Sets the value of the entry with the OccupiedEntry's key,
    /// and returns the entry's old value.
    pub fn insert(&mut self, key: String, val: Bson) -> Option<Bson> {
        let key_slice = &key[..];

        if self.contains_key(key_slice) {
            let position = self.position(key_slice).unwrap();
            self.keys.remove(position);
        }

        self.keys.push(key.to_owned());
        self.document.insert(key.to_owned(), val.to_owned())
    }

    /// Takes the value of the entry out of the document, and returns it.
    pub fn remove(&mut self, key: &str) -> Option<Bson> {
        let position = self.position(key);
        if position.is_some() {
            self.keys.remove(position.unwrap());
        }
        self.document.remove(key)
    }
}

#[cfg(test)]
mod test {
    use super::OrderedDocument;
    use bson::Bson;

    #[test]
    fn ordered_insert() {
        let mut doc = OrderedDocument::new();
        doc.insert("first".to_owned(), Bson::I32(1));
        doc.insert("second".to_owned(), Bson::String("foo".to_owned()));
        doc.insert("alphanumeric".to_owned(), Bson::String("bar".to_owned()));

        let expected_keys = vec!(
            "first".to_owned(),
            "second".to_owned(),
            "alphanumeric".to_owned(),
        );

        let keys: Vec<_> = doc.iter().map(|(key, _)| key.to_owned()).collect();
        assert_eq!(expected_keys, keys);
    }

    #[test]
    fn remove() {
        let mut doc = OrderedDocument::new();
        doc.insert("first".to_owned(), Bson::I32(1));
        doc.insert("second".to_owned(), Bson::String("foo".to_owned()));
        doc.insert("alphanumeric".to_owned(), Bson::String("bar".to_owned()));

        assert!(doc.remove("second").is_some());
        assert!(doc.remove("none").is_none());

        let expected_keys = vec!(
            "first",
            "alphanumeric",
        );

        let keys: Vec<_> = doc.iter().map(|(key, _)| key.to_owned()).collect();
        assert_eq!(expected_keys, keys);
    }
}
<|endoftext|>


<commit_before>use std::{fmt, iter, ptr};

// ////////////////////////////////////////////////////////////////////////////
// SkipNode
// ////////////////////////////////////////////////////////////////////////////

/// SkipNodes are make up the SkipList.  The SkipList owns the first head-node
/// (which has no value) and each node has ownership of the next node through
/// `next`.
///
/// The node has a `level` which corresponds to how 'high' the node reaches, and
/// this should be equal to the length of the vector of links.  There is a
/// corresponding vector of link lengths which is used to reach a particular
/// index.
///
/// Lastly, each node contains a link to the immediately previous node in case
/// one needs to parse the list backwards.
///
/// In cases where the value is not applicable, `None` should be used.  In
/// particular, as there is no tail node, the value of `next` in the last node
/// should be `None`.
#[derive(Clone, Debug)]
pub struct SkipNode<V> {
    // key and value should never be None, with the sole exception being the
    // head node.
    pub value: Option<V>,
    // how high the node reaches.  This should be equal to the vector length.
    pub level: usize,
    // The immediately next element (and owns that next node).
    pub next: Option<Box<SkipNode<V>>>,
    // The immediately previous element.
    pub prev: Option<*mut SkipNode<V>>,
    // Vector of links to the next node at the respective level.  This vector
    // *must* be of length `self.level + 1`.  links[0] stores a pointer to the
    // same node as next.
    pub links: Vec<Option<*mut SkipNode<V>>>,
    // The corresponding length of each link
    pub links_len: Vec<usize>,
}

// ///////////////////////////////////////////////
// Inherent methods
// ///////////////////////////////////////////////

impl<V> SkipNode<V> {
    /// Create a new head node.
    pub fn head(total_levels: usize) -> Self {
        SkipNode {
            value: None,
            level: total_levels - 1,
            next: None,
            prev: None,
            links: iter::repeat(None).take(total_levels).collect(),
            links_len: iter::repeat(0).take(total_levels).collect(),
        }
    }

    /// Create a new SkipNode with the given value.  The values of `prev` and
    /// `next` will all be `None` and have to be adjusted.
    pub fn new(value: V, level: usize) -> Self {
        SkipNode {
            value: Some(value),
            level,
            next: None,
            prev: None,
            links: iter::repeat(None).take(level + 1).collect(),
            links_len: iter::repeat(0).take(level + 1).collect(),
        }
    }

    /// Consumes the node returning the value it contains.
    pub fn into_inner(mut self) -> Option<V> {
        self.value.take()
    }

    /// Returns `true` is the node is a head-node.
    pub fn is_head(&self) -> bool {
        self.prev.is_none()
    }
}

impl<V> Drop for SkipNode<V> {
    fn drop(&mut self) {
        while let Some(mut node) = self.next.take() {
            self.next = node.next.take();
        }
    }
}

// ///////////////////////////////////////////////
// Trait implementation
// ///////////////////////////////////////////////

impl<V> fmt::Display for SkipNode<V>
where
    V: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref v) = self.value {
            write!(f, "{}", v)
        } else {
            Ok(())
        }
    }
}

// /////////////////////////////////
// Iterators
// /////////////////////////////////
// Since Iterators (currently) only pop from front and back,
// they can be shared by some data structures.
// There's no need for a dummy head (that contains no value) in the iterator.
// so the members are named first and last instaed of head/end to avoid confusion.

/// Consuming iterator.
pub struct IntoIter<T> {
    pub(crate) first: Option<Box<SkipNode<T>>>,
    pub(crate) last: *mut SkipNode<T>,
    pub(crate) size: usize,
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let mut popped_node = self.first.take()?;
        self.size -= 1;
        self.first = popped_node.next.take().map(|mut node| {
            node.prev = None;
            node
        });
        if self.first.is_none() {
            self.last = ptr::null_mut();
        }
        popped_node.into_inner()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.size, Some(self.size))
    }
}

impl<T> DoubleEndedIterator for IntoIter<T> {
    fn next_back(&mut self) -> Option<T> {
        if self.first.is_none() {
            return None;
        }
        assert!(!self.last.is_null());
        unsafe {
            let new_last = (*self.last).prev;
            let popped_node = match new_last {
                Some(new_last) => {
                    self.last = new_last;
                    (*new_last).next.take().unwrap()
                }
                None => self.first.take().unwrap(),
            };
            self.size -= 1;
            popped_node.into_inner()
        }
    }
}
<commit_msg>Add safe methods for pointer manupilation in SkipNode.<commit_after>use std::{fmt, iter, ptr};

pub enum NextNotFound {
    NoMoreElements,
    NoSuchElement,
}

// ////////////////////////////////////////////////////////////////////////////
// SkipNode
// ////////////////////////////////////////////////////////////////////////////

/// SkipNodes are make up the SkipList.  The SkipList owns the first head-node
/// (which has no value) and each node has ownership of the next node through
/// `next`.
///
/// The node has a `level` which corresponds to how 'high' the node reaches, and
/// this should be equal to the length of the vector of links.  There is a
/// corresponding vector of link lengths which is used to reach a particular
/// index.
///
/// Lastly, each node contains a link to the immediately previous node in case
/// one needs to parse the list backwards.
///
/// In cases where the value is not applicable, `None` should be used.  In
/// particular, as there is no tail node, the value of `next` in the last node
/// should be `None`.
#[derive(Clone, Debug)]
pub struct SkipNode<V> {
    // key and value should never be None, with the sole exception being the
    // head node.
    pub value: Option<V>,
    // how high the node reaches.  This should be equal to the vector length.
    pub level: usize,
    // The immediately next element (and owns that next node).
    pub next: Option<Box<SkipNode<V>>>,
    // The immediately previous element.
    pub prev: Option<*mut SkipNode<V>>,
    // Vector of links to the next node at the respective level.  This vector
    // *must* be of length `self.level + 1`.  links[0] stores a pointer to the
    // same node as next.
    pub links: Vec<Option<*mut SkipNode<V>>>,
    // The corresponding length of each link
    pub links_len: Vec<usize>,
}

// ///////////////////////////////////////////////
// Inherent methods
// ///////////////////////////////////////////////

impl<V> SkipNode<V> {
    /// Create a new head node.
    pub fn head(total_levels: usize) -> Self {
        SkipNode {
            value: None,
            level: total_levels - 1,
            next: None,
            prev: None,
            links: iter::repeat(None).take(total_levels).collect(),
            links_len: iter::repeat(0).take(total_levels).collect(),
        }
    }

    /// Create a new SkipNode with the given value.  The values of `prev` and
    /// `next` will all be `None` and have to be adjusted.
    pub fn new(value: V, level: usize) -> Self {
        SkipNode {
            value: Some(value),
            level,
            next: None,
            prev: None,
            links: iter::repeat(None).take(level + 1).collect(),
            links_len: iter::repeat(0).take(level + 1).collect(),
        }
    }

    /// Consumes the node returning the value it contains.
    pub fn into_inner(mut self) -> Option<V> {
        self.value.take()
    }

    /// Returns `true` is the node is a head-node.
    pub fn is_head(&self) -> bool {
        self.prev.is_none()
    }

    pub fn advance_if(
        &self,
        level: usize,
        predicate: impl FnOnce(&Self, &Self) -> bool,
    ) -> Result<&Self, NextNotFound> {
        let next = unsafe {
            self.links[level]
                .and_then(|ptr| ptr.as_ref())
                .ok_or(NextNotFound::NoMoreElements)?
        };
        if predicate(self, next) {
            Ok(next)
        } else {
            Err(NextNotFound::NoSuchElement)
        }
    }

    pub fn advance_if_mut(
        &mut self,
        level: usize,
        predicate: impl FnOnce(&Self, &Self) -> bool,
    ) -> Result<&mut Self, NextNotFound> {
        let next = unsafe {
            self.links[level]
                .and_then(|ptr| ptr.as_mut())
                .ok_or(NextNotFound::NoMoreElements)?
        };
        if predicate(self, next) {
            Ok(next)
        } else {
            Err(NextNotFound::NoSuchElement)
        }
    }
}

impl<V> Drop for SkipNode<V> {
    fn drop(&mut self) {
        while let Some(mut node) = self.next.take() {
            self.next = node.next.take();
        }
    }
}

// ///////////////////////////////////////////////
// Trait implementation
// ///////////////////////////////////////////////

impl<V> fmt::Display for SkipNode<V>
where
    V: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref v) = self.value {
            write!(f, "{}", v)
        } else {
            Ok(())
        }
    }
}

// /////////////////////////////////
// Iterators
// /////////////////////////////////
// Since Iterators (currently) only pop from front and back,
// they can be shared by some data structures.
// There's no need for a dummy head (that contains no value) in the iterator.
// so the members are named first and last instaed of head/end to avoid confusion.

/// Consuming iterator.
pub struct IntoIter<T> {
    pub(crate) first: Option<Box<SkipNode<T>>>,
    pub(crate) last: *mut SkipNode<T>,
    pub(crate) size: usize,
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let mut popped_node = self.first.take()?;
        self.size -= 1;
        self.first = popped_node.next.take().map(|mut node| {
            node.prev = None;
            node
        });
        if self.first.is_none() {
            self.last = ptr::null_mut();
        }
        popped_node.into_inner()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.size, Some(self.size))
    }
}

impl<T> DoubleEndedIterator for IntoIter<T> {
    fn next_back(&mut self) -> Option<T> {
        if self.first.is_none() {
            return None;
        }
        assert!(!self.last.is_null());
        unsafe {
            let new_last = (*self.last).prev;
            let popped_node = match new_last {
                Some(new_last) => {
                    self.last = new_last;
                    (*new_last).next.take().unwrap()
                }
                None => self.first.take().unwrap(),
            };
            self.size -= 1;
            popped_node.into_inner()
        }
    }
}
<|endoftext|>


<commit_before>extern crate bincode;
extern crate rustc_serialize;

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use rustc_serialize::{Encodable, Decodable};

use ::{KeyType, ValueType};

use std::collections::{BTreeMap, BTreeSet};
use std::collections::btree_map::Iter as MIter;
use std::collections::btree_set::Iter as SIter;
use std::error::Error;
use std::fs::{File, OpenOptions};
use std::io::{Read, Write, ErrorKind, Seek, SeekFrom};
use std::io::Error as IOError;
use std::iter::Peekable;
use std::marker::PhantomData;

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
pub struct KeyValuePair<K: KeyType, V: ValueType> {
    pub key: K,
    pub value: V,
}

pub struct WALFile<K: KeyType, V: ValueType> {
    fd: File,  // the WAL file
    key_size: usize,
    value_size: usize,
    _k_marker: PhantomData<K>,
    _v_marker: PhantomData<V>
}

pub struct WALIterator<'a, K: KeyType + 'a, V: ValueType + 'a> {
    wal_file: &'a mut WALFile<K,V>,  // the WAL file
    _k_marker: PhantomData<K>,
    _v_marker: PhantomData<V>
}

impl <K: KeyType, V: ValueType> WALFile<K,V> {
    pub fn new(wal_file_path: String, key_size: usize, value_size: usize) -> Result<WALFile<K,V>, Box<Error>> {
        let mut wal_file = try!(OpenOptions::new().read(true).write(true).create(true).open(wal_file_path));

        return Ok(WALFile{fd: wal_file,
                          key_size: key_size,
                          value_size: value_size,
                          _k_marker: PhantomData,
                          _v_marker: PhantomData});
    }

    pub fn is_new(&self) -> Result<bool, Box<Error>> {
        Ok(try!(self.fd.metadata()).len() == 0)
    }

    pub fn len(&self) -> Result<u64, Box<Error>> {
        Ok(try!(self.fd.metadata()).len())
    }

    pub fn write_record(&mut self, kv: &KeyValuePair<K,V>) -> Result<(), Box<Error>> {
        // encode the record
        let record_size = self.key_size + self.value_size;
        let mut buff = try!(encode(&kv, SizeLimit::Bounded(record_size as u64)));

        // padd it out to the max size
        if buff.len() > self.key_size + self.value_size {
            return Err(From::from(IOError::new(ErrorKind::InvalidData, "Key and value size are too large")));
        } else {
            let diff = (self.key_size + self.value_size) - buff.len();
            buff.extend(vec![0; diff]);
        }

        match self.fd.write_all(&buff) {
            Ok(_) => Ok( () ),
            Err(e) => Err(From::from(e))
        }
    }
}

impl <'a, K: KeyType, V: ValueType> IntoIterator for &'a mut WALFile<K,V> {
    type Item = KeyValuePair<K,V>;
    type IntoIter = WALIterator<'a, K,V>;

    fn into_iter(self) -> Self::IntoIter {
        let k_marker = self._k_marker;
        let v_marker = self._v_marker;

        // seek back to the start
        self.fd.seek(SeekFrom::Start(0));

        // create our iterator
        WALIterator{wal_file: self, _k_marker: k_marker, _v_marker: v_marker}
    }
}

impl <'a, K: KeyType, V: ValueType> Iterator for WALIterator<'a,K,V> {
    type Item = KeyValuePair<K,V>;

    fn next(&mut self) -> Option<Self::Item> {
        let total_size = self.wal_file.key_size + self.wal_file.value_size;
        let mut buff = vec![0; total_size];

        println!("Creating buffer: {}", total_size);

        // attempt to read a buffer's worth and decode
        match self.wal_file.fd.read_exact(&mut buff) {
            Ok(_) => {
                match decode(&buff) {
                    Ok(record) => Some(record),
                    Err(_) => None
                }
            },
            Err(e) => {
                println!("ERROR: {}", e);
                None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use tests::gen_temp_name;
    use std::fs;
    use std::fs::OpenOptions;
    use wal_file::{WALFile, KeyValuePair};

    #[test]
    fn test_iterator() {
        let temp_path = gen_temp_name();
        let file_path = temp_path.to_owned() + ".wal";

        // create a new blank file
        OpenOptions::new().read(true).write(true).truncate(true).create(true).open(&file_path);

        let mut wal_file = WALFile::new(file_path, 20, 20).unwrap();

        let kv1 = KeyValuePair{key: "hello".to_owned(), value: "world".to_owned()};
        let kv2 = KeyValuePair{key: "foo".to_owned(), value: "bar".to_owned()};

        wal_file.write_record(&kv1).unwrap();
        wal_file.write_record(&kv2).unwrap();

        let mut wal_it = wal_file.into_iter();

        let it_kv1 = wal_it.next().unwrap();

        assert!(kv1.key == it_kv1.key);
        assert!(kv1.value == it_kv1.value);

        let it_kv2 = wal_it.next().unwrap();

        assert!(kv2.key == it_kv2.key);
        assert!(kv2.value == it_kv2.value);
    }
}
<commit_msg>Cleaned up wal_file iterator<commit_after>extern crate bincode;
extern crate rustc_serialize;

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};

use ::{KeyType, ValueType};

use std::error::Error;
use std::fs::{File, OpenOptions};
use std::io::{Read, Write, ErrorKind, Seek, SeekFrom};
use std::io::Error as IOError;
use std::marker::PhantomData;

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
pub struct KeyValuePair<K: KeyType, V: ValueType> {
    pub key: K,
    pub value: V,
}

pub struct WALFile<K: KeyType, V: ValueType> {
    fd: File,  // the WAL file
    key_size: usize,
    value_size: usize,
    _k_marker: PhantomData<K>,
    _v_marker: PhantomData<V>
}

pub struct WALIterator<'a, K: KeyType + 'a, V: ValueType + 'a> {
    wal_file: &'a mut WALFile<K,V>,  // the WAL file
}

impl <K: KeyType, V: ValueType> WALFile<K,V> {
    pub fn new(wal_file_path: String, key_size: usize, value_size: usize) -> Result<WALFile<K,V>, Box<Error>> {
        let wal_file = try!(OpenOptions::new().read(true).write(true).create(true).open(wal_file_path));

        return Ok(WALFile{fd: wal_file,
                          key_size: key_size,
                          value_size: value_size,
                          _k_marker: PhantomData,
                          _v_marker: PhantomData});
    }

    pub fn is_new(&self) -> Result<bool, Box<Error>> {
        Ok(try!(self.fd.metadata()).len() == 0)
    }

    pub fn len(&self) -> Result<u64, Box<Error>> {
        Ok(try!(self.fd.metadata()).len())
    }

    pub fn write_record(&mut self, kv: &KeyValuePair<K,V>) -> Result<(), Box<Error>> {
        // encode the record
        let record_size = self.key_size + self.value_size;
        let mut buff = try!(encode(&kv, SizeLimit::Bounded(record_size as u64)));

        // padd it out to the max size
        if buff.len() > self.key_size + self.value_size {
            return Err(From::from(IOError::new(ErrorKind::InvalidData, "Key and value size are too large")));
        } else {
            let diff = (self.key_size + self.value_size) - buff.len();
            buff.extend(vec![0; diff]);
        }

        match self.fd.write_all(&buff) {
            Ok(_) => Ok( () ),
            Err(e) => Err(From::from(e))
        }
    }
}

impl <'a, K: KeyType, V: ValueType> IntoIterator for &'a mut WALFile<K,V> {
    type Item = KeyValuePair<K,V>;
    type IntoIter = WALIterator<'a, K,V>;

    fn into_iter(self) -> Self::IntoIter {
        // seek back to the start
        self.fd.seek(SeekFrom::Start(0));

        // create our iterator
        WALIterator{wal_file: self}
    }
}

impl <'a, K: KeyType, V: ValueType> Iterator for WALIterator<'a,K,V> {
    type Item = KeyValuePair<K,V>;

    fn next(&mut self) -> Option<Self::Item> {
        let total_size = self.wal_file.key_size + self.wal_file.value_size;
        let mut buff = vec![0; total_size];

        println!("Creating buffer: {}", total_size);

        // attempt to read a buffer's worth and decode
        match self.wal_file.fd.read_exact(&mut buff) {
            Ok(_) => {
                match decode(&buff) {
                    Ok(record) => Some(record),
                    Err(_) => None
                }
            },
            Err(e) => {
                println!("ERROR: {}", e);
                None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use tests::gen_temp_name;
    use std::fs;
    use std::fs::OpenOptions;
    use wal_file::{WALFile, KeyValuePair};

    #[test]
    fn test_iterator() {
        let temp_path = gen_temp_name();
        let file_path = temp_path.to_owned() + ".wal";

        // create a new blank file
        OpenOptions::new().read(true).write(true).truncate(true).create(true).open(&file_path);

        let mut wal_file = WALFile::new(file_path, 20, 20).unwrap();

        let kv1 = KeyValuePair{key: "hello".to_owned(), value: "world".to_owned()};
        let kv2 = KeyValuePair{key: "foo".to_owned(), value: "bar".to_owned()};

        wal_file.write_record(&kv1).unwrap();
        wal_file.write_record(&kv2).unwrap();

        let mut wal_it = wal_file.into_iter();

        let it_kv1 = wal_it.next().unwrap();

        assert!(kv1.key == it_kv1.key);
        assert!(kv1.value == it_kv1.value);

        let it_kv2 = wal_it.next().unwrap();

        assert!(kv2.key == it_kv2.key);
        assert!(kv2.value == it_kv2.value);
    }
}
<|endoftext|>


<commit_before>//! Game world generation
//!
//! Unlike the game-agnostic `mapgen` module, this module does have gameplay-specific logic.

use calx::{hex_neighbors, seeded_rng, Dijkstra};
use euclid::{point2, vec2};
use location::{Location, Portal, Sector};
use mapgen::{self, MapGen, Size2D, Vault, VaultCell};
use rand::distributions::{Distribution, Standard};
use rand::{seq, Rng};
use serde;
use spec;
use std::collections::{BTreeSet, HashMap};
use std::iter::FromIterator;
use std::slice;
use terrain::Terrain;
use world::Loadout;

/// Static generated world.
pub struct Worldgen {
    seed: u32,
    terrain: HashMap<Location, Terrain>,
    portals: HashMap<Location, Portal>,
    spawns: Vec<(Location, Loadout)>,
    player_entry: Location,
}

impl Worldgen {
    pub fn new(seed: u32) -> Worldgen {
        let mut ret = Worldgen {
            seed,
            terrain: HashMap::new(),
            portals: HashMap::new(),
            spawns: Vec::new(),
            player_entry: Location::new(0, 0, 0),
        };

        let mut rng = seeded_rng(&seed);

        let mut cave_entrance = Location::new(0, 0, 0);
        for depth in 1..11 {
            let up_stairs;
            let down_stairs;
            let mut spawns = Vec::new();

            {
                let mut digger = SectorDigger::new(&mut ret, Sector::new(0, 0, depth as i16));
                let domain = digger.domain();

                if depth < 5 {
                    mapgen::RoomsAndCorridors.dig(&mut rng, &mut digger, domain);
                } else {
                    mapgen::Caves.dig(&mut rng, &mut digger, domain);
                }

                up_stairs = digger.up_portal.expect("Mapgen didn't create stairs up");
                down_stairs = digger
                    .down_portal
                    .expect("Mapgen didn't create stairs down");

                // Spawn
                digger.clear_spawns_near_entrance(12);

                let mut spawn_locs =
                    seq::sample_iter(&mut rng, digger.spawn_region.iter(), 20).unwrap();
                let n_spawns = spawn_locs.len();

                for &loc in spawn_locs.drain(0..n_spawns / 2) {
                    spawns.push((
                        loc,
                        spec::pick(&mut rng, depth, spec::ITEM_SPECS.as_slice())
                            .expect("No item spawn"),
                    ))
                }

                for &loc in spawn_locs {
                    spawns.push((
                        loc,
                        spec::pick(&mut rng, depth, (&*spec::MOB_SPECS).as_slice())
                            .expect("No mob spawn"),
                    ))
                }
            }

            ret.spawns.extend(spawns.into_iter());
            if depth > 1 {
                ret.portal(cave_entrance, up_stairs + vec2(1, 1));
                ret.portal(up_stairs, cave_entrance - vec2(1, 1));
            } else {
                ret.player_entry = up_stairs;
            }

            cave_entrance = down_stairs;

            // TODO: Generator needs an option to not generate stairs down on bottom level
        }

        ret
    }

    pub fn seed(&self) -> u32 { self.seed }

    pub fn get_terrain(&self, loc: Location) -> Terrain {
        if let Some(&t) = self.terrain.get(&loc) {
            t
        } else {
            self.default_terrain(loc)
        }
    }

    fn default_terrain(&self, loc: Location) -> Terrain {
        use Terrain::*;
        if loc.z == 0 {
            /*
            match loc.noise() {
                n if n > 0.8 => Tree,
                n if n > -0.8 => Grass,
                _ => Water,
            }
            */
            Ground
        } else {
            Rock
        }
    }

    pub fn get_portal(&self, loc: Location) -> Option<Location> {
        self.portals.get(&loc).map(|&p| loc + p)
    }

    pub fn spawns(&self) -> slice::Iter<(Location, Loadout)> { self.spawns.iter() }

    pub fn player_entry(&self) -> Location { self.player_entry }

    /// Punch a (one-way) portal between two points.
    fn portal(&mut self, origin: Location, destination: Location) {
        self.portals
            .insert(origin, Portal::new(origin, destination));
    }
}

impl serde::Serialize for Worldgen {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        self.seed.serialize(s)
    }
}

impl<'a> serde::Deserialize<'a> for Worldgen {
    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
        Ok(Worldgen::new(serde::Deserialize::deserialize(d)?))
    }
}

struct SectorDigger<'a> {
    worldgen: &'a mut Worldgen,
    sector: Sector,
    up_portal: Option<Location>,
    down_portal: Option<Location>,
    spawn_region: BTreeSet<Location>,
}

impl<'a> SectorDigger<'a> {
    fn new(worldgen: &'a mut Worldgen, sector: Sector) -> SectorDigger<'a> {
        SectorDigger {
            worldgen,
            sector,
            up_portal: None,
            down_portal: None,
            spawn_region: BTreeSet::new(),
        }
    }

    fn domain(&self) -> impl Iterator<Item = mapgen::Point2D> {
        let sector_origin = self.sector.origin();
        let mapgen_origin = mapgen::Point2D::zero();

        self.sector.iter().filter_map(move |loc| {
            let pos = mapgen_origin
                + sector_origin
                    .v2_at(loc)
                    .expect("Sector points are not Euclidean");

            // Okay, this part is a bit hairy, hang on.
            // Sectors are arranged in a rectangular grid, so there are some cells where you can
            // step diagonally across two sector boundaries. However, we want to pretend that
            // sectors are only connected in the four cardinal directions, so we omit these cells
            // from the domain to prevent surprise holes between two diagonally adjacent shafts.
            if Self::is_next_to_diagonal_sector(loc) {
                None
            } else {
                Some(pos)
            }
        })
    }

    fn loc(&self, pos: mapgen::Point2D) -> Location { self.sector.origin() + pos.to_vector() }

    fn set(&mut self, loc: Location, terrain: Terrain) {
        if terrain != Terrain::Gate {
            debug_assert_ne!(self.up_portal, Some(loc), "mapgen overwriting exit");
            debug_assert_ne!(self.down_portal, Some(loc), "mapgen overwriting exit");
        }

        // FIXME: Mapgen currently trips this with exit placement
        // debug_assert!(
        //     !Self::is_next_to_diagonal_sector(loc) &&
        //     self.sector.iter().find(|x| x == &loc).is_some(),
        //     "Setting a location outsides sector area");

        self.worldgen.terrain.insert(loc, terrain);
    }

    /// Set with some extra heuristics
    fn smart_set(&mut self, loc: Location, terrain: Terrain) {
        if let Some(&previous) = self.worldgen.terrain.get(&loc) {
            // Never turn terrain already set as passable as impassable.
            if !previous.blocks_walk() && terrain.blocks_walk() {
                return;
            }

            // Don't clobber doors, period.
            if previous == Terrain::Door {
                return;
            }
        }
        self.set(loc, terrain);
    }

    fn clear_spawns_near_entrance(&mut self, range: u32) {
        for loc in Dijkstra::new(
            vec![self.up_portal.expect("No entrance generated")],
            |loc| {
                self.worldgen
                    .terrain
                    .get(loc)
                    .cloned()
                    .unwrap_or(Terrain::Empty)
                    .is_open()
            },
            10_000,
        ).weights
            .iter()
            .filter_map(|(loc, &w)| if w <= range { Some(loc) } else { None })
        {
            self.spawn_region.remove(loc);
        }
    }

    fn is_next_to_diagonal_sector(loc: Location) -> bool {
        hex_neighbors(loc)
            .map(|x| x.sector().taxicab_distance(loc.sector()))
            .any(|d| d > 1)
    }
}

impl<'a> mapgen::Dungeon for SectorDigger<'a> {
    type Vault = Room;

    fn sample_vault<R: Rng>(&mut self, rng: &mut R) -> Self::Vault { rng.gen() }

    fn dig_chamber<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, area: I) {
        for pos in area {
            let loc = self.loc(pos);
            self.spawn_region.insert(loc);
            self.set(loc, Terrain::Ground);
        }
    }

    fn dig_corridor<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, path: I) {
        for pos in path {
            let loc = self.loc(pos);
            // Do not spawn things in corridors.
            self.spawn_region.remove(&loc);
            self.set(loc, Terrain::Ground);
        }
    }

    fn place_vault(&mut self, vault: &Self::Vault, pos: mapgen::Point2D) {
        use mapgen::VaultCell::*;
        let points: Vec<(mapgen::Point2D, VaultCell)> = vault.get_shape();
        for &(p, c) in &points {
            let loc = self.loc(pos + p.to_vector());

            match c {
                Interior => {
                    self.set(loc, Terrain::Ground);
                    self.spawn_region.insert(loc);
                }
                DiggableWall | UndiggableWall => self.smart_set(loc, Terrain::Wall),
            }
        }
    }

    fn add_door(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Door);
    }

    fn add_up_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.up_portal = Some(loc);
        self.spawn_region.remove(&loc);
        // TODO: Add debug assertion that up-stairs are dug in the upper wall of a room
        self.set(loc, Terrain::Gate);
    }

    fn add_down_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.down_portal = Some(loc);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Gate);

        // Visual hack to make the down-stairs show up better, carve out the rock blob that would
        // be drawn partially in front of it on screen.
        self.set(loc + vec2(1, 1), Terrain::Empty);
        // TODO: Add a debug assertion here to check that the exit location has proper enclosure to
        // begin with, both it and the rock cell in front of it need to be surrounded by undug
        // front tiles.
        self.spawn_region.remove(&(loc + vec2(1, 1))); // Just in case
    }
}

#[derive(Debug)]
pub struct Room {
    size: Size2D,
}

impl mapgen::Vault for Room {
    fn get_shape<T: FromIterator<(mapgen::Point2D, VaultCell)>>(&self) -> T {
        (-1..self.size.height + 1)
            .flat_map(move |y| {
                (-1..self.size.width + 1).map(move |x| {
                    let x_wall = x == -1 || x == self.size.width;
                    let y_wall = y == -1 || y == self.size.height;
                    let p = point2(x, y);

                    if x_wall && y_wall {
                        (p, VaultCell::UndiggableWall)
                    } else if x_wall || y_wall {
                        (p, VaultCell::DiggableWall)
                    } else {
                        (p, VaultCell::Interior)
                    }
                })
            })
            .collect()
    }
}

impl Distribution<Room> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Room {
        Room {
            size: Size2D::new(rng.gen_range(2, 8), rng.gen_range(2, 8)),
        }
    }
}

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

    #[test]
    fn test_determinism() {
        use rand;

        let mut rng = rand::thread_rng();

        let seed: u32 = rng.gen();
        let mut gen = Worldgen::new(seed);

        // Build the value repeatedly using the same seed and see that they are all equal.
        for _ in 1..4 {
            let second = Worldgen::new(seed);

            assert_eq!(gen.seed, second.seed);
            // These can make huge printouts so don't use assert_eq that would try to print them to
            // stdout
            assert!(gen.terrain == second.terrain);
            assert!(gen.portals == second.portals);
            assert_eq!(gen.player_entry, second.player_entry);

            gen = second;
        }
    }
}
<commit_msg>All earth is rock now that overworld is shelved<commit_after>//! Game world generation
//!
//! Unlike the game-agnostic `mapgen` module, this module does have gameplay-specific logic.

use calx::{hex_neighbors, seeded_rng, Dijkstra};
use euclid::{point2, vec2};
use location::{Location, Portal, Sector};
use mapgen::{self, MapGen, Size2D, Vault, VaultCell};
use rand::distributions::{Distribution, Standard};
use rand::{seq, Rng};
use serde;
use spec;
use std::collections::{BTreeSet, HashMap};
use std::iter::FromIterator;
use std::slice;
use terrain::Terrain;
use world::Loadout;

/// Static generated world.
pub struct Worldgen {
    seed: u32,
    terrain: HashMap<Location, Terrain>,
    portals: HashMap<Location, Portal>,
    spawns: Vec<(Location, Loadout)>,
    player_entry: Location,
}

impl Worldgen {
    pub fn new(seed: u32) -> Worldgen {
        let mut ret = Worldgen {
            seed,
            terrain: HashMap::new(),
            portals: HashMap::new(),
            spawns: Vec::new(),
            player_entry: Location::new(0, 0, 0),
        };

        let mut rng = seeded_rng(&seed);

        let mut cave_entrance = Location::new(0, 0, 0);
        for depth in 1..11 {
            let up_stairs;
            let down_stairs;
            let mut spawns = Vec::new();

            {
                let mut digger = SectorDigger::new(&mut ret, Sector::new(0, 0, depth as i16));
                let domain = digger.domain();

                if depth < 5 {
                    mapgen::RoomsAndCorridors.dig(&mut rng, &mut digger, domain);
                } else {
                    mapgen::Caves.dig(&mut rng, &mut digger, domain);
                }

                up_stairs = digger.up_portal.expect("Mapgen didn't create stairs up");
                down_stairs = digger
                    .down_portal
                    .expect("Mapgen didn't create stairs down");

                // Spawn
                digger.clear_spawns_near_entrance(12);

                let mut spawn_locs =
                    seq::sample_iter(&mut rng, digger.spawn_region.iter(), 20).unwrap();
                let n_spawns = spawn_locs.len();

                for &loc in spawn_locs.drain(0..n_spawns / 2) {
                    spawns.push((
                        loc,
                        spec::pick(&mut rng, depth, spec::ITEM_SPECS.as_slice())
                            .expect("No item spawn"),
                    ))
                }

                for &loc in spawn_locs {
                    spawns.push((
                        loc,
                        spec::pick(&mut rng, depth, (&*spec::MOB_SPECS).as_slice())
                            .expect("No mob spawn"),
                    ))
                }
            }

            ret.spawns.extend(spawns.into_iter());
            if depth > 1 {
                ret.portal(cave_entrance, up_stairs + vec2(1, 1));
                ret.portal(up_stairs, cave_entrance - vec2(1, 1));
            } else {
                ret.player_entry = up_stairs;
            }

            cave_entrance = down_stairs;

            // TODO: Generator needs an option to not generate stairs down on bottom level
        }

        ret
    }

    pub fn seed(&self) -> u32 { self.seed }

    pub fn get_terrain(&self, loc: Location) -> Terrain {
        if let Some(&t) = self.terrain.get(&loc) {
            t
        } else {
            self.default_terrain(loc)
        }
    }

    fn default_terrain(&self, _loc: Location) -> Terrain { Terrain::Rock }

    pub fn get_portal(&self, loc: Location) -> Option<Location> {
        self.portals.get(&loc).map(|&p| loc + p)
    }

    pub fn spawns(&self) -> slice::Iter<(Location, Loadout)> { self.spawns.iter() }

    pub fn player_entry(&self) -> Location { self.player_entry }

    /// Punch a (one-way) portal between two points.
    fn portal(&mut self, origin: Location, destination: Location) {
        self.portals
            .insert(origin, Portal::new(origin, destination));
    }
}

impl serde::Serialize for Worldgen {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        self.seed.serialize(s)
    }
}

impl<'a> serde::Deserialize<'a> for Worldgen {
    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
        Ok(Worldgen::new(serde::Deserialize::deserialize(d)?))
    }
}

struct SectorDigger<'a> {
    worldgen: &'a mut Worldgen,
    sector: Sector,
    up_portal: Option<Location>,
    down_portal: Option<Location>,
    spawn_region: BTreeSet<Location>,
}

impl<'a> SectorDigger<'a> {
    fn new(worldgen: &'a mut Worldgen, sector: Sector) -> SectorDigger<'a> {
        SectorDigger {
            worldgen,
            sector,
            up_portal: None,
            down_portal: None,
            spawn_region: BTreeSet::new(),
        }
    }

    fn domain(&self) -> impl Iterator<Item = mapgen::Point2D> {
        let sector_origin = self.sector.origin();
        let mapgen_origin = mapgen::Point2D::zero();

        self.sector.iter().filter_map(move |loc| {
            let pos = mapgen_origin
                + sector_origin
                    .v2_at(loc)
                    .expect("Sector points are not Euclidean");

            // Okay, this part is a bit hairy, hang on.
            // Sectors are arranged in a rectangular grid, so there are some cells where you can
            // step diagonally across two sector boundaries. However, we want to pretend that
            // sectors are only connected in the four cardinal directions, so we omit these cells
            // from the domain to prevent surprise holes between two diagonally adjacent shafts.
            if Self::is_next_to_diagonal_sector(loc) {
                None
            } else {
                Some(pos)
            }
        })
    }

    fn loc(&self, pos: mapgen::Point2D) -> Location { self.sector.origin() + pos.to_vector() }

    fn set(&mut self, loc: Location, terrain: Terrain) {
        if terrain != Terrain::Gate {
            debug_assert_ne!(self.up_portal, Some(loc), "mapgen overwriting exit");
            debug_assert_ne!(self.down_portal, Some(loc), "mapgen overwriting exit");
        }

        // FIXME: Mapgen currently trips this with exit placement
        // debug_assert!(
        //     !Self::is_next_to_diagonal_sector(loc) &&
        //     self.sector.iter().find(|x| x == &loc).is_some(),
        //     "Setting a location outsides sector area");

        self.worldgen.terrain.insert(loc, terrain);
    }

    /// Set with some extra heuristics
    fn smart_set(&mut self, loc: Location, terrain: Terrain) {
        if let Some(&previous) = self.worldgen.terrain.get(&loc) {
            // Never turn terrain already set as passable as impassable.
            if !previous.blocks_walk() && terrain.blocks_walk() {
                return;
            }

            // Don't clobber doors, period.
            if previous == Terrain::Door {
                return;
            }
        }
        self.set(loc, terrain);
    }

    fn clear_spawns_near_entrance(&mut self, range: u32) {
        for loc in Dijkstra::new(
            vec![self.up_portal.expect("No entrance generated")],
            |loc| {
                self.worldgen
                    .terrain
                    .get(loc)
                    .cloned()
                    .unwrap_or(Terrain::Empty)
                    .is_open()
            },
            10_000,
        ).weights
            .iter()
            .filter_map(|(loc, &w)| if w <= range { Some(loc) } else { None })
        {
            self.spawn_region.remove(loc);
        }
    }

    fn is_next_to_diagonal_sector(loc: Location) -> bool {
        hex_neighbors(loc)
            .map(|x| x.sector().taxicab_distance(loc.sector()))
            .any(|d| d > 1)
    }
}

impl<'a> mapgen::Dungeon for SectorDigger<'a> {
    type Vault = Room;

    fn sample_vault<R: Rng>(&mut self, rng: &mut R) -> Self::Vault { rng.gen() }

    fn dig_chamber<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, area: I) {
        for pos in area {
            let loc = self.loc(pos);
            self.spawn_region.insert(loc);
            self.set(loc, Terrain::Ground);
        }
    }

    fn dig_corridor<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, path: I) {
        for pos in path {
            let loc = self.loc(pos);
            // Do not spawn things in corridors.
            self.spawn_region.remove(&loc);
            self.set(loc, Terrain::Ground);
        }
    }

    fn place_vault(&mut self, vault: &Self::Vault, pos: mapgen::Point2D) {
        use mapgen::VaultCell::*;
        let points: Vec<(mapgen::Point2D, VaultCell)> = vault.get_shape();
        for &(p, c) in &points {
            let loc = self.loc(pos + p.to_vector());

            match c {
                Interior => {
                    self.set(loc, Terrain::Ground);
                    self.spawn_region.insert(loc);
                }
                DiggableWall | UndiggableWall => self.smart_set(loc, Terrain::Wall),
            }
        }
    }

    fn add_door(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Door);
    }

    fn add_up_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.up_portal = Some(loc);
        self.spawn_region.remove(&loc);
        // TODO: Add debug assertion that up-stairs are dug in the upper wall of a room
        self.set(loc, Terrain::Gate);
    }

    fn add_down_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.down_portal = Some(loc);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Gate);

        // Visual hack to make the down-stairs show up better, carve out the rock blob that would
        // be drawn partially in front of it on screen.
        self.set(loc + vec2(1, 1), Terrain::Empty);
        // TODO: Add a debug assertion here to check that the exit location has proper enclosure to
        // begin with, both it and the rock cell in front of it need to be surrounded by undug
        // front tiles.
        self.spawn_region.remove(&(loc + vec2(1, 1))); // Just in case
    }
}

#[derive(Debug)]
pub struct Room {
    size: Size2D,
}

impl mapgen::Vault for Room {
    fn get_shape<T: FromIterator<(mapgen::Point2D, VaultCell)>>(&self) -> T {
        (-1..self.size.height + 1)
            .flat_map(move |y| {
                (-1..self.size.width + 1).map(move |x| {
                    let x_wall = x == -1 || x == self.size.width;
                    let y_wall = y == -1 || y == self.size.height;
                    let p = point2(x, y);

                    if x_wall && y_wall {
                        (p, VaultCell::UndiggableWall)
                    } else if x_wall || y_wall {
                        (p, VaultCell::DiggableWall)
                    } else {
                        (p, VaultCell::Interior)
                    }
                })
            })
            .collect()
    }
}

impl Distribution<Room> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Room {
        Room {
            size: Size2D::new(rng.gen_range(2, 8), rng.gen_range(2, 8)),
        }
    }
}

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

    #[test]
    fn test_determinism() {
        use rand;

        let mut rng = rand::thread_rng();

        let seed: u32 = rng.gen();
        let mut gen = Worldgen::new(seed);

        // Build the value repeatedly using the same seed and see that they are all equal.
        for _ in 1..4 {
            let second = Worldgen::new(seed);

            assert_eq!(gen.seed, second.seed);
            // These can make huge printouts so don't use assert_eq that would try to print them to
            // stdout
            assert!(gen.terrain == second.terrain);
            assert!(gen.portals == second.portals);
            assert_eq!(gen.player_entry, second.player_entry);

            gen = second;
        }
    }
}
<|endoftext|>