<commit_before>use std::collections::HashSet;
use std::iter::FromIterator;
use {Trie, TrieCommon};
use keys::TrieKey;

const TEST_DATA: [(&'static str, u32); 7] = [("abcdefgh", 19),
                                             ("abcdef", 18),
                                             ("abcd", 17),
                                             ("ab", 16),
                                             ("a", 15),
                                             ("acbdef", 30),
                                             ("bcdefgh", 29)];

fn test_trie() -> Trie<&'static str, u32> {
    let mut trie = Trie::new();

    for &(key, val) in &TEST_DATA {
        trie.insert(key, val);
        assert!(trie.check_integrity());
    }

    trie
}

#[test]
fn get_nonexistant() {
    let trie = test_trie();
    assert!(trie.get(&"nonexistant").is_none());
    assert!(trie.get(&"").is_none());
}

#[test]
fn subtrie_nonexistant() {
    let trie = test_trie();
    assert!(trie.subtrie(&"nonexistant").is_none());
    assert!(trie.subtrie(&"").is_some());
}

#[test]
fn subtrie() {
    let mut trie = Trie::new();
    trie.insert("hello", 55);
    assert!(trie.subtrie(&"h").is_none());
    assert!(trie.subtrie(&"hello").is_some());
}
#[test]
fn subtrie_string() {
    let mut trie = Trie::new();
    trie.insert("hello".to_string(), 55);

    let h = "h".to_string();
    let hello = "hello".to_string();

    assert!(trie.subtrie(&h).is_none());
    assert!(trie.subtrie(&hello).is_some());
}

#[test]
fn empty_key() {
    let mut trie = test_trie();
    trie.insert(&"", 99);
    assert_eq!(*trie.get(&"").unwrap(), 99);
    assert_eq!(trie.remove(&""), Some(99));
}

#[test]
fn insert() {
    let trie = test_trie();

    for &(key, val) in &TEST_DATA {
        assert_eq!(*trie.get(&key).unwrap(), val);
    }

    assert!(trie.check_integrity());
    assert_eq!(trie.len(), TEST_DATA.len());
}

#[test]
fn insert_replace() {
    let mut trie = Trie::new();
    assert_eq!(trie.insert("haskell", 18), None);
    let length = trie.len();
    assert_eq!(trie.insert("haskell", 36), Some(18));
    assert_eq!(trie.len(), length);
}

#[test]
fn map_with_default() {
    let mut trie = test_trie();
    trie.map_with_default(&"abcd",
                          {
                              |x| *x = *x + 1
                          },
                          42);
    assert_eq!(*trie.get(&"abcd").unwrap(), 17 + 1);
    trie.map_with_default(&"zzz",
                          {
                              |x| *x = *x + 1
                          },
                          42);
    assert_eq!(*trie.get(&"zzz").unwrap(), 42);
}

#[test]
fn remove() {
    let mut trie = test_trie();

    // Remove.
    for &(key, val) in &TEST_DATA {
        let res = trie.remove(&key);
        assert_eq!(res, Some(val));
        assert!(trie.check_integrity());
    }

    // Check non-existance.
    for &(key, _) in &TEST_DATA {
        assert!(trie.get(&key).is_none());
    }
}

#[test]
fn remove_simple() {
    let mut trie = Trie::new();

    trie.insert("HELL", 66);
    trie.insert("HELLO", 77);
    let val = trie.remove(&"HELLO");
    assert_eq!(val, Some(77));
}

#[test]
fn nearest_ancestor_root() {
    let mut trie = Trie::new();
    trie.insert("", 55);
    assert_eq!(trie.get_ancestor_value(&""), Some(&55));
}

#[test]
fn nearest_ancestor() {
    let trie = test_trie();
    assert_eq!(trie.get_ancestor_value(&""), None);

    // Test identity prefixes.
    for &(key, val) in &TEST_DATA {
        assert_eq!(trie.get_ancestor_value(&key), Some(&val));
    }

    assert_eq!(trie.get_ancestor_value(&"abcdefg"), trie.get(&"abcdef"));
    assert_eq!(trie.get_ancestor_value(&"abcde"), trie.get(&"abcd"));
    assert_eq!(trie.get_ancestor_value(&"aauksdjk"), trie.get(&"a"));
}

#[test]
fn nearest_ancestor_no_child_fn() {
    let mut t = Trie::new();
    t.insert("ab", 5);
    let anc = t.get_ancestor(&"abc").unwrap();
    assert_eq!(*anc.value().unwrap(), 5);
}

#[test]
fn raw_ancestor() {
    let mut t = Trie::new();

    for &(key, _) in &TEST_DATA {
        assert_eq!(t.get_raw_ancestor(&key).key(), t.key());
    }

    t.insert("wow", 0);
    t.insert("hella", 1);
    t.insert("hellb", 2);

    // Ancestor should be "hell" node.
    let anc = t.get_raw_ancestor(&"hello");
    assert_eq!(anc.len(), 2);
}

// Check that the subtrie prefix is correct for raw_ancestor.
#[test]
fn raw_ancestor_prefix() {
    let mut t = Trie::new();

    t.insert("abac", ());
    t.insert("abaz", ());

    let anc = t.get_raw_ancestor(&"aba");
    assert_eq!(anc.prefix, "aba".encode());
}

#[test]
fn iter() {
    type Set = HashSet<(&'static str, u32)>;
    let trie = test_trie();
    let expected = TEST_DATA.iter().map(|&x| x).collect::<Set>();
    let observed = trie.iter().map(|(&k, &v)| (k, v)).collect::<Set>();
    assert_eq!(expected, observed);
}

#[test]
fn get_raw_descendant() {
    let trie = test_trie();
    assert_eq!(trie.get_raw_descendant(&"abcdefgh").and_then(|t| t.value()),
               Some(&19));
    assert_eq!(trie.get_raw_descendant(&"abcdefg").and_then(|t| t.value()),
               Some(&19));
    assert!(trie.get_raw_descendant(&"acbg").is_none());
}

#[test]
fn raw_descendant_prefix() {
    let mut t = Trie::new();

    t.insert("abczzzz", ());
    t.insert("abcaaaa", ());

    assert_eq!(t.get_raw_descendant(&"a").unwrap().prefix, "abc".encode());
    assert_eq!(t.get_raw_descendant(&"abc").unwrap().prefix, "abc".encode());
    assert_eq!(t.get_raw_descendant(&"abca").unwrap().prefix, "abcaaaa".encode());
}

#[test]
fn get_prefix_bug() {
    let mut trie = Trie::new();
    trie.insert("abdc", 5);
    trie.insert("abde", 6);
    assert!(trie.get(&"abc").is_none());
}

#[test]
fn get_ancestor_bug() {
    let mut trie = Trie::new();
    trie.insert("abc", 1);
    trie.insert("abcde", 2);
    assert_eq!(trie.get_ancestor_value(&"abcdz"), Some(&1));
}

#[test]
fn root_replace_bug() {
    let mut trie = Trie::new();
    trie.insert("a", ());
    trie.insert("p", ());
    trie.remove(&"a");
    assert_eq!(trie.len(), 1);
    trie.remove(&"p");
    assert_eq!(trie.len(), 0);
}

#[test]
fn subtrie_insert() {
    let mut trie = Trie::new();
    trie.insert("abc", 3);
    {
        let mut subtrie = trie.subtrie_mut(&"abc").unwrap();
        assert_eq!(subtrie.insert("somerandomshit", 666), Err(()));
        assert_eq!(subtrie.insert("abcdef", 6), Ok(None));
        assert_eq!(subtrie.insert("abc", 9), Ok(Some(3)));
    }
    assert_eq!(trie.get(&"abc"), Some(&9));
    assert_eq!(trie.get(&"abcdef"), Some(&6));
    assert_eq!(trie.len(), 2);
}

#[test]
fn subtrie_len() {
    let trie = test_trie();
    assert_eq!(trie.subtrie(&"ab").unwrap().len(), 4);
    assert_eq!(trie.subtrie(&"a").unwrap().len(), 6);
    assert_eq!(trie.subtrie(&"").unwrap().len(), trie.len());
    assert_eq!(trie.subtrie(&"bcdefgh").unwrap().len(), 1);
}

// Subtrie functions that return references should return references valid for
// the lifetime of the *original trie* from which they were borrowed, NOT
// just the lifetime of the subtrie (which may be shorter).
#[test]
fn subtrie_lifetime() {
    let trie = test_trie();
    let subtrie_value = {
        let subtrie = trie.subtrie(&"ab").unwrap();
        subtrie.value()
    };
    assert_eq!(*subtrie_value.unwrap(), 16);
}

#[test]
fn subtrie_mut_lifetime() {
    let mut trie = test_trie();
    let subtrie_value = {
        let mut subtrie = trie.subtrie_mut(&"ab").unwrap();
        *subtrie.value_mut().unwrap() = 999;
        subtrie.value()
    };
    assert_eq!(*subtrie_value.unwrap(), 999);
}

#[test]
fn ancestor_key() {
    let trie = test_trie();
    let subtrie = trie.get_ancestor(&"abcde").unwrap();
    assert_eq!(*subtrie.key().unwrap(), "abcd");
    assert_eq!(subtrie.prefix, "abcd".encode());
    assert_eq!(*subtrie.get(&"abcdef").unwrap().unwrap(), 18);
    assert_eq!(*subtrie.get(&"abcdefgh").unwrap().unwrap(), 19);
}

#[test]
fn child_subtrie_keys() {
    let trie = test_trie();
    let subtrie = trie.subtrie(&"abcd").unwrap();
    for subsubtrie in subtrie.children() {
        // This subtrie should be for "abcde".
        assert_eq!(*subsubtrie.get(&"abcdef").unwrap().unwrap(), 18);
        assert_eq!(*subsubtrie.get(&"abcdefgh").unwrap().unwrap(), 19);
    }
}

#[test]
fn int_keys() {
    let mut trie = Trie::new();
    trie.insert(0x00ffu64, "asdf");
    trie.insert(0xdeadbeefu64, "asdf");
    assert!(trie.check_integrity());
}

#[test]
fn from_iter() {
    let trie: Trie<&str, u32> = Trie::from_iter(vec![
        ("test", 10), ("hello", 12)
    ]);
    assert_eq!(*trie.get(&"test").unwrap(), 10);
    assert_eq!(*trie.get(&"hello").unwrap(), 12);
    assert_eq!(trie.len(), 2);
}

#[test]
fn test_get_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    assert_eq!(*trie.get("/boot").unwrap(), "dir");
}

#[test]
fn test_get_mut_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    assert_eq!(*trie.get_mut("/boot").unwrap(), "dir");
}

#[test]
fn test_remove_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    assert_eq!(trie.remove("/boot").unwrap(), "dir");
}

#[test]
fn test_subtrie_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.subtrie("/boot").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "dir");
}

#[test]
fn test_subtrie_mut_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.subtrie_mut("/boot").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "dir");
}

#[test]
fn test_get_ancestor_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "ancestor");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.get_ancestor("/boot/lo").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "ancestor");
}

#[test]
fn test_get_ancestor_value_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "ancestor");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let ancestor_val = trie.get_ancestor_value("/boot/lo").unwrap();
    assert_eq!(*ancestor_val, "ancestor");
}

#[test]
fn test_get_raw_ancestor_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "boot");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.get_raw_ancestor("/boot/lo");
    assert_eq!(*subtrie.value().unwrap(), "boot");
}

#[test]
fn test_get_raw_descendant_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    trie.insert("/boot/lol".to_string(), "lol");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.get_raw_descendant("/boot/lo").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "lol");
}
<commit_msg>add test for unicode support<commit_after>use std::collections::HashSet;
use std::iter::FromIterator;
use {Trie, TrieCommon};
use keys::TrieKey;

const TEST_DATA: [(&'static str, u32); 7] = [("abcdefgh", 19),
                                             ("abcdef", 18),
                                             ("abcd", 17),
                                             ("ab", 16),
                                             ("a", 15),
                                             ("acbdef", 30),
                                             ("bcdefgh", 29)];

fn test_trie() -> Trie<&'static str, u32> {
    let mut trie = Trie::new();

    for &(key, val) in &TEST_DATA {
        trie.insert(key, val);
        assert!(trie.check_integrity());
    }

    trie
}

#[test]
fn get_nonexistant() {
    let trie = test_trie();
    assert!(trie.get(&"nonexistant").is_none());
    assert!(trie.get(&"").is_none());
}

#[test]
fn subtrie_nonexistant() {
    let trie = test_trie();
    assert!(trie.subtrie(&"nonexistant").is_none());
    assert!(trie.subtrie(&"").is_some());
}

#[test]
fn unicode() {
    let mut trie = Trie::new();
    trie.insert("bär", 1);
    trie.insert("bären", 2);

    assert_eq!(*trie.get("bär").unwrap(), 1);
    let values = trie.get_raw_descendant("bä").unwrap().values().collect::<HashSet<_>>();
    assert_eq!([1,2].iter().collect::<HashSet<_>>(), values);
}

#[test]
fn subtrie() {
    let mut trie = Trie::new();
    trie.insert("hello", 55);
    assert!(trie.subtrie(&"h").is_none());
    assert!(trie.subtrie(&"hello").is_some());
}
#[test]
fn subtrie_string() {
    let mut trie = Trie::new();
    trie.insert("hello".to_string(), 55);

    let h = "h".to_string();
    let hello = "hello".to_string();

    assert!(trie.subtrie(&h).is_none());
    assert!(trie.subtrie(&hello).is_some());
}

#[test]
fn empty_key() {
    let mut trie = test_trie();
    trie.insert(&"", 99);
    assert_eq!(*trie.get(&"").unwrap(), 99);
    assert_eq!(trie.remove(&""), Some(99));
}

#[test]
fn insert() {
    let trie = test_trie();

    for &(key, val) in &TEST_DATA {
        assert_eq!(*trie.get(&key).unwrap(), val);
    }

    assert!(trie.check_integrity());
    assert_eq!(trie.len(), TEST_DATA.len());
}

#[test]
fn insert_replace() {
    let mut trie = Trie::new();
    assert_eq!(trie.insert("haskell", 18), None);
    let length = trie.len();
    assert_eq!(trie.insert("haskell", 36), Some(18));
    assert_eq!(trie.len(), length);
}

#[test]
fn map_with_default() {
    let mut trie = test_trie();
    trie.map_with_default(&"abcd",
                          {
                              |x| *x = *x + 1
                          },
                          42);
    assert_eq!(*trie.get(&"abcd").unwrap(), 17 + 1);
    trie.map_with_default(&"zzz",
                          {
                              |x| *x = *x + 1
                          },
                          42);
    assert_eq!(*trie.get(&"zzz").unwrap(), 42);
}

#[test]
fn remove() {
    let mut trie = test_trie();

    // Remove.
    for &(key, val) in &TEST_DATA {
        let res = trie.remove(&key);
        assert_eq!(res, Some(val));
        assert!(trie.check_integrity());
    }

    // Check non-existance.
    for &(key, _) in &TEST_DATA {
        assert!(trie.get(&key).is_none());
    }
}

#[test]
fn remove_simple() {
    let mut trie = Trie::new();

    trie.insert("HELL", 66);
    trie.insert("HELLO", 77);
    let val = trie.remove(&"HELLO");
    assert_eq!(val, Some(77));
}

#[test]
fn nearest_ancestor_root() {
    let mut trie = Trie::new();
    trie.insert("", 55);
    assert_eq!(trie.get_ancestor_value(&""), Some(&55));
}

#[test]
fn nearest_ancestor() {
    let trie = test_trie();
    assert_eq!(trie.get_ancestor_value(&""), None);

    // Test identity prefixes.
    for &(key, val) in &TEST_DATA {
        assert_eq!(trie.get_ancestor_value(&key), Some(&val));
    }

    assert_eq!(trie.get_ancestor_value(&"abcdefg"), trie.get(&"abcdef"));
    assert_eq!(trie.get_ancestor_value(&"abcde"), trie.get(&"abcd"));
    assert_eq!(trie.get_ancestor_value(&"aauksdjk"), trie.get(&"a"));
}

#[test]
fn nearest_ancestor_no_child_fn() {
    let mut t = Trie::new();
    t.insert("ab", 5);
    let anc = t.get_ancestor(&"abc").unwrap();
    assert_eq!(*anc.value().unwrap(), 5);
}

#[test]
fn raw_ancestor() {
    let mut t = Trie::new();

    for &(key, _) in &TEST_DATA {
        assert_eq!(t.get_raw_ancestor(&key).key(), t.key());
    }

    t.insert("wow", 0);
    t.insert("hella", 1);
    t.insert("hellb", 2);

    // Ancestor should be "hell" node.
    let anc = t.get_raw_ancestor(&"hello");
    assert_eq!(anc.len(), 2);
}

// Check that the subtrie prefix is correct for raw_ancestor.
#[test]
fn raw_ancestor_prefix() {
    let mut t = Trie::new();

    t.insert("abac", ());
    t.insert("abaz", ());

    let anc = t.get_raw_ancestor(&"aba");
    assert_eq!(anc.prefix, "aba".encode());
}

#[test]
fn iter() {
    type Set = HashSet<(&'static str, u32)>;
    let trie = test_trie();
    let expected = TEST_DATA.iter().map(|&x| x).collect::<Set>();
    let observed = trie.iter().map(|(&k, &v)| (k, v)).collect::<Set>();
    assert_eq!(expected, observed);
}

#[test]
fn get_raw_descendant() {
    let trie = test_trie();
    assert_eq!(trie.get_raw_descendant(&"abcdefgh").and_then(|t| t.value()),
               Some(&19));
    assert_eq!(trie.get_raw_descendant(&"abcdefg").and_then(|t| t.value()),
               Some(&19));
    assert!(trie.get_raw_descendant(&"acbg").is_none());
}

#[test]
fn raw_descendant_prefix() {
    let mut t = Trie::new();

    t.insert("abczzzz", ());
    t.insert("abcaaaa", ());

    assert_eq!(t.get_raw_descendant(&"a").unwrap().prefix, "abc".encode());
    assert_eq!(t.get_raw_descendant(&"abc").unwrap().prefix, "abc".encode());
    assert_eq!(t.get_raw_descendant(&"abca").unwrap().prefix, "abcaaaa".encode());
}

#[test]
fn get_prefix_bug() {
    let mut trie = Trie::new();
    trie.insert("abdc", 5);
    trie.insert("abde", 6);
    assert!(trie.get(&"abc").is_none());
}

#[test]
fn get_ancestor_bug() {
    let mut trie = Trie::new();
    trie.insert("abc", 1);
    trie.insert("abcde", 2);
    assert_eq!(trie.get_ancestor_value(&"abcdz"), Some(&1));
}

#[test]
fn root_replace_bug() {
    let mut trie = Trie::new();
    trie.insert("a", ());
    trie.insert("p", ());
    trie.remove(&"a");
    assert_eq!(trie.len(), 1);
    trie.remove(&"p");
    assert_eq!(trie.len(), 0);
}

#[test]
fn subtrie_insert() {
    let mut trie = Trie::new();
    trie.insert("abc", 3);
    {
        let mut subtrie = trie.subtrie_mut(&"abc").unwrap();
        assert_eq!(subtrie.insert("somerandomshit", 666), Err(()));
        assert_eq!(subtrie.insert("abcdef", 6), Ok(None));
        assert_eq!(subtrie.insert("abc", 9), Ok(Some(3)));
    }
    assert_eq!(trie.get(&"abc"), Some(&9));
    assert_eq!(trie.get(&"abcdef"), Some(&6));
    assert_eq!(trie.len(), 2);
}

#[test]
fn subtrie_len() {
    let trie = test_trie();
    assert_eq!(trie.subtrie(&"ab").unwrap().len(), 4);
    assert_eq!(trie.subtrie(&"a").unwrap().len(), 6);
    assert_eq!(trie.subtrie(&"").unwrap().len(), trie.len());
    assert_eq!(trie.subtrie(&"bcdefgh").unwrap().len(), 1);
}

// Subtrie functions that return references should return references valid for
// the lifetime of the *original trie* from which they were borrowed, NOT
// just the lifetime of the subtrie (which may be shorter).
#[test]
fn subtrie_lifetime() {
    let trie = test_trie();
    let subtrie_value = {
        let subtrie = trie.subtrie(&"ab").unwrap();
        subtrie.value()
    };
    assert_eq!(*subtrie_value.unwrap(), 16);
}

#[test]
fn subtrie_mut_lifetime() {
    let mut trie = test_trie();
    let subtrie_value = {
        let mut subtrie = trie.subtrie_mut(&"ab").unwrap();
        *subtrie.value_mut().unwrap() = 999;
        subtrie.value()
    };
    assert_eq!(*subtrie_value.unwrap(), 999);
}

#[test]
fn ancestor_key() {
    let trie = test_trie();
    let subtrie = trie.get_ancestor(&"abcde").unwrap();
    assert_eq!(*subtrie.key().unwrap(), "abcd");
    assert_eq!(subtrie.prefix, "abcd".encode());
    assert_eq!(*subtrie.get(&"abcdef").unwrap().unwrap(), 18);
    assert_eq!(*subtrie.get(&"abcdefgh").unwrap().unwrap(), 19);
}

#[test]
fn child_subtrie_keys() {
    let trie = test_trie();
    let subtrie = trie.subtrie(&"abcd").unwrap();
    for subsubtrie in subtrie.children() {
        // This subtrie should be for "abcde".
        assert_eq!(*subsubtrie.get(&"abcdef").unwrap().unwrap(), 18);
        assert_eq!(*subsubtrie.get(&"abcdefgh").unwrap().unwrap(), 19);
    }
}

#[test]
fn int_keys() {
    let mut trie = Trie::new();
    trie.insert(0x00ffu64, "asdf");
    trie.insert(0xdeadbeefu64, "asdf");
    assert!(trie.check_integrity());
}

#[test]
fn from_iter() {
    let trie: Trie<&str, u32> = Trie::from_iter(vec![
        ("test", 10), ("hello", 12)
    ]);
    assert_eq!(*trie.get(&"test").unwrap(), 10);
    assert_eq!(*trie.get(&"hello").unwrap(), 12);
    assert_eq!(trie.len(), 2);
}

#[test]
fn test_get_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    assert_eq!(*trie.get("/boot").unwrap(), "dir");
}

#[test]
fn test_get_mut_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    assert_eq!(*trie.get_mut("/boot").unwrap(), "dir");
}

#[test]
fn test_remove_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    assert_eq!(trie.remove("/boot").unwrap(), "dir");
}

#[test]
fn test_subtrie_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.subtrie("/boot").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "dir");
}

#[test]
fn test_subtrie_mut_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.subtrie_mut("/boot").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "dir");
}

#[test]
fn test_get_ancestor_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "ancestor");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.get_ancestor("/boot/lo").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "ancestor");
}

#[test]
fn test_get_ancestor_value_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "ancestor");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let ancestor_val = trie.get_ancestor_value("/boot/lo").unwrap();
    assert_eq!(*ancestor_val, "ancestor");
}

#[test]
fn test_get_raw_ancestor_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "boot");
    trie.insert("/boot/lol".to_string(), "dir");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.get_raw_ancestor("/boot/lo");
    assert_eq!(*subtrie.value().unwrap(), "boot");
}

#[test]
fn test_get_raw_descendant_borrow() {
    let mut trie = Trie::new();
    trie.insert("/boot".to_string(), "dir");
    trie.insert("/boot/lol".to_string(), "lol");
    trie.insert("/bleh".to_string(), "ohi");
    let subtrie = trie.get_raw_descendant("/boot/lo").unwrap();
    assert_eq!(*subtrie.value().unwrap(), "lol");
}
<|endoftext|>


<commit_before>extern crate num_traits;

use std::f64;
use std::cmp::Ordering;
use std::fmt;
use std::str::FromStr;
use std::ops::{Add, Mul};
use self::num_traits::{One, Zero};

#[derive(PartialOrd, Debug, Default, Clone, Copy)]
pub struct LogProb {
    value: f64
}

impl LogProb {
    pub fn new(value: f64) -> Result<Self, String> {
        if 0.0 <= value && value <= 1.0 {
            Ok(LogProb { value: - value.log2() })
        } else {
            Err(format!("Value {} is not a probability (i.e. not in the interval [0,1]).", value))
        }
    }

    pub fn probability(&self) -> f64 {
        (- &self.value).exp2()
    }
}

/// Flipped instance to reflect that a probability of ``0.25`` is "more expensive" than ``0.5``.
impl Ord for LogProb {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.partial_cmp(&other) {
            Some(ordering) => ordering,
            None => {
                if self.value.is_nan() {
                    if other.value.is_nan() {
                        Ordering::Equal
                    } else {
                        Ordering::Greater
                    }
                } else {
                    Ordering::Less
                }
            }
        }
    }
}

impl PartialEq for LogProb {
    fn eq(&self, other: &Self) -> bool {
        if self.value.is_nan() {
            if other.value.is_nan() { true } else { false }
        } else if other.value.is_nan() {
            false
        } else {
            self.value == other.value
        }
    }
}

impl Eq for LogProb {}

impl Mul for LogProb {
    type Output = Self;

    fn mul(self, other: Self) -> Self {
        LogProb { value: self.value.add(other.value) }
    }
}

impl Add for LogProb {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        match LogProb::new(self.probability().add(other.probability())) {
            Ok(p) => p,
            Err(e) => panic!(e)
        }
    }
}

impl Zero for LogProb {
    fn zero() -> LogProb {
        LogProb { value: f64::INFINITY }
    }

    fn is_zero(&self) -> bool {
        self.value.is_infinite()
    }
}

impl One for LogProb {
    fn one() -> LogProb {
        LogProb { value: 0.0 }
    }
}

impl FromStr for LogProb {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.parse() {
            Ok(p) => LogProb::new(p),
            Err(e) => Err(e.to_string())
        }
    }
}

impl fmt::Display for LogProb {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.probability())
    }
}
<commit_msg>src/util/log_prob.rs:  improved impl Add and added impl {Sub,Div}<commit_after>extern crate num_traits;

use std::f64;
use std::cmp::Ordering;
use std::fmt;
use std::str::FromStr;
use std::ops::{Add, Sub, Mul, Div};
use self::num_traits::{One, Zero};

#[derive(PartialOrd, Debug, Default, Clone, Copy)]
pub struct LogProb {
    value: f64
}

impl LogProb {
    pub fn new(value: f64) -> Result<Self, String> {
        if 0.0 <= value && value <= 1.0 {
            Ok(LogProb { value: - value.ln() })
        } else {
            Err(format!("{} is not a probability (i.e. not in the interval [0,1]).", value))
        }
    }

    pub fn probability(&self) -> f64 {
        (- &self.value).exp()
    }
}

/// Flipped instance to reflect that a probability of ``0.25`` is "more expensive" than ``0.5``.
impl Ord for LogProb {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.partial_cmp(&other) {
            Some(ordering) => ordering,
            None => {
                if self.value.is_nan() {
                    if other.value.is_nan() {
                        Ordering::Equal
                    } else {
                        Ordering::Greater
                    }
                } else {
                    Ordering::Less
                }
            }
        }
    }
}

impl PartialEq for LogProb {
    fn eq(&self, other: &Self) -> bool {
        if self.value.is_nan() {
            if other.value.is_nan() { true } else { false }
        } else if other.value.is_nan() {
            false
        } else {
            self.value - other.value < f64::EPSILON
        }
    }
}

impl Eq for LogProb {}

impl Add for LogProb {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        let (a, b) = (self.value, other.value);

        let (x, y) = if a > b {
            (a, b)
        } else {
            (b, a)
        };

        match x - (x - y).exp().ln_1p() {
            s if s >= 0.0 => LogProb { value: s },
            _             => panic!(format!("exp(-{}) + exp(-{}) is not a probability, i.e. not in the interval [0,1].", x, y)),
        }
    }
}

impl Sub for LogProb {
    type Output = Self ;

    fn sub(self, other: Self) -> Self {
        match (self.value, other.value) {
            (x, y) if x <= y => LogProb { value: (x - (-(x - y).exp_m1()).ln()) },
            (x, y) if x >  y => panic!(format!("exp(-{}) - exp(-{}) is not a probability, i.e. not on the interval [0,1].", x, y)),
            _                => unreachable!(),
        }
    }
}

#[test]
fn test_add_sub() {
    match (LogProb::new(0.5), LogProb::new(0.25), LogProb::new(0.75)) {
        (Ok(x), Ok(y), Ok(z)) => {
            assert_eq!(x + y, z);
            assert_eq!(y + x, z);
            assert_eq!(z - x, y);
            assert_eq!(z - y, x);
        },
        _ => panic!(),
    }
}

impl Mul for LogProb {
    type Output = Self;

    fn mul(self, other: Self) -> Self {
        LogProb { value: self.value.add(other.value) }
    }
}

impl Div for LogProb {
    type Output = Self;

    fn div(self, other: Self) -> Self {
        match self.value.sub(other.value) {
            x if x >= 0.0
                => LogProb { value: x },
            x if x <  0.0
                => panic!(format!("{} is not a probability (i.e. not in the interval [0,1]).", x)),
            _   => unreachable!(),
        }
    }
}

impl Zero for LogProb {
    fn zero() -> LogProb {
        LogProb { value: f64::INFINITY }
    }

    fn is_zero(&self) -> bool {
        self.value.is_infinite()
    }
}

impl One for LogProb {
    fn one() -> LogProb {
        LogProb { value: 0.0 }
    }
}

impl FromStr for LogProb {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.parse() {
            Ok(p) => LogProb::new(p),
            Err(e) => Err(e.to_string())
        }
    }
}

impl fmt::Display for LogProb {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.probability())
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::logging::{LogEntry, LogSchema};
use anyhow::format_err;
use diem_crypto::HashValue;
use diem_logger::prelude::*;
use diem_state_view::StateViewId;
use diem_types::{
    ledger_info::LedgerInfoWithSignatures,
    transaction::{
        Transaction, TransactionOutput, TransactionStatus, TransactionToCommit, Version,
    },
};
use diem_vm::VMExecutor;
use executor_types::{BlockExecutor, Error, StateComputeResult};
use fail::fail_point;

use crate::{
    metrics::{
        DIEM_EXECUTOR_COMMIT_BLOCKS_SECONDS, DIEM_EXECUTOR_EXECUTE_BLOCK_SECONDS,
        DIEM_EXECUTOR_SAVE_TRANSACTIONS_SECONDS, DIEM_EXECUTOR_TRANSACTIONS_SAVED,
        DIEM_EXECUTOR_VM_EXECUTE_BLOCK_SECONDS,
    },
    types::ProcessedVMOutput,
    Executor,
};
use diem_types::protocol_spec::ProtocolSpec;

impl<PS, V> BlockExecutor for Executor<PS, V>
where
    PS: ProtocolSpec,
    V: VMExecutor,
{
    fn committed_block_id(&self) -> anyhow::Result<HashValue, Error> {
        Ok(Self::committed_block_id(self))
    }

    fn reset(&self) -> anyhow::Result<(), Error> {
        self.reset_cache()
    }

    fn execute_block(
        &self,
        block: (HashValue, Vec<Transaction>),
        parent_block_id: HashValue,
    ) -> anyhow::Result<StateComputeResult, Error> {
        let (block_id, mut transactions) = block;
        let read_lock = self.cache.read();

        // Reconfiguration rule - if a block is a child of pending reconfiguration, it needs to be empty
        // So we roll over the executed state until it's committed and we start new epoch.
        let (output, state_compute_result) = if parent_block_id != read_lock.committed_block_id()
            && read_lock
                .get_block(&parent_block_id)?
                .lock()
                .output()
                .has_reconfiguration()
        {
            let parent = read_lock.get_block(&parent_block_id)?;
            drop(read_lock);
            let parent_block = parent.lock();
            let parent_output = parent_block.output();

            info!(
                LogSchema::new(LogEntry::BlockExecutor).block_id(block_id),
                "reconfig_descendant_block_received"
            );

            let output = ProcessedVMOutput::new(
                vec![],
                parent_output.executed_trees().clone(),
                parent_output.epoch_state().clone(),
            );

            let parent_accu = parent_output.executed_trees().txn_accumulator();
            let state_compute_result = output.compute_result(
                parent_accu.frozen_subtree_roots().clone(),
                parent_accu.num_leaves(),
            );

            // Reset the reconfiguration suffix transactions to empty list.
            transactions = vec![];

            (output, state_compute_result)
        } else {
            info!(
                LogSchema::new(LogEntry::BlockExecutor).block_id(block_id),
                "execute_block"
            );

            let _timer = DIEM_EXECUTOR_EXECUTE_BLOCK_SECONDS.start_timer();

            let parent_block_executed_trees =
                Self::get_executed_trees_from_lock(&read_lock, parent_block_id)?;

            // Hold a ref to the current base smt, so that all in-mem state in between the
            // currently committed version and the end version of the parent block won't go away
            // during execution.
            let _base_smt = read_lock.committed_trees().state_tree().clone();

            let state_view = self.get_executed_state_view_from_lock(
                &read_lock,
                StateViewId::BlockExecution { block_id },
                &parent_block_executed_trees,
            );
            drop(read_lock);

            let vm_outputs = {
                let _timer = DIEM_EXECUTOR_VM_EXECUTE_BLOCK_SECONDS.start_timer();
                fail_point!("executor::vm_execute_block", |_| {
                    Err(Error::from(anyhow::anyhow!(
                        "Injected error in vm_execute_block"
                    )))
                });
                V::execute_block(transactions.clone(), &state_view).map_err(anyhow::Error::from)?
            };

            let status: Vec<_> = vm_outputs
                .iter()
                .map(TransactionOutput::status)
                .cloned()
                .collect();
            if !status.is_empty() {
                trace!("Execution status: {:?}", status);
            }

            let (account_to_state, account_to_proof) = state_view.into();
            let output = Self::process_vm_outputs(
                account_to_state,
                account_to_proof,
                &transactions,
                vm_outputs,
                &parent_block_executed_trees,
            )
            .map_err(|err| format_err!("Failed to execute block: {}", err))?;

            let parent_accu = parent_block_executed_trees.txn_accumulator();

            let state_compute_result = output.compute_result(
                parent_accu.frozen_subtree_roots().clone(),
                parent_accu.num_leaves(),
            );
            (output, state_compute_result)
        };

        // Add the output to the speculation_output_tree
        self.cache
            .write()
            .add_block(parent_block_id, (block_id, transactions, output))?;

        Ok(state_compute_result)
    }

    fn commit_blocks(
        &self,
        block_ids: Vec<HashValue>,
        ledger_info_with_sigs: LedgerInfoWithSignatures,
    ) -> anyhow::Result<(), Error> {
        let _timer = DIEM_EXECUTOR_COMMIT_BLOCKS_SECONDS.start_timer();
        let block_id_to_commit = ledger_info_with_sigs.ledger_info().consensus_block_id();
        let read_lock = self.cache.read();

        info!(
            LogSchema::new(LogEntry::BlockExecutor).block_id(block_id_to_commit),
            "commit_block"
        );

        let version = ledger_info_with_sigs.ledger_info().version();

        let num_txns_in_li = version
            .checked_add(1)
            .ok_or_else(|| format_err!("version + 1 overflows"))?;
        let num_persistent_txns = read_lock.synced_trees().txn_accumulator().num_leaves();

        if num_txns_in_li < num_persistent_txns {
            return Err(Error::InternalError {
                error: format!(
                    "Try to commit stale transactions with the last version as {}",
                    version
                ),
            });
        }

        if num_txns_in_li == num_persistent_txns {
            return Ok(());
        }

        // All transactions that need to go to storage. In the above example, this means all the
        // transactions in A, B and C whose status == TransactionStatus::Keep.
        // This must be done before calculate potential skipping of transactions in idempotent commit.
        let mut txns_to_keep = vec![];
        let arc_blocks = block_ids
            .iter()
            .map(|id| read_lock.get_block(id))
            .collect::<anyhow::Result<Vec<_>, Error>>()?;
        let blocks = arc_blocks.iter().map(|b| b.lock()).collect::<Vec<_>>();
        for (txn, txn_data) in blocks.iter().flat_map(|block| {
            itertools::zip_eq(block.transactions(), block.output().transaction_data())
        }) {
            if let TransactionStatus::Keep(recorded_status) = txn_data.status() {
                txns_to_keep.push(TransactionToCommit::new(
                    txn.clone(),
                    txn_data.account_blobs().clone(),
                    Some(txn_data.jf_node_hashes().clone()),
                    txn_data.write_set().clone(),
                    txn_data.events().to_vec(),
                    txn_data.gas_used(),
                    recorded_status.clone(),
                ));
            }
        }

        let last_block = blocks
            .last()
            .ok_or_else(|| format_err!("CommittableBlockBatch is empty"))?;

        // Check that the version in ledger info (computed by consensus) matches the version
        // computed by us.
        let num_txns_in_speculative_accumulator = last_block
            .output()
            .executed_trees()
            .txn_accumulator()
            .num_leaves();
        assert_eq!(
            num_txns_in_li, num_txns_in_speculative_accumulator as Version,
            "Number of transactions in ledger info ({}) does not match number of transactions \
             in accumulator ({}).",
            num_txns_in_li, num_txns_in_speculative_accumulator,
        );
        drop(blocks);
        drop(read_lock);

        let num_txns_to_keep = txns_to_keep.len() as u64;

        // Skip txns that are already committed to allow failures in state sync process.
        let first_version_to_keep = num_txns_in_li - num_txns_to_keep;
        assert!(
            first_version_to_keep <= num_persistent_txns,
            "first_version {} in the blocks to commit cannot exceed # of committed txns: {}.",
            first_version_to_keep,
            num_persistent_txns
        );

        let num_txns_to_skip = num_persistent_txns - first_version_to_keep;
        let first_version_to_commit = first_version_to_keep + num_txns_to_skip;

        if num_txns_to_skip != 0 {
            debug!(
                LogSchema::new(LogEntry::BlockExecutor)
                    .latest_synced_version(num_persistent_txns - 1)
                    .first_version_to_keep(first_version_to_keep)
                    .num_txns_to_keep(num_txns_to_keep)
                    .first_version_to_commit(first_version_to_commit),
                "skip_transactions_when_committing"
            );
        }

        // Skip duplicate txns that are already persistent.
        let txns_to_commit = &txns_to_keep[num_txns_to_skip as usize..];
        let num_txns_to_commit = txns_to_commit.len() as u64;
        {
            let _timer = DIEM_EXECUTOR_SAVE_TRANSACTIONS_SECONDS.start_timer();
            DIEM_EXECUTOR_TRANSACTIONS_SAVED.observe(num_txns_to_commit as f64);

            assert_eq!(first_version_to_commit, num_txns_in_li - num_txns_to_commit);
            fail_point!("executor::commit_blocks", |_| {
                Err(Error::from(anyhow::anyhow!(
                    "Injected error in commit_blocks"
                )))
            });

            self.db.writer.save_transactions(
                txns_to_commit,
                first_version_to_commit,
                Some(&ledger_info_with_sigs),
            )?;
        }

        self.cache
            .write()
            .prune(ledger_info_with_sigs.ledger_info())?;

        // Now that the blocks are persisted successfully, we can reply to consensus
        Ok(())
    }
}
<commit_msg>[executor] fix refs to blocks leaking to outside of the lock<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::logging::{LogEntry, LogSchema};
use anyhow::format_err;
use diem_crypto::HashValue;
use diem_logger::prelude::*;
use diem_state_view::StateViewId;
use diem_types::{
    ledger_info::LedgerInfoWithSignatures,
    transaction::{
        Transaction, TransactionOutput, TransactionStatus, TransactionToCommit, Version,
    },
};
use diem_vm::VMExecutor;
use executor_types::{BlockExecutor, Error, StateComputeResult};
use fail::fail_point;

use crate::{
    metrics::{
        DIEM_EXECUTOR_COMMIT_BLOCKS_SECONDS, DIEM_EXECUTOR_EXECUTE_BLOCK_SECONDS,
        DIEM_EXECUTOR_SAVE_TRANSACTIONS_SECONDS, DIEM_EXECUTOR_TRANSACTIONS_SAVED,
        DIEM_EXECUTOR_VM_EXECUTE_BLOCK_SECONDS,
    },
    types::ProcessedVMOutput,
    Executor,
};
use diem_types::{proof::definition::LeafCount, protocol_spec::ProtocolSpec};

impl<PS, V> BlockExecutor for Executor<PS, V>
where
    PS: ProtocolSpec,
    V: VMExecutor,
{
    fn committed_block_id(&self) -> anyhow::Result<HashValue, Error> {
        Ok(Self::committed_block_id(self))
    }

    fn reset(&self) -> anyhow::Result<(), Error> {
        self.reset_cache()
    }

    fn execute_block(
        &self,
        block: (HashValue, Vec<Transaction>),
        parent_block_id: HashValue,
    ) -> anyhow::Result<StateComputeResult, Error> {
        let (block_id, mut transactions) = block;
        let read_lock = self.cache.read();

        // Reconfiguration rule - if a block is a child of pending reconfiguration, it needs to be empty
        // So we roll over the executed state until it's committed and we start new epoch.
        let (output, state_compute_result) = if parent_block_id != read_lock.committed_block_id()
            && read_lock
                .get_block(&parent_block_id)?
                .lock()
                .output()
                .has_reconfiguration()
        {
            let parent = read_lock.get_block(&parent_block_id)?;
            drop(read_lock);
            let parent_block = parent.lock();
            let parent_output = parent_block.output();

            info!(
                LogSchema::new(LogEntry::BlockExecutor).block_id(block_id),
                "reconfig_descendant_block_received"
            );

            let output = ProcessedVMOutput::new(
                vec![],
                parent_output.executed_trees().clone(),
                parent_output.epoch_state().clone(),
            );

            let parent_accu = parent_output.executed_trees().txn_accumulator();
            let state_compute_result = output.compute_result(
                parent_accu.frozen_subtree_roots().clone(),
                parent_accu.num_leaves(),
            );

            // Reset the reconfiguration suffix transactions to empty list.
            transactions = vec![];

            (output, state_compute_result)
        } else {
            info!(
                LogSchema::new(LogEntry::BlockExecutor).block_id(block_id),
                "execute_block"
            );

            let _timer = DIEM_EXECUTOR_EXECUTE_BLOCK_SECONDS.start_timer();

            let parent_block_executed_trees =
                Self::get_executed_trees_from_lock(&read_lock, parent_block_id)?;

            // Hold a ref to the current base smt, so that all in-mem state in between the
            // currently committed version and the end version of the parent block won't go away
            // during execution.
            let _base_smt = read_lock.committed_trees().state_tree().clone();

            let state_view = self.get_executed_state_view_from_lock(
                &read_lock,
                StateViewId::BlockExecution { block_id },
                &parent_block_executed_trees,
            );
            drop(read_lock);

            let vm_outputs = {
                let _timer = DIEM_EXECUTOR_VM_EXECUTE_BLOCK_SECONDS.start_timer();
                fail_point!("executor::vm_execute_block", |_| {
                    Err(Error::from(anyhow::anyhow!(
                        "Injected error in vm_execute_block"
                    )))
                });
                V::execute_block(transactions.clone(), &state_view).map_err(anyhow::Error::from)?
            };

            let status: Vec<_> = vm_outputs
                .iter()
                .map(TransactionOutput::status)
                .cloned()
                .collect();
            if !status.is_empty() {
                trace!("Execution status: {:?}", status);
            }

            let (account_to_state, account_to_proof) = state_view.into();
            let output = Self::process_vm_outputs(
                account_to_state,
                account_to_proof,
                &transactions,
                vm_outputs,
                &parent_block_executed_trees,
            )
            .map_err(|err| format_err!("Failed to execute block: {}", err))?;

            let parent_accu = parent_block_executed_trees.txn_accumulator();

            let state_compute_result = output.compute_result(
                parent_accu.frozen_subtree_roots().clone(),
                parent_accu.num_leaves(),
            );
            (output, state_compute_result)
        };

        // Add the output to the speculation_output_tree
        self.cache
            .write()
            .add_block(parent_block_id, (block_id, transactions, output))?;

        Ok(state_compute_result)
    }

    fn commit_blocks(
        &self,
        block_ids: Vec<HashValue>,
        ledger_info_with_sigs: LedgerInfoWithSignatures,
    ) -> anyhow::Result<(), Error> {
        let _timer = DIEM_EXECUTOR_COMMIT_BLOCKS_SECONDS.start_timer();
        let block_id_to_commit = ledger_info_with_sigs.ledger_info().consensus_block_id();
        info!(
            LogSchema::new(LogEntry::BlockExecutor).block_id(block_id_to_commit),
            "commit_block"
        );
        let version = ledger_info_with_sigs.ledger_info().version();
        let num_txns_in_li = version
            .checked_add(1)
            .ok_or_else(|| format_err!("version + 1 overflows"))?;

        if let Some((num_persistent_txns, txns_to_keep)) =
            self.lock_and_get_txns_to_keep(block_ids, version, num_txns_in_li)?
        {
            self.save_to_db(
                &ledger_info_with_sigs,
                num_txns_in_li,
                num_persistent_txns,
                &txns_to_keep,
            )?;

            self.cache
                .write()
                .prune(ledger_info_with_sigs.ledger_info())?;
        }

        // Now that the blocks are persisted successfully, we can reply to consensus
        Ok(())
    }
}

impl<PS, V> Executor<PS, V> {
    fn lock_and_get_txns_to_keep(
        &self,
        block_ids: Vec<HashValue>,
        version: Version,
        num_txns_in_li: u64,
    ) -> anyhow::Result<Option<(LeafCount, Vec<TransactionToCommit>)>, Error> {
        let read_lock = self.cache.read();
        let num_persistent_txns = read_lock.synced_trees().txn_accumulator().num_leaves();

        if num_txns_in_li < num_persistent_txns {
            return Err(Error::InternalError {
                error: format!(
                    "Try to commit stale transactions with the last version as {}",
                    version
                ),
            });
        }

        if num_txns_in_li == num_persistent_txns {
            return Ok(None);
        }

        // All transactions that need to go to storage. In the above example, this means all the
        // transactions in A, B and C whose status == TransactionStatus::Keep.
        // This must be done before calculate potential skipping of transactions in idempotent commit.
        let mut txns_to_keep = vec![];
        let arc_blocks = block_ids
            .iter()
            .map(|id| read_lock.get_block(id))
            .collect::<anyhow::Result<Vec<_>, Error>>()?;
        let blocks = arc_blocks.iter().map(|b| b.lock()).collect::<Vec<_>>();
        for (txn, txn_data) in blocks.iter().flat_map(|block| {
            itertools::zip_eq(block.transactions(), block.output().transaction_data())
        }) {
            if let TransactionStatus::Keep(recorded_status) = txn_data.status() {
                txns_to_keep.push(TransactionToCommit::new(
                    txn.clone(),
                    txn_data.account_blobs().clone(),
                    Some(txn_data.jf_node_hashes().clone()),
                    txn_data.write_set().clone(),
                    txn_data.events().to_vec(),
                    txn_data.gas_used(),
                    recorded_status.clone(),
                ));
            }
        }

        let last_block = blocks
            .last()
            .ok_or_else(|| format_err!("CommittableBlockBatch is empty"))?;

        // Check that the version in ledger info (computed by consensus) matches the version
        // computed by us.
        let num_txns_in_speculative_accumulator = last_block
            .output()
            .executed_trees()
            .txn_accumulator()
            .num_leaves();
        assert_eq!(
            num_txns_in_li, num_txns_in_speculative_accumulator as Version,
            "Number of transactions in ledger info ({}) does not match number of transactions \
             in accumulator ({}).",
            num_txns_in_li, num_txns_in_speculative_accumulator,
        );
        Ok(Some((num_persistent_txns, txns_to_keep)))
    }

    fn save_to_db(
        &self,
        ledger_info_with_sigs: &LedgerInfoWithSignatures,
        num_txns_in_li: u64,
        num_persistent_txns: u64,
        txns_to_keep: &[TransactionToCommit],
    ) -> anyhow::Result<(), Error> {
        let num_txns_to_keep = txns_to_keep.len() as u64;

        // Skip txns that are already committed to allow failures in state sync process.
        let first_version_to_keep = num_txns_in_li - num_txns_to_keep;
        assert!(
            first_version_to_keep <= num_persistent_txns,
            "first_version {} in the blocks to commit cannot exceed # of committed txns: {}.",
            first_version_to_keep,
            num_persistent_txns
        );

        let num_txns_to_skip = num_persistent_txns - first_version_to_keep;
        let first_version_to_commit = first_version_to_keep + num_txns_to_skip;
        if num_txns_to_skip != 0 {
            debug!(
                LogSchema::new(LogEntry::BlockExecutor)
                    .latest_synced_version(num_persistent_txns - 1)
                    .first_version_to_keep(first_version_to_keep)
                    .num_txns_to_keep(num_txns_to_keep)
                    .first_version_to_commit(first_version_to_commit),
                "skip_transactions_when_committing"
            );
        }

        // Skip duplicate txns that are already persistent.
        let txns_to_commit = &txns_to_keep[num_txns_to_skip as usize..];
        let num_txns_to_commit = txns_to_commit.len() as u64;
        {
            let _timer = DIEM_EXECUTOR_SAVE_TRANSACTIONS_SECONDS.start_timer();
            DIEM_EXECUTOR_TRANSACTIONS_SAVED.observe(num_txns_to_commit as f64);

            assert_eq!(first_version_to_commit, num_txns_in_li - num_txns_to_commit);
            fail_point!("executor::commit_blocks", |_| {
                Err(Error::from(anyhow::anyhow!(
                    "Injected error in commit_blocks"
                )))
            });

            self.db.writer.save_transactions(
                txns_to_commit,
                first_version_to_commit,
                Some(ledger_info_with_sigs),
            )?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::convert::Into;

/// The result of computing the index of a tree node's sibling.
pub enum SiblingIndex {
    /// Numerical underflow in computing the index.
    Underflow,
    /// Numerical overflow in computing the index.
    Overflow,
    /// The computed index is out of range, with the second value giving the
    /// number of siblings.
    OutOfRange(usize, usize),
    /// A successfully computed index value.
    Valid(usize),
}

impl SiblingIndex {
    pub fn of(sibling_count: usize,
              here_index: usize,
              offset: isize) -> Self {
        let offset_abs = offset.abs();
        if offset_abs < 0 {
            // offset is Int::min_value().
            let mut new_index = match here_index.checked_sub(1) {
                Some(x) => x,
                None => return SiblingIndex::Underflow,
            };
            new_index = match new_index.checked_sub((offset_abs + 1isize).abs() as usize) {
                Some(x) => x,
                None => return SiblingIndex::Underflow,
            };
            SiblingIndex::Valid(new_index)
        } else if offset_abs == 0 {
            SiblingIndex::Valid(here_index)
        } else {
            let new_index = match here_index.checked_add(offset_abs as usize) {
                Some(x) => x,
                None => return SiblingIndex::Overflow,
            };
            if new_index >= sibling_count {
                return SiblingIndex::OutOfRange(new_index, sibling_count);
            }
            SiblingIndex::Valid(new_index)
        }
    }

    /// Safely computes the index of a tree node's sibling.
    ///
    /// For `sibling_count` siblings and the current node at `here_index`, the
    /// index of the node that is the given offset from `here_index` is computed
    /// using checked arithmetic.
    pub fn compute(sibling_count: usize,
                   here_index: usize,
                   offset: isize) -> Option<usize> {
        SiblingIndex::of(sibling_count, here_index, offset).into()
    }
}

impl Into<Option<usize>> for SiblingIndex {
    /// Unwraps the index to get its value, or panics with an error message if
    /// `self` is not `SiblingIndex::Valid`.
    fn into(self) -> Option<usize> {
        match self {
            SiblingIndex::Underflow => panic!["numerical underflow computing sibling offset"],
            SiblingIndex::Overflow => panic!["numerical overflow computing sibling offset"],
            SiblingIndex::OutOfRange(_, _) => None,
            SiblingIndex::Valid(new_index) => Some(new_index),
        }
    }
}

/// The result of computing the index of a child.
pub enum ChildIndex {
    /// The computed index is out of range, with the second value giving the
    /// number of children.
    OutOfRange(usize, usize),
    /// A successfully computed index value.
    Valid(usize),
}

impl ChildIndex {
    /// Validates that a tree node has a child at the given index.
    pub fn of(child_count: usize, index: usize) -> Self {
        if index >= child_count {
            ChildIndex::OutOfRange(index, child_count)
        } else {
            ChildIndex::Valid(index)
        }
    }

    pub fn compute(child_count: usize, index: usize) -> Option<usize> {
        ChildIndex::of(child_count, index).into()
    }
}

impl Into<Option<usize>> for ChildIndex {
    /// Unwraps the index to get its value, or panics with an error message if
    /// `self` is not `ChildIndex::Valid`.
    fn into(self) -> Option<usize> {
        match self {
            ChildIndex::OutOfRange(_, _) => None,
            ChildIndex::Valid(new_index) => Some(new_index),
        }
     }
}
<commit_msg>Documentation bump.<commit_after>use std::convert::Into;

/// The result of computing the index of a nonroot tree node's sibling.
pub enum SiblingIndex {
    /// Numerical underflow in computing the index.
    Underflow,
    /// Numerical overflow in computing the index.
    Overflow,
    /// The computed index is out of range, with the second value giving the
    /// number of siblings.
    OutOfRange(usize, usize),
    /// A successfully computed index value.
    Valid(usize),
}

impl SiblingIndex {
    pub fn of(sibling_count: usize,
              here_index: usize,
              offset: isize) -> Self {
        let offset_abs = offset.abs();
        if offset_abs < 0 {
            // offset is Int::min_value().
            let mut new_index = match here_index.checked_sub(1) {
                Some(x) => x,
                None => return SiblingIndex::Underflow,
            };
            new_index = match new_index.checked_sub((offset_abs + 1isize).abs() as usize) {
                Some(x) => x,
                None => return SiblingIndex::Underflow,
            };
            SiblingIndex::Valid(new_index)
        } else if offset_abs == 0 {
            SiblingIndex::Valid(here_index)
        } else {
            let new_index = match here_index.checked_add(offset_abs as usize) {
                Some(x) => x,
                None => return SiblingIndex::Overflow,
            };
            if new_index >= sibling_count {
                return SiblingIndex::OutOfRange(new_index, sibling_count);
            }
            SiblingIndex::Valid(new_index)
        }
    }

    /// Safely computes the index of a tree node's sibling.
    ///
    /// For `sibling_count` siblings and the current node at `here_index`, the
    /// index of the node that is the given offset from `here_index` is computed
    /// using checked arithmetic.
    pub fn compute(sibling_count: usize,
                   here_index: usize,
                   offset: isize) -> Option<usize> {
        SiblingIndex::of(sibling_count, here_index, offset).into()
    }
}

impl Into<Option<usize>> for SiblingIndex {
    /// Unwraps the index to get its value, or panics with an error message if
    /// `self` is not `SiblingIndex::Valid`.
    fn into(self) -> Option<usize> {
        match self {
            SiblingIndex::Underflow => panic!["numerical underflow computing sibling offset"],
            SiblingIndex::Overflow => panic!["numerical overflow computing sibling offset"],
            SiblingIndex::OutOfRange(_, _) => None,
            SiblingIndex::Valid(new_index) => Some(new_index),
        }
    }
}

/// The result of computing the index of a child.
pub enum ChildIndex {
    /// The computed index is out of range, with the second value giving the
    /// actual number of children.
    OutOfRange(usize, usize),
    /// A successfully computed index value.
    Valid(usize),
}

impl ChildIndex {
    /// Validates that a tree node has a child at the given index.
    pub fn of(child_count: usize, index: usize) -> Self {
        if index >= child_count {
            ChildIndex::OutOfRange(index, child_count)
        } else {
            ChildIndex::Valid(index)
        }
    }

    pub fn compute(child_count: usize, index: usize) -> Option<usize> {
        ChildIndex::of(child_count, index).into()
    }
}

impl Into<Option<usize>> for ChildIndex {
    /// Unwraps the index to get its value, or panics with an error message if
    /// `self` is not `ChildIndex::Valid`.
    fn into(self) -> Option<usize> {
        match self {
            ChildIndex::OutOfRange(_, _) => None,
            ChildIndex::Valid(new_index) => Some(new_index),
        }
     }
}
<|endoftext|>


<commit_before>use syntax::ast;
use std::borrow::ToOwned;
use std::ascii::AsciiExt;

pub fn simple_path(p: &ast::Path) -> Option<String> {
    match &p.segments[..] {
        [ref single] if single.parameters.is_empty()
            => Some(single.identifier.as_str().to_owned()),
        _ => None,
    }
}

pub fn pat_to_var(p: &ast::Pat) -> Option<String> {
    match p.node {
        ast::PatIdent(ast::BindByValue(_), id, None)
            => Some(id.node.as_str().to_owned()),
        _ => None,
    }
}

pub fn pat_to_ident(p: &ast::Pat) -> Option<ast::Ident> {
    match p.node {
        ast::PatIdent(_, id, None) => Some(id.node),
        _ => None
    }
}

pub fn convert_type(orig: &String) -> String {
    println!("input type: {}", orig);
    let name = match &orig[..] {
        "f32" => "float",
        "f64" => "double",
        "i32" => "int",
        "u32" => "uint",
        name => name,
    };
    // convert Uppercase+Rust types to lowercase+glsl
    match name.to_ascii_lowercase() {
        ref n if n.len() == 4 && (n.starts_with("pnt") || n.starts_with("rot")) => {
            if !(*n.as_bytes().last().unwrap() as char).is_numeric() {
                name.to_string()
            }
            else {
                let mut start = "vec".to_string();
                start.push(*n.as_bytes().last().unwrap() as char);
                start
            }
        },
        ref n if n.len() == 4 && (n.starts_with("vec") || n.starts_with("mat")) => {
            if !(*n.as_bytes().last().unwrap() as char).is_numeric() {
                name.to_string()
            }
            else {
                n.clone()
            }
        },
        _ => name.to_string(),
    }
}<commit_msg>Remove annoying print<commit_after>use syntax::ast;
use std::borrow::ToOwned;
use std::ascii::AsciiExt;

pub fn simple_path(p: &ast::Path) -> Option<String> {
    match &p.segments[..] {
        [ref single] if single.parameters.is_empty()
            => Some(single.identifier.as_str().to_owned()),
        _ => None,
    }
}

pub fn pat_to_var(p: &ast::Pat) -> Option<String> {
    match p.node {
        ast::PatIdent(ast::BindByValue(_), id, None)
            => Some(id.node.as_str().to_owned()),
        _ => None,
    }
}

pub fn pat_to_ident(p: &ast::Pat) -> Option<ast::Ident> {
    match p.node {
        ast::PatIdent(_, id, None) => Some(id.node),
        _ => None
    }
}

pub fn convert_type(orig: &String) -> String {
    let name = match &orig[..] {
        "f32" => "float",
        "f64" => "double",
        "i32" => "int",
        "u32" => "uint",
        name => name,
    };
    // convert Uppercase+Rust types to lowercase+glsl
    match name.to_ascii_lowercase() {
        ref n if n.len() == 4 && (n.starts_with("pnt") || n.starts_with("rot")) => {
            if !(*n.as_bytes().last().unwrap() as char).is_numeric() {
                name.to_string()
            }
            else {
                let mut start = "vec".to_string();
                start.push(*n.as_bytes().last().unwrap() as char);
                start
            }
        },
        ref n if n.len() == 4 && (n.starts_with("vec") || n.starts_with("mat")) => {
            if !(*n.as_bytes().last().unwrap() as char).is_numeric() {
                name.to_string()
            }
            else {
                n.clone()
            }
        },
        _ => name.to_string(),
    }
}<|endoftext|>


<commit_before>// Copyright (c) 2017 fd developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0>
// or the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use crate::exec;
use crate::exit_codes::{merge_exitcodes, ExitCode};
use crate::fshelper;
use crate::internal::{opts::FdOptions, osstr_to_bytes, MAX_BUFFER_LENGTH};
use crate::output;

use std::borrow::Cow;
use std::convert::TryFrom;
use std::error::Error;
use std::ffi::OsStr;
use std::fs;
use std::io;
use std::path::PathBuf;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

use ignore::overrides::OverrideBuilder;
use ignore::{self, WalkBuilder};
use regex::bytes::Regex;

/// The receiver thread can either be buffering results or directly streaming to the console.
enum ReceiverMode {
    /// Receiver is still buffering in order to sort the results, if the search finishes fast
    /// enough.
    Buffering,

    /// Receiver is directly printing results to the output.
    Streaming,
}

/// The Worker threads can result in a valid entry having PathBuf or an error.
pub enum WorkerResult {
    Entry(PathBuf),
    Error(ignore::Error),
}

/// Converts a WorkerResult to a PathBuf. This is required as in some cases,
/// the underlying `ignore` libary will return an error when we actually
/// still want to show the path.
///
/// Currently, the cases supported are:
///   - a broken symlink should still be found and the path returned
impl TryFrom<WorkerResult> for PathBuf {
    type Error = ignore::Error;

    fn try_from(worker_result: WorkerResult) -> Result<PathBuf, Self::Error> {
        match worker_result {
            WorkerResult::Entry(value) => return Ok(value),
            WorkerResult::Error(ignore_error) => {
                if let ignore::Error::WithPath { path, err } = &ignore_error {
                    if let ignore::Error::Io(ref io_err) = **err {
                        if io_err.kind() == io::ErrorKind::NotFound
                            && fs::symlink_metadata(&path)
                                .map_or(false, |m| m.file_type().is_symlink())
                        {
                            return Ok(path.to_path_buf());
                        }
                    }
                }
                return Err(ignore_error);
            }
        }
    }
}

/// Recursively scan the given search path for files / pathnames matching the pattern.
///
/// If the `--exec` argument was supplied, this will create a thread pool for executing
/// jobs in parallel from a given command line and the discovered paths. Otherwise, each
/// path will simply be written to standard output.
pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<FdOptions>) -> ExitCode {
    let mut path_iter = path_vec.iter();
    let first_path_buf = path_iter
        .next()
        .expect("Error: Path vector can not be empty");
    let (tx, rx) = channel();

    let mut override_builder = OverrideBuilder::new(first_path_buf.as_path());

    for pattern in &config.exclude_patterns {
        let res = override_builder.add(pattern);
        if res.is_err() {
            print_error_and_exit!("Malformed exclude pattern '{}'", pattern);
        }
    }
    let overrides = override_builder.build().unwrap_or_else(|_| {
        print_error_and_exit!("Mismatch in exclude patterns");
    });

    let mut walker = WalkBuilder::new(first_path_buf.as_path());
    walker
        .hidden(config.ignore_hidden)
        .ignore(config.read_fdignore)
        .parents(config.read_fdignore || config.read_vcsignore)
        .git_ignore(config.read_vcsignore)
        .git_global(config.read_vcsignore)
        .git_exclude(config.read_vcsignore)
        .overrides(overrides)
        .follow_links(config.follow_links)
        // No need to check for supported platforms, option is unavailable on unsupported ones
        .same_file_system(config.one_file_system)
        .max_depth(config.max_depth);

    if config.read_fdignore {
        walker.add_custom_ignore_filename(".fdignore");
    }

    for ignore_file in &config.ignore_files {
        let result = walker.add_ignore(ignore_file);
        match result {
            Some(ignore::Error::Partial(_)) => (),
            Some(err) => {
                print_error!(
                    "{}",
                    format!(
                        "Malformed pattern in custom ignore file '{}': {}.",
                        ignore_file.to_string_lossy(),
                        err.description()
                    )
                );
            }
            None => (),
        }
    }

    for path_entry in path_iter {
        walker.add(path_entry.as_path());
    }

    let parallel_walker = walker.threads(config.threads).build_parallel();

    let wants_to_quit = Arc::new(AtomicBool::new(false));
    if config.ls_colors.is_some() && config.command.is_none() {
        let wq = Arc::clone(&wants_to_quit);
        ctrlc::set_handler(move || {
            if wq.load(Ordering::Relaxed) {
                // Ctrl-C has been pressed twice, exit NOW
                process::exit(ExitCode::KilledBySigint.into());
            } else {
                wq.store(true, Ordering::Relaxed);
            }
        })
        .unwrap();
    }

    // Spawn the thread that receives all results through the channel.
    let receiver_thread = spawn_receiver(&config, &wants_to_quit, rx);

    // Spawn the sender threads.
    spawn_senders(&config, &wants_to_quit, pattern, parallel_walker, tx);

    // Wait for the receiver thread to print out all results.
    let exit_code = receiver_thread.join().unwrap();

    if wants_to_quit.load(Ordering::Relaxed) {
        process::exit(ExitCode::KilledBySigint.into());
    }

    exit_code
}

fn spawn_receiver(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    rx: Receiver<WorkerResult>,
) -> thread::JoinHandle<ExitCode> {
    let config = Arc::clone(config);
    let wants_to_quit = Arc::clone(wants_to_quit);

    let show_filesystem_errors = config.show_filesystem_errors;
    let threads = config.threads;

    thread::spawn(move || {
        // This will be set to `Some` if the `--exec` argument was supplied.
        if let Some(ref cmd) = config.command {
            if cmd.in_batch_mode() {
                exec::batch(rx, cmd, show_filesystem_errors)
            } else {
                let shared_rx = Arc::new(Mutex::new(rx));

                let out_perm = Arc::new(Mutex::new(()));

                // Each spawned job will store it's thread handle in here.
                let mut handles = Vec::with_capacity(threads);
                for _ in 0..threads {
                    let rx = Arc::clone(&shared_rx);
                    let cmd = Arc::clone(cmd);
                    let out_perm = Arc::clone(&out_perm);

                    // Spawn a job thread that will listen for and execute inputs.
                    let handle =
                        thread::spawn(move || exec::job(rx, cmd, out_perm, show_filesystem_errors));

                    // Push the handle of the spawned thread into the vector for later joining.
                    handles.push(handle);
                }

                // Wait for all threads to exit before exiting the program.
                let mut results: Vec<ExitCode> = Vec::new();
                for h in handles {
                    results.push(h.join().unwrap());
                }

                merge_exitcodes(results)
            }
        } else {
            let start = time::Instant::now();

            let mut buffer = vec![];

            // Start in buffering mode
            let mut mode = ReceiverMode::Buffering;

            // Maximum time to wait before we start streaming to the console.
            let max_buffer_time = config
                .max_buffer_time
                .unwrap_or_else(|| time::Duration::from_millis(100));

            let stdout = io::stdout();
            let mut stdout = stdout.lock();

            for worker_result in rx {
                match PathBuf::try_from(worker_result) {
                    Ok(value) => {
                        match mode {
                            ReceiverMode::Buffering => {
                                buffer.push(value);

                                // Have we reached the maximum buffer size or maximum buffering time?
                                if buffer.len() > MAX_BUFFER_LENGTH
                                    || time::Instant::now() - start > max_buffer_time
                                {
                                    // Flush the buffer
                                    for v in &buffer {
                                        output::print_entry(
                                            &mut stdout,
                                            v,
                                            &config,
                                            &wants_to_quit,
                                        );
                                    }
                                    buffer.clear();

                                    // Start streaming
                                    mode = ReceiverMode::Streaming;
                                }
                            }
                            ReceiverMode::Streaming => {
                                output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                            }
                        }
                    }
                    Err(err) => {
                        if show_filesystem_errors {
                            print_error!("{}", err);
                        }
                    }
                }
            }

            // If we have finished fast enough (faster than max_buffer_time), we haven't streamed
            // anything to the console, yet. In this case, sort the results and print them:
            if !buffer.is_empty() {
                buffer.sort();
                for value in buffer {
                    output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                }
            }

            ExitCode::Success
        }
    })
}

fn spawn_senders(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    pattern: Arc<Regex>,
    parallel_walker: ignore::WalkParallel,
    tx: Sender<WorkerResult>,
) {
    parallel_walker.run(|| {
        let config = Arc::clone(config);
        let pattern = Arc::clone(&pattern);
        let tx_thread = tx.clone();
        let wants_to_quit = Arc::clone(wants_to_quit);

        Box::new(move |entry_o| {
            if wants_to_quit.load(Ordering::Relaxed) {
                return ignore::WalkState::Quit;
            }

            let entry = match entry_o {
                Ok(e) => e,
                Err(err) => {
                    tx_thread.send(WorkerResult::Error(err)).unwrap();
                    return ignore::WalkState::Continue;
                }
            };

            if entry.depth() == 0 {
                return ignore::WalkState::Continue;
            }

            // Check the name first, since it doesn't require metadata
            let entry_path = entry.path();

            let search_str: Cow<OsStr> = if config.search_full_path {
                match fshelper::path_absolute_form(entry_path) {
                    Ok(path_abs_buf) => Cow::Owned(path_abs_buf.as_os_str().to_os_string()),
                    Err(_) => {
                        print_error_and_exit!("Unable to retrieve absolute path.");
                    }
                }
            } else {
                match entry_path.file_name() {
                    Some(filename) => Cow::Borrowed(filename),
                    None => unreachable!(
                        "Encountered file system entry without a file name. This should only \
                         happen for paths like 'foo/bar/..' or '/' which are not supposed to \
                         appear in a file system traversal."
                    ),
                }
            };

            if !pattern.is_match(&osstr_to_bytes(search_str.as_ref())) {
                return ignore::WalkState::Continue;
            }

            // Filter out unwanted extensions.
            if let Some(ref exts_regex) = config.extensions {
                if let Some(path_str) = entry_path.file_name() {
                    if !exts_regex.is_match(&osstr_to_bytes(path_str)) {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted file types.
            if let Some(ref file_types) = config.file_types {
                if let Some(ref entry_type) = entry.file_type() {
                    if (!file_types.files && entry_type.is_file())
                        || (!file_types.directories && entry_type.is_dir())
                        || (!file_types.symlinks && entry_type.is_symlink())
                        || (file_types.executables_only
                            && !entry
                                .metadata()
                                .map(|m| fshelper::is_executable(&m))
                                .unwrap_or(false))
                        || (file_types.empty_only && !fshelper::is_empty(&entry))
                        || !(entry_type.is_file() || entry_type.is_dir() || entry_type.is_symlink())
                    {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted sizes if it is a file and we have been given size constraints.
            if !config.size_constraints.is_empty() {
                if entry_path.is_file() {
                    if let Ok(metadata) = entry_path.metadata() {
                        let file_size = metadata.len();
                        if config
                            .size_constraints
                            .iter()
                            .any(|sc| !sc.is_within(file_size))
                        {
                            return ignore::WalkState::Continue;
                        }
                    } else {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted modification times
            if !config.time_constraints.is_empty() {
                let mut matched = false;
                if let Ok(metadata) = entry_path.metadata() {
                    if let Ok(modified) = metadata.modified() {
                        matched = config
                            .time_constraints
                            .iter()
                            .all(|tf| tf.applies_to(&modified));
                    }
                }
                if !matched {
                    return ignore::WalkState::Continue;
                }
            }

            let send_result = tx_thread.send(WorkerResult::Entry(entry_path.to_owned()));

            if !send_result.is_ok() {
                return ignore::WalkState::Quit;
            }

            ignore::WalkState::Continue
        })
    });
}
<commit_msg>Add comment for broken symlinks<commit_after>// Copyright (c) 2017 fd developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0>
// or the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use crate::exec;
use crate::exit_codes::{merge_exitcodes, ExitCode};
use crate::fshelper;
use crate::internal::{opts::FdOptions, osstr_to_bytes, MAX_BUFFER_LENGTH};
use crate::output;

use std::borrow::Cow;
use std::convert::TryFrom;
use std::error::Error;
use std::ffi::OsStr;
use std::fs;
use std::io;
use std::path::PathBuf;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

use ignore::overrides::OverrideBuilder;
use ignore::{self, WalkBuilder};
use regex::bytes::Regex;

/// The receiver thread can either be buffering results or directly streaming to the console.
enum ReceiverMode {
    /// Receiver is still buffering in order to sort the results, if the search finishes fast
    /// enough.
    Buffering,

    /// Receiver is directly printing results to the output.
    Streaming,
}

/// The Worker threads can result in a valid entry having PathBuf or an error.
pub enum WorkerResult {
    Entry(PathBuf),
    Error(ignore::Error),
}

/// Converts a WorkerResult to a PathBuf. This is required as in some cases,
/// the underlying `ignore` libary will return an error when we actually
/// still want to show the path.
///
/// Currently, the cases supported are:
///   - a broken symlink should still be found and the path returned
impl TryFrom<WorkerResult> for PathBuf {
    type Error = ignore::Error;

    fn try_from(worker_result: WorkerResult) -> Result<PathBuf, Self::Error> {
        match worker_result {
            WorkerResult::Entry(value) => return Ok(value),
            WorkerResult::Error(ignore_error) => {
                if let ignore::Error::WithPath { path, err } = &ignore_error {
                    if let ignore::Error::Io(ref io_err) = **err {
                        if io_err.kind() == io::ErrorKind::NotFound
                            && fs::symlink_metadata(&path)
                                .map_or(false, |m| m.file_type().is_symlink())
                        {
                            // Broken symlink
                            return Ok(path.to_path_buf());
                        }
                    }
                }
                return Err(ignore_error);
            }
        }
    }
}

/// Recursively scan the given search path for files / pathnames matching the pattern.
///
/// If the `--exec` argument was supplied, this will create a thread pool for executing
/// jobs in parallel from a given command line and the discovered paths. Otherwise, each
/// path will simply be written to standard output.
pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<FdOptions>) -> ExitCode {
    let mut path_iter = path_vec.iter();
    let first_path_buf = path_iter
        .next()
        .expect("Error: Path vector can not be empty");
    let (tx, rx) = channel();

    let mut override_builder = OverrideBuilder::new(first_path_buf.as_path());

    for pattern in &config.exclude_patterns {
        let res = override_builder.add(pattern);
        if res.is_err() {
            print_error_and_exit!("Malformed exclude pattern '{}'", pattern);
        }
    }
    let overrides = override_builder.build().unwrap_or_else(|_| {
        print_error_and_exit!("Mismatch in exclude patterns");
    });

    let mut walker = WalkBuilder::new(first_path_buf.as_path());
    walker
        .hidden(config.ignore_hidden)
        .ignore(config.read_fdignore)
        .parents(config.read_fdignore || config.read_vcsignore)
        .git_ignore(config.read_vcsignore)
        .git_global(config.read_vcsignore)
        .git_exclude(config.read_vcsignore)
        .overrides(overrides)
        .follow_links(config.follow_links)
        // No need to check for supported platforms, option is unavailable on unsupported ones
        .same_file_system(config.one_file_system)
        .max_depth(config.max_depth);

    if config.read_fdignore {
        walker.add_custom_ignore_filename(".fdignore");
    }

    for ignore_file in &config.ignore_files {
        let result = walker.add_ignore(ignore_file);
        match result {
            Some(ignore::Error::Partial(_)) => (),
            Some(err) => {
                print_error!(
                    "{}",
                    format!(
                        "Malformed pattern in custom ignore file '{}': {}.",
                        ignore_file.to_string_lossy(),
                        err.description()
                    )
                );
            }
            None => (),
        }
    }

    for path_entry in path_iter {
        walker.add(path_entry.as_path());
    }

    let parallel_walker = walker.threads(config.threads).build_parallel();

    let wants_to_quit = Arc::new(AtomicBool::new(false));
    if config.ls_colors.is_some() && config.command.is_none() {
        let wq = Arc::clone(&wants_to_quit);
        ctrlc::set_handler(move || {
            if wq.load(Ordering::Relaxed) {
                // Ctrl-C has been pressed twice, exit NOW
                process::exit(ExitCode::KilledBySigint.into());
            } else {
                wq.store(true, Ordering::Relaxed);
            }
        })
        .unwrap();
    }

    // Spawn the thread that receives all results through the channel.
    let receiver_thread = spawn_receiver(&config, &wants_to_quit, rx);

    // Spawn the sender threads.
    spawn_senders(&config, &wants_to_quit, pattern, parallel_walker, tx);

    // Wait for the receiver thread to print out all results.
    let exit_code = receiver_thread.join().unwrap();

    if wants_to_quit.load(Ordering::Relaxed) {
        process::exit(ExitCode::KilledBySigint.into());
    }

    exit_code
}

fn spawn_receiver(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    rx: Receiver<WorkerResult>,
) -> thread::JoinHandle<ExitCode> {
    let config = Arc::clone(config);
    let wants_to_quit = Arc::clone(wants_to_quit);

    let show_filesystem_errors = config.show_filesystem_errors;
    let threads = config.threads;

    thread::spawn(move || {
        // This will be set to `Some` if the `--exec` argument was supplied.
        if let Some(ref cmd) = config.command {
            if cmd.in_batch_mode() {
                exec::batch(rx, cmd, show_filesystem_errors)
            } else {
                let shared_rx = Arc::new(Mutex::new(rx));

                let out_perm = Arc::new(Mutex::new(()));

                // Each spawned job will store it's thread handle in here.
                let mut handles = Vec::with_capacity(threads);
                for _ in 0..threads {
                    let rx = Arc::clone(&shared_rx);
                    let cmd = Arc::clone(cmd);
                    let out_perm = Arc::clone(&out_perm);

                    // Spawn a job thread that will listen for and execute inputs.
                    let handle =
                        thread::spawn(move || exec::job(rx, cmd, out_perm, show_filesystem_errors));

                    // Push the handle of the spawned thread into the vector for later joining.
                    handles.push(handle);
                }

                // Wait for all threads to exit before exiting the program.
                let mut results: Vec<ExitCode> = Vec::new();
                for h in handles {
                    results.push(h.join().unwrap());
                }

                merge_exitcodes(results)
            }
        } else {
            let start = time::Instant::now();

            let mut buffer = vec![];

            // Start in buffering mode
            let mut mode = ReceiverMode::Buffering;

            // Maximum time to wait before we start streaming to the console.
            let max_buffer_time = config
                .max_buffer_time
                .unwrap_or_else(|| time::Duration::from_millis(100));

            let stdout = io::stdout();
            let mut stdout = stdout.lock();

            for worker_result in rx {
                match PathBuf::try_from(worker_result) {
                    Ok(value) => {
                        match mode {
                            ReceiverMode::Buffering => {
                                buffer.push(value);

                                // Have we reached the maximum buffer size or maximum buffering time?
                                if buffer.len() > MAX_BUFFER_LENGTH
                                    || time::Instant::now() - start > max_buffer_time
                                {
                                    // Flush the buffer
                                    for v in &buffer {
                                        output::print_entry(
                                            &mut stdout,
                                            v,
                                            &config,
                                            &wants_to_quit,
                                        );
                                    }
                                    buffer.clear();

                                    // Start streaming
                                    mode = ReceiverMode::Streaming;
                                }
                            }
                            ReceiverMode::Streaming => {
                                output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                            }
                        }
                    }
                    Err(err) => {
                        if show_filesystem_errors {
                            print_error!("{}", err);
                        }
                    }
                }
            }

            // If we have finished fast enough (faster than max_buffer_time), we haven't streamed
            // anything to the console, yet. In this case, sort the results and print them:
            if !buffer.is_empty() {
                buffer.sort();
                for value in buffer {
                    output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                }
            }

            ExitCode::Success
        }
    })
}

fn spawn_senders(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    pattern: Arc<Regex>,
    parallel_walker: ignore::WalkParallel,
    tx: Sender<WorkerResult>,
) {
    parallel_walker.run(|| {
        let config = Arc::clone(config);
        let pattern = Arc::clone(&pattern);
        let tx_thread = tx.clone();
        let wants_to_quit = Arc::clone(wants_to_quit);

        Box::new(move |entry_o| {
            if wants_to_quit.load(Ordering::Relaxed) {
                return ignore::WalkState::Quit;
            }

            let entry = match entry_o {
                Ok(e) => e,
                Err(err) => {
                    tx_thread.send(WorkerResult::Error(err)).unwrap();
                    return ignore::WalkState::Continue;
                }
            };

            if entry.depth() == 0 {
                return ignore::WalkState::Continue;
            }

            // Check the name first, since it doesn't require metadata
            let entry_path = entry.path();

            let search_str: Cow<OsStr> = if config.search_full_path {
                match fshelper::path_absolute_form(entry_path) {
                    Ok(path_abs_buf) => Cow::Owned(path_abs_buf.as_os_str().to_os_string()),
                    Err(_) => {
                        print_error_and_exit!("Unable to retrieve absolute path.");
                    }
                }
            } else {
                match entry_path.file_name() {
                    Some(filename) => Cow::Borrowed(filename),
                    None => unreachable!(
                        "Encountered file system entry without a file name. This should only \
                         happen for paths like 'foo/bar/..' or '/' which are not supposed to \
                         appear in a file system traversal."
                    ),
                }
            };

            if !pattern.is_match(&osstr_to_bytes(search_str.as_ref())) {
                return ignore::WalkState::Continue;
            }

            // Filter out unwanted extensions.
            if let Some(ref exts_regex) = config.extensions {
                if let Some(path_str) = entry_path.file_name() {
                    if !exts_regex.is_match(&osstr_to_bytes(path_str)) {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted file types.
            if let Some(ref file_types) = config.file_types {
                if let Some(ref entry_type) = entry.file_type() {
                    if (!file_types.files && entry_type.is_file())
                        || (!file_types.directories && entry_type.is_dir())
                        || (!file_types.symlinks && entry_type.is_symlink())
                        || (file_types.executables_only
                            && !entry
                                .metadata()
                                .map(|m| fshelper::is_executable(&m))
                                .unwrap_or(false))
                        || (file_types.empty_only && !fshelper::is_empty(&entry))
                        || !(entry_type.is_file() || entry_type.is_dir() || entry_type.is_symlink())
                    {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted sizes if it is a file and we have been given size constraints.
            if !config.size_constraints.is_empty() {
                if entry_path.is_file() {
                    if let Ok(metadata) = entry_path.metadata() {
                        let file_size = metadata.len();
                        if config
                            .size_constraints
                            .iter()
                            .any(|sc| !sc.is_within(file_size))
                        {
                            return ignore::WalkState::Continue;
                        }
                    } else {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted modification times
            if !config.time_constraints.is_empty() {
                let mut matched = false;
                if let Ok(metadata) = entry_path.metadata() {
                    if let Ok(modified) = metadata.modified() {
                        matched = config
                            .time_constraints
                            .iter()
                            .all(|tf| tf.applies_to(&modified));
                    }
                }
                if !matched {
                    return ignore::WalkState::Continue;
                }
            }

            let send_result = tx_thread.send(WorkerResult::Entry(entry_path.to_owned()));

            if !send_result.is_ok() {
                return ignore::WalkState::Quit;
            }

            ignore::WalkState::Continue
        })
    });
}
<|endoftext|>