<commit_before>use nalgebra::{Point3, Scalar};
use num::Float;
use num::traits::FloatConst;
use std::cmp;
use std::marker::PhantomData;

use super::generate::{ConjointPointGenerator, IndexPolygonGenerator, PolygonGenerator};
use super::primitive::{Polygon, Triangle, Quad};

#[derive(Clone)]
pub struct UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    nu: usize, // Meridians.
    nv: usize, // Parallels.
    phantom: PhantomData<T>,
}

impl<T> UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    pub fn with_unit_radius(nu: usize, nv: usize) -> Self {
        let nu = cmp::max(3, nu);
        let nv = cmp::max(2, nv);
        UVSphere {
            nu: nu,
            nv: nv,
            phantom: PhantomData,
        }
    }

    fn point(&self, u: usize, v: usize) -> Point3<T> {
        let u = (T::from(u).unwrap() / T::from(self.nu).unwrap()) * T::PI() * (T::one() + T::one());
        let v = (T::from(v).unwrap() / T::from(self.nv).unwrap()) * T::PI();
        Point3::new(u.cos() * v.sin(), u.sin() * v.sin(), v.cos())
    }

    fn index_point(&self, u: usize, v: usize) -> usize {
        if v == 0 {
            0
        }
        else if v == self.nv {
            ((self.nv - 1) * self.nu) + 1
        }
        else {
            ((v - 1) * self.nu) + (u % self.nu) + 1
        }
    }

    fn map_polygon_index(&self, index: usize) -> (usize, usize) {
        (index % self.nu, index / self.nu)
    }
}

impl<T> ConjointPointGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Point3<T>;

    fn conjoint_point(&self, index: usize) -> Self::Output {
        if index == 0 {
            self.point(0, 0)
        }
        else if index == self.conjoint_point_count() - 1 {
            self.point(0, self.nv)
        }
        else {
            let index = index - 1;
            self.point(index % self.nu, (index / self.nv) + 1)
        }
    }

    fn conjoint_point_count(&self) -> usize {
        (self.nv - 1) * self.nu + 2
    }
}

impl<T> PolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<Point3<T>>;

    fn polygon(&self, index: usize) -> Self::Output {
        let (u, v) = self.map_polygon_index(index);

        // Generate the points at the requested meridian and parallel. The
        // upper and lower bounds of (u, v) are always used, so generate them
        // in advance (`low` and `high`). Emit triangles at the poles,
        // otherwise quads.
        let low = self.point(u, v);
        let high = self.point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.point(u, v + 1),
                high,
                self.point(u + 1, v),
            ))
        }
    }

    fn polygon_count(&self) -> usize {
        self.nu * self.nv
    }
}

impl<T> IndexPolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<usize>;

    fn index_polygon(&self, index: usize) -> <Self as IndexPolygonGenerator>::Output {
        let (u, v) = self.map_polygon_index(index);

        let low = self.index_point(u, v);
        let high = self.index_point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.index_point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.index_point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.index_point(u, v + 1),
                high,
                self.index_point(u + 1, v),
            ))
        }
    }
}
<commit_msg>Add sanity tests to `mesh::sphere` module.<commit_after>use nalgebra::{Point3, Scalar};
use num::Float;
use num::traits::FloatConst;
use std::cmp;
use std::marker::PhantomData;

use super::generate::{ConjointPointGenerator, IndexPolygonGenerator, PolygonGenerator};
use super::primitive::{Polygon, Triangle, Quad};

#[derive(Clone)]
pub struct UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    nu: usize, // Meridians.
    nv: usize, // Parallels.
    phantom: PhantomData<T>,
}

impl<T> UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    pub fn with_unit_radius(nu: usize, nv: usize) -> Self {
        let nu = cmp::max(3, nu);
        let nv = cmp::max(2, nv);
        UVSphere {
            nu: nu,
            nv: nv,
            phantom: PhantomData,
        }
    }

    fn point(&self, u: usize, v: usize) -> Point3<T> {
        let u = (T::from(u).unwrap() / T::from(self.nu).unwrap()) * T::PI() * (T::one() + T::one());
        let v = (T::from(v).unwrap() / T::from(self.nv).unwrap()) * T::PI();
        Point3::new(u.cos() * v.sin(), u.sin() * v.sin(), v.cos())
    }

    fn index_point(&self, u: usize, v: usize) -> usize {
        if v == 0 {
            0
        }
        else if v == self.nv {
            ((self.nv - 1) * self.nu) + 1
        }
        else {
            ((v - 1) * self.nu) + (u % self.nu) + 1
        }
    }

    fn map_polygon_index(&self, index: usize) -> (usize, usize) {
        (index % self.nu, index / self.nu)
    }
}

impl<T> ConjointPointGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Point3<T>;

    fn conjoint_point(&self, index: usize) -> Self::Output {
        if index == 0 {
            self.point(0, 0)
        }
        else if index == self.conjoint_point_count() - 1 {
            self.point(0, self.nv)
        }
        else {
            let index = index - 1;
            self.point(index % self.nu, (index / self.nv) + 1)
        }
    }

    fn conjoint_point_count(&self) -> usize {
        (self.nv - 1) * self.nu + 2
    }
}

impl<T> PolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<Point3<T>>;

    fn polygon(&self, index: usize) -> Self::Output {
        let (u, v) = self.map_polygon_index(index);

        // Generate the points at the requested meridian and parallel. The
        // upper and lower bounds of (u, v) are always used, so generate them
        // in advance (`low` and `high`). Emit triangles at the poles,
        // otherwise quads.
        let low = self.point(u, v);
        let high = self.point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.point(u, v + 1),
                high,
                self.point(u + 1, v),
            ))
        }
    }

    fn polygon_count(&self) -> usize {
        self.nu * self.nv
    }
}

impl<T> IndexPolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<usize>;

    fn index_polygon(&self, index: usize) -> <Self as IndexPolygonGenerator>::Output {
        let (u, v) = self.map_polygon_index(index);

        let low = self.index_point(u, v);
        let high = self.index_point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.index_point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.index_point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.index_point(u, v + 1),
                high,
                self.index_point(u + 1, v),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;
    use std::iter::FromIterator;

    use super::super::*;

    #[test]
    fn conjoint_point_count() {
        assert_eq!(
            5,
            sphere::UVSphere::<f32>::with_unit_radius(3, 2)
                .conjoint_points() // 5 conjoint points.
                .count()
        );
    }

    #[test]
    fn polygon_point_count() {
        assert_eq!(
            18,
            sphere::UVSphere::<f32>::with_unit_radius(3, 2)
                .polygons() // 6 triangles, 18 points.
                .points()
                .count()
        );
    }

    #[test]
    fn index_to_conjoint_point_mapping() {
        assert_eq!(
            5,
            BTreeSet::from_iter(
                sphere::UVSphere::<f32>::with_unit_radius(3, 2)
                    .index_polygons() // 18 points, 5 indeces.
                    .points()
            ).len()
        )
    }
}
<|endoftext|>


<commit_before>use postgres;
use uuid::Uuid;
use std::fmt;
use chrono::{NaiveDateTime, Utc};

use apple_music;
use spotify;
use error::Error;
use super::{conn, Model};
use model::provider::Provider;
use model::state::State;
use model::enclosure::Enclosure;
use model::track::Track;
use model::artist::Artist;

static PROPS: [&'static str; 14]  = ["id",
                                     "provider",
                                     "identifier",
                                     "owner_id",
                                     "owner_name",
                                     "url",
                                     "title",
                                     "description",
                                     "thumbnail_url",
                                     "artwork_url",
                                     "published_at",
                                     "created_at",
                                     "updated_at",
                                     "state"];

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Album {
    pub id:            Uuid,
    pub provider:      Provider,
    pub identifier:    String,
    pub owner_id:      Option<String>,
    pub owner_name:    Option<String>,
    pub url:           String,
    pub title:         String,
    pub description:   Option<String>,
    pub thumbnail_url: Option<String>,
    pub artwork_url:   Option<String>,
    pub published_at:  NaiveDateTime,
    pub created_at:    NaiveDateTime,
    pub updated_at:    NaiveDateTime,
    pub state:         State,
    pub tracks:        Vec<Track>,
    pub artists:       Option<Vec<Artist>>,
}

impl PartialEq for Album {
    fn eq(&self, p: &Album) -> bool {
        return self.identifier == p.identifier && self.provider == p.provider
    }
}

impl fmt::Display for Album {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.provider, self.identifier)
    }
}

impl<'a> Model<'a> for Album {
    fn table_name() -> String { "albums".to_string() }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }

    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Album> {
        let mut albums = Vec::new();
        for row in rows.iter() {
            let album = Album {
                id:            row.get(0),
                provider:      Provider::new(row.get(1)),
                identifier:    row.get(2),
                owner_id:      row.get(3),
                owner_name:    row.get(4),
                url:           row.get(5),
                title:         row.get(6),
                description:   row.get(7),
                thumbnail_url: row.get(8),
                artwork_url:   row.get(9),
                published_at:  row.get(10),
                created_at:    row.get(11),
                updated_at:    row.get(12),
                state:         State::new(row.get(13)),
                tracks:        vec![],
                artists:       None,
            };
            albums.push(album)
        }
        albums
    }

    fn create(&self) -> Result<Album, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO albums (provider, identifier, url, title)
                                      VALUES ($1, $2, $3, $4) RETURNING id"));
        let rows = try!(stmt.query(&[&self.provider.to_string(), &self.identifier, &self.url, &self.title]));
        let mut album = self.clone();
        for row in rows.iter() {
            album.id = row.get(0);
        }
        Ok(album)
    }

    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = Utc::now().naive_utc();
        let conn = try!(conn());
        let stmt = try!(conn.prepare("UPDATE albums SET
                                      provider      = $2,
                                      identifier    = $3,
                                      owner_id      = $4,
                                      owner_name    = $5,
                                      url           = $6,
                                      title         = $7,
                                      description   = $8,
                                      thumbnail_url = $9,
                                      artwork_url   = $10,
                                      published_at  = $11,
                                      created_at    = $12,
                                      updated_at    = $13,
                                      state         = $14
                                      WHERE id = $1"));
        let result = stmt.query(&[&self.id,
                                  &self.provider.to_string(),
                                  &self.identifier,
                                  &self.owner_id,
                                  &self.owner_name,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.thumbnail_url,
                                  &self.artwork_url,
                                  &self.published_at,
                                  &self.created_at,
                                  &self.updated_at,
                                  &self.state.to_string(),
        ]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected),
        }
    }
}

impl<'a> Enclosure<'a> for Album {
    fn new(provider: Provider, identifier: String) -> Album {
        Album {
            id:            Uuid::new_v4(),
            provider:      provider,
            identifier:    identifier,
            owner_id:      None,
            owner_name:    None,
            url:           "".to_string(),
            title:         "".to_string(),
            description:   None,
            thumbnail_url: None,
            artwork_url:   None,
            published_at:  Utc::now().naive_utc(),
            created_at:    Utc::now().naive_utc(),
            updated_at:    Utc::now().naive_utc(),
            state:         State::Alive,
            tracks:        vec![],
            artists:       None,
        }
    }
    fn set_url(&mut self, url: String) -> &mut Album {
        self.url = url;
        self
    }
    fn set_owner_id(&mut self, owner_id: Option<String>) -> &mut Album {
        self.owner_id = owner_id;
        self
    }
    fn fetch_props(&mut self) -> Result<(), Error> {
        match self.provider {
            Provider::AppleMusic => {
                let country = apple_music::country(&self.url);
                match apple_music::fetch_album(&self.identifier, &country) {
                    Ok(album) => self.update_with_am_album(&album),
                    Err(_)    => self.disable(),
                }
            }
            Provider::Spotify => match spotify::fetch_album(&self.identifier) {
                Ok(album) => self.update_with_sp_album(&album),
                Err(_)    => self.disable(),
            },
            _ => self,
        };
        match self.state {
            State::Alive => Ok(()),
            State::Dead  => Err(Error::NotFound),
        }
    }
    fn find_by_entry_id(entry_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums p LEFT JOIN album_entries pe
                        ON p.id = pe.album_id
                        WHERE pe.entry_id = $1
                        ORDER BY p.published_at DESC",
                     Album::props_str("p."))).unwrap();
        let rows = stmt.query(&[&entry_id]).unwrap();
        Album::rows_to_items(rows)
    }
}

impl Album {
    pub fn find_all() -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums ORDER BY albums.published_at DESC",
                     Album::props_str(""))).unwrap();
        let rows = stmt.query(&[]).unwrap();
        Album::rows_to_items(rows)
    }
    fn add_track(&mut self, track: &Track) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_tracks (track_id, album_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&track.id, &self.id]));
        Ok(())
    }
    fn add_artist(&mut self, artist: Artist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_artists (album_id, artist_id) VALUES ($1, $2)"));
        try!(stmt.query(&[&self.id, &artist.id]));
        match self.artists {
            Some(ref mut artists) => artists.push(artist),
            None                  => self.artists = Some(vec![artist]),

        }
        Ok(())
    }

    pub fn from_sp_album(album: &spotify::Album) -> Album {
        Album::new(Provider::Spotify, (*album).id.to_string())
            .update_with_sp_album(album)
            .clone()
    }

    pub fn from_am_album(album: &apple_music::Album) -> Album {
        Album::new(Provider::AppleMusic, (*album).id.to_string())
            .update_with_am_album(album)
            .clone()
    }

    fn add_tracks(&mut self, tracks: Vec<Track>) {
        self.tracks = tracks.iter().map(|t| {
            let mut t = t.clone();
            if let Ok(new_track) = Track::find_or_create(t.provider,
                                                         t.identifier.to_string()) {
                t.id      = new_track.id;
                let _     = self.add_track(&t);
            };
            t
        }).collect::<Vec<_>>();
    }

    pub fn update_with_sp_album(&mut self, album: &spotify::Album) -> &mut Album {
        self.provider       = Provider::Spotify;
        self.identifier     = album.id.to_string();
        if album.artists.len() > 0 {
            self.owner_id   = Some(album.artists[0].id.clone());
            self.owner_name = Some(album.artists[0].name.clone());
        }
        self.url            = album.uri.clone();
        self.title          = album.name.clone();
        self.description    = None;
        self.state          = State::Alive;
        self.published_at   = Utc::now().naive_utc();
        if album.images.len() > 0 {
            self.artwork_url   = Some(album.images[0].url.clone());
            self.thumbnail_url = Some(album.images[0].url.clone());
        }
        if album.images.len() > 1 {
            self.thumbnail_url = Some(album.images[1].url.clone());
        }
        let artists = album.artists
            .iter()
            .map(|artist| {
                Artist::find_or_create(self.provider, artist.id.to_string())
                    .map(|mut a| a.update_with_sp_artist(&artist).clone())
            })
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<Artist>>();
        for mut a in artists {
            let _ = a.save();
            let _ = self.add_artist(a);
        }
        let tracks = album.tracks.clone()
            .map(|t| t.items).unwrap_or(vec![]).iter()
            .map(|ref t| Track::from_sp_track(t))
            .map(|ref mut t| t.update_with_sp_album(&album).clone())
            .collect::<Vec<_>>();
        self.add_tracks(tracks);
        self
    }

    pub fn update_with_am_album(&mut self, album: &apple_music::Album) -> &mut Album {
        self.provider      = Provider::AppleMusic;
        self.identifier    = album.id.to_string();
        self.owner_id      = Some(album.album_artist.to_string());
        self.owner_name    = Some(album.album_artist.to_string());
        self.url           = album.music_url.to_string();
        self.title         = album.title.to_string();
        self.description   = None;
        self.thumbnail_url = Some(album.artwork_url.to_string());
        self.artwork_url   = Some(album.artwork_url.to_string());
        self.state         = State::Alive;
        if let Ok(mut artist) = Artist::find_or_create(self.provider,
                                                       album.album_artist.to_string()) {
            artist.name  = album.album_artist.to_string();
            let _ = artist.save();
            let _ = self.add_artist(artist);
        }
        self
    }

    pub fn disable(&mut self) -> &mut Album {
        self.state = State::Dead;
        self
    }

    pub fn delete(&self) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("DELETE FROM albums WHERE id=$1"));
        let result = stmt.query(&[&self.id]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected)
        }
    }
}

#[cfg(test)]
mod test {
    use model::enclosure::Enclosure;
    use model::Model;
    use super::Album;
    use Provider;
    #[test]
    fn test_new() {
        let album = Album::new(Provider::YouTube,
                               "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string());
        assert_eq!(album.provider, Provider::YouTube);
        assert_eq!(&album.identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX")
    }
    #[test]
    fn test_find_of_create() {
        let identifier = "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string();
        let result     = Album::find_or_create(Provider::YouTube, identifier);
        assert!(result.is_ok());
    }
    #[test]
    fn test_delete() {
        let identifier = "test_delete".to_string();
        let album      = Album::find_or_create(Provider::YouTube, identifier).unwrap();
        let result     = album.delete();
        assert!(result.is_ok());
    }
    #[test]
    fn test_save() {
        let id = "test_save";
        let mut album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        album.title   = "title".to_string();
        let result    = album.save();
        assert!(result.is_ok());
        let album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        assert_eq!(&album.title, "title");
    }
}
<commit_msg>Implement Album::find_by_artist<commit_after>use postgres;
use uuid::Uuid;
use std::fmt;
use chrono::{NaiveDateTime, Utc};

use apple_music;
use spotify;
use error::Error;
use super::{conn, Model};
use model::provider::Provider;
use model::state::State;
use model::enclosure::Enclosure;
use model::track::Track;
use model::artist::Artist;

static PROPS: [&'static str; 14]  = ["id",
                                     "provider",
                                     "identifier",
                                     "owner_id",
                                     "owner_name",
                                     "url",
                                     "title",
                                     "description",
                                     "thumbnail_url",
                                     "artwork_url",
                                     "published_at",
                                     "created_at",
                                     "updated_at",
                                     "state"];

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Album {
    pub id:            Uuid,
    pub provider:      Provider,
    pub identifier:    String,
    pub owner_id:      Option<String>,
    pub owner_name:    Option<String>,
    pub url:           String,
    pub title:         String,
    pub description:   Option<String>,
    pub thumbnail_url: Option<String>,
    pub artwork_url:   Option<String>,
    pub published_at:  NaiveDateTime,
    pub created_at:    NaiveDateTime,
    pub updated_at:    NaiveDateTime,
    pub state:         State,
    pub tracks:        Vec<Track>,
    pub artists:       Option<Vec<Artist>>,
}

impl PartialEq for Album {
    fn eq(&self, p: &Album) -> bool {
        return self.identifier == p.identifier && self.provider == p.provider
    }
}

impl fmt::Display for Album {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.provider, self.identifier)
    }
}

impl<'a> Model<'a> for Album {
    fn table_name() -> String { "albums".to_string() }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }

    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Album> {
        let mut albums = Vec::new();
        for row in rows.iter() {
            let album = Album {
                id:            row.get(0),
                provider:      Provider::new(row.get(1)),
                identifier:    row.get(2),
                owner_id:      row.get(3),
                owner_name:    row.get(4),
                url:           row.get(5),
                title:         row.get(6),
                description:   row.get(7),
                thumbnail_url: row.get(8),
                artwork_url:   row.get(9),
                published_at:  row.get(10),
                created_at:    row.get(11),
                updated_at:    row.get(12),
                state:         State::new(row.get(13)),
                tracks:        vec![],
                artists:       None,
            };
            albums.push(album)
        }
        albums
    }

    fn create(&self) -> Result<Album, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO albums (provider, identifier, url, title)
                                      VALUES ($1, $2, $3, $4) RETURNING id"));
        let rows = try!(stmt.query(&[&self.provider.to_string(), &self.identifier, &self.url, &self.title]));
        let mut album = self.clone();
        for row in rows.iter() {
            album.id = row.get(0);
        }
        Ok(album)
    }

    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = Utc::now().naive_utc();
        let conn = try!(conn());
        let stmt = try!(conn.prepare("UPDATE albums SET
                                      provider      = $2,
                                      identifier    = $3,
                                      owner_id      = $4,
                                      owner_name    = $5,
                                      url           = $6,
                                      title         = $7,
                                      description   = $8,
                                      thumbnail_url = $9,
                                      artwork_url   = $10,
                                      published_at  = $11,
                                      created_at    = $12,
                                      updated_at    = $13,
                                      state         = $14
                                      WHERE id = $1"));
        let result = stmt.query(&[&self.id,
                                  &self.provider.to_string(),
                                  &self.identifier,
                                  &self.owner_id,
                                  &self.owner_name,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.thumbnail_url,
                                  &self.artwork_url,
                                  &self.published_at,
                                  &self.created_at,
                                  &self.updated_at,
                                  &self.state.to_string(),
        ]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected),
        }
    }
}

impl<'a> Enclosure<'a> for Album {
    fn new(provider: Provider, identifier: String) -> Album {
        Album {
            id:            Uuid::new_v4(),
            provider:      provider,
            identifier:    identifier,
            owner_id:      None,
            owner_name:    None,
            url:           "".to_string(),
            title:         "".to_string(),
            description:   None,
            thumbnail_url: None,
            artwork_url:   None,
            published_at:  Utc::now().naive_utc(),
            created_at:    Utc::now().naive_utc(),
            updated_at:    Utc::now().naive_utc(),
            state:         State::Alive,
            tracks:        vec![],
            artists:       None,
        }
    }
    fn set_url(&mut self, url: String) -> &mut Album {
        self.url = url;
        self
    }
    fn set_owner_id(&mut self, owner_id: Option<String>) -> &mut Album {
        self.owner_id = owner_id;
        self
    }
    fn fetch_props(&mut self) -> Result<(), Error> {
        match self.provider {
            Provider::AppleMusic => {
                let country = apple_music::country(&self.url);
                match apple_music::fetch_album(&self.identifier, &country) {
                    Ok(album) => self.update_with_am_album(&album),
                    Err(_)    => self.disable(),
                }
            }
            Provider::Spotify => match spotify::fetch_album(&self.identifier) {
                Ok(album) => self.update_with_sp_album(&album),
                Err(_)    => self.disable(),
            },
            _ => self,
        };
        match self.state {
            State::Alive => Ok(()),
            State::Dead  => Err(Error::NotFound),
        }
    }
    fn find_by_entry_id(entry_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums p LEFT JOIN album_entries pe
                        ON p.id = pe.album_id
                        WHERE pe.entry_id = $1
                        ORDER BY p.published_at DESC",
                     Album::props_str("p."))).unwrap();
        let rows = stmt.query(&[&entry_id]).unwrap();
        Album::rows_to_items(rows)
    }
}

impl Album {
    pub fn find_all() -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums ORDER BY albums.published_at DESC",
                     Album::props_str(""))).unwrap();
        let rows = stmt.query(&[]).unwrap();
        Album::rows_to_items(rows)
    }
    fn add_track(&mut self, track: &Track) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_tracks (track_id, album_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&track.id, &self.id]));
        Ok(())
    }
    fn add_artist(&mut self, artist: Artist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_artists (album_id, artist_id) VALUES ($1, $2)"));
        try!(stmt.query(&[&self.id, &artist.id]));
        match self.artists {
            Some(ref mut artists) => artists.push(artist),
            None                  => self.artists = Some(vec![artist]),

        }
        Ok(())
    }
    pub fn find_by_artist(artist_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums
                      LEFT OUTER JOIN album_artists ON album_artists.album_id = albums.id
                      WHERE album_artists.artist_id = $1 ORDER BY albums.created_at DESC",
                     Album::props_str("albums."))).unwrap();
        let rows = stmt.query(&[&artist_id]).unwrap();
        Album::rows_to_items(rows)
    }
    pub fn from_sp_album(album: &spotify::Album) -> Album {
        Album::new(Provider::Spotify, (*album).id.to_string())
            .update_with_sp_album(album)
            .clone()
    }

    pub fn from_am_album(album: &apple_music::Album) -> Album {
        Album::new(Provider::AppleMusic, (*album).id.to_string())
            .update_with_am_album(album)
            .clone()
    }

    fn add_tracks(&mut self, tracks: Vec<Track>) {
        self.tracks = tracks.iter().map(|t| {
            let mut t = t.clone();
            if let Ok(new_track) = Track::find_or_create(t.provider,
                                                         t.identifier.to_string()) {
                t.id      = new_track.id;
                let _     = self.add_track(&t);
            };
            t
        }).collect::<Vec<_>>();
    }

    pub fn update_with_sp_album(&mut self, album: &spotify::Album) -> &mut Album {
        self.provider       = Provider::Spotify;
        self.identifier     = album.id.to_string();
        if album.artists.len() > 0 {
            self.owner_id   = Some(album.artists[0].id.clone());
            self.owner_name = Some(album.artists[0].name.clone());
        }
        self.url            = album.uri.clone();
        self.title          = album.name.clone();
        self.description    = None;
        self.state          = State::Alive;
        self.published_at   = Utc::now().naive_utc();
        if album.images.len() > 0 {
            self.artwork_url   = Some(album.images[0].url.clone());
            self.thumbnail_url = Some(album.images[0].url.clone());
        }
        if album.images.len() > 1 {
            self.thumbnail_url = Some(album.images[1].url.clone());
        }
        let artists = album.artists
            .iter()
            .map(|artist| {
                Artist::find_or_create(self.provider, artist.id.to_string())
                    .map(|mut a| a.update_with_sp_artist(&artist).clone())
            })
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<Artist>>();
        for mut a in artists {
            let _ = a.save();
            let _ = self.add_artist(a);
        }
        let tracks = album.tracks.clone()
            .map(|t| t.items).unwrap_or(vec![]).iter()
            .map(|ref t| Track::from_sp_track(t))
            .map(|ref mut t| t.update_with_sp_album(&album).clone())
            .collect::<Vec<_>>();
        self.add_tracks(tracks);
        self
    }

    pub fn update_with_am_album(&mut self, album: &apple_music::Album) -> &mut Album {
        self.provider      = Provider::AppleMusic;
        self.identifier    = album.id.to_string();
        self.owner_id      = Some(album.album_artist.to_string());
        self.owner_name    = Some(album.album_artist.to_string());
        self.url           = album.music_url.to_string();
        self.title         = album.title.to_string();
        self.description   = None;
        self.thumbnail_url = Some(album.artwork_url.to_string());
        self.artwork_url   = Some(album.artwork_url.to_string());
        self.state         = State::Alive;
        if let Ok(mut artist) = Artist::find_or_create(self.provider,
                                                       album.album_artist.to_string()) {
            artist.name  = album.album_artist.to_string();
            let _ = artist.save();
            let _ = self.add_artist(artist);
        }
        self
    }

    pub fn disable(&mut self) -> &mut Album {
        self.state = State::Dead;
        self
    }

    pub fn delete(&self) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("DELETE FROM albums WHERE id=$1"));
        let result = stmt.query(&[&self.id]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected)
        }
    }
}

#[cfg(test)]
mod test {
    use model::enclosure::Enclosure;
    use model::Model;
    use super::Album;
    use Provider;
    #[test]
    fn test_new() {
        let album = Album::new(Provider::YouTube,
                               "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string());
        assert_eq!(album.provider, Provider::YouTube);
        assert_eq!(&album.identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX")
    }
    #[test]
    fn test_find_of_create() {
        let identifier = "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string();
        let result     = Album::find_or_create(Provider::YouTube, identifier);
        assert!(result.is_ok());
    }
    #[test]
    fn test_delete() {
        let identifier = "test_delete".to_string();
        let album      = Album::find_or_create(Provider::YouTube, identifier).unwrap();
        let result     = album.delete();
        assert!(result.is_ok());
    }
    #[test]
    fn test_save() {
        let id = "test_save";
        let mut album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        album.title   = "title".to_string();
        let result    = album.save();
        assert!(result.is_ok());
        let album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        assert_eq!(&album.title, "title");
    }
}
<|endoftext|>


<commit_before>use super::prelude::*;

use conduit_cookie::RequestSession;
use diesel::prelude::*;

use crate::db::RequestTransaction;
use crate::util::errors::{std_error, AppResult, ChainError, Unauthorized};

use crate::models::ApiToken;
use crate::models::user::{UserNoEmailType, ALL_COLUMNS};
use crate::models::User;
use crate::schema::users;

#[derive(Debug, Clone, Copy)]
pub struct CurrentUser;

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum AuthenticationSource {
    SessionCookie,
    ApiToken { api_token_id: i32 },
}

impl Middleware for CurrentUser {
    fn before(&self, req: &mut dyn Request) -> Result<(), Box<dyn Error + Send>> {
        // Check if the request has a session cookie with a `user_id` property inside
        let id = {
            req.session()
                .get("user_id")
                .and_then(|s| s.parse::<i32>().ok())
        };

        let conn = req.db_conn().map_err(std_error)?;

        if let Some(id) = id {
            // If it did, look for a user in the database with the given `user_id`
            let maybe_user = users::table
                .select(ALL_COLUMNS)
                .find(id)
                .first::<UserNoEmailType>(&*conn);
            drop(conn);
            if let Ok(user) = maybe_user {
                let user = User::from(user);
                // Attach the `User` model from the database to the request
                req.mut_extensions().insert(user);
                req.mut_extensions()
                    .insert(AuthenticationSource::SessionCookie);
            }
        } else {
            // Otherwise, look for an `Authorization` header on the request
            // and try to find a user in the database with a matching API token
            let user_auth = if let Some(headers) = req.headers().find("Authorization") {
                ApiToken::find_by_api_token(&conn, headers[0])
                    .and_then(|api_token| {
                        User::find(&conn, api_token.user_id).map(|user| {
                            (
                                AuthenticationSource::ApiToken {
                                    api_token_id: api_token.id,
                                },
                                user,
                            )
                        })
                    })
                    .optional()
                    .map_err(|e| Box::new(e) as Box<dyn Error + Send>)?
            } else {
                None
            };

            drop(conn);

            if let Some((api_token, user)) = user_auth {
                // Attach the `User` model from the database and the API token to the request
                req.mut_extensions().insert(user);
                req.mut_extensions().insert(api_token);
            }
        }

        Ok(())
    }
}

pub trait RequestUser {
    fn user(&self) -> AppResult<&User>;
    fn authentication_source(&self) -> AppResult<AuthenticationSource>;
}

impl<'a> RequestUser for dyn Request + 'a {
    fn user(&self) -> AppResult<&User> {
        self.extensions()
            .find::<User>()
            .chain_error(|| Unauthorized)
    }

    fn authentication_source(&self) -> AppResult<AuthenticationSource> {
        self.extensions()
            .find::<AuthenticationSource>()
            .cloned()
            .chain_error(|| Unauthorized)
    }
}

impl AuthenticationSource {
    pub fn api_token_id(self) -> Option<i32> {
        match self {
            AuthenticationSource::SessionCookie => None,
            AuthenticationSource::ApiToken { api_token_id } => Some(api_token_id),
        }
    }
}
<commit_msg>cargo fmt<commit_after>use super::prelude::*;

use conduit_cookie::RequestSession;
use diesel::prelude::*;

use crate::db::RequestTransaction;
use crate::util::errors::{std_error, AppResult, ChainError, Unauthorized};

use crate::models::user::{UserNoEmailType, ALL_COLUMNS};
use crate::models::ApiToken;
use crate::models::User;
use crate::schema::users;

#[derive(Debug, Clone, Copy)]
pub struct CurrentUser;

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum AuthenticationSource {
    SessionCookie,
    ApiToken { api_token_id: i32 },
}

impl Middleware for CurrentUser {
    fn before(&self, req: &mut dyn Request) -> Result<(), Box<dyn Error + Send>> {
        // Check if the request has a session cookie with a `user_id` property inside
        let id = {
            req.session()
                .get("user_id")
                .and_then(|s| s.parse::<i32>().ok())
        };

        let conn = req.db_conn().map_err(std_error)?;

        if let Some(id) = id {
            // If it did, look for a user in the database with the given `user_id`
            let maybe_user = users::table
                .select(ALL_COLUMNS)
                .find(id)
                .first::<UserNoEmailType>(&*conn);
            drop(conn);
            if let Ok(user) = maybe_user {
                let user = User::from(user);
                // Attach the `User` model from the database to the request
                req.mut_extensions().insert(user);
                req.mut_extensions()
                    .insert(AuthenticationSource::SessionCookie);
            }
        } else {
            // Otherwise, look for an `Authorization` header on the request
            // and try to find a user in the database with a matching API token
            let user_auth = if let Some(headers) = req.headers().find("Authorization") {
                ApiToken::find_by_api_token(&conn, headers[0])
                    .and_then(|api_token| {
                        User::find(&conn, api_token.user_id).map(|user| {
                            (
                                AuthenticationSource::ApiToken {
                                    api_token_id: api_token.id,
                                },
                                user,
                            )
                        })
                    })
                    .optional()
                    .map_err(|e| Box::new(e) as Box<dyn Error + Send>)?
            } else {
                None
            };

            drop(conn);

            if let Some((api_token, user)) = user_auth {
                // Attach the `User` model from the database and the API token to the request
                req.mut_extensions().insert(user);
                req.mut_extensions().insert(api_token);
            }
        }

        Ok(())
    }
}

pub trait RequestUser {
    fn user(&self) -> AppResult<&User>;
    fn authentication_source(&self) -> AppResult<AuthenticationSource>;
}

impl<'a> RequestUser for dyn Request + 'a {
    fn user(&self) -> AppResult<&User> {
        self.extensions()
            .find::<User>()
            .chain_error(|| Unauthorized)
    }

    fn authentication_source(&self) -> AppResult<AuthenticationSource> {
        self.extensions()
            .find::<AuthenticationSource>()
            .cloned()
            .chain_error(|| Unauthorized)
    }
}

impl AuthenticationSource {
    pub fn api_token_id(self) -> Option<i32> {
        match self {
            AuthenticationSource::SessionCookie => None,
            AuthenticationSource::ApiToken { api_token_id } => Some(api_token_id),
        }
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt;

use cranelift::codegen::{
    entity::SecondaryMap,
    ir::{self, entities::AnyEntity, function::DisplayFunctionAnnotations},
    write::{FuncWriter, PlainWriter},
    ValueLabelsRanges,
};

use crate::prelude::*;

/// This module provides the [CommentWriter] which makes it possible
/// to add comments to the written cranelift ir.
///
/// # Example
///
/// ```clif
/// test compile
/// target x86_64
///
/// function u0:0(i64, i64, i64) system_v {
/// ; symbol _ZN119_$LT$example..IsNotEmpty$u20$as$u20$mini_core..FnOnce$LT$$LP$$RF$$u27$a$u20$$RF$$u27$b$u20$$u5b$u16$u5d$$C$$RP$$GT$$GT$9call_once17he85059d5e6a760a0E
/// ; instance Instance { def: Item(DefId(0/0:29 ~ example[8787]::{{impl}}[0]::call_once[0])), substs: [ReErased, ReErased] }
/// ; sig ([IsNotEmpty, (&&[u16],)]; c_variadic: false)->(u8, u8)
///
/// ; ssa {_2: NOT_SSA, _4: NOT_SSA, _0: NOT_SSA, _3: (empty), _1: NOT_SSA}
/// ; msg   loc.idx    param    pass mode            ssa flags  ty
/// ; ret    _0      = v0       ByRef                NOT_SSA    (u8, u8)
/// ; arg    _1      = v1       ByRef                NOT_SSA    IsNotEmpty
/// ; arg    _2.0    = v2       ByVal(types::I64)    NOT_SSA    &&[u16]
///
///     ss0 = explicit_slot 0 ; _1: IsNotEmpty size=0 align=1,8
///     ss1 = explicit_slot 8 ; _2: (&&[u16],) size=8 align=8,8
///     ss2 = explicit_slot 8 ; _4: (&&[u16],) size=8 align=8,8
///     sig0 = (i64, i64, i64) system_v
///     sig1 = (i64, i64, i64) system_v
///     fn0 = colocated u0:6 sig1 ; Instance { def: Item(DefId(0/0:31 ~ example[8787]::{{impl}}[1]::call_mut[0])), substs: [ReErased, ReErased] }
///
/// ebb0(v0: i64, v1: i64, v2: i64):
///     v3 = stack_addr.i64 ss0
///     v4 = stack_addr.i64 ss1
///     store v2, v4
///     v5 = stack_addr.i64 ss2
///     jump ebb1
///
/// ebb1:
///     nop
/// ; _3 = &mut _1
/// ; _4 = _2
///     v6 = load.i64 v4
///     store v6, v5
/// ;
/// ; _0 = const mini_core::FnMut::call_mut(move _3, move _4)
///     v7 = load.i64 v5
///     call fn0(v0, v3, v7)
///     jump ebb2
///
/// ebb2:
///     nop
/// ;
/// ; return
///     return
/// }
/// ```

#[derive(Debug)]
pub struct CommentWriter {
    global_comments: Vec<String>,
    entity_comments: HashMap<AnyEntity, String>,
    inst_comments: HashMap<Inst, String>,
}

impl CommentWriter {
    pub fn new<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Self {
        CommentWriter {
            global_comments: vec![
                format!("symbol {}", tcx.symbol_name(instance).name.as_str()),
                format!("instance {:?}", instance),
                format!(
                    "sig {:?}",
                    tcx.normalize_erasing_late_bound_regions(
                        ParamEnv::reveal_all(),
                        &instance.fn_sig(tcx)
                    )
                ),
                String::new(),
            ],
            entity_comments: HashMap::new(),
            inst_comments: HashMap::new(),
        }
    }
}

impl FuncWriter for &'_ CommentWriter {
    fn write_preamble(
        &mut self,
        w: &mut dyn fmt::Write,
        func: &Function,
        reg_info: Option<&isa::RegInfo>,
    ) -> Result<bool, fmt::Error> {
        for comment in &self.global_comments {
            if !comment.is_empty() {
                writeln!(w, "; {}", comment)?;
            } else {
                writeln!(w, "")?;
            }
        }
        if !self.global_comments.is_empty() {
            writeln!(w, "")?;
        }

        self.super_preamble(w, func, reg_info)
    }

    fn write_entity_definition(
        &mut self,
        w: &mut dyn fmt::Write,
        _func: &Function,
        entity: AnyEntity,
        value: &dyn fmt::Display,
    ) -> fmt::Result {
        write!(w, "    {} = {}", entity, value)?;

        if let Some(comment) = self.entity_comments.get(&entity) {
            writeln!(w, " ; {}", comment.replace('\n', "\n; "))
        } else {
            writeln!(w, "")
        }
    }

    fn write_ebb_header(
        &mut self,
        w: &mut dyn fmt::Write,
        func: &Function,
        isa: Option<&dyn isa::TargetIsa>,
        ebb: Ebb,
        indent: usize,
    ) -> fmt::Result {
        PlainWriter.write_ebb_header(w, func, isa, ebb, indent)
    }

    fn write_instruction(
        &mut self,
        w: &mut dyn fmt::Write,
        func: &Function,
        aliases: &SecondaryMap<Value, Vec<Value>>,
        isa: Option<&dyn isa::TargetIsa>,
        inst: Inst,
        indent: usize,
    ) -> fmt::Result {
        PlainWriter.write_instruction(w, func, aliases, isa, inst, indent)?;
        if let Some(comment) = self.inst_comments.get(&inst) {
            writeln!(w, "; {}", comment.replace('\n', "\n; "))?;
        }
        Ok(())
    }
}

#[cfg(debug_assertions)]
impl<'a, 'tcx, B: Backend + 'static> FunctionCx<'_, 'tcx, B> {
    pub fn add_global_comment<S: Into<String>>(&mut self, comment: S) {
        self.clif_comments.global_comments.push(comment.into());
    }

    pub fn add_entity_comment<'s, S: Into<Cow<'s, str>>, E: Into<AnyEntity>>(
        &mut self,
        entity: E,
        comment: S,
    ) {
        use std::collections::hash_map::Entry;
        match self.clif_comments.entity_comments.entry(entity.into()) {
            Entry::Occupied(mut occ) => {
                occ.get_mut().push('\n');
                occ.get_mut().push_str(comment.into().as_ref());
            }
            Entry::Vacant(vac) => {
                vac.insert(comment.into().into_owned());
            }
        }
    }

    pub fn add_comment<'s, S: Into<Cow<'s, str>>>(&mut self, inst: Inst, comment: S) {
        use std::collections::hash_map::Entry;
        match self.clif_comments.inst_comments.entry(inst) {
            Entry::Occupied(mut occ) => {
                occ.get_mut().push('\n');
                occ.get_mut().push_str(comment.into().as_ref());
            }
            Entry::Vacant(vac) => {
                vac.insert(comment.into().into_owned());
            }
        }
    }
}

pub fn write_clif_file<'tcx>(
    tcx: TyCtxt<'tcx>,
    postfix: &str,
    instance: Instance<'tcx>,
    func: &ir::Function,
    mut clif_comments: &CommentWriter,
    value_ranges: Option<&ValueLabelsRanges>,
) {
    use std::io::Write;

    let symbol_name = tcx.symbol_name(instance).name.as_str();
    let clif_file_name = format!(
        "{}/{}__{}.{}.clif",
        concat!(env!("CARGO_MANIFEST_DIR"), "/target/out/clif"),
        tcx.crate_name(LOCAL_CRATE),
        symbol_name,
        postfix,
    );

    let mut clif = String::new();
    cranelift::codegen::write::decorate_function(
        &mut clif_comments,
        &mut clif,
        &func,
        &DisplayFunctionAnnotations {
            isa: Some(&*crate::build_isa(
                tcx.sess, true, /* PIC doesn't matter here */
            )),
            value_ranges,
        },
    )
    .unwrap();

    match ::std::fs::File::create(clif_file_name) {
        Ok(mut file) => {
            let target_triple = crate::target_triple(tcx.sess);
            writeln!(file, "test compile").unwrap();
            writeln!(file, "set is_pic").unwrap();
            writeln!(file, "target {}", target_triple).unwrap();
            writeln!(file, "").unwrap();
            file.write(clif.as_bytes()).unwrap();
        }
        Err(e) => {
            tcx.sess.warn(&format!("err opening clif file: {:?}", e));
        }
    }
}

impl<'a, 'tcx, B: Backend + 'static> fmt::Debug for FunctionCx<'_, 'tcx, B> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "{:?}", self.instance.substs)?;
        writeln!(f, "{:?}", self.local_map)?;

        let mut clif = String::new();
        ::cranelift::codegen::write::decorate_function(
            &mut &self.clif_comments,
            &mut clif,
            &self.bcx.func,
            // FIXME use DisplayFunctionAnnotations::default() instead
            &DisplayFunctionAnnotations {
                isa: None,
                value_ranges: None,
            },
        )
        .unwrap();
        writeln!(f, "\n{}", clif)
    }
}
<commit_msg>Use DisplayFunctionAnnotations::default()<commit_after>use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt;

use cranelift::codegen::{
    entity::SecondaryMap,
    ir::{self, entities::AnyEntity, function::DisplayFunctionAnnotations},
    write::{FuncWriter, PlainWriter},
    ValueLabelsRanges,
};

use crate::prelude::*;

/// This module provides the [CommentWriter] which makes it possible
/// to add comments to the written cranelift ir.
///
/// # Example
///
/// ```clif
/// test compile
/// target x86_64
///
/// function u0:0(i64, i64, i64) system_v {
/// ; symbol _ZN119_$LT$example..IsNotEmpty$u20$as$u20$mini_core..FnOnce$LT$$LP$$RF$$u27$a$u20$$RF$$u27$b$u20$$u5b$u16$u5d$$C$$RP$$GT$$GT$9call_once17he85059d5e6a760a0E
/// ; instance Instance { def: Item(DefId(0/0:29 ~ example[8787]::{{impl}}[0]::call_once[0])), substs: [ReErased, ReErased] }
/// ; sig ([IsNotEmpty, (&&[u16],)]; c_variadic: false)->(u8, u8)
///
/// ; ssa {_2: NOT_SSA, _4: NOT_SSA, _0: NOT_SSA, _3: (empty), _1: NOT_SSA}
/// ; msg   loc.idx    param    pass mode            ssa flags  ty
/// ; ret    _0      = v0       ByRef                NOT_SSA    (u8, u8)
/// ; arg    _1      = v1       ByRef                NOT_SSA    IsNotEmpty
/// ; arg    _2.0    = v2       ByVal(types::I64)    NOT_SSA    &&[u16]
///
///     ss0 = explicit_slot 0 ; _1: IsNotEmpty size=0 align=1,8
///     ss1 = explicit_slot 8 ; _2: (&&[u16],) size=8 align=8,8
///     ss2 = explicit_slot 8 ; _4: (&&[u16],) size=8 align=8,8
///     sig0 = (i64, i64, i64) system_v
///     sig1 = (i64, i64, i64) system_v
///     fn0 = colocated u0:6 sig1 ; Instance { def: Item(DefId(0/0:31 ~ example[8787]::{{impl}}[1]::call_mut[0])), substs: [ReErased, ReErased] }
///
/// ebb0(v0: i64, v1: i64, v2: i64):
///     v3 = stack_addr.i64 ss0
///     v4 = stack_addr.i64 ss1
///     store v2, v4
///     v5 = stack_addr.i64 ss2
///     jump ebb1
///
/// ebb1:
///     nop
/// ; _3 = &mut _1
/// ; _4 = _2
///     v6 = load.i64 v4
///     store v6, v5
/// ;
/// ; _0 = const mini_core::FnMut::call_mut(move _3, move _4)
///     v7 = load.i64 v5
///     call fn0(v0, v3, v7)
///     jump ebb2
///
/// ebb2:
///     nop
/// ;
/// ; return
///     return
/// }
/// ```

#[derive(Debug)]
pub struct CommentWriter {
    global_comments: Vec<String>,
    entity_comments: HashMap<AnyEntity, String>,
    inst_comments: HashMap<Inst, String>,
}

impl CommentWriter {
    pub fn new<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Self {
        CommentWriter {
            global_comments: vec![
                format!("symbol {}", tcx.symbol_name(instance).name.as_str()),
                format!("instance {:?}", instance),
                format!(
                    "sig {:?}",
                    tcx.normalize_erasing_late_bound_regions(
                        ParamEnv::reveal_all(),
                        &instance.fn_sig(tcx)
                    )
                ),
                String::new(),
            ],
            entity_comments: HashMap::new(),
            inst_comments: HashMap::new(),
        }
    }
}

impl FuncWriter for &'_ CommentWriter {
    fn write_preamble(
        &mut self,
        w: &mut dyn fmt::Write,
        func: &Function,
        reg_info: Option<&isa::RegInfo>,
    ) -> Result<bool, fmt::Error> {
        for comment in &self.global_comments {
            if !comment.is_empty() {
                writeln!(w, "; {}", comment)?;
            } else {
                writeln!(w, "")?;
            }
        }
        if !self.global_comments.is_empty() {
            writeln!(w, "")?;
        }

        self.super_preamble(w, func, reg_info)
    }

    fn write_entity_definition(
        &mut self,
        w: &mut dyn fmt::Write,
        _func: &Function,
        entity: AnyEntity,
        value: &dyn fmt::Display,
    ) -> fmt::Result {
        write!(w, "    {} = {}", entity, value)?;

        if let Some(comment) = self.entity_comments.get(&entity) {
            writeln!(w, " ; {}", comment.replace('\n', "\n; "))
        } else {
            writeln!(w, "")
        }
    }

    fn write_ebb_header(
        &mut self,
        w: &mut dyn fmt::Write,
        func: &Function,
        isa: Option<&dyn isa::TargetIsa>,
        ebb: Ebb,
        indent: usize,
    ) -> fmt::Result {
        PlainWriter.write_ebb_header(w, func, isa, ebb, indent)
    }

    fn write_instruction(
        &mut self,
        w: &mut dyn fmt::Write,
        func: &Function,
        aliases: &SecondaryMap<Value, Vec<Value>>,
        isa: Option<&dyn isa::TargetIsa>,
        inst: Inst,
        indent: usize,
    ) -> fmt::Result {
        PlainWriter.write_instruction(w, func, aliases, isa, inst, indent)?;
        if let Some(comment) = self.inst_comments.get(&inst) {
            writeln!(w, "; {}", comment.replace('\n', "\n; "))?;
        }
        Ok(())
    }
}

#[cfg(debug_assertions)]
impl<'a, 'tcx, B: Backend + 'static> FunctionCx<'_, 'tcx, B> {
    pub fn add_global_comment<S: Into<String>>(&mut self, comment: S) {
        self.clif_comments.global_comments.push(comment.into());
    }

    pub fn add_entity_comment<'s, S: Into<Cow<'s, str>>, E: Into<AnyEntity>>(
        &mut self,
        entity: E,
        comment: S,
    ) {
        use std::collections::hash_map::Entry;
        match self.clif_comments.entity_comments.entry(entity.into()) {
            Entry::Occupied(mut occ) => {
                occ.get_mut().push('\n');
                occ.get_mut().push_str(comment.into().as_ref());
            }
            Entry::Vacant(vac) => {
                vac.insert(comment.into().into_owned());
            }
        }
    }

    pub fn add_comment<'s, S: Into<Cow<'s, str>>>(&mut self, inst: Inst, comment: S) {
        use std::collections::hash_map::Entry;
        match self.clif_comments.inst_comments.entry(inst) {
            Entry::Occupied(mut occ) => {
                occ.get_mut().push('\n');
                occ.get_mut().push_str(comment.into().as_ref());
            }
            Entry::Vacant(vac) => {
                vac.insert(comment.into().into_owned());
            }
        }
    }
}

pub fn write_clif_file<'tcx>(
    tcx: TyCtxt<'tcx>,
    postfix: &str,
    instance: Instance<'tcx>,
    func: &ir::Function,
    mut clif_comments: &CommentWriter,
    value_ranges: Option<&ValueLabelsRanges>,
) {
    use std::io::Write;

    let symbol_name = tcx.symbol_name(instance).name.as_str();
    let clif_file_name = format!(
        "{}/{}__{}.{}.clif",
        concat!(env!("CARGO_MANIFEST_DIR"), "/target/out/clif"),
        tcx.crate_name(LOCAL_CRATE),
        symbol_name,
        postfix,
    );

    let mut clif = String::new();
    cranelift::codegen::write::decorate_function(
        &mut clif_comments,
        &mut clif,
        &func,
        &DisplayFunctionAnnotations {
            isa: Some(&*crate::build_isa(
                tcx.sess, true, /* PIC doesn't matter here */
            )),
            value_ranges,
        },
    )
    .unwrap();

    match ::std::fs::File::create(clif_file_name) {
        Ok(mut file) => {
            let target_triple = crate::target_triple(tcx.sess);
            writeln!(file, "test compile").unwrap();
            writeln!(file, "set is_pic").unwrap();
            writeln!(file, "target {}", target_triple).unwrap();
            writeln!(file, "").unwrap();
            file.write(clif.as_bytes()).unwrap();
        }
        Err(e) => {
            tcx.sess.warn(&format!("err opening clif file: {:?}", e));
        }
    }
}

impl<'a, 'tcx, B: Backend + 'static> fmt::Debug for FunctionCx<'_, 'tcx, B> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "{:?}", self.instance.substs)?;
        writeln!(f, "{:?}", self.local_map)?;

        let mut clif = String::new();
        ::cranelift::codegen::write::decorate_function(
            &mut &self.clif_comments,
            &mut clif,
            &self.bcx.func,
            &DisplayFunctionAnnotations::default(),
        )
        .unwrap();
        writeln!(f, "\n{}", clif)
    }
}
<|endoftext|>


<commit_before>use super::RawOutcome;
use crate::core::{FilterList, LowercaseText};
use crate::utils::is_stop_char;
use crate::{Lang, Script};

use once_cell::sync::Lazy;
use std::collections::HashMap;

const AFR: &str = "abcdefghijklmnopqrstuvwxyzáèéêëíîïóôúû";
const AKA: &str = "abdefghiklmnoprstuwyɔɛ";
const AZE: &str = "abcdefghijklmnopqrstuvxyzçöüğışə̇";
const CAT: &str = "abcdefghijklmnopqrstuvwxyz·àçèéíïòóúü";
const CES: &str = "abcdefghijklmnopqrstuvwxyzáéóúýčďěňřšťůž";
const DAN: &str = "abcdefghijklmnopqrstuvwxyzåæø";
const DEU: &str = "abcdefghijklmnopqrstuvwxyzßäöü";
const ENG: &str = "abcdefghijklmnopqrstuvwxyz";
const EPO: &str = "abcdefghijklmnoprstuvzĉĝĥĵŝŭ";
const EST: &str = "abcdefghijklmnopqrstuvwxyzäõöü";
const FIN: &str = "abcdefghijklmnopqrstuvwxyzäöšž";
const FRA: &str = "abcdefghijklmnopqrstuvwxyzàâçèéêëîïôùûüÿœ";
const HRV: &str = "abcdefghijklmnopqrstuvwxyzćčđšž";
const HUN: &str = "abcdefghijklmnopqrstuvwxyzáéíóöúüőű";
const IND: &str = "abcdefghijklmnopqrstuvwxyz";
const ITA: &str = "abcdefghijklmnopqrstuvwxyzàèéìòù";
const JAV: &str = "abcdefghijklmnopqrstuvwxyzèé";
const LAT: &str = "abcdefghijklmnopqrstuvwxyz";
const LAV: &str = "abcdefghijklmnopqrstuvwxyzāčēģīķļņōŗšūž";
const LIT: &str = "abcdefghijklmnopqrstuvwxyząčėęįšūųž";
const NLD: &str = "abcdefghijklmnopqrstuvwxyzàèéëïĳ";
const NOB: &str = "abcdefghijklmnopqrstuvwxyzåæø";
const POL: &str = "abcdefghijklmnopqrstuvwxyzóąćęłńśźż";
const POR: &str = "abcdefghijklmnopqrstuvwxyzàáâãçéêíóôõú";
const RON: &str = "abcdefghijklmnopqrstuvwxyzâîăşţ";
const SLK: &str = "abcdefghijklmnopqrstuvwxyzáäéíóôúýčďĺľňŕšťž";
const SLV: &str = "abcdefghijklmnopqrstuvwxyzčšž";
const SNA: &str = "abcdefghijklmnopqrstuvwxyz";
const SPA: &str = "abcdefghijklmnopqrstuvwxyz¡¿áéíñóúü";
const SWE: &str = "abcdefghijklmnopqrstuvwxyzäåö";
const TGL: &str = "abcdefghijklmnopqrstuvwxyzáéíñóú";
const TUK: &str = "abdefghijklmnoprstuwyzäçöüýňşž";
const TUR: &str = "abcdefghijklmnopqrstuvwxyzçöüğış̇";
const UZB: &str = "abcdefghijklmnopqrstuvxyzʻ";
const VIE: &str =
    "abcdefghijklmnopqrstuvwxyzàáâãèéêìíòóôõùúýăđĩũơưạảấầẩẫậắằẳẵặẹẻẽếềểễệỉịọỏốồổỗộớờởỡợụủứừửữựỳỵỷỹ";
const ZUL: &str = "abcdefghijklmnopqrstuvwxyz";

const LATIN_ALPHABETS: &[(Lang, &str)] = &[
    (Lang::Afr, AFR),
    (Lang::Aka, AKA),
    (Lang::Aze, AZE),
    (Lang::Cat, CAT),
    (Lang::Ces, CES),
    (Lang::Dan, DAN),
    (Lang::Deu, DEU),
    (Lang::Eng, ENG),
    (Lang::Epo, EPO),
    (Lang::Est, EST),
    (Lang::Fin, FIN),
    (Lang::Fra, FRA),
    (Lang::Hrv, HRV),
    (Lang::Hun, HUN),
    (Lang::Ind, IND),
    (Lang::Ita, ITA),
    (Lang::Jav, JAV),
    (Lang::Lat, LAT),
    (Lang::Lav, LAV),
    (Lang::Lit, LIT),
    (Lang::Nld, NLD),
    (Lang::Nob, NOB),
    (Lang::Pol, POL),
    (Lang::Por, POR),
    (Lang::Ron, RON),
    (Lang::Slk, SLK),
    (Lang::Slv, SLV),
    (Lang::Sna, SNA),
    (Lang::Spa, SPA),
    (Lang::Swe, SWE),
    (Lang::Tgl, TGL),
    (Lang::Tuk, TUK),
    (Lang::Tur, TUR),
    (Lang::Uzb, UZB),
    (Lang::Vie, VIE),
    (Lang::Zul, ZUL),
];

/// Inverted map binding a character to a set of languages.
pub static ALPHABET_LANG_MAP: Lazy<Vec<(char, Vec<Lang>)>> = Lazy::new(|| {
    let mut map = HashMap::new();

    for (lang, alphabet) in LATIN_ALPHABETS {
        for c in alphabet.chars() {
            let entry = map.entry(c).or_insert_with(Vec::new);
            entry.push(*lang);
        }
    }

    let mut char_lang = Vec::new();

    for (ch, languages) in map {
        char_lang.push((ch, languages));
    }

    char_lang
});

pub fn alphabet_calculate_scores(text: &LowercaseText, filter_list: &FilterList) -> RawOutcome {
    let char_lang = &*ALPHABET_LANG_MAP;
    let max_raw_score = text.chars().filter(|&ch| !is_stop_char(ch)).count();

    let mut raw_scores: Vec<(Lang, usize)> = Script::Latin
        .langs()
        .iter()
        .filter(|&&l| filter_list.is_allowed(l))
        .map(|&l| (l, 0))
        .collect();

    let mut common_score = 0;
    for ch in text.chars() {
        if is_stop_char(ch) {
            continue;
        }

        if let Some((_, languages)) = char_lang.iter().find(|(c, _)| *c == ch) {
            // if this char is common to all Languages, add 2 to a common counter
            // instead of iterating over all Languages counters.
            if languages.len() == LATIN_ALPHABETS.len() {
                common_score += 2;
            } else {
                for lang in languages {
                    if let Some((_, score)) = raw_scores.iter_mut().find(|(l, _)| l == lang) {
                        *score += 2;
                    }
                }
            }
        }
    }

    raw_scores.sort_unstable_by(|a, b| b.1.cmp(&a.1));

    let mut normalized_scores = vec![];

    for &(lang, raw_score) in &raw_scores {
        let normalized_score =
            (raw_score + common_score).saturating_sub(max_raw_score) as f64 / max_raw_score as f64;
        normalized_scores.push((lang, normalized_score));
    }

    RawOutcome {
        count: max_raw_score,
        raw_scores,
        scores: normalized_scores,
    }
}
<commit_msg>Use binary search instead of iter find<commit_after>use super::RawOutcome;
use crate::core::{FilterList, LowercaseText};
use crate::utils::is_stop_char;
use crate::{Lang, Script};

use once_cell::sync::Lazy;
use std::collections::HashMap;

const AFR: &str = "abcdefghijklmnopqrstuvwxyzáèéêëíîïóôúû";
const AKA: &str = "abdefghiklmnoprstuwyɔɛ";
const AZE: &str = "abcdefghijklmnopqrstuvxyzçöüğışə̇";
const CAT: &str = "abcdefghijklmnopqrstuvwxyz·àçèéíïòóúü";
const CES: &str = "abcdefghijklmnopqrstuvwxyzáéóúýčďěňřšťůž";
const DAN: &str = "abcdefghijklmnopqrstuvwxyzåæø";
const DEU: &str = "abcdefghijklmnopqrstuvwxyzßäöü";
const ENG: &str = "abcdefghijklmnopqrstuvwxyz";
const EPO: &str = "abcdefghijklmnoprstuvzĉĝĥĵŝŭ";
const EST: &str = "abcdefghijklmnopqrstuvwxyzäõöü";
const FIN: &str = "abcdefghijklmnopqrstuvwxyzäöšž";
const FRA: &str = "abcdefghijklmnopqrstuvwxyzàâçèéêëîïôùûüÿœ";
const HRV: &str = "abcdefghijklmnopqrstuvwxyzćčđšž";
const HUN: &str = "abcdefghijklmnopqrstuvwxyzáéíóöúüőű";
const IND: &str = "abcdefghijklmnopqrstuvwxyz";
const ITA: &str = "abcdefghijklmnopqrstuvwxyzàèéìòù";
const JAV: &str = "abcdefghijklmnopqrstuvwxyzèé";
const LAT: &str = "abcdefghijklmnopqrstuvwxyz";
const LAV: &str = "abcdefghijklmnopqrstuvwxyzāčēģīķļņōŗšūž";
const LIT: &str = "abcdefghijklmnopqrstuvwxyząčėęįšūųž";
const NLD: &str = "abcdefghijklmnopqrstuvwxyzàèéëïĳ";
const NOB: &str = "abcdefghijklmnopqrstuvwxyzåæø";
const POL: &str = "abcdefghijklmnopqrstuvwxyzóąćęłńśźż";
const POR: &str = "abcdefghijklmnopqrstuvwxyzàáâãçéêíóôõú";
const RON: &str = "abcdefghijklmnopqrstuvwxyzâîăşţ";
const SLK: &str = "abcdefghijklmnopqrstuvwxyzáäéíóôúýčďĺľňŕšťž";
const SLV: &str = "abcdefghijklmnopqrstuvwxyzčšž";
const SNA: &str = "abcdefghijklmnopqrstuvwxyz";
const SPA: &str = "abcdefghijklmnopqrstuvwxyz¡¿áéíñóúü";
const SWE: &str = "abcdefghijklmnopqrstuvwxyzäåö";
const TGL: &str = "abcdefghijklmnopqrstuvwxyzáéíñóú";
const TUK: &str = "abdefghijklmnoprstuwyzäçöüýňşž";
const TUR: &str = "abcdefghijklmnopqrstuvwxyzçöüğış̇";
const UZB: &str = "abcdefghijklmnopqrstuvxyzʻ";
const VIE: &str =
    "abcdefghijklmnopqrstuvwxyzàáâãèéêìíòóôõùúýăđĩũơưạảấầẩẫậắằẳẵặẹẻẽếềểễệỉịọỏốồổỗộớờởỡợụủứừửữựỳỵỷỹ";
const ZUL: &str = "abcdefghijklmnopqrstuvwxyz";

const LATIN_ALPHABETS: &[(Lang, &str)] = &[
    (Lang::Afr, AFR),
    (Lang::Aka, AKA),
    (Lang::Aze, AZE),
    (Lang::Cat, CAT),
    (Lang::Ces, CES),
    (Lang::Dan, DAN),
    (Lang::Deu, DEU),
    (Lang::Eng, ENG),
    (Lang::Epo, EPO),
    (Lang::Est, EST),
    (Lang::Fin, FIN),
    (Lang::Fra, FRA),
    (Lang::Hrv, HRV),
    (Lang::Hun, HUN),
    (Lang::Ind, IND),
    (Lang::Ita, ITA),
    (Lang::Jav, JAV),
    (Lang::Lat, LAT),
    (Lang::Lav, LAV),
    (Lang::Lit, LIT),
    (Lang::Nld, NLD),
    (Lang::Nob, NOB),
    (Lang::Pol, POL),
    (Lang::Por, POR),
    (Lang::Ron, RON),
    (Lang::Slk, SLK),
    (Lang::Slv, SLV),
    (Lang::Sna, SNA),
    (Lang::Spa, SPA),
    (Lang::Swe, SWE),
    (Lang::Tgl, TGL),
    (Lang::Tuk, TUK),
    (Lang::Tur, TUR),
    (Lang::Uzb, UZB),
    (Lang::Vie, VIE),
    (Lang::Zul, ZUL),
];

/// Inverted map binding a character to a set of languages.
pub static ALPHABET_LANG_MAP: Lazy<Vec<(char, Vec<Lang>)>> = Lazy::new(|| {
    let mut map = HashMap::new();

    for (lang, alphabet) in LATIN_ALPHABETS {
        for c in alphabet.chars() {
            let entry = map.entry(c).or_insert_with(Vec::new);
            entry.push(*lang);
        }
    }

    let mut char_lang = Vec::new();

    for (ch, languages) in map {
        char_lang.push((ch, languages));
    }

    char_lang.sort_unstable_by_key(|(c, _)| *c);
    char_lang
});

pub fn alphabet_calculate_scores(text: &LowercaseText, filter_list: &FilterList) -> RawOutcome {
    let char_lang = &*ALPHABET_LANG_MAP;
    let max_raw_score = text.chars().filter(|&ch| !is_stop_char(ch)).count();

    let mut raw_scores: Vec<(Lang, usize)> = Script::Latin
        .langs()
        .iter()
        .filter(|&&l| filter_list.is_allowed(l))
        .map(|&l| (l, 0))
        .collect();

    let mut common_score = 0;
    for ch in text.chars() {
        if is_stop_char(ch) {
            continue;
        }

        if let Ok(position) = char_lang.binary_search_by_key(&ch, |(c, _)| *c) {
            let (_, languages) = &char_lang[position];
            // if this char is common to all Languages, add 2 to a common counter
            // instead of iterating over all Languages counters.
            if languages.len() == LATIN_ALPHABETS.len() {
                common_score += 2;
            } else {
                for lang in languages {
                    if let Some((_, score)) = raw_scores.iter_mut().find(|(l, _)| l == lang) {
                        *score += 2;
                    }
                }
            }
        }
    }

    raw_scores.sort_unstable_by(|a, b| b.1.cmp(&a.1));

    let mut normalized_scores = vec![];

    for &(lang, raw_score) in &raw_scores {
        let normalized_score =
            (raw_score + common_score).saturating_sub(max_raw_score) as f64 / max_raw_score as f64;
        normalized_scores.push((lang, normalized_score));
    }

    RawOutcome {
        count: max_raw_score,
        raw_scores,
        scores: normalized_scores,
    }
}
<|endoftext|>


<commit_before>pub mod tag_command;
pub mod autotag;
pub mod update_tags;

pub use self::tag_command::TagCommand;
pub use self::autotag::AutoTagCommand;
pub use self::update_tags::UpdateTagsCommand;

use LalafellBot;
use database::AutotagUser;

use error::*;

use discord;
use discord::model::{UserId, LiveServer, Role, RoleId};

use serde_json;

use std::sync::{Arc, Mutex};
use std::collections::HashSet;

lazy_static! {
  // Limbo roles are roles that may or may not be added to the Discord bot state.
  static ref LIMBO_ROLES: Mutex<Vec<Role>> = Mutex::default();
}

pub struct Tagger;

impl Tagger {
  pub fn search_tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, server: &str, character_name: &str, ignore_verified: bool) -> Result<Option<String>> {
    let params = &[
      ("one", "characters"),
      ("strict", "on"),
      ("server|et", server)
    ];

    let res = bot.xivdb.search(character_name, params).chain_err(|| "could not query XIVDB")?;

    let search_chars = res.characters.unwrap().results;
    if search_chars.is_empty() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    let char_id = match search_chars[0]["id"].as_u64() {
      Some(u) => u,
      None => return Err("character ID was not a u64".into())
    };

    let name = match search_chars[0]["name"].as_str() {
      Some(s) => s,
      None => return Err("character name was not a string".into())
    };

    if name.to_lowercase() != character_name.to_lowercase() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    Tagger::tag(bot, who, on, char_id, ignore_verified)
  }

  fn find_or_create_role(bot: &LalafellBot, server: &LiveServer, name: &str, add_roles: &mut Vec<Role>, created_roles: &mut Vec<Role>) -> Result<()> {
    let lower_name = name.to_lowercase();
    match server.roles.iter().find(|x| x.name.to_lowercase() == lower_name) {
      Some(r) => add_roles.push(r.clone()),
      None => {
        let role = bot.discord.create_role(server.id, Some(&lower_name), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    Ok(())
  }

  pub fn tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, char_id: u64, ignore_verified: bool) -> Result<Option<String>> {
    let is_verified = match bot.database.read().unwrap().autotags.users.iter().find(|u| u.user_id == who.0 && u.server_id == on.id.0) {
      Some(u) => {
        if u.verification.verified && !ignore_verified && char_id != u.character_id {
          return Ok(Some(format!("{} is verified as {} on {}, so they cannot switch to another account.", who.mention(), u.character, u.server)));
        }
        u.verification.verified
      },
      None => false
    };

    let member = match bot.discord.get_member(on.id, who) {
      Ok(m) => m,
      Err(discord::Error::Status(_, Some(discord_error))) => {
        let error: DiscordNotFoundError = serde_json::from_value(discord_error)
          .chain_err(|| "could not get member for tagging and could not parse error")?;
        if error.code == 10013 {
          let mut database = bot.database.write().unwrap();
          match database.autotags.users.iter().position(|u| u.user_id == who.0) {
            Some(id) => {
              database.autotags.users.remove(id);
              return Err(format!("could not find user {} on server {}: removing from database", who.0, on.id.0).into());
            },
            _ => return Err("could not find user {} on server {}, but was not in database".into())
          }
        } else {
          let message = error.message.map(|x| format!(" ({})", x)).unwrap_or_default();
          return Err(format!("could not get member for tagging with unknown error code: {}{}", error.code, message).into());
        }
      },
      Err(e) => return Err(e).chain_err(|| "could not get member for tagging")
    };

    let character = bot.xivdb.character(char_id).chain_err(|| "could not look up character")?;

    bot.database.write().unwrap().autotags.update_or_add(AutotagUser::new(
      who.0,
      on.id.0,
      character.lodestone_id,
      &character.name,
      &character.server
    ));

    // Get a copy of the roles on the server.
    let mut roles = on.roles.clone();
    // Check for existing limbo roles.
    {
      let limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for role in &roles {
        // If the server has updated to contain the limbo role, remove it.
        if let Some(i) = limbo.iter().position(|x| x.id == role.id) {
          limbo.remove(i);
        }
      }
    }
    // Get a copy of the limbo roles.
    let limbo = LIMBO_ROLES.lock().unwrap().clone();
    // Extend the server roles with the limbo roles.
    roles.extend(limbo);

    // Find or create the necessary roles
    let mut created_roles = Vec::new();
    let mut add_roles = Vec::new();
    Tagger::find_or_create_role(&bot, on, &character.data.race, &mut add_roles, &mut created_roles)?;
    Tagger::find_or_create_role(&bot, on, &character.data.gender, &mut add_roles, &mut created_roles)?;
    Tagger::find_or_create_role(&bot, on, &character.server, &mut add_roles, &mut created_roles)?;

    if is_verified {
      Tagger::find_or_create_role(&bot, on, "verified", &mut add_roles, &mut created_roles)?;
    }

    // If we created any roles, the server may or may not update with them fast enough, so store a copy in the limbo
    // roles.
    {
      let mut limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for created in &created_roles {
        limbo.push(created.clone());
      }
    }

    debug!("Created the following roles:\n:{:#?}", created_roles);

    debug!("Existing roles:\n{:#?}", add_roles);

    // Extend the roles to add with the roles we created.
    add_roles.extend(created_roles);
    // Get all the roles that are part of groups
    let all_group_roles: Vec<String> = bot.config.roles.groups.iter().flat_map(|x| x).map(|x| x.to_lowercase()).collect();
    // Filter all roles on the server to only the roles the member has
    let keep: Vec<&Role> = roles.iter().filter(|x| member.roles.contains(&x.id)).collect();
    // Filter all the roles the member has, keeping the ones not in a group. These roles will not be touched when
    // updating the tag.
    let keep: Vec<&Role> = keep.into_iter().filter(|x| !all_group_roles.contains(&x.name.to_lowercase())).collect();
    debug!("Roles to keep:\n{:#?}", keep);
    // Combine the two sets of roles and map them to IDs
    let mut role_set: Vec<RoleId> = add_roles.iter().map(|r| r.id).chain(keep.into_iter().map(|r| r.id)).collect();
    // Sort the IDs so we can dedup them
    role_set.sort();
    // Remove the duplicate roles, if any
    role_set.dedup();

    debug!("Final role set:\n{:#?}", role_set);

    // Only update the roles if they are different
    let different = {
      let member_roles: HashSet<u64> = member.roles.iter().map(|x| x.0).collect();
      let actual_role_set: HashSet<u64> = role_set.iter().map(|x| x.0).collect();
      member_roles != actual_role_set
    };
    if different {
      bot.discord.edit_member_roles(on.id, who, &role_set).chain_err(|| "could not add roles")?;
    }

    // cannot edit nickname of those with a higher role
    let nick = match member.nick {
      Some(n) => n,
      None => Default::default()
    };
    if nick != character.name {
      bot.discord.edit_member(on.id, who, |e| e.nickname(&character.name)).ok();
    }
    Ok(None)
  }
}

#[derive(Debug, Deserialize)]
struct DiscordNotFoundError {
  code: u64,
  message: Option<String>
}
<commit_msg>fix: check more error codes<commit_after>pub mod tag_command;
pub mod autotag;
pub mod update_tags;

pub use self::tag_command::TagCommand;
pub use self::autotag::AutoTagCommand;
pub use self::update_tags::UpdateTagsCommand;

use LalafellBot;
use database::AutotagUser;

use error::*;

use discord;
use discord::model::{UserId, LiveServer, Role, RoleId};

use serde_json;

use std::sync::{Arc, Mutex};
use std::collections::HashSet;

lazy_static! {
  // Limbo roles are roles that may or may not be added to the Discord bot state.
  static ref LIMBO_ROLES: Mutex<Vec<Role>> = Mutex::default();
}

pub struct Tagger;

impl Tagger {
  pub fn search_tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, server: &str, character_name: &str, ignore_verified: bool) -> Result<Option<String>> {
    let params = &[
      ("one", "characters"),
      ("strict", "on"),
      ("server|et", server)
    ];

    let res = bot.xivdb.search(character_name, params).chain_err(|| "could not query XIVDB")?;

    let search_chars = res.characters.unwrap().results;
    if search_chars.is_empty() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    let char_id = match search_chars[0]["id"].as_u64() {
      Some(u) => u,
      None => return Err("character ID was not a u64".into())
    };

    let name = match search_chars[0]["name"].as_str() {
      Some(s) => s,
      None => return Err("character name was not a string".into())
    };

    if name.to_lowercase() != character_name.to_lowercase() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    Tagger::tag(bot, who, on, char_id, ignore_verified)
  }

  fn find_or_create_role(bot: &LalafellBot, server: &LiveServer, name: &str, add_roles: &mut Vec<Role>, created_roles: &mut Vec<Role>) -> Result<()> {
    let lower_name = name.to_lowercase();
    match server.roles.iter().find(|x| x.name.to_lowercase() == lower_name) {
      Some(r) => add_roles.push(r.clone()),
      None => {
        let role = bot.discord.create_role(server.id, Some(&lower_name), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    Ok(())
  }

  pub fn tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, char_id: u64, ignore_verified: bool) -> Result<Option<String>> {
    let is_verified = match bot.database.read().unwrap().autotags.users.iter().find(|u| u.user_id == who.0 && u.server_id == on.id.0) {
      Some(u) => {
        if u.verification.verified && !ignore_verified && char_id != u.character_id {
          return Ok(Some(format!("{} is verified as {} on {}, so they cannot switch to another account.", who.mention(), u.character, u.server)));
        }
        u.verification.verified
      },
      None => false
    };

    // FIXME: This is a goddamn disaster
    let member = match bot.discord.get_member(on.id, who) {
      Ok(m) => m,
      Err(discord::Error::Status(_, Some(discord_error))) => {
        let error: DiscordNotFoundError = serde_json::from_value(discord_error)
          .chain_err(|| "could not get member for tagging and could not parse error")?;
          let mut database = bot.database.write().unwrap();
          if error.code == 10013 || error.code == 10007 {
            match database.autotags.users.iter().position(|u| u.user_id == who.0) {
              Some(id) => {
                database.autotags.users.remove(id);
                let message = match error.code {
                  // Unknown user
                  10013 => format!("could not find user {}: removing from database", who.0),
                  // User not a member on this server
                  10007 => format!("user {} is not on server {}: removing from database", who.0, on.id.0),
                  _ => unreachable!()
                };
                return Err(message.into());
              },
              _ => return Err("could not find user {} on server {}, but was not in database".into())
            }
          } else {
            let message = error.message.map(|x| format!(" ({})", x)).unwrap_or_default();
            return Err(format!("could not get member for tagging with unknown error code: {}{}", error.code, message).into());
          }
      },
      Err(e) => return Err(e).chain_err(|| "could not get member for tagging")
    };

    let character = bot.xivdb.character(char_id).chain_err(|| "could not look up character")?;

    bot.database.write().unwrap().autotags.update_or_add(AutotagUser::new(
      who.0,
      on.id.0,
      character.lodestone_id,
      &character.name,
      &character.server
    ));

    // Get a copy of the roles on the server.
    let mut roles = on.roles.clone();
    // Check for existing limbo roles.
    {
      let limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for role in &roles {
        // If the server has updated to contain the limbo role, remove it.
        if let Some(i) = limbo.iter().position(|x| x.id == role.id) {
          limbo.remove(i);
        }
      }
    }
    // Get a copy of the limbo roles.
    let limbo = LIMBO_ROLES.lock().unwrap().clone();
    // Extend the server roles with the limbo roles.
    roles.extend(limbo);

    // Find or create the necessary roles
    let mut created_roles = Vec::new();
    let mut add_roles = Vec::new();
    Tagger::find_or_create_role(&bot, on, &character.data.race, &mut add_roles, &mut created_roles)?;
    Tagger::find_or_create_role(&bot, on, &character.data.gender, &mut add_roles, &mut created_roles)?;
    Tagger::find_or_create_role(&bot, on, &character.server, &mut add_roles, &mut created_roles)?;

    if is_verified {
      Tagger::find_or_create_role(&bot, on, "verified", &mut add_roles, &mut created_roles)?;
    }

    // If we created any roles, the server may or may not update with them fast enough, so store a copy in the limbo
    // roles.
    {
      let mut limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for created in &created_roles {
        limbo.push(created.clone());
      }
    }

    debug!("Created the following roles:\n:{:#?}", created_roles);

    debug!("Existing roles:\n{:#?}", add_roles);

    // Extend the roles to add with the roles we created.
    add_roles.extend(created_roles);
    // Get all the roles that are part of groups
    let all_group_roles: Vec<String> = bot.config.roles.groups.iter().flat_map(|x| x).map(|x| x.to_lowercase()).collect();
    // Filter all roles on the server to only the roles the member has
    let keep: Vec<&Role> = roles.iter().filter(|x| member.roles.contains(&x.id)).collect();
    // Filter all the roles the member has, keeping the ones not in a group. These roles will not be touched when
    // updating the tag.
    let keep: Vec<&Role> = keep.into_iter().filter(|x| !all_group_roles.contains(&x.name.to_lowercase())).collect();
    debug!("Roles to keep:\n{:#?}", keep);
    // Combine the two sets of roles and map them to IDs
    let mut role_set: Vec<RoleId> = add_roles.iter().map(|r| r.id).chain(keep.into_iter().map(|r| r.id)).collect();
    // Sort the IDs so we can dedup them
    role_set.sort();
    // Remove the duplicate roles, if any
    role_set.dedup();

    debug!("Final role set:\n{:#?}", role_set);

    // Only update the roles if they are different
    let different = {
      let member_roles: HashSet<u64> = member.roles.iter().map(|x| x.0).collect();
      let actual_role_set: HashSet<u64> = role_set.iter().map(|x| x.0).collect();
      member_roles != actual_role_set
    };
    if different {
      bot.discord.edit_member_roles(on.id, who, &role_set).chain_err(|| "could not add roles")?;
    }

    // cannot edit nickname of those with a higher role
    let nick = match member.nick {
      Some(n) => n,
      None => Default::default()
    };
    if nick != character.name {
      bot.discord.edit_member(on.id, who, |e| e.nickname(&character.name)).ok();
    }
    Ok(None)
  }
}

#[derive(Debug, Deserialize)]
struct DiscordNotFoundError {
  code: u64,
  message: Option<String>
}
<|endoftext|>


<commit_before>
use std::ffi;
use std::ffi::AsOsStr;
use std::convert::AsRef;
use std::path::Path;
use std::slice;

pub mod record;

use htslib;

pub struct HeaderView {
    inner: *mut htslib::vcf::bcf_hdr_t,
}

impl HeaderView {
    fn new(inner: *mut htslib::vcf::bcf_hdr_t) -> Self {
        HeaderView { inner: inner }
    }

    #[inline]
    fn inner(&self) -> htslib::vcf::bcf_hdr_t {
        unsafe { (*self.inner) }
    }

    pub fn sample_count(&self) -> u32 {
        self.inner().n[htslib::vcf::BCF_DT_SAMPLE as usize] as u32
    }

    pub fn samples(&self) -> Vec<&[u8]> {
        let names = unsafe { slice::from_raw_parts(self.inner().samples, self.sample_count() as usize) };
        names.iter().map(|name| unsafe { ffi::CStr::from_ptr(*name).to_bytes() }).collect()
    }
}

pub struct Reader {
    inner: *mut htslib::vcf::htsFile,
    header: HeaderView,
}

impl Reader {
   pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let htsfile = bcf_open(path, b"r");
        let header = unsafe { htslib::vcf::bcf_hdr_read(htsfile) };
        Reader { inner: htsfile, header: HeaderView::new(header) }
    }

    pub fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::vcf::bcf_read(self.inner, self.header.inner, record.inner) } {
            0  => {
                record.header = self.header.inner;
                Ok(())
            },
            -1 => Err(ReadError::NoMoreRecord),
            _  => Err(ReadError::Invalid),
        }
    }

    pub fn records(&self) -> Records {
        Records { reader: self }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Records<'a> {
    reader: &'a Reader,
}

impl<'a> Iterator for Records<'a> {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.reader.read(&mut record) {
            Err(ReadError::NoMoreRecord) => None,
            Err(e)                       => Some(Err(e)),
            Ok(())                       => Some(Ok(record)),
        }
    }
}

/// Wrapper for opening a BCF file.
fn bcf_open<P: AsRef<Path>>(path: &P, mode: &[u8]) -> *mut htslib::vcf::htsFile {
    unsafe {
        htslib::vcf::hts_open(
            path.as_ref().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    NoMoreRecord,
}

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

    #[test]
    fn test_read() {
        let bcf = Reader::new(&"test.bcf");
        assert_eq!(bcf.header.samples(), [b"NA12878.subsample-0.25-0"]);

        for (i, rec) in bcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");

            assert_eq!(record.rid().expect("Error reading rid."), 0);
            assert_eq!(record.pos(), 10021 + i as u32);
            assert_eq!(record.qual(), 0f32);
            assert_eq!(record.info(b"MQ0F").float().ok().expect("Error reading info."), [1.0]);
            if i == 59 {
                assert_eq!(record.info(b"SGB").float().ok().expect("Error reading info."), [-0.379885]);
            }
            // the artificial "not observed" allele is present in each record.
            assert_eq!(record.alleles().iter().last().unwrap(), b"<X>");

            let mut fmt = record.format(b"PL");
            let pl = fmt.integer().ok().expect("Error reading format.");
            assert_eq!(pl.len(), 1);
            if i == 59 {
                assert_eq!(pl[0].len(), 6);
            }
            else {
                assert_eq!(pl[0].len(), 3);
            }
        }
    }
}
<commit_msg>BCF writer.<commit_after>
use std::ffi;
use std::ffi::AsOsStr;
use std::convert::AsRef;
use std::path::Path;
use std::slice;

pub mod record;

use htslib;

pub struct HeaderView {
    inner: *mut htslib::vcf::bcf_hdr_t,
}

impl HeaderView {
    fn new(inner: *mut htslib::vcf::bcf_hdr_t) -> Self {
        HeaderView { inner: inner }
    }

    #[inline]
    fn inner(&self) -> htslib::vcf::bcf_hdr_t {
        unsafe { (*self.inner) }
    }

    pub fn sample_count(&self) -> u32 {
        self.inner().n[htslib::vcf::BCF_DT_SAMPLE as usize] as u32
    }

    pub fn samples(&self) -> Vec<&[u8]> {
        let names = unsafe { slice::from_raw_parts(self.inner().samples, self.sample_count() as usize) };
        names.iter().map(|name| unsafe { ffi::CStr::from_ptr(*name).to_bytes() }).collect()
    }
}

pub struct Reader {
    inner: *mut htslib::vcf::htsFile,
    header: HeaderView,
}

impl Reader {
   pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let htsfile = bcf_open(path, b"r");
        let header = unsafe { htslib::vcf::bcf_hdr_read(htsfile) };
        Reader { inner: htsfile, header: HeaderView::new(header) }
    }

    pub fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::vcf::bcf_read(self.inner, self.header.inner, record.inner) } {
            0  => {
                record.header = self.header.inner;
                Ok(())
            },
            -1 => Err(ReadError::NoMoreRecord),
            _  => Err(ReadError::Invalid),
        }
    }

    pub fn records(&self) -> Records {
        Records { reader: self }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Writer<'a> {
    inner: *mut htslib::vcf::htsFile,
    header: &'a HeaderView,
}

impl<'a> Writer<'a> {
    pub fn with_template<P: AsRef<Path>>(template: &'a Reader, path: &P) -> Self {
        let htsfile = bcf_open(path, b"w");
        unsafe { htslib::vcf::bcf_hdr_write(htsfile, template.header.inner) };
        Writer { inner: htsfile, header: &template.header }
    }

    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::vcf::bcf_write(self.inner, self.header.inner, record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl<'a> Drop for Writer<'a> {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Records<'a> {
    reader: &'a Reader,
}

impl<'a> Iterator for Records<'a> {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.reader.read(&mut record) {
            Err(ReadError::NoMoreRecord) => None,
            Err(e)                       => Some(Err(e)),
            Ok(())                       => Some(Ok(record)),
        }
    }
}

/// Wrapper for opening a BCF file.
fn bcf_open<P: AsRef<Path>>(path: &P, mode: &[u8]) -> *mut htslib::vcf::htsFile {
    unsafe {
        htslib::vcf::hts_open(
            path.as_ref().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    NoMoreRecord,
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use std::path::Path;

    fn _test_read<P: AsRef<Path>>(path: &P) {
        let bcf = Reader::new(path);
        assert_eq!(bcf.header.samples(), [b"NA12878.subsample-0.25-0"]);

        for (i, rec) in bcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");

            assert_eq!(record.rid().expect("Error reading rid."), 0);
            assert_eq!(record.pos(), 10021 + i as u32);
            assert_eq!(record.qual(), 0f32);
            assert_eq!(record.info(b"MQ0F").float().ok().expect("Error reading info."), [1.0]);
            if i == 59 {
                assert_eq!(record.info(b"SGB").float().ok().expect("Error reading info."), [-0.379885]);
            }
            // the artificial "not observed" allele is present in each record.
            assert_eq!(record.alleles().iter().last().unwrap(), b"<X>");

            let mut fmt = record.format(b"PL");
            let pl = fmt.integer().ok().expect("Error reading format.");
            assert_eq!(pl.len(), 1);
            if i == 59 {
                assert_eq!(pl[0].len(), 6);
            }
            else {
                assert_eq!(pl[0].len(), 3);
            }
        }
    }

    #[test]
    fn test_read() {
        _test_read(&"test.bcf");
    }

    #[test]
    fn test_write() {
        let bcf = Reader::new(&"test.bcf");
        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bcfpath = tmp.path().join("test.bcf");
        println!("{:?}", bcfpath);
        {
            let mut writer = Writer::with_template(&bcf, &bcfpath);
            for record in bcf.records() {
                writer.write(&record.ok().expect("Error reading record.")).ok().expect("Error writing record");
            }
        }
        {
            _test_read(&bcfpath);
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<|endoftext|>