<commit_before>use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::slice::Iter as SliceIter;
use std::rc::Rc;

use chrono::{
    DateTime,
    FixedOffset,
    Offset,
    UTC,
};

use super::rfc2047::decode_rfc2047;
use super::rfc822::Rfc822DateParser;

/// Trait for converting from RFC822 Header values into
/// Rust types.
#[stable]
pub trait FromHeader {
    /// Parse the `value` of the header.
    ///
    /// Returns None if the value failed to be parsed
    fn from_header(value: String) -> Option<Self>;
}

/// Trait for converting from a Rust type into a Header value.
#[stable]
pub trait ToHeader {
    /// Turn the `value` into a String suitable for being used in
    /// a message header.
    ///
    /// Returns None if the value cannot be stringified.
    fn to_header(value: Self) -> Option<String>;
}

/// Trait for converting from a Rust time into a Header value
/// that handles its own folding.
///
/// Be mindful that this trait does not mean that the value will
/// not be folded later, rather that the type returns a value that
/// should not be folded, given that the header value starts so far
/// in to a line.
#[unstable]
pub trait ToFoldedHeader {
    fn to_folded_header(start_pos: uint, value: Self) -> Option<String>;
}

impl<T: ToHeader> ToFoldedHeader for T {
    fn to_folded_header(_: uint, value: T) -> Option<String> {
        // We ignore the start_position because the thing will fold anyway.
        ToHeader::to_header(value)
    }
}

impl FromHeader for String {
    fn from_header(value: String) -> Option<String> {
        #[derive(Show,Copy)]
        enum ParseState {
            Normal(uint),
            SeenEquals(uint),
            SeenQuestion(uint, uint),
        }

        let mut state = ParseState::Normal(0u);
        let mut decoded = String::new();
        let mut pos = 0u;

        let value_slice = value.as_slice();

        while pos < value.len() {
            let ch_range = value_slice.char_range_at(pos);
            let c = ch_range.ch;

            state = match (state, c) {
                (ParseState::SeenQuestion(start_pos, 4), '=') => {
                    // Go to decode if we've seen enough ?
                    let part_decoded = decode_rfc2047(value_slice.slice(start_pos, ch_range.next));
                    let to_push = match part_decoded {
                        Some(ref s) => s.as_slice(),
                        // Decoding failed, push the undecoded string in.
                        None => value_slice.slice(start_pos, pos),
                    };
                    decoded.push_str(to_push);
                    // Revert us to normal state, but starting at the next character.
                    ParseState::Normal(ch_range.next)
                },
                (ParseState::SeenQuestion(start_pos, count), '?') => {
                    ParseState::SeenQuestion(start_pos, count + 1)
                },
                (ParseState::SeenQuestion(start_pos, count), _) => {
                    if count > 4 {
                        // This isn't a RFC2047 sequence, so go back to a normal.
                        ParseState::Normal(start_pos)
                    } else {
                        state
                    }
                }
                (ParseState::SeenEquals(start_pos), '?') => {
                    ParseState::SeenQuestion(start_pos, 1)
                },
                (ParseState::SeenEquals(start_pos), _) => {
                    // This isn't a RFC2047 sequence, so go back to a normal.
                    ParseState::Normal(start_pos)
                }
                (ParseState::Normal(start_pos), '=') => {
                    if start_pos != pos {
                        // Push all up to the =, if there is stuff to push.
                        decoded.push_str(value_slice.slice(start_pos, pos));
                    }
                    ParseState::SeenEquals(pos)
                },
                (ParseState::Normal(_), _) => state,
            };

            pos = ch_range.next;
        }

        // Don't forget to push on whatever we have left
        let last_start = match state {
            ParseState::Normal(start_pos) => start_pos,
            ParseState::SeenEquals(start_pos) => start_pos,
            ParseState::SeenQuestion(start_pos, _) => start_pos,
        };
        decoded.push_str(value_slice.slice_from(last_start));

        Some(decoded)
    }
}

impl FromHeader for DateTime<FixedOffset> {
    fn from_header(value: String) -> Option<DateTime<FixedOffset>> {
        let mut parser = Rfc822DateParser::new(value.as_slice());
        parser.consume_datetime()
    }
}

impl FromHeader for DateTime<UTC> {
    fn from_header(value: String) -> Option<DateTime<UTC>> {
        let dt: Option<DateTime<FixedOffset>> = FromHeader::from_header(value);
        dt.map(|i| i.with_offset(UTC))
    }
}

impl ToHeader for String {
    fn to_header(value: String) -> Option<String> {
        Some(value)
    }
}

impl<'a> ToHeader for &'a str {
    fn to_header(value: &'a str) -> Option<String> {
        Some(value.to_string())
    }
}

/// Represents an RFC 822 Header
#[derive(PartialEq, Eq, Clone, Hash)]
#[unstable]
pub struct Header {
    /// The name of this header
    pub name: String,
    value: String,
}

impl Header {
    /// Creates a new Header for the given `name` and `value`
    #[unstable]
    pub fn new(name: String, value: String) -> Header {
        Header {
            name: name,
            value: value,
        }
    }

    /// Creates a new Header for the given `name` and `value`,
    /// as converted through the `ToHeader` or `ToFoldedHeader` trait.
    ///
    /// Returns None if the value failed to be converted.
    #[unstable]
    pub fn new_with_value<T: ToFoldedHeader>(name: String, value: T) -> Option<Header> {
        let header_len = name.len() + 2;
        ToFoldedHeader::to_folded_header(header_len, value).map(|val| { Header::new(name.clone(), val) })
    }

    /// Get the value represented by this header, as parsed
    /// into whichever type `T`
    #[unstable]
    pub fn get_value<T: FromHeader>(&self) -> Option<T> {
        FromHeader::from_header(self.value.clone())
    }
}

impl fmt::Show for Header {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}: {}", self.name, self.value)
    }
}

#[unstable]
pub struct HeaderIter<'s> {
    iter: SliceIter<'s, Rc<Header>>
}

impl<'s> HeaderIter<'s> {
    #[unstable]
    fn new(iter: SliceIter<'s, Rc<Header>>) -> HeaderIter<'s> {
        HeaderIter {
            iter: iter
        }
    }
}

impl<'s> Iterator<&'s Header> for HeaderIter<'s> {
    fn next(&mut self) -> Option<&'s Header> {
        match self.iter.next() {
            Some(s) => Some(s.deref()),
            None => None,
        }
    }
}

/// A collection of Headers
#[derive(Eq,PartialEq)]
#[unstable]
pub struct HeaderMap {
    // We store headers "twice" inside the HeaderMap.
    //
    // The first is as an ordered list of headers,
    // which is used to iterate over.
    ordered_headers: Vec<Rc<Header>>,
    // The second is as a mapping between header names
    // and all of the headers with that name.
    //
    // This allows quick retrival of a header by name.
    headers: HashMap<String, Vec<Rc<Header>>>,
}

impl HeaderMap {
    #[unstable]
    pub fn new() -> HeaderMap {
        HeaderMap {
            ordered_headers: Vec::new(),
            headers: HashMap::new(),
        }
    }

    /// Adds a header to the collection
    #[unstable]
    pub fn insert(&mut self, header: Header) {
        let header_name = header.name.clone();
        let rc = Rc::new(header);
        // Add to the ordered list of headers
        self.ordered_headers.push(rc.clone());

        // and to the mapping between header names and values.
        match self.headers.entry(header_name) {
            Entry::Occupied(mut entry) => {
                entry.get_mut().push(rc.clone());
            },
            Entry::Vacant(entry) => {
                // There haven't been any headers with this name
                // as of yet, so make a new list and push it in.
                let mut header_list = Vec::new();
                header_list.push(rc.clone());
                entry.set(header_list);
            },
        };
    }

    /// Get an Iterator over the collection of headers.
    #[unstable]
    pub fn iter(&self) -> HeaderIter {
        HeaderIter::new(self.ordered_headers.iter())
    }

    /// Get the last value of the header with `name`
    #[unstable]
    pub fn get(&self, name: String) -> Option<&Header> {
        self.headers.get(&name).map(|headers| { headers.last().unwrap() })
                               .map(|rc| { rc.deref() })
    }

    /// Get the last value of the header with `name`, as a decoded type.
    #[unstable]
    pub fn get_value<T: FromHeader>(&self, name: String) -> Option<T> {
        match self.get(name) {
            Some(ref header) => header.get_value(),
            None => None,
        }
    }

    #[unstable]
    /// Get the number of headers within this map.
    pub fn len(&self) -> uint {
        self.ordered_headers.len()
    }

    #[unstable]
    /// Find a list of headers of `name`, `None` if there
    /// are no headers with that name.
    pub fn find(&self, name: &String) -> Option<Vec<&Header>> {
        let headers_rcs = self.headers.get(name);
        if headers_rcs.is_some() {
            Some(headers_rcs.unwrap().iter().map(|rc| { rc.deref() }).collect())
        } else {
            None
        }
    }
}

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

    use chrono::{
        DateTime,
        FixedOffset,
        Offset,
        UTC,
    };

    static SAMPLE_HEADERS: [(&'static str, &'static str); 4] = [
        ("Test", "Value"),
        ("Test", "Value 2"),
        ("Test-2", "Value 3"),
        ("Test-Multiline", "Foo\nBar"),
    ];

    fn make_sample_headers() -> Vec<Header> {
        SAMPLE_HEADERS.iter().map(|&(name, value)| {
            Header::new(name.to_string(), value.to_string())
        }).collect()
    }

    #[test]
    fn test_header_to_string() {
        let header = Header::new("Test".to_string(), "Value".to_string());
        assert_eq!(header.to_string(), "Test: Value".to_string());
    }

    #[test]
    fn test_string_get_value() {
        struct HeaderTest<'s> {
            input: &'s str,
            result: Option<&'s str>,
        }

        let tests = vec![
            HeaderTest {
                input: "Value",
                result: Some("Value"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Test=20text?=",
                result: Some("Test text"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Multiple?= =?utf-8?b?ZW5jb2Rpbmdz?=",
                result: Some("Multiple encodings"),
            },
            HeaderTest {
                input: "Some things with =?utf-8?b?ZW5jb2Rpbmdz?=, other things without.",
                result: Some("Some things with encodings, other things without."),
            },
            HeaderTest {
                input: "Encoding =?utf-8?q?fail",
                result: Some("Encoding =?utf-8?q?fail"),
            },
        ];

        for test in tests.into_iter() {
            let header = Header::new("Test".to_string(), test.input.to_string());
            let string_value = header.get_value::<String>();
            assert_eq!(string_value, test.result.map(|s| { s.to_string() }));
        }
    }

    #[test]
    fn test_datetime_get_value() {
        let header = Header::new("Date".to_string(), "Wed, 17 Dec 2014 09:35:07 +0100".to_string());
        let dt_value = header.get_value::<DateTime<FixedOffset>>();
        assert_eq!(dt_value, Some(FixedOffset::east(3600).ymd(2014, 12, 17).and_hms(9, 35, 7)));
    }

    #[test]
    fn test_datetime_utc_get_value() {
        let header = Header::new("Date".to_string(), "Wed, 17 Dec 2014 09:35:07 +0100".to_string());
        let dt_value = header.get_value::<DateTime<UTC>>();
        assert_eq!(dt_value, Some(UTC.ymd(2014, 12, 17).and_hms(8, 35, 7)));
    }

    #[test]
    fn test_to_header_string() {
        let header = Header::new_with_value("Test".to_string(), "Value".to_string()).unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_to_header_str() {
        let header = Header::new_with_value("Test".to_string(), "Value").unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_header_map_len() {
        let mut headers = HeaderMap::new();
        for (i, header) in make_sample_headers().into_iter().enumerate() {
            headers.insert(header);
            assert_eq!(headers.len(), i + 1);
        }
    }
    #[test]
    fn test_header_map_iter() {
        let mut headers = HeaderMap::new();
        let mut expected_headers = HashSet::new();
        for header in make_sample_headers().into_iter() {
            headers.insert(header.clone());
            expected_headers.insert(header);
        }

        let mut count = 0u;
        // Ensure all the headers returned are expected
        for header in headers.iter() {
            assert!(expected_headers.contains(header));
            count += 1;
        }
        // And that there is the right number of them
        assert_eq!(count, expected_headers.len());
    }
}
<commit_msg>Imported std::ops::Deref, which left the prelude<commit_after>use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::ops::Deref;
use std::slice::Iter as SliceIter;
use std::rc::Rc;

use chrono::{
    DateTime,
    FixedOffset,
    Offset,
    UTC,
};

use super::rfc2047::decode_rfc2047;
use super::rfc822::Rfc822DateParser;

/// Trait for converting from RFC822 Header values into
/// Rust types.
#[stable]
pub trait FromHeader {
    /// Parse the `value` of the header.
    ///
    /// Returns None if the value failed to be parsed
    fn from_header(value: String) -> Option<Self>;
}

/// Trait for converting from a Rust type into a Header value.
#[stable]
pub trait ToHeader {
    /// Turn the `value` into a String suitable for being used in
    /// a message header.
    ///
    /// Returns None if the value cannot be stringified.
    fn to_header(value: Self) -> Option<String>;
}

/// Trait for converting from a Rust time into a Header value
/// that handles its own folding.
///
/// Be mindful that this trait does not mean that the value will
/// not be folded later, rather that the type returns a value that
/// should not be folded, given that the header value starts so far
/// in to a line.
#[unstable]
pub trait ToFoldedHeader {
    fn to_folded_header(start_pos: uint, value: Self) -> Option<String>;
}

impl<T: ToHeader> ToFoldedHeader for T {
    fn to_folded_header(_: uint, value: T) -> Option<String> {
        // We ignore the start_position because the thing will fold anyway.
        ToHeader::to_header(value)
    }
}

impl FromHeader for String {
    fn from_header(value: String) -> Option<String> {
        #[derive(Show,Copy)]
        enum ParseState {
            Normal(uint),
            SeenEquals(uint),
            SeenQuestion(uint, uint),
        }

        let mut state = ParseState::Normal(0u);
        let mut decoded = String::new();
        let mut pos = 0u;

        let value_slice = value.as_slice();

        while pos < value.len() {
            let ch_range = value_slice.char_range_at(pos);
            let c = ch_range.ch;

            state = match (state, c) {
                (ParseState::SeenQuestion(start_pos, 4), '=') => {
                    // Go to decode if we've seen enough ?
                    let part_decoded = decode_rfc2047(value_slice.slice(start_pos, ch_range.next));
                    let to_push = match part_decoded {
                        Some(ref s) => s.as_slice(),
                        // Decoding failed, push the undecoded string in.
                        None => value_slice.slice(start_pos, pos),
                    };
                    decoded.push_str(to_push);
                    // Revert us to normal state, but starting at the next character.
                    ParseState::Normal(ch_range.next)
                },
                (ParseState::SeenQuestion(start_pos, count), '?') => {
                    ParseState::SeenQuestion(start_pos, count + 1)
                },
                (ParseState::SeenQuestion(start_pos, count), _) => {
                    if count > 4 {
                        // This isn't a RFC2047 sequence, so go back to a normal.
                        ParseState::Normal(start_pos)
                    } else {
                        state
                    }
                }
                (ParseState::SeenEquals(start_pos), '?') => {
                    ParseState::SeenQuestion(start_pos, 1)
                },
                (ParseState::SeenEquals(start_pos), _) => {
                    // This isn't a RFC2047 sequence, so go back to a normal.
                    ParseState::Normal(start_pos)
                }
                (ParseState::Normal(start_pos), '=') => {
                    if start_pos != pos {
                        // Push all up to the =, if there is stuff to push.
                        decoded.push_str(value_slice.slice(start_pos, pos));
                    }
                    ParseState::SeenEquals(pos)
                },
                (ParseState::Normal(_), _) => state,
            };

            pos = ch_range.next;
        }

        // Don't forget to push on whatever we have left
        let last_start = match state {
            ParseState::Normal(start_pos) => start_pos,
            ParseState::SeenEquals(start_pos) => start_pos,
            ParseState::SeenQuestion(start_pos, _) => start_pos,
        };
        decoded.push_str(value_slice.slice_from(last_start));

        Some(decoded)
    }
}

impl FromHeader for DateTime<FixedOffset> {
    fn from_header(value: String) -> Option<DateTime<FixedOffset>> {
        let mut parser = Rfc822DateParser::new(value.as_slice());
        parser.consume_datetime()
    }
}

impl FromHeader for DateTime<UTC> {
    fn from_header(value: String) -> Option<DateTime<UTC>> {
        let dt: Option<DateTime<FixedOffset>> = FromHeader::from_header(value);
        dt.map(|i| i.with_offset(UTC))
    }
}

impl ToHeader for String {
    fn to_header(value: String) -> Option<String> {
        Some(value)
    }
}

impl<'a> ToHeader for &'a str {
    fn to_header(value: &'a str) -> Option<String> {
        Some(value.to_string())
    }
}

/// Represents an RFC 822 Header
#[derive(PartialEq, Eq, Clone, Hash)]
#[unstable]
pub struct Header {
    /// The name of this header
    pub name: String,
    value: String,
}

impl Header {
    /// Creates a new Header for the given `name` and `value`
    #[unstable]
    pub fn new(name: String, value: String) -> Header {
        Header {
            name: name,
            value: value,
        }
    }

    /// Creates a new Header for the given `name` and `value`,
    /// as converted through the `ToHeader` or `ToFoldedHeader` trait.
    ///
    /// Returns None if the value failed to be converted.
    #[unstable]
    pub fn new_with_value<T: ToFoldedHeader>(name: String, value: T) -> Option<Header> {
        let header_len = name.len() + 2;
        ToFoldedHeader::to_folded_header(header_len, value).map(|val| { Header::new(name.clone(), val) })
    }

    /// Get the value represented by this header, as parsed
    /// into whichever type `T`
    #[unstable]
    pub fn get_value<T: FromHeader>(&self) -> Option<T> {
        FromHeader::from_header(self.value.clone())
    }
}

impl fmt::Show for Header {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}: {}", self.name, self.value)
    }
}

#[unstable]
pub struct HeaderIter<'s> {
    iter: SliceIter<'s, Rc<Header>>
}

impl<'s> HeaderIter<'s> {
    #[unstable]
    fn new(iter: SliceIter<'s, Rc<Header>>) -> HeaderIter<'s> {
        HeaderIter {
            iter: iter
        }
    }
}

impl<'s> Iterator<&'s Header> for HeaderIter<'s> {
    fn next(&mut self) -> Option<&'s Header> {
        match self.iter.next() {
            Some(s) => Some(s.deref()),
            None => None,
        }
    }
}

/// A collection of Headers
#[derive(Eq,PartialEq)]
#[unstable]
pub struct HeaderMap {
    // We store headers "twice" inside the HeaderMap.
    //
    // The first is as an ordered list of headers,
    // which is used to iterate over.
    ordered_headers: Vec<Rc<Header>>,
    // The second is as a mapping between header names
    // and all of the headers with that name.
    //
    // This allows quick retrival of a header by name.
    headers: HashMap<String, Vec<Rc<Header>>>,
}

impl HeaderMap {
    #[unstable]
    pub fn new() -> HeaderMap {
        HeaderMap {
            ordered_headers: Vec::new(),
            headers: HashMap::new(),
        }
    }

    /// Adds a header to the collection
    #[unstable]
    pub fn insert(&mut self, header: Header) {
        let header_name = header.name.clone();
        let rc = Rc::new(header);
        // Add to the ordered list of headers
        self.ordered_headers.push(rc.clone());

        // and to the mapping between header names and values.
        match self.headers.entry(header_name) {
            Entry::Occupied(mut entry) => {
                entry.get_mut().push(rc.clone());
            },
            Entry::Vacant(entry) => {
                // There haven't been any headers with this name
                // as of yet, so make a new list and push it in.
                let mut header_list = Vec::new();
                header_list.push(rc.clone());
                entry.set(header_list);
            },
        };
    }

    /// Get an Iterator over the collection of headers.
    #[unstable]
    pub fn iter(&self) -> HeaderIter {
        HeaderIter::new(self.ordered_headers.iter())
    }

    /// Get the last value of the header with `name`
    #[unstable]
    pub fn get(&self, name: String) -> Option<&Header> {
        self.headers.get(&name).map(|headers| { headers.last().unwrap() })
                               .map(|rc| { rc.deref() })
    }

    /// Get the last value of the header with `name`, as a decoded type.
    #[unstable]
    pub fn get_value<T: FromHeader>(&self, name: String) -> Option<T> {
        match self.get(name) {
            Some(ref header) => header.get_value(),
            None => None,
        }
    }

    #[unstable]
    /// Get the number of headers within this map.
    pub fn len(&self) -> uint {
        self.ordered_headers.len()
    }

    #[unstable]
    /// Find a list of headers of `name`, `None` if there
    /// are no headers with that name.
    pub fn find(&self, name: &String) -> Option<Vec<&Header>> {
        let headers_rcs = self.headers.get(name);
        if headers_rcs.is_some() {
            Some(headers_rcs.unwrap().iter().map(|rc| { rc.deref() }).collect())
        } else {
            None
        }
    }
}

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

    use chrono::{
        DateTime,
        FixedOffset,
        Offset,
        UTC,
    };

    static SAMPLE_HEADERS: [(&'static str, &'static str); 4] = [
        ("Test", "Value"),
        ("Test", "Value 2"),
        ("Test-2", "Value 3"),
        ("Test-Multiline", "Foo\nBar"),
    ];

    fn make_sample_headers() -> Vec<Header> {
        SAMPLE_HEADERS.iter().map(|&(name, value)| {
            Header::new(name.to_string(), value.to_string())
        }).collect()
    }

    #[test]
    fn test_header_to_string() {
        let header = Header::new("Test".to_string(), "Value".to_string());
        assert_eq!(header.to_string(), "Test: Value".to_string());
    }

    #[test]
    fn test_string_get_value() {
        struct HeaderTest<'s> {
            input: &'s str,
            result: Option<&'s str>,
        }

        let tests = vec![
            HeaderTest {
                input: "Value",
                result: Some("Value"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Test=20text?=",
                result: Some("Test text"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Multiple?= =?utf-8?b?ZW5jb2Rpbmdz?=",
                result: Some("Multiple encodings"),
            },
            HeaderTest {
                input: "Some things with =?utf-8?b?ZW5jb2Rpbmdz?=, other things without.",
                result: Some("Some things with encodings, other things without."),
            },
            HeaderTest {
                input: "Encoding =?utf-8?q?fail",
                result: Some("Encoding =?utf-8?q?fail"),
            },
        ];

        for test in tests.into_iter() {
            let header = Header::new("Test".to_string(), test.input.to_string());
            let string_value = header.get_value::<String>();
            assert_eq!(string_value, test.result.map(|s| { s.to_string() }));
        }
    }

    #[test]
    fn test_datetime_get_value() {
        let header = Header::new("Date".to_string(), "Wed, 17 Dec 2014 09:35:07 +0100".to_string());
        let dt_value = header.get_value::<DateTime<FixedOffset>>();
        assert_eq!(dt_value, Some(FixedOffset::east(3600).ymd(2014, 12, 17).and_hms(9, 35, 7)));
    }

    #[test]
    fn test_datetime_utc_get_value() {
        let header = Header::new("Date".to_string(), "Wed, 17 Dec 2014 09:35:07 +0100".to_string());
        let dt_value = header.get_value::<DateTime<UTC>>();
        assert_eq!(dt_value, Some(UTC.ymd(2014, 12, 17).and_hms(8, 35, 7)));
    }

    #[test]
    fn test_to_header_string() {
        let header = Header::new_with_value("Test".to_string(), "Value".to_string()).unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_to_header_str() {
        let header = Header::new_with_value("Test".to_string(), "Value").unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_header_map_len() {
        let mut headers = HeaderMap::new();
        for (i, header) in make_sample_headers().into_iter().enumerate() {
            headers.insert(header);
            assert_eq!(headers.len(), i + 1);
        }
    }
    #[test]
    fn test_header_map_iter() {
        let mut headers = HeaderMap::new();
        let mut expected_headers = HashSet::new();
        for header in make_sample_headers().into_iter() {
            headers.insert(header.clone());
            expected_headers.insert(header);
        }

        let mut count = 0u;
        // Ensure all the headers returned are expected
        for header in headers.iter() {
            assert!(expected_headers.contains(header));
            count += 1;
        }
        // And that there is the right number of them
        assert_eq!(count, expected_headers.len());
    }
}
<|endoftext|>


<commit_before>#![allow(unused_variables)]

use super::color::{White, Black};
use super::piece::{mod, Piece, WK, WR, BK, BR};
use super::square::{Square,File,Rank};
use super::moves::{Move,Plies};
use super::pos::{Position,Kingside,Queenside};

pub fn make_move(mut p: Position, m: &Move) -> Position {
    make_move_mut(&mut p, m);
    p
}

pub fn make_move_mut(p: &mut Position, m: &Move) {
    let from = m.from();
    let to = m.to();
    let curr_piece = p.at(from).unwrap();

    if m.is_capture() || m.is_en_passant() || curr_piece.piece_type() == piece::Type::Pawn {
        p.set_ply_count_mut(Plies(0));
    } else {
        let Plies(temp) = p.ply_count();
        p.set_ply_count_mut(Plies(temp+1));
    }

    if m.is_castle() {

        //no en passant
        p.set_en_passant_mut(None);

        let castle_color = match from.rank() {
            Rank(0) => White,
            Rank(7) => Black,
            _ => panic!(),
        };
        let castle_side = match to.file() {
            File(2) => Queenside,
            File(6) => Kingside,
            _ => panic!(),
        };

        let (rook_from, rook_to) = match (castle_color, castle_side) {
            (White, Kingside ) => (Square::new(File(7),Rank(0)), Square::new(File(5),Rank(0))),
            (White, Queenside) => (Square::new(File(0),Rank(0)), Square::new(File(3),Rank(0))),
            (Black, Kingside ) => (Square::new(File(7),Rank(7)), Square::new(File(5),Rank(7))),
            (Black, Queenside) => (Square::new(File(0),Rank(7)), Square::new(File(3),Rank(7))),
        };

        p.remove_at_mut(from);
        p.set_at_mut(to, curr_piece);
        let curr_rook = p.at(rook_from).unwrap();
        p.remove_at_mut(rook_from);
        p.set_at_mut(rook_to, curr_rook);

    } else if m.is_en_passant() {
        //TODO implement en passant
        unimplemented!();
    } else if m.is_promote() {
        //no en passant
        p.set_en_passant_mut(None);

        //change the board and promote
        if m.is_capture() {
            p.remove_at_mut(to);
        }
        p.remove_at_mut(from);
        p.set_at_mut(to, Piece::new(curr_piece.color(), m.promote().unwrap()));
    } else {
        //TODO set enpassant

        //change the board
        if m.is_capture() {
            p.remove_at_mut(to);
        }
        p.remove_at_mut(from);
        p.set_at_mut(to, curr_piece);
    }

    //castling
    if p.at(Square::new(File(4),Rank(0))) != Some(WK) {
        p.set_castle_mut(Queenside, White, false);
        p.set_castle_mut(Kingside , White, false);
    }
    if p.at(Square::new(File(0),Rank(0))) != Some(WR) {
        p.set_castle_mut(Queenside, White, false);
    }
    if p.at(Square::new(File(7),Rank(0))) != Some(WR) {
        p.set_castle_mut(Kingside , White, false);
    }
    if p.at(Square::new(File(4),Rank(7))) != Some(BK) {
        p.set_castle_mut(Queenside, Black, false);
        p.set_castle_mut(Kingside , Black, false);
    }
    if p.at(Square::new(File(0),Rank(7))) != Some(BR) {
        p.set_castle_mut(Queenside, Black, false);
    }
    if p.at(Square::new(File(7),Rank(7))) != Some(BR) {
        p.set_castle_mut(Kingside , Black, false);
    }

    let next_color = p.side_to_move().invert();
    p.set_side_to_move_mut(next_color);
}
<commit_msg>Implement en passant in make_move()<commit_after>#![allow(unused_variables)]

use super::color::{White, Black};
use super::piece::{mod, Piece, WK, WR, BK, BR};
use super::square::{Square,File,Rank};
use super::moves::{Move,Plies};
use super::pos::{Position,Kingside,Queenside};

pub fn make_move(mut p: Position, m: &Move) -> Position {
    make_move_mut(&mut p, m);
    p
}

pub fn make_move_mut(p: &mut Position, m: &Move) {
    let from = m.from();
    let to = m.to();
    let curr_piece = p.at(from).unwrap();

    if m.is_capture() || m.is_en_passant() || curr_piece.piece_type() == piece::Type::Pawn {
        p.set_ply_count_mut(Plies(0));
    } else {
        let Plies(temp) = p.ply_count();
        p.set_ply_count_mut(Plies(temp+1));
    }

    if m.is_castle() {

        //no en passant
        p.set_en_passant_mut(None);

        let castle_color = match from.rank() {
            Rank(0) => White,
            Rank(7) => Black,
            _ => panic!(),
        };
        let castle_side = match to.file() {
            File(2) => Queenside,
            File(6) => Kingside,
            _ => panic!(),
        };

        let (rook_from, rook_to) = match (castle_color, castle_side) {
            (White, Kingside ) => (Square::new(File(7),Rank(0)), Square::new(File(5),Rank(0))),
            (White, Queenside) => (Square::new(File(0),Rank(0)), Square::new(File(3),Rank(0))),
            (Black, Kingside ) => (Square::new(File(7),Rank(7)), Square::new(File(5),Rank(7))),
            (Black, Queenside) => (Square::new(File(0),Rank(7)), Square::new(File(3),Rank(7))),
        };

        p.remove_at_mut(from);
        p.set_at_mut(to, curr_piece);
        let curr_rook = p.at(rook_from).unwrap();
        p.remove_at_mut(rook_from);
        p.set_at_mut(rook_to, curr_rook);

    } else if m.is_en_passant() {

        //no en passant after this
        p.set_en_passant_mut(None);

        let captured = Square::new(to.file(), from.rank());

        p.remove_at_mut(captured);
        p.remove_at_mut(from);
        p.set_at_mut(to, curr_piece);

    } else if m.is_promote() {
        //no en passant
        p.set_en_passant_mut(None);

        //change the board and promote
        if m.is_capture() {
            p.remove_at_mut(to);
        }
        p.remove_at_mut(from);
        p.set_at_mut(to, Piece::new(curr_piece.color(), m.promote().unwrap()));
    } else {
        //TODO set enpassant

        //change the board
        if m.is_capture() {
            p.remove_at_mut(to);
        }
        p.remove_at_mut(from);
        p.set_at_mut(to, curr_piece);
    }

    //castling
    if p.at(Square::new(File(4),Rank(0))) != Some(WK) {
        p.set_castle_mut(Queenside, White, false);
        p.set_castle_mut(Kingside , White, false);
    }
    if p.at(Square::new(File(0),Rank(0))) != Some(WR) {
        p.set_castle_mut(Queenside, White, false);
    }
    if p.at(Square::new(File(7),Rank(0))) != Some(WR) {
        p.set_castle_mut(Kingside , White, false);
    }
    if p.at(Square::new(File(4),Rank(7))) != Some(BK) {
        p.set_castle_mut(Queenside, Black, false);
        p.set_castle_mut(Kingside , Black, false);
    }
    if p.at(Square::new(File(0),Rank(7))) != Some(BR) {
        p.set_castle_mut(Queenside, Black, false);
    }
    if p.at(Square::new(File(7),Rank(7))) != Some(BR) {
        p.set_castle_mut(Kingside , Black, false);
    }

    let next_color = p.side_to_move().invert();
    p.set_side_to_move_mut(next_color);
}
<|endoftext|>


<commit_before>use gfx_app::{ColorFormat, DepthFormat};
use gfx_app::renderer::EncoderQueue;
use gfx;
use terrain;
use character;
use specs;
use std::time::Instant;
use character::character::CharacterSprite;

pub type Delta = f32;

pub struct DrawSystem<D: gfx::Device> {
  render_target_view: gfx::handle::RenderTargetView<D::Resources, ColorFormat>,
  depth_stencil_view: gfx::handle::DepthStencilView<D::Resources, DepthFormat>,
  terrain_system: terrain::DrawSystem<D::Resources>,
  character_system: character::DrawSystem<D::Resources>,
  encoder_queue: EncoderQueue<D>,
  game_time: Instant,
  frames: u32,
  cool_down: f32,
}

impl<D: gfx::Device> DrawSystem<D> {
  pub fn new<F>(factory: &mut F,
                rtv: gfx::handle::RenderTargetView<D::Resources, ColorFormat>,
                dsv: gfx::handle::DepthStencilView<D::Resources, DepthFormat>,
                encoder_queue: EncoderQueue<D>)
                -> DrawSystem<D>
    where F: gfx::Factory<D::Resources>
  {
    DrawSystem {
      render_target_view: rtv.clone(),
      depth_stencil_view: dsv.clone(),
      terrain_system: terrain::DrawSystem::new(factory, rtv.clone(), dsv.clone()),
      character_system: character::DrawSystem::new(factory, rtv.clone(), dsv.clone()),
      encoder_queue,
      game_time: Instant::now(),
      frames: 0,
      cool_down: 1.0
    }
  }
}

impl<D> specs::System<Delta> for DrawSystem<D>
  where D: gfx::Device,
        D::CommandBuffer: Send,
{
  fn run(&mut self, arg: specs::RunArg, delta: Delta) {
    use specs::Join;
    let mut encoder = self.encoder_queue.receiver.recv().unwrap();
    let (mut terrain, mut character, mut sprite) = arg.fetch(|w| {
      if self.cool_down == 0.0 {
        self.cool_down = self.cool_down + 0.07;
      }
      self.cool_down = (self.cool_down - delta).max(0.0);
      (w.read::<terrain::Drawable>(),
       w.write::<character::CharacterDrawable>(),
       w.write::<CharacterSprite>())
    });

    let current_time = Instant::now();
    self.frames = self.frames + 1;
    if cfg!(feature = "dev") && (current_time.duration_since(self.game_time ).as_secs()) >= 1 {
      println!("{:?} ms/frames", 1000.0 / self.frames as f32);
      self.frames = 0;
      self.game_time = Instant::now();
    }

    encoder.clear(&self.render_target_view, [16.0 / 256.0, 14.0 / 256.0, 22.0 / 256.0, 1.0]);
    encoder.clear_depth(&self.depth_stencil_view, 1.0);

    for t in (&mut terrain).join() {
      self.terrain_system.draw(t, &mut encoder);
    }

    for (c, s) in (&mut character, &mut sprite).join() {
      if self.cool_down == 0.0 {
        s.update();
      }
      self.character_system.draw(c, s, &mut encoder);
    }

    if let Err(e) = self.encoder_queue.sender.send(encoder) {
      println!("Disconnected, cannot return encoder to mpsc: {}", e);
    };
  }
}
<commit_msg>Use gray background color<commit_after>use gfx_app::{ColorFormat, DepthFormat};
use gfx_app::renderer::EncoderQueue;
use gfx;
use terrain;
use character;
use specs;
use std::time::Instant;
use character::character::CharacterSprite;

pub type Delta = f32;

pub struct DrawSystem<D: gfx::Device> {
  render_target_view: gfx::handle::RenderTargetView<D::Resources, ColorFormat>,
  depth_stencil_view: gfx::handle::DepthStencilView<D::Resources, DepthFormat>,
  terrain_system: terrain::DrawSystem<D::Resources>,
  character_system: character::DrawSystem<D::Resources>,
  encoder_queue: EncoderQueue<D>,
  game_time: Instant,
  frames: u32,
  cool_down: f32,
}

impl<D: gfx::Device> DrawSystem<D> {
  pub fn new<F>(factory: &mut F,
                rtv: gfx::handle::RenderTargetView<D::Resources, ColorFormat>,
                dsv: gfx::handle::DepthStencilView<D::Resources, DepthFormat>,
                encoder_queue: EncoderQueue<D>)
                -> DrawSystem<D>
    where F: gfx::Factory<D::Resources>
  {
    DrawSystem {
      render_target_view: rtv.clone(),
      depth_stencil_view: dsv.clone(),
      terrain_system: terrain::DrawSystem::new(factory, rtv.clone(), dsv.clone()),
      character_system: character::DrawSystem::new(factory, rtv.clone(), dsv.clone()),
      encoder_queue,
      game_time: Instant::now(),
      frames: 0,
      cool_down: 1.0
    }
  }
}

impl<D> specs::System<Delta> for DrawSystem<D>
  where D: gfx::Device,
        D::CommandBuffer: Send,
{
  fn run(&mut self, arg: specs::RunArg, delta: Delta) {
    use specs::Join;
    let mut encoder = self.encoder_queue.receiver.recv().unwrap();
    let (mut terrain, mut character, mut sprite) = arg.fetch(|w| {
      if self.cool_down == 0.0 {
        self.cool_down = self.cool_down + 0.07;
      }
      self.cool_down = (self.cool_down - delta).max(0.0);
      (w.read::<terrain::Drawable>(),
       w.write::<character::CharacterDrawable>(),
       w.write::<CharacterSprite>())
    });

    let current_time = Instant::now();
    self.frames = self.frames + 1;
    if cfg!(feature = "dev") && (current_time.duration_since(self.game_time ).as_secs()) >= 1 {
      println!("{:?} ms/frames", 1000.0 / self.frames as f32);
      self.frames = 0;
      self.game_time = Instant::now();
    }

    encoder.clear(&self.render_target_view, [16.0 / 256.0, 16.0 / 256.0, 20.0 / 256.0, 1.0]);
    encoder.clear_depth(&self.depth_stencil_view, 1.0);

    for t in (&mut terrain).join() {
      self.terrain_system.draw(t, &mut encoder);
    }

    for (c, s) in (&mut character, &mut sprite).join() {
      if self.cool_down == 0.0 {
        s.update();
      }
      self.character_system.draw(c, s, &mut encoder);
    }

    if let Err(e) = self.encoder_queue.sender.send(encoder) {
      println!("Disconnected, cannot return encoder to mpsc: {}", e);
    };
  }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::path;

#[rustfmt::skip]
use ::image;

use crate::context::{Context, DebugId};
use crate::error::GameError;
use crate::error::GameResult;
use crate::filesystem;
use crate::graphics;
use crate::graphics::shader::*;
use crate::graphics::*;

/// Generic in-GPU-memory image data available to be drawn on the screen.
/// You probably just want to look at the `Image` type.
#[derive(Clone, PartialEq)]
pub struct ImageGeneric<B>
where
    B: BackendSpec,
{
    pub(crate) texture: gfx::handle::RawShaderResourceView<B::Resources>,
    pub(crate) texture_handle: gfx::handle::RawTexture<B::Resources>,
    pub(crate) sampler_info: gfx::texture::SamplerInfo,
    pub(crate) blend_mode: Option<BlendMode>,
    pub(crate) width: u16,
    pub(crate) height: u16,

    pub(crate) debug_id: DebugId,
}

impl<B> ImageGeneric<B>
where
    B: BackendSpec,
{
    /// A helper function that just takes a factory directly so we can make an image
    /// without needing the full context object, so we can create an Image while still
    /// creating the GraphicsContext.
    pub(crate) fn make_raw(
        factory: &mut <B as BackendSpec>::Factory,
        sampler_info: &texture::SamplerInfo,
        width: u16,
        height: u16,
        rgba: &[u8],
        color_format: gfx::format::Format,
        debug_id: DebugId,
    ) -> GameResult<Self> {
        if width == 0 || height == 0 {
            let msg = format!(
                "Tried to create a texture of size {}x{}, each dimension must
                be >0",
                width, height
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        // Check for overflow, which might happen on 32-bit systems.
        // Textures can be max u16*u16, pixels, but then have 4 bytes per pixel.
        let uwidth = width as usize;
        let uheight = height as usize;
        let expected_bytes = uwidth
            .checked_mul(uheight)
            .and_then(|size| size.checked_mul(4))
            .ok_or_else(|| {
                let msg = format!(
                    "Integer overflow in Image::make_raw, image size: {} {}",
                    uwidth, uheight
                );
                GameError::ResourceLoadError(msg)
            })?;
        if expected_bytes != rgba.len() {
            let msg = format!(
                "Tried to create a texture of size {}x{}, but gave {} bytes of data (expected {})",
                width,
                height,
                rgba.len(),
                expected_bytes
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        let kind = gfx::texture::Kind::D2(width, height, gfx::texture::AaMode::Single);
        use gfx::memory::Bind;
        let gfx::format::Format(surface_format, channel_type) = color_format;
        let texinfo = gfx::texture::Info {
            kind,
            levels: 1,
            format: surface_format,
            bind: Bind::SHADER_RESOURCE
                | Bind::RENDER_TARGET
                | Bind::TRANSFER_SRC
                | Bind::TRANSFER_DST,
            usage: gfx::memory::Usage::Dynamic,
        };
        let raw_tex = factory.create_texture_raw(
            texinfo,
            Some(channel_type),
            Some((&[rgba], gfx::texture::Mipmap::Provided)),
        )?;
        let resource_desc = gfx::texture::ResourceDesc {
            channel: channel_type,
            layer: None,
            min: 0,
            max: raw_tex.get_info().levels - 1,
            swizzle: gfx::format::Swizzle::new(),
        };
        let raw_view = factory.view_texture_as_shader_resource_raw(&raw_tex, resource_desc)?;
        Ok(Self {
            texture: raw_view,
            texture_handle: raw_tex,
            sampler_info: *sampler_info,
            blend_mode: None,
            width,
            height,
            debug_id,
        })
    }

    /// A helper function to get the raw gfx texture handle
    pub fn get_raw_texture_handle(&self) -> gfx::handle::RawTexture<B::Resources> {
        self.texture_handle.clone()
    }

    /// A helper function to get the raw gfx texture view
    pub fn get_raw_texture_view(&self) -> gfx::handle::RawShaderResourceView<B::Resources> {
        self.texture.clone()
    }
}

/// In-GPU-memory image data available to be drawn on the screen,
/// using the OpenGL backend.
///
/// Under the hood this is just an `Arc`'ed texture handle and
/// some metadata, so cloning it is fairly cheap; it doesn't
/// make another copy of the underlying image data.
pub type Image = ImageGeneric<GlBackendSpec>;

/// The supported formats for saving an image.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ImageFormat {
    /// .png image format (defaults to RGBA with 8-bit channels.)
    Png,
}

impl Image {
    /// Load a new image from the file at the given path. The documentation for the
    /// [`filesystem`](../filesystem/index.html) module explains how the path must be specified.
    pub fn new<P: AsRef<path::Path>>(context: &mut Context, path: P) -> GameResult<Self> {
        let mut buf = Vec::new();
        let mut reader = context.filesystem.open(path)?;
        let _ = reader.read_to_end(&mut buf)?;
        Self::from_bytes(context, &buf)
    }

    /// Creates a new `Image` from the given buffer, which should contain an image encoded
    /// in a supported image file format.
    pub fn from_bytes(context: &mut Context, bytes: &[u8]) -> GameResult<Self> {
        let img = image::load_from_memory(&bytes)?.to_rgba8();
        let (width, height) = img.dimensions();
        Self::from_rgba8(context, width as u16, height as u16, &img)
    }

    /// Creates a new `Image` from the given buffer of `u8` RGBA values.
    ///
    /// The pixel layout is row-major.  That is,
    /// the first 4 `u8` values make the top-left pixel in the `Image`, the
    /// next 4 make the next pixel in the same row, and so on to the end of
    /// the row.  The next `width * 4` values make up the second row, and so
    /// on.
    pub fn from_rgba8(
        context: &mut Context,
        width: u16,
        height: u16,
        rgba: &[u8],
    ) -> GameResult<Self> {
        let debug_id = DebugId::get(context);
        let color_format = context.gfx_context.color_format();
        Self::make_raw(
            &mut *context.gfx_context.factory,
            &context.gfx_context.default_sampler_info,
            width,
            height,
            rgba,
            color_format,
            debug_id,
        )
    }

    /// Dumps the `Image`'s data to a `Vec` of `u8` RGBA values.
    pub fn to_rgba8(&self, ctx: &mut Context) -> GameResult<Vec<u8>> {
        use gfx::memory::Typed;
        use gfx::traits::FactoryExt;

        let gfx = &mut ctx.gfx_context;
        let w = self.width;
        let h = self.height;

        // Note: In the GFX example, the download buffer is created ahead of time
        // and updated on screen resize events. This may be preferable, but then
        // the buffer also needs to be updated when we switch to/from a canvas.
        // Unsure of the performance impact of creating this as it is needed.
        // Probably okay for now though, since this probably won't be a super
        // common operation.
        let dl_buffer = gfx
            .factory
            .create_download_buffer::<[u8; 4]>(w as usize * h as usize)?;

        let factory = &mut *gfx.factory;
        let mut local_encoder = GlBackendSpec::encoder(factory);

        local_encoder.copy_texture_to_buffer_raw(
            &self.texture_handle,
            None,
            gfx::texture::RawImageInfo {
                xoffset: 0,
                yoffset: 0,
                zoffset: 0,
                width: w as u16,
                height: h as u16,
                depth: 0,
                format: gfx.color_format(),
                mipmap: 0,
            },
            dl_buffer.raw(),
            0,
        )?;
        local_encoder.flush(&mut *gfx.device);

        let reader = gfx.factory.read_mapping(&dl_buffer)?;

        // intermediary buffer to avoid casting.
        let mut data = Vec::with_capacity(self.width as usize * self.height as usize * 4);
        // Assuming OpenGL backend whose typical readback option (glReadPixels) has origin at bottom left.
        // Image formats on the other hand usually deal with top right.
        for y in (0..self.height as usize).rev() {
            data.extend(
                reader
                    .iter()
                    .skip(y * self.width as usize)
                    .take(self.width as usize)
                    .flatten(),
            );
        }
        Ok(data)
    }

    /// Encode the `Image` to the given file format and
    /// write it out to the given path.
    ///
    /// See the [`filesystem`](../filesystem/index.html) module docs for where exactly
    /// the file will end up.
    pub fn encode<P: AsRef<path::Path>>(
        &self,
        ctx: &mut Context,
        format: ImageFormat,
        path: P,
    ) -> GameResult {
        use std::io;
        let data = self.to_rgba8(ctx)?;
        let f = filesystem::create(ctx, path)?;
        let writer = &mut io::BufWriter::new(f);
        let color_format = image::ColorType::Rgba8;
        match format {
            ImageFormat::Png => image::png::PngEncoder::new(writer)
                .encode(
                    &data,
                    u32::from(self.width),
                    u32::from(self.height),
                    color_format,
                )
                .map_err(Into::into),
        }
    }

    /// A little helper function that creates a new `Image` that is just
    /// a solid square of the given size and color.  Mainly useful for
    /// debugging.
    pub fn solid(context: &mut Context, size: u16, color: Color) -> GameResult<Self> {
        let (r, g, b, a) = color.into();
        let pixel_array: [u8; 4] = [r, g, b, a];
        let size_squared = size as usize * size as usize;
        let mut buffer = Vec::with_capacity(size_squared);
        for _i in 0..size_squared {
            buffer.extend(&pixel_array[..]);
        }
        Image::from_rgba8(context, size, size, &buffer)
    }

    /// Return the width of the image.
    pub fn width(&self) -> u16 {
        self.width
    }

    /// Return the height of the image.
    pub fn height(&self) -> u16 {
        self.height
    }

    /// Get the filter mode for the image.
    pub fn filter(&self) -> FilterMode {
        self.sampler_info.filter.into()
    }

    /// Set the filter mode for the image.
    pub fn set_filter(&mut self, mode: FilterMode) {
        self.sampler_info.filter = mode.into();
    }

    /// Returns the dimensions of the image.
    pub fn dimensions(&self) -> Rect {
        Rect::new(0.0, 0.0, f32::from(self.width()), f32::from(self.height()))
    }

    /// Gets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn wrap(&self) -> (WrapMode, WrapMode) {
        (self.sampler_info.wrap_mode.0, self.sampler_info.wrap_mode.1)
    }

    /// Sets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn set_wrap(&mut self, wrap_x: WrapMode, wrap_y: WrapMode) {
        self.sampler_info.wrap_mode.0 = wrap_x;
        self.sampler_info.wrap_mode.1 = wrap_y;
    }
}

impl fmt::Debug for Image {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "<Image: {}x{}, {:p}, texture address {:p}, sampler: {:?}>",
            self.width(),
            self.height(),
            self,
            &self.texture,
            &self.sampler_info
        )
    }
}

impl Drawable for Image {
    fn draw(&self, ctx: &mut Context, param: DrawParam) -> GameResult {
        self.debug_id.assert(ctx);

        let gfx = &mut ctx.gfx_context;
        let src_width = param.src.w;
        let src_height = param.src.h;
        // We have to mess with the scale to make everything
        // be its-unit-size-in-pixels.
        let new_param = match param.trans {
            Transform::Values { scale, .. } => {
                let new_scale = mint::Vector2 {
                    x: scale.x * src_width * f32::from(self.width),
                    y: scale.y * src_height * f32::from(self.height),
                };
                param.scale(new_scale)
            }
            Transform::Matrix(_) => param,
        };

        gfx.update_instance_properties(new_param)?;
        let sampler = gfx
            .samplers
            .get_or_insert(self.sampler_info, gfx.factory.as_mut());
        gfx.data.vbuf = gfx.quad_vertex_buffer.clone();
        let typed_thingy = gfx
            .backend_spec
            .raw_to_typed_shader_resource(self.texture.clone());
        gfx.data.tex = (typed_thingy, sampler);
        let previous_mode: Option<BlendMode> = if let Some(mode) = self.blend_mode {
            let current_mode = gfx.blend_mode();
            if current_mode != mode {
                gfx.set_blend_mode(mode)?;
                Some(current_mode)
            } else {
                None
            }
        } else {
            None
        };

        gfx.draw(None)?;
        if let Some(mode) = previous_mode {
            gfx.set_blend_mode(mode)?;
        }
        Ok(())
    }

    fn dimensions(&self, _: &mut Context) -> Option<graphics::Rect> {
        Some(self.dimensions())
    }

    fn set_blend_mode(&mut self, mode: Option<BlendMode>) {
        self.blend_mode = mode;
    }

    fn blend_mode(&self) -> Option<BlendMode> {
        self.blend_mode
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ContextBuilder;
    #[test]
    fn test_invalid_image_size() {
        let (ctx, _) = &mut ContextBuilder::new("unittest", "unittest").build().unwrap();
        let _i = assert!(Image::from_rgba8(ctx, 0, 0, &[]).is_err());
        let _i = assert!(Image::from_rgba8(ctx, 3432, 432, &[]).is_err());
        let _i = Image::from_rgba8(ctx, 2, 2, &[99; 16]).unwrap();
    }
}
<commit_msg>Scale Image correctly when using Transform::Matrix<commit_after>use std::io::Read;
use std::path;

#[rustfmt::skip]
use ::image;

use crate::context::{Context, DebugId};
use crate::error::GameError;
use crate::error::GameResult;
use crate::filesystem;
use crate::graphics;
use crate::graphics::shader::*;
use crate::graphics::*;

/// Generic in-GPU-memory image data available to be drawn on the screen.
/// You probably just want to look at the `Image` type.
#[derive(Clone, PartialEq)]
pub struct ImageGeneric<B>
where
    B: BackendSpec,
{
    pub(crate) texture: gfx::handle::RawShaderResourceView<B::Resources>,
    pub(crate) texture_handle: gfx::handle::RawTexture<B::Resources>,
    pub(crate) sampler_info: gfx::texture::SamplerInfo,
    pub(crate) blend_mode: Option<BlendMode>,
    pub(crate) width: u16,
    pub(crate) height: u16,

    pub(crate) debug_id: DebugId,
}

impl<B> ImageGeneric<B>
where
    B: BackendSpec,
{
    /// A helper function that just takes a factory directly so we can make an image
    /// without needing the full context object, so we can create an Image while still
    /// creating the GraphicsContext.
    pub(crate) fn make_raw(
        factory: &mut <B as BackendSpec>::Factory,
        sampler_info: &texture::SamplerInfo,
        width: u16,
        height: u16,
        rgba: &[u8],
        color_format: gfx::format::Format,
        debug_id: DebugId,
    ) -> GameResult<Self> {
        if width == 0 || height == 0 {
            let msg = format!(
                "Tried to create a texture of size {}x{}, each dimension must
                be >0",
                width, height
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        // Check for overflow, which might happen on 32-bit systems.
        // Textures can be max u16*u16, pixels, but then have 4 bytes per pixel.
        let uwidth = width as usize;
        let uheight = height as usize;
        let expected_bytes = uwidth
            .checked_mul(uheight)
            .and_then(|size| size.checked_mul(4))
            .ok_or_else(|| {
                let msg = format!(
                    "Integer overflow in Image::make_raw, image size: {} {}",
                    uwidth, uheight
                );
                GameError::ResourceLoadError(msg)
            })?;
        if expected_bytes != rgba.len() {
            let msg = format!(
                "Tried to create a texture of size {}x{}, but gave {} bytes of data (expected {})",
                width,
                height,
                rgba.len(),
                expected_bytes
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        let kind = gfx::texture::Kind::D2(width, height, gfx::texture::AaMode::Single);
        use gfx::memory::Bind;
        let gfx::format::Format(surface_format, channel_type) = color_format;
        let texinfo = gfx::texture::Info {
            kind,
            levels: 1,
            format: surface_format,
            bind: Bind::SHADER_RESOURCE
                | Bind::RENDER_TARGET
                | Bind::TRANSFER_SRC
                | Bind::TRANSFER_DST,
            usage: gfx::memory::Usage::Dynamic,
        };
        let raw_tex = factory.create_texture_raw(
            texinfo,
            Some(channel_type),
            Some((&[rgba], gfx::texture::Mipmap::Provided)),
        )?;
        let resource_desc = gfx::texture::ResourceDesc {
            channel: channel_type,
            layer: None,
            min: 0,
            max: raw_tex.get_info().levels - 1,
            swizzle: gfx::format::Swizzle::new(),
        };
        let raw_view = factory.view_texture_as_shader_resource_raw(&raw_tex, resource_desc)?;
        Ok(Self {
            texture: raw_view,
            texture_handle: raw_tex,
            sampler_info: *sampler_info,
            blend_mode: None,
            width,
            height,
            debug_id,
        })
    }

    /// A helper function to get the raw gfx texture handle
    pub fn get_raw_texture_handle(&self) -> gfx::handle::RawTexture<B::Resources> {
        self.texture_handle.clone()
    }

    /// A helper function to get the raw gfx texture view
    pub fn get_raw_texture_view(&self) -> gfx::handle::RawShaderResourceView<B::Resources> {
        self.texture.clone()
    }
}

/// In-GPU-memory image data available to be drawn on the screen,
/// using the OpenGL backend.
///
/// Under the hood this is just an `Arc`'ed texture handle and
/// some metadata, so cloning it is fairly cheap; it doesn't
/// make another copy of the underlying image data.
pub type Image = ImageGeneric<GlBackendSpec>;

/// The supported formats for saving an image.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ImageFormat {
    /// .png image format (defaults to RGBA with 8-bit channels.)
    Png,
}

impl Image {
    /// Load a new image from the file at the given path. The documentation for the
    /// [`filesystem`](../filesystem/index.html) module explains how the path must be specified.
    pub fn new<P: AsRef<path::Path>>(context: &mut Context, path: P) -> GameResult<Self> {
        let mut buf = Vec::new();
        let mut reader = context.filesystem.open(path)?;
        let _ = reader.read_to_end(&mut buf)?;
        Self::from_bytes(context, &buf)
    }

    /// Creates a new `Image` from the given buffer, which should contain an image encoded
    /// in a supported image file format.
    pub fn from_bytes(context: &mut Context, bytes: &[u8]) -> GameResult<Self> {
        let img = image::load_from_memory(&bytes)?.to_rgba8();
        let (width, height) = img.dimensions();
        Self::from_rgba8(context, width as u16, height as u16, &img)
    }

    /// Creates a new `Image` from the given buffer of `u8` RGBA values.
    ///
    /// The pixel layout is row-major.  That is,
    /// the first 4 `u8` values make the top-left pixel in the `Image`, the
    /// next 4 make the next pixel in the same row, and so on to the end of
    /// the row.  The next `width * 4` values make up the second row, and so
    /// on.
    pub fn from_rgba8(
        context: &mut Context,
        width: u16,
        height: u16,
        rgba: &[u8],
    ) -> GameResult<Self> {
        let debug_id = DebugId::get(context);
        let color_format = context.gfx_context.color_format();
        Self::make_raw(
            &mut *context.gfx_context.factory,
            &context.gfx_context.default_sampler_info,
            width,
            height,
            rgba,
            color_format,
            debug_id,
        )
    }

    /// Dumps the `Image`'s data to a `Vec` of `u8` RGBA values.
    pub fn to_rgba8(&self, ctx: &mut Context) -> GameResult<Vec<u8>> {
        use gfx::memory::Typed;
        use gfx::traits::FactoryExt;

        let gfx = &mut ctx.gfx_context;
        let w = self.width;
        let h = self.height;

        // Note: In the GFX example, the download buffer is created ahead of time
        // and updated on screen resize events. This may be preferable, but then
        // the buffer also needs to be updated when we switch to/from a canvas.
        // Unsure of the performance impact of creating this as it is needed.
        // Probably okay for now though, since this probably won't be a super
        // common operation.
        let dl_buffer = gfx
            .factory
            .create_download_buffer::<[u8; 4]>(w as usize * h as usize)?;

        let factory = &mut *gfx.factory;
        let mut local_encoder = GlBackendSpec::encoder(factory);

        local_encoder.copy_texture_to_buffer_raw(
            &self.texture_handle,
            None,
            gfx::texture::RawImageInfo {
                xoffset: 0,
                yoffset: 0,
                zoffset: 0,
                width: w as u16,
                height: h as u16,
                depth: 0,
                format: gfx.color_format(),
                mipmap: 0,
            },
            dl_buffer.raw(),
            0,
        )?;
        local_encoder.flush(&mut *gfx.device);

        let reader = gfx.factory.read_mapping(&dl_buffer)?;

        // intermediary buffer to avoid casting.
        let mut data = Vec::with_capacity(self.width as usize * self.height as usize * 4);
        // Assuming OpenGL backend whose typical readback option (glReadPixels) has origin at bottom left.
        // Image formats on the other hand usually deal with top right.
        for y in (0..self.height as usize).rev() {
            data.extend(
                reader
                    .iter()
                    .skip(y * self.width as usize)
                    .take(self.width as usize)
                    .flatten(),
            );
        }
        Ok(data)
    }

    /// Encode the `Image` to the given file format and
    /// write it out to the given path.
    ///
    /// See the [`filesystem`](../filesystem/index.html) module docs for where exactly
    /// the file will end up.
    pub fn encode<P: AsRef<path::Path>>(
        &self,
        ctx: &mut Context,
        format: ImageFormat,
        path: P,
    ) -> GameResult {
        use std::io;
        let data = self.to_rgba8(ctx)?;
        let f = filesystem::create(ctx, path)?;
        let writer = &mut io::BufWriter::new(f);
        let color_format = image::ColorType::Rgba8;
        match format {
            ImageFormat::Png => image::png::PngEncoder::new(writer)
                .encode(
                    &data,
                    u32::from(self.width),
                    u32::from(self.height),
                    color_format,
                )
                .map_err(Into::into),
        }
    }

    /// A little helper function that creates a new `Image` that is just
    /// a solid square of the given size and color.  Mainly useful for
    /// debugging.
    pub fn solid(context: &mut Context, size: u16, color: Color) -> GameResult<Self> {
        let (r, g, b, a) = color.into();
        let pixel_array: [u8; 4] = [r, g, b, a];
        let size_squared = size as usize * size as usize;
        let mut buffer = Vec::with_capacity(size_squared);
        for _i in 0..size_squared {
            buffer.extend(&pixel_array[..]);
        }
        Image::from_rgba8(context, size, size, &buffer)
    }

    /// Return the width of the image.
    pub fn width(&self) -> u16 {
        self.width
    }

    /// Return the height of the image.
    pub fn height(&self) -> u16 {
        self.height
    }

    /// Get the filter mode for the image.
    pub fn filter(&self) -> FilterMode {
        self.sampler_info.filter.into()
    }

    /// Set the filter mode for the image.
    pub fn set_filter(&mut self, mode: FilterMode) {
        self.sampler_info.filter = mode.into();
    }

    /// Returns the dimensions of the image.
    pub fn dimensions(&self) -> Rect {
        Rect::new(0.0, 0.0, f32::from(self.width()), f32::from(self.height()))
    }

    /// Gets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn wrap(&self) -> (WrapMode, WrapMode) {
        (self.sampler_info.wrap_mode.0, self.sampler_info.wrap_mode.1)
    }

    /// Sets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn set_wrap(&mut self, wrap_x: WrapMode, wrap_y: WrapMode) {
        self.sampler_info.wrap_mode.0 = wrap_x;
        self.sampler_info.wrap_mode.1 = wrap_y;
    }
}

impl fmt::Debug for Image {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "<Image: {}x{}, {:p}, texture address {:p}, sampler: {:?}>",
            self.width(),
            self.height(),
            self,
            &self.texture,
            &self.sampler_info
        )
    }
}

impl Drawable for Image {
    fn draw(&self, ctx: &mut Context, param: DrawParam) -> GameResult {
        self.debug_id.assert(ctx);

        let gfx = &mut ctx.gfx_context;
        let src_width = param.src.w;
        let src_height = param.src.h;
        // We have to mess with the scale to make everything
        // be its-unit-size-in-pixels.
        let scale_x = src_width * f32::from(self.width);
        let scale_y = src_height * f32::from(self.height);
        let new_param = match param.trans {
            Transform::Values { scale, .. } => {
                let new_scale = mint::Vector2 {
                    x: scale.x * scale_x,
                    y: scale.y * scale_y,
                };
                param.scale(new_scale)
            }
            Transform::Matrix(m) => {
                let m_scale = Matrix4::from_scale((scale_x, scale_y, 1.0).into());
                param.transform(Matrix4::from(m) * m_scale)
            }
        };

        gfx.update_instance_properties(new_param)?;
        let sampler = gfx
            .samplers
            .get_or_insert(self.sampler_info, gfx.factory.as_mut());
        gfx.data.vbuf = gfx.quad_vertex_buffer.clone();
        let typed_thingy = gfx
            .backend_spec
            .raw_to_typed_shader_resource(self.texture.clone());
        gfx.data.tex = (typed_thingy, sampler);
        let previous_mode: Option<BlendMode> = if let Some(mode) = self.blend_mode {
            let current_mode = gfx.blend_mode();
            if current_mode != mode {
                gfx.set_blend_mode(mode)?;
                Some(current_mode)
            } else {
                None
            }
        } else {
            None
        };

        gfx.draw(None)?;
        if let Some(mode) = previous_mode {
            gfx.set_blend_mode(mode)?;
        }
        Ok(())
    }

    fn dimensions(&self, _: &mut Context) -> Option<graphics::Rect> {
        Some(self.dimensions())
    }

    fn set_blend_mode(&mut self, mode: Option<BlendMode>) {
        self.blend_mode = mode;
    }

    fn blend_mode(&self) -> Option<BlendMode> {
        self.blend_mode
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ContextBuilder;
    #[test]
    fn test_invalid_image_size() {
        let (ctx, _) = &mut ContextBuilder::new("unittest", "unittest").build().unwrap();
        let _i = assert!(Image::from_rgba8(ctx, 0, 0, &[]).is_err());
        let _i = assert!(Image::from_rgba8(ctx, 3432, 432, &[]).is_err());
        let _i = Image::from_rgba8(ctx, 2, 2, &[99; 16]).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry::*;

use util::*;

/// A Postings map (doc id => highlights) for a single term.
/// Records which Documents contain the term, and at which locations in the documents.
pub type PostingsMap = BTreeMap<String, Vec<(usize, usize)>>;

/// An extension trait for iterables over PostingsMaps
/// that enables the computation of the union of Postingsmaps
pub trait PostingsMerge {
    /// Computes the map containing the union of the maps in self
    fn merge_postings(self) -> PostingsMap;
}

impl<'a, Iter> PostingsMerge for Iter
    where Iter: IntoIterator,
          Iter::Item: Borrow<PostingsMap> {
    fn merge_postings(self) -> PostingsMap {
        let mut map = PostingsMap::new();
        for tree in self {
            let tree = tree.borrow();
            for (doc_id, highlights) in tree.iter() {
                match map.entry(doc_id.clone()) {
                    Vacant(entry) => {
                        entry.insert(highlights.clone());
                    }
                    Occupied(mut entry) => {
                        let entry = entry.get_mut();
                        let mut last_search = 0;
                        for &highlight in highlights {
                            last_search = entry.search_coalesce(last_search, highlight);
                        }
                    }
                }
            }
        }
        map
    }
}

/// An extension trait for slices of PostingsMaps,
/// that enables the computation of the intersection
/// of PostingsMaps
pub trait PostingsIntersect {
    /// Computes the map containing the intersection of the the maps in self
    fn intersect_postings(self) -> PostingsMap;
}

impl<'a> PostingsIntersect for &'a [PostingsMap] {
    fn intersect_postings(self) -> PostingsMap {
        match self {
            [] => PostingsMap::new(),
            [ref posting] => posting.clone(),
            [ref posting0, rest..] => {
                self.intersection()
                    .map(|doc_id| {
                        let mut highlights = posting0[doc_id].clone();
                        for posting in rest {
                            let mut last_search = 0;
                            for &highlight in &posting[doc_id] {
                                last_search = highlights.search_coalesce(last_search, highlight);
                            }
                        }
                        (doc_id.clone(), highlights)
                    })
                    .collect()
            }
        }
    }
}

<commit_msg>Remove unnecessary .iter()<commit_after>use std::borrow::Borrow;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry::*;

use util::*;

/// A Postings map (doc id => highlights) for a single term.
/// Records which Documents contain the term, and at which locations in the documents.
pub type PostingsMap = BTreeMap<String, Vec<(usize, usize)>>;

/// An extension trait for iterables over PostingsMaps
/// that enables the computation of the union of Postingsmaps
pub trait PostingsMerge {
    /// Computes the map containing the union of the maps in self
    fn merge_postings(self) -> PostingsMap;
}

impl<'a, Iter> PostingsMerge for Iter
    where Iter: IntoIterator,
          Iter::Item: Borrow<PostingsMap> {
    fn merge_postings(self) -> PostingsMap {
        let mut map = PostingsMap::new();
        for tree in self {
            for (doc_id, highlights) in tree.borrow() {
                match map.entry(doc_id.clone()) {
                    Vacant(entry) => {
                        entry.insert(highlights.clone());
                    }
                    Occupied(mut entry) => {
                        let entry = entry.get_mut();
                        let mut last_search = 0;
                        for &highlight in highlights {
                            last_search = entry.search_coalesce(last_search, highlight);
                        }
                    }
                }
            }
        }
        map
    }
}

/// An extension trait for slices of PostingsMaps,
/// that enables the computation of the intersection
/// of PostingsMaps
pub trait PostingsIntersect {
    /// Computes the map containing the intersection of the the maps in self
    fn intersect_postings(self) -> PostingsMap;
}

impl<'a> PostingsIntersect for &'a [PostingsMap] {
    fn intersect_postings(self) -> PostingsMap {
        match self {
            [] => PostingsMap::new(),
            [ref posting] => posting.clone(),
            [ref posting0, rest..] => {
                self.intersection()
                    .map(|doc_id| {
                        let mut highlights = posting0[doc_id].clone();
                        for posting in rest {
                            let mut last_search = 0;
                            for &highlight in &posting[doc_id] {
                                last_search = highlights.search_coalesce(last_search, highlight);
                            }
                        }
                        (doc_id.clone(), highlights)
                    })
                    .collect()
            }
        }
    }
}

<|endoftext|>


<commit_before>mod idt;

macro_rules! handler {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(
                          *const ExceptionStackFrame) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

macro_rules! handler_with_error_code {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("pop rsi // pop error code into rsi
                      mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(
                          *const ExceptionStackFrame, u64) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

lazy_static! {
    static ref IDT: idt::Idt = {
        let mut idt = idt::Idt::new();

        idt.set_handler(0, handler!(divide_by_zero_handler));
        idt.set_handler(6, handler!(invalid_opcode_handler));
        idt.set_handler(14, handler_with_error_code!(page_fault_handler));

        idt
    };
}

pub fn init() {
    IDT.load();
}

#[derive(Debug)]
#[repr(C)]
struct ExceptionStackFrame {
    instruction_pointer: u64,
    code_segment: u64,
    cpu_flags: u64,
    stack_pointer: u64,
    stack_segment: u64,
}

use vga_buffer::print_error;

extern "C" fn divide_by_zero_handler(stack_frame: *const ExceptionStackFrame)
    -> !
{
    unsafe {
        print_error(format_args!("EXCEPTION: DIVIDE BY ZERO\n{:#?}",
            *stack_frame));
    }
    loop {}
}

extern "C" fn invalid_opcode_handler(stack_frame: *const ExceptionStackFrame)
    -> !
{
    unsafe {
        print_error(format_args!("EXCEPTION: INVALID OPCODE at {:#x}\n{:#?}",
            (*stack_frame).instruction_pointer, *stack_frame));
    }
    loop {}
}

extern "C" fn page_fault_handler(stack_frame: *const ExceptionStackFrame,
                                 error_code: u64) -> !
{
    unsafe {
        print_error(format_args!(
            "EXCEPTION: PAGE FAULT with error code {:?}\n{:#?}",
            error_code, *stack_frame));
    }
    loop {}
}
<commit_msg>Translate error code and print cr2 register<commit_after>mod idt;

macro_rules! handler {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(
                          *const ExceptionStackFrame) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

macro_rules! handler_with_error_code {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("pop rsi // pop error code into rsi
                      mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(
                          *const ExceptionStackFrame, u64) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

lazy_static! {
    static ref IDT: idt::Idt = {
        let mut idt = idt::Idt::new();

        idt.set_handler(0, handler!(divide_by_zero_handler));
        idt.set_handler(6, handler!(invalid_opcode_handler));
        idt.set_handler(14, handler_with_error_code!(page_fault_handler));

        idt
    };
}

pub fn init() {
    IDT.load();
}

#[derive(Debug)]
#[repr(C)]
struct ExceptionStackFrame {
    instruction_pointer: u64,
    code_segment: u64,
    cpu_flags: u64,
    stack_pointer: u64,
    stack_segment: u64,
}

use vga_buffer::print_error;

extern "C" fn divide_by_zero_handler(stack_frame: *const ExceptionStackFrame)
    -> !
{
    unsafe {
        print_error(format_args!("EXCEPTION: DIVIDE BY ZERO\n{:#?}",
            *stack_frame));
    }
    loop {}
}

extern "C" fn invalid_opcode_handler(stack_frame: *const ExceptionStackFrame)
    -> !
{
    unsafe {
        print_error(format_args!("EXCEPTION: INVALID OPCODE at {:#x}\n{:#?}",
            (*stack_frame).instruction_pointer, *stack_frame));
    }
    loop {}
}

bitflags! {
    flags PageFaultErrorCode: u64 {
        const PROTECTION_VIOLATION = 1 << 0,
        const CAUSED_BY_WRITE = 1 << 1,
        const USER_MODE = 1 << 2,
        const MALFORMED_TABLE = 1 << 3,
        const INSTRUCTION_FETCH = 1 << 4,
    }
}

extern "C" fn page_fault_handler(stack_frame: *const ExceptionStackFrame,
                                 error_code: u64) -> !
{
    use x86::controlregs;
    unsafe {
        print_error(format_args!(
            "EXCEPTION: PAGE FAULT while accessing {:#x}\
            \nerror code: {:?}\n{:#?}",
            controlregs::cr2(),
            PageFaultErrorCode::from_bits(error_code).unwrap(),
            *stack_frame));
    }
    loop {}
}
<|endoftext|>