<commit_before>use std::iter::{Iterator, Peekable};
use std::collections::btree_map::{BTreeMap, Range};
use std::cmp::Ordering::*;

use super::Result;

use self::NextIterValue::*;


pub type Patch = BTreeMap<Vec<u8>, Change>;
pub type Iter<'a> = Box<Iterator<Item = (&'a [u8], &'a [u8])> + 'a>;

#[derive(Debug, Clone)]
pub enum Change {
    Put(Vec<u8>),
    Delete,
}

pub struct Fork {
    snapshot: Box<Snapshot>,
    changes: Patch,
    changelog: Vec<(Vec<u8>, Option<Change>)>,
}

pub struct ForkIter<'a> {
    snapshot: Peekable<Iter<'a>>,
    changes: Peekable<Range<'a, Vec<u8>, Change>>,
}

#[derive(Debug)]
enum NextIterValue<'a> {
    Stored(&'a [u8], &'a [u8]),
    Replaced(&'a [u8], &'a [u8]),
    Inserted(&'a [u8], &'a [u8]),
    Deleted,
    MissDeleted,
}

pub trait Database: Send + Sync + 'static {
    fn clone(&self) -> Box<Database>;
    fn snapshot(&self) -> Box<Snapshot>;
    fn fork(&self) -> Fork {
        Fork {
            snapshot: self.snapshot(),
            changes: Patch::new(),
            changelog: Vec::new(),
        }
    }
    fn merge(&mut self, patch: Patch) -> Result<()>;
}

pub trait Snapshot {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>>;
    fn contains(&self, key: &[u8]) -> bool {
        self.get(key).is_some()
    }
    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a>;
}

impl Snapshot for Fork {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(ref v) => Some(v.clone()),
                    Change::Delete => None,
                }
            }
            None => self.snapshot.get(key),
        }
    }

    fn contains(&self, key: &[u8]) -> bool {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(..) => true,
                    Change::Delete => false,
                }
            }
            None => self.snapshot.get(key).is_some(),
        }
    }

    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a> {
        use std::collections::Bound::*;
        let range = (Included(from), Unbounded);
        Box::new(ForkIter {
                     snapshot: self.snapshot.iter(from).peekable(),
                     changes: self.changes.range::<[u8], _>(range).peekable(),
                 })
    }
}

impl Fork {
    pub fn checkpoint(&mut self) {
        self.changelog.clear()
    }

    pub fn rollback(&mut self) {
        for (k, c) in self.changelog.drain(..).rev() {
            match c {
                Some(change) => self.changes.insert(k, change),
                None => self.changes.remove(&k),
            };
        }
    }

    pub fn put(&mut self, key: Vec<u8>, value: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Put(value))))
    }

    pub fn remove(&mut self, key: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Delete)));
    }

    pub fn remove_by_prefix(&mut self, prefix: &[u8]) {
        for (k, _) in self.snapshot.iter(prefix) {
            if !k.starts_with(prefix) {
                return;
            }
            self.changes.insert(k.to_vec(), Change::Delete);
        }
    }

    pub fn into_patch(self) -> Patch {
        self.changes
    }

    pub fn merge(&mut self, patch: Patch) {
        self.changes.extend(patch)
    }
}

// TODO: Does we needed this AsRef / AsMut impls?

impl AsRef<Snapshot> for Snapshot {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        self
    }
}

impl AsRef<Snapshot + 'static> for Fork {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        &*self
    }
}

impl ::std::fmt::Debug for Fork {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "Fork(..)")
    }
}

impl<'a> NextIterValue<'a> {
    fn skip_changes(&self) -> bool {
        match *self {
            Replaced(..) | Inserted(..) | Deleted | MissDeleted => true,
            Stored(..) => false,
        }
    }

    fn skip_snapshot(&self) -> bool {
        match *self {
            Stored(..) | Replaced(..) | Deleted => true,
            Inserted(..) | MissDeleted => false,
        }
    }

    fn value(&self) -> Option<(&'a [u8], &'a [u8])> {
        match *self {
            Stored(k, v) => Some((k, v)),
            Replaced(k, v) => Some((k, v)),
            Inserted(k, v) => Some((k, v)),
            Deleted | MissDeleted => None,
        }
    }
}

impl<'a> Iterator for ForkIter<'a> {
    type Item = (&'a [u8], &'a [u8]);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let next = match self.changes.peek() {
                Some(&(k, ref change)) => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => {
                            match **change {
                                Change::Put(ref v) => {
                                    match k[..].cmp(key) {
                                        Equal => Replaced(k, v),
                                        Less => Inserted(k, v),
                                        Greater => Stored(key, value),
                                    }
                                }
                                Change::Delete => {
                                    match k[..].cmp(key) {
                                        Equal => Deleted,
                                        Less => MissDeleted,
                                        Greater => Stored(key, value),
                                    }
                                }
                            }
                        }
                        None => {
                            match **change {
                                Change::Put(ref v) => Inserted(k, v),
                                Change::Delete => MissDeleted,
                            }
                        }
                    }
                }
                None => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => Stored(key, value),
                        None => return None,
                    }
                }
            };
            if next.skip_changes() {
                self.changes.next();
            }
            if next.skip_snapshot() {
                self.snapshot.next();
            }
            if let Some(value) = next.value() {
                return Some(value);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::{MemoryDB, Database, Snapshot, Fork};

    #[test]
    fn fork_iter() {
        let mut db = MemoryDB::new();
        let mut fork = db.fork();

        fork.put(vec![10], vec![10]);
        fork.put(vec![20], vec![20]);
        fork.put(vec![30], vec![30]);

        db.merge(fork.into_patch()).unwrap();

        let mut fork = db.fork();

        fn assert_iter(fork: &Fork, from: u8, assumed: &[(u8, u8)]) {
            assert_eq!(fork.iter(&[from])
                           .map(|(k, v)| (k[0], v[0]))
                           .collect::<Vec<_>>(),
                       assumed);
        }

        // Stored
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 5, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 10, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 11, &[(20, 20), (30, 30)]);
        assert_iter(&fork, 31, &[]);

        // Inserted
        fork.put(vec![5], vec![5]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (30, 30)]);
        fork.put(vec![25], vec![25]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30)]);
        fork.put(vec![35], vec![35]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30), (35, 35)]);

        // Double inserted
        fork.put(vec![25], vec![23]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (30, 30), (35, 35)]);
        fork.put(vec![26], vec![26]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (26, 26), (30, 30), (35, 35)]);

        // Replaced
        let mut fork = db.fork();
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 30)]);
        fork.put(vec![30], vec![31]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 31)]);

        // Deleted
        let mut fork = db.fork();
        fork.remove(vec![20]);
        assert_iter(&fork, 0, &[(10, 10), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);

        // MissDeleted
        let mut fork = db.fork();
        fork.remove(vec![5]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![15]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![35]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
    }
}
<commit_msg>Panic when trying to merge into unfinalized fork<commit_after>use std::iter::{Iterator, Peekable};
use std::collections::btree_map::{BTreeMap, Range};
use std::cmp::Ordering::*;

use super::Result;

use self::NextIterValue::*;


pub type Patch = BTreeMap<Vec<u8>, Change>;
pub type Iter<'a> = Box<Iterator<Item = (&'a [u8], &'a [u8])> + 'a>;

#[derive(Debug, Clone)]
pub enum Change {
    Put(Vec<u8>),
    Delete,
}

pub struct Fork {
    snapshot: Box<Snapshot>,
    changes: Patch,
    changelog: Vec<(Vec<u8>, Option<Change>)>,
}

pub struct ForkIter<'a> {
    snapshot: Peekable<Iter<'a>>,
    changes: Peekable<Range<'a, Vec<u8>, Change>>,
}

#[derive(Debug)]
enum NextIterValue<'a> {
    Stored(&'a [u8], &'a [u8]),
    Replaced(&'a [u8], &'a [u8]),
    Inserted(&'a [u8], &'a [u8]),
    Deleted,
    MissDeleted,
}

pub trait Database: Send + Sync + 'static {
    fn clone(&self) -> Box<Database>;
    fn snapshot(&self) -> Box<Snapshot>;
    fn fork(&self) -> Fork {
        Fork {
            snapshot: self.snapshot(),
            changes: Patch::new(),
            changelog: Vec::new(),
        }
    }
    fn merge(&mut self, patch: Patch) -> Result<()>;
}

pub trait Snapshot {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>>;
    fn contains(&self, key: &[u8]) -> bool {
        self.get(key).is_some()
    }
    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a>;
}

impl Snapshot for Fork {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(ref v) => Some(v.clone()),
                    Change::Delete => None,
                }
            }
            None => self.snapshot.get(key),
        }
    }

    fn contains(&self, key: &[u8]) -> bool {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(..) => true,
                    Change::Delete => false,
                }
            }
            None => self.snapshot.get(key).is_some(),
        }
    }

    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a> {
        use std::collections::Bound::*;
        let range = (Included(from), Unbounded);
        Box::new(ForkIter {
                     snapshot: self.snapshot.iter(from).peekable(),
                     changes: self.changes.range::<[u8], _>(range).peekable(),
                 })
    }
}

impl Fork {
    pub fn checkpoint(&mut self) {
        self.changelog.clear()
    }

    pub fn rollback(&mut self) {
        for (k, c) in self.changelog.drain(..).rev() {
            match c {
                Some(change) => self.changes.insert(k, change),
                None => self.changes.remove(&k),
            };
        }
    }

    pub fn put(&mut self, key: Vec<u8>, value: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Put(value))))
    }

    pub fn remove(&mut self, key: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Delete)));
    }

    pub fn remove_by_prefix(&mut self, prefix: &[u8]) {
        for (k, _) in self.snapshot.iter(prefix) {
            if !k.starts_with(prefix) {
                return;
            }
            self.changes.insert(k.to_vec(), Change::Delete);
        }
    }

    pub fn into_patch(self) -> Patch {
        self.changes
    }

    pub fn merge(&mut self, patch: Patch) {
        if !self.changelog.is_empty() {
            panic!("merge into a fork is impossible because it has unfinalized changes");
        }
        self.changes.extend(patch)
    }
}

// TODO: Does we needed this AsRef / AsMut impls?

impl AsRef<Snapshot> for Snapshot {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        self
    }
}

impl AsRef<Snapshot + 'static> for Fork {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        &*self
    }
}

impl ::std::fmt::Debug for Fork {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "Fork(..)")
    }
}

impl<'a> NextIterValue<'a> {
    fn skip_changes(&self) -> bool {
        match *self {
            Replaced(..) | Inserted(..) | Deleted | MissDeleted => true,
            Stored(..) => false,
        }
    }

    fn skip_snapshot(&self) -> bool {
        match *self {
            Stored(..) | Replaced(..) | Deleted => true,
            Inserted(..) | MissDeleted => false,
        }
    }

    fn value(&self) -> Option<(&'a [u8], &'a [u8])> {
        match *self {
            Stored(k, v) => Some((k, v)),
            Replaced(k, v) => Some((k, v)),
            Inserted(k, v) => Some((k, v)),
            Deleted | MissDeleted => None,
        }
    }
}

impl<'a> Iterator for ForkIter<'a> {
    type Item = (&'a [u8], &'a [u8]);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let next = match self.changes.peek() {
                Some(&(k, ref change)) => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => {
                            match **change {
                                Change::Put(ref v) => {
                                    match k[..].cmp(key) {
                                        Equal => Replaced(k, v),
                                        Less => Inserted(k, v),
                                        Greater => Stored(key, value),
                                    }
                                }
                                Change::Delete => {
                                    match k[..].cmp(key) {
                                        Equal => Deleted,
                                        Less => MissDeleted,
                                        Greater => Stored(key, value),
                                    }
                                }
                            }
                        }
                        None => {
                            match **change {
                                Change::Put(ref v) => Inserted(k, v),
                                Change::Delete => MissDeleted,
                            }
                        }
                    }
                }
                None => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => Stored(key, value),
                        None => return None,
                    }
                }
            };
            if next.skip_changes() {
                self.changes.next();
            }
            if next.skip_snapshot() {
                self.snapshot.next();
            }
            if let Some(value) = next.value() {
                return Some(value);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::{MemoryDB, Database, Snapshot, Fork};

    #[test]
    fn fork_iter() {
        let mut db = MemoryDB::new();
        let mut fork = db.fork();

        fork.put(vec![10], vec![10]);
        fork.put(vec![20], vec![20]);
        fork.put(vec![30], vec![30]);

        db.merge(fork.into_patch()).unwrap();

        let mut fork = db.fork();

        fn assert_iter(fork: &Fork, from: u8, assumed: &[(u8, u8)]) {
            assert_eq!(fork.iter(&[from])
                           .map(|(k, v)| (k[0], v[0]))
                           .collect::<Vec<_>>(),
                       assumed);
        }

        // Stored
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 5, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 10, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 11, &[(20, 20), (30, 30)]);
        assert_iter(&fork, 31, &[]);

        // Inserted
        fork.put(vec![5], vec![5]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (30, 30)]);
        fork.put(vec![25], vec![25]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30)]);
        fork.put(vec![35], vec![35]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30), (35, 35)]);

        // Double inserted
        fork.put(vec![25], vec![23]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (30, 30), (35, 35)]);
        fork.put(vec![26], vec![26]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (26, 26), (30, 30), (35, 35)]);

        // Replaced
        let mut fork = db.fork();
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 30)]);
        fork.put(vec![30], vec![31]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 31)]);

        // Deleted
        let mut fork = db.fork();
        fork.remove(vec![20]);
        assert_iter(&fork, 0, &[(10, 10), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);

        // MissDeleted
        let mut fork = db.fork();
        fork.remove(vec![5]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![15]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![35]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
    }
}

{'fraction_non_alphanumeric': 0.1849499920622321, 'fraction_numerical': 0.0358787109064931, 'mean_word_length': 1.451420417693605, 'pattern_counts': {'":': 0, '<': 91, '<?xml version=': 0, '>': 195, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6925631', 'n_tokens_mistral': 6595, 'n_tokens_neox': 6033, 'n_words': 1348}