<commit_before>use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use dovahkiin::types::custom_types::bytes::SmallBytes;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use index::btree::Ordering;
use index::EntryKey;
use index::Slice;
use index::{id_from_key, key_with_id};
use itertools::Itertools;
use parking_lot::Mutex;
use parking_lot::RwLock;
use ram::cell::Cell;
use ram::schema::Field;
use ram::schema::Schema;
use ram::types::*;
use std::cell::RefCell;
use std::collections::btree_map::Range;
use std::collections::hash_map::DefaultHasher;
use std::collections::{BTreeMap, BTreeSet, LinkedList};
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem;
use std::sync::Arc;
use utils::lru_cache::LRUCache;

// LevelTree items cannot been added or removed individually
// Items must been merged from higher level in bulk
// Deletion will be performed when merging to or from higher level tree
// Because tree update will not to be performed in parallel. Unlike memtable, a single r/w lock
// should be sufficient. Thus concurrency control will be simple and efficient.

type Tombstones = BTreeSet<EntryKey>;
type PageCache<S> = Arc<Mutex<LRUCache<Id, Arc<SSPage<S>>>>>;

const PAGE_SCHEMA: &'static str = "NEB_SSTABLE_PAGE";

lazy_static! {
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    neb_client: Arc<AsyncClient>,
    index: BTreeMap<EntryKey, Id>,
    tombstones: Tombstones,
    pages: PageCache<S>,
    marker: PhantomData<S>,
}
struct SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    id: Id,
    slice: S,
    len: usize,
}

impl<S> SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn from_cell(cell: Cell) -> Self {
        let mut slice = S::init();
        let slice_len = slice.as_slice().len();
        let cell_id = cell.id();
        let keys = cell.data;
        let keys_len = keys.len().unwrap();
        debug_assert!(slice_len == keys_len);
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        for (i, val) in keys_array.iter().enumerate() {
            slice.as_slice()[i] = EntryKey::from(val.as_slice());
        }
        Self {
            id: cell_id,
            slice,
            len: keys_array.len(),
        }
    }

    pub fn to_cell(&self) -> Cell
    {
        let value = self
            .slice
            .as_slice_immute()
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .take(self.len)
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }

    pub fn new() -> Self {
        SSPage {
            slice: S::init(),
            id: Id::unit_id(),
            len: 0,
        }
    }
}

impl<S> LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        let client = neb_client.clone();
        let mut index = BTreeMap::new();
        Self {
            neb_client: neb_client.clone(),
            index,
            tombstones: Tombstones::new(),
            // TODO: carefully set the capacity
            pages: Arc::new(Mutex::new(LRUCache::new(
                100,
                move |id| {
                    wait(client.read_cell(*id))
                        .unwrap()
                        .map(|cell| Arc::new(SSPage::from_cell(cell)))
                        .ok()
                },
                |_, _| {
                    // will do nothing on evict for pages will be write to back storage once updated
                },
            ))),
            marker: PhantomData {},
        }
    }

    pub fn seek(&self, key: &EntryKey, ordering: Ordering) -> RTCursor<S> {
        debug!("Seeking for {:?}, in index len {}", key, self.index.len());
        let mut range = match ordering {
            Ordering::Forward => self.index.range::<EntryKey, _>(key..),
            Ordering::Backward => self.index.range::<EntryKey, _>(..=key),
        };
        let first = match ordering {
            Ordering::Forward => range.next(),
            Ordering::Backward => range.next_back(),
        }
        .map(|(_, page_id)| {
            debug!("First page id is {:?}", page_id);
            self.pages.lock().get_or_fetch(page_id).unwrap().clone()
        });

        let (pos, page_len) = if let Some(ref page) = first {
            (
                page.slice
                    .as_slice_immute()
                    .binary_search(key)
                    .unwrap_or_else(|i| i),
                page.slice.as_slice_immute().len(),
            )
        } else {
            (0, 0)
        };
        let mut cursor = RTCursor {
            ordering,
            range,
            page_len,
            current: first,
            index: pos,
            cache: self.pages.clone(),
        };

        if cursor.current.is_some() {
            match ordering {
                Ordering::Forward => {
                    if pos >= page_len {
                        cursor.next();
                    }
                }
                Ordering::Backward => {}
            }
        }

        cursor
    }

    pub fn merge(&mut self, mut source: Vec<EntryKey>, tombstones: &mut Tombstones) {
        let mut pages_cache = self.pages.lock();
        let (keys_to_removed, mut merged, mut ids_to_reuse) = {
            let (target_keys, target_pages): (Vec<_>, Vec<_>) = {
                let first = source.first().unwrap();
                let last = source.last().unwrap();
                self.index
                    .range::<EntryKey, _>(first..=last)
                    .map(|(k, id)| (k, pages_cache.get_or_fetch(id).unwrap().clone()))
                    .unzip()
            };
            let mut merged: Vec<SSPage<S>> = vec![SSPage::new()];
            let page_size = merged[0].slice.len();
            let target_page_ids = target_pages.iter().map(|p| p.id).collect::<LinkedList<_>>();
            let mut target = target_pages
                .iter()
                .map(|page| &page.slice.as_slice_immute()[..page.len])
                .flat_map(|x| x)
                .map(|x| (*x).clone())
                .collect_vec();

            let src_len = source.len();
            let tgt_len = target.len();

            debug!("Src len {}, target len {}", src_len, tgt_len);

            let mut srci = 0;
            let mut tgti = 0;

            while srci < src_len || tgti < tgt_len {
                let mut from_source = false;
                let mut lowest: &mut EntryKey = if srci >= src_len {
                    // source overflowed
                    tgti += 1;
                    &mut target[tgti - 1]
                } else if tgti >= tgt_len {
                    // target overflowed
                    srci += 1;
                    from_source = true;
                    &mut source[srci - 1]
                } else {
                    // normal case, compare
                    let t = &mut target[tgti];
                    let s = &mut source[srci];
                    if t <= s {
                        tgti += 1;
                        t
                    } else {
                        srci += 1;
                        from_source = true;
                        s
                    }
                };

                if from_source {
                    if tombstones.remove(lowest) {
                        continue;
                    }
                } else {
                    if self.tombstones.remove(lowest) {
                        continue;
                    }
                }

                let merging_page_len = {
                    let merging_page = merged.last_mut().unwrap();
                    let page_len = merging_page.len;
                    debug!("i is {}, picking {:?}", page_len, lowest);
                    mem::swap(&mut merging_page.slice.as_slice()[page_len], lowest);
                    merging_page.len += 1;
                    merging_page.len
                };

                if merging_page_len >= page_size {
                    merged.push(SSPage::new());
                }
            }
            if merged.last().unwrap().len == 0 {
                let merged_len = merged.len();
                merged.remove(merged_len - 1);
            }
            (
                target_keys.into_iter().map(|x| x.clone()).collect_vec(),
                merged,
                target_page_ids,
            )
        };

        debug!("Merged {} pages", merged.len());

        for k in keys_to_removed {
            debug!("Remove stale page index {:?}", k);
            self.index.remove(&k).unwrap();
        }

        for id in &ids_to_reuse {
            debug!("Remove stale page cache {:?}", id);
            pages_cache.remove(id);
        }

        for mut page in merged {
            let reuse_id = ids_to_reuse.pop_back();
            let id = reuse_id.unwrap_or_else(|| Id::rand());
            let index_key = page.slice.as_slice_immute()[0].clone();
            page.id = id;
            let cell = page.to_cell();
            if reuse_id.is_some() {
                debug!("Reuse page id by updating {:?}", cell.id());
                wait(self.neb_client.update_cell(cell)).unwrap();
            } else {
                debug!("Inserting page with id {:?}", cell.id());
                wait(self.neb_client.write_cell(cell)).unwrap();
            }

            debug!("Insert page index key {:?}, id {:?}", index_key, id);
            self.index.insert(index_key, id);
        }
    }
}

pub trait SortableEntrySlice: Sized + Slice<Item = EntryKey> {}

pub struct RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    index: usize,
    ordering: Ordering,
    page_len: usize,
    range: Range<'a, EntryKey, Id>,
    current: Option<Arc<SSPage<S>>>,
    cache: PageCache<S>,
}

impl<'a, S> RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn next(&mut self) -> bool {
        if self.current.is_some() {
            match self.ordering {
                Ordering::Forward => {
                    if self.index + 1 >= self.page_len {
                        if let Some((_, page_id)) = self.range.next() {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.index = 0;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.index += 1;
                    }
                }
                Ordering::Backward => {
                    if self.index == 0 {
                        if let Some((_, page_id)) = self.range.next_back() {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.index = self.page_len - 1;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.index -= 1;
                    }
                }
            }
            true
        } else {
            false
        }
    }

    pub fn current(&self) -> Option<&EntryKey> {
        if self.index <= 0 || self.index >= self.page_len {
            return None;
        }
        if let Some(ref page) = self.current {
            Some(&page.slice.as_slice_immute()[self.index])
        } else {
            None
        }
    }

    fn get_page(&self, id: &Id) -> Arc<SSPage<S>> {
        self.cache.lock().get_or_fetch(id).unwrap().clone()
    }
}

pub fn get_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new("*", type_id_of(Type::SmallBytes), false, true, None),
    }
}

macro_rules! impl_sspage_slice {
    ($t: ty, $et: ty, $n: expr) => {
        impl_slice_ops!($t, $et, $n);
        impl SortableEntrySlice for $t {}
    };
}

#[cfg(test)]
mod test {
    use super::*;
    use client;
    use index::sstable::LevelTree;
    use server::NebServer;
    use server::ServerOptions;
    use std::ptr;
    use std::sync::Arc;

    type SmallPage = [EntryKey; 5];

    impl_sspage_slice!(SmallPage, EntryKey, 5);

    pub fn init() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5200");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 16 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(super::get_schema());

        let mut tree: LevelTree<SmallPage> = LevelTree::new(&client);
        let id = Id::unit_id();
        let key = smallvec![1, 2, 3, 4, 5, 6];
        let mut tombstones = BTreeSet::new();
        info!("test insertion");

        let mut key_id = key.clone();
        key_with_id(&mut key_id, &id);

        tree.merge(vec![key_id], &mut tombstones);
        let mut cursor = tree.seek(&key, Ordering::Forward);
        assert_eq!(id_from_key(cursor.current().unwrap()), id);
    }
}
<commit_msg>sstable init test passed<commit_after>use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use dovahkiin::types::custom_types::bytes::SmallBytes;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use index::btree::Ordering;
use index::EntryKey;
use index::Slice;
use index::{id_from_key, key_with_id};
use itertools::Itertools;
use parking_lot::Mutex;
use parking_lot::RwLock;
use ram::cell::Cell;
use ram::schema::Field;
use ram::schema::Schema;
use ram::types::*;
use std::cell::RefCell;
use std::collections::btree_map::Range;
use std::collections::hash_map::DefaultHasher;
use std::collections::{BTreeMap, BTreeSet, LinkedList};
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem;
use std::sync::Arc;
use utils::lru_cache::LRUCache;

// LevelTree items cannot been added or removed individually
// Items must been merged from higher level in bulk
// Deletion will be performed when merging to or from higher level tree
// Because tree update will not to be performed in parallel. Unlike memtable, a single r/w lock
// should be sufficient. Thus concurrency control will be simple and efficient.

type Tombstones = BTreeSet<EntryKey>;
type PageCache<S> = Arc<Mutex<LRUCache<Id, Arc<SSPage<S>>>>>;

const PAGE_SCHEMA: &'static str = "NEB_SSTABLE_PAGE";

lazy_static! {
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    neb_client: Arc<AsyncClient>,
    index: BTreeMap<EntryKey, Id>,
    tombstones: Tombstones,
    pages: PageCache<S>,
    marker: PhantomData<S>,
}
struct SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    id: Id,
    slice: S,
    len: usize,
}

impl<S> SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn from_cell(cell: Cell) -> Self {
        let mut slice = S::init();
        let slice_len = slice.as_slice().len();
        let cell_id = cell.id();
        let keys = cell.data;
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        for (i, val) in keys_array.iter().enumerate() {
            slice.as_slice()[i] = EntryKey::from(val.as_slice());
        }
        Self {
            id: cell_id,
            slice,
            len: keys_array.len(),
        }
    }

    pub fn to_cell(&self) -> Cell
    {
        let value = self
            .slice
            .as_slice_immute()
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .take(self.len)
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }

    pub fn new() -> Self {
        SSPage {
            slice: S::init(),
            id: Id::unit_id(),
            len: 0,
        }
    }
}

impl<S> LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        let client = neb_client.clone();
        let mut index = BTreeMap::new();
        Self {
            neb_client: neb_client.clone(),
            index,
            tombstones: Tombstones::new(),
            // TODO: carefully set the capacity
            pages: Arc::new(Mutex::new(LRUCache::new(
                100,
                move |id| {
                    wait(client.read_cell(*id))
                        .unwrap()
                        .map(|cell| Arc::new(SSPage::from_cell(cell)))
                        .ok()
                },
                |_, _| {
                    // will do nothing on evict for pages will be write to back storage once updated
                },
            ))),
            marker: PhantomData {},
        }
    }

    pub fn seek(&self, key: &EntryKey, ordering: Ordering) -> RTCursor<S> {
        debug!("Seeking for {:?}, in index len {}", key, self.index.len());
        let mut range = match ordering {
            Ordering::Forward => self.index.range::<EntryKey, _>(key..),
            Ordering::Backward => self.index.range::<EntryKey, _>(..=key),
        };
        let first = match ordering {
            Ordering::Forward => range.next(),
            Ordering::Backward => range.next_back(),
        }
        .map(|(_, page_id)| {
            debug!("First page id is {:?}", page_id);
            self.pages.lock().get_or_fetch(page_id).unwrap().clone()
        });

        let (pos, page_len) = if let Some(ref page) = first {
            debug!("First page has {} elements", page.len);
            (
                page.slice
                    .as_slice_immute()[..page.len]
                    .binary_search(key)
                    .unwrap_or_else(|i| i),
                page.len,
            )
        } else {
            (0, 0)
        };

        debug!("First page search pos is {}", pos);

        let mut cursor = RTCursor {
            ordering,
            range,
            page_len,
            current: first,
            index: pos,
            cache: self.pages.clone(),
        };

        if cursor.current.is_some() {
            match ordering {
                Ordering::Forward => {
                    if pos >= page_len {
                        cursor.next();
                    }
                }
                Ordering::Backward => {}
            }
        }

        cursor
    }

    pub fn merge(&mut self, mut source: Vec<EntryKey>, tombstones: &mut Tombstones) {
        let mut pages_cache = self.pages.lock();
        let (keys_to_removed, mut merged, mut ids_to_reuse) = {
            let (target_keys, target_pages): (Vec<_>, Vec<_>) = {
                let first = source.first().unwrap();
                let last = source.last().unwrap();
                self.index
                    .range::<EntryKey, _>(first..=last)
                    .map(|(k, id)| (k, pages_cache.get_or_fetch(id).unwrap().clone()))
                    .unzip()
            };
            let mut merged: Vec<SSPage<S>> = vec![SSPage::new()];
            let page_size = merged[0].slice.len();
            let target_page_ids = target_pages.iter().map(|p| p.id).collect::<LinkedList<_>>();
            let mut target = target_pages
                .iter()
                .map(|page| &page.slice.as_slice_immute()[..page.len])
                .flat_map(|x| x)
                .map(|x| (*x).clone())
                .collect_vec();

            let src_len = source.len();
            let tgt_len = target.len();

            debug!("Src len {}, target len {}", src_len, tgt_len);

            let mut srci = 0;
            let mut tgti = 0;

            while srci < src_len || tgti < tgt_len {
                let mut from_source = false;
                let mut lowest: &mut EntryKey = if srci >= src_len {
                    // source overflowed
                    tgti += 1;
                    &mut target[tgti - 1]
                } else if tgti >= tgt_len {
                    // target overflowed
                    srci += 1;
                    from_source = true;
                    &mut source[srci - 1]
                } else {
                    // normal case, compare
                    let t = &mut target[tgti];
                    let s = &mut source[srci];
                    if t <= s {
                        tgti += 1;
                        t
                    } else {
                        srci += 1;
                        from_source = true;
                        s
                    }
                };

                if from_source {
                    if tombstones.remove(lowest) {
                        continue;
                    }
                } else {
                    if self.tombstones.remove(lowest) {
                        continue;
                    }
                }

                let merging_page_len = {
                    let merging_page = merged.last_mut().unwrap();
                    let page_len = merging_page.len;
                    debug!("i is {}, picking {:?}", page_len, lowest);
                    mem::swap(&mut merging_page.slice.as_slice()[page_len], lowest);
                    merging_page.len += 1;
                    merging_page.len
                };

                if merging_page_len >= page_size {
                    merged.push(SSPage::new());
                }
            }
            if merged.last().unwrap().len == 0 {
                let merged_len = merged.len();
                merged.remove(merged_len - 1);
            }
            (
                target_keys.into_iter().map(|x| x.clone()).collect_vec(),
                merged,
                target_page_ids,
            )
        };

        debug!("Merged {} pages", merged.len());

        for k in keys_to_removed {
            debug!("Remove stale page index {:?}", k);
            self.index.remove(&k).unwrap();
        }

        for id in &ids_to_reuse {
            debug!("Remove stale page cache {:?}", id);
            pages_cache.remove(id);
        }

        for mut page in merged {
            let reuse_id = ids_to_reuse.pop_back();
            let id = reuse_id.unwrap_or_else(|| Id::rand());
            let index_key = page.slice.as_slice_immute()[0].clone();
            page.id = id;
            let cell = page.to_cell();
            if reuse_id.is_some() {
                debug!("Reuse page id by updating {:?}", cell.id());
                wait(self.neb_client.update_cell(cell)).unwrap();
            } else {
                debug!("Inserting page with id {:?}", cell.id());
                wait(self.neb_client.write_cell(cell)).unwrap();
            }

            debug!("Insert page index key {:?}, id {:?}", index_key, id);
            self.index.insert(index_key, id);
        }
    }
}

pub trait SortableEntrySlice: Sized + Slice<Item = EntryKey> {}

pub struct RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    index: usize,
    ordering: Ordering,
    page_len: usize,
    range: Range<'a, EntryKey, Id>,
    current: Option<Arc<SSPage<S>>>,
    cache: PageCache<S>,
}

impl<'a, S> RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn next(&mut self) -> bool {
        if self.current.is_some() {
            match self.ordering {
                Ordering::Forward => {
                    if self.index + 1 >= self.page_len {
                        if let Some((_, page_id)) = self.range.next() {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.index = 0;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.index += 1;
                    }
                }
                Ordering::Backward => {
                    if self.index == 0 {
                        if let Some((_, page_id)) = self.range.next_back() {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.index = self.page_len - 1;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.index -= 1;
                    }
                }
            }
            true
        } else {
            false
        }
    }

    pub fn current(&self) -> Option<&EntryKey> {
        if self.index < 0 || self.index >= self.page_len {
            return None;
        }
        if let Some(ref page) = self.current {
            Some(&page.slice.as_slice_immute()[self.index])
        } else {
            None
        }
    }

    fn get_page(&self, id: &Id) -> Arc<SSPage<S>> {
        self.cache.lock().get_or_fetch(id).unwrap().clone()
    }
}

pub fn get_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new("*", type_id_of(Type::SmallBytes), false, true, None),
    }
}

macro_rules! impl_sspage_slice {
    ($t: ty, $et: ty, $n: expr) => {
        impl_slice_ops!($t, $et, $n);
        impl SortableEntrySlice for $t {}
    };
}

#[cfg(test)]
mod test {
    use super::*;
    use client;
    use index::sstable::LevelTree;
    use server::NebServer;
    use server::ServerOptions;
    use std::ptr;
    use std::sync::Arc;
    use futures::prelude::*;

    type SmallPage = [EntryKey; 5];

    impl_sspage_slice!(SmallPage, EntryKey, 5);

    #[test]
    pub fn init() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5200");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 16 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(super::get_schema()).wait();

        let mut tree: LevelTree<SmallPage> = LevelTree::new(&client);
        let id = Id::unit_id();
        let key = smallvec![1, 2, 3, 4, 5, 6];
        let mut tombstones = BTreeSet::new();
        info!("test insertion");

        let mut key_id = key.clone();
        key_with_id(&mut key_id, &id);

        tree.merge(vec![key_id], &mut tombstones);
        let mut cursor = tree.seek(&key, Ordering::Forward);
        assert_eq!(id_from_key(cursor.current().unwrap()), id);
    }
}
<|endoftext|>


<commit_before>use core::mem::*;

#[test]
fn size_of_basic() {
    assert_eq!(size_of::<u8>(), 1);
    assert_eq!(size_of::<u16>(), 2);
    assert_eq!(size_of::<u32>(), 4);
    assert_eq!(size_of::<u64>(), 8);
}

#[test]
#[cfg(target_pointer_width = "16")]
fn size_of_16() {
    assert_eq!(size_of::<usize>(), 2);
    assert_eq!(size_of::<*const usize>(), 2);
}

#[test]
#[cfg(target_pointer_width = "32")]
fn size_of_32() {
    assert_eq!(size_of::<usize>(), 4);
    assert_eq!(size_of::<*const usize>(), 4);
}

#[test]
#[cfg(target_pointer_width = "64")]
fn size_of_64() {
    assert_eq!(size_of::<usize>(), 8);
    assert_eq!(size_of::<*const usize>(), 8);
}

#[test]
fn size_of_val_basic() {
    assert_eq!(size_of_val(&1u8), 1);
    assert_eq!(size_of_val(&1u16), 2);
    assert_eq!(size_of_val(&1u32), 4);
    assert_eq!(size_of_val(&1u64), 8);
}

#[test]
fn align_of_basic() {
    assert_eq!(align_of::<u8>(), 1);
    assert_eq!(align_of::<u16>(), 2);
    assert_eq!(align_of::<u32>(), 4);
}

#[test]
#[cfg(target_pointer_width = "16")]
fn align_of_16() {
    assert_eq!(align_of::<usize>(), 2);
    assert_eq!(align_of::<*const usize>(), 2);
}

#[test]
#[cfg(target_pointer_width = "32")]
fn align_of_32() {
    assert_eq!(align_of::<usize>(), 4);
    assert_eq!(align_of::<*const usize>(), 4);
}

#[test]
#[cfg(target_pointer_width = "64")]
fn align_of_64() {
    assert_eq!(align_of::<usize>(), 8);
    assert_eq!(align_of::<*const usize>(), 8);
}

#[test]
fn align_of_val_basic() {
    assert_eq!(align_of_val(&1u8), 1);
    assert_eq!(align_of_val(&1u16), 2);
    assert_eq!(align_of_val(&1u32), 4);
}

#[test]
fn test_swap() {
    let mut x = 31337;
    let mut y = 42;
    swap(&mut x, &mut y);
    assert_eq!(x, 42);
    assert_eq!(y, 31337);
}

#[test]
fn test_replace() {
    let mut x = Some("test".to_string());
    let y = replace(&mut x, None);
    assert!(x.is_none());
    assert!(y.is_some());
}

#[test]
fn test_transmute_copy() {
    assert_eq!(1, unsafe { transmute_copy(&1) });
}

#[test]
fn test_transmute() {
    trait Foo {
        fn dummy(&self) {}
    }
    impl Foo for isize {}

    let a = box 100isize as Box<dyn Foo>;
    unsafe {
        let x: ::core::raw::TraitObject = transmute(a);
        assert!(*(x.data as *const isize) == 100);
        let _x: Box<dyn Foo> = transmute(x);
    }

    unsafe {
        assert_eq!(transmute::<_, Vec<u8>>("L".to_string()), [76]);
    }
}

#[test]
#[allow(dead_code)]
fn test_discriminant_send_sync() {
    enum Regular {
        A,
        B(i32),
    }
    enum NotSendSync {
        A(*const i32),
    }

    fn is_send_sync<T: Send + Sync>() {}

    is_send_sync::<Discriminant<Regular>>();
    is_send_sync::<Discriminant<NotSendSync>>();
}

#[test]
fn test_const_forget() {
    const fn test_const_forget<T>(x: T) {
        forget(x);
    }

    // Writing this function signature without const-forget
    // triggers compiler errors:
    // 1) That we use a non-const fn inside a const fn
    // 2) without the forget, it complains about the destructor of Box
    const fn const_forget_box<T>(mut x: Box<T>) {
        forget(x);
    }

    const _: () = test_const_forget(0i32);
    const _: () = test_const_forget(Vec::<Vec<Box<i32>>>::new());
}<commit_msg>An enter as last character pleases tidy it seems<commit_after>use core::mem::*;

#[test]
fn size_of_basic() {
    assert_eq!(size_of::<u8>(), 1);
    assert_eq!(size_of::<u16>(), 2);
    assert_eq!(size_of::<u32>(), 4);
    assert_eq!(size_of::<u64>(), 8);
}

#[test]
#[cfg(target_pointer_width = "16")]
fn size_of_16() {
    assert_eq!(size_of::<usize>(), 2);
    assert_eq!(size_of::<*const usize>(), 2);
}

#[test]
#[cfg(target_pointer_width = "32")]
fn size_of_32() {
    assert_eq!(size_of::<usize>(), 4);
    assert_eq!(size_of::<*const usize>(), 4);
}

#[test]
#[cfg(target_pointer_width = "64")]
fn size_of_64() {
    assert_eq!(size_of::<usize>(), 8);
    assert_eq!(size_of::<*const usize>(), 8);
}

#[test]
fn size_of_val_basic() {
    assert_eq!(size_of_val(&1u8), 1);
    assert_eq!(size_of_val(&1u16), 2);
    assert_eq!(size_of_val(&1u32), 4);
    assert_eq!(size_of_val(&1u64), 8);
}

#[test]
fn align_of_basic() {
    assert_eq!(align_of::<u8>(), 1);
    assert_eq!(align_of::<u16>(), 2);
    assert_eq!(align_of::<u32>(), 4);
}

#[test]
#[cfg(target_pointer_width = "16")]
fn align_of_16() {
    assert_eq!(align_of::<usize>(), 2);
    assert_eq!(align_of::<*const usize>(), 2);
}

#[test]
#[cfg(target_pointer_width = "32")]
fn align_of_32() {
    assert_eq!(align_of::<usize>(), 4);
    assert_eq!(align_of::<*const usize>(), 4);
}

#[test]
#[cfg(target_pointer_width = "64")]
fn align_of_64() {
    assert_eq!(align_of::<usize>(), 8);
    assert_eq!(align_of::<*const usize>(), 8);
}

#[test]
fn align_of_val_basic() {
    assert_eq!(align_of_val(&1u8), 1);
    assert_eq!(align_of_val(&1u16), 2);
    assert_eq!(align_of_val(&1u32), 4);
}

#[test]
fn test_swap() {
    let mut x = 31337;
    let mut y = 42;
    swap(&mut x, &mut y);
    assert_eq!(x, 42);
    assert_eq!(y, 31337);
}

#[test]
fn test_replace() {
    let mut x = Some("test".to_string());
    let y = replace(&mut x, None);
    assert!(x.is_none());
    assert!(y.is_some());
}

#[test]
fn test_transmute_copy() {
    assert_eq!(1, unsafe { transmute_copy(&1) });
}

#[test]
fn test_transmute() {
    trait Foo {
        fn dummy(&self) {}
    }
    impl Foo for isize {}

    let a = box 100isize as Box<dyn Foo>;
    unsafe {
        let x: ::core::raw::TraitObject = transmute(a);
        assert!(*(x.data as *const isize) == 100);
        let _x: Box<dyn Foo> = transmute(x);
    }

    unsafe {
        assert_eq!(transmute::<_, Vec<u8>>("L".to_string()), [76]);
    }
}

#[test]
#[allow(dead_code)]
fn test_discriminant_send_sync() {
    enum Regular {
        A,
        B(i32),
    }
    enum NotSendSync {
        A(*const i32),
    }

    fn is_send_sync<T: Send + Sync>() {}

    is_send_sync::<Discriminant<Regular>>();
    is_send_sync::<Discriminant<NotSendSync>>();
}

#[test]
fn test_const_forget() {
    const fn test_const_forget<T>(x: T) {
        forget(x);
    }

    // Writing this function signature without const-forget
    // triggers compiler errors:
    // 1) That we use a non-const fn inside a const fn
    // 2) without the forget, it complains about the destructor of Box
    const fn const_forget_box<T>(mut x: Box<T>) {
        forget(x);
    }

    const _: () = test_const_forget(0i32);
    const _: () = test_const_forget(Vec::<Vec<Box<i32>>>::new());
}
<|endoftext|>


<commit_before>use nom::{ConsumerState, FileProducer, Producer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use utility::resize_producer;

use metadata::{
  Metadata, Data,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
};
use metadata::types::MetaDataConsumer;

// Will return true when the unwrapped value of `option` and `other` match
// or `option` is `Option::None`, otherwise false.
#[inline]
pub fn optional_eq<T: Eq>(option: Option<T>, other: T) -> bool {
  option.map_or(true, |value| value == other)
}

// With the given filename, return all metadata blocks available.
//
// This function expects a flac file, but will return a proper `Result::Err`
// when things go wrong.
//
// # Failures
//
// * `ErrorKind::NotFound` is returned when the given filename isn't found.
// * `ErrorKind::InvalidData` is returned when the data within the file
//   isn't valid FLAC data.
pub fn get_metadata(filename: &str) -> Result<Vec<Metadata>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer    = MetaDataConsumer::new();
    let mut buffer_size = 1024;
    let mut is_error    = false;

    loop {
      match *producer.apply(&mut consumer) {
        ConsumerState::Done(_, _)      => break,
        ConsumerState::Continue(await) => {
          let result = resize_producer(&mut producer, &await, buffer_size);

          if let Some(size) = result {
            buffer_size = size;
          }

          continue;
        }
        ConsumerState::Error(_)        => {
          is_error = true;

          break;
        }
      }
    }

    if is_error {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    } else {
      Ok(consumer.data)
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_stream_info("path/to/file.flac") {
///   Ok(stream_info) => {
///     // Use the stream_info variable...
///   }
///   Err(error)      => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let stream_info = metadata::get_stream_info("path/to/file.flac").unwrap();
/// ```
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let Data::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_vorbis_comment("path/to/file.flac") {
///   Ok(vorbis_comment) => {
///     // Use the vorbis_comment variable...
///   }
///   Err(error)         => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let vorbis_comment =
///   metadata::get_vorbis_comment("path/to/file.flac").unwrap();
/// ```
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let Data::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_cue_sheet("path/to/file.flac") {
///   Ok(cue_sheet) => {
///     // Use the cue_sheet variable...
///   }
///   Err(error)    => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors
///
/// ```no_run
/// use flac::metadata;
///
/// let cue_sheet = metadata::get_cue_sheet("path/to/file.flac").unwrap();
/// ```
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let Data::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file and this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// Putting `None` into any of the optional constraints conveys that you
/// want any of that parameter. Otherwise it will try to look for the image
/// that matches within the given constraints.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let result = metadata::get_picture("path/to/file.flac",
///                                    Some(PictureType::FileIcon),
///                                    Some("image/gif"),
///                                    None, None, None,
///                                    None, None);
///
/// match result {
///   Ok(picture) => {
///     // Use the picture variable...
///   }
///   Err(error)  => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let picture = metadata::get_picture("path/to/file.flac",
///                                     Some(PictureType::FileIcon),
///                                     Some("image/gif"),
///                                     None, None, None,
///                                     None, None).unwrap();
/// ```
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let Data::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq(picture_type, picture.picture_type) &&
           optional_eq(mime_type, &picture.mime_type) &&
           optional_eq(description, &picture.description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

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

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq(Some(1), 0));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq(None, 0), "Should always return true when None");
    assert!(optional_eq(Some(0), 0), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<commit_msg>Rewrite `get_metadata` to use `ReadStream`<commit_after>use nom::{ConsumerState, FileProducer, Producer};
use std::io::{Error, ErrorKind, Result};
use std::u32;
use std::fs::File;

use utility::{ErrorKind, ReadStream};

use metadata::{
  Metadata, Data,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
};
use metadata::types::Consumer;

// Will return true when the unwrapped value of `option` and `other` match
// or `option` is `Option::None`, otherwise false.
#[inline]
pub fn optional_eq<T: Eq>(option: Option<T>, other: T) -> bool {
  option.map_or(true, |value| value == other)
}

// With the given filename, return all metadata blocks available.
//
// This function expects a flac file, but will return a proper `Result::Err`
// when things go wrong.
//
// # Failures
//
// * `ErrorKind::NotFound` is returned when the given filename isn't found.
// * `ErrorKind::InvalidData` is returned when the data within the file
//   isn't valid FLAC data.
pub fn get_metadata(filename: &str) -> Result<Vec<Metadata>> {
  File::open(filename).and_then(|file| {
    let mut stream   = ReadStream::new(file);
    let mut consumer = Consumer::new();
    let mut is_error = false;

    loop {
      match consumer.handle(&mut stream) {
        Ok(_)                         => break,
        Err(ErrorKind::Consumed(_))   => continue,
        Err(ErrorKind::Incomplete(_)) => continue,
        Err(_)                        => {
          is_error = true;

          break;
        }
      }
    }

    if is_error {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    } else {
      Ok(consumer.data)
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_stream_info("path/to/file.flac") {
///   Ok(stream_info) => {
///     // Use the stream_info variable...
///   }
///   Err(error)      => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let stream_info = metadata::get_stream_info("path/to/file.flac").unwrap();
/// ```
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let Data::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_vorbis_comment("path/to/file.flac") {
///   Ok(vorbis_comment) => {
///     // Use the vorbis_comment variable...
///   }
///   Err(error)         => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let vorbis_comment =
///   metadata::get_vorbis_comment("path/to/file.flac").unwrap();
/// ```
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let Data::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_cue_sheet("path/to/file.flac") {
///   Ok(cue_sheet) => {
///     // Use the cue_sheet variable...
///   }
///   Err(error)    => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors
///
/// ```no_run
/// use flac::metadata;
///
/// let cue_sheet = metadata::get_cue_sheet("path/to/file.flac").unwrap();
/// ```
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let Data::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file and this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// Putting `None` into any of the optional constraints conveys that you
/// want any of that parameter. Otherwise it will try to look for the image
/// that matches within the given constraints.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let result = metadata::get_picture("path/to/file.flac",
///                                    Some(PictureType::FileIcon),
///                                    Some("image/gif"),
///                                    None, None, None,
///                                    None, None);
///
/// match result {
///   Ok(picture) => {
///     // Use the picture variable...
///   }
///   Err(error)  => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let picture = metadata::get_picture("path/to/file.flac",
///                                     Some(PictureType::FileIcon),
///                                     Some("image/gif"),
///                                     None, None, None,
///                                     None, None).unwrap();
/// ```
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let Data::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq(picture_type, picture.picture_type) &&
           optional_eq(mime_type, &picture.mime_type) &&
           optional_eq(description, &picture.description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

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

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq(Some(1), 0));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq(None, 0), "Should always return true when None");
    assert!(optional_eq(Some(0), 0), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<|endoftext|>


<commit_before>//! `core_arch`

#[macro_use]
mod macros;

#[cfg(any(target_arch = "arm", target_arch = "aarch64", dox))]
mod acle;

mod simd;

#[doc(include = "core_arch_docs.md")]
#[stable(feature = "simd_arch", since = "1.27.0")]
pub mod arch {
    /// Platform-specific intrinsics for the `x86` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "x86", dox))]
    #[doc(cfg(target_arch = "x86"))]
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub mod x86 {
        #[stable(feature = "simd_x86", since = "1.27.0")]
        pub use crate::core_arch::x86::*;
    }

    /// Platform-specific intrinsics for the `x86_64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "x86_64", dox))]
    #[doc(cfg(target_arch = "x86_64"))]
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub mod x86_64 {
        #[stable(feature = "simd_x86", since = "1.27.0")]
        pub use crate::core_arch::x86::*;
        #[stable(feature = "simd_x86", since = "1.27.0")]
        pub use crate::core_arch::x86_64::*;
    }

    /// Platform-specific intrinsics for the `arm` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "arm", dox))]
    #[doc(cfg(target_arch = "arm"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod arm {
        pub use crate::core_arch::arm::*;
    }

    /// Platform-specific intrinsics for the `aarch64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "aarch64", dox))]
    #[doc(cfg(target_arch = "aarch64"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod aarch64 {
        pub use crate::core_arch::{aarch64::*, arm::*};
    }

    /// Platform-specific intrinsics for the `wasm32` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "wasm32", dox))]
    #[doc(cfg(target_arch = "wasm32"))]
    #[stable(feature = "simd_wasm32", since = "1.33.0")]
    pub mod wasm32 {
        #[stable(feature = "simd_wasm32", since = "1.33.0")]
        pub use crate::core_arch::wasm32::*;
    }

    /// Platform-specific intrinsics for the `mips` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "mips", dox))]
    #[doc(cfg(target_arch = "mips"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod mips {
        pub use crate::core_arch::mips::*;
    }

    /// Platform-specific intrinsics for the `mips64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "mips64", dox))]
    #[doc(cfg(target_arch = "mips64"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod mips64 {
        pub use crate::core_arch::mips::*;
    }

    /// Platform-specific intrinsics for the `PowerPC` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "powerpc", dox))]
    #[doc(cfg(target_arch = "powerpc"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod powerpc {
        pub use crate::core_arch::powerpc::*;
    }

    /// Platform-specific intrinsics for the `PowerPC64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "powerpc64", dox))]
    #[doc(cfg(target_arch = "powerpc64"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod powerpc64 {
        pub use crate::core_arch::powerpc64::*;
    }

    /// Platform-specific intrinsics for the `NVPTX` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "nvptx", target_arch = "nvptx64", dox))]
    #[doc(cfg(any(target_arch = "nvptx", target_arch = "nvptx64")))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod nvptx {
        pub use crate::core_arch::nvptx::*;
    }
}

mod simd_llvm;

#[cfg(any(target_arch = "x86", target_arch = "x86_64", dox))]
#[doc(cfg(any(target_arch = "x86", target_arch = "x86_64")))]
mod x86;
#[cfg(any(target_arch = "x86_64", dox))]
#[doc(cfg(target_arch = "x86_64"))]
mod x86_64;

#[cfg(any(target_arch = "aarch64", dox))]
#[doc(cfg(target_arch = "aarch64"))]
mod aarch64;
#[cfg(any(target_arch = "arm", target_arch = "aarch64", dox))]
#[doc(cfg(any(target_arch = "arm", target_arch = "aarch64")))]
mod arm;

#[cfg(any(target_arch = "wasm32", dox))]
#[doc(cfg(target_arch = "wasm32"))]
mod wasm32;

#[cfg(any(target_arch = "mips", target_arch = "mips64", dox))]
#[doc(cfg(any(target_arch = "mips", target_arch = "mips64")))]
mod mips;

#[cfg(any(target_arch = "powerpc", target_arch = "powerpc64", dox))]
#[doc(cfg(any(target_arch = "powerpc", target_arch = "powerpc64")))]
mod powerpc;

#[cfg(any(target_arch = "powerpc64", dox))]
#[doc(cfg(target_arch = "powerpc64"))]
mod powerpc64;

#[cfg(any(target_arch = "nvptx", target_arch = "nvptx64", dox))]
#[doc(cfg(any(target_arch = "nvptx", target_arch = "nvptx64")))]
mod nvptx;
<commit_msg>Clarify documentation about wasm32 target_feature gates<commit_after>//! `core_arch`

#[macro_use]
mod macros;

#[cfg(any(target_arch = "arm", target_arch = "aarch64", dox))]
mod acle;

mod simd;

#[doc(include = "core_arch_docs.md")]
#[stable(feature = "simd_arch", since = "1.27.0")]
pub mod arch {
    /// Platform-specific intrinsics for the `x86` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "x86", dox))]
    #[doc(cfg(target_arch = "x86"))]
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub mod x86 {
        #[stable(feature = "simd_x86", since = "1.27.0")]
        pub use crate::core_arch::x86::*;
    }

    /// Platform-specific intrinsics for the `x86_64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "x86_64", dox))]
    #[doc(cfg(target_arch = "x86_64"))]
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub mod x86_64 {
        #[stable(feature = "simd_x86", since = "1.27.0")]
        pub use crate::core_arch::x86::*;
        #[stable(feature = "simd_x86", since = "1.27.0")]
        pub use crate::core_arch::x86_64::*;
    }

    /// Platform-specific intrinsics for the `arm` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "arm", dox))]
    #[doc(cfg(target_arch = "arm"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod arm {
        pub use crate::core_arch::arm::*;
    }

    /// Platform-specific intrinsics for the `aarch64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "aarch64", dox))]
    #[doc(cfg(target_arch = "aarch64"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod aarch64 {
        pub use crate::core_arch::{aarch64::*, arm::*};
    }

    /// Platform-specific intrinsics for the `wasm32` platform.
    ///

    /// # Availability
    ///
    /// Note that intrinsics gated by `target_feature = "atomics"` or `target_feature = "simd128"`
    /// are only available **when the standard library itself is compiled with the the respective
    /// target feature**. This version of the standard library is not obtainable via `rustup`,
    /// but rather will require the standard library to be compiled from source.
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "wasm32", dox))]
    #[doc(cfg(target_arch = "wasm32"))]
    #[stable(feature = "simd_wasm32", since = "1.33.0")]
    pub mod wasm32 {
        #[stable(feature = "simd_wasm32", since = "1.33.0")]
        pub use crate::core_arch::wasm32::*;
    }

    /// Platform-specific intrinsics for the `mips` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "mips", dox))]
    #[doc(cfg(target_arch = "mips"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod mips {
        pub use crate::core_arch::mips::*;
    }

    /// Platform-specific intrinsics for the `mips64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "mips64", dox))]
    #[doc(cfg(target_arch = "mips64"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod mips64 {
        pub use crate::core_arch::mips::*;
    }

    /// Platform-specific intrinsics for the `PowerPC` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "powerpc", dox))]
    #[doc(cfg(target_arch = "powerpc"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod powerpc {
        pub use crate::core_arch::powerpc::*;
    }

    /// Platform-specific intrinsics for the `PowerPC64` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "powerpc64", dox))]
    #[doc(cfg(target_arch = "powerpc64"))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod powerpc64 {
        pub use crate::core_arch::powerpc64::*;
    }

    /// Platform-specific intrinsics for the `NVPTX` platform.
    ///
    /// See the [module documentation](../index.html) for more details.
    #[cfg(any(target_arch = "nvptx", target_arch = "nvptx64", dox))]
    #[doc(cfg(any(target_arch = "nvptx", target_arch = "nvptx64")))]
    #[unstable(feature = "stdsimd", issue = "27731")]
    pub mod nvptx {
        pub use crate::core_arch::nvptx::*;
    }
}

mod simd_llvm;

#[cfg(any(target_arch = "x86", target_arch = "x86_64", dox))]
#[doc(cfg(any(target_arch = "x86", target_arch = "x86_64")))]
mod x86;
#[cfg(any(target_arch = "x86_64", dox))]
#[doc(cfg(target_arch = "x86_64"))]
mod x86_64;

#[cfg(any(target_arch = "aarch64", dox))]
#[doc(cfg(target_arch = "aarch64"))]
mod aarch64;
#[cfg(any(target_arch = "arm", target_arch = "aarch64", dox))]
#[doc(cfg(any(target_arch = "arm", target_arch = "aarch64")))]
mod arm;

#[cfg(any(target_arch = "wasm32", dox))]
#[doc(cfg(target_arch = "wasm32"))]
mod wasm32;

#[cfg(any(target_arch = "mips", target_arch = "mips64", dox))]
#[doc(cfg(any(target_arch = "mips", target_arch = "mips64")))]
mod mips;

#[cfg(any(target_arch = "powerpc", target_arch = "powerpc64", dox))]
#[doc(cfg(any(target_arch = "powerpc", target_arch = "powerpc64")))]
mod powerpc;

#[cfg(any(target_arch = "powerpc64", dox))]
#[doc(cfg(target_arch = "powerpc64"))]
mod powerpc64;

#[cfg(any(target_arch = "nvptx", target_arch = "nvptx64", dox))]
#[doc(cfg(any(target_arch = "nvptx", target_arch = "nvptx64")))]
mod nvptx;
<|endoftext|>


<commit_before>//! Trait solving using Chalk.
use std::sync::Arc;

use base_db::CrateId;
use chalk_ir::cast::Cast;
use chalk_solve::{logging_db::LoggingRustIrDatabase, Solver};
use hir_def::{lang_item::LangItemTarget, TraitId};

use crate::{db::HirDatabase, DebruijnIndex, Substs};

use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};

use self::chalk::{from_chalk, Interner, ToChalk};

pub(crate) mod chalk;

// This controls the maximum size of types Chalk considers. If we set this too
// high, we can run into slow edge cases; if we set it too low, Chalk won't
// find some solutions.
// FIXME this is currently hardcoded in the recursive solver
// const CHALK_SOLVER_MAX_SIZE: usize = 10;

/// This controls how much 'time' we give the Chalk solver before giving up.
const CHALK_SOLVER_FUEL: i32 = 100;

#[derive(Debug, Copy, Clone)]
struct ChalkContext<'a> {
    db: &'a dyn HirDatabase,
    krate: CrateId,
}

fn create_chalk_solver() -> chalk_recursive::RecursiveSolver<Interner> {
    let overflow_depth = 100;
    let caching_enabled = true;
    chalk_recursive::RecursiveSolver::new(overflow_depth, caching_enabled)
}

/// A set of clauses that we assume to be true. E.g. if we are inside this function:
/// ```rust
/// fn foo<T: Default>(t: T) {}
/// ```
/// we assume that `T: Default`.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TraitEnvironment {
    pub predicates: Vec<GenericPredicate>,
}

impl TraitEnvironment {
    /// Returns trait refs with the given self type which are supposed to hold
    /// in this trait env. E.g. if we are in `foo<T: SomeTrait>()`, this will
    /// find that `T: SomeTrait` if we call it for `T`.
    pub(crate) fn trait_predicates_for_self_ty<'a>(
        &'a self,
        ty: &'a Ty,
    ) -> impl Iterator<Item = &'a TraitRef> + 'a {
        self.predicates.iter().filter_map(move |pred| match pred {
            GenericPredicate::Implemented(tr) if tr.self_ty() == ty => Some(tr),
            _ => None,
        })
    }
}

/// Something (usually a goal), along with an environment.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct InEnvironment<T> {
    pub environment: Arc<TraitEnvironment>,
    pub value: T,
}

impl<T> InEnvironment<T> {
    pub fn new(environment: Arc<TraitEnvironment>, value: T) -> InEnvironment<T> {
        InEnvironment { environment, value }
    }
}

/// Something that needs to be proven (by Chalk) during type checking, e.g. that
/// a certain type implements a certain trait. Proving the Obligation might
/// result in additional information about inference variables.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Obligation {
    /// Prove that a certain type implements a trait (the type is the `Self` type
    /// parameter to the `TraitRef`).
    Trait(TraitRef),
    Projection(ProjectionPredicate),
}

impl Obligation {
    pub fn from_predicate(predicate: GenericPredicate) -> Option<Obligation> {
        match predicate {
            GenericPredicate::Implemented(trait_ref) => Some(Obligation::Trait(trait_ref)),
            GenericPredicate::Projection(projection_pred) => {
                Some(Obligation::Projection(projection_pred))
            }
            GenericPredicate::Error => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ProjectionPredicate {
    pub projection_ty: ProjectionTy,
    pub ty: Ty,
}

impl TypeWalk for ProjectionPredicate {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.projection_ty.walk(f);
        self.ty.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.projection_ty.walk_mut_binders(f, binders);
        self.ty.walk_mut_binders(f, binders);
    }
}

/// Solve a trait goal using Chalk.
pub(crate) fn trait_solve_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: Canonical<InEnvironment<Obligation>>,
) -> Option<Solution> {
    let _p = profile::span("trait_solve_query").detail(|| match &goal.value.value {
        Obligation::Trait(it) => db.trait_data(it.trait_).name.to_string(),
        Obligation::Projection(_) => "projection".to_string(),
    });
    log::info!("trait_solve_query({})", goal.value.value.display(db));

    if let Obligation::Projection(pred) = &goal.value.value {
        if let Ty::Bound(_) = &pred.projection_ty.parameters[0] {
            // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
            return Some(Solution::Ambig(Guidance::Unknown));
        }
    }

    let canonical = goal.to_chalk(db).cast(&Interner);

    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
    let solution = solve(db, krate, &u_canonical);
    solution.map(|solution| solution_from_chalk(db, solution))
}

fn solve(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
) -> Option<chalk_solve::Solution<Interner>> {
    let context = ChalkContext { db, krate };
    log::debug!("solve goal: {:?}", goal);
    let mut solver = create_chalk_solver();

    let fuel = std::cell::Cell::new(CHALK_SOLVER_FUEL);

    let should_continue = || {
        context.db.check_canceled();
        let remaining = fuel.get();
        fuel.set(remaining - 1);
        if remaining == 0 {
            log::debug!("fuel exhausted");
        }
        remaining > 0
    };

    let solution = if is_chalk_debug() {
        let logging_db = LoggingRustIrDatabase::new(context);
        let solve = || {
            let solution = solver.solve_limited(&logging_db, goal, should_continue);
            log::debug!("chalk program:\n{}", logging_db);
            solution
        };

        // don't set the TLS for Chalk unless Chalk debugging is active, to make
        // extra sure we only use it for debugging
        chalk::tls::set_current_program(db, solve)
    } else {
        solver.solve_limited(&context, goal, should_continue)
    };

    log::debug!("solve({:?}) => {:?}", goal, solution);

    solution
}

fn is_chalk_debug() -> bool {
    std::env::var("CHALK_DEBUG").is_ok()
}

fn solution_from_chalk(
    db: &dyn HirDatabase,
    solution: chalk_solve::Solution<Interner>,
) -> Solution {
    let convert_subst = |subst: chalk_ir::Canonical<chalk_ir::Substitution<Interner>>| {
        let result = from_chalk(db, subst);
        SolutionVariables(result)
    };
    match solution {
        chalk_solve::Solution::Unique(constr_subst) => {
            let subst = chalk_ir::Canonical {
                value: constr_subst.value.subst,
                binders: constr_subst.binders,
            };
            Solution::Unique(convert_subst(subst))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Definite(subst)) => {
            Solution::Ambig(Guidance::Definite(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Suggested(subst)) => {
            Solution::Ambig(Guidance::Suggested(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Unknown) => {
            Solution::Ambig(Guidance::Unknown)
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SolutionVariables(pub Canonical<Substs>);

#[derive(Clone, Debug, PartialEq, Eq)]
/// A (possible) solution for a proposed goal.
pub enum Solution {
    /// The goal indeed holds, and there is a unique value for all existential
    /// variables.
    Unique(SolutionVariables),

    /// The goal may be provable in multiple ways, but regardless we may have some guidance
    /// for type inference. In this case, we don't return any lifetime
    /// constraints, since we have not "committed" to any particular solution
    /// yet.
    Ambig(Guidance),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// When a goal holds ambiguously (e.g., because there are multiple possible
/// solutions), we issue a set of *guidance* back to type inference.
pub enum Guidance {
    /// The existential variables *must* have the given values if the goal is
    /// ever to hold, but that alone isn't enough to guarantee the goal will
    /// actually hold.
    Definite(SolutionVariables),

    /// There are multiple plausible values for the existentials, but the ones
    /// here are suggested as the preferred choice heuristically. These should
    /// be used for inference fallback only.
    Suggested(SolutionVariables),

    /// There's no useful information to feed back to type inference
    Unknown,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum FnTrait {
    FnOnce,
    FnMut,
    Fn,
}

impl FnTrait {
    fn lang_item_name(self) -> &'static str {
        match self {
            FnTrait::FnOnce => "fn_once",
            FnTrait::FnMut => "fn_mut",
            FnTrait::Fn => "fn",
        }
    }

    pub fn get_id(&self, db: &dyn HirDatabase, krate: CrateId) -> Option<TraitId> {
        let target = db.lang_item(krate, self.lang_item_name().into())?;
        match target {
            LangItemTarget::TraitId(t) => Some(t),
            _ => None,
        }
    }
}
<commit_msg>Only print chalk programs with CHALK_PRINT<commit_after>//! Trait solving using Chalk.
use std::sync::Arc;

use base_db::CrateId;
use chalk_ir::cast::Cast;
use chalk_solve::{logging_db::LoggingRustIrDatabase, Solver};
use hir_def::{lang_item::LangItemTarget, TraitId};

use crate::{db::HirDatabase, DebruijnIndex, Substs};

use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};

use self::chalk::{from_chalk, Interner, ToChalk};

pub(crate) mod chalk;

// This controls the maximum size of types Chalk considers. If we set this too
// high, we can run into slow edge cases; if we set it too low, Chalk won't
// find some solutions.
// FIXME this is currently hardcoded in the recursive solver
// const CHALK_SOLVER_MAX_SIZE: usize = 10;

/// This controls how much 'time' we give the Chalk solver before giving up.
const CHALK_SOLVER_FUEL: i32 = 100;

#[derive(Debug, Copy, Clone)]
struct ChalkContext<'a> {
    db: &'a dyn HirDatabase,
    krate: CrateId,
}

fn create_chalk_solver() -> chalk_recursive::RecursiveSolver<Interner> {
    let overflow_depth = 100;
    let caching_enabled = true;
    chalk_recursive::RecursiveSolver::new(overflow_depth, caching_enabled)
}

/// A set of clauses that we assume to be true. E.g. if we are inside this function:
/// ```rust
/// fn foo<T: Default>(t: T) {}
/// ```
/// we assume that `T: Default`.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TraitEnvironment {
    pub predicates: Vec<GenericPredicate>,
}

impl TraitEnvironment {
    /// Returns trait refs with the given self type which are supposed to hold
    /// in this trait env. E.g. if we are in `foo<T: SomeTrait>()`, this will
    /// find that `T: SomeTrait` if we call it for `T`.
    pub(crate) fn trait_predicates_for_self_ty<'a>(
        &'a self,
        ty: &'a Ty,
    ) -> impl Iterator<Item = &'a TraitRef> + 'a {
        self.predicates.iter().filter_map(move |pred| match pred {
            GenericPredicate::Implemented(tr) if tr.self_ty() == ty => Some(tr),
            _ => None,
        })
    }
}

/// Something (usually a goal), along with an environment.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct InEnvironment<T> {
    pub environment: Arc<TraitEnvironment>,
    pub value: T,
}

impl<T> InEnvironment<T> {
    pub fn new(environment: Arc<TraitEnvironment>, value: T) -> InEnvironment<T> {
        InEnvironment { environment, value }
    }
}

/// Something that needs to be proven (by Chalk) during type checking, e.g. that
/// a certain type implements a certain trait. Proving the Obligation might
/// result in additional information about inference variables.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Obligation {
    /// Prove that a certain type implements a trait (the type is the `Self` type
    /// parameter to the `TraitRef`).
    Trait(TraitRef),
    Projection(ProjectionPredicate),
}

impl Obligation {
    pub fn from_predicate(predicate: GenericPredicate) -> Option<Obligation> {
        match predicate {
            GenericPredicate::Implemented(trait_ref) => Some(Obligation::Trait(trait_ref)),
            GenericPredicate::Projection(projection_pred) => {
                Some(Obligation::Projection(projection_pred))
            }
            GenericPredicate::Error => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ProjectionPredicate {
    pub projection_ty: ProjectionTy,
    pub ty: Ty,
}

impl TypeWalk for ProjectionPredicate {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.projection_ty.walk(f);
        self.ty.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.projection_ty.walk_mut_binders(f, binders);
        self.ty.walk_mut_binders(f, binders);
    }
}

/// Solve a trait goal using Chalk.
pub(crate) fn trait_solve_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: Canonical<InEnvironment<Obligation>>,
) -> Option<Solution> {
    let _p = profile::span("trait_solve_query").detail(|| match &goal.value.value {
        Obligation::Trait(it) => db.trait_data(it.trait_).name.to_string(),
        Obligation::Projection(_) => "projection".to_string(),
    });
    log::info!("trait_solve_query({})", goal.value.value.display(db));

    if let Obligation::Projection(pred) = &goal.value.value {
        if let Ty::Bound(_) = &pred.projection_ty.parameters[0] {
            // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
            return Some(Solution::Ambig(Guidance::Unknown));
        }
    }

    let canonical = goal.to_chalk(db).cast(&Interner);

    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
    let solution = solve(db, krate, &u_canonical);
    solution.map(|solution| solution_from_chalk(db, solution))
}

fn solve(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
) -> Option<chalk_solve::Solution<Interner>> {
    let context = ChalkContext { db, krate };
    log::debug!("solve goal: {:?}", goal);
    let mut solver = create_chalk_solver();

    let fuel = std::cell::Cell::new(CHALK_SOLVER_FUEL);

    let should_continue = || {
        context.db.check_canceled();
        let remaining = fuel.get();
        fuel.set(remaining - 1);
        if remaining == 0 {
            log::debug!("fuel exhausted");
        }
        remaining > 0
    };

    let mut solve = || {
        if is_chalk_print() {
            let logging_db = LoggingRustIrDatabase::new(context);
            let solution = solver.solve_limited(&logging_db, goal, should_continue);
            log::debug!("chalk program:\n{}", logging_db);
            solution
        } else {
            solver.solve_limited(&context, goal, should_continue)
        }
    };

    // don't set the TLS for Chalk unless Chalk debugging is active, to make
    // extra sure we only use it for debugging
    let solution =
        if is_chalk_debug() { chalk::tls::set_current_program(db, solve) } else { solve() };

    log::debug!("solve({:?}) => {:?}", goal, solution);

    solution
}

fn is_chalk_debug() -> bool {
    std::env::var("CHALK_DEBUG").is_ok()
}

fn is_chalk_print() -> bool {
    std::env::var("CHALK_PRINT").is_ok()
}

fn solution_from_chalk(
    db: &dyn HirDatabase,
    solution: chalk_solve::Solution<Interner>,
) -> Solution {
    let convert_subst = |subst: chalk_ir::Canonical<chalk_ir::Substitution<Interner>>| {
        let result = from_chalk(db, subst);
        SolutionVariables(result)
    };
    match solution {
        chalk_solve::Solution::Unique(constr_subst) => {
            let subst = chalk_ir::Canonical {
                value: constr_subst.value.subst,
                binders: constr_subst.binders,
            };
            Solution::Unique(convert_subst(subst))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Definite(subst)) => {
            Solution::Ambig(Guidance::Definite(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Suggested(subst)) => {
            Solution::Ambig(Guidance::Suggested(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Unknown) => {
            Solution::Ambig(Guidance::Unknown)
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SolutionVariables(pub Canonical<Substs>);

#[derive(Clone, Debug, PartialEq, Eq)]
/// A (possible) solution for a proposed goal.
pub enum Solution {
    /// The goal indeed holds, and there is a unique value for all existential
    /// variables.
    Unique(SolutionVariables),

    /// The goal may be provable in multiple ways, but regardless we may have some guidance
    /// for type inference. In this case, we don't return any lifetime
    /// constraints, since we have not "committed" to any particular solution
    /// yet.
    Ambig(Guidance),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// When a goal holds ambiguously (e.g., because there are multiple possible
/// solutions), we issue a set of *guidance* back to type inference.
pub enum Guidance {
    /// The existential variables *must* have the given values if the goal is
    /// ever to hold, but that alone isn't enough to guarantee the goal will
    /// actually hold.
    Definite(SolutionVariables),

    /// There are multiple plausible values for the existentials, but the ones
    /// here are suggested as the preferred choice heuristically. These should
    /// be used for inference fallback only.
    Suggested(SolutionVariables),

    /// There's no useful information to feed back to type inference
    Unknown,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum FnTrait {
    FnOnce,
    FnMut,
    Fn,
}

impl FnTrait {
    fn lang_item_name(self) -> &'static str {
        match self {
            FnTrait::FnOnce => "fn_once",
            FnTrait::FnMut => "fn_mut",
            FnTrait::Fn => "fn",
        }
    }

    pub fn get_id(&self, db: &dyn HirDatabase, krate: CrateId) -> Option<TraitId> {
        let target = db.lang_item(krate, self.lang_item_name().into())?;
        match target {
            LangItemTarget::TraitId(t) => Some(t),
            _ => None,
        }
    }
}
<|endoftext|>