<commit_before>use crate::App;
use anyhow::Error;
use dashmap::{DashMap, SharedValue};
use diesel::{pg::upsert::excluded, prelude::*};
use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};

/// crates.io receives a lot of download requests, and we can't execute a write query to the
/// database during each connection for performance reasons. To reduce the write load, this struct
/// collects the pending updates from the current process and writes in batch.
///
/// To avoid locking the whole data structure behind a RwLock, which could potentially delay
/// requests, this uses the dashmap crate. A DashMap has the same public API as an HashMap, but
/// stores the items into `num_cpus()*4` individually locked shards. This approach reduces the
/// likelyhood of a request encountering a locked shard.
///
/// Persisting the download counts in the database also takes advantage of the inner sharding of
/// DashMaps: to avoid locking all the download requests at the same time each iteration only
/// persists a single shard at the time.
///
/// The disadvantage of this approach is that download counts are stored in memory until they're
/// persisted, so it's possible to lose some of them if the process exits ungracefully. While
/// that's far from ideal, the advantage of batching database updates far outweights potentially
/// losing some download counts.
#[derive(Debug)]
pub struct DownloadsCounter {
    /// Inner storage for the download counts.
    inner: DashMap<i32, AtomicUsize>,
    /// Index of the next shard that should be persisted by `persist_next_shard`.
    shard_idx: AtomicUsize,
    /// Number of downloads that are not yet persisted on the database. This is just used as a
    /// metric included in log lines, and it's not guaranteed to be accurate.
    pending_count: AtomicI64,
}

impl DownloadsCounter {
    pub(crate) fn new() -> Self {
        Self {
            inner: DashMap::new(),
            shard_idx: AtomicUsize::new(0),
            pending_count: AtomicI64::new(0),
        }
    }

    pub(crate) fn increment(&self, version_id: i32) {
        self.pending_count.fetch_add(1, Ordering::SeqCst);

        if let Some(counter) = self.inner.get(&version_id) {
            // The version is already recorded in the DashMap, so we don't need to lock the whole
            // shard in write mode. The shard is instead locked in read mode, which allows an
            // unbounded number of readers as long as there are no write locks.
            counter.value().fetch_add(1, Ordering::SeqCst);
        } else {
            // The version is not in the DashMap, so we need to lock the whole shard in write mode
            // and insert the version into it. This has worse performance than the above case.
            self.inner
                .entry(version_id)
                .and_modify(|counter| {
                    // Handle the version being inserted by another thread while we were waiting
                    // for the write lock on the shard.
                    counter.fetch_add(1, Ordering::SeqCst);
                })
                .or_insert_with(|| AtomicUsize::new(1));
        }
    }

    pub fn persist_all_shards(&self, app: &App) -> Result<PersistStats, Error> {
        let conn = app.primary_database.get()?;
        self.persist_all_shards_with_conn(&conn)
    }

    pub fn persist_next_shard(&self, app: &App) -> Result<PersistStats, Error> {
        let conn = app.primary_database.get()?;
        self.persist_next_shard_with_conn(&conn)
    }

    fn persist_all_shards_with_conn(&self, conn: &PgConnection) -> Result<PersistStats, Error> {
        let mut stats = PersistStats::default();
        for shard in self.inner.shards() {
            let shard = std::mem::take(&mut *shard.write());
            stats = stats.merge(self.persist_shard(&conn, shard)?);
        }

        Ok(stats)
    }

    fn persist_next_shard_with_conn(&self, conn: &PgConnection) -> Result<PersistStats, Error> {
        // Replace the next shard in the ring with an empty HashMap (clearing it), and return the
        // previous contents for processing. The fetch_add method wraps around on overflow, so it's
        // fine to keep incrementing it without resetting.
        let shards = self.inner.shards();
        let idx = self.shard_idx.fetch_add(1, Ordering::SeqCst) % shards.len();
        let shard = std::mem::take(&mut *shards[idx].write());

        let mut stats = self.persist_shard(&conn, shard)?;
        stats.shard = Some(idx);
        Ok(stats)
    }

    fn persist_shard(
        &self,
        conn: &PgConnection,
        shard: HashMap<i32, SharedValue<AtomicUsize>>,
    ) -> Result<PersistStats, Error> {
        use crate::schema::version_downloads::dsl::*;

        let mut counted_downloads = 0;
        let mut counted_versions = 0;
        let mut to_insert = Vec::new();
        for (key, atomic) in shard.iter() {
            let count = atomic.get().load(Ordering::SeqCst);
            counted_downloads += count;
            counted_versions += 1;

            to_insert.push((*key, count));
        }

        if !to_insert.is_empty() {
            // The rows we're about to insert need to be sorted to avoid deadlocks when multiple
            // instances of crates.io are running at the same time.
            //
            // In PostgreSQL a transaction modifying a row locks that row until the transaction is
            // committed. Multiple transactions inserting rows into a table could end up
            // deadlocking each other though: PostgreSQL will detect that deadlock, abort one of
            // the transactions and allow the other one to continue. We don't want that to happen,
            // as we'd lose the downloads from the aborted transaction.
            //
            // Ensuring the rows are inserted in a consistent order (in our case by sorting them by
            // the version ID) will prevent deadlocks from occuring. For more information:
            //
            //     https://www.postgresql.org/docs/11/explicit-locking.html#LOCKING-DEADLOCKS
            //
            to_insert.sort_by_key(|(key, _)| *key);

            let to_insert = to_insert
                .into_iter()
                .map(|(key, count)| (version_id.eq(key), downloads.eq(count as i32)))
                .collect::<Vec<_>>();

            diesel::insert_into(version_downloads)
                .values(&to_insert)
                .on_conflict((version_id, date))
                .do_update()
                .set(downloads.eq(downloads + excluded(downloads)))
                .execute(conn)?;
        }

        let old_pending = self
            .pending_count
            .fetch_sub(counted_downloads as i64, Ordering::SeqCst);

        Ok(PersistStats {
            shard: None,
            counted_downloads,
            counted_versions,
            pending_downloads: old_pending - counted_downloads as i64,
        })
    }

    pub fn shards_count(&self) -> usize {
        self.inner.shards().len()
    }
}

#[derive(Debug, Default, Copy, Clone)]
pub struct PersistStats {
    shard: Option<usize>,
    counted_downloads: usize,
    counted_versions: usize,
    pending_downloads: i64,
}

impl PersistStats {
    fn merge(self, other: PersistStats) -> Self {
        Self {
            shard: if self.shard == other.shard {
                other.shard
            } else {
                None
            },
            counted_downloads: self.counted_downloads + other.counted_downloads,
            counted_versions: self.counted_versions + other.counted_versions,
            pending_downloads: other.pending_downloads,
        }
    }

    pub fn log(&self) {
        println!(
            "downloads_counter shard={} counted_versions={} counted_downloads={} pending_downloads={}",
            self.shard.map(|s| s.to_string()).unwrap_or_else(|| "all".into()),
            self.counted_versions,
            self.counted_downloads,
            self.pending_downloads,
        );
    }
}
<commit_msg>downloads_counter: add tests for persist_* methods<commit_after>use crate::App;
use anyhow::Error;
use dashmap::{DashMap, SharedValue};
use diesel::{pg::upsert::excluded, prelude::*};
use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};

/// crates.io receives a lot of download requests, and we can't execute a write query to the
/// database during each connection for performance reasons. To reduce the write load, this struct
/// collects the pending updates from the current process and writes in batch.
///
/// To avoid locking the whole data structure behind a RwLock, which could potentially delay
/// requests, this uses the dashmap crate. A DashMap has the same public API as an HashMap, but
/// stores the items into `num_cpus()*4` individually locked shards. This approach reduces the
/// likelyhood of a request encountering a locked shard.
///
/// Persisting the download counts in the database also takes advantage of the inner sharding of
/// DashMaps: to avoid locking all the download requests at the same time each iteration only
/// persists a single shard at the time.
///
/// The disadvantage of this approach is that download counts are stored in memory until they're
/// persisted, so it's possible to lose some of them if the process exits ungracefully. While
/// that's far from ideal, the advantage of batching database updates far outweights potentially
/// losing some download counts.
#[derive(Debug)]
pub struct DownloadsCounter {
    /// Inner storage for the download counts.
    inner: DashMap<i32, AtomicUsize>,
    /// Index of the next shard that should be persisted by `persist_next_shard`.
    shard_idx: AtomicUsize,
    /// Number of downloads that are not yet persisted on the database. This is just used as a
    /// metric included in log lines, and it's not guaranteed to be accurate.
    pending_count: AtomicI64,
}

impl DownloadsCounter {
    pub(crate) fn new() -> Self {
        Self {
            inner: DashMap::new(),
            shard_idx: AtomicUsize::new(0),
            pending_count: AtomicI64::new(0),
        }
    }

    pub(crate) fn increment(&self, version_id: i32) {
        self.pending_count.fetch_add(1, Ordering::SeqCst);

        if let Some(counter) = self.inner.get(&version_id) {
            // The version is already recorded in the DashMap, so we don't need to lock the whole
            // shard in write mode. The shard is instead locked in read mode, which allows an
            // unbounded number of readers as long as there are no write locks.
            counter.value().fetch_add(1, Ordering::SeqCst);
        } else {
            // The version is not in the DashMap, so we need to lock the whole shard in write mode
            // and insert the version into it. This has worse performance than the above case.
            self.inner
                .entry(version_id)
                .and_modify(|counter| {
                    // Handle the version being inserted by another thread while we were waiting
                    // for the write lock on the shard.
                    counter.fetch_add(1, Ordering::SeqCst);
                })
                .or_insert_with(|| AtomicUsize::new(1));
        }
    }

    pub fn persist_all_shards(&self, app: &App) -> Result<PersistStats, Error> {
        let conn = app.primary_database.get()?;
        self.persist_all_shards_with_conn(&conn)
    }

    pub fn persist_next_shard(&self, app: &App) -> Result<PersistStats, Error> {
        let conn = app.primary_database.get()?;
        self.persist_next_shard_with_conn(&conn)
    }

    fn persist_all_shards_with_conn(&self, conn: &PgConnection) -> Result<PersistStats, Error> {
        let mut stats = PersistStats::default();
        for shard in self.inner.shards() {
            let shard = std::mem::take(&mut *shard.write());
            stats = stats.merge(self.persist_shard(&conn, shard)?);
        }

        Ok(stats)
    }

    fn persist_next_shard_with_conn(&self, conn: &PgConnection) -> Result<PersistStats, Error> {
        // Replace the next shard in the ring with an empty HashMap (clearing it), and return the
        // previous contents for processing. The fetch_add method wraps around on overflow, so it's
        // fine to keep incrementing it without resetting.
        let shards = self.inner.shards();
        let idx = self.shard_idx.fetch_add(1, Ordering::SeqCst) % shards.len();
        let shard = std::mem::take(&mut *shards[idx].write());

        let mut stats = self.persist_shard(&conn, shard)?;
        stats.shard = Some(idx);
        Ok(stats)
    }

    fn persist_shard(
        &self,
        conn: &PgConnection,
        shard: HashMap<i32, SharedValue<AtomicUsize>>,
    ) -> Result<PersistStats, Error> {
        use crate::schema::version_downloads::dsl::*;

        let mut counted_downloads = 0;
        let mut counted_versions = 0;
        let mut to_insert = Vec::new();
        for (key, atomic) in shard.iter() {
            let count = atomic.get().load(Ordering::SeqCst);
            counted_downloads += count;
            counted_versions += 1;

            to_insert.push((*key, count));
        }

        if !to_insert.is_empty() {
            // The rows we're about to insert need to be sorted to avoid deadlocks when multiple
            // instances of crates.io are running at the same time.
            //
            // In PostgreSQL a transaction modifying a row locks that row until the transaction is
            // committed. Multiple transactions inserting rows into a table could end up
            // deadlocking each other though: PostgreSQL will detect that deadlock, abort one of
            // the transactions and allow the other one to continue. We don't want that to happen,
            // as we'd lose the downloads from the aborted transaction.
            //
            // Ensuring the rows are inserted in a consistent order (in our case by sorting them by
            // the version ID) will prevent deadlocks from occuring. For more information:
            //
            //     https://www.postgresql.org/docs/11/explicit-locking.html#LOCKING-DEADLOCKS
            //
            to_insert.sort_by_key(|(key, _)| *key);

            let to_insert = to_insert
                .into_iter()
                .map(|(key, count)| (version_id.eq(key), downloads.eq(count as i32)))
                .collect::<Vec<_>>();

            diesel::insert_into(version_downloads)
                .values(&to_insert)
                .on_conflict((version_id, date))
                .do_update()
                .set(downloads.eq(downloads + excluded(downloads)))
                .execute(conn)?;
        }

        let old_pending = self
            .pending_count
            .fetch_sub(counted_downloads as i64, Ordering::SeqCst);

        Ok(PersistStats {
            shard: None,
            counted_downloads,
            counted_versions,
            pending_downloads: old_pending - counted_downloads as i64,
        })
    }

    pub fn shards_count(&self) -> usize {
        self.inner.shards().len()
    }
}

#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
pub struct PersistStats {
    shard: Option<usize>,
    counted_downloads: usize,
    counted_versions: usize,
    pending_downloads: i64,
}

impl PersistStats {
    fn merge(self, other: PersistStats) -> Self {
        Self {
            shard: if self.shard == other.shard {
                other.shard
            } else {
                None
            },
            counted_downloads: self.counted_downloads + other.counted_downloads,
            counted_versions: self.counted_versions + other.counted_versions,
            pending_downloads: other.pending_downloads,
        }
    }

    pub fn log(&self) {
        println!(
            "downloads_counter shard={} counted_versions={} counted_downloads={} pending_downloads={}",
            self.shard.map(|s| s.to_string()).unwrap_or_else(|| "all".into()),
            self.counted_versions,
            self.counted_downloads,
            self.pending_downloads,
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::{Crate, NewCrate, NewUser, NewVersion, User};
    use diesel::PgConnection;
    use semver::Version;

    #[test]
    fn test_increment_and_persist_all() {
        let counter = DownloadsCounter::new();
        let conn = crate::db::test_conn();
        let mut state = State::new(&conn);

        let v1 = state.new_version(&conn);
        let v2 = state.new_version(&conn);
        let v3 = state.new_version(&conn);

        // Add 15 downloads between v1 and v2, and no downloads for v3.
        for _ in 0..10 {
            counter.increment(v1);
        }
        for _ in 0..5 {
            counter.increment(v2);
        }
        assert_eq!(15, counter.pending_count.load(Ordering::SeqCst));

        // Persist everything to the database
        let stats = counter
            .persist_all_shards_with_conn(&conn)
            .expect("failed to persist all shards");

        // Ensure the stats are accurate
        assert_eq!(
            stats,
            PersistStats {
                shard: None,
                counted_downloads: 15,
                counted_versions: 2,
                pending_downloads: 0,
            }
        );

        // Ensure the download counts in the database are what we expect.
        state.assert_downloads_count(&conn, v1, 10);
        state.assert_downloads_count(&conn, v2, 5);
        state.assert_downloads_count(&conn, v3, 0);
    }

    #[test]
    fn test_increment_and_persist_shard() {
        let counter = DownloadsCounter::new();
        let conn = crate::db::test_conn();
        let mut state = State::new(&conn);

        let v1 = state.new_version(&conn);
        let v1_shard = counter.inner.determine_map(&v1);

        // For this test to work we need the two versions to be stored in different shards.
        let mut v2 = state.new_version(&conn);
        while counter.inner.determine_map(&v2) == v1_shard {
            v2 = state.new_version(&conn);
        }
        let v2_shard = counter.inner.determine_map(&v2);

        // Add 15 downloads between v1 and v2.
        for _ in 0..10 {
            counter.increment(v1);
        }
        for _ in 0..5 {
            counter.increment(v2);
        }
        assert_eq!(15, counter.pending_count.load(Ordering::SeqCst));

        // Persist one shard at the time and ensure the stats returned for each shard are expected.
        let mut pending = 15;
        for shard in 0..counter.shards_count() {
            let stats = counter
                .persist_next_shard_with_conn(&conn)
                .expect("failed to persist shard");

            if shard == v1_shard {
                pending -= 10;
                assert_eq!(
                    stats,
                    PersistStats {
                        shard: Some(shard),
                        counted_downloads: 10,
                        counted_versions: 1,
                        pending_downloads: pending,
                    }
                );
                state.assert_downloads_count(&conn, v1, 10);
            } else if shard == v2_shard {
                pending -= 5;
                assert_eq!(
                    stats,
                    PersistStats {
                        shard: Some(shard),
                        counted_downloads: 5,
                        counted_versions: 1,
                        pending_downloads: pending,
                    }
                );
                state.assert_downloads_count(&conn, v2, 5);
            } else {
                assert_eq!(
                    stats,
                    PersistStats {
                        shard: Some(shard),
                        counted_downloads: 0,
                        counted_versions: 0,
                        pending_downloads: pending,
                    }
                );
            };
        }
        assert_eq!(pending, 0);

        // Finally ensure that the download counts in the database are what we expect.
        state.assert_downloads_count(&conn, v1, 10);
        state.assert_downloads_count(&conn, v2, 5);
    }

    struct State {
        user: User,
        krate: Crate,
        next_version: u32,
    }

    impl State {
        fn new(conn: &PgConnection) -> Self {
            let user = NewUser {
                gh_id: 0,
                gh_login: "ghost",
                ..NewUser::default()
            }
            .create_or_update(None, conn)
            .expect("failed to create user");

            let krate = NewCrate {
                name: "foo",
                ..NewCrate::default()
            }
            .create_or_update(conn, user.id, None)
            .expect("failed to create crate");

            Self {
                user,
                krate,
                next_version: 1,
            }
        }

        fn new_version(&mut self, conn: &PgConnection) -> i32 {
            let version = NewVersion::new(
                self.krate.id,
                &Version::parse(&format!("{}.0.0", self.next_version)).unwrap(),
                &HashMap::new(),
                None,
                None,
                0,
                self.user.id,
            )
            .expect("failed to create version")
            .save(conn, &[], "ghost@example.com")
            .expect("failed to save version");

            self.next_version += 1;
            version.id
        }

        fn assert_downloads_count(&self, conn: &PgConnection, version: i32, expected: i64) {
            use crate::schema::version_downloads::dsl::*;
            use diesel::dsl::*;

            let actual: Option<i64> = version_downloads
                .select(sum(downloads))
                .filter(version_id.eq(version))
                .first(conn)
                .unwrap();
            assert_eq!(actual.unwrap_or(0), expected);
        }
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::Cursor;
use std::fmt;

use http::{Status, ContentType};
use response::{Response, Stream};

/// Trait implemented by types that send a response to clients.
///
/// Types that implement this trait can be used as the return type of a handler,
/// as illustrated below:
///
/// ```rust,ignore
/// #[get("/")]
/// fn index() -> T { ... }
/// ```
///
/// In this example, `T` can be any type that implements `Responder`.
///
/// # Return Value
///
/// A `Responder` returns an `Ok(Response)` or an `Err(Status)`:
///
///   * An `Ok` variant means that the `Responder` was successful in generating
///     a `Response`. The `Response` will be written out to the client.
///
///   * An `Err` variant means that the `Responder` could not or did not
///     generate a `Response`. The contained `Status` will be used to find the
///     relevant error catcher which then generates an error response.
///
/// # Provided Implementations
///
/// Rocket implements `Responder` for several standard library types. Their
/// behavior is documented here. Note that the `Result` implementation is
/// overloaded, allowing for two `Responder`s to be used at once, depending on
/// the variant.
///
///   * **&str**
///
///     Sets the `Content-Type` to `text/plain`. The string is used as the body
///     of the response, which is fixed size and not streamed. To stream a raw
///     string, use `Stream::from(Cursor::new(string))`.
///
///   * **String**
///
///     Sets the `Content-Type`t to `text/html`. The string is used as the body
///     of the response, which is fixed size and not streamed. To stream a
///     string, use `Stream::from(Cursor::new(string))`.
///
///   * **File**
///
///     Streams the `File` to the client. This is essentially an alias to
///     `Stream::from(file)`.
///
///   * **()**
///
///     Responds with an empty body. No Content-Type is set.
///
///   * **Option<T>**
///
///     If the `Option` is `Some`, the wrapped responder is used to respond to
///     the client. Otherwise, an `Err` with status **404 Not Found** is
///     returned and a warning is printed to the console.
///
///   * **Result<T, E>** _where_ **E: Debug**
///
///     If the `Result` is `Ok`, the wrapped responder is used to respond to the
///     client. Otherwise, an `Err` with status **500 Internal Server Error** is
///     returned and the error is printed to the console using the `Debug`
///     implementation.
///
///   * **Result<T, E>** _where_ **E: Debug + Responder**
///
///     If the `Result` is `Ok`, the wrapped `Ok` responder is used to respond
///     to the client. If the `Result` is `Err`, the wrapped `Err` responder is
///     used to respond to the client.
///
/// # Implementation Tips
///
/// This section describes a few best practices to take into account when
/// implementing `Responder`.
///
/// ## Debug
///
/// A type implementing `Responder` should implement the `Debug` trait when
/// possible. This is because the `Responder` implementation for `Result`
/// requires its `Err` type to implement `Debug`. Therefore, a type implementing
/// `Debug` can more easily be composed.
///
/// ## Joining and Merging
///
/// When chaining/wrapping other `Responder`s, use the
/// [merge](/rocket/struct.Response.html#method.merge) or
/// [join](/rocket/struct.Response.html#method.join) methods on the `Response`
/// struct. Ensure that you document the merging or joining behavior
/// appropriately.
///
/// # Example
///
/// Say that you have a custom type, `Person`:
///
/// ```rust
/// struct Person {
///     name: String,
///     age: u16
/// }
/// ```
///
/// You'd like to use `Person` as a `Responder` so that you can return a
/// `Person` directly from a handler:
///
/// ```rust,ignore
/// #[get("/person/<id>")]
/// fn person(id: usize) -> Option<Person> {
///     Person::from_id(id)
/// }
/// ```
///
/// You want the `Person` responder to set two header fields: `X-Person-Name`
/// and `X-Person-Age` as well as supply a custom representation of the object
/// (`Content-Type: application/x-person`) in the body of the response. The
/// following `Responder` implementation accomplishes this:
///
/// ```rust
/// # #[derive(Debug)]
/// # struct Person { name: String, age: u16 }
/// #
/// use std::io::Cursor;
///
/// use rocket::response::{self, Response, Responder};
/// use rocket::http::ContentType;
///
/// impl<'r> Responder<'r> for Person {
///     fn respond(self) -> response::Result<'r> {
///         Response::build()
///             .sized_body(Cursor::new(format!("{:?}", self)))
///             .raw_header("X-Person-Name", self.name)
///             .raw_header("X-Person-Age", self.age.to_string())
///             .header(ContentType::new("application", "x-person"))
///             .ok()
///     }
/// }
/// ```
pub trait Responder<'r> {
    /// Attempts to write a response to `res`.
    ///
    /// If writing the response successfully completes, an outcome of `Success`
    /// is returned. If writing the response begins but fails, an outcome of
    /// `Failure` is returned. If writing a response fails before writing
    /// anything out, an outcome of `Forward` can be returned, which causes the
    /// response to be written by the appropriate error catcher instead.
    fn respond(self) -> Result<Response<'r>, Status>;
}

/// Sets the `Content-Type`t to `text/plain` if it is not already set. Sends the
/// string as the body of the response. Never fails.
///
/// # Example
///
/// ```rust
/// use rocket::response::Responder;
/// use rocket::http::ContentType;
///
/// let mut response = "Hello".respond().unwrap();
///
/// let body_string = response.body().and_then(|b| b.into_string());
/// assert_eq!(body_string, Some("Hello".to_string()));
///
/// let content_type: Vec<_> = response.header_values("Content-Type").collect();
/// assert_eq!(content_type.len(), 1);
/// assert_eq!(content_type[0], ContentType::Plain.to_string());
/// ```
impl<'r> Responder<'r> for &'r str {
    fn respond(self) -> Result<Response<'r>, Status> {
        Response::build()
            .header(ContentType::Plain)
            .sized_body(Cursor::new(self))
            .ok()
    }
}

impl Responder<'static> for String {
    fn respond(self) -> Result<Response<'static>, Status> {
        Response::build()
            .header(ContentType::HTML)
            .sized_body(Cursor::new(self))
            .ok()
    }
}

/// Essentially aliases Stream<File>.
impl Responder<'static> for File {
    fn respond(self) -> Result<Response<'static>, Status> {
        Stream::from(self).respond()
    }
}

/// Empty response.
impl Responder<'static> for () {
    fn respond(self) -> Result<Response<'static>, Status> {
        Ok(Response::new())
    }
}

impl<'r, R: Responder<'r>> Responder<'r> for Option<R> {
    fn respond(self) -> Result<Response<'r>, Status> {
        self.map_or_else(|| {
            warn_!("Response was `None`.");
            Err(Status::NotFound)
        }, |r| r.respond())
    }
}

impl<'r, R: Responder<'r>, E: fmt::Debug> Responder<'r> for Result<R, E> {
    default fn respond(self) -> Result<Response<'r>, Status> {
        self.map(|r| r.respond()).unwrap_or_else(|e| {
            warn_!("Response was `Err`: {:?}.", e);
            Err(Status::InternalServerError)
        })
    }
}

impl<'r, R: Responder<'r>, E: Responder<'r> + fmt::Debug> Responder<'r> for Result<R, E> {
    fn respond(self) -> Result<Response<'r>, Status> {
        match self {
            Ok(responder) => responder.respond(),
            Err(responder) => responder.respond(),
        }
    }
}
<commit_msg>Cleanup Responder documentation.<commit_after>use std::fs::File;
use std::io::Cursor;
use std::fmt;

use http::{Status, ContentType};
use response::{Response, Stream};

/// Trait implemented by types that send a response to clients.
///
/// Types that implement this trait can be used as the return type of a handler,
/// as illustrated below:
///
/// ```rust,ignore
/// #[get("/")]
/// fn index() -> T { ... }
/// ```
///
/// In this example, `T` can be any type that implements `Responder`.
///
/// # Return Value
///
/// A `Responder` returns an `Ok(Response)` or an `Err(Status)`:
///
///   * An `Ok` variant means that the `Responder` was successful in generating
///     a `Response`. The `Response` will be written out to the client.
///
///   * An `Err` variant means that the `Responder` could not or did not
///     generate a `Response`. The contained `Status` will be used to find the
///     relevant error catcher which then generates an error response.
///
/// # Provided Implementations
///
/// Rocket implements `Responder` for several standard library types. Their
/// behavior is documented here. Note that the `Result` implementation is
/// overloaded, allowing for two `Responder`s to be used at once, depending on
/// the variant.
///
///   * **&str**
///
///     Sets the `Content-Type` to `text/plain`. The string is used as the body
///     of the response, which is fixed size and not streamed. To stream a raw
///     string, use `Stream::from(Cursor::new(string))`.
///
///   * **String**
///
///     Sets the `Content-Type`t to `text/html`. The string is used as the body
///     of the response, which is fixed size and not streamed. To stream a
///     string, use `Stream::from(Cursor::new(string))`.
///
///   * **File**
///
///     Streams the `File` to the client. This is essentially an alias to
///     `Stream::from(file)`.
///
///   * **()**
///
///     Responds with an empty body. No Content-Type is set.
///
///   * **Option<T>**
///
///     If the `Option` is `Some`, the wrapped responder is used to respond to
///     the client. Otherwise, an `Err` with status **404 Not Found** is
///     returned and a warning is printed to the console.
///
///   * **Result<T, E>** _where_ **E: Debug**
///
///     If the `Result` is `Ok`, the wrapped responder is used to respond to the
///     client. Otherwise, an `Err` with status **500 Internal Server Error** is
///     returned and the error is printed to the console using the `Debug`
///     implementation.
///
///   * **Result<T, E>** _where_ **E: Debug + Responder**
///
///     If the `Result` is `Ok`, the wrapped `Ok` responder is used to respond
///     to the client. If the `Result` is `Err`, the wrapped `Err` responder is
///     used to respond to the client.
///
/// # Implementation Tips
///
/// This section describes a few best practices to take into account when
/// implementing `Responder`.
///
/// ## Debug
///
/// A type implementing `Responder` should implement the `Debug` trait when
/// possible. This is because the `Responder` implementation for `Result`
/// requires its `Err` type to implement `Debug`. Therefore, a type implementing
/// `Debug` can more easily be composed.
///
/// ## Joining and Merging
///
/// When chaining/wrapping other `Responder`s, use the
/// [merge](/rocket/struct.Response.html#method.merge) or
/// [join](/rocket/struct.Response.html#method.join) methods on the `Response`
/// or `ResponseBuilder` struct. Ensure that you document the merging or joining
/// behavior appropriately.
///
/// # Example
///
/// Say that you have a custom type, `Person`:
///
/// ```rust
/// struct Person {
///     name: String,
///     age: u16
/// }
/// ```
///
/// You'd like to use `Person` as a `Responder` so that you can return a
/// `Person` directly from a handler:
///
/// ```rust,ignore
/// #[get("/person/<id>")]
/// fn person(id: usize) -> Option<Person> {
///     Person::from_id(id)
/// }
/// ```
///
/// You want the `Person` responder to set two header fields: `X-Person-Name`
/// and `X-Person-Age` as well as supply a custom representation of the object
/// (`Content-Type: application/x-person`) in the body of the response. The
/// following `Responder` implementation accomplishes this:
///
/// ```rust
/// # #[derive(Debug)]
/// # struct Person { name: String, age: u16 }
/// #
/// use std::io::Cursor;
///
/// use rocket::response::{self, Response, Responder};
/// use rocket::http::ContentType;
///
/// impl<'r> Responder<'r> for Person {
///     fn respond(self) -> response::Result<'r> {
///         Response::build()
///             .sized_body(Cursor::new(format!("{:?}", self)))
///             .raw_header("X-Person-Name", self.name)
///             .raw_header("X-Person-Age", self.age.to_string())
///             .header(ContentType::new("application", "x-person"))
///             .ok()
///     }
/// }
/// ```
pub trait Responder<'r> {
    /// Returns `Ok` if a `Response` could be generated successfully. Otherwise,
    /// returns an `Err` with a failing `Status`.
    ///
    /// When using Rocket's code generation, if an `Ok(Response)` is returned,
    /// the response will be written out to the client. If an `Err(Status)` is
    /// returned, the error catcher for the given status is retrieved and called
    /// to generate a final error response, which is then written out to the
    /// client.
    fn respond(self) -> Result<Response<'r>, Status>;
}

/// Returns a response with Content-Type `text/plain` and a fixed-size body
/// containing the string `self`. Always returns `Ok`.
///
/// # Example
///
/// ```rust
/// use rocket::response::Responder;
/// use rocket::http::ContentType;
///
/// let mut response = "Hello".respond().unwrap();
///
/// let body_string = response.body().and_then(|b| b.into_string());
/// assert_eq!(body_string, Some("Hello".to_string()));
///
/// let content_type: Vec<_> = response.header_values("Content-Type").collect();
/// assert_eq!(content_type.len(), 1);
/// assert_eq!(content_type[0], ContentType::Plain.to_string());
/// ```
impl<'r> Responder<'r> for &'r str {
    fn respond(self) -> Result<Response<'r>, Status> {
        Response::build()
            .header(ContentType::Plain)
            .sized_body(Cursor::new(self))
            .ok()
    }
}

/// Returns a response with Content-Type `text/html` and a fized-size body
/// containing the string `self`. Always returns `Ok`.
impl Responder<'static> for String {
    fn respond(self) -> Result<Response<'static>, Status> {
        Response::build()
            .header(ContentType::HTML)
            .sized_body(Cursor::new(self))
            .ok()
    }
}

/// Aliases Stream<File>: `Stream::from(self)`.
impl Responder<'static> for File {
    fn respond(self) -> Result<Response<'static>, Status> {
        Stream::from(self).respond()
    }
}

/// Returns an empty, default `Response`. Always returns `Ok`.
impl Responder<'static> for () {
    fn respond(self) -> Result<Response<'static>, Status> {
        Ok(Response::new())
    }
}

/// If `self` is `Some`, responds with the wrapped `Responder`. Otherwise prints
/// a warning message and returns an `Err` of `Status::NotFound`.
impl<'r, R: Responder<'r>> Responder<'r> for Option<R> {
    fn respond(self) -> Result<Response<'r>, Status> {
        self.map_or_else(|| {
            warn_!("Response was `None`.");
            Err(Status::NotFound)
        }, |r| r.respond())
    }
}

/// If `self` is `Ok`, responds with the wrapped `Responder`. Otherwise prints a
/// warning message with the `Err` value returns an `Err` of
/// `Status::InternalServerError`.
impl<'r, R: Responder<'r>, E: fmt::Debug> Responder<'r> for Result<R, E> {
    default fn respond(self) -> Result<Response<'r>, Status> {
        self.map(|r| r.respond()).unwrap_or_else(|e| {
            warn_!("Response was `Err`: {:?}.", e);
            Err(Status::InternalServerError)
        })
    }
}

/// Responds with the wrapped `Responder` in `self`, whether it is `Ok` or
/// `Err`.
impl<'r, R: Responder<'r>, E: Responder<'r> + fmt::Debug> Responder<'r> for Result<R, E> {
    fn respond(self) -> Result<Response<'r>, Status> {
        match self {
            Ok(responder) => responder.respond(),
            Err(responder) => responder.respond(),
        }
    }
}
<|endoftext|>


<commit_before>use wheel;
use streaming::{self, StreamingSieve};

#[cfg(target_pointer_width = "32")]
const SQRT: usize = 1 << 16;
#[cfg(target_pointer_width = "64")]
const SQRT: usize = 1 << 32;

enum Early {
    Two,
    Three,
    Five,
    Done,
}

/// An iterator over all primes.
///
/// This will yield primes indefinitely (bits in `usize`
/// permitting). If there is an known upper bound, sieving first with
/// `Sieve` and using its `primes_from` method may be more efficient,
/// especially if the bound is small.
///
/// This requires *O(sqrt(p))* memory to yield prime `p`, where `X` is
/// the maximum value of `usize`.
///
/// # Examples
///
/// ```rust
/// # extern crate primal;
/// let count = primal::Primes::all().take_while(|p| *p < 1_000_000).count();
/// println!("{}", count);
/// ```
pub struct Primes {
    early: Early,
    base: usize,
    ones: primal_bit::IntoOnes,
    streaming: StreamingSieve,
    sieving_primes: Option<Box<Primes>>,
    left_over: Option<usize>,
}

impl Primes {
    /// The sequence `2, 3, 5, 7, 11, ...`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # extern crate primal;
    /// // print the first 20 primes
    /// for p in primal::Primes::all().take(20) {
    ///     println!("{}", p);
    /// }
    /// ```
    pub fn all() -> Primes {
        Primes::sqrt(SQRT)
    }

    fn sqrt(sqrt: usize) -> Primes {
        let (sieving, limit) = if sqrt < streaming::isqrt(streaming::SEG_LEN) {
            (None, sqrt * sqrt)
        } else {
            (Some(Box::new(Primes::sqrt(streaming::isqrt(sqrt)))),
             streaming::SEG_LEN)
        };
        let mut streaming = StreamingSieve::new(limit);

        let ones = {
            let (_, bits) = streaming.next().unwrap();
            bits.clone().into_ones()
        };

        // we manually add the primes
        streaming.small = None;
        // go to the end.
        streaming.limit = !0;

        Primes {
            early: Early::Two,
            base: 0,
            ones,
            streaming,
            sieving_primes: sieving,
            left_over: None,
        }
    }

    #[inline]
    fn from_bit_index(&self, i: usize) -> Option<usize> {
        let w = wheel::from_bit_index(i);
        self.base.checked_add(w)
    }

    fn advance_ones(&mut self) -> bool {
        let low = self.streaming.low;
        let high = low.saturating_add(streaming::SEG_LEN);

        for q in self.left_over.into_iter().chain(self.sieving_primes.as_mut().unwrap()) {
            if q.saturating_mul(q) >= high {
                self.left_over = Some(q);
                break
            }
            if q > streaming::isqrt(streaming::SEG_LEN) {
                self.streaming.add_sieving_prime(q, low);
            }
        }

        match self.streaming.next() {
            Some((_, bits)) => {
                self.base = low;
                self.ones = bits.clone().into_ones();
                true
            },
            None => false,
        }
    }
}

impl Iterator for Primes {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<usize> {
        match self.early {
            Early::Done => {}
            Early::Two => {
                self.early = Early::Three;
                return Some(2)
            }
            Early::Three => {
                self.early = Early::Five;
                return Some(3)
            }
            Early::Five => {
                self.early = Early::Done;
                return Some(5)
            }
        }

        loop {
            if let Some(i) = self.ones.next() {
                return self.from_bit_index(i);
            }
            if !self.advance_ones() {
                return None;
            }
        }
    }

    fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
    where
        F: FnMut(Acc, Self::Item) -> Acc
    {
        match self.early {
            Early::Done => {}
            Early::Two => {
                acc = f(acc, 2);
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Three => {
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Five => {
                acc = f(acc, 5);
            }
        }
        loop {
            while let Some(i) = self.ones.next() {
                match self.from_bit_index(i) {
                    Some(p) => acc = f(acc, p),
                    None => return acc,
                }
            }
            if !self.advance_ones() {
                return acc;
            }
        }
    }
}

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

    fn check_equality(limit: usize) {
        let sieve = Sieve::new(limit);

        let real = sieve.primes_from(0).take_while(|x| *x < limit);
        let computed = Primes::all().take_while(|x| *x < limit);

        let mut i = 0;
        for (r, c) in real.zip(computed) {
            assert_eq!(c, r);
            i += 1;
        }
        assert_eq!(sieve.prime_pi(limit), i);
    }

    #[test]
    fn equality() {
        check_equality(20_000_000);
    }

    #[test]
    fn equality_huge() {
        let limit = if cfg!(feature = "slow_tests") {
            // This takes a minute or so in debug mode, but it does work!
            ::std::u32::MAX as usize
        } else {
            100_000_000
        };
        check_equality(limit);
    }
}
<commit_msg>Test Primes::fold<commit_after>use wheel;
use streaming::{self, StreamingSieve};

#[cfg(target_pointer_width = "32")]
const SQRT: usize = 1 << 16;
#[cfg(target_pointer_width = "64")]
const SQRT: usize = 1 << 32;

enum Early {
    Two,
    Three,
    Five,
    Done,
}

/// An iterator over all primes.
///
/// This will yield primes indefinitely (bits in `usize`
/// permitting). If there is an known upper bound, sieving first with
/// `Sieve` and using its `primes_from` method may be more efficient,
/// especially if the bound is small.
///
/// This requires *O(sqrt(p))* memory to yield prime `p`, where `X` is
/// the maximum value of `usize`.
///
/// # Examples
///
/// ```rust
/// # extern crate primal;
/// let count = primal::Primes::all().take_while(|p| *p < 1_000_000).count();
/// println!("{}", count);
/// ```
pub struct Primes {
    early: Early,
    base: usize,
    ones: primal_bit::IntoOnes,
    streaming: StreamingSieve,
    sieving_primes: Option<Box<Primes>>,
    left_over: Option<usize>,
}

impl Primes {
    /// The sequence `2, 3, 5, 7, 11, ...`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # extern crate primal;
    /// // print the first 20 primes
    /// for p in primal::Primes::all().take(20) {
    ///     println!("{}", p);
    /// }
    /// ```
    pub fn all() -> Primes {
        Primes::sqrt(SQRT)
    }

    fn sqrt(sqrt: usize) -> Primes {
        let (sieving, limit) = if sqrt < streaming::isqrt(streaming::SEG_LEN) {
            (None, sqrt * sqrt)
        } else {
            (Some(Box::new(Primes::sqrt(streaming::isqrt(sqrt)))),
             streaming::SEG_LEN)
        };
        let mut streaming = StreamingSieve::new(limit);

        let ones = {
            let (_, bits) = streaming.next().unwrap();
            bits.clone().into_ones()
        };

        // we manually add the primes
        streaming.small = None;
        // go to the end.
        streaming.limit = !0;

        Primes {
            early: Early::Two,
            base: 0,
            ones,
            streaming,
            sieving_primes: sieving,
            left_over: None,
        }
    }

    #[inline]
    fn from_bit_index(&self, i: usize) -> Option<usize> {
        let w = wheel::from_bit_index(i);
        self.base.checked_add(w)
    }

    fn advance_ones(&mut self) -> bool {
        let low = self.streaming.low;
        let high = low.saturating_add(streaming::SEG_LEN);

        for q in self.left_over.into_iter().chain(self.sieving_primes.as_mut().unwrap()) {
            if q.saturating_mul(q) >= high {
                self.left_over = Some(q);
                break
            }
            if q > streaming::isqrt(streaming::SEG_LEN) {
                self.streaming.add_sieving_prime(q, low);
            }
        }

        match self.streaming.next() {
            Some((_, bits)) => {
                self.base = low;
                self.ones = bits.clone().into_ones();
                true
            },
            None => false,
        }
    }
}

impl Iterator for Primes {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<usize> {
        match self.early {
            Early::Done => {}
            Early::Two => {
                self.early = Early::Three;
                return Some(2)
            }
            Early::Three => {
                self.early = Early::Five;
                return Some(3)
            }
            Early::Five => {
                self.early = Early::Done;
                return Some(5)
            }
        }

        loop {
            if let Some(i) = self.ones.next() {
                return self.from_bit_index(i);
            }
            if !self.advance_ones() {
                return None;
            }
        }
    }

    fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
    where
        F: FnMut(Acc, Self::Item) -> Acc
    {
        match self.early {
            Early::Done => {}
            Early::Two => {
                acc = f(acc, 2);
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Three => {
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Five => {
                acc = f(acc, 5);
            }
        }
        loop {
            while let Some(i) = self.ones.next() {
                match self.from_bit_index(i) {
                    Some(p) => acc = f(acc, p),
                    None => return acc,
                }
            }
            if !self.advance_ones() {
                return acc;
            }
        }
    }
}

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

    fn check_equality(limit: usize) {
        let sieve = Sieve::new(limit);

        let real = sieve.primes_from(0).take_while(|x| *x < limit);
        let computed = Primes::all().take_while(|x| *x < limit);

        let mut i = 0;
        for (r, c) in real.zip(computed) {
            assert_eq!(c, r);
            i += 1;
        }
        assert_eq!(sieve.prime_pi(limit), i);
    }

    #[test]
    fn equality() {
        check_equality(20_000_000);
    }

    #[test]
    fn equality_huge() {
        let limit = if cfg!(feature = "slow_tests") {
            // This takes a minute or so in debug mode, but it does work!
            ::std::u32::MAX as usize
        } else {
            100_000_000
        };
        check_equality(limit);
    }

    #[test]
    #[should_panic = "123456791"]
    fn fold() {
        // There's no termination until we exceed `usize::MAX`, which
        // will take too long, but we can cut it short by unwinding.
        Primes::all().fold((), |(), p| {
            if p > 123456789 {
                panic!(format!("{}", p));
            }
        })
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use num_derive::FromPrimitive;
use arg_enum_proc_macro::ArgEnum;

/// Sample position for subsampled chroma
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSamplePosition {
  /// The source video transfer function must be signaled
  /// outside the AV1 bitstream.
  Unknown,
  /// Horizontally co-located with (0, 0) luma sample, vertically positioned
  /// in the middle between two luma samples.
  Vertical,
  /// Co-located with (0, 0) luma sample.
  Colocated
}

impl Default for ChromaSamplePosition {
  fn default() -> Self {
    ChromaSamplePosition::Unknown
  }
}

/// Chroma subsampling format
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSampling {
  /// Both vertically and horizontally subsampled.
  Cs420,
  /// Horizontally subsampled.
  Cs422,
  /// Not subsampled.
  Cs444,
  /// Monochrome.
  Cs400
}

impl Default for ChromaSampling {
  fn default() -> Self {
    ChromaSampling::Cs420
  }
}

impl ChromaSampling {
  /// Provides the amount to right shift the luma plane dimensions to get the
  ///  chroma plane dimensions.
  /// Only values 0 or 1 are ever returned.
  /// The plane dimensions must also be rounded up to accomodate odd luma plane
  ///  sizes.
  /// Cs400 returns None, as there are no chroma planes.
  pub fn get_decimation(self) -> Option<(usize, usize)> {
    use self::ChromaSampling::*;
    match self {
      Cs420 => Some((1, 1)),
      Cs422 => Some((1, 0)),
      Cs444 => Some((0, 0)),
      Cs400 => None,
    }
  }

  pub fn get_chroma_dimensions(self, luma_width: usize, luma_height: usize)
   -> (usize, usize) {
    if let Some((ss_x, ss_y)) = self.get_decimation() {
      ((luma_width + ss_x) >> ss_x, (luma_height + ss_y) >> ss_y)
    }
    else {
      (0, 0)
    }
  }
}

/// Supported Color Primaries
///
/// As defined by “Color primaries” section of ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ColorPrimaries {
  BT709 = 1,
  Unspecified,
  BT470M = 4,
  BT470BG,
  ST170M,
  ST240M,
  Film,
  BT2020,
  ST428,
  P3DCI,
  P3Display,
  Tech3213 = 22
}

impl Default for ColorPrimaries {
  fn default() -> Self {
    ColorPrimaries::Unspecified
  }
}

/// Supported Transfer Characteristics
///
/// As defined by “Transfer characteristics” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum TransferCharacteristics {
  BT1886 = 1,
  Unspecified,
  BT470M = 4,
  BT470BG,
  ST170M,
  ST240M,
  Linear,
  Logarithmic100,
  Logarithmic316,
  XVYCC,
  BT1361E,
  SRGB,
  BT2020Ten,
  BT2020Twelve,
  PerceptualQuantizer,
  ST428,
  HybridLogGamma
}

impl Default for TransferCharacteristics {
  fn default() -> Self {
    TransferCharacteristics::Unspecified
  }
}

/// Matrix coefficients
///
/// As defined by the “Matrix coefficients” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum MatrixCoefficients {
  Identity = 0,
  BT709,
  Unspecified,
  BT470M = 4,
  BT470BG,
  ST170M,
  ST240M,
  YCgCo,
  BT2020NonConstantLuminance,
  BT2020ConstantLuminance,
  ST2085,
  ChromaticityDerivedNonConstantLuminance,
  ChromaticityDerivedConstantLuminance,
  ICtCp
}

impl Default for MatrixCoefficients {
  fn default() -> Self {
    MatrixCoefficients::Unspecified
  }
}

/// Signal the content color description
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
  pub color_primaries: ColorPrimaries,
  pub transfer_characteristics: TransferCharacteristics,
  pub matrix_coefficients: MatrixCoefficients
}

/// Allowed pixel value range
///
/// C.f. VideoFullRangeFlag variable specified in ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum PixelRange {
  /// Studio swing representation
  Limited,
  /// Full swing representation
  Full
}

impl Default for PixelRange {
  fn default() -> Self {
    PixelRange::Limited
  }
}

/// High dynamic range content light level
///
/// As defined by CEA-861.3, Appendix A.
#[derive(Copy, Clone, Debug)]
pub struct ContentLight {
  /// Maximum content light level
  pub max_content_light_level: u16,
  /// Maximum frame-average light level
  pub max_frame_average_light_level: u16
}

/// Chromaticity coordinates expressed as 0.16 fixed-point values
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Point {
  pub x: u16,
  pub y: u16
}

/// High dynamic range mastering display color volume
///
/// As defined by CIE 1931
#[derive(Copy, Clone, Debug)]
pub struct MasteringDisplay {
  /// Chromaticity coordinates in Red, Green, Blue order
  /// expressed as 0.16 fixed-point
  pub primaries: [Point; 3],
  /// Chromaticity coordinates expressed as 0.16 fixed-point
  pub white_point: Point,
  /// 24.8 fixed-point maximum luminance in candelas per square meter
  pub max_luminance: u32,
  /// 18.14 fixed-point minimum luminance in candelas per square meter
  pub min_luminance: u32
}

<commit_msg>Document color config enum members<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use num_derive::FromPrimitive;
use arg_enum_proc_macro::ArgEnum;

/// Sample position for subsampled chroma
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSamplePosition {
  /// The source video transfer function must be signaled
  /// outside the AV1 bitstream.
  Unknown,
  /// Horizontally co-located with (0, 0) luma sample, vertically positioned
  /// in the middle between two luma samples.
  Vertical,
  /// Co-located with (0, 0) luma sample.
  Colocated
}

impl Default for ChromaSamplePosition {
  fn default() -> Self {
    ChromaSamplePosition::Unknown
  }
}

/// Chroma subsampling format
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSampling {
  /// Both vertically and horizontally subsampled.
  Cs420,
  /// Horizontally subsampled.
  Cs422,
  /// Not subsampled.
  Cs444,
  /// Monochrome.
  Cs400
}

impl Default for ChromaSampling {
  fn default() -> Self {
    ChromaSampling::Cs420
  }
}

impl ChromaSampling {
  /// Provides the amount to right shift the luma plane dimensions to get the
  ///  chroma plane dimensions.
  /// Only values 0 or 1 are ever returned.
  /// The plane dimensions must also be rounded up to accomodate odd luma plane
  ///  sizes.
  /// Cs400 returns None, as there are no chroma planes.
  pub fn get_decimation(self) -> Option<(usize, usize)> {
    use self::ChromaSampling::*;
    match self {
      Cs420 => Some((1, 1)),
      Cs422 => Some((1, 0)),
      Cs444 => Some((0, 0)),
      Cs400 => None,
    }
  }

  /// Calculates the size of a chroma plane for this sampling type, given the luma plane dimensions.
  pub fn get_chroma_dimensions(self, luma_width: usize, luma_height: usize)
   -> (usize, usize) {
    if let Some((ss_x, ss_y)) = self.get_decimation() {
      ((luma_width + ss_x) >> ss_x, (luma_height + ss_y) >> ss_y)
    }
    else {
      (0, 0)
    }
  }
}

/// Supported Color Primaries
///
/// As defined by “Color primaries” section of ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ColorPrimaries {
  /// BT.709
  BT709 = 1,
  /// Unspecified, must be signaled or inferred outside of the bitstream
  Unspecified,
  /// BT.470 System M (historical)
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  ST170M,
  /// SMPTE 240M (historical)
  ST240M,
  /// Generic film
  Film,
  /// BT.2020, BT.2100
  BT2020,
  /// SMPTE 248 (CIE 1921 XYZ)
  ST428,
  /// SMPTE RP 431-2
  P3DCI,
  /// SMPTE EG 432-1
  P3Display,
  /// EBU Tech. 3213-E
  Tech3213 = 22
}

impl Default for ColorPrimaries {
  fn default() -> Self {
    ColorPrimaries::Unspecified
  }
}

/// Supported Transfer Characteristics
///
/// As defined by “Transfer characteristics” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum TransferCharacteristics {
  /// BT.709
  BT1886 = 1,
  /// Unspecified, must be signaled or inferred outside of the bitstream
  Unspecified,
  /// BT.470 System M (historical)
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  ST170M,
  /// SMPTE 240 M
  ST240M,
  /// Linear
  Linear,
  /// Logarithmic (100:1 range)
  Logarithmic100,
  /// Logarithmic ((100 * √10):1 range)
  Logarithmic316,
  /// IEC 61966-2-4
  XVYCC,
  /// BT.1361 extended color gamut system (historical)
  BT1361E,
  /// sRGB or sYCC
  SRGB,
  /// BT.2020 10-bit systems
  BT2020Ten,
  /// BT.2020 12-bit systems
  BT2020Twelve,
  /// SMPTE ST 2084, ITU BT.2100 PQ
  PerceptualQuantizer,
  /// SMPTE ST 428
  ST428,
  /// BT.2100 HLG, ARIB STD-B67
  HybridLogGamma
}

impl Default for TransferCharacteristics {
  fn default() -> Self {
    TransferCharacteristics::Unspecified
  }
}

/// Matrix coefficients
///
/// As defined by the “Matrix coefficients” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum MatrixCoefficients {
  /// Identity matrix
  Identity = 0,
  /// BT.709
  BT709,
  /// Unspecified, must be signaled or inferred outside of the bitstream.
  Unspecified,
  /// US FCC 73.628
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  ST170M,
  /// SMPTE 240 M
  ST240M,
  /// YCgCo
  YCgCo,
  /// BT.2020 non-constant luminance, BT.2100 YCbCr
  BT2020NonConstantLuminance,
  /// BT.2020 constant luminance
  BT2020ConstantLuminance,
  /// SMPTE ST 2085 YDzDx
  ST2085,
  /// Chromaticity-derived non-constant luminance
  ChromaticityDerivedNonConstantLuminance,
  /// Chromaticity-derived constant luminance
  ChromaticityDerivedConstantLuminance,
  /// BT.2020 ICtCp
  ICtCp
}

impl Default for MatrixCoefficients {
  fn default() -> Self {
    MatrixCoefficients::Unspecified
  }
}

/// Signal the content color description
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
  pub color_primaries: ColorPrimaries,
  pub transfer_characteristics: TransferCharacteristics,
  pub matrix_coefficients: MatrixCoefficients
}

/// Allowed pixel value range
///
/// C.f. VideoFullRangeFlag variable specified in ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum PixelRange {
  /// Studio swing representation
  Limited,
  /// Full swing representation
  Full
}

impl Default for PixelRange {
  fn default() -> Self {
    PixelRange::Limited
  }
}

/// High dynamic range content light level
///
/// As defined by CEA-861.3, Appendix A.
#[derive(Copy, Clone, Debug)]
pub struct ContentLight {
  /// Maximum content light level
  pub max_content_light_level: u16,
  /// Maximum frame-average light level
  pub max_frame_average_light_level: u16
}

/// Chromaticity coordinates expressed as 0.16 fixed-point values
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Point {
  pub x: u16,
  pub y: u16
}

/// High dynamic range mastering display color volume
///
/// As defined by CIE 1931
#[derive(Copy, Clone, Debug)]
pub struct MasteringDisplay {
  /// Chromaticity coordinates in Red, Green, Blue order
  /// expressed as 0.16 fixed-point
  pub primaries: [Point; 3],
  /// Chromaticity coordinates expressed as 0.16 fixed-point
  pub white_point: Point,
  /// 24.8 fixed-point maximum luminance in candelas per square meter
  pub max_luminance: u32,
  /// 18.14 fixed-point minimum luminance in candelas per square meter
  pub min_luminance: u32
}

<|endoftext|>


<commit_before>use crate::db;
use futures::future;
use hyper::{body, Client, Uri};
use hyper_tls::HttpsConnector;
use image::imageops::{overlay, FilterType};
use image::{DynamicImage, ImageBuffer, ImageFormat, RgbaImage};
use serde_json;
use std::fs::OpenOptions;
use std::io::prelude::*;
use webp;

pub struct ImageAndIDType {
	id: String,
	image: DynamicImage,
}

async fn get_icon(clan: db::ClanStats, resize_to: u32) -> Result<ImageAndIDType, String> {
	println!("1");

	let mut icon_to_get = clan.emblems.x195_portal.as_str();
	if icon_to_get.len() == 0 {
		icon_to_get = clan.emblems.x256_wowp.as_str();
		if icon_to_get.len() == 0 {
			return Err(format!("Clan {} has no usable icons to fetch", clan.tag));
		}
	}

	println!("2");

	if !icon_to_get.contains("http://") && !icon_to_get.contains("https://") {
		return Err(format!(
			"Clan {} icon has an invalid url: {} , make sure the url starts with http(s)://",
			clan.tag, icon_to_get
		));
	}

	println!("3");

	let https = HttpsConnector::new();
	let client = Client::builder().build::<_, hyper::Body>(https);

	let uri: Uri = match icon_to_get.parse() {
		Ok(v) => v,
		Err(e) => return Err(format!("Clan {} icon has an invalid url: {}", clan.tag, e)),
	};
	let resp = match client.get(uri).await {
		Ok(v) => v,
		Err(e) => {
			return Err(format!(
				"Can't fetch image {}, error: {}, for clan: {}",
				icon_to_get, e, clan.tag
			))
		}
	};

	let resp_bytes = match body::to_bytes(resp).await {
		Ok(v) => v,
		Err(e) => {
			return Err(format!(
				"Unable to fetch image {}, error: {}, for clan: {}",
				icon_to_get, e, clan.tag
			))
		}
	};

	println!("4");

	let body_bytes_ref = resp_bytes.as_ref();

	println!("5");

	let parsed_image = image::load_from_memory(body_bytes_ref).or_else(|e| {
		Err(format!(
			"Unable to parse image {}, error: {}, for clan: {}",
			icon_to_get, e, clan.tag,
		))
	})?;

	println!("6");

	Ok(ImageAndIDType {
		id: clan.id,
		image: parsed_image.resize(resize_to, resize_to, FilterType::CatmullRom),
	})
}

// GetIcons fetches all clan icons and creates a grid
pub async fn get() -> Result<(), String> {
	let img_size: u32 = 60;
	let imgs_in_a_row: usize = 30;

	println!("Getting clan icons...");

	let mut img_and_id: Vec<ImageAndIDType> = Vec::new();

	let current_stats = db::get_current_stats().clone();

	// for clan in current_stats.iter() {
	// 	match get_icon(clan.clone().1.clone(), img_size).await {
	// 		Ok(v) => img_and_id.push(v),
	// 		Err(e) => println!("{}", e),
	// 	};
	// }

	println!("0.1");
	let awaiting = current_stats
		.iter()
		.map(|clan| get_icon(clan.clone().1.clone(), img_size));
	println!("0.2");
	let awaiting_res = future::join_all(awaiting).await;
	println!("0.3");

	for res in awaiting_res {
		match res {
			Ok(v) => img_and_id.push(v),
			Err(e) => println!("{}", e),
		};
	}

	if img_and_id.len() == 0 {
		return Ok(());
	}

	let mut ids: Vec<Vec<String>> = vec![Vec::new()];
	let mut output_img: RgbaImage = ImageBuffer::new(img_size * imgs_in_a_row as u32, img_size);

	for img_obj in img_and_id {
		let id = img_obj.id;
		if ids.last().unwrap().len() == imgs_in_a_row {
			ids.push(Vec::new());

			// increase the size of the output_img
			let mut new_output_img: RgbaImage =
				ImageBuffer::new(img_size * imgs_in_a_row as u32, img_size * ids.len() as u32);
			overlay(&mut new_output_img, &output_img, 0, 0);
			output_img = new_output_img;
		}

		let ids_last = ids.last_mut().unwrap();
		ids_last.push(id);

		let from_top = (ids.len() - 1) as u32 * img_size;
		let from_left = (ids[ids.len() - 1].len() - 1) as u32 * img_size;

		overlay(&mut output_img, &img_obj.image, from_left, from_top);
	}

	output_img
		.save_with_format("./icons/allIcons.png", ImageFormat::Png)
		.or_else(|e| {
			Err(format!(
				"Unable to to save image in png format, error: {}",
				e,
			))
		})?;

	let output_img_dynamic = DynamicImage::ImageRgba8(output_img);
	let rgba_bytes = output_img_dynamic.to_rgba8();

	let encoded_webp = webp::Encoder::from_rgba(
		&rgba_bytes,
		img_size * imgs_in_a_row as u32, // Width
		img_size * ids.len() as u32,     // Height
	)
	.encode(50.0);

	let encoded_webp_bytes = encoded_webp.to_vec();

	let mut file = OpenOptions::new()
		.write(true)
		.create(true)
		.truncate(true)
		.open("./icons/allIcons.webp")
		.or_else(|e| {
			Err(format!(
				"Unable to to save image in webp format, error: {}",
				e,
			))
		})?;

	file.write(&encoded_webp_bytes).or_else(|e| {
		Err(format!(
			"Unable to to write webp data to image file, error: {}",
			e,
		))
	})?;

	let clan_location_in_icons = serde_json::to_vec(&ids).or_else(|e| {
		Err(format!(
			"Unable to to parse clan icons locations to file, error: {}",
			e,
		))
	})?;

	file = OpenOptions::new()
		.write(true)
		.create(true)
		.truncate(true)
		.open("./icons/allIcons.json")
		.or_else(|e| {
			Err(format!(
				"Unable to to create clan icons locations to file, error: {}",
				e,
			))
		})?;

	file.write_all(&clan_location_in_icons).or_else(|e| {
		Err(format!(
			"Unable to to write clan icons locations to file, error: {}",
			e,
		))
	})?;

	Ok(())
}
<commit_msg>Removed debugging<commit_after>use crate::db;
use futures::future;
use hyper::{body, Client, Uri};
use hyper_tls::HttpsConnector;
use image::imageops::{overlay, FilterType};
use image::{DynamicImage, ImageBuffer, ImageFormat, RgbaImage};
use serde_json;
use std::fs::OpenOptions;
use std::io::prelude::*;
use webp;

pub struct ImageAndIDType {
	id: String,
	image: DynamicImage,
}

async fn get_icon(clan: db::ClanStats, resize_to: u32) -> Result<ImageAndIDType, String> {
	let mut icon_to_get = clan.emblems.x195_portal.as_str();
	if icon_to_get.len() == 0 {
		icon_to_get = clan.emblems.x256_wowp.as_str();
		if icon_to_get.len() == 0 {
			return Err(format!("Clan {} has no usable icons to fetch", clan.tag));
		}
	}

	if !icon_to_get.contains("http://") && !icon_to_get.contains("https://") {
		return Err(format!(
			"Clan {} icon has an invalid url: {} , make sure the url starts with http(s)://",
			clan.tag, icon_to_get
		));
	}

	let https = HttpsConnector::new();
	let client = Client::builder().build::<_, hyper::Body>(https);

	let uri: Uri = match icon_to_get.parse() {
		Ok(v) => v,
		Err(e) => return Err(format!("Clan {} icon has an invalid url: {}", clan.tag, e)),
	};
	let resp = match client.get(uri).await {
		Ok(v) => v,
		Err(e) => {
			return Err(format!(
				"Can't fetch image {}, error: {}, for clan: {}",
				icon_to_get, e, clan.tag
			))
		}
	};

	let resp_bytes = match body::to_bytes(resp).await {
		Ok(v) => v,
		Err(e) => {
			return Err(format!(
				"Unable to fetch image {}, error: {}, for clan: {}",
				icon_to_get, e, clan.tag
			))
		}
	};

	let body_bytes_ref = resp_bytes.as_ref();

	let parsed_image = image::load_from_memory(body_bytes_ref).or_else(|e| {
		Err(format!(
			"Unable to parse image {}, error: {}, for clan: {}",
			icon_to_get, e, clan.tag,
		))
	})?;

	Ok(ImageAndIDType {
		id: clan.id,
		image: parsed_image.resize(resize_to, resize_to, FilterType::CatmullRom),
	})
}

// GetIcons fetches all clan icons and creates a grid
pub async fn get() -> Result<(), String> {
	let img_size: u32 = 60;
	let imgs_in_a_row: usize = 30;

	println!("Getting clan icons...");

	let mut img_and_id: Vec<ImageAndIDType> = Vec::new();

	let current_stats = db::get_current_stats().clone();

	let awaiting = current_stats
		.iter()
		.map(|clan| get_icon(clan.clone().1.clone(), img_size));
	let awaiting_res = future::join_all(awaiting).await;

	for res in awaiting_res {
		match res {
			Ok(v) => img_and_id.push(v),
			Err(e) => println!("{}", e),
		};
	}

	if img_and_id.len() == 0 {
		return Ok(());
	}

	let mut ids: Vec<Vec<String>> = vec![Vec::new()];
	let mut output_img: RgbaImage = ImageBuffer::new(img_size * imgs_in_a_row as u32, img_size);

	for img_obj in img_and_id {
		let id = img_obj.id;
		if ids.last().unwrap().len() == imgs_in_a_row {
			ids.push(Vec::new());

			// increase the size of the output_img
			let mut new_output_img: RgbaImage =
				ImageBuffer::new(img_size * imgs_in_a_row as u32, img_size * ids.len() as u32);
			overlay(&mut new_output_img, &output_img, 0, 0);
			output_img = new_output_img;
		}

		let ids_last = ids.last_mut().unwrap();
		ids_last.push(id);

		let from_top = (ids.len() - 1) as u32 * img_size;
		let from_left = (ids[ids.len() - 1].len() - 1) as u32 * img_size;

		overlay(&mut output_img, &img_obj.image, from_left, from_top);
	}

	output_img
		.save_with_format("./icons/allIcons.png", ImageFormat::Png)
		.or_else(|e| {
			Err(format!(
				"Unable to to save image in png format, error: {}",
				e,
			))
		})?;

	let output_img_dynamic = DynamicImage::ImageRgba8(output_img);
	let rgba_bytes = output_img_dynamic.to_rgba8();

	let encoded_webp = webp::Encoder::from_rgba(
		&rgba_bytes,
		img_size * imgs_in_a_row as u32, // Width
		img_size * ids.len() as u32,     // Height
	)
	.encode(50.0);

	let encoded_webp_bytes = encoded_webp.to_vec();

	let mut file = OpenOptions::new()
		.write(true)
		.create(true)
		.truncate(true)
		.open("./icons/allIcons.webp")
		.or_else(|e| {
			Err(format!(
				"Unable to to save image in webp format, error: {}",
				e,
			))
		})?;

	file.write(&encoded_webp_bytes).or_else(|e| {
		Err(format!(
			"Unable to to write webp data to image file, error: {}",
			e,
		))
	})?;

	let clan_location_in_icons = serde_json::to_vec(&ids).or_else(|e| {
		Err(format!(
			"Unable to to parse clan icons locations to file, error: {}",
			e,
		))
	})?;

	file = OpenOptions::new()
		.write(true)
		.create(true)
		.truncate(true)
		.open("./icons/allIcons.json")
		.or_else(|e| {
			Err(format!(
				"Unable to to create clan icons locations to file, error: {}",
				e,
			))
		})?;

	file.write_all(&clan_location_in_icons).or_else(|e| {
		Err(format!(
			"Unable to to write clan icons locations to file, error: {}",
			e,
		))
	})?;

	Ok(())
}
<|endoftext|>


<commit_before>#![crate_name = "hostname"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Alan Andrade <alan.andradec@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 *
 * Synced with:
 *
 * https://www.opensource.apple.com/source/shell_cmds/shell_cmds-170/hostname/hostname.c?txt
 */

extern crate getopts;
extern crate libc;

#[macro_use]
extern crate uucore;

use getopts::Options;
use std::collections::hash_set::HashSet;
use std::iter::repeat;
use std::str;
use std::io::Write;
use std::net::ToSocketAddrs;

static NAME: &'static str = "hostname";

extern {
    fn gethostname(name: *mut libc::c_char, namelen: libc::size_t) -> libc::c_int;
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
extern {
    fn sethostname(name: *const libc::c_char, namelen: libc::c_int) -> libc::c_int;
}

#[cfg(target_os = "linux")]
extern {
    fn sethostname(name: *const libc::c_char, namelen: libc::size_t) -> libc::c_int;
}

pub fn uumain(args: Vec<String>) -> i32 {
    let program = &args[0];

    let mut opts = Options::new();
    opts.optflag("d", "domain", "Display the name of the DNS domain if possible");
    opts.optflag("i", "ip-address", "Display the network address(es) of the host");
    opts.optflag("f", "fqdn", "Display the FQDN (Fully Qualified Domain Name) (default)");   // TODO: support --long
    opts.optflag("s", "short", "Display the short hostname (the portion before the first dot) if possible");
    opts.optflag("h", "help", "Show help");
    opts.optflag("V", "version", "Show program's version");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        _ => { help_menu(program, opts); return 0; }
    };

    if matches.opt_present("h") {
        help_menu(program, opts);
        return 0
    }
    if matches.opt_present("V") { version(); return 0 }

    match matches.free.len() {
        0 => {
            let hostname = xgethostname();

            if matches.opt_present("i") {
                match hostname.to_socket_addrs() {
                    Ok(addresses) => {
                        let mut hashset = HashSet::new();
                        let mut output = String::new();
                        for addr in addresses {
                            // XXX: not sure why this is necessary...
                            if !hashset.contains(&addr) {
                                output.push_str(&format!("{}", addr));
                                output.push_str(" ");
                                hashset.insert(addr.clone());
                            }
                        }
                        let len = output.len();
                        if len > 0 {
                            println!("{}", &output[0 .. len - 1]);
                        }
                    }
                    Err(f) => {
                        show_error!("{}", f);
                        return 1;
                    }
                }
            } else {
                if matches.opt_present("s") {
                    let mut it = hostname.char_indices().filter(|&ci| ci.1 == '.');
                    let ci = it.next();
                    if ci.is_some() {
                        println!("{}", &hostname[0 .. ci.unwrap().0]);
                        return 0;
                    }
                } else if matches.opt_present("d") {
                    let mut it = hostname.char_indices().filter(|&ci| ci.1 == '.');
                    let ci = it.next();
                    if ci.is_some() {
                        println!("{}", &hostname[ci.unwrap().0 + 1 .. ]);
                        return 0;
                    }
                }

                println!("{}", hostname);
            }
        }
        1 => xsethostname(matches.free.last().unwrap()),
        _ => help_menu(program, opts)
    };

    0
}

fn version() {
    println!("hostname 1.0.0");
}

fn help_menu(program: &str, options: Options) {
    version();
    println!("");
    println!("Usage:");
    println!("  {} [OPTION]... [HOSTNAME]", program);
    println!("");
    print!("{}", options.usage("Print or set the system's host name."));
}

fn xgethostname() -> String {
    let namelen = 256usize;
    let mut name : Vec<u8> = repeat(0).take(namelen).collect();
    let err = unsafe {
        gethostname (name.as_mut_ptr() as *mut libc::c_char,
                                        namelen as libc::size_t)
    };

    if err != 0 {
        panic!("Cannot determine hostname");
    }

    let last_char = name.iter().position(|byte| *byte == 0).unwrap_or(namelen);

    str::from_utf8(&name[..last_char]).unwrap().to_string()
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
fn xsethostname(name: &str) {
    let vec_name: Vec<libc::c_char> = name.bytes().map(|c| c as libc::c_char).collect();

    let err = unsafe {
        sethostname (vec_name.as_ptr(), vec_name.len() as libc::c_int)
    };

    if err != 0 {
        println!("Cannot set hostname to {}", name);
    }
}

#[cfg(target_os = "linux")]
fn xsethostname(name: &str) {
    let vec_name: Vec<libc::c_char> = name.bytes().map(|c| c as libc::c_char).collect();

    let err = unsafe {
        sethostname (vec_name.as_ptr(), vec_name.len() as libc::size_t)
    };

    if err != 0 {
        println!("Cannot set hostname to {}", name);
    }
}

#[allow(dead_code)]
fn main() {
    std::process::exit(uumain(std::env::args().collect()));
}
<commit_msg>Use name/version when showing hostname's usage.<commit_after>#![crate_name = "hostname"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Alan Andrade <alan.andradec@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 *
 * Synced with:
 *
 * https://www.opensource.apple.com/source/shell_cmds/shell_cmds-170/hostname/hostname.c?txt
 */

extern crate getopts;
extern crate libc;

#[macro_use]
extern crate uucore;

use getopts::Options;
use std::collections::hash_set::HashSet;
use std::iter::repeat;
use std::str;
use std::io::Write;
use std::net::ToSocketAddrs;

static NAME: &'static str = "hostname";
static VERSION: &'static str = "1.0.0";

extern {
    fn gethostname(name: *mut libc::c_char, namelen: libc::size_t) -> libc::c_int;
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
extern {
    fn sethostname(name: *const libc::c_char, namelen: libc::c_int) -> libc::c_int;
}

#[cfg(target_os = "linux")]
extern {
    fn sethostname(name: *const libc::c_char, namelen: libc::size_t) -> libc::c_int;
}

pub fn uumain(args: Vec<String>) -> i32 {
    let program = &args[0];

    let mut opts = Options::new();
    opts.optflag("d", "domain", "Display the name of the DNS domain if possible");
    opts.optflag("i", "ip-address", "Display the network address(es) of the host");
    opts.optflag("f", "fqdn", "Display the FQDN (Fully Qualified Domain Name) (default)");   // TODO: support --long
    opts.optflag("s", "short", "Display the short hostname (the portion before the first dot) if possible");
    opts.optflag("h", "help", "Show help");
    opts.optflag("V", "version", "Show program's version");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        _ => { help_menu(program, opts); return 0; }
    };

    if matches.opt_present("h") {
        help_menu(program, opts);
        return 0
    }
    if matches.opt_present("V") { version(); return 0 }

    match matches.free.len() {
        0 => {
            let hostname = xgethostname();

            if matches.opt_present("i") {
                match hostname.to_socket_addrs() {
                    Ok(addresses) => {
                        let mut hashset = HashSet::new();
                        let mut output = String::new();
                        for addr in addresses {
                            // XXX: not sure why this is necessary...
                            if !hashset.contains(&addr) {
                                output.push_str(&format!("{}", addr));
                                output.push_str(" ");
                                hashset.insert(addr.clone());
                            }
                        }
                        let len = output.len();
                        if len > 0 {
                            println!("{}", &output[0 .. len - 1]);
                        }
                    }
                    Err(f) => {
                        show_error!("{}", f);
                        return 1;
                    }
                }
            } else {
                if matches.opt_present("s") {
                    let mut it = hostname.char_indices().filter(|&ci| ci.1 == '.');
                    let ci = it.next();
                    if ci.is_some() {
                        println!("{}", &hostname[0 .. ci.unwrap().0]);
                        return 0;
                    }
                } else if matches.opt_present("d") {
                    let mut it = hostname.char_indices().filter(|&ci| ci.1 == '.');
                    let ci = it.next();
                    if ci.is_some() {
                        println!("{}", &hostname[ci.unwrap().0 + 1 .. ]);
                        return 0;
                    }
                }

                println!("{}", hostname);
            }
        }
        1 => xsethostname(matches.free.last().unwrap()),
        _ => help_menu(program, opts)
    };

    0
}

fn version() {
    println!("{} {}", NAME, VERSION);
}

fn help_menu(program: &str, options: Options) {
    version();
    println!("");
    println!("Usage:");
    println!("  {} [OPTION]... [HOSTNAME]", program);
    println!("");
    print!("{}", options.usage("Print or set the system's host name."));
}

fn xgethostname() -> String {
    let namelen = 256usize;
    let mut name : Vec<u8> = repeat(0).take(namelen).collect();
    let err = unsafe {
        gethostname (name.as_mut_ptr() as *mut libc::c_char,
                                        namelen as libc::size_t)
    };

    if err != 0 {
        panic!("Cannot determine hostname");
    }

    let last_char = name.iter().position(|byte| *byte == 0).unwrap_or(namelen);

    str::from_utf8(&name[..last_char]).unwrap().to_string()
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
fn xsethostname(name: &str) {
    let vec_name: Vec<libc::c_char> = name.bytes().map(|c| c as libc::c_char).collect();

    let err = unsafe {
        sethostname (vec_name.as_ptr(), vec_name.len() as libc::c_int)
    };

    if err != 0 {
        println!("Cannot set hostname to {}", name);
    }
}

#[cfg(target_os = "linux")]
fn xsethostname(name: &str) {
    let vec_name: Vec<libc::c_char> = name.bytes().map(|c| c as libc::c_char).collect();

    let err = unsafe {
        sethostname (vec_name.as_ptr(), vec_name.len() as libc::size_t)
    };

    if err != 0 {
        println!("Cannot set hostname to {}", name);
    }
}

#[allow(dead_code)]
fn main() {
    std::process::exit(uumain(std::env::args().collect()));
}
<|endoftext|>