<commit_before>use std::cmp::{max, min};
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;
use utils::calculate_ratio;

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
pub struct Match {
    pub first_start: usize,
    pub second_start: usize,
    pub size: usize,
}

impl Match {
    fn new(first_start: usize, second_start: usize, size: usize) -> Match {
        Match {
            first_start: first_start,
            second_start: second_start,
            size: size,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Opcode {
    pub tag: String,
    pub first_start: usize,
    pub first_end: usize,
    pub second_start: usize,
    pub second_end: usize,
}

impl Opcode {
    fn new(
        tag: String,
        first_start: usize,
        first_end: usize,
        second_start: usize,
        second_end: usize,
    ) -> Opcode {
        Opcode {
            tag: tag,
            first_start: first_start,
            first_end: first_end,
            second_start: second_start,
            second_end: second_end,
        }
    }
}

pub trait Sequence: Debug {
    fn len(&self) -> usize;
    fn at_index(&self, index: usize) -> Option<&str>;
}

impl Sequence for str {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index > self.len() {
            return None;
        }
        unsafe { Some(self.slice_unchecked(index, index + 1)) }
    }
}

impl<'a> Sequence for Vec<&'a str> {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index < self.len() {
            return Some(self[index]);
        }
        None
    }
}

impl Sequence for String {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index > self.len() {
            return None;
        }
        unsafe { Some(self.slice_unchecked(index, index + 1)) }
    }
}

impl<'a> Sequence for Vec<String> {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index < self.len() {
            return Some(&self[index]);
        }
        None
    }
}

pub struct SequenceMatcher<'a, T: 'a + ?Sized + Sequence> {
    first_sequence: &'a T,
    second_sequence: &'a T,
    matching_blocks: Option<Vec<Match>>,
    opcodes: Option<Vec<Opcode>>,
    is_junk: Option<fn(&str) -> bool>,
    second_sequence_elements: HashMap<&'a str, Vec<usize>>,
    second_sequence_popular: HashSet<&'a str>,
}

impl<'a, T: ?Sized + Sequence> SequenceMatcher<'a, T> {
    pub fn new(first_sequence: &'a T, second_sequence: &'a T) -> SequenceMatcher<'a, T> {
        let mut matcher = SequenceMatcher {
            first_sequence: first_sequence,
            second_sequence: second_sequence,
            matching_blocks: None,
            opcodes: None,
            is_junk: None,
            second_sequence_elements: HashMap::new(),
            second_sequence_popular: HashSet::new(),
        };
        matcher.set_seqs(first_sequence, second_sequence);
        matcher
    }

    pub fn set_is_junk(&mut self, is_junk: Option<fn(&str) -> bool>) {
        self.is_junk = is_junk;
        self.set_second_seq(self.second_sequence);
    }

    pub fn set_seqs(&mut self, first_sequence: &'a T, second_sequence: &'a T) {
        self.set_first_seq(first_sequence);
        self.set_second_seq(second_sequence);
    }

    pub fn set_first_seq(&mut self, sequence: &'a T) {
        self.first_sequence = sequence;
        self.matching_blocks = None;
        self.opcodes = None;
    }

    pub fn set_second_seq(&mut self, sequence: &'a T) {
        self.second_sequence = sequence;
        self.matching_blocks = None;
        self.opcodes = None;
        self.chain_second_seq();
    }

    fn chain_second_seq(&mut self) {
        let second_sequence = self.second_sequence;
        let mut second_sequence_elements = HashMap::new();
        for i in 0..second_sequence.len() {
            let mut counter = second_sequence_elements
                .entry(second_sequence.at_index(i).unwrap())
                .or_insert(Vec::new());
            counter.push(i);
        }
        if self.is_junk.is_some() {
            let mut junk = Vec::new();
            for element in second_sequence_elements.keys() {
                if (self.is_junk.unwrap())(element) {
                    junk.push(element.clone());
                }
            }
            for element in &junk {
                second_sequence_elements.remove(element);
            }
        }
        let mut popular = HashSet::new();
        let len = second_sequence.len();
        if len >= 200 {
            let test_len = (len as f32 / 100.0).floor() as usize + 1;
            for (element, indexes) in second_sequence_elements.iter() {
                if indexes.len() > test_len {
                    popular.insert(element.clone());
                }
            }
            for element in &popular {
                second_sequence_elements.remove(element);
            }
        }
        self.second_sequence_elements = second_sequence_elements;
        self.second_sequence_popular = popular;
    }

    pub fn find_longest_match(
        &self,
        first_start: usize,
        first_end: usize,
        second_start: usize,
        second_end: usize,
    ) -> Match {
        let first_sequence = &self.first_sequence;
        let second_sequence = &self.second_sequence;
        let second_sequence_elements = &self.second_sequence_elements;
        let (mut best_i, mut best_j, mut best_size) = (first_start, second_start, 0);
        let mut j2len: HashMap<usize, usize> = HashMap::new();
        for i in first_start..first_end {
            let mut new_j2len: HashMap<usize, usize> = HashMap::new();
            match second_sequence_elements.get(first_sequence.at_index(i).unwrap()) {
                Some(indexes) => {
                    for j in indexes {
                        let j = j.clone();
                        if j < second_start {
                            continue;
                        };
                        if j >= second_end {
                            break;
                        };
                        let mut size = 0;
                        if j > 0 {
                            match j2len.get(&(j - 1)) {
                                Some(k) => {
                                    size = k.clone();
                                }
                                None => {}
                            }
                        }
                        size += 1;
                        new_j2len.insert(j, size);
                        if size > best_size {
                            best_i = i + 1 - size;
                            best_j = j + 1 - size;
                            best_size = size;
                        }
                    }
                }
                None => {}
            }
            j2len = new_j2len;
        }
        for _ in 0..2 {
            while best_i > first_start
                && best_j > second_start
                && first_sequence.at_index(best_i - 1) == second_sequence.at_index(best_j - 1)
            {
                best_i = best_i - 1;
                best_j = best_j - 1;
                best_size = best_size + 1;
            }
            while best_i + best_size < first_end && best_j + best_size < second_end
                && first_sequence.at_index(best_i + best_size)
                    == second_sequence.at_index(best_j + best_size)
            {
                best_size += 1;
            }
        }
        Match::new(best_i, best_j, best_size)
    }

    pub fn get_matching_blocks(&mut self) -> Vec<Match> {
        if self.matching_blocks.as_ref().is_some() {
            return self.matching_blocks.as_ref().unwrap().clone();
        }
        let (first_length, second_length) = (self.first_sequence.len(), self.second_sequence.len());
        let mut matches = Vec::new();
        let mut queue = vec![(0, first_length, 0, second_length)];
        while !queue.is_empty() {
            let (first_start, first_end, second_start, second_end) = queue.pop().unwrap();
            let m = self.find_longest_match(first_start, first_end, second_start, second_end);
            match m.size {
                0 => {}
                _ => {
                    if first_start < m.first_start && second_start < m.second_start {
                        queue.push((first_start, m.first_start, second_start, m.second_start));
                    }
                    if m.first_start + m.size < first_end && m.second_start + m.size < second_end {
                        queue.push((
                            m.first_start + m.size,
                            first_end,
                            m.second_start + m.size,
                            second_end,
                        ));
                    }
                    matches.push(m);
                }
            }
        }
        matches.sort_by(|a, b| a.cmp(b));
        let (mut first_start, mut second_start, mut size) = (0, 0, 0);
        let mut non_adjacent = Vec::new();
        for m in &matches {
            if first_start + size == m.first_start && second_start + size == m.second_start {
                size += m.size
            } else {
                if size != 0 {
                    non_adjacent.push(Match::new(first_start, second_start, size));
                }
                first_start = m.first_start;
                second_start = m.second_start;
                size = m.size;
            }
        }
        if size != 0 {
            non_adjacent.push(Match::new(first_start, second_start, size));
        }
        non_adjacent.push(Match::new(first_length, second_length, 0));
        self.matching_blocks = Some(non_adjacent);
        self.matching_blocks.as_ref().unwrap().clone()
    }

    pub fn get_opcodes(&mut self) -> Vec<Opcode> {
        if self.opcodes.as_ref().is_some() {
            return self.opcodes.as_ref().unwrap().clone();
        }
        let mut opcodes = Vec::new();
        let (mut i, mut j) = (0, 0);
        for m in self.get_matching_blocks() {
            let mut tag = String::new();
            if i < m.first_start && j < m.second_start {
                tag = String::from("replace");
            } else if i < m.first_start {
                tag = String::from("delete");
            } else if j < m.second_start {
                tag = String::from("insert");
            }
            if !tag.is_empty() {
                opcodes.push(Opcode::new(tag, i, m.first_start, j, m.second_start));
            }
            i = m.first_start + m.size;
            j = m.second_start + m.size;
            if m.size != 0 {
                opcodes.push(Opcode::new(
                    String::from("equal"),
                    m.first_start,
                    i,
                    m.second_start,
                    j,
                ));
            }
        }
        self.opcodes = Some(opcodes);
        return self.opcodes.as_ref().unwrap().clone();
    }

    pub fn get_grouped_opcodes(&mut self, n: usize) -> Vec<Vec<Opcode>> {
        let mut res = Vec::new();
        let mut codes = self.get_opcodes();
        if codes.is_empty() {
            codes.push(Opcode::new("equal".to_string(), 0, 1, 0, 1));
        }

        if codes.first().unwrap().tag == "equal" {
            let opcode = codes.first_mut().unwrap();
            opcode.first_start = max(opcode.first_start, opcode.first_end.saturating_sub(n));
            opcode.second_start = max(opcode.second_start, opcode.second_end.saturating_sub(n));
        }
        if codes.last().unwrap().tag == "equal" {
            let opcode = codes.last_mut().unwrap();
            opcode.first_end = min(opcode.first_start + n, opcode.first_end);
            opcode.second_end = min(opcode.second_start + n, opcode.second_end);
        }
        let nn = n + n;
        let mut group = Vec::new();
        for code in &codes {
            let (mut first_start, mut second_start) = (code.first_start, code.second_start);
            if code.tag == "equal" && code.first_end - code.first_start > nn {
                group.push(Opcode::new(
                    code.tag.clone(),
                    code.first_start,
                    min(code.first_end, code.first_start + n),
                    code.second_start,
                    min(code.second_end, code.second_start + n),
                ));
                res.push(group.clone());
                group.clear();
                first_start = max(first_start, code.first_end.saturating_sub(n));
                second_start = max(second_start, code.second_end.saturating_sub(n));
            }
            group.push(Opcode::new(
                code.tag.clone(),
                first_start,
                code.first_end,
                second_start,
                code.second_end,
            ));
        }
        if !group.is_empty() && !(group.len() == 1 && group.first().unwrap().tag == "equal") {
            res.push(group.clone());
        }
        res
    }

    pub fn ratio(&mut self) -> f32 {
        let matches = self.get_matching_blocks()
            .iter()
            .fold(0, |res, &m| res + m.size);
        calculate_ratio(
            matches,
            self.first_sequence.len() + self.second_sequence.len(),
        )
    }
}
<commit_msg>Remove second_sequence_popular from the SequenceMatcher struct.<commit_after>use std::cmp::{max, min};
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;
use utils::calculate_ratio;

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
pub struct Match {
    pub first_start: usize,
    pub second_start: usize,
    pub size: usize,
}

impl Match {
    fn new(first_start: usize, second_start: usize, size: usize) -> Match {
        Match {
            first_start: first_start,
            second_start: second_start,
            size: size,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Opcode {
    pub tag: String,
    pub first_start: usize,
    pub first_end: usize,
    pub second_start: usize,
    pub second_end: usize,
}

impl Opcode {
    fn new(
        tag: String,
        first_start: usize,
        first_end: usize,
        second_start: usize,
        second_end: usize,
    ) -> Opcode {
        Opcode {
            tag: tag,
            first_start: first_start,
            first_end: first_end,
            second_start: second_start,
            second_end: second_end,
        }
    }
}

pub trait Sequence: Debug {
    fn len(&self) -> usize;
    fn at_index(&self, index: usize) -> Option<&str>;
}

impl Sequence for str {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index > self.len() {
            return None;
        }
        unsafe { Some(self.slice_unchecked(index, index + 1)) }
    }
}

impl<'a> Sequence for Vec<&'a str> {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index < self.len() {
            return Some(self[index]);
        }
        None
    }
}

impl Sequence for String {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index > self.len() {
            return None;
        }
        unsafe { Some(self.slice_unchecked(index, index + 1)) }
    }
}

impl<'a> Sequence for Vec<String> {
    fn len(&self) -> usize {
        self.len()
    }

    fn at_index(&self, index: usize) -> Option<&str> {
        if index < self.len() {
            return Some(&self[index]);
        }
        None
    }
}

pub struct SequenceMatcher<'a, T: 'a + ?Sized + Sequence> {
    first_sequence: &'a T,
    second_sequence: &'a T,
    matching_blocks: Option<Vec<Match>>,
    opcodes: Option<Vec<Opcode>>,
    is_junk: Option<fn(&str) -> bool>,
    second_sequence_elements: HashMap<&'a str, Vec<usize>>,
}

impl<'a, T: ?Sized + Sequence> SequenceMatcher<'a, T> {
    pub fn new(first_sequence: &'a T, second_sequence: &'a T) -> SequenceMatcher<'a, T> {
        let mut matcher = SequenceMatcher {
            first_sequence: first_sequence,
            second_sequence: second_sequence,
            matching_blocks: None,
            opcodes: None,
            is_junk: None,
            second_sequence_elements: HashMap::new(),
        };
        matcher.set_seqs(first_sequence, second_sequence);
        matcher
    }

    pub fn set_is_junk(&mut self, is_junk: Option<fn(&str) -> bool>) {
        self.is_junk = is_junk;
        self.set_second_seq(self.second_sequence);
    }

    pub fn set_seqs(&mut self, first_sequence: &'a T, second_sequence: &'a T) {
        self.set_first_seq(first_sequence);
        self.set_second_seq(second_sequence);
    }

    pub fn set_first_seq(&mut self, sequence: &'a T) {
        self.first_sequence = sequence;
        self.matching_blocks = None;
        self.opcodes = None;
    }

    pub fn set_second_seq(&mut self, sequence: &'a T) {
        self.second_sequence = sequence;
        self.matching_blocks = None;
        self.opcodes = None;
        self.chain_second_seq();
    }

    fn chain_second_seq(&mut self) {
        let second_sequence = self.second_sequence;
        let mut second_sequence_elements = HashMap::new();
        for i in 0..second_sequence.len() {
            let mut counter = second_sequence_elements
                .entry(second_sequence.at_index(i).unwrap())
                .or_insert(Vec::new());
            counter.push(i);
        }
        if self.is_junk.is_some() {
            let mut junk = Vec::new();
            for element in second_sequence_elements.keys() {
                if (self.is_junk.unwrap())(element) {
                    junk.push(element.clone());
                }
            }
            for element in &junk {
                second_sequence_elements.remove(element);
            }
        }
        let mut popular = HashSet::new();
        let len = second_sequence.len();
        if len >= 200 {
            let test_len = (len as f32 / 100.0).floor() as usize + 1;
            for (element, indexes) in second_sequence_elements.iter() {
                if indexes.len() > test_len {
                    popular.insert(element.clone());
                }
            }
            for element in &popular {
                second_sequence_elements.remove(element);
            }
        }
        self.second_sequence_elements = second_sequence_elements;
    }

    pub fn find_longest_match(
        &self,
        first_start: usize,
        first_end: usize,
        second_start: usize,
        second_end: usize,
    ) -> Match {
        let first_sequence = &self.first_sequence;
        let second_sequence = &self.second_sequence;
        let second_sequence_elements = &self.second_sequence_elements;
        let (mut best_i, mut best_j, mut best_size) = (first_start, second_start, 0);
        let mut j2len: HashMap<usize, usize> = HashMap::new();
        for i in first_start..first_end {
            let mut new_j2len: HashMap<usize, usize> = HashMap::new();
            match second_sequence_elements.get(first_sequence.at_index(i).unwrap()) {
                Some(indexes) => {
                    for j in indexes {
                        let j = j.clone();
                        if j < second_start {
                            continue;
                        };
                        if j >= second_end {
                            break;
                        };
                        let mut size = 0;
                        if j > 0 {
                            match j2len.get(&(j - 1)) {
                                Some(k) => {
                                    size = k.clone();
                                }
                                None => {}
                            }
                        }
                        size += 1;
                        new_j2len.insert(j, size);
                        if size > best_size {
                            best_i = i + 1 - size;
                            best_j = j + 1 - size;
                            best_size = size;
                        }
                    }
                }
                None => {}
            }
            j2len = new_j2len;
        }
        for _ in 0..2 {
            while best_i > first_start
                && best_j > second_start
                && first_sequence.at_index(best_i - 1) == second_sequence.at_index(best_j - 1)
            {
                best_i = best_i - 1;
                best_j = best_j - 1;
                best_size = best_size + 1;
            }
            while best_i + best_size < first_end && best_j + best_size < second_end
                && first_sequence.at_index(best_i + best_size)
                    == second_sequence.at_index(best_j + best_size)
            {
                best_size += 1;
            }
        }
        Match::new(best_i, best_j, best_size)
    }

    pub fn get_matching_blocks(&mut self) -> Vec<Match> {
        if self.matching_blocks.as_ref().is_some() {
            return self.matching_blocks.as_ref().unwrap().clone();
        }
        let (first_length, second_length) = (self.first_sequence.len(), self.second_sequence.len());
        let mut matches = Vec::new();
        let mut queue = vec![(0, first_length, 0, second_length)];
        while !queue.is_empty() {
            let (first_start, first_end, second_start, second_end) = queue.pop().unwrap();
            let m = self.find_longest_match(first_start, first_end, second_start, second_end);
            match m.size {
                0 => {}
                _ => {
                    if first_start < m.first_start && second_start < m.second_start {
                        queue.push((first_start, m.first_start, second_start, m.second_start));
                    }
                    if m.first_start + m.size < first_end && m.second_start + m.size < second_end {
                        queue.push((
                            m.first_start + m.size,
                            first_end,
                            m.second_start + m.size,
                            second_end,
                        ));
                    }
                    matches.push(m);
                }
            }
        }
        matches.sort_by(|a, b| a.cmp(b));
        let (mut first_start, mut second_start, mut size) = (0, 0, 0);
        let mut non_adjacent = Vec::new();
        for m in &matches {
            if first_start + size == m.first_start && second_start + size == m.second_start {
                size += m.size
            } else {
                if size != 0 {
                    non_adjacent.push(Match::new(first_start, second_start, size));
                }
                first_start = m.first_start;
                second_start = m.second_start;
                size = m.size;
            }
        }
        if size != 0 {
            non_adjacent.push(Match::new(first_start, second_start, size));
        }
        non_adjacent.push(Match::new(first_length, second_length, 0));
        self.matching_blocks = Some(non_adjacent);
        self.matching_blocks.as_ref().unwrap().clone()
    }

    pub fn get_opcodes(&mut self) -> Vec<Opcode> {
        if self.opcodes.as_ref().is_some() {
            return self.opcodes.as_ref().unwrap().clone();
        }
        let mut opcodes = Vec::new();
        let (mut i, mut j) = (0, 0);
        for m in self.get_matching_blocks() {
            let mut tag = String::new();
            if i < m.first_start && j < m.second_start {
                tag = String::from("replace");
            } else if i < m.first_start {
                tag = String::from("delete");
            } else if j < m.second_start {
                tag = String::from("insert");
            }
            if !tag.is_empty() {
                opcodes.push(Opcode::new(tag, i, m.first_start, j, m.second_start));
            }
            i = m.first_start + m.size;
            j = m.second_start + m.size;
            if m.size != 0 {
                opcodes.push(Opcode::new(
                    String::from("equal"),
                    m.first_start,
                    i,
                    m.second_start,
                    j,
                ));
            }
        }
        self.opcodes = Some(opcodes);
        return self.opcodes.as_ref().unwrap().clone();
    }

    pub fn get_grouped_opcodes(&mut self, n: usize) -> Vec<Vec<Opcode>> {
        let mut res = Vec::new();
        let mut codes = self.get_opcodes();
        if codes.is_empty() {
            codes.push(Opcode::new("equal".to_string(), 0, 1, 0, 1));
        }

        if codes.first().unwrap().tag == "equal" {
            let opcode = codes.first_mut().unwrap();
            opcode.first_start = max(opcode.first_start, opcode.first_end.saturating_sub(n));
            opcode.second_start = max(opcode.second_start, opcode.second_end.saturating_sub(n));
        }
        if codes.last().unwrap().tag == "equal" {
            let opcode = codes.last_mut().unwrap();
            opcode.first_end = min(opcode.first_start + n, opcode.first_end);
            opcode.second_end = min(opcode.second_start + n, opcode.second_end);
        }
        let nn = n + n;
        let mut group = Vec::new();
        for code in &codes {
            let (mut first_start, mut second_start) = (code.first_start, code.second_start);
            if code.tag == "equal" && code.first_end - code.first_start > nn {
                group.push(Opcode::new(
                    code.tag.clone(),
                    code.first_start,
                    min(code.first_end, code.first_start + n),
                    code.second_start,
                    min(code.second_end, code.second_start + n),
                ));
                res.push(group.clone());
                group.clear();
                first_start = max(first_start, code.first_end.saturating_sub(n));
                second_start = max(second_start, code.second_end.saturating_sub(n));
            }
            group.push(Opcode::new(
                code.tag.clone(),
                first_start,
                code.first_end,
                second_start,
                code.second_end,
            ));
        }
        if !group.is_empty() && !(group.len() == 1 && group.first().unwrap().tag == "equal") {
            res.push(group.clone());
        }
        res
    }

    pub fn ratio(&mut self) -> f32 {
        let matches = self.get_matching_blocks()
            .iter()
            .fold(0, |res, &m| res + m.size);
        calculate_ratio(
            matches,
            self.first_sequence.len() + self.second_sequence.len(),
        )
    }
}
<|endoftext|>


<commit_before>use backend::*;
use backend::collecting::Collecting;
use backend::for_context::ForContext;
use backend::package_processor::PackageProcessor;
use codeviz::rust::*;
use core::*;
use errors::*;
use naming::{self, FromNaming};
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::rc::Rc;

const EXT: &str = "rs";
const RUST_CONTEXT: &str = "rust";

pub trait Listeners {
    fn configure(&self, _processor: &mut ProcessorOptions) -> Result<()> {
        Ok(())
    }
}

/// A vector of listeners is a valid listener.
impl Listeners for Vec<Box<Listeners>> {
    fn configure(&self, processor: &mut ProcessorOptions) -> Result<()> {
        for listeners in self {
            listeners.configure(processor)?;
        }

        Ok(())
    }
}

pub struct ProcessorOptions {
}

impl ProcessorOptions {
    pub fn new() -> ProcessorOptions {
        ProcessorOptions {}
    }
}

pub struct Processor {
    env: Environment,
    out_path: PathBuf,
    id_converter: Option<Box<naming::Naming>>,
    package_prefix: Option<RpPackage>,
    listeners: Box<Listeners>,
    to_lower_snake: Box<naming::Naming>,
    hash_map: ImportedName,
    json_value: ImportedName,
}

impl Processor {
    pub fn new(_options: ProcessorOptions,
               env: Environment,
               out_path: PathBuf,
               id_converter: Option<Box<naming::Naming>>,
               package_prefix: Option<RpPackage>,
               listeners: Box<Listeners>)
               -> Processor {
        Processor {
            env: env,
            out_path: out_path,
            id_converter: id_converter,
            package_prefix: package_prefix,
            listeners: listeners,
            to_lower_snake: naming::SnakeCase::new().to_lower_snake(),
            hash_map: Name::imported("std::collections", "HashMap"),
            json_value: Name::imported_alias("serde_json", "Value", "json"),
        }
    }

    fn ident(&self, name: &str) -> String {
        if let Some(ref id_converter) = self.id_converter {
            id_converter.convert(name)
        } else {
            name.to_owned()
        }
    }

    fn convert_type_id(&self, pos: &RpPos, type_id: &RpTypeId) -> Result<Name> {
        let (package, registered) = self.env
            .lookup(&type_id.package, &type_id.name)
            .map_err(|e| Error::pos(e.description().to_owned(), pos.clone()))?;

        let name = registered.name().join(".");

        if let Some(ref prefix) = type_id.name.prefix {
            let package_name = self.package(package).parts.join("::");
            return Ok(Name::Imported(Name::imported_alias(&package_name, &name, prefix)));
        }

        Ok(Name::Local(Name::local(&name)))
    }

    fn into_type(&self, type_id: &RpTypeId, field: &RpLoc<RpField>) -> Result<Statement> {
        let stmt = self.into_rust_type(type_id, field.pos(), &field.ty)?;

        if field.is_optional() {
            return Ok(stmt!["Option<", stmt, ">"]);
        }

        Ok(stmt)
    }

    pub fn into_rust_type(&self,
                          type_id: &RpTypeId,
                          pos: &RpPos,
                          ty: &RpType)
                          -> Result<Statement> {
        let ty = match *ty {
            RpType::String => stmt!["String"],
            RpType::Signed { ref size } => {
                if size.map(|s| s <= 32usize).unwrap_or(true) {
                    stmt!["i32"]
                } else {
                    stmt!["i64"]
                }
            }
            RpType::Unsigned { ref size } => {
                if size.map(|s| s <= 32usize).unwrap_or(true) {
                    stmt!["u32"]
                } else {
                    stmt!["u64"]
                }
            }
            RpType::Float => stmt!["f32"],
            RpType::Double => stmt!["f64"],
            RpType::Boolean => stmt!["bool"],
            RpType::Array { ref inner } => {
                let argument = self.into_rust_type(type_id, pos, inner)?;
                stmt!["Vec<", argument, ">"]
            }
            RpType::Name { ref name } => {
                let type_id = type_id.with_name(name.clone());
                stmt![self.convert_type_id(pos, &type_id)?]
            }
            RpType::Map { ref key, ref value } => {
                let key = self.into_rust_type(type_id, pos, key)?;
                let value = self.into_rust_type(type_id, pos, value)?;
                stmt![&self.hash_map, "<", key, ", ", value, ">"]
            }
            RpType::Any => stmt![&self.json_value],
            ref t => {
                return Err(Error::pos(format!("unsupported type: {:?}", t), pos.clone()));
            }
        };

        Ok(ty)
    }

    // Build the corresponding element out of a field declaration.
    fn field_element(&self, type_id: &RpTypeId, field: &RpLoc<RpField>) -> Result<Element> {
        let mut elements = Elements::new();

        let ident = self.ident(field.ident());
        let type_spec = self.into_type(type_id, field)?;

        if field.is_optional() {
            elements.push(stmt!["#[serde(skip_serializing_if=\"Option::is_none\")]"]);
        }

        if field.name() != ident {
            elements.push(stmt!["#[serde(rename = ", Variable::String(field.name().to_owned()), ")]"]);
        }

        elements.push(stmt![ident, ": ", type_spec, ","]);

        Ok(elements.into())
    }

    fn write_mod_files(&self, files: &BTreeMap<&RpVersionedPackage, FileSpec>) -> Result<()> {
        let mut packages: BTreeMap<PathBuf, BTreeSet<String>> = BTreeMap::new();

        for (key, _) in files {
            let mut current = self.out_path().to_owned();

            let mut it = self.package(key).parts.into_iter().peekable();

            while let Some(part) = it.next() {
                current = current.join(part);

                if let Some(next) = it.peek() {
                    let mut full_path = current.clone();

                    full_path = full_path.join("mod");
                    full_path.set_extension(self.ext());

                    packages.entry(full_path)
                        .or_insert_with(BTreeSet::new)
                        .insert(next.clone());
                }
            }
        }

        for (full_path, children) in packages {
            if let Some(parent) = full_path.parent() {
                if !parent.is_dir() {
                    debug!("+dir: {}", parent.display());
                    fs::create_dir_all(parent)?;
                }
            }

            if !full_path.is_file() {
                debug!("+mod: {}", full_path.display());
                let mut f = File::create(full_path)?;

                for child in children {
                    writeln!(f, "pub mod {};", child)?;
                }
            }
        }

        Ok(())
    }
}

impl Backend for Processor {
    fn process(&self) -> Result<()> {
        let files = self.populate_files()?;
        self.write_mod_files(&files)?;
        self.write_files(files)
    }

    fn verify(&self) -> Result<Vec<Error>> {
        Ok(vec![])
    }
}

impl Collecting for FileSpec {
    type Processor = Processor;

    fn new() -> Self {
        FileSpec::new()
    }

    fn into_bytes(self, _: &Self::Processor) -> Result<Vec<u8>> {
        let mut out = String::new();
        self.format(&mut out)?;
        Ok(out.into_bytes())
    }
}

impl PackageProcessor for Processor {
    type Out = FileSpec;

    fn ext(&self) -> &str {
        EXT
    }

    fn env(&self) -> &Environment {
        &self.env
    }

    fn package_prefix(&self) -> &Option<RpPackage> {
        &self.package_prefix
    }

    fn out_path(&self) -> &Path {
        &self.out_path
    }

    fn default_process(&self, _out: &mut Self::Out, _type_id: &RpTypeId, _: &RpPos) -> Result<()> {
        Ok(())
    }

    fn process_tuple(&self,
                     out: &mut Self::Out,
                     type_id: &RpTypeId,
                     _: &RpPos,
                     body: Rc<RpTupleBody>)
                     -> Result<()> {
        let mut fields = Statement::new();

        for field in &body.fields {
            fields.push(self.into_type(type_id, field)?);
        }

        let mut elements = Elements::new();
        elements.push("#[derive(Serialize, Deserialize, Debug)]");
        elements.push(stmt!["struct ", &body.name, "(", fields.join(", "), ");"]);

        out.push(elements);
        Ok(())
    }

    fn process_enum(&self,
                    out: &mut Self::Out,
                    _: &RpTypeId,
                    _: &RpPos,
                    body: Rc<RpEnumBody>)
                    -> Result<()> {
        let mut enum_spec = EnumSpec::new(&body.name);
        enum_spec.public();

        for code in body.codes.for_context(RUST_CONTEXT) {
            enum_spec.push(code.move_inner().lines);
        }

        out.push(enum_spec);
        Ok(())
    }

    fn process_type(&self,
                    out: &mut Self::Out,
                    type_id: &RpTypeId,
                    _: &RpPos,
                    body: Rc<RpTypeBody>)
                    -> Result<()> {
        let mut fields = Elements::new();

        for field in &body.fields {
            fields.push(self.field_element(type_id, field)?);
        }

        let mut struct_spec = StructSpec::new(&body.name);
        struct_spec.public();

        struct_spec.push_attribute("#[derive(Serialize, Deserialize, Debug)]");
        struct_spec.push(fields);

        for code in body.codes.for_context(RUST_CONTEXT) {
            struct_spec.push(code.move_inner().lines);
        }

        out.push(struct_spec);
        Ok(())
    }

    fn process_interface(&self,
                         out: &mut Self::Out,
                         type_id: &RpTypeId,
                         _: &RpPos,
                         body: Rc<RpInterfaceBody>)
                         -> Result<()> {
        let mut enum_spec = EnumSpec::new(&body.name);
        enum_spec.public();

        enum_spec.push_attribute("#[derive(Serialize, Deserialize, Debug)]");
        enum_spec.push_attribute("#[serde(tag = \"type\")]");

        for code in body.codes.for_context(RUST_CONTEXT) {
            enum_spec.push(code.move_inner().lines);
        }

        for (_, ref sub_type) in &body.sub_types {
            let mut elements = Elements::new();

            if let Some(name) = sub_type.names.first() {
                elements.push(stmt!["#[serde(rename = ",
                                    Variable::String((**name).to_owned()),
                                    ")]"]);
            }

            elements.push(stmt![&sub_type.name, " {"]);

            for field in body.fields.iter().chain(sub_type.fields.iter()) {
                elements.push_nested(self.field_element(type_id, field)?);
            }

            elements.push("},");

            enum_spec.push(elements);
        }

        out.push(enum_spec);

        Ok(())
    }
}
<commit_msg>also build root modules<commit_after>use backend::*;
use backend::collecting::Collecting;
use backend::for_context::ForContext;
use backend::package_processor::PackageProcessor;
use codeviz::rust::*;
use core::*;
use errors::*;
use naming::{self, FromNaming};
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::rc::Rc;

const MOD: &str = "mod";
const EXT: &str = "rs";
const RUST_CONTEXT: &str = "rust";

pub trait Listeners {
    fn configure(&self, _processor: &mut ProcessorOptions) -> Result<()> {
        Ok(())
    }
}

/// A vector of listeners is a valid listener.
impl Listeners for Vec<Box<Listeners>> {
    fn configure(&self, processor: &mut ProcessorOptions) -> Result<()> {
        for listeners in self {
            listeners.configure(processor)?;
        }

        Ok(())
    }
}

pub struct ProcessorOptions {
}

impl ProcessorOptions {
    pub fn new() -> ProcessorOptions {
        ProcessorOptions {}
    }
}

pub struct Processor {
    env: Environment,
    out_path: PathBuf,
    id_converter: Option<Box<naming::Naming>>,
    package_prefix: Option<RpPackage>,
    listeners: Box<Listeners>,
    to_lower_snake: Box<naming::Naming>,
    hash_map: ImportedName,
    json_value: ImportedName,
}

impl Processor {
    pub fn new(_options: ProcessorOptions,
               env: Environment,
               out_path: PathBuf,
               id_converter: Option<Box<naming::Naming>>,
               package_prefix: Option<RpPackage>,
               listeners: Box<Listeners>)
               -> Processor {
        Processor {
            env: env,
            out_path: out_path,
            id_converter: id_converter,
            package_prefix: package_prefix,
            listeners: listeners,
            to_lower_snake: naming::SnakeCase::new().to_lower_snake(),
            hash_map: Name::imported("std::collections", "HashMap"),
            json_value: Name::imported_alias("serde_json", "Value", "json"),
        }
    }

    fn ident(&self, name: &str) -> String {
        if let Some(ref id_converter) = self.id_converter {
            id_converter.convert(name)
        } else {
            name.to_owned()
        }
    }

    fn convert_type_id(&self, pos: &RpPos, type_id: &RpTypeId) -> Result<Name> {
        let (package, registered) = self.env
            .lookup(&type_id.package, &type_id.name)
            .map_err(|e| Error::pos(e.description().to_owned(), pos.clone()))?;

        let name = registered.name().join(".");

        if let Some(ref prefix) = type_id.name.prefix {
            let package_name = self.package(package).parts.join("::");
            return Ok(Name::Imported(Name::imported_alias(&package_name, &name, prefix)));
        }

        Ok(Name::Local(Name::local(&name)))
    }

    fn into_type(&self, type_id: &RpTypeId, field: &RpLoc<RpField>) -> Result<Statement> {
        let stmt = self.into_rust_type(type_id, field.pos(), &field.ty)?;

        if field.is_optional() {
            return Ok(stmt!["Option<", stmt, ">"]);
        }

        Ok(stmt)
    }

    pub fn into_rust_type(&self,
                          type_id: &RpTypeId,
                          pos: &RpPos,
                          ty: &RpType)
                          -> Result<Statement> {
        let ty = match *ty {
            RpType::String => stmt!["String"],
            RpType::Signed { ref size } => {
                if size.map(|s| s <= 32usize).unwrap_or(true) {
                    stmt!["i32"]
                } else {
                    stmt!["i64"]
                }
            }
            RpType::Unsigned { ref size } => {
                if size.map(|s| s <= 32usize).unwrap_or(true) {
                    stmt!["u32"]
                } else {
                    stmt!["u64"]
                }
            }
            RpType::Float => stmt!["f32"],
            RpType::Double => stmt!["f64"],
            RpType::Boolean => stmt!["bool"],
            RpType::Array { ref inner } => {
                let argument = self.into_rust_type(type_id, pos, inner)?;
                stmt!["Vec<", argument, ">"]
            }
            RpType::Name { ref name } => {
                let type_id = type_id.with_name(name.clone());
                stmt![self.convert_type_id(pos, &type_id)?]
            }
            RpType::Map { ref key, ref value } => {
                let key = self.into_rust_type(type_id, pos, key)?;
                let value = self.into_rust_type(type_id, pos, value)?;
                stmt![&self.hash_map, "<", key, ", ", value, ">"]
            }
            RpType::Any => stmt![&self.json_value],
            ref t => {
                return Err(Error::pos(format!("unsupported type: {:?}", t), pos.clone()));
            }
        };

        Ok(ty)
    }

    // Build the corresponding element out of a field declaration.
    fn field_element(&self, type_id: &RpTypeId, field: &RpLoc<RpField>) -> Result<Element> {
        let mut elements = Elements::new();

        let ident = self.ident(field.ident());
        let type_spec = self.into_type(type_id, field)?;

        if field.is_optional() {
            elements.push(stmt!["#[serde(skip_serializing_if=\"Option::is_none\")]"]);
        }

        if field.name() != ident {
            elements.push(stmt!["#[serde(rename = ", Variable::String(field.name().to_owned()), ")]"]);
        }

        elements.push(stmt![ident, ": ", type_spec, ","]);

        Ok(elements.into())
    }

    fn write_mod_files(&self, files: &BTreeMap<&RpVersionedPackage, FileSpec>) -> Result<()> {
        let mut packages: BTreeMap<PathBuf, BTreeSet<String>> = BTreeMap::new();
        let mut root_names = BTreeSet::new();

        for (key, _) in files {
            let mut current = self.out_path().to_owned();

            let mut it = self.package(key).parts.into_iter().peekable();

            if let Some(root) = it.peek() {
                root_names.insert(root.to_owned());
            }

            while let Some(part) = it.next() {
                current = current.join(part);

                if let Some(next) = it.peek() {
                    let mut full_path = current.join(MOD);
                    full_path.set_extension(self.ext());

                    packages.entry(full_path)
                        .or_insert_with(BTreeSet::new)
                        .insert(next.clone());
                }
            }
        }

        let mut root_mod = self.out_path().join(MOD);
        root_mod.set_extension(self.ext());
        packages.insert(root_mod, root_names);

        for (full_path, children) in packages {
            if let Some(parent) = full_path.parent() {
                if !parent.is_dir() {
                    debug!("+dir: {}", parent.display());
                    fs::create_dir_all(parent)?;
                }
            }

            if !full_path.is_file() {
                debug!("+mod: {}", full_path.display());
                let mut f = File::create(full_path)?;

                for child in children {
                    writeln!(f, "pub mod {};", child)?;
                }
            }
        }

        Ok(())
    }
}

impl Backend for Processor {
    fn process(&self) -> Result<()> {
        let files = self.populate_files()?;
        self.write_mod_files(&files)?;
        self.write_files(files)
    }

    fn verify(&self) -> Result<Vec<Error>> {
        Ok(vec![])
    }
}

impl Collecting for FileSpec {
    type Processor = Processor;

    fn new() -> Self {
        FileSpec::new()
    }

    fn into_bytes(self, _: &Self::Processor) -> Result<Vec<u8>> {
        let mut out = String::new();
        self.format(&mut out)?;
        Ok(out.into_bytes())
    }
}

impl PackageProcessor for Processor {
    type Out = FileSpec;

    fn ext(&self) -> &str {
        EXT
    }

    fn env(&self) -> &Environment {
        &self.env
    }

    fn package_prefix(&self) -> &Option<RpPackage> {
        &self.package_prefix
    }

    fn out_path(&self) -> &Path {
        &self.out_path
    }

    fn default_process(&self, _out: &mut Self::Out, _type_id: &RpTypeId, _: &RpPos) -> Result<()> {
        Ok(())
    }

    fn process_tuple(&self,
                     out: &mut Self::Out,
                     type_id: &RpTypeId,
                     _: &RpPos,
                     body: Rc<RpTupleBody>)
                     -> Result<()> {
        let mut fields = Statement::new();

        for field in &body.fields {
            fields.push(self.into_type(type_id, field)?);
        }

        let mut elements = Elements::new();
        elements.push("#[derive(Serialize, Deserialize, Debug)]");
        elements.push(stmt!["struct ", &body.name, "(", fields.join(", "), ");"]);

        out.push(elements);
        Ok(())
    }

    fn process_enum(&self,
                    out: &mut Self::Out,
                    _: &RpTypeId,
                    _: &RpPos,
                    body: Rc<RpEnumBody>)
                    -> Result<()> {
        let mut enum_spec = EnumSpec::new(&body.name);
        enum_spec.public();

        for code in body.codes.for_context(RUST_CONTEXT) {
            enum_spec.push(code.move_inner().lines);
        }

        out.push(enum_spec);
        Ok(())
    }

    fn process_type(&self,
                    out: &mut Self::Out,
                    type_id: &RpTypeId,
                    _: &RpPos,
                    body: Rc<RpTypeBody>)
                    -> Result<()> {
        let mut fields = Elements::new();

        for field in &body.fields {
            fields.push(self.field_element(type_id, field)?);
        }

        let mut struct_spec = StructSpec::new(&body.name);
        struct_spec.public();

        struct_spec.push_attribute("#[derive(Serialize, Deserialize, Debug)]");
        struct_spec.push(fields);

        for code in body.codes.for_context(RUST_CONTEXT) {
            struct_spec.push(code.move_inner().lines);
        }

        out.push(struct_spec);
        Ok(())
    }

    fn process_interface(&self,
                         out: &mut Self::Out,
                         type_id: &RpTypeId,
                         _: &RpPos,
                         body: Rc<RpInterfaceBody>)
                         -> Result<()> {
        let mut enum_spec = EnumSpec::new(&body.name);
        enum_spec.public();

        enum_spec.push_attribute("#[derive(Serialize, Deserialize, Debug)]");
        enum_spec.push_attribute("#[serde(tag = \"type\")]");

        for code in body.codes.for_context(RUST_CONTEXT) {
            enum_spec.push(code.move_inner().lines);
        }

        for (_, ref sub_type) in &body.sub_types {
            let mut elements = Elements::new();

            if let Some(name) = sub_type.names.first() {
                elements.push(stmt!["#[serde(rename = ",
                                    Variable::String((**name).to_owned()),
                                    ")]"]);
            }

            elements.push(stmt![&sub_type.name, " {"]);

            for field in body.fields.iter().chain(sub_type.fields.iter()) {
                elements.push_nested(self.field_element(type_id, field)?);
            }

            elements.push("},");

            enum_spec.push(elements);
        }

        out.push(enum_spec);

        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate postgres;
extern crate getopts;

use postgres::{Connection, TlsMode};

use std::env;
use getopts::Options;
use std::process;

extern crate pombase;

use pombase::db::*;
use pombase::web::config::*;
use pombase::web::data_build::*;
use pombase::interpro::parse_interpro;

const PKG_NAME: &str = env!("CARGO_PKG_NAME");
const VERSION: &str = env!("CARGO_PKG_VERSION");

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    print!("{} v{}\n", PKG_NAME, VERSION);

    let args: Vec<String> = env::args().collect();
    let mut opts = Options::new();

    opts.optflag("h", "help", "print this help message");
    opts.optopt("c", "config-file", "Configuration file name", "CONFIG");
    opts.optopt("p", "postgresql-connection-string",
                "PostgresSQL connection string like: postgres://user:pass@host/db_name",
                "CONN_STR");
    opts.optopt("i", "interpro-data-file",
                "The name of the file generated by 'pombase-interpro'", "FILE");
    opts.optopt("d", "output-directory",
                "Destination directory for the output", "DIR");
    opts.optflag("j", "store-json",
                 "optionally create a 'web_json' schema to store the generated JSON in the database");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!("Invalid options\n{}", f)
    };

    let program = args[0].clone();

    if matches.opt_present("help") {
        print_usage(&program, opts);
        process::exit(0);
    }

    if !matches.opt_present("config-file") {
        print!("no -c|--config-file option\n");
        print_usage(&program, opts);
        process::exit(1);
    }
    if !matches.opt_present("postgresql-connection-string") {
        print!("no -p|--postgresql-connection-string option\n");
        print_usage(&program, opts);
        process::exit(1);
    }
    if !matches.opt_present("output-directory") {
        print!("no -d|--output-directory option\n");
        print_usage(&program, opts);
        process::exit(1);
    }

    let config = Config::read(&matches.opt_str("c").unwrap());
    let connection_string = matches.opt_str("p").unwrap();
    let interpro_json = matches.opt_str("i").unwrap();
    let output_dir = matches.opt_str("d").unwrap();

    let conn = match Connection::connect(connection_string.as_str(), TlsMode::None) {
        Ok(conn) => conn,
        Err(err) => panic!("failed to connect using: {}, err: {}", connection_string, err)
    };

    let raw = Raw::new(&conn);
    let interpro_data = parse_interpro(&config, &interpro_json);
    let web_data_build = WebDataBuild::new(&raw, &interpro_data, &config);
    let web_data = web_data_build.get_web_data();

    match web_data.write(&config, &output_dir) {
        Ok(_) => (),
        Err(e) => {
            panic!("error while writing: {}", e);
        },
    }

    if matches.opt_present("store-json") {
        conn.execute("DROP SCHEMA IF EXISTS web_json CASCADE", &[]).unwrap();
        conn.execute("CREATE SCHEMA web_json", &[]).unwrap();
        conn.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;", &[]).unwrap();
        conn.execute("CREATE TABLE web_json.gene (uniquename TEXT, data JSONB)", &[]).unwrap();
        conn.execute("CREATE TABLE web_json.term (termid TEXT, data JSONB)", &[]).unwrap();
        conn.execute("CREATE TABLE web_json.reference (uniquename TEXT, data JSONB)", &[]).unwrap();

        web_data.store_jsonb(&conn);

        print!("stored results as JSONB using {}\n", &connection_string);
    }
}
<commit_msg>Return Result from main()<commit_after>extern crate postgres;
extern crate getopts;

use postgres::{Connection, TlsMode};

use std::env;
use getopts::Options;
use std::process;

extern crate pombase;

use pombase::db::*;
use pombase::web::config::*;
use pombase::web::data_build::*;
use pombase::interpro::parse_interpro;

const PKG_NAME: &str = env!("CARGO_PKG_NAME");
const VERSION: &str = env!("CARGO_PKG_VERSION");

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn main() -> Result<(), std::io::Error> {
    print!("{} v{}\n", PKG_NAME, VERSION);

    let args: Vec<String> = env::args().collect();
    let mut opts = Options::new();

    opts.optflag("h", "help", "print this help message");
    opts.optopt("c", "config-file", "Configuration file name", "CONFIG");
    opts.optopt("p", "postgresql-connection-string",
                "PostgresSQL connection string like: postgres://user:pass@host/db_name",
                "CONN_STR");
    opts.optopt("i", "interpro-data-file",
                "The name of the file generated by 'pombase-interpro'", "FILE");
    opts.optopt("d", "output-directory",
                "Destination directory for the output", "DIR");
    opts.optflag("j", "store-json",
                 "optionally create a 'web_json' schema to store the generated JSON in the database");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!("Invalid options\n{}", f)
    };

    let program = args[0].clone();

    if matches.opt_present("help") {
        print_usage(&program, opts);
        process::exit(0);
    }

    if !matches.opt_present("config-file") {
        print!("no -c|--config-file option\n");
        print_usage(&program, opts);
        process::exit(1);
    }
    if !matches.opt_present("postgresql-connection-string") {
        print!("no -p|--postgresql-connection-string option\n");
        print_usage(&program, opts);
        process::exit(1);
    }
    if !matches.opt_present("output-directory") {
        print!("no -d|--output-directory option\n");
        print_usage(&program, opts);
        process::exit(1);
    }

    let config = Config::read(&matches.opt_str("c").unwrap());
    let connection_string = matches.opt_str("p").unwrap();
    let interpro_json = matches.opt_str("i").unwrap();
    let output_dir = matches.opt_str("d").unwrap();

    let conn = match Connection::connect(connection_string.as_str(), TlsMode::None) {
        Ok(conn) => conn,
        Err(err) => panic!("failed to connect using: {}, err: {}", connection_string, err)
    };

    let raw = Raw::new(&conn);
    let interpro_data = parse_interpro(&config, &interpro_json);
    let web_data_build = WebDataBuild::new(&raw, &interpro_data, &config);
    let web_data = web_data_build.get_web_data();

    match web_data.write(&config, &output_dir) {
        Ok(_) => (),
        Err(e) => {
            panic!("error while writing: {}", e);
        },
    }

    if matches.opt_present("store-json") {
        conn.execute("DROP SCHEMA IF EXISTS web_json CASCADE", &[])?;
        conn.execute("CREATE SCHEMA web_json", &[])?;
        conn.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;", &[])?;
        conn.execute("CREATE TABLE web_json.gene (uniquename TEXT, data JSONB)", &[])?;
        conn.execute("CREATE TABLE web_json.term (termid TEXT, data JSONB)", &[])?;
        conn.execute("CREATE TABLE web_json.reference (uniquename TEXT, data JSONB)", &[])?;

        web_data.store_jsonb(&conn);

        print!("stored results as JSONB using {}\n", &connection_string);
    }

    Ok(())
}
<|endoftext|>


<commit_before>use primal_bit::BitVec;
use std::{cmp, ptr};

use wheel;
use super::{StreamingSieve};

const MINIMUM_PRESIEVE: usize = 2 * 3 * 5;
const PRESIEVE_PRIMES: &'static [usize] = &[7, 11, 13, 17, 19, 23, 29];

#[derive(Debug)]
pub struct Presieve {
    sieve: BitVec,
    presieve_prod: usize,
    presieve_idx: usize,
}
impl Presieve {
    #[inline(never)]
    pub fn new(limit_bits: usize) -> Presieve {
        let mut prod = MINIMUM_PRESIEVE;
        let mut idx = 0;
        for (i, &x) in PRESIEVE_PRIMES.iter().enumerate() {
            let new_prod = prod * x;
            if wheel::bits_for(new_prod) > limit_bits {
                break
            }
            prod = new_prod;
            idx = i;
        }

        let len = cmp::min(wheel::bits_for(prod), limit_bits);

        if idx == 0 {
            Presieve {
                sieve: BitVec::new(),
                presieve_prod: prod,
                presieve_idx: idx,
            }
        } else {
            let mut sievers = vec![];
            for &x in &PRESIEVE_PRIMES[..idx] {
                let (use_, _idx) = wheel::bit_index(x);
                if use_ {
                    sievers.push(wheel::State::new(wheel::Wheel30, x, prod));
                }
            }
            let mut sieve =  BitVec::from_elem(len, true);
            StreamingSieve::small_primes_sieve(&mut sieve, &mut sievers);

            Presieve {
                sieve: sieve,
                presieve_prod: prod,
                presieve_idx: idx,
            }
        }
    }
    pub fn smallest_unincluded_prime(&self) -> usize {
        PRESIEVE_PRIMES[self.presieve_idx]
    }
    pub fn mark_small_primes(&self, sieve: &mut BitVec) {
        for &x in &PRESIEVE_PRIMES[..self.presieve_idx] {
            let (use_, idx) = wheel::bit_index(x);
            if use_ {
                sieve.set(idx, true)
            }
        }
    }
    pub fn apply(&self, sieve: &mut BitVec, low: usize) {
        if self.sieve.len() == 0 {
            return
        }
        let offset = (low % self.presieve_prod) * wheel::BYTE_SIZE / wheel::BYTE_MODULO / 8;

        copy_all(sieve.as_bytes_mut(),
                 self.sieve.as_bytes(),
                 offset);
        fn copy_all(mut dst: &mut [u8], src: &[u8], init_offset: usize) {
            let mut pos = 0;
            // pre-fill data at the start, as a rotation of `src`.
            pos += memcpy(&mut dst[pos..], &src[init_offset..]);
            if pos >= dst.len() { return }

            pos += memcpy(&mut dst[pos..], &src[..init_offset]);
            if pos >= dst.len() { return }

            // progressively copy the prefix to the rest of the
            // vector, doubling each time.
            while pos < dst.len() {
                let (filled, unfilled) = dst.split_at_mut(pos);
                pos += memcpy(unfilled, filled);
            }
        }
        fn memcpy<'d>(dst: &'d mut [u8], src: &[u8]) -> usize {
            let l = cmp::min(dst.len(), src.len());
            unsafe {
                ptr::copy(src.as_ptr(), dst.as_mut_ptr(), l);
            }
            l
        }
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use test::Bencher;
    use super::Presieve;
    fn run_presieve(b: &mut Bencher, n: usize) {
        b.iter(|| Presieve::new(::wheel::bits_for(n)))
    }
    #[bench]
    fn presieve_small(b: &mut Bencher) {
        run_presieve(b, 100)
    }
    #[bench]
    fn presieve_medium(b: &mut Bencher) {
        run_presieve(b, 10_000)
    }
    #[bench]
    fn presieve_large(b: &mut Bencher) {
        run_presieve(b, 100_000)
    }
    #[bench]
    fn presieve_larger(b: &mut Bencher) {
        run_presieve(b, 1_000_000)
    }
    #[bench]
    fn presieve_huge(b: &mut Bencher) {
        run_presieve(b, 10_000_000)
    }
}
<commit_msg>sieve: fix an unused mut<commit_after>use primal_bit::BitVec;
use std::{cmp, ptr};

use wheel;
use super::{StreamingSieve};

const MINIMUM_PRESIEVE: usize = 2 * 3 * 5;
const PRESIEVE_PRIMES: &'static [usize] = &[7, 11, 13, 17, 19, 23, 29];

#[derive(Debug)]
pub struct Presieve {
    sieve: BitVec,
    presieve_prod: usize,
    presieve_idx: usize,
}
impl Presieve {
    #[inline(never)]
    pub fn new(limit_bits: usize) -> Presieve {
        let mut prod = MINIMUM_PRESIEVE;
        let mut idx = 0;
        for (i, &x) in PRESIEVE_PRIMES.iter().enumerate() {
            let new_prod = prod * x;
            if wheel::bits_for(new_prod) > limit_bits {
                break
            }
            prod = new_prod;
            idx = i;
        }

        let len = cmp::min(wheel::bits_for(prod), limit_bits);

        if idx == 0 {
            Presieve {
                sieve: BitVec::new(),
                presieve_prod: prod,
                presieve_idx: idx,
            }
        } else {
            let mut sievers = vec![];
            for &x in &PRESIEVE_PRIMES[..idx] {
                let (use_, _idx) = wheel::bit_index(x);
                if use_ {
                    sievers.push(wheel::State::new(wheel::Wheel30, x, prod));
                }
            }
            let mut sieve =  BitVec::from_elem(len, true);
            StreamingSieve::small_primes_sieve(&mut sieve, &mut sievers);

            Presieve {
                sieve: sieve,
                presieve_prod: prod,
                presieve_idx: idx,
            }
        }
    }
    pub fn smallest_unincluded_prime(&self) -> usize {
        PRESIEVE_PRIMES[self.presieve_idx]
    }
    pub fn mark_small_primes(&self, sieve: &mut BitVec) {
        for &x in &PRESIEVE_PRIMES[..self.presieve_idx] {
            let (use_, idx) = wheel::bit_index(x);
            if use_ {
                sieve.set(idx, true)
            }
        }
    }
    pub fn apply(&self, sieve: &mut BitVec, low: usize) {
        if self.sieve.len() == 0 {
            return
        }
        let offset = (low % self.presieve_prod) * wheel::BYTE_SIZE / wheel::BYTE_MODULO / 8;

        copy_all(sieve.as_bytes_mut(),
                 self.sieve.as_bytes(),
                 offset);
        fn copy_all(dst: &mut [u8], src: &[u8], init_offset: usize) {
            let mut pos = 0;
            // pre-fill data at the start, as a rotation of `src`.
            pos += memcpy(&mut dst[pos..], &src[init_offset..]);
            if pos >= dst.len() { return }

            pos += memcpy(&mut dst[pos..], &src[..init_offset]);
            if pos >= dst.len() { return }

            // progressively copy the prefix to the rest of the
            // vector, doubling each time.
            while pos < dst.len() {
                let (filled, unfilled) = dst.split_at_mut(pos);
                pos += memcpy(unfilled, filled);
            }
        }
        fn memcpy<'d>(dst: &'d mut [u8], src: &[u8]) -> usize {
            let l = cmp::min(dst.len(), src.len());
            unsafe {
                ptr::copy(src.as_ptr(), dst.as_mut_ptr(), l);
            }
            l
        }
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use test::Bencher;
    use super::Presieve;
    fn run_presieve(b: &mut Bencher, n: usize) {
        b.iter(|| Presieve::new(::wheel::bits_for(n)))
    }
    #[bench]
    fn presieve_small(b: &mut Bencher) {
        run_presieve(b, 100)
    }
    #[bench]
    fn presieve_medium(b: &mut Bencher) {
        run_presieve(b, 10_000)
    }
    #[bench]
    fn presieve_large(b: &mut Bencher) {
        run_presieve(b, 100_000)
    }
    #[bench]
    fn presieve_larger(b: &mut Bencher) {
        run_presieve(b, 1_000_000)
    }
    #[bench]
    fn presieve_huge(b: &mut Bencher) {
        run_presieve(b, 10_000_000)
    }
}
<|endoftext|>


<commit_before>pub fn float_to_fixed(before: int, f: float) -> i32 {
    (1i32 << before) * (f as i32)
}

pub fn fixed_to_float(before: int, f: i32) -> float {
    f as float * 1.0f / ((1i32 << before) as float)
}

pub fn fixed_to_rounded_int(before: int, f: i32) -> int {
    let half = 1i32 << (before-1);
    if f > 0i32 {
        ((half + f) >> before) as int
    } else {
       -((half - f) >> before) as int
    }
}

/* Generate a 32-bit TrueType tag from its 4 characters */
pub fn true_type_tag(a: char, b: char, c: char, d: char) -> u32 {
    (a << 24 | b << 16 | c << 8 | d) as u32
}

#[test]
fn test_true_type_tag() {
    assert true_type_tag('c', 'm', 'a', 'p') == 0x_63_6D_61_70_u32;
}
<commit_msg>Add first-cut implementation of transform_text<commit_after>enum CompressionMode {
    CompressNone,
    CompressWhitespace,
    CompressWhitespaceNewline,
    DiscardNewline
}

impl CompressionMode : cmp::Eq {
    pure fn eq(other: &CompressionMode) -> bool {
        match (self, *other) {
            (CompressNone, CompressNone) => true,
            (CompressWhitespace, CompressWhitespace) => true,
            (CompressWhitespaceNewline, CompressWhitespaceNewline) => true,
            (DiscardNewline, DiscardNewline) => true,
            _ => false
        }
    }
    pure fn ne(other: &CompressionMode) -> bool {
        !self.eq(other)
    }
}

// ported from Gecko's nsTextFrameUtils::TransformText.
//
// High level TODOs:
// * consider incoming text state (preceding spaces, arabic, etc)
// * send outgoing text state (dual of above)
// * record skipped and kept chars for mapping original to new text
// * various edge cases for bidi, CJK, combining char seqs, etc.
pub fn transform_text(text: &str, mode: CompressionMode) -> ~str {
    let out_str: ~str = ~"";
    match mode {
        CompressNone | DiscardNewline => {
            do str::each_char(text) |ch: char| {
                if is_discardable_char(ch, mode) {
                    // TODO: record skipped char
                } else {
                    // TODO: record kept char
                    if ch == '\t' {
                        // TODO: set "has tab" flag
                    }
                    str::push_char(&out_str, ch);
                }

                true
            }
        },

        CompressWhitespace | CompressWhitespaceNewline => {
            let mut in_whitespace: bool = false;
            do str::each_char(text) |ch: char| {
                // TODO: discard newlines between CJK chars
                let mut next_in_whitespace: bool = match (ch, mode) {
                    // TODO: check for following char that may create
                    // a Unicode combining-character sequence with a
                    // space, in which case it shouldn't be  compressed.
                    (' ', _)  => true,
                    ('\t', _) => true,
                    ('\n', CompressWhitespaceNewline) => true,
                    (_, _)    => false
                };

                if next_in_whitespace {
                    if is_always_discardable_char(ch) {
                        // revert whitespace setting, since this char was discarded
                        next_in_whitespace = in_whitespace;
                        // TODO: record skipped char
                    } else {
                        // TODO: record kept char
                        str::push_char(&out_str, ch);
                    }
                } else {
                    if in_whitespace {
                        // TODO: record skipped char
                    } else {
                        // TODO: record kept char
                        str::push_char(&out_str, ch);
                    }
                }
                // save whitespace context for next char
                in_whitespace = next_in_whitespace;
                true
            } /* /do str::each_chari */
        }
    }

    return out_str;

    fn is_discardable_char(ch: char, mode: CompressionMode) -> bool {
        if is_always_discardable_char(ch) {
            return true;
        }
        match mode {
            DiscardNewline | CompressWhitespaceNewline => ch == '\n',
            _ => false
        }
    }

    fn is_always_discardable_char(_ch: char) -> bool {
        // TODO: check for bidi control chars, soft hyphens.
        false
    }
}

pub fn float_to_fixed(before: int, f: float) -> i32 {
    (1i32 << before) * (f as i32)
}

pub fn fixed_to_float(before: int, f: i32) -> float {
    f as float * 1.0f / ((1i32 << before) as float)
}

pub fn fixed_to_rounded_int(before: int, f: i32) -> int {
    let half = 1i32 << (before-1);
    if f > 0i32 {
        ((half + f) >> before) as int
    } else {
       -((half - f) >> before) as int
    }
}

/* Generate a 32-bit TrueType tag from its 4 characters */
pub fn true_type_tag(a: char, b: char, c: char, d: char) -> u32 {
    (a << 24 | b << 16 | c << 8 | d) as u32
}

#[test]
fn test_true_type_tag() {
    assert true_type_tag('c', 'm', 'a', 'p') == 0x_63_6D_61_70_u32;
}
<|endoftext|>


<commit_before>//! UDP for IOCP
//!
//! Note that most of this module is quite similar to the TCP module, so if
//! something seems odd you may also want to try the docs over there.

use std::fmt;
use std::io::prelude::*;
use std::io;
use std::mem;
use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::{Mutex, MutexGuard};

#[allow(unused_imports)]
use net2::{UdpBuilder, UdpSocketExt};
use winapi::*;
use miow::iocp::CompletionStatus;
use miow::net::SocketAddrBuf;
use miow::net::UdpSocketExt as MiowUdpSocketExt;

use {poll, Ready, Poll, PollOpt, Token};
use event::Evented;
use sys::windows::from_raw_arc::FromRawArc;
use sys::windows::selector::{Overlapped, ReadyBinding};

pub struct UdpSocket {
    imp: Imp,
    registration: Mutex<Option<poll::Registration>>,
}

#[derive(Clone)]
struct Imp {
    inner: FromRawArc<Io>,
}

struct Io {
    read: Overlapped,
    write: Overlapped,
    socket: net::UdpSocket,
    inner: Mutex<Inner>,
}

struct Inner {
    iocp: ReadyBinding,
    read: State<Vec<u8>, Vec<u8>>,
    write: State<Vec<u8>, (Vec<u8>, usize)>,
    read_buf: SocketAddrBuf,
}

enum State<T, U> {
    Empty,
    Pending(T),
    Ready(U),
    Error(io::Error),
}

impl UdpSocket {
    pub fn new(socket: net::UdpSocket) -> io::Result<UdpSocket> {
        Ok(UdpSocket {
            registration: Mutex::new(None),
            imp: Imp {
                inner: FromRawArc::new(Io {
                    read: Overlapped::new(recv_done),
                    write: Overlapped::new(send_done),
                    socket: socket,
                    inner: Mutex::new(Inner {
                        iocp: ReadyBinding::new(),
                        read: State::Empty,
                        write: State::Empty,
                        read_buf: SocketAddrBuf::new(),
                    }),
                }),
            },
        })
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.imp.inner.socket.local_addr()
    }

    pub fn try_clone(&self) -> io::Result<UdpSocket> {
        self.imp.inner.socket.try_clone().and_then(UdpSocket::new)
    }

    /// Note that unlike `TcpStream::write` this function will not attempt to
    /// continue writing `buf` until its entirely written.
    ///
    /// TODO: This... may be wrong in the long run. We're reporting that we
    ///       successfully wrote all of the bytes in `buf` but it's possible
    ///       that we don't actually end up writing all of them!
    pub fn send_to(&self, buf: &[u8], target: &SocketAddr)
                   -> io::Result<usize> {
        let mut me = self.inner();
        let me = &mut *me;

        match me.write {
            State::Empty => {}
            _ => return Err(would_block()),
        }

        if !me.iocp.registered() {
            return Err(would_block())
        }

        let interest = me.iocp.readiness();
        me.iocp.set_readiness(interest & !Ready::writable());

        let mut owned_buf = me.iocp.get_buffer(64 * 1024);
        let amt = owned_buf.write(buf)?;
        unsafe {
            trace!("scheduling a send");
            self.imp.inner.socket.send_to_overlapped(&owned_buf, target,
                                                     self.imp.inner.write.as_mut_ptr())
        }?;
        me.write = State::Pending(owned_buf);
        mem::forget(self.imp.clone());
        Ok(amt)
    }

    /// Note that unlike `TcpStream::write` this function will not attempt to
    /// continue writing `buf` until its entirely written.
    ///
    /// TODO: This... may be wrong in the long run. We're reporting that we
    ///       successfully wrote all of the bytes in `buf` but it's possible
    ///       that we don't actually end up writing all of them!
    pub fn send(&self, buf: &[u8]) -> io::Result<usize> {
        let mut me = self.inner();
        let me = &mut *me;

        match me.write {
            State::Empty => {}
            _ => return Err(would_block()),
        }

        if !me.iocp.registered() {
            return Err(would_block())
        }

        let interest = me.iocp.readiness();
        me.iocp.set_readiness(interest & !Ready::writable());

        let mut owned_buf = me.iocp.get_buffer(64 * 1024);
        let amt = owned_buf.write(buf)?;
        unsafe {
            trace!("scheduling a send");
            self.imp.inner.socket.send_overlapped(&owned_buf, self.imp.inner.write.as_mut_ptr())

        }?;
        me.write = State::Pending(owned_buf);
        mem::forget(self.imp.clone());
        Ok(amt)
    }

    pub fn recv_from(&self, mut buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
        let mut me = self.inner();
        match mem::replace(&mut me.read, State::Empty) {
            State::Empty => Err(would_block()),
            State::Pending(b) => { me.read = State::Pending(b); Err(would_block()) }
            State::Ready(data) => {
                // If we weren't provided enough space to receive the message
                // then don't actually read any data, just return an error.
                if buf.len() < data.len() {
                    me.read = State::Ready(data);
                    Err(io::Error::from_raw_os_error(WSAEMSGSIZE as i32))
                } else {
                    let r = if let Some(addr) = me.read_buf.to_socket_addr() {
                        buf.write(&data).unwrap();
                        Ok((data.len(), addr))
                    } else {
                        Err(io::Error::new(io::ErrorKind::Other,
                                           "failed to parse socket address"))
                    };
                    me.iocp.put_buffer(data);
                    self.imp.schedule_read_from(&mut me);
                    r
                }
            }
            State::Error(e) => {
                self.imp.schedule_read_from(&mut me);
                Err(e)
            }
        }
    }

    pub fn recv(&self, mut buf: &mut [u8])
                     -> io::Result<usize> {
        //Since recv_from can be used on connected sockets just call it and drop the address.
        self.recv_from(buf).map(|(size,_)| size)
    }

    pub fn connect(&self, addr: SocketAddr) -> io::Result<()> {
        self.imp.inner.socket.connect(addr)
    }

    pub fn broadcast(&self) -> io::Result<bool> {
        self.imp.inner.socket.broadcast()
    }

    pub fn set_broadcast(&self, on: bool) -> io::Result<()> {
        self.imp.inner.socket.set_broadcast(on)
    }

    pub fn multicast_loop_v4(&self) -> io::Result<bool> {
        self.imp.inner.socket.multicast_loop_v4()
    }

    pub fn set_multicast_loop_v4(&self, on: bool) -> io::Result<()> {
        self.imp.inner.socket.set_multicast_loop_v4(on)
    }

    pub fn multicast_ttl_v4(&self) -> io::Result<u32> {
        self.imp.inner.socket.multicast_ttl_v4()
    }

    pub fn set_multicast_ttl_v4(&self, ttl: u32) -> io::Result<()> {
        self.imp.inner.socket.set_multicast_ttl_v4(ttl)
    }

    pub fn multicast_loop_v6(&self) -> io::Result<bool> {
        self.imp.inner.socket.multicast_loop_v6()
    }

    pub fn set_multicast_loop_v6(&self, on: bool) -> io::Result<()> {
        self.imp.inner.socket.set_multicast_loop_v6(on)
    }

    pub fn ttl(&self) -> io::Result<u32> {
        self.imp.inner.socket.ttl()
    }

    pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
        self.imp.inner.socket.set_ttl(ttl)
    }

    pub fn join_multicast_v4(&self,
                             multiaddr: &Ipv4Addr,
                             interface: &Ipv4Addr) -> io::Result<()> {
        self.imp.inner.socket.join_multicast_v4(multiaddr, interface)
    }

    pub fn join_multicast_v6(&self,
                             multiaddr: &Ipv6Addr,
                             interface: u32) -> io::Result<()> {
        self.imp.inner.socket.join_multicast_v6(multiaddr, interface)
    }

    pub fn leave_multicast_v4(&self,
                              multiaddr: &Ipv4Addr,
                              interface: &Ipv4Addr) -> io::Result<()> {
        self.imp.inner.socket.leave_multicast_v4(multiaddr, interface)
    }

    pub fn leave_multicast_v6(&self,
                              multiaddr: &Ipv6Addr,
                              interface: u32) -> io::Result<()> {
        self.imp.inner.socket.leave_multicast_v6(multiaddr, interface)
    }

    pub fn set_only_v6(&self, only_v6: bool) -> io::Result<()> {
        self.imp.inner.socket.set_only_v6(only_v6)
    }

    pub fn only_v6(&self) -> io::Result<bool> {
        self.imp.inner.socket.only_v6()
    }

    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
        self.imp.inner.socket.take_error()
    }

    fn inner(&self) -> MutexGuard<Inner> {
        self.imp.inner()
    }

    fn post_register(&self, interest: Ready, me: &mut Inner) {
        if interest.is_readable() {
            //We use recv_from here since it is well specified for both
            //connected and non-connected sockets and we can discard the address
            //when calling recv().
            self.imp.schedule_read_from(me);
        }
        // See comments in TcpSocket::post_register for what's going on here
        if interest.is_writable() {
            if let State::Empty = me.write {
                self.imp.add_readiness(me, Ready::writable());
            }
        }
    }
}

impl Imp {
    fn inner(&self) -> MutexGuard<Inner> {
        self.inner.inner.lock().unwrap()
    }

    fn schedule_read_from(&self, me: &mut Inner) {
        match me.read {
            State::Empty => {}
            _ => return,
        }

        let interest = me.iocp.readiness();
        me.iocp.set_readiness(interest & !Ready::readable());

        let mut buf = me.iocp.get_buffer(64 * 1024);
        let res = unsafe {
            trace!("scheduling a read");
            let cap = buf.capacity();
            buf.set_len(cap);
            self.inner.socket.recv_from_overlapped(&mut buf, &mut me.read_buf,
                                                   self.inner.read.as_mut_ptr())
        };
        match res {
            Ok(_) => {
                me.read = State::Pending(buf);
                mem::forget(self.clone());
            }
            Err(e) => {
                me.read = State::Error(e);
                self.add_readiness(me, Ready::readable());
                me.iocp.put_buffer(buf);
            }
        }
    }

    // See comments in tcp::StreamImp::push
    fn add_readiness(&self, me: &Inner, set: Ready) {
        me.iocp.set_readiness(set | me.iocp.readiness());
    }
}

impl Evented for UdpSocket {
    fn register(&self, poll: &Poll, token: Token,
                interest: Ready, opts: PollOpt) -> io::Result<()> {
        let mut me = self.inner();
        me.iocp.register_socket(&self.imp.inner.socket,
                                     poll, token, interest, opts,
                                     &self.registration)?;
        self.post_register(interest, &mut me);
        Ok(())
    }

    fn reregister(&self, poll: &Poll, token: Token,
                  interest: Ready, opts: PollOpt) -> io::Result<()> {
        let mut me = self.inner();
        me.iocp.reregister_socket(&self.imp.inner.socket,
                                       poll, token, interest,
                                       opts, &self.registration)?;
        self.post_register(interest, &mut me);
        Ok(())
    }

    fn deregister(&self, poll: &Poll) -> io::Result<()> {
        self.inner().iocp.deregister(&self.imp.inner.socket,
                                     poll, &self.registration)
    }
}

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

impl Drop for UdpSocket {
    fn drop(&mut self) {
        let inner = self.inner();

        // If we're still internally reading, we're no longer interested. Note
        // though that we don't cancel any writes which may have been issued to
        // preserve the same semantics as Unix.
        unsafe {
            match inner.read {
                State::Pending(_) => {
                    drop(super::cancel(&self.imp.inner.socket,
                                       &self.imp.inner.read));
                }
                State::Empty |
                State::Ready(_) |
                State::Error(_) => {}
            }
        }
    }
}

fn send_done(status: &OVERLAPPED_ENTRY) {
    let status = CompletionStatus::from_entry(status);
    trace!("finished a send {}", status.bytes_transferred());
    let me2 = Imp {
        inner: unsafe { overlapped2arc!(status.overlapped(), Io, write) },
    };
    let mut me = me2.inner();
    me.write = State::Empty;
    me2.add_readiness(&mut me, Ready::writable());
}

fn recv_done(status: &OVERLAPPED_ENTRY) {
    let status = CompletionStatus::from_entry(status);
    trace!("finished a recv {}", status.bytes_transferred());
    let me2 = Imp {
        inner: unsafe { overlapped2arc!(status.overlapped(), Io, read) },
    };
    let mut me = me2.inner();
    let mut buf = match mem::replace(&mut me.read, State::Empty) {
        State::Pending(buf) => buf,
        _ => unreachable!(),
    };
    unsafe {
        buf.set_len(status.bytes_transferred() as usize);
    }
    me.read = State::Ready(buf);
    me2.add_readiness(&mut me, Ready::readable());
}

// TODO: Use std's allocation free io::Error
const WOULDBLOCK: i32 = ::winapi::winerror::WSAEWOULDBLOCK as i32;

/// Returns a std `WouldBlock` error without allocating
pub fn would_block() -> ::std::io::Error {
    ::std::io::Error::from_raw_os_error(WOULDBLOCK)
}
<commit_msg>Remove unused mutable argument on windows (#645)<commit_after>//! UDP for IOCP
//!
//! Note that most of this module is quite similar to the TCP module, so if
//! something seems odd you may also want to try the docs over there.

use std::fmt;
use std::io::prelude::*;
use std::io;
use std::mem;
use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::{Mutex, MutexGuard};

#[allow(unused_imports)]
use net2::{UdpBuilder, UdpSocketExt};
use winapi::*;
use miow::iocp::CompletionStatus;
use miow::net::SocketAddrBuf;
use miow::net::UdpSocketExt as MiowUdpSocketExt;

use {poll, Ready, Poll, PollOpt, Token};
use event::Evented;
use sys::windows::from_raw_arc::FromRawArc;
use sys::windows::selector::{Overlapped, ReadyBinding};

pub struct UdpSocket {
    imp: Imp,
    registration: Mutex<Option<poll::Registration>>,
}

#[derive(Clone)]
struct Imp {
    inner: FromRawArc<Io>,
}

struct Io {
    read: Overlapped,
    write: Overlapped,
    socket: net::UdpSocket,
    inner: Mutex<Inner>,
}

struct Inner {
    iocp: ReadyBinding,
    read: State<Vec<u8>, Vec<u8>>,
    write: State<Vec<u8>, (Vec<u8>, usize)>,
    read_buf: SocketAddrBuf,
}

enum State<T, U> {
    Empty,
    Pending(T),
    Ready(U),
    Error(io::Error),
}

impl UdpSocket {
    pub fn new(socket: net::UdpSocket) -> io::Result<UdpSocket> {
        Ok(UdpSocket {
            registration: Mutex::new(None),
            imp: Imp {
                inner: FromRawArc::new(Io {
                    read: Overlapped::new(recv_done),
                    write: Overlapped::new(send_done),
                    socket: socket,
                    inner: Mutex::new(Inner {
                        iocp: ReadyBinding::new(),
                        read: State::Empty,
                        write: State::Empty,
                        read_buf: SocketAddrBuf::new(),
                    }),
                }),
            },
        })
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.imp.inner.socket.local_addr()
    }

    pub fn try_clone(&self) -> io::Result<UdpSocket> {
        self.imp.inner.socket.try_clone().and_then(UdpSocket::new)
    }

    /// Note that unlike `TcpStream::write` this function will not attempt to
    /// continue writing `buf` until its entirely written.
    ///
    /// TODO: This... may be wrong in the long run. We're reporting that we
    ///       successfully wrote all of the bytes in `buf` but it's possible
    ///       that we don't actually end up writing all of them!
    pub fn send_to(&self, buf: &[u8], target: &SocketAddr)
                   -> io::Result<usize> {
        let mut me = self.inner();
        let me = &mut *me;

        match me.write {
            State::Empty => {}
            _ => return Err(would_block()),
        }

        if !me.iocp.registered() {
            return Err(would_block())
        }

        let interest = me.iocp.readiness();
        me.iocp.set_readiness(interest & !Ready::writable());

        let mut owned_buf = me.iocp.get_buffer(64 * 1024);
        let amt = owned_buf.write(buf)?;
        unsafe {
            trace!("scheduling a send");
            self.imp.inner.socket.send_to_overlapped(&owned_buf, target,
                                                     self.imp.inner.write.as_mut_ptr())
        }?;
        me.write = State::Pending(owned_buf);
        mem::forget(self.imp.clone());
        Ok(amt)
    }

    /// Note that unlike `TcpStream::write` this function will not attempt to
    /// continue writing `buf` until its entirely written.
    ///
    /// TODO: This... may be wrong in the long run. We're reporting that we
    ///       successfully wrote all of the bytes in `buf` but it's possible
    ///       that we don't actually end up writing all of them!
    pub fn send(&self, buf: &[u8]) -> io::Result<usize> {
        let mut me = self.inner();
        let me = &mut *me;

        match me.write {
            State::Empty => {}
            _ => return Err(would_block()),
        }

        if !me.iocp.registered() {
            return Err(would_block())
        }

        let interest = me.iocp.readiness();
        me.iocp.set_readiness(interest & !Ready::writable());

        let mut owned_buf = me.iocp.get_buffer(64 * 1024);
        let amt = owned_buf.write(buf)?;
        unsafe {
            trace!("scheduling a send");
            self.imp.inner.socket.send_overlapped(&owned_buf, self.imp.inner.write.as_mut_ptr())

        }?;
        me.write = State::Pending(owned_buf);
        mem::forget(self.imp.clone());
        Ok(amt)
    }

    pub fn recv_from(&self, mut buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
        let mut me = self.inner();
        match mem::replace(&mut me.read, State::Empty) {
            State::Empty => Err(would_block()),
            State::Pending(b) => { me.read = State::Pending(b); Err(would_block()) }
            State::Ready(data) => {
                // If we weren't provided enough space to receive the message
                // then don't actually read any data, just return an error.
                if buf.len() < data.len() {
                    me.read = State::Ready(data);
                    Err(io::Error::from_raw_os_error(WSAEMSGSIZE as i32))
                } else {
                    let r = if let Some(addr) = me.read_buf.to_socket_addr() {
                        buf.write(&data).unwrap();
                        Ok((data.len(), addr))
                    } else {
                        Err(io::Error::new(io::ErrorKind::Other,
                                           "failed to parse socket address"))
                    };
                    me.iocp.put_buffer(data);
                    self.imp.schedule_read_from(&mut me);
                    r
                }
            }
            State::Error(e) => {
                self.imp.schedule_read_from(&mut me);
                Err(e)
            }
        }
    }

    pub fn recv(&self, buf: &mut [u8])
                     -> io::Result<usize> {
        //Since recv_from can be used on connected sockets just call it and drop the address.
        self.recv_from(buf).map(|(size,_)| size)
    }

    pub fn connect(&self, addr: SocketAddr) -> io::Result<()> {
        self.imp.inner.socket.connect(addr)
    }

    pub fn broadcast(&self) -> io::Result<bool> {
        self.imp.inner.socket.broadcast()
    }

    pub fn set_broadcast(&self, on: bool) -> io::Result<()> {
        self.imp.inner.socket.set_broadcast(on)
    }

    pub fn multicast_loop_v4(&self) -> io::Result<bool> {
        self.imp.inner.socket.multicast_loop_v4()
    }

    pub fn set_multicast_loop_v4(&self, on: bool) -> io::Result<()> {
        self.imp.inner.socket.set_multicast_loop_v4(on)
    }

    pub fn multicast_ttl_v4(&self) -> io::Result<u32> {
        self.imp.inner.socket.multicast_ttl_v4()
    }

    pub fn set_multicast_ttl_v4(&self, ttl: u32) -> io::Result<()> {
        self.imp.inner.socket.set_multicast_ttl_v4(ttl)
    }

    pub fn multicast_loop_v6(&self) -> io::Result<bool> {
        self.imp.inner.socket.multicast_loop_v6()
    }

    pub fn set_multicast_loop_v6(&self, on: bool) -> io::Result<()> {
        self.imp.inner.socket.set_multicast_loop_v6(on)
    }

    pub fn ttl(&self) -> io::Result<u32> {
        self.imp.inner.socket.ttl()
    }

    pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
        self.imp.inner.socket.set_ttl(ttl)
    }

    pub fn join_multicast_v4(&self,
                             multiaddr: &Ipv4Addr,
                             interface: &Ipv4Addr) -> io::Result<()> {
        self.imp.inner.socket.join_multicast_v4(multiaddr, interface)
    }

    pub fn join_multicast_v6(&self,
                             multiaddr: &Ipv6Addr,
                             interface: u32) -> io::Result<()> {
        self.imp.inner.socket.join_multicast_v6(multiaddr, interface)
    }

    pub fn leave_multicast_v4(&self,
                              multiaddr: &Ipv4Addr,
                              interface: &Ipv4Addr) -> io::Result<()> {
        self.imp.inner.socket.leave_multicast_v4(multiaddr, interface)
    }

    pub fn leave_multicast_v6(&self,
                              multiaddr: &Ipv6Addr,
                              interface: u32) -> io::Result<()> {
        self.imp.inner.socket.leave_multicast_v6(multiaddr, interface)
    }

    pub fn set_only_v6(&self, only_v6: bool) -> io::Result<()> {
        self.imp.inner.socket.set_only_v6(only_v6)
    }

    pub fn only_v6(&self) -> io::Result<bool> {
        self.imp.inner.socket.only_v6()
    }

    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
        self.imp.inner.socket.take_error()
    }

    fn inner(&self) -> MutexGuard<Inner> {
        self.imp.inner()
    }

    fn post_register(&self, interest: Ready, me: &mut Inner) {
        if interest.is_readable() {
            //We use recv_from here since it is well specified for both
            //connected and non-connected sockets and we can discard the address
            //when calling recv().
            self.imp.schedule_read_from(me);
        }
        // See comments in TcpSocket::post_register for what's going on here
        if interest.is_writable() {
            if let State::Empty = me.write {
                self.imp.add_readiness(me, Ready::writable());
            }
        }
    }
}

impl Imp {
    fn inner(&self) -> MutexGuard<Inner> {
        self.inner.inner.lock().unwrap()
    }

    fn schedule_read_from(&self, me: &mut Inner) {
        match me.read {
            State::Empty => {}
            _ => return,
        }

        let interest = me.iocp.readiness();
        me.iocp.set_readiness(interest & !Ready::readable());

        let mut buf = me.iocp.get_buffer(64 * 1024);
        let res = unsafe {
            trace!("scheduling a read");
            let cap = buf.capacity();
            buf.set_len(cap);
            self.inner.socket.recv_from_overlapped(&mut buf, &mut me.read_buf,
                                                   self.inner.read.as_mut_ptr())
        };
        match res {
            Ok(_) => {
                me.read = State::Pending(buf);
                mem::forget(self.clone());
            }
            Err(e) => {
                me.read = State::Error(e);
                self.add_readiness(me, Ready::readable());
                me.iocp.put_buffer(buf);
            }
        }
    }

    // See comments in tcp::StreamImp::push
    fn add_readiness(&self, me: &Inner, set: Ready) {
        me.iocp.set_readiness(set | me.iocp.readiness());
    }
}

impl Evented for UdpSocket {
    fn register(&self, poll: &Poll, token: Token,
                interest: Ready, opts: PollOpt) -> io::Result<()> {
        let mut me = self.inner();
        me.iocp.register_socket(&self.imp.inner.socket,
                                     poll, token, interest, opts,
                                     &self.registration)?;
        self.post_register(interest, &mut me);
        Ok(())
    }

    fn reregister(&self, poll: &Poll, token: Token,
                  interest: Ready, opts: PollOpt) -> io::Result<()> {
        let mut me = self.inner();
        me.iocp.reregister_socket(&self.imp.inner.socket,
                                       poll, token, interest,
                                       opts, &self.registration)?;
        self.post_register(interest, &mut me);
        Ok(())
    }

    fn deregister(&self, poll: &Poll) -> io::Result<()> {
        self.inner().iocp.deregister(&self.imp.inner.socket,
                                     poll, &self.registration)
    }
}

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

impl Drop for UdpSocket {
    fn drop(&mut self) {
        let inner = self.inner();

        // If we're still internally reading, we're no longer interested. Note
        // though that we don't cancel any writes which may have been issued to
        // preserve the same semantics as Unix.
        unsafe {
            match inner.read {
                State::Pending(_) => {
                    drop(super::cancel(&self.imp.inner.socket,
                                       &self.imp.inner.read));
                }
                State::Empty |
                State::Ready(_) |
                State::Error(_) => {}
            }
        }
    }
}

fn send_done(status: &OVERLAPPED_ENTRY) {
    let status = CompletionStatus::from_entry(status);
    trace!("finished a send {}", status.bytes_transferred());
    let me2 = Imp {
        inner: unsafe { overlapped2arc!(status.overlapped(), Io, write) },
    };
    let mut me = me2.inner();
    me.write = State::Empty;
    me2.add_readiness(&mut me, Ready::writable());
}

fn recv_done(status: &OVERLAPPED_ENTRY) {
    let status = CompletionStatus::from_entry(status);
    trace!("finished a recv {}", status.bytes_transferred());
    let me2 = Imp {
        inner: unsafe { overlapped2arc!(status.overlapped(), Io, read) },
    };
    let mut me = me2.inner();
    let mut buf = match mem::replace(&mut me.read, State::Empty) {
        State::Pending(buf) => buf,
        _ => unreachable!(),
    };
    unsafe {
        buf.set_len(status.bytes_transferred() as usize);
    }
    me.read = State::Ready(buf);
    me2.add_readiness(&mut me, Ready::readable());
}

// TODO: Use std's allocation free io::Error
const WOULDBLOCK: i32 = ::winapi::winerror::WSAEWOULDBLOCK as i32;

/// Returns a std `WouldBlock` error without allocating
pub fn would_block() -> ::std::io::Error {
    ::std::io::Error::from_raw_os_error(WOULDBLOCK)
}
<|endoftext|>