<commit_before>use super::CrateVersion;
use serde_json;
use std::path::Path;

use git2::{
    build::RepoBuilder, Delta, DiffFormat, Error as GitError, ErrorClass, Object, ObjectType, Oid,
    Reference, Repository, Tree,
};
use std::str;

static INDEX_GIT_URL: &str = "https://github.com/rust-lang/crates.io-index";
static LAST_SEEN_REFNAME: &str = "refs/heads/crates-index-diff_last-seen";
static EMPTY_TREE_HASH: &str = "4b825dc642cb6eb9a060e54bf8d69288fbee4904";
static LINE_ADDED_INDICATOR: char = '+';

/// A wrapper for a repository of the crates.io index.
pub struct Index {
    /// The name and path of the reference used to keep track of the last seen state of the
    /// crates.io repository. The default value is `refs/heads/crates-index-diff_last-seen`.
    pub seen_ref_name: &'static str,
    /// The crates.io repository.
    repo: Repository,
}

/// Options for use in `Index::from_path_or_cloned_with_options`
#[non_exhaustive]
pub struct CloneOptions<'a> {
    pub repository_url: String,
    pub fetch_options: Option<git2::FetchOptions<'a>>,
}

impl<'a> Default for CloneOptions<'a> {
    fn default() -> Self {
        CloneOptions {
            repository_url: INDEX_GIT_URL.into(),
            fetch_options: None,
        }
    }
}

impl Index {
    /// Return the crates.io repository.
    pub fn repository(&self) -> &Repository {
        &self.repo
    }

    /// Return the reference pointing to the state we have seen after calling `fetch_changes()`.
    pub fn last_seen_reference(&self) -> Result<Reference, GitError> {
        self.repo.find_reference(self.seen_ref_name)
    }

    /// Return a new `Index` instance from the given `path`, which should contain a bare or non-bare
    /// clone of the `crates.io` index.
    /// If the directory does not contain the repository or does not exist, it will be cloned from
    /// the official location automatically (with complete history).
    ///
    /// An error will occour if the repository exists and the remote URL does not match the given repository URL.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use crates_index_diff::{CloneOptions, Index};
    ///
    /// # let path = tempdir::TempDir::new("index").unwrap();
    /// let mut options = CloneOptions {
    ///   repository_url: "https://github.com/rust-lang/staging.crates.io-index".into(),
    ///   ..Default::default()
    /// };
    ///
    ///
    /// let index = Index::from_path_or_cloned_with_options(path, options)?;
    /// # Ok::<(), git2::Error>(())
    /// ```
    /// Or to access a private repository, use fetch options.
    ///
    /// ```no_run
    /// use crates_index_diff::{CloneOptions, Index};
    /// let fo = {
    ///     let mut fo = git2::FetchOptions::new();
    ///     fo.remote_callbacks({
    ///         let mut callbacks = git2::RemoteCallbacks::new();
    ///         callbacks.credentials(|_url, username_from_url, _allowed_types| {
    ///             git2::Cred::ssh_key_from_memory(
    ///                 username_from_url.unwrap(),
    ///                 None,
    ///                 &std::env::var("PRIVATE_KEY").unwrap(),
    ///                 None,
    ///             )
    ///         });
    ///         callbacks
    ///     });
    ///     fo
    /// };
    /// Index::from_path_or_cloned_with_options(
    ///     "index",
    ///     CloneOptions {
    ///         repository_url: "git@github.com:private-index/goes-here.git".into(),
    ///         fetch_options: Some(fo),
    ///     },
    /// ).unwrap();
    /// ```
    pub fn from_path_or_cloned_with_options(
        path: impl AsRef<Path>,
        CloneOptions {
            repository_url,
            fetch_options,
        }: CloneOptions,
    ) -> Result<Index, GitError> {
        let mut repo_did_exist = true;
        let repo = Repository::open(path.as_ref()).or_else(|err| {
            if err.class() == ErrorClass::Repository {
                repo_did_exist = false;
                let mut builder = RepoBuilder::new();
                if let Some(fo) = fetch_options {
                    builder.fetch_options(fo);
                }
                builder.bare(true).clone(&repository_url, path.as_ref())
            } else {
                Err(err)
            }
        })?;

        if repo_did_exist {
            let remote = repo.find_remote("origin")?;
            let actual_remote_url = remote
                .url()
                .ok_or_else(|| GitError::from_str("did not obtain URL of remote named 'origin'"))?;
            if actual_remote_url != repository_url {
                return Err(GitError::from_str(&format!(
                    "Actual 'origin' remote url {:#?} did not match desired one at {:#?}",
                    actual_remote_url, repository_url
                )));
            }
        }

        Ok(Index {
            repo,
            seen_ref_name: LAST_SEEN_REFNAME,
        })
    }

    /// Return a new `Index` instance from the given `path`, which should contain a bare or non-bare
    /// clone of the `crates.io` index.
    /// If the directory does not contain the repository or does not exist, it will be cloned from
    /// the official location automatically (with complete history).
    pub fn from_path_or_cloned(path: impl AsRef<Path>) -> Result<Index, GitError> {
        Index::from_path_or_cloned_with_options(path, CloneOptions::default())
    }

    /// As `peek_changes_with_options`, but without the options.
    pub fn peek_changes(&self) -> Result<(Vec<CrateVersion>, git2::Oid), GitError> {
        self.peek_changes_with_options(None)
    }

    /// Return all `CrateVersion`s that are observed between the last time `fetch_changes(…)` was called
    /// and the latest state of the `crates.io` index repository, which is obtained by fetching
    /// the remote called `origin`.
    /// The `last_seen_reference()` will not be created or updated.
    /// The second field in the returned tuple is the commit object to which the changes were provided.
    /// If one would set the `last_seen_reference()` to that object, the effect is exactly the same
    /// as if `fetch_changes(…)` had been called.
    ///
    /// # Resource Usage
    ///
    /// As this method fetches the git repository, loose objects or small packs may be created. Over time,
    /// these will accumulate and either slow down subsequent operations, or cause them to fail due to exhaustion
    /// of the maximum number of open file handles as configured with `ulimit`.
    ///
    /// Thus it is advised for the caller to run `git gc` occasionally based on their own requirements and usage patterns.
    pub fn peek_changes_with_options(
        &self,
        options: Option<&mut git2::FetchOptions<'_>>,
    ) -> Result<(Vec<CrateVersion>, git2::Oid), GitError> {
        let from = self
            .last_seen_reference()
            .and_then(|r| {
                r.target().ok_or_else(|| {
                    GitError::from_str("last-seen reference did not have a valid target")
                })
            })
            .or_else(|_| Oid::from_str(EMPTY_TREE_HASH))?;
        let to = {
            self.repo.find_remote("origin").and_then(|mut r| {
                r.fetch(&["refs/heads/*:refs/remotes/origin/*"], options, None)
            })?;
            let latest_fetched_commit_oid =
                self.repo.refname_to_id("refs/remotes/origin/master")?;
            latest_fetched_commit_oid
        };

        Ok((
            self.changes_from_objects(
                &self.repo.find_object(from, None)?,
                &self.repo.find_object(to, None)?,
            )?,
            to,
        ))
    }

    /// As `fetch_changes_with_options`, but without the options.
    pub fn fetch_changes(&self) -> Result<Vec<CrateVersion>, GitError> {
        self.fetch_changes_with_options(None)
    }

    /// Return all `CrateVersion`s that are observed between the last time this method was called
    /// and the latest state of the `crates.io` index repository, which is obtained by fetching
    /// the remote called `origin`.
    /// The `last_seen_reference()` will be created or adjusted to point to the latest fetched
    /// state, which causes this method to have a different result each time it is called.
    ///
    /// # Resource Usage
    ///
    /// As this method fetches the git repository, loose objects or small packs may be created. Over time,
    /// these will accumulate and either slow down subsequent operations, or cause them to fail due to exhaustion
    /// of the maximum number of open file handles as configured with `ulimit`.
    ///
    /// Thus it is advised for the caller to run `git gc` occasionally based on their own requirements and usage patterns.
    pub fn fetch_changes_with_options(
        &self,
        options: Option<&mut git2::FetchOptions<'_>>,
    ) -> Result<Vec<CrateVersion>, GitError> {
        let (changes, to) = self.peek_changes_with_options(options)?;
        self.set_last_seen_reference(to)?;
        Ok(changes)
    }

    /// Set the last seen reference to the given Oid. It will be created if it does not yet exists.
    pub fn set_last_seen_reference(&self, to: Oid) -> Result<(), GitError> {
        self.last_seen_reference()
            .and_then(|mut seen_ref| {
                seen_ref.set_target(to, "updating seen-ref head to latest fetched commit")
            })
            .or_else(|_err| {
                self.repo.reference(
                    self.seen_ref_name,
                    to,
                    true,
                    "creating seen-ref at latest fetched commit",
                )
            })?;
        Ok(())
    }

    /// Return all `CreateVersion`s observed between `from` and `to`. Both parameter are ref-specs
    /// pointing to either a commit or a tree.
    /// Learn more about specifying revisions
    /// in the
    /// [official documentation](https://www.kernel.org/pub/software/scm/git/docs/gitrevisions.html)
    pub fn changes(
        &self,
        from: impl AsRef<str>,
        to: impl AsRef<str>,
    ) -> Result<Vec<CrateVersion>, GitError> {
        self.changes_from_objects(
            &self.repo.revparse_single(from.as_ref())?,
            &self.repo.revparse_single(to.as_ref())?,
        )
    }

    /// Similar to `changes()`, but requires `from` and `to` objects to be provided. They may point
    /// to either `Commit`s or `Tree`s.
    pub fn changes_from_objects(
        &self,
        from: &Object,
        to: &Object,
    ) -> Result<Vec<CrateVersion>, GitError> {
        fn into_tree<'a>(repo: &'a Repository, obj: &Object) -> Result<Tree<'a>, GitError> {
            repo.find_tree(match obj.kind() {
                Some(ObjectType::Commit) => obj
                    .as_commit()
                    .expect("object of kind commit yields commit")
                    .tree_id(),
                _ =>
                /* let it possibly fail later */
                {
                    obj.id()
                }
            })
        }
        let diff = self.repo.diff_tree_to_tree(
            Some(&into_tree(&self.repo, from)?),
            Some(&into_tree(&self.repo, to)?),
            None,
        )?;
        let mut res: Vec<CrateVersion> = Vec::new();
        diff.print(DiffFormat::Patch, |delta, _, diffline| {
            if diffline.origin() != LINE_ADDED_INDICATOR {
                return true;
            }

            if !match delta.status() {
                Delta::Added | Delta::Modified => true,
                _ => false,
            } {
                return true;
            }

            if let Ok(c) = serde_json::from_slice(diffline.content()) {
                res.push(c)
            }
            true
        })
        .map(|_| res)
    }
}
<commit_msg>remove non-exhaustive attribute…<commit_after>use super::CrateVersion;
use serde_json;
use std::path::Path;

use git2::{
    build::RepoBuilder, Delta, DiffFormat, Error as GitError, ErrorClass, Object, ObjectType, Oid,
    Reference, Repository, Tree,
};
use std::str;

static INDEX_GIT_URL: &str = "https://github.com/rust-lang/crates.io-index";
static LAST_SEEN_REFNAME: &str = "refs/heads/crates-index-diff_last-seen";
static EMPTY_TREE_HASH: &str = "4b825dc642cb6eb9a060e54bf8d69288fbee4904";
static LINE_ADDED_INDICATOR: char = '+';

/// A wrapper for a repository of the crates.io index.
pub struct Index {
    /// The name and path of the reference used to keep track of the last seen state of the
    /// crates.io repository. The default value is `refs/heads/crates-index-diff_last-seen`.
    pub seen_ref_name: &'static str,
    /// The crates.io repository.
    repo: Repository,
}

/// Options for use in `Index::from_path_or_cloned_with_options`
pub struct CloneOptions<'a> {
    pub repository_url: String,
    pub fetch_options: Option<git2::FetchOptions<'a>>,
}

impl<'a> Default for CloneOptions<'a> {
    fn default() -> Self {
        CloneOptions {
            repository_url: INDEX_GIT_URL.into(),
            fetch_options: None,
        }
    }
}

impl Index {
    /// Return the crates.io repository.
    pub fn repository(&self) -> &Repository {
        &self.repo
    }

    /// Return the reference pointing to the state we have seen after calling `fetch_changes()`.
    pub fn last_seen_reference(&self) -> Result<Reference, GitError> {
        self.repo.find_reference(self.seen_ref_name)
    }

    /// Return a new `Index` instance from the given `path`, which should contain a bare or non-bare
    /// clone of the `crates.io` index.
    /// If the directory does not contain the repository or does not exist, it will be cloned from
    /// the official location automatically (with complete history).
    ///
    /// An error will occour if the repository exists and the remote URL does not match the given repository URL.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use crates_index_diff::{CloneOptions, Index};
    ///
    /// # let path = tempdir::TempDir::new("index").unwrap();
    /// let mut options = CloneOptions {
    ///   repository_url: "https://github.com/rust-lang/staging.crates.io-index".into(),
    ///   ..Default::default()
    /// };
    ///
    ///
    /// let index = Index::from_path_or_cloned_with_options(path, options)?;
    /// # Ok::<(), git2::Error>(())
    /// ```
    /// Or to access a private repository, use fetch options.
    ///
    /// ```no_run
    /// use crates_index_diff::{CloneOptions, Index};
    /// let fo = {
    ///     let mut fo = git2::FetchOptions::new();
    ///     fo.remote_callbacks({
    ///         let mut callbacks = git2::RemoteCallbacks::new();
    ///         callbacks.credentials(|_url, username_from_url, _allowed_types| {
    ///             git2::Cred::ssh_key_from_memory(
    ///                 username_from_url.unwrap(),
    ///                 None,
    ///                 &std::env::var("PRIVATE_KEY").unwrap(),
    ///                 None,
    ///             )
    ///         });
    ///         callbacks
    ///     });
    ///     fo
    /// };
    /// Index::from_path_or_cloned_with_options(
    ///     "index",
    ///     CloneOptions {
    ///         repository_url: "git@github.com:private-index/goes-here.git".into(),
    ///         fetch_options: Some(fo),
    ///     },
    /// ).unwrap();
    /// ```
    pub fn from_path_or_cloned_with_options(
        path: impl AsRef<Path>,
        CloneOptions {
            repository_url,
            fetch_options,
        }: CloneOptions,
    ) -> Result<Index, GitError> {
        let mut repo_did_exist = true;
        let repo = Repository::open(path.as_ref()).or_else(|err| {
            if err.class() == ErrorClass::Repository {
                repo_did_exist = false;
                let mut builder = RepoBuilder::new();
                if let Some(fo) = fetch_options {
                    builder.fetch_options(fo);
                }
                builder.bare(true).clone(&repository_url, path.as_ref())
            } else {
                Err(err)
            }
        })?;

        if repo_did_exist {
            let remote = repo.find_remote("origin")?;
            let actual_remote_url = remote
                .url()
                .ok_or_else(|| GitError::from_str("did not obtain URL of remote named 'origin'"))?;
            if actual_remote_url != repository_url {
                return Err(GitError::from_str(&format!(
                    "Actual 'origin' remote url {:#?} did not match desired one at {:#?}",
                    actual_remote_url, repository_url
                )));
            }
        }

        Ok(Index {
            repo,
            seen_ref_name: LAST_SEEN_REFNAME,
        })
    }

    /// Return a new `Index` instance from the given `path`, which should contain a bare or non-bare
    /// clone of the `crates.io` index.
    /// If the directory does not contain the repository or does not exist, it will be cloned from
    /// the official location automatically (with complete history).
    pub fn from_path_or_cloned(path: impl AsRef<Path>) -> Result<Index, GitError> {
        Index::from_path_or_cloned_with_options(path, CloneOptions::default())
    }

    /// As `peek_changes_with_options`, but without the options.
    pub fn peek_changes(&self) -> Result<(Vec<CrateVersion>, git2::Oid), GitError> {
        self.peek_changes_with_options(None)
    }

    /// Return all `CrateVersion`s that are observed between the last time `fetch_changes(…)` was called
    /// and the latest state of the `crates.io` index repository, which is obtained by fetching
    /// the remote called `origin`.
    /// The `last_seen_reference()` will not be created or updated.
    /// The second field in the returned tuple is the commit object to which the changes were provided.
    /// If one would set the `last_seen_reference()` to that object, the effect is exactly the same
    /// as if `fetch_changes(…)` had been called.
    ///
    /// # Resource Usage
    ///
    /// As this method fetches the git repository, loose objects or small packs may be created. Over time,
    /// these will accumulate and either slow down subsequent operations, or cause them to fail due to exhaustion
    /// of the maximum number of open file handles as configured with `ulimit`.
    ///
    /// Thus it is advised for the caller to run `git gc` occasionally based on their own requirements and usage patterns.
    pub fn peek_changes_with_options(
        &self,
        options: Option<&mut git2::FetchOptions<'_>>,
    ) -> Result<(Vec<CrateVersion>, git2::Oid), GitError> {
        let from = self
            .last_seen_reference()
            .and_then(|r| {
                r.target().ok_or_else(|| {
                    GitError::from_str("last-seen reference did not have a valid target")
                })
            })
            .or_else(|_| Oid::from_str(EMPTY_TREE_HASH))?;
        let to = {
            self.repo.find_remote("origin").and_then(|mut r| {
                r.fetch(&["refs/heads/*:refs/remotes/origin/*"], options, None)
            })?;
            let latest_fetched_commit_oid =
                self.repo.refname_to_id("refs/remotes/origin/master")?;
            latest_fetched_commit_oid
        };

        Ok((
            self.changes_from_objects(
                &self.repo.find_object(from, None)?,
                &self.repo.find_object(to, None)?,
            )?,
            to,
        ))
    }

    /// As `fetch_changes_with_options`, but without the options.
    pub fn fetch_changes(&self) -> Result<Vec<CrateVersion>, GitError> {
        self.fetch_changes_with_options(None)
    }

    /// Return all `CrateVersion`s that are observed between the last time this method was called
    /// and the latest state of the `crates.io` index repository, which is obtained by fetching
    /// the remote called `origin`.
    /// The `last_seen_reference()` will be created or adjusted to point to the latest fetched
    /// state, which causes this method to have a different result each time it is called.
    ///
    /// # Resource Usage
    ///
    /// As this method fetches the git repository, loose objects or small packs may be created. Over time,
    /// these will accumulate and either slow down subsequent operations, or cause them to fail due to exhaustion
    /// of the maximum number of open file handles as configured with `ulimit`.
    ///
    /// Thus it is advised for the caller to run `git gc` occasionally based on their own requirements and usage patterns.
    pub fn fetch_changes_with_options(
        &self,
        options: Option<&mut git2::FetchOptions<'_>>,
    ) -> Result<Vec<CrateVersion>, GitError> {
        let (changes, to) = self.peek_changes_with_options(options)?;
        self.set_last_seen_reference(to)?;
        Ok(changes)
    }

    /// Set the last seen reference to the given Oid. It will be created if it does not yet exists.
    pub fn set_last_seen_reference(&self, to: Oid) -> Result<(), GitError> {
        self.last_seen_reference()
            .and_then(|mut seen_ref| {
                seen_ref.set_target(to, "updating seen-ref head to latest fetched commit")
            })
            .or_else(|_err| {
                self.repo.reference(
                    self.seen_ref_name,
                    to,
                    true,
                    "creating seen-ref at latest fetched commit",
                )
            })?;
        Ok(())
    }

    /// Return all `CreateVersion`s observed between `from` and `to`. Both parameter are ref-specs
    /// pointing to either a commit or a tree.
    /// Learn more about specifying revisions
    /// in the
    /// [official documentation](https://www.kernel.org/pub/software/scm/git/docs/gitrevisions.html)
    pub fn changes(
        &self,
        from: impl AsRef<str>,
        to: impl AsRef<str>,
    ) -> Result<Vec<CrateVersion>, GitError> {
        self.changes_from_objects(
            &self.repo.revparse_single(from.as_ref())?,
            &self.repo.revparse_single(to.as_ref())?,
        )
    }

    /// Similar to `changes()`, but requires `from` and `to` objects to be provided. They may point
    /// to either `Commit`s or `Tree`s.
    pub fn changes_from_objects(
        &self,
        from: &Object,
        to: &Object,
    ) -> Result<Vec<CrateVersion>, GitError> {
        fn into_tree<'a>(repo: &'a Repository, obj: &Object) -> Result<Tree<'a>, GitError> {
            repo.find_tree(match obj.kind() {
                Some(ObjectType::Commit) => obj
                    .as_commit()
                    .expect("object of kind commit yields commit")
                    .tree_id(),
                _ =>
                /* let it possibly fail later */
                {
                    obj.id()
                }
            })
        }
        let diff = self.repo.diff_tree_to_tree(
            Some(&into_tree(&self.repo, from)?),
            Some(&into_tree(&self.repo, to)?),
            None,
        )?;
        let mut res: Vec<CrateVersion> = Vec::new();
        diff.print(DiffFormat::Patch, |delta, _, diffline| {
            if diffline.origin() != LINE_ADDED_INDICATOR {
                return true;
            }

            if !match delta.status() {
                Delta::Added | Delta::Modified => true,
                _ => false,
            } {
                return true;
            }

            if let Ok(c) = serde_json::from_slice(diffline.content()) {
                res.push(c)
            }
            true
        })
        .map(|_| res)
    }
}
<|endoftext|>


<commit_before>use crate::convert::*;
use crate::folded_multiply::*;
use core::hash::Hasher;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [`AHasher`] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: [u64; 2],
    sum: [u64; 2],
    key: u128,
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    /// # Example
    ///
    /// ```
    /// use std::hash::Hasher;
    /// use ahash::AHasher;
    ///
    /// let mut hasher = AHasher::new_with_keys(123, 456);
    ///
    /// hasher.write_u32(1989);
    /// hasher.write_u8(11);
    /// hasher.write_u8(9);
    /// hasher.write(b"Huh?");
    ///
    /// println!("Hash is {:x}!", hasher.finish());
    /// ```
    #[inline]
    pub fn new_with_keys(key0: u64, key1: u64) -> Self {
        Self {
            buffer: [key0, key1],
            sum: [key1, key0],
            key: [key0, key1].convert(),
        }
    }

    #[cfg(test)]
    pub(crate) fn test_with_keys(key1: u64, key2: u64) -> AHasher {
        use crate::random_state::scramble_keys;
        let (k1, k2) = scramble_keys(key1, key2);
        AHasher {
            buffer: [k1, k2],
            sum: [k2, k1],
            key: [k1, k2].convert(),
        }
    }

    #[inline(always)]
    fn hash_in(&mut self, new_value: u128) {
        self.buffer = aeshash(self.buffer.convert(), new_value).convert();
        self.sum = add_by_64s(self.sum.convert(), self.buffer.convert()).convert();
    }

    #[inline(always)]
    fn hash_in_2(&mut self, v1: u128, v2: u128) {
        self.hash_in(v1);
        self.hash_in(v2);
    }
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {
    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        self.hash_in(i);
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //This will be scrambled by the first AES round in any branch.
        self.buffer[1] = self.buffer[1].wrapping_add(length);
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() <= 8 {
            let value: [u64; 2] = if data.len() >= 2 {
                if data.len() >= 4 {
                    //len 4-8
                    [data.read_u32().0 as u64, data.read_last_u32() as u64]
                } else {
                    //len 2-3
                    [data.read_u16().0 as u64, data[data.len() - 1] as u64]
                }
            } else {
                if data.len() > 0 {
                    [data[0] as u64, 0]
                } else {
                    [0, 0]
                }
            };
            self.hash_in(value.convert());
        } else {
            if data.len() > 32 {
                if data.len() > 64 {
                    let tail = data.read_last_u128x4();
                    let mut current: [u128; 2] = [self.buffer.convert(), self.key];
                    let mut sum: [u128; 2] = [0, 0];
                    current[0] = aeshash(current[0], tail[0]);
                    current[1] = aeshash(current[1], tail[1]);
                    sum[0] = add_by_64s(sum[0], current[0]);
                    sum[1] = add_by_64s(sum[1], current[1]);
                    current[0] = aeshash(current[0], tail[2]);
                    current[1] = aeshash(current[1], tail[3]);
                    sum[0] = add_by_64s(sum[0], current[0]);
                    sum[1] = add_by_64s(sum[1], current[1]);
                    while data.len() > 64 {
                        let (blocks, rest) = data.read_u128x4();
                        current[0] = aeshash(current[0], blocks[0]);
                        current[1] = aeshash(current[1], blocks[1]);
                        sum[0] = add_by_64s(sum[0], current[0]);
                        sum[1] = add_by_64s(sum[1], current[1]);
                        current[0] = aeshash(current[0], blocks[2]);
                        current[1] = aeshash(current[1], blocks[3]);
                        sum[0] = add_by_64s(sum[0], current[0]);
                        sum[1] = add_by_64s(sum[1], current[1]);
                        data = rest;
                    }
                    self.buffer = current[0].convert();
                    self.hash_in(current[1]);
                    self.hash_in_2(sum[0], sum[1]);
                } else {
                    //len 33-64
                    let (head, _) = data.read_u128x2();
                    let tail = data.read_last_u128x2();
                    self.hash_in_2(head[0], head[1]);
                    self.hash_in_2(tail[0], tail[1]);
                }
            } else {
                if data.len() > 16 {
                    //len 17-32
                    self.hash_in_2(data.read_u128().0, data.read_last_u128());
                } else {
                    //len 9-16
                    let value: [u64; 2] = [data.read_u64().0, data.read_last_u64()];
                    self.hash_in(value.convert());
                }
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        let result: [u64; 2] = aeshashx2(self.sum.convert(), self.key, self.buffer.convert()).convert();
        result[0] //.wrapping_add(result[1])
    }
}

#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)))]
#[inline(always)]
fn aeshash(value: u128, xor: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        transmute(_mm_aesdec_si128(value, transmute(xor)))
    }
}
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)))]
#[inline(always)]
fn aeshashx2(value: u128, k1: u128, k2: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        let value = _mm_aesdec_si128(value, transmute(k1));
        transmute(_mm_aesdec_si128(value, transmute(k2)))
    }
}

#[cfg(test)]
mod tests {
    use crate::aes_hash::*;
    use crate::convert::Convert;
    use std::collections::HashMap;
    use std::hash::BuildHasherDefault;

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_builder() {
        let mut map = HashMap::<u32, u64, BuildHasherDefault<AHasher>>::default();
        map.insert(1, 3);
    }

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_default() {
        let hasher_a = AHasher::default();
        assert_ne!(0, hasher_a.buffer[0]);
        assert_ne!(0, hasher_a.buffer[1]);
        assert_ne!(hasher_a.buffer[0], hasher_a.buffer[1]);
        let hasher_b = AHasher::default();
        assert_eq!(hasher_a.buffer[0], hasher_b.buffer[0]);
        assert_eq!(hasher_a.buffer[1], hasher_b.buffer[1]);
    }

    #[test]
    fn test_hash() {
        let mut result: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value: [u64; 2] = [1 << 32, 0xFEDCBA9876543210];
        result = aeshash(value.convert(), result.convert()).convert();
        result = aeshash(result.convert(), result.convert()).convert();
        let mut result2: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value2: [u64; 2] = [1, 0xFEDCBA9876543210];
        result2 = aeshash(value2.convert(), result2.convert()).convert();
        result2 = aeshash(result2.convert(), result.convert()).convert();
        let result: [u8; 16] = result.convert();
        let result2: [u8; 16] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<commit_msg>Speed up long string tail with AES<commit_after>use crate::convert::*;
use crate::folded_multiply::*;
use core::hash::Hasher;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [`AHasher`] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: [u64; 2],
    sum: [u64; 2],
    key: u128,
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    /// # Example
    ///
    /// ```
    /// use std::hash::Hasher;
    /// use ahash::AHasher;
    ///
    /// let mut hasher = AHasher::new_with_keys(123, 456);
    ///
    /// hasher.write_u32(1989);
    /// hasher.write_u8(11);
    /// hasher.write_u8(9);
    /// hasher.write(b"Huh?");
    ///
    /// println!("Hash is {:x}!", hasher.finish());
    /// ```
    #[inline]
    pub fn new_with_keys(key0: u64, key1: u64) -> Self {
        Self {
            buffer: [key0, key1],
            sum: [key1, key0],
            key: [key0, key1].convert(),
        }
    }

    #[cfg(test)]
    pub(crate) fn test_with_keys(key1: u64, key2: u64) -> AHasher {
        use crate::random_state::scramble_keys;
        let (k1, k2) = scramble_keys(key1, key2);
        AHasher {
            buffer: [k1, k2],
            sum: [k2, k1],
            key: [k1, k2].convert(),
        }
    }

    #[inline(always)]
    fn hash_in(&mut self, new_value: u128) {
        self.buffer = aeshash(self.buffer.convert(), new_value).convert();
        self.sum = add_by_64s(self.sum.convert(), self.buffer.convert()).convert();
    }

    #[inline(always)]
    fn hash_in_2(&mut self, v1: u128, v2: u128) {
        self.hash_in(v1);
        self.hash_in(v2);
    }
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {
    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        self.hash_in(i);
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //This will be scrambled by the first AES round in any branch.
        self.buffer[1] = self.buffer[1].wrapping_add(length);
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() <= 8 {
            let value: [u64; 2] = if data.len() >= 2 {
                if data.len() >= 4 {
                    //len 4-8
                    [data.read_u32().0 as u64, data.read_last_u32() as u64]
                } else {
                    //len 2-3
                    [data.read_u16().0 as u64, data[data.len() - 1] as u64]
                }
            } else {
                if data.len() > 0 {
                    [data[0] as u64, 0]
                } else {
                    [0, 0]
                }
            };
            self.hash_in(value.convert());
        } else {
            if data.len() > 32 {
                if data.len() > 64 {
                    let tail = data.read_last_u128x4();
                    let mut current: [u128; 2] = [self.buffer.convert(), self.key];
                    let mut sum: [u128; 2] = [0, 0];
                    current[0] = aeshash(current[0], tail[0]);
                    current[1] = aeshash(current[1], tail[1]);
                    sum[0] = add_by_64s(sum[0], current[0]);
                    sum[1] = add_by_64s(sum[1], current[1]);
                    current[0] = aeshash(current[0], tail[2]);
                    current[1] = aeshash(current[1], tail[3]);
                    sum[0] = add_by_64s(sum[0], current[0]);
                    sum[1] = add_by_64s(sum[1], current[1]);
                    while data.len() > 64 {
                        let (blocks, rest) = data.read_u128x4();
                        current[0] = aeshash(current[0], blocks[0]);
                        current[1] = aeshash(current[1], blocks[1]);
                        sum[0] = add_by_64s(sum[0], current[0]);
                        sum[1] = add_by_64s(sum[1], current[1]);
                        current[0] = aeshash(current[0], blocks[2]);
                        current[1] = aeshash(current[1], blocks[3]);
                        sum[0] = add_by_64s(sum[0], current[0]);
                        sum[1] = add_by_64s(sum[1], current[1]);
                        data = rest;
                    }
                    self.buffer = current[0].convert();
                    self.hash_in(current[1]);
                    self.sum = add_by_64s(self.sum.convert(), sum[0]).convert();
                    self.sum = add_by_64s(self.sum.convert(), sum[1]).convert();
                } else {
                    //len 33-64
                    let (head, _) = data.read_u128x2();
                    let tail = data.read_last_u128x2();
                    self.hash_in_2(head[0], head[1]);
                    self.hash_in_2(tail[0], tail[1]);
                }
            } else {
                if data.len() > 16 {
                    //len 17-32
                    self.hash_in_2(data.read_u128().0, data.read_last_u128());
                } else {
                    //len 9-16
                    let value: [u64; 2] = [data.read_u64().0, data.read_last_u64()];
                    self.hash_in(value.convert());
                }
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        let result: [u64; 2] = aeshashx2(self.sum.convert(), self.key, self.buffer.convert()).convert();
        result[0] //.wrapping_add(result[1])
    }
}

#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)))]
#[inline(always)]
fn aeshash(value: u128, xor: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        transmute(_mm_aesdec_si128(value, transmute(xor)))
    }
}
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)))]
#[inline(always)]
fn aeshashx2(value: u128, k1: u128, k2: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        let value = _mm_aesdec_si128(value, transmute(k1));
        transmute(_mm_aesdec_si128(value, transmute(k2)))
    }
}

#[cfg(test)]
mod tests {
    use crate::aes_hash::*;
    use crate::convert::Convert;
    use std::collections::HashMap;
    use std::hash::BuildHasherDefault;

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_builder() {
        let mut map = HashMap::<u32, u64, BuildHasherDefault<AHasher>>::default();
        map.insert(1, 3);
    }

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_default() {
        let hasher_a = AHasher::default();
        assert_ne!(0, hasher_a.buffer[0]);
        assert_ne!(0, hasher_a.buffer[1]);
        assert_ne!(hasher_a.buffer[0], hasher_a.buffer[1]);
        let hasher_b = AHasher::default();
        assert_eq!(hasher_a.buffer[0], hasher_b.buffer[0]);
        assert_eq!(hasher_a.buffer[1], hasher_b.buffer[1]);
    }

    #[test]
    fn test_hash() {
        let mut result: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value: [u64; 2] = [1 << 32, 0xFEDCBA9876543210];
        result = aeshash(value.convert(), result.convert()).convert();
        result = aeshash(result.convert(), result.convert()).convert();
        let mut result2: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value2: [u64; 2] = [1, 0xFEDCBA9876543210];
        result2 = aeshash(value2.convert(), result2.convert()).convert();
        result2 = aeshash(result2.convert(), result.convert()).convert();
        let result: [u8; 16] = result.convert();
        let result2: [u8; 16] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::fmt;
use std::fmt::{ Display, Formatter };
use std::io;
use std::io::{ BufReader, ErrorKind, Read };
use std::result::Result;

/// Wrapper for a Reader, providing convenience methods to read the stream bit-by-bit.
///
/// # Examples
///
/// extern crate compression;
///
/// use compression::bitreader::BitReader;
/// use std::fs::File;
///
/// let f = try!(File::open("filename"));
/// let mut br = BitReader::new(f);
/// let byte: u8 = br.read_u8().unwrap();
#[derive(Debug)]
pub struct BitReader<R: Read> {
	inner: BufReader<R>,
	bit_pos: u8,
	current_byte: Option<u8>,
}

impl<R: Read> BitReader<R> {
	pub fn new(inner: R) -> BitReader<R> {
		BitReader{
			inner: BufReader::new(inner),
			bit_pos: 0,
			current_byte: None,
		}
	}

	fn read_exact(&mut self, mut buf: &mut [u8]) -> io::Result<()> {
		while !buf.is_empty() {
			match self.inner.read(buf) {
				Ok(0) => break,
				Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
				Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
				Err(e) => return Err(e),
			}
		}
		if !buf.is_empty() {
			Err(io::Error::new(ErrorKind::Other, "failed to fill whole buffer"))
		} else {
			Ok(())
		}
	}

	pub fn read_u8(&mut self) -> Result<u8, BitReaderError> {
		let mut buf = &mut [0u8];

		match (self.current_byte, self.read_exact(buf)) {
			(Some(byte), Ok(())) => {
				self.current_byte = Some(buf[0]);
				Ok((byte >> self.bit_pos) | (buf[0] << (8 - self.bit_pos)))
			},
			(None, Ok(())) => Ok(buf[0]),
			(_, Err(_)) => Err(BitReaderError::Unspecified),
		}
	}

	pub fn read_u16(&mut self) -> Result<u16, BitReaderError> {
		let mut buf = &mut [0u8; 2];

		match (self.current_byte, self.read_exact(buf)) {
			(Some(byte), Ok(())) => {
				self.current_byte = Some(buf[1]);
				Ok(((byte as u16) >> self.bit_pos) | ((buf[0] as u16) << (8 - self.bit_pos)) | ((buf[1] as u16) << (16 - self.bit_pos)))
			},
			(None, Ok(())) => Ok(((buf[1] as u16) << 8) | (buf[0] as u16)),
			(_, Err(_)) => Err(BitReaderError::Unspecified),
		}
	}

	pub fn read_u32(&mut self) -> Result<u32, BitReaderError> {
		match (self.read_u16(), self.read_u16()) {
			(Ok(my_u16_0), Ok(my_u16_1)) => Ok((my_u16_1 as u32) << 16 | (my_u16_0 as u32)),
			(_, _) => Err(BitReaderError::Unspecified),
		}
	}

	pub fn read_bit(&mut self) -> Result<bool, BitReaderError> {
		match (self.current_byte, self.bit_pos) {
			(Some(byte), bit_pos) => {
				self.bit_pos = (self.bit_pos + 1) % 8;
				if self.bit_pos == 0 {
					self.current_byte = None;
				}
				Ok(byte >> bit_pos & 1 == 1)
			},
			(None, _) => {
				let mut buf = &mut [0u8; 1];
				match self.read_exact(buf) {
					Ok(()) => {
						self.current_byte = Some(buf[0]);
						self.bit_pos = 1;
						Ok(buf[0] & 1 == 1)
					},
					Err(_) => Err(BitReaderError::Unspecified),
				}
			}
		}
	}

	pub fn read_n_bits(&mut self, n: usize) -> Result<Vec<bool>, BitReaderError> {
		unimplemented!()
	}

	pub fn read_zero_terminated_string(&mut self) -> Result<Vec<u8>, BitReaderError> {
		let mut my_string = Vec::with_capacity(16);

		loop {
			match self.read_u8() {
				Ok(0) => break,
				Ok(byte) => my_string.push(byte),
				Err(_) => return Err(BitReaderError::Unspecified),
			}
		}

		Ok(my_string)
	}

	pub fn read_fixed_length_string(&mut self, len: usize) -> Result<Vec<u8>, BitReaderError> {
		let mut my_string = Vec::with_capacity(16);

		for _ in 0..len {
			match self.read_u8() {
				Ok(byte) => my_string.push(byte),
				Err(_) => return Err(BitReaderError::Unspecified),
			}
		}

		Ok(my_string)
	}
}

#[derive(Debug, Clone, PartialEq)]
pub enum BitReaderError {
	Unspecified
}

impl Display for BitReaderError {
	fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
		fmt.write_str(self.description())
	}
}

impl Error for BitReaderError {
	fn description(&self) -> &str {
		match self {
			_ => "Generic error",
		}
	}
}

mod tests {
	#[test]
	fn should_read_one_u8() {
		use super::*;
		use std::io::{ Cursor };

		let expected = 0x1f;
		let mut br = BitReader::new(Cursor::new(vec![expected, 0x8b]));

		match br.read_u8() {
			Ok(byte) => assert_eq!(expected, byte),
			Err(_) => panic!("Should have read one byte"),
		}
	}

	#[test]
	fn should_read_two_u8() {
		use super::*;
		use std::io::{ Cursor };

		let (expected0, expected1) = (0x1f, 0x8b);
		let mut br = BitReader::new(Cursor::new(vec![expected0, expected1]));

		match (br.read_u8(), br.read_u8()) {
			(Ok(my_u8_0), Ok(my_u8_1)) => assert_eq!((expected0, expected1), (my_u8_0, my_u8_1)),
			_ => panic!("Should have read two bytes"),
		}
	}

	#[test]
	fn should_read_one_u16() {
		use super::*;
		use std::io::{ Cursor };

		let expected = 0x8b1f;
		let mut br = BitReader::new(Cursor::new(vec![(expected & 0x00FF) as u8, (expected >> 8) as u8]));

		match br.read_u16() {
			Ok(my_u16) => assert_eq!(expected, my_u16),
			_ => panic!("Should have read one u16"),
		}
	}

	#[test]
	fn should_read_one_set_bit() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![3]));

		match br.read_bit() {
			Ok(my_bit) => assert!(my_bit),
			_ => panic!("Should have read one set bit"),
		}
	}

	#[test]
	fn should_read_some_bits() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![134, 1]));

		match (br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit(),
		       br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit()) {
			(Ok(my_bit_0), Ok(my_bit_1), Ok(my_bit_2), Ok(my_bit_3), Ok(my_bit_4),
			 Ok(my_bit_5), Ok(my_bit_6), Ok(my_bit_7), Ok(my_bit_8), Ok(my_bit_9)) => {
				assert!(!my_bit_0);
				assert!(my_bit_1);
				assert!(my_bit_2);
				assert!(!my_bit_3);
				assert!(!my_bit_4);
				assert!(!my_bit_5);
				assert!(!my_bit_6);
				assert!(my_bit_7);
				assert!(my_bit_8);
				assert!(!my_bit_9);
			},
			_ => panic!("Should have read 10 bits"),
		}
	}

	#[test]
	fn should_read_u8_after_bit() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0b10001101, 0b00010101]));

		match (br.read_bit(), br.read_u8()) {
			(Ok(my_bit), Ok(my_u8)) => {
				assert!(my_bit);
				assert_eq!(0b11000110, my_u8);
			},
			_ => panic!("Should have read one set bit and one u8"),
		}
	}

	#[test]
	fn should_read_u16_after_bit() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0b10001101, 0b00010101, 0b00010101]));

		match (br.read_bit(), br.read_u16()) {
			(Ok(my_bit), Ok(my_u16)) => {
				assert!(my_bit);
				assert_eq!(0b1000101011000110, my_u16);
			},
			_ => panic!("Should have read one set bit and one u16"),
		}
	}

	#[test]
	fn should_read_u32() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0x8e, 0x30, 0x04, 0x56]));

		match br.read_u32() {
			Ok(my_u32) => {
				assert_eq!(0x5604308e, my_u32);
			},
			_ => panic!("Should have read one u32"),
		}
	}

	#[test]
	fn should_read_zero_terminated_string() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
		]));

		match br.read_zero_terminated_string() {
			Ok(my_vec) => assert_eq!("xxxxxyyyyy.txt", &(String::from_utf8(my_vec).unwrap())),
			_ => panic!("Should have read zero-terminated string"),
		}
	}

	#[test]
	fn should_read_fixed_length_string() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
		]));

		match br.read_fixed_length_string(14) {
			Ok(my_vec) => assert_eq!("xxxxxyyyyy.txt", &(String::from_utf8(my_vec).unwrap())),
			_ => panic!("Should have read fixed-length string"),
		}
	}

	#[test]
	fn should_read_eight_bits() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0x08]));

		match br.read_n_bits(8) {
			Ok(bits) => assert_eq!(vec![false, false, false, false, true, false, false, false], bits),
			_ => panic!("Should have read 8 bits"),
		}
	}
}
<commit_msg>MOD: read n bits<commit_after>use std::error::Error;
use std::fmt;
use std::fmt::{ Display, Formatter };
use std::io;
use std::io::{ BufReader, ErrorKind, Read };
use std::result::Result;

/// Wrapper for a Reader, providing convenience methods to read the stream bit-by-bit.
///
/// # Examples
///
/// extern crate compression;
///
/// use compression::bitreader::BitReader;
/// use std::fs::File;
///
/// let f = try!(File::open("filename"));
/// let mut br = BitReader::new(f);
/// let byte: u8 = br.read_u8().unwrap();
#[derive(Debug)]
pub struct BitReader<R: Read> {
	inner: BufReader<R>,
	bit_pos: u8,
	current_byte: Option<u8>,
}

impl<R: Read> BitReader<R> {
	pub fn new(inner: R) -> BitReader<R> {
		BitReader{
			inner: BufReader::new(inner),
			bit_pos: 0,
			current_byte: None,
		}
	}

	fn read_exact(&mut self, mut buf: &mut [u8]) -> io::Result<()> {
		while !buf.is_empty() {
			match self.inner.read(buf) {
				Ok(0) => break,
				Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
				Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
				Err(e) => return Err(e),
			}
		}
		if !buf.is_empty() {
			Err(io::Error::new(ErrorKind::Other, "failed to fill whole buffer"))
		} else {
			Ok(())
		}
	}

	pub fn read_u8(&mut self) -> Result<u8, BitReaderError> {
		let mut buf = &mut [0u8];

		match (self.current_byte, self.read_exact(buf)) {
			(Some(byte), Ok(())) => {
				self.current_byte = Some(buf[0]);
				Ok((byte >> self.bit_pos) | (buf[0] << (8 - self.bit_pos)))
			},
			(None, Ok(())) => Ok(buf[0]),
			(_, Err(_)) => Err(BitReaderError::Unspecified),
		}
	}

	pub fn read_u16(&mut self) -> Result<u16, BitReaderError> {
		let mut buf = &mut [0u8; 2];

		match (self.current_byte, self.read_exact(buf)) {
			(Some(byte), Ok(())) => {
				self.current_byte = Some(buf[1]);
				Ok(((byte as u16) >> self.bit_pos) | ((buf[0] as u16) << (8 - self.bit_pos)) | ((buf[1] as u16) << (16 - self.bit_pos)))
			},
			(None, Ok(())) => Ok(((buf[1] as u16) << 8) | (buf[0] as u16)),
			(_, Err(_)) => Err(BitReaderError::Unspecified),
		}
	}

	pub fn read_u32(&mut self) -> Result<u32, BitReaderError> {
		match (self.read_u16(), self.read_u16()) {
			(Ok(my_u16_0), Ok(my_u16_1)) => Ok((my_u16_1 as u32) << 16 | (my_u16_0 as u32)),
			(_, _) => Err(BitReaderError::Unspecified),
		}
	}

	pub fn read_bit(&mut self) -> Result<bool, BitReaderError> {
		match (self.current_byte, self.bit_pos) {
			(Some(byte), bit_pos) => {
				self.bit_pos = (self.bit_pos + 1) % 8;
				if self.bit_pos == 0 {
					self.current_byte = None;
				}
				Ok(byte >> bit_pos & 1 == 1)
			},
			(None, _) => {
				let mut buf = &mut [0u8; 1];
				match self.read_exact(buf) {
					Ok(()) => {
						self.current_byte = Some(buf[0]);
						self.bit_pos = 1;
						Ok(buf[0] & 1 == 1)
					},
					Err(_) => Err(BitReaderError::Unspecified),
				}
			}
		}
	}

	pub fn read_n_bits(&mut self, n: usize) -> Result<Vec<bool>, BitReaderError> {
		let mut v = Vec::with_capacity(n);

		for _ in 0..n {
			match self.read_bit() {
				Ok(bit) => v.push(bit),
				Err(_) => return Err(BitReaderError::Unspecified),
			}
		}

		Ok(v)
	}

	pub fn read_zero_terminated_string(&mut self) -> Result<Vec<u8>, BitReaderError> {
		let mut my_string = Vec::with_capacity(16);

		loop {
			match self.read_u8() {
				Ok(0) => break,
				Ok(byte) => my_string.push(byte),
				Err(_) => return Err(BitReaderError::Unspecified),
			}
		}

		Ok(my_string)
	}

	pub fn read_fixed_length_string(&mut self, len: usize) -> Result<Vec<u8>, BitReaderError> {
		let mut my_string = Vec::with_capacity(16);

		for _ in 0..len {
			match self.read_u8() {
				Ok(byte) => my_string.push(byte),
				Err(_) => return Err(BitReaderError::Unspecified),
			}
		}

		Ok(my_string)
	}
}

#[derive(Debug, Clone, PartialEq)]
pub enum BitReaderError {
	Unspecified
}

impl Display for BitReaderError {
	fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
		fmt.write_str(self.description())
	}
}

impl Error for BitReaderError {
	fn description(&self) -> &str {
		match self {
			_ => "Generic error",
		}
	}
}

mod tests {
	#[test]
	fn should_read_one_u8() {
		use super::*;
		use std::io::{ Cursor };

		let expected = 0x1f;
		let mut br = BitReader::new(Cursor::new(vec![expected, 0x8b]));

		match br.read_u8() {
			Ok(byte) => assert_eq!(expected, byte),
			Err(_) => panic!("Should have read one byte"),
		}
	}

	#[test]
	fn should_read_two_u8() {
		use super::*;
		use std::io::{ Cursor };

		let (expected0, expected1) = (0x1f, 0x8b);
		let mut br = BitReader::new(Cursor::new(vec![expected0, expected1]));

		match (br.read_u8(), br.read_u8()) {
			(Ok(my_u8_0), Ok(my_u8_1)) => assert_eq!((expected0, expected1), (my_u8_0, my_u8_1)),
			_ => panic!("Should have read two bytes"),
		}
	}

	#[test]
	fn should_read_one_u16() {
		use super::*;
		use std::io::{ Cursor };

		let expected = 0x8b1f;
		let mut br = BitReader::new(Cursor::new(vec![(expected & 0x00FF) as u8, (expected >> 8) as u8]));

		match br.read_u16() {
			Ok(my_u16) => assert_eq!(expected, my_u16),
			_ => panic!("Should have read one u16"),
		}
	}

	#[test]
	fn should_read_one_set_bit() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![3]));

		match br.read_bit() {
			Ok(my_bit) => assert!(my_bit),
			_ => panic!("Should have read one set bit"),
		}
	}

	#[test]
	fn should_read_some_bits() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![134, 1]));

		match (br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit(),
		       br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit(), br.read_bit()) {
			(Ok(my_bit_0), Ok(my_bit_1), Ok(my_bit_2), Ok(my_bit_3), Ok(my_bit_4),
			 Ok(my_bit_5), Ok(my_bit_6), Ok(my_bit_7), Ok(my_bit_8), Ok(my_bit_9)) => {
				assert!(!my_bit_0);
				assert!(my_bit_1);
				assert!(my_bit_2);
				assert!(!my_bit_3);
				assert!(!my_bit_4);
				assert!(!my_bit_5);
				assert!(!my_bit_6);
				assert!(my_bit_7);
				assert!(my_bit_8);
				assert!(!my_bit_9);
			},
			_ => panic!("Should have read 10 bits"),
		}
	}

	#[test]
	fn should_read_u8_after_bit() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0b10001101, 0b00010101]));

		match (br.read_bit(), br.read_u8()) {
			(Ok(my_bit), Ok(my_u8)) => {
				assert!(my_bit);
				assert_eq!(0b11000110, my_u8);
			},
			_ => panic!("Should have read one set bit and one u8"),
		}
	}

	#[test]
	fn should_read_u16_after_bit() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0b10001101, 0b00010101, 0b00010101]));

		match (br.read_bit(), br.read_u16()) {
			(Ok(my_bit), Ok(my_u16)) => {
				assert!(my_bit);
				assert_eq!(0b1000101011000110, my_u16);
			},
			_ => panic!("Should have read one set bit and one u16"),
		}
	}

	#[test]
	fn should_read_u32() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0x8e, 0x30, 0x04, 0x56]));

		match br.read_u32() {
			Ok(my_u32) => {
				assert_eq!(0x5604308e, my_u32);
			},
			_ => panic!("Should have read one u32"),
		}
	}

	#[test]
	fn should_read_zero_terminated_string() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
		]));

		match br.read_zero_terminated_string() {
			Ok(my_vec) => assert_eq!("xxxxxyyyyy.txt", &(String::from_utf8(my_vec).unwrap())),
			_ => panic!("Should have read zero-terminated string"),
		}
	}

	#[test]
	fn should_read_fixed_length_string() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
			0x78, 0x78, 0x78, 0x78, 0x78,
			0x79, 0x79, 0x79, 0x79, 0x79,
			0x2e, 0x74, 0x78, 0x74, 0x00,
		]));

		match br.read_fixed_length_string(14) {
			Ok(my_vec) => assert_eq!("xxxxxyyyyy.txt", &(String::from_utf8(my_vec).unwrap())),
			_ => panic!("Should have read fixed-length string"),
		}
	}

	#[test]
	fn should_read_eight_bits() {
		use super::*;
		use std::io::{ Cursor };

		let mut br = BitReader::new(Cursor::new(vec![0x08]));

		match br.read_n_bits(8) {
			Ok(bits) => assert_eq!(vec![false, false, false, true, false, false, false, false], bits),
			_ => panic!("Should have read 8 bits"),
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Documentation generation for rustbuild.
//!
//! This module implements generation for all bits and pieces of documentation
//! for the Rust project. This notably includes suites like the rust book, the
//! nomicon, standalone documentation, etc.
//!
//! Everything here is basically just a shim around calling either `rustbook` or
//! `rustdoc`.

use std::fs::{self, File};
use std::io::prelude::*;
use std::io;
use std::path::Path;
use std::process::Command;

use {Build, Compiler, Mode};
use util::{cp_r, symlink_dir};
use build_helper::up_to_date;

/// Invoke `rustbook` as compiled in `stage` for `target` for the doc book
/// `name` into the `out` path.
///
/// This will not actually generate any documentation if the documentation has
/// already been generated.
pub fn rustbook(build: &Build, target: &str, name: &str) {
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));

    let out = out.join(name);
    let compiler = Compiler::new(0, &build.config.build);
    let src = build.src.join("src/doc").join(name);
    let index = out.join("index.html");
    let rustbook = build.tool(&compiler, "rustbook");
    if up_to_date(&src, &index) && up_to_date(&rustbook, &index) {
        return
    }
    println!("Rustbook ({}) - {}", target, name);
    let _ = fs::remove_dir_all(&out);
    build.run(build.tool_cmd(&compiler, "rustbook")
                   .arg("build")
                   .arg(&src)
                   .arg("-d")
                   .arg(out));
}

/// Build the book and associated stuff.
///
/// We need to build:
///
/// * Book (first edition)
/// * Book (second edition)
/// * Index page
/// * Redirect pages
pub fn book(build: &Build, target: &str, name: &str) {
    // build book first edition
    rustbook(build, target, &format!("{}/first-edition", name));

    // build book second edition
    rustbook(build, target, &format!("{}/second-edition", name));

    // build the index page
    let index = format!("{}/index.md", name);
    println!("Documenting book index ({})", target);
    invoke_rustdoc(build, target, &index);

    // build the redirect pages
    println!("Documenting book redirect pages ({})", target);
    for file in t!(fs::read_dir(build.src.join("src/doc/book/redirects"))) {
        let file = t!(file);
        let path = file.path();
        let path = path.to_str().unwrap();

        invoke_rustdoc(build, target, path);
    }
}

fn invoke_rustdoc(build: &Build, target: &str, markdown: &str) {
    let out = build.doc_out(target);

    let compiler = Compiler::new(0, &build.config.build);

    let path = build.src.join("src/doc").join(markdown);

    let rustdoc = build.rustdoc(&compiler);

    let favicon = build.src.join("src/doc/favicon.inc");
    let footer = build.src.join("src/doc/footer.inc");

    let version_input = build.src.join("src/doc/version_info.html.template");
    let version_info = out.join("version_info.html");

    if !up_to_date(&version_input, &version_info) {
        let mut info = String::new();
        t!(t!(File::open(&version_input)).read_to_string(&mut info));
        let info = info.replace("VERSION", &build.rust_release())
                       .replace("SHORT_HASH", build.rust_info.sha_short().unwrap_or(""))
                       .replace("STAMP", build.rust_info.sha().unwrap_or(""));
        t!(t!(File::create(&version_info)).write_all(info.as_bytes()));
    }

    let mut cmd = Command::new(&rustdoc);

    build.add_rustc_lib_path(&compiler, &mut cmd);

    let out = out.join("book");

    t!(fs::copy(build.src.join("src/doc/rust.css"), out.join("rust.css")));

    cmd.arg("--html-after-content").arg(&footer)
        .arg("--html-before-content").arg(&version_info)
        .arg("--html-in-header").arg(&favicon)
        .arg("--markdown-playground-url")
        .arg("https://play.rust-lang.org/")
        .arg("-o").arg(&out)
        .arg(&path)
        .arg("--markdown-css")
        .arg("rust.css");

    build.run(&mut cmd);
}

/// Generates all standalone documentation as compiled by the rustdoc in `stage`
/// for the `target` into `out`.
///
/// This will list all of `src/doc` looking for markdown files and appropriately
/// perform transformations like substituting `VERSION`, `SHORT_HASH`, and
/// `STAMP` alongw ith providing the various header/footer HTML we've cutomized.
///
/// In the end, this is just a glorified wrapper around rustdoc!
pub fn standalone(build: &Build, target: &str) {
    println!("Documenting standalone ({})", target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));

    let compiler = Compiler::new(0, &build.config.build);

    let favicon = build.src.join("src/doc/favicon.inc");
    let footer = build.src.join("src/doc/footer.inc");
    let full_toc = build.src.join("src/doc/full-toc.inc");
    t!(fs::copy(build.src.join("src/doc/rust.css"), out.join("rust.css")));

    let version_input = build.src.join("src/doc/version_info.html.template");
    let version_info = out.join("version_info.html");

    if !up_to_date(&version_input, &version_info) {
        let mut info = String::new();
        t!(t!(File::open(&version_input)).read_to_string(&mut info));
        let info = info.replace("VERSION", &build.rust_release())
                       .replace("SHORT_HASH", build.rust_info.sha_short().unwrap_or(""))
                       .replace("STAMP", build.rust_info.sha().unwrap_or(""));
        t!(t!(File::create(&version_info)).write_all(info.as_bytes()));
    }

    for file in t!(fs::read_dir(build.src.join("src/doc"))) {
        let file = t!(file);
        let path = file.path();
        let filename = path.file_name().unwrap().to_str().unwrap();
        if !filename.ends_with(".md") || filename == "README.md" {
            continue
        }

        let html = out.join(filename).with_extension("html");
        let rustdoc = build.rustdoc(&compiler);
        if up_to_date(&path, &html) &&
           up_to_date(&footer, &html) &&
           up_to_date(&favicon, &html) &&
           up_to_date(&full_toc, &html) &&
           up_to_date(&version_info, &html) &&
           up_to_date(&rustdoc, &html) {
            continue
        }

        let mut cmd = Command::new(&rustdoc);
        build.add_rustc_lib_path(&compiler, &mut cmd);
        cmd.arg("--html-after-content").arg(&footer)
           .arg("--html-before-content").arg(&version_info)
           .arg("--html-in-header").arg(&favicon)
           .arg("--markdown-playground-url")
           .arg("https://play.rust-lang.org/")
           .arg("-o").arg(&out)
           .arg(&path);

        if filename == "not_found.md" {
            cmd.arg("--markdown-no-toc")
               .arg("--markdown-css")
               .arg("https://doc.rust-lang.org/rust.css");
        } else {
            cmd.arg("--markdown-css").arg("rust.css");
        }
        build.run(&mut cmd);
    }
}

/// Compile all standard library documentation.
///
/// This will generate all documentation for the standard library and its
/// dependencies. This is largely just a wrapper around `cargo doc`.
pub fn std(build: &Build, stage: u32, target: &str) {
    println!("Documenting stage{} std ({})", stage, target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(stage, &build.config.build);
    let compiler = if build.force_use_stage1(&compiler, target) {
        Compiler::new(1, compiler.host)
    } else {
        compiler
    };
    let out_dir = build.stage_out(&compiler, Mode::Libstd)
                       .join(target).join("doc");
    let rustdoc = build.rustdoc(&compiler);

    // Here what we're doing is creating a *symlink* (directory junction on
    // Windows) to the final output location. This is not done as an
    // optimization but rather for correctness. We've got three trees of
    // documentation, one for std, one for test, and one for rustc. It's then
    // our job to merge them all together.
    //
    // Unfortunately rustbuild doesn't know nearly as well how to merge doc
    // trees as rustdoc does itself, so instead of actually having three
    // separate trees we just have rustdoc output to the same location across
    // all of them.
    //
    // This way rustdoc generates output directly into the output, and rustdoc
    // will also directly handle merging.
    let my_out = build.crate_doc_out(target);
    build.clear_if_dirty(&my_out, &rustdoc);
    t!(symlink_dir_force(&my_out, &out_dir));

    let mut cargo = build.cargo(&compiler, Mode::Libstd, target, "doc");
    cargo.arg("--manifest-path")
         .arg(build.src.join("src/libstd/Cargo.toml"))
         .arg("--features").arg(build.std_features());

    // We don't want to build docs for internal std dependencies unless
    // in compiler-docs mode. When not in that mode, we whitelist the crates
    // for which docs must be built.
    if !build.config.compiler_docs {
        cargo.arg("--no-deps");
        for krate in &["alloc", "collections", "core", "std", "std_unicode"] {
            cargo.arg("-p").arg(krate);
            // Create all crate output directories first to make sure rustdoc uses
            // relative links.
            // FIXME: Cargo should probably do this itself.
            t!(fs::create_dir_all(out_dir.join(krate)));
        }
    }

    build.run(&mut cargo);
    cp_r(&my_out, &out);
}

/// Compile all libtest documentation.
///
/// This will generate all documentation for libtest and its dependencies. This
/// is largely just a wrapper around `cargo doc`.
pub fn test(build: &Build, stage: u32, target: &str) {
    println!("Documenting stage{} test ({})", stage, target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(stage, &build.config.build);
    let compiler = if build.force_use_stage1(&compiler, target) {
        Compiler::new(1, compiler.host)
    } else {
        compiler
    };
    let out_dir = build.stage_out(&compiler, Mode::Libtest)
                       .join(target).join("doc");
    let rustdoc = build.rustdoc(&compiler);

    // See docs in std above for why we symlink
    let my_out = build.crate_doc_out(target);
    build.clear_if_dirty(&my_out, &rustdoc);
    t!(symlink_dir_force(&my_out, &out_dir));

    let mut cargo = build.cargo(&compiler, Mode::Libtest, target, "doc");
    cargo.arg("--manifest-path")
         .arg(build.src.join("src/libtest/Cargo.toml"));
    build.run(&mut cargo);
    cp_r(&my_out, &out);
}

/// Generate all compiler documentation.
///
/// This will generate all documentation for the compiler libraries and their
/// dependencies. This is largely just a wrapper around `cargo doc`.
pub fn rustc(build: &Build, stage: u32, target: &str) {
    println!("Documenting stage{} compiler ({})", stage, target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(stage, &build.config.build);
    let compiler = if build.force_use_stage1(&compiler, target) {
        Compiler::new(1, compiler.host)
    } else {
        compiler
    };
    let out_dir = build.stage_out(&compiler, Mode::Librustc)
                       .join(target).join("doc");
    let rustdoc = build.rustdoc(&compiler);

    // See docs in std above for why we symlink
    let my_out = build.crate_doc_out(target);
    build.clear_if_dirty(&my_out, &rustdoc);
    t!(symlink_dir_force(&my_out, &out_dir));

    let mut cargo = build.cargo(&compiler, Mode::Librustc, target, "doc");
    cargo.arg("--manifest-path")
         .arg(build.src.join("src/rustc/Cargo.toml"))
         .arg("--features").arg(build.rustc_features());

    if build.config.compiler_docs {
        // src/rustc/Cargo.toml contains bin crates called rustc and rustdoc
        // which would otherwise overwrite the docs for the real rustc and
        // rustdoc lib crates.
        cargo.arg("-p").arg("rustc_driver")
             .arg("-p").arg("rustdoc");
    } else {
        // Like with libstd above if compiler docs aren't enabled then we're not
        // documenting internal dependencies, so we have a whitelist.
        cargo.arg("--no-deps");
        for krate in &["proc_macro"] {
            cargo.arg("-p").arg(krate);
        }
    }

    build.run(&mut cargo);
    cp_r(&my_out, &out);
}

/// Generates the HTML rendered error-index by running the
/// `error_index_generator` tool.
pub fn error_index(build: &Build, target: &str) {
    println!("Documenting error index ({})", target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(0, &build.config.build);
    let mut index = build.tool_cmd(&compiler, "error_index_generator");
    index.arg("html");
    index.arg(out.join("error-index.html"));

    // FIXME: shouldn't have to pass this env var
    index.env("CFG_BUILD", &build.config.build);

    build.run(&mut index);
}

fn symlink_dir_force(src: &Path, dst: &Path) -> io::Result<()> {
    if let Ok(m) = fs::symlink_metadata(dst) {
        if m.file_type().is_dir() {
            try!(fs::remove_dir_all(dst));
        } else {
            // handle directory junctions on windows by falling back to
            // `remove_dir`.
            try!(fs::remove_file(dst).or_else(|_| {
                fs::remove_dir(dst)
            }));
        }
    }

    symlink_dir(src, dst)
}
<commit_msg>fix whitespace<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Documentation generation for rustbuild.
//!
//! This module implements generation for all bits and pieces of documentation
//! for the Rust project. This notably includes suites like the rust book, the
//! nomicon, standalone documentation, etc.
//!
//! Everything here is basically just a shim around calling either `rustbook` or
//! `rustdoc`.

use std::fs::{self, File};
use std::io::prelude::*;
use std::io;
use std::path::Path;
use std::process::Command;

use {Build, Compiler, Mode};
use util::{cp_r, symlink_dir};
use build_helper::up_to_date;

/// Invoke `rustbook` as compiled in `stage` for `target` for the doc book
/// `name` into the `out` path.
///
/// This will not actually generate any documentation if the documentation has
/// already been generated.
pub fn rustbook(build: &Build, target: &str, name: &str) {
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));

    let out = out.join(name);
    let compiler = Compiler::new(0, &build.config.build);
    let src = build.src.join("src/doc").join(name);
    let index = out.join("index.html");
    let rustbook = build.tool(&compiler, "rustbook");
    if up_to_date(&src, &index) && up_to_date(&rustbook, &index) {
        return
    }
    println!("Rustbook ({}) - {}", target, name);
    let _ = fs::remove_dir_all(&out);
    build.run(build.tool_cmd(&compiler, "rustbook")
                   .arg("build")
                   .arg(&src)
                   .arg("-d")
                   .arg(out));
}

/// Build the book and associated stuff.
///
/// We need to build:
///
/// * Book (first edition)
/// * Book (second edition)
/// * Index page
/// * Redirect pages
pub fn book(build: &Build, target: &str, name: &str) {
    // build book first edition
    rustbook(build, target, &format!("{}/first-edition", name));

    // build book second edition
    rustbook(build, target, &format!("{}/second-edition", name));

    // build the index page
    let index = format!("{}/index.md", name);
    println!("Documenting book index ({})", target);
    invoke_rustdoc(build, target, &index);

    // build the redirect pages
    println!("Documenting book redirect pages ({})", target);
    for file in t!(fs::read_dir(build.src.join("src/doc/book/redirects"))) {
        let file = t!(file);
        let path = file.path();
        let path = path.to_str().unwrap();

        invoke_rustdoc(build, target, path);
    }
}

fn invoke_rustdoc(build: &Build, target: &str, markdown: &str) {
    let out = build.doc_out(target);

    let compiler = Compiler::new(0, &build.config.build);

    let path = build.src.join("src/doc").join(markdown);

    let rustdoc = build.rustdoc(&compiler);

    let favicon = build.src.join("src/doc/favicon.inc");
    let footer = build.src.join("src/doc/footer.inc");

    let version_input = build.src.join("src/doc/version_info.html.template");
    let version_info = out.join("version_info.html");

    if !up_to_date(&version_input, &version_info) {
        let mut info = String::new();
        t!(t!(File::open(&version_input)).read_to_string(&mut info));
        let info = info.replace("VERSION", &build.rust_release())
                       .replace("SHORT_HASH", build.rust_info.sha_short().unwrap_or(""))
                       .replace("STAMP", build.rust_info.sha().unwrap_or(""));
        t!(t!(File::create(&version_info)).write_all(info.as_bytes()));
    }

    let mut cmd = Command::new(&rustdoc);

    build.add_rustc_lib_path(&compiler, &mut cmd);

    let out = out.join("book");

    t!(fs::copy(build.src.join("src/doc/rust.css"), out.join("rust.css")));

    cmd.arg("--html-after-content").arg(&footer)
        .arg("--html-before-content").arg(&version_info)
        .arg("--html-in-header").arg(&favicon)
        .arg("--markdown-playground-url")
        .arg("https://play.rust-lang.org/")
        .arg("-o").arg(&out)
        .arg(&path)
        .arg("--markdown-css")
        .arg("rust.css");

    build.run(&mut cmd);
}

/// Generates all standalone documentation as compiled by the rustdoc in `stage`
/// for the `target` into `out`.
///
/// This will list all of `src/doc` looking for markdown files and appropriately
/// perform transformations like substituting `VERSION`, `SHORT_HASH`, and
/// `STAMP` alongw ith providing the various header/footer HTML we've cutomized.
///
/// In the end, this is just a glorified wrapper around rustdoc!
pub fn standalone(build: &Build, target: &str) {
    println!("Documenting standalone ({})", target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));

    let compiler = Compiler::new(0, &build.config.build);

    let favicon = build.src.join("src/doc/favicon.inc");
    let footer = build.src.join("src/doc/footer.inc");
    let full_toc = build.src.join("src/doc/full-toc.inc");
    t!(fs::copy(build.src.join("src/doc/rust.css"), out.join("rust.css")));

    let version_input = build.src.join("src/doc/version_info.html.template");
    let version_info = out.join("version_info.html");

    if !up_to_date(&version_input, &version_info) {
        let mut info = String::new();
        t!(t!(File::open(&version_input)).read_to_string(&mut info));
        let info = info.replace("VERSION", &build.rust_release())
                       .replace("SHORT_HASH", build.rust_info.sha_short().unwrap_or(""))
                       .replace("STAMP", build.rust_info.sha().unwrap_or(""));
        t!(t!(File::create(&version_info)).write_all(info.as_bytes()));
    }

    for file in t!(fs::read_dir(build.src.join("src/doc"))) {
        let file = t!(file);
        let path = file.path();
        let filename = path.file_name().unwrap().to_str().unwrap();
        if !filename.ends_with(".md") || filename == "README.md" {
            continue
        }

        let html = out.join(filename).with_extension("html");
        let rustdoc = build.rustdoc(&compiler);
        if up_to_date(&path, &html) &&
           up_to_date(&footer, &html) &&
           up_to_date(&favicon, &html) &&
           up_to_date(&full_toc, &html) &&
           up_to_date(&version_info, &html) &&
           up_to_date(&rustdoc, &html) {
            continue
        }

        let mut cmd = Command::new(&rustdoc);
        build.add_rustc_lib_path(&compiler, &mut cmd);
        cmd.arg("--html-after-content").arg(&footer)
           .arg("--html-before-content").arg(&version_info)
           .arg("--html-in-header").arg(&favicon)
           .arg("--markdown-playground-url")
           .arg("https://play.rust-lang.org/")
           .arg("-o").arg(&out)
           .arg(&path);

        if filename == "not_found.md" {
            cmd.arg("--markdown-no-toc")
               .arg("--markdown-css")
               .arg("https://doc.rust-lang.org/rust.css");
        } else {
            cmd.arg("--markdown-css").arg("rust.css");
        }
        build.run(&mut cmd);
    }
}

/// Compile all standard library documentation.
///
/// This will generate all documentation for the standard library and its
/// dependencies. This is largely just a wrapper around `cargo doc`.
pub fn std(build: &Build, stage: u32, target: &str) {
    println!("Documenting stage{} std ({})", stage, target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(stage, &build.config.build);
    let compiler = if build.force_use_stage1(&compiler, target) {
        Compiler::new(1, compiler.host)
    } else {
        compiler
    };
    let out_dir = build.stage_out(&compiler, Mode::Libstd)
                       .join(target).join("doc");
    let rustdoc = build.rustdoc(&compiler);

    // Here what we're doing is creating a *symlink* (directory junction on
    // Windows) to the final output location. This is not done as an
    // optimization but rather for correctness. We've got three trees of
    // documentation, one for std, one for test, and one for rustc. It's then
    // our job to merge them all together.
    //
    // Unfortunately rustbuild doesn't know nearly as well how to merge doc
    // trees as rustdoc does itself, so instead of actually having three
    // separate trees we just have rustdoc output to the same location across
    // all of them.
    //
    // This way rustdoc generates output directly into the output, and rustdoc
    // will also directly handle merging.
    let my_out = build.crate_doc_out(target);
    build.clear_if_dirty(&my_out, &rustdoc);
    t!(symlink_dir_force(&my_out, &out_dir));

    let mut cargo = build.cargo(&compiler, Mode::Libstd, target, "doc");
    cargo.arg("--manifest-path")
         .arg(build.src.join("src/libstd/Cargo.toml"))
         .arg("--features").arg(build.std_features());

    // We don't want to build docs for internal std dependencies unless
    // in compiler-docs mode. When not in that mode, we whitelist the crates
    // for which docs must be built.
    if !build.config.compiler_docs {
        cargo.arg("--no-deps");
        for krate in &["alloc", "collections", "core", "std", "std_unicode"] {
            cargo.arg("-p").arg(krate);
            // Create all crate output directories first to make sure rustdoc uses
            // relative links.
            // FIXME: Cargo should probably do this itself.
            t!(fs::create_dir_all(out_dir.join(krate)));
        }
    }

    build.run(&mut cargo);
    cp_r(&my_out, &out);
}

/// Compile all libtest documentation.
///
/// This will generate all documentation for libtest and its dependencies. This
/// is largely just a wrapper around `cargo doc`.
pub fn test(build: &Build, stage: u32, target: &str) {
    println!("Documenting stage{} test ({})", stage, target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(stage, &build.config.build);
    let compiler = if build.force_use_stage1(&compiler, target) {
        Compiler::new(1, compiler.host)
    } else {
        compiler
    };
    let out_dir = build.stage_out(&compiler, Mode::Libtest)
                       .join(target).join("doc");
    let rustdoc = build.rustdoc(&compiler);

    // See docs in std above for why we symlink
    let my_out = build.crate_doc_out(target);
    build.clear_if_dirty(&my_out, &rustdoc);
    t!(symlink_dir_force(&my_out, &out_dir));

    let mut cargo = build.cargo(&compiler, Mode::Libtest, target, "doc");
    cargo.arg("--manifest-path")
         .arg(build.src.join("src/libtest/Cargo.toml"));
    build.run(&mut cargo);
    cp_r(&my_out, &out);
}

/// Generate all compiler documentation.
///
/// This will generate all documentation for the compiler libraries and their
/// dependencies. This is largely just a wrapper around `cargo doc`.
pub fn rustc(build: &Build, stage: u32, target: &str) {
    println!("Documenting stage{} compiler ({})", stage, target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(stage, &build.config.build);
    let compiler = if build.force_use_stage1(&compiler, target) {
        Compiler::new(1, compiler.host)
    } else {
        compiler
    };
    let out_dir = build.stage_out(&compiler, Mode::Librustc)
                       .join(target).join("doc");
    let rustdoc = build.rustdoc(&compiler);

    // See docs in std above for why we symlink
    let my_out = build.crate_doc_out(target);
    build.clear_if_dirty(&my_out, &rustdoc);
    t!(symlink_dir_force(&my_out, &out_dir));

    let mut cargo = build.cargo(&compiler, Mode::Librustc, target, "doc");
    cargo.arg("--manifest-path")
         .arg(build.src.join("src/rustc/Cargo.toml"))
         .arg("--features").arg(build.rustc_features());

    if build.config.compiler_docs {
        // src/rustc/Cargo.toml contains bin crates called rustc and rustdoc
        // which would otherwise overwrite the docs for the real rustc and
        // rustdoc lib crates.
        cargo.arg("-p").arg("rustc_driver")
             .arg("-p").arg("rustdoc");
    } else {
        // Like with libstd above if compiler docs aren't enabled then we're not
        // documenting internal dependencies, so we have a whitelist.
        cargo.arg("--no-deps");
        for krate in &["proc_macro"] {
            cargo.arg("-p").arg(krate);
        }
    }

    build.run(&mut cargo);
    cp_r(&my_out, &out);
}

/// Generates the HTML rendered error-index by running the
/// `error_index_generator` tool.
pub fn error_index(build: &Build, target: &str) {
    println!("Documenting error index ({})", target);
    let out = build.doc_out(target);
    t!(fs::create_dir_all(&out));
    let compiler = Compiler::new(0, &build.config.build);
    let mut index = build.tool_cmd(&compiler, "error_index_generator");
    index.arg("html");
    index.arg(out.join("error-index.html"));

    // FIXME: shouldn't have to pass this env var
    index.env("CFG_BUILD", &build.config.build);

    build.run(&mut index);
}

fn symlink_dir_force(src: &Path, dst: &Path) -> io::Result<()> {
    if let Ok(m) = fs::symlink_metadata(dst) {
        if m.file_type().is_dir() {
            try!(fs::remove_dir_all(dst));
        } else {
            // handle directory junctions on windows by falling back to
            // `remove_dir`.
            try!(fs::remove_file(dst).or_else(|_| {
                fs::remove_dir(dst)
            }));
        }
    }

    symlink_dir(src, dst)
}
<|endoftext|>