<commit_before>#![allow(unstable)]

extern crate "recursive_sync" as rs;

use rs::RMutex;
use std::thread::Thread;
use std::sync::Arc;

#[test]
fn recursive_test() {
    let mutex = RMutex::new(0i32);
    {
        let mut outer_lock = mutex.lock();
        {
            let mut inner_lock = mutex.lock();
            *inner_lock = 1;
        }
        *outer_lock = 2;
    }
    assert_eq!(*mutex.lock(), 2);
}

#[test]
fn test_guarding() {
    let count = 10000;
    let mutex = Arc::new(RMutex::new(0i32));
    let mut guards = Vec::new();

    for _ in (0..count) {
        let mutex = mutex.clone();
        guards.push(Thread::scoped(move || {
            *mutex.lock() += 1;
        }));
    }

    drop(guards);

    assert_eq!(*mutex.lock(), count);
}
<commit_msg>Better test for guarding<commit_after>#![allow(unstable)]

extern crate "recursive_sync" as rs;

use rs::RMutex;
use std::thread::Thread;
use std::sync::Arc;
use std::io::timer::sleep;
use std::time::duration::Duration;

#[test]
fn recursive_test() {
    let mutex = RMutex::new(0i32);
    {
        let mut outer_lock = mutex.lock();
        {
            let mut inner_lock = mutex.lock();
            *inner_lock = 1;
        }
        *outer_lock = 2;
    }
    assert_eq!(*mutex.lock(), 2);
}

#[test]
fn test_guarding() {
    let count = 1000;
    let mutex = Arc::new(RMutex::new(0i32));
    let mut guards = Vec::new();

    for _ in (0..count) {
        let mutex = mutex.clone();
        guards.push(Thread::scoped(move || {
            let mut value_ref = mutex.lock();
            let value = *value_ref;

            sleep(Duration::milliseconds(1));

            *value_ref = value + 1;
        }));
    }

    drop(guards);

    assert_eq!(*mutex.lock(), count);
}
<|endoftext|>


<commit_before>extern crate xml_oxide;

use xml_oxide::{parser::OxideParser, sax::Event};

use std::char;

use std::time::{Duration, Instant};

// let mut my_sax_handler = MySaxHandler {
//     counter: 0,
//     end_counter: 0,
//     char_counter: 0,
// };
// {
//     let mut sax_parser = SaxParser::new();
//     sax_parser.set_content_handler(&mut my_sax_handler);
//     sax_parser.parse(&mut f);
// }
// println!("START EVENT COUNT:{}", my_sax_handler.counter);
// println!("END EVENT COUNT:{}", my_sax_handler.end_counter);
// println!("TOTAL CHARS:{}", my_sax_handler.char_counter);

struct MyCollectorSaxHandler {
    start_counter: usize,
    end_counter: usize,
    char_counter: usize,
    start_el_name_vec: Vec<String>,
    end_el_name_vec: Vec<String>,
    // characters should be collected because SAX parser can send them splitted for various reasons.
    characters_collected_vec: Vec<String>,
    characters_buf: String,
}

fn collect_with_parser<R: std::io::Read>(f: R) -> MyCollectorSaxHandler {
    let mut data = MyCollectorSaxHandler {
        start_counter: 0,
        end_counter: 0,
        char_counter: 0,
        start_el_name_vec: Vec::new(),
        end_el_name_vec: Vec::new(),
        characters_collected_vec: Vec::new(),
        characters_buf: String::new(),
    };

    let mut p = OxideParser::start(f);

    loop {
        let res = p.read_event();

        match res {
            Event::StartElement(el) => {
                data.start_counter = data.start_counter + 1;
                data.start_el_name_vec.push(el.name.to_owned());
                for attr in el.attributes.iter() {}
            }
            Event::EndElement(el) => {
                data.end_counter += 1;
                data.end_el_name_vec.push(el.name.to_owned());
            }
            Event::EndDocument => {
                break;
            }
            Event::Characters(chars) => {
                data.characters_buf.push_str(chars);
            }

            _ => {}
        }
    }

    data
}

#[test]
fn test_basic() {
    let mut s = String::from("<rootEl><value>5</value></rootEl>");

    let data = collect_with_parser(s.as_bytes());

    assert_eq!(data.start_counter, 2);
    assert_eq!(data.end_counter, 2);
    assert_eq!(data.start_el_name_vec.get(0).unwrap(), "rootEl");
    println!("{:?}", data.characters_collected_vec);
    assert_eq!(data.end_el_name_vec.get(0).unwrap(), "value");

    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "5");
    // assert_eq!(my_sax_handler.end, );
}

#[test]
fn test_66_EntityRef() {
    let c = char::from_u32(60).unwrap();
    // println!("{}", c);
    assert_eq!(c, '<'); //8898  ⋂    60 <

    let c2 = char::from_u32(u32::from_str_radix("0003C", 16).unwrap()).unwrap();
    // println!("{}", c2);
    assert_eq!(c2, '<'); //8898  ⋂    60 <

    let mut s = String::from("<rootEl><value>1<2<3<4</value></rootEl>");

    let data = collect_with_parser(s.as_bytes());

    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "1<2<3<4");
    // assert_eq!(my_sax_handler.end, );
}

#[test]
fn test_18_CDATA() {
    let mut s = String::from(
        "<rootEl>1<2<3<4<![CDATA[ \
         1<2<3<4]]><![CDATA[]]></rootEl>",
    );
    let data = collect_with_parser(s.as_bytes());

    println!("{}", data.characters_collected_vec.get(0).unwrap());
    assert_eq!(
        data.characters_collected_vec.get(0).unwrap(),
        "1<2<3<4 1<2<3<4"
    );
}

// comments should be ignored in content handler
#[test]
fn test_15_Comment() {
    let mut s = String::from("<rootEl>comments<!--are ignored--><!---->.</rootEl>");
    let data = collect_with_parser(s.as_bytes());

    println!("{}", data.characters_collected_vec.get(0).unwrap());
    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "comments.");
}

#[test]
#[should_panic]
fn test_15_Comment_not_well_formed() {
    let mut s = String::from(
        "<rootEl>comments<!-- are not well formed with 3 \
         hyphen at the end unless it is empty---><!---->.</rootEl>",
    );
    let data = collect_with_parser(s.as_bytes());

    println!("{}", data.characters_collected_vec.get(0).unwrap());
    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "comments.");
}
<commit_msg>cargo fix<commit_after>extern crate xml_oxide;

use xml_oxide::{parser::OxideParser, sax::Event};

use std::char;

// let mut my_sax_handler = MySaxHandler {
//     counter: 0,
//     end_counter: 0,
//     char_counter: 0,
// };
// {
//     let mut sax_parser = SaxParser::new();
//     sax_parser.set_content_handler(&mut my_sax_handler);
//     sax_parser.parse(&mut f);
// }
// println!("START EVENT COUNT:{}", my_sax_handler.counter);
// println!("END EVENT COUNT:{}", my_sax_handler.end_counter);
// println!("TOTAL CHARS:{}", my_sax_handler.char_counter);

struct MyCollectorSaxHandler {
    start_counter: usize,
    end_counter: usize,
    char_counter: usize,
    start_el_name_vec: Vec<String>,
    end_el_name_vec: Vec<String>,
    // characters should be collected because SAX parser can send them splitted for various reasons.
    characters_collected_vec: Vec<String>,
    characters_buf: String,
}

fn collect_with_parser<R: std::io::Read>(f: R) -> MyCollectorSaxHandler {
    let mut data = MyCollectorSaxHandler {
        start_counter: 0,
        end_counter: 0,
        char_counter: 0,
        start_el_name_vec: Vec::new(),
        end_el_name_vec: Vec::new(),
        characters_collected_vec: Vec::new(),
        characters_buf: String::new(),
    };

    let mut p = OxideParser::start(f);

    loop {
        let res = p.read_event();

        match res {
            Event::StartElement(el) => {
                data.start_counter = data.start_counter + 1;
                data.start_el_name_vec.push(el.name.to_owned());
                for _attr in el.attributes.iter() {}
            }
            Event::EndElement(el) => {
                data.end_counter += 1;
                data.end_el_name_vec.push(el.name.to_owned());
            }
            Event::EndDocument => {
                break;
            }
            Event::Characters(chars) => {
                data.characters_buf.push_str(chars);
            }

            _ => {}
        }
    }

    data
}

#[test]
fn test_basic() {
    let s = String::from("<rootEl><value>5</value></rootEl>");

    let data = collect_with_parser(s.as_bytes());

    assert_eq!(data.start_counter, 2);
    assert_eq!(data.end_counter, 2);
    assert_eq!(data.start_el_name_vec.get(0).unwrap(), "rootEl");
    println!("{:?}", data.characters_collected_vec);
    assert_eq!(data.end_el_name_vec.get(0).unwrap(), "value");

    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "5");
    // assert_eq!(my_sax_handler.end, );
}

#[test]
fn test_66_EntityRef() {
    let c = char::from_u32(60).unwrap();
    // println!("{}", c);
    assert_eq!(c, '<'); //8898  ⋂    60 <

    let c2 = char::from_u32(u32::from_str_radix("0003C", 16).unwrap()).unwrap();
    // println!("{}", c2);
    assert_eq!(c2, '<'); //8898  ⋂    60 <

    let s = String::from("<rootEl><value>1<2<3<4</value></rootEl>");

    let data = collect_with_parser(s.as_bytes());

    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "1<2<3<4");
    // assert_eq!(my_sax_handler.end, );
}

#[test]
fn test_18_CDATA() {
    let s = String::from(
        "<rootEl>1<2<3<4<![CDATA[ \
         1<2<3<4]]><![CDATA[]]></rootEl>",
    );
    let data = collect_with_parser(s.as_bytes());

    println!("{}", data.characters_collected_vec.get(0).unwrap());
    assert_eq!(
        data.characters_collected_vec.get(0).unwrap(),
        "1<2<3<4 1<2<3<4"
    );
}

// comments should be ignored in content handler
#[test]
fn test_15_Comment() {
    let s = String::from("<rootEl>comments<!--are ignored--><!---->.</rootEl>");
    let data = collect_with_parser(s.as_bytes());

    println!("{}", data.characters_collected_vec.get(0).unwrap());
    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "comments.");
}

#[test]
#[should_panic]
fn test_15_Comment_not_well_formed() {
    let s = String::from(
        "<rootEl>comments<!-- are not well formed with 3 \
         hyphen at the end unless it is empty---><!---->.</rootEl>",
    );
    let data = collect_with_parser(s.as_bytes());

    println!("{}", data.characters_collected_vec.get(0).unwrap());
    assert_eq!(data.characters_collected_vec.get(0).unwrap(), "comments.");
}
<|endoftext|>


<commit_before>
extern crate chainbuf;

#[cfg(test)]
mod test {
    use chainbuf::{CHB_MIN_SIZE, Chain};
    use std::rand::{task_rng, Rng};
    use std::str;

    #[test]
    fn test_append_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld";
        let ls = s.len();
        chain.append_bytes(s.as_bytes());
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_prepend_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.prepend_bytes(s);
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_from_foreign_moves_all_data() {
        let mut orig = Chain::new();
        orig.append_bytes("HelloWorld".as_bytes());
        let sz = orig.len();
        let new = Chain::from_foreign(orig);
        assert_eq!(new.len(), sz);
    }

    #[test]
    fn test_pullup_returns_none_on_empty_chain() {
        let mut chain = Chain::new();
        assert!(chain.pullup(1).is_none());
    }

    #[test]
    fn test_pullup_returns_what_has_been_appended() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let res = chain.pullup(ls);
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s);
    }

    #[test]
    fn test_pullup_does_not_change_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let olds = chain.len();
        chain.pullup(ls / 2);
        assert_eq!(chain.len(), olds);
    }

    #[test]
    fn test_pullup_works_on_large_sequences() {
        let mut chain = Chain::new();
        let total = 2048u;
        let mut t = total;
        let one_seq = 128u;
        let mut buf = Vec::new();
        while t > 0 {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            buf.extend(b.iter().map(|x| x.clone()));
            chain.append_bytes(b);
            t -= one_seq;
        }
        {
            let ret = chain.pullup(total);
            assert!(ret.is_some());
            assert_eq!(ret.unwrap(), buf.as_slice());
        }
        assert_eq!(chain.len(), total);
    }

    #[test]
    fn test_pullup_all_returns_all_data() {
        let mut chain = Chain::new();
        let s = "helloworld".as_bytes();
        chain.append_bytes(s);
        let res = chain.pullup_all();
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s);
    }

    #[test]
    fn test_concat_increase_dst_length() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain1.append_bytes("HelloWorld".as_bytes());
        let l1 = chain1.len();
        chain2.append_bytes("HelloWorld".as_bytes());
        let l2 = chain2.len();
        chain1.concat(chain2);
        assert_eq!(chain1.len(), l1+l2);
    }

    #[test]
    fn test_concat_appends_content() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let b = "HelloWorld".as_bytes();
        let bl = b.len();
        chain2.append_bytes(b);
        chain1.concat(chain2);
        let res = chain1.pullup(bl);

        assert!(res.is_some());
        assert_eq!(res.unwrap(), b);
    }

    #[test]
    fn test_reset_empties_chain() {
        let mut chain = Chain::new();
        chain.append_bytes("HelloWorld".as_bytes());
        chain.reset();
        assert!(chain.pullup(1).is_none());
        assert_eq!(chain.len(), 0);
    }

    #[test]
    fn test_append_copies_data() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let s = "HelloWorld";
        let b = s.as_bytes();
        let lb = b.len();
        let mut ss = String::from_str(s);
        ss.push_str(s);
        chain1.append_bytes(b);
        chain2.append_bytes(b);
        chain1.append(&mut chain2);
        {
            let res = chain1.pullup(2 * lb);
            assert!(res.is_some());
            assert_eq!(res.unwrap(), ss.as_bytes());
        }
        assert_eq!(chain1.len(), 2 * lb);
    }

    #[test]
    fn test_move_from_moves_data() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let s = "HelloWorld";
        let b = s.as_bytes();
        let lb = b.len();
        let hlb = lb / 2;
        chain1.append_bytes(b);
        chain2.append_bytes(b);
        chain1.move_from(&mut chain2, hlb);
        assert_eq!(chain1.len(), lb + hlb);
        assert_eq!(chain2.len(), hlb);
        {
            let mut ss = String::from_str(s);
            ss.push_str(s.slice(0, hlb));
            let r = b.slice_from(hlb);
            let r1 = chain1.pullup(lb + hlb);
            let r2 = chain2.pullup(hlb);
            assert!(r1.is_some());
            assert!(r2.is_some());
            assert_eq!(r1.unwrap(), ss.as_bytes());
            assert_eq!(r2.unwrap(), r);
        }
    }

    #[test]
    fn test_move_from_move_on_node_edge() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let s:String = task_rng().gen_ascii_chars().take(CHB_MIN_SIZE).collect();
        let sb = s.as_bytes();
        chain2.append_bytes(sb);
        chain2.append_bytes(sb);
        chain2.append_bytes(sb);
        chain2.append_bytes(sb);
        chain1.move_from(&mut chain2, CHB_MIN_SIZE * 2);
        assert_eq!(chain1.len(), chain2.len());
    }

    #[test]
    fn test_move_from_returns_number_of_bytes_moved() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain2.append_bytes("helloworld".as_bytes());
        let orig_size = chain2.len();
        let moved = chain1.move_from(&mut chain2, 3);
        let new_size = chain2.len();
        assert_eq!(moved, 3);
        assert_eq!(orig_size - moved, chain2.len());
        let moved_some_more = chain1.move_from(&mut chain2, orig_size);
        assert_eq!(moved_some_more, new_size);
        assert_eq!(chain2.len(), 0);
    }

    #[test]
    fn test_reserve_returns_buffer_of_requested_size() {
        let mut chain = Chain::new();
        let buf = chain.reserve(10);
        assert_eq!(buf.len(), 10);
    }

    #[test]
    fn test_reserve_returns_free_buffer() {
        let mut chain = Chain::new();
        chain.append_bytes("helloworld".as_bytes());
        let buf = chain.reserve(10);
        let pat = Vec::from_elem(10, 0u8);
        assert_eq!(buf.as_slice(), pat.as_slice());
    }

    #[test]
    fn test_reserve_and_written_modifies_chain() {
        let mut chain = Chain::new();
        let s = "helloworld".as_bytes();
        let sl = s.len();
        {
            let buf = chain.reserve(10);
            for (i, c) in s.iter().enumerate() {
                buf[i] = *c as u8;
            }
        }
        chain.written(sl);
        assert_eq!(chain.len(), sl);
        assert_eq!(chain.pullup(sl).unwrap(), s);
    }

    #[test]
    fn test_drain_changes_chain_length() {
        let mut chain = Chain::new();
        let s = "helloworld".as_bytes();
        let hsl = s.len() / 2;
        chain.append_bytes(s);
        let was_l = chain.len();
        chain.drain(hsl);
        let new_l = chain.len();
        assert!(new_l < was_l);
        assert_eq!(new_l, was_l - hsl);
    }

    #[test]
    fn test_to_utf8_str_returns_none_on_non_utf8() {
        let mut chain = Chain::new();
        let b = [0xf0_u8, 0xff_u8, 0xff_u8, 0x10_u8];
        chain.append_bytes(b);
        let res = chain.to_utf8_str();
        assert!(res.is_none());
    }

    #[test]
    fn test_to_utf8_returns_correct_string() {
        let mut chain = Chain::new();
        let s:String = task_rng().gen_ascii_chars().take(CHB_MIN_SIZE * 4).collect();
        chain.append_bytes(s.as_bytes());
        let res = chain.to_utf8_str();
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s.as_slice());
    }

    #[test]
    fn test_find_returns_none_on_empty_chain() {
        let chain = Chain::new();
        let res = chain.find("helloworld".as_bytes());
        assert!(res.is_none());
    }

    #[test]
    fn test_find_returns_zero_on_empty_needle() {
        let mut chain = Chain::new();
        chain.append_bytes("helloworld".as_bytes());
        let res = chain.find("".as_bytes());
        assert!(res.is_some());
        assert_eq!(res.unwrap(), 0);
    }

    #[test]
    fn test_find_returns_none_if_not_found() {
        let mut chain = Chain::new();
        let needle = [1u8, 2u8, 3u8];
        let one_seq = 128u;
        for _ in range(0u, 20) {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            chain.append_bytes(b);
        }
        let res = chain.find(needle);
        assert!(res.is_none());
    }

    #[test]
    fn test_find_returns_correct_offset() {
        let mut chain = Chain::new();
        let mut offs = 0;
        let mut needle = String::from_str("");
        for i in range(0u, 20) {
            let mut int_rng = task_rng();
            let s:String = task_rng().gen_ascii_chars().take(int_rng.gen_range(50, 100)).collect();
            let bytes = s.as_bytes();

            chain.append_bytes(bytes);
            if i < 11 {
                offs += bytes.len();
            }
            if i == 11 {
                let m = bytes.len();
                let b = int_rng.gen_range(5, m/2);
                let e = int_rng.gen_range(b, m);
                needle = String::from_str(str::from_utf8(bytes.slice(b, e)).unwrap());
                offs += b;
            }
        }
        let res = chain.find(needle.as_bytes());
        assert!(res.is_some());
        assert_eq!(res.unwrap(), offs);
    }

    #[test]
    fn test_chains_with_same_content_are_equal() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let total = 2048u;
        let mut t = total;
        let one_seq = 128u;
        while t > 0 {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            chain1.append_bytes(b);
            chain2.append_bytes(b);
            t -= one_seq;
        }
        assert!(chain1 == chain2);
        let res1 = chain1.pullup(total).unwrap();
        let res2 = chain2.pullup(total).unwrap();
        assert_eq!(res1, res2);
    }

    #[test]
    fn test_chains_with_different_content_are_not_equal() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain1.append_bytes("hello".as_bytes());
        chain2.append_bytes("world".as_bytes());
        assert!(chain1 != chain2);
    }
}
<commit_msg>[fix] remove undetermenism from some tests<commit_after>
extern crate chainbuf;

#[cfg(test)]
mod test {
    use chainbuf::{CHB_MIN_SIZE, Chain};
    use std::rand::{task_rng, Rng};

    #[test]
    fn test_append_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld";
        let ls = s.len();
        chain.append_bytes(s.as_bytes());
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_prepend_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.prepend_bytes(s);
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_from_foreign_moves_all_data() {
        let mut orig = Chain::new();
        orig.append_bytes("HelloWorld".as_bytes());
        let sz = orig.len();
        let new = Chain::from_foreign(orig);
        assert_eq!(new.len(), sz);
    }

    #[test]
    fn test_pullup_returns_none_on_empty_chain() {
        let mut chain = Chain::new();
        assert!(chain.pullup(1).is_none());
    }

    #[test]
    fn test_pullup_returns_what_has_been_appended() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let res = chain.pullup(ls);
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s);
    }

    #[test]
    fn test_pullup_does_not_change_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let olds = chain.len();
        chain.pullup(ls / 2);
        assert_eq!(chain.len(), olds);
    }

    #[test]
    fn test_pullup_works_on_large_sequences() {
        let mut chain = Chain::new();
        let total = 2048u;
        let mut t = total;
        let one_seq = 128u;
        let mut buf = Vec::new();
        while t > 0 {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            buf.extend(b.iter().map(|x| x.clone()));
            chain.append_bytes(b);
            t -= one_seq;
        }
        {
            let ret = chain.pullup(total);
            assert!(ret.is_some());
            assert_eq!(ret.unwrap(), buf.as_slice());
        }
        assert_eq!(chain.len(), total);
    }

    #[test]
    fn test_pullup_all_returns_all_data() {
        let mut chain = Chain::new();
        let s = "helloworld".as_bytes();
        chain.append_bytes(s);
        let res = chain.pullup_all();
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s);
    }

    #[test]
    fn test_concat_increase_dst_length() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain1.append_bytes("HelloWorld".as_bytes());
        let l1 = chain1.len();
        chain2.append_bytes("HelloWorld".as_bytes());
        let l2 = chain2.len();
        chain1.concat(chain2);
        assert_eq!(chain1.len(), l1+l2);
    }

    #[test]
    fn test_concat_appends_content() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let b = "HelloWorld".as_bytes();
        let bl = b.len();
        chain2.append_bytes(b);
        chain1.concat(chain2);
        let res = chain1.pullup(bl);

        assert!(res.is_some());
        assert_eq!(res.unwrap(), b);
    }

    #[test]
    fn test_reset_empties_chain() {
        let mut chain = Chain::new();
        chain.append_bytes("HelloWorld".as_bytes());
        chain.reset();
        assert!(chain.pullup(1).is_none());
        assert_eq!(chain.len(), 0);
    }

    #[test]
    fn test_append_copies_data() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let s = "HelloWorld";
        let b = s.as_bytes();
        let lb = b.len();
        let mut ss = String::from_str(s);
        ss.push_str(s);
        chain1.append_bytes(b);
        chain2.append_bytes(b);
        chain1.append(&mut chain2);
        {
            let res = chain1.pullup(2 * lb);
            assert!(res.is_some());
            assert_eq!(res.unwrap(), ss.as_bytes());
        }
        assert_eq!(chain1.len(), 2 * lb);
    }

    #[test]
    fn test_move_from_moves_data() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let s = "HelloWorld";
        let b = s.as_bytes();
        let lb = b.len();
        let hlb = lb / 2;
        chain1.append_bytes(b);
        chain2.append_bytes(b);
        chain1.move_from(&mut chain2, hlb);
        assert_eq!(chain1.len(), lb + hlb);
        assert_eq!(chain2.len(), hlb);
        {
            let mut ss = String::from_str(s);
            ss.push_str(s.slice(0, hlb));
            let r = b.slice_from(hlb);
            let r1 = chain1.pullup(lb + hlb);
            let r2 = chain2.pullup(hlb);
            assert!(r1.is_some());
            assert!(r2.is_some());
            assert_eq!(r1.unwrap(), ss.as_bytes());
            assert_eq!(r2.unwrap(), r);
        }
    }

    #[test]
    fn test_move_from_move_on_node_edge() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let s:String = task_rng().gen_ascii_chars().take(CHB_MIN_SIZE).collect();
        let sb = s.as_bytes();
        chain2.append_bytes(sb);
        chain2.append_bytes(sb);
        chain2.append_bytes(sb);
        chain2.append_bytes(sb);
        chain1.move_from(&mut chain2, CHB_MIN_SIZE * 2);
        assert_eq!(chain1.len(), chain2.len());
    }

    #[test]
    fn test_move_from_returns_number_of_bytes_moved() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain2.append_bytes("helloworld".as_bytes());
        let orig_size = chain2.len();
        let moved = chain1.move_from(&mut chain2, 3);
        let new_size = chain2.len();
        assert_eq!(moved, 3);
        assert_eq!(orig_size - moved, chain2.len());
        let moved_some_more = chain1.move_from(&mut chain2, orig_size);
        assert_eq!(moved_some_more, new_size);
        assert_eq!(chain2.len(), 0);
    }

    #[test]
    fn test_reserve_returns_buffer_of_requested_size() {
        let mut chain = Chain::new();
        let buf = chain.reserve(10);
        assert_eq!(buf.len(), 10);
    }

    #[test]
    fn test_reserve_returns_free_buffer() {
        let mut chain = Chain::new();
        chain.append_bytes("helloworld".as_bytes());
        let buf = chain.reserve(10);
        let pat = Vec::from_elem(10, 0u8);
        assert_eq!(buf.as_slice(), pat.as_slice());
    }

    #[test]
    fn test_reserve_and_written_modifies_chain() {
        let mut chain = Chain::new();
        let s = "helloworld".as_bytes();
        let sl = s.len();
        {
            let buf = chain.reserve(10);
            for (i, c) in s.iter().enumerate() {
                buf[i] = *c as u8;
            }
        }
        chain.written(sl);
        assert_eq!(chain.len(), sl);
        assert_eq!(chain.pullup(sl).unwrap(), s);
    }

    #[test]
    fn test_drain_changes_chain_length() {
        let mut chain = Chain::new();
        let s = "helloworld".as_bytes();
        let hsl = s.len() / 2;
        chain.append_bytes(s);
        let was_l = chain.len();
        chain.drain(hsl);
        let new_l = chain.len();
        assert!(new_l < was_l);
        assert_eq!(new_l, was_l - hsl);
    }

    #[test]
    fn test_to_utf8_str_returns_none_on_non_utf8() {
        let mut chain = Chain::new();
        let b = [0xf0_u8, 0xff_u8, 0xff_u8, 0x10_u8];
        chain.append_bytes(b);
        let res = chain.to_utf8_str();
        assert!(res.is_none());
    }

    #[test]
    fn test_to_utf8_returns_correct_string() {
        let mut chain = Chain::new();
        let s:String = task_rng().gen_ascii_chars().take(CHB_MIN_SIZE * 4).collect();
        chain.append_bytes(s.as_bytes());
        let res = chain.to_utf8_str();
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s.as_slice());
    }

    #[test]
    fn test_find_returns_none_on_empty_chain() {
        let chain = Chain::new();
        let res = chain.find("helloworld".as_bytes());
        assert!(res.is_none());
    }

    #[test]
    fn test_find_returns_zero_on_empty_needle() {
        let mut chain = Chain::new();
        chain.append_bytes("helloworld".as_bytes());
        let res = chain.find("".as_bytes());
        assert!(res.is_some());
        assert_eq!(res.unwrap(), 0);
    }

    #[test]
    fn test_find_returns_none_if_not_found() {
        let mut chain = Chain::new();
        let needle = [1u8, 2u8, 3u8];
        let one_seq = 128u;
        for _ in range(0u, 20) {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            chain.append_bytes(b);
        }
        let res = chain.find(needle);
        assert!(res.is_none());
    }

    #[test]
    fn test_find_returns_correct_offset() {
        let mut chain = Chain::new();
        let mut offs = 0;
        let needle = "the quick brown fox jumps over the lazy dog";
        for i in range(0u, 20) {
            let mut int_rng = task_rng();
            let s:String = task_rng().gen_ascii_chars().take(int_rng.gen_range(50, 100)).collect();
            let bytes = s.as_bytes();

            chain.append_bytes(bytes);
            if i <= 11 {
                offs += bytes.len();
            }
            if i == 11 {
                chain.append_bytes(needle.as_bytes());
            }
        }
        let res = chain.find(needle.as_bytes());
        assert!(res.is_some());
        assert_eq!(res.unwrap(), offs);
    }

    #[test]
    fn test_chains_with_same_content_are_equal() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let total = 2048u;
        let mut t = total;
        let one_seq = 128u;
        while t > 0 {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            chain1.append_bytes(b);
            chain2.append_bytes(b);
            t -= one_seq;
        }
        assert!(chain1 == chain2);
        let res1 = chain1.pullup(total).unwrap();
        let res2 = chain2.pullup(total).unwrap();
        assert_eq!(res1, res2);
    }

    #[test]
    fn test_chains_with_different_content_are_not_equal() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain1.append_bytes("hello".as_bytes());
        chain2.append_bytes("world".as_bytes());
        assert!(chain1 != chain2);
    }
}
<|endoftext|>


<commit_before>//! Upper confidence bound (UCB1) algorithm for graph search.

use crate::backprop::BackpropSelector;
use crate::game::{Game, State, Statistics};
use crate::graph::{EdgeData, VertexData};
use crate::rollout::RolloutSelector;
use log::{error, trace};
use rand::Rng;
use search_graph;

use std::cmp::Ordering;
use std::error::Error;
use std::f64;
use std::fmt;
use std::result::Result;

/// Represents success when computing the UCB score for a child.
pub enum UcbSuccess<'id> {
  /// No (finite) value can be computed, but the UCB policy indicates that
  /// this child should be selected. E.g., the child has not yet been visited.
  Select(search_graph::view::EdgeRef<'id>),
  /// A value is computed.
  Value(search_graph::view::EdgeRef<'id>, f64),
}

/// Represents error conditions when computing the UCB score for a child.
#[derive(Clone, Debug)]
pub enum UcbError {
  /// An error was encountered during computation of UCB score (such as
  /// encountering a `None` result when `PartialCmp`'ing two floating-point
  /// values).
  InvalidComputation,
}

unsafe impl Send for UcbError {}

/// Lazy iterator over UCB scores for a series of edges.
pub struct EdgeUcbIter<'a, 'b, 'id, G, I>
where
  G: Game,
  I: Iterator<Item = search_graph::view::EdgeRef<'id>>,
{
  log_parent_visits: f64,
  explore_bias: f64,
  graph: &'b search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  edges: I,
}

impl fmt::Display for UcbError {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    match *self {
      UcbError::InvalidComputation => write!(f, "Numerical error when computing UCB score"),
    }
  }
}

impl Error for UcbError {
  fn description(&self) -> &str {
    match *self {
      UcbError::InvalidComputation => "invalid computation",
    }
  }
}

impl<'a, 'b, 'id, G, I> EdgeUcbIter<'a, 'b, 'id, G, I>
where
  G: Game,
  I: Iterator<Item = search_graph::view::EdgeRef<'id>>,
{
  /// Constructs a new `EdgeUcbIter` that will compute UCB scores using the
  /// given constants. All floating-point values are assumed to be valid
  /// floating-point values and positive.
  ///
  ///  - `log_parent_visits`: ln(visits to parent vertex).
  ///  - `explore_bias`: scalar bias controlling tradeoff between search width
  ///    and search depth (lower = wider, higher = deeper).
  ///  - `role`: the game role whose score to maximize.
  ///  - `edges`: an iterator over edge handles for which to compute UCB
  ///    scores. This should usually be a list of child edges which share a
  ///    parent vertex.
  pub fn new(
    log_parent_visits: f64,
    explore_bias: f64,
    graph: &'b search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
    edges: I,
  ) -> Self {
    EdgeUcbIter {
      log_parent_visits: log_parent_visits,
      explore_bias: explore_bias,
      graph,
      edges: edges,
    }
  }
}

impl<'a, 'b, 'id, G, I> Iterator for EdgeUcbIter<'a, 'b, 'id, G, I>
where
  G: Game,
  I: Iterator<Item = search_graph::view::EdgeRef<'id>>,
{
  type Item = Result<UcbSuccess<'id>, UcbError>;

  fn next(&mut self) -> Option<Result<UcbSuccess<'id>, UcbError>> {
    self.edges.next().map(|e| {
      if self.graph.edge_data(e).statistics.visits() == 0 {
        trace!("EdgeUcbIter selects unvisited action");
        Ok(UcbSuccess::Select(e))
      } else {
        trace!("EdgeUcbIter calls child_score with log_parent_visits = {}, explore_bias = {}, edge = {:?}",
               self.log_parent_visits,
               self.explore_bias,
               e);
        Ok(child_score(
          self.log_parent_visits,
          self.explore_bias,
          self.graph,
          e,
        ))
      }
    })
  }

  fn size_hint(&self) -> (usize, Option<usize>) {
    self.edges.size_hint()
  }
}

/// Returns the UCB policy result for the given values.
pub fn child_score<'a, 'id, G: Game>(
  log_parent_visits: f64,
  explore_bias: f64,
  graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  child: search_graph::view::EdgeRef<'id>,
) -> UcbSuccess<'id> {
  let statistics = &graph.edge_data(child).statistics;
  if statistics.visits() == 0 {
    UcbSuccess::Select(child)
  } else {
    let child_visits = statistics.visits() as f64;
    let child_score =
      statistics.score(graph.node_state(graph.edge_source(child)).active_player()) as f64;
    let ucb =
      child_score / child_visits + explore_bias * f64::sqrt(log_parent_visits / child_visits);
    UcbSuccess::Value(child, ucb)
  }
}

/// Returns `true` iff `e` could be selected by the UCB policy during rollout
/// from its parent vertex. Assumes we haven't yet altered the parent vertex
/// statistics. (Callers should ensure that this is not called more than once
/// for a given parent vertex.)
///
/// "Could be" reflects how this is different from calling `find_best_child`: it
/// is possible (and common) for multiple child edges of a vertex to be best
/// children (as when there are multiple children that have not yet been
/// explored). This is dealt with in `find_best_child` by randomly choosing one
/// of the best children. But when doing backpropagation on a full game state
/// graph (not just a tree), we want to know all of the parent edges which could
/// have rolled out to a given child.
pub fn is_best_child<'a, 'id, G: Game>(
  graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  e: search_graph::view::EdgeRef<'id>,
  explore_bias: f64,
) -> bool {
  let statistics = &graph.edge_data(e).statistics;
  // trace!("is_best_child: edge {} has {} visits", e.get_id(), stats.visits);
  if statistics.visits() == 0 {
    // Edge has been visited, but statistics aren't yet updated.
    // trace!("is_best_child: edge {} is a best child because stats.visits == 0", e.get_id());
    return true;
  }
  let parent = graph.edge_source(e);
  // trace!("is_best_child: edge {} (from node {}) has {} siblings", e.get_id(), parent.get_id(), siblings.len());
  let log_parent_visits = {
    let mut parent_visits = 0;
    for child_edge in graph.children(parent) {
      parent_visits += graph.edge_data(child_edge).statistics.visits();
    }
    f64::ln(parent_visits as f64)
  };
  let mut edge_ucb = None;
  let mut best_ucb = ::std::f64::MIN;
  let ucb_iter = EdgeUcbIter::new(
    log_parent_visits,
    explore_bias,
    graph,
    graph.children(parent),
  );
  // Scan through siblings to find the maximum UCB score. This is
  // short-circuited using a lazy iterator to ameliorate the O(n) running
  // time.
  for ucb in ucb_iter {
    match ucb {
      Ok(UcbSuccess::Select(sibling)) if e == sibling => {
        // trace!("is_best_child: edge {} is best child by fiat of score computation", e.get_id());
        return true;
      }
      Ok(UcbSuccess::Select(_)) => {
        if let Some(_) = edge_ucb {
          // trace!("is_best_child: edge {} has a ucb of {}, but sibling edge {} is best child by fiat of score computation", e.get_id(), u, sibling.get_id());
          return false;
        }
      }
      Ok(UcbSuccess::Value(sibling, score)) => {
        if sibling == e {
          if best_ucb > score {
            // trace!("is_best_child: found ucb {} for edge {}, but a sibling has a higher ucb of {}", score, e.get_id(), best_ucb);
            return false;
          }
          edge_ucb = Some(score)
        } else if let Some(u) = edge_ucb {
          if score > u {
            // trace!("is_best_child: found ucb {} for edge {}, but a sibling has a higher ucb of {}", score, e.get_id(), best_ucb);
            return false;
          }
        }
        if score > best_ucb {
          best_ucb = score;
        }
      }
      Err(e) => panic!("error {:?} computing ucb for best child search", e),
    }
  }
  match edge_ucb {
    Some(u) if u >= best_ucb => {
      // Target edge has a UCB score which matches the maximum we found.
      // trace!("is_best_child: edge {} has a ucb of {}, which matches the max sibling ucb of {}", e.get_id(), u, best_ucb);
      true
    }
    _ => {
      // ThudEdges are not best children by default.
      // trace!("is_best_child: edge {} has a ucb of {:?}, which does not match the max sibling ucb of {}", e.get_id(), edge_ucb, best_ucb);
      false
    }
  }
}

/// Returns the child edge of `parent` that is best according to the UCB1
/// criterion.
///
/// This function will panic if `parent` has no children.
pub fn find_best_child<'a, 'id, G, R>(
  graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  parent: search_graph::view::NodeRef<'id>,
  explore_bias: f64,
  rng: &mut R,
) -> Result<search_graph::view::EdgeRef<'id>, UcbError>
where
  G: Game,
  R: Rng,
{
  let log_parent_visits = {
    let mut parent_visits = 0;
    for child in graph.children(parent) {
      parent_visits += graph.edge_data(child).statistics.visits();
    }
    if parent_visits == 0 {
      // When we visit a vertex for the first time, it will have zero visits.
      0.0
    } else {
      // Otherwise, it should be positive.
      f64::ln(parent_visits as f64)
    }
  };
  let mut sampling_count = 0u32;
  let mut ucb_iter = EdgeUcbIter::new(
    log_parent_visits,
    explore_bias,
    graph,
    graph.children(parent),
  );
  let mut best: search_graph::view::EdgeRef<'id>;
  let mut best_ucb: f64;
  match ucb_iter.next().expect("vertex has no children")? {
    UcbSuccess::Select(e) => {
      best = e;
      best_ucb = f64::MAX;
    }
    UcbSuccess::Value(e, v) => {
      best = e;
      best_ucb = v;
    }
  }

  for ucb in ucb_iter {
    let (edge, value) = match ucb? {
      UcbSuccess::Select(e) => (e, f64::MAX),
      UcbSuccess::Value(e, v) => (e, v),
    };
    match value.partial_cmp(&best_ucb) {
      None => {
        error!("find_best_child: invalid floating-point comparison");
        return Err(UcbError::InvalidComputation);
      }
      Some(Ordering::Greater) => {
        trace!("find_best_child_edge: new best action with score {}", value);
        best = edge;
        best_ucb = value;
        sampling_count = 1;
      }
      Some(Ordering::Equal) => {
        // We use reservoir sampling to break ties.
        trace!("find_best_child: found action with identical score; sampling to break tie");
        sampling_count += 1;
        if rng.gen_ratio(1, sampling_count) {
          trace!("find_best_child: broke tie by selecting new action");
          best = edge;
        } else {
          trace!("find_best_child: broke tie by keeping old action");
        }
      }
      Some(Ordering::Less) => (),
    }
  }
  Ok(best)
}

/// [Rollout selector](../rollout/trait.RolloutSelector.html) that chooses a
/// child with the highest the UCB1 score.
///
/// If more than one child has the same rollout score, chooses one such child at
/// random.
pub struct Rollout {
  explore_bias: f64,
}

impl<'a> From<&'a crate::SearchSettings> for Rollout {
  fn from(settings: &'a crate::SearchSettings) -> Self {
    Rollout {
      explore_bias: settings.explore_bias,
    }
  }
}

impl RolloutSelector for Rollout {
  type Error = UcbError;

  fn select<'a, 'id, G: Game, R: Rng>(
    &self,
    graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
    parent: search_graph::view::NodeRef<'id>,
    rng: &mut R,
  ) -> Result<search_graph::view::EdgeRef<'id>, UcbError> {
    find_best_child(graph, parent, self.explore_bias, rng)
  }
}

/// [Backprop selector](../backprop/trait.BackpropSelector.html) that traverses
/// upward edges that would have been selected by the [UCB1
/// rollout](struct.Rollout.html) policy.
pub struct BestParentBackprop {
  explore_bias: f64,
}

impl<'a> From<&'a crate::SearchSettings> for BestParentBackprop {
  fn from(settings: &'a crate::SearchSettings) -> Self {
    BestParentBackprop {
      explore_bias: settings.explore_bias,
    }
  }
}

impl<'id> BackpropSelector<'id> for BestParentBackprop {
  // TODO: this requires an allocation. We have everything in place for this
  // selector to be allocation-free, except Rust doesn't yet support
  // ATCs/HKTs. We need ATC support because this iterator type will have its
  // lifetime constrained by the borrow of `graph` in the select method, but
  // that lifetime isn't known statically.
  type Items = std::vec::IntoIter<search_graph::view::EdgeRef<'id>>;

  fn select<G: Game, R: Rng>(
    &self,
    graph: &search_graph::view::View<'_, 'id, G::State, VertexData, EdgeData<G>>,
    node: search_graph::view::NodeRef<'id>,
    _payoff: &G::Payoff,
    _rng: &mut R,
  ) -> Self::Items {
    let result: Vec<search_graph::view::EdgeRef<'id>> = graph
      .parents(node)
      .filter(|&parent_edge| is_best_child(graph, parent_edge, self.explore_bias))
      .collect();
    result.into_iter()
  }
}
<commit_msg>Delete docs for function parameter that doesn't exist anymore.<commit_after>//! Upper confidence bound (UCB1) algorithm for graph search.

use crate::backprop::BackpropSelector;
use crate::game::{Game, State, Statistics};
use crate::graph::{EdgeData, VertexData};
use crate::rollout::RolloutSelector;
use log::{error, trace};
use rand::Rng;
use search_graph;

use std::cmp::Ordering;
use std::error::Error;
use std::f64;
use std::fmt;
use std::result::Result;

/// Represents success when computing the UCB score for a child.
pub enum UcbSuccess<'id> {
  /// No (finite) value can be computed, but the UCB policy indicates that
  /// this child should be selected. E.g., the child has not yet been visited.
  Select(search_graph::view::EdgeRef<'id>),
  /// A value is computed.
  Value(search_graph::view::EdgeRef<'id>, f64),
}

/// Represents error conditions when computing the UCB score for a child.
#[derive(Clone, Debug)]
pub enum UcbError {
  /// An error was encountered during computation of UCB score (such as
  /// encountering a `None` result when `PartialCmp`'ing two floating-point
  /// values).
  InvalidComputation,
}

unsafe impl Send for UcbError {}

/// Lazy iterator over UCB scores for a series of edges.
pub struct EdgeUcbIter<'a, 'b, 'id, G, I>
where
  G: Game,
  I: Iterator<Item = search_graph::view::EdgeRef<'id>>,
{
  log_parent_visits: f64,
  explore_bias: f64,
  graph: &'b search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  edges: I,
}

impl fmt::Display for UcbError {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    match *self {
      UcbError::InvalidComputation => write!(f, "Numerical error when computing UCB score"),
    }
  }
}

impl Error for UcbError {
  fn description(&self) -> &str {
    match *self {
      UcbError::InvalidComputation => "invalid computation",
    }
  }
}

impl<'a, 'b, 'id, G, I> EdgeUcbIter<'a, 'b, 'id, G, I>
where
  G: Game,
  I: Iterator<Item = search_graph::view::EdgeRef<'id>>,
{
  /// Constructs a new `EdgeUcbIter` that will compute UCB scores using the
  /// given constants. All floating-point values are assumed to be valid
  /// floating-point values and positive.
  ///
  ///  - `log_parent_visits`: ln(visits to parent vertex).
  ///  - `explore_bias`: scalar bias controlling tradeoff between search width
  ///    and search depth (lower = wider, higher = deeper).
  ///  - `edges`: an iterator over edge handles for which to compute UCB
  ///    scores. This should usually be a list of child edges which share a
  ///    parent vertex.
  pub fn new(
    log_parent_visits: f64,
    explore_bias: f64,
    graph: &'b search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
    edges: I,
  ) -> Self {
    EdgeUcbIter {
      log_parent_visits: log_parent_visits,
      explore_bias: explore_bias,
      graph,
      edges: edges,
    }
  }
}

impl<'a, 'b, 'id, G, I> Iterator for EdgeUcbIter<'a, 'b, 'id, G, I>
where
  G: Game,
  I: Iterator<Item = search_graph::view::EdgeRef<'id>>,
{
  type Item = Result<UcbSuccess<'id>, UcbError>;

  fn next(&mut self) -> Option<Result<UcbSuccess<'id>, UcbError>> {
    self.edges.next().map(|e| {
      if self.graph.edge_data(e).statistics.visits() == 0 {
        trace!("EdgeUcbIter selects unvisited action");
        Ok(UcbSuccess::Select(e))
      } else {
        trace!("EdgeUcbIter calls child_score with log_parent_visits = {}, explore_bias = {}, edge = {:?}",
               self.log_parent_visits,
               self.explore_bias,
               e);
        Ok(child_score(
          self.log_parent_visits,
          self.explore_bias,
          self.graph,
          e,
        ))
      }
    })
  }

  fn size_hint(&self) -> (usize, Option<usize>) {
    self.edges.size_hint()
  }
}

/// Returns the UCB policy result for the given values.
pub fn child_score<'a, 'id, G: Game>(
  log_parent_visits: f64,
  explore_bias: f64,
  graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  child: search_graph::view::EdgeRef<'id>,
) -> UcbSuccess<'id> {
  let statistics = &graph.edge_data(child).statistics;
  if statistics.visits() == 0 {
    UcbSuccess::Select(child)
  } else {
    let child_visits = statistics.visits() as f64;
    let child_score =
      statistics.score(graph.node_state(graph.edge_source(child)).active_player()) as f64;
    let ucb =
      child_score / child_visits + explore_bias * f64::sqrt(log_parent_visits / child_visits);
    UcbSuccess::Value(child, ucb)
  }
}

/// Returns `true` iff `e` could be selected by the UCB policy during rollout
/// from its parent vertex. Assumes we haven't yet altered the parent vertex
/// statistics. (Callers should ensure that this is not called more than once
/// for a given parent vertex.)
///
/// "Could be" reflects how this is different from calling `find_best_child`: it
/// is possible (and common) for multiple child edges of a vertex to be best
/// children (as when there are multiple children that have not yet been
/// explored). This is dealt with in `find_best_child` by randomly choosing one
/// of the best children. But when doing backpropagation on a full game state
/// graph (not just a tree), we want to know all of the parent edges which could
/// have rolled out to a given child.
pub fn is_best_child<'a, 'id, G: Game>(
  graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  e: search_graph::view::EdgeRef<'id>,
  explore_bias: f64,
) -> bool {
  let statistics = &graph.edge_data(e).statistics;
  // trace!("is_best_child: edge {} has {} visits", e.get_id(), stats.visits);
  if statistics.visits() == 0 {
    // Edge has been visited, but statistics aren't yet updated.
    // trace!("is_best_child: edge {} is a best child because stats.visits == 0", e.get_id());
    return true;
  }
  let parent = graph.edge_source(e);
  // trace!("is_best_child: edge {} (from node {}) has {} siblings", e.get_id(), parent.get_id(), siblings.len());
  let log_parent_visits = {
    let mut parent_visits = 0;
    for child_edge in graph.children(parent) {
      parent_visits += graph.edge_data(child_edge).statistics.visits();
    }
    f64::ln(parent_visits as f64)
  };
  let mut edge_ucb = None;
  let mut best_ucb = ::std::f64::MIN;
  let ucb_iter = EdgeUcbIter::new(
    log_parent_visits,
    explore_bias,
    graph,
    graph.children(parent),
  );
  // Scan through siblings to find the maximum UCB score. This is
  // short-circuited using a lazy iterator to ameliorate the O(n) running
  // time.
  for ucb in ucb_iter {
    match ucb {
      Ok(UcbSuccess::Select(sibling)) if e == sibling => {
        // trace!("is_best_child: edge {} is best child by fiat of score computation", e.get_id());
        return true;
      }
      Ok(UcbSuccess::Select(_)) => {
        if let Some(_) = edge_ucb {
          // trace!("is_best_child: edge {} has a ucb of {}, but sibling edge {} is best child by fiat of score computation", e.get_id(), u, sibling.get_id());
          return false;
        }
      }
      Ok(UcbSuccess::Value(sibling, score)) => {
        if sibling == e {
          if best_ucb > score {
            // trace!("is_best_child: found ucb {} for edge {}, but a sibling has a higher ucb of {}", score, e.get_id(), best_ucb);
            return false;
          }
          edge_ucb = Some(score)
        } else if let Some(u) = edge_ucb {
          if score > u {
            // trace!("is_best_child: found ucb {} for edge {}, but a sibling has a higher ucb of {}", score, e.get_id(), best_ucb);
            return false;
          }
        }
        if score > best_ucb {
          best_ucb = score;
        }
      }
      Err(e) => panic!("error {:?} computing ucb for best child search", e),
    }
  }
  match edge_ucb {
    Some(u) if u >= best_ucb => {
      // Target edge has a UCB score which matches the maximum we found.
      // trace!("is_best_child: edge {} has a ucb of {}, which matches the max sibling ucb of {}", e.get_id(), u, best_ucb);
      true
    }
    _ => {
      // ThudEdges are not best children by default.
      // trace!("is_best_child: edge {} has a ucb of {:?}, which does not match the max sibling ucb of {}", e.get_id(), edge_ucb, best_ucb);
      false
    }
  }
}

/// Returns the child edge of `parent` that is best according to the UCB1
/// criterion.
///
/// This function will panic if `parent` has no children.
pub fn find_best_child<'a, 'id, G, R>(
  graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
  parent: search_graph::view::NodeRef<'id>,
  explore_bias: f64,
  rng: &mut R,
) -> Result<search_graph::view::EdgeRef<'id>, UcbError>
where
  G: Game,
  R: Rng,
{
  let log_parent_visits = {
    let mut parent_visits = 0;
    for child in graph.children(parent) {
      parent_visits += graph.edge_data(child).statistics.visits();
    }
    if parent_visits == 0 {
      // When we visit a vertex for the first time, it will have zero visits.
      0.0
    } else {
      // Otherwise, it should be positive.
      f64::ln(parent_visits as f64)
    }
  };
  let mut sampling_count = 0u32;
  let mut ucb_iter = EdgeUcbIter::new(
    log_parent_visits,
    explore_bias,
    graph,
    graph.children(parent),
  );
  let mut best: search_graph::view::EdgeRef<'id>;
  let mut best_ucb: f64;
  match ucb_iter.next().expect("vertex has no children")? {
    UcbSuccess::Select(e) => {
      best = e;
      best_ucb = f64::MAX;
    }
    UcbSuccess::Value(e, v) => {
      best = e;
      best_ucb = v;
    }
  }

  for ucb in ucb_iter {
    let (edge, value) = match ucb? {
      UcbSuccess::Select(e) => (e, f64::MAX),
      UcbSuccess::Value(e, v) => (e, v),
    };
    match value.partial_cmp(&best_ucb) {
      None => {
        error!("find_best_child: invalid floating-point comparison");
        return Err(UcbError::InvalidComputation);
      }
      Some(Ordering::Greater) => {
        trace!("find_best_child_edge: new best action with score {}", value);
        best = edge;
        best_ucb = value;
        sampling_count = 1;
      }
      Some(Ordering::Equal) => {
        // We use reservoir sampling to break ties.
        trace!("find_best_child: found action with identical score; sampling to break tie");
        sampling_count += 1;
        if rng.gen_ratio(1, sampling_count) {
          trace!("find_best_child: broke tie by selecting new action");
          best = edge;
        } else {
          trace!("find_best_child: broke tie by keeping old action");
        }
      }
      Some(Ordering::Less) => (),
    }
  }
  Ok(best)
}

/// [Rollout selector](../rollout/trait.RolloutSelector.html) that chooses a
/// child with the highest the UCB1 score.
///
/// If more than one child has the same rollout score, chooses one such child at
/// random.
pub struct Rollout {
  explore_bias: f64,
}

impl<'a> From<&'a crate::SearchSettings> for Rollout {
  fn from(settings: &'a crate::SearchSettings) -> Self {
    Rollout {
      explore_bias: settings.explore_bias,
    }
  }
}

impl RolloutSelector for Rollout {
  type Error = UcbError;

  fn select<'a, 'id, G: Game, R: Rng>(
    &self,
    graph: &search_graph::view::View<'a, 'id, G::State, VertexData, EdgeData<G>>,
    parent: search_graph::view::NodeRef<'id>,
    rng: &mut R,
  ) -> Result<search_graph::view::EdgeRef<'id>, UcbError> {
    find_best_child(graph, parent, self.explore_bias, rng)
  }
}

/// [Backprop selector](../backprop/trait.BackpropSelector.html) that traverses
/// upward edges that would have been selected by the [UCB1
/// rollout](struct.Rollout.html) policy.
pub struct BestParentBackprop {
  explore_bias: f64,
}

impl<'a> From<&'a crate::SearchSettings> for BestParentBackprop {
  fn from(settings: &'a crate::SearchSettings) -> Self {
    BestParentBackprop {
      explore_bias: settings.explore_bias,
    }
  }
}

impl<'id> BackpropSelector<'id> for BestParentBackprop {
  // TODO: this requires an allocation. We have everything in place for this
  // selector to be allocation-free, except Rust doesn't yet support
  // ATCs/HKTs. We need ATC support because this iterator type will have its
  // lifetime constrained by the borrow of `graph` in the select method, but
  // that lifetime isn't known statically.
  type Items = std::vec::IntoIter<search_graph::view::EdgeRef<'id>>;

  fn select<G: Game, R: Rng>(
    &self,
    graph: &search_graph::view::View<'_, 'id, G::State, VertexData, EdgeData<G>>,
    node: search_graph::view::NodeRef<'id>,
    _payoff: &G::Payoff,
    _rng: &mut R,
  ) -> Self::Items {
    let result: Vec<search_graph::view::EdgeRef<'id>> = graph
      .parents(node)
      .filter(|&parent_edge| is_best_child(graph, parent_edge, self.explore_bias))
      .collect();
    result.into_iter()
  }
}
<|endoftext|>


<commit_before>use crate::interface::{Compiler, Result};
use crate::passes::{self, BoxedResolver, QueryContext};

use rustc_ast::{self, ast};
use rustc_codegen_ssa::traits::CodegenBackend;
use rustc_data_structures::sync::{Lrc, OnceCell, WorkerLocal};
use rustc_errors::ErrorReported;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_hir::Crate;
use rustc_incremental::DepGraphFuture;
use rustc_lint::LintStore;
use rustc_middle::arena::Arena;
use rustc_middle::dep_graph::DepGraph;
use rustc_middle::ty::steal::Steal;
use rustc_middle::ty::{GlobalCtxt, ResolverOutputs, TyCtxt};
use rustc_session::config::{OutputFilenames, OutputType};
use rustc_session::{output::find_crate_name, Session};
use rustc_span::symbol::sym;
use std::any::Any;
use std::cell::{Ref, RefCell, RefMut};
use std::mem;
use std::rc::Rc;

/// Represent the result of a query.
/// This result can be stolen with the `take` method and generated with the `compute` method.
pub struct Query<T> {
    result: RefCell<Option<Result<T>>>,
}

impl<T> Query<T> {
    fn compute<F: FnOnce() -> Result<T>>(&self, f: F) -> Result<&Query<T>> {
        let mut result = self.result.borrow_mut();
        if result.is_none() {
            *result = Some(f());
        }
        result.as_ref().unwrap().as_ref().map(|_| self).map_err(|err| *err)
    }

    /// Takes ownership of the query result. Further attempts to take or peek the query
    /// result will panic unless it is generated by calling the `compute` method.
    pub fn take(&self) -> T {
        self.result.borrow_mut().take().expect("missing query result").unwrap()
    }

    /// Borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek(&self) -> Ref<'_, T> {
        Ref::map(self.result.borrow(), |r| {
            r.as_ref().unwrap().as_ref().expect("missing query result")
        })
    }

    /// Mutably borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek_mut(&self) -> RefMut<'_, T> {
        RefMut::map(self.result.borrow_mut(), |r| {
            r.as_mut().unwrap().as_mut().expect("missing query result")
        })
    }
}

impl<T> Default for Query<T> {
    fn default() -> Self {
        Query { result: RefCell::new(None) }
    }
}

pub struct Queries<'tcx> {
    compiler: &'tcx Compiler,
    gcx: OnceCell<GlobalCtxt<'tcx>>,

    arena: WorkerLocal<Arena<'tcx>>,
    hir_arena: WorkerLocal<rustc_ast_lowering::Arena<'tcx>>,

    dep_graph_future: Query<Option<DepGraphFuture>>,
    parse: Query<ast::Crate>,
    crate_name: Query<String>,
    register_plugins: Query<(ast::Crate, Lrc<LintStore>)>,
    expansion: Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>,
    dep_graph: Query<DepGraph>,
    lower_to_hir: Query<(&'tcx Crate<'tcx>, Steal<ResolverOutputs>)>,
    prepare_outputs: Query<OutputFilenames>,
    global_ctxt: Query<QueryContext<'tcx>>,
    ongoing_codegen: Query<Box<dyn Any>>,
}

impl<'tcx> Queries<'tcx> {
    pub fn new(compiler: &'tcx Compiler) -> Queries<'tcx> {
        Queries {
            compiler,
            gcx: OnceCell::new(),
            arena: WorkerLocal::new(|_| Arena::default()),
            hir_arena: WorkerLocal::new(|_| rustc_ast_lowering::Arena::default()),
            dep_graph_future: Default::default(),
            parse: Default::default(),
            crate_name: Default::default(),
            register_plugins: Default::default(),
            expansion: Default::default(),
            dep_graph: Default::default(),
            lower_to_hir: Default::default(),
            prepare_outputs: Default::default(),
            global_ctxt: Default::default(),
            ongoing_codegen: Default::default(),
        }
    }

    fn session(&self) -> &Lrc<Session> {
        &self.compiler.sess
    }
    fn codegen_backend(&self) -> &Lrc<Box<dyn CodegenBackend>> {
        &self.compiler.codegen_backend()
    }

    pub fn dep_graph_future(&self) -> Result<&Query<Option<DepGraphFuture>>> {
        self.dep_graph_future.compute(|| {
            Ok(self
                .session()
                .opts
                .build_dep_graph()
                .then(|| rustc_incremental::load_dep_graph(self.session())))
        })
    }

    pub fn parse(&self) -> Result<&Query<ast::Crate>> {
        self.parse.compute(|| {
            passes::parse(self.session(), &self.compiler.input).map_err(|mut parse_error| {
                parse_error.emit();
                ErrorReported
            })
        })
    }

    pub fn register_plugins(&self) -> Result<&Query<(ast::Crate, Lrc<LintStore>)>> {
        self.register_plugins.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let krate = self.parse()?.take();

            let empty: &(dyn Fn(&Session, &mut LintStore) + Sync + Send) = &|_, _| {};
            let result = passes::register_plugins(
                self.session(),
                &*self.codegen_backend().metadata_loader(),
                self.compiler.register_lints.as_deref().unwrap_or_else(|| empty),
                krate,
                &crate_name,
            );

            // Compute the dependency graph (in the background). We want to do
            // this as early as possible, to give the DepGraph maximum time to
            // load before dep_graph() is called, but it also can't happen
            // until after rustc_incremental::prepare_session_directory() is
            // called, which happens within passes::register_plugins().
            self.dep_graph_future().ok();

            result
        })
    }

    pub fn crate_name(&self) -> Result<&Query<String>> {
        self.crate_name.compute(|| {
            Ok(match self.compiler.crate_name {
                Some(ref crate_name) => crate_name.clone(),
                None => {
                    let parse_result = self.parse()?;
                    let krate = parse_result.peek();
                    find_crate_name(Some(self.session()), &krate.attrs, &self.compiler.input)
                }
            })
        })
    }

    pub fn expansion(
        &self,
    ) -> Result<&Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>> {
        log::trace!("expansion");
        self.expansion.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let (krate, lint_store) = self.register_plugins()?.take();
            let _timer = self.session().timer("configure_and_expand");
            passes::configure_and_expand(
                self.session().clone(),
                lint_store.clone(),
                self.codegen_backend().metadata_loader(),
                krate,
                &crate_name,
            )
            .map(|(krate, resolver)| {
                (krate, Steal::new(Rc::new(RefCell::new(resolver))), lint_store)
            })
        })
    }

    pub fn dep_graph(&self) -> Result<&Query<DepGraph>> {
        self.dep_graph.compute(|| {
            Ok(match self.dep_graph_future()?.take() {
                None => DepGraph::new_disabled(),
                Some(future) => {
                    let (prev_graph, prev_work_products) =
                        self.session().time("blocked_on_dep_graph_loading", || {
                            future
                                .open()
                                .unwrap_or_else(|e| rustc_incremental::LoadResult::Error {
                                    message: format!("could not decode incremental cache: {:?}", e),
                                })
                                .open(self.session())
                        });
                    DepGraph::new(prev_graph, prev_work_products)
                }
            })
        })
    }

    pub fn lower_to_hir(&'tcx self) -> Result<&Query<(&'tcx Crate<'tcx>, Steal<ResolverOutputs>)>> {
        self.lower_to_hir.compute(|| {
            let expansion_result = self.expansion()?;
            let peeked = expansion_result.peek();
            let krate = &peeked.0;
            let resolver = peeked.1.steal();
            let lint_store = &peeked.2;
            let hir = resolver.borrow_mut().access(|resolver| {
                Ok(passes::lower_to_hir(
                    self.session(),
                    lint_store,
                    resolver,
                    &*self.dep_graph()?.peek(),
                    &krate,
                    &self.hir_arena,
                ))
            })?;
            let hir = self.hir_arena.alloc(hir);
            Ok((hir, Steal::new(BoxedResolver::to_resolver_outputs(resolver))))
        })
    }

    pub fn prepare_outputs(&self) -> Result<&Query<OutputFilenames>> {
        self.prepare_outputs.compute(|| {
            let expansion_result = self.expansion()?;
            let (krate, boxed_resolver, _) = &*expansion_result.peek();
            let crate_name = self.crate_name()?;
            let crate_name = crate_name.peek();
            passes::prepare_outputs(
                self.session(),
                self.compiler,
                &krate,
                &boxed_resolver,
                &crate_name,
            )
        })
    }

    pub fn global_ctxt(&'tcx self) -> Result<&Query<QueryContext<'tcx>>> {
        self.global_ctxt.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let outputs = self.prepare_outputs()?.peek().clone();
            let lint_store = self.expansion()?.peek().2.clone();
            let hir = self.lower_to_hir()?.peek();
            let dep_graph = self.dep_graph()?.peek().clone();
            let (ref krate, ref resolver_outputs) = &*hir;
            let _timer = self.session().timer("create_global_ctxt");
            Ok(passes::create_global_ctxt(
                self.compiler,
                lint_store,
                krate,
                dep_graph,
                resolver_outputs.steal(),
                outputs,
                &crate_name,
                &self.gcx,
                &self.arena,
            ))
        })
    }

    pub fn ongoing_codegen(&'tcx self) -> Result<&Query<Box<dyn Any>>> {
        self.ongoing_codegen.compute(|| {
            let outputs = self.prepare_outputs()?;
            self.global_ctxt()?.peek_mut().enter(|tcx| {
                tcx.analysis(LOCAL_CRATE).ok();

                // Don't do code generation if there were any errors
                self.session().compile_status()?;

                // Hook for compile-fail tests.
                Self::check_for_rustc_errors_attr(tcx);

                Ok(passes::start_codegen(&***self.codegen_backend(), tcx, &*outputs.peek()))
            })
        })
    }

    /// Check for the `#[rustc_error]` annotation, which forces an error in codegen. This is used
    /// to write compile-fail tests that actually test that compilation succeeds without reporting
    /// an error.
    fn check_for_rustc_errors_attr(tcx: TyCtxt<'_>) {
        let def_id = match tcx.entry_fn(LOCAL_CRATE) {
            Some((def_id, _)) => def_id,
            _ => return,
        };

        let attrs = &*tcx.get_attrs(def_id.to_def_id());
        let attrs = attrs.iter().filter(|attr| attr.check_name(sym::rustc_error));
        for attr in attrs {
            match attr.meta_item_list() {
                // Check if there is a `#[rustc_error(delay_span_bug_from_inside_query)]`.
                Some(list)
                    if list.iter().any(|list_item| {
                        matches!(
                            list_item.ident().map(|i| i.name),
                            Some(sym::delay_span_bug_from_inside_query)
                        )
                    }) =>
                {
                    tcx.ensure().trigger_delay_span_bug(def_id);
                }

                // Bare `#[rustc_error]`.
                None => {
                    tcx.sess.span_fatal(
                        tcx.def_span(def_id),
                        "fatal error triggered by #[rustc_error]",
                    );
                }

                // Some other attribute.
                Some(_) => {
                    tcx.sess.span_warn(
                        tcx.def_span(def_id),
                        "unexpected annotation used with `#[rustc_error(...)]!",
                    );
                }
            }
        }
    }

    pub fn linker(&'tcx self) -> Result<Linker> {
        let dep_graph = self.dep_graph()?;
        let prepare_outputs = self.prepare_outputs()?;
        let ongoing_codegen = self.ongoing_codegen()?;

        let sess = self.session().clone();
        let codegen_backend = self.codegen_backend().clone();

        Ok(Linker {
            sess,
            dep_graph: dep_graph.peek().clone(),
            prepare_outputs: prepare_outputs.take(),
            ongoing_codegen: ongoing_codegen.take(),
            codegen_backend,
        })
    }
}

pub struct Linker {
    sess: Lrc<Session>,
    dep_graph: DepGraph,
    prepare_outputs: OutputFilenames,
    ongoing_codegen: Box<dyn Any>,
    codegen_backend: Lrc<Box<dyn CodegenBackend>>,
}

impl Linker {
    pub fn link(self) -> Result<()> {
        let codegen_results =
            self.codegen_backend.join_codegen(self.ongoing_codegen, &self.sess, &self.dep_graph)?;
        let prof = self.sess.prof.clone();
        let dep_graph = self.dep_graph;
        prof.generic_activity("drop_dep_graph").run(move || drop(dep_graph));

        if !self
            .sess
            .opts
            .output_types
            .keys()
            .any(|&i| i == OutputType::Exe || i == OutputType::Metadata)
        {
            return Ok(());
        }
        self.codegen_backend.link(&self.sess, codegen_results, &self.prepare_outputs)
    }
}

impl Compiler {
    pub fn enter<F, T>(&self, f: F) -> T
    where
        F: for<'tcx> FnOnce(&'tcx Queries<'tcx>) -> T,
    {
        let mut _timer = None;
        let queries = Queries::new(&self);
        let ret = f(&queries);

        if self.session().opts.debugging_opts.query_stats {
            if let Ok(gcx) = queries.global_ctxt() {
                gcx.peek_mut().print_stats();
            }
        }

        _timer = Some(self.session().timer("free_global_ctxt"));

        ret
    }

    // This method is different to all the other methods in `Compiler` because
    // it lacks a `Queries` entry. It's also not currently used. It does serve
    // as an example of how `Compiler` can be used, with additional steps added
    // between some passes. And see `rustc_driver::run_compiler` for a more
    // complex example.
    pub fn compile(&self) -> Result<()> {
        let linker = self.enter(|queries| {
            queries.prepare_outputs()?;

            if self.session().opts.output_types.contains_key(&OutputType::DepInfo)
                && self.session().opts.output_types.len() == 1
            {
                return Ok(None);
            }

            queries.global_ctxt()?;

            // Drop AST after creating GlobalCtxt to free memory.
            mem::drop(queries.expansion()?.take());

            queries.ongoing_codegen()?;

            let linker = queries.linker()?;
            Ok(Some(linker))
        })?;

        if let Some(linker) = linker {
            linker.link()?
        }

        Ok(())
    }
}
<commit_msg>Remove `Compiler::compile()`.<commit_after>use crate::interface::{Compiler, Result};
use crate::passes::{self, BoxedResolver, QueryContext};

use rustc_ast::{self, ast};
use rustc_codegen_ssa::traits::CodegenBackend;
use rustc_data_structures::sync::{Lrc, OnceCell, WorkerLocal};
use rustc_errors::ErrorReported;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_hir::Crate;
use rustc_incremental::DepGraphFuture;
use rustc_lint::LintStore;
use rustc_middle::arena::Arena;
use rustc_middle::dep_graph::DepGraph;
use rustc_middle::ty::steal::Steal;
use rustc_middle::ty::{GlobalCtxt, ResolverOutputs, TyCtxt};
use rustc_session::config::{OutputFilenames, OutputType};
use rustc_session::{output::find_crate_name, Session};
use rustc_span::symbol::sym;
use std::any::Any;
use std::cell::{Ref, RefCell, RefMut};
use std::rc::Rc;

/// Represent the result of a query.
/// This result can be stolen with the `take` method and generated with the `compute` method.
pub struct Query<T> {
    result: RefCell<Option<Result<T>>>,
}

impl<T> Query<T> {
    fn compute<F: FnOnce() -> Result<T>>(&self, f: F) -> Result<&Query<T>> {
        let mut result = self.result.borrow_mut();
        if result.is_none() {
            *result = Some(f());
        }
        result.as_ref().unwrap().as_ref().map(|_| self).map_err(|err| *err)
    }

    /// Takes ownership of the query result. Further attempts to take or peek the query
    /// result will panic unless it is generated by calling the `compute` method.
    pub fn take(&self) -> T {
        self.result.borrow_mut().take().expect("missing query result").unwrap()
    }

    /// Borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek(&self) -> Ref<'_, T> {
        Ref::map(self.result.borrow(), |r| {
            r.as_ref().unwrap().as_ref().expect("missing query result")
        })
    }

    /// Mutably borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek_mut(&self) -> RefMut<'_, T> {
        RefMut::map(self.result.borrow_mut(), |r| {
            r.as_mut().unwrap().as_mut().expect("missing query result")
        })
    }
}

impl<T> Default for Query<T> {
    fn default() -> Self {
        Query { result: RefCell::new(None) }
    }
}

pub struct Queries<'tcx> {
    compiler: &'tcx Compiler,
    gcx: OnceCell<GlobalCtxt<'tcx>>,

    arena: WorkerLocal<Arena<'tcx>>,
    hir_arena: WorkerLocal<rustc_ast_lowering::Arena<'tcx>>,

    dep_graph_future: Query<Option<DepGraphFuture>>,
    parse: Query<ast::Crate>,
    crate_name: Query<String>,
    register_plugins: Query<(ast::Crate, Lrc<LintStore>)>,
    expansion: Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>,
    dep_graph: Query<DepGraph>,
    lower_to_hir: Query<(&'tcx Crate<'tcx>, Steal<ResolverOutputs>)>,
    prepare_outputs: Query<OutputFilenames>,
    global_ctxt: Query<QueryContext<'tcx>>,
    ongoing_codegen: Query<Box<dyn Any>>,
}

impl<'tcx> Queries<'tcx> {
    pub fn new(compiler: &'tcx Compiler) -> Queries<'tcx> {
        Queries {
            compiler,
            gcx: OnceCell::new(),
            arena: WorkerLocal::new(|_| Arena::default()),
            hir_arena: WorkerLocal::new(|_| rustc_ast_lowering::Arena::default()),
            dep_graph_future: Default::default(),
            parse: Default::default(),
            crate_name: Default::default(),
            register_plugins: Default::default(),
            expansion: Default::default(),
            dep_graph: Default::default(),
            lower_to_hir: Default::default(),
            prepare_outputs: Default::default(),
            global_ctxt: Default::default(),
            ongoing_codegen: Default::default(),
        }
    }

    fn session(&self) -> &Lrc<Session> {
        &self.compiler.sess
    }
    fn codegen_backend(&self) -> &Lrc<Box<dyn CodegenBackend>> {
        &self.compiler.codegen_backend()
    }

    pub fn dep_graph_future(&self) -> Result<&Query<Option<DepGraphFuture>>> {
        self.dep_graph_future.compute(|| {
            Ok(self
                .session()
                .opts
                .build_dep_graph()
                .then(|| rustc_incremental::load_dep_graph(self.session())))
        })
    }

    pub fn parse(&self) -> Result<&Query<ast::Crate>> {
        self.parse.compute(|| {
            passes::parse(self.session(), &self.compiler.input).map_err(|mut parse_error| {
                parse_error.emit();
                ErrorReported
            })
        })
    }

    pub fn register_plugins(&self) -> Result<&Query<(ast::Crate, Lrc<LintStore>)>> {
        self.register_plugins.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let krate = self.parse()?.take();

            let empty: &(dyn Fn(&Session, &mut LintStore) + Sync + Send) = &|_, _| {};
            let result = passes::register_plugins(
                self.session(),
                &*self.codegen_backend().metadata_loader(),
                self.compiler.register_lints.as_deref().unwrap_or_else(|| empty),
                krate,
                &crate_name,
            );

            // Compute the dependency graph (in the background). We want to do
            // this as early as possible, to give the DepGraph maximum time to
            // load before dep_graph() is called, but it also can't happen
            // until after rustc_incremental::prepare_session_directory() is
            // called, which happens within passes::register_plugins().
            self.dep_graph_future().ok();

            result
        })
    }

    pub fn crate_name(&self) -> Result<&Query<String>> {
        self.crate_name.compute(|| {
            Ok(match self.compiler.crate_name {
                Some(ref crate_name) => crate_name.clone(),
                None => {
                    let parse_result = self.parse()?;
                    let krate = parse_result.peek();
                    find_crate_name(Some(self.session()), &krate.attrs, &self.compiler.input)
                }
            })
        })
    }

    pub fn expansion(
        &self,
    ) -> Result<&Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>> {
        log::trace!("expansion");
        self.expansion.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let (krate, lint_store) = self.register_plugins()?.take();
            let _timer = self.session().timer("configure_and_expand");
            passes::configure_and_expand(
                self.session().clone(),
                lint_store.clone(),
                self.codegen_backend().metadata_loader(),
                krate,
                &crate_name,
            )
            .map(|(krate, resolver)| {
                (krate, Steal::new(Rc::new(RefCell::new(resolver))), lint_store)
            })
        })
    }

    pub fn dep_graph(&self) -> Result<&Query<DepGraph>> {
        self.dep_graph.compute(|| {
            Ok(match self.dep_graph_future()?.take() {
                None => DepGraph::new_disabled(),
                Some(future) => {
                    let (prev_graph, prev_work_products) =
                        self.session().time("blocked_on_dep_graph_loading", || {
                            future
                                .open()
                                .unwrap_or_else(|e| rustc_incremental::LoadResult::Error {
                                    message: format!("could not decode incremental cache: {:?}", e),
                                })
                                .open(self.session())
                        });
                    DepGraph::new(prev_graph, prev_work_products)
                }
            })
        })
    }

    pub fn lower_to_hir(&'tcx self) -> Result<&Query<(&'tcx Crate<'tcx>, Steal<ResolverOutputs>)>> {
        self.lower_to_hir.compute(|| {
            let expansion_result = self.expansion()?;
            let peeked = expansion_result.peek();
            let krate = &peeked.0;
            let resolver = peeked.1.steal();
            let lint_store = &peeked.2;
            let hir = resolver.borrow_mut().access(|resolver| {
                Ok(passes::lower_to_hir(
                    self.session(),
                    lint_store,
                    resolver,
                    &*self.dep_graph()?.peek(),
                    &krate,
                    &self.hir_arena,
                ))
            })?;
            let hir = self.hir_arena.alloc(hir);
            Ok((hir, Steal::new(BoxedResolver::to_resolver_outputs(resolver))))
        })
    }

    pub fn prepare_outputs(&self) -> Result<&Query<OutputFilenames>> {
        self.prepare_outputs.compute(|| {
            let expansion_result = self.expansion()?;
            let (krate, boxed_resolver, _) = &*expansion_result.peek();
            let crate_name = self.crate_name()?;
            let crate_name = crate_name.peek();
            passes::prepare_outputs(
                self.session(),
                self.compiler,
                &krate,
                &boxed_resolver,
                &crate_name,
            )
        })
    }

    pub fn global_ctxt(&'tcx self) -> Result<&Query<QueryContext<'tcx>>> {
        self.global_ctxt.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let outputs = self.prepare_outputs()?.peek().clone();
            let lint_store = self.expansion()?.peek().2.clone();
            let hir = self.lower_to_hir()?.peek();
            let dep_graph = self.dep_graph()?.peek().clone();
            let (ref krate, ref resolver_outputs) = &*hir;
            let _timer = self.session().timer("create_global_ctxt");
            Ok(passes::create_global_ctxt(
                self.compiler,
                lint_store,
                krate,
                dep_graph,
                resolver_outputs.steal(),
                outputs,
                &crate_name,
                &self.gcx,
                &self.arena,
            ))
        })
    }

    pub fn ongoing_codegen(&'tcx self) -> Result<&Query<Box<dyn Any>>> {
        self.ongoing_codegen.compute(|| {
            let outputs = self.prepare_outputs()?;
            self.global_ctxt()?.peek_mut().enter(|tcx| {
                tcx.analysis(LOCAL_CRATE).ok();

                // Don't do code generation if there were any errors
                self.session().compile_status()?;

                // Hook for compile-fail tests.
                Self::check_for_rustc_errors_attr(tcx);

                Ok(passes::start_codegen(&***self.codegen_backend(), tcx, &*outputs.peek()))
            })
        })
    }

    /// Check for the `#[rustc_error]` annotation, which forces an error in codegen. This is used
    /// to write compile-fail tests that actually test that compilation succeeds without reporting
    /// an error.
    fn check_for_rustc_errors_attr(tcx: TyCtxt<'_>) {
        let def_id = match tcx.entry_fn(LOCAL_CRATE) {
            Some((def_id, _)) => def_id,
            _ => return,
        };

        let attrs = &*tcx.get_attrs(def_id.to_def_id());
        let attrs = attrs.iter().filter(|attr| attr.check_name(sym::rustc_error));
        for attr in attrs {
            match attr.meta_item_list() {
                // Check if there is a `#[rustc_error(delay_span_bug_from_inside_query)]`.
                Some(list)
                    if list.iter().any(|list_item| {
                        matches!(
                            list_item.ident().map(|i| i.name),
                            Some(sym::delay_span_bug_from_inside_query)
                        )
                    }) =>
                {
                    tcx.ensure().trigger_delay_span_bug(def_id);
                }

                // Bare `#[rustc_error]`.
                None => {
                    tcx.sess.span_fatal(
                        tcx.def_span(def_id),
                        "fatal error triggered by #[rustc_error]",
                    );
                }

                // Some other attribute.
                Some(_) => {
                    tcx.sess.span_warn(
                        tcx.def_span(def_id),
                        "unexpected annotation used with `#[rustc_error(...)]!",
                    );
                }
            }
        }
    }

    pub fn linker(&'tcx self) -> Result<Linker> {
        let dep_graph = self.dep_graph()?;
        let prepare_outputs = self.prepare_outputs()?;
        let ongoing_codegen = self.ongoing_codegen()?;

        let sess = self.session().clone();
        let codegen_backend = self.codegen_backend().clone();

        Ok(Linker {
            sess,
            dep_graph: dep_graph.peek().clone(),
            prepare_outputs: prepare_outputs.take(),
            ongoing_codegen: ongoing_codegen.take(),
            codegen_backend,
        })
    }
}

pub struct Linker {
    sess: Lrc<Session>,
    dep_graph: DepGraph,
    prepare_outputs: OutputFilenames,
    ongoing_codegen: Box<dyn Any>,
    codegen_backend: Lrc<Box<dyn CodegenBackend>>,
}

impl Linker {
    pub fn link(self) -> Result<()> {
        let codegen_results =
            self.codegen_backend.join_codegen(self.ongoing_codegen, &self.sess, &self.dep_graph)?;
        let prof = self.sess.prof.clone();
        let dep_graph = self.dep_graph;
        prof.generic_activity("drop_dep_graph").run(move || drop(dep_graph));

        if !self
            .sess
            .opts
            .output_types
            .keys()
            .any(|&i| i == OutputType::Exe || i == OutputType::Metadata)
        {
            return Ok(());
        }
        self.codegen_backend.link(&self.sess, codegen_results, &self.prepare_outputs)
    }
}

impl Compiler {
    pub fn enter<F, T>(&self, f: F) -> T
    where
        F: for<'tcx> FnOnce(&'tcx Queries<'tcx>) -> T,
    {
        let mut _timer = None;
        let queries = Queries::new(&self);
        let ret = f(&queries);

        if self.session().opts.debugging_opts.query_stats {
            if let Ok(gcx) = queries.global_ctxt() {
                gcx.peek_mut().print_stats();
            }
        }

        _timer = Some(self.session().timer("free_global_ctxt"));

        ret
    }
}
<|endoftext|>