<commit_before>use perlin::index::boolean_index::{IndexBuilder, BooleanIndex};
use perlin::storage::RamStorage;
use rand;
use rand::{XorShiftRng, Rng};

#[allow(dead_code)]
pub fn prepare_index(documents: usize, document_size: usize) -> BooleanIndex<usize> {
    let collection: Vec<Vec<usize>> = (0..documents)
        .map(|_| ZipfGenerator::new(voc_size(45, 0.5, documents * document_size)).take(document_size).collect())
        .collect();
    println!("Preparing index with {} documents each {} terms",
             documents,
             document_size);
    let index = IndexBuilder::<_, RamStorage<_>>::new()
        .create(collection.iter().map(|i| i.iter().cloned()));
    index.unwrap()
}

// Implementation of Heaps' Law
pub fn voc_size(k: usize, b: f64, tokens: usize) -> usize {
    ((k as f64) * (tokens as f64).powf(b)) as usize
}

#[derive(Clone)]
pub struct ZipfGenerator {
    voc_size: usize,
    factor: f32,
    acc_probs: Vec<f32>,
    rng: XorShiftRng,
}

impl ZipfGenerator {
    pub fn new(voc_size: usize) -> Self {
        let mut res = ZipfGenerator {
            voc_size: voc_size,
            factor: (1.78 * voc_size as f32).ln(),
            acc_probs: Vec::with_capacity(voc_size),
            rng: rand::weak_rng(),
        };
        let mut acc = 0.0;
        for i in 1..voc_size {
            acc += 1.0 / (i as f32 * res.factor);
            res.acc_probs.push(acc);
        }
        res.acc_probs.push(1f32);
        res
    }
}

impl<'a> Iterator for &'a mut ZipfGenerator {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let dice = self.rng.next_f32();
        let result = match self.acc_probs.binary_search_by(|v| v.partial_cmp(&dice).unwrap()) {
            Ok(index) => index,
            Err(index) => index
        };
        Some(result)
    }
}
<commit_msg>clippy for bench<commit_after>use perlin::index::boolean_index::{IndexBuilder, BooleanIndex};
use perlin::storage::RamStorage;
use rand;
use rand::{XorShiftRng, Rng};

#[allow(dead_code)]
pub fn prepare_index(documents: usize, document_size: usize) -> BooleanIndex<usize> {
    let collection: Vec<Vec<usize>> = (0..documents)
        .map(|_| ZipfGenerator::new(voc_size(45, 0.5, documents * document_size)).take(document_size).collect())
        .collect();
    println!("Preparing index with {} documents each {} terms",
             documents,
             document_size);
    let index = IndexBuilder::<_, RamStorage<_>>::new()
        .create(collection.iter().map(|i| i.iter().cloned()));
    index.unwrap()
}

// Implementation of Heaps' Law
pub fn voc_size(k: usize, b: f64, tokens: usize) -> usize {
    ((k as f64) * (tokens as f64).powf(b)) as usize
}

#[derive(Clone)]
pub struct ZipfGenerator {
    voc_size: usize,
    factor: f32,
    acc_probs: Vec<f32>,
    rng: XorShiftRng,
}

impl ZipfGenerator {
    pub fn new(voc_size: usize) -> Self {
        let mut res = ZipfGenerator {
            voc_size: voc_size,
            factor: (1.78 * voc_size as f32).ln(),
            acc_probs: Vec::with_capacity(voc_size),
            rng: rand::weak_rng(),
        };
        let mut acc = 0.0;
        for i in 1..voc_size {
            acc += 1.0 / (i as f32 * res.factor);
            res.acc_probs.push(acc);
        }
        res.acc_probs.push(1f32);
        res
    }
}

impl<'a> Iterator for &'a mut ZipfGenerator {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let dice = self.rng.next_f32();
        let result = match self.acc_probs.binary_search_by(|v| v.partial_cmp(&dice).unwrap()) {
            Ok(index) | Err(index) => index
        };
        Some(result)
    }
}
<|endoftext|>


<commit_before>use std::any::Any;
use std::fmt;
use std::io::Read;
use std::io::Write;

#[cfg(feature = "bytes")]
use bytes::Bytes;

use clear::Clear;
use error::ProtobufError;
use error::ProtobufResult;
use reflect::MessageDescriptor;
use reflect::ProtobufValue;
use stream::with_coded_output_stream_to_bytes;
use stream::CodedInputStream;
use stream::CodedOutputStream;
use stream::WithCodedInputStream;
use stream::WithCodedOutputStream;
use unknown::UnknownFields;

/// Trait implemented for all generated structs for protobuf messages.
/// Also, generated messages implement `Clone + Default + PartialEq`
pub trait Message: fmt::Debug + Clear + Send + Sync + ProtobufValue {
    /// Message descriptor for this message, used for reflection.
    fn descriptor(&self) -> &'static MessageDescriptor;

    /// True iff all required fields are initialized.
    /// Always returns `true` for protobuf 3.
    fn is_initialized(&self) -> bool;

    /// Update this message object with fields read from given stream.
    fn merge_from(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()>;

    /// Write message to the stream.
    ///
    /// Sizes of this messages and nested messages must be cached
    /// by calling `compute_size` prior to this call.
    fn write_to_with_cached_sizes(&self, os: &mut CodedOutputStream) -> ProtobufResult<()>;

    /// Compute and cache size of this message and all nested messages
    fn compute_size(&self) -> u32;

    /// Get size previously computed by `compute_size`.
    fn get_cached_size(&self) -> u32;

    /// Write the message to the stream.
    ///
    /// Results in error if message is not fully initialized.
    fn write_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.check_initialized()?;

        // cache sizes
        self.compute_size();
        // TODO: reserve additional
        self.write_to_with_cached_sizes(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the stream prepending the message with message length
    /// encoded as varint.
    fn write_length_delimited_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        let size = self.compute_size();
        os.write_raw_varint32(size)?;
        self.write_to_with_cached_sizes(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_vec(&self, vec: &mut Vec<u8>) -> ProtobufResult<()> {
        let mut os = CodedOutputStream::vec(vec);
        self.write_length_delimited_to(&mut os)?;
        os.flush()?;
        Ok(())
    }

    /// Update this message object with fields read from given stream.
    fn merge_from_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
        let mut is = CodedInputStream::from_bytes(bytes);
        self.merge_from(&mut is)
    }

    /// Check if all required fields of this object are initialized.
    fn check_initialized(&self) -> ProtobufResult<()> {
        if !self.is_initialized() {
            Err(ProtobufError::message_not_initialized(
                self.descriptor().name(),
            ))
        } else {
            Ok(())
        }
    }

    /// Write the message to the writer.
    fn write_to_writer(&self, w: &mut Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    fn write_to_vec(&self, v: &mut Vec<u8>) -> ProtobufResult<()> {
        v.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.

    /// > **Note**: You can use parse_from_bytes to do the reverse.
    fn write_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        self.check_initialized()?;

        let size = self.compute_size() as usize;
        let mut v = Vec::with_capacity(size);
        // skip zerofill
        unsafe {
            v.set_len(size);
        }
        {
            let mut os = CodedOutputStream::bytes(&mut v);
            self.write_to_with_cached_sizes(&mut os)?;
            os.check_eof();
        }
        Ok(v)
    }

    /// Write the message to the writer, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_writer(&self, w: &mut Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_length_delimited_to(os))
    }

    /// Write the message to the bytes vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        with_coded_output_stream_to_bytes(|os| self.write_length_delimited_to(os))
    }

    /// Get a reference to unknown fields.
    fn get_unknown_fields(&self) -> &UnknownFields;
    /// Get a mutable reference to unknown fields.
    fn mut_unknown_fields(&mut self) -> &mut UnknownFields;

    /// Create an empty message object.
    fn new() -> Self
    where
        Self: Sized;

    /// Get message descriptor for message type.
    fn descriptor_static() -> &'static MessageDescriptor
    where
        Self: Sized,
    {
        panic!(
            "descriptor_static is not implemented for message, \
             LITE_RUNTIME must be used"
        );
    }

    /// Return a pointer to default immutable message.
    fn default_instance() -> &'static Self
    where
        Self: Sized;
}

impl Message {
    pub fn downcast_box<T: Any>(self: Box<Self>) -> Result<Box<T>, Box<Message>> {
        if self.as_any().is::<T>() {
            unsafe {
                let raw: *mut Message = Box::into_raw(self);
                Ok(Box::from_raw(raw as *mut T))
            }
        } else {
            Err(self)
        }
    }
}

impl Clone for Box<Message> {
    fn clone(&self) -> Self {
        use std::ops::Deref;
        self.descriptor().clone(self.deref())
    }
}

#[cfg(off)] // don't need it
impl PartialEq for Box<Message> {
    fn eq(&self, other: &Box<Message>) -> bool {
        use std::ops::Deref;
        self.descriptor() == other.descriptor() && self.descriptor().eq(self.deref(), other.deref())
    }
}

pub fn message_down_cast<'a, M: Message + 'a>(m: &'a Message) -> &'a M {
    m.as_any().downcast_ref::<M>().unwrap()
}

pub fn message_down_cast_mut<'a, M: Message + 'a>(m: &'a mut Message) -> &'a mut M {
    m.as_any_mut().downcast_mut::<M>().unwrap()
}

/// Parse message from stream.
pub fn parse_from<M: Message>(is: &mut CodedInputStream) -> ProtobufResult<M> {
    let mut r: M = Message::new();
    r.merge_from(is)?;
    r.check_initialized()?;
    Ok(r)
}

/// Parse message from reader.
/// Parse stops on EOF or when error encountered.
pub fn parse_from_reader<M: Message>(reader: &mut Read) -> ProtobufResult<M> {
    reader.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from byte array.
pub fn parse_from_bytes<M: Message>(bytes: &[u8]) -> ProtobufResult<M> {
    bytes.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from `Bytes` object.
/// Resulting message may share references to the passed bytes object.
#[cfg(feature = "bytes")]
pub fn parse_from_carllerche_bytes<M: Message>(bytes: &Bytes) -> ProtobufResult<M> {
    // Call trait explicitly to avoid accidental construction from `&[u8]`
    WithCodedInputStream::with_coded_input_stream(bytes, |is| parse_from::<M>(is))
}
<commit_msg>Fix doc comment<commit_after>use std::any::Any;
use std::fmt;
use std::io::Read;
use std::io::Write;

#[cfg(feature = "bytes")]
use bytes::Bytes;

use clear::Clear;
use error::ProtobufError;
use error::ProtobufResult;
use reflect::MessageDescriptor;
use reflect::ProtobufValue;
use stream::with_coded_output_stream_to_bytes;
use stream::CodedInputStream;
use stream::CodedOutputStream;
use stream::WithCodedInputStream;
use stream::WithCodedOutputStream;
use unknown::UnknownFields;

/// Trait implemented for all generated structs for protobuf messages.
/// Also, generated messages implement `Clone + Default + PartialEq`
pub trait Message: fmt::Debug + Clear + Send + Sync + ProtobufValue {
    /// Message descriptor for this message, used for reflection.
    fn descriptor(&self) -> &'static MessageDescriptor;

    /// True iff all required fields are initialized.
    /// Always returns `true` for protobuf 3.
    fn is_initialized(&self) -> bool;

    /// Update this message object with fields read from given stream.
    fn merge_from(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()>;

    /// Write message to the stream.
    ///
    /// Sizes of this messages and nested messages must be cached
    /// by calling `compute_size` prior to this call.
    fn write_to_with_cached_sizes(&self, os: &mut CodedOutputStream) -> ProtobufResult<()>;

    /// Compute and cache size of this message and all nested messages
    fn compute_size(&self) -> u32;

    /// Get size previously computed by `compute_size`.
    fn get_cached_size(&self) -> u32;

    /// Write the message to the stream.
    ///
    /// Results in error if message is not fully initialized.
    fn write_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.check_initialized()?;

        // cache sizes
        self.compute_size();
        // TODO: reserve additional
        self.write_to_with_cached_sizes(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the stream prepending the message with message length
    /// encoded as varint.
    fn write_length_delimited_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        let size = self.compute_size();
        os.write_raw_varint32(size)?;
        self.write_to_with_cached_sizes(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_vec(&self, vec: &mut Vec<u8>) -> ProtobufResult<()> {
        let mut os = CodedOutputStream::vec(vec);
        self.write_length_delimited_to(&mut os)?;
        os.flush()?;
        Ok(())
    }

    /// Update this message object with fields read from given stream.
    fn merge_from_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
        let mut is = CodedInputStream::from_bytes(bytes);
        self.merge_from(&mut is)
    }

    /// Check if all required fields of this object are initialized.
    fn check_initialized(&self) -> ProtobufResult<()> {
        if !self.is_initialized() {
            Err(ProtobufError::message_not_initialized(
                self.descriptor().name(),
            ))
        } else {
            Ok(())
        }
    }

    /// Write the message to the writer.
    fn write_to_writer(&self, w: &mut Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    fn write_to_vec(&self, v: &mut Vec<u8>) -> ProtobufResult<()> {
        v.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    ///
    /// > **Note**: You can use `parse_from_bytes` to do the reverse.
    fn write_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        self.check_initialized()?;

        let size = self.compute_size() as usize;
        let mut v = Vec::with_capacity(size);
        // skip zerofill
        unsafe {
            v.set_len(size);
        }
        {
            let mut os = CodedOutputStream::bytes(&mut v);
            self.write_to_with_cached_sizes(&mut os)?;
            os.check_eof();
        }
        Ok(v)
    }

    /// Write the message to the writer, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_writer(&self, w: &mut Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_length_delimited_to(os))
    }

    /// Write the message to the bytes vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        with_coded_output_stream_to_bytes(|os| self.write_length_delimited_to(os))
    }

    /// Get a reference to unknown fields.
    fn get_unknown_fields(&self) -> &UnknownFields;
    /// Get a mutable reference to unknown fields.
    fn mut_unknown_fields(&mut self) -> &mut UnknownFields;

    /// Create an empty message object.
    fn new() -> Self
    where
        Self: Sized;

    /// Get message descriptor for message type.
    fn descriptor_static() -> &'static MessageDescriptor
    where
        Self: Sized,
    {
        panic!(
            "descriptor_static is not implemented for message, \
             LITE_RUNTIME must be used"
        );
    }

    /// Return a pointer to default immutable message.
    fn default_instance() -> &'static Self
    where
        Self: Sized;
}

impl Message {
    pub fn downcast_box<T: Any>(self: Box<Self>) -> Result<Box<T>, Box<Message>> {
        if self.as_any().is::<T>() {
            unsafe {
                let raw: *mut Message = Box::into_raw(self);
                Ok(Box::from_raw(raw as *mut T))
            }
        } else {
            Err(self)
        }
    }
}

impl Clone for Box<Message> {
    fn clone(&self) -> Self {
        use std::ops::Deref;
        self.descriptor().clone(self.deref())
    }
}

#[cfg(off)] // don't need it
impl PartialEq for Box<Message> {
    fn eq(&self, other: &Box<Message>) -> bool {
        use std::ops::Deref;
        self.descriptor() == other.descriptor() && self.descriptor().eq(self.deref(), other.deref())
    }
}

pub fn message_down_cast<'a, M: Message + 'a>(m: &'a Message) -> &'a M {
    m.as_any().downcast_ref::<M>().unwrap()
}

pub fn message_down_cast_mut<'a, M: Message + 'a>(m: &'a mut Message) -> &'a mut M {
    m.as_any_mut().downcast_mut::<M>().unwrap()
}

/// Parse message from stream.
pub fn parse_from<M: Message>(is: &mut CodedInputStream) -> ProtobufResult<M> {
    let mut r: M = Message::new();
    r.merge_from(is)?;
    r.check_initialized()?;
    Ok(r)
}

/// Parse message from reader.
/// Parse stops on EOF or when error encountered.
pub fn parse_from_reader<M: Message>(reader: &mut Read) -> ProtobufResult<M> {
    reader.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from byte array.
pub fn parse_from_bytes<M: Message>(bytes: &[u8]) -> ProtobufResult<M> {
    bytes.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from `Bytes` object.
/// Resulting message may share references to the passed bytes object.
#[cfg(feature = "bytes")]
pub fn parse_from_carllerche_bytes<M: Message>(bytes: &Bytes) -> ProtobufResult<M> {
    // Call trait explicitly to avoid accidental construction from `&[u8]`
    WithCodedInputStream::with_coded_input_stream(bytes, |is| parse_from::<M>(is))
}
<|endoftext|>


<commit_before>//! Finding the dominators in a control-flow graph.
//!
//! Algorithm based on Keith D. Cooper, Timothy J. Harvey, and Ken Kennedy,
//! "A Simple, Fast Dominance Algorithm",
//! Rice Computer Science TS-06-33870,
//! <https://www.cs.rice.edu/~keith/EMBED/dom.pdf>.

use super::iterate::reverse_post_order;
use super::ControlFlowGraph;
use rustc_index::vec::{Idx, IndexVec};
use std::borrow::BorrowMut;
use std::cmp::Ordering;

#[cfg(test)]
mod tests;

pub fn dominators<G: ControlFlowGraph>(graph: G) -> Dominators<G::Node> {
    let start_node = graph.start_node();
    let rpo = reverse_post_order(&graph, start_node);
    dominators_given_rpo(graph, &rpo)
}

fn dominators_given_rpo<G: ControlFlowGraph + BorrowMut<G>>(
    mut graph: G,
    rpo: &[G::Node],
) -> Dominators<G::Node> {
    let start_node = graph.borrow().start_node();
    assert_eq!(rpo[0], start_node);

    // compute the post order index (rank) for each node
    let mut post_order_rank: IndexVec<G::Node, usize> =
        (0..graph.borrow().num_nodes()).map(|_| 0).collect();
    for (index, node) in rpo.iter().rev().cloned().enumerate() {
        post_order_rank[node] = index;
    }

    let mut immediate_dominators: IndexVec<G::Node, Option<G::Node>> =
        (0..graph.borrow().num_nodes()).map(|_| None).collect();
    immediate_dominators[start_node] = Some(start_node);

    let mut changed = true;
    while changed {
        changed = false;

        for &node in &rpo[1..] {
            let mut new_idom = None;
            for pred in graph.borrow_mut().predecessors(node) {
                if immediate_dominators[pred].is_some() {
                    // (*) dominators for `pred` have been calculated
                    new_idom = Some(if let Some(new_idom) = new_idom {
                        intersect(&post_order_rank, &immediate_dominators, new_idom, pred)
                    } else {
                        pred
                    });
                }
            }

            if new_idom != immediate_dominators[node] {
                immediate_dominators[node] = new_idom;
                changed = true;
            }
        }
    }

    Dominators { post_order_rank, immediate_dominators }
}

fn intersect<Node: Idx>(
    post_order_rank: &IndexVec<Node, usize>,
    immediate_dominators: &IndexVec<Node, Option<Node>>,
    mut node1: Node,
    mut node2: Node,
) -> Node {
    while node1 != node2 {
        while post_order_rank[node1] < post_order_rank[node2] {
            node1 = immediate_dominators[node1].unwrap();
        }

        while post_order_rank[node2] < post_order_rank[node1] {
            node2 = immediate_dominators[node2].unwrap();
        }
    }

    node1
}

#[derive(Clone, Debug)]
pub struct Dominators<N: Idx> {
    post_order_rank: IndexVec<N, usize>,
    immediate_dominators: IndexVec<N, Option<N>>,
}

impl<Node: Idx> Dominators<Node> {
    pub fn is_reachable(&self, node: Node) -> bool {
        self.immediate_dominators[node].is_some()
    }

    pub fn immediate_dominator(&self, node: Node) -> Node {
        assert!(self.is_reachable(node), "node {:?} is not reachable", node);
        self.immediate_dominators[node].unwrap()
    }

    pub fn dominators(&self, node: Node) -> Iter<'_, Node> {
        assert!(self.is_reachable(node), "node {:?} is not reachable", node);
        Iter { dominators: self, node: Some(node) }
    }

    pub fn is_dominated_by(&self, node: Node, dom: Node) -> bool {
        // FIXME -- could be optimized by using post-order-rank
        self.dominators(node).any(|n| n == dom)
    }

    /// Provide deterministic ordering of nodes such that, if any two nodes have a dominator
    /// relationship, the dominator will always precede the dominated. (The relative ordering
    /// of two unrelated nodes will also be consistent, but otherwise the order has no
    /// meaning.) This method cannot be used to determine if either Node dominates the other.
    pub fn rank_partial_cmp(&self, lhs: Node, rhs: Node) -> Option<Ordering> {
        self.post_order_rank[lhs].partial_cmp(&self.post_order_rank[rhs])
    }
}

pub struct Iter<'dom, Node: Idx> {
    dominators: &'dom Dominators<Node>,
    node: Option<Node>,
}

impl<'dom, Node: Idx> Iterator for Iter<'dom, Node> {
    type Item = Node;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(node) = self.node {
            let dom = self.dominators.immediate_dominator(node);
            if dom == node {
                self.node = None; // reached the root
            } else {
                self.node = Some(dom);
            }
            Some(node)
        } else {
            None
        }
    }
}
<commit_msg>Clean up dominators_given_rpo<commit_after>//! Finding the dominators in a control-flow graph.
//!
//! Algorithm based on Keith D. Cooper, Timothy J. Harvey, and Ken Kennedy,
//! "A Simple, Fast Dominance Algorithm",
//! Rice Computer Science TS-06-33870,
//! <https://www.cs.rice.edu/~keith/EMBED/dom.pdf>.

use super::iterate::reverse_post_order;
use super::ControlFlowGraph;
use rustc_index::vec::{Idx, IndexVec};
use std::cmp::Ordering;

#[cfg(test)]
mod tests;

pub fn dominators<G: ControlFlowGraph>(graph: G) -> Dominators<G::Node> {
    let start_node = graph.start_node();
    let rpo = reverse_post_order(&graph, start_node);
    dominators_given_rpo(graph, &rpo)
}

fn dominators_given_rpo<G: ControlFlowGraph>(graph: G, rpo: &[G::Node]) -> Dominators<G::Node> {
    let start_node = graph.start_node();
    assert_eq!(rpo[0], start_node);

    // compute the post order index (rank) for each node
    let mut post_order_rank = IndexVec::from_elem_n(0, graph.num_nodes());
    for (index, node) in rpo.iter().rev().cloned().enumerate() {
        post_order_rank[node] = index;
    }

    let mut immediate_dominators = IndexVec::from_elem_n(None, graph.num_nodes());
    immediate_dominators[start_node] = Some(start_node);

    let mut changed = true;
    while changed {
        changed = false;

        for &node in &rpo[1..] {
            let mut new_idom = None;
            for pred in graph.predecessors(node) {
                if immediate_dominators[pred].is_some() {
                    // (*) dominators for `pred` have been calculated
                    new_idom = Some(if let Some(new_idom) = new_idom {
                        intersect(&post_order_rank, &immediate_dominators, new_idom, pred)
                    } else {
                        pred
                    });
                }
            }

            if new_idom != immediate_dominators[node] {
                immediate_dominators[node] = new_idom;
                changed = true;
            }
        }
    }

    Dominators { post_order_rank, immediate_dominators }
}

fn intersect<Node: Idx>(
    post_order_rank: &IndexVec<Node, usize>,
    immediate_dominators: &IndexVec<Node, Option<Node>>,
    mut node1: Node,
    mut node2: Node,
) -> Node {
    while node1 != node2 {
        while post_order_rank[node1] < post_order_rank[node2] {
            node1 = immediate_dominators[node1].unwrap();
        }

        while post_order_rank[node2] < post_order_rank[node1] {
            node2 = immediate_dominators[node2].unwrap();
        }
    }

    node1
}

#[derive(Clone, Debug)]
pub struct Dominators<N: Idx> {
    post_order_rank: IndexVec<N, usize>,
    immediate_dominators: IndexVec<N, Option<N>>,
}

impl<Node: Idx> Dominators<Node> {
    pub fn is_reachable(&self, node: Node) -> bool {
        self.immediate_dominators[node].is_some()
    }

    pub fn immediate_dominator(&self, node: Node) -> Node {
        assert!(self.is_reachable(node), "node {:?} is not reachable", node);
        self.immediate_dominators[node].unwrap()
    }

    pub fn dominators(&self, node: Node) -> Iter<'_, Node> {
        assert!(self.is_reachable(node), "node {:?} is not reachable", node);
        Iter { dominators: self, node: Some(node) }
    }

    pub fn is_dominated_by(&self, node: Node, dom: Node) -> bool {
        // FIXME -- could be optimized by using post-order-rank
        self.dominators(node).any(|n| n == dom)
    }

    /// Provide deterministic ordering of nodes such that, if any two nodes have a dominator
    /// relationship, the dominator will always precede the dominated. (The relative ordering
    /// of two unrelated nodes will also be consistent, but otherwise the order has no
    /// meaning.) This method cannot be used to determine if either Node dominates the other.
    pub fn rank_partial_cmp(&self, lhs: Node, rhs: Node) -> Option<Ordering> {
        self.post_order_rank[lhs].partial_cmp(&self.post_order_rank[rhs])
    }
}

pub struct Iter<'dom, Node: Idx> {
    dominators: &'dom Dominators<Node>,
    node: Option<Node>,
}

impl<'dom, Node: Idx> Iterator for Iter<'dom, Node> {
    type Item = Node;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(node) = self.node {
            let dom = self.dominators.immediate_dominator(node);
            if dom == node {
                self.node = None; // reached the root
            } else {
                self.node = Some(dom);
            }
            Some(node)
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>//! A variant of `SortedMap` that preserves insertion order.

use std::hash::{Hash, Hasher};
use std::iter::FromIterator;

use crate::stable_hasher::{HashStable, StableHasher};
use rustc_index::vec::{Idx, IndexVec};

/// An indexed multi-map that preserves insertion order while permitting both *O*(log *n*) lookup of
/// an item by key and *O*(1) lookup by index.
///
/// This data structure is a hybrid of an [`IndexVec`] and a [`SortedMap`]. Like `IndexVec`,
/// `SortedIndexMultiMap` assigns a typed index to each item while preserving insertion order.
/// Like `SortedMap`, `SortedIndexMultiMap` has efficient lookup of items by key. However, this
/// is accomplished by sorting an array of item indices instead of the items themselves.
///
/// Unlike `SortedMap`, this data structure can hold multiple equivalent items at once, so the
/// `get_by_key` method and its variants return an iterator instead of an `Option`. Equivalent
/// items will be yielded in insertion order.
///
/// Unlike a general-purpose map like `BTreeSet` or `HashSet`, `SortedMap` and
/// `SortedIndexMultiMap` require *O*(*n*) time to insert a single item. This is because we may need
/// to insert into the middle of the sorted array. Users should avoid mutating this data structure
/// in-place.
///
/// [`SortedMap`]: super::SortedMap
#[derive(Clone, Debug)]
pub struct SortedIndexMultiMap<I: Idx, K, V> {
    /// The elements of the map in insertion order.
    items: IndexVec<I, (K, V)>,

    /// Indices of the items in the set, sorted by the item's key.
    idx_sorted_by_item_key: Vec<I>,
}

impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
    pub fn new() -> Self {
        SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
    }

    pub fn len(&self) -> usize {
        self.items.len()
    }

    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }

    /// Returns an iterator over the items in the map in insertion order.
    pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> {
        self.items.into_iter()
    }

    /// Returns an iterator over the items in the map in insertion order along with their indices.
    pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> {
        self.items.into_iter_enumerated()
    }

    /// Returns an iterator over the items in the map in insertion order.
    pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
        self.items.iter().map(|(ref k, ref v)| (k, v))
    }

    /// Returns an iterator over the items in the map in insertion order along with their indices.
    pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> {
        self.items.iter_enumerated().map(|(i, (ref k, ref v))| (i, (k, v)))
    }

    /// Returns the item in the map with the given index.
    pub fn get(&self, idx: I) -> Option<&(K, V)> {
        self.items.get(idx)
    }

    /// Returns an iterator over the items in the map that are equal to `key`.
    ///
    /// If there are multiple items that are equivalent to `key`, they will be yielded in
    /// insertion order.
    pub fn get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V> {
        self.get_by_key_enumerated(key).map(|(_, v)| v)
    }

    /// Returns an iterator over the items in the map that are equal to `key` along with their
    /// indices.
    ///
    /// If there are multiple items that are equivalent to `key`, they will be yielded in
    /// insertion order.
    pub fn get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)> {
        let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key);
        self.idx_sorted_by_item_key[lower_bound..].iter().map_while(move |&i| {
            let (k, v) = &self.items[i];
            (k == &key).then_some((i, v))
        })
    }
}

impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {}
impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> {
    fn eq(&self, other: &Self) -> bool {
        // No need to compare the sorted index. If the items are the same, the index will be too.
        self.items == other.items
    }
}

impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
where
    K: Hash,
    V: Hash,
{
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.items.hash(hasher)
    }
}
impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
where
    K: HashStable<C>,
    V: HashStable<C>,
{
    fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
        self.items.hash_stable(ctx, hasher)
    }
}

impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
    fn from_iter<J>(iter: J) -> Self
    where
        J: IntoIterator<Item = (K, V)>,
    {
        let items = IndexVec::from_iter(iter);
        let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();

        // `sort_by_key` is stable, so insertion order is preserved for duplicate items.
        idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0);

        SortedIndexMultiMap { items, idx_sorted_by_item_key }
    }
}

impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
    type Output = V;

    fn index(&self, idx: I) -> &Self::Output {
        &self.items[idx].1
    }
}

#[cfg(tests)]
mod tests;
<commit_msg>Add `#[inline]`s to `SortedIndexMultiMap`<commit_after>//! A variant of `SortedMap` that preserves insertion order.

use std::hash::{Hash, Hasher};
use std::iter::FromIterator;

use crate::stable_hasher::{HashStable, StableHasher};
use rustc_index::vec::{Idx, IndexVec};

/// An indexed multi-map that preserves insertion order while permitting both *O*(log *n*) lookup of
/// an item by key and *O*(1) lookup by index.
///
/// This data structure is a hybrid of an [`IndexVec`] and a [`SortedMap`]. Like `IndexVec`,
/// `SortedIndexMultiMap` assigns a typed index to each item while preserving insertion order.
/// Like `SortedMap`, `SortedIndexMultiMap` has efficient lookup of items by key. However, this
/// is accomplished by sorting an array of item indices instead of the items themselves.
///
/// Unlike `SortedMap`, this data structure can hold multiple equivalent items at once, so the
/// `get_by_key` method and its variants return an iterator instead of an `Option`. Equivalent
/// items will be yielded in insertion order.
///
/// Unlike a general-purpose map like `BTreeSet` or `HashSet`, `SortedMap` and
/// `SortedIndexMultiMap` require *O*(*n*) time to insert a single item. This is because we may need
/// to insert into the middle of the sorted array. Users should avoid mutating this data structure
/// in-place.
///
/// [`SortedMap`]: super::SortedMap
#[derive(Clone, Debug)]
pub struct SortedIndexMultiMap<I: Idx, K, V> {
    /// The elements of the map in insertion order.
    items: IndexVec<I, (K, V)>,

    /// Indices of the items in the set, sorted by the item's key.
    idx_sorted_by_item_key: Vec<I>,
}

impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
    #[inline]
    pub fn new() -> Self {
        SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.items.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }

    /// Returns an iterator over the items in the map in insertion order.
    #[inline]
    pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> {
        self.items.into_iter()
    }

    /// Returns an iterator over the items in the map in insertion order along with their indices.
    #[inline]
    pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> {
        self.items.into_iter_enumerated()
    }

    /// Returns an iterator over the items in the map in insertion order.
    #[inline]
    pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
        self.items.iter().map(|(ref k, ref v)| (k, v))
    }

    /// Returns an iterator over the items in the map in insertion order along with their indices.
    #[inline]
    pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> {
        self.items.iter_enumerated().map(|(i, (ref k, ref v))| (i, (k, v)))
    }

    /// Returns the item in the map with the given index.
    #[inline]
    pub fn get(&self, idx: I) -> Option<&(K, V)> {
        self.items.get(idx)
    }

    /// Returns an iterator over the items in the map that are equal to `key`.
    ///
    /// If there are multiple items that are equivalent to `key`, they will be yielded in
    /// insertion order.
    #[inline]
    pub fn get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V> {
        self.get_by_key_enumerated(key).map(|(_, v)| v)
    }

    /// Returns an iterator over the items in the map that are equal to `key` along with their
    /// indices.
    ///
    /// If there are multiple items that are equivalent to `key`, they will be yielded in
    /// insertion order.
    #[inline]
    pub fn get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)> {
        let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key);
        self.idx_sorted_by_item_key[lower_bound..].iter().map_while(move |&i| {
            let (k, v) = &self.items[i];
            (k == &key).then_some((i, v))
        })
    }
}

impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {}
impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> {
    fn eq(&self, other: &Self) -> bool {
        // No need to compare the sorted index. If the items are the same, the index will be too.
        self.items == other.items
    }
}

impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
where
    K: Hash,
    V: Hash,
{
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.items.hash(hasher)
    }
}
impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
where
    K: HashStable<C>,
    V: HashStable<C>,
{
    fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
        self.items.hash_stable(ctx, hasher)
    }
}

impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
    fn from_iter<J>(iter: J) -> Self
    where
        J: IntoIterator<Item = (K, V)>,
    {
        let items = IndexVec::from_iter(iter);
        let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();

        // `sort_by_key` is stable, so insertion order is preserved for duplicate items.
        idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0);

        SortedIndexMultiMap { items, idx_sorted_by_item_key }
    }
}

impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
    type Output = V;

    fn index(&self, idx: I) -> &Self::Output {
        &self.items[idx].1
    }
}

#[cfg(tests)]
mod tests;
<|endoftext|>


<commit_before>use resize::Pixel::Gray8;
use resize::Type::Triangle;
use resize::Type::Lanczos3;
use std::fs;
use std::path::Path;

fn load_png(mut data: &[u8]) -> (usize, usize, Vec<u8>) {
    let decoder = png::Decoder::new(&mut data);
    let (info, mut reader) = decoder.read_info().unwrap();
    let mut src = vec![0; info.buffer_size()];
    reader.next_frame(&mut src).unwrap();

    (info.width as usize, info.height as usize, src)
}

fn write_png(path: &Path, w2: usize, h2: usize, pixels: &[u8]) {
    let outfh = fs::File::create(path).unwrap();
    let encoder = png::Encoder::new(outfh, w2 as u32, h2 as u32);
    encoder.write_header().unwrap().write_image_data(pixels).unwrap();
}

fn img_diff(a: &[u8], b: &[u8]) -> f64 {
    assert_eq!(a.len(), b.len());
    let sum = a.iter().cloned().zip(b.iter().cloned()).map(|(a,b)| {
        ((a as i32 - b as i32) as u32).pow(2)
    }).sum::<u32>();
    sum as f64 / a.len() as f64
}

fn assert_equals(img: &[u8], w2: usize, h2: usize, expected_filename: &str) {
    let (_, _, expected) = load_png(&fs::read(&expected_filename).expect(expected_filename));

    let diff = img_diff(&img, &expected);
    if diff > 0.00001 {
        let bad_file = Path::new(expected_filename).with_extension("failed-test.png");
        write_png(&bad_file, w2, h2, img);
        panic!("Test failed: {} differs by {}; see {}", expected_filename, diff, bad_file.display());
    }
}

#[test]
fn qa() {
    let tiger = &include_bytes!("../examples/tiger.png")[..];

    let (w1, h1, src) = load_png(&tiger);
    let mut res1 = vec![];
    let mut res2 = vec![];
    for w2 in [1, 2, 10, 100, 1000, 2000].iter().copied() {
        for h2 in [1, 2, 9, 99, 999, 1555].iter().copied() {
            res1.clear();
            res1.resize(w2 * h2, 0);

            resize::new(w1, h1, w2, h2, Gray8, Lanczos3).resize(&src, &mut res1);
            assert_equals(&res1, w2, h2, &format!("tests/t{}x{}.png", w2, h2));

            res2.clear();
            res2.resize(100 * 100, 255);

            resize::new(w2, h2, 100, 100, Gray8, Triangle).resize(&res1, &mut res2);
            assert_equals(&res2, 100, 100, &format!("tests/t{}x{}-100.png", w2, h2));
        }
    }
}
<commit_msg>Improved test<commit_after>use resize::Pixel::Gray8;
use resize::Type::Triangle;
use resize::Type::Lanczos3;
use std::fs;
use std::path::Path;

fn load_png(mut data: &[u8]) -> (usize, usize, Vec<u8>) {
    let decoder = png::Decoder::new(&mut data);
    let (info, mut reader) = decoder.read_info().unwrap();
    let mut src = vec![0; info.buffer_size()];
    reader.next_frame(&mut src).unwrap();

    (info.width as usize, info.height as usize, src)
}

fn write_png(path: &Path, w2: usize, h2: usize, pixels: &[u8]) {
    let outfh = fs::File::create(path).unwrap();
    let encoder = png::Encoder::new(outfh, w2 as u32, h2 as u32);
    encoder.write_header().unwrap().write_image_data(pixels).unwrap();
}

fn img_diff(a: &[u8], b: &[u8]) -> f64 {
    assert_eq!(a.len(), b.len());
    let sum = a.iter().cloned().zip(b.iter().cloned()).map(|(a,b)| {
        (a as i32 - b as i32).pow(2) as u32
    }).sum::<u32>();
    sum as f64 / a.len() as f64
}

fn assert_equals(img: &[u8], w2: usize, h2: usize, expected_filename: &str) {
    let (_, _, expected) = load_png(&fs::read(&expected_filename).expect(expected_filename));

    let diff = img_diff(&img, &expected);
    if diff > 0.0004 {
        let bad_file = Path::new(expected_filename).with_extension("failed-test.png");
        write_png(&bad_file, w2, h2, img);
        panic!("Test failed: {} differs by {}; see {} ", expected_filename, diff, bad_file.display());
    }
}

#[test]
fn qa() {
    let tiger = &include_bytes!("../examples/tiger.png")[..];

    let (w1, h1, src) = load_png(&tiger);
    let mut res1 = vec![];
    let mut res2 = vec![];
    for w2 in [1, 2, 10, 100, 1000, 2000].iter().copied() {
        for h2 in [1, 2, 9, 99, 999, 1555].iter().copied() {
            res1.clear();
            res1.resize(w2 * h2, 0);

            resize::new(w1, h1, w2, h2, Gray8, Lanczos3).resize(&src, &mut res1);
            assert_equals(&res1, w2, h2, &format!("tests/t{}x{}.png", w2, h2));

            res2.clear();
            res2.resize(100 * 100, 255);

            resize::new(w2, h2, 100, 100, Gray8, Triangle).resize(&res1, &mut res2);
            assert_equals(&res2, 100, 100, &format!("tests/t{}x{}-100.png", w2, h2));
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::net::udp::UdpSocket;
use std::io::net::ip::{Ipv4Addr, SocketAddr};

struct UtpPacketHeader {
    ver_type: u8, // ver: u4, type: u4
    extension: u8,
    connection_id: u16,
    timestamp_microseconds: u32,
    timestamp_difference_microseconds: u32,
    wnd_size: u32,
    seq_nr: u16,
    ack_nr: u16,
}

impl UtpPacketHeader {
}

fn connection_handler(stream: TcpStream) {
    let mut sock = stream;
    println!("New client {}", sock.peer_name().unwrap());

    // Read request
    let mut buf = [0, ..512];
    let mut count: uint = 0;

    match sock.read(buf) {
        Err(e) => { println!("Error detected! {}", e) }
        Ok(c) => {
            count = c;
            println!("Received {} bytes:", c);
            println!("{}", std::str::from_utf8(buf).unwrap());
        }
    }

    // Answer request
    match sock.write(buf.slice(0, count)) {
        _ => {}
    }

    println!("Gone!");
    drop(sock);

}

fn main() {
    let header = UtpPacketHeader {
        ver_type: 0 | 0 << 4,
        extension: 0,
        connection_id: 0,
        timestamp_microseconds: 0,
        timestamp_difference_microseconds: 0,
        wnd_size: 0,
        seq_nr: 0,
        ack_nr: 0,
    };

    let mut buf = [0, ..512];
    let mut sock = UdpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 }).unwrap();
    match sock.recvfrom(buf) {
        Ok((_, src)) => {
            unsafe {
                let buf: &[u8, ..20] = std::mem::transmute(&header);
                let _ = sock.sendto(*buf, src);
            }
        }
        Err(_) => {}
    }
    drop(sock);
    return;
    // Create socket
    let socket = TcpListener::bind("127.0.0.1", 8080);

    // Listen for new connections
    let mut acceptor = socket.listen();

    // Spawn a new process for handling each incoming connection
    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("{}", e) }
            Ok(stream) => spawn(proc() { connection_handler(stream) })
        }
    }

    drop(acceptor);
}
<commit_msg>Hide implementation.<commit_after>use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::net::udp::UdpSocket;
use std::io::net::ip::{Ipv4Addr, SocketAddr};

struct UtpPacketHeader {
    ver_type: u8, // ver: u4, type: u4
    extension: u8,
    connection_id: u16,
    timestamp_microseconds: u32,
    timestamp_difference_microseconds: u32,
    wnd_size: u32,
    seq_nr: u16,
    ack_nr: u16,
}

impl UtpPacketHeader {
    fn bytes(&self) -> &[u8] {
        unsafe {
            let buf: &[u8, ..20] = std::mem::transmute(self);
            return buf.as_slice();
        }
    }
}

fn connection_handler(stream: TcpStream) {
    let mut sock = stream;
    println!("New client {}", sock.peer_name().unwrap());

    // Read request
    let mut buf = [0, ..512];
    let mut count: uint = 0;

    match sock.read(buf) {
        Err(e) => { println!("Error detected! {}", e) }
        Ok(c) => {
            count = c;
            println!("Received {} bytes:", c);
            println!("{}", std::str::from_utf8(buf).unwrap());
        }
    }

    // Answer request
    match sock.write(buf.slice(0, count)) {
        _ => {}
    }

    println!("Gone!");
    drop(sock);

}

fn main() {
    let header = UtpPacketHeader {
        ver_type: 0 | 0 << 4,
        extension: 0,
        connection_id: 0,
        timestamp_microseconds: 0,
        timestamp_difference_microseconds: 0,
        wnd_size: 0,
        seq_nr: 0,
        ack_nr: 0,
    };

    let mut buf = [0, ..512];
    let mut sock = UdpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 }).unwrap();
    match sock.recvfrom(buf) {
        Ok((_, src)) => {
            let _ = sock.sendto(header.bytes(), src);
        }
        Err(_) => {}
    }
    drop(sock);
    return;
    // Create socket
    let socket = TcpListener::bind("127.0.0.1", 8080);

    // Listen for new connections
    let mut acceptor = socket.listen();

    // Spawn a new process for handling each incoming connection
    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("{}", e) }
            Ok(stream) => spawn(proc() { connection_handler(stream) })
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::transmute;
    use std::rand::random;
    use std::fmt;

    #[allow(dead_code,non_camel_case_types)]
    #[deriving(PartialEq,Eq)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    impl fmt::Show for UtpPacketType {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let s = match *self {
                ST_DATA  => "ST_DATA",
                ST_FIN   => "ST_FIN",
                ST_STATE => "ST_STATE",
                ST_RESET => "ST_RESET",
                ST_SYN   => "ST_SYN",
            };
            write!(f, "{}", s)
        }
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.type_ver;
            self.type_ver = t as u8 << 4 | version;
        }

        fn get_type(self) -> UtpPacketType {
            let t: UtpPacketType = unsafe { transmute(self.type_ver >> 4) };
            t
        }

        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }

        fn decode(buf: &[u8]) -> UtpPacketHeader {
            UtpPacketHeader {
                type_ver: buf[0],
                extension: buf[1],
                connection_id: buf[2] as u16 << 8 | buf[3] as u16,
                timestamp_microseconds: buf[4] as u32 << 24 | buf[5] as u32 << 16 | buf[6] as u32 << 8 | buf[7] as u32,
                timestamp_difference_microseconds: buf[8] as u32 << 24 | buf[9] as u32 << 16 | buf[10] as u32 << 8 | buf[11] as u32,
                wnd_size: buf[12] as u32 << 24 | buf[13] as u32 << 16 | buf[14] as u32 << 8 | buf[15] as u32,
                seq_nr: buf[16] as u16 << 8 | buf[17] as u16,
                ack_nr: buf[18] as u16 << 8 | buf[19] as u16,
            }
        }
    }

    impl fmt::Show for UtpPacketHeader {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let t: UtpPacketType = self.get_type();
            write!(f, "(type: {}, version: {}, extension: {}, connection_id: {}, timestamp_microseconds: {}, timestamp_difference_microseconds: {}, wnd_size: {}, seq_nr: {}, ack_nr: {})", t, self.type_ver & 0x0F, self.extension, self.connection_id, self.timestamp_microseconds, self.timestamp_difference_microseconds, self.wnd_size, self.seq_nr, self.ack_nr)
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            self.header.set_type(t);
        }

        fn get_type(self) -> UtpPacketType {
            self.header.get_type()
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }

        fn decode(buf: &[u8]) -> UtpPacket {
            let header = UtpPacketHeader::decode(buf);
            UtpPacket { header: header, payload: Vec::from_slice(buf.slice(20, buf.len())) }
        }
    }

    #[allow(non_camel_case_types)]
    enum UtpSocketState {
        CS_NEW,
        CS_CONNECTED,
        CS_SYN_SENT,
    }

    pub struct UtpSocket {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
        state: UtpSocketState,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            let r: u16 = random();
            match skt {
                Ok(x)  => Ok(UtpSocket {
                    socket: x,
                    connected_to: addr,
                    connection_id: r.to_be(),
                    seq_nr: 1,
                    ack_nr: 0,
                    state: CS_NEW,
                }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(mut self, other: SocketAddr) -> UtpSocket {
            self.connected_to = other;
            /*
            let stream = UtpStream::new(s, other);
            stream.connect()
            */
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.connected_to;
            let _result = self.socket.sendto(packet.bytes().as_slice(), dst);

            self.state = CS_SYN_SENT;

            let mut buf = [0, ..512];
            let (len, addr) = self.socket.recvfrom(buf).unwrap();
            println!("{} {}", addr, self.connected_to);
            assert!(addr == self.connected_to);
            let packet = UtpPacket::decode(buf.slice(0, len));
            //assert!(packet.get_type() == ST_STATE);
            println!("{}", packet.header);

            self.seq_nr += 1;

            self.state = CS_CONNECTED;

            self
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            let mut b = [0, ..512];
            let response = self.socket.recvfrom(b);

            let _src: SocketAddr;
            let read;
            match response {
                Ok((nread, src)) => { read = nread; _src = src },
                Err(e) => return Err(e),
            };

            let packet = UtpPacket::decode(b);
            let x = packet.header;
            println!("{}", packet.header);
            match packet.get_type() {
                ST_SYN => { // Respond with an ACK
                    self.seq_nr = random();
                    let mut resp = UtpPacket::new();
                    resp.set_type(ST_STATE);
                    let t = time::get_time();
                    let self_t_micro: u32 = (t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32;
                    let other_t_micro: u32 = x.timestamp_microseconds;
                    resp.header.timestamp_microseconds = self_t_micro.to_be();
                    resp.header.timestamp_difference_microseconds = (self_t_micro.to_le() - other_t_micro.to_le()).to_be();
                    resp.header.connection_id = self.connection_id;
                    resp.header.seq_nr = self.seq_nr;
                    resp.header.ack_nr = x.seq_nr.to_le().to_be();
                    self.socket.sendto(resp.bytes().as_slice(), _src);

                    self.seq_nr += 1;

                    self.state = CS_CONNECTED;
                }
                _ => {}
            };

            for i in range(0u, ::std::cmp::min(buf.len(), read-20)) {
                buf[i] = b[i+20];
            }

            println!("{}", Vec::from_slice(buf.slice(0,read-20)));

            response
        }

        pub fn sendto(&mut self, buf: &[u8], dst: SocketAddr) -> IoResult<()> {
            self.socket.sendto(buf, dst)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            let r: u16 = random();
            UtpSocket {
                socket: self.socket.clone(),
                connected_to: self.connected_to,
                connection_id: r.to_be(),
                seq_nr: 1,
                ack_nr: 0,
                state: CS_NEW,
            }
        }
    }

    /*
    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UtpSocket,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UtpSocket, conn: SocketAddr) -> UtpStream {
            UtpStream {
                socket: socket,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.socket.connection_id;
            packet.header.seq_nr = self.socket.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.socket.connected_to;
            let _result = self.socket.sendto(packet.bytes().as_slice(), dst);

            let mut buf = [0, ..512];
            let (len, addr) = self.socket.recvfrom(buf).unwrap();
            println!("{} {}", addr, self.socket.connected_to);
            assert!(addr == self.socket.connected_to);
            let packet = UtpPacket::decode(buf.slice(0, len));
            //assert!(packet.get_type() == ST_STATE);
            println!("{}", packet.header);

            self.socket.seq_nr += 1;

            self
        }

        pub fn terminate(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_FIN);
            packet.header.connection_id = self.socket.connection_id;
            packet.header.seq_nr = self.socket.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.socket.connected_to;
            let _result = self.socket.sendto(packet.bytes().as_slice(), dst);

            self.socket.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UtpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.socket.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e)
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.socket.connection_id;
            packet.header.seq_nr = self.socket.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.socket.connected_to;
            let result = self.socket.sendto(packet.bytes().as_slice(), dst);

            self.socket.seq_nr += 1;
            result
        }
    }
    */
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::UtpSocket;
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_nread, _src)) => {
                        spawn(proc() {
                            /*
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                            */
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let stream = sock.connect(addr);
            //let _ = stream.write([0]);
            //let mut buf = [0, ..512];
            //stream.read(buf);
            //let stream = stream.terminate();
            drop(stream);
        }
        _ => usage(),
    }
}
<commit_msg>Encapsulate sendto data into packets.<commit_after>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::transmute;
    use std::rand::random;
    use std::fmt;

    #[allow(dead_code,non_camel_case_types)]
    #[deriving(PartialEq,Eq)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    impl fmt::Show for UtpPacketType {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let s = match *self {
                ST_DATA  => "ST_DATA",
                ST_FIN   => "ST_FIN",
                ST_STATE => "ST_STATE",
                ST_RESET => "ST_RESET",
                ST_SYN   => "ST_SYN",
            };
            write!(f, "{}", s)
        }
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.type_ver;
            self.type_ver = t as u8 << 4 | version;
        }

        fn get_type(self) -> UtpPacketType {
            let t: UtpPacketType = unsafe { transmute(self.type_ver >> 4) };
            t
        }

        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }

        fn decode(buf: &[u8]) -> UtpPacketHeader {
            UtpPacketHeader {
                type_ver: buf[0],
                extension: buf[1],
                connection_id: buf[2] as u16 << 8 | buf[3] as u16,
                timestamp_microseconds: buf[4] as u32 << 24 | buf[5] as u32 << 16 | buf[6] as u32 << 8 | buf[7] as u32,
                timestamp_difference_microseconds: buf[8] as u32 << 24 | buf[9] as u32 << 16 | buf[10] as u32 << 8 | buf[11] as u32,
                wnd_size: buf[12] as u32 << 24 | buf[13] as u32 << 16 | buf[14] as u32 << 8 | buf[15] as u32,
                seq_nr: buf[16] as u16 << 8 | buf[17] as u16,
                ack_nr: buf[18] as u16 << 8 | buf[19] as u16,
            }
        }
    }

    impl fmt::Show for UtpPacketHeader {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let t: UtpPacketType = self.get_type();
            write!(f, "(type: {}, version: {}, extension: {}, connection_id: {}, timestamp_microseconds: {}, timestamp_difference_microseconds: {}, wnd_size: {}, seq_nr: {}, ack_nr: {})", t, self.type_ver & 0x0F, self.extension, self.connection_id, self.timestamp_microseconds, self.timestamp_difference_microseconds, self.wnd_size, self.seq_nr, self.ack_nr)
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            self.header.set_type(t);
        }

        fn get_type(self) -> UtpPacketType {
            self.header.get_type()
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }

        fn decode(buf: &[u8]) -> UtpPacket {
            let header = UtpPacketHeader::decode(buf);
            UtpPacket { header: header, payload: Vec::from_slice(buf.slice(20, buf.len())) }
        }
    }

    #[allow(non_camel_case_types)]
    enum UtpSocketState {
        CS_NEW,
        CS_CONNECTED,
        CS_SYN_SENT,
    }

    pub struct UtpSocket {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
        state: UtpSocketState,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            let r: u16 = random();
            match skt {
                Ok(x)  => Ok(UtpSocket {
                    socket: x,
                    connected_to: addr,
                    connection_id: r.to_be(),
                    seq_nr: 1,
                    ack_nr: 0,
                    state: CS_NEW,
                }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(mut self, other: SocketAddr) -> UtpSocket {
            self.connected_to = other;
            /*
            let stream = UtpStream::new(s, other);
            stream.connect()
            */
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.connected_to;
            let _result = self.socket.sendto(packet.bytes().as_slice(), dst);

            self.state = CS_SYN_SENT;

            let mut buf = [0, ..512];
            let (len, addr) = self.socket.recvfrom(buf).unwrap();
            println!("{} {}", addr, self.connected_to);
            assert!(addr == self.connected_to);
            let packet = UtpPacket::decode(buf.slice(0, len));
            //assert!(packet.get_type() == ST_STATE);
            println!("{}", packet.header);

            self.seq_nr += 1;

            self.state = CS_CONNECTED;

            self
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            let mut b = [0, ..512];
            let response = self.socket.recvfrom(b);

            let _src: SocketAddr;
            let read;
            match response {
                Ok((nread, src)) => { read = nread; _src = src },
                Err(e) => return Err(e),
            };

            let packet = UtpPacket::decode(b);
            let x = packet.header;
            println!("{}", packet.header);
            match packet.get_type() {
                ST_SYN => { // Respond with an ACK
                    self.seq_nr = random();
                    let mut resp = UtpPacket::new();
                    resp.set_type(ST_STATE);
                    let t = time::get_time();
                    let self_t_micro: u32 = (t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32;
                    let other_t_micro: u32 = x.timestamp_microseconds;
                    resp.header.timestamp_microseconds = self_t_micro.to_be();
                    resp.header.timestamp_difference_microseconds = (self_t_micro.to_le() - other_t_micro.to_le()).to_be();
                    resp.header.connection_id = self.connection_id;
                    resp.header.seq_nr = self.seq_nr;
                    resp.header.ack_nr = x.seq_nr.to_le().to_be();
                    self.socket.sendto(resp.bytes().as_slice(), _src);

                    self.seq_nr += 1;

                    self.state = CS_CONNECTED;
                }
                _ => {}
            };

            for i in range(0u, ::std::cmp::min(buf.len(), read-20)) {
                buf[i] = b[i+20];
            }

            println!("{}", Vec::from_slice(buf.slice(0,read-20)));

            response
        }

        pub fn sendto(&mut self, buf: &[u8], dst: SocketAddr) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();
            packet.header.seq_nr = self.seq_nr.to_be();
            packet.header.connection_id = self.connection_id;

            self.seq_nr += 1;
            self.socket.sendto(packet.bytes().as_slice(), dst)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            let r: u16 = random();
            UtpSocket {
                socket: self.socket.clone(),
                connected_to: self.connected_to,
                connection_id: r.to_be(),
                seq_nr: 1,
                ack_nr: 0,
                state: CS_NEW,
            }
        }
    }

    /*
    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UtpSocket,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UtpSocket, conn: SocketAddr) -> UtpStream {
            UtpStream {
                socket: socket,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.socket.connection_id;
            packet.header.seq_nr = self.socket.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.socket.connected_to;
            let _result = self.socket.sendto(packet.bytes().as_slice(), dst);

            let mut buf = [0, ..512];
            let (len, addr) = self.socket.recvfrom(buf).unwrap();
            println!("{} {}", addr, self.socket.connected_to);
            assert!(addr == self.socket.connected_to);
            let packet = UtpPacket::decode(buf.slice(0, len));
            //assert!(packet.get_type() == ST_STATE);
            println!("{}", packet.header);

            self.socket.seq_nr += 1;

            self
        }

        pub fn terminate(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_FIN);
            packet.header.connection_id = self.socket.connection_id;
            packet.header.seq_nr = self.socket.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.socket.connected_to;
            let _result = self.socket.sendto(packet.bytes().as_slice(), dst);

            self.socket.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UtpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.socket.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e)
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.socket.connection_id;
            packet.header.seq_nr = self.socket.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let dst = self.socket.connected_to;
            let result = self.socket.sendto(packet.bytes().as_slice(), dst);

            self.socket.seq_nr += 1;
            result
        }
    }
    */
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::UtpSocket;
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_nread, _src)) => {
                        spawn(proc() {
                            /*
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                            */
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let stream = sock.connect(addr);
            //let _ = stream.write([0]);
            //let mut buf = [0, ..512];
            //stream.read(buf);
            //let stream = stream.terminate();
            drop(stream);
        }
        _ => usage(),
    }
}
<|endoftext|>