<commit_before>use futures::{Async, Future, Poll, Stream};
use h2::Reason;
use h2::server::{Builder, Connection, Handshake, SendResponse};
use tokio_io::{AsyncRead, AsyncWrite};

use ::headers::content_length_parse_all;
use ::body::Payload;
use body::internal::FullDataArg;
use ::common::exec::H2Exec;
use ::headers;
use ::service::Service;
use ::proto::Dispatched;
use super::{PipeToSendStream, SendBuf};

use ::{Body, Response};

pub(crate) struct Server<T, S, B, E>
where
    S: Service,
    B: Payload,
{
    exec: E,
    service: S,
    state: State<T, B>,
}

enum State<T, B>
where
    B: Payload,
{
    Handshaking(Handshake<T, SendBuf<B::Data>>),
    Serving(Serving<T, B>),
    Closed,
}

struct Serving<T, B>
where
    B: Payload,
{
    conn: Connection<T, SendBuf<B::Data>>,
}

impl<T, S, B, E> Server<T, S, B, E>
where
    T: AsyncRead + AsyncWrite,
    S: Service<ReqBody=Body, ResBody=B>,
    S::Error: Into<Box<::std::error::Error + Send + Sync>>,
    //S::Future: Send + 'static,
    B: Payload,
    E: H2Exec<S::Future, B>,
{
    pub(crate) fn new(io: T, service: S, exec: E) -> Server<T, S, B, E> {
        let handshake = Builder::new()
            .handshake(io);
        Server {
            exec,
            state: State::Handshaking(handshake),
            service,
        }
    }

    pub fn graceful_shutdown(&mut self) {
        trace!("graceful_shutdown");
        match self.state {
            State::Handshaking(..) => {
                // fall-through, to replace state with Closed
            },
            State::Serving(ref mut srv) => {
                srv.conn.graceful_shutdown();
                return;
            },
            State::Closed => {
                return;
            }
        }
        self.state = State::Closed;
    }
}

impl<T, S, B, E> Future for Server<T, S, B, E>
where
    T: AsyncRead + AsyncWrite,
    S: Service<ReqBody=Body, ResBody=B>,
    S::Error: Into<Box<::std::error::Error + Send + Sync>>,
    //S::Future: Send + 'static,
    B: Payload,
    E: H2Exec<S::Future, B>,
{
    type Item = Dispatched;
    type Error = ::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let next = match self.state {
                State::Handshaking(ref mut h) => {
                    let conn = try_ready!(h.poll().map_err(::Error::new_h2));
                    State::Serving(Serving {
                        conn: conn,
                    })
                },
                State::Serving(ref mut srv) => {
                    try_ready!(srv.poll_server(&mut self.service, &self.exec));
                    return Ok(Async::Ready(Dispatched::Shutdown));
                }
                State::Closed => {
                    // graceful_shutdown was called before handshaking finished,
                    // nothing to do here...
                    return Ok(Async::Ready(Dispatched::Shutdown));
                }
            };
            self.state = next;
        }
    }
}

impl<T, B> Serving<T, B>
where
    T: AsyncRead + AsyncWrite,
    B: Payload,
{
    fn poll_server<S, E>(&mut self, service: &mut S, exec: &E) -> Poll<(), ::Error>
    where
        S: Service<
            ReqBody=Body,
            ResBody=B,
        >,
        S::Error: Into<Box<::std::error::Error + Send + Sync>>,
        E: H2Exec<S::Future, B>,
    {
        loop {
            // At first, polls the readiness of supplied service.
            match service.poll_ready() {
                Ok(Async::Ready(())) => (),
                Ok(Async::NotReady) => {
                    // use `poll_close` instead of `poll`, in order to avoid accepting a request.
                    try_ready!(self.conn.poll_close().map_err(::Error::new_h2));
                    trace!("incoming connection complete");
                    return Ok(Async::Ready(()));
                }
                Err(err) => {
                    trace!("service closed");
                    return Err(::Error::new_user_service(err));
                }
            }

            // When the service is ready, accepts an incoming request.
            if let Some((req, respond)) = try_ready!(self.conn.poll().map_err(::Error::new_h2)) {
                trace!("incoming request");
                let content_length = content_length_parse_all(req.headers());
                let req = req.map(|stream| {
                    ::Body::h2(stream, content_length)
                });
                let fut = H2Stream::new(service.call(req), respond);
                exec.execute_h2stream(fut)?;
            } else {
                // no more incoming streams...
                trace!("incoming connection complete");
                return Ok(Async::Ready(()))
            }
        }
    }
}

#[allow(missing_debug_implementations)]
pub struct H2Stream<F, B>
where
    B: Payload,
{
    reply: SendResponse<SendBuf<B::Data>>,
    state: H2StreamState<F, B>,
}

enum H2StreamState<F, B>
where
    B: Payload,
{
    Service(F),
    Body(PipeToSendStream<B>),
}

impl<F, B> H2Stream<F, B>
where
    F: Future<Item=Response<B>>,
    F::Error: Into<Box<::std::error::Error + Send + Sync>>,
    B: Payload,
{
    fn new(fut: F, respond: SendResponse<SendBuf<B::Data>>) -> H2Stream<F, B> {
        H2Stream {
            reply: respond,
            state: H2StreamState::Service(fut),
        }
    }

    fn poll2(&mut self) -> Poll<(), ::Error> {
        loop {
            let next = match self.state {
                H2StreamState::Service(ref mut h) => {
                    let res = match h.poll() {
                        Ok(Async::Ready(r)) => r,
                        Ok(Async::NotReady) => {
                            // Body is not yet ready, so we want to check if the client has sent a
                            // RST_STREAM frame which would cancel the current request.
                            if let Async::Ready(reason) =
                                self.reply.poll_reset().map_err(|e| ::Error::new_h2(e))?
                            {
                                debug!("stream received RST_STREAM: {:?}", reason);
                                return Err(::Error::new_h2(reason.into()));
                            }
                            return Ok(Async::NotReady);
                        }
                        Err(e) => return Err(::Error::new_user_service(e)),
                    };

                    let (head, mut body) = res.into_parts();
                    let mut res = ::http::Response::from_parts(head, ());
                    super::strip_connection_headers(res.headers_mut(), false);

                    // set Date header if it isn't already set...
                    res
                        .headers_mut()
                        .entry(::http::header::DATE)
                        .expect("DATE is a valid HeaderName")
                        .or_insert_with(::proto::h1::date::update_and_header_value);

                    macro_rules! reply {
                        ($eos:expr) => ({
                            match self.reply.send_response(res, $eos) {
                                Ok(tx) => tx,
                                Err(e) => {
                                    trace!("send response error: {}", e);
                                    self.reply.send_reset(Reason::INTERNAL_ERROR);
                                    return Err(::Error::new_h2(e));
                                }
                            }
                        })
                    }

                    // automatically set Content-Length from body...
                    if let Some(len) = body.content_length() {
                        headers::set_content_length_if_missing(res.headers_mut(), len);
                    }

                    if let Some(full) = body.__hyper_full_data(FullDataArg(())).0 {
                        let mut body_tx = reply!(false);
                        let buf = SendBuf(Some(full));
                        body_tx
                            .send_data(buf, true)
                            .map_err(::Error::new_body_write)?;
                        return Ok(Async::Ready(()));
                    }

                    if !body.is_end_stream() {
                        let body_tx = reply!(false);
                        H2StreamState::Body(PipeToSendStream::new(body, body_tx))
                    } else {
                        reply!(true);
                        return Ok(Async::Ready(()));
                    }
                },
                H2StreamState::Body(ref mut pipe) => {
                    return pipe.poll();
                }
            };
            self.state = next;
        }
    }
}

impl<F, B> Future for H2Stream<F, B>
where
    F: Future<Item=Response<B>>,
    F::Error: Into<Box<::std::error::Error + Send + Sync>>,
    B: Payload,
{
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.poll2()
            .map_err(|e| debug!("stream error: {}", e))
    }
}

<commit_msg>fix(http2): send INTERNAL_ERROR when user's Service errors<commit_after>use futures::{Async, Future, Poll, Stream};
use h2::Reason;
use h2::server::{Builder, Connection, Handshake, SendResponse};
use tokio_io::{AsyncRead, AsyncWrite};

use ::headers::content_length_parse_all;
use ::body::Payload;
use body::internal::FullDataArg;
use ::common::exec::H2Exec;
use ::headers;
use ::service::Service;
use ::proto::Dispatched;
use super::{PipeToSendStream, SendBuf};

use ::{Body, Response};

pub(crate) struct Server<T, S, B, E>
where
    S: Service,
    B: Payload,
{
    exec: E,
    service: S,
    state: State<T, B>,
}

enum State<T, B>
where
    B: Payload,
{
    Handshaking(Handshake<T, SendBuf<B::Data>>),
    Serving(Serving<T, B>),
    Closed,
}

struct Serving<T, B>
where
    B: Payload,
{
    conn: Connection<T, SendBuf<B::Data>>,
}

impl<T, S, B, E> Server<T, S, B, E>
where
    T: AsyncRead + AsyncWrite,
    S: Service<ReqBody=Body, ResBody=B>,
    S::Error: Into<Box<::std::error::Error + Send + Sync>>,
    //S::Future: Send + 'static,
    B: Payload,
    E: H2Exec<S::Future, B>,
{
    pub(crate) fn new(io: T, service: S, exec: E) -> Server<T, S, B, E> {
        let handshake = Builder::new()
            .handshake(io);
        Server {
            exec,
            state: State::Handshaking(handshake),
            service,
        }
    }

    pub fn graceful_shutdown(&mut self) {
        trace!("graceful_shutdown");
        match self.state {
            State::Handshaking(..) => {
                // fall-through, to replace state with Closed
            },
            State::Serving(ref mut srv) => {
                srv.conn.graceful_shutdown();
                return;
            },
            State::Closed => {
                return;
            }
        }
        self.state = State::Closed;
    }
}

impl<T, S, B, E> Future for Server<T, S, B, E>
where
    T: AsyncRead + AsyncWrite,
    S: Service<ReqBody=Body, ResBody=B>,
    S::Error: Into<Box<::std::error::Error + Send + Sync>>,
    //S::Future: Send + 'static,
    B: Payload,
    E: H2Exec<S::Future, B>,
{
    type Item = Dispatched;
    type Error = ::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let next = match self.state {
                State::Handshaking(ref mut h) => {
                    let conn = try_ready!(h.poll().map_err(::Error::new_h2));
                    State::Serving(Serving {
                        conn: conn,
                    })
                },
                State::Serving(ref mut srv) => {
                    try_ready!(srv.poll_server(&mut self.service, &self.exec));
                    return Ok(Async::Ready(Dispatched::Shutdown));
                }
                State::Closed => {
                    // graceful_shutdown was called before handshaking finished,
                    // nothing to do here...
                    return Ok(Async::Ready(Dispatched::Shutdown));
                }
            };
            self.state = next;
        }
    }
}

impl<T, B> Serving<T, B>
where
    T: AsyncRead + AsyncWrite,
    B: Payload,
{
    fn poll_server<S, E>(&mut self, service: &mut S, exec: &E) -> Poll<(), ::Error>
    where
        S: Service<
            ReqBody=Body,
            ResBody=B,
        >,
        S::Error: Into<Box<::std::error::Error + Send + Sync>>,
        E: H2Exec<S::Future, B>,
    {
        loop {
            // At first, polls the readiness of supplied service.
            match service.poll_ready() {
                Ok(Async::Ready(())) => (),
                Ok(Async::NotReady) => {
                    // use `poll_close` instead of `poll`, in order to avoid accepting a request.
                    try_ready!(self.conn.poll_close().map_err(::Error::new_h2));
                    trace!("incoming connection complete");
                    return Ok(Async::Ready(()));
                }
                Err(err) => {
                    trace!("service closed");
                    return Err(::Error::new_user_service(err));
                }
            }

            // When the service is ready, accepts an incoming request.
            if let Some((req, respond)) = try_ready!(self.conn.poll().map_err(::Error::new_h2)) {
                trace!("incoming request");
                let content_length = content_length_parse_all(req.headers());
                let req = req.map(|stream| {
                    ::Body::h2(stream, content_length)
                });
                let fut = H2Stream::new(service.call(req), respond);
                exec.execute_h2stream(fut)?;
            } else {
                // no more incoming streams...
                trace!("incoming connection complete");
                return Ok(Async::Ready(()))
            }
        }
    }
}

#[allow(missing_debug_implementations)]
pub struct H2Stream<F, B>
where
    B: Payload,
{
    reply: SendResponse<SendBuf<B::Data>>,
    state: H2StreamState<F, B>,
}

enum H2StreamState<F, B>
where
    B: Payload,
{
    Service(F),
    Body(PipeToSendStream<B>),
}

impl<F, B> H2Stream<F, B>
where
    F: Future<Item=Response<B>>,
    F::Error: Into<Box<::std::error::Error + Send + Sync>>,
    B: Payload,
{
    fn new(fut: F, respond: SendResponse<SendBuf<B::Data>>) -> H2Stream<F, B> {
        H2Stream {
            reply: respond,
            state: H2StreamState::Service(fut),
        }
    }

    fn poll2(&mut self) -> Poll<(), ::Error> {
        loop {
            let next = match self.state {
                H2StreamState::Service(ref mut h) => {
                    let res = match h.poll() {
                        Ok(Async::Ready(r)) => r,
                        Ok(Async::NotReady) => {
                            // Body is not yet ready, so we want to check if the client has sent a
                            // RST_STREAM frame which would cancel the current request.
                            if let Async::Ready(reason) =
                                self.reply.poll_reset().map_err(|e| ::Error::new_h2(e))?
                            {
                                debug!("stream received RST_STREAM: {:?}", reason);
                                return Err(::Error::new_h2(reason.into()));
                            }
                            return Ok(Async::NotReady);
                        }
                        Err(e) => {
                            let err = ::Error::new_user_service(e);
                            warn!("http2 service errored: {}", err);
                            self.reply.send_reset(Reason::INTERNAL_ERROR);
                            return Err(err);
                        },
                    };

                    let (head, mut body) = res.into_parts();
                    let mut res = ::http::Response::from_parts(head, ());
                    super::strip_connection_headers(res.headers_mut(), false);

                    // set Date header if it isn't already set...
                    res
                        .headers_mut()
                        .entry(::http::header::DATE)
                        .expect("DATE is a valid HeaderName")
                        .or_insert_with(::proto::h1::date::update_and_header_value);

                    macro_rules! reply {
                        ($eos:expr) => ({
                            match self.reply.send_response(res, $eos) {
                                Ok(tx) => tx,
                                Err(e) => {
                                    trace!("send response error: {}", e);
                                    self.reply.send_reset(Reason::INTERNAL_ERROR);
                                    return Err(::Error::new_h2(e));
                                }
                            }
                        })
                    }

                    // automatically set Content-Length from body...
                    if let Some(len) = body.content_length() {
                        headers::set_content_length_if_missing(res.headers_mut(), len);
                    }

                    if let Some(full) = body.__hyper_full_data(FullDataArg(())).0 {
                        let mut body_tx = reply!(false);
                        let buf = SendBuf(Some(full));
                        body_tx
                            .send_data(buf, true)
                            .map_err(::Error::new_body_write)?;
                        return Ok(Async::Ready(()));
                    }

                    if !body.is_end_stream() {
                        let body_tx = reply!(false);
                        H2StreamState::Body(PipeToSendStream::new(body, body_tx))
                    } else {
                        reply!(true);
                        return Ok(Async::Ready(()));
                    }
                },
                H2StreamState::Body(ref mut pipe) => {
                    return pipe.poll();
                }
            };
            self.state = next;
        }
    }
}

impl<F, B> Future for H2Stream<F, B>
where
    F: Future<Item=Response<B>>,
    F::Error: Into<Box<::std::error::Error + Send + Sync>>,
    B: Payload,
{
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.poll2()
            .map_err(|e| debug!("stream error: {}", e))
    }
}

<|endoftext|>


<commit_before>use std::ascii::AsciiExt;
use std::collections::HashMap;
use time::Timespec;

use conduit::{Request, Response};
use conduit_router::RequestParams;
use diesel::pg::PgConnection;
use diesel::prelude::*;
use diesel;
use pg::GenericConnection;
use pg::rows::Row;

use {Model, Crate};
use db::RequestTransaction;
use schema::*;
use util::{RequestUtils, CargoResult, ChainError, internal};

#[derive(Clone, Identifiable, Queryable)]
pub struct Keyword {
    pub id: i32,
    pub keyword: String,
    pub crates_cnt: i32,
    pub created_at: Timespec,
}

#[derive(Associations, Insertable, Identifiable)]
#[belongs_to(Keyword)]
#[belongs_to(Crate)]
#[table_name="crates_keywords"]
#[primary_key(crate_id, keyword_id)]
pub struct CrateKeyword {
    crate_id: i32,
    keyword_id: i32,
}

#[derive(RustcEncodable, RustcDecodable)]
pub struct EncodableKeyword {
    pub id: String,
    pub keyword: String,
    pub created_at: String,
    pub crates_cnt: i32,
}

impl Keyword {
    pub fn find_by_keyword(conn: &PgConnection, name: &str)
                           -> QueryResult<Keyword> {
        keywords::table
            .filter(keywords::keyword.eq(::lower(name)))
            .first(&*conn)
    }

    pub fn find_or_create_all(conn: &PgConnection, names: &[&str]) -> QueryResult<Vec<Keyword>> {
        use diesel::pg::upsert::*;
        use diesel::expression::dsl::any;

        #[derive(Insertable)]
        #[table_name="keywords"]
        struct NewKeyword<'a> {
            keyword: &'a str,
        }

        let (lowercase_names, new_keywords): (Vec<_>, Vec<_>) = names.iter()
            .map(|s| (s.to_lowercase(), NewKeyword { keyword: *s }))
            .unzip();

        // https://github.com/diesel-rs/diesel/issues/797
        if !new_keywords.is_empty() {
            diesel::insert(&new_keywords.on_conflict_do_nothing()).into(keywords::table)
                .execute(conn)?;
        }
        keywords::table.filter(::lower(keywords::keyword).eq(any(lowercase_names)))
            .load(conn)
    }

    pub fn find_or_insert(conn: &GenericConnection, name: &str)
                          -> CargoResult<Keyword> {
        // TODO: racy (the select then insert is not atomic)
        let stmt = conn.prepare("SELECT * FROM keywords
                                      WHERE keyword = LOWER($1)")?;
        for row in stmt.query(&[&name])?.iter() {
            return Ok(Model::from_row(&row))
        }

        let stmt = conn.prepare("INSERT INTO keywords (keyword) VALUES (LOWER($1))
                                      RETURNING *")?;
        let rows = stmt.query(&[&name])?;
        Ok(Model::from_row(&rows.iter().next().chain_error(|| {
            internal("no version returned")
        })?))
    }

    pub fn valid_name(name: &str) -> bool {
        if name.is_empty() { return false }
        name.chars().next().unwrap().is_alphanumeric() &&
            name.chars().all(|c| c.is_alphanumeric() || c == '_' || c == '-') &&
            name.chars().all(|c| c.is_ascii())
    }

    pub fn encodable(self) -> EncodableKeyword {
        let Keyword { crates_cnt, keyword, created_at, .. } = self;
        EncodableKeyword {
            id: keyword.clone(),
            created_at: ::encode_time(created_at),
            crates_cnt: crates_cnt,
            keyword: keyword,
        }
    }

    pub fn update_crate(conn: &PgConnection,
                        krate: &Crate,
                        keywords: &[&str]) -> QueryResult<()> {
        conn.transaction(|| {
            let keywords = Keyword::find_or_create_all(conn, keywords)?;
            diesel::delete(CrateKeyword::belonging_to(krate))
                .execute(conn)?;
            let crate_keywords = keywords.into_iter().map(|kw| {
                CrateKeyword { crate_id: krate.id, keyword_id: kw.id }
            }).collect::<Vec<_>>();
            diesel::insert(&crate_keywords).into(crates_keywords::table)
                .execute(conn)?;
            Ok(())
        })
    }

    pub fn update_crate_old(conn: &GenericConnection,
                        krate: &Crate,
                        keywords: &[String]) -> CargoResult<()> {
        let old_kws = krate.keywords(conn)?;
        let old_kws = old_kws.iter().map(|kw| {
            (&kw.keyword[..], kw)
        }).collect::<HashMap<_, _>>();
        let new_kws = keywords.iter().map(|k| {
            let kw = Keyword::find_or_insert(conn, k)?;
            Ok((k.as_str(), kw))
        }).collect::<CargoResult<HashMap<_, _>>>()?;

        let to_rm = old_kws.iter().filter(|&(kw, _)| {
            !new_kws.contains_key(kw)
        }).map(|(_, v)| v.id).collect::<Vec<_>>();
        let to_add = new_kws.iter().filter(|&(kw, _)| {
            !old_kws.contains_key(kw)
        }).map(|(_, v)| v.id).collect::<Vec<_>>();

        if !to_rm.is_empty() {
            conn.execute("DELETE FROM crates_keywords
                                WHERE keyword_id = ANY($1)
                                  AND crate_id = $2",
                         &[&to_rm, &krate.id])?;
        }

        if !to_add.is_empty() {
            let insert = to_add.iter().map(|id| {
                let crate_id: i32 = krate.id;
                let id: i32 = *id;
                format!("({}, {})", crate_id,  id)
            }).collect::<Vec<_>>().join(", ");
            conn.execute(&format!("INSERT INTO crates_keywords
                                        (crate_id, keyword_id) VALUES {}",
                                  insert),
                         &[])?;
        }

        Ok(())
    }
}

impl Model for Keyword {
    fn from_row(row: &Row) -> Keyword {
        Keyword {
            id: row.get("id"),
            created_at: row.get("created_at"),
            crates_cnt: row.get("crates_cnt"),
            keyword: row.get("keyword"),
        }
    }
    fn table_name(_: Option<Keyword>) -> &'static str { "keywords" }
}

/// Handles the `GET /keywords` route.
pub fn index(req: &mut Request) -> CargoResult<Response> {
    use diesel::expression::dsl::sql;
    use diesel::types::BigInt;
    use schema::keywords;

    let conn = req.db_conn()?;
    let (offset, limit) = req.pagination(10, 100)?;
    let query = req.query();
    let sort = query.get("sort").map(|s| &s[..]).unwrap_or("alpha");

    let mut query = keywords::table
        .select((keywords::all_columns, sql::<BigInt>("COUNT(*) OVER ()")))
        .limit(limit)
        .offset(offset)
        .into_boxed();

    if sort == "crates" {
        query = query.order(keywords::crates_cnt.desc());
    } else {
        query = query.order(keywords::keyword.asc());
    }

    let data = query.load::<(Keyword, i64)>(&*conn)?;
    let total = data.get(0).map(|&(_, t)| t).unwrap_or(0);
    let kws = data.into_iter()
        .map(|(k, _)| k.encodable()).collect::<Vec<_>>();

    #[derive(RustcEncodable)]
    struct R { keywords: Vec<EncodableKeyword>, meta: Meta }
    #[derive(RustcEncodable)]
    struct Meta { total: i64 }

    Ok(req.json(&R {
        keywords: kws,
        meta: Meta { total: total },
    }))
}

/// Handles the `GET /keywords/:keyword_id` route.
pub fn show(req: &mut Request) -> CargoResult<Response> {
    let name = &req.params()["keyword_id"];
    let conn = req.db_conn()?;

    let kw = Keyword::find_by_keyword(&conn, name)?;

    #[derive(RustcEncodable)]
    struct R { keyword: EncodableKeyword }
    Ok(req.json(&R { keyword: kw.encodable() }))
}
<commit_msg>Lowercase keywords before inserting<commit_after>use std::ascii::AsciiExt;
use std::collections::HashMap;
use time::Timespec;

use conduit::{Request, Response};
use conduit_router::RequestParams;
use diesel::pg::PgConnection;
use diesel::prelude::*;
use diesel;
use pg::GenericConnection;
use pg::rows::Row;

use {Model, Crate};
use db::RequestTransaction;
use schema::*;
use util::{RequestUtils, CargoResult, ChainError, internal};

#[derive(Clone, Identifiable, Queryable)]
pub struct Keyword {
    pub id: i32,
    pub keyword: String,
    pub crates_cnt: i32,
    pub created_at: Timespec,
}

#[derive(Associations, Insertable, Identifiable)]
#[belongs_to(Keyword)]
#[belongs_to(Crate)]
#[table_name="crates_keywords"]
#[primary_key(crate_id, keyword_id)]
pub struct CrateKeyword {
    crate_id: i32,
    keyword_id: i32,
}

#[derive(RustcEncodable, RustcDecodable)]
pub struct EncodableKeyword {
    pub id: String,
    pub keyword: String,
    pub created_at: String,
    pub crates_cnt: i32,
}

impl Keyword {
    pub fn find_by_keyword(conn: &PgConnection, name: &str)
                           -> QueryResult<Keyword> {
        keywords::table
            .filter(keywords::keyword.eq(::lower(name)))
            .first(&*conn)
    }

    pub fn find_or_create_all(conn: &PgConnection, names: &[&str]) -> QueryResult<Vec<Keyword>> {
        use diesel::pg::upsert::*;
        use diesel::expression::dsl::any;

        #[derive(Insertable)]
        #[table_name="keywords"]
        struct NewKeyword<'a> {
            keyword: &'a str,
        }

        let lowercase_names: Vec<_> = names.iter()
            .map(|s| s.to_lowercase())
            .collect();

        let new_keywords: Vec<_> = lowercase_names.iter()
            .map(|s| NewKeyword { keyword: s })
            .collect();

        // https://github.com/diesel-rs/diesel/issues/797
        if !new_keywords.is_empty() {
            diesel::insert(&new_keywords.on_conflict_do_nothing())
                .into(keywords::table)
                .execute(conn)?;
        }
        keywords::table
            .filter(::lower(keywords::keyword).eq(any(&lowercase_names)))
            .load(conn)
    }

    pub fn find_or_insert(conn: &GenericConnection, name: &str)
                          -> CargoResult<Keyword> {
        // TODO: racy (the select then insert is not atomic)
        let stmt = conn.prepare("SELECT * FROM keywords
                                      WHERE keyword = LOWER($1)")?;
        for row in stmt.query(&[&name])?.iter() {
            return Ok(Model::from_row(&row))
        }

        let stmt = conn.prepare("INSERT INTO keywords (keyword) VALUES (LOWER($1))
                                      RETURNING *")?;
        let rows = stmt.query(&[&name])?;
        Ok(Model::from_row(&rows.iter().next().chain_error(|| {
            internal("no version returned")
        })?))
    }

    pub fn valid_name(name: &str) -> bool {
        if name.is_empty() { return false }
        name.chars().next().unwrap().is_alphanumeric() &&
            name.chars().all(|c| c.is_alphanumeric() || c == '_' || c == '-') &&
            name.chars().all(|c| c.is_ascii())
    }

    pub fn encodable(self) -> EncodableKeyword {
        let Keyword { crates_cnt, keyword, created_at, .. } = self;
        EncodableKeyword {
            id: keyword.clone(),
            created_at: ::encode_time(created_at),
            crates_cnt: crates_cnt,
            keyword: keyword,
        }
    }

    pub fn update_crate(conn: &PgConnection,
                        krate: &Crate,
                        keywords: &[&str]) -> QueryResult<()> {
        conn.transaction(|| {
            let keywords = Keyword::find_or_create_all(conn, keywords)?;
            diesel::delete(CrateKeyword::belonging_to(krate))
                .execute(conn)?;
            let crate_keywords = keywords.into_iter().map(|kw| {
                CrateKeyword { crate_id: krate.id, keyword_id: kw.id }
            }).collect::<Vec<_>>();
            diesel::insert(&crate_keywords).into(crates_keywords::table)
                .execute(conn)?;
            Ok(())
        })
    }

    pub fn update_crate_old(conn: &GenericConnection,
                        krate: &Crate,
                        keywords: &[String]) -> CargoResult<()> {
        let old_kws = krate.keywords(conn)?;
        let old_kws = old_kws.iter().map(|kw| {
            (&kw.keyword[..], kw)
        }).collect::<HashMap<_, _>>();
        let new_kws = keywords.iter().map(|k| {
            let kw = Keyword::find_or_insert(conn, k)?;
            Ok((k.as_str(), kw))
        }).collect::<CargoResult<HashMap<_, _>>>()?;

        let to_rm = old_kws.iter().filter(|&(kw, _)| {
            !new_kws.contains_key(kw)
        }).map(|(_, v)| v.id).collect::<Vec<_>>();
        let to_add = new_kws.iter().filter(|&(kw, _)| {
            !old_kws.contains_key(kw)
        }).map(|(_, v)| v.id).collect::<Vec<_>>();

        if !to_rm.is_empty() {
            conn.execute("DELETE FROM crates_keywords
                                WHERE keyword_id = ANY($1)
                                  AND crate_id = $2",
                         &[&to_rm, &krate.id])?;
        }

        if !to_add.is_empty() {
            let insert = to_add.iter().map(|id| {
                let crate_id: i32 = krate.id;
                let id: i32 = *id;
                format!("({}, {})", crate_id,  id)
            }).collect::<Vec<_>>().join(", ");
            conn.execute(&format!("INSERT INTO crates_keywords
                                        (crate_id, keyword_id) VALUES {}",
                                  insert),
                         &[])?;
        }

        Ok(())
    }
}

impl Model for Keyword {
    fn from_row(row: &Row) -> Keyword {
        Keyword {
            id: row.get("id"),
            created_at: row.get("created_at"),
            crates_cnt: row.get("crates_cnt"),
            keyword: row.get("keyword"),
        }
    }
    fn table_name(_: Option<Keyword>) -> &'static str { "keywords" }
}

/// Handles the `GET /keywords` route.
pub fn index(req: &mut Request) -> CargoResult<Response> {
    use diesel::expression::dsl::sql;
    use diesel::types::BigInt;
    use schema::keywords;

    let conn = req.db_conn()?;
    let (offset, limit) = req.pagination(10, 100)?;
    let query = req.query();
    let sort = query.get("sort").map(|s| &s[..]).unwrap_or("alpha");

    let mut query = keywords::table
        .select((keywords::all_columns, sql::<BigInt>("COUNT(*) OVER ()")))
        .limit(limit)
        .offset(offset)
        .into_boxed();

    if sort == "crates" {
        query = query.order(keywords::crates_cnt.desc());
    } else {
        query = query.order(keywords::keyword.asc());
    }

    let data = query.load::<(Keyword, i64)>(&*conn)?;
    let total = data.get(0).map(|&(_, t)| t).unwrap_or(0);
    let kws = data.into_iter()
        .map(|(k, _)| k.encodable()).collect::<Vec<_>>();

    #[derive(RustcEncodable)]
    struct R { keywords: Vec<EncodableKeyword>, meta: Meta }
    #[derive(RustcEncodable)]
    struct Meta { total: i64 }

    Ok(req.json(&R {
        keywords: kws,
        meta: Meta { total: total },
    }))
}

/// Handles the `GET /keywords/:keyword_id` route.
pub fn show(req: &mut Request) -> CargoResult<Response> {
    let name = &req.params()["keyword_id"];
    let conn = req.db_conn()?;

    let kw = Keyword::find_by_keyword(&conn, name)?;

    #[derive(RustcEncodable)]
    struct R { keyword: EncodableKeyword }
    Ok(req.json(&R { keyword: kw.encodable() }))
}
<|endoftext|>


<commit_before>//! Types for the [`m.room.join_rules`] event.
//!
//! [`m.room.join_rules`]: https://spec.matrix.org/v1.2/client-server-api/#mroomjoin_rules

use std::{borrow::Cow, collections::BTreeMap};

use ruma_macros::EventContent;
use serde::{
    de::{Deserializer, Error},
    Deserialize, Serialize,
};
use serde_json::{value::RawValue as RawJsonValue, Value as JsonValue};

use crate::{events::EmptyStateKey, serde::from_raw_json_value, OwnedRoomId, PrivOwnedStr};

/// The content of an `m.room.join_rules` event.
///
/// Describes how users are allowed to join the room.
#[derive(Clone, Debug, Serialize, EventContent)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[ruma_event(type = "m.room.join_rules", kind = State, state_key_type = EmptyStateKey)]
pub struct RoomJoinRulesEventContent {
    /// The type of rules used for users wishing to join this room.
    #[ruma_event(skip_redaction)]
    #[serde(flatten)]
    pub join_rule: JoinRule,
}

impl RoomJoinRulesEventContent {
    /// Creates a new `RoomJoinRulesEventContent` with the given rule.
    pub fn new(join_rule: JoinRule) -> Self {
        Self { join_rule }
    }

    /// Creates a new `RoomJoinRulesEventContent` with the restricted rule and the given set of
    /// allow rules.
    pub fn restricted(allow: Vec<AllowRule>) -> Self {
        Self { join_rule: JoinRule::Restricted(Restricted::new(allow)) }
    }

    /// Creates a new `RoomJoinRulesEventContent` with the knock restricted rule and the given set
    /// of allow rules.
    pub fn knock_restricted(allow: Vec<AllowRule>) -> Self {
        Self { join_rule: JoinRule::KnockRestricted(Restricted::new(allow)) }
    }
}

impl<'de> Deserialize<'de> for RoomJoinRulesEventContent {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let join_rule = JoinRule::deserialize(deserializer)?;
        Ok(RoomJoinRulesEventContent { join_rule })
    }
}

impl RoomJoinRulesEvent {
    /// Obtain the join rule, regardless of whether this event is redacted.
    pub fn join_rule(&self) -> &JoinRule {
        match self {
            Self::Original(ev) => &ev.content.join_rule,
            Self::Redacted(ev) => &ev.content.join_rule,
        }
    }
}

impl SyncRoomJoinRulesEvent {
    /// Obtain the join rule, regardless of whether this event is redacted.
    pub fn join_rule(&self) -> &JoinRule {
        match self {
            Self::Original(ev) => &ev.content.join_rule,
            Self::Redacted(ev) => &ev.content.join_rule,
        }
    }
}

/// The rule used for users wishing to join this room.
///
/// This type can hold an arbitrary string. To check for values that are not available as a
/// documented variant here, use its string representation, obtained through `.as_str()`.
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[serde(tag = "join_rule")]
pub enum JoinRule {
    /// A user who wishes to join the room must first receive an invite to the room from someone
    /// already inside of the room.
    #[serde(rename = "invite")]
    Invite,

    /// Users can join the room if they are invited, or they can request an invite to the room.
    ///
    /// They can be allowed (invited) or denied (kicked/banned) access.
    #[serde(rename = "knock")]
    Knock,

    /// Reserved but not yet implemented by the Matrix specification.
    #[serde(rename = "private")]
    Private,

    /// Users can join the room if they are invited, or if they meet any of the conditions
    /// described in a set of [`AllowRule`]s.
    #[serde(rename = "restricted")]
    Restricted(Restricted),

    /// Users can join the room if they are invited, or if they meet any of the conditions
    /// described in a set of [`AllowRule`]s, or they can request an invite to the room.
    #[serde(rename = "knock_restricted")]
    KnockRestricted(Restricted),

    /// Anyone can join the room without any prior action.
    #[serde(rename = "public")]
    Public,

    #[doc(hidden)]
    #[serde(skip_serializing)]
    _Custom(PrivOwnedStr),
}

impl JoinRule {
    /// Returns the string name of this `JoinRule`
    pub fn as_str(&self) -> &str {
        match self {
            JoinRule::Invite => "invite",
            JoinRule::Knock => "knock",
            JoinRule::Private => "private",
            JoinRule::Restricted(_) => "restricted",
            JoinRule::KnockRestricted(_) => "knock_restricted",
            JoinRule::Public => "public",
            JoinRule::_Custom(rule) => &rule.0,
        }
    }
}

impl<'de> Deserialize<'de> for JoinRule {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let json: Box<RawJsonValue> = Box::deserialize(deserializer)?;

        #[derive(Deserialize)]
        struct ExtractType<'a> {
            #[serde(borrow)]
            join_rule: Option<Cow<'a, str>>,
        }

        let join_rule = serde_json::from_str::<ExtractType<'_>>(json.get())
            .map_err(serde::de::Error::custom)?
            .join_rule
            .ok_or_else(|| D::Error::missing_field("join_rule"))?;

        match join_rule.as_ref() {
            "invite" => Ok(Self::Invite),
            "knock" => Ok(Self::Knock),
            "private" => Ok(Self::Private),
            "restricted" => from_raw_json_value(&json).map(Self::Restricted),
            "knock_restricted" => from_raw_json_value(&json).map(Self::KnockRestricted),
            "public" => Ok(Self::Public),
            _ => Ok(Self::_Custom(PrivOwnedStr(join_rule.into()))),
        }
    }
}

/// Configuration of the `Restricted` join rule.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct Restricted {
    /// Allow rules which describe conditions that allow joining a room.
    pub allow: Vec<AllowRule>,
}

impl Restricted {
    /// Constructs a new rule set for restricted rooms with the given rules.
    pub fn new(allow: Vec<AllowRule>) -> Self {
        Self { allow }
    }
}

/// An allow rule which defines a condition that allows joining a room.
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[serde(tag = "type")]
pub enum AllowRule {
    /// Joining is allowed if a user is already a member of the room with the id `room_id`.
    #[serde(rename = "m.room_membership")]
    RoomMembership(RoomMembership),

    #[doc(hidden)]
    _Custom(CustomAllowRule),
}

impl AllowRule {
    /// Constructs an `AllowRule` with membership of the room with the given id as its predicate.
    pub fn room_membership(room_id: OwnedRoomId) -> Self {
        Self::RoomMembership(RoomMembership::new(room_id))
    }
}

/// Allow rule which grants permission to join based on the membership of another room.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct RoomMembership {
    /// The id of the room which being a member of grants permission to join another room.
    pub room_id: OwnedRoomId,
}

impl RoomMembership {
    /// Constructs a new room membership rule for the given room id.
    pub fn new(room_id: OwnedRoomId) -> Self {
        Self { room_id }
    }
}

#[doc(hidden)]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct CustomAllowRule {
    #[serde(rename = "type")]
    rule_type: String,
    #[serde(flatten)]
    extra: BTreeMap<String, JsonValue>,
}

impl<'de> Deserialize<'de> for AllowRule {
    fn deserialize<D>(deserializer: D) -> Result<AllowRule, D::Error>
    where
        D: Deserializer<'de>,
    {
        let json: Box<RawJsonValue> = Box::deserialize(deserializer)?;

        // Extracts the `type` value.
        #[derive(Deserialize)]
        struct ExtractType<'a> {
            #[serde(borrow, rename = "type")]
            rule_type: Option<Cow<'a, str>>,
        }

        // Get the value of `type` if present.
        let rule_type = serde_json::from_str::<ExtractType<'_>>(json.get())
            .map_err(serde::de::Error::custom)?
            .rule_type;

        match rule_type.as_deref() {
            Some("m.room_membership") => from_raw_json_value(&json).map(Self::RoomMembership),
            Some(_) => from_raw_json_value(&json).map(Self::_Custom),
            None => Err(D::Error::missing_field("type")),
        }
    }
}

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

    use super::{AllowRule, JoinRule, OriginalSyncRoomJoinRulesEvent, RoomJoinRulesEventContent};
    use crate::room_id;

    #[test]
    fn deserialize() {
        let json = r#"{"join_rule": "public"}"#;
        let event: RoomJoinRulesEventContent = serde_json::from_str(json).unwrap();
        assert_matches!(event, RoomJoinRulesEventContent { join_rule: JoinRule::Public });
    }

    #[test]
    fn deserialize_restricted() {
        let json = r#"{
            "join_rule": "restricted",
            "allow": [
                {
                    "type": "m.room_membership",
                    "room_id": "!mods:example.org"
                },
                {
                    "type": "m.room_membership",
                    "room_id": "!users:example.org"
                }
            ]
        }"#;
        let event: RoomJoinRulesEventContent = serde_json::from_str(json).unwrap();
        match event.join_rule {
            JoinRule::Restricted(restricted) => assert_eq!(
                restricted.allow,
                &[
                    AllowRule::room_membership(room_id!("!mods:example.org").to_owned()),
                    AllowRule::room_membership(room_id!("!users:example.org").to_owned())
                ]
            ),
            rule => panic!("Deserialized to wrong variant: {rule:?}"),
        }
    }

    #[test]
    fn deserialize_restricted_event() {
        let json = r#"{
            "type": "m.room.join_rules",
            "sender": "@admin:community.rs",
            "content": {
                "join_rule": "restricted",
                "allow":[
                    { "type": "m.room_membership","room_id": "!KqeUnzmXPIhHRaWMTs:mccarty.io" }
                ]
            },
            "state_key": "",
            "origin_server_ts":1630508835342,
            "unsigned": {
                "age":4165521871
            },
            "event_id": "$0ACb9KSPlT3al3kikyRYvFhMqXPP9ZcQOBrsdIuh58U"
        }"#;

        assert_matches!(serde_json::from_str::<OriginalSyncRoomJoinRulesEvent>(json), Ok(_));
    }
}
<commit_msg>events: Use rename_all for JoinRule variants<commit_after>//! Types for the [`m.room.join_rules`] event.
//!
//! [`m.room.join_rules`]: https://spec.matrix.org/v1.2/client-server-api/#mroomjoin_rules

use std::{borrow::Cow, collections::BTreeMap};

use ruma_macros::EventContent;
use serde::{
    de::{Deserializer, Error},
    Deserialize, Serialize,
};
use serde_json::{value::RawValue as RawJsonValue, Value as JsonValue};

use crate::{events::EmptyStateKey, serde::from_raw_json_value, OwnedRoomId, PrivOwnedStr};

/// The content of an `m.room.join_rules` event.
///
/// Describes how users are allowed to join the room.
#[derive(Clone, Debug, Serialize, EventContent)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[ruma_event(type = "m.room.join_rules", kind = State, state_key_type = EmptyStateKey)]
pub struct RoomJoinRulesEventContent {
    /// The type of rules used for users wishing to join this room.
    #[ruma_event(skip_redaction)]
    #[serde(flatten)]
    pub join_rule: JoinRule,
}

impl RoomJoinRulesEventContent {
    /// Creates a new `RoomJoinRulesEventContent` with the given rule.
    pub fn new(join_rule: JoinRule) -> Self {
        Self { join_rule }
    }

    /// Creates a new `RoomJoinRulesEventContent` with the restricted rule and the given set of
    /// allow rules.
    pub fn restricted(allow: Vec<AllowRule>) -> Self {
        Self { join_rule: JoinRule::Restricted(Restricted::new(allow)) }
    }

    /// Creates a new `RoomJoinRulesEventContent` with the knock restricted rule and the given set
    /// of allow rules.
    pub fn knock_restricted(allow: Vec<AllowRule>) -> Self {
        Self { join_rule: JoinRule::KnockRestricted(Restricted::new(allow)) }
    }
}

impl<'de> Deserialize<'de> for RoomJoinRulesEventContent {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let join_rule = JoinRule::deserialize(deserializer)?;
        Ok(RoomJoinRulesEventContent { join_rule })
    }
}

impl RoomJoinRulesEvent {
    /// Obtain the join rule, regardless of whether this event is redacted.
    pub fn join_rule(&self) -> &JoinRule {
        match self {
            Self::Original(ev) => &ev.content.join_rule,
            Self::Redacted(ev) => &ev.content.join_rule,
        }
    }
}

impl SyncRoomJoinRulesEvent {
    /// Obtain the join rule, regardless of whether this event is redacted.
    pub fn join_rule(&self) -> &JoinRule {
        match self {
            Self::Original(ev) => &ev.content.join_rule,
            Self::Redacted(ev) => &ev.content.join_rule,
        }
    }
}

/// The rule used for users wishing to join this room.
///
/// This type can hold an arbitrary string. To check for values that are not available as a
/// documented variant here, use its string representation, obtained through `.as_str()`.
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[serde(tag = "join_rule", rename_all = "snake_case")]
pub enum JoinRule {
    /// A user who wishes to join the room must first receive an invite to the room from someone
    /// already inside of the room.
    Invite,

    /// Users can join the room if they are invited, or they can request an invite to the room.
    ///
    /// They can be allowed (invited) or denied (kicked/banned) access.
    Knock,

    /// Reserved but not yet implemented by the Matrix specification.
    Private,

    /// Users can join the room if they are invited, or if they meet any of the conditions
    /// described in a set of [`AllowRule`]s.
    Restricted(Restricted),

    /// Users can join the room if they are invited, or if they meet any of the conditions
    /// described in a set of [`AllowRule`]s, or they can request an invite to the room.
    KnockRestricted(Restricted),

    /// Anyone can join the room without any prior action.
    Public,

    #[doc(hidden)]
    #[serde(skip_serializing)]
    _Custom(PrivOwnedStr),
}

impl JoinRule {
    /// Returns the string name of this `JoinRule`
    pub fn as_str(&self) -> &str {
        match self {
            JoinRule::Invite => "invite",
            JoinRule::Knock => "knock",
            JoinRule::Private => "private",
            JoinRule::Restricted(_) => "restricted",
            JoinRule::KnockRestricted(_) => "knock_restricted",
            JoinRule::Public => "public",
            JoinRule::_Custom(rule) => &rule.0,
        }
    }
}

impl<'de> Deserialize<'de> for JoinRule {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let json: Box<RawJsonValue> = Box::deserialize(deserializer)?;

        #[derive(Deserialize)]
        struct ExtractType<'a> {
            #[serde(borrow)]
            join_rule: Option<Cow<'a, str>>,
        }

        let join_rule = serde_json::from_str::<ExtractType<'_>>(json.get())
            .map_err(serde::de::Error::custom)?
            .join_rule
            .ok_or_else(|| D::Error::missing_field("join_rule"))?;

        match join_rule.as_ref() {
            "invite" => Ok(Self::Invite),
            "knock" => Ok(Self::Knock),
            "private" => Ok(Self::Private),
            "restricted" => from_raw_json_value(&json).map(Self::Restricted),
            "knock_restricted" => from_raw_json_value(&json).map(Self::KnockRestricted),
            "public" => Ok(Self::Public),
            _ => Ok(Self::_Custom(PrivOwnedStr(join_rule.into()))),
        }
    }
}

/// Configuration of the `Restricted` join rule.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct Restricted {
    /// Allow rules which describe conditions that allow joining a room.
    pub allow: Vec<AllowRule>,
}

impl Restricted {
    /// Constructs a new rule set for restricted rooms with the given rules.
    pub fn new(allow: Vec<AllowRule>) -> Self {
        Self { allow }
    }
}

/// An allow rule which defines a condition that allows joining a room.
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[serde(tag = "type")]
pub enum AllowRule {
    /// Joining is allowed if a user is already a member of the room with the id `room_id`.
    #[serde(rename = "m.room_membership")]
    RoomMembership(RoomMembership),

    #[doc(hidden)]
    _Custom(CustomAllowRule),
}

impl AllowRule {
    /// Constructs an `AllowRule` with membership of the room with the given id as its predicate.
    pub fn room_membership(room_id: OwnedRoomId) -> Self {
        Self::RoomMembership(RoomMembership::new(room_id))
    }
}

/// Allow rule which grants permission to join based on the membership of another room.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct RoomMembership {
    /// The id of the room which being a member of grants permission to join another room.
    pub room_id: OwnedRoomId,
}

impl RoomMembership {
    /// Constructs a new room membership rule for the given room id.
    pub fn new(room_id: OwnedRoomId) -> Self {
        Self { room_id }
    }
}

#[doc(hidden)]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct CustomAllowRule {
    #[serde(rename = "type")]
    rule_type: String,
    #[serde(flatten)]
    extra: BTreeMap<String, JsonValue>,
}

impl<'de> Deserialize<'de> for AllowRule {
    fn deserialize<D>(deserializer: D) -> Result<AllowRule, D::Error>
    where
        D: Deserializer<'de>,
    {
        let json: Box<RawJsonValue> = Box::deserialize(deserializer)?;

        // Extracts the `type` value.
        #[derive(Deserialize)]
        struct ExtractType<'a> {
            #[serde(borrow, rename = "type")]
            rule_type: Option<Cow<'a, str>>,
        }

        // Get the value of `type` if present.
        let rule_type = serde_json::from_str::<ExtractType<'_>>(json.get())
            .map_err(serde::de::Error::custom)?
            .rule_type;

        match rule_type.as_deref() {
            Some("m.room_membership") => from_raw_json_value(&json).map(Self::RoomMembership),
            Some(_) => from_raw_json_value(&json).map(Self::_Custom),
            None => Err(D::Error::missing_field("type")),
        }
    }
}

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

    use super::{AllowRule, JoinRule, OriginalSyncRoomJoinRulesEvent, RoomJoinRulesEventContent};
    use crate::room_id;

    #[test]
    fn deserialize() {
        let json = r#"{"join_rule": "public"}"#;
        let event: RoomJoinRulesEventContent = serde_json::from_str(json).unwrap();
        assert_matches!(event, RoomJoinRulesEventContent { join_rule: JoinRule::Public });
    }

    #[test]
    fn deserialize_restricted() {
        let json = r#"{
            "join_rule": "restricted",
            "allow": [
                {
                    "type": "m.room_membership",
                    "room_id": "!mods:example.org"
                },
                {
                    "type": "m.room_membership",
                    "room_id": "!users:example.org"
                }
            ]
        }"#;
        let event: RoomJoinRulesEventContent = serde_json::from_str(json).unwrap();
        match event.join_rule {
            JoinRule::Restricted(restricted) => assert_eq!(
                restricted.allow,
                &[
                    AllowRule::room_membership(room_id!("!mods:example.org").to_owned()),
                    AllowRule::room_membership(room_id!("!users:example.org").to_owned())
                ]
            ),
            rule => panic!("Deserialized to wrong variant: {rule:?}"),
        }
    }

    #[test]
    fn deserialize_restricted_event() {
        let json = r#"{
            "type": "m.room.join_rules",
            "sender": "@admin:community.rs",
            "content": {
                "join_rule": "restricted",
                "allow":[
                    { "type": "m.room_membership","room_id": "!KqeUnzmXPIhHRaWMTs:mccarty.io" }
                ]
            },
            "state_key": "",
            "origin_server_ts":1630508835342,
            "unsigned": {
                "age":4165521871
            },
            "event_id": "$0ACb9KSPlT3al3kikyRYvFhMqXPP9ZcQOBrsdIuh58U"
        }"#;

        assert_matches!(serde_json::from_str::<OriginalSyncRoomJoinRulesEvent>(json), Ok(_));
    }
}
<|endoftext|>


<commit_before>extern crate luthor;

use std::rc::Rc;
use std::cell::RefCell;
use std::old_io::{File, Open, Read, Write};
use std::old_io::IoError;
use std::old_io::IoResult;
use super::GapBuffer;
use super::gap_buffer;
use super::Position;
use super::Range;
use super::Cursor;
use super::type_detection;
use self::luthor::token::{Token, Category};
use self::luthor::lexers;

/// A UTF-8 buffer with bounds-checked cursor management and persistence.
pub struct Buffer {
    data: Rc<RefCell<GapBuffer>>,
    lexer: Option<fn(&str) -> Vec<Token>>,
    pub path: Option<Path>,
    pub cursor: Cursor,
}

impl Buffer {
    /// Returns the contents of the buffer as a string.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn data(&self) -> String {
        self.data.borrow().to_string()
    }

    /// Writes the contents of the buffer to its path.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::old_io::{File, Open, Read};
    ///
    /// // Set up a buffer and point it to a path.
    /// let mut buffer = scribe::buffer::new();
    /// let write_path = Path::new("my_doc");
    /// buffer.path = Some(write_path.clone());
    ///
    /// // Put some data into the buffer and save it.
    /// buffer.insert("scribe");
    /// buffer.save();
    ///
    /// # let saved_data = File::open_mode(&write_path, Open, Read)
    /// #   .unwrap().read_to_string().unwrap();
    /// # assert_eq!(saved_data, "scribe");
    ///
    /// # std::old_io::fs::unlink(&write_path);
    /// ```
    pub fn save(&self) -> Option<IoError> {
        let path = match self.path.clone() {
            Some(p) => p,
            None => Path::new(""),
        };

        // Try to open and write to the file, returning any errors encountered.
        let mut file = match File::open_mode(&path, Open, Write) {
            Ok(f) => f,
            Err(error) => return Some(error),
        };

        // We use to_string here because we don't want to write the gap contents.
        match file.write(self.data().to_string().as_bytes()) {
            Ok(_) => (),
            Err(error) => return Some(error),
        }

        return None
    }

    /// Inserts `data` into the buffer at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn insert(&mut self, data: &str) {
        self.data.borrow_mut().insert(data, &self.cursor);
    }

    /// Deletes a character at the cursor position. If at the end
    /// of the current line, it'll try to delete a newline character
    /// (joining the lines), succeeding if there's a line below.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// buffer.delete();
    /// assert_eq!(buffer.data(), "cribe");
    /// ```
    pub fn delete(&mut self) {
        // We need to specify a range to delete, so start at
        // the current offset and delete the character to the right.
        let mut end = Position{ line: self.cursor.line, offset: self.cursor.offset + 1 };

        // If there isn't a character to the right,
        // delete the newline by jumping to the start
        // of the next line. If it doesn't exist, that's okay;
        // these values are bounds-checked by delete() anyway.
        if !self.data.borrow().in_bounds(&end) {
            end.line += 1;
            end.offset = 0;
        }

        self.data.borrow_mut().delete(&Range{ start: *self.cursor, end: end});
    }

    /// Produces a set of tokens based on the buffer data
    /// suitable for colorized display, using a lexer for the
    /// buffer data's language and/or format. If a lexer is not
    /// available, the set will consist of a single text-category token.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    ///
    /// // Build the buffer data string back by combining its token lexemes.
    /// let mut data = String::new();
    /// for token in buffer.tokens().iter() {
    ///     data.push_str(token.lexeme.as_slice());
    /// }
    /// assert_eq!(data, "scribe");
    /// ```
    pub fn tokens(&self) -> Vec<Token> {
        match self.lexer {
            Some(lexer) => lexer(&self.data()),
            None => vec![Token{ lexeme: self.data(), category: Category::Text }],
        }
    }
}

/// Creates a new empty buffer. The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::new();
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn new() -> Buffer {
    let data = Rc::new(RefCell::new(gap_buffer::new(String::new())));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    Buffer{ data: data.clone(), path: None, cursor: cursor, lexer: None }
}

/// Creates a new buffer by reading the UTF-8 interpreted file contents of the specified path.
/// The buffer's cursor is set to the beginning of the buffer. The buffer data's type will be
/// inferred based on its extension, and an appropriate lexer will be used, if available (see
/// tokens method for further information on why this happens).
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::from_file(Path::new("tests/sample/file")).unwrap();
/// assert_eq!(buffer.data(), "it works!\n");
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn from_file(path: Path) -> IoResult<Buffer> {
    // Try to open and read the file, returning any errors encountered.
    let mut file = match File::open_mode(&path, Open, Read) {
        Ok(f) => f,
        Err(error) => return Err(error),
    };
    let mut data = match file.read_to_string() {
        Ok(d) => d,
        Err(error) => return Err(error),
    };

    let data = Rc::new(RefCell::new(gap_buffer::new(data)));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    // Detect the file type and use its corresponding lexer, if available.
    let lexer = match type_detection::from_path(&path) {
        Some(type_detection::Type::JSON) => Some(lexers::json::lex as fn(&str) -> Vec<Token>),
        Some(type_detection::Type::XML) => Some(lexers::xml::lex as fn(&str) -> Vec<Token>),
        _ => None,
    };

    // Create a new buffer using the loaded data, path, and other defaults.
    Ok(Buffer{ data: data.clone(), path: Some(path), cursor: cursor, lexer: lexer })
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::luthor::token::{Token, Category};

    fn placeholder_lexer(_: &str) -> Vec<Token> {
        vec![Token{ lexeme: "lexer".to_string(), category: Category::Text }]
    }

    #[test]
    fn tokens_returns_one_text_token_when_no_lexer_is_set() {
        let mut buffer = new();
        buffer.insert("scribe");
        let expected_tokens = vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }];
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn tokens_returns_result_of_lexer_when_set() {
        let mut buffer = new();
        let expected_tokens = placeholder_lexer("scribe");
        buffer.lexer = Some(placeholder_lexer);
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn delete_joins_lines_when_invoked_at_end_of_line() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe library");
    }

    #[test]
    fn delete_does_nothing_when_invoked_at_the_end_of_the_document() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_down();
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe\n library");
    }
}
<commit_msg>Implement filename method for the buffer type.<commit_after>extern crate luthor;

use std::rc::Rc;
use std::cell::RefCell;
use std::str::from_utf8;
use std::old_io::{File, Open, Read, Write};
use std::old_io::IoError;
use std::old_io::IoResult;
use super::GapBuffer;
use super::gap_buffer;
use super::Position;
use super::Range;
use super::Cursor;
use super::type_detection;
use self::luthor::token::{Token, Category};
use self::luthor::lexers;

/// A UTF-8 buffer with bounds-checked cursor management and persistence.
pub struct Buffer {
    data: Rc<RefCell<GapBuffer>>,
    lexer: Option<fn(&str) -> Vec<Token>>,
    pub path: Option<Path>,
    pub cursor: Cursor,
}

impl Buffer {
    /// Returns the contents of the buffer as a string.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn data(&self) -> String {
        self.data.borrow().to_string()
    }

    /// Writes the contents of the buffer to its path.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::old_io::{File, Open, Read};
    ///
    /// // Set up a buffer and point it to a path.
    /// let mut buffer = scribe::buffer::new();
    /// let write_path = Path::new("my_doc");
    /// buffer.path = Some(write_path.clone());
    ///
    /// // Put some data into the buffer and save it.
    /// buffer.insert("scribe");
    /// buffer.save();
    ///
    /// # let saved_data = File::open_mode(&write_path, Open, Read)
    /// #   .unwrap().read_to_string().unwrap();
    /// # assert_eq!(saved_data, "scribe");
    ///
    /// # std::old_io::fs::unlink(&write_path);
    /// ```
    pub fn save(&self) -> Option<IoError> {
        let path = match self.path.clone() {
            Some(p) => p,
            None => Path::new(""),
        };

        // Try to open and write to the file, returning any errors encountered.
        let mut file = match File::open_mode(&path, Open, Write) {
            Ok(f) => f,
            Err(error) => return Some(error),
        };

        // We use to_string here because we don't want to write the gap contents.
        match file.write(self.data().to_string().as_bytes()) {
            Ok(_) => (),
            Err(error) => return Some(error),
        }

        return None
    }

    /// Inserts `data` into the buffer at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn insert(&mut self, data: &str) {
        self.data.borrow_mut().insert(data, &self.cursor);
    }

    /// Deletes a character at the cursor position. If at the end
    /// of the current line, it'll try to delete a newline character
    /// (joining the lines), succeeding if there's a line below.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// buffer.delete();
    /// assert_eq!(buffer.data(), "cribe");
    /// ```
    pub fn delete(&mut self) {
        // We need to specify a range to delete, so start at
        // the current offset and delete the character to the right.
        let mut end = Position{ line: self.cursor.line, offset: self.cursor.offset + 1 };

        // If there isn't a character to the right,
        // delete the newline by jumping to the start
        // of the next line. If it doesn't exist, that's okay;
        // these values are bounds-checked by delete() anyway.
        if !self.data.borrow().in_bounds(&end) {
            end.line += 1;
            end.offset = 0;
        }

        self.data.borrow_mut().delete(&Range{ start: *self.cursor, end: end});
    }

    /// Produces a set of tokens based on the buffer data
    /// suitable for colorized display, using a lexer for the
    /// buffer data's language and/or format. If a lexer is not
    /// available, the set will consist of a single text-category token.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    ///
    /// // Build the buffer data string back by combining its token lexemes.
    /// let mut data = String::new();
    /// for token in buffer.tokens().iter() {
    ///     data.push_str(token.lexeme.as_slice());
    /// }
    /// assert_eq!(data, "scribe");
    /// ```
    pub fn tokens(&self) -> Vec<Token> {
        match self.lexer {
            Some(lexer) => lexer(&self.data()),
            None => vec![Token{ lexeme: self.data(), category: Category::Text }],
        }
    }

    /// Returns the filename portion of the buffer's path, if
    /// the path is set and its filename is a valid UTF-8 sequence.
    ///
    /// # Examples
    ///
    /// ```
    /// let buffer = scribe::buffer::from_file(Path::new("tests/sample/file")).unwrap();
    /// assert_eq!(buffer.filename().unwrap(), "file");
    /// ```
    pub fn filename(&self) -> Option<String> {
        match self.path {
            Some(ref path) => {
                match path.filename() {
                    Some(filename) => {
                        match from_utf8(filename) {
                            Ok(utf8_filename) => Some(utf8_filename.to_string()),
                            Err(_) => None,
                        }
                    },
                    None => None,
                }
            },
            None => None,
        }
    }
}

/// Creates a new empty buffer. The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::new();
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn new() -> Buffer {
    let data = Rc::new(RefCell::new(gap_buffer::new(String::new())));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    Buffer{ data: data.clone(), path: None, cursor: cursor, lexer: None }
}

/// Creates a new buffer by reading the UTF-8 interpreted file contents of the specified path.
/// The buffer's cursor is set to the beginning of the buffer. The buffer data's type will be
/// inferred based on its extension, and an appropriate lexer will be used, if available (see
/// tokens method for further information on why this happens).
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::from_file(Path::new("tests/sample/file")).unwrap();
/// assert_eq!(buffer.data(), "it works!\n");
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn from_file(path: Path) -> IoResult<Buffer> {
    // Try to open and read the file, returning any errors encountered.
    let mut file = match File::open_mode(&path, Open, Read) {
        Ok(f) => f,
        Err(error) => return Err(error),
    };
    let mut data = match file.read_to_string() {
        Ok(d) => d,
        Err(error) => return Err(error),
    };

    let data = Rc::new(RefCell::new(gap_buffer::new(data)));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    // Detect the file type and use its corresponding lexer, if available.
    let lexer = match type_detection::from_path(&path) {
        Some(type_detection::Type::JSON) => Some(lexers::json::lex as fn(&str) -> Vec<Token>),
        Some(type_detection::Type::XML) => Some(lexers::xml::lex as fn(&str) -> Vec<Token>),
        _ => None,
    };

    // Create a new buffer using the loaded data, path, and other defaults.
    Ok(Buffer{ data: data.clone(), path: Some(path), cursor: cursor, lexer: lexer })
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::luthor::token::{Token, Category};

    fn placeholder_lexer(_: &str) -> Vec<Token> {
        vec![Token{ lexeme: "lexer".to_string(), category: Category::Text }]
    }

    #[test]
    fn tokens_returns_one_text_token_when_no_lexer_is_set() {
        let mut buffer = new();
        buffer.insert("scribe");
        let expected_tokens = vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }];
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn tokens_returns_result_of_lexer_when_set() {
        let mut buffer = new();
        let expected_tokens = placeholder_lexer("scribe");
        buffer.lexer = Some(placeholder_lexer);
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn delete_joins_lines_when_invoked_at_end_of_line() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe library");
    }

    #[test]
    fn delete_does_nothing_when_invoked_at_the_end_of_the_document() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_down();
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe\n library");
    }
}
<|endoftext|>


<commit_before>use crate::{Assist, AssistCtx, AssistId};

use ast::LoopBodyOwner;
use ra_fmt::unwrap_trivial_block;
use ra_syntax::{ast, AstNode};

// Assist: unwrap_block
//
// This assist removes if...else, for, while and loop control statements to just keep the body.
//
// ```
// fn foo() {
//     if true {<|>
//         println!("foo");
//     }
// }
// ```
// ->
// ```
// fn foo() {
//     println!("foo");
// }
// ```
pub(crate) fn unwrap_block(ctx: AssistCtx) -> Option<Assist> {
    let res = if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
        // if expression
        let mut expr_to_unwrap: Option<ast::Expr> = None;
        for block_expr in if_expr.blocks() {
            if let Some(block) = block_expr.block() {
                let cursor_in_range =
                    block.l_curly_token()?.text_range().contains_range(ctx.frange.range);

                if cursor_in_range {
                    let exprto = unwrap_trivial_block(block_expr);
                    expr_to_unwrap = Some(exprto);
                    break;
                }
            }
        }
        let expr_to_unwrap = expr_to_unwrap?;
        // Find if we are in a else if block
        let ancestor = ctx
            .sema
            .ancestors_with_macros(if_expr.syntax().clone())
            .skip(1)
            .find_map(ast::IfExpr::cast);

        if let Some(ancestor) = ancestor {
            Some((ast::Expr::IfExpr(ancestor), expr_to_unwrap))
        } else {
            Some((ast::Expr::IfExpr(if_expr), expr_to_unwrap))
        }
    } else if let Some(for_expr) = ctx.find_node_at_offset::<ast::ForExpr>() {
        // for expression
        let block_expr = for_expr.loop_body()?;
        let block = block_expr.block()?;
        let cursor_in_range = block.l_curly_token()?.text_range().contains_range(ctx.frange.range);

        if cursor_in_range {
            let expr_to_unwrap = unwrap_trivial_block(block_expr);

            Some((ast::Expr::ForExpr(for_expr), expr_to_unwrap))
        } else {
            None
        }
    } else if let Some(while_expr) = ctx.find_node_at_offset::<ast::WhileExpr>() {
        // while expression
        let block_expr = while_expr.loop_body()?;
        let block = block_expr.block()?;
        let cursor_in_range = block.l_curly_token()?.text_range().contains_range(ctx.frange.range);

        if cursor_in_range {
            let expr_to_unwrap = unwrap_trivial_block(block_expr);

            Some((ast::Expr::WhileExpr(while_expr), expr_to_unwrap))
        } else {
            None
        }
    } else if let Some(loop_expr) = ctx.find_node_at_offset::<ast::LoopExpr>() {
        // loop expression
        let block_expr = loop_expr.loop_body()?;
        let block = block_expr.block()?;
        let cursor_in_range = block.l_curly_token()?.text_range().contains_range(ctx.frange.range);

        if cursor_in_range {
            let expr_to_unwrap = unwrap_trivial_block(block_expr);

            Some((ast::Expr::LoopExpr(loop_expr), expr_to_unwrap))
        } else {
            None
        }
    } else {
        None
    };

    let (expr, expr_to_unwrap) = res?;
    ctx.add_assist(AssistId("unwrap_block"), "Unwrap block", |edit| {
        edit.set_cursor(expr.syntax().text_range().start());
        edit.target(expr_to_unwrap.syntax().text_range());

        let pat_start: &[_] = &[' ', '{', '\n'];
        let expr_to_unwrap = expr_to_unwrap.to_string();
        let expr_string = expr_to_unwrap.trim_start_matches(pat_start);
        let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
        expr_string_lines.pop(); // Delete last line

        let expr_string = expr_string_lines
            .into_iter()
            .map(|line| line.replacen("    ", "", 1)) // Delete indentation
            .collect::<Vec<String>>()
            .join("\n");

        edit.replace(expr.syntax().text_range(), expr_string);
    })
}

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

    use super::*;

    #[test]
    fn simple_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {<|>
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>foo();

                //comment
                bar();
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                } else {<|>
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {<|>
                    println!("bar");
                } else {
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                <|>println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                bar();<|>
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn issue_example_with_if() {
        check_assist(
            unwrap_block,
            r#"
            fn complete_enum_variants(acc: &mut Completions, ctx: &CompletionContext, ty: &Type) {
                if let Some(ty) = &ctx.expected_type {<|>
                    if let Some(Adt::Enum(enum_data)) = ty.as_adt() {
                        let variants = enum_data.variants(ctx.db);

                        let module = if let Some(module) = ctx.scope().module() {
                            // Compute path from the completion site if available.
                            module
                        } else {
                            // Otherwise fall back to the enum's definition site.
                            enum_data.module(ctx.db)
                        };

                        for variant in variants {
                            if let Some(path) = module.find_use_path(ctx.db, ModuleDef::from(variant)) {
                                // Variants with trivial paths are already added by the existing completion logic,
                                // so we should avoid adding these twice
                                if path.segments.len() > 1 {
                                    acc.add_enum_variant(ctx, variant, Some(path.to_string()));
                                }
                            }
                        }
                    }
                }
            }
            "#,
            r#"
            fn complete_enum_variants(acc: &mut Completions, ctx: &CompletionContext, ty: &Type) {
                <|>if let Some(Adt::Enum(enum_data)) = ty.as_adt() {
                    let variants = enum_data.variants(ctx.db);

                    let module = if let Some(module) = ctx.scope().module() {
                        // Compute path from the completion site if available.
                        module
                    } else {
                        // Otherwise fall back to the enum's definition site.
                        enum_data.module(ctx.db)
                    };

                    for variant in variants {
                        if let Some(path) = module.find_use_path(ctx.db, ModuleDef::from(variant)) {
                            // Variants with trivial paths are already added by the existing completion logic,
                            // so we should avoid adding these twice
                            if path.segments.len() > 1 {
                                acc.add_enum_variant(ctx, variant, Some(path.to_string()));
                            }
                        }
                    }
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {
                    if true {<|>
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                for i in 0..5 {
                    <|>foo();

                    //comment
                    bar();
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_loop() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                loop {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_while() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                while true {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_while_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                while true {
                    if true {
                        foo();<|>

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
        );
    }
}
<commit_msg>Add unwrap block assist #4156<commit_after>use crate::{Assist, AssistCtx, AssistId};

use ast::{BlockExpr, Expr, LoopBodyOwner};
use ra_fmt::unwrap_trivial_block;
use ra_syntax::{ast, AstNode, TextRange};

// Assist: unwrap_block
//
// This assist removes if...else, for, while and loop control statements to just keep the body.
//
// ```
// fn foo() {
//     if true {<|>
//         println!("foo");
//     }
// }
// ```
// ->
// ```
// fn foo() {
//     println!("foo");
// }
// ```
pub(crate) fn unwrap_block(ctx: AssistCtx) -> Option<Assist> {
    let res = if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
        // if expression
        let mut expr_to_unwrap: Option<ast::Expr> = None;
        for block_expr in if_expr.blocks() {
            if let Some(expr) = excract_expr(ctx.frange.range, block_expr) {
                expr_to_unwrap = Some(expr);
                break;
            }
        }
        let expr_to_unwrap = expr_to_unwrap?;
        // Find if we are in a else if block
        let ancestor = if_expr.syntax().ancestors().skip(1).find_map(ast::IfExpr::cast);

        if let Some(ancestor) = ancestor {
            Some((ast::Expr::IfExpr(ancestor), expr_to_unwrap))
        } else {
            Some((ast::Expr::IfExpr(if_expr), expr_to_unwrap))
        }
    } else if let Some(for_expr) = ctx.find_node_at_offset::<ast::ForExpr>() {
        // for expression
        let block_expr = for_expr.loop_body()?;
        excract_expr(ctx.frange.range, block_expr)
            .map(|expr_to_unwrap| (ast::Expr::ForExpr(for_expr), expr_to_unwrap))
    } else if let Some(while_expr) = ctx.find_node_at_offset::<ast::WhileExpr>() {
        // while expression
        let block_expr = while_expr.loop_body()?;
        excract_expr(ctx.frange.range, block_expr)
            .map(|expr_to_unwrap| (ast::Expr::WhileExpr(while_expr), expr_to_unwrap))
    } else if let Some(loop_expr) = ctx.find_node_at_offset::<ast::LoopExpr>() {
        // loop expression
        let block_expr = loop_expr.loop_body()?;
        excract_expr(ctx.frange.range, block_expr)
            .map(|expr_to_unwrap| (ast::Expr::LoopExpr(loop_expr), expr_to_unwrap))
    } else {
        None
    };

    let (expr, expr_to_unwrap) = res?;
    ctx.add_assist(AssistId("unwrap_block"), "Unwrap block", |edit| {
        edit.set_cursor(expr.syntax().text_range().start());
        edit.target(expr_to_unwrap.syntax().text_range());

        let pat_start: &[_] = &[' ', '{', '\n'];
        let expr_to_unwrap = expr_to_unwrap.to_string();
        let expr_string = expr_to_unwrap.trim_start_matches(pat_start);
        let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
        expr_string_lines.pop(); // Delete last line

        let expr_string = expr_string_lines
            .into_iter()
            .map(|line| line.replacen("    ", "", 1)) // Delete indentation
            .collect::<Vec<String>>()
            .join("\n");

        edit.replace(expr.syntax().text_range(), expr_string);
    })
}

fn excract_expr(cursor_range: TextRange, block_expr: BlockExpr) -> Option<Expr> {
    let block = block_expr.block()?;
    let cursor_in_range = block.l_curly_token()?.text_range().contains_range(cursor_range);

    if cursor_in_range {
        Some(unwrap_trivial_block(block_expr))
    } else {
        None
    }
}

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

    use super::*;

    #[test]
    fn simple_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {<|>
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>foo();

                //comment
                bar();
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                } else {<|>
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {<|>
                    println!("bar");
                } else {
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                <|>println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                bar();<|>
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn issue_example_with_if() {
        check_assist(
            unwrap_block,
            r#"
            fn complete_enum_variants(acc: &mut Completions, ctx: &CompletionContext, ty: &Type) {
                if let Some(ty) = &ctx.expected_type {<|>
                    if let Some(Adt::Enum(enum_data)) = ty.as_adt() {
                        let variants = enum_data.variants(ctx.db);

                        let module = if let Some(module) = ctx.scope().module() {
                            // Compute path from the completion site if available.
                            module
                        } else {
                            // Otherwise fall back to the enum's definition site.
                            enum_data.module(ctx.db)
                        };

                        for variant in variants {
                            if let Some(path) = module.find_use_path(ctx.db, ModuleDef::from(variant)) {
                                // Variants with trivial paths are already added by the existing completion logic,
                                // so we should avoid adding these twice
                                if path.segments.len() > 1 {
                                    acc.add_enum_variant(ctx, variant, Some(path.to_string()));
                                }
                            }
                        }
                    }
                }
            }
            "#,
            r#"
            fn complete_enum_variants(acc: &mut Completions, ctx: &CompletionContext, ty: &Type) {
                <|>if let Some(Adt::Enum(enum_data)) = ty.as_adt() {
                    let variants = enum_data.variants(ctx.db);

                    let module = if let Some(module) = ctx.scope().module() {
                        // Compute path from the completion site if available.
                        module
                    } else {
                        // Otherwise fall back to the enum's definition site.
                        enum_data.module(ctx.db)
                    };

                    for variant in variants {
                        if let Some(path) = module.find_use_path(ctx.db, ModuleDef::from(variant)) {
                            // Variants with trivial paths are already added by the existing completion logic,
                            // so we should avoid adding these twice
                            if path.segments.len() > 1 {
                                acc.add_enum_variant(ctx, variant, Some(path.to_string()));
                            }
                        }
                    }
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {
                    if true {<|>
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                for i in 0..5 {
                    <|>foo();

                    //comment
                    bar();
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_loop() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                loop {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_while() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                while true {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_while_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                while true {
                    if true {
                        foo();<|>

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
        );
    }
}
<|endoftext|>