<commit_before>use crate::runtime::task::{Id, RawTask};

use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::panic::{RefUnwindSafe, UnwindSafe};
use std::pin::Pin;
use std::task::{Context, Poll, Waker};

cfg_rt! {
    /// An owned permission to join on a task (await its termination).
    ///
    /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] for
    /// a task rather than a thread.
    ///
    /// A `JoinHandle` *detaches* the associated task when it is dropped, which
    /// means that there is no longer any handle to the task, and no way to `join`
    /// on it.
    ///
    /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`]
    /// functions.
    ///
    /// # Examples
    ///
    /// Creation from [`task::spawn`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn(async {
    ///     // some work here
    /// });
    /// # }
    /// ```
    ///
    /// Creation from [`task::spawn_blocking`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn_blocking(|| {
    ///     // some blocking work here
    /// });
    /// # }
    /// ```
    ///
    /// The generic parameter `T` in `JoinHandle<T>` is the return type of the spawned task.
    /// If the return value is an i32, the join handle has type `JoinHandle<i32>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<i32> = task::spawn(async {
    ///     5 + 3
    /// });
    /// # }
    ///
    /// ```
    ///
    /// If the task does not have a return value, the join handle has type `JoinHandle<()>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<()> = task::spawn(async {
    ///     println!("I return nothing.");
    /// });
    /// # }
    /// ```
    ///
    /// Note that `handle.await` doesn't give you the return type directly. It is wrapped in a
    /// `Result` because panics in the spawned task are caught by Tokio. The `?` operator has
    /// to be double chained to extract the returned value:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         Ok(5 + 3)
    ///     });
    ///
    ///     let result = join_handle.await??;
    ///     assert_eq!(result, 8);
    ///     Ok(())
    /// }
    /// ```
    ///
    /// If the task panics, the error is a [`JoinError`] that contains the panic:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    /// use std::panic;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         panic!("boom");
    ///     });
    ///
    ///     let err = join_handle.await.unwrap_err();
    ///     assert!(err.is_panic());
    ///     Ok(())
    /// }
    ///
    /// ```
    /// Child being detached and outliving its parent:
    ///
    /// ```no_run
    /// use tokio::task;
    /// use tokio::time;
    /// use std::time::Duration;
    ///
    /// # #[tokio::main] async fn main() {
    /// let original_task = task::spawn(async {
    ///     let _detached_task = task::spawn(async {
    ///         // Here we sleep to make sure that the first task returns before.
    ///         time::sleep(Duration::from_millis(10)).await;
    ///         // This will be called, even though the JoinHandle is dropped.
    ///         println!("♫ Still alive ♫");
    ///     });
    /// });
    ///
    /// original_task.await.expect("The task being joined has panicked");
    /// println!("Original task is joined.");
    ///
    /// // We make sure that the new task has time to run, before the main
    /// // task returns.
    ///
    /// time::sleep(Duration::from_millis(1000)).await;
    /// # }
    /// ```
    ///
    /// [`task::spawn`]: crate::task::spawn()
    /// [`task::spawn_blocking`]: crate::task::spawn_blocking
    /// [`std::thread::JoinHandle`]: std::thread::JoinHandle
    /// [`JoinError`]: crate::task::JoinError
    pub struct JoinHandle<T> {
        raw: Option<RawTask>,
        id: Id,
        _p: PhantomData<T>,
    }
}

unsafe impl<T: Send> Send for JoinHandle<T> {}
unsafe impl<T: Send> Sync for JoinHandle<T> {}

impl<T> UnwindSafe for JoinHandle<T> {}
impl<T> RefUnwindSafe for JoinHandle<T> {}

impl<T> JoinHandle<T> {
    pub(super) fn new(raw: RawTask, id: Id) -> JoinHandle<T> {
        JoinHandle {
            raw: Some(raw),
            id,
            _p: PhantomData,
        }
    }

    /// Abort the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will fail with a [cancelled] `JoinError`.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       true
    ///    }));
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       false
    ///    }));
    ///
    ///    for handle in &handles {
    ///        handle.abort();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    /// [cancelled]: method@super::error::JoinError::is_cancelled
    pub fn abort(&self) {
        if let Some(raw) = self.raw {
            raw.remote_abort();
        }
    }

    /// Checks if the task associated with this `JoinHandle` has finished.
    ///
    /// Please note that this method can return `false` even if [`abort`] has been
    /// called on the task. This is because the cancellation process may take
    /// some time, and this method does not return `true` until it has
    /// completed.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// # #[tokio::main(flavor = "current_thread")]
    /// # async fn main() {
    /// # time::pause();
    /// let handle1 = tokio::spawn(async {
    ///     // do some stuff here
    /// });
    /// let handle2 = tokio::spawn(async {
    ///     // do some other stuff here
    ///     time::sleep(time::Duration::from_secs(10)).await;
    /// });
    /// // Wait for the task to finish
    /// handle2.abort();
    /// time::sleep(time::Duration::from_secs(1)).await;
    /// assert!(handle1.is_finished());
    /// assert!(handle2.is_finished());
    /// # }
    /// ```
    /// [`abort`]: method@JoinHandle::abort
    pub fn is_finished(&self) -> bool {
        if let Some(raw) = self.raw {
            let state = raw.header().state.load();
            state.is_complete()
        } else {
            true
        }
    }

    /// Set the waker that is notified when the task completes.
    pub(crate) fn set_join_waker(&mut self, waker: &Waker) {
        if let Some(raw) = self.raw {
            if raw.try_set_join_waker(waker) {
                // In this case the task has already completed. We wake the waker immediately.
                waker.wake_by_ref();
            }
        }
    }

    /// Returns a new `AbortHandle` that can be used to remotely abort this task.
    #[cfg(any(tokio_unstable, test))]
    pub(crate) fn abort_handle(&self) -> super::AbortHandle {
        let raw = self.raw.map(|raw| {
            raw.ref_inc();
            raw
        });
        super::AbortHandle::new(raw, self.id.clone())
    }

    /// Returns a [task ID] that uniquely identifies this task relative to other
    /// currently spawned tasks.
    ///
    /// **Note**: This is an [unstable API][unstable]. The public API of this type
    /// may break in 1.x releases. See [the documentation on unstable
    /// features][unstable] for details.
    ///
    /// [task ID]: crate::task::Id
    /// [unstable]: crate#unstable-features
    #[cfg(tokio_unstable)]
    #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
    pub fn id(&self) -> super::Id {
        self.id.clone()
    }
}

impl<T> Unpin for JoinHandle<T> {}

impl<T> Future for JoinHandle<T> {
    type Output = super::Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut ret = Poll::Pending;

        // Keep track of task budget
        let coop = ready!(crate::coop::poll_proceed(cx));

        // Raw should always be set. If it is not, this is due to polling after
        // completion
        let raw = self
            .raw
            .as_ref()
            .expect("polling after `JoinHandle` already completed");

        // Try to read the task output. If the task is not yet complete, the
        // waker is stored and is notified once the task does complete.
        //
        // The function must go via the vtable, which requires erasing generic
        // types. To do this, the function "return" is placed on the stack
        // **before** calling the function and is passed into the function using
        // `*mut ()`.
        //
        // Safety:
        //
        // The type of `T` must match the task's output type.
        unsafe {
            raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker());
        }

        if ret.is_ready() {
            coop.made_progress();
        }

        ret
    }
}

impl<T> Drop for JoinHandle<T> {
    fn drop(&mut self) {
        if let Some(raw) = self.raw.take() {
            if raw.header().state.drop_join_handle_fast().is_ok() {
                return;
            }

            raw.drop_join_handle_slow();
        }
    }
}

impl<T> fmt::Debug for JoinHandle<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("JoinHandle")
            .field("id", &self.id)
            .finish()
    }
}
<commit_msg>task: add cancel safety docs to `JoinHandle` (#4901)<commit_after>use crate::runtime::task::{Id, RawTask};

use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::panic::{RefUnwindSafe, UnwindSafe};
use std::pin::Pin;
use std::task::{Context, Poll, Waker};

cfg_rt! {
    /// An owned permission to join on a task (await its termination).
    ///
    /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] for
    /// a task rather than a thread.
    ///
    /// A `JoinHandle` *detaches* the associated task when it is dropped, which
    /// means that there is no longer any handle to the task, and no way to `join`
    /// on it.
    ///
    /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`]
    /// functions.
    ///
    /// # Cancel safety
    ///
    /// The `&mut JoinHandle<T>` type is cancel safe. If it is used as the event
    /// in a `tokio::select!` statement and some other branch completes first,
    /// then it is guaranteed that the output of the task is not lost.
    ///
    /// If a `JoinHandle` is dropped, then the task continues running in the
    /// background and its return value is lost.
    ///
    /// # Examples
    ///
    /// Creation from [`task::spawn`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn(async {
    ///     // some work here
    /// });
    /// # }
    /// ```
    ///
    /// Creation from [`task::spawn_blocking`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn_blocking(|| {
    ///     // some blocking work here
    /// });
    /// # }
    /// ```
    ///
    /// The generic parameter `T` in `JoinHandle<T>` is the return type of the spawned task.
    /// If the return value is an i32, the join handle has type `JoinHandle<i32>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<i32> = task::spawn(async {
    ///     5 + 3
    /// });
    /// # }
    ///
    /// ```
    ///
    /// If the task does not have a return value, the join handle has type `JoinHandle<()>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<()> = task::spawn(async {
    ///     println!("I return nothing.");
    /// });
    /// # }
    /// ```
    ///
    /// Note that `handle.await` doesn't give you the return type directly. It is wrapped in a
    /// `Result` because panics in the spawned task are caught by Tokio. The `?` operator has
    /// to be double chained to extract the returned value:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         Ok(5 + 3)
    ///     });
    ///
    ///     let result = join_handle.await??;
    ///     assert_eq!(result, 8);
    ///     Ok(())
    /// }
    /// ```
    ///
    /// If the task panics, the error is a [`JoinError`] that contains the panic:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    /// use std::panic;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         panic!("boom");
    ///     });
    ///
    ///     let err = join_handle.await.unwrap_err();
    ///     assert!(err.is_panic());
    ///     Ok(())
    /// }
    ///
    /// ```
    /// Child being detached and outliving its parent:
    ///
    /// ```no_run
    /// use tokio::task;
    /// use tokio::time;
    /// use std::time::Duration;
    ///
    /// # #[tokio::main] async fn main() {
    /// let original_task = task::spawn(async {
    ///     let _detached_task = task::spawn(async {
    ///         // Here we sleep to make sure that the first task returns before.
    ///         time::sleep(Duration::from_millis(10)).await;
    ///         // This will be called, even though the JoinHandle is dropped.
    ///         println!("♫ Still alive ♫");
    ///     });
    /// });
    ///
    /// original_task.await.expect("The task being joined has panicked");
    /// println!("Original task is joined.");
    ///
    /// // We make sure that the new task has time to run, before the main
    /// // task returns.
    ///
    /// time::sleep(Duration::from_millis(1000)).await;
    /// # }
    /// ```
    ///
    /// [`task::spawn`]: crate::task::spawn()
    /// [`task::spawn_blocking`]: crate::task::spawn_blocking
    /// [`std::thread::JoinHandle`]: std::thread::JoinHandle
    /// [`JoinError`]: crate::task::JoinError
    pub struct JoinHandle<T> {
        raw: Option<RawTask>,
        id: Id,
        _p: PhantomData<T>,
    }
}

unsafe impl<T: Send> Send for JoinHandle<T> {}
unsafe impl<T: Send> Sync for JoinHandle<T> {}

impl<T> UnwindSafe for JoinHandle<T> {}
impl<T> RefUnwindSafe for JoinHandle<T> {}

impl<T> JoinHandle<T> {
    pub(super) fn new(raw: RawTask, id: Id) -> JoinHandle<T> {
        JoinHandle {
            raw: Some(raw),
            id,
            _p: PhantomData,
        }
    }

    /// Abort the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will fail with a [cancelled] `JoinError`.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       true
    ///    }));
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       false
    ///    }));
    ///
    ///    for handle in &handles {
    ///        handle.abort();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    /// [cancelled]: method@super::error::JoinError::is_cancelled
    pub fn abort(&self) {
        if let Some(raw) = self.raw {
            raw.remote_abort();
        }
    }

    /// Checks if the task associated with this `JoinHandle` has finished.
    ///
    /// Please note that this method can return `false` even if [`abort`] has been
    /// called on the task. This is because the cancellation process may take
    /// some time, and this method does not return `true` until it has
    /// completed.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// # #[tokio::main(flavor = "current_thread")]
    /// # async fn main() {
    /// # time::pause();
    /// let handle1 = tokio::spawn(async {
    ///     // do some stuff here
    /// });
    /// let handle2 = tokio::spawn(async {
    ///     // do some other stuff here
    ///     time::sleep(time::Duration::from_secs(10)).await;
    /// });
    /// // Wait for the task to finish
    /// handle2.abort();
    /// time::sleep(time::Duration::from_secs(1)).await;
    /// assert!(handle1.is_finished());
    /// assert!(handle2.is_finished());
    /// # }
    /// ```
    /// [`abort`]: method@JoinHandle::abort
    pub fn is_finished(&self) -> bool {
        if let Some(raw) = self.raw {
            let state = raw.header().state.load();
            state.is_complete()
        } else {
            true
        }
    }

    /// Set the waker that is notified when the task completes.
    pub(crate) fn set_join_waker(&mut self, waker: &Waker) {
        if let Some(raw) = self.raw {
            if raw.try_set_join_waker(waker) {
                // In this case the task has already completed. We wake the waker immediately.
                waker.wake_by_ref();
            }
        }
    }

    /// Returns a new `AbortHandle` that can be used to remotely abort this task.
    #[cfg(any(tokio_unstable, test))]
    pub(crate) fn abort_handle(&self) -> super::AbortHandle {
        let raw = self.raw.map(|raw| {
            raw.ref_inc();
            raw
        });
        super::AbortHandle::new(raw, self.id.clone())
    }

    /// Returns a [task ID] that uniquely identifies this task relative to other
    /// currently spawned tasks.
    ///
    /// **Note**: This is an [unstable API][unstable]. The public API of this type
    /// may break in 1.x releases. See [the documentation on unstable
    /// features][unstable] for details.
    ///
    /// [task ID]: crate::task::Id
    /// [unstable]: crate#unstable-features
    #[cfg(tokio_unstable)]
    #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
    pub fn id(&self) -> super::Id {
        self.id.clone()
    }
}

impl<T> Unpin for JoinHandle<T> {}

impl<T> Future for JoinHandle<T> {
    type Output = super::Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut ret = Poll::Pending;

        // Keep track of task budget
        let coop = ready!(crate::coop::poll_proceed(cx));

        // Raw should always be set. If it is not, this is due to polling after
        // completion
        let raw = self
            .raw
            .as_ref()
            .expect("polling after `JoinHandle` already completed");

        // Try to read the task output. If the task is not yet complete, the
        // waker is stored and is notified once the task does complete.
        //
        // The function must go via the vtable, which requires erasing generic
        // types. To do this, the function "return" is placed on the stack
        // **before** calling the function and is passed into the function using
        // `*mut ()`.
        //
        // Safety:
        //
        // The type of `T` must match the task's output type.
        unsafe {
            raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker());
        }

        if ret.is_ready() {
            coop.made_progress();
        }

        ret
    }
}

impl<T> Drop for JoinHandle<T> {
    fn drop(&mut self) {
        if let Some(raw) = self.raw.take() {
            if raw.header().state.drop_join_handle_fast().is_ok() {
                return;
            }

            raw.drop_join_handle_slow();
        }
    }
}

impl<T> fmt::Debug for JoinHandle<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("JoinHandle")
            .field("id", &self.id)
            .finish()
    }
}
<|endoftext|>


<commit_before>use super::namespaces;
use crate::{library::Library, nameutil::crate_name, version::Version};
use std::cmp::Ordering;
use std::collections::btree_map::BTreeMap;
use std::ops::{Deref, DerefMut};
use std::vec::IntoIter;

fn is_first_char_up(s: &str) -> bool {
    s.chars().next().unwrap().is_uppercase()
}

fn check_up_eq(a: &str, b: &str) -> Ordering {
    let is_a_up = is_first_char_up(a);
    let is_b_up = is_first_char_up(b);
    if is_a_up != is_b_up {
        if is_a_up {
            return Ordering::Greater;
        }
        return Ordering::Less;
    }
    Ordering::Equal
}

/// This function is used by the `Imports` type to generate output like `cargo fmt` would.
///
/// For example:
///
/// ```text
/// use gdk; // lowercases come first.
/// use Window;
///
/// use gdk::foo; // lowercases come first here as well.
/// use gdk::Foo;
/// ```
#[allow(clippy::trivially_copy_pass_by_ref)] // because of the way this is used, the refs are needed
fn compare_imports(a: &(&String, &ImportConditions), b: &(&String, &ImportConditions)) -> Ordering {
    let s = check_up_eq(a.0, b.0);
    if s != Ordering::Equal {
        return s;
    }
    let mut a = a.0.split("::");
    let mut b = b.0.split("::");
    loop {
        match (a.next(), b.next()) {
            (Some(a), Some(b)) => {
                let s = check_up_eq(a, b);
                if s != Ordering::Equal {
                    break s;
                }
                let s = a.partial_cmp(b).unwrap();
                if s != Ordering::Equal {
                    break s;
                }
            }
            (Some(_), None) => break Ordering::Greater,
            (None, Some(_)) => break Ordering::Less,
            (None, None) => break Ordering::Equal,
        }
    }
}

/// Provides assistance in generating use declarations.
///
/// It takes into account that use declaration referring to names within the
/// same crate will look differently. It also avoids generating spurious
/// declarations referring to names from within the same module as the one we
/// are generating code for.
#[derive(Clone, Debug, Default)]
pub struct Imports {
    /// Name of the current crate.
    crate_name: String,
    /// Name defined within current module. It doesn't need use declaration.
    ///
    /// NOTE: Currently we don't need to support more than one such name.
    defined: Option<String>,
    defaults: ImportConditions,
    map: BTreeMap<String, ImportConditions>,
}

impl Imports {
    pub fn new(gir: &Library) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: None,
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    pub fn with_defined(gir: &Library, name: &str) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: Some(name.to_owned()),
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    #[must_use = "ImportsWithDefault must live while defaults are needed"]
    pub fn with_defaults(
        &mut self,
        version: Option<Version>,
        constraint: &Option<String>,
    ) -> ImportsWithDefault<'_> {
        let constraints = if let Some(constraint) = constraint {
            vec![constraint.clone()]
        } else {
            vec![]
        };
        self.defaults = ImportConditions {
            version,
            constraints,
        };

        ImportsWithDefault::new(self)
    }

    fn reset_defaults(&mut self) {
        self.defaults.clear();
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    /// Uses defaults.
    pub fn add(&mut self, name: &str) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let defaults = &self.defaults;
            let entry = self
                .map
                .entry(name.to_owned())
                .or_insert_with(|| defaults.clone());
            entry.update_version(self.defaults.version);
            entry.update_constraints(self.defaults.constraints.clone());
        }
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_version(&mut self, name: &str, version: Option<Version>) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                version,
                constraints: Vec::new(),
            });
            entry.update_version(version);
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            entry.constraints.clear();
        }
    }

    /// Declares that name should be available through its last path component and provides
    /// an optional feature constraint.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_constraint(
        &mut self,
        name: &str,
        version: Option<Version>,
        constraint: Option<&str>,
    ) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = if let Some(constraint) = constraint {
                let constraint = String::from(constraint);
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: vec![constraint.clone()],
                });
                entry.add_constraint(constraint);
                entry
            } else {
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: Vec::new(),
                });
                // Since there is no constraint on this import, if any constraint
                // is present, we can just remove it.
                entry.constraints.clear();
                entry
            };
            entry.update_version(version);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type(&mut self, used_type: &str) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add(&used_type[2..]);
            } else {
                self.add(&used_type[..i]);
            }
        } else {
            self.add(used_type);
        }
    }

    pub fn add_used_types(&mut self, used_types: &[String]) {
        for s in used_types {
            self.add_used_type(s);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type_with_version(&mut self, used_type: &str, version: Option<Version>) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add_with_version(&used_type[2..], version);
            } else {
                self.add_with_version(&used_type[..i], version);
            }
        } else {
            self.add_with_version(used_type, version);
        }
    }

    /// Tries to strip crate name prefix from given name.
    ///
    /// Returns `None` if name matches crate name exactly. Otherwise returns
    /// name with crate name prefix stripped or full name if there was no match.
    fn strip_crate_name<'a>(&self, name: &'a str) -> Option<&'a str> {
        let prefix = &self.crate_name;
        if !name.starts_with(prefix) {
            return Some(name);
        }
        let rest = &name[prefix.len()..];
        if rest.is_empty() {
            None
        } else if rest.starts_with("::") {
            Some(&rest["::".len()..])
        } else {
            // It was false positive, return the whole name.
            Some(name)
        }
    }

    pub fn iter(&self) -> IntoIter<(&String, &ImportConditions)> {
        let mut imports = self.map.iter().collect::<Vec<_>>();
        imports.sort_by(compare_imports);
        imports.into_iter()
    }
}

pub struct ImportsWithDefault<'a> {
    imports: &'a mut Imports,
}

impl<'a> ImportsWithDefault<'a> {
    fn new(imports: &'a mut Imports) -> Self {
        Self { imports }
    }
}

impl Drop for ImportsWithDefault<'_> {
    fn drop(&mut self) {
        self.imports.reset_defaults()
    }
}

impl Deref for ImportsWithDefault<'_> {
    type Target = Imports;
    fn deref(&self) -> &Self::Target {
        self.imports
    }
}

impl DerefMut for ImportsWithDefault<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.imports
    }
}

#[derive(Clone, Debug, Default)]
pub struct ImportConditions {
    pub version: Option<Version>,
    pub constraints: Vec<String>,
}

impl ImportConditions {
    fn clear(&mut self) {
        *self = ImportConditions::default();
    }

    fn update_version(&mut self, version: Option<Version>) {
        if version < self.version {
            self.version = version;
        }
    }

    fn add_constraint(&mut self, constraint: String) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        // Otherwise, we just check if the constraint
        // is already present or not before adding it.
        if !self.constraints.iter().any(|x| x == &constraint) {
            self.constraints.push(constraint);
        }
    }

    fn update_constraints(&mut self, constraints: Vec<String>) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        if constraints.is_empty() {
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            self.constraints.clear();
        } else {
            // Otherwise, we just check if the constraint
            // is already present or not before adding it.
            for constraint in constraints {
                if !self.constraints.iter().any(|x| x == &constraint) {
                    self.constraints.push(constraint.clone());
                }
            }
        }
    }
}

fn make_crate_name(gir: &Library) -> String {
    let name = gir.namespace(namespaces::MAIN).name.as_str();
    if name == "GObject" {
        crate_name("GLib")
    } else {
        crate_name(name)
    }
}
<commit_msg>Remove unneeded clone on constraints<commit_after>use super::namespaces;
use crate::{library::Library, nameutil::crate_name, version::Version};
use std::cmp::Ordering;
use std::collections::btree_map::BTreeMap;
use std::ops::{Deref, DerefMut};
use std::vec::IntoIter;

fn is_first_char_up(s: &str) -> bool {
    s.chars().next().unwrap().is_uppercase()
}

fn check_up_eq(a: &str, b: &str) -> Ordering {
    let is_a_up = is_first_char_up(a);
    let is_b_up = is_first_char_up(b);
    if is_a_up != is_b_up {
        if is_a_up {
            return Ordering::Greater;
        }
        return Ordering::Less;
    }
    Ordering::Equal
}

/// This function is used by the `Imports` type to generate output like `cargo fmt` would.
///
/// For example:
///
/// ```text
/// use gdk; // lowercases come first.
/// use Window;
///
/// use gdk::foo; // lowercases come first here as well.
/// use gdk::Foo;
/// ```
#[allow(clippy::trivially_copy_pass_by_ref)] // because of the way this is used, the refs are needed
fn compare_imports(a: &(&String, &ImportConditions), b: &(&String, &ImportConditions)) -> Ordering {
    let s = check_up_eq(a.0, b.0);
    if s != Ordering::Equal {
        return s;
    }
    let mut a = a.0.split("::");
    let mut b = b.0.split("::");
    loop {
        match (a.next(), b.next()) {
            (Some(a), Some(b)) => {
                let s = check_up_eq(a, b);
                if s != Ordering::Equal {
                    break s;
                }
                let s = a.partial_cmp(b).unwrap();
                if s != Ordering::Equal {
                    break s;
                }
            }
            (Some(_), None) => break Ordering::Greater,
            (None, Some(_)) => break Ordering::Less,
            (None, None) => break Ordering::Equal,
        }
    }
}

/// Provides assistance in generating use declarations.
///
/// It takes into account that use declaration referring to names within the
/// same crate will look differently. It also avoids generating spurious
/// declarations referring to names from within the same module as the one we
/// are generating code for.
#[derive(Clone, Debug, Default)]
pub struct Imports {
    /// Name of the current crate.
    crate_name: String,
    /// Name defined within current module. It doesn't need use declaration.
    ///
    /// NOTE: Currently we don't need to support more than one such name.
    defined: Option<String>,
    defaults: ImportConditions,
    map: BTreeMap<String, ImportConditions>,
}

impl Imports {
    pub fn new(gir: &Library) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: None,
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    pub fn with_defined(gir: &Library, name: &str) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: Some(name.to_owned()),
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    #[must_use = "ImportsWithDefault must live while defaults are needed"]
    pub fn with_defaults(
        &mut self,
        version: Option<Version>,
        constraint: &Option<String>,
    ) -> ImportsWithDefault<'_> {
        let constraints = if let Some(constraint) = constraint {
            vec![constraint.clone()]
        } else {
            vec![]
        };
        self.defaults = ImportConditions {
            version,
            constraints,
        };

        ImportsWithDefault::new(self)
    }

    fn reset_defaults(&mut self) {
        self.defaults.clear();
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    /// Uses defaults.
    pub fn add(&mut self, name: &str) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let defaults = &self.defaults;
            let entry = self
                .map
                .entry(name.to_owned())
                .or_insert_with(|| defaults.clone());
            entry.update_version(self.defaults.version);
            entry.update_constraints(&self.defaults.constraints);
        }
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_version(&mut self, name: &str, version: Option<Version>) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                version,
                constraints: Vec::new(),
            });
            entry.update_version(version);
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            entry.constraints.clear();
        }
    }

    /// Declares that name should be available through its last path component and provides
    /// an optional feature constraint.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_constraint(
        &mut self,
        name: &str,
        version: Option<Version>,
        constraint: Option<&str>,
    ) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = if let Some(constraint) = constraint {
                let constraint = String::from(constraint);
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: vec![constraint.clone()],
                });
                entry.add_constraint(constraint);
                entry
            } else {
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: Vec::new(),
                });
                // Since there is no constraint on this import, if any constraint
                // is present, we can just remove it.
                entry.constraints.clear();
                entry
            };
            entry.update_version(version);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type(&mut self, used_type: &str) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add(&used_type[2..]);
            } else {
                self.add(&used_type[..i]);
            }
        } else {
            self.add(used_type);
        }
    }

    pub fn add_used_types(&mut self, used_types: &[String]) {
        for s in used_types {
            self.add_used_type(s);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type_with_version(&mut self, used_type: &str, version: Option<Version>) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add_with_version(&used_type[2..], version);
            } else {
                self.add_with_version(&used_type[..i], version);
            }
        } else {
            self.add_with_version(used_type, version);
        }
    }

    /// Tries to strip crate name prefix from given name.
    ///
    /// Returns `None` if name matches crate name exactly. Otherwise returns
    /// name with crate name prefix stripped or full name if there was no match.
    fn strip_crate_name<'a>(&self, name: &'a str) -> Option<&'a str> {
        let prefix = &self.crate_name;
        if !name.starts_with(prefix) {
            return Some(name);
        }
        let rest = &name[prefix.len()..];
        if rest.is_empty() {
            None
        } else if rest.starts_with("::") {
            Some(&rest["::".len()..])
        } else {
            // It was false positive, return the whole name.
            Some(name)
        }
    }

    pub fn iter(&self) -> IntoIter<(&String, &ImportConditions)> {
        let mut imports = self.map.iter().collect::<Vec<_>>();
        imports.sort_by(compare_imports);
        imports.into_iter()
    }
}

pub struct ImportsWithDefault<'a> {
    imports: &'a mut Imports,
}

impl<'a> ImportsWithDefault<'a> {
    fn new(imports: &'a mut Imports) -> Self {
        Self { imports }
    }
}

impl Drop for ImportsWithDefault<'_> {
    fn drop(&mut self) {
        self.imports.reset_defaults()
    }
}

impl Deref for ImportsWithDefault<'_> {
    type Target = Imports;
    fn deref(&self) -> &Self::Target {
        self.imports
    }
}

impl DerefMut for ImportsWithDefault<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.imports
    }
}

#[derive(Clone, Debug, Default)]
pub struct ImportConditions {
    pub version: Option<Version>,
    pub constraints: Vec<String>,
}

impl ImportConditions {
    fn clear(&mut self) {
        self.version = None;
        self.constraints.clear();
    }

    fn update_version(&mut self, version: Option<Version>) {
        if version < self.version {
            self.version = version;
        }
    }

    fn add_constraint(&mut self, constraint: String) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        // Otherwise, we just check if the constraint
        // is already present or not before adding it.
        if !self.constraints.iter().any(|x| x == &constraint) {
            self.constraints.push(constraint);
        }
    }

    fn update_constraints(&mut self, constraints: &[String]) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        if constraints.is_empty() {
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            self.constraints.clear();
        } else {
            // Otherwise, we just check if the constraint
            // is already present or not before adding it.
            for constraint in constraints {
                if !self.constraints.iter().any(|x| x == constraint) {
                    self.constraints.push(constraint.clone());
                }
            }
        }
    }
}

fn make_crate_name(gir: &Library) -> String {
    let name = gir.namespace(namespaces::MAIN).name.as_str();
    if name == "GObject" {
        crate_name("GLib")
    } else {
        crate_name(name)
    }
}
<|endoftext|>


<commit_before>use futures::future::ok;

use std::rc::Rc;

use super::{BoxedNewPeerFuture, Peer};
use super::{ConstructParams, PeerConstructor, Specifier};

use std::io::Read;
use tokio_io::AsyncRead;

use std::io::Error as IoError;

#[derive(Debug)]
pub struct Message2Line<T: Specifier>(pub T);
impl<T: Specifier> Specifier for Message2Line<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p| packet2line_peer(p))
    }
    specifier_boilerplate!(typ=Line noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = Message2LineClass,
    target = Message2Line,
    prefixes = ["msg2line:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
Line filter: ensure each message (a chunk from one read call from underlying specifier)
contains no inner newlines and terminates with one newline.

Reverse of the `line2msg:`.

Replaces both newlines (\x0A) and carrige returns (\x0D) with spaces (\x20) for each read.

Does not affect writing at all. Use this specifier on both ends to get bi-directional behaviour.

Automatically inserted by --line option on top of the stack containing a websocket.

Example: TODO
"#
);

#[derive(Debug)]
pub struct Line2Message<T: Specifier>(pub T);
impl<T: Specifier> Specifier for Line2Message<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let retain_newlines = !cp.program_options.linemode_strip_newlines;
        let inner = self.0.construct(cp.clone());
        inner.map(move |p| line2packet_peer(p, retain_newlines))
    }
    specifier_boilerplate!(typ=Line noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name=Line2MessageClass,
    target=Line2Message,
    prefixes=["line2msg:"],
    arg_handling=subspec,
    overlay = true,
    MessageOriented,
    MulticonnectnessDependsOnInnerType,
    help=r#"
Line filter: encure that each message (a successful read call) is obtained from a line
coming from underlying specifier, buffering up or splitting content as needed.

Reverse of the `msg2line:`.

Does not affect writing at all. Use this specifier on both ends to get bi-directional behaviour.

Automatically inserted by --line option at the top of the stack opposite to websocket-containing stack.

Example: TODO
"#
);

pub fn packet2line_peer(inner_peer: Peer) -> BoxedNewPeerFuture {
    let filtered = Packet2LineWrapper(inner_peer.0);
    let thepeer = Peer::new(filtered, inner_peer.1);
    Box::new(ok(thepeer)) as BoxedNewPeerFuture
}
struct Packet2LineWrapper(Box<AsyncRead>);

impl Read for Packet2LineWrapper {
    fn read(&mut self, b: &mut [u8]) -> Result<usize, IoError> {
        let l = b.len();
        assert!(l > 1);
        let mut n = match self.0.read(&mut b[..(l - 1)]) {
            Ok(x) => x,
            Err(e) => return Err(e),
        };
        if n == 0 {
            return Ok(n);
        }
        // chomp away \n or \r\n
        if n > 0 && b[n - 1] == b'\n' {
            n -= 1;
        }
        if n > 0 && b[n - 1] == b'\r' {
            n -= 1;
        }
        // replace those with spaces
        for c in b.iter_mut().take(n) {
            if *c == b'\n' || *c == b'\r' {
                *c = b' ';
            }
        }
        // add back one \n
        b[n] = b'\n';
        n += 1;

        Ok(n)
    }
}
impl AsyncRead for Packet2LineWrapper {}

pub fn line2packet_peer(inner_peer: Peer, retain_newlines: bool) -> BoxedNewPeerFuture {
    let filtered = Line2PacketWrapper {
        inner: inner_peer.0,
        queue: vec![],
        retain_newlines,
    };
    let thepeer = Peer::new(filtered, inner_peer.1);
    Box::new(ok(thepeer)) as BoxedNewPeerFuture
}
struct Line2PacketWrapper {
    inner: Box<AsyncRead>,
    queue: Vec<u8>,
    retain_newlines: bool,
}

impl Read for Line2PacketWrapper {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, IoError> {
        //eprint!("ql={} ", self.queue.len());
        let mut queued_line_len = None;
        for i in 0..self.queue.len() {
            if self.queue[i] == b'\n' {
                queued_line_len = Some(i);
                break;
            }
        }
        //eprint!("qll={:?} ", queued_line_len);

        if let Some(mut n) = queued_line_len {
            n += 1;
            buf[0..n].copy_from_slice(&self.queue[0..n]);
            ::std::mem::drop(self.queue.drain(0..n));
            if !self.retain_newlines {
                if n > 0 && (buf[n - 1] == b'\n') {
                    n -= 1
                }
                if n > 0 && (buf[n - 1] == b'\r') {
                    n -= 1
                }
            }
            //eprintln!("n={}", n);
            Ok(n)
        } else {
            let mut n = match self.inner.read(buf) {
                Ok(x) => x,
                Err(e) => return Err(e),
            };

            if n == 0 {
                if self.queue.is_empty() {
                    warn!(
                        "Throwing away {} bytes of incomplete line",
                        self.queue.len()
                    );
                }
                return Ok(0);
            }

            let mut happy_case =
                self.queue.is_empty() && (!buf[0..(n - 1)].contains(&b'\n')) && buf[n - 1] == b'\n';

            if happy_case {
                // Specifically to avoid allocations when data is already nice
                if !self.retain_newlines {
                    if n > 0 && (buf[n - 1] == b'\n') {
                        n -= 1
                    }
                    if n > 0 && (buf[n - 1] == b'\r') {
                        n -= 1
                    }
                }
                //eprintln!("happy n={}", n);
                Ok(n)
            } else {
                // Just queue up and recurse
                self.queue.extend_from_slice(&buf[0..n]);
                //eprintln!(" recurse");
                self.read(buf)
            }
        }
    }
}
impl AsyncRead for Line2PacketWrapper {}
<commit_msg>Fix line_peer's warning<commit_after>use futures::future::ok;

use std::rc::Rc;

use super::{BoxedNewPeerFuture, Peer};
use super::{ConstructParams, PeerConstructor, Specifier};

use std::io::Read;
use tokio_io::AsyncRead;

use std::io::Error as IoError;

#[derive(Debug)]
pub struct Message2Line<T: Specifier>(pub T);
impl<T: Specifier> Specifier for Message2Line<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p| packet2line_peer(p))
    }
    specifier_boilerplate!(typ=Line noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = Message2LineClass,
    target = Message2Line,
    prefixes = ["msg2line:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
Line filter: ensure each message (a chunk from one read call from underlying specifier)
contains no inner newlines and terminates with one newline.

Reverse of the `line2msg:`.

Replaces both newlines (\x0A) and carrige returns (\x0D) with spaces (\x20) for each read.

Does not affect writing at all. Use this specifier on both ends to get bi-directional behaviour.

Automatically inserted by --line option on top of the stack containing a websocket.

Example: TODO
"#
);

#[derive(Debug)]
pub struct Line2Message<T: Specifier>(pub T);
impl<T: Specifier> Specifier for Line2Message<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let retain_newlines = !cp.program_options.linemode_strip_newlines;
        let inner = self.0.construct(cp.clone());
        inner.map(move |p| line2packet_peer(p, retain_newlines))
    }
    specifier_boilerplate!(typ=Line noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name=Line2MessageClass,
    target=Line2Message,
    prefixes=["line2msg:"],
    arg_handling=subspec,
    overlay = true,
    MessageOriented,
    MulticonnectnessDependsOnInnerType,
    help=r#"
Line filter: encure that each message (a successful read call) is obtained from a line
coming from underlying specifier, buffering up or splitting content as needed.

Reverse of the `msg2line:`.

Does not affect writing at all. Use this specifier on both ends to get bi-directional behaviour.

Automatically inserted by --line option at the top of the stack opposite to websocket-containing stack.

Example: TODO
"#
);

pub fn packet2line_peer(inner_peer: Peer) -> BoxedNewPeerFuture {
    let filtered = Packet2LineWrapper(inner_peer.0);
    let thepeer = Peer::new(filtered, inner_peer.1);
    Box::new(ok(thepeer)) as BoxedNewPeerFuture
}
struct Packet2LineWrapper(Box<AsyncRead>);

impl Read for Packet2LineWrapper {
    fn read(&mut self, b: &mut [u8]) -> Result<usize, IoError> {
        let l = b.len();
        assert!(l > 1);
        let mut n = match self.0.read(&mut b[..(l - 1)]) {
            Ok(x) => x,
            Err(e) => return Err(e),
        };
        if n == 0 {
            return Ok(n);
        }
        // chomp away \n or \r\n
        if n > 0 && b[n - 1] == b'\n' {
            n -= 1;
        }
        if n > 0 && b[n - 1] == b'\r' {
            n -= 1;
        }
        // replace those with spaces
        for c in b.iter_mut().take(n) {
            if *c == b'\n' || *c == b'\r' {
                *c = b' ';
            }
        }
        // add back one \n
        b[n] = b'\n';
        n += 1;

        Ok(n)
    }
}
impl AsyncRead for Packet2LineWrapper {}

pub fn line2packet_peer(inner_peer: Peer, retain_newlines: bool) -> BoxedNewPeerFuture {
    let filtered = Line2PacketWrapper {
        inner: inner_peer.0,
        queue: vec![],
        retain_newlines,
    };
    let thepeer = Peer::new(filtered, inner_peer.1);
    Box::new(ok(thepeer)) as BoxedNewPeerFuture
}
struct Line2PacketWrapper {
    inner: Box<AsyncRead>,
    queue: Vec<u8>,
    retain_newlines: bool,
}

impl Read for Line2PacketWrapper {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, IoError> {
        //eprint!("ql={} ", self.queue.len());
        let mut queued_line_len = None;
        for i in 0..self.queue.len() {
            if self.queue[i] == b'\n' {
                queued_line_len = Some(i);
                break;
            }
        }
        //eprint!("qll={:?} ", queued_line_len);

        if let Some(mut n) = queued_line_len {
            n += 1;
            buf[0..n].copy_from_slice(&self.queue[0..n]);
            ::std::mem::drop(self.queue.drain(0..n));
            if !self.retain_newlines {
                if n > 0 && (buf[n - 1] == b'\n') {
                    n -= 1
                }
                if n > 0 && (buf[n - 1] == b'\r') {
                    n -= 1
                }
            }
            //eprintln!("n={}", n);
            Ok(n)
        } else {
            let mut n = match self.inner.read(buf) {
                Ok(x) => x,
                Err(e) => return Err(e),
            };

            if n == 0 {
                if !self.queue.is_empty() {
                    warn!(
                        "Throwing away {} bytes of incomplete line",
                        self.queue.len()
                    );
                }
                return Ok(0);
            }

            let mut happy_case =
                self.queue.is_empty() && (!buf[0..(n - 1)].contains(&b'\n')) && buf[n - 1] == b'\n';

            if happy_case {
                // Specifically to avoid allocations when data is already nice
                if !self.retain_newlines {
                    if n > 0 && (buf[n - 1] == b'\n') {
                        n -= 1
                    }
                    if n > 0 && (buf[n - 1] == b'\r') {
                        n -= 1
                    }
                }
                //eprintln!("happy n={}", n);
                Ok(n)
            } else {
                // Just queue up and recurse
                self.queue.extend_from_slice(&buf[0..n]);
                //eprintln!(" recurse");
                self.read(buf)
            }
        }
    }
}
impl AsyncRead for Line2PacketWrapper {}
<|endoftext|>


<commit_before>use std::collections::HashSet;
use std::sync::Arc;
use super::types::{
    AttestationRecord,
    AttesterMap,
};
use super::attestation_parent_hashes::{
    attestation_parent_hashes,
    ParentHashesError,
};
use super::db::{
    ClientDB,
    DBError
};
use super::db::stores::{
    BlockStore,
    ValidatorStore,
};
use super::types::{
    Hash256,
};
use super::message_generation::generate_signed_message;
use super::signature_verification::{
    verify_aggregate_signature_for_indices,
    SignatureVerificationError,
};

#[derive(Debug,PartialEq)]
pub enum AttestationValidationError {
    SlotTooHigh,
    SlotTooLow,
    JustifiedSlotIncorrect,
    UnknownJustifiedBlock,
    TooManyObliqueHashes,
    BadCurrentHashes,
    BadObliqueHashes,
    BadAttesterMap,
    IntWrapping,
    PublicKeyCorrupt,
    NoPublicKeyForValidator,
    BadBitfieldLength,
    InvalidBitfield,
    InvalidBitfieldEndBits,
    NoSignatures,
    NonZeroTrailingBits,
    BadAggregateSignature,
    DBError(String),
}

/// The context against which some attestation should be validated.
pub struct AttestationValidationContext<T>
    where T: ClientDB + Sized
{
    /// The slot as determined by the system time.
    pub block_slot: u64,
    /// The cycle_length as determined by the chain configuration.
    pub cycle_length: u8,
    /// The last justified slot as per the client's view of the canonical chain.
    pub last_justified_slot: u64,
    /// A vec of the hashes of the blocks preceeding the present slot.
    pub parent_hashes: Arc<Vec<Hash256>>,
    /// The store containing block information.
    pub block_store: Arc<BlockStore<T>>,
    /// The store containing validator information.
    pub validator_store: Arc<ValidatorStore<T>>,
    /// A map of (slot, shard_id) to the attestation set of validation indices.
    pub attester_map: Arc<AttesterMap>,
}

impl<T> AttestationValidationContext<T>
    where T: ClientDB
{
    /// Validate a (fully deserialized) AttestationRecord against this context.
    ///
    /// The function will return a HashSet of validator indices (canonical validator indices not
    /// attestation indices) if the validation passed successfully, or an error otherwise.
    ///
    /// The attestation's aggregate signature will be verified, therefore the function must able to
    /// access all required validation public keys via the `validator_store`.
    pub fn validate_attestation(&self, a: &AttestationRecord)
        -> Result<HashSet<usize>, AttestationValidationError>
    {
        /*
         * The attesation slot must not be higher than the block that contained it.
         */
        if a.slot > self.block_slot {
            return Err(AttestationValidationError::SlotTooHigh);
        }

        /*
         * The slot of this attestation must not be more than cycle_length + 1 distance
         * from the block that contained it.
         */
        if a.slot < self.block_slot
            .saturating_sub(u64::from(self.cycle_length).saturating_add(1)) {
            return Err(AttestationValidationError::SlotTooLow);
        }

        /*
         * The attestation must indicate that its last justified slot is the same as the last justified
         * slot known to us.
         */
        if a.justified_slot != self.last_justified_slot {
            return Err(AttestationValidationError::JustifiedSlotIncorrect);
        }

        /*
         * There is no need to include more oblique parents hashes than there are blocks
         * in a cycle.
         */
        if a.oblique_parent_hashes.len() > usize::from(self.cycle_length) {
            return Err(AttestationValidationError::TooManyObliqueHashes);
        }

        /*
         * Retrieve the set of attestation indices for this slot and shard id.
         *
         * This is an array mapping the order that validators will appear in the bitfield to the
         * canonincal index of a validator.
         */
        let attestation_indices = self.attester_map.get(&(a.slot, a.shard_id))
            .ok_or(AttestationValidationError::BadAttesterMap)?;

        /*
         * The bitfield must be no longer than the minimum required to represent each validator in the
         * attestation indicies for this slot and shard id.
         */
        if a.attester_bitfield.num_bytes() !=
            bytes_for_bits(attestation_indices.len())
        {
            return Err(AttestationValidationError::BadBitfieldLength);
       }

        /*
         * If there are excess bits in the bitfield because the number of a validators in not a
         * multiple of 8, reject this attestation record.
         *
         * Allow extra set bits would permit mutliple different byte layouts (and therefore hashes) to
         * refer to the same AttesationRecord.
         */
        if a.attester_bitfield.len() > attestation_indices.len() {
            return Err(AttestationValidationError::InvalidBitfieldEndBits)
        }

        /*
         * The specified justified block hash must be known to us
         */
        if !self.block_store.block_exists(&a.justified_block_hash)? {
            return Err(AttestationValidationError::UnknownJustifiedBlock)
        }

        let signed_message = {
            let parent_hashes = attestation_parent_hashes(
                self.cycle_length,
                self.block_slot,
                a.slot,
                &self.parent_hashes,
                &a.oblique_parent_hashes)?;
            generate_signed_message(
                a.slot,
                &parent_hashes,
                a.shard_id,
                &a.shard_block_hash,
                a.justified_slot)
        };

        let voted_hashmap =
            verify_aggregate_signature_for_indices(
                &signed_message,
                &a.aggregate_sig,
                &attestation_indices,
                &a.attester_bitfield,
                &self.validator_store)?;

        /*
         * If the hashmap of voters is None, the signature verification failed.
         */
        match voted_hashmap {
            None => Err(AttestationValidationError::BadAggregateSignature),
            Some(hashmap) => Ok(hashmap),
        }
    }
}

fn bytes_for_bits(bits: usize) -> usize {
    (bits.saturating_sub(1) / 8) + 1
}

impl From<ParentHashesError> for AttestationValidationError {
    fn from(e: ParentHashesError) -> Self {
        match e {
            ParentHashesError::BadCurrentHashes
                => AttestationValidationError::BadCurrentHashes,
            ParentHashesError::BadObliqueHashes
                => AttestationValidationError::BadObliqueHashes,
            ParentHashesError::SlotTooLow
                => AttestationValidationError::SlotTooLow,
            ParentHashesError::SlotTooHigh
                => AttestationValidationError::SlotTooHigh,
            ParentHashesError::IntWrapping
                => AttestationValidationError::IntWrapping
        }
    }
}

impl From<DBError> for AttestationValidationError {
    fn from(e: DBError) -> Self {
        AttestationValidationError::DBError(e.message)
    }
}

impl From<SignatureVerificationError> for AttestationValidationError {
    fn from(e: SignatureVerificationError) -> Self {
        match e {
            SignatureVerificationError::BadValidatorIndex
                => AttestationValidationError::BadAttesterMap,
            SignatureVerificationError::PublicKeyCorrupt
                => AttestationValidationError::PublicKeyCorrupt,
            SignatureVerificationError::NoPublicKeyForValidator
                => AttestationValidationError::NoPublicKeyForValidator,
            SignatureVerificationError::DBError(s)
                => AttestationValidationError::DBError(s),
        }
    }
}
<commit_msg>Rename voted_hashmap to voted_hashset<commit_after>use std::collections::HashSet;
use std::sync::Arc;
use super::types::{
    AttestationRecord,
    AttesterMap,
};
use super::attestation_parent_hashes::{
    attestation_parent_hashes,
    ParentHashesError,
};
use super::db::{
    ClientDB,
    DBError
};
use super::db::stores::{
    BlockStore,
    ValidatorStore,
};
use super::types::{
    Hash256,
};
use super::message_generation::generate_signed_message;
use super::signature_verification::{
    verify_aggregate_signature_for_indices,
    SignatureVerificationError,
};

#[derive(Debug,PartialEq)]
pub enum AttestationValidationError {
    SlotTooHigh,
    SlotTooLow,
    JustifiedSlotIncorrect,
    UnknownJustifiedBlock,
    TooManyObliqueHashes,
    BadCurrentHashes,
    BadObliqueHashes,
    BadAttesterMap,
    IntWrapping,
    PublicKeyCorrupt,
    NoPublicKeyForValidator,
    BadBitfieldLength,
    InvalidBitfield,
    InvalidBitfieldEndBits,
    NoSignatures,
    NonZeroTrailingBits,
    BadAggregateSignature,
    DBError(String),
}

/// The context against which some attestation should be validated.
pub struct AttestationValidationContext<T>
    where T: ClientDB + Sized
{
    /// The slot as determined by the system time.
    pub block_slot: u64,
    /// The cycle_length as determined by the chain configuration.
    pub cycle_length: u8,
    /// The last justified slot as per the client's view of the canonical chain.
    pub last_justified_slot: u64,
    /// A vec of the hashes of the blocks preceeding the present slot.
    pub parent_hashes: Arc<Vec<Hash256>>,
    /// The store containing block information.
    pub block_store: Arc<BlockStore<T>>,
    /// The store containing validator information.
    pub validator_store: Arc<ValidatorStore<T>>,
    /// A map of (slot, shard_id) to the attestation set of validation indices.
    pub attester_map: Arc<AttesterMap>,
}

impl<T> AttestationValidationContext<T>
    where T: ClientDB
{
    /// Validate a (fully deserialized) AttestationRecord against this context.
    ///
    /// The function will return a HashSet of validator indices (canonical validator indices not
    /// attestation indices) if the validation passed successfully, or an error otherwise.
    ///
    /// The attestation's aggregate signature will be verified, therefore the function must able to
    /// access all required validation public keys via the `validator_store`.
    pub fn validate_attestation(&self, a: &AttestationRecord)
        -> Result<HashSet<usize>, AttestationValidationError>
    {
        /*
         * The attesation slot must not be higher than the block that contained it.
         */
        if a.slot > self.block_slot {
            return Err(AttestationValidationError::SlotTooHigh);
        }

        /*
         * The slot of this attestation must not be more than cycle_length + 1 distance
         * from the block that contained it.
         */
        if a.slot < self.block_slot
            .saturating_sub(u64::from(self.cycle_length).saturating_add(1)) {
            return Err(AttestationValidationError::SlotTooLow);
        }

        /*
         * The attestation must indicate that its last justified slot is the same as the last justified
         * slot known to us.
         */
        if a.justified_slot != self.last_justified_slot {
            return Err(AttestationValidationError::JustifiedSlotIncorrect);
        }

        /*
         * There is no need to include more oblique parents hashes than there are blocks
         * in a cycle.
         */
        if a.oblique_parent_hashes.len() > usize::from(self.cycle_length) {
            return Err(AttestationValidationError::TooManyObliqueHashes);
        }

        /*
         * Retrieve the set of attestation indices for this slot and shard id.
         *
         * This is an array mapping the order that validators will appear in the bitfield to the
         * canonincal index of a validator.
         */
        let attestation_indices = self.attester_map.get(&(a.slot, a.shard_id))
            .ok_or(AttestationValidationError::BadAttesterMap)?;

        /*
         * The bitfield must be no longer than the minimum required to represent each validator in the
         * attestation indicies for this slot and shard id.
         */
        if a.attester_bitfield.num_bytes() !=
            bytes_for_bits(attestation_indices.len())
        {
            return Err(AttestationValidationError::BadBitfieldLength);
       }

        /*
         * If there are excess bits in the bitfield because the number of a validators in not a
         * multiple of 8, reject this attestation record.
         *
         * Allow extra set bits would permit mutliple different byte layouts (and therefore hashes) to
         * refer to the same AttesationRecord.
         */
        if a.attester_bitfield.len() > attestation_indices.len() {
            return Err(AttestationValidationError::InvalidBitfieldEndBits)
        }

        /*
         * The specified justified block hash must be known to us
         */
        if !self.block_store.block_exists(&a.justified_block_hash)? {
            return Err(AttestationValidationError::UnknownJustifiedBlock)
        }

        let signed_message = {
            let parent_hashes = attestation_parent_hashes(
                self.cycle_length,
                self.block_slot,
                a.slot,
                &self.parent_hashes,
                &a.oblique_parent_hashes)?;
            generate_signed_message(
                a.slot,
                &parent_hashes,
                a.shard_id,
                &a.shard_block_hash,
                a.justified_slot)
        };

        let voted_hashset =
            verify_aggregate_signature_for_indices(
                &signed_message,
                &a.aggregate_sig,
                &attestation_indices,
                &a.attester_bitfield,
                &self.validator_store)?;

        /*
         * If the hashmap of voters is None, the signature verification failed.
         */
        match voted_hashset {
            None => Err(AttestationValidationError::BadAggregateSignature),
            Some(hashmap) => Ok(hashmap),
        }
    }
}

fn bytes_for_bits(bits: usize) -> usize {
    (bits.saturating_sub(1) / 8) + 1
}

impl From<ParentHashesError> for AttestationValidationError {
    fn from(e: ParentHashesError) -> Self {
        match e {
            ParentHashesError::BadCurrentHashes
                => AttestationValidationError::BadCurrentHashes,
            ParentHashesError::BadObliqueHashes
                => AttestationValidationError::BadObliqueHashes,
            ParentHashesError::SlotTooLow
                => AttestationValidationError::SlotTooLow,
            ParentHashesError::SlotTooHigh
                => AttestationValidationError::SlotTooHigh,
            ParentHashesError::IntWrapping
                => AttestationValidationError::IntWrapping
        }
    }
}

impl From<DBError> for AttestationValidationError {
    fn from(e: DBError) -> Self {
        AttestationValidationError::DBError(e.message)
    }
}

impl From<SignatureVerificationError> for AttestationValidationError {
    fn from(e: SignatureVerificationError) -> Self {
        match e {
            SignatureVerificationError::BadValidatorIndex
                => AttestationValidationError::BadAttesterMap,
            SignatureVerificationError::PublicKeyCorrupt
                => AttestationValidationError::PublicKeyCorrupt,
            SignatureVerificationError::NoPublicKeyForValidator
                => AttestationValidationError::NoPublicKeyForValidator,
            SignatureVerificationError::DBError(s)
                => AttestationValidationError::DBError(s),
        }
    }
}
<|endoftext|>


<commit_before>use indexmap::IndexMap;
use semver::Version;
use serde::{Deserialize, Deserializer, Serialize};

use std::collections::{BTreeMap, BTreeSet};

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Definitions {
    /// The Syn version used to generate the introspection file.
    pub version: Version,
    pub types: Vec<Node>,
    pub tokens: BTreeMap<String, String>,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Node {
    pub ident: String,
    pub features: Features,
    #[serde(
        flatten,
        skip_serializing_if = "is_private",
        deserialize_with = "de_data"
    )]
    pub data: Data,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum Data {
    Private,
    #[serde(rename = "fields")]
    Struct(Fields),
    #[serde(rename = "variants")]
    Enum(Variants),
}

pub type Fields = IndexMap<String, Type>;
pub type Variants = IndexMap<String, Vec<Type>>;

#[derive(Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Type {
    /// Type defined by `syn`
    Syn(String),

    /// Type defined in `std`.
    Std(String),

    /// Type external to `syn`
    #[serde(rename = "proc_macro2")]
    Ext(String),

    /// Token type
    Token(String),

    /// Token group
    Group(String),

    /// Punctuated list
    Punctuated(Punctuated),

    Option(Box<Type>),
    Box(Box<Type>),
    Vec(Box<Type>),
    Tuple(Vec<Type>),
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Punctuated {
    pub element: Box<Type>,
    pub punct: String,
}

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct Features {
    pub any: BTreeSet<String>,
}

fn is_private(data: &Data) -> bool {
    match data {
        Data::Private => true,
        Data::Struct(_) | Data::Enum(_) => false,
    }
}

fn de_data<'de, D>(deserializer: D) -> Result<Data, D::Error>
where
    D: Deserializer<'de>,
{
    let option = Option::deserialize(deserializer)?;
    Ok(option.unwrap_or(Data::Private))
}
<commit_msg>Better name for deserializing Private data if absent<commit_after>use indexmap::IndexMap;
use semver::Version;
use serde::{Deserialize, Deserializer, Serialize};

use std::collections::{BTreeMap, BTreeSet};

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Definitions {
    /// The Syn version used to generate the introspection file.
    pub version: Version,
    pub types: Vec<Node>,
    pub tokens: BTreeMap<String, String>,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Node {
    pub ident: String,
    pub features: Features,
    #[serde(
        flatten,
        skip_serializing_if = "is_private",
        deserialize_with = "private_if_absent"
    )]
    pub data: Data,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum Data {
    Private,
    #[serde(rename = "fields")]
    Struct(Fields),
    #[serde(rename = "variants")]
    Enum(Variants),
}

pub type Fields = IndexMap<String, Type>;
pub type Variants = IndexMap<String, Vec<Type>>;

#[derive(Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Type {
    /// Type defined by `syn`
    Syn(String),

    /// Type defined in `std`.
    Std(String),

    /// Type external to `syn`
    #[serde(rename = "proc_macro2")]
    Ext(String),

    /// Token type
    Token(String),

    /// Token group
    Group(String),

    /// Punctuated list
    Punctuated(Punctuated),

    Option(Box<Type>),
    Box(Box<Type>),
    Vec(Box<Type>),
    Tuple(Vec<Type>),
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Punctuated {
    pub element: Box<Type>,
    pub punct: String,
}

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct Features {
    pub any: BTreeSet<String>,
}

fn is_private(data: &Data) -> bool {
    match data {
        Data::Private => true,
        Data::Struct(_) | Data::Enum(_) => false,
    }
}

fn private_if_absent<'de, D>(deserializer: D) -> Result<Data, D::Error>
where
    D: Deserializer<'de>,
{
    let option = Option::deserialize(deserializer)?;
    Ok(option.unwrap_or(Data::Private))
}
<|endoftext|>


<commit_before>use ::{KeyType, ValueType};

use wal_file::KeyValuePair;

use std::iter::{Peekable, empty, Map};
use std::collections::{BTreeMap, BTreeSet};
use std::collections::btree_map;
use std::collections::btree_set;

pub struct MultiMap<K: KeyType, V: ValueType> {
    multi_map: BTreeMap<K, BTreeSet<V>>,
    count: usize  // total number of KV pairs
}

pub struct MultiMapIterator<'a, K: KeyType + 'a, V: ValueType + 'a> {
    cur_key: Option<&'a K>,
    key_it: btree_map::Iter<'a,K,BTreeSet<V>>,
    value_it: Option<btree_set::Iter<'a,V>>,
}

impl <'a, K: KeyType, V: ValueType> MultiMap<K,V> {
    pub fn new() -> MultiMap<K,V> {
        return MultiMap{multi_map: BTreeMap::<K,BTreeSet<V>>::new(), count: 0};
    }

    pub fn insert(&mut self, key: K, value: V) -> usize {
        self.count += 1;

        if let Some(set) = self.multi_map.get_mut(&key) {
            set.insert(value);
            return self.count;
        }

        let mut set = BTreeSet::<V>::new();

        set.insert(value);

        self.multi_map.insert(key, set);

        return self.count;
    }
}

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

    fn into_iter(self) -> Self::IntoIter {
        let mut key_it = self.multi_map.iter();
        let cur_entry = key_it.next();

        // check to see if our map is empty
        if cur_entry.is_none() {
            return MultiMapIterator{cur_key: None, key_it: key_it, value_it: None};
        }

        // safe to call unwrap as we tested above
        let (cur_key, cur_set) = cur_entry.unwrap();

        return MultiMapIterator{cur_key: Some(cur_key), key_it: key_it, value_it: Some(cur_set.iter())};
    }
}

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

    fn next(&mut self) -> Option<Self::Item> {
        // this is our invariant, when it's None we've gone through everything
        if self.cur_key.is_none() {
            return None
        }

        // should be safe to call unwrap here, because we checked for None above
        let mut cur_val = self.value_it.as_mut().unwrap().next();

        // check to see if we've gone through everything in the set
        if cur_val.is_none() {
            let cur_entry = self.key_it.next(); // increment our key iterator

            if cur_entry.is_none() {
                self.cur_key = None; // set our invariant
                return None;
            }

            // safe to call unwrap because we checked it above
            let (cur_key, cur_set) = cur_entry.unwrap();

            self.cur_key = Some(cur_key); // set our key
            self.value_it = Some(cur_set.iter()); // set our value iterator
            cur_val = self.value_it.as_mut().unwrap().next(); // set our current value
        }

        return Some(KeyValuePair{key: self.cur_key.unwrap().clone(), value: cur_val.unwrap().clone()});
    }
}

#[cfg(test)]
mod tests {
    use multi_map::{MultiMap, MultiMapIterator};
    use wal_file::KeyValuePair;

    #[test]
    fn test_insert() {
        let mut mmap = MultiMap::<i32,String>::new();

        mmap.insert(12, String::from("abc"));
        mmap.insert(23, String::from("abc"));
        mmap.insert(23, String::from("def"));

        let mut it = mmap.into_iter();

        let e1 = it.next().unwrap();
        assert!(12 == e1.key);
        assert!(String::from("abc") == e1.value);

        let e2 = it.next().unwrap();
        assert!(23 == e2.key);
        assert!(String::from("abc") == e2.value);

        let e3 = it.next().unwrap();
        assert!(23 == e3.key);
        assert!(String::from("def") == e3.value);
    }
}

<commit_msg>Added check for return of insert<commit_after>use ::{KeyType, ValueType};

use wal_file::KeyValuePair;

use std::iter::{Peekable, empty, Map};
use std::collections::{BTreeMap, BTreeSet};
use std::collections::btree_map;
use std::collections::btree_set;

pub struct MultiMap<K: KeyType, V: ValueType> {
    multi_map: BTreeMap<K, BTreeSet<V>>,
    count: usize  // total number of KV pairs
}

pub struct MultiMapIterator<'a, K: KeyType + 'a, V: ValueType + 'a> {
    cur_key: Option<&'a K>,
    key_it: btree_map::Iter<'a,K,BTreeSet<V>>,
    value_it: Option<btree_set::Iter<'a,V>>,
}

impl <'a, K: KeyType, V: ValueType> MultiMap<K,V> {
    pub fn new() -> MultiMap<K,V> {
        return MultiMap{multi_map: BTreeMap::<K,BTreeSet<V>>::new(), count: 0};
    }

    pub fn insert(&mut self, key: K, value: V) -> usize {
        self.count += 1;

        if let Some(set) = self.multi_map.get_mut(&key) {
            set.insert(value);
            return self.count;
        }

        let mut set = BTreeSet::<V>::new();

        set.insert(value);

        self.multi_map.insert(key, set);

        return self.count;
    }
}

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

    fn into_iter(self) -> Self::IntoIter {
        let mut key_it = self.multi_map.iter();
        let cur_entry = key_it.next();

        // check to see if our map is empty
        if cur_entry.is_none() {
            return MultiMapIterator{cur_key: None, key_it: key_it, value_it: None};
        }

        // safe to call unwrap as we tested above
        let (cur_key, cur_set) = cur_entry.unwrap();

        return MultiMapIterator{cur_key: Some(cur_key), key_it: key_it, value_it: Some(cur_set.iter())};
    }
}

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

    fn next(&mut self) -> Option<Self::Item> {
        // this is our invariant, when it's None we've gone through everything
        if self.cur_key.is_none() {
            return None
        }

        // should be safe to call unwrap here, because we checked for None above
        let mut cur_val = self.value_it.as_mut().unwrap().next();

        // check to see if we've gone through everything in the set
        if cur_val.is_none() {
            let cur_entry = self.key_it.next(); // increment our key iterator

            if cur_entry.is_none() {
                self.cur_key = None; // set our invariant
                return None;
            }

            // safe to call unwrap because we checked it above
            let (cur_key, cur_set) = cur_entry.unwrap();

            self.cur_key = Some(cur_key); // set our key
            self.value_it = Some(cur_set.iter()); // set our value iterator
            cur_val = self.value_it.as_mut().unwrap().next(); // set our current value
        }

        return Some(KeyValuePair{key: self.cur_key.unwrap().clone(), value: cur_val.unwrap().clone()});
    }
}

#[cfg(test)]
mod tests {
    use multi_map::{MultiMap, MultiMapIterator};
    use wal_file::KeyValuePair;

    #[test]
    fn test_insert() {
        let mut mmap = MultiMap::<i32,String>::new();

        assert!(mmap.insert(12, String::from("abc")) == 1);
        assert!(mmap.insert(23, String::from("abc")) == 2);
        assert!(mmap.insert(23, String::from("def")) == 3);

        let mut it = mmap.into_iter();

        let e1 = it.next().unwrap();
        assert!(12 == e1.key);
        assert!(String::from("abc") == e1.value);

        let e2 = it.next().unwrap();
        assert!(23 == e2.key);
        assert!(String::from("abc") == e2.value);

        let e3 = it.next().unwrap();
        assert!(23 == e3.key);
        assert!(String::from("def") == e3.value);
    }
}

<|endoftext|>