<commit_before>use std::collections::HashMap;
#[cfg(not(feature = "model"))]
use std::marker::PhantomData;

use crate::json::{from_number, Value};
#[cfg(feature = "model")]
use crate::model::channel::AttachmentType;
use crate::model::channel::MessageFlags;

/// A builder to create the inner content of a [`Webhook`]'s execution.
///
/// This is a structured way of cleanly creating the inner execution payload,
/// to reduce potential argument counts.
///
/// Refer to the documentation for [`execute_webhook`] on restrictions with
/// execution payloads and its fields.
///
/// # Examples
///
/// Creating two embeds, and then sending them as part of the delivery
/// payload of [`Webhook::execute`]:
///
/// ```rust,no_run
/// use serenity::http::Http;
/// use serenity::model::channel::Embed;
/// use serenity::utils::Colour;
///
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// # let http = Http::default();
/// let url = "https://discord.com/api/webhooks/245037420704169985/ig5AO-wdVWpCBtUUMxmgsWryqgsW3DChbKYOINftJ4DCrUbnkedoYZD0VOH1QLr-S3sV";
/// let webhook = http.get_webhook_from_url(url).await?;
///
/// let website = Embed::fake(|e| {
///     e.title("The Rust Language Website")
///         .description("Rust is a systems programming language.")
///         .colour(Colour::from_rgb(222, 165, 132))
/// });
///
/// let resources = Embed::fake(|e| {
///     e.title("Rust Resources")
///         .description("A few resources to help with learning Rust")
///         .colour(0xDEA584)
///         .field("The Rust Book", "A comprehensive resource for Rust.", false)
///         .field("Rust by Example", "A collection of Rust examples", false)
/// });
///
/// webhook
///     .execute(&http, false, |w| {
///         w.content("Here's some information on Rust:").embeds(vec![website, resources])
///     })
///     .await?;
/// #     Ok(())
/// # }
/// ```
///
/// [`Webhook`]: crate::model::webhook::Webhook
/// [`Webhook::execute`]: crate::model::webhook::Webhook::execute
/// [`execute_webhook`]: crate::http::client::Http::execute_webhook
#[derive(Clone, Debug)]
pub struct ExecuteWebhook<'a>(
    pub HashMap<&'static str, Value>,
    #[cfg(feature = "model")] pub Vec<AttachmentType<'a>>,
    #[cfg(not(feature = "model"))] PhantomData<&'a ()>,
);

impl<'a> ExecuteWebhook<'a> {
    /// Override the default avatar of the webhook with an image URL.
    ///
    /// # Examples
    ///
    /// Overriding the default avatar:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let avatar_url = "https://i.imgur.com/KTs6whd.jpg";
    ///
    /// webhook.execute(&http, false, |w| w.avatar_url(avatar_url).content("Here's a webhook")).await?;
    /// #     Ok(())
    /// # }
    /// ```
    pub fn avatar_url<S: ToString>(&mut self, avatar_url: S) -> &mut Self {
        self.0.insert("avatar_url", Value::from(avatar_url.to_string()));
        self
    }

    /// Set the content of the message.
    ///
    /// Note that when setting at least one embed via [`Self::embeds`], this may be
    /// omitted.
    ///
    /// # Examples
    ///
    /// Sending a webhook with a content of `"foo"`:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook.execute(&http, false, |w| w.content("foo")).await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn content<S: ToString>(&mut self, content: S) -> &mut Self {
        self.0.insert("content", Value::from(content.to_string()));
        self
    }

    /// Appends a file to the webhook message.
    #[cfg(feature = "model")]
    pub fn add_file<T: Into<AttachmentType<'a>>>(&mut self, file: T) -> &mut Self {
        self.1.push(file.into());
        self
    }

    /// Appends a list of files to the webhook message.
    #[cfg(feature = "model")]
    pub fn add_files<T: Into<AttachmentType<'a>>, It: IntoIterator<Item = T>>(
        &mut self,
        files: It,
    ) -> &mut Self {
        self.1.extend(files.into_iter().map(|f| f.into()));
        self
    }

    /// Sets a list of files to include in the webhook message.
    ///
    /// Calling this multiple times will overwrite the file list.
    /// To append files, call [`Self::add_file`] or [`Self::add_files`] instead.
    #[cfg(feature = "model")]
    pub fn files<T: Into<AttachmentType<'a>>, It: IntoIterator<Item = T>>(
        &mut self,
        files: It,
    ) -> &mut Self {
        self.1 = files.into_iter().map(|f| f.into()).collect();
        self
    }

    /// Set the embeds associated with the message.
    ///
    /// This should be used in combination with [`Embed::fake`], creating one
    /// or more fake embeds to send to the API.
    ///
    /// # Examples
    ///
    /// Refer to the [struct-level documentation] for an example on how to use
    /// embeds.
    ///
    /// [`Embed::fake`]: crate::model::channel::Embed::fake
    /// [`Webhook::execute`]: crate::model::webhook::Webhook::execute
    /// [struct-level documentation]: #examples
    pub fn embeds(&mut self, embeds: Vec<Value>) -> &mut Self {
        self.0.insert("embeds", Value::from(embeds));
        self
    }

    /// Whether the message is a text-to-speech message.
    ///
    /// # Examples
    ///
    /// Sending a webhook with text-to-speech enabled:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook.execute(&http, false, |w| w.content("hello").tts(true)).await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn tts(&mut self, tts: bool) -> &mut Self {
        self.0.insert("tts", Value::from(tts));
        self
    }

    /// Override the default username of the webhook.
    ///
    /// # Examples
    ///
    /// Overriding the username to `"hakase"`:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook.execute(&http, false, |w| w.content("hello").username("hakase")).await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn username<S: ToString>(&mut self, username: S) -> &mut Self {
        self.0.insert("username", Value::from(username.to_string()));
        self
    }

    /// Sets the flags for the message.
    ///
    /// # Examples
    ///
    /// Suppressing an embed on the message.
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// # use serenity::model::channel::MessageFlags;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook
    ///     .execute(&http, false, |w| {
    ///         w.content("https://docs.rs/serenity/latest/serenity/")
    ///             .flags(MessageFlags::SUPPRESS_EMBEDS)
    ///     })
    ///     .await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn flags(&mut self, flags: MessageFlags) -> &mut Self {
        self.0.insert("flags", from_number(flags.bits()));
        self
    }
}

impl<'a> Default for ExecuteWebhook<'a> {
    /// Returns a default set of values for a [`Webhook`] execution.
    ///
    /// The only default value is [`Self::tts`] being set to `false`.
    ///
    /// # Examples
    ///
    /// Creating an [`ExecuteWebhook`] builder:
    ///
    /// ```rust
    /// use serenity::builder::ExecuteWebhook;
    ///
    /// let executor = ExecuteWebhook::default();
    /// ```
    ///
    /// [`Webhook`]: crate::model::webhook::Webhook
    fn default() -> ExecuteWebhook<'a> {
        let mut map = HashMap::new();
        map.insert("tts", Value::from(false));

        ExecuteWebhook(map, Default::default())
    }
}
<commit_msg>Add support for `allowed_mentions` in the `ExecuteWebhook` builder  (#1665)<commit_after>use std::collections::HashMap;
#[cfg(not(feature = "model"))]
use std::marker::PhantomData;

use super::CreateAllowedMentions;
use crate::json::{self, from_number, Value};
#[cfg(feature = "model")]
use crate::model::channel::AttachmentType;
use crate::model::channel::MessageFlags;

/// A builder to create the inner content of a [`Webhook`]'s execution.
///
/// This is a structured way of cleanly creating the inner execution payload,
/// to reduce potential argument counts.
///
/// Refer to the documentation for [`execute_webhook`] on restrictions with
/// execution payloads and its fields.
///
/// # Examples
///
/// Creating two embeds, and then sending them as part of the delivery
/// payload of [`Webhook::execute`]:
///
/// ```rust,no_run
/// use serenity::http::Http;
/// use serenity::model::channel::Embed;
/// use serenity::utils::Colour;
///
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// # let http = Http::default();
/// let url = "https://discord.com/api/webhooks/245037420704169985/ig5AO-wdVWpCBtUUMxmgsWryqgsW3DChbKYOINftJ4DCrUbnkedoYZD0VOH1QLr-S3sV";
/// let webhook = http.get_webhook_from_url(url).await?;
///
/// let website = Embed::fake(|e| {
///     e.title("The Rust Language Website")
///         .description("Rust is a systems programming language.")
///         .colour(Colour::from_rgb(222, 165, 132))
/// });
///
/// let resources = Embed::fake(|e| {
///     e.title("Rust Resources")
///         .description("A few resources to help with learning Rust")
///         .colour(0xDEA584)
///         .field("The Rust Book", "A comprehensive resource for Rust.", false)
///         .field("Rust by Example", "A collection of Rust examples", false)
/// });
///
/// webhook
///     .execute(&http, false, |w| {
///         w.content("Here's some information on Rust:").embeds(vec![website, resources])
///     })
///     .await?;
/// #     Ok(())
/// # }
/// ```
///
/// [`Webhook`]: crate::model::webhook::Webhook
/// [`Webhook::execute`]: crate::model::webhook::Webhook::execute
/// [`execute_webhook`]: crate::http::client::Http::execute_webhook
#[derive(Clone, Debug)]
pub struct ExecuteWebhook<'a>(
    pub HashMap<&'static str, Value>,
    #[cfg(feature = "model")] pub Vec<AttachmentType<'a>>,
    #[cfg(not(feature = "model"))] PhantomData<&'a ()>,
);

impl<'a> ExecuteWebhook<'a> {
    /// Override the default avatar of the webhook with an image URL.
    ///
    /// # Examples
    ///
    /// Overriding the default avatar:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let avatar_url = "https://i.imgur.com/KTs6whd.jpg";
    ///
    /// webhook.execute(&http, false, |w| w.avatar_url(avatar_url).content("Here's a webhook")).await?;
    /// #     Ok(())
    /// # }
    /// ```
    pub fn avatar_url<S: ToString>(&mut self, avatar_url: S) -> &mut Self {
        self.0.insert("avatar_url", Value::from(avatar_url.to_string()));
        self
    }

    /// Set the content of the message.
    ///
    /// Note that when setting at least one embed via [`Self::embeds`], this may be
    /// omitted.
    ///
    /// # Examples
    ///
    /// Sending a webhook with a content of `"foo"`:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook.execute(&http, false, |w| w.content("foo")).await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn content<S: ToString>(&mut self, content: S) -> &mut Self {
        self.0.insert("content", Value::from(content.to_string()));
        self
    }

    /// Appends a file to the webhook message.
    #[cfg(feature = "model")]
    pub fn add_file<T: Into<AttachmentType<'a>>>(&mut self, file: T) -> &mut Self {
        self.1.push(file.into());
        self
    }

    /// Appends a list of files to the webhook message.
    #[cfg(feature = "model")]
    pub fn add_files<T: Into<AttachmentType<'a>>, It: IntoIterator<Item = T>>(
        &mut self,
        files: It,
    ) -> &mut Self {
        self.1.extend(files.into_iter().map(|f| f.into()));
        self
    }

    /// Sets a list of files to include in the webhook message.
    ///
    /// Calling this multiple times will overwrite the file list.
    /// To append files, call [`Self::add_file`] or [`Self::add_files`] instead.
    #[cfg(feature = "model")]
    pub fn files<T: Into<AttachmentType<'a>>, It: IntoIterator<Item = T>>(
        &mut self,
        files: It,
    ) -> &mut Self {
        self.1 = files.into_iter().map(|f| f.into()).collect();
        self
    }

    /// Set the allowed mentions for the message.
    pub fn allowed_mentions<F>(&mut self, f: F) -> &mut Self
    where
        F: FnOnce(&mut CreateAllowedMentions) -> &mut CreateAllowedMentions,
    {
        let mut allowed_mentions = CreateAllowedMentions::default();
        f(&mut allowed_mentions);
        let map = json::hashmap_to_json_map(allowed_mentions.0);
        let allowed_mentions = Value::from(map);

        self.0.insert("allowed_mentions", allowed_mentions);
        self
    }

    /// Set the embeds associated with the message.
    ///
    /// This should be used in combination with [`Embed::fake`], creating one
    /// or more fake embeds to send to the API.
    ///
    /// # Examples
    ///
    /// Refer to the [struct-level documentation] for an example on how to use
    /// embeds.
    ///
    /// [`Embed::fake`]: crate::model::channel::Embed::fake
    /// [`Webhook::execute`]: crate::model::webhook::Webhook::execute
    /// [struct-level documentation]: #examples
    pub fn embeds(&mut self, embeds: Vec<Value>) -> &mut Self {
        self.0.insert("embeds", Value::from(embeds));
        self
    }

    /// Whether the message is a text-to-speech message.
    ///
    /// # Examples
    ///
    /// Sending a webhook with text-to-speech enabled:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook.execute(&http, false, |w| w.content("hello").tts(true)).await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn tts(&mut self, tts: bool) -> &mut Self {
        self.0.insert("tts", Value::from(tts));
        self
    }

    /// Override the default username of the webhook.
    ///
    /// # Examples
    ///
    /// Overriding the username to `"hakase"`:
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook.execute(&http, false, |w| w.content("hello").username("hakase")).await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn username<S: ToString>(&mut self, username: S) -> &mut Self {
        self.0.insert("username", Value::from(username.to_string()));
        self
    }

    /// Sets the flags for the message.
    ///
    /// # Examples
    ///
    /// Suppressing an embed on the message.
    ///
    /// ```rust,no_run
    /// # use serenity::http::Http;
    /// # use serenity::model::channel::MessageFlags;
    /// #
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// # let http = Http::default();
    /// # let webhook = http.get_webhook_with_token(0, "").await?;
    /// #
    /// let execution = webhook
    ///     .execute(&http, false, |w| {
    ///         w.content("https://docs.rs/serenity/latest/serenity/")
    ///             .flags(MessageFlags::SUPPRESS_EMBEDS)
    ///     })
    ///     .await;
    ///
    /// if let Err(why) = execution {
    ///     println!("Err sending webhook: {:?}", why);
    /// }
    /// #     Ok(())
    /// # }
    /// ```
    pub fn flags(&mut self, flags: MessageFlags) -> &mut Self {
        self.0.insert("flags", from_number(flags.bits()));
        self
    }
}

impl<'a> Default for ExecuteWebhook<'a> {
    /// Returns a default set of values for a [`Webhook`] execution.
    ///
    /// The only default value is [`Self::tts`] being set to `false`.
    ///
    /// # Examples
    ///
    /// Creating an [`ExecuteWebhook`] builder:
    ///
    /// ```rust
    /// use serenity::builder::ExecuteWebhook;
    ///
    /// let executor = ExecuteWebhook::default();
    /// ```
    ///
    /// [`Webhook`]: crate::model::webhook::Webhook
    fn default() -> ExecuteWebhook<'a> {
        let mut map = HashMap::new();
        map.insert("tts", Value::from(false));

        ExecuteWebhook(map, Default::default())
    }
}
<|endoftext|>


<commit_before>use std::fmt::{mod, Show, Formatter};
use std::io::{USER_DIR};
use std::io::fs::{mkdir_recursive, rmdir_recursive, PathExtensions};
use serialize::{Encodable, Encoder};
use url::Url;
use git2;

use util::{CargoResult, ChainError, human, ToUrl, internal, Require};

#[deriving(PartialEq,Clone,Encodable)]
pub enum GitReference {
    Master,
    Other(String)
}

#[deriving(PartialEq,Clone,Encodable)]
pub struct GitRevision(String);

impl GitReference {
    pub fn for_str<S: Str>(string: S) -> GitReference {
        if string.as_slice() == "master" {
            GitReference::Master
        } else {
            GitReference::Other(string.as_slice().to_string())
        }
    }
}

impl Str for GitReference {
    fn as_slice(&self) -> &str {
        match *self {
            GitReference::Master => "master",
            GitReference::Other(ref string) => string.as_slice()
        }
    }
}

impl Show for GitReference {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

impl Str for GitRevision {
    fn as_slice(&self) -> &str {
        let GitRevision(ref me) = *self;
        me.as_slice()
    }
}

impl Show for GitRevision {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

/// GitRemote represents a remote repository. It gets cloned into a local
/// GitDatabase.
#[deriving(PartialEq,Clone,Show)]
pub struct GitRemote {
    url: Url,
}

#[deriving(PartialEq,Clone,Encodable)]
struct EncodableGitRemote {
    url: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitRemote {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitRemote {
            url: self.url.to_string()
        }.encode(s)
    }
}

/// GitDatabase is a local clone of a remote repository's database. Multiple
/// GitCheckouts can be cloned from this GitDatabase.
pub struct GitDatabase {
    remote: GitRemote,
    path: Path,
    repo: git2::Repository,
}

#[deriving(Encodable)]
pub struct EncodableGitDatabase {
    remote: GitRemote,
    path: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitDatabase {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitDatabase {
            remote: self.remote.clone(),
            path: self.path.display().to_string()
        }.encode(s)
    }
}

/// GitCheckout is a local checkout of a particular revision. Calling
/// `clone_into` with a reference will resolve the reference into a revision,
/// and return a CargoError if no revision for that reference was found.
pub struct GitCheckout<'a> {
    database: &'a GitDatabase,
    location: Path,
    revision: GitRevision,
    repo: git2::Repository,
}

#[deriving(Encodable)]
pub struct EncodableGitCheckout {
    database: EncodableGitDatabase,
    location: String,
    revision: String,
}

impl<'a, E, S: Encoder<E>> Encodable<S, E> for GitCheckout<'a> {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitCheckout {
            location: self.location.display().to_string(),
            revision: self.revision.to_string(),
            database: EncodableGitDatabase {
                remote: self.database.remote.clone(),
                path: self.database.path.display().to_string(),
            },
        }.encode(s)
    }
}

// Implementations

impl GitRemote {
    pub fn new(url: &Url) -> GitRemote {
        GitRemote { url: url.clone() }
    }

    pub fn get_url(&self) -> &Url {
        &self.url
    }

    pub fn rev_for<S: Str>(&self, path: &Path, reference: S)
                           -> CargoResult<GitRevision> {
        let db = try!(self.db_at(path));
        db.rev_for(reference)
    }

    pub fn checkout(&self, into: &Path) -> CargoResult<GitDatabase> {
        let repo = match git2::Repository::open(into) {
            Ok(repo) => {
                try!(self.fetch_into(&repo).chain_error(|| {
                    internal(format!("failed to fetch into {}", into.display()))
                }));
                repo
            }
            Err(..) => {
                try!(self.clone_into(into).chain_error(|| {
                    internal(format!("failed to clone into: {}", into.display()))
                }))
            }
        };

        Ok(GitDatabase { remote: self.clone(), path: into.clone(), repo: repo })
    }

    pub fn db_at(&self, db_path: &Path) -> CargoResult<GitDatabase> {
        let repo = try!(git2::Repository::open(db_path));
        Ok(GitDatabase {
            remote: self.clone(),
            path: db_path.clone(),
            repo: repo,
        })
    }

    fn fetch_into(&self, dst: &git2::Repository) -> CargoResult<()> {
        // Create a local anonymous remote in the repository to fetch the url
        let url = self.url.to_string();
        let refspec = "refs/heads/*:refs/heads/*";
        fetch(dst, url.as_slice(), refspec)
    }

    fn clone_into(&self, dst: &Path) -> CargoResult<git2::Repository> {
        let url = self.url.to_string();
        if dst.exists() {
            try!(rmdir_recursive(dst));
        }
        try!(mkdir_recursive(dst, USER_DIR));
        let repo = try!(git2::Repository::init_bare(dst));
        try!(fetch(&repo, url.as_slice(), "refs/heads/*:refs/heads/*"));
        Ok(repo)
    }
}

impl GitDatabase {
    fn get_path<'a>(&'a self) -> &'a Path {
        &self.path
    }

    pub fn copy_to(&self, rev: GitRevision, dest: &Path)
                   -> CargoResult<GitCheckout> {
        let checkout = match git2::Repository::open(dest) {
            Ok(repo) => {
                let checkout = GitCheckout::new(dest, self, rev, repo);
                if !checkout.is_fresh() {
                    try!(checkout.fetch());
                    try!(checkout.reset());
                    assert!(checkout.is_fresh());
                }
                checkout
            }
            Err(..) => try!(GitCheckout::clone_into(dest, self, rev)),
        };
        try!(checkout.update_submodules().chain_error(|| {
            internal("failed to update submodules")
        }));
        Ok(checkout)
    }

    pub fn rev_for<S: Str>(&self, reference: S) -> CargoResult<GitRevision> {
        let rev = try!(self.repo.revparse_single(reference.as_slice()));
        Ok(GitRevision(rev.id().to_string()))
    }

    pub fn has_ref<S: Str>(&self, reference: S) -> CargoResult<()> {
        try!(self.repo.revparse_single(reference.as_slice()));
        Ok(())
    }
}

impl<'a> GitCheckout<'a> {
    fn new<'a>(path: &Path, database: &'a GitDatabase, revision: GitRevision,
               repo: git2::Repository)
               -> GitCheckout<'a>
    {
        GitCheckout {
            location: path.clone(),
            database: database,
            revision: revision,
            repo: repo,
        }
    }

    fn clone_into<'a>(into: &Path, database: &'a GitDatabase,
                      revision: GitRevision)
                      -> CargoResult<GitCheckout<'a>>
    {
        let repo = try!(GitCheckout::clone_repo(database.get_path(), into));
        let checkout = GitCheckout::new(into, database, revision, repo);
        try!(checkout.reset());
        Ok(checkout)
    }

    pub fn get_rev(&self) -> &str {
        self.revision.as_slice()
    }

    fn clone_repo(source: &Path, into: &Path) -> CargoResult<git2::Repository> {
        let dirname = into.dir_path();

        try!(mkdir_recursive(&dirname, USER_DIR).chain_error(|| {
            human(format!("Couldn't mkdir {}", dirname.display()))
        }));

        if into.exists() {
            try!(rmdir_recursive(into).chain_error(|| {
                human(format!("Couldn't rmdir {}", into.display()))
            }));
        }

        let url = try!(source.to_url().map_err(human));
        let url = url.to_string();
        let repo = try!(git2::Repository::clone(url.as_slice(),
                                                into).chain_error(|| {
            internal(format!("failed to clone {} into {}", source.display(),
                             into.display()))
        }));
        Ok(repo)
    }

    fn is_fresh(&self) -> bool {
        match self.repo.revparse_single("HEAD") {
            Ok(head) => head.id().to_string() == self.revision.to_string(),
            _ => false,
        }
    }

    fn fetch(&self) -> CargoResult<()> {
        info!("fetch {}", self.repo.path().display());
        let url = try!(self.database.path.to_url().map_err(human));
        let url = url.to_string();
        let refspec = "refs/heads/*:refs/heads/*";
        try!(fetch(&self.repo, url.as_slice(), refspec));
        Ok(())
    }

    fn reset(&self) -> CargoResult<()> {
        info!("reset {} to {}", self.repo.path().display(),
              self.revision.as_slice());
        let oid = try!(git2::Oid::from_str(self.revision.as_slice()));
        let object = try!(self.repo.find_object(oid, None));
        try!(self.repo.reset(&object, git2::Hard, None, None));
        Ok(())
    }

    fn update_submodules(&self) -> CargoResult<()> {
        return update_submodules(&self.repo);

        fn update_submodules(repo: &git2::Repository) -> CargoResult<()> {
            info!("update submodules for: {}", repo.path().display());

            for mut child in try!(repo.submodules()).into_iter() {
                try!(child.init(false));
                let url = try!(child.url().require(|| {
                    internal("non-utf8 url for submodule")
                }));

                // A submodule which is listed in .gitmodules but not actually
                // checked out will not have a head id, so we should ignore it.
                let head = match child.head_id() {
                    Some(head) => head,
                    None => continue,
                };

                // If the submodule hasn't been checked out yet, we need to
                // clone it. If it has been checked out and the head is the same
                // as the submodule's head, then we can bail out and go to the
                // next submodule.
                let repo = match child.open() {
                    Ok(repo) => {
                        if child.head_id() == try!(repo.head()).target() {
                            continue
                        }
                        repo
                    }
                    Err(..) => {
                        let path = repo.path().dir_path().join(child.path());
                        try!(git2::Repository::clone(url, &path))
                    }
                };

                // Fetch data from origin and reset to the head commit
                let refspec = "refs/heads/*:refs/heads/*";
                try!(fetch(&repo, url, refspec).chain_error(|| {
                    internal(format!("failed to fetch submodule `{}` from {}",
                                     child.name().unwrap_or(""), url))
                }));

                let obj = try!(repo.find_object(head, None));
                try!(repo.reset(&obj, git2::Hard, None, None));
                try!(update_submodules(&repo));
            }
            Ok(())
        }
    }
}

fn with_authentication<T>(url: &str,
                          cfg: &git2::Config,
                          f: |git2::Credentials| -> CargoResult<T>)
                          -> CargoResult<T> {
    // Prepare the authentication callbacks.
    //
    // We check the `allowed` types of credentials, and we try to do as much as
    // possible based on that:
    //
    // * Prioritize SSH keys from the local ssh agent as they're likely the most
    //   reliable. The username here is prioritized from the credential
    //   callback, then from whatever is configured in git itself, and finally
    //   we fall back to the generic user of `git`.
    //
    // * If a username/password is allowed, then we fallback to git2-rs's
    //   implementation of the credential helper. This is what is configured
    //   with `credential.helper` in git, and is the interface for the OSX
    //   keychain, for example.
    //
    // * After the above two have failed, we just kinda grapple attempting to
    //   return *something*.
    let mut cred_helper = git2::CredentialHelper::new(url);
    cred_helper.config(cfg);
    let mut cred_error = false;
    let ret = f(|url, username, allowed| {
        let creds = if allowed.contains(git2::SSH_KEY) {
            let user = username.map(|s| s.to_string())
                               .or_else(|| cred_helper.username.clone())
                               .unwrap_or("git".to_string());
            git2::Cred::ssh_key_from_agent(user.as_slice())
        } else if allowed.contains(git2::USER_PASS_PLAINTEXT) {
            git2::Cred::credential_helper(cfg, url, username)
        } else if allowed.contains(git2::DEFAULT) {
            git2::Cred::default()
        } else {
            Err(git2::Error::from_str("no authentication available"))
        };
        cred_error = creds.is_err();
        creds
    });
    if cred_error {
        ret.chain_error(|| {
            human("Failed to authenticate when downloading repository")
        })
    } else {
        ret
    }
}

pub fn fetch(repo: &git2::Repository, url: &str,
             refspec: &str) -> CargoResult<()> {
    // Create a local anonymous remote in the repository to fetch the url

    with_authentication(url, &try!(repo.config()), |f| {
        let mut cb = git2::RemoteCallbacks::new()
                                       .credentials(f);
        let mut remote = try!(repo.remote_anonymous(url.as_slice(), refspec));
        try!(remote.add_fetch("refs/tags/*:refs/tags/*"));
        remote.set_callbacks(&mut cb);
        try!(remote.fetch(None, None));
        Ok(())
    })
}
<commit_msg>Don't abort on possibly-corrupt submodules<commit_after>use std::fmt::{mod, Show, Formatter};
use std::io::{USER_DIR};
use std::io::fs::{mkdir_recursive, rmdir_recursive, PathExtensions};
use serialize::{Encodable, Encoder};
use url::Url;
use git2;

use util::{CargoResult, ChainError, human, ToUrl, internal, Require};

#[deriving(PartialEq,Clone,Encodable)]
pub enum GitReference {
    Master,
    Other(String)
}

#[deriving(PartialEq,Clone,Encodable)]
pub struct GitRevision(String);

impl GitReference {
    pub fn for_str<S: Str>(string: S) -> GitReference {
        if string.as_slice() == "master" {
            GitReference::Master
        } else {
            GitReference::Other(string.as_slice().to_string())
        }
    }
}

impl Str for GitReference {
    fn as_slice(&self) -> &str {
        match *self {
            GitReference::Master => "master",
            GitReference::Other(ref string) => string.as_slice()
        }
    }
}

impl Show for GitReference {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

impl Str for GitRevision {
    fn as_slice(&self) -> &str {
        let GitRevision(ref me) = *self;
        me.as_slice()
    }
}

impl Show for GitRevision {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

/// GitRemote represents a remote repository. It gets cloned into a local
/// GitDatabase.
#[deriving(PartialEq,Clone,Show)]
pub struct GitRemote {
    url: Url,
}

#[deriving(PartialEq,Clone,Encodable)]
struct EncodableGitRemote {
    url: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitRemote {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitRemote {
            url: self.url.to_string()
        }.encode(s)
    }
}

/// GitDatabase is a local clone of a remote repository's database. Multiple
/// GitCheckouts can be cloned from this GitDatabase.
pub struct GitDatabase {
    remote: GitRemote,
    path: Path,
    repo: git2::Repository,
}

#[deriving(Encodable)]
pub struct EncodableGitDatabase {
    remote: GitRemote,
    path: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitDatabase {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitDatabase {
            remote: self.remote.clone(),
            path: self.path.display().to_string()
        }.encode(s)
    }
}

/// GitCheckout is a local checkout of a particular revision. Calling
/// `clone_into` with a reference will resolve the reference into a revision,
/// and return a CargoError if no revision for that reference was found.
pub struct GitCheckout<'a> {
    database: &'a GitDatabase,
    location: Path,
    revision: GitRevision,
    repo: git2::Repository,
}

#[deriving(Encodable)]
pub struct EncodableGitCheckout {
    database: EncodableGitDatabase,
    location: String,
    revision: String,
}

impl<'a, E, S: Encoder<E>> Encodable<S, E> for GitCheckout<'a> {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitCheckout {
            location: self.location.display().to_string(),
            revision: self.revision.to_string(),
            database: EncodableGitDatabase {
                remote: self.database.remote.clone(),
                path: self.database.path.display().to_string(),
            },
        }.encode(s)
    }
}

// Implementations

impl GitRemote {
    pub fn new(url: &Url) -> GitRemote {
        GitRemote { url: url.clone() }
    }

    pub fn get_url(&self) -> &Url {
        &self.url
    }

    pub fn rev_for<S: Str>(&self, path: &Path, reference: S)
                           -> CargoResult<GitRevision> {
        let db = try!(self.db_at(path));
        db.rev_for(reference)
    }

    pub fn checkout(&self, into: &Path) -> CargoResult<GitDatabase> {
        let repo = match git2::Repository::open(into) {
            Ok(repo) => {
                try!(self.fetch_into(&repo).chain_error(|| {
                    internal(format!("failed to fetch into {}", into.display()))
                }));
                repo
            }
            Err(..) => {
                try!(self.clone_into(into).chain_error(|| {
                    internal(format!("failed to clone into: {}", into.display()))
                }))
            }
        };

        Ok(GitDatabase { remote: self.clone(), path: into.clone(), repo: repo })
    }

    pub fn db_at(&self, db_path: &Path) -> CargoResult<GitDatabase> {
        let repo = try!(git2::Repository::open(db_path));
        Ok(GitDatabase {
            remote: self.clone(),
            path: db_path.clone(),
            repo: repo,
        })
    }

    fn fetch_into(&self, dst: &git2::Repository) -> CargoResult<()> {
        // Create a local anonymous remote in the repository to fetch the url
        let url = self.url.to_string();
        let refspec = "refs/heads/*:refs/heads/*";
        fetch(dst, url.as_slice(), refspec)
    }

    fn clone_into(&self, dst: &Path) -> CargoResult<git2::Repository> {
        let url = self.url.to_string();
        if dst.exists() {
            try!(rmdir_recursive(dst));
        }
        try!(mkdir_recursive(dst, USER_DIR));
        let repo = try!(git2::Repository::init_bare(dst));
        try!(fetch(&repo, url.as_slice(), "refs/heads/*:refs/heads/*"));
        Ok(repo)
    }
}

impl GitDatabase {
    fn get_path<'a>(&'a self) -> &'a Path {
        &self.path
    }

    pub fn copy_to(&self, rev: GitRevision, dest: &Path)
                   -> CargoResult<GitCheckout> {
        let checkout = match git2::Repository::open(dest) {
            Ok(repo) => {
                let checkout = GitCheckout::new(dest, self, rev, repo);
                if !checkout.is_fresh() {
                    try!(checkout.fetch());
                    try!(checkout.reset());
                    assert!(checkout.is_fresh());
                }
                checkout
            }
            Err(..) => try!(GitCheckout::clone_into(dest, self, rev)),
        };
        try!(checkout.update_submodules().chain_error(|| {
            internal("failed to update submodules")
        }));
        Ok(checkout)
    }

    pub fn rev_for<S: Str>(&self, reference: S) -> CargoResult<GitRevision> {
        let rev = try!(self.repo.revparse_single(reference.as_slice()));
        Ok(GitRevision(rev.id().to_string()))
    }

    pub fn has_ref<S: Str>(&self, reference: S) -> CargoResult<()> {
        try!(self.repo.revparse_single(reference.as_slice()));
        Ok(())
    }
}

impl<'a> GitCheckout<'a> {
    fn new<'a>(path: &Path, database: &'a GitDatabase, revision: GitRevision,
               repo: git2::Repository)
               -> GitCheckout<'a>
    {
        GitCheckout {
            location: path.clone(),
            database: database,
            revision: revision,
            repo: repo,
        }
    }

    fn clone_into<'a>(into: &Path, database: &'a GitDatabase,
                      revision: GitRevision)
                      -> CargoResult<GitCheckout<'a>>
    {
        let repo = try!(GitCheckout::clone_repo(database.get_path(), into));
        let checkout = GitCheckout::new(into, database, revision, repo);
        try!(checkout.reset());
        Ok(checkout)
    }

    pub fn get_rev(&self) -> &str {
        self.revision.as_slice()
    }

    fn clone_repo(source: &Path, into: &Path) -> CargoResult<git2::Repository> {
        let dirname = into.dir_path();

        try!(mkdir_recursive(&dirname, USER_DIR).chain_error(|| {
            human(format!("Couldn't mkdir {}", dirname.display()))
        }));

        if into.exists() {
            try!(rmdir_recursive(into).chain_error(|| {
                human(format!("Couldn't rmdir {}", into.display()))
            }));
        }

        let url = try!(source.to_url().map_err(human));
        let url = url.to_string();
        let repo = try!(git2::Repository::clone(url.as_slice(),
                                                into).chain_error(|| {
            internal(format!("failed to clone {} into {}", source.display(),
                             into.display()))
        }));
        Ok(repo)
    }

    fn is_fresh(&self) -> bool {
        match self.repo.revparse_single("HEAD") {
            Ok(head) => head.id().to_string() == self.revision.to_string(),
            _ => false,
        }
    }

    fn fetch(&self) -> CargoResult<()> {
        info!("fetch {}", self.repo.path().display());
        let url = try!(self.database.path.to_url().map_err(human));
        let url = url.to_string();
        let refspec = "refs/heads/*:refs/heads/*";
        try!(fetch(&self.repo, url.as_slice(), refspec));
        Ok(())
    }

    fn reset(&self) -> CargoResult<()> {
        info!("reset {} to {}", self.repo.path().display(),
              self.revision.as_slice());
        let oid = try!(git2::Oid::from_str(self.revision.as_slice()));
        let object = try!(self.repo.find_object(oid, None));
        try!(self.repo.reset(&object, git2::Hard, None, None));
        Ok(())
    }

    fn update_submodules(&self) -> CargoResult<()> {
        return update_submodules(&self.repo);

        fn update_submodules(repo: &git2::Repository) -> CargoResult<()> {
            info!("update submodules for: {}", repo.path().display());

            for mut child in try!(repo.submodules()).into_iter() {
                try!(child.init(false));
                let url = try!(child.url().require(|| {
                    internal("non-utf8 url for submodule")
                }));

                // A submodule which is listed in .gitmodules but not actually
                // checked out will not have a head id, so we should ignore it.
                let head = match child.head_id() {
                    Some(head) => head,
                    None => continue,
                };

                // If the submodule hasn't been checked out yet, we need to
                // clone it. If it has been checked out and the head is the same
                // as the submodule's head, then we can bail out and go to the
                // next submodule.
                let head_and_repo = child.open().and_then(|repo| {
                    Ok((try!(repo.head()).target(), repo))
                });
                let repo = match head_and_repo {
                    Ok((head, repo)) => {
                        if child.head_id() == head {
                            continue
                        }
                        repo
                    }
                    Err(..) => {
                        let path = repo.path().dir_path().join(child.path());
                        try!(git2::Repository::clone(url, &path))
                    }
                };

                // Fetch data from origin and reset to the head commit
                let refspec = "refs/heads/*:refs/heads/*";
                try!(fetch(&repo, url, refspec).chain_error(|| {
                    internal(format!("failed to fetch submodule `{}` from {}",
                                     child.name().unwrap_or(""), url))
                }));

                let obj = try!(repo.find_object(head, None));
                try!(repo.reset(&obj, git2::Hard, None, None));
                try!(update_submodules(&repo));
            }
            Ok(())
        }
    }
}

fn with_authentication<T>(url: &str,
                          cfg: &git2::Config,
                          f: |git2::Credentials| -> CargoResult<T>)
                          -> CargoResult<T> {
    // Prepare the authentication callbacks.
    //
    // We check the `allowed` types of credentials, and we try to do as much as
    // possible based on that:
    //
    // * Prioritize SSH keys from the local ssh agent as they're likely the most
    //   reliable. The username here is prioritized from the credential
    //   callback, then from whatever is configured in git itself, and finally
    //   we fall back to the generic user of `git`.
    //
    // * If a username/password is allowed, then we fallback to git2-rs's
    //   implementation of the credential helper. This is what is configured
    //   with `credential.helper` in git, and is the interface for the OSX
    //   keychain, for example.
    //
    // * After the above two have failed, we just kinda grapple attempting to
    //   return *something*.
    let mut cred_helper = git2::CredentialHelper::new(url);
    cred_helper.config(cfg);
    let mut cred_error = false;
    let ret = f(|url, username, allowed| {
        let creds = if allowed.contains(git2::SSH_KEY) {
            let user = username.map(|s| s.to_string())
                               .or_else(|| cred_helper.username.clone())
                               .unwrap_or("git".to_string());
            git2::Cred::ssh_key_from_agent(user.as_slice())
        } else if allowed.contains(git2::USER_PASS_PLAINTEXT) {
            git2::Cred::credential_helper(cfg, url, username)
        } else if allowed.contains(git2::DEFAULT) {
            git2::Cred::default()
        } else {
            Err(git2::Error::from_str("no authentication available"))
        };
        cred_error = creds.is_err();
        creds
    });
    if cred_error {
        ret.chain_error(|| {
            human("Failed to authenticate when downloading repository")
        })
    } else {
        ret
    }
}

pub fn fetch(repo: &git2::Repository, url: &str,
             refspec: &str) -> CargoResult<()> {
    // Create a local anonymous remote in the repository to fetch the url

    with_authentication(url, &try!(repo.config()), |f| {
        let mut cb = git2::RemoteCallbacks::new()
                                       .credentials(f);
        let mut remote = try!(repo.remote_anonymous(url.as_slice(), refspec));
        try!(remote.add_fetch("refs/tags/*:refs/tags/*"));
        remote.set_callbacks(&mut cb);
        try!(remote.fetch(None, None));
        Ok(())
    })
}
<|endoftext|>


<commit_before>use diesel::mysql::{Mysql, MysqlConnection};
use diesel::*;
use heck::ToUpperCamelCase;
use std::{borrow::Cow, error::Error};

use super::data_structures::*;
use super::information_schema::DefaultSchema;
use super::table_data::TableName;
use crate::print_schema::ColumnSorting;

diesel::sql_function! {
    #[sql_name = "NULLIF"]
    fn null_if_text(lhs: sql_types::Text, rhs: sql_types::Text) -> sql_types::Nullable<sql_types::Text>
}

pub fn get_table_data<'a>(
    conn: &mut MysqlConnection,
    table: &'a TableName,
    column_sorting: &ColumnSorting,
) -> QueryResult<Vec<ColumnInformation>> {
    use self::information_schema::columns::dsl::*;

    let schema_name = match table.schema {
        Some(ref name) => Cow::Borrowed(name),
        None => Cow::Owned(Mysql::default_schema(conn)?),
    };

    let type_schema = None::<String>.into_sql();
    let query = columns
        .select((
            column_name,
            column_type,
            type_schema,
            __is_nullable,
            // MySQL comments are not nullable and are empty strings if not set
            null_if_text(column_comment, ""),
        ))
        .filter(table_name.eq(&table.sql_name))
        .filter(table_schema.eq(schema_name));
    match column_sorting {
        ColumnSorting::OrdinalPosition => query.order(ordinal_position).load(conn),
        ColumnSorting::Name => query.order(column_name).load(conn),
    }
}

mod information_schema {
    use diesel::prelude::{allow_tables_to_appear_in_same_query, table};

    table! {
        information_schema.tables (table_schema, table_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            table_comment -> VarChar,
        }
    }

    table! {
        information_schema.table_constraints (constraint_schema, constraint_name) {
            table_schema -> VarChar,
            constraint_schema -> VarChar,
            constraint_name -> VarChar,
            constraint_type -> VarChar,
        }
    }

    table! {
        information_schema.key_column_usage (constraint_schema, constraint_name) {
            constraint_schema -> VarChar,
            constraint_name -> VarChar,
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            referenced_table_schema -> VarChar,
            referenced_table_name -> VarChar,
            referenced_column_name -> VarChar,
        }
    }

    table! {
        information_schema.columns (table_schema, table_name, column_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            #[sql_name = "is_nullable"]
            __is_nullable -> VarChar,
            ordinal_position -> BigInt,
            udt_name -> VarChar,
            udt_schema -> VarChar,
            column_type -> VarChar,
            column_comment -> VarChar,
        }
    }

    allow_tables_to_appear_in_same_query!(table_constraints, key_column_usage);
}

/// Even though this is using `information_schema`, MySQL needs non-ANSI columns
/// in order to do this.
pub fn load_foreign_key_constraints(
    connection: &mut MysqlConnection,
    schema_name: Option<&str>,
) -> QueryResult<Vec<ForeignKeyConstraint>> {
    use self::information_schema::key_column_usage as kcu;
    use self::information_schema::table_constraints as tc;

    let default_schema = Mysql::default_schema(connection)?;
    let schema_name = match schema_name {
        Some(name) => name,
        None => &default_schema,
    };

    let constraints = tc::table
        .filter(tc::constraint_type.eq("FOREIGN KEY"))
        .filter(tc::table_schema.eq(schema_name))
        .filter(kcu::referenced_column_name.is_not_null())
        .inner_join(
            kcu::table.on(tc::constraint_schema
                .eq(kcu::constraint_schema)
                .and(tc::constraint_name.eq(kcu::constraint_name))),
        )
        .select((
            (kcu::table_name, kcu::table_schema),
            (kcu::referenced_table_name, kcu::referenced_table_schema),
            kcu::column_name,
            kcu::referenced_column_name,
        ))
        .load::<(TableName, TableName, String, _)>(connection)?
        .into_iter()
        .map(
            |(mut child_table, mut parent_table, foreign_key, primary_key)| {
                child_table.strip_schema_if_matches(&default_schema);
                parent_table.strip_schema_if_matches(&default_schema);

                ForeignKeyConstraint {
                    child_table,
                    parent_table,
                    foreign_key: foreign_key.clone(),
                    foreign_key_rust_name: foreign_key,
                    primary_key,
                }
            },
        )
        .collect();
    Ok(constraints)
}

pub fn determine_column_type(
    attr: &ColumnInformation,
) -> Result<ColumnType, Box<dyn Error + Send + Sync + 'static>> {
    let tpe = determine_type_name(&attr.type_name)?;
    let unsigned = determine_unsigned(&attr.type_name);

    Ok(ColumnType {
        schema: None,
        sql_name: tpe.trim().to_lowercase(),
        rust_name: tpe.trim().to_upper_camel_case(),
        is_array: false,
        is_nullable: attr.nullable,
        is_unsigned: unsigned,
    })
}

pub fn get_table_comment(
    conn: &mut MysqlConnection,
    table: &TableName,
) -> QueryResult<Option<String>> {
    use self::information_schema::tables::dsl::*;

    let schema_name = match table.schema {
        Some(ref name) => Cow::Borrowed(name),
        None => Cow::Owned(Mysql::default_schema(conn)?),
    };

    let comment: String = tables
        .select(table_comment)
        .filter(table_name.eq(&table.sql_name))
        .filter(table_schema.eq(schema_name))
        .get_result(conn)?;

    if comment.is_empty() {
        Ok(None)
    } else {
        Ok(Some(comment))
    }
}

fn determine_type_name(
    sql_type_name: &str,
) -> Result<String, Box<dyn Error + Send + Sync + 'static>> {
    let result = if sql_type_name == "tinyint(1)" {
        "bool"
    } else if sql_type_name.starts_with("int") {
        "integer"
    } else if let Some(idx) = sql_type_name.find('(') {
        &sql_type_name[..idx]
    } else {
        sql_type_name
    };

    if determine_unsigned(result) {
        Ok(result
            .to_lowercase()
            .replace("unsigned", "")
            .trim()
            .to_owned())
    } else if result.contains(' ') {
        Err(format!("unrecognized type {:?}", result).into())
    } else {
        Ok(result.to_owned())
    }
}

fn determine_unsigned(sql_type_name: &str) -> bool {
    sql_type_name.to_lowercase().contains("unsigned")
}

#[test]
fn values_which_already_map_to_type_are_returned_unchanged() {
    assert_eq!("text", determine_type_name("text").unwrap());
    assert_eq!("integer", determine_type_name("integer").unwrap());
    assert_eq!("biginteger", determine_type_name("biginteger").unwrap());
}

#[test]
fn trailing_parenthesis_are_stripped() {
    assert_eq!("varchar", determine_type_name("varchar(255)").unwrap());
    assert_eq!("decimal", determine_type_name("decimal(10, 2)").unwrap());
    assert_eq!("float", determine_type_name("float(1)").unwrap());
}

#[test]
fn tinyint_is_bool_if_limit_1() {
    assert_eq!("bool", determine_type_name("tinyint(1)").unwrap());
    assert_eq!("tinyint", determine_type_name("tinyint(2)").unwrap());
}

#[test]
fn int_is_treated_as_integer() {
    assert_eq!("integer", determine_type_name("int").unwrap());
    assert_eq!("integer", determine_type_name("int(11)").unwrap());
}

#[test]
fn unsigned_types_are_supported() {
    assert!(determine_unsigned("float unsigned"));
    assert!(determine_unsigned("UNSIGNED INT"));
    assert!(determine_unsigned("unsigned bigint"));
    assert!(!determine_unsigned("bigint"));
    assert!(!determine_unsigned("FLOAT"));
    assert_eq!("float", determine_type_name("float unsigned").unwrap());
    assert_eq!("int", determine_type_name("UNSIGNED INT").unwrap());
    assert_eq!("bigint", determine_type_name("unsigned bigint").unwrap());
}

#[test]
fn types_with_space_are_not_supported() {
    assert!(determine_type_name("lol wat").is_err());
}

#[cfg(test)]
mod test {
    extern crate dotenvy;

    use self::dotenvy::dotenv;
    use super::*;
    use std::env;

    fn connection() -> MysqlConnection {
        dotenv().ok();

        let connection_url = env::var("MYSQL_DATABASE_URL")
            .or_else(|_| env::var("DATABASE_URL"))
            .expect("DATABASE_URL must be set in order to run tests");
        let mut connection = MysqlConnection::establish(&connection_url).unwrap();
        connection.begin_test_transaction().unwrap();
        connection
    }

    #[test]
    fn get_table_data_loads_column_information() {
        let mut connection = connection();

        diesel::sql_query(
            "DROP TABLE IF EXISTS diesel_test.table_1",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query(
            "CREATE TABLE diesel_test.table_1 \
            (id SERIAL PRIMARY KEY COMMENT 'column comment') \
            COMMENT 'table comment'",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query(
            "DROP TABLE IF EXISTS diesel_test.table_2",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query(
            "CREATE TABLE diesel_test.table_2 (id SERIAL PRIMARY KEY)",
        )
        .execute(&mut connection)
        .unwrap();

        let table_1 = TableName::new("table_1", "diesel_test");
        let table_2 = TableName::new("table_2", "diesel_test");

        let id_with_comment = ColumnInformation::new(
            "id",
            "bigint unsigned",
            None,
            false,
            Some("column comment".to_string()),
        );
        let id_without_comment = ColumnInformation::new("id", "bigint unsigned", None, false, None);
        assert_eq!(
            Ok(vec![id_with_comment]),
            get_table_data(&mut connection, &table_1, &ColumnSorting::OrdinalPosition)
        );
        assert_eq!(
            Ok(vec![id_without_comment]),
            get_table_data(&mut connection, &table_2, &ColumnSorting::OrdinalPosition)
        );
    }

    #[test]
    fn gets_table_comment() {
        let mut connection = connection();

        diesel::sql_query("DROP TABLE IF EXISTS diesel_test.table_1")
            .execute(&mut connection)
            .unwrap();
        diesel::sql_query(
            "CREATE TABLE diesel_test.table_1 (id SERIAL PRIMARY KEY) COMMENT 'table comment'",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query("DROP TABLE IF EXISTS diesel_test.table_2")
            .execute(&mut connection)
            .unwrap();
        diesel::sql_query("CREATE TABLE diesel_test.table_2 (id SERIAL PRIMARY KEY)")
            .execute(&mut connection)
            .unwrap();

        let table_1 = TableName::new("table_1", "diesel_test");
        let table_2 = TableName::new("table_2", "diesel_test");
        assert_eq!(
            Ok(Some("table comment".to_string())),
            get_table_comment(&mut connection, &table_1)
        );
        assert_eq!(Ok(None), get_table_comment(&mut connection, &table_2));
    }
}
<commit_msg>Do not use SERIAL in mysql column information test<commit_after>use diesel::mysql::{Mysql, MysqlConnection};
use diesel::*;
use heck::ToUpperCamelCase;
use std::{borrow::Cow, error::Error};

use super::data_structures::*;
use super::information_schema::DefaultSchema;
use super::table_data::TableName;
use crate::print_schema::ColumnSorting;

diesel::sql_function! {
    #[sql_name = "NULLIF"]
    fn null_if_text(lhs: sql_types::Text, rhs: sql_types::Text) -> sql_types::Nullable<sql_types::Text>
}

pub fn get_table_data<'a>(
    conn: &mut MysqlConnection,
    table: &'a TableName,
    column_sorting: &ColumnSorting,
) -> QueryResult<Vec<ColumnInformation>> {
    use self::information_schema::columns::dsl::*;

    let schema_name = match table.schema {
        Some(ref name) => Cow::Borrowed(name),
        None => Cow::Owned(Mysql::default_schema(conn)?),
    };

    let type_schema = None::<String>.into_sql();
    let query = columns
        .select((
            column_name,
            column_type,
            type_schema,
            __is_nullable,
            // MySQL comments are not nullable and are empty strings if not set
            null_if_text(column_comment, ""),
        ))
        .filter(table_name.eq(&table.sql_name))
        .filter(table_schema.eq(schema_name));
    match column_sorting {
        ColumnSorting::OrdinalPosition => query.order(ordinal_position).load(conn),
        ColumnSorting::Name => query.order(column_name).load(conn),
    }
}

mod information_schema {
    use diesel::prelude::{allow_tables_to_appear_in_same_query, table};

    table! {
        information_schema.tables (table_schema, table_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            table_comment -> VarChar,
        }
    }

    table! {
        information_schema.table_constraints (constraint_schema, constraint_name) {
            table_schema -> VarChar,
            constraint_schema -> VarChar,
            constraint_name -> VarChar,
            constraint_type -> VarChar,
        }
    }

    table! {
        information_schema.key_column_usage (constraint_schema, constraint_name) {
            constraint_schema -> VarChar,
            constraint_name -> VarChar,
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            referenced_table_schema -> VarChar,
            referenced_table_name -> VarChar,
            referenced_column_name -> VarChar,
        }
    }

    table! {
        information_schema.columns (table_schema, table_name, column_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            #[sql_name = "is_nullable"]
            __is_nullable -> VarChar,
            ordinal_position -> BigInt,
            udt_name -> VarChar,
            udt_schema -> VarChar,
            column_type -> VarChar,
            column_comment -> VarChar,
        }
    }

    allow_tables_to_appear_in_same_query!(table_constraints, key_column_usage);
}

/// Even though this is using `information_schema`, MySQL needs non-ANSI columns
/// in order to do this.
pub fn load_foreign_key_constraints(
    connection: &mut MysqlConnection,
    schema_name: Option<&str>,
) -> QueryResult<Vec<ForeignKeyConstraint>> {
    use self::information_schema::key_column_usage as kcu;
    use self::information_schema::table_constraints as tc;

    let default_schema = Mysql::default_schema(connection)?;
    let schema_name = match schema_name {
        Some(name) => name,
        None => &default_schema,
    };

    let constraints = tc::table
        .filter(tc::constraint_type.eq("FOREIGN KEY"))
        .filter(tc::table_schema.eq(schema_name))
        .filter(kcu::referenced_column_name.is_not_null())
        .inner_join(
            kcu::table.on(tc::constraint_schema
                .eq(kcu::constraint_schema)
                .and(tc::constraint_name.eq(kcu::constraint_name))),
        )
        .select((
            (kcu::table_name, kcu::table_schema),
            (kcu::referenced_table_name, kcu::referenced_table_schema),
            kcu::column_name,
            kcu::referenced_column_name,
        ))
        .load::<(TableName, TableName, String, _)>(connection)?
        .into_iter()
        .map(
            |(mut child_table, mut parent_table, foreign_key, primary_key)| {
                child_table.strip_schema_if_matches(&default_schema);
                parent_table.strip_schema_if_matches(&default_schema);

                ForeignKeyConstraint {
                    child_table,
                    parent_table,
                    foreign_key: foreign_key.clone(),
                    foreign_key_rust_name: foreign_key,
                    primary_key,
                }
            },
        )
        .collect();
    Ok(constraints)
}

pub fn determine_column_type(
    attr: &ColumnInformation,
) -> Result<ColumnType, Box<dyn Error + Send + Sync + 'static>> {
    let tpe = determine_type_name(&attr.type_name)?;
    let unsigned = determine_unsigned(&attr.type_name);

    Ok(ColumnType {
        schema: None,
        sql_name: tpe.trim().to_lowercase(),
        rust_name: tpe.trim().to_upper_camel_case(),
        is_array: false,
        is_nullable: attr.nullable,
        is_unsigned: unsigned,
    })
}

pub fn get_table_comment(
    conn: &mut MysqlConnection,
    table: &TableName,
) -> QueryResult<Option<String>> {
    use self::information_schema::tables::dsl::*;

    let schema_name = match table.schema {
        Some(ref name) => Cow::Borrowed(name),
        None => Cow::Owned(Mysql::default_schema(conn)?),
    };

    let comment: String = tables
        .select(table_comment)
        .filter(table_name.eq(&table.sql_name))
        .filter(table_schema.eq(schema_name))
        .get_result(conn)?;

    if comment.is_empty() {
        Ok(None)
    } else {
        Ok(Some(comment))
    }
}

fn determine_type_name(
    sql_type_name: &str,
) -> Result<String, Box<dyn Error + Send + Sync + 'static>> {
    let result = if sql_type_name == "tinyint(1)" {
        "bool"
    } else if sql_type_name.starts_with("int") {
        "integer"
    } else if let Some(idx) = sql_type_name.find('(') {
        &sql_type_name[..idx]
    } else {
        sql_type_name
    };

    if determine_unsigned(result) {
        Ok(result
            .to_lowercase()
            .replace("unsigned", "")
            .trim()
            .to_owned())
    } else if result.contains(' ') {
        Err(format!("unrecognized type {:?}", result).into())
    } else {
        Ok(result.to_owned())
    }
}

fn determine_unsigned(sql_type_name: &str) -> bool {
    sql_type_name.to_lowercase().contains("unsigned")
}

#[test]
fn values_which_already_map_to_type_are_returned_unchanged() {
    assert_eq!("text", determine_type_name("text").unwrap());
    assert_eq!("integer", determine_type_name("integer").unwrap());
    assert_eq!("biginteger", determine_type_name("biginteger").unwrap());
}

#[test]
fn trailing_parenthesis_are_stripped() {
    assert_eq!("varchar", determine_type_name("varchar(255)").unwrap());
    assert_eq!("decimal", determine_type_name("decimal(10, 2)").unwrap());
    assert_eq!("float", determine_type_name("float(1)").unwrap());
}

#[test]
fn tinyint_is_bool_if_limit_1() {
    assert_eq!("bool", determine_type_name("tinyint(1)").unwrap());
    assert_eq!("tinyint", determine_type_name("tinyint(2)").unwrap());
}

#[test]
fn int_is_treated_as_integer() {
    assert_eq!("integer", determine_type_name("int").unwrap());
    assert_eq!("integer", determine_type_name("int(11)").unwrap());
}

#[test]
fn unsigned_types_are_supported() {
    assert!(determine_unsigned("float unsigned"));
    assert!(determine_unsigned("UNSIGNED INT"));
    assert!(determine_unsigned("unsigned bigint"));
    assert!(!determine_unsigned("bigint"));
    assert!(!determine_unsigned("FLOAT"));
    assert_eq!("float", determine_type_name("float unsigned").unwrap());
    assert_eq!("int", determine_type_name("UNSIGNED INT").unwrap());
    assert_eq!("bigint", determine_type_name("unsigned bigint").unwrap());
}

#[test]
fn types_with_space_are_not_supported() {
    assert!(determine_type_name("lol wat").is_err());
}

#[cfg(test)]
mod test {
    extern crate dotenvy;

    use self::dotenvy::dotenv;
    use super::*;
    use std::env;

    fn connection() -> MysqlConnection {
        dotenv().ok();

        let connection_url = env::var("MYSQL_DATABASE_URL")
            .or_else(|_| env::var("DATABASE_URL"))
            .expect("DATABASE_URL must be set in order to run tests");
        let mut connection = MysqlConnection::establish(&connection_url).unwrap();
        connection.begin_test_transaction().unwrap();
        connection
    }

    #[test]
    fn get_table_data_loads_column_information() {
        let mut connection = connection();

        diesel::sql_query(
            "DROP TABLE IF EXISTS diesel_test.table_1",
        )
        .execute(&mut connection)
        .unwrap();
        // uses VARCHAR(255) as the type because SERIAL returned bigint on most platforms and bigint(20) on MacOS
        diesel::sql_query(
            "CREATE TABLE diesel_test.table_1 \
            (id VARCHAR(255) PRIMARY KEY COMMENT 'column comment') \
            COMMENT 'table comment'",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query(
            "DROP TABLE IF EXISTS diesel_test.table_2",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query(
            "CREATE TABLE diesel_test.table_2 (id VARCHAR(255) PRIMARY KEY)",
        )
        .execute(&mut connection)
        .unwrap();

        let table_1 = TableName::new("table_1", "diesel_test");
        let table_2 = TableName::new("table_2", "diesel_test");

        let id_with_comment = ColumnInformation::new(
            "id",
            "varchar(255)",
            None,
            false,
            Some("column comment".to_string()),
        );
        let id_without_comment = ColumnInformation::new("id", "varchar(255)", None, false, None);
        assert_eq!(
            Ok(vec![id_with_comment]),
            get_table_data(&mut connection, &table_1, &ColumnSorting::OrdinalPosition)
        );
        assert_eq!(
            Ok(vec![id_without_comment]),
            get_table_data(&mut connection, &table_2, &ColumnSorting::OrdinalPosition)
        );
    }

    #[test]
    fn gets_table_comment() {
        let mut connection = connection();

        diesel::sql_query("DROP TABLE IF EXISTS diesel_test.table_1")
            .execute(&mut connection)
            .unwrap();
        diesel::sql_query(
            "CREATE TABLE diesel_test.table_1 (id SERIAL PRIMARY KEY) COMMENT 'table comment'",
        )
        .execute(&mut connection)
        .unwrap();
        diesel::sql_query("DROP TABLE IF EXISTS diesel_test.table_2")
            .execute(&mut connection)
            .unwrap();
        diesel::sql_query("CREATE TABLE diesel_test.table_2 (id SERIAL PRIMARY KEY)")
            .execute(&mut connection)
            .unwrap();

        let table_1 = TableName::new("table_1", "diesel_test");
        let table_2 = TableName::new("table_2", "diesel_test");
        assert_eq!(
            Ok(Some("table comment".to_string())),
            get_table_comment(&mut connection, &table_1)
        );
        assert_eq!(Ok(None), get_table_comment(&mut connection, &table_2));
    }
}
<|endoftext|>


<commit_before>// lock-free stack
use std::{
    fmt::{self, Debug},
    ops::Deref,
    sync::atomic::Ordering::{Relaxed, SeqCst},
};

use sled_sync::{
    debug_delay, pin, unprotected, Atomic, Guard, Owned, Shared,
};

/// A node in the lock-free `Stack`.
#[derive(Debug)]
pub(crate) struct Node<T: Send + 'static> {
    inner: T,
    next: Atomic<Node<T>>,
}

impl<T: Send + 'static> Drop for Node<T> {
    fn drop(&mut self) {
        unsafe {
            let next =
                self.next.load(Relaxed, unprotected()).as_raw();
            if !next.is_null() {
                drop(Box::from_raw(next as *mut Node<T>));
            }
        }
    }
}

/// A simple lock-free stack, with the ability to atomically
/// append or entirely swap-out entries.
pub(crate) struct Stack<T: Send + 'static> {
    head: Atomic<Node<T>>,
}

impl<T: Send + 'static> Default for Stack<T> {
    fn default() -> Stack<T> {
        Stack {
            head: Atomic::null(),
        }
    }
}

impl<T: Send + 'static> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            let curr =
                self.head.load(Relaxed, unprotected()).as_raw();
            if !curr.is_null() {
                drop(Box::from_raw(curr as *mut Node<T>));
            }
        }
    }
}

impl<T> Debug for Stack<T>
where
    T: Debug + Send + 'static + Sync,
{
    fn fmt(
        &self,
        formatter: &mut fmt::Formatter<'_>,
    ) -> Result<(), fmt::Error> {
        let guard = pin();
        let head = self.head(&guard);
        let iter = StackIter::from_ptr(head, &guard);

        formatter.write_str("Stack [")?;
        let mut written = false;
        for node in iter {
            if written {
                formatter.write_str(", ")?;
            }
            formatter
                .write_str(&*format!("({:?}) ", &node as *const _))?;
            node.fmt(formatter)?;
            written = true;
        }
        formatter.write_str("]")?;
        Ok(())
    }
}

impl<T: Send + 'static> Deref for Node<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T: Send + Sync + 'static> Stack<T> {
    /// Add an item to the stack, spinning until successful.
    pub(crate) fn push(&self, inner: T) {
        debug_delay();
        let node = Owned::new(Node {
            inner,
            next: Atomic::null(),
        });

        unsafe {
            let node = node.into_shared(unprotected());

            loop {
                let head = self.head(unprotected());
                node.deref().next.store(head, SeqCst);
                if self
                    .head
                    .compare_and_set(
                        head,
                        node,
                        SeqCst,
                        unprotected(),
                    )
                    .is_ok()
                {
                    return;
                }
            }
        }
    }

    /// Pop the next item off the stack. Returns None if nothing is there.
    fn _pop<'g>(&self, guard: &'g Guard) -> Option<T> {
        use std::ptr;
        debug_delay();
        let mut head = self.head(&guard);
        loop {
            match unsafe { head.as_ref() } {
                Some(h) => {
                    let next = h.next.load(SeqCst, &guard);
                    match self
                        .head
                        .compare_and_set(head, next, SeqCst, &guard)
                    {
                        Ok(_) => unsafe {
                            let head_owned = head.into_owned();
                            guard.defer(move || head_owned);
                            return Some(ptr::read(&h.inner));
                        },
                        Err(h) => head = h.current,
                    }
                }
                None => return None,
            }
        }
    }

    /// compare and push
    pub(crate) fn cap<'g>(
        &self,
        old: Shared<'_, Node<T>>,
        new: T,
        guard: &'g Guard,
    ) -> Result<Shared<'g, Node<T>>, Shared<'g, Node<T>>> {
        debug_delay();
        let node = Owned::new(Node {
            inner: new,
            next: Atomic::from(old),
        });

        let node = node.into_shared(guard);

        let res = self.head.compare_and_set(old, node, SeqCst, guard);

        match res {
            Err(e) => {
                unsafe {
                    // we want to set next to null to prevent
                    // the current shared head from being
                    // dropped when we drop this node.
                    node.deref().next.store(Shared::null(), SeqCst);
                    let node_owned = node.into_owned();
                    guard.defer(move || node_owned);
                }
                Err(e.current)
            }
            Ok(_) => Ok(node),
        }
    }

    /// compare and swap
    pub(crate) fn cas<'g>(
        &self,
        old: Shared<'g, Node<T>>,
        new: Shared<'g, Node<T>>,
        guard: &'g Guard,
    ) -> Result<Shared<'g, Node<T>>, Shared<'g, Node<T>>> {
        debug_delay();
        let res = self.head.compare_and_set(old, new, SeqCst, guard);

        match res {
            Ok(_) => {
                if !old.is_null() {
                    unsafe {
                        let old_owned = old.into_owned();
                        guard.defer(move || old_owned)
                    };
                }
                Ok(new)
            }
            Err(e) => {
                if !new.is_null() {
                    unsafe {
                        let new_owned = new.into_owned();
                        guard.defer(move || new_owned)
                    };
                }

                Err(e.current)
            }
        }
    }

    /// Returns the current head pointer of the stack, which can
    /// later be used as the key for cas and cap operations.
    pub(crate) fn head<'g>(
        &self,
        guard: &'g Guard,
    ) -> Shared<'g, Node<T>> {
        self.head.load(SeqCst, guard)
    }
}

/// An iterator over nodes in a lock-free stack.
pub(crate) struct StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    inner: Shared<'a, Node<T>>,
    guard: &'a Guard,
}

impl<'a, T> StackIter<'a, T>
where
    T: 'a + Send + 'static + Sync,
{
    /// Creates a StackIter from a pointer to one.
    pub(crate) fn from_ptr<'b>(
        ptr: Shared<'b, Node<T>>,
        guard: &'b Guard,
    ) -> StackIter<'b, T> {
        StackIter {
            inner: ptr,
            guard,
        }
    }
}

impl<'a, T> Iterator for StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        debug_delay();
        if self.inner.is_null() {
            None
        } else {
            unsafe {
                let ret = &self.inner.deref().inner;
                self.inner =
                    self.inner.deref().next.load(SeqCst, self.guard);
                Some(ret)
            }
        }
    }
}

/// Turns a vector of elements into a lock-free stack
/// of them, and returns the head of the stack.
pub(crate) fn node_from_frag_vec<T>(from: Vec<T>) -> Owned<Node<T>>
where
    T: Send + 'static + Sync,
{
    let mut last = None;

    for item in from.into_iter().rev() {
        last = if let Some(last) = last {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::from(last),
            }))
        } else {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::null(),
            }))
        }
    }

    last.expect("at least one frag was provided in the from Vec")
}

#[test]
fn basic_functionality() {
    use std::sync::Arc;
    use std::thread;

    let guard = pin();
    let ll = Arc::new(Stack::default());
    assert_eq!(ll._pop(&guard), None);
    ll.push(1);
    let ll2 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        ll2.push(2);
        ll2.push(3);
        ll2.push(4);
    });
    t.join().unwrap();
    ll.push(5);
    assert_eq!(ll._pop(&guard), Some(5));
    assert_eq!(ll._pop(&guard), Some(4));
    let ll3 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll3._pop(&guard), Some(3));
        assert_eq!(ll3._pop(&guard), Some(2));
    });
    t.join().unwrap();
    assert_eq!(ll._pop(&guard), Some(1));
    let ll4 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll4._pop(&guard), None);
    });
    t.join().unwrap();
}
<commit_msg>Immediately drop nodes from failed cap/cas on Stack<commit_after>// lock-free stack
use std::{
    fmt::{self, Debug},
    ops::Deref,
    sync::atomic::Ordering::{Relaxed, SeqCst},
};

use sled_sync::{
    debug_delay, pin, unprotected, Atomic, Guard, Owned, Shared,
};

/// A node in the lock-free `Stack`.
#[derive(Debug)]
pub(crate) struct Node<T: Send + 'static> {
    inner: T,
    next: Atomic<Node<T>>,
}

impl<T: Send + 'static> Drop for Node<T> {
    fn drop(&mut self) {
        unsafe {
            let next =
                self.next.load(Relaxed, unprotected()).as_raw();
            if !next.is_null() {
                drop(Box::from_raw(next as *mut Node<T>));
            }
        }
    }
}

/// A simple lock-free stack, with the ability to atomically
/// append or entirely swap-out entries.
pub(crate) struct Stack<T: Send + 'static> {
    head: Atomic<Node<T>>,
}

impl<T: Send + 'static> Default for Stack<T> {
    fn default() -> Stack<T> {
        Stack {
            head: Atomic::null(),
        }
    }
}

impl<T: Send + 'static> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            let curr =
                self.head.load(Relaxed, unprotected()).as_raw();
            if !curr.is_null() {
                drop(Box::from_raw(curr as *mut Node<T>));
            }
        }
    }
}

impl<T> Debug for Stack<T>
where
    T: Debug + Send + 'static + Sync,
{
    fn fmt(
        &self,
        formatter: &mut fmt::Formatter<'_>,
    ) -> Result<(), fmt::Error> {
        let guard = pin();
        let head = self.head(&guard);
        let iter = StackIter::from_ptr(head, &guard);

        formatter.write_str("Stack [")?;
        let mut written = false;
        for node in iter {
            if written {
                formatter.write_str(", ")?;
            }
            formatter
                .write_str(&*format!("({:?}) ", &node as *const _))?;
            node.fmt(formatter)?;
            written = true;
        }
        formatter.write_str("]")?;
        Ok(())
    }
}

impl<T: Send + 'static> Deref for Node<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T: Send + Sync + 'static> Stack<T> {
    /// Add an item to the stack, spinning until successful.
    pub(crate) fn push(&self, inner: T) {
        debug_delay();
        let node = Owned::new(Node {
            inner,
            next: Atomic::null(),
        });

        unsafe {
            let node = node.into_shared(unprotected());

            loop {
                let head = self.head(unprotected());
                node.deref().next.store(head, SeqCst);
                if self
                    .head
                    .compare_and_set(
                        head,
                        node,
                        SeqCst,
                        unprotected(),
                    )
                    .is_ok()
                {
                    return;
                }
            }
        }
    }

    /// Pop the next item off the stack. Returns None if nothing is there.
    fn _pop<'g>(&self, guard: &'g Guard) -> Option<T> {
        use std::ptr;
        debug_delay();
        let mut head = self.head(&guard);
        loop {
            match unsafe { head.as_ref() } {
                Some(h) => {
                    let next = h.next.load(SeqCst, &guard);
                    match self
                        .head
                        .compare_and_set(head, next, SeqCst, &guard)
                    {
                        Ok(_) => unsafe {
                            let head_owned = head.into_owned();
                            guard.defer(move || head_owned);
                            return Some(ptr::read(&h.inner));
                        },
                        Err(h) => head = h.current,
                    }
                }
                None => return None,
            }
        }
    }

    /// compare and push
    pub(crate) fn cap<'g>(
        &self,
        old: Shared<'_, Node<T>>,
        new: T,
        guard: &'g Guard,
    ) -> Result<Shared<'g, Node<T>>, Shared<'g, Node<T>>> {
        debug_delay();
        let node = Owned::new(Node {
            inner: new,
            next: Atomic::from(old),
        });

        let node = node.into_shared(guard);

        let res = self.head.compare_and_set(old, node, SeqCst, guard);

        match res {
            Err(e) => {
                unsafe {
                    // we want to set next to null to prevent
                    // the current shared head from being
                    // dropped when we drop this node.
                    node.deref().next.store(Shared::null(), SeqCst);
                    let node_owned = node.into_owned();
                    drop(node_owned)
                }
                Err(e.current)
            }
            Ok(_) => Ok(node),
        }
    }

    /// compare and swap
    pub(crate) fn cas<'g>(
        &self,
        old: Shared<'g, Node<T>>,
        new: Shared<'g, Node<T>>,
        guard: &'g Guard,
    ) -> Result<Shared<'g, Node<T>>, Shared<'g, Node<T>>> {
        debug_delay();
        let res = self.head.compare_and_set(old, new, SeqCst, guard);

        match res {
            Ok(_) => {
                if !old.is_null() {
                    unsafe {
                        let old_owned = old.into_owned();
                        guard.defer(move || old_owned)
                    };
                }
                Ok(new)
            }
            Err(e) => {
                if !new.is_null() {
                    unsafe {
                        let new_owned = new.into_owned();
                        drop(new_owned)
                    };
                }

                Err(e.current)
            }
        }
    }

    /// Returns the current head pointer of the stack, which can
    /// later be used as the key for cas and cap operations.
    pub(crate) fn head<'g>(
        &self,
        guard: &'g Guard,
    ) -> Shared<'g, Node<T>> {
        self.head.load(SeqCst, guard)
    }
}

/// An iterator over nodes in a lock-free stack.
pub(crate) struct StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    inner: Shared<'a, Node<T>>,
    guard: &'a Guard,
}

impl<'a, T> StackIter<'a, T>
where
    T: 'a + Send + 'static + Sync,
{
    /// Creates a StackIter from a pointer to one.
    pub(crate) fn from_ptr<'b>(
        ptr: Shared<'b, Node<T>>,
        guard: &'b Guard,
    ) -> StackIter<'b, T> {
        StackIter { inner: ptr, guard }
    }
}

impl<'a, T> Iterator for StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        debug_delay();
        if self.inner.is_null() {
            None
        } else {
            unsafe {
                let ret = &self.inner.deref().inner;
                self.inner =
                    self.inner.deref().next.load(SeqCst, self.guard);
                Some(ret)
            }
        }
    }
}

/// Turns a vector of elements into a lock-free stack
/// of them, and returns the head of the stack.
pub(crate) fn node_from_frag_vec<T>(from: Vec<T>) -> Owned<Node<T>>
where
    T: Send + 'static + Sync,
{
    let mut last = None;

    for item in from.into_iter().rev() {
        last = if let Some(last) = last {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::from(last),
            }))
        } else {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::null(),
            }))
        }
    }

    last.expect("at least one frag was provided in the from Vec")
}

#[test]
fn basic_functionality() {
    use std::sync::Arc;
    use std::thread;

    let guard = pin();
    let ll = Arc::new(Stack::default());
    assert_eq!(ll._pop(&guard), None);
    ll.push(1);
    let ll2 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        ll2.push(2);
        ll2.push(3);
        ll2.push(4);
    });
    t.join().unwrap();
    ll.push(5);
    assert_eq!(ll._pop(&guard), Some(5));
    assert_eq!(ll._pop(&guard), Some(4));
    let ll3 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll3._pop(&guard), Some(3));
        assert_eq!(ll3._pop(&guard), Some(2));
    });
    t.join().unwrap();
    assert_eq!(ll._pop(&guard), Some(1));
    let ll4 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll4._pop(&guard), None);
    });
    t.join().unwrap();
}
<|endoftext|>