<commit_before>//! This module defines a cheaply-copyable cursor into a TokenStream's data.
//!
//! It does this by copying the data into a stably-addressed structured buffer,
//! and holding raw pointers into that buffer to allow walking through delimited
//! groups cheaply.
//!
//! This module is heavily commented as it contains the only unsafe code in
//! `syn`, and caution should be made when editing it. It provides a safe
//! interface, but is fragile internally.

use proc_macro2::*;

use std::ptr;
use std::marker::PhantomData;

#[cfg(all(feature = "verbose-trace", not(feature = "all-features")))]
use std::fmt::{self, Debug};

/// Internal type which is used instead of `TokenTree` to represent a single
/// `TokenTree` within a `SynomBuffer`.
enum Entry {
    /// Mimicing types from proc-macro.
    Group(Span, Delimiter, SynomBuffer),
    Term(Span, Term),
    Op(Span, char, Spacing),
    Literal(Span, Literal),
    /// End entries contain a raw pointer to the entry from the containing
    /// TokenTree.
    End(*const Entry),
}

/// A buffer of data which contains a structured representation of the input
/// `TokenStream` object.
pub struct SynomBuffer {
    // NOTE: Do not derive clone on this - there are raw pointers inside which
    // will be messed up. Moving the `SynomBuffer` itself is safe as the actual
    // backing slices won't be moved.
    data: Box<[Entry]>,
}

impl SynomBuffer {
    // NOTE: DO NOT MUTATE THE `Vec` RETURNED FROM THIS FUNCTION ONCE IT
    // RETURNS, THE ADDRESS OF ITS BACKING MEMORY MUST REMAIN STABLE.
    fn inner_new(stream: TokenStream, up: *const Entry) -> SynomBuffer {
        // Build up the entries list, recording the locations of any Groups
        // in the list to be processed later.
        let mut entries = Vec::new();
        let mut seqs = Vec::new();
        for tt in stream {
            match tt.kind {
                TokenNode::Term(sym) => {
                    entries.push(Entry::Term(tt.span, sym));
                }
                TokenNode::Op(chr, ok) => {
                    entries.push(Entry::Op(tt.span, chr, ok));
                }
                TokenNode::Literal(lit) => {
                    entries.push(Entry::Literal(tt.span, lit));
                }
                TokenNode::Group(delim, seq_stream) => {
                    // Record the index of the interesting entry, and store an
                    // `End(null)` there temporarially.
                    seqs.push((entries.len(), tt.span, delim, seq_stream));
                    entries.push(Entry::End(ptr::null()));
                }
            }
        }
        // Add an `End` entry to the end with a reference to the enclosing token
        // stream which was passed in.
        entries.push(Entry::End(up));

        // NOTE: This is done to ensure that we don't accidentally modify the
        // length of the backing buffer. The backing buffer must remain at a
        // constant address after this point, as we are going to store a raw
        // pointer into it.
        let mut entries = entries.into_boxed_slice();
        for (idx, span, delim, seq_stream) in seqs {
            // We know that this index refers to one of the temporary
            // `End(null)` entries, and we know that the last entry is
            // `End(up)`, so the next index is also valid.
            let seq_up = &entries[idx + 1] as *const Entry;

            // The end entry stored at the end of this Entry::Group should
            // point to the Entry which follows the Group in the list.
            let inner = Self::inner_new(seq_stream, seq_up);
            entries[idx] = Entry::Group(span, delim, inner);
        }

        SynomBuffer { data: entries }
    }

    /// Create a new SynomBuffer, which contains the data from the given
    /// TokenStream.
    pub fn new(stream: TokenStream) -> SynomBuffer {
        Self::inner_new(stream, ptr::null())
    }

    /// Create a cursor referencing the first token in the input.
    pub fn begin(&self) -> Cursor {
        unsafe { Cursor::create(&self.data[0], &self.data[self.data.len() - 1]) }
    }
}

/// A cursor into an input `TokenStream`'s data. This cursor holds a reference
/// into the immutable data which is used internally to represent a
/// `TokenStream`, and can be efficiently manipulated and copied around.
///
/// An empty `Cursor` can be created directly, or one may create a `SynomBuffer`
/// object and get a cursor to its first token with `begin()`.
///
/// Two cursors are equal if they have the same location in the same input
/// stream, and have the same scope.
#[derive(Copy, Clone, Eq, PartialEq)]
pub struct Cursor<'a> {
    /// The current entry which the `Cursor` is pointing at.
    ptr: *const Entry,
    /// This is the only `Entry::End(..)` object which this cursor is allowed to
    /// point at. All other `End` objects are skipped over in `Cursor::create`.
    scope: *const Entry,
    /// This uses the &'a reference which guarantees that these pointers are
    /// still valid.
    marker: PhantomData<&'a Entry>,
}

impl<'a> Cursor<'a> {
    /// Create a cursor referencing a static empty TokenStream.
    pub fn empty() -> Self {
        // It's safe in this situation for us to put an `Entry` object in global
        // storage, despite it not actually being safe to send across threads
        // (`Term` is a reference into a thread-local table). This is because
        // this entry never includes a `Term` object.
        //
        // This wrapper struct allows us to break the rules and put a `Sync`
        // object in global storage.
        struct UnsafeSyncEntry(Entry);
        unsafe impl Sync for UnsafeSyncEntry {}
        static EMPTY_ENTRY: UnsafeSyncEntry = UnsafeSyncEntry(Entry::End(0 as *const Entry));

        Cursor {
            ptr: &EMPTY_ENTRY.0,
            scope: &EMPTY_ENTRY.0,
            marker: PhantomData,
        }
    }

    /// This create method intelligently exits non-explicitly-entered
    /// `None`-delimited scopes when the cursor reaches the end of them,
    /// allowing for them to be treated transparently.
    unsafe fn create(mut ptr: *const Entry, scope: *const Entry) -> Self {
        // NOTE: If we're looking at a `End(..)`, we want to advance the cursor
        // past it, unless `ptr == scope`, which means that we're at the edge of
        // our cursor's scope. We should only have `ptr != scope` at the exit
        // from None-delimited groups entered with `ignore_none`.
        while let Entry::End(exit) = *ptr {
            if ptr == scope {
                break;
            }
            ptr = exit;
        }

        Cursor {
            ptr: ptr,
            scope: scope,
            marker: PhantomData,
        }
    }

    /// Get the current entry.
    fn entry(self) -> &'a Entry {
        unsafe { &*self.ptr }
    }

    /// Bump the cursor to point at the next token after the current one. This
    /// is undefined behavior if the cursor is currently looking at an
    /// `Entry::End`.
    unsafe fn bump(self) -> Cursor<'a> {
        Cursor::create(self.ptr.offset(1), self.scope)
    }

    /// If the cursor is looking at a `None`-delimited group, move it to look at
    /// the first token inside instead. If the group is empty, this will move
    /// the cursor past the `None`-delimited group.
    ///
    /// WARNING: This mutates its argument.
    fn ignore_none(&mut self) {
        if let Entry::Group(_, Delimiter::None, ref buf) = *self.entry() {
            // NOTE: We call `Cursor::create` here to make sure that situations
            // where we should immediately exit the span after entering it are
            // handled correctly.
            unsafe {
                *self = Cursor::create(&buf.data[0], self.scope);
            }
        }
    }

    /// Check if the cursor is currently pointing at the end of its valid scope.
    #[inline]
    pub fn eof(self) -> bool {
        // We're at eof if we're at the end of our scope.
        self.ptr == self.scope
    }

    /// If the cursor is pointing at a Group with the given `Delimiter`, return
    /// a cursor into that group, and one pointing to the next `TokenTree`.
    pub fn group(mut self, delim: Delimiter) -> Option<(Cursor<'a>, Span, Cursor<'a>)> {
        // If we're not trying to enter a none-delimited group, we want to
        // ignore them. We have to make sure to _not_ ignore them when we want
        // to enter them, of course. For obvious reasons.
        if delim != Delimiter::None {
            self.ignore_none();
        }

        if let Entry::Group(span, group_delim, ref buf) = *self.entry() {
            if group_delim == delim {
                return Some((buf.begin(), span, unsafe { self.bump() }));
            }
        }

        None
    }

    /// If the cursor is pointing at a Term, return it and a cursor pointing at
    /// the next `TokenTree`.
    pub fn term(mut self) -> Option<(Span, Term, Cursor<'a>)> {
        self.ignore_none();
        match *self.entry() {
            Entry::Term(span, term) => Some((span, term, unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at an Op, return it and a cursor pointing
    /// at the next `TokenTree`.
    pub fn op(mut self) -> Option<(Span, char, Spacing, Cursor<'a>)> {
        self.ignore_none();
        match *self.entry() {
            Entry::Op(span, op, spacing) => Some((span, op, spacing, unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a Literal, return it and a cursor pointing
    /// at the next `TokenTree`.
    pub fn literal(mut self) -> Option<(Span, Literal, Cursor<'a>)> {
        self.ignore_none();
        match *self.entry() {
            Entry::Literal(span, ref lit) => Some((span, lit.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// Copy all remaining tokens visible from this cursor into a `TokenStream`.
    pub fn token_stream(self) -> TokenStream {
        let mut tts = Vec::new();
        let mut cursor = self;
        while let Some((tt, rest)) = cursor.token_tree() {
            tts.push(tt);
            cursor = rest;
        }
        tts.into_iter().collect()
    }

    /// If the cursor is looking at a `TokenTree`, returns it along with a
    /// cursor pointing to the next token in the group, otherwise returns
    /// `None`.
    ///
    /// This method does not treat `None`-delimited groups as invisible, and
    /// will return a `Group(None, ..)` if the cursor is looking at one.
    pub fn token_tree(self) -> Option<(TokenTree, Cursor<'a>)> {
        let tree = match *self.entry() {
            Entry::Group(span, delim, ref buf) => {
                let stream = buf.begin().token_stream();
                TokenTree {
                    span: span,
                    kind: TokenNode::Group(delim, stream),
                }
            }
            Entry::Literal(span, ref lit) => TokenTree {
                span: span,
                kind: TokenNode::Literal(lit.clone()),
            },
            Entry::Term(span, sym) => TokenTree {
                span: span,
                kind: TokenNode::Term(sym),
            },
            Entry::Op(span, chr, spacing) => TokenTree {
                span: span,
                kind: TokenNode::Op(chr, spacing),
            },
            Entry::End(..) => {
                return None;
            }
        };

        Some((tree, unsafe { self.bump() }))
    }
}

// We do a custom implementation for `Debug` as the default implementation is
// pretty useless.
#[cfg(all(feature = "verbose-trace", not(feature = "all-features")))]
impl<'a> Debug for Cursor<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Print what the cursor is currently looking at.
        // This will look like Cursor("some remaining tokens here")
        f.debug_tuple("Cursor")
            .field(&self.token_stream().to_string())
            .finish()
    }
}
<commit_msg>Expose Span pointed to by Cursor<commit_after>//! This module defines a cheaply-copyable cursor into a TokenStream's data.
//!
//! It does this by copying the data into a stably-addressed structured buffer,
//! and holding raw pointers into that buffer to allow walking through delimited
//! groups cheaply.
//!
//! This module is heavily commented as it contains the only unsafe code in
//! `syn`, and caution should be made when editing it. It provides a safe
//! interface, but is fragile internally.

use proc_macro2::*;

use std::ptr;
use std::marker::PhantomData;

#[cfg(all(feature = "verbose-trace", not(feature = "all-features")))]
use std::fmt::{self, Debug};

/// Internal type which is used instead of `TokenTree` to represent a single
/// `TokenTree` within a `SynomBuffer`.
enum Entry {
    /// Mimicing types from proc-macro.
    Group(Span, Delimiter, SynomBuffer),
    Term(Span, Term),
    Op(Span, char, Spacing),
    Literal(Span, Literal),
    /// End entries contain a raw pointer to the entry from the containing
    /// TokenTree.
    End(*const Entry),
}

/// A buffer of data which contains a structured representation of the input
/// `TokenStream` object.
pub struct SynomBuffer {
    // NOTE: Do not derive clone on this - there are raw pointers inside which
    // will be messed up. Moving the `SynomBuffer` itself is safe as the actual
    // backing slices won't be moved.
    data: Box<[Entry]>,
}

impl SynomBuffer {
    // NOTE: DO NOT MUTATE THE `Vec` RETURNED FROM THIS FUNCTION ONCE IT
    // RETURNS, THE ADDRESS OF ITS BACKING MEMORY MUST REMAIN STABLE.
    fn inner_new(stream: TokenStream, up: *const Entry) -> SynomBuffer {
        // Build up the entries list, recording the locations of any Groups
        // in the list to be processed later.
        let mut entries = Vec::new();
        let mut seqs = Vec::new();
        for tt in stream {
            match tt.kind {
                TokenNode::Term(sym) => {
                    entries.push(Entry::Term(tt.span, sym));
                }
                TokenNode::Op(chr, ok) => {
                    entries.push(Entry::Op(tt.span, chr, ok));
                }
                TokenNode::Literal(lit) => {
                    entries.push(Entry::Literal(tt.span, lit));
                }
                TokenNode::Group(delim, seq_stream) => {
                    // Record the index of the interesting entry, and store an
                    // `End(null)` there temporarially.
                    seqs.push((entries.len(), tt.span, delim, seq_stream));
                    entries.push(Entry::End(ptr::null()));
                }
            }
        }
        // Add an `End` entry to the end with a reference to the enclosing token
        // stream which was passed in.
        entries.push(Entry::End(up));

        // NOTE: This is done to ensure that we don't accidentally modify the
        // length of the backing buffer. The backing buffer must remain at a
        // constant address after this point, as we are going to store a raw
        // pointer into it.
        let mut entries = entries.into_boxed_slice();
        for (idx, span, delim, seq_stream) in seqs {
            // We know that this index refers to one of the temporary
            // `End(null)` entries, and we know that the last entry is
            // `End(up)`, so the next index is also valid.
            let seq_up = &entries[idx + 1] as *const Entry;

            // The end entry stored at the end of this Entry::Group should
            // point to the Entry which follows the Group in the list.
            let inner = Self::inner_new(seq_stream, seq_up);
            entries[idx] = Entry::Group(span, delim, inner);
        }

        SynomBuffer { data: entries }
    }

    /// Create a new SynomBuffer, which contains the data from the given
    /// TokenStream.
    pub fn new(stream: TokenStream) -> SynomBuffer {
        Self::inner_new(stream, ptr::null())
    }

    /// Create a cursor referencing the first token in the input.
    pub fn begin(&self) -> Cursor {
        unsafe { Cursor::create(&self.data[0], &self.data[self.data.len() - 1]) }
    }
}

/// A cursor into an input `TokenStream`'s data. This cursor holds a reference
/// into the immutable data which is used internally to represent a
/// `TokenStream`, and can be efficiently manipulated and copied around.
///
/// An empty `Cursor` can be created directly, or one may create a `SynomBuffer`
/// object and get a cursor to its first token with `begin()`.
///
/// Two cursors are equal if they have the same location in the same input
/// stream, and have the same scope.
#[derive(Copy, Clone, Eq, PartialEq)]
pub struct Cursor<'a> {
    /// The current entry which the `Cursor` is pointing at.
    ptr: *const Entry,
    /// This is the only `Entry::End(..)` object which this cursor is allowed to
    /// point at. All other `End` objects are skipped over in `Cursor::create`.
    scope: *const Entry,
    /// This uses the &'a reference which guarantees that these pointers are
    /// still valid.
    marker: PhantomData<&'a Entry>,
}

impl<'a> Cursor<'a> {
    /// Create a cursor referencing a static empty TokenStream.
    pub fn empty() -> Self {
        // It's safe in this situation for us to put an `Entry` object in global
        // storage, despite it not actually being safe to send across threads
        // (`Term` is a reference into a thread-local table). This is because
        // this entry never includes a `Term` object.
        //
        // This wrapper struct allows us to break the rules and put a `Sync`
        // object in global storage.
        struct UnsafeSyncEntry(Entry);
        unsafe impl Sync for UnsafeSyncEntry {}
        static EMPTY_ENTRY: UnsafeSyncEntry = UnsafeSyncEntry(Entry::End(0 as *const Entry));

        Cursor {
            ptr: &EMPTY_ENTRY.0,
            scope: &EMPTY_ENTRY.0,
            marker: PhantomData,
        }
    }

    /// This create method intelligently exits non-explicitly-entered
    /// `None`-delimited scopes when the cursor reaches the end of them,
    /// allowing for them to be treated transparently.
    unsafe fn create(mut ptr: *const Entry, scope: *const Entry) -> Self {
        // NOTE: If we're looking at a `End(..)`, we want to advance the cursor
        // past it, unless `ptr == scope`, which means that we're at the edge of
        // our cursor's scope. We should only have `ptr != scope` at the exit
        // from None-delimited groups entered with `ignore_none`.
        while let Entry::End(exit) = *ptr {
            if ptr == scope {
                break;
            }
            ptr = exit;
        }

        Cursor {
            ptr: ptr,
            scope: scope,
            marker: PhantomData,
        }
    }

    /// Get the current entry.
    fn entry(self) -> &'a Entry {
        unsafe { &*self.ptr }
    }

    /// Bump the cursor to point at the next token after the current one. This
    /// is undefined behavior if the cursor is currently looking at an
    /// `Entry::End`.
    unsafe fn bump(self) -> Cursor<'a> {
        Cursor::create(self.ptr.offset(1), self.scope)
    }

    /// If the cursor is looking at a `None`-delimited group, move it to look at
    /// the first token inside instead. If the group is empty, this will move
    /// the cursor past the `None`-delimited group.
    ///
    /// WARNING: This mutates its argument.
    fn ignore_none(&mut self) {
        if let Entry::Group(_, Delimiter::None, ref buf) = *self.entry() {
            // NOTE: We call `Cursor::create` here to make sure that situations
            // where we should immediately exit the span after entering it are
            // handled correctly.
            unsafe {
                *self = Cursor::create(&buf.data[0], self.scope);
            }
        }
    }

    /// Check if the cursor is currently pointing at the end of its valid scope.
    #[inline]
    pub fn eof(self) -> bool {
        // We're at eof if we're at the end of our scope.
        self.ptr == self.scope
    }

    /// If the cursor is pointing at a Group with the given `Delimiter`, return
    /// a cursor into that group, and one pointing to the next `TokenTree`.
    pub fn group(mut self, delim: Delimiter) -> Option<(Cursor<'a>, Span, Cursor<'a>)> {
        // If we're not trying to enter a none-delimited group, we want to
        // ignore them. We have to make sure to _not_ ignore them when we want
        // to enter them, of course. For obvious reasons.
        if delim != Delimiter::None {
            self.ignore_none();
        }

        if let Entry::Group(span, group_delim, ref buf) = *self.entry() {
            if group_delim == delim {
                return Some((buf.begin(), span, unsafe { self.bump() }));
            }
        }

        None
    }

    /// If the cursor is pointing at a Term, return it and a cursor pointing at
    /// the next `TokenTree`.
    pub fn term(mut self) -> Option<(Span, Term, Cursor<'a>)> {
        self.ignore_none();
        match *self.entry() {
            Entry::Term(span, term) => Some((span, term, unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at an Op, return it and a cursor pointing
    /// at the next `TokenTree`.
    pub fn op(mut self) -> Option<(Span, char, Spacing, Cursor<'a>)> {
        self.ignore_none();
        match *self.entry() {
            Entry::Op(span, op, spacing) => Some((span, op, spacing, unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a Literal, return it and a cursor pointing
    /// at the next `TokenTree`.
    pub fn literal(mut self) -> Option<(Span, Literal, Cursor<'a>)> {
        self.ignore_none();
        match *self.entry() {
            Entry::Literal(span, ref lit) => Some((span, lit.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// Copy all remaining tokens visible from this cursor into a `TokenStream`.
    pub fn token_stream(self) -> TokenStream {
        let mut tts = Vec::new();
        let mut cursor = self;
        while let Some((tt, rest)) = cursor.token_tree() {
            tts.push(tt);
            cursor = rest;
        }
        tts.into_iter().collect()
    }

    /// If the cursor is looking at a `TokenTree`, returns it along with a
    /// cursor pointing to the next token in the group, otherwise returns
    /// `None`.
    ///
    /// This method does not treat `None`-delimited groups as invisible, and
    /// will return a `Group(None, ..)` if the cursor is looking at one.
    pub fn token_tree(self) -> Option<(TokenTree, Cursor<'a>)> {
        let tree = match *self.entry() {
            Entry::Group(span, delim, ref buf) => {
                let stream = buf.begin().token_stream();
                TokenTree {
                    span: span,
                    kind: TokenNode::Group(delim, stream),
                }
            }
            Entry::Literal(span, ref lit) => TokenTree {
                span: span,
                kind: TokenNode::Literal(lit.clone()),
            },
            Entry::Term(span, sym) => TokenTree {
                span: span,
                kind: TokenNode::Term(sym),
            },
            Entry::Op(span, chr, spacing) => TokenTree {
                span: span,
                kind: TokenNode::Op(chr, spacing),
            },
            Entry::End(..) => {
                return None;
            }
        };

        Some((tree, unsafe { self.bump() }))
    }

    /// Returns the `Span` of the current token, or `Span::call_site()` if this
    /// cursor points to eof.
    pub fn span(self) -> Span {
        match *self.entry() {
            Entry::Group(span, ..) => span,
            Entry::Literal(span, ..) => span,
            Entry::Term(span, ..) => span,
            Entry::Op(span, ..) => span,
            Entry::End(..) => Span::call_site(),
        }
    }
}

// We do a custom implementation for `Debug` as the default implementation is
// pretty useless.
#[cfg(all(feature = "verbose-trace", not(feature = "all-features")))]
impl<'a> Debug for Cursor<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Print what the cursor is currently looking at.
        // This will look like Cursor("some remaining tokens here")
        f.debug_tuple("Cursor")
            .field(&self.token_stream().to_string())
            .finish()
    }
}
<|endoftext|>


<commit_before>#[derive(Debug,Clone,RustcEncodable)]
pub struct Image {
    pub id: i32,
    pub name: String,
    pub tags: Vec<String>,
    pub got_from: Option<String>,
    pub original_link: Option<String>,
    pub uploader: Option<String>,
    pub rating: Option<char>,
    pub score: i32
}

/// Image voting error
pub enum VoteImageError {
    /// Voting up/down twice
    Already,
    /// No such image
    NoImage
}

/// Partial tags, e.g. `Some*`
#[derive(Debug,Clone)]
enum AnyWith {
    After(String),
    Before(String)
}

/// Ascending/descending sort
#[derive(Debug,Clone)]
enum AscDesc {
    Asc,
    Desc
}

/// Sorting, e.g. `sort:asc/desc:id/score`
#[derive(Debug,Clone)]
enum OrderBy {
    Id,
    Score
}

#[derive(Debug,Clone)]
enum Tag {
    /// Just a tag
    Include(String),
    /// Exlude tag
    Exclude(String),
    /// Rating, e.g. rating:s for safe
    Rating(Vec<String>),
    /// Partial tag
    AnyWith(AnyWith),
    /// from:derpibooru,konachan etc.
    From(Vec<String>),
    /// Uploader search
    Uploader(Vec<String>),
    /// Sorting
    OrderBy(OrderBy, AscDesc)
}

/// `pub` is used to switch DBs, though only postgres works TODO: fix sqlite sometime
pub mod postgres;
mod sqlite;

#[cfg(feature = "sqlite")]
pub use self::sqlite::Db;

#[cfg(feature = "postgresql")]
pub use self::postgres::Db;

/// Tag parsing
fn parse_tag(tag: &str) -> Tag {
    let tag = tag.replace('\'', "''").replace(r'\\', r"\\");
    let all = tag.split(':').collect::<Vec<_>>();
    match all.len() {
        1 => {
            let tag = all[0];
            if tag.starts_with('-') {
                Tag::Exclude(tag[1..].to_string())
            } else if tag.starts_with('*') {
                Tag::AnyWith(AnyWith::After(tag[1..].to_string()))
            } else if tag.ends_with('*') {
                let mut n = tag.to_string();
                n.pop();
                Tag::AnyWith(AnyWith::Before(n))
            } else {
                Tag::Include(tag.to_string())
            }
        },
        2 => {
            let kind = all[0];
            let values = all[1].split(',').map(String::from).collect::<Vec<_>>();
            match kind {
                "rating"    => Tag::Rating(values),
                "from"      => Tag::From(values),
                "uploader"  => Tag::Uploader(values),
                _           => Tag::Include(tag.to_string()) // Probably shouldn't be anything there?
            }
        },
        3 => { // Probably sort, but should check
            if all[0] == "sort" {
                let aod = match all[1] {
                    "asc" => AscDesc::Asc, // - -> +
                    "desc" | _ => AscDesc::Desc, // + -> -
                };
                let by = match all[2] {
                    "score" => OrderBy::Score,
                    "id" | _ => OrderBy::Id,
                };

                Tag::OrderBy(by, aod)
            } else {
                Tag::Include(tag.to_string())
            }
        },
        _   => { // Shouldn't happen
            Tag::Include(tag.to_string())
        }
    }
}

/// Strcture to ease image addition
pub struct ImageBuilder {
    name: String,
    tags: Vec<String>,
    got_from: Option<String>,
    original_link: Option<String>,
    uploader: Option<String>,
    score: i32,
    rating: Option<char>
}

impl ImageBuilder {
    pub fn new(name: &str, tags: &[String]) -> Self {
        ImageBuilder{
            name: name.to_string(),
            tags: tags.to_owned(),
            got_from: None,
            original_link: None,
            uploader: None,
            score: 0,
            rating: None
        }
    }

    pub fn got_from(mut self, got_from: &str) -> Self {
        self.got_from = Some(got_from.to_string());
        self
    }

    pub fn original_link(mut self, original_link: &str) -> Self {
        self.original_link = Some(original_link.to_string());
        self
    }

    pub fn uploader(mut self, uploader: &str) -> Self {
        self.uploader = Some(uploader.to_string());
        self
    }

    pub fn score(mut self, score: i32) -> Self {
        self.score = score;
        self
    }

    pub fn rating(mut self, rating: char) -> Self {
        self.rating = Some(rating);
        self
    }

    pub fn finalize(self) -> Self { self }
}
<commit_msg>whoops<commit_after>#[derive(Debug,Clone,RustcEncodable)]
pub struct Image {
    pub id: i32,
    pub name: String,
    pub tags: Vec<String>,
    pub got_from: Option<String>,
    pub original_link: Option<String>,
    pub uploader: Option<String>,
    pub rating: Option<char>,
    pub score: i32
}

/// Image voting error
pub enum VoteImageError {
    /// Voting up/down twice
    Already,
    /// No such image
    NoImage
}

/// Partial tags, e.g. `Some*`
#[derive(Debug,Clone)]
enum AnyWith {
    After(String),
    Before(String)
}

/// Ascending/descending sort
#[derive(Debug,Clone)]
enum AscDesc {
    Asc,
    Desc
}

/// Sorting, e.g. `sort:asc/desc:id/score`
#[derive(Debug,Clone)]
enum OrderBy {
    Id,
    Score
}

#[derive(Debug,Clone)]
enum Tag {
    /// Just a tag
    Include(String),
    /// Exlude tag
    Exclude(String),
    /// Rating, e.g. rating:s for safe
    Rating(Vec<String>),
    /// Partial tag
    AnyWith(AnyWith),
    /// from:derpibooru,konachan etc.
    From(Vec<String>),
    /// Uploader search
    Uploader(Vec<String>),
    /// Sorting
    OrderBy(OrderBy, AscDesc)
}

/// `pub` is used to switch DBs, though only postgres works TODO: fix sqlite sometime
pub mod postgres;
mod sqlite;

#[cfg(feature = "sqlite")]
pub use self::sqlite::Db;

#[cfg(feature = "postgresql")]
pub use self::postgres::Db;

/// Tag parsing
fn parse_tag(tag: &str) -> Tag {
    let tag = tag.replace('\'', "''").replace('\\', r"\\");
    let all = tag.split(':').collect::<Vec<_>>();
    match all.len() {
        1 => {
            let tag = all[0];
            if tag.starts_with('-') {
                Tag::Exclude(tag[1..].to_string())
            } else if tag.starts_with('*') {
                Tag::AnyWith(AnyWith::After(tag[1..].to_string()))
            } else if tag.ends_with('*') {
                let mut n = tag.to_string();
                n.pop();
                Tag::AnyWith(AnyWith::Before(n))
            } else {
                Tag::Include(tag.to_string())
            }
        },
        2 => {
            let kind = all[0];
            let values = all[1].split(',').map(String::from).collect::<Vec<_>>();
            match kind {
                "rating"    => Tag::Rating(values),
                "from"      => Tag::From(values),
                "uploader"  => Tag::Uploader(values),
                _           => Tag::Include(tag.to_string()) // Probably shouldn't be anything there?
            }
        },
        3 => { // Probably sort, but should check
            if all[0] == "sort" {
                let aod = match all[1] {
                    "asc" => AscDesc::Asc, // - -> +
                    "desc" | _ => AscDesc::Desc, // + -> -
                };
                let by = match all[2] {
                    "score" => OrderBy::Score,
                    "id" | _ => OrderBy::Id,
                };

                Tag::OrderBy(by, aod)
            } else {
                Tag::Include(tag.to_string())
            }
        },
        _   => { // Shouldn't happen
            Tag::Include(tag.to_string())
        }
    }
}

/// Strcture to ease image addition
pub struct ImageBuilder {
    name: String,
    tags: Vec<String>,
    got_from: Option<String>,
    original_link: Option<String>,
    uploader: Option<String>,
    score: i32,
    rating: Option<char>
}

impl ImageBuilder {
    pub fn new(name: &str, tags: &[String]) -> Self {
        ImageBuilder{
            name: name.to_string(),
            tags: tags.to_owned(),
            got_from: None,
            original_link: None,
            uploader: None,
            score: 0,
            rating: None
        }
    }

    pub fn got_from(mut self, got_from: &str) -> Self {
        self.got_from = Some(got_from.to_string());
        self
    }

    pub fn original_link(mut self, original_link: &str) -> Self {
        self.original_link = Some(original_link.to_string());
        self
    }

    pub fn uploader(mut self, uploader: &str) -> Self {
        self.uploader = Some(uploader.to_string());
        self
    }

    pub fn score(mut self, score: i32) -> Self {
        self.score = score;
        self
    }

    pub fn rating(mut self, rating: char) -> Self {
        self.rating = Some(rating);
        self
    }

    pub fn finalize(self) -> Self { self }
}
<|endoftext|>


<commit_before>use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, TextRange, TextUnit};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let mut unescaped = String::with_capacity(inside_str.len());
        let mut error = Ok(());
        rustc_lexer::unescape::unescape_str(inside_str, &mut |_, unescaped_char| {
            match unescaped_char {
                Ok(c) => unescaped.push(c),
                Err(_) => error = Err(()),
            }
        });
        if error.is_err() {
            eprintln!("Error unescaping string");
        } else {
            edit.replace(literal.syntax().text_range(), format!("r#\"{}\"#", unescaped));
        }
    });
    ctx.build()
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(literal.syntax().text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(literal.syntax().text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(literal.syntax().text_range());
        edit.insert(literal.syntax().text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(literal.syntax().text_range().end(), "#");
    });
    ctx.build()
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(literal.syntax().text_range(), format!("r{}", result));
    });
    ctx.build()
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }
}
<commit_msg>feat(assists): Be smart about hashes<commit_after>use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, TextRange, TextUnit};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let mut unescaped = String::with_capacity(inside_str.len());
        let mut error = Ok(());
        rustc_lexer::unescape::unescape_str(inside_str, &mut |_, unescaped_char| {
            match unescaped_char {
                Ok(c) => unescaped.push(c),
                Err(_) => error = Err(()),
            }
        });
        if error.is_err() {
            eprintln!("Error unescaping string");
        } else {
            let mut max_hash_streak = 0;
            unescaped.chars().fold(0, |acc, c| {
                if c == '#' {
                    acc + 1
                } else {
                    if acc > max_hash_streak {
                        max_hash_streak = acc;
                    }
                    0
                }
            });
            let mut hashes = String::with_capacity(max_hash_streak + 1);
            for _ in 0..hashes.capacity() {
                hashes.push('#');
            }
            edit.replace(
                literal.syntax().text_range(),
                format!("r{}\"{}\"{}", hashes, unescaped, hashes),
            );
        }
    });
    ctx.build()
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(literal.syntax().text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(literal.syntax().text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(literal.syntax().text_range());
        edit.insert(literal.syntax().text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(literal.syntax().text_range().end(), "#");
    });
    ctx.build()
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(literal.syntax().text_range(), format!("r{}", result));
    });
    ctx.build()
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random##
string"###;
            }
            "####,
        )
    }
    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "unstable-msc3488")]

use assert_matches::assert_matches;
use assign::assign;
use js_int::uint;
use ruma_common::{
    event_id,
    events::{
        location::{
            AssetContent, AssetType, LocationContent, LocationEventContent, ZoomLevel,
            ZoomLevelError,
        },
        message::MessageContent,
        room::message::{
            InReplyTo, LocationMessageEventContent, MessageType, Relation, RoomMessageEventContent,
        },
        AnyMessageLikeEvent, MessageLikeEvent, MessageLikeUnsigned, OriginalMessageLikeEvent,
    },
    room_id, user_id, MilliSecondsSinceUnixEpoch,
};
use serde_json::{from_value as from_json_value, json, to_value as to_json_value};

#[test]
fn plain_content_serialization() {
    let event_content = LocationEventContent::plain(
        "Alice was at geo:51.5008,0.1247;u=35",
        LocationContent::new("geo:51.5008,0.1247;u=35".to_owned()),
    );

    assert_eq!(
        to_json_value(&event_content).unwrap(),
        json!({
            "org.matrix.msc1767.text": "Alice was at geo:51.5008,0.1247;u=35",
            "m.location": {
                "uri": "geo:51.5008,0.1247;u=35",
            },
        })
    );
}

#[test]
fn event_serialization() {
    let event = OriginalMessageLikeEvent {
        content: assign!(
            LocationEventContent::with_message(
                MessageContent::html(
                    "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                    "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                ),
                assign!(
                    LocationContent::new("geo:51.5008,0.1247;u=35".to_owned()),
                    {
                        description: Some("Alice's whereabouts".into()),
                        zoom_level: Some(ZoomLevel::new(4).unwrap())
                    }
                )
            ),
            {
                ts: Some(MilliSecondsSinceUnixEpoch(uint!(1_636_829_458))),
                relates_to: Some(Relation::Reply {
                    in_reply_to: InReplyTo::new(event_id!("$replyevent:example.com").to_owned()),
                }),
            }
        ),
        event_id: event_id!("$event:notareal.hs").to_owned(),
        sender: user_id!("@user:notareal.hs").to_owned(),
        origin_server_ts: MilliSecondsSinceUnixEpoch(uint!(134_829_848)),
        room_id: room_id!("!roomid:notareal.hs").to_owned(),
        unsigned: MessageLikeUnsigned::default(),
    };

    assert_eq!(
        to_json_value(&event).unwrap(),
        json!({
            "content": {
                "org.matrix.msc1767.message": [
                    {
                        "body": "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                        "mimetype": "text/html",
                    },
                    {
                        "body": "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                        "mimetype": "text/plain",
                    },
                ],
                "m.location": {
                    "uri": "geo:51.5008,0.1247;u=35",
                    "description": "Alice's whereabouts",
                    "zoom_level": 4,
                },
                "m.ts": 1_636_829_458,
                "m.relates_to": {
                    "m.in_reply_to": {
                        "event_id": "$replyevent:example.com",
                    },
                },
            },
            "event_id": "$event:notareal.hs",
            "origin_server_ts": 134_829_848,
            "room_id": "!roomid:notareal.hs",
            "sender": "@user:notareal.hs",
            "type": "m.location",
        })
    );
}

#[test]
fn plain_content_deserialization() {
    let json_data = json!({
        "m.text": "Alice was at geo:51.5008,0.1247;u=35",
        "m.location": {
            "uri": "geo:51.5008,0.1247;u=35",
        },
    });

    assert_matches!(
        from_json_value::<LocationEventContent>(json_data)
            .unwrap(),
        LocationEventContent {
            message,
            location: LocationContent {
                uri,
                description: None,
                zoom_level: None,
                ..
            },
            asset: AssetContent {
                type_: AssetType::Self_,
                ..
            },
            ts: None,
            ..
        }
        if message.find_plain() == Some("Alice was at geo:51.5008,0.1247;u=35")
            && message.find_html().is_none()
            && uri == "geo:51.5008,0.1247;u=35"
    );
}

#[test]
fn zoomlevel_deserialization_pass() {
    let json_data = json!({
        "uri": "geo:51.5008,0.1247;u=35",
        "zoom_level": 16,
    });

    assert_matches!(
        from_json_value::<LocationContent>(json_data).unwrap(),
        LocationContent {
            zoom_level: Some(zoom_level),
            ..
        } if zoom_level.get() == uint!(16)
    );
}

#[test]
fn zoomlevel_deserialization_too_high() {
    let json_data = json!({
        "uri": "geo:51.5008,0.1247;u=35",
        "zoom_level": 30,
    });

    assert_matches!(
        from_json_value::<LocationContent>(json_data),
        Err(err)
            if err.is_data()
            && format!("{err}") == format!("{}", ZoomLevelError::TooHigh)
    );
}

#[test]
fn message_event_deserialization() {
    let json_data = json!({
        "content": {
            "org.matrix.msc1767.message": [
                { "body": "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021" },
            ],
            "m.location": {
                "uri": "geo:51.5008,0.1247;u=35",
                "description": "Alice's whereabouts",
                "zoom_level": 4,
            },
            "m.ts": 1_636_829_458,
            "m.relates_to": {
                "m.in_reply_to": {
                    "event_id": "$replyevent:example.com",
                },
            },
        },
        "event_id": "$event:notareal.hs",
        "origin_server_ts": 134_829_848,
        "room_id": "!roomid:notareal.hs",
        "sender": "@user:notareal.hs",
        "type": "m.location",
    });

    assert_matches!(
        from_json_value::<AnyMessageLikeEvent>(json_data).unwrap(),
        AnyMessageLikeEvent::Location(MessageLikeEvent::Original(OriginalMessageLikeEvent {
            content: LocationEventContent {
                message,
                location: LocationContent {
                    uri,
                    description: Some(description),
                    zoom_level: Some(zoom_level),
                    ..
                },
                asset: AssetContent {
                    type_: AssetType::Self_,
                    ..
                },
                ts: Some(ts),
                ..
            },
            event_id,
            origin_server_ts,
            room_id,
            sender,
            unsigned
        })) if event_id == event_id!("$event:notareal.hs")
            && message.find_plain() == Some("Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021")
            && message.find_html().is_none()
            && uri == "geo:51.5008,0.1247;u=35"
            && description == "Alice's whereabouts"
            && zoom_level.get() == uint!(4)
            && ts == MilliSecondsSinceUnixEpoch(uint!(1_636_829_458))
            && origin_server_ts == MilliSecondsSinceUnixEpoch(uint!(134_829_848))
            && room_id == room_id!("!roomid:notareal.hs")
            && sender == user_id!("@user:notareal.hs")
            && unsigned.is_empty()
    );
}

#[test]
fn room_message_serialization() {
    let message_event_content =
        RoomMessageEventContent::new(MessageType::Location(LocationMessageEventContent::new(
            "Alice was at geo:51.5008,0.1247;u=35".to_owned(),
            "geo:51.5008,0.1247;u=35".to_owned(),
        )));

    assert_eq!(
        to_json_value(&message_event_content).unwrap(),
        json!({
            "body": "Alice was at geo:51.5008,0.1247;u=35",
            "geo_uri": "geo:51.5008,0.1247;u=35",
            "msgtype": "m.location",
            "org.matrix.msc1767.text": "Alice was at geo:51.5008,0.1247;u=35",
            "org.matrix.msc3488.location": {
                "uri": "geo:51.5008,0.1247;u=35",
            },
        })
    );
}

#[test]
fn room_message_stable_deserialization() {
    let json_data = json!({
        "body": "Alice was at geo:51.5008,0.1247;u=35",
        "geo_uri": "geo:51.5008,0.1247;u=35",
        "msgtype": "m.location",
        "m.text": "Alice was at geo:51.5008,0.1247;u=35",
        "m.location": {
            "uri": "geo:51.5008,0.1247;u=35",
        },
    });

    let event_content = from_json_value::<RoomMessageEventContent>(json_data).unwrap();
    assert_matches!(event_content.msgtype, MessageType::Location(_));
    if let MessageType::Location(content) = event_content.msgtype {
        assert_eq!(content.body, "Alice was at geo:51.5008,0.1247;u=35");
        assert_eq!(content.geo_uri, "geo:51.5008,0.1247;u=35");
        let message = content.message.unwrap();
        assert_eq!(message.len(), 1);
        assert_eq!(message[0].body, "Alice was at geo:51.5008,0.1247;u=35");
        assert_eq!(content.location.unwrap().uri, "geo:51.5008,0.1247;u=35");
    }
}

#[test]
fn room_message_unstable_deserialization() {
    let json_data = json!({
        "body": "Alice was at geo:51.5008,0.1247;u=35",
        "geo_uri": "geo:51.5008,0.1247;u=35",
        "msgtype": "m.location",
        "org.matrix.msc1767.text": "Alice was at geo:51.5008,0.1247;u=35",
        "org.matrix.msc3488.location": {
            "uri": "geo:51.5008,0.1247;u=35",
        },
    });

    let event_content = from_json_value::<RoomMessageEventContent>(json_data).unwrap();
    assert_matches!(event_content.msgtype, MessageType::Location(_));
    if let MessageType::Location(content) = event_content.msgtype {
        assert_eq!(content.body, "Alice was at geo:51.5008,0.1247;u=35");
        assert_eq!(content.geo_uri, "geo:51.5008,0.1247;u=35");
        let message = content.message.unwrap();
        assert_eq!(message.len(), 1);
        assert_eq!(message[0].body, "Alice was at geo:51.5008,0.1247;u=35");
        assert_eq!(content.location.unwrap().uri, "geo:51.5008,0.1247;u=35");
    }
}
<commit_msg>events: Clean up location event tests<commit_after>#![cfg(feature = "unstable-msc3488")]

use assert_matches::assert_matches;
use assign::assign;
use js_int::uint;
use ruma_common::{
    event_id,
    events::{
        location::{AssetType, LocationContent, LocationEventContent, ZoomLevel, ZoomLevelError},
        message::MessageContent,
        room::message::{
            InReplyTo, LocationMessageEventContent, MessageType, Relation, RoomMessageEventContent,
        },
        AnyMessageLikeEvent, MessageLikeEvent, MessageLikeUnsigned, OriginalMessageLikeEvent,
    },
    room_id, user_id, MilliSecondsSinceUnixEpoch,
};
use serde_json::{from_value as from_json_value, json, to_value as to_json_value};

#[test]
fn plain_content_serialization() {
    let event_content = LocationEventContent::plain(
        "Alice was at geo:51.5008,0.1247;u=35",
        LocationContent::new("geo:51.5008,0.1247;u=35".to_owned()),
    );

    assert_eq!(
        to_json_value(&event_content).unwrap(),
        json!({
            "org.matrix.msc1767.text": "Alice was at geo:51.5008,0.1247;u=35",
            "m.location": {
                "uri": "geo:51.5008,0.1247;u=35",
            },
        })
    );
}

#[test]
fn event_serialization() {
    let event = OriginalMessageLikeEvent {
        content: assign!(
            LocationEventContent::with_message(
                MessageContent::html(
                    "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                    "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                ),
                assign!(
                    LocationContent::new("geo:51.5008,0.1247;u=35".to_owned()),
                    {
                        description: Some("Alice's whereabouts".into()),
                        zoom_level: Some(ZoomLevel::new(4).unwrap())
                    }
                )
            ),
            {
                ts: Some(MilliSecondsSinceUnixEpoch(uint!(1_636_829_458))),
                relates_to: Some(Relation::Reply {
                    in_reply_to: InReplyTo::new(event_id!("$replyevent:example.com").to_owned()),
                }),
            }
        ),
        event_id: event_id!("$event:notareal.hs").to_owned(),
        sender: user_id!("@user:notareal.hs").to_owned(),
        origin_server_ts: MilliSecondsSinceUnixEpoch(uint!(134_829_848)),
        room_id: room_id!("!roomid:notareal.hs").to_owned(),
        unsigned: MessageLikeUnsigned::default(),
    };

    assert_eq!(
        to_json_value(&event).unwrap(),
        json!({
            "content": {
                "org.matrix.msc1767.message": [
                    {
                        "body": "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                        "mimetype": "text/html",
                    },
                    {
                        "body": "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021",
                        "mimetype": "text/plain",
                    },
                ],
                "m.location": {
                    "uri": "geo:51.5008,0.1247;u=35",
                    "description": "Alice's whereabouts",
                    "zoom_level": 4,
                },
                "m.ts": 1_636_829_458,
                "m.relates_to": {
                    "m.in_reply_to": {
                        "event_id": "$replyevent:example.com",
                    },
                },
            },
            "event_id": "$event:notareal.hs",
            "origin_server_ts": 134_829_848,
            "room_id": "!roomid:notareal.hs",
            "sender": "@user:notareal.hs",
            "type": "m.location",
        })
    );
}

#[test]
fn plain_content_deserialization() {
    let json_data = json!({
        "m.text": "Alice was at geo:51.5008,0.1247;u=35",
        "m.location": {
            "uri": "geo:51.5008,0.1247;u=35",
        },
    });

    let ev = from_json_value::<LocationEventContent>(json_data).unwrap();

    assert_eq!(ev.message.find_plain(), Some("Alice was at geo:51.5008,0.1247;u=35"));
    assert_eq!(ev.message.find_html(), None);
    assert_eq!(ev.location.uri, "geo:51.5008,0.1247;u=35");
    assert_eq!(ev.location.description, None);
    assert_matches!(ev.location.zoom_level, None);
    assert_eq!(ev.asset.type_, AssetType::Self_);
    assert_eq!(ev.ts, None);
}

#[test]
fn zoomlevel_deserialization_pass() {
    let json_data = json!({
        "uri": "geo:51.5008,0.1247;u=35",
        "zoom_level": 16,
    });

    assert_matches!(
        from_json_value::<LocationContent>(json_data).unwrap(),
        LocationContent {
            zoom_level: Some(zoom_level),
            ..
        } if zoom_level.get() == uint!(16)
    );
}

#[test]
fn zoomlevel_deserialization_too_high() {
    let json_data = json!({
        "uri": "geo:51.5008,0.1247;u=35",
        "zoom_level": 30,
    });

    let err = from_json_value::<LocationContent>(json_data).unwrap_err();
    assert!(err.is_data());
    assert_eq!(err.to_string(), ZoomLevelError::TooHigh.to_string());
}

#[test]
fn message_event_deserialization() {
    let json_data = json!({
        "content": {
            "org.matrix.msc1767.message": [
                { "body": "Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021" },
            ],
            "m.location": {
                "uri": "geo:51.5008,0.1247;u=35",
                "description": "Alice's whereabouts",
                "zoom_level": 4,
            },
            "m.ts": 1_636_829_458,
            "m.relates_to": {
                "m.in_reply_to": {
                    "event_id": "$replyevent:example.com",
                },
            },
        },
        "event_id": "$event:notareal.hs",
        "origin_server_ts": 134_829_848,
        "room_id": "!roomid:notareal.hs",
        "sender": "@user:notareal.hs",
        "type": "m.location",
    });

    let ev = from_json_value::<AnyMessageLikeEvent>(json_data).unwrap();

    let ev =
        assert_matches!(ev, AnyMessageLikeEvent::Location(MessageLikeEvent::Original(ev)) => ev);

    assert_eq!(
        ev.content.message.find_plain(),
        Some("Alice was at geo:51.5008,0.1247;u=35 as of Sat Nov 13 18:50:58 2021")
    );
    assert_eq!(ev.content.message.find_html(), None);
    assert_eq!(ev.content.location.uri, "geo:51.5008,0.1247;u=35");
    assert_eq!(ev.content.location.description.as_deref(), Some("Alice's whereabouts"));
    assert_eq!(ev.content.location.zoom_level.unwrap().get(), uint!(4));
    assert_eq!(ev.content.asset.type_, AssetType::Self_);
    assert_eq!(ev.content.ts, Some(MilliSecondsSinceUnixEpoch(uint!(1_636_829_458))));

    assert_eq!(ev.event_id, event_id!("$event:notareal.hs"));
    assert_eq!(ev.origin_server_ts, MilliSecondsSinceUnixEpoch(uint!(134_829_848)));
    assert_eq!(ev.room_id, room_id!("!roomid:notareal.hs"));
    assert_eq!(ev.sender, user_id!("@user:notareal.hs"));
    assert!(ev.unsigned.is_empty());
}

#[test]
fn room_message_serialization() {
    let message_event_content =
        RoomMessageEventContent::new(MessageType::Location(LocationMessageEventContent::new(
            "Alice was at geo:51.5008,0.1247;u=35".to_owned(),
            "geo:51.5008,0.1247;u=35".to_owned(),
        )));

    assert_eq!(
        to_json_value(&message_event_content).unwrap(),
        json!({
            "body": "Alice was at geo:51.5008,0.1247;u=35",
            "geo_uri": "geo:51.5008,0.1247;u=35",
            "msgtype": "m.location",
            "org.matrix.msc1767.text": "Alice was at geo:51.5008,0.1247;u=35",
            "org.matrix.msc3488.location": {
                "uri": "geo:51.5008,0.1247;u=35",
            },
        })
    );
}

#[test]
fn room_message_stable_deserialization() {
    let json_data = json!({
        "body": "Alice was at geo:51.5008,0.1247;u=35",
        "geo_uri": "geo:51.5008,0.1247;u=35",
        "msgtype": "m.location",
        "m.text": "Alice was at geo:51.5008,0.1247;u=35",
        "m.location": {
            "uri": "geo:51.5008,0.1247;u=35",
        },
    });

    let event_content = from_json_value::<RoomMessageEventContent>(json_data).unwrap();
    let content = assert_matches!(event_content.msgtype, MessageType::Location(c) => c);

    assert_eq!(content.body, "Alice was at geo:51.5008,0.1247;u=35");
    assert_eq!(content.geo_uri, "geo:51.5008,0.1247;u=35");
    let message = content.message.unwrap();
    assert_eq!(message.len(), 1);
    assert_eq!(message[0].body, "Alice was at geo:51.5008,0.1247;u=35");
    assert_eq!(content.location.unwrap().uri, "geo:51.5008,0.1247;u=35");
}

#[test]
fn room_message_unstable_deserialization() {
    let json_data = json!({
        "body": "Alice was at geo:51.5008,0.1247;u=35",
        "geo_uri": "geo:51.5008,0.1247;u=35",
        "msgtype": "m.location",
        "org.matrix.msc1767.text": "Alice was at geo:51.5008,0.1247;u=35",
        "org.matrix.msc3488.location": {
            "uri": "geo:51.5008,0.1247;u=35",
        },
    });

    let event_content = from_json_value::<RoomMessageEventContent>(json_data).unwrap();
    let content = assert_matches!(event_content.msgtype, MessageType::Location(c) => c);

    assert_eq!(content.body, "Alice was at geo:51.5008,0.1247;u=35");
    assert_eq!(content.geo_uri, "geo:51.5008,0.1247;u=35");
    let message = content.message.unwrap();
    assert_eq!(message.len(), 1);
    assert_eq!(message[0].body, "Alice was at geo:51.5008,0.1247;u=35");
    assert_eq!(content.location.unwrap().uri, "geo:51.5008,0.1247;u=35");
}
<|endoftext|>


<commit_before>use nom::alphanumeric;

named!(pub symbol<&str, &str>,
    call!(alphanumeric)
);

#[cfg(test)]
mod tests {
    #[test]
    fn test_symbol() {
        //".text"
        //"*"
        //"hello*.o"
        //"\"spaces are ok, just quote the identifier\""
        //"this+is-another*crazy[example]"
    }
}
<commit_msg>add symbols and patterns<commit_after>use nom::{IResult, Needed, ErrorKind};

named!(quoted<&str, &str>, delimited!(
    tag!("\""),
    take_until!("\""),
    tag!("\"")
));

fn simple(input: &str) -> IResult<&str, &str> {
    let mut iter = input.char_indices();
    if let Some((_, c)) = iter.next() {
        // starts with a letter, underscore, or period
        if !(c.is_alphabetic() || c == '_' || c == '.') {
            return IResult::Error(ErrorKind::Char);
        }
    } else {
        return IResult::Incomplete(Needed::Size(1));
    }
    for (i, c) in iter {
        // may include letters, digits, underscores, periods, and hyphens
        if !(c.is_alphanumeric() || c == '_' || c == '.' || c == '-') {
            return IResult::Done(&input[i..], &input[..i]);
        }
    }
    IResult::Done(&input[input.len()..], &input[..])
}

named!(pub symbol<&str, &str>, alt_complete!(
     quoted | simple
));

fn is_pattern(c: char) -> bool {
    c.is_alphanumeric() || "_.$/\\~=+[]*?-!<>^:".contains(c)
}

named!(pub pattern<&str, &str>, alt_complete!(
    quoted | take_while!(is_pattern)
));

#[cfg(test)]
mod tests {
    use idents::symbol;
    use idents::pattern;

    #[test]
    fn test_symbol() {
        assert_done!(symbol(".0"), ".0");
        assert_done!(symbol(".text"), ".text");
        assert_done!(symbol("a-b"), "a-b");
        assert_done!(symbol("\"spaces are ok, just quote the identifier\""),
                     "spaces are ok, just quote the identifier");
    }

    #[test]
    fn test_pattern() {
        assert_done!(pattern("0"), "0");
        assert_done!(pattern(".text"), ".text");
        assert_done!(pattern("hello*.o"), "hello*.o");
        assert_done!(pattern("\"spaces are ok, just quote the identifier\""),
                     "spaces are ok, just quote the identifier");
        assert_done!(pattern("this+is-another*crazy[example]"),
                     "this+is-another*crazy[example]");
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Execution plan for reading Parquet files

use std::fs::File;
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::{fmt, thread};

use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::common;
use crate::execution::physical_plan::{ExecutionPlan, Partition};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::{RecordBatch, RecordBatchReader};
use parquet::file::reader::SerializedFileReader;

use crossbeam::channel::{bounded, Receiver, RecvError, Sender};
use fmt::{Debug, Formatter};
use parquet::arrow::{ArrowReader, ParquetFileArrowReader};

/// Execution plan for scanning a Parquet file
#[derive(Debug)]
pub struct ParquetExec {
    /// Path to directory containing partitioned Parquet files with the same schema
    filenames: Vec<String>,
    /// Schema after projection is applied
    schema: SchemaRef,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
}

impl ParquetExec {
    /// Create a new Parquet reader execution plan
    pub fn try_new(
        path: &str,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Result<Self> {
        let mut filenames: Vec<String> = vec![];
        common::build_file_list(path, &mut filenames, ".parquet")?;
        if filenames.is_empty() {
            Err(ExecutionError::General("No files found".to_string()))
        } else {
            let file = File::open(&filenames[0])?;
            let file_reader = Rc::new(SerializedFileReader::new(file)?);
            let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
            let schema = arrow_reader.get_schema()?;

            let projection = match projection {
                Some(p) => p,
                None => (0..schema.fields().len()).collect(),
            };

            let projected_schema = Schema::new(
                projection
                    .iter()
                    .map(|i| schema.field(*i).clone())
                    .collect(),
            );

            Ok(Self {
                filenames,
                schema: Arc::new(projected_schema),
                projection,
                batch_size,
            })
        }
    }
}

impl ExecutionPlan for ParquetExec {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn partitions(&self) -> Result<Vec<Arc<dyn Partition>>> {
        let partitions = self
            .filenames
            .iter()
            .map(|filename| {
                Arc::new(ParquetPartition::new(
                    &filename,
                    self.projection.clone(),
                    self.schema.clone(),
                    self.batch_size,
                )) as Arc<dyn Partition>
            })
            .collect();
        Ok(partitions)
    }
}

struct ParquetPartition {
    iterator: Arc<Mutex<dyn RecordBatchReader + Send + Sync>>,
}

impl Debug for ParquetPartition {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("ParquetPartition").finish()
    }
}

impl ParquetPartition {
    /// Create a new Parquet partition
    pub fn new(
        filename: &str,
        projection: Vec<usize>,
        schema: SchemaRef,
        batch_size: usize,
    ) -> Self {
        // because the parquet implementation is not thread-safe, it is necessary to execute
        // on a thread and communicate with channels
        let (response_tx, response_rx): (
            Sender<ArrowResult<Option<RecordBatch>>>,
            Receiver<ArrowResult<Option<RecordBatch>>>,
        ) = bounded(2);

        let filename = filename.to_string();

        thread::spawn(move || {
            if let Err(e) = read_file(&filename, projection, batch_size, response_tx) {
                println!("Parquet reader thread terminated due to error: {:?}", e);
            }
        });

        let iterator = Arc::new(Mutex::new(ParquetIterator {
            schema,
            response_rx,
        }));

        Self { iterator }
    }
}

fn send_result(
    response_tx: &Sender<ArrowResult<Option<RecordBatch>>>,
    result: ArrowResult<Option<RecordBatch>>,
) -> Result<()> {
    response_tx
        .send(result)
        .map_err(|e| ExecutionError::ExecutionError(format!("{:?}", e)))?;
    Ok(())
}

fn read_file(
    filename: &str,
    projection: Vec<usize>,
    batch_size: usize,
    response_tx: Sender<ArrowResult<Option<RecordBatch>>>,
) -> Result<()> {
    let file = File::open(&filename)?;
    let file_reader = Rc::new(SerializedFileReader::new(file)?);
    let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
    let mut batch_reader =
        arrow_reader.get_record_reader_by_columns(projection.clone(), batch_size)?;
    loop {
        match batch_reader.next_batch() {
            Ok(Some(batch)) => send_result(&response_tx, Ok(Some(batch)))?,
            Ok(None) => {
                // finished reading file
                send_result(&response_tx, Ok(None))?;
                break;
            }
            Err(e) => {
                let err_msg = format!("Error reading batch from {}: {:?}", filename, e);
                // send error to operator
                send_result(
                    &response_tx,
                    Err(ArrowError::ParquetError(err_msg.clone())),
                )?;
                // terminate thread with error
                return Err(ExecutionError::ExecutionError(err_msg));
            }
        }
    }
    Ok(())
}

impl Partition for ParquetPartition {
    fn execute(&self) -> Result<Arc<Mutex<dyn RecordBatchReader + Send + Sync>>> {
        Ok(self.iterator.clone())
    }
}

struct ParquetIterator {
    schema: SchemaRef,
    response_rx: Receiver<ArrowResult<Option<RecordBatch>>>,
}

impl RecordBatchReader for ParquetIterator {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn next_batch(&mut self) -> ArrowResult<Option<RecordBatch>> {
        match self.response_rx.recv() {
            Ok(batch) => batch,
            // RecvError means receiver has exited and closed the channel
            Err(RecvError) => Ok(None),
        }
    }
}

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

    #[test]
    fn test() -> Result<()> {
        let testdata =
            env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
        let filename = format!("{}/alltypes_plain.parquet", testdata);
        let parquet_exec = ParquetExec::try_new(&filename, Some(vec![0, 1, 2]), 1024)?;
        let partitions = parquet_exec.partitions()?;
        assert_eq!(partitions.len(), 1);

        let results = partitions[0].execute()?;
        let mut results = results.lock().unwrap();
        let batch = results.next_batch()?.unwrap();

        assert_eq!(8, batch.num_rows());
        assert_eq!(3, batch.num_columns());

        let schema = batch.schema();
        let field_names: Vec<&str> =
            schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        Ok(())
    }
}
<commit_msg>ARROW-9726: [Rust] [DataFusion] Do not create parquet reader thread until execute is called<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Execution plan for reading Parquet files

use std::fs::File;
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::{fmt, thread};

use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::common;
use crate::execution::physical_plan::{ExecutionPlan, Partition};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::{RecordBatch, RecordBatchReader};
use parquet::file::reader::SerializedFileReader;

use crossbeam::channel::{bounded, Receiver, RecvError, Sender};
use fmt::{Debug, Formatter};
use parquet::arrow::{ArrowReader, ParquetFileArrowReader};

/// Execution plan for scanning a Parquet file
#[derive(Debug)]
pub struct ParquetExec {
    /// Path to directory containing partitioned Parquet files with the same schema
    filenames: Vec<String>,
    /// Schema after projection is applied
    schema: SchemaRef,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
}

impl ParquetExec {
    /// Create a new Parquet reader execution plan
    pub fn try_new(
        path: &str,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Result<Self> {
        let mut filenames: Vec<String> = vec![];
        common::build_file_list(path, &mut filenames, ".parquet")?;
        if filenames.is_empty() {
            Err(ExecutionError::General("No files found".to_string()))
        } else {
            let file = File::open(&filenames[0])?;
            let file_reader = Rc::new(SerializedFileReader::new(file)?);
            let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
            let schema = arrow_reader.get_schema()?;

            let projection = match projection {
                Some(p) => p,
                None => (0..schema.fields().len()).collect(),
            };

            let projected_schema = Schema::new(
                projection
                    .iter()
                    .map(|i| schema.field(*i).clone())
                    .collect(),
            );

            Ok(Self {
                filenames,
                schema: Arc::new(projected_schema),
                projection,
                batch_size,
            })
        }
    }
}

impl ExecutionPlan for ParquetExec {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn partitions(&self) -> Result<Vec<Arc<dyn Partition>>> {
        let partitions = self
            .filenames
            .iter()
            .map(|filename| {
                Arc::new(ParquetPartition::new(
                    &filename,
                    self.projection.clone(),
                    self.schema.clone(),
                    self.batch_size,
                )) as Arc<dyn Partition>
            })
            .collect();
        Ok(partitions)
    }
}

struct ParquetPartition {
    /// Filename for this partition
    filename: String,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
    schema: SchemaRef,
}

impl Debug for ParquetPartition {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("ParquetPartition").finish()
    }
}

impl ParquetPartition {
    /// Create a new Parquet partition
    pub fn new(
        filename: &str,
        projection: Vec<usize>,
        schema: SchemaRef,
        batch_size: usize,
    ) -> Self {
        Self {
            filename: filename.to_owned(),
            projection,
            schema,
            batch_size,
        }
    }
}

fn send_result(
    response_tx: &Sender<ArrowResult<Option<RecordBatch>>>,
    result: ArrowResult<Option<RecordBatch>>,
) -> Result<()> {
    response_tx
        .send(result)
        .map_err(|e| ExecutionError::ExecutionError(e.to_string()))?;
    Ok(())
}

fn read_file(
    filename: &str,
    projection: Vec<usize>,
    batch_size: usize,
    response_tx: Sender<ArrowResult<Option<RecordBatch>>>,
) -> Result<()> {
    let file = File::open(&filename)?;
    let file_reader = Rc::new(SerializedFileReader::new(file)?);
    let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
    let mut batch_reader =
        arrow_reader.get_record_reader_by_columns(projection.clone(), batch_size)?;
    loop {
        match batch_reader.next_batch() {
            Ok(Some(batch)) => send_result(&response_tx, Ok(Some(batch)))?,
            Ok(None) => {
                // finished reading file
                send_result(&response_tx, Ok(None))?;
                break;
            }
            Err(e) => {
                let err_msg =
                    format!("Error reading batch from {}: {}", filename, e.to_string());
                // send error to operator
                send_result(
                    &response_tx,
                    Err(ArrowError::ParquetError(err_msg.clone())),
                )?;
                // terminate thread with error
                return Err(ExecutionError::ExecutionError(err_msg));
            }
        }
    }
    Ok(())
}

impl Partition for ParquetPartition {
    fn execute(&self) -> Result<Arc<Mutex<dyn RecordBatchReader + Send + Sync>>> {
        // because the parquet implementation is not thread-safe, it is necessary to execute
        // on a thread and communicate with channels
        let (response_tx, response_rx): (
            Sender<ArrowResult<Option<RecordBatch>>>,
            Receiver<ArrowResult<Option<RecordBatch>>>,
        ) = bounded(2);

        let filename = self.filename.clone();
        let projection = self.projection.clone();
        let batch_size = self.batch_size;

        thread::spawn(move || {
            if let Err(e) = read_file(&filename, projection, batch_size, response_tx) {
                println!("Parquet reader thread terminated due to error: {:?}", e);
            }
        });

        let iterator = Arc::new(Mutex::new(ParquetIterator {
            schema: self.schema.clone(),
            response_rx,
        }));

        Ok(iterator)
    }
}

struct ParquetIterator {
    schema: SchemaRef,
    response_rx: Receiver<ArrowResult<Option<RecordBatch>>>,
}

impl RecordBatchReader for ParquetIterator {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn next_batch(&mut self) -> ArrowResult<Option<RecordBatch>> {
        match self.response_rx.recv() {
            Ok(batch) => batch,
            // RecvError means receiver has exited and closed the channel
            Err(RecvError) => Ok(None),
        }
    }
}

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

    #[test]
    fn test() -> Result<()> {
        let testdata =
            env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
        let filename = format!("{}/alltypes_plain.parquet", testdata);
        let parquet_exec = ParquetExec::try_new(&filename, Some(vec![0, 1, 2]), 1024)?;
        let partitions = parquet_exec.partitions()?;
        assert_eq!(partitions.len(), 1);

        let results = partitions[0].execute()?;
        let mut results = results.lock().unwrap();
        let batch = results.next_batch()?.unwrap();

        assert_eq!(8, batch.num_rows());
        assert_eq!(3, batch.num_columns());

        let schema = batch.schema();
        let field_names: Vec<&str> =
            schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        Ok(())
    }
}
<|endoftext|>


<commit_before>#![feature(std_misc)]

use std::ascii::AsciiExt;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::thread;
use std::sync::mpsc::channel;

/// Compute the frequency of each letter (technically of each unicode codepoint) using the given
/// number of worker threads.
pub fn frequency(texts: &[&str], num_workers: usize) -> HashMap<char, usize> {
    assert!(num_workers > 0);
    let part_size_floor = texts.len() / num_workers;
    let rem = texts.len() % num_workers;
    let part_size = if rem > 0 { part_size_floor + 1 } else { part_size_floor };
    let mut parts: Vec<Vec<String>> = Vec::with_capacity(part_size);
    for _ in 0 .. num_workers {
        parts.push(Vec::with_capacity(part_size));
    }
    let mut i = 0;
    for line in texts.iter() {
        // We'll need to clone those strings in order to satisfy some lifetime guarantees. Basically
        // it's hard for the system to be sure that the threads spawned don't outlive the strings.
        parts[i].push(line.to_string());
        i = (i + 1) % num_workers;
    }

    let (tx, rx) = channel();

    for part in parts {
        let tx = tx.clone();
        thread::spawn(move || {
            tx.send(count(part)).unwrap();
        });
    }

    let mut results: HashMap<char, usize> = HashMap::new();
    for _ in 0..num_workers {
        let part_results = rx.recv().unwrap();
        for (c, n) in part_results.into_iter() {
            match results.entry(c) {
                Entry::Vacant(view) => { view.insert(n); },
                Entry::Occupied(mut view) => {
                    *view.get_mut() += n;
                }
            }
        }
    }
    results
}

fn count(lines: Vec<String>) -> HashMap<char, usize> {
    let mut results: HashMap<char, usize> = HashMap::new();
    for line in lines.iter() {
        for c in line.chars() {
            if c.is_alphabetic() {
                match results.entry(c.to_lowercase().next().unwrap()) {
                    Entry::Vacant(view) => { view.insert(1); },
                    Entry::Occupied(mut view) => {
                        *view.get_mut() += 1;
                    }
                }
            }
        }
    }
    results
}
<commit_msg>Remove feature flag and unused import<commit_after>use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::thread;
use std::sync::mpsc::channel;

/// Compute the frequency of each letter (technically of each unicode codepoint) using the given
/// number of worker threads.
pub fn frequency(texts: &[&str], num_workers: usize) -> HashMap<char, usize> {
    assert!(num_workers > 0);
    let part_size_floor = texts.len() / num_workers;
    let rem = texts.len() % num_workers;
    let part_size = if rem > 0 { part_size_floor + 1 } else { part_size_floor };
    let mut parts: Vec<Vec<String>> = Vec::with_capacity(part_size);
    for _ in 0 .. num_workers {
        parts.push(Vec::with_capacity(part_size));
    }
    let mut i = 0;
    for line in texts.iter() {
        // We'll need to clone those strings in order to satisfy some lifetime guarantees. Basically
        // it's hard for the system to be sure that the threads spawned don't outlive the strings.
        parts[i].push(line.to_string());
        i = (i + 1) % num_workers;
    }

    let (tx, rx) = channel();

    for part in parts {
        let tx = tx.clone();
        thread::spawn(move || {
            tx.send(count(part)).unwrap();
        });
    }

    let mut results: HashMap<char, usize> = HashMap::new();
    for _ in 0..num_workers {
        let part_results = rx.recv().unwrap();
        for (c, n) in part_results.into_iter() {
            match results.entry(c) {
                Entry::Vacant(view) => { view.insert(n); },
                Entry::Occupied(mut view) => {
                    *view.get_mut() += n;
                }
            }
        }
    }
    results
}

fn count(lines: Vec<String>) -> HashMap<char, usize> {
    let mut results: HashMap<char, usize> = HashMap::new();
    for line in lines.iter() {
        for c in line.chars() {
            if c.is_alphabetic() {
                match results.entry(c.to_lowercase().next().unwrap()) {
                    Entry::Vacant(view) => { view.insert(1); },
                    Entry::Occupied(mut view) => {
                        *view.get_mut() += 1;
                    }
                }
            }
        }
    }
    results
}
<|endoftext|>