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

//! A library for passing arbitrary file descriptors when spawning child processes.
//!
//! # Example
//!
//! ```rust
//! use command_fds::{CommandFdExt, FdMapping};
//! use std::fs::File;
//! use std::os::unix::io::AsRawFd;
//! use std::process::Command;
//!
//! // Open a file.
//! let file = File::open("Cargo.toml").unwrap();
//!
//! // Prepare to run `ls -l /proc/self/fd` with some FDs mapped.
//! let mut command = Command::new("ls");
//! command.arg("-l").arg("/proc/self/fd");
//! command
//!     .fd_mappings(vec![
//!         // Map `file` as FD 3 in the child process.
//!         FdMapping {
//!             parent_fd: file.as_raw_fd(),
//!             child_fd: 3,
//!         },
//!         // Map this process's stdin as FD 5 in the child process.
//!         FdMapping {
//!             parent_fd: 0,
//!             child_fd: 5,
//!         },
//!     ])
//!     .unwrap();
//!
//! // Spawn the child process.
//! let mut child = command.spawn().unwrap();
//! child.wait().unwrap();
//! ```

use nix::fcntl::{fcntl, FcntlArg, FdFlag};
use nix::unistd::dup2;
use std::cmp::max;
use std::io::{self, ErrorKind};
use std::os::unix::io::RawFd;
use std::os::unix::process::CommandExt;
use std::process::Command;
use thiserror::Error;

/// A mapping from a file descriptor in the parent to a file descriptor in the child, to be applied
/// when spawning a child process.
///
/// The parent_fd must be kept open until after the child is spawned.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FdMapping {
    pub parent_fd: RawFd,
    pub child_fd: RawFd,
}

/// Error setting up FD mappings, because there were two or more mappings for the same child FD.
#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)]
#[error("Two or more mappings for the same child FD")]
pub struct FdMappingCollision;

/// Extension to add file descriptor mappings to a [`Command`].
pub trait CommandFdExt {
    /// Adds the given set of file descriptor to the command.
    ///
    /// Calling this more than once on the same command may result in unexpected behaviour.
    fn fd_mappings(&mut self, mappings: Vec<FdMapping>) -> Result<&mut Self, FdMappingCollision>;
}

impl CommandFdExt for Command {
    fn fd_mappings(&mut self, mappings: Vec<FdMapping>) -> Result<&mut Self, FdMappingCollision> {
        // Validate that there are no conflicting mappings to the same child FD.
        let mut child_fds: Vec<RawFd> = mappings.iter().map(|mapping| mapping.child_fd).collect();
        child_fds.sort_unstable();
        child_fds.dedup();
        if child_fds.len() != mappings.len() {
            return Err(FdMappingCollision);
        }

        // Register the callback to apply the mappings after forking but before execing.
        unsafe {
            self.pre_exec(move || map_fds(&mappings));
        }

        Ok(self)
    }
}

fn map_fds(mappings: &[FdMapping]) -> io::Result<()> {
    if mappings.is_empty() {
        // No need to do anything, and finding first_unused_fd would fail.
        return Ok(());
    }

    // Find the first FD which is higher than any parent or child FD in the mapping, so we can
    // safely use it and higher FDs as temporary FDs. There may be other files open with these FDs,
    // so we still need to ensure we don't conflict with them.
    let first_safe_fd = mappings
        .iter()
        .map(|mapping| max(mapping.parent_fd, mapping.child_fd))
        .max()
        .unwrap()
        + 1;

    // If any parent FDs conflict with child FDs, then first duplicate them to a temporary FD which
    // is clear of either range. Mappings to the same FD are fine though, we can handle them by just
    // removing the FD_CLOEXEC flag.
    let child_fds: Vec<RawFd> = mappings.iter().map(|mapping| mapping.child_fd).collect();
    let mappings = mappings
        .iter()
        .map(|mapping| {
            Ok(
                if child_fds.contains(&mapping.parent_fd) && mapping.parent_fd != mapping.child_fd {
                    let temporary_fd =
                        fcntl(mapping.parent_fd, FcntlArg::F_DUPFD_CLOEXEC(first_safe_fd))?;
                    FdMapping {
                        parent_fd: temporary_fd,
                        child_fd: mapping.child_fd,
                    }
                } else {
                    mapping.to_owned()
                },
            )
        })
        .collect::<nix::Result<Vec<_>>>()
        .map_err(nix_to_io_error)?;

    // Now we can actually duplicate FDs to the desired child FDs.
    for mapping in mappings {
        if mapping.child_fd == mapping.parent_fd {
            // Remove the FD_CLOEXEC flag, so the FD will be kept open when exec is called for the
            // child.
            fcntl(mapping.parent_fd, FcntlArg::F_SETFD(FdFlag::empty()))
                .map_err(nix_to_io_error)?;
        } else {
            // This closes child_fd if it is already open as something else, and clears the
            // FD_CLOEXEC flag on child_fd.
            dup2(mapping.parent_fd, mapping.child_fd).map_err(nix_to_io_error)?;
        }
    }

    Ok(())
}

/// Convert a [`nix::Error`] to a [`std::io::Error`].
fn nix_to_io_error(error: nix::Error) -> io::Error {
    if let nix::Error::Sys(errno) = error {
        io::Error::from_raw_os_error(errno as i32)
    } else {
        io::Error::new(ErrorKind::Other, error)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nix::unistd::close;
    use std::collections::HashSet;
    use std::fs::{read_dir, File};
    use std::os::unix::io::AsRawFd;
    use std::process::Output;
    use std::str;
    use std::sync::Once;

    static SETUP: Once = Once::new();

    #[test]
    fn conflicting_mappings() {
        setup();

        let mut command = Command::new("ls");

        // The same mapping can't be included twice.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
            ])
            .is_err());

        // Mapping two different FDs to the same FD isn't allowed either.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 6,
                },
            ])
            .is_err());
    }

    #[test]
    fn no_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        assert!(command.fd_mappings(vec![]).is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3], 0);
    }

    #[test]
    fn one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file an otherwise unused FD.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 5,
            },])
            .is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3, 5], 0);
    }

    #[test]
    fn swap_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        let fd2 = file2.as_raw_fd();
        // Map files to each other's FDs, to ensure that the temporary FD logic works.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    parent_fd: fd1,
                    child_fd: fd2,
                },
                FdMapping {
                    parent_fd: fd2,
                    child_fd: fd1,
                },
            ])
            .is_ok(),);

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1, fd2], 1);
    }

    #[test]
    fn one_to_one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        // Map files to each other's FDs, to ensure that the temporary FD logic works.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: fd1,
                child_fd: fd1,
            }])
            .is_ok());

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1], 1);

        // Keep file2 open until the end, to ensure that it's not passed to the child.
        drop(file2);
    }

    #[test]
    fn map_stdin() {
        setup();

        let mut command = Command::new("cat");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file to stdin.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 0,
            },])
            .is_ok());

        let output = command.output().unwrap();
        assert!(output.status.success());
        assert_eq!(output.stdout, b"test 1");
    }

    /// Parse the output of ls into a set of filenames
    fn parse_ls_output(output: &[u8]) -> HashSet<String> {
        str::from_utf8(output)
            .unwrap()
            .split_terminator("\n")
            .map(str::to_owned)
            .collect()
    }

    /// Check that the output of `ls /proc/self/fd` contains the expected set of FDs, plus exactly
    /// `extra` extra FDs.
    fn expect_fds(output: &Output, expected_fds: &[RawFd], extra: usize) {
        assert!(output.status.success());
        let expected_fds: HashSet<String> = expected_fds.iter().map(RawFd::to_string).collect();
        let fds = parse_ls_output(&output.stdout);
        if extra == 0 {
            assert_eq!(fds, expected_fds);
        } else {
            assert!(expected_fds.is_subset(&fds));
            assert_eq!(fds.len(), expected_fds.len() + extra);
        }
    }

    fn setup() {
        SETUP.call_once(close_excess_fds);
    }

    /// Close all file descriptors apart from stdin, stdout and stderr.
    ///
    /// This is necessary because GitHub Actions opens a bunch of others for some reason.
    fn close_excess_fds() {
        let dir = read_dir("/proc/self/fd").unwrap();
        for entry in dir {
            let entry = entry.unwrap();
            let fd: RawFd = entry.file_name().to_str().unwrap().parse().unwrap();
            if fd > 3 {
                close(fd).unwrap();
            }
        }
    }
}
<commit_msg>Fix typo in comment.<commit_after>// Copyright 2021, The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! A library for passing arbitrary file descriptors when spawning child processes.
//!
//! # Example
//!
//! ```rust
//! use command_fds::{CommandFdExt, FdMapping};
//! use std::fs::File;
//! use std::os::unix::io::AsRawFd;
//! use std::process::Command;
//!
//! // Open a file.
//! let file = File::open("Cargo.toml").unwrap();
//!
//! // Prepare to run `ls -l /proc/self/fd` with some FDs mapped.
//! let mut command = Command::new("ls");
//! command.arg("-l").arg("/proc/self/fd");
//! command
//!     .fd_mappings(vec![
//!         // Map `file` as FD 3 in the child process.
//!         FdMapping {
//!             parent_fd: file.as_raw_fd(),
//!             child_fd: 3,
//!         },
//!         // Map this process's stdin as FD 5 in the child process.
//!         FdMapping {
//!             parent_fd: 0,
//!             child_fd: 5,
//!         },
//!     ])
//!     .unwrap();
//!
//! // Spawn the child process.
//! let mut child = command.spawn().unwrap();
//! child.wait().unwrap();
//! ```

use nix::fcntl::{fcntl, FcntlArg, FdFlag};
use nix::unistd::dup2;
use std::cmp::max;
use std::io::{self, ErrorKind};
use std::os::unix::io::RawFd;
use std::os::unix::process::CommandExt;
use std::process::Command;
use thiserror::Error;

/// A mapping from a file descriptor in the parent to a file descriptor in the child, to be applied
/// when spawning a child process.
///
/// The parent_fd must be kept open until after the child is spawned.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FdMapping {
    pub parent_fd: RawFd,
    pub child_fd: RawFd,
}

/// Error setting up FD mappings, because there were two or more mappings for the same child FD.
#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)]
#[error("Two or more mappings for the same child FD")]
pub struct FdMappingCollision;

/// Extension to add file descriptor mappings to a [`Command`].
pub trait CommandFdExt {
    /// Adds the given set of file descriptors to the command.
    ///
    /// Calling this more than once on the same command may result in unexpected behaviour.
    fn fd_mappings(&mut self, mappings: Vec<FdMapping>) -> Result<&mut Self, FdMappingCollision>;
}

impl CommandFdExt for Command {
    fn fd_mappings(&mut self, mappings: Vec<FdMapping>) -> Result<&mut Self, FdMappingCollision> {
        // Validate that there are no conflicting mappings to the same child FD.
        let mut child_fds: Vec<RawFd> = mappings.iter().map(|mapping| mapping.child_fd).collect();
        child_fds.sort_unstable();
        child_fds.dedup();
        if child_fds.len() != mappings.len() {
            return Err(FdMappingCollision);
        }

        // Register the callback to apply the mappings after forking but before execing.
        unsafe {
            self.pre_exec(move || map_fds(&mappings));
        }

        Ok(self)
    }
}

fn map_fds(mappings: &[FdMapping]) -> io::Result<()> {
    if mappings.is_empty() {
        // No need to do anything, and finding first_unused_fd would fail.
        return Ok(());
    }

    // Find the first FD which is higher than any parent or child FD in the mapping, so we can
    // safely use it and higher FDs as temporary FDs. There may be other files open with these FDs,
    // so we still need to ensure we don't conflict with them.
    let first_safe_fd = mappings
        .iter()
        .map(|mapping| max(mapping.parent_fd, mapping.child_fd))
        .max()
        .unwrap()
        + 1;

    // If any parent FDs conflict with child FDs, then first duplicate them to a temporary FD which
    // is clear of either range. Mappings to the same FD are fine though, we can handle them by just
    // removing the FD_CLOEXEC flag.
    let child_fds: Vec<RawFd> = mappings.iter().map(|mapping| mapping.child_fd).collect();
    let mappings = mappings
        .iter()
        .map(|mapping| {
            Ok(
                if child_fds.contains(&mapping.parent_fd) && mapping.parent_fd != mapping.child_fd {
                    let temporary_fd =
                        fcntl(mapping.parent_fd, FcntlArg::F_DUPFD_CLOEXEC(first_safe_fd))?;
                    FdMapping {
                        parent_fd: temporary_fd,
                        child_fd: mapping.child_fd,
                    }
                } else {
                    mapping.to_owned()
                },
            )
        })
        .collect::<nix::Result<Vec<_>>>()
        .map_err(nix_to_io_error)?;

    // Now we can actually duplicate FDs to the desired child FDs.
    for mapping in mappings {
        if mapping.child_fd == mapping.parent_fd {
            // Remove the FD_CLOEXEC flag, so the FD will be kept open when exec is called for the
            // child.
            fcntl(mapping.parent_fd, FcntlArg::F_SETFD(FdFlag::empty()))
                .map_err(nix_to_io_error)?;
        } else {
            // This closes child_fd if it is already open as something else, and clears the
            // FD_CLOEXEC flag on child_fd.
            dup2(mapping.parent_fd, mapping.child_fd).map_err(nix_to_io_error)?;
        }
    }

    Ok(())
}

/// Convert a [`nix::Error`] to a [`std::io::Error`].
fn nix_to_io_error(error: nix::Error) -> io::Error {
    if let nix::Error::Sys(errno) = error {
        io::Error::from_raw_os_error(errno as i32)
    } else {
        io::Error::new(ErrorKind::Other, error)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nix::unistd::close;
    use std::collections::HashSet;
    use std::fs::{read_dir, File};
    use std::os::unix::io::AsRawFd;
    use std::process::Output;
    use std::str;
    use std::sync::Once;

    static SETUP: Once = Once::new();

    #[test]
    fn conflicting_mappings() {
        setup();

        let mut command = Command::new("ls");

        // The same mapping can't be included twice.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
            ])
            .is_err());

        // Mapping two different FDs to the same FD isn't allowed either.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 6,
                },
            ])
            .is_err());
    }

    #[test]
    fn no_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        assert!(command.fd_mappings(vec![]).is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3], 0);
    }

    #[test]
    fn one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file an otherwise unused FD.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 5,
            },])
            .is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3, 5], 0);
    }

    #[test]
    fn swap_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        let fd2 = file2.as_raw_fd();
        // Map files to each other's FDs, to ensure that the temporary FD logic works.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    parent_fd: fd1,
                    child_fd: fd2,
                },
                FdMapping {
                    parent_fd: fd2,
                    child_fd: fd1,
                },
            ])
            .is_ok(),);

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1, fd2], 1);
    }

    #[test]
    fn one_to_one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        // Map files to each other's FDs, to ensure that the temporary FD logic works.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: fd1,
                child_fd: fd1,
            }])
            .is_ok());

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1], 1);

        // Keep file2 open until the end, to ensure that it's not passed to the child.
        drop(file2);
    }

    #[test]
    fn map_stdin() {
        setup();

        let mut command = Command::new("cat");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file to stdin.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 0,
            },])
            .is_ok());

        let output = command.output().unwrap();
        assert!(output.status.success());
        assert_eq!(output.stdout, b"test 1");
    }

    /// Parse the output of ls into a set of filenames
    fn parse_ls_output(output: &[u8]) -> HashSet<String> {
        str::from_utf8(output)
            .unwrap()
            .split_terminator("\n")
            .map(str::to_owned)
            .collect()
    }

    /// Check that the output of `ls /proc/self/fd` contains the expected set of FDs, plus exactly
    /// `extra` extra FDs.
    fn expect_fds(output: &Output, expected_fds: &[RawFd], extra: usize) {
        assert!(output.status.success());
        let expected_fds: HashSet<String> = expected_fds.iter().map(RawFd::to_string).collect();
        let fds = parse_ls_output(&output.stdout);
        if extra == 0 {
            assert_eq!(fds, expected_fds);
        } else {
            assert!(expected_fds.is_subset(&fds));
            assert_eq!(fds.len(), expected_fds.len() + extra);
        }
    }

    fn setup() {
        SETUP.call_once(close_excess_fds);
    }

    /// Close all file descriptors apart from stdin, stdout and stderr.
    ///
    /// This is necessary because GitHub Actions opens a bunch of others for some reason.
    fn close_excess_fds() {
        let dir = read_dir("/proc/self/fd").unwrap();
        for entry in dir {
            let entry = entry.unwrap();
            let fd: RawFd = entry.file_name().to_str().unwrap().parse().unwrap();
            if fd > 3 {
                close(fd).unwrap();
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Runtime support for the async/await syntax for futures.
//!
//! This crate serves as a masquerade over the `futures` crate itself,
//! reexporting all of its contents. It's intended that you'll do:
//!
//! ```
//! extern crate futures_await as futures;
//! ```
//!
//! This crate adds a `prelude` module which contains various traits as well as
//! the `async` and `await` macros you'll likely want to use.
//!
//! See the crates's README for more information about usage.

#![feature(generator_trait)]
#![feature(use_extern_macros)]

extern crate futures_async_macro; // the compiler lies that this has no effect
extern crate futures_await_macro;
extern crate futures;

pub use futures::*;

pub mod prelude {
    pub use {Future, Stream, Sink, Poll, Async, AsyncSink, StartSend};
    pub use {IntoFuture};
    pub use futures_async_macro::*; // TODO: can we selectively import just `async`?
    pub use futures_await_macro::await;
}

/// A hidden module that's the "runtime support" for the async/await syntax.
///
/// The `async` attribute and the `await` macro both assume that they can find
/// this module and use its contents. All of their dependencies are defined or
/// reexported here in one way shape or form.
///
/// This module has absolutely not stability at all. All contents may change at
/// any time without notice. Do not use this module in your code if you wish
/// your code to be stable.
#[doc(hidden)]
pub mod __rt {
    pub use std::result::Result::{Ok, Err};
    pub use std::option::Option::{Some, None};
    pub use std::boxed::Box;
    pub use std::ops::Generator;

    use futures::Poll;
    use futures::{Future, Async};
    use std::ops::State;

    /// Convenience trait to project from `Result` and get the item/error types
    ///
    /// This is how we work with type aliases like `io::Result` without knowing
    /// whether you're using a type alias.
    pub trait FutureType {
        type Item;
        type Error;

        fn into_result(self) -> Result<Self::Item, Self::Error>;
    }

    impl<T, E> FutureType for Result<T, E> {
        type Item = T;
        type Error = E;

        fn into_result(self) -> Result<Self::Item, Self::Error> {
            self
        }
    }

    /// Random hack for this causing problems in the compiler's typechecking
    /// pass. Ideally this trait and impl would not be needed.
    ///
    /// ```ignore
    /// -> impl Future<Item = <T as FutureType>::Item,
    ///                Error = <T as FutureType>::Error>
    /// ```
    pub trait MyFuture<T: FutureType>: Future<Item=T::Item, Error=T::Error> {}
    impl<F: Future + ?Sized> MyFuture<Result<F::Item, F::Error>> for F {}

    /// Small shim to translate from a generator to a future.
    ///
    /// This is the translation layer from the generator/coroutine protocol to
    /// the futures protocol.
    pub struct GenFuture<T>(T);

    pub fn gen<T>(t: T) -> GenFuture<T>
        where T: Generator<Yield = ()>,
              T::Return: FutureType,
    {
        GenFuture(t)
    }

    impl<T> Future for GenFuture<T>
        where T: Generator<Yield = ()>,
              T::Return: FutureType,
    {
        type Item = <T::Return as FutureType>::Item;
        type Error = <T::Return as FutureType>::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            match self.0.resume(()) {
                State::Yielded(()) => Ok(Async::NotReady),
                State::Complete(e) => e.into_result().map(Async::Ready),
            }
        }
    }
}
<commit_msg>Minor tweaks:<commit_after>//! Runtime support for the async/await syntax for futures.
//!
//! This crate serves as a masquerade over the `futures` crate itself,
//! reexporting all of its contents. It's intended that you'll do:
//!
//! ```
//! extern crate futures_await as futures;
//! ```
//!
//! This crate adds a `prelude` module which contains various traits as well as
//! the `async` and `await` macros you'll likely want to use.
//!
//! See the crates's README for more information about usage.

#![feature(conservative_impl_trait)]
#![feature(generator_trait)]
#![feature(try_trait)]
#![feature(use_extern_macros)]

extern crate futures_async_macro; // the compiler lies that this has no effect
extern crate futures_await_macro;
extern crate futures;

pub use futures::*;

pub mod prelude {
    pub use {Future, Stream, Sink, Poll, Async, AsyncSink, StartSend};
    pub use {IntoFuture};
    pub use futures_async_macro::*; // TODO: can we selectively import just `async`?
    pub use futures_await_macro::await;
}

/// A hidden module that's the "runtime support" for the async/await syntax.
///
/// The `async` attribute and the `await` macro both assume that they can find
/// this module and use its contents. All of their dependencies are defined or
/// reexported here in one way shape or form.
///
/// This module has absolutely not stability at all. All contents may change at
/// any time without notice. Do not use this module in your code if you wish
/// your code to be stable.
#[doc(hidden)]
pub mod __rt {
    pub use std::boxed::Box;
    pub use std::ops::Try;
    pub use std::option::Option::{Some, None};
    pub use std::result::Result::{Ok, Err};
    pub use std::ops::Generator;

    use futures::Poll;
    use futures::{Future, Async};
    use std::ops::State;

    /// Random hack for this causing problems in the compiler's typechecking
    /// pass. Ideally this trait and impl would not be needed.
    ///
    /// ```ignore
    /// -> impl Future<Item = <T as Try>::Ok,
    ///                Error = <T as Try>::Error>
    /// ```
    pub trait MyFuture<T: Try>: Future<Item=T::Ok, Error=T::Error> {}

    impl<F, T> MyFuture<T> for F
        where F: Future<Item = T::Ok, Error = T::Error> + ?Sized,
              T: Try,
    {}

    /// Small shim to translate from a generator to a future.
    ///
    /// This is the translation layer from the generator/coroutine protocol to
    /// the futures protocol.
    struct GenFuture<T>(T);

    pub fn gen<T>(t: T) -> impl Future<Item = <T::Return as Try>::Ok,
                                       Error = <T::Return as Try>::Error>
        where T: Generator<Yield = ()>,
              T::Return: Try,
    {
        GenFuture(t)
    }

    impl<T> Future for GenFuture<T>
        where T: Generator<Yield = ()>,
              T::Return: Try,
    {
        type Item = <T::Return as Try>::Ok;
        type Error = <T::Return as Try>::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            match self.0.resume() {
                State::Yielded(()) => Ok(Async::NotReady),
                State::Complete(e) => e.into_result().map(Async::Ready),
            }
        }
    }
}
<|endoftext|>


<commit_before>#![deny(missing_docs)]
#![unstable]

//! A library for setting current values for stack scope,
//! such as application structure.

use std::rc::Rc;
use std::cell::RefCell;

impl<'a, T, U: GetFrom<T>> GetFrom<&'a RefCell<T>> for U {
    #[inline(always)]
    fn get_from(obj: & &'a RefCell<T>) -> U {
        GetFrom::get_from(obj.borrow().deref())
    }
}

impl<T, U: GetFrom<T>> GetFrom<Rc<RefCell<T>>> for U {
    #[inline(always)]
    fn get_from(obj: &Rc<RefCell<T>>) -> U {
        GetFrom::get_from(obj.borrow().deref())
    }
}

impl<'a, F, T: SetAt<F>> SetAt<&'a RefCell<F>> for T {
    #[inline(always)]
    fn set_at(self, obj: &mut &'a RefCell<F>) {
        self.set_at(obj.borrow_mut().deref_mut())
    }
}

impl<F, T: SetAt<F>> SetAt<Rc<RefCell<F>>> for T {
    #[inline(always)]
    fn set_at(self, obj: &mut Rc<RefCell<F>>) {
        self.set_at(obj.borrow_mut().deref_mut())
    }
}

impl<'a, F, T: ActOn<F, U>, U> ActOn<&'a RefCell<F>, U> for T {
    #[inline(always)]
    fn act_on(self, obj: &mut &'a RefCell<F>) -> U {
        self.act_on(obj.borrow_mut().deref_mut())
    }
}

impl<F, T: ActOn<F, U>, U> ActOn<Rc<RefCell<F>>, U> for T {
    #[inline(always)]
    fn act_on(self, obj: &mut Rc<RefCell<F>>) -> U {
        self.act_on(obj.borrow_mut().deref_mut())
    }
}

/// Something that can be set at an object.
#[unstable]
pub trait SetAt<F> {
    /// Modify `F` with self.
    fn set_at(self, &mut F);
}

/// Automatically implemented through the `SetAt` trait.
#[unstable]
pub trait Set<T> {
    /// Set value.
    fn set(mut self, val: T) -> Self;

    /// Set value through mutable reference.
    fn set_mut(&mut self, val: T) -> &mut Self;
}

impl<T, U: SetAt<T>> Set<U> for T {
    #[inline(always)]
    fn set(mut self, val: U) -> T {
        val.set_at(&mut self);
        self
    }

    #[inline(always)]
    fn set_mut(&mut self, val: U) -> &mut T {
        val.set_at(self);
        self
    }
}

/// Something that can be retrieved from another object.
#[unstable]
pub trait GetFrom<T> {
    /// Gets value from object.
    fn get_from(obj: &T) -> Self;
}

/// Automatically implemented through the `GetFrom` trait.
#[unstable]
pub trait Get<T> {
    /// Returns new value.
    fn get(&self) -> T;
}

impl<T, U: GetFrom<T>> Get<U> for T {
    #[inline(always)]
    fn get(&self) -> U {
        GetFrom::get_from(self)
    }
}

/// Does something to an object.
#[unstable]
pub trait ActOn<T, U> {
    /// Does something to an object.
    fn act_on(self, &mut T) -> U;
}

/// Automatically implemented through the `ActOn` trait.
#[unstable]
pub trait Action<T, U> {
    /// Does something.
    fn action(&mut self, val: T) -> U;
}

impl<T, U: ActOn<T, V>, V> Action<U, V> for T {
    #[inline(always)]
    fn action(&mut self, val: U) -> V {
        val.act_on(self)
    }
}

<commit_msg>Upgrade to latest Rust<commit_after>#![deny(missing_docs)]
#![unstable]

//! A library for setting current values for stack scope,
//! such as application structure.

use std::rc::Rc;
use std::cell::RefCell;
use std::ops::{ Deref, DerefMut };

impl<'a, T, U: GetFrom<T>> GetFrom<&'a RefCell<T>> for U {
    #[inline(always)]
    fn get_from(obj: & &'a RefCell<T>) -> U {
        GetFrom::get_from(obj.borrow().deref())
    }
}

impl<T, U: GetFrom<T>> GetFrom<Rc<RefCell<T>>> for U {
    #[inline(always)]
    fn get_from(obj: &Rc<RefCell<T>>) -> U {
        GetFrom::get_from(obj.borrow().deref())
    }
}

impl<'a, F, T: SetAt<F>> SetAt<&'a RefCell<F>> for T {
    #[inline(always)]
    fn set_at(self, obj: &mut &'a RefCell<F>) {
        self.set_at(obj.borrow_mut().deref_mut())
    }
}

impl<F, T: SetAt<F>> SetAt<Rc<RefCell<F>>> for T {
    #[inline(always)]
    fn set_at(self, obj: &mut Rc<RefCell<F>>) {
        self.set_at(obj.borrow_mut().deref_mut())
    }
}

impl<'a, F, T: ActOn<F, U>, U> ActOn<&'a RefCell<F>, U> for T {
    #[inline(always)]
    fn act_on(self, obj: &mut &'a RefCell<F>) -> U {
        self.act_on(obj.borrow_mut().deref_mut())
    }
}

impl<F, T: ActOn<F, U>, U> ActOn<Rc<RefCell<F>>, U> for T {
    #[inline(always)]
    fn act_on(self, obj: &mut Rc<RefCell<F>>) -> U {
        self.act_on(obj.borrow_mut().deref_mut())
    }
}

/// Something that can be set at an object.
#[unstable]
pub trait SetAt<F> {
    /// Modify `F` with self.
    fn set_at(self, &mut F);
}

/// Automatically implemented through the `SetAt` trait.
#[unstable]
pub trait Set<T> {
    /// Set value.
    fn set(mut self, val: T) -> Self;

    /// Set value through mutable reference.
    fn set_mut(&mut self, val: T) -> &mut Self;
}

impl<T, U: SetAt<T>> Set<U> for T {
    #[inline(always)]
    fn set(mut self, val: U) -> T {
        val.set_at(&mut self);
        self
    }

    #[inline(always)]
    fn set_mut(&mut self, val: U) -> &mut T {
        val.set_at(self);
        self
    }
}

/// Something that can be retrieved from another object.
#[unstable]
pub trait GetFrom<T> {
    /// Gets value from object.
    fn get_from(obj: &T) -> Self;
}

/// Automatically implemented through the `GetFrom` trait.
#[unstable]
pub trait Get<T> {
    /// Returns new value.
    fn get(&self) -> T;
}

impl<T, U: GetFrom<T>> Get<U> for T {
    #[inline(always)]
    fn get(&self) -> U {
        GetFrom::get_from(self)
    }
}

/// Does something to an object.
#[unstable]
pub trait ActOn<T, U> {
    /// Does something to an object.
    fn act_on(self, &mut T) -> U;
}

/// Automatically implemented through the `ActOn` trait.
#[unstable]
pub trait Action<T, U> {
    /// Does something.
    fn action(&mut self, val: T) -> U;
}

impl<T, U: ActOn<T, V>, V> Action<U, V> for T {
    #[inline(always)]
    fn action(&mut self, val: U) -> V {
        val.act_on(self)
    }
}

<|endoftext|>


<commit_before>use Prefab;
use calx::{self, hex_neighbors, Dijkstra};
use euclid::{point2, vec2};
use form::{self, Form};
use image::{self, GenericImage, SubImage};
use location::{Location, Portal, Sector};
use mapgen::{self, MapGen, Size2D, VaultCell};
use rand::{self, Rand, Rng, SeedableRng};
use serde;
use std::collections::{BTreeSet, HashMap};
use std::io::Cursor;
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: seed,
            terrain: HashMap::new(),
            portals: HashMap::new(),
            spawns: Vec::new(),
            player_entry: Location::new(25, 0, 0),
        };

        ret.load_map_bitmap(
            Location::new(0, 0, 0),
            include_bytes!("../assets/overland.png"),
        );

        let mut rng: ::Rng = SeedableRng::from_seed([seed, seed, seed, seed]);

        let mut cave_entrance = Location::new(40, 12, 0);
        ret.cave_entrance(cave_entrance);
        ret.terrain.insert(cave_entrance, Terrain::Gate);
        ret.terrain
            .insert(cave_entrance + vec2(1, 1), Terrain::Empty);

        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();
                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 = rand::sample(&mut rng, digger.spawn_region.iter(), 20);
                let n_spawns = spawn_locs.len();

                let items = Form::filter(|f| f.is_item() && f.at_depth(depth));
                for &loc in spawn_locs.drain(0..n_spawns / 2) {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &items)
                            .expect("No item spawn")
                            .loadout
                            .clone(),
                    ))
                }

                let mobs = Form::filter(|f| f.is_mob() && f.at_depth(depth));
                for &loc in spawn_locs {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &mobs)
                            .expect("No mob spawn")
                            .loadout
                            .clone(),
                    ))
                }
            }

            ret.spawns.extend(spawns.into_iter());
            ret.portal(cave_entrance, up_stairs + vec2(1, 1));
            ret.portal(up_stairs, cave_entrance - vec2(1, 1));
            cave_entrance = down_stairs;

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

        ret
    }

    fn load_prefab(&mut self, origin: Location, prefab: &Prefab) {
        for (&p, &(ref terrain, ref entities)) in prefab.iter() {
            let loc = origin + p;

            self.terrain.insert(loc, *terrain);

            for spawn in entities.iter() {
                if spawn == "player" {
                    self.player_entry = loc;
                } else {
                    let form = Form::named(spawn)
                        .expect(&format!("Bad prefab: Form '{}' not found!", spawn));
                    self.spawns.push((loc, form.loadout.clone()));
                }
            }
        }
    }

    fn load_map_bitmap(&mut self, origin: Location, data: &[u8]) {
        let mut image = image::load(Cursor::new(data), image::ImageFormat::PNG).unwrap();
        // Skip the bottom horizontal line, it's used to store metadata pixels.
        let (w, h) = (image.width(), image.height());
        let input_map = SubImage::new(&mut image, 0, 0, w, h - 1);

        let prefab = calx::Prefab::parse(input_map).expect("Invalid overworld image map");

        self.terrain.extend(
            prefab
                .into_iter()
                .filter_map(|(p, c)| Terrain::from_color(c).map(|t| (origin + p, t))),
        );
    }

    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 }

    /// Make a cave entrance going down.
    fn cave_entrance(&mut self, loc: Location) {
        const DOWNBOUND_ENCLOSURE: [(i32, i32); 5] = [(1, 0), (0, 1), (2, 1), (1, 2), (2, 2)];

        for &v in &DOWNBOUND_ENCLOSURE {
            let loc = loc + vec2(v.0, v.1);
            self.terrain.insert(loc, Terrain::Rock);
        }

        self.terrain.insert(loc, Terrain::Ground);
        self.terrain.insert(loc + vec2(1, 1), Terrain::Ground);
    }

    /// 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(),
        }
    }

    // TODO return impl
    fn domain(&self) -> Vec<mapgen::Point2D> {
        let mut ret = Vec::new();

        let sector_origin = self.sector.origin();
        let mapgen_origin = mapgen::Point2D::zero();

        for loc in self.sector.iter() {
            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) {
                continue;
            }

            ret.push(pos);
        }

        ret
    }

    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);
    }

    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) {
        // TODO
        unimplemented!();
    }

    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
    }
}

pub struct Room {
    size: Size2D,
}

impl mapgen::Vault for Room {
    fn get_shape<T: FromIterator<(mapgen::Point2D, VaultCell)>>(&self) -> T {
        (-1..self.size.height)
            .flat_map(move |y| {
                (-1..self.size.width).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 Rand for Room {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        Room {
            size: Size2D::new(rng.gen_range(3, 10), rng.gen_range(3, 10)),
        }
    }
}

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

    #[test]
    fn test_determinism() {
        use rand::{self, Rand};

        let mut rng = rand::thread_rng();

        let seed = u32::rand(&mut rng);
        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>Fix room code, generate smaller rooms<commit_after>use Prefab;
use calx::{self, hex_neighbors, Dijkstra};
use euclid::{point2, vec2};
use form::{self, Form};
use image::{self, GenericImage, SubImage};
use location::{Location, Portal, Sector};
use mapgen::{self, MapGen, Size2D, VaultCell};
use rand::{self, Rand, Rng, SeedableRng};
use serde;
use std::collections::{BTreeSet, HashMap};
use std::io::Cursor;
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: seed,
            terrain: HashMap::new(),
            portals: HashMap::new(),
            spawns: Vec::new(),
            player_entry: Location::new(25, 0, 0),
        };

        ret.load_map_bitmap(
            Location::new(0, 0, 0),
            include_bytes!("../assets/overland.png"),
        );

        let mut rng: ::Rng = SeedableRng::from_seed([seed, seed, seed, seed]);

        let mut cave_entrance = Location::new(40, 12, 0);
        ret.cave_entrance(cave_entrance);
        ret.terrain.insert(cave_entrance, Terrain::Gate);
        ret.terrain
            .insert(cave_entrance + vec2(1, 1), Terrain::Empty);

        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();
                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 = rand::sample(&mut rng, digger.spawn_region.iter(), 20);
                let n_spawns = spawn_locs.len();

                let items = Form::filter(|f| f.is_item() && f.at_depth(depth));
                for &loc in spawn_locs.drain(0..n_spawns / 2) {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &items)
                            .expect("No item spawn")
                            .loadout
                            .clone(),
                    ))
                }

                let mobs = Form::filter(|f| f.is_mob() && f.at_depth(depth));
                for &loc in spawn_locs {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &mobs)
                            .expect("No mob spawn")
                            .loadout
                            .clone(),
                    ))
                }
            }

            ret.spawns.extend(spawns.into_iter());
            ret.portal(cave_entrance, up_stairs + vec2(1, 1));
            ret.portal(up_stairs, cave_entrance - vec2(1, 1));
            cave_entrance = down_stairs;

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

        ret
    }

    fn load_prefab(&mut self, origin: Location, prefab: &Prefab) {
        for (&p, &(ref terrain, ref entities)) in prefab.iter() {
            let loc = origin + p;

            self.terrain.insert(loc, *terrain);

            for spawn in entities.iter() {
                if spawn == "player" {
                    self.player_entry = loc;
                } else {
                    let form = Form::named(spawn)
                        .expect(&format!("Bad prefab: Form '{}' not found!", spawn));
                    self.spawns.push((loc, form.loadout.clone()));
                }
            }
        }
    }

    fn load_map_bitmap(&mut self, origin: Location, data: &[u8]) {
        let mut image = image::load(Cursor::new(data), image::ImageFormat::PNG).unwrap();
        // Skip the bottom horizontal line, it's used to store metadata pixels.
        let (w, h) = (image.width(), image.height());
        let input_map = SubImage::new(&mut image, 0, 0, w, h - 1);

        let prefab = calx::Prefab::parse(input_map).expect("Invalid overworld image map");

        self.terrain.extend(
            prefab
                .into_iter()
                .filter_map(|(p, c)| Terrain::from_color(c).map(|t| (origin + p, t))),
        );
    }

    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 }

    /// Make a cave entrance going down.
    fn cave_entrance(&mut self, loc: Location) {
        const DOWNBOUND_ENCLOSURE: [(i32, i32); 5] = [(1, 0), (0, 1), (2, 1), (1, 2), (2, 2)];

        for &v in &DOWNBOUND_ENCLOSURE {
            let loc = loc + vec2(v.0, v.1);
            self.terrain.insert(loc, Terrain::Rock);
        }

        self.terrain.insert(loc, Terrain::Ground);
        self.terrain.insert(loc + vec2(1, 1), Terrain::Ground);
    }

    /// 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(),
        }
    }

    // TODO return impl
    fn domain(&self) -> Vec<mapgen::Point2D> {
        let mut ret = Vec::new();

        let sector_origin = self.sector.origin();
        let mapgen_origin = mapgen::Point2D::zero();

        for loc in self.sector.iter() {
            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) {
                continue;
            }

            ret.push(pos);
        }

        ret
    }

    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);
    }

    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) {
        // TODO
        unimplemented!();
    }

    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
    }
}

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 Rand for Room {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        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::{self, Rand};

        let mut rng = rand::thread_rng();

        let seed = u32::rand(&mut rng);
        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|>


<commit_before>//! This is the Buffer module for Xt.

// This file is part of Xt.

// This is the Xt text editor; it edits text.
// Copyright (C) 2016-2017 The Xt Developers

// This program is free software: you can redistribute it and/or
// modify it under the terms of the GNU General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.

// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see
// <http://www.gnu.org/licenses/>.

extern crate slog;

use logging::init_logger;
use server::gapbuffer::GapBuffer;
use std::path::PathBuf;
use utils::types::BufferID;
use utils::uuid::get_uuid_buffer;

/// Struct for a Buffer in Xt
/// This stores metadata about a buffer.
#[derive(Debug)]
pub struct Buffer {
    /// UUID of a buffer.
    pub buffer_uuid: BufferID,
    /// File path of a buffer.
    pub file_path: Option<PathBuf>,
    /// Active status of a buffer.
    pub active: bool,
    /// Temporary status of a buffer.
    pub temporary: bool,
    /// If a buffer is r/o.
    /// If this is `true`, then r/w is `false.`
    /// Likewise, if this is `false`, then r/w is `true.
    pub read_only: bool,
    /// Major mode of a buffer
    pub major_mode: &'static str,
    /// Vector of Minor modes in a buffer
    pub minor_modes: Vec<&'static str>,
    /// Dirty status of a buffer
    pub dirty: bool,
    /// Contents of a buffer.
    pub text: GapBuffer<u8>,
    /// Logger instance.
    pub logger: slog::Logger,
}

impl Buffer {
    /// Return a new instance of `Buffer`.
    pub fn new() -> Buffer {
        let buffer_uuid = get_uuid_buffer();

        let logger = init_logger().new(o!("buffer_uuid" =>
                                          buffer_uuid.clone()));
        Buffer {
            buffer_uuid: buffer_uuid,
            file_path: None,
            active: false,
            temporary: false,
            read_only: false,
            major_mode: "fundamental-mode",
            minor_modes: Vec::new(),
            dirty: false,
            text: GapBuffer::new(),
            logger: logger,
        }
    }

    /// Return true if a buffer is active.
    /// That is to say, a 'focused' buffer, or otherwise in-use, and
    /// *not* inactive.
    /// If it's not a active buffer, then return false.
    pub fn is_active(&self) -> bool {
        self.active
    }

    /// Return true if a buffer is temporary.
    /// If it's not a temporary buffer, then return false.
    /// A temporary buffer is defined by the characteristic of having
    /// no 'target' to save to.
    pub fn is_temporary(&self) -> bool {
        self.temporary
    }

    /// Return true if a buffer is read only.

    /// If it's not read-only, then it returns false, and the buffer
    /// will be presumed to be read-write.

    /// A read-only buffer is defined by the characteristic of the FS
    /// denying write access, or being set manually.
    pub fn is_ro(&self) -> bool {
        self.read_only
    }

    /// Return true if a buffer is dirty.
    /// If it's not dirty, then return false.
    /// A dirty buffer is defined by the characteristic of the buffer
    /// having changed since the last save.
    /// Temporary buffers are *exempt* from this field.
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Return the current major mode of a buffer.
    pub fn get_major_mode(&self) -> &str {
        self.major_mode
    }

    /// Return a `Vec<&'static str>` of minor modes.
    pub fn get_minor_modes(&self) -> Vec<&'static str> {
        self.minor_modes.clone()
    }

    /// Return the length of a buffer.
    pub fn get_buffer_length(&self) -> usize {
        self.text.len()
    }
}

impl Default for Buffer {
    fn default() -> Buffer {
        Buffer::new()
    }
}

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

    #[test]
    fn test_buffer_default_values() {
        let buf = Buffer::new();

        assert_eq!(None, buf.file_path);
        assert_eq!(false, buf.active);
        assert_eq!(false, buf.is_active());

        assert_eq!(false, buf.temporary);
        assert_eq!(false, buf.is_temporary());

        assert_eq!(false, buf.read_only);
        assert_eq!(false, buf.is_ro());

        assert_eq!("fundamental-mode", buf.major_mode);
        assert_eq!(false, buf.dirty);
    }

    #[test]
    fn test_buffer_major_minor_modes() {
        let mut buf = Buffer::new();
        buf.major_mode = "rust-mode";
        buf.minor_modes.push("auto-wrap-mode");
        buf.minor_modes.push("spell-check-mode");
        buf.minor_modes.push("complete-me-mode");
        buf.minor_modes.push("language-server-mode");

        assert_eq!(4, buf.minor_modes.len());
        assert_eq!("rust-mode", buf.major_mode);
    }
}
<commit_msg>Refactor Buffer modulex<commit_after>//! This is the Buffer module for Xt.

// This file is part of Xt.

// This is the Xt text editor; it edits text.
// Copyright (C) 2016-2017 The Xt Developers

// This program is free software: you can redistribute it and/or
// modify it under the terms of the GNU General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.

// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see
// <http://www.gnu.org/licenses/>.

extern crate gapbuffer;

use self::gapbuffer::GapBuffer;
use std::path::PathBuf;

/// Struct for a Buffer object in Xt.
/// Stores buffer state & metadata.
#[derive(Debug)]
pub struct Buffer {
    /// File path of a buffer.
    pub file_path: Option<PathBuf>,
    /// Active status of a buffer.
    pub active: bool,
    /// Temporary status of a buffer.
    pub temporary: bool,
    /// If a buffer is r/o.
    /// If this is `true`, then r/w is `false.`
    /// Likewise, if this is `false`, then r/w is `true.
    pub read_only: bool,
    /// Major mode of a buffer
    pub major_mode: &'static str,
    /// Vector of Minor modes in a buffer
    pub minor_modes: Vec<&'static str>,
    /// Dirty status of a buffer
    pub dirty: bool,
    /// Contents of a buffer.
    pub text: GapBuffer<u8>,
}

impl Buffer {
    /// Return a new instance of `Buffer`.
    pub fn new() -> Buffer {
        Buffer {
            file_path: None,
            active: false,
            temporary: false,
            read_only: false,
            major_mode: "fundamental-mode",
            minor_modes: Vec::new(),
            dirty: false,
            text: GapBuffer::new()
        }
    }

    /// Return true if a buffer is active.
    /// That is to say, a 'focused' buffer, or otherwise in-use, and
    /// *not* inactive.
    /// If it's not a active buffer, then return false.
    pub fn is_active(&self) -> bool {
        self.active
    }

    /// Return true if a buffer is temporary.
    /// If it's not a temporary buffer, then return false.
    /// A temporary buffer is defined by the characteristic of having
    /// no 'target' to save to.
    pub fn is_temporary(&self) -> bool {
        self.temporary
    }

    /// Return true if a buffer is read only.

    /// If it's not read-only, then it returns false, and the buffer
    /// will be presumed to be read-write.

    /// A read-only buffer is defined by the characteristic of the FS
    /// denying write access, or being set manually.
    pub fn is_ro(&self) -> bool {
        self.read_only
    }

    /// Return true if a buffer is dirty.
    /// If it's not dirty, then return false.
    /// A dirty buffer is defined by the characteristic of the buffer
    /// having changed since the last save.
    /// Temporary buffers are *exempt* from this field.
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Return the current major mode of a buffer.
    pub fn get_major_mode(&self) -> &str {
        self.major_mode
    }

    /// Return a `Vec<&'static str>` of minor modes.
    pub fn get_minor_modes(&self) -> Vec<&'static str> {
        self.minor_modes.clone()
    }

    /// Return the length of a buffer.
    pub fn get_buffer_length(&self) -> usize {
        self.text.len()
    }
}

impl Default for Buffer {
    fn default() -> Buffer {
        Buffer::new()
    }
}

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

    #[test]
    fn test_buffer_default_values() {
        let buf = Buffer::new();

        assert_eq!(None, buf.file_path);
        assert_eq!(false, buf.active);
        assert_eq!(false, buf.is_active());

        assert_eq!(false, buf.temporary);
        assert_eq!(false, buf.is_temporary());

        assert_eq!(false, buf.read_only);
        assert_eq!(false, buf.is_ro());

        assert_eq!("fundamental-mode", buf.major_mode);
        assert_eq!(false, buf.dirty);
    }

    #[test]
    fn test_buffer_major_minor_modes() {
        let mut buf = Buffer::new();
        buf.major_mode = "rust-mode";
        buf.minor_modes.push("auto-wrap-mode");
        buf.minor_modes.push("spell-check-mode");
        buf.minor_modes.push("complete-me-mode");
        buf.minor_modes.push("language-server-mode");

        assert_eq!(4, buf.minor_modes.len());
        assert_eq!("rust-mode", buf.major_mode);
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ascii::AsciiExt;

use sys_info;

use util::collections::HashMap;
use util::config::{self, ReadableSize};

use super::Result;

pub use raftstore::store::Config as RaftStoreConfig;
pub use storage::Config as StorageConfig;

pub const DEFAULT_CLUSTER_ID: u64 = 0;
pub const DEFAULT_LISTENING_ADDR: &'static str = "127.0.0.1:20160";
const DEFAULT_ADVERTISE_LISTENING_ADDR: &'static str = "";
const DEFAULT_NOTIFY_CAPACITY: usize = 40960;
const DEFAULT_GRPC_CONCURRENCY: usize = 4;
const DEFAULT_GRPC_CONCURRENT_STREAM: usize = 1024;
const DEFAULT_GRPC_RAFT_CONN_NUM: usize = 10;
const DEFAULT_GRPC_STREAM_INITIAL_WINDOW_SIZE: u64 = 2 * 1024 * 1024;
const DEFAULT_MESSAGES_PER_TICK: usize = 4096;

// Assume a request can be finished in 1ms, a request at position x will wait about
// 0.001 * x secs to be actual started. A server-is-busy error will trigger 2 seconds
// backoff. So when it needs to wait for more than 2 seconds, return error won't causse
// larger latency.
pub const DEFAULT_MAX_RUNNING_TASK_COUNT: usize = 2 as usize * 1000;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct Config {
    #[serde(skip)]
    pub cluster_id: u64,

    // Server listening address.
    pub addr: String,

    // Server advertise listening address for outer communication.
    // If not set, we will use listening address instead.
    pub advertise_addr: String,
    pub notify_capacity: usize,
    pub messages_per_tick: usize,
    pub grpc_concurrency: usize,
    pub grpc_concurrent_stream: usize,
    pub grpc_raft_conn_num: usize,
    pub grpc_stream_initial_window_size: ReadableSize,
    pub end_point_concurrency: usize,
    pub end_point_max_tasks: usize,
    // Server labels to specify some attributes about this server.
    #[serde(with = "config::order_map_serde")]
    pub labels: HashMap<String, String>,
}

impl Default for Config {
    fn default() -> Config {
        let cpu_num = sys_info::cpu_num().unwrap();
        let concurrency = if cpu_num > 8 {
            (cpu_num as f64 * 0.8) as usize
        } else {
            4
        };
        Config {
            cluster_id: DEFAULT_CLUSTER_ID,
            addr: DEFAULT_LISTENING_ADDR.to_owned(),
            labels: HashMap::default(),
            advertise_addr: DEFAULT_ADVERTISE_LISTENING_ADDR.to_owned(),
            notify_capacity: DEFAULT_NOTIFY_CAPACITY,
            messages_per_tick: DEFAULT_MESSAGES_PER_TICK,
            grpc_concurrency: DEFAULT_GRPC_CONCURRENCY,
            grpc_concurrent_stream: DEFAULT_GRPC_CONCURRENT_STREAM,
            grpc_raft_conn_num: DEFAULT_GRPC_RAFT_CONN_NUM,
            grpc_stream_initial_window_size: ReadableSize(DEFAULT_GRPC_STREAM_INITIAL_WINDOW_SIZE),
            end_point_concurrency: concurrency,
            end_point_max_tasks: DEFAULT_MAX_RUNNING_TASK_COUNT,
        }
    }
}

impl Config {
    pub fn validate(&mut self) -> Result<()> {
        box_try!(config::check_addr(&self.addr));
        if !self.advertise_addr.is_empty() {
            box_try!(config::check_addr(&self.advertise_addr));
        } else {
            info!("no advertise-addr is specified, fall back to addr.");
            self.advertise_addr = self.addr.clone();
        }
        if self.advertise_addr.starts_with("0.") {
            return Err(box_err!(
                "invalid advertise-addr: {:?}",
                self.advertise_addr
            ));
        }

        if self.end_point_concurrency == 0 {
            return Err(box_err!("server.end-point-concurrency should not be 0."));
        }

        if self.end_point_max_tasks == 0 {
            return Err(box_err!("server.end-point-max-tasks should not be 0."));
        }

        for (k, v) in &self.labels {
            validate_label(k, "key")?;
            validate_label(v, "value")?;
        }

        Ok(())
    }
}

fn validate_label(s: &str, tp: &str) -> Result<()> {
    let report_err = || {
        box_err!(
            "store label {}: {:?} not match ^[a-z0-9]([a-z0-9-._]*[a-z0-9])?",
            tp,
            s
        )
    };
    if s.is_empty() {
        return Err(report_err());
    }
    let mut chrs = s.chars();
    let first_char = chrs.next().unwrap();
    if !first_char.is_ascii_lowercase() && !first_char.is_ascii_digit() {
        return Err(report_err());
    }
    let last_char = match chrs.next_back() {
        None => return Ok(()),
        Some(c) => c,
    };
    if !last_char.is_ascii_lowercase() && !last_char.is_ascii_digit() {
        return Err(report_err());
    }
    for c in chrs {
        if !c.is_ascii_lowercase() && !c.is_ascii_digit() && !"-._".contains(c) {
            return Err(report_err());
        }
    }
    Ok(())
}

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

    #[test]
    fn test_config_validate() {
        let mut cfg = Config::default();
        assert!(cfg.advertise_addr.is_empty());
        cfg.validate().unwrap();
        assert_eq!(cfg.addr, cfg.advertise_addr);

        let mut invalid_cfg = cfg.clone();
        invalid_cfg.end_point_concurrency = 0;
        assert!(invalid_cfg.validate().is_err());

        let mut invalid_cfg = cfg.clone();
        invalid_cfg.end_point_max_tasks = 0;
        assert!(invalid_cfg.validate().is_err());

        invalid_cfg = Config::default();
        invalid_cfg.addr = "0.0.0.0:1000".to_owned();
        assert!(invalid_cfg.validate().is_err());
        invalid_cfg.advertise_addr = "127.0.0.1:1000".to_owned();
        invalid_cfg.validate().unwrap();

        cfg.labels.insert("k1".to_owned(), "v1".to_owned());
        cfg.validate().unwrap();
        cfg.labels.insert("k2".to_owned(), "v2?".to_owned());
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_store_labels() {
        let invalid_cases = vec!["", "123*", ".123", "Cab", "abC", "💖"];

        for case in invalid_cases {
            assert!(validate_label(case, "dummy").is_err());
        }

        let valid_cases = vec!["a", "0", "a.1-2", "b_1.2", "cab-012", "3ac.8b2"];

        for case in valid_cases {
            validate_label(case, "dummy").unwrap();
        }
    }
}
<commit_msg>config: adjust label checks (#2372)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ascii::AsciiExt;

use sys_info;

use util::collections::HashMap;
use util::config::{self, ReadableSize};

use super::Result;

pub use raftstore::store::Config as RaftStoreConfig;
pub use storage::Config as StorageConfig;

pub const DEFAULT_CLUSTER_ID: u64 = 0;
pub const DEFAULT_LISTENING_ADDR: &'static str = "127.0.0.1:20160";
const DEFAULT_ADVERTISE_LISTENING_ADDR: &'static str = "";
const DEFAULT_NOTIFY_CAPACITY: usize = 40960;
const DEFAULT_GRPC_CONCURRENCY: usize = 4;
const DEFAULT_GRPC_CONCURRENT_STREAM: usize = 1024;
const DEFAULT_GRPC_RAFT_CONN_NUM: usize = 10;
const DEFAULT_GRPC_STREAM_INITIAL_WINDOW_SIZE: u64 = 2 * 1024 * 1024;
const DEFAULT_MESSAGES_PER_TICK: usize = 4096;

// Assume a request can be finished in 1ms, a request at position x will wait about
// 0.001 * x secs to be actual started. A server-is-busy error will trigger 2 seconds
// backoff. So when it needs to wait for more than 2 seconds, return error won't causse
// larger latency.
pub const DEFAULT_MAX_RUNNING_TASK_COUNT: usize = 2 as usize * 1000;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct Config {
    #[serde(skip)]
    pub cluster_id: u64,

    // Server listening address.
    pub addr: String,

    // Server advertise listening address for outer communication.
    // If not set, we will use listening address instead.
    pub advertise_addr: String,
    pub notify_capacity: usize,
    pub messages_per_tick: usize,
    pub grpc_concurrency: usize,
    pub grpc_concurrent_stream: usize,
    pub grpc_raft_conn_num: usize,
    pub grpc_stream_initial_window_size: ReadableSize,
    pub end_point_concurrency: usize,
    pub end_point_max_tasks: usize,
    // Server labels to specify some attributes about this server.
    #[serde(with = "config::order_map_serde")]
    pub labels: HashMap<String, String>,
}

impl Default for Config {
    fn default() -> Config {
        let cpu_num = sys_info::cpu_num().unwrap();
        let concurrency = if cpu_num > 8 {
            (cpu_num as f64 * 0.8) as usize
        } else {
            4
        };
        Config {
            cluster_id: DEFAULT_CLUSTER_ID,
            addr: DEFAULT_LISTENING_ADDR.to_owned(),
            labels: HashMap::default(),
            advertise_addr: DEFAULT_ADVERTISE_LISTENING_ADDR.to_owned(),
            notify_capacity: DEFAULT_NOTIFY_CAPACITY,
            messages_per_tick: DEFAULT_MESSAGES_PER_TICK,
            grpc_concurrency: DEFAULT_GRPC_CONCURRENCY,
            grpc_concurrent_stream: DEFAULT_GRPC_CONCURRENT_STREAM,
            grpc_raft_conn_num: DEFAULT_GRPC_RAFT_CONN_NUM,
            grpc_stream_initial_window_size: ReadableSize(DEFAULT_GRPC_STREAM_INITIAL_WINDOW_SIZE),
            end_point_concurrency: concurrency,
            end_point_max_tasks: DEFAULT_MAX_RUNNING_TASK_COUNT,
        }
    }
}

impl Config {
    pub fn validate(&mut self) -> Result<()> {
        box_try!(config::check_addr(&self.addr));
        if !self.advertise_addr.is_empty() {
            box_try!(config::check_addr(&self.advertise_addr));
        } else {
            info!("no advertise-addr is specified, fall back to addr.");
            self.advertise_addr = self.addr.clone();
        }
        if self.advertise_addr.starts_with("0.") {
            return Err(box_err!(
                "invalid advertise-addr: {:?}",
                self.advertise_addr
            ));
        }

        if self.end_point_concurrency == 0 {
            return Err(box_err!("server.end-point-concurrency should not be 0."));
        }

        if self.end_point_max_tasks == 0 {
            return Err(box_err!("server.end-point-max-tasks should not be 0."));
        }

        for (k, v) in &self.labels {
            validate_label(k, "key")?;
            validate_label(v, "value")?;
        }

        Ok(())
    }
}

fn validate_label(s: &str, tp: &str) -> Result<()> {
    let report_err = || {
        box_err!(
            "store label {}: {:?} not match ^[a-zA-Z0-9]([a-zA-Z0-9-._]*[a-zA-Z0-9])?",
            tp,
            s
        )
    };
    if s.is_empty() {
        return Err(report_err());
    }
    let mut chrs = s.chars();
    let first_char = chrs.next().unwrap();
    if !first_char.is_ascii_alphanumeric() {
        return Err(report_err());
    }
    let last_char = match chrs.next_back() {
        None => return Ok(()),
        Some(c) => c,
    };
    if !last_char.is_ascii_alphanumeric() {
        return Err(report_err());
    }
    for c in chrs {
        if !c.is_ascii_alphanumeric() && !"-._".contains(c) {
            return Err(report_err());
        }
    }
    Ok(())
}

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

    #[test]
    fn test_config_validate() {
        let mut cfg = Config::default();
        assert!(cfg.advertise_addr.is_empty());
        cfg.validate().unwrap();
        assert_eq!(cfg.addr, cfg.advertise_addr);

        let mut invalid_cfg = cfg.clone();
        invalid_cfg.end_point_concurrency = 0;
        assert!(invalid_cfg.validate().is_err());

        let mut invalid_cfg = cfg.clone();
        invalid_cfg.end_point_max_tasks = 0;
        assert!(invalid_cfg.validate().is_err());

        invalid_cfg = Config::default();
        invalid_cfg.addr = "0.0.0.0:1000".to_owned();
        assert!(invalid_cfg.validate().is_err());
        invalid_cfg.advertise_addr = "127.0.0.1:1000".to_owned();
        invalid_cfg.validate().unwrap();

        cfg.labels.insert("k1".to_owned(), "v1".to_owned());
        cfg.validate().unwrap();
        cfg.labels.insert("k2".to_owned(), "v2?".to_owned());
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_store_labels() {
        let invalid_cases = vec!["", "123*", ".123", "💖"];

        for case in invalid_cases {
            assert!(validate_label(case, "dummy").is_err());
        }

        let valid_cases = vec![
            "a",
            "0",
            "a.1-2",
            "Cab",
            "abC",
            "b_1.2",
            "cab-012",
            "3ac.8b2",
        ];

        for case in valid_cases {
            validate_label(case, "dummy").unwrap();
        }
    }
}
<|endoftext|>