<commit_before>extern crate itertools;

use std::collections::{HashSet,HashMap};
use self::itertools::Itertools;        // fold_results
use std::io;                           // Error
use std::io::{Read,BufReader,BufRead}; // read_to_string
use std::path::Path;
use std::num;                          // ParseIntError
use std::fmt;                          // Formatter, format!, Display, Debug, write!
use std::error;
use std::fs::File;                     // File, open

// TODO "readme.mk"
// TODO documentation
#[derive(Debug)]
pub enum DFAError {
    MissingStartingState,
    MissingFinalStates,
    IncompleteTransition(usize),
    IllformedTransition(usize),
    DuplicatedTransition(usize),
    Io(io::Error),
    Parse(num::ParseIntError,Option<usize>),
}

impl fmt::Display for DFAError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            DFAError::Io(ref err) => write!(f, "IO error: {}", err),
            DFAError::MissingStartingState => write!(f, "The file is empty or only contains white characters."),
            DFAError::MissingFinalStates => write!(f, "The file does not specify the list of final states."),
            DFAError::IncompleteTransition(ref line) => write!(f, "Line {}: missing the src or the dest state.", line),
            DFAError::IllformedTransition(ref line) => write!(f, "Line {}: too much elements.", line),
            DFAError::DuplicatedTransition(ref line) => write!(f, "Line {}: duplicated transition.", line),
            DFAError::Parse(ref err,ref line) => write!(f, "Parse error on line {:?}: {}", line, err),
        }
    }
}

impl error::Error for DFAError {
    fn description(&self) -> &str {
        match *self {
            DFAError::Io(ref err) => err.description(),
            DFAError::MissingStartingState => "The file is empty or only contains white characters.",
            DFAError::MissingFinalStates => "The file does not specify the list of final states.",
            DFAError::IncompleteTransition(_) => "Missing the src or the dest state.",
            DFAError::IllformedTransition(_) => "Too much elements.",
            DFAError::DuplicatedTransition(_) => "Duplicated transition.",
            DFAError::Parse(ref err,_) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            DFAError::Io(ref err) => Some(err),
            DFAError::Parse(ref err,_) => Some(err),
            _ => None,
        }
    }
}

impl From<io::Error> for DFAError {
    fn from(err: io::Error) -> DFAError {
        DFAError::Io(err)
    }
}

impl From<num::ParseIntError> for DFAError {
    fn from(err: num::ParseIntError) -> DFAError {
        DFAError::Parse(err,None)
    }
}

#[derive(Debug)]
pub struct DFA {
    transitions : HashMap<(char,usize),usize>,
    start       : usize,
    finals      : HashSet<usize>,
}

#[derive(Debug)]
pub struct DFABuilder {
    transitions : HashMap<(char,usize),usize>,
    start       : usize,
    finals      : HashSet<usize>,
}

pub trait DFABuilding {
    fn add_start(mut self, state: usize) -> Self;
    fn add_final(mut self, state: usize) -> Self;
    fn add_transition(mut self, symb: char, src: usize, dest: usize) -> Self;
    fn finalize(self) -> Result<DFA, DFAError>;
}

impl DFABuilder {
    pub fn new() -> Result<DFABuilder,DFAError> {
        Ok(DFABuilder{transitions: HashMap::new(), start: 0, finals: HashSet::new()})
    }
}

impl DFABuilding for Result<DFABuilder,DFAError> {
    fn add_start(self, state: usize) -> Result<DFABuilder,DFAError> {
        self.and_then(|mut dfa| {
            dfa.start = state;
            Ok(dfa)
        })
    }

    fn add_final(self, state: usize) -> Result<DFABuilder,DFAError> {
        self.and_then(|mut dfa| {
            dfa.finals.insert(state);
            Ok(dfa)
        })
    }

    fn add_transition(self, symb: char, src: usize, dest: usize) -> Result<DFABuilder,DFAError> {
        self.and_then(|mut dfa| {
            if dfa.transitions.insert((symb,src), dest).is_some() {
                return Err(DFAError::DuplicatedTransition(0));
            }
            Ok(dfa)
        })
    }

    fn finalize(self) -> Result<DFA, DFAError> {
        self.and_then(|dfa| {
            Ok(DFA{transitions: dfa.transitions, start: dfa.start, finals: dfa.finals})
        })
    }
}

impl DFA {
    pub fn new() -> DFA {
        DFA{transitions: HashMap::new(), start: 0, finals: HashSet::new()}
    }

    fn parse_dfa_error(contents: &str, line: usize) -> Result<usize, DFAError> {
            contents.parse::<usize>()
                    .map_err(|e| DFAError::Parse(e,Some(line)))
    }

    pub fn new_from_file<P: AsRef<Path>>(file_path: P) -> Result<DFA, DFAError> {
        let file = try!(File::open(file_path));
        let file = BufReader::new(file);
        DFA::new_from_lines(&mut file.lines())
    }

    fn new_from_lines(lines : &mut Iterator<Item=io::Result<String>>) -> Result<DFA, DFAError> {
        let mut dfa = DFA::new();
        let mut lines = lines
            .map(|line| {
                line.and_then(|contents| Ok(contents.split('#').nth(0).unwrap().trim().to_owned()))
            })
            .enumerate().map(|(nline,line)| (nline+1,line))
            .filter(|&(_,ref line)| {
                // Mandatory otherwise unwrap will take the ownership of the String
                let line = line.as_ref();
                line.is_err() || !line.unwrap().is_empty()
            });
        // Starting state
        let (nline,line) = try!(lines.next().ok_or(DFAError::MissingStartingState));
        let line = try!(line);
        dfa.start = try!(DFA::parse_dfa_error(&line,nline));
        // Final states
        let (nline,line) = try!(lines.next().ok_or(DFAError::MissingFinalStates));
        let line = try!(line);
        dfa.finals = try!(line
            .split_whitespace()
            .map(|token| DFA::parse_dfa_error(token,nline))
            .fold_results(HashSet::new(), |mut acc, elt| {
                acc.insert(elt);
                acc
            }));
        for (nline,line) in lines {
            let line = try!(line);
            let mut tokens = line.split_whitespace();
            // can't fail because lines iterates over the non-empty line
            let mut symbs = tokens.next().unwrap().chars();
            let symb = symbs.nth(0).unwrap();
            if symbs.next().is_some() {
                return Err(DFAError::IllformedTransition(nline));
            }
            let src = try!(tokens
                .next()
                .ok_or(DFAError::IncompleteTransition(nline))
                .and_then(|contents| DFA::parse_dfa_error(contents,nline)));
            let dest = try!(tokens
                .next()
                .ok_or(DFAError::IncompleteTransition(nline))
                .and_then(|contents| DFA::parse_dfa_error(contents,nline)));
            if tokens.next().is_some() {
                return Err(DFAError::IllformedTransition(nline));
            }
            if dfa.transitions.insert((symb,src), dest).is_some() {
                return Err(DFAError::DuplicatedTransition(nline));
            }
        }
        Ok(dfa)
    }

    pub fn new_from_string(file: &str) -> Result<DFA, DFAError> {
        DFA::new_from_lines(&mut file.lines().map(|line| Ok(line.to_string())))
    }

    // TODO return the position of the first match
    //      maybe create an another function to do that
    pub fn run(&self, input: &str) -> bool {
        let f = input
            .chars()
            .fold(Some(self.start), |state,c| {
                match state {
                    Some(n) => self.transitions.get(&(c,n)).map(|v| *v),
                    None => None,
                }
            });
        match f {
            Some(n) => self.finals.contains(&n),
            None => false
        }
    }
}

impl fmt::Display for DFA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(writeln!(f, "START: {}", self.start));
        try!(writeln!(f, "FINALS:"));
        for fi in self.finals.iter() {
            try!(writeln!(f,"  {}", fi));
        }
        try!(writeln!(f, "TRANSITIONS:"));
        for (tr,d) in self.transitions.iter() {
            let (c,s) = *tr;
            try!(writeln!(f, "  ({},{}) => {}", c, s, d));
        }
        write!(f, "")
    }
}

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

    #[test]
    fn test_dfa() {
        let model =
            "0\n\
             3\n\
             a 0 1\n\
             c 0 3\n\
             b 1 2\n\
             a 2 1\n\
             c 2 3";
        let dfa = DFA::new_from_string(&model).unwrap();
        let samples =
            vec![("ababac", false),
                 ("ababc", true),
                 ("", false),
                 ("abc", true),
                 ("c", true),
                 ("ac", false),
                 ("ababababababababababababababababababababc", true),];

        for (input,expected_result) in samples {
            assert!(dfa.run(input) == expected_result, "input false for: \"{}\"", input);
        }
    }

    #[test]
    fn test_empty_file() {
        let model =
            "";
        match DFA::new_from_string(model) {
            Err(DFAError::MissingStartingState) => assert!(true),
            _ => assert!(false, "Missing state expected."),
        }
    }

    #[test]
    fn test_start_not_a_number() {
        let model =
            "a";
        match DFA::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 1),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_many_starts() {
        let model =
            "0 1\n\
             3\n\
             a 0 1\n\
             c 0 3\n\
             b 1 2\n\
             a 2 1\n\
             c 2 3";
        match DFA::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 1),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_no_finals() {
        let model =
            "1\n\
            ";
        match DFA::new_from_string(model) {
            Err(DFAError::MissingFinalStates) => assert!(true),
            _ => assert!(false, "Missing final states expected."),
        }
    }

    #[test]
    fn test_finals_not_a_number() {
        let model =
            "1\n\
             2 a 3";
        match DFA::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 2),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_no_transistions() {
        let model =
            "0\n\
             3";
        let _dfa = DFA::new_from_string(&model).unwrap();
    }

    #[test]
    fn test_transitions_with_at_least_four_elements() {
        let model =
            "0\n\
             3\n\
             a 0 1 8";
        match DFA::new_from_string(model) {
            Err(DFAError::IllformedTransition(line)) => assert!(line == 3),
            _ => assert!(false, "IllformedTransition expected."),
        }
    }

    #[test]
    #[should_panic]
    fn test_transitions_start_with_at_least_two_chars() {
        let model =
            "0\n\
             3\n\
             ab 2 3";
        let _dfa = DFA::new_from_string(&model).unwrap();
    }

    #[test]
    fn test_transitions_with_src_not_a_number() {
        let model =
            "0\n\
             3\n\
             c b 3";
        match DFA::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 3),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_transitions_with_dest_not_a_number() {
        let model =
            "0\n\
             3\n\
             c 2 b";
        match DFA::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 3),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_duplicated_transition() {
        let model =
            "0\n\
             3\n\
             c 2 3\n\
             c 2 4";
        match DFA::new_from_string(model) {
            Err(DFAError::DuplicatedTransition(line)) => assert!(line == 4),
            _ => assert!(false, "DuplicatedTransition expected."),
        }
    }

    #[test]
    fn test_read_from_fake_file() {
        let file = "fake.txt";
        match DFA::new_from_file(file) {
            Err(DFAError::Io(_)) => assert!(true),
            _ => assert!(false, "Io::Error expected."),
        }
    }
}
<commit_msg>Add DFAReader<commit_after>extern crate itertools;

use std::collections::{HashSet,HashMap};
use self::itertools::Itertools;        // fold_results
use std::io;                           // Error
use std::io::{Read,BufReader,BufRead}; // read_to_string
use std::path::Path;
use std::num;                          // ParseIntError
use std::fmt;                          // Formatter, format!, Display, Debug, write!
use std::error;
use std::fs::File;                     // File, open

// TODO "readme.mk"
// TODO documentation
#[derive(Debug)]
pub enum DFAError {
    MissingStartingState,
    MissingFinalStates,
    IncompleteTransition(usize),
    IllformedTransition(usize),
    DuplicatedTransition(usize),
    Io(io::Error),
    Parse(num::ParseIntError,Option<usize>),
}

impl fmt::Display for DFAError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            DFAError::Io(ref err) => write!(f, "IO error: {}", err),
            DFAError::MissingStartingState => write!(f, "The file is empty or only contains white characters."),
            DFAError::MissingFinalStates => write!(f, "The file does not specify the list of final states."),
            DFAError::IncompleteTransition(ref line) => write!(f, "Line {}: missing the src or the dest state.", line),
            DFAError::IllformedTransition(ref line) => write!(f, "Line {}: too much elements.", line),
            DFAError::DuplicatedTransition(ref line) => write!(f, "Line {}: duplicated transition.", line),
            DFAError::Parse(ref err,ref line) => write!(f, "Parse error on line {:?}: {}", line, err),
        }
    }
}

impl error::Error for DFAError {
    fn description(&self) -> &str {
        match *self {
            DFAError::Io(ref err) => err.description(),
            DFAError::MissingStartingState => "The file is empty or only contains white characters.",
            DFAError::MissingFinalStates => "The file does not specify the list of final states.",
            DFAError::IncompleteTransition(_) => "Missing the src or the dest state.",
            DFAError::IllformedTransition(_) => "Too much elements.",
            DFAError::DuplicatedTransition(_) => "Duplicated transition.",
            DFAError::Parse(ref err,_) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            DFAError::Io(ref err) => Some(err),
            DFAError::Parse(ref err,_) => Some(err),
            _ => None,
        }
    }
}

impl From<io::Error> for DFAError {
    fn from(err: io::Error) -> DFAError {
        DFAError::Io(err)
    }
}

impl From<num::ParseIntError> for DFAError {
    fn from(err: num::ParseIntError) -> DFAError {
        DFAError::Parse(err,None)
    }
}

#[derive(Debug)]
pub struct DFA {
    transitions : HashMap<(char,usize),usize>,
    start       : usize,
    finals      : HashSet<usize>,
}

#[derive(Debug)]
pub struct DFABuilder {
    transitions : HashMap<(char,usize),usize>,
    start       : usize,
    finals      : HashSet<usize>,
}

pub trait DFABuilding {
    fn add_start(mut self, state: usize) -> Self;
    fn add_final(mut self, state: usize) -> Self;
    fn add_transition(mut self, symb: char, src: usize, dest: usize) -> Self;
    fn finalize(self) -> Result<DFA, DFAError>;
}

impl DFABuilder {
    pub fn new() -> Result<DFABuilder,DFAError> {
        Ok(DFABuilder{transitions: HashMap::new(), start: 0, finals: HashSet::new()})
    }
}

impl DFABuilding for Result<DFABuilder,DFAError> {
    fn add_start(self, state: usize) -> Result<DFABuilder,DFAError> {
        self.and_then(|mut dfa| {
            dfa.start = state;
            Ok(dfa)
        })
    }

    fn add_final(self, state: usize) -> Result<DFABuilder,DFAError> {
        self.and_then(|mut dfa| {
            dfa.finals.insert(state);
            Ok(dfa)
        })
    }

    fn add_transition(self, symb: char, src: usize, dest: usize) -> Result<DFABuilder,DFAError> {
        self.and_then(|mut dfa| {
            if dfa.transitions.insert((symb,src), dest).is_some() {
                return Err(DFAError::DuplicatedTransition(0));
            }
            Ok(dfa)
        })
    }

    fn finalize(self) -> Result<DFA, DFAError> {
        self.and_then(|dfa| {
            Ok(DFA{transitions: dfa.transitions, start: dfa.start, finals: dfa.finals})
        })
    }
}
pub struct DFAReader;

impl DFAReader {
    fn parse_dfa_error(contents: &str, line: usize) -> Result<usize, DFAError> {
            contents.parse::<usize>()
                    .map_err(|e| DFAError::Parse(e,Some(line)))
    }

    pub fn new_from_file<P: AsRef<Path>>(file_path: P) -> Result<DFA, DFAError> {
        let file = try!(File::open(file_path));
        let file = BufReader::new(file);
        DFAReader::new_from_lines(&mut file.lines())
    }

    fn new_from_lines(lines : &mut Iterator<Item=io::Result<String>>) -> Result<DFA, DFAError> {
        let mut dfa = DFABuilder::new().finalize().unwrap();
        //let mut dfa = DFABuilder::new();
        let mut lines = lines
            .map(|line| {
                line.and_then(|contents| Ok(contents.split('#').nth(0).unwrap().trim().to_owned()))
            })
            .enumerate().map(|(nline,line)| (nline+1,line))
            .filter(|&(_,ref line)| {
                // Mandatory otherwise unwrap will take the ownership of the String
                let line = line.as_ref();
                line.is_err() || !line.unwrap().is_empty()
            });
        // Starting state
        let (nline,line) = try!(lines.next().ok_or(DFAError::MissingStartingState));
        let line = try!(line);
        dfa.start = try!(DFAReader::parse_dfa_error(&line,nline));
        //let start = try!(DFAReader::parse_dfa_error(&line,nline));
        //try!(dfa.add_start(start).map_err(|e| );
        // Final states
        let (nline,line) = try!(lines.next().ok_or(DFAError::MissingFinalStates));
        let line = try!(line);
        dfa.finals = try!(line
            .split_whitespace()
            .map(|token| DFAReader::parse_dfa_error(token,nline))
            .fold_results(HashSet::new(), |mut acc, elt| {
                acc.insert(elt);
                acc
            }));
        //try!(line
            //.split_whitespace()
            //.map(|token| DFAReader::parse_dfa_error(token,nline))
            //.fold_results(dfa, |mut acc, elt| acc.add_final(elt)));
        //try!(dfa.map_err(|e| ));
        for (nline,line) in lines {
            let line = try!(line);
            let mut tokens = line.split_whitespace();
            // can't fail because lines iterates over the non-empty line
            let mut symbs = tokens.next().unwrap().chars();
            let symb = symbs.nth(0).unwrap();
            if symbs.next().is_some() {
                return Err(DFAError::IllformedTransition(nline));
            }
            let src = try!(tokens
                .next()
                .ok_or(DFAError::IncompleteTransition(nline))
                .and_then(|contents| DFAReader::parse_dfa_error(contents,nline)));
            let dest = try!(tokens
                .next()
                .ok_or(DFAError::IncompleteTransition(nline))
                .and_then(|contents| DFAReader::parse_dfa_error(contents,nline)));
            if tokens.next().is_some() {
                return Err(DFAError::IllformedTransition(nline));
            }
            if dfa.transitions.insert((symb,src), dest).is_some() {
                return Err(DFAError::DuplicatedTransition(nline));
            }
            //try!(dfa.add_transition(symb,src,dest).map_err(|e| ));
        }
        Ok(dfa)
    }

    pub fn new_from_string(file: &str) -> Result<DFA, DFAError> {
        DFAReader::new_from_lines(&mut file.lines().map(|line| Ok(line.to_string())))
    }
}

impl DFA {
    // TODO return the position of the first match
    //      maybe create an another function to do that
    pub fn run(&self, input: &str) -> bool {
        let f = input
            .chars()
            .fold(Some(self.start), |state,c| {
                match state {
                    Some(n) => self.transitions.get(&(c,n)).map(|v| *v),
                    None => None,
                }
            });
        match f {
            Some(n) => self.finals.contains(&n),
            None => false
        }
    }
}

impl fmt::Display for DFA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(writeln!(f, "START: {}", self.start));
        try!(writeln!(f, "FINALS:"));
        for fi in self.finals.iter() {
            try!(writeln!(f,"  {}", fi));
        }
        try!(writeln!(f, "TRANSITIONS:"));
        for (tr,d) in self.transitions.iter() {
            let (c,s) = *tr;
            try!(writeln!(f, "  ({},{}) => {}", c, s, d));
        }
        write!(f, "")
    }
}

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

    #[test]
    fn test_dfa() {
        let model =
            "0\n\
             3\n\
             a 0 1\n\
             c 0 3\n\
             b 1 2\n\
             a 2 1\n\
             c 2 3";
        let dfa = DFAReader::new_from_string(&model).unwrap();
        let samples =
            vec![("ababac", false),
                 ("ababc", true),
                 ("", false),
                 ("abc", true),
                 ("c", true),
                 ("ac", false),
                 ("ababababababababababababababababababababc", true),];

        for (input,expected_result) in samples {
            assert!(dfa.run(input) == expected_result, "input false for: \"{}\"", input);
        }
    }

    #[test]
    fn test_empty_file() {
        let model =
            "";
        match DFAReader::new_from_string(model) {
            Err(DFAError::MissingStartingState) => assert!(true),
            _ => assert!(false, "Missing state expected."),
        }
    }

    #[test]
    fn test_start_not_a_number() {
        let model =
            "a";
        match DFAReader::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 1),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_many_starts() {
        let model =
            "0 1\n\
             3\n\
             a 0 1\n\
             c 0 3\n\
             b 1 2\n\
             a 2 1\n\
             c 2 3";
        match DFAReader::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 1),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_no_finals() {
        let model =
            "1\n\
            ";
        match DFAReader::new_from_string(model) {
            Err(DFAError::MissingFinalStates) => assert!(true),
            _ => assert!(false, "Missing final states expected."),
        }
    }

    #[test]
    fn test_finals_not_a_number() {
        let model =
            "1\n\
             2 a 3";
        match DFAReader::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 2),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_no_transistions() {
        let model =
            "0\n\
             3";
        let _dfa = DFAReader::new_from_string(&model).unwrap();
    }

    #[test]
    fn test_transitions_with_at_least_four_elements() {
        let model =
            "0\n\
             3\n\
             a 0 1 8";
        match DFAReader::new_from_string(model) {
            Err(DFAError::IllformedTransition(line)) => assert!(line == 3),
            _ => assert!(false, "IllformedTransition expected."),
        }
    }

    #[test]
    #[should_panic]
    fn test_transitions_start_with_at_least_two_chars() {
        let model =
            "0\n\
             3\n\
             ab 2 3";
        let _dfa = DFAReader::new_from_string(&model).unwrap();
    }

    #[test]
    fn test_transitions_with_src_not_a_number() {
        let model =
            "0\n\
             3\n\
             c b 3";
        match DFAReader::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 3),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_transitions_with_dest_not_a_number() {
        let model =
            "0\n\
             3\n\
             c 2 b";
        match DFAReader::new_from_string(model) {
            Err(DFAError::Parse(_,line)) => assert!(line.unwrap() == 3),
            _ => assert!(false, "Parsing error."),
        }
    }

    #[test]
    fn test_duplicated_transition() {
        let model =
            "0\n\
             3\n\
             c 2 3\n\
             c 2 4";
        match DFAReader::new_from_string(model) {
            Err(DFAError::DuplicatedTransition(line)) => assert!(line == 4),
            _ => assert!(false, "DuplicatedTransition expected."),
        }
    }

    #[test]
    fn test_read_from_fake_file() {
        let file = "fake.txt";
        match DFAReader::new_from_file(file) {
            Err(DFAError::Io(_)) => assert!(true),
            _ => assert!(false, "Io::Error expected."),
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;
use std::io;
use std::str;

use bytes::{Bytes, BytesMut};
use bytes::{IntoBuf, Buf, BufMut, BigEndian};
use tokio_io::codec::{Encoder, Decoder};
use tokio_proto::streaming::pipeline::Frame;

use constants::*;
use packet::{PacketMagic, PTYPES};

#[derive(Debug)]
pub struct PacketHeader {
    pub magic: PacketMagic,
    pub ptype: u32,
    pub psize: u32,
}

pub struct PacketCodec {
    pub data_todo: Option<usize>,
}

type PacketItem = Frame<PacketHeader, BytesMut, io::Error>;

impl PacketHeader {
    pub fn admin_decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        let newline = buf[..].iter().position(|b| *b == b'\n');
        if let Some(n) = newline {
            let line = buf.split_to(n);
            buf.split_to(1); // drop the newline itself
            let data_str = match str::from_utf8(&line[..]) {
                Ok(s) => s,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "invalid string")),
            };
            info!("admin command data: {:?}", data_str);
            let command = match data_str.trim() {
                "version" => ADMIN_VERSION,
                "status" => ADMIN_STATUS,
                _ => ADMIN_UNKNOWN,
            };
            return Ok(Some(Frame::Message {
                message: PacketHeader {
                    magic: PacketMagic::TEXT,
                    ptype: command,
                    psize: 0,
                },
                body: false,
            }));
        }
        Ok(None) // Wait for more data
    }

    pub fn decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        debug!("Decoding {:?}", buf);
        // Peek at first 4
        // Is this a req/res
        if buf.len() < 4 {
            return Ok(None);
        }
        let mut magic_buf: [u8; 4] = [0; 4];
        magic_buf.clone_from_slice(&buf[0..4]);
        let magic = match magic_buf {
            REQ => PacketMagic::REQ,
            RES => PacketMagic::RES,
            // TEXT/ADMIN protocol
            _ => PacketMagic::TEXT,
        };
        debug!("Magic is {:?}", magic);
        if magic == PacketMagic::TEXT {
            debug!("admin protocol detected");
            return PacketHeader::admin_decode(buf);
        }
        if buf.len() < 12 {
            return Ok(None);
        }
        buf.split_to(4);
        // Now get the type
        let ptype = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("We got a {}", &PTYPES[ptype as usize].name);
        // Now the length
        let psize = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("Data section is {} bytes", psize);
        Ok(Some(Frame::Message {
            message: PacketHeader {
                magic: magic,
                ptype: ptype,
                psize: psize,
            },
            body: true, // TODO: false for 0 psize?
        }))
    }

    pub fn to_bytes(&self) -> Bytes {
        let magic = match self.magic {
            PacketMagic::UNKNOWN => panic!("Unknown packet magic cannot be sent"),
            PacketMagic::REQ => REQ,
            PacketMagic::RES => RES,
            PacketMagic::TEXT => {
                return Bytes::from_static(b"");
            }
        };
        let mut buf = BytesMut::with_capacity(12);
        buf.extend(magic.iter());
        buf.put_u32::<BigEndian>(self.ptype);
        buf.put_u32::<BigEndian>(self.psize);
        buf.freeze()
    }

    pub fn new_text_res(body: &BytesMut) -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: body.len() as u32,
        }
    }

}

impl Decoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
        match self.data_todo {
            None => {
                match PacketHeader::decode(buf)? {
                    Some(Frame::Message { message, body }) => {
                        self.data_todo = Some(message.psize as usize);
                        Ok(Some(Frame::Message {
                            message: message,
                            body: body,
                        }))
                    }
                    Some(_) => panic!("Expecting Frame::Message, got something else"),
                    None => Ok(None),
                }
            }
            Some(0) => {
                self.data_todo = None;
                Ok(Some(Frame::Body { chunk: None }))
            }
            Some(data_todo) => {
                let chunk_size = min(buf.len(), data_todo);
                self.data_todo = Some(data_todo - chunk_size);
                Ok(Some(Frame::Body { chunk: Some(buf.split_to(chunk_size)) }))
            }
        }
    }
}

impl Encoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
        debug!("Encoding {:?}", msg);
        match msg {
            Frame::Message { message, body } => {
                if body {
                    debug!("body follows")
                }
                buf.extend(message.to_bytes())
            }
            Frame::Body { chunk } => {
                match chunk {
                    Some(chunk) => buf.extend_from_slice(&chunk[..]),
                    None => {}
                }
            }
            Frame::Error { error } => {
                error!("Sending error frame. {}", error);
                buf.extend("ERR UNKNOWN_COMMAND Unknown+server+command".bytes())
            }
        }
        Ok(())
    }
}
<commit_msg>Fix formatting<commit_after>use std::cmp::min;
use std::io;
use std::str;

use bytes::{Bytes, BytesMut};
use bytes::{IntoBuf, Buf, BufMut, BigEndian};
use tokio_io::codec::{Encoder, Decoder};
use tokio_proto::streaming::pipeline::Frame;

use constants::*;
use packet::{PacketMagic, PTYPES};

#[derive(Debug)]
pub struct PacketHeader {
    pub magic: PacketMagic,
    pub ptype: u32,
    pub psize: u32,
}

pub struct PacketCodec {
    pub data_todo: Option<usize>,
}

type PacketItem = Frame<PacketHeader, BytesMut, io::Error>;

impl PacketHeader {
    pub fn admin_decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        let newline = buf[..].iter().position(|b| *b == b'\n');
        if let Some(n) = newline {
            let line = buf.split_to(n);
            buf.split_to(1); // drop the newline itself
            let data_str = match str::from_utf8(&line[..]) {
                Ok(s) => s,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "invalid string")),
            };
            info!("admin command data: {:?}", data_str);
            let command = match data_str.trim() {
                "version" => ADMIN_VERSION,
                "status" => ADMIN_STATUS,
                _ => ADMIN_UNKNOWN,
            };
            return Ok(Some(Frame::Message {
                message: PacketHeader {
                    magic: PacketMagic::TEXT,
                    ptype: command,
                    psize: 0,
                },
                body: false,
            }));
        }
        Ok(None) // Wait for more data
    }

    pub fn decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        debug!("Decoding {:?}", buf);
        // Peek at first 4
        // Is this a req/res
        if buf.len() < 4 {
            return Ok(None);
        }
        let mut magic_buf: [u8; 4] = [0; 4];
        magic_buf.clone_from_slice(&buf[0..4]);
        let magic = match magic_buf {
            REQ => PacketMagic::REQ,
            RES => PacketMagic::RES,
            // TEXT/ADMIN protocol
            _ => PacketMagic::TEXT,
        };
        debug!("Magic is {:?}", magic);
        if magic == PacketMagic::TEXT {
            debug!("admin protocol detected");
            return PacketHeader::admin_decode(buf);
        }
        if buf.len() < 12 {
            return Ok(None);
        }
        buf.split_to(4);
        // Now get the type
        let ptype = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("We got a {}", &PTYPES[ptype as usize].name);
        // Now the length
        let psize = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("Data section is {} bytes", psize);
        Ok(Some(Frame::Message {
            message: PacketHeader {
                magic: magic,
                ptype: ptype,
                psize: psize,
            },
            body: true, // TODO: false for 0 psize?
        }))
    }

    pub fn to_bytes(&self) -> Bytes {
        let magic = match self.magic {
            PacketMagic::UNKNOWN => panic!("Unknown packet magic cannot be sent"),
            PacketMagic::REQ => REQ,
            PacketMagic::RES => RES,
            PacketMagic::TEXT => {
                return Bytes::from_static(b"");
            }
        };
        let mut buf = BytesMut::with_capacity(12);
        buf.extend(magic.iter());
        buf.put_u32::<BigEndian>(self.ptype);
        buf.put_u32::<BigEndian>(self.psize);
        buf.freeze()
    }

    pub fn new_text_res(body: &BytesMut) -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: body.len() as u32,
        }
    }
}

impl Decoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
        match self.data_todo {
            None => {
                match PacketHeader::decode(buf)? {
                    Some(Frame::Message { message, body }) => {
                        self.data_todo = Some(message.psize as usize);
                        Ok(Some(Frame::Message {
                            message: message,
                            body: body,
                        }))
                    }
                    Some(_) => panic!("Expecting Frame::Message, got something else"),
                    None => Ok(None),
                }
            }
            Some(0) => {
                self.data_todo = None;
                Ok(Some(Frame::Body { chunk: None }))
            }
            Some(data_todo) => {
                let chunk_size = min(buf.len(), data_todo);
                self.data_todo = Some(data_todo - chunk_size);
                Ok(Some(Frame::Body { chunk: Some(buf.split_to(chunk_size)) }))
            }
        }
    }
}

impl Encoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
        debug!("Encoding {:?}", msg);
        match msg {
            Frame::Message { message, body } => {
                if body {
                    debug!("body follows")
                }
                buf.extend(message.to_bytes())
            }
            Frame::Body { chunk } => {
                match chunk {
                    Some(chunk) => buf.extend_from_slice(&chunk[..]),
                    None => {}
                }
            }
            Frame::Error { error } => {
                error!("Sending error frame. {}", error);
                buf.extend("ERR UNKNOWN_COMMAND Unknown+server+command".bytes())
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>const DATA: &'static str = include_str!("../data/input_25.txt");
const SEED: u64 = 20151125;
const BASE: u64 = 252533;
const MODULUS: u64 = 33554393;

pub fn main() -> Vec<String> {
    if let Some((row, col)) = parse_input(DATA.trim()) {
        let index = coord_to_index(row, col);
        let code = SEED * mod_exp(BASE, index, MODULUS) % MODULUS;
        vec![code.to_string()]
    } else {
        vec![]
    }
}

fn parse_input(input: &str) -> Option<(u64, u64)> {
    match &input.split_whitespace().collect::<Vec<_>>()[..] {
        [.., "row", r, "column", c] => {
            r.trim_matches(',')
             .parse::<u64>()
             .and_then(|r| {
                 c.trim_matches('.')
                  .parse::<u64>()
                  .map(|c| (r, c))
             })
             .ok()
        }
        _ => None,
    }
}

fn coord_to_index(row: u64, col: u64) -> u64 {
    (1..(row + col - 1)).sum::<u64>() + col - 1
}

// modular exponentiation by repeated squaring
fn mod_exp(base: u64, exponent: u64, modulus: u64) -> u64 {
    // tail recursive version, with accumulating multiplier argument
    fn mod_exp_inner(base: u64, exponent: u64, modulus: u64, mult: u64) -> u64 {
        match exponent {
            0 => mult,
            1 => base * mult % modulus,
            e if e % 2 == 0 => mod_exp_inner(base * base % modulus, e >> 1, modulus, mult),
            e => {
                mod_exp_inner(base * base % modulus,
                              (e - 1) >> 1,
                              modulus,
                              base * mult % modulus)
            }
        }
    }
    mod_exp_inner(base, exponent, modulus, 1)
}

// this was my original iterator approach, much slower but generates all of the intermediate values
#[cfg(day25_iterator)]
mod iterator {
    use super::{SEED, BASE, MODULUS};

    struct CodeGen {
        running: bool,
        last: u64,
    }

    impl CodeGen {
        fn new() -> CodeGen {
            CodeGen {
                running: false,
                last: SEED,
            }
        }
    }

    impl Iterator for CodeGen {
        type Item = u64;

        fn next(&mut self) -> Option<u64> {
            if self.running {
                self.last = (self.last * BASE) % MODULUS;
            } else {
                self.running = true;
            }
            Some(self.last)
        }
    }
}
<commit_msg>day 25: add tests<commit_after>const DATA: &'static str = include_str!("../data/input_25.txt");
const SEED: u64 = 20151125;
const BASE: u64 = 252533;
const MODULUS: u64 = 33554393;

pub fn main() -> Vec<String> {
    if let Some((row, col)) = parse_input(DATA.trim()) {
        vec![code_lookup(row, col).to_string()]
    } else {
        vec![]
    }
}

fn parse_input(input: &str) -> Option<(u64, u64)> {
    match &input.split_whitespace().collect::<Vec<_>>()[..] {
        [.., "row", r, "column", c] => {
            r.trim_matches(',')
             .parse::<u64>()
             .and_then(|r| {
                 c.trim_matches('.')
                  .parse::<u64>()
                  .map(|c| (r, c))
             })
             .ok()
        }
        _ => None,
    }
}

fn code_lookup(row: u64, col: u64) -> u64 {
    let index = coord_to_index(row, col);
    SEED * mod_exp(BASE, index, MODULUS) % MODULUS
}

fn coord_to_index(row: u64, col: u64) -> u64 {
    (1..(row + col - 1)).sum::<u64>() + col - 1
}

// modular exponentiation by repeated squaring
fn mod_exp(base: u64, exponent: u64, modulus: u64) -> u64 {
    // tail recursive version, with accumulating multiplier argument
    fn mod_exp_inner(base: u64, exponent: u64, modulus: u64, mult: u64) -> u64 {
        match exponent {
            0 => mult,
            1 => base * mult % modulus,
            e if e % 2 == 0 => mod_exp_inner(base * base % modulus, e >> 1, modulus, mult),
            e => {
                mod_exp_inner(base * base % modulus,
                              (e - 1) >> 1,
                              modulus,
                              base * mult % modulus)
            }
        }
    }
    mod_exp_inner(base, exponent, modulus, 1)
}

#[cfg(test)]
mod test {
    use super::{coord_to_index, code_lookup};

    #[test]
    fn examples() {
        assert_eq!(coord_to_index(1, 1), 0);
        assert_eq!(code_lookup(1, 1), 20151125);
        assert_eq!(coord_to_index(2, 1), 1);
        assert_eq!(code_lookup(2, 1), 31916031);
        assert_eq!(coord_to_index(1, 2), 2);
        assert_eq!(code_lookup(1, 2), 18749137);
    }
}

// this was my original iterator approach, much slower but generates all of the intermediate values
#[cfg(day25_iterator)]
mod iterator {
    use super::{SEED, BASE, MODULUS};

    struct CodeGen {
        running: bool,
        last: u64,
    }

    impl CodeGen {
        fn new() -> CodeGen {
            CodeGen {
                running: false,
                last: SEED,
            }
        }
    }

    impl Iterator for CodeGen {
        type Item = u64;

        fn next(&mut self) -> Option<u64> {
            if self.running {
                self.last = (self.last * BASE) % MODULUS;
            } else {
                self.running = true;
            }
            Some(self.last)
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::rc::Rc;
use std::hash;
use digit_slice::DigitSlice;

pub enum Shape<'a> {
    /// A natural number represented as a little-endian sequence of 8-bit
    /// digits.
    Atom(&'a [u8]),
    /// An ordered pair of nouns.
    Cell(&'a Noun, &'a Noun),
}

/// A Nock noun, the basic unit of representation.
///
/// A noun is an atom or a cell. An atom is any natural number. A cell is any
/// ordered pair of nouns.
#[derive(Clone, PartialEq, Eq)]
pub struct Noun(Inner);

#[derive(Clone, PartialEq, Eq)]
enum Inner {
    Atom(Rc<Vec<u8>>),
    Cell(Rc<Noun>, Rc<Noun>),
}

impl Noun {
    fn get<'a>(&'a self) -> Shape<'a> {
        match self.0 {
            Inner::Atom(ref v) => Shape::Atom(&v),
            Inner::Cell(ref a, ref b) => Shape::Cell(&*a, &*b),
        }
    }

    /// Memory address or other unique identifier for the noun.
    fn addr(&self) -> usize {
        &*self as *const _ as usize
    }

    /// Run a memoizing fold over the noun
    fn fold<'a, F, G, T: Clone>(&'a self, mut leaf: F, mut branch: G) -> T
        where F: FnMut(&[u8]) -> T,
              G: FnMut(&T, &T) -> T
    {
        fn h<'a, F, G, T>(noun: &'a Noun,
                          memo: &mut HashMap<usize, T>,
                          leaf: &mut F,
                          branch: &mut G)
                          -> T
            where F: FnMut(&[u8]) -> T,
                  G: FnMut(&T, &T) -> T,
                  T: Clone
        {
            let key = noun.addr();

            if memo.contains_key(&key) {
                memo.get(&key).unwrap().clone()
            } else {
                let ret = match noun.get() {
                    Shape::Atom(x) => leaf(x),
                    Shape::Cell(ref a, ref b) => {
                        let a = h(*a, memo, leaf, branch);
                        let b = h(*b, memo, leaf, branch);
                        branch(&a, &b)
                    }
                };
                memo.insert(key, ret.clone());
                ret
            }
        }

        h(self, &mut HashMap::new(), &mut leaf, &mut branch)
    }

    /// Build a new atom noun from a little-endian 8-bit digit sequence.
    fn atom(digits: &[u8]) -> Noun {
        Noun(Inner::Atom(Rc::new(digits.to_vec())))
    }

    /// Build a new cell noun from two existing nouns.
    fn cell(&mut self, a: &Noun, b: &Noun) -> Noun {
        Noun(Inner::Cell(Rc::new(a.clone()), Rc::new(b.clone())))
    }

    fn from<T: ToNoun>(item: T) -> Noun {
        item.to_noun()
    }
}

/// Trait for types that can convert themselves to a noun.
pub trait ToNoun {
    fn to_noun(&self) -> Noun;
}

impl<T> ToNoun for T where T: DigitSlice
{
    fn to_noun(&self) -> Noun {
        Noun::atom(self.as_digits())
    }
}

#[cfg(test)]
mod tests {
    use super::Noun;

    #[test]
    fn scratch() {
        let x = Noun::from(123u32);
        assert!(x == Noun::from(123u8));
    }
}
<commit_msg>Generic Shape, use for folding<commit_after>use std::collections::HashMap;
use std::rc::Rc;
use std::hash;
use digit_slice::DigitSlice;

/// A wrapper for referencing Noun-like patterns.
pub enum Shape<A, N> {
    Atom(A),
    Cell(N, N),
}

/// A Nock noun, the basic unit of representation.
///
/// A noun is an atom or a cell. An atom is any natural number. A cell is any
/// ordered pair of nouns.
///
/// Atoms are represented by a little-endian byte array of 8-bit digits.
#[derive(Clone, PartialEq, Eq)]
pub struct Noun(Inner);

#[derive(Clone, PartialEq, Eq)]
enum Inner {
    Atom(Rc<Vec<u8>>),
    Cell(Rc<Noun>, Rc<Noun>),
}

impl Noun {
    fn get<'a>(&'a self) -> Shape<&'a [u8], &'a Noun> {
        match self.0 {
            Inner::Atom(ref v) => Shape::Atom(&v),
            Inner::Cell(ref a, ref b) => Shape::Cell(&*a, &*b),
        }
    }

    /// Memory address or other unique identifier for the noun.
    fn addr(&self) -> usize {
        &*self as *const _ as usize
    }

    /// Run a memoizing fold over the noun
    fn fold<'a, F, T>(&'a self, mut f: F) -> T
        where F: FnMut(Shape<&'a [u8], T>) -> T,
              T: Clone
    {
        fn h<'a, F, T>(noun: &'a Noun,
                       memo: &mut HashMap<usize, T>,
                       f: &mut F)
                       -> T
            where F: FnMut(Shape<&'a [u8], T>) -> T,
                  T: Clone
        {
            let key = noun.addr();

            if memo.contains_key(&key) {
                memo.get(&key).unwrap().clone()
            } else {
                let ret = match noun.get() {
                    Shape::Atom(x) => f(Shape::Atom(x)),
                    Shape::Cell(ref a, ref b) => {
                        let a = h(*a, memo, f);
                        let b = h(*b, memo, f);
                        let ret = f(Shape::Cell(a, b));
                        ret
                    }
                };
                memo.insert(key, ret.clone());
                ret
            }
        }

        h(self, &mut HashMap::new(), &mut f)
    }

    /// Build a new atom noun from a little-endian 8-bit digit sequence.
    pub fn atom(digits: &[u8]) -> Noun {
        Noun(Inner::Atom(Rc::new(digits.to_vec())))
    }

    /// Build a new cell noun from two existing nouns.
    pub fn cell(&mut self, a: &Noun, b: &Noun) -> Noun {
        Noun(Inner::Cell(Rc::new(a.clone()), Rc::new(b.clone())))
    }

    pub fn from<T: ToNoun>(item: T) -> Noun {
        item.to_noun()
    }
}

impl hash::Hash for Noun {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        fn f<H: hash::Hasher>(state: &mut H, shape: Shape<&[u8], u64>) -> u64 {
            match shape {
                Shape::Atom(x) => x.hash(state),
                Shape::Cell(a, b) => {
                    a.hash(state);
                    b.hash(state);
                }
            }
            state.finish()
        }
        self.fold(|x| f(state, x))
            .hash(state);
    }
}

/// Trait for types that can convert themselves to a noun.
pub trait ToNoun {
    fn to_noun(&self) -> Noun;
}

impl<T> ToNoun for T where T: DigitSlice
{
    fn to_noun(&self) -> Noun {
        Noun::atom(self.as_digits())
    }
}

#[cfg(test)]
mod tests {
    use super::Noun;

    #[test]
    fn scratch() {
        let x = Noun::from(123u32);
        assert!(x == Noun::from(123u8));
    }
}
<|endoftext|>


<commit_before>use crate::npyffi::{NpyTypes, PyArray_Descr, NPY_TYPES, PY_ARRAY_API};
use cfg_if::cfg_if;
use pyo3::{ffi, prelude::*, pyobject_native_type_core, types::PyType, AsPyPointer, PyNativeType};
use std::os::raw::c_int;

pub use num_complex::Complex32 as c32;
pub use num_complex::Complex64 as c64;

/// Binding of [`numpy.dtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.html).
///
/// # Example
/// ```
/// use pyo3::types::IntoPyDict;
/// pyo3::Python::with_gil(|py| {
///     let locals = [("np", numpy::get_array_module(py).unwrap())].into_py_dict(py);
///     let dtype: &numpy::PyArrayDescr = py
///         .eval("np.array([1, 2, 3.0]).dtype", Some(locals), None)
///         .unwrap()
///         .downcast()
///         .unwrap();
///     assert_eq!(dtype.get_datatype().unwrap(), numpy::DataType::Float64);
/// });
/// ```
pub struct PyArrayDescr(PyAny);

pyobject_native_type_core!(
    PyArrayDescr,
    *PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    #module=Some("numpy"),
    #checkfunction=arraydescr_check
);

unsafe fn arraydescr_check(op: *mut ffi::PyObject) -> c_int {
    ffi::PyObject_TypeCheck(
        op,
        PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    )
}

impl PyArrayDescr {
    /// Returns `self` as `*mut PyArray_Descr`.
    pub fn as_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.as_ptr() as _
    }

    /// Returns `self` as `*mut PyArray_Descr` while increasing the reference count.
    ///
    /// Useful in cases where the descriptor is stolen by the API.
    pub fn into_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.into_ptr() as _
    }

    /// Returns the internal `PyType` that this `dtype` holds.
    ///
    /// # Example
    /// ```
    /// pyo3::Python::with_gil(|py| {
    ///    let array = numpy::PyArray::from_vec(py, vec![0.0, 1.0, 2.0f64]);
    ///    let dtype = array.dtype();
    ///    assert_eq!(dtype.get_type().name().unwrap().to_string(), "float64");
    /// });
    /// ```
    pub fn get_type(&self) -> &PyType {
        let dtype_type_ptr = unsafe { *self.as_dtype_ptr() }.typeobj;
        unsafe { PyType::from_type_ptr(self.py(), dtype_type_ptr) }
    }

    /// Returns the data type as `DataType` enum.
    pub fn get_datatype(&self) -> Option<DataType> {
        DataType::from_typenum(self.get_typenum())
    }

    /// Shortcut for creating a descriptor of 'object' type.
    pub fn object(py: Python) -> &Self {
        Self::from_npy_type(py, NPY_TYPES::NPY_OBJECT)
    }

    fn from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
        unsafe {
            let descr = PY_ARRAY_API.PyArray_DescrFromType(npy_type as i32);
            py.from_owned_ptr(descr as _)
        }
    }

    pub(crate) fn get_typenum(&self) -> std::os::raw::c_int {
        unsafe { *self.as_dtype_ptr() }.type_num
    }
}

/// Represents numpy data type.
///
/// This is an incomplete counterpart of
/// [Enumerated Types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
/// in numpy C-API.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DataType {
    Bool,
    Int8,
    Int16,
    Int32,
    Int64,
    Uint8,
    Uint16,
    Uint32,
    Uint64,
    Float32,
    Float64,
    Complex32,
    Complex64,
    Object,
}

impl DataType {
    /// Construct `DataType` from
    /// [Enumerated Types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types).
    pub fn from_typenum(typenum: c_int) -> Option<Self> {
        Some(match typenum {
            x if x == NPY_TYPES::NPY_BOOL as i32 => DataType::Bool,
            x if x == NPY_TYPES::NPY_BYTE as i32 => DataType::Int8,
            x if x == NPY_TYPES::NPY_SHORT as i32 => DataType::Int16,
            x if x == NPY_TYPES::NPY_INT as i32 => DataType::Int32,
            x if x == NPY_TYPES::NPY_LONG as i32 => return DataType::from_clong(false),
            x if x == NPY_TYPES::NPY_LONGLONG as i32 => DataType::Int64,
            x if x == NPY_TYPES::NPY_UBYTE as i32 => DataType::Uint8,
            x if x == NPY_TYPES::NPY_USHORT as i32 => DataType::Uint16,
            x if x == NPY_TYPES::NPY_UINT as i32 => DataType::Uint32,
            x if x == NPY_TYPES::NPY_ULONG as i32 => return DataType::from_clong(true),
            x if x == NPY_TYPES::NPY_ULONGLONG as i32 => DataType::Uint64,
            x if x == NPY_TYPES::NPY_FLOAT as i32 => DataType::Float32,
            x if x == NPY_TYPES::NPY_DOUBLE as i32 => DataType::Float64,
            x if x == NPY_TYPES::NPY_CFLOAT as i32 => DataType::Complex32,
            x if x == NPY_TYPES::NPY_CDOUBLE as i32 => DataType::Complex64,
            x if x == NPY_TYPES::NPY_OBJECT as i32 => DataType::Object,
            _ => return None,
        })
    }

    /// Convert `self` into
    /// [Enumerated Types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types).
    pub fn into_ctype(self) -> NPY_TYPES {
        match self {
            DataType::Bool => NPY_TYPES::NPY_BOOL,
            DataType::Int8 => NPY_TYPES::NPY_BYTE,
            DataType::Int16 => NPY_TYPES::NPY_SHORT,
            DataType::Int32 => NPY_TYPES::NPY_INT,
            #[cfg(all(target_pointer_width = "64", not(windows)))]
            DataType::Int64 => NPY_TYPES::NPY_LONG,
            #[cfg(any(target_pointer_width = "32", windows))]
            DataType::Int64 => NPY_TYPES::NPY_LONGLONG,
            DataType::Uint8 => NPY_TYPES::NPY_UBYTE,
            DataType::Uint16 => NPY_TYPES::NPY_USHORT,
            DataType::Uint32 => NPY_TYPES::NPY_UINT,
            DataType::Uint64 => NPY_TYPES::NPY_ULONGLONG,
            DataType::Float32 => NPY_TYPES::NPY_FLOAT,
            DataType::Float64 => NPY_TYPES::NPY_DOUBLE,
            DataType::Complex32 => NPY_TYPES::NPY_CFLOAT,
            DataType::Complex64 => NPY_TYPES::NPY_CDOUBLE,
            DataType::Object => NPY_TYPES::NPY_OBJECT,
        }
    }

    #[inline(always)]
    fn from_clong(is_usize: bool) -> Option<Self> {
        if cfg!(any(target_pointer_width = "32", windows)) {
            Some(if is_usize {
                DataType::Uint32
            } else {
                DataType::Int32
            })
        } else if cfg!(all(target_pointer_width = "64", not(windows))) {
            Some(if is_usize {
                DataType::Uint64
            } else {
                DataType::Int64
            })
        } else {
            None
        }
    }
}

/// Represents that a type can be an element of `PyArray`.
///
/// Currently, only integer/float/complex types are supported.
/// If you come up with a nice implementation for some other types, we're happy to receive your PR :)
/// You may refer to the [numpy document](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
/// for all types that numpy supports.
///
/// # Safety
///
/// A type `T` that implements this trait should be safe when managed in numpy array,
/// thus implementing this trait is marked unsafe.
/// This means that all data types except for `DataType::Object` are assumed to be trivially copyable.
/// Furthermore, it is assumed that for `DataType::Object` the elements are pointers into the Python heap
/// and that the corresponding `Clone` implemenation will never panic as it only increases the reference count.
///
/// # Custom element types
///
/// You can implement this trait to manage arrays of custom element types, but they still need to be stored
/// on Python's heap using PyO3's [Py](pyo3::Py) type.
///
/// ```
/// use numpy::{ndarray::Array2, DataType, Element, PyArray, PyArrayDescr, ToPyArray};
/// use pyo3::{pyclass, Py, Python};
///
/// #[pyclass]
/// pub struct CustomElement;
///
/// // The transparent wrapper is necessary as one cannot implement
/// // a foreign trait (`Element`) on a foreign type (`Py`) directly.
/// #[derive(Clone)]
/// #[repr(transparent)]
/// pub struct Wrapper(pub Py<CustomElement>);
///
/// unsafe impl Element for Wrapper {
///     const DATA_TYPE: DataType = DataType::Object;
///
///     fn is_same_type(dtype: &PyArrayDescr) -> bool {
///         dtype.get_datatype() == Some(DataType::Object)
///     }
///
///     fn get_dtype(py: Python) -> &PyArrayDescr {
///         PyArrayDescr::object(py)
///     }
/// }
///
/// Python::with_gil(|py| {
///     let array = Array2::<Wrapper>::from_shape_fn((2, 3), |(_i, _j)| {
///         Wrapper(Py::new(py, CustomElement).unwrap())
///     });
///
///     let _array: &PyArray<Wrapper, _> = array.to_pyarray(py);
/// });
/// ```
pub unsafe trait Element: Clone + Send {
    /// `DataType` corresponding to this type.
    const DATA_TYPE: DataType;

    /// Returns if the give `dtype` is convertible to `Self` in Rust.
    fn is_same_type(dtype: &PyArrayDescr) -> bool;

    /// Create `dtype`.
    fn get_dtype(py: Python) -> &PyArrayDescr;
}

macro_rules! impl_num_element {
    ($t:ty, $npy_dat_t:ident $(,$npy_types: ident)+) => {
        unsafe impl Element for $t {
            const DATA_TYPE: DataType = DataType::$npy_dat_t;
            fn is_same_type(dtype: &PyArrayDescr) -> bool {
                $(dtype.get_typenum() == NPY_TYPES::$npy_types as i32 ||)+ false
            }
            fn get_dtype(py: Python) -> &PyArrayDescr {
                PyArrayDescr::from_npy_type(py, DataType::$npy_dat_t.into_ctype())
            }
        }
    };
}

impl_num_element!(bool, Bool, NPY_BOOL);
impl_num_element!(i8, Int8, NPY_BYTE);
impl_num_element!(i16, Int16, NPY_SHORT);
impl_num_element!(u8, Uint8, NPY_UBYTE);
impl_num_element!(u16, Uint16, NPY_USHORT);
impl_num_element!(f32, Float32, NPY_FLOAT);
impl_num_element!(f64, Float64, NPY_DOUBLE);
impl_num_element!(c32, Complex32, NPY_CFLOAT);
impl_num_element!(c64, Complex64, NPY_CDOUBLE);

cfg_if! {
    if #[cfg(all(target_pointer_width = "64", windows))] {
            impl_num_element!(usize, Uint64, NPY_ULONGLONG);
    } else if #[cfg(all(target_pointer_width = "64", not(windows)))] {
            impl_num_element!(usize, Uint64, NPY_ULONG, NPY_ULONGLONG);
    } else if #[cfg(all(target_pointer_width = "32", windows))] {
            impl_num_element!(usize, Uint32, NPY_UINT, NPY_ULONG);
    } else if #[cfg(all(target_pointer_width = "32", not(windows)))] {
            impl_num_element!(usize, Uint32, NPY_UINT);
    }
}
cfg_if! {
    if #[cfg(any(target_pointer_width = "32", windows))] {
        impl_num_element!(i32, Int32, NPY_INT, NPY_LONG);
        impl_num_element!(u32, Uint32, NPY_UINT, NPY_ULONG);
        impl_num_element!(i64, Int64, NPY_LONGLONG);
        impl_num_element!(u64, Uint64, NPY_ULONGLONG);
    } else if #[cfg(all(target_pointer_width = "64", not(windows)))] {
        impl_num_element!(i32, Int32, NPY_INT);
        impl_num_element!(u32, Uint32, NPY_UINT);
        impl_num_element!(i64, Int64, NPY_LONG, NPY_LONGLONG);
        impl_num_element!(u64, Uint64, NPY_ULONG, NPY_ULONGLONG);
    }
}

unsafe impl Element for PyObject {
    const DATA_TYPE: DataType = DataType::Object;
    fn is_same_type(dtype: &PyArrayDescr) -> bool {
        dtype.get_typenum() == NPY_TYPES::NPY_OBJECT as i32
    }
    fn get_dtype(py: Python) -> &PyArrayDescr {
        PyArrayDescr::object(py)
    }
}
<commit_msg>Add `PyArrayDescr::of<T>()`<commit_after>use crate::npyffi::{NpyTypes, PyArray_Descr, NPY_TYPES, PY_ARRAY_API};
use cfg_if::cfg_if;
use pyo3::{ffi, prelude::*, pyobject_native_type_core, types::PyType, AsPyPointer, PyNativeType};
use std::os::raw::c_int;

pub use num_complex::Complex32 as c32;
pub use num_complex::Complex64 as c64;

/// Binding of [`numpy.dtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.html).
///
/// # Example
/// ```
/// use pyo3::types::IntoPyDict;
/// pyo3::Python::with_gil(|py| {
///     let locals = [("np", numpy::get_array_module(py).unwrap())].into_py_dict(py);
///     let dtype: &numpy::PyArrayDescr = py
///         .eval("np.array([1, 2, 3.0]).dtype", Some(locals), None)
///         .unwrap()
///         .downcast()
///         .unwrap();
///     assert_eq!(dtype.get_datatype().unwrap(), numpy::DataType::Float64);
/// });
/// ```
pub struct PyArrayDescr(PyAny);

pyobject_native_type_core!(
    PyArrayDescr,
    *PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    #module=Some("numpy"),
    #checkfunction=arraydescr_check
);

unsafe fn arraydescr_check(op: *mut ffi::PyObject) -> c_int {
    ffi::PyObject_TypeCheck(
        op,
        PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    )
}

impl PyArrayDescr {
    /// Returns `self` as `*mut PyArray_Descr`.
    pub fn as_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.as_ptr() as _
    }

    /// Returns `self` as `*mut PyArray_Descr` while increasing the reference count.
    ///
    /// Useful in cases where the descriptor is stolen by the API.
    pub fn into_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.into_ptr() as _
    }

    /// Returns the internal `PyType` that this `dtype` holds.
    ///
    /// # Example
    /// ```
    /// pyo3::Python::with_gil(|py| {
    ///    let array = numpy::PyArray::from_vec(py, vec![0.0, 1.0, 2.0f64]);
    ///    let dtype = array.dtype();
    ///    assert_eq!(dtype.get_type().name().unwrap().to_string(), "float64");
    /// });
    /// ```
    pub fn get_type(&self) -> &PyType {
        let dtype_type_ptr = unsafe { *self.as_dtype_ptr() }.typeobj;
        unsafe { PyType::from_type_ptr(self.py(), dtype_type_ptr) }
    }

    /// Returns the data type as `DataType` enum.
    pub fn get_datatype(&self) -> Option<DataType> {
        DataType::from_typenum(self.get_typenum())
    }

    /// Shortcut for creating a descriptor of 'object' type.
    pub fn object(py: Python) -> &Self {
        Self::from_npy_type(py, NPY_TYPES::NPY_OBJECT)
    }

    /// Returns the type descriptor for a registered type.
    pub fn of<T: Element>(py: Python) -> &Self {
        T::get_dtype(py)
    }

    fn from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
        unsafe {
            let descr = PY_ARRAY_API.PyArray_DescrFromType(npy_type as i32);
            py.from_owned_ptr(descr as _)
        }
    }

    pub(crate) fn get_typenum(&self) -> std::os::raw::c_int {
        unsafe { *self.as_dtype_ptr() }.type_num
    }
}

/// Represents numpy data type.
///
/// This is an incomplete counterpart of
/// [Enumerated Types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
/// in numpy C-API.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DataType {
    Bool,
    Int8,
    Int16,
    Int32,
    Int64,
    Uint8,
    Uint16,
    Uint32,
    Uint64,
    Float32,
    Float64,
    Complex32,
    Complex64,
    Object,
}

impl DataType {
    /// Construct `DataType` from
    /// [Enumerated Types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types).
    pub fn from_typenum(typenum: c_int) -> Option<Self> {
        Some(match typenum {
            x if x == NPY_TYPES::NPY_BOOL as i32 => DataType::Bool,
            x if x == NPY_TYPES::NPY_BYTE as i32 => DataType::Int8,
            x if x == NPY_TYPES::NPY_SHORT as i32 => DataType::Int16,
            x if x == NPY_TYPES::NPY_INT as i32 => DataType::Int32,
            x if x == NPY_TYPES::NPY_LONG as i32 => return DataType::from_clong(false),
            x if x == NPY_TYPES::NPY_LONGLONG as i32 => DataType::Int64,
            x if x == NPY_TYPES::NPY_UBYTE as i32 => DataType::Uint8,
            x if x == NPY_TYPES::NPY_USHORT as i32 => DataType::Uint16,
            x if x == NPY_TYPES::NPY_UINT as i32 => DataType::Uint32,
            x if x == NPY_TYPES::NPY_ULONG as i32 => return DataType::from_clong(true),
            x if x == NPY_TYPES::NPY_ULONGLONG as i32 => DataType::Uint64,
            x if x == NPY_TYPES::NPY_FLOAT as i32 => DataType::Float32,
            x if x == NPY_TYPES::NPY_DOUBLE as i32 => DataType::Float64,
            x if x == NPY_TYPES::NPY_CFLOAT as i32 => DataType::Complex32,
            x if x == NPY_TYPES::NPY_CDOUBLE as i32 => DataType::Complex64,
            x if x == NPY_TYPES::NPY_OBJECT as i32 => DataType::Object,
            _ => return None,
        })
    }

    /// Convert `self` into
    /// [Enumerated Types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types).
    pub fn into_ctype(self) -> NPY_TYPES {
        match self {
            DataType::Bool => NPY_TYPES::NPY_BOOL,
            DataType::Int8 => NPY_TYPES::NPY_BYTE,
            DataType::Int16 => NPY_TYPES::NPY_SHORT,
            DataType::Int32 => NPY_TYPES::NPY_INT,
            #[cfg(all(target_pointer_width = "64", not(windows)))]
            DataType::Int64 => NPY_TYPES::NPY_LONG,
            #[cfg(any(target_pointer_width = "32", windows))]
            DataType::Int64 => NPY_TYPES::NPY_LONGLONG,
            DataType::Uint8 => NPY_TYPES::NPY_UBYTE,
            DataType::Uint16 => NPY_TYPES::NPY_USHORT,
            DataType::Uint32 => NPY_TYPES::NPY_UINT,
            DataType::Uint64 => NPY_TYPES::NPY_ULONGLONG,
            DataType::Float32 => NPY_TYPES::NPY_FLOAT,
            DataType::Float64 => NPY_TYPES::NPY_DOUBLE,
            DataType::Complex32 => NPY_TYPES::NPY_CFLOAT,
            DataType::Complex64 => NPY_TYPES::NPY_CDOUBLE,
            DataType::Object => NPY_TYPES::NPY_OBJECT,
        }
    }

    #[inline(always)]
    fn from_clong(is_usize: bool) -> Option<Self> {
        if cfg!(any(target_pointer_width = "32", windows)) {
            Some(if is_usize {
                DataType::Uint32
            } else {
                DataType::Int32
            })
        } else if cfg!(all(target_pointer_width = "64", not(windows))) {
            Some(if is_usize {
                DataType::Uint64
            } else {
                DataType::Int64
            })
        } else {
            None
        }
    }
}

/// Represents that a type can be an element of `PyArray`.
///
/// Currently, only integer/float/complex types are supported.
/// If you come up with a nice implementation for some other types, we're happy to receive your PR :)
/// You may refer to the [numpy document](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
/// for all types that numpy supports.
///
/// # Safety
///
/// A type `T` that implements this trait should be safe when managed in numpy array,
/// thus implementing this trait is marked unsafe.
/// This means that all data types except for `DataType::Object` are assumed to be trivially copyable.
/// Furthermore, it is assumed that for `DataType::Object` the elements are pointers into the Python heap
/// and that the corresponding `Clone` implemenation will never panic as it only increases the reference count.
///
/// # Custom element types
///
/// You can implement this trait to manage arrays of custom element types, but they still need to be stored
/// on Python's heap using PyO3's [Py](pyo3::Py) type.
///
/// ```
/// use numpy::{ndarray::Array2, DataType, Element, PyArray, PyArrayDescr, ToPyArray};
/// use pyo3::{pyclass, Py, Python};
///
/// #[pyclass]
/// pub struct CustomElement;
///
/// // The transparent wrapper is necessary as one cannot implement
/// // a foreign trait (`Element`) on a foreign type (`Py`) directly.
/// #[derive(Clone)]
/// #[repr(transparent)]
/// pub struct Wrapper(pub Py<CustomElement>);
///
/// unsafe impl Element for Wrapper {
///     const DATA_TYPE: DataType = DataType::Object;
///
///     fn is_same_type(dtype: &PyArrayDescr) -> bool {
///         dtype.get_datatype() == Some(DataType::Object)
///     }
///
///     fn get_dtype(py: Python) -> &PyArrayDescr {
///         PyArrayDescr::object(py)
///     }
/// }
///
/// Python::with_gil(|py| {
///     let array = Array2::<Wrapper>::from_shape_fn((2, 3), |(_i, _j)| {
///         Wrapper(Py::new(py, CustomElement).unwrap())
///     });
///
///     let _array: &PyArray<Wrapper, _> = array.to_pyarray(py);
/// });
/// ```
pub unsafe trait Element: Clone + Send {
    /// `DataType` corresponding to this type.
    const DATA_TYPE: DataType;

    /// Returns if the give `dtype` is convertible to `Self` in Rust.
    fn is_same_type(dtype: &PyArrayDescr) -> bool;

    /// Create `dtype`.
    fn get_dtype(py: Python) -> &PyArrayDescr;
}

macro_rules! impl_num_element {
    ($t:ty, $npy_dat_t:ident $(,$npy_types: ident)+) => {
        unsafe impl Element for $t {
            const DATA_TYPE: DataType = DataType::$npy_dat_t;
            fn is_same_type(dtype: &PyArrayDescr) -> bool {
                $(dtype.get_typenum() == NPY_TYPES::$npy_types as i32 ||)+ false
            }
            fn get_dtype(py: Python) -> &PyArrayDescr {
                PyArrayDescr::from_npy_type(py, DataType::$npy_dat_t.into_ctype())
            }
        }
    };
}

impl_num_element!(bool, Bool, NPY_BOOL);
impl_num_element!(i8, Int8, NPY_BYTE);
impl_num_element!(i16, Int16, NPY_SHORT);
impl_num_element!(u8, Uint8, NPY_UBYTE);
impl_num_element!(u16, Uint16, NPY_USHORT);
impl_num_element!(f32, Float32, NPY_FLOAT);
impl_num_element!(f64, Float64, NPY_DOUBLE);
impl_num_element!(c32, Complex32, NPY_CFLOAT);
impl_num_element!(c64, Complex64, NPY_CDOUBLE);

cfg_if! {
    if #[cfg(all(target_pointer_width = "64", windows))] {
            impl_num_element!(usize, Uint64, NPY_ULONGLONG);
    } else if #[cfg(all(target_pointer_width = "64", not(windows)))] {
            impl_num_element!(usize, Uint64, NPY_ULONG, NPY_ULONGLONG);
    } else if #[cfg(all(target_pointer_width = "32", windows))] {
            impl_num_element!(usize, Uint32, NPY_UINT, NPY_ULONG);
    } else if #[cfg(all(target_pointer_width = "32", not(windows)))] {
            impl_num_element!(usize, Uint32, NPY_UINT);
    }
}
cfg_if! {
    if #[cfg(any(target_pointer_width = "32", windows))] {
        impl_num_element!(i32, Int32, NPY_INT, NPY_LONG);
        impl_num_element!(u32, Uint32, NPY_UINT, NPY_ULONG);
        impl_num_element!(i64, Int64, NPY_LONGLONG);
        impl_num_element!(u64, Uint64, NPY_ULONGLONG);
    } else if #[cfg(all(target_pointer_width = "64", not(windows)))] {
        impl_num_element!(i32, Int32, NPY_INT);
        impl_num_element!(u32, Uint32, NPY_UINT);
        impl_num_element!(i64, Int64, NPY_LONG, NPY_LONGLONG);
        impl_num_element!(u64, Uint64, NPY_ULONG, NPY_ULONGLONG);
    }
}

unsafe impl Element for PyObject {
    const DATA_TYPE: DataType = DataType::Object;
    fn is_same_type(dtype: &PyArrayDescr) -> bool {
        dtype.get_typenum() == NPY_TYPES::NPY_OBJECT as i32
    }
    fn get_dtype(py: Python) -> &PyArrayDescr {
        PyArrayDescr::object(py)
    }
}
<|endoftext|>