<commit_before>//! Efficient container for locations annotated across a set of named
//! reference sequences.
use std::collections::HashMap;
use std::hash::Hash;

use crate::data_structures::interval_tree;
use crate::data_structures::interval_tree::{IntervalTree, IntervalTreeIterator};
use crate::utils::Interval;
use bio_types::annot::loc::Loc;

/// Efficient container for querying annotations, using `HashMap` and
/// `IntervalTree`.
///
/// The container is parameterized over the type of the reference
/// sequence names `R` (which is often a `String`) and the type of the
/// contained objects `T`.
///
/// The container finds annotations that overlap a specific query
/// location. Overlaps are identified without regard for strandedness
/// and without regard for e.g. spliced-out introns within the
/// annotation or the query.
///
/// Thus, the overlapping annotations identified by querying a
/// `AnnotMap` may need further filtering.
#[derive(Debug, Clone)]
pub struct AnnotMap<R, T>
where
    R: Hash + Eq,
{
    refid_itrees: HashMap<R, IntervalTree<isize, T>>,
}

impl<R, T> Default for AnnotMap<R, T>
where
    R: Eq + Hash,
{
    fn default() -> Self {
        AnnotMap {
            refid_itrees: HashMap::new(),
        }
    }
}

impl<R, T> AnnotMap<R, T>
where
    R: Eq + Hash,
{
    /// Creates a new, empty `AnnotMap`.
    ///
    /// ```
    /// extern crate bio;
    /// use bio::data_structures::annot_map::AnnotMap;
    /// let mut genes: AnnotMap<String, String> = AnnotMap::new();
    /// ```
    pub fn new() -> Self {
        Default::default()
    }

    /// Inserts an object into the container at a specified location.
    ///
    /// # Arguments
    ///
    /// `data` is the data item to be inserted into the annotation map
    ///
    /// `location` is the location associated with the data item.
    ///
    /// ```
    /// extern crate bio_types;
    /// extern crate bio;
    /// use bio::data_structures::annot_map::AnnotMap;
    /// use bio_types::annot::contig::Contig;
    /// use bio_types::strand::ReqStrand;
    /// let mut genes: AnnotMap<String, String> = AnnotMap::new();
    /// let tma22 = Contig::new(
    ///     "chrX".to_owned(),
    ///     461829,
    ///     462426 - 461829,
    ///     ReqStrand::Forward,
    /// );
    /// genes.insert_at("TMA22".to_owned(), &tma22);
    /// let tma19 = Contig::new(
    ///     "chrXI".to_owned(),
    ///     334412,
    ///     (334916 - 334412),
    ///     ReqStrand::Reverse,
    /// );
    /// genes.insert_at("TMA19".to_owned(), &tma19);
    /// ```
    pub fn insert_at<L>(&mut self, data: T, location: &L)
    where
        R: Eq + Hash + Clone,
        L: Loc<RefID = R>,
    {
        let itree = self
            .refid_itrees
            .entry(location.refid().clone())
            .or_insert_with(IntervalTree::new);
        let rng = location.start()..(location.start() + (location.length() as isize));
        itree.insert(rng, data);
    }

    /// Creates an `Iterator` that will visit all entries that overlap
    /// a query location.
    ///
    ///
    /// # Arguments
    ///
    /// `location` is the annotation location to be searched for
    /// overlapping data items.
    ///
    /// ```
    /// extern crate bio_types;
    /// extern crate bio;
    /// use bio::data_structures::annot_map::AnnotMap;
    /// use bio_types::annot::contig::Contig;
    /// use bio_types::strand::ReqStrand;
    /// let mut genes: AnnotMap<String, String> = AnnotMap::new();
    /// let tma22 = Contig::new(
    ///     "chrX".to_owned(),
    ///     461829,
    ///     462426 - 461829,
    ///     ReqStrand::Forward,
    /// );
    /// genes.insert_at("TMA22".to_owned(), &tma22);
    /// let tma19 = Contig::new(
    ///     "chrXI".to_owned(),
    ///     334412,
    ///     (334916 - 334412),
    ///     ReqStrand::Reverse,
    /// );
    /// genes.insert_at("TMA19".to_owned(), &tma19);
    /// let query = Contig::new("chrX".to_owned(), 462400, 100, ReqStrand::Forward);
    /// let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
    /// assert_eq!(hits, vec!["TMA22"]);
    /// let query = Contig::new("chrXI".to_owned(), 334400, 100, ReqStrand::Forward);
    /// let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
    /// assert_eq!(hits, vec!["TMA19"]);
    /// ```
    pub fn find<'a, L>(&'a self, location: &'a L) -> AnnotMapIterator<'a, R, T>
    where
        L: Loc<RefID = R>,
    {
        if let Some(itree) = self.refid_itrees.get(location.refid()) {
            let interval = location.start()..(location.start() + (location.length() as isize));
            let itree_iter = itree.find(interval);
            AnnotMapIterator {
                itree_iter: Some(itree_iter),
                refid: location.refid(),
            }
        } else {
            AnnotMapIterator {
                itree_iter: None,
                refid: location.refid(),
            }
        }
    }
}

impl<R, T> AnnotMap<R, T>
where
    R: Eq + Hash + Clone,
    T: Loc<RefID = R>,
{
    /// Inserts an object with the `Loc` trait into the container at
    /// its location.
    ///
    /// # Argument
    ///
    /// `data` is the data to be inserted based on its location.
    ///
    /// Equivalent to inserting `data` at `data.contig()`.
    pub fn insert_loc(&mut self, data: T) {
        let itree = self
            .refid_itrees
            .entry(data.refid().clone())
            .or_insert_with(IntervalTree::new);
        let rng = data.start()..(data.start() + (data.length() as isize));
        itree.insert(rng, data);
    }
}

/// A view of one annotation in a `AnnotMap` container.
#[derive(Debug, Clone)]
pub struct Entry<'a, R, T>
where
    R: Eq + Hash,
{
    itree_entry: interval_tree::Entry<'a, isize, T>,
    refid: &'a R,
}

impl<'a, R, T> Entry<'a, R, T>
where
    R: Eq + Hash,
{
    /// Returns a reference to the data value in the `AnnotMap`.
    pub fn data(&self) -> &'a T {
        self.itree_entry.data()
    }

    /// Returns a reference to the interval spanned by the annotation.
    pub fn interval(&self) -> &'a Interval<isize> {
        self.itree_entry.interval()
    }

    /// Returns a reference to the identifier of the annotated reference sequence.
    pub fn refid(&self) -> &'a R {
        self.refid
    }
}

/// An iterator over annotation entries (of type `Entry`) in a
/// `AnnotMap`.
///
/// This struct is created by the `find` function on `AnnotMap`.
pub struct AnnotMapIterator<'a, R, T>
where
    R: Eq + Hash,
{
    itree_iter: Option<IntervalTreeIterator<'a, isize, T>>,
    refid: &'a R,
}

impl<'a, R, T> Iterator for AnnotMapIterator<'a, R, T>
where
    R: 'a + Eq + Hash,
    T: 'a,
{
    type Item = Entry<'a, R, T>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.itree_iter {
            Some(ref mut iter) => match iter.next() {
                Some(next_itree) => Some(Entry {
                    itree_entry: next_itree,
                    refid: self.refid,
                }),
                None => None,
            },
            None => None,
        }
    }
}

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

    use bio_types::annot::contig::Contig;
    use bio_types::strand::ReqStrand;

    #[test]
    fn lookup() {
        let mut genes: AnnotMap<String, String> = AnnotMap::new();
        genes.insert_at(
            "TMA22".to_owned(),
            &Contig::new(
                "chrX".to_owned(),
                461829,
                462426 - 461829,
                ReqStrand::Forward,
            ),
        );
        genes.insert_at(
            "TMA19".to_owned(),
            &Contig::new(
                "chrXI".to_owned(),
                334412,
                334916 - 334412,
                ReqStrand::Reverse,
            ),
        );

        let query = Contig::new("chrX".to_owned(), 462400, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert_eq!(hits, vec!["TMA22"]);

        let query = Contig::new("chrXI".to_owned(), 334400, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert_eq!(hits, vec!["TMA19"]);

        let query = Contig::new("chrXI".to_owned(), 334916, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());

        let query = Contig::new("chrX".to_owned(), 461729, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());

        let query = Contig::new("chrXI".to_owned(), 462400, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());

        let query = Contig::new("NotFound".to_owned(), 0, 0, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());
    }

    #[test]
    fn overlaps() {
        let mut genes: AnnotMap<String, String> = AnnotMap::new();

        genes.insert_at(
            "a".to_owned(),
            &Contig::new("chr01".to_owned(), 1000, 1000, ReqStrand::Forward),
        );
        genes.insert_at(
            "b".to_owned(),
            &Contig::new("chr01".to_owned(), 1300, 1000, ReqStrand::Forward),
        );
        genes.insert_at(
            "c".to_owned(),
            &Contig::new("chr01".to_owned(), 1700, 1000, ReqStrand::Forward),
        );
        genes.insert_at(
            "d".to_owned(),
            &Contig::new("chr01".to_owned(), 2200, 1000, ReqStrand::Forward),
        );

        let query = Contig::new("chr01".to_owned(), 1050, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["a"]);

        let query = Contig::new("chr01".to_owned(), 1450, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["a", "b"]);

        let query = Contig::new("chr01".to_owned(), 1850, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["a", "b", "c"]);

        let query = Contig::new("chr01".to_owned(), 2250, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["b", "c", "d"]);

        let query = Contig::new("chr01".to_owned(), 2650, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["c", "d"]);
    }
}
<commit_msg>Improve docs: data_structures::annot_map (#314)<commit_after>//! Efficient container for locations annotated across a set of named
//! reference sequences.
//!
//! # Example
//!
//! ```
//! extern crate bio_types;
//! use bio::data_structures::annot_map::AnnotMap;
//! use bio_types::annot::contig::Contig;
//! use bio_types::strand::ReqStrand;
//!
//! // Insert a String annotation into the annotation map at a specified location.
//! let mut genes: AnnotMap<String, String> = AnnotMap::new();
//! let tma22 = Contig::new(
//!     "chrX".to_owned(),
//!     461829,
//!     462426 - 461829,
//!     ReqStrand::Forward,
//! );
//! genes.insert_at("TMA22".to_owned(), &tma22);
//!
//! // Find annotations that overlap a specific query
//! let query = Contig::new("chrX".to_owned(), 462400, 100, ReqStrand::Forward);
//! let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
//! assert_eq!(hits, vec!["TMA22"]);
//! ```

use std::collections::HashMap;
use std::hash::Hash;

use crate::data_structures::interval_tree;
use crate::data_structures::interval_tree::{IntervalTree, IntervalTreeIterator};
use crate::utils::Interval;
use bio_types::annot::loc::Loc;

/// Efficient container for querying annotations, using `HashMap` and
/// `IntervalTree`.
///
/// The container is parameterized over the type of the reference
/// sequence names `R` (which is often a `String`) and the type of the
/// contained objects `T`.
///
/// The container finds annotations that overlap a specific query
/// location. Overlaps are identified without regard for strandedness
/// and without regard for e.g. spliced-out introns within the
/// annotation or the query.
///
/// Thus, the overlapping annotations identified by querying a
/// `AnnotMap` may need further filtering.
#[derive(Debug, Clone)]
pub struct AnnotMap<R, T>
where
    R: Hash + Eq,
{
    refid_itrees: HashMap<R, IntervalTree<isize, T>>,
}

impl<R, T> Default for AnnotMap<R, T>
where
    R: Eq + Hash,
{
    fn default() -> Self {
        AnnotMap {
            refid_itrees: HashMap::new(),
        }
    }
}

impl<R, T> AnnotMap<R, T>
where
    R: Eq + Hash,
{
    /// Create a new, empty `AnnotMap`. Used in conjunction with `insert_at`
    /// or `insert_loc`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Insert an object into the container at a specified location (`Loc`).
    ///
    /// # Arguments
    ///
    /// * `data` - any type of data to be inserted at the location / region
    /// * `location` - any object with the `Loc` trait implemented, determining
    ///   the Range at which to insert the `data`
    ///
    /// # Example
    ///
    /// ```
    /// extern crate bio_types;
    /// use bio::data_structures::annot_map::AnnotMap;
    /// use bio_types::annot::contig::Contig;
    /// use bio_types::strand::ReqStrand;
    ///
    /// let mut genes: AnnotMap<String, String> = AnnotMap::new();
    /// let tma22 = Contig::new(
    ///     "chrX".to_owned(),
    ///     461829,
    ///     462426 - 461829,
    ///     ReqStrand::Forward,
    /// );
    /// genes.insert_at("TMA22".to_owned(), &tma22);
    /// ```
    pub fn insert_at<L>(&mut self, data: T, location: &L)
    where
        R: Eq + Hash + Clone,
        L: Loc<RefID = R>,
    {
        let itree = self
            .refid_itrees
            .entry(location.refid().clone())
            .or_insert_with(IntervalTree::new);
        let rng = location.start()..(location.start() + (location.length() as isize));
        itree.insert(rng, data);
    }

    /// Create an `Iterator` that will visit all entries that overlap
    /// a query location.
    pub fn find<'a, L>(&'a self, location: &'a L) -> AnnotMapIterator<'a, R, T>
    where
        L: Loc<RefID = R>,
    {
        if let Some(itree) = self.refid_itrees.get(location.refid()) {
            let interval = location.start()..(location.start() + (location.length() as isize));
            let itree_iter = itree.find(interval);
            AnnotMapIterator {
                itree_iter: Some(itree_iter),
                refid: location.refid(),
            }
        } else {
            AnnotMapIterator {
                itree_iter: None,
                refid: location.refid(),
            }
        }
    }
}

impl<R, T> AnnotMap<R, T>
where
    R: Eq + Hash + Clone,
    T: Loc<RefID = R>,
{
    /// Insert an object with the `Loc` trait into the container at
    /// its location.
    ///
    /// This inserts all of `data` at the Range of length `data.length()`
    /// that starts at `data.start()`.
    ///
    /// # Example
    ///
    /// ```
    /// extern crate bio_types;
    /// use bio::data_structures::annot_map::AnnotMap;
    /// use bio_types::annot::contig::Contig;
    /// use bio_types::strand::ReqStrand;
    ///
    /// let mut gene_locs = AnnotMap::new();
    /// let tma19 = Contig::new(
    ///     String::from("chrXI"),
    ///     334412,
    ///     (334916 - 334412),
    ///     ReqStrand::Reverse,
    /// );
    /// let assert_copy = tma19.clone();
    /// gene_locs.insert_loc(tma19);
    /// // Find annotations that overlap a specific query
    /// let query = Contig::new(String::from("chrXI"), 334400, 100, ReqStrand::Reverse);
    /// let hits: Vec<&Contig<String, ReqStrand>> = gene_locs.find(&query).map(|e| e.data()).collect();
    /// assert_eq!(hits, vec![&assert_copy]);
    /// ```
    pub fn insert_loc(&mut self, data: T) {
        let itree = self
            .refid_itrees
            .entry(data.refid().clone())
            .or_insert_with(IntervalTree::new);
        let rng = data.start()..(data.start() + (data.length() as isize));
        itree.insert(rng, data);
    }
}

/// A view of one annotation in a `AnnotMap` container.
#[derive(Debug, Clone)]
pub struct Entry<'a, R, T>
where
    R: Eq + Hash,
{
    itree_entry: interval_tree::Entry<'a, isize, T>,
    refid: &'a R,
}

impl<'a, R, T> Entry<'a, R, T>
where
    R: Eq + Hash,
{
    /// Return a reference to the data value in the `AnnotMap`.
    pub fn data(&self) -> &'a T {
        self.itree_entry.data()
    }

    /// Return a reference to the interval spanned by the annotation.
    pub fn interval(&self) -> &'a Interval<isize> {
        self.itree_entry.interval()
    }

    /// Return a reference to the identifier of the annotated reference sequence.
    pub fn refid(&self) -> &'a R {
        self.refid
    }
}

/// An iterator over annotation entries (of type `Entry`) in a
/// `AnnotMap`.
///
/// This struct is created by the `find` function on `AnnotMap`.
pub struct AnnotMapIterator<'a, R, T>
where
    R: Eq + Hash,
{
    itree_iter: Option<IntervalTreeIterator<'a, isize, T>>,
    refid: &'a R,
}

impl<'a, R, T> Iterator for AnnotMapIterator<'a, R, T>
where
    R: 'a + Eq + Hash,
    T: 'a,
{
    type Item = Entry<'a, R, T>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.itree_iter {
            Some(ref mut iter) => match iter.next() {
                Some(next_itree) => Some(Entry {
                    itree_entry: next_itree,
                    refid: self.refid,
                }),
                None => None,
            },
            None => None,
        }
    }
}

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

    use bio_types::annot::contig::Contig;
    use bio_types::strand::ReqStrand;

    #[test]
    fn lookup() {
        let mut genes: AnnotMap<String, String> = AnnotMap::new();
        genes.insert_at(
            "TMA22".to_owned(),
            &Contig::new(
                "chrX".to_owned(),
                461829,
                462426 - 461829,
                ReqStrand::Forward,
            ),
        );
        genes.insert_at(
            "TMA19".to_owned(),
            &Contig::new(
                "chrXI".to_owned(),
                334412,
                334916 - 334412,
                ReqStrand::Reverse,
            ),
        );

        let query = Contig::new("chrX".to_owned(), 462400, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert_eq!(hits, vec!["TMA22"]);

        let query = Contig::new("chrXI".to_owned(), 334400, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert_eq!(hits, vec!["TMA19"]);

        let query = Contig::new("chrXI".to_owned(), 334916, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());

        let query = Contig::new("chrX".to_owned(), 461729, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());

        let query = Contig::new("chrXI".to_owned(), 462400, 100, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());

        let query = Contig::new("NotFound".to_owned(), 0, 0, ReqStrand::Forward);
        let hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        assert!(hits.is_empty());
    }

    #[test]
    fn overlaps() {
        let mut genes: AnnotMap<String, String> = AnnotMap::new();

        genes.insert_at(
            "a".to_owned(),
            &Contig::new("chr01".to_owned(), 1000, 1000, ReqStrand::Forward),
        );
        genes.insert_at(
            "b".to_owned(),
            &Contig::new("chr01".to_owned(), 1300, 1000, ReqStrand::Forward),
        );
        genes.insert_at(
            "c".to_owned(),
            &Contig::new("chr01".to_owned(), 1700, 1000, ReqStrand::Forward),
        );
        genes.insert_at(
            "d".to_owned(),
            &Contig::new("chr01".to_owned(), 2200, 1000, ReqStrand::Forward),
        );

        let query = Contig::new("chr01".to_owned(), 1050, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["a"]);

        let query = Contig::new("chr01".to_owned(), 1450, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["a", "b"]);

        let query = Contig::new("chr01".to_owned(), 1850, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["a", "b", "c"]);

        let query = Contig::new("chr01".to_owned(), 2250, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["b", "c", "d"]);

        let query = Contig::new("chr01".to_owned(), 2650, 100, ReqStrand::Forward);
        let mut hits: Vec<&String> = genes.find(&query).map(|e| e.data()).collect();
        hits.sort();
        assert_eq!(hits, vec!["c", "d"]);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Carl-Erwin Griffith
//
// Permission is hereby granted, free of charge, to any
// person obtaining a copy of this software and associated
// documentation files (the "Software"), to deal in the
// Software without restriction, including without
// limitation the rights to use, copy, modify, merge,
// publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software
// is furnished to do so, subject to the following
// conditions:
//
// The above copyright notice and this permission notice
// shall be included in all copies or substantial portions
// of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER

pub mod line;
use crate::core::codepointinfo::CodepointInfo;
use std::time::Duration;

use self::line::Line;
use self::line::LineCellIndex;

pub type LineIndex = usize;

/// A Screen is composed of Line(s).

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Screen {
    /// the underlying lines storage
    pub line: Vec<Line>,
    /// the index of the line filled with the push method
    pub current_line_index: LineCellIndex,
    /// maximum number of elements the screen line can hold
    pub width: usize,
    /// maximum number of lines the screen can hold
    pub height: usize,
    /// the number of elements pushed in the screen
    pub nb_push: usize,
    /// placeholder to record the offset of the first pushed CodepointInfo (used by View)
    pub first_offset: u64,
    /// placeholder to record the offset of the last pushed CodepointInfo (used by View)
    pub last_offset: u64,
    /// placeholder to record the maximum offset of the document (eof)
    pub doc_max_offset: u64,
    /// time spent to generate the screen content
    pub time_to_build: Duration,

    pub input_size: usize,
}

impl Screen {
    pub fn new(width: usize, height: usize) -> Screen {
        let mut line: Vec<Line> = Vec::new();
        for _ in 0..height {
            line.push(Line::new(width));
        }

        Screen {
            line,
            current_line_index: 0,
            width,
            height,
            nb_push: 0,
            first_offset: 0,
            last_offset: 0,
            doc_max_offset: 0,
            time_to_build: Duration::new(0, 0),
            input_size: 0,
        }
    }

    pub fn resize(&mut self, width: usize, height: usize) {
        self.line.resize(height, Line::new(width));
        for i in 0..height {
            self.line[i].resize(width);
        }
        self.width = width;
        self.height = height;
        self.current_line_index = 0;
        self.nb_push = 0;
        self.first_offset = 0;
        self.last_offset = 0;
        self.doc_max_offset = 0;
        self.input_size = 0;
    }

    /// append
    pub fn push(&mut self, cpi: CodepointInfo) -> (bool, usize) {
        if self.current_line_index == self.height {
            return (false, self.current_line_index);
        }

        if self.line[self.current_line_index].read_only {
            self.current_line_index += 1;
        }

        if self.current_line_index == self.height {
            return (false, self.current_line_index);
        }

        let cp = cpi.cp;
        let line = &mut self.line[self.current_line_index];
        let (ok, _) = line.push(cpi);

        if ok {
            self.nb_push += 1;
            if cp == '\n' || cp == '\r' {
                line.read_only = true;
            }
        }
        (ok, self.current_line_index)
    }

    pub fn clear(&mut self) {
        for h in 0..self.height {
            self.line[h].clear();
        }
        self.current_line_index = 0;
        self.nb_push = 0;
        self.first_offset = 0;
        self.last_offset = 0;
        self.doc_max_offset = 0;
        self.input_size = 0;
    }

    pub fn get_mut_line(&mut self, index: usize) -> Option<&mut Line> {
        if index < self.height {
            Some(&mut self.line[index])
        } else {
            None
        }
    }

    pub fn get_line(&self, index: usize) -> Option<&Line> {
        if index < self.height {
            Some(&self.line[index])
        } else {
            None
        }
    }

    pub fn get_mut_used_line(&mut self, index: usize) -> Option<&mut Line> {
        if index < self.current_line_index {
            Some(&mut self.line[index])
        } else {
            None
        }
    }

    pub fn get_used_line_clipped(&mut self, index: usize) -> (Option<&Line>, LineIndex) {
        let index = ::std::cmp::min(index, self.current_line_index);
        (self.get_used_line(index), index)
    }

    pub fn get_used_line(&self, index: usize) -> Option<&Line> {
        if index < self.current_line_index {
            Some(&self.line[index])
        } else {
            None
        }
    }

    /// there must be 2 line a least
    pub fn get_first_used_line(&self) -> Option<&Line> {
        if 0 < self.current_line_index {
            Some(&self.line[0])
        } else {
            None
        }
    }

    /// there must be 2 line a least
    pub fn get_last_used_line(&self) -> Option<&Line> {
        if self.current_line_index > 0 {
            Some(&self.line[self.current_line_index - 1])
        } else {
            None
        }
    }

    pub fn get_cpinfo(&self, x: usize, y: usize) -> Option<&CodepointInfo> {
        match self.get_line(y) {
            None => None,
            Some(l) => l.get_cpi(x),
        }
    }

    pub fn get_mut_cpinfo(&mut self, x: usize, y: usize) -> Option<&mut CodepointInfo> {
        match self.get_mut_line(y) {
            None => None,
            Some(l) => l.get_mut_cpi(x),
        }
    }

    pub fn get_used_cpinfo(&self, x: usize, y: usize) -> Option<&CodepointInfo> {
        match self.get_used_line(y) {
            None => None,
            Some(l) => l.get_used_cpi(x),
        }
    }

    pub fn get_mut_used_cpinfo(&mut self, x: usize, y: usize) -> Option<&mut CodepointInfo> {
        match self.get_mut_used_line(y) {
            None => None,
            Some(l) => l.get_mut_used_cpi(x),
        }
    }

    pub fn get_mut_first_cpinfo(&mut self) -> (Option<&mut CodepointInfo>, usize, usize) {
        match self.get_mut_used_line(0) {
            None => (None, 0, 0),
            Some(l) => (l.get_mut_used_cpi(0), 0, 0),
        }
    }

    pub fn get_mut_last_cpinfo(&mut self) -> (Option<&mut CodepointInfo>, usize, usize) {
        let y = self.current_line_index;
        match self.get_mut_used_line(y) {
            None => (None, 0, 0),
            Some(l) => {
                let x = l.nb_cells;
                (l.get_mut_used_cpi(0), x, y)
            }
        }
    }

    pub fn get_first_cpinfo(&self) -> (Option<&CodepointInfo>, usize, usize) {
        match self.get_used_line(0) {
            None => (None, 0, 0),
            Some(l) => (l.get_used_cpi(0), 0, 0),
        }
    }

    pub fn get_last_cpinfo(&self) -> (Option<&CodepointInfo>, usize, usize) {
        let y = self.current_line_index;
        match self.get_used_line(y) {
            None => (None, 0, 0),
            Some(l) => {
                if l.nb_cells > 0 {
                    let x = l.nb_cells - 1;
                    (l.get_used_cpi(0), x, y)
                } else {
                    (None, 0, 0)
                }
            }
        }
    }

    pub fn get_used_cpinfo_clipped(
        &mut self,
        x: usize,
        y: usize,
    ) -> (Option<&CodepointInfo>, LineCellIndex, LineIndex) {
        match self.get_used_line_clipped(y) {
            (None, li) => (None, x, li),
            (Some(l), li) => match l.get_used_cpi_clipped(x) {
                (optcpi, lci) => (optcpi, lci, li),
            },
        }
    }

    pub fn find_cpi_by_offset(&self, offset: u64) -> (Option<&CodepointInfo>, usize, usize) {
        // TODO: use dichotomic search

        if offset < self.first_offset || offset > self.last_offset {
            return (None, 0, 0);
        }

        for y in 0..self.height {
            let l = self.get_line(y).unwrap();
            if l.nb_cells == 0 {
                continue;
            }

            for x in 0..l.width {
                let cpi = l.get_cpi(x).unwrap();
                if cpi.offset == offset {
                    return (Some(cpi), x, y);
                }
            }
        }
        (None, 0, 0)
    }

    pub fn contains_offset(&self, offset: u64) -> bool {
        if offset < self.first_offset || offset > self.last_offset {
            return false;
        }

        let (cpi, _, _) = self.find_cpi_by_offset(offset);
        match cpi {
            Some(_) => true,
            _ => false,
        }
    }
}

#[test]
fn test_screen() {
    let mut scr = Screen::new(640, 480);
    assert_eq!(640, scr.width);
    assert_eq!(480, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());

    scr.resize(800, 600);
    assert_eq!(800, scr.width);
    assert_eq!(600, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());

    scr.resize(1024, 768);
    assert_eq!(1024, scr.width);
    assert_eq!(768, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());

    scr.resize(640, 480);
    assert_eq!(640, scr.width);
    assert_eq!(480, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());
}
<commit_msg>wrong name used (even if the aliased type is the same)<commit_after>// Copyright (c) Carl-Erwin Griffith
//
// Permission is hereby granted, free of charge, to any
// person obtaining a copy of this software and associated
// documentation files (the "Software"), to deal in the
// Software without restriction, including without
// limitation the rights to use, copy, modify, merge,
// publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software
// is furnished to do so, subject to the following
// conditions:
//
// The above copyright notice and this permission notice
// shall be included in all copies or substantial portions
// of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER

pub mod line;
use crate::core::codepointinfo::CodepointInfo;
use std::time::Duration;

use self::line::Line;
use self::line::LineCellIndex;

pub type LineIndex = usize;

/// A Screen is composed of Line(s).

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Screen {
    /// the underlying lines storage
    pub line: Vec<Line>,
    /// the index of the line filled with the push method
    pub current_line_index: LineIndex,
    /// maximum number of elements the screen line can hold
    pub width: usize,
    /// maximum number of lines the screen can hold
    pub height: usize,
    /// the number of elements pushed in the screen
    pub nb_push: usize,
    /// placeholder to record the offset of the first pushed CodepointInfo (used by View)
    pub first_offset: u64,
    /// placeholder to record the offset of the last pushed CodepointInfo (used by View)
    pub last_offset: u64,
    /// placeholder to record the maximum offset of the document (eof)
    pub doc_max_offset: u64,
    /// time spent to generate the screen content
    pub time_to_build: Duration,

    pub input_size: usize,
}

impl Screen {
    pub fn new(width: usize, height: usize) -> Screen {
        let mut line: Vec<Line> = Vec::new();
        for _ in 0..height {
            line.push(Line::new(width));
        }

        Screen {
            line,
            current_line_index: 0,
            width,
            height,
            nb_push: 0,
            first_offset: 0,
            last_offset: 0,
            doc_max_offset: 0,
            time_to_build: Duration::new(0, 0),
            input_size: 0,
        }
    }

    pub fn resize(&mut self, width: usize, height: usize) {
        self.line.resize(height, Line::new(width));
        for i in 0..height {
            self.line[i].resize(width);
        }
        self.width = width;
        self.height = height;
        self.current_line_index = 0;
        self.nb_push = 0;
        self.first_offset = 0;
        self.last_offset = 0;
        self.doc_max_offset = 0;
        self.input_size = 0;
    }

    /// append
    pub fn push(&mut self, cpi: CodepointInfo) -> (bool, usize) {
        if self.current_line_index == self.height {
            return (false, self.current_line_index);
        }

        if self.line[self.current_line_index].read_only {
            self.current_line_index += 1;
        }

        if self.current_line_index == self.height {
            return (false, self.current_line_index);
        }

        let cp = cpi.cp;
        let line = &mut self.line[self.current_line_index];
        let (ok, _) = line.push(cpi);

        if ok {
            self.nb_push += 1;
            if cp == '\n' || cp == '\r' {
                line.read_only = true;
            }
        }
        (ok, self.current_line_index)
    }

    pub fn clear(&mut self) {
        for h in 0..self.height {
            self.line[h].clear();
        }
        self.current_line_index = 0;
        self.nb_push = 0;
        self.first_offset = 0;
        self.last_offset = 0;
        self.doc_max_offset = 0;
        self.input_size = 0;
    }

    pub fn get_mut_line(&mut self, index: usize) -> Option<&mut Line> {
        if index < self.height {
            Some(&mut self.line[index])
        } else {
            None
        }
    }

    pub fn get_line(&self, index: usize) -> Option<&Line> {
        if index < self.height {
            Some(&self.line[index])
        } else {
            None
        }
    }

    pub fn get_mut_used_line(&mut self, index: usize) -> Option<&mut Line> {
        if index < self.current_line_index {
            Some(&mut self.line[index])
        } else {
            None
        }
    }

    pub fn get_used_line_clipped(&mut self, index: usize) -> (Option<&Line>, LineIndex) {
        let index = ::std::cmp::min(index, self.current_line_index);
        (self.get_used_line(index), index)
    }

    pub fn get_used_line(&self, index: usize) -> Option<&Line> {
        if index < self.current_line_index {
            Some(&self.line[index])
        } else {
            None
        }
    }

    /// there must be 2 line a least
    pub fn get_first_used_line(&self) -> Option<&Line> {
        if 0 < self.current_line_index {
            Some(&self.line[0])
        } else {
            None
        }
    }

    /// there must be 2 line a least
    pub fn get_last_used_line(&self) -> Option<&Line> {
        if self.current_line_index > 0 {
            Some(&self.line[self.current_line_index - 1])
        } else {
            None
        }
    }

    pub fn get_cpinfo(&self, x: usize, y: usize) -> Option<&CodepointInfo> {
        match self.get_line(y) {
            None => None,
            Some(l) => l.get_cpi(x),
        }
    }

    pub fn get_mut_cpinfo(&mut self, x: usize, y: usize) -> Option<&mut CodepointInfo> {
        match self.get_mut_line(y) {
            None => None,
            Some(l) => l.get_mut_cpi(x),
        }
    }

    pub fn get_used_cpinfo(&self, x: usize, y: usize) -> Option<&CodepointInfo> {
        match self.get_used_line(y) {
            None => None,
            Some(l) => l.get_used_cpi(x),
        }
    }

    pub fn get_mut_used_cpinfo(&mut self, x: usize, y: usize) -> Option<&mut CodepointInfo> {
        match self.get_mut_used_line(y) {
            None => None,
            Some(l) => l.get_mut_used_cpi(x),
        }
    }

    pub fn get_mut_first_cpinfo(&mut self) -> (Option<&mut CodepointInfo>, usize, usize) {
        match self.get_mut_used_line(0) {
            None => (None, 0, 0),
            Some(l) => (l.get_mut_used_cpi(0), 0, 0),
        }
    }

    pub fn get_mut_last_cpinfo(&mut self) -> (Option<&mut CodepointInfo>, usize, usize) {
        let y = self.current_line_index;
        match self.get_mut_used_line(y) {
            None => (None, 0, 0),
            Some(l) => {
                let x = l.nb_cells;
                (l.get_mut_used_cpi(0), x, y)
            }
        }
    }

    pub fn get_first_cpinfo(&self) -> (Option<&CodepointInfo>, usize, usize) {
        match self.get_used_line(0) {
            None => (None, 0, 0),
            Some(l) => (l.get_used_cpi(0), 0, 0),
        }
    }

    pub fn get_last_cpinfo(&self) -> (Option<&CodepointInfo>, usize, usize) {
        let y = self.current_line_index;
        match self.get_used_line(y) {
            None => (None, 0, 0),
            Some(l) => {
                if l.nb_cells > 0 {
                    let x = l.nb_cells - 1;
                    (l.get_used_cpi(0), x, y)
                } else {
                    (None, 0, 0)
                }
            }
        }
    }

    pub fn get_used_cpinfo_clipped(
        &mut self,
        x: usize,
        y: usize,
    ) -> (Option<&CodepointInfo>, LineCellIndex, LineIndex) {
        match self.get_used_line_clipped(y) {
            (None, li) => (None, x, li),
            (Some(l), li) => match l.get_used_cpi_clipped(x) {
                (optcpi, lci) => (optcpi, lci, li),
            },
        }
    }

    pub fn find_cpi_by_offset(&self, offset: u64) -> (Option<&CodepointInfo>, usize, usize) {
        // TODO: use dichotomic search

        if offset < self.first_offset || offset > self.last_offset {
            return (None, 0, 0);
        }

        for y in 0..self.height {
            let l = self.get_line(y).unwrap();
            if l.nb_cells == 0 {
                continue;
            }

            for x in 0..l.width {
                let cpi = l.get_cpi(x).unwrap();
                if cpi.offset == offset {
                    return (Some(cpi), x, y);
                }
            }
        }
        (None, 0, 0)
    }

    pub fn contains_offset(&self, offset: u64) -> bool {
        if offset < self.first_offset || offset > self.last_offset {
            return false;
        }

        let (cpi, _, _) = self.find_cpi_by_offset(offset);
        match cpi {
            Some(_) => true,
            _ => false,
        }
    }
}

#[test]
fn test_screen() {
    let mut scr = Screen::new(640, 480);
    assert_eq!(640, scr.width);
    assert_eq!(480, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());

    scr.resize(800, 600);
    assert_eq!(800, scr.width);
    assert_eq!(600, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());

    scr.resize(1024, 768);
    assert_eq!(1024, scr.width);
    assert_eq!(768, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());

    scr.resize(640, 480);
    assert_eq!(640, scr.width);
    assert_eq!(480, scr.height);
    assert_eq!(scr.height, scr.line.len());
    assert_eq!(scr.width, scr.line[0].cells.len());
}
<|endoftext|>


<commit_before>use hir::HirDisplay;
use syntax::{ast, AstNode, SyntaxToken, TextSize};
use test_utils::mark;

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: infer_function_return_type
//
// Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return
// type specified.
//
// ```
// fn foo() { 4<|>2i32 }
// ```
// ->
// ```
// fn foo() -> i32 { 42i32 }
// ```
pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let expr = ctx.find_node_at_offset::<ast::Expr>()?;
    let (tail_expr, insert_pos) = extract_tail(expr)?;
    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
    let ty = ctx.sema.type_of_expr(&tail_expr).filter(|ty| !ty.is_unit())?;
    let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;

    acc.add(
        AssistId("change_return_type_to_result", AssistKind::RefactorRewrite),
        "Wrap return type in Result",
        tail_expr.syntax().text_range(),
        |builder| {
            let insert_pos = insert_pos.text_range().end() + TextSize::from(1);
            builder.insert(insert_pos, &format!("-> {} ", ty));
        },
    )
}

fn extract_tail(expr: ast::Expr) -> Option<(ast::Expr, SyntaxToken)> {
    let (ret_ty, tail_expr, insert_pos) =
        if let Some(closure) = expr.syntax().ancestors().find_map(ast::ClosureExpr::cast) {
            let tail_expr = match closure.body()? {
                ast::Expr::BlockExpr(block) => block.expr()?,
                body => body,
            };
            let ret_ty = closure.ret_type();
            let rpipe = closure.param_list()?.syntax().last_token()?;
            (ret_ty, tail_expr, rpipe)
        } else {
            let func = expr.syntax().ancestors().find_map(ast::Fn::cast)?;
            let tail_expr = func.body()?.expr()?;
            let ret_ty = func.ret_type();
            let rparen = func.param_list()?.r_paren_token()?;
            (ret_ty, tail_expr, rparen)
        };
    if ret_ty.is_some() {
        mark::hit!(existing_ret_type);
        mark::hit!(existing_ret_type_closure);
        return None;
    }
    // check whether the expr we were at is indeed the tail expression
    if !tail_expr.syntax().text_range().contains_range(expr.syntax().text_range()) {
        mark::hit!(not_tail_expr);
        return None;
    }
    Some((tail_expr, insert_pos))
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn infer_return_type() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                45<|>
            }"#,
            r#"fn foo() -> i32 {
                45
            }"#,
        );
    }

    #[test]
    fn infer_return_type_nested() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                if true {
                    3<|>
                } else {
                    5
                }
            }"#,
            r#"fn foo() -> i32 {
                if true {
                    3
                } else {
                    5
                }
            }"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified() {
        mark::check!(existing_ret_type);
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() -> i32 {
                ( 45<|> + 32 ) * 123
            }"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr() {
        mark::check!(not_tail_expr);
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                let x = <|>3;
                ( 45 + 32 ) * 123
            }"#,
        );
    }

    #[test]
    fn not_applicable_unit_return_type() {
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                (<|>)
            }"#,
        );
    }

    #[test]
    fn infer_return_type_closure_block() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                |x: i32| {
                    x<|>
                };
            }"#,
            r#"fn foo() {
                |x: i32| -> i32 {
                    x
                };
            }"#,
        );
    }

    #[test]
    fn infer_return_type_closure() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                |x: i32| x<|>;
            }"#,
            r#"fn foo() {
                |x: i32| -> i32 x;
            }"#,
        );
    }

    #[test]
    fn infer_return_type_nested_closure() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                || {
                    if true {
                        3<|>
                    } else {
                        5
                    }
                }
            }"#,
            r#"fn foo() {
                || -> i32 {
                    if true {
                        3
                    } else {
                        5
                    }
                }
            }"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified_closure() {
        mark::check!(existing_ret_type_closure);
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                || -> i32 { 3<|> }
            }"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr_closure() {
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                || -> i32 {
                    let x = 3<|>;
                    6
                }
            }"#,
        );
    }
}
<commit_msg>Wrap non-block expressions in closures with a block<commit_after>use hir::HirDisplay;
use syntax::{ast, AstNode, SyntaxToken, TextSize};
use test_utils::mark;

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: infer_function_return_type
//
// Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return
// type specified.
//
// ```
// fn foo() { 4<|>2i32 }
// ```
// ->
// ```
// fn foo() -> i32 { 42i32 }
// ```
pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let expr = ctx.find_node_at_offset::<ast::Expr>()?;
    let (tail_expr, insert_pos, wrap_expr) = extract_tail(expr)?;
    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
    let ty = ctx.sema.type_of_expr(&tail_expr).filter(|ty| !ty.is_unit())?;
    let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;

    acc.add(
        AssistId("infer_function_return_type", AssistKind::RefactorRewrite),
        "Add this function's return type",
        tail_expr.syntax().text_range(),
        |builder| {
            let insert_pos = insert_pos.text_range().end() + TextSize::from(1);
            builder.insert(insert_pos, &format!("-> {} ", ty));
            if wrap_expr {
                mark::hit!(wrap_closure_non_block_expr);
                // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
                builder.replace(tail_expr.syntax().text_range(), &format!("{{{}}}", tail_expr));
            }
        },
    )
}

fn extract_tail(expr: ast::Expr) -> Option<(ast::Expr, SyntaxToken, bool)> {
    let (ret_ty, tail_expr, insert_pos, wrap_expr) =
        if let Some(closure) = expr.syntax().ancestors().find_map(ast::ClosureExpr::cast) {
            let (tail_expr, wrap_expr) = match closure.body()? {
                ast::Expr::BlockExpr(block) => (block.expr()?, false),
                body => (body, true),
            };
            let ret_ty = closure.ret_type();
            let rpipe = closure.param_list()?.syntax().last_token()?;
            (ret_ty, tail_expr, rpipe, wrap_expr)
        } else {
            let func = expr.syntax().ancestors().find_map(ast::Fn::cast)?;
            let tail_expr = func.body()?.expr()?;
            let ret_ty = func.ret_type();
            let rparen = func.param_list()?.r_paren_token()?;
            (ret_ty, tail_expr, rparen, false)
        };
    if ret_ty.is_some() {
        mark::hit!(existing_ret_type);
        mark::hit!(existing_ret_type_closure);
        return None;
    }
    // check whether the expr we were at is indeed the tail expression
    if !tail_expr.syntax().text_range().contains_range(expr.syntax().text_range()) {
        mark::hit!(not_tail_expr);
        return None;
    }
    Some((tail_expr, insert_pos, wrap_expr))
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn infer_return_type() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                45<|>
            }"#,
            r#"fn foo() -> i32 {
                45
            }"#,
        );
    }

    #[test]
    fn infer_return_type_nested() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                if true {
                    3<|>
                } else {
                    5
                }
            }"#,
            r#"fn foo() -> i32 {
                if true {
                    3
                } else {
                    5
                }
            }"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified() {
        mark::check!(existing_ret_type);
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() -> i32 {
                ( 45<|> + 32 ) * 123
            }"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr() {
        mark::check!(not_tail_expr);
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                let x = <|>3;
                ( 45 + 32 ) * 123
            }"#,
        );
    }

    #[test]
    fn not_applicable_unit_return_type() {
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                (<|>)
            }"#,
        );
    }

    #[test]
    fn infer_return_type_closure_block() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                |x: i32| {
                    x<|>
                };
            }"#,
            r#"fn foo() {
                |x: i32| -> i32 {
                    x
                };
            }"#,
        );
    }

    #[test]
    fn infer_return_type_closure() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                |x: i32| { x<|> };
            }"#,
            r#"fn foo() {
                |x: i32| -> i32 { x };
            }"#,
        );
    }

    #[test]
    fn infer_return_type_closure_wrap() {
        mark::check!(wrap_closure_non_block_expr);
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                |x: i32| x<|>;
            }"#,
            r#"fn foo() {
                |x: i32| -> i32 {x};
            }"#,
        );
    }

    #[test]
    fn infer_return_type_nested_closure() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                || {
                    if true {
                        3<|>
                    } else {
                        5
                    }
                }
            }"#,
            r#"fn foo() {
                || -> i32 {
                    if true {
                        3
                    } else {
                        5
                    }
                }
            }"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified_closure() {
        mark::check!(existing_ret_type_closure);
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                || -> i32 { 3<|> }
            }"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr_closure() {
        check_assist_not_applicable(
            infer_function_return_type,
            r#"fn foo() {
                || -> i32 {
                    let x = 3<|>;
                    6
                }
            }"#,
        );
    }
}
<|endoftext|>


<commit_before>mod page;
pub use self::page::Page;

use std::collections::BTreeMap;
use std::default::Default;
use std::env;
use std::str::FromStr;
use std::time::Duration;

use html5ever::tokenizer::BufferQueue;
use html5ever::tokenizer::Tokenizer;
use reqwest::header::{CONTENT_TYPE, USER_AGENT};
use reqwest::Client as HttpClient;
use reqwest::Response as HttpResponse;
use reqwest::StatusCode;
use robotparser::RobotFileParser;
use sidekiq::Client as SidekiqClient;
use sidekiq::ClientOpts as SidekiqClientOpts;
use sidekiq::RedisPool;
use url::Url;

const MAMAN_ENV: &str = "MAMAN_ENV";
const MAMAN_ENV_DEFAULT: &str = "development";

pub struct Spider<'a> {
    pub base_url: Url,
    pub visited_urls: Vec<Url>,
    pub unvisited_urls: Vec<Url>,
    pub env: String,
    pub limit: isize,
    pub mime_types: Vec<mime::Mime>,
    sidekiq: SidekiqClient,
    robot_parser: RobotFileParser<'a>,
}

impl<'a> Spider<'a> {
    pub fn new(
        redis_pool: RedisPool,
        base_url: Url,
        limit: isize,
        mime_types: Vec<mime::Mime>,
    ) -> Spider<'a> {
        let maman_env =
            env::var(&MAMAN_ENV.to_owned()).unwrap_or_else(|_| MAMAN_ENV_DEFAULT.to_owned());
        let robots_txt = base_url.join("/robots.txt").unwrap();
        let robot_file_parser = RobotFileParser::new(robots_txt);
        let client_opts = SidekiqClientOpts {
            namespace: Some(maman_env.to_string()),
        };
        let sidekiq = SidekiqClient::new(redis_pool, client_opts);
        Spider {
            base_url,
            visited_urls: Vec::new(),
            unvisited_urls: Vec::new(),
            sidekiq,
            env: maman_env,
            robot_parser: robot_file_parser,
            limit,
            mime_types,
        }
    }

    pub fn visit_page(&mut self, page: Page) {
        self.visited_urls.push(page.url.clone());
        for u in &page.urls {
            self.unvisited_urls.push(u.clone());
        }
        if let Err(err) = self.sidekiq.push(page.to_job()) {
            error!("SidekiqClient push failed: {}", err);
        }
    }

    pub fn crawl(&mut self) {
        self.robot_parser.read();
        let base_url = self.base_url.clone();
        if let Some(response) = Spider::load_url(self.base_url.as_ref(), &self.mime_types) {
            self.visit(&base_url, response);
            while let Some(url) = self.unvisited_urls.pop() {
                if self.continue_to_crawl() {
                    if !self.visited_urls.contains(&url) {
                        if let Some(response) = Spider::load_url(url.as_ref(), &self.mime_types) {
                            self.visit(&url, response);
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    pub fn read_page(page: Page, document: &str) -> Tokenizer<Page> {
        let mut tok = Tokenizer::new(page, Default::default());
        let mut input = BufferQueue::new();
        input.push_back(String::from(document).into());
        let _ = tok.feed(&mut input);
        tok.end();
        tok
    }

    fn visit(&mut self, page_url: &Url, response: HttpResponse) {
        if self.can_visit(page_url) {
            info!("{}", page_url);
            if let Some(page) = Spider::read_response(page_url, response) {
                self.visit_page(page);
            }
        }
    }

    fn continue_to_crawl(&self) -> bool {
        self.limit == 0 || (self.visited_urls.len() as isize) < self.limit
    }

    fn can_visit(&self, page_url: &Url) -> bool {
        self.robot_parser.can_fetch(maman_name!(), page_url.path())
    }

    fn read_response(page_url: &Url, mut response: HttpResponse) -> Option<Page> {
        let mut headers = BTreeMap::new();
        {
            for (key, value) in response.headers().iter() {
                headers.insert(
                    key.as_str().to_string(),
                    value.to_str().unwrap_or("").to_string(),
                );
            }
        }
        match response.text() {
            Ok(content) => {
                let page = Page::new(
                    page_url.clone(),
                    content.to_string(),
                    headers,
                    response.status().to_string(),
                );
                let read = Spider::read_page(page, &content);
                Some(read.sink)
            }
            _ => None,
        }
    }

    fn load_url(url: &str, mime_types: &[mime::Mime]) -> Option<HttpResponse> {
        let client = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build()
            .expect("HttpClient failed to construct");
        match client
            .get(url)
            .header(USER_AGENT, maman_user_agent!())
            .send()
        {
            Err(_) => None,
            Ok(response) => match response.status() {
                StatusCode::OK | StatusCode::NOT_MODIFIED => {
                    if mime_types.is_empty() {
                        Some(response)
                    } else {
                        let content_type = response
                            .headers()
                            .get(CONTENT_TYPE)
                            .and_then(|value| value.to_str().ok())
                            .and_then(|value| value.parse::<mime::Mime>().ok())
                            .and_then(|value| {
                                let (type_, subtype) = (value.type_(), value.subtype());
                                {
                                    let mut text = type_.to_string();
                                    text.push_str("/");
                                    text.push_str(subtype.as_ref());
                                    mime::Mime::from_str(&text).ok()
                                }
                            });
                        match content_type {
                            Some(ct) => {
                                if mime_types.contains(&ct) {
                                    Some(response)
                                } else {
                                    None
                                }
                            }
                            None => None,
                        }
                    }
                }
                _ => None,
            },
        }
    }
}
<commit_msg>clippy: calling `push_str()` using a single-character string literal<commit_after>mod page;
pub use self::page::Page;

use std::collections::BTreeMap;
use std::default::Default;
use std::env;
use std::str::FromStr;
use std::time::Duration;

use html5ever::tokenizer::BufferQueue;
use html5ever::tokenizer::Tokenizer;
use reqwest::header::{CONTENT_TYPE, USER_AGENT};
use reqwest::Client as HttpClient;
use reqwest::Response as HttpResponse;
use reqwest::StatusCode;
use robotparser::RobotFileParser;
use sidekiq::Client as SidekiqClient;
use sidekiq::ClientOpts as SidekiqClientOpts;
use sidekiq::RedisPool;
use url::Url;

const MAMAN_ENV: &str = "MAMAN_ENV";
const MAMAN_ENV_DEFAULT: &str = "development";

pub struct Spider<'a> {
    pub base_url: Url,
    pub visited_urls: Vec<Url>,
    pub unvisited_urls: Vec<Url>,
    pub env: String,
    pub limit: isize,
    pub mime_types: Vec<mime::Mime>,
    sidekiq: SidekiqClient,
    robot_parser: RobotFileParser<'a>,
}

impl<'a> Spider<'a> {
    pub fn new(
        redis_pool: RedisPool,
        base_url: Url,
        limit: isize,
        mime_types: Vec<mime::Mime>,
    ) -> Spider<'a> {
        let maman_env =
            env::var(&MAMAN_ENV.to_owned()).unwrap_or_else(|_| MAMAN_ENV_DEFAULT.to_owned());
        let robots_txt = base_url.join("/robots.txt").unwrap();
        let robot_file_parser = RobotFileParser::new(robots_txt);
        let client_opts = SidekiqClientOpts {
            namespace: Some(maman_env.to_string()),
        };
        let sidekiq = SidekiqClient::new(redis_pool, client_opts);
        Spider {
            base_url,
            visited_urls: Vec::new(),
            unvisited_urls: Vec::new(),
            sidekiq,
            env: maman_env,
            robot_parser: robot_file_parser,
            limit,
            mime_types,
        }
    }

    pub fn visit_page(&mut self, page: Page) {
        self.visited_urls.push(page.url.clone());
        for u in &page.urls {
            self.unvisited_urls.push(u.clone());
        }
        if let Err(err) = self.sidekiq.push(page.to_job()) {
            error!("SidekiqClient push failed: {}", err);
        }
    }

    pub fn crawl(&mut self) {
        self.robot_parser.read();
        let base_url = self.base_url.clone();
        if let Some(response) = Spider::load_url(self.base_url.as_ref(), &self.mime_types) {
            self.visit(&base_url, response);
            while let Some(url) = self.unvisited_urls.pop() {
                if self.continue_to_crawl() {
                    if !self.visited_urls.contains(&url) {
                        if let Some(response) = Spider::load_url(url.as_ref(), &self.mime_types) {
                            self.visit(&url, response);
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    pub fn read_page(page: Page, document: &str) -> Tokenizer<Page> {
        let mut tok = Tokenizer::new(page, Default::default());
        let mut input = BufferQueue::new();
        input.push_back(String::from(document).into());
        let _ = tok.feed(&mut input);
        tok.end();
        tok
    }

    fn visit(&mut self, page_url: &Url, response: HttpResponse) {
        if self.can_visit(page_url) {
            info!("{}", page_url);
            if let Some(page) = Spider::read_response(page_url, response) {
                self.visit_page(page);
            }
        }
    }

    fn continue_to_crawl(&self) -> bool {
        self.limit == 0 || (self.visited_urls.len() as isize) < self.limit
    }

    fn can_visit(&self, page_url: &Url) -> bool {
        self.robot_parser.can_fetch(maman_name!(), page_url.path())
    }

    fn read_response(page_url: &Url, mut response: HttpResponse) -> Option<Page> {
        let mut headers = BTreeMap::new();
        {
            for (key, value) in response.headers().iter() {
                headers.insert(
                    key.as_str().to_string(),
                    value.to_str().unwrap_or("").to_string(),
                );
            }
        }
        match response.text() {
            Ok(content) => {
                let page = Page::new(
                    page_url.clone(),
                    content.to_string(),
                    headers,
                    response.status().to_string(),
                );
                let read = Spider::read_page(page, &content);
                Some(read.sink)
            }
            _ => None,
        }
    }

    fn load_url(url: &str, mime_types: &[mime::Mime]) -> Option<HttpResponse> {
        let client = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build()
            .expect("HttpClient failed to construct");
        match client
            .get(url)
            .header(USER_AGENT, maman_user_agent!())
            .send()
        {
            Err(_) => None,
            Ok(response) => match response.status() {
                StatusCode::OK | StatusCode::NOT_MODIFIED => {
                    if mime_types.is_empty() {
                        Some(response)
                    } else {
                        let content_type = response
                            .headers()
                            .get(CONTENT_TYPE)
                            .and_then(|value| value.to_str().ok())
                            .and_then(|value| value.parse::<mime::Mime>().ok())
                            .and_then(|value| {
                                let (type_, subtype) = (value.type_(), value.subtype());
                                {
                                    let mut text = type_.to_string();
                                    text.push('/');
                                    text.push_str(subtype.as_ref());
                                    mime::Mime::from_str(&text).ok()
                                }
                            });
                        match content_type {
                            Some(ct) => {
                                if mime_types.contains(&ct) {
                                    Some(response)
                                } else {
                                    None
                                }
                            }
                            None => None,
                        }
                    }
                }
                _ => None,
            },
        }
    }
}
<|endoftext|>


<commit_before>//! A simple wait-free, grow-only pagetable, assumes a dense keyspace.

use std::{
    alloc::{alloc_zeroed, Layout},
    mem::{align_of, size_of},
    sync::atomic::Ordering::{Relaxed, SeqCst},
};

use super::*;

const FANFACTOR: u64 = 18;
const FANOUT: u64 = 1 << FANFACTOR;
const FAN_MASK: u64 = FANOUT - 1;

#[doc(hidden)]
pub const PAGETABLE_NODE_SZ: usize = size_of::<Node1<()>>();

pub type PageId = u64;

#[inline(always)]
fn split_fanout(i: u64) -> (u64, u64) {
    // right shift 32 on 32-bit pointer systems panics
    #[cfg(target_pointer_width = "64")]
    assert!(
        i <= 1 << (FANFACTOR * 2),
        "trying to access key of {}, which is \
         higher than 2 ^ {}",
        i,
        (FANFACTOR * 2)
    );

    let left = i >> FANFACTOR;
    let right = i & FAN_MASK;
    (left, right)
}

struct Node1<T: Send + 'static> {
    children: [Atomic<Node2<T>>; FANOUT as usize],
}

struct Node2<T: Send + 'static> {
    children: [Atomic<T>; FANOUT as usize],
}

impl<T: Send + 'static> Node1<T> {
    fn new() -> Box<Node1<T>> {
        let size = size_of::<Node1<T>>();
        let align = align_of::<Node1<T>>();

        unsafe {
            let layout = Layout::from_size_align_unchecked(size, align);

            #[allow(clippy::cast_ptr_alignment)]
            let ptr = alloc_zeroed(layout) as *mut Node1<T>;

            Box::from_raw(ptr)
        }
    }
}

impl<T: Send + 'static> Node2<T> {
    fn new() -> Owned<Node2<T>> {
        let size = size_of::<Node2<T>>();
        let align = align_of::<Node2<T>>();

        unsafe {
            let layout = Layout::from_size_align_unchecked(size, align);

            #[allow(clippy::cast_ptr_alignment)]
            let ptr = alloc_zeroed(layout) as *mut Node2<T>;

            Owned::from_raw(ptr)
        }
    }
}

/// A simple lock-free radix tree.
pub struct PageTable<T>
where
    T: 'static + Send + Sync,
{
    head: Atomic<Node1<T>>,
}

impl<T> Default for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn default() -> PageTable<T> {
        let head = Node1::new();
        PageTable {
            head: Atomic::from(head),
        }
    }
}

impl<T> PageTable<T>
where
    T: 'static + Send + Sync,
{
    /// Atomically swap the previous value in a tree with a new one.
    pub fn swap<'g>(
        &self,
        pid: PageId,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> Shared<'g, T> {
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);
        debug_delay();
        tip.swap(new, SeqCst, guard)
    }

    /// Compare and swap an old value to a new one.
    pub fn cas<'g>(
        &self,
        pid: PageId,
        old: Shared<'g, T>,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> std::result::Result<Shared<'g, T>, Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        debug_delay();

        let _ = tip
            .compare_and_set(old, new, SeqCst, guard)
            .map_err(|e| e.current)?;

        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
        }
        Ok(new)
    }

    /// Try to get a value from the tree.
    pub fn get<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        let res = tip.load(SeqCst, guard);
        if res.is_null() {
            None
        } else {
            Some(res)
        }
    }

    /// Delete a value from the tree, returning the old value if it was set.
    pub fn del<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let old = self.swap(pid, Shared::null(), guard);
        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
            Some(old)
        } else {
            None
        }
    }
}

fn traverse<'g, T: 'static + Send>(
    head: Shared<'g, Node1<T>>,
    k: PageId,
    guard: &'g Guard,
) -> &'g Atomic<T> {
    let (l1k, l2k) = split_fanout(k);

    debug_delay();
    let l1 = unsafe { &head.deref().children };

    debug_delay();
    let mut l2_ptr = l1[usize::try_from(l1k).unwrap()].load(SeqCst, guard);

    if l2_ptr.is_null() {
        let next_child = Node2::new().into_shared(guard);

        debug_delay();
        let ret = l1[usize::try_from(l1k).unwrap()]
            .compare_and_set(l2_ptr, next_child, SeqCst, guard);

        match ret {
            Ok(_) => {
                // CAS worked
                l2_ptr = next_child;
            }
            Err(e) => {
                // another thread beat us, drop unused created
                // child and use what is already set
                l2_ptr = e.current;
            }
        }
    }

    debug_delay();
    let l2 = unsafe { &l2_ptr.deref().children };

    &l2[usize::try_from(l2k).unwrap()]
}

impl<T> Drop for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn drop(&mut self) {
        unsafe {
            let head = self.head.load(Relaxed, &unprotected()).into_owned();
            drop(head);
        }
    }
}

impl<T: Send + 'static> Drop for Node1<T> {
    fn drop(&mut self) {
        unsafe {
            let children: Vec<*const Node2<T>> = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null())
                .collect();

            for child in children {
                drop(Box::from_raw(child as *mut Node2<T>));
            }
        }
    }
}

impl<T: Send + 'static> Drop for Node2<T> {
    fn drop(&mut self) {
        unsafe {
            let children: Vec<*const T> = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null())
                .collect();

            for child in children {
                drop(Box::from_raw(child as *mut T));
            }
        }
    }
}

#[test]
fn test_split_fanout() {
    assert_eq!(
        split_fanout(0b11_1111_1111_1111_1111),
        (0, 0b11_1111_1111_1111_1111)
    );
    assert_eq!(
        split_fanout(0b111_1111_1111_1111_1111),
        (0b1, 0b11_1111_1111_1111_1111)
    );
}

#[test]
fn basic_functionality() {
    unsafe {
        let guard = pin();
        let rt = PageTable::default();
        let v1 = Owned::new(5).into_shared(&guard);
        rt.cas(0, Shared::null(), v1, &guard).unwrap();
        let ptr = rt.get(0, &guard).unwrap();
        assert_eq!(ptr.deref(), &5);
        rt.cas(0, ptr, Owned::new(6).into_shared(&guard), &guard)
            .unwrap();
        assert_eq!(rt.get(0, &guard).unwrap().deref(), &6);
        rt.del(0, &guard);
        assert!(rt.get(0, &guard).is_none());

        let k2 = 321 << FANFACTOR;
        let k3 = k2 + 1;

        let v2 = Owned::new(2).into_shared(&guard);
        rt.cas(k2, Shared::null(), v2, &guard).unwrap();
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
        assert!(rt.get(k3, &guard).is_none());
        let v3 = Owned::new(3).into_shared(&guard);
        rt.cas(k3, Shared::null(), v3, &guard).unwrap();
        assert_eq!(rt.get(k3, &guard).unwrap().deref(), &3);
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
    }
}

#[test]
#[ignore]
fn test_model() {
    use self::Shared as EpochShared;
    use model::{model, prop_oneof};

    model! {
        Model => let mut m = std::collections::HashMap::new(),
        Implementation => let i = PageTable::default(),
        Insert((u64, u64))((k, new) in (0u64..4, 0u64..4)) => {
            if !m.contains_key(&k) {
                m.insert(k, new);

                let guard = pin();
                let v = Owned::new(new).into_shared(&guard);
                i.cas(k, EpochShared::null(), v, &guard ).expect("should be able to insert a value");
            }
        },
        Get(u64)(k in 0u64..4) => {
            let guard = pin();
            let expected = m.get(&k);
            let actual = i.get(k, &guard).map(|s| unsafe { s.deref() });
            assert_eq!(expected, actual);
        },
        Cas((u64, u64, u64))((k, old, new) in (0u64..4, 0u64..4, 0u64..4)) => {
            let guard = pin();
            let expected_current = m.get(&k).cloned();
            let actual_current = i.get(k, &guard);
            assert_eq!(expected_current, actual_current.map(|s| unsafe { *s.deref() }));
            if expected_current.is_none() {
                continue;
            }

            let new_v = Owned::new(new).into_shared(&guard);

            if expected_current == Some(old) {
                m.insert(k, new);
                let cas_res = i.cas(k, actual_current.unwrap(), new_v, &guard);
                assert!(cas_res.is_ok());
            };
        }
    }
}

#[test]
#[ignore]
fn test_linearizability() {
    use self::Shared as EpochShared;
    use model::{linearizable, prop_oneof, Shared};

    linearizable! {
        Implementation => let i = Shared::new(PageTable::default()),
        Get(u64)(k in 0u64..4) -> Option<u64> {
            let guard = pin();
            unsafe {
                i.get(k, &guard).map(|s| *s.deref())
            }
        },
        Insert((u64, u64))((k, new) in (0u64..4, 0u64..4)) -> bool {
            let guard = pin();
            let v = Owned::new(new).into_shared(&guard);
            i.cas(k, EpochShared::null(), v, &guard).is_err()
        },
        Cas((u64, u64, u64))((k, old, new)
            in (0u64..4, 0u64..4, 0u64..4))
            -> std::result::Result<(), u64> {
            let guard = pin();
            i.cas(k, Owned::new(old).into_shared(&guard), Owned::new(new).into_shared(&guard), &guard)
                .map(|_| ())
                .map_err(|s| if s.is_null() {
                    0
                } else {
                    unsafe { *s.deref() }
                })
        }
    }
}
<commit_msg>don't turn filter iterator into a vector before using it when dropping the pagetable<commit_after>//! A simple wait-free, grow-only pagetable, assumes a dense keyspace.

use std::{
    alloc::{alloc_zeroed, Layout},
    mem::{align_of, size_of},
    sync::atomic::Ordering::{Relaxed, SeqCst},
};

use super::*;

const FANFACTOR: u64 = 18;
const FANOUT: u64 = 1 << FANFACTOR;
const FAN_MASK: u64 = FANOUT - 1;

#[doc(hidden)]
pub const PAGETABLE_NODE_SZ: usize = size_of::<Node1<()>>();

pub type PageId = u64;

#[inline(always)]
fn split_fanout(i: u64) -> (u64, u64) {
    // right shift 32 on 32-bit pointer systems panics
    #[cfg(target_pointer_width = "64")]
    assert!(
        i <= 1 << (FANFACTOR * 2),
        "trying to access key of {}, which is \
         higher than 2 ^ {}",
        i,
        (FANFACTOR * 2)
    );

    let left = i >> FANFACTOR;
    let right = i & FAN_MASK;
    (left, right)
}

struct Node1<T: Send + 'static> {
    children: [Atomic<Node2<T>>; FANOUT as usize],
}

struct Node2<T: Send + 'static> {
    children: [Atomic<T>; FANOUT as usize],
}

impl<T: Send + 'static> Node1<T> {
    fn new() -> Box<Node1<T>> {
        let size = size_of::<Node1<T>>();
        let align = align_of::<Node1<T>>();

        unsafe {
            let layout = Layout::from_size_align_unchecked(size, align);

            #[allow(clippy::cast_ptr_alignment)]
            let ptr = alloc_zeroed(layout) as *mut Node1<T>;

            Box::from_raw(ptr)
        }
    }
}

impl<T: Send + 'static> Node2<T> {
    fn new() -> Owned<Node2<T>> {
        let size = size_of::<Node2<T>>();
        let align = align_of::<Node2<T>>();

        unsafe {
            let layout = Layout::from_size_align_unchecked(size, align);

            #[allow(clippy::cast_ptr_alignment)]
            let ptr = alloc_zeroed(layout) as *mut Node2<T>;

            Owned::from_raw(ptr)
        }
    }
}

impl<T: Send + 'static> Drop for Node1<T> {
    fn drop(&mut self) {
        unsafe {
            let children = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null());

            for child in children {
                drop(Box::from_raw(child as *mut Node2<T>));
            }
        }
    }
}

impl<T: Send + 'static> Drop for Node2<T> {
    fn drop(&mut self) {
        unsafe {
            let children = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null());

            for child in children {
                drop(Box::from_raw(child as *mut T));
            }
        }
    }
}

/// A simple lock-free radix tree.
pub struct PageTable<T>
where
    T: 'static + Send + Sync,
{
    head: Atomic<Node1<T>>,
}

impl<T> Default for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn default() -> PageTable<T> {
        let head = Node1::new();
        PageTable {
            head: Atomic::from(head),
        }
    }
}

impl<T> PageTable<T>
where
    T: 'static + Send + Sync,
{
    /// Atomically swap the previous value in a tree with a new one.
    pub fn swap<'g>(
        &self,
        pid: PageId,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> Shared<'g, T> {
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);
        debug_delay();
        tip.swap(new, SeqCst, guard)
    }

    /// Compare and swap an old value to a new one.
    pub fn cas<'g>(
        &self,
        pid: PageId,
        old: Shared<'g, T>,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> std::result::Result<Shared<'g, T>, Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        debug_delay();

        let _ = tip
            .compare_and_set(old, new, SeqCst, guard)
            .map_err(|e| e.current)?;

        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
        }
        Ok(new)
    }

    /// Try to get a value from the tree.
    pub fn get<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        let res = tip.load(SeqCst, guard);
        if res.is_null() {
            None
        } else {
            Some(res)
        }
    }

    /// Delete a value from the tree, returning the old value if it was set.
    pub fn del<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let old = self.swap(pid, Shared::null(), guard);
        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
            Some(old)
        } else {
            None
        }
    }
}

fn traverse<'g, T: 'static + Send>(
    head: Shared<'g, Node1<T>>,
    k: PageId,
    guard: &'g Guard,
) -> &'g Atomic<T> {
    let (l1k, l2k) = split_fanout(k);

    debug_delay();
    let l1 = unsafe { &head.deref().children };

    debug_delay();
    let mut l2_ptr = l1[usize::try_from(l1k).unwrap()].load(SeqCst, guard);

    if l2_ptr.is_null() {
        let next_child = Node2::new().into_shared(guard);

        debug_delay();
        let ret = l1[usize::try_from(l1k).unwrap()]
            .compare_and_set(l2_ptr, next_child, SeqCst, guard);

        match ret {
            Ok(_) => {
                // CAS worked
                l2_ptr = next_child;
            }
            Err(e) => {
                // another thread beat us, drop unused created
                // child and use what is already set
                l2_ptr = e.current;
            }
        }
    }

    debug_delay();
    let l2 = unsafe { &l2_ptr.deref().children };

    &l2[usize::try_from(l2k).unwrap()]
}

impl<T> Drop for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn drop(&mut self) {
        unsafe {
            let head = self.head.load(Relaxed, &unprotected()).into_owned();
            drop(head);
        }
    }
}

#[test]
fn test_split_fanout() {
    assert_eq!(
        split_fanout(0b11_1111_1111_1111_1111),
        (0, 0b11_1111_1111_1111_1111)
    );
    assert_eq!(
        split_fanout(0b111_1111_1111_1111_1111),
        (0b1, 0b11_1111_1111_1111_1111)
    );
}

#[test]
fn basic_functionality() {
    unsafe {
        let guard = pin();
        let rt = PageTable::default();
        let v1 = Owned::new(5).into_shared(&guard);
        rt.cas(0, Shared::null(), v1, &guard).unwrap();
        let ptr = rt.get(0, &guard).unwrap();
        assert_eq!(ptr.deref(), &5);
        rt.cas(0, ptr, Owned::new(6).into_shared(&guard), &guard)
            .unwrap();
        assert_eq!(rt.get(0, &guard).unwrap().deref(), &6);
        rt.del(0, &guard);
        assert!(rt.get(0, &guard).is_none());

        let k2 = 321 << FANFACTOR;
        let k3 = k2 + 1;

        let v2 = Owned::new(2).into_shared(&guard);
        rt.cas(k2, Shared::null(), v2, &guard).unwrap();
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
        assert!(rt.get(k3, &guard).is_none());
        let v3 = Owned::new(3).into_shared(&guard);
        rt.cas(k3, Shared::null(), v3, &guard).unwrap();
        assert_eq!(rt.get(k3, &guard).unwrap().deref(), &3);
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
    }
}

#[test]
#[ignore]
fn test_model() {
    use self::Shared as EpochShared;
    use model::{model, prop_oneof};

    model! {
        Model => let mut m = std::collections::HashMap::new(),
        Implementation => let i = PageTable::default(),
        Insert((u64, u64))((k, new) in (0u64..4, 0u64..4)) => {
            if !m.contains_key(&k) {
                m.insert(k, new);

                let guard = pin();
                let v = Owned::new(new).into_shared(&guard);
                i.cas(k, EpochShared::null(), v, &guard ).expect("should be able to insert a value");
            }
        },
        Get(u64)(k in 0u64..4) => {
            let guard = pin();
            let expected = m.get(&k);
            let actual = i.get(k, &guard).map(|s| unsafe { s.deref() });
            assert_eq!(expected, actual);
        },
        Cas((u64, u64, u64))((k, old, new) in (0u64..4, 0u64..4, 0u64..4)) => {
            let guard = pin();
            let expected_current = m.get(&k).cloned();
            let actual_current = i.get(k, &guard);
            assert_eq!(expected_current, actual_current.map(|s| unsafe { *s.deref() }));
            if expected_current.is_none() {
                continue;
            }

            let new_v = Owned::new(new).into_shared(&guard);

            if expected_current == Some(old) {
                m.insert(k, new);
                let cas_res = i.cas(k, actual_current.unwrap(), new_v, &guard);
                assert!(cas_res.is_ok());
            };
        }
    }
}

#[test]
#[ignore]
fn test_linearizability() {
    use self::Shared as EpochShared;
    use model::{linearizable, prop_oneof, Shared};

    linearizable! {
        Implementation => let i = Shared::new(PageTable::default()),
        Get(u64)(k in 0u64..4) -> Option<u64> {
            let guard = pin();
            unsafe {
                i.get(k, &guard).map(|s| *s.deref())
            }
        },
        Insert((u64, u64))((k, new) in (0u64..4, 0u64..4)) -> bool {
            let guard = pin();
            let v = Owned::new(new).into_shared(&guard);
            i.cas(k, EpochShared::null(), v, &guard).is_err()
        },
        Cas((u64, u64, u64))((k, old, new)
            in (0u64..4, 0u64..4, 0u64..4))
            -> std::result::Result<(), u64> {
            let guard = pin();
            i.cas(k, Owned::new(old).into_shared(&guard), Owned::new(new).into_shared(&guard), &guard)
                .map(|_| ())
                .map_err(|s| if s.is_null() {
                    0
                } else {
                    unsafe { *s.deref() }
                })
        }
    }
}
<|endoftext|>