<commit_before>#![crate_id = "seax-svm"]
#![crate_type="lib"]
#![feature(box_syntax)]
#![allow(dead_code)]

/// Contains the Seax Virtual Machine (SVM) and miscellaneous
/// support code.
pub mod svm {
    use svm::slist::List;
    use svm::slist::Stack;
    use std::iter::IteratorExt;

    /// Singly-linked list and stack implementations. `List<T>` is a
    /// singly-linked cons list with boxed items. `Stack<T>` is basically
    /// just a struct containing a boxed pointer to the head of a list,
    /// and some methods.
    pub mod slist {

        use svm::slist::List::{Cons,Nil};
        use std::fmt::Show;

        /// A stack implementation wrapping a `List<T>`
        ///
        /// This is essentially just a struct containing a boxed pointer
        /// to the head of a `List<T>`, so that when items are pushed to or
        /// popped from the stack, the pointer is changed to point to the
        /// new head item. There may be saner ways of doing this.
        pub struct Stack<T> {
            /// The head item of the list.
            head: Box<List<T>>
        }

        impl<T> Stack<T> {
            /// Push an item to the top of the stack, returning a new stack
            pub fn push(self, it: T) -> Stack<T> {
                Stack { head: Box::new(self.head.prepend(it)) }
            }

            /// Peak at the top item of the stack.
            ///
            /// Returns Some<T> if there is an item on top of the stack,
            /// and None if the stack is empty.
            pub fn peek(&self) -> Option<&T> {
                match *self.head {
                    Nil => None,
                    Cons(ref it,_) => Some(it)
                }
            }

            /// Returns an empty stack.
            pub fn empty() -> Stack<T> {
                Stack { head: box Nil }
            }

            /// Wraps a list into a stack.
            pub fn new(l: List<T>) -> Stack<T> {
                Stack { head: box l }
            }
        }

        /// Singly-linked cons list.
        ///
        /// This is used internally to represent list primitives in the machine.
        pub enum List<T> {
            Cons(T, Box<List<T>>),
            Nil,
        }

        /// Public implementation for List.
        impl<T> List<T> {

            /// Creates a new empty list
            pub fn new() -> List<T> {
                Nil
            }

            /// Prepends the given item to the list, returning the new head item.
            pub fn prepend(self, it: T) -> List<T> {
                Cons(it, box self)
            }

            /// Returns the length of the list.
            pub fn length (&self) -> isize {
                match *self {
                    Cons(_, ref tail) => 1 + tail.length(),
                    Nil => 0
                }
            }
        }

        impl<T> List<T> where T: Show {
            /// Return a string representation of the list
            fn to_string(&self) -> String {
                match *self {
                    Cons(ref head, ref tail) => format!("({:?}, {})", head, tail.to_string()),
                    Nil => format!("nil")
                }
            }

        }

        /// Convenience macro for making lists.
        ///
        /// Usage: `list!(1i32, 2i32, 3i32);` expands to
        /// `Cons(1i32, Box::new(Cons(2i32, Box::new(Cons(3i32, Box::new(Nil))))));`.
        macro_rules! list(
            ( $e:expr, $($rest:expr),+ ) => ( Cons($e, Box::new(list!( $( $rest ),+ )) ));
            ( $e:expr ) => ( Cons($e, Box::new(Nil)) );
            () => ( @Empty )
        );

        #[cfg(test)]
        mod tests {
            use super::{List, Stack};
            use super::List::{Cons,Nil};

            #[test]
            fn test_list_length() {
                let full_list: List<i32> = list!(1i32, 2i32, 3i32);
                let empty_list: List<i32> = List::new();
                assert_eq!(full_list.length(), 3);
                assert_eq!(empty_list.length(), 0);
            }

            #[test]
            fn test_list_to_string() {
                let l: List<i32> = Cons(1, Box::new(Cons(2, Box::new(Cons(3, Box::new(Nil))))));
                assert_eq!(l.to_string(), "(1i32, (2i32, (3i32, nil)))");
            }

            #[test]
            fn test_stack_peek() {
                let s: Stack<i32> = Stack::new(list!(1i32, 2i32, 3i32));
                assert_eq!(s.peek(), Some(&1));
            }

            #[test]
            fn test_stack_push() {
                let mut s: Stack<i32> = Stack::empty();
                assert_eq!(s.peek(), None);
                s = s.push(1);
                assert_eq!(s.peek(), Some(&1));
                s = s.push(6);
                assert_eq!(s.peek(), Some(&6));
            }

            #[test]
            fn test_stack_pop() {
                let mut s: Stack<i32> = Stack::empty();
                assert_eq!(s.peek(), None); // TODO: implement
                s = s.push(1);
                assert_eq!(s.peek(), Some(&1));
                s = s.push(6);
                assert_eq!(s.peek(), Some(&6));
            }

            #[test]
            fn test_list_macro() {
                let l: List<i32> = list!(1i32, 2i32, 3i32);
                assert_eq!(l.to_string(), "(1i32, (2i32, (3i32, nil)))")
            }
        }
    }

    /// SVM item types
    pub enum SVMCell {
        Atom,
        ListCell(Box<List<SVMCell>>)
    }

    pub enum Atom {
        AtomUInt(usize),
        AtomSInt(isize),
        AtomFloat(f64),
        AtomStr(String)
    }

    /// SVM instruction types
    pub enum SVMInstruction {
        /// `nil`
        ///
        /// Pushes an empty list (nil) onto the stack
        InstNIL,
        /// `ldc`: `L`oa`d` `C`onstant. Loads a constant (atom)
        InstLDC(Atom),
        /// `ld`: `L`oa`d`. Pushes a variable onto the stack.
        ///
        /// The variable is indicated by the argument, a pair.
        /// The pair's `car` specifies the level, the `cdr` the position.
        /// So `(1 . 3)` gives the current function's (level 1) third
        /// parameter.
        InstLD,
        /// `ldf`: `L`oa`d` `F`unction.
        ///
        ///  Takes one list argument representing a function and constructs
        ///  a closure (a pair containing the function and the current
        ///  environment) and pushes that onto the stack.
        InstLDF,
        /// `join`
        ///
        /// Pops a list reference from the dump and makes this the new value
        /// of `C`. This instruction occurs at the end of both alternatives of
        ///  a `sel`.
        InstJOIN,
        /// `ap`: `Ap`ply.
        ///
        /// Pops a closure and a list of parameter values from the stack.
        /// The closure is applied to the parameters by installing its
        /// environment as the current one, pushing the parameter list
        /// in front of that, clearing the stack, and setting `C` to the
        /// closure's function pointer. The previous values of `S`, `E`,
        ///  and the next value of `C` are saved on the dump.
        InstAP,
        /// `ret`: `Ret`urn.
        ///
        /// Pops one return value from the stack, restores
        /// `S`, `E`, and `C` from the dump, and pushes
        /// the return value onto the now-current stack.
        InstRET,
        /// `dum`: `Dum`my.
        ///
        /// Pops a dummy environment (an empty list) onto the `E` stack.
        InstDUM,
        /// `rap`: `R`ecursive `Ap`ply.
        /// Works like `ap`, only that it replaces an occurrence of a
        /// dummy environment with the current one, thus making recursive
        ///  functions possible.
        InstRAP,
        /// `sel`: `Sel`ect
        ///
        /// Expects two list arguments, and pops a value from the stack.
        /// The first list is executed if the popped value was non-nil,
        /// the second list otherwise. Before one of these list pointers
        ///  is made the new `C`, a pointer to the instruction following
        ///  `sel` is saved on the dump.
        InstSEL,
        /// `add`
        ///
        /// Pops two numbers off of the stack and adds them, pushing the
        /// result onto the stack. This will up-convert integers to floating
        /// point if necessary.
        ///
        /// TODO: figure out what happens when you try to add things that aren't
        /// numbers (maybe the compiler won't let this happen?).
        InstADD,
        /// `sub`: `Sub`tract
        ///
        /// Pops two numbers off of the stack and subtracts the first from the
        /// second, pushing the result onto the stack. This will up-convert
        /// integers to floating point if necessary.
        ///
        /// TODO: figure out what happens when you try to subtract things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstSUB,
        /// `mul`: `Mul`tiply
        ///
        /// Pops two numbers off of the stack and multiplies them, pushing the
        /// result onto the stack. This will up-convert integers to floating
        /// point if necessary.
        ///
        /// TODO: figure out what happens when you try to multiply things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstMUL,
        /// `div`: `Div`ide
        ///
        /// Pops two numbers off of the stack and divides the first by the second,
        /// pushing the result onto the stack. This performs integer division.
        ///
        /// TODO: figure out what happens when you try to divide things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstDIV,
        /// `fdiv`: `F`loating-point `div`ide
        ///
        /// Pops two numbers off of the stack and divides the first by the second,
        /// pushing the result onto the stack. This performs float division.
        ///
        /// TODO: figure out what happens when you try to divide things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        ///
        /// TODO: Not sure if there should be separate float and int divide words
        /// I guess the compiler can figure this out
        InstFDIV,
        /// `mod`: `Mod`ulo
        ///
        /// Pops two numbers off of the stack and divides the first by the second,
        /// pushing the remainder onto the stack.
        ///
        /// TODO: figure out what happens when you try to modulo things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstMOD
        // TODO: add some hardcoded I/O instructions here so that you can
        //  do I/O without farming everything out to `stdio`
        // TODO: add `cons` and `cdr` words
    }

    /// Represents a SVM machine state
    pub struct State {
        stack: Stack<SVMCell>,
        env: Stack<SVMCell>,
        control: Stack<SVMCell>,
        dump: Stack<SVMCell>
    }

    impl State {

        /// Creates a new empty state
        fn new() -> State {
            State {
                stack: Stack::empty(),
                env: Stack::empty(),
                control: Stack::empty(),
                dump: Stack::empty()
            }
        }

        /// Evaluates an instruction.
        ///
        /// Evaluates an instruction against a state, returning a new state.
        pub fn eval(self, inst: SVMInstruction) -> State {
            match inst {
                SVMInstruction::InstNIL => State {
                    stack: self.stack.push(SVMCell::ListCell(box List::new())),
                    env: self.env,
                    control: self.control,
                    dump: self.dump
                },
                _ => { unimplemented!() }
            }
        }
    }

    /*
    /// Evaluates a program.
    ///
    /// Evaluates a program represented as an `Iterator` of `SVMInstruction`s.
    /// Returns the final machine state at the end of execution

    pub fn evalProgram(insts: Iterator<Item=SVMInstruction>) -> State {
        insts.fold(State::new(), |last_state: State, inst: SVMInstruction| last_state.eval(inst));
    }*/

}
<commit_msg>Add tests for State bits<commit_after>#![crate_id = "seax-svm"]
#![crate_type="lib"]
#![feature(box_syntax)]

/// Contains the Seax Virtual Machine (SVM) and miscellaneous
/// support code.
pub mod svm {
    use svm::slist::List;
    use svm::slist::Stack;
    use std::iter::IteratorExt;
    use std::fmt;

    /// Singly-linked list and stack implementations. `List<T>` is a
    /// singly-linked cons list with boxed items. `Stack<T>` is basically
    /// just a struct containing a boxed pointer to the head of a list,
    /// and some methods.
    pub mod slist {

        use svm::slist::List::{Cons,Nil};
        use std::fmt;

        /// A stack implementation wrapping a `List<T>`
        ///
        /// This is essentially just a struct containing a boxed pointer
        /// to the head of a `List<T>`, so that when items are pushed to or
        /// popped from the stack, the pointer is changed to point to the
        /// new head item. There may be saner ways of doing this.
        pub struct Stack<T> {
            /// The head item of the list.
            head: Box<List<T>>
        }

        impl<T> Stack<T> {
            /// Push an item to the top of the stack, returning a new stack
            pub fn push(self, it: T) -> Stack<T> {
                Stack { head: Box::new(self.head.prepend(it)) }
            }

            /// Peak at the top item of the stack.
            ///
            /// Returns Some<T> if there is an item on top of the stack,
            /// and None if the stack is empty.
            pub fn peek(&self) -> Option<&T> {
                match *self.head {
                    Nil => None,
                    Cons(ref it,_) => Some(it)
                }
            }

            /// Returns an empty stack.
            pub fn empty() -> Stack<T> {
                Stack { head: box Nil }
            }

            /// Wraps a list into a stack.
            pub fn new(l: List<T>) -> Stack<T> {
                Stack { head: box l }
            }

            /// Returns the length of the stack. This just calls
            /// `List::length()` on the wrapped list.
            pub fn length(&self) -> isize {
                self.head.length()
            }
        }

        /// Singly-linked cons list.
        ///
        /// This is used internally to represent list primitives in the
        /// machine.
        #[deriving(Show)]
        #[deriving(PartialEq)]
        pub enum List<T> {
            Cons(T, Box<List<T>>),
            Nil,
        }

        /// Public implementation for List.
        impl<T> List<T> {

            /// Creates a new empty list
            pub fn new() -> List<T> {
                Nil
            }

            /// Prepends the given item to the list, returning the new head item.
            pub fn prepend(self, it: T) -> List<T> {
                Cons(it, box self)
            }

            /// Returns the length of the list.
            pub fn length (&self) -> isize {
                match *self {
                    Cons(_, ref tail) => 1 + tail.length(),
                    Nil => 0
                }
            }
        }

        impl<T> List<T> where T: fmt::Show {
            /// Return a string representation of the list
            fn to_string(&self) -> String {
                match *self {
                    Cons(ref head, ref tail) => format!("({:?}, {})", head, tail.to_string()),
                    Nil => format!("nil")
                }
            }
        }

        impl<T> fmt::Show for List<T> {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                write!(f, "{}", self.to_string)
            }
        }

        /// Convenience macro for making lists.
        ///
        /// Usage: `list!(1i32, 2i32, 3i32);` expands to
        /// `Cons(1i32, Box::new(Cons(2i32, Box::new(Cons(3i32, Box::new(Nil))))));`.
        macro_rules! list(
            ( $e:expr, $($rest:expr),+ ) => ( Cons($e, Box::new(list!( $( $rest ),+ )) ));
            ( $e:expr ) => ( Cons($e, Box::new(Nil)) );
            () => ( @Empty )
        );

        #[cfg(test)]
        mod tests {
            use super::{List, Stack};
            use super::List::{Cons,Nil};

            #[test]
            fn test_list_length() {
                let full_list: List<i32> = list!(1i32, 2i32, 3i32);
                let empty_list: List<i32> = List::new();
                assert_eq!(full_list.length(), 3);
                assert_eq!(empty_list.length(), 0);
            }

            #[test]
            fn test_list_to_string() {
                let l: List<i32> = Cons(1, Box::new(Cons(2, Box::new(Cons(3, Box::new(Nil))))));
                assert_eq!(l.to_string(), "(1i32, (2i32, (3i32, nil)))");
            }

            #[test]
            fn test_stack_length() {
                let full_stack: Stack<i32> = Stack::new(list!(1i32, 2i32, 3i32));
                let empty_stack: Stack<i32> = Stack::empty();
                assert_eq!(full_stack.length(), 3);
                assert_eq!(empty_stack.length(), 0);
            }

            #[test]
            fn test_stack_peek() {
                let full_stack: Stack<i32> = Stack::new(list!(1i32, 2i32, 3i32));
                let empty_stack: Stack<i32> = Stack::empty();
                assert_eq!(full_stack.peek(), Some(&1));
                assert_eq!(empty_stack.peek(), None);
            }

            #[test]
            fn test_stack_push() {
                let mut s: Stack<i32> = Stack::empty();
                assert_eq!(s.peek(), None);
                s = s.push(1);
                assert_eq!(s.peek(), Some(&1));
                s = s.push(6);
                assert_eq!(s.peek(), Some(&6));
            }

            #[test]
            fn test_stack_pop() {
                let mut s: Stack<i32> = Stack::empty();
                assert_eq!(s.peek(), None); // TODO: implement
                s = s.push(1);
                assert_eq!(s.peek(), Some(&1));
                s = s.push(6);
                assert_eq!(s.peek(), Some(&6));
            }

            #[test]
            fn test_list_macro() {
                let l: List<i32> = list!(1i32, 2i32, 3i32);
                assert_eq!(l.to_string(), "(1i32, (2i32, (3i32, nil)))")
            }
        }
    }

    /// SVM item types
    #[deriving(Show)]
    #[deriving(PartialEq)]
    pub enum SVMCell {
        Atom,
        ListCell(Box<List<SVMCell>>)
    }

    pub enum Atom {
        AtomUInt(usize),
        AtomSInt(isize),
        AtomFloat(f64),
        AtomStr(String)
    }

    /// SVM instruction types
    pub enum SVMInstruction {
        /// `nil`
        ///
        /// Pushes an empty list (nil) onto the stack
        InstNIL,
        /// `ldc`: `L`oa`d` `C`onstant. Loads a constant (atom)
        InstLDC(Atom),
        /// `ld`: `L`oa`d`. Pushes a variable onto the stack.
        ///
        /// The variable is indicated by the argument, a pair.
        /// The pair's `car` specifies the level, the `cdr` the position.
        /// So `(1 . 3)` gives the current function's (level 1) third
        /// parameter.
        InstLD,
        /// `ldf`: `L`oa`d` `F`unction.
        ///
        ///  Takes one list argument representing a function and constructs
        ///  a closure (a pair containing the function and the current
        ///  environment) and pushes that onto the stack.
        InstLDF,
        /// `join`
        ///
        /// Pops a list reference from the dump and makes this the new value
        /// of `C`. This instruction occurs at the end of both alternatives of
        ///  a `sel`.
        InstJOIN,
        /// `ap`: `Ap`ply.
        ///
        /// Pops a closure and a list of parameter values from the stack.
        /// The closure is applied to the parameters by installing its
        /// environment as the current one, pushing the parameter list
        /// in front of that, clearing the stack, and setting `C` to the
        /// closure's function pointer. The previous values of `S`, `E`,
        ///  and the next value of `C` are saved on the dump.
        InstAP,
        /// `ret`: `Ret`urn.
        ///
        /// Pops one return value from the stack, restores
        /// `S`, `E`, and `C` from the dump, and pushes
        /// the return value onto the now-current stack.
        InstRET,
        /// `dum`: `Dum`my.
        ///
        /// Pops a dummy environment (an empty list) onto the `E` stack.
        InstDUM,
        /// `rap`: `R`ecursive `Ap`ply.
        /// Works like `ap`, only that it replaces an occurrence of a
        /// dummy environment with the current one, thus making recursive
        ///  functions possible.
        InstRAP,
        /// `sel`: `Sel`ect
        ///
        /// Expects two list arguments, and pops a value from the stack.
        /// The first list is executed if the popped value was non-nil,
        /// the second list otherwise. Before one of these list pointers
        ///  is made the new `C`, a pointer to the instruction following
        ///  `sel` is saved on the dump.
        InstSEL,
        /// `add`
        ///
        /// Pops two numbers off of the stack and adds them, pushing the
        /// result onto the stack. This will up-convert integers to floating
        /// point if necessary.
        ///
        /// TODO: figure out what happens when you try to add things that aren't
        /// numbers (maybe the compiler won't let this happen?).
        InstADD,
        /// `sub`: `Sub`tract
        ///
        /// Pops two numbers off of the stack and subtracts the first from the
        /// second, pushing the result onto the stack. This will up-convert
        /// integers to floating point if necessary.
        ///
        /// TODO: figure out what happens when you try to subtract things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstSUB,
        /// `mul`: `Mul`tiply
        ///
        /// Pops two numbers off of the stack and multiplies them, pushing the
        /// result onto the stack. This will up-convert integers to floating
        /// point if necessary.
        ///
        /// TODO: figure out what happens when you try to multiply things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstMUL,
        /// `div`: `Div`ide
        ///
        /// Pops two numbers off of the stack and divides the first by the second,
        /// pushing the result onto the stack. This performs integer division.
        ///
        /// TODO: figure out what happens when you try to divide things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstDIV,
        /// `fdiv`: `F`loating-point `div`ide
        ///
        /// Pops two numbers off of the stack and divides the first by the second,
        /// pushing the result onto the stack. This performs float division.
        ///
        /// TODO: figure out what happens when you try to divide things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        ///
        /// TODO: Not sure if there should be separate float and int divide words
        /// I guess the compiler can figure this out
        InstFDIV,
        /// `mod`: `Mod`ulo
        ///
        /// Pops two numbers off of the stack and divides the first by the second,
        /// pushing the remainder onto the stack.
        ///
        /// TODO: figure out what happens when you try to modulo things that
        /// aren't numbers (maybe the compiler won't let this happen?).
        InstMOD
        // TODO: add some hardcoded I/O instructions here so that you can
        //  do I/O without farming everything out to `stdio`
        // TODO: add `cons` and `cdr` words
    }

    /// Represents a SVM machine state
    pub struct State {
        stack: Stack<SVMCell>,
        env: Stack<SVMCell>,
        control: Stack<SVMCell>,
        dump: Stack<SVMCell>
    }

    impl State {

        /// Creates a new empty state
        fn new() -> State {
            State {
                stack: Stack::empty(),
                env: Stack::empty(),
                control: Stack::empty(),
                dump: Stack::empty()
            }
        }

        /// Evaluates an instruction.
        ///
        /// Evaluates an instruction against a state, returning a new state.
        pub fn eval(self, inst: SVMInstruction) -> State {
            match inst {
                SVMInstruction::InstNIL => State {
                    stack: self.stack.push(SVMCell::ListCell(box List::new())),
                    env: self.env,
                    control: self.control,
                    dump: self.dump
                },
                _ => { unimplemented!() }
            }
        }
    }

    /*
    /// Evaluates a program.
    ///
    /// Evaluates a program represented as an `Iterator` of `SVMInstruction`s.
    /// Returns the final machine state at the end of execution

    pub fn evalProgram(insts: Iterator<Item=SVMInstruction>) -> State {
        insts.fold(State::new(), |last_state: State, inst: SVMInstruction| last_state.eval(inst));
    }*/

    #[cfg(test)]
    mod tests {
        use super::State;
        use super::SVMInstruction;
        use super::slist::List::{Cons,Nil};

        #[test]
        fn test_empty_state() {
            let state = State::new();
            assert_eq!(state.stack.length(), 0);
            assert_eq!(state.env.length(), 0);
            assert_eq!(state.control.length(), 0);
            assert_eq!(state.dump.length(), 0);
        }

        #[test]
        fn test_eval_nil () {
            let mut state = State::new();
            assert_eq!(state.stack.peek(), None);
            state = state.eval(SVMInstruction::InstNIL);
            assert_eq!(state.stack.peek(), Some(&Nil));
        }
    }

}
<|endoftext|>


<commit_before>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

fn serialize(saddr: &SocketAddr) -> [u8; 18] {
    let mut result = [0; 18];
    result[..16].copy_from_slice(&match saddr.ip() {
        IpAddr::V4(ref ip) => &ip.to_ipv6_mapped(),
        IpAddr::V6(ref ip) => ip,
    }.octets());
    result[16..].copy_from_slice(&saddr.port().to_ne_bytes());
    result
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = serialize(&src);
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(serialize(&src)) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<commit_msg>Fix build<commit_after>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

fn serialize(saddr: &SocketAddr) -> [u8; 18] {
    let mut result = [0; 18];
    result[..16].copy_from_slice(&match saddr.ip() {
        IpAddr::V4(ref ip) => ip.to_ipv6_mapped().octets(),
        IpAddr::V6(ref ip) => ip.octets(),
    });
    result[16..].copy_from_slice(&saddr.port().to_ne_bytes());
    result
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = serialize(&src);
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(serialize(&src)) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<|endoftext|>


<commit_before>
use std::cmp;
use std::mem;
use std::iter::*;
use image_data::*;
use rayon::prelude::*;

fn minf(a: f32, b: f32) -> f32 { if a < b { a } else { b } }
fn maxf(a: f32, b: f32) -> f32 { if a > b { a } else { b } }
fn saturate(a: f32) -> f32 { maxf(minf(a, 1.0), 0.0) }
fn fnorm(a: u8) -> f32 { a as f32 / 255.0 }
fn bnorm(a: f32) -> u8 { (saturate(a) * 255.0) as u8 }
fn to_rgb(a: &Rgba) -> Rgb { [a[0], a[1], a[2]] }
fn to_rgba(a: &Rgb) -> Rgba { [a[0], a[1], a[2], 255u8] }

const ALPHA_THRESHOLD: u8 = 128;

fn pixel_dist(a: &Rgba, b: &Rgba) -> u32 {
	fn dist_one(a: u8, b: u8) -> u32 {
		let x = a as i32 - b as i32;
		(x * x) as u32
	}
	dist_one(a[0], b[0]) +
	dist_one(a[1], b[1]) +
	dist_one(a[2], b[2])
}

fn pixel_dist_alpha(a: &Rgba, b: &Rgba) -> u32 {
	if a[3] < ALPHA_THRESHOLD && b[3] < ALPHA_THRESHOLD {
		0
	} else {
		pixel_dist(a, b) + 1
	}
}

fn block_dist<T: Fn(&Rgba, &Rgba) -> u32> (table: &[Rgba; 4], pixels: &Block, dist: T) -> u32 {
	let mut sum = 0u32;
	let mut max = 0u32;
	for pix in pixels.into_iter().rev() {
		let e = (0..4).into_iter().map(|x| dist(&table[x], pix)).min().unwrap();
		sum += e;
		max = cmp::max(max, e);
	}
	sum + max * 16
}

fn build_table(ends: &(Rgb, Rgb)) -> [Rgba; 4] {
	fn build_entry(ends: &(Rgb, Rgb), index: usize) -> Rgba {
		let b = if index < 2 { index * 3 } else { index - 1 } as u16;
		let a = 3 - b;
		[((b * ends.0[0] as u16 + a * ends.1[0] as u16) / 3) as u8,
		 ((b * ends.0[1] as u16 + a * ends.1[1] as u16) / 3) as u8,
		 ((b * ends.0[2] as u16 + a * ends.1[2] as u16) / 3) as u8,
		 255]
	}
	let mut interps = [[0u8; 4]; 4];
	for i in 0..4 {
		interps[i] = build_entry(ends, i);
	}
	interps
}

fn build_table_alpha(ends: &(Rgb, Rgb)) -> [Rgba; 4] {
	let mut mid = [0u8; 3];
	for i in 0..3 {
		mid[i] = (((ends.0[i] as u16) + (ends.1[i] as u16)) / 2) as u8;
	}
	[to_rgba(&ends.0), to_rgba(&ends.1), to_rgba(&mid), [0u8; 4]]
}

fn minmax(pixels: &Block) -> (Rgb, Rgb) {
	let mut min = [255u8; 3];
	let mut max = [0u8; 3];
	for p in pixels.into_iter() {
		for i in 0..3 {
			min[i] = cmp::min(min[i], p[i]);
			max[i] = cmp::max(max[i], p[i]);
		}
	}
	(min, max)
}

fn quantize_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	fn quantize(end: &Rgb) -> Rgb {
		[end[0] & 0xF8, end[1] & 0xFC, end[2] & 0xF8]
	}
	(quantize(&ends.0), quantize(&ends.1))
}

fn extrapolate_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	let mut a = ends.0;
	let mut b = ends.1;
	for i in 0..3 {
		let f = (fnorm(ends.0[i]) - fnorm(ends.1[i])) * 1.5;
		a[i] = bnorm(fnorm(ends.0[i]) - f);
		b[i] = bnorm(fnorm(ends.1[i]) + f);
	}
	(a, b)
}

fn endpoints_sorted(ends: &(Rgb, Rgb)) -> bool {
	encode_endpoint(&ends.0) < encode_endpoint(&ends.1)
}

fn sort_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	if endpoints_sorted(ends) {
		*ends
	} else {
		(ends.1, ends.0)
	}
}

fn sort_quantize_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	quantize_endpoints(&sort_endpoints(ends))
}

fn build_endpoints(pixels: &Block, quality: u8) -> Vec<(Rgb, Rgb)> {
	let (min, max) = quantize_endpoints(&minmax(pixels));

	let mut out = Vec::new();
	out.push((min, max));

	if quality == 255 {
		for i in 0..3 {
			for o in (min[i] + 1)..max[i] {
				let mut min = min;
				min[i] = o;

				for i in 0..3 {
					for o in (min[i] + 1)..max[i] {
						let mut max = max;
						max[i] = o;

						out.push((min, max));
					}
				}
			}
		}
	} if quality > 0 && min != max {
		for i in 0..(BLOCK_PIXELS - 1) {
			for j in (i + 1)..BLOCK_PIXELS {

				let px = (to_rgb(&pixels[i]), to_rgb(&pixels[j]));
				let qpx = sort_quantize_endpoints(&px);
				out.push(qpx);

				if quality > 1 {
					let extrapolated = sort_quantize_endpoints(&extrapolate_endpoints(&px));
					out.push(extrapolated);

					if quality > 2 {
						let e = (extrapolated.0, max);
						if endpoints_sorted(&e) {
							out.push(e);
						}

						let e = (min, extrapolated.1);
						if endpoints_sorted(&e) {
							out.push(e);
						}

						if quality > 3 {
							out.push(sort_endpoints(&(extrapolated.0, qpx.1)));
							out.push(sort_endpoints(&(qpx.0, extrapolated.1)));
						}
					}
				}

			}
		}
	}

	out
}

fn encode_endpoint(end: &Rgb) -> u16 {
	((end[0] as u16 >> 3) << 11) |
	((end[1] as u16 >> 2) << 5) |
	(end[2] as u16 >> 3)
}

fn encode_pixels<T: Fn(&Rgba, &Rgba) -> u32>(table: &[Rgba; 4], pixels: &Block, dist: T) -> u32 {
	let mut mask = 0u32;
	for pix in pixels.into_iter().rev() {
		let best = (0..4).into_iter().min_by_key(|x| dist(&table[*x], pix)).unwrap() as u32;
		mask = (mask << 2) | (best & 0x03);
	}
	mask
}

fn encode_block(pixels: &Block, quality: u8) -> u64 {
	let alpha = pixels.into_iter().any(|p| p[3] < ALPHA_THRESHOLD);

	if alpha {
		let (ends, table) =
			build_endpoints(pixels, quality).iter()
			.map(|ends| (*ends, build_table_alpha(ends)))
			.min_by_key(|tbl| block_dist(&tbl.1, pixels, pixel_dist_alpha)).unwrap();

		let min = encode_endpoint(&ends.1) as u64;
		let max = encode_endpoint(&ends.0) as u64;

		let encoded = encode_pixels(&table, &pixels, pixel_dist_alpha) as u64;
		(encoded << 32) | (min << 16) | max

	} else {
		let (ends, table) =
			build_endpoints(pixels, quality).iter()
			.map(|ends| (*ends, build_table(ends)))
			.min_by_key(|tbl| block_dist(&tbl.1, pixels, pixel_dist)).unwrap();

		let min = encode_endpoint(&ends.0) as u64;
		let max = encode_endpoint(&ends.1) as u64;

		let encoded = encode_pixels(&table, &pixels, pixel_dist) as u64;
		(encoded << 32) | (min << 16) | max
	}
}

pub fn encode(image: &ImageData, quality: u8) -> Result<Vec<u8>, ()> {
	if image.size.0 % BLOCK_SIZE != 0 || image.size.1 % BLOCK_SIZE != 0 {
		return Err(());
	}

	let blocks = image.blocks().count();

    let mut out = Vec::with_capacity(blocks);
	(0..blocks).into_par_iter()
		.map(|i| {
			encode_block(&image.blocks().nth(i).unwrap(), quality)
		}).collect_into(&mut out);

	unsafe {
		let mut out: Vec<u8> = mem::transmute(out);
		out.set_len(blocks * 8);
		Ok(out)
	}
}
<commit_msg>Better quantization.<commit_after>
use std::cmp;
use std::mem;
use std::iter::*;
use image_data::*;
use rayon::prelude::*;

fn minf(a: f32, b: f32) -> f32 { if a < b { a } else { b } }
fn maxf(a: f32, b: f32) -> f32 { if a > b { a } else { b } }
fn saturate(a: f32) -> f32 { maxf(minf(a, 1.0), 0.0) }
fn fnorm(a: u8) -> f32 { a as f32 / 255.0 }
fn bnorm(a: f32) -> u8 { (saturate(a) * 255.0) as u8 }
fn to_rgb(a: &Rgba) -> Rgb { [a[0], a[1], a[2]] }
fn to_rgba(a: &Rgb) -> Rgba { [a[0], a[1], a[2], 255u8] }

const ALPHA_THRESHOLD: u8 = 128;

fn pixel_dist(a: &Rgba, b: &Rgba) -> u32 {
	fn dist_one(a: u8, b: u8) -> u32 {
		let x = a as i32 - b as i32;
		(x * x) as u32
	}
	dist_one(a[0], b[0]) +
	dist_one(a[1], b[1]) +
	dist_one(a[2], b[2])
}

fn pixel_dist_alpha(a: &Rgba, b: &Rgba) -> u32 {
	if a[3] < ALPHA_THRESHOLD && b[3] < ALPHA_THRESHOLD {
		0
	} else {
		pixel_dist(a, b) + 1
	}
}

fn block_dist<T: Fn(&Rgba, &Rgba) -> u32>(table: &[Rgba; 4], pixels: &Block, dist: T) -> u32 {
	let mut sum = 0u32;
	let mut max = 0u32;
	for pix in pixels.into_iter().rev() {
		let e = (0..4).into_iter().map(|x| dist(&table[x], pix)).min().unwrap();
		sum += e;
		max = cmp::max(max, e);
	}
	sum + max * 16
}

fn build_table(ends: &(Rgb, Rgb)) -> [Rgba; 4] {
	fn build_entry(ends: &(Rgb, Rgb), index: usize) -> Rgba {
		let b = if index < 2 { index * 3 } else { index - 1 } as u16;
		let a = 3 - b;
		[((b * ends.0[0] as u16 + a * ends.1[0] as u16) / 3) as u8,
		 ((b * ends.0[1] as u16 + a * ends.1[1] as u16) / 3) as u8,
		 ((b * ends.0[2] as u16 + a * ends.1[2] as u16) / 3) as u8,
		 255]
	}
	let mut interps = [[0u8; 4]; 4];
	for i in 0..4 {
		interps[i] = build_entry(ends, i);
	}
	interps
}

fn build_table_alpha(ends: &(Rgb, Rgb)) -> [Rgba; 4] {
	let mut mid = [0u8; 3];
	for i in 0..3 {
		mid[i] = (((ends.0[i] as u16) + (ends.1[i] as u16)) / 2) as u8;
	}
	[to_rgba(&ends.0), to_rgba(&ends.1), to_rgba(&mid), [0u8; 4]]
}

fn minmax(pixels: &Block) -> (Rgb, Rgb) {
	let mut min = [255u8; 3];
	let mut max = [0u8; 3];
	for p in pixels.into_iter() {
		for i in 0..3 {
			min[i] = cmp::min(min[i], p[i]);
			max[i] = cmp::max(max[i], p[i]);
		}
	}
	(min, max)
}

fn quantize_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	fn quantize(end: &Rgb) -> Rgb {
		//[end[0] & 0xF8, end[1] & 0xFC, end[2] & 0xF8]
		let bits = [3, 2, 3];
		let mut e = *end;
		for i in 0..3 {
			let q = (e[i] >> bits[i]) << bits[i];
			e[i] = q + ((e[i] >> (bits[i] - 1)) & 0x01);
		}
		e
	}
	(quantize(&ends.0), quantize(&ends.1))
}

fn extrapolate_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	let mut a = ends.0;
	let mut b = ends.1;
	for i in 0..3 {
		let f = (fnorm(ends.0[i]) - fnorm(ends.1[i])) * 1.5;
		a[i] = bnorm(fnorm(ends.0[i]) - f);
		b[i] = bnorm(fnorm(ends.1[i]) + f);
	}
	(a, b)
}

fn endpoints_sorted(ends: &(Rgb, Rgb)) -> bool {
	encode_endpoint(&ends.0) < encode_endpoint(&ends.1)
}

fn sort_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	if endpoints_sorted(ends) {
		*ends
	} else {
		(ends.1, ends.0)
	}
}

fn sort_quantize_endpoints(ends: &(Rgb, Rgb)) -> (Rgb, Rgb) {
	quantize_endpoints(&sort_endpoints(ends))
}

fn build_endpoints(pixels: &Block, quality: u8) -> Vec<(Rgb, Rgb)> {
	let (min, max) = quantize_endpoints(&minmax(pixels));

	let mut out = Vec::new();
	out.push((min, max));

	if quality == 255 {
		for i in 0..3 {
			for o in (min[i] + 1)..max[i] {
				let mut min = min;
				min[i] = o;

				for i in 0..3 {
					for o in (min[i] + 1)..max[i] {
						let mut max = max;
						max[i] = o;

						out.push((min, max));
					}
				}
			}
		}
	} if quality > 0 && min != max {
		for i in 0..(BLOCK_PIXELS - 1) {
			for j in (i + 1)..BLOCK_PIXELS {

				let px = (to_rgb(&pixels[i]), to_rgb(&pixels[j]));
				let qpx = sort_quantize_endpoints(&px);
				out.push(qpx);

				if quality > 1 {
					let extrapolated = sort_quantize_endpoints(&extrapolate_endpoints(&px));
					out.push(extrapolated);

					if quality > 2 {
						let e = (extrapolated.0, max);
						if endpoints_sorted(&e) {
							out.push(e);
						}

						let e = (min, extrapolated.1);
						if endpoints_sorted(&e) {
							out.push(e);
						}

						if quality > 3 {
							out.push(sort_endpoints(&(extrapolated.0, qpx.1)));
							out.push(sort_endpoints(&(qpx.0, extrapolated.1)));
						}
					}
				}

			}
		}
	}

	out
}

fn encode_endpoint(end: &Rgb) -> u16 {
	((end[0] as u16 >> 3) << 11) |
	((end[1] as u16 >> 2) << 5) |
	(end[2] as u16 >> 3)
}

fn encode_pixels<T: Fn(&Rgba, &Rgba) -> u32>(table: &[Rgba; 4], pixels: &Block, dist: T) -> u32 {
	let mut mask = 0u32;
	for pix in pixels.into_iter().rev() {
		let best = (0..4).into_iter().min_by_key(|x| dist(&table[*x], pix)).unwrap() as u32;
		mask = (mask << 2) | (best & 0x03);
	}
	mask
}

fn encode_block(pixels: &Block, quality: u8) -> u64 {
	let alpha = pixels.into_iter().any(|p| p[3] < ALPHA_THRESHOLD);

	if alpha {
		let (ends, table) =
			build_endpoints(pixels, quality).iter()
			.map(|ends| (*ends, build_table_alpha(ends)))
			.min_by_key(|tbl| block_dist(&tbl.1, pixels, pixel_dist_alpha)).unwrap();

		let min = encode_endpoint(&ends.1) as u64;
		let max = encode_endpoint(&ends.0) as u64;

		let encoded = encode_pixels(&table, &pixels, pixel_dist_alpha) as u64;
		(encoded << 32) | (min << 16) | max

	} else {
		let (ends, table) =
			build_endpoints(pixels, quality).iter()
			.map(|ends| (*ends, build_table(ends)))
			.min_by_key(|tbl| block_dist(&tbl.1, pixels, pixel_dist)).unwrap();

		let min = encode_endpoint(&ends.0) as u64;
		let max = encode_endpoint(&ends.1) as u64;

		let encoded = encode_pixels(&table, &pixels, pixel_dist) as u64;
		(encoded << 32) | (min << 16) | max
	}
}

pub fn encode(image: &ImageData, quality: u8) -> Result<Vec<u8>, ()> {
	if image.size.0 % BLOCK_SIZE != 0 || image.size.1 % BLOCK_SIZE != 0 {
		return Err(());
	}

	let blocks = image.blocks().count();

    let mut out = Vec::with_capacity(blocks);
	(0..blocks).into_par_iter()
		.map(|i| {
			encode_block(&image.blocks().nth(i).unwrap(), quality)
		}).collect_into(&mut out);

	unsafe {
		let mut out: Vec<u8> = mem::transmute(out);
		out.set_len(blocks * 8);
		Ok(out)
	}
}
<|endoftext|>


<commit_before>#![deny(missing_docs)]

use varint::consume_single_varint;
use stream::*;
use std::io::Read;

/// A consumer for a byte stream
pub struct ByteConsumer<'a> {
    read: &'a mut Read,
    type_: StreamType,
}

impl<'a> ByteConsumer<'a> {
    /// Return a ByteConsumer from for single messages, varint or leb128-delimited
    pub fn new(read: &'a mut Read, type_: StreamType) -> ByteConsumer {
        ByteConsumer { read, type_ }
    }
}

impl<'a> Iterator for ByteConsumer<'a> {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.type_ {
            StreamType::Leb128 | StreamType::Varint => consume_single_varint(self.read),
            StreamType::Single => {
                let ret: Option<Vec<u8>>;
                let mut buf = Vec::new();
                match self.read.read_to_end(&mut buf) {
                    Ok(x) => {
                        if x > 0 {
                            ret = Some(buf);
                        } else {
                            ret = None
                        }
                    }
                    Err(_) => ret = None,
                }
                ret
            }
        }
    }
}
<commit_msg>Make ByteConsumer own its reader<commit_after>#![deny(missing_docs)]

use varint::consume_single_varint;
use stream::*;
use std::io::Read;

/// A consumer for a byte stream
pub struct ByteConsumer<T: Read> {
    read: T,
    type_: StreamType,
}

impl <T: Read> ByteConsumer<T> {
    /// Return a ByteConsumer from for single messages, varint or leb128-delimited
    pub fn new(read: T, type_: StreamType) -> ByteConsumer<T> {
        ByteConsumer { read, type_ }
    }
}

impl <T: Read> Iterator for ByteConsumer<T> {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.type_ {
            StreamType::Leb128 | StreamType::Varint => consume_single_varint(&mut self.read),
            StreamType::Single => {
                let ret: Option<Vec<u8>>;
                let mut buf = Vec::new();
                match self.read.read_to_end(&mut buf) {
                    Ok(x) => {
                        if x > 0 {
                            ret = Some(buf);
                        } else {
                            ret = None
                        }
                    }
                    Err(_) => ret = None,
                }
                ret
            }
        }
    }
}
<|endoftext|>


<commit_before>use clap::{crate_version, App, Arg};
use csv;
use io::Write;
use std::io;
use std::process;
use table::ascii_table_writer;

const DESCRIPTION: &str = "A command to print ASCII table from stdin";
const USAGE_TEMPLATE: &str = r#"
Usage:
    {usage}

Flags:
{flags}
"#;

fn main() {
    let matches = App::new("table")
        .version(crate_version!())
        .author("Naoto Kaneko <naoty.k@gmail.com>")
        .about(DESCRIPTION)
        .template(USAGE_TEMPLATE.trim())
        .version_short("v")
        .arg(
            Arg::with_name("has_headers")
                .short("H")
                .long("header")
                .help("Prints table with headers"),
        )
        .get_matches();

    if let Err(error) = start(matches.is_present("has_headers")) {
        eprintln!("{}", error);
        process::exit(1);
    }
}

fn start(has_headers: bool) -> io::Result<()> {
    let mut reader = csv::ReaderBuilder::new()
        .delimiter(b'\t')
        .has_headers(false)
        .from_reader(io::stdin());

    let mut writer = ascii_table_writer::new(io::stdout(), has_headers);

    for result in reader.records() {
        if let Err(error) = result {
            eprintln!("failed to read TSV: {}", error);
            process::exit(1);
        }

        let record = result.unwrap();

        for field in record.iter() {
            writer.write(field.as_bytes())?;
        }

        writer.write(b"\n")?;
    }

    writer.flush()
}
<commit_msg>Refactor some codes<commit_after>use clap::{crate_version, App, Arg};
use csv;
use io::Write;
use std::io;
use std::process;
use table::ascii_table_writer;

const DESCRIPTION: &str = "A command to print ASCII table from stdin";
const USAGE_TEMPLATE: &str = r#"
Usage:
    {usage}

Flags:
{flags}
"#;

fn main() {
    let matches = App::new("table")
        .version(crate_version!())
        .author("Naoto Kaneko <naoty.k@gmail.com>")
        .about(DESCRIPTION)
        .template(USAGE_TEMPLATE.trim())
        .version_short("v")
        .arg(
            Arg::with_name("header")
                .short("H")
                .long("header")
                .help("Prints table with headers"),
        )
        .get_matches();

    if let Err(error) = start(matches.is_present("header")) {
        eprintln!("{}", error);
        process::exit(1);
    }
}

fn start(has_headers: bool) -> io::Result<()> {
    let mut reader = csv::ReaderBuilder::new()
        .delimiter(b'\t')
        .has_headers(false)
        .from_reader(io::stdin());

    let mut writer = ascii_table_writer::new(io::stdout(), has_headers);

    for result in reader.records() {
        match result {
            Ok(record) => {
                for field in record.iter() {
                    writer.write(field.as_bytes())?;
                }
                writer.write(b"\n")?;
            }
            Err(error) => {
                eprintln!("failed to read record: {}", error);
                process::exit(1);
            }
        }
    }

    writer.flush()
}
<|endoftext|>


<commit_before>#![feature(collections)]
#![feature(unicode)]

//modules

mod config;
mod word;

//extern crates

extern crate rand;
extern crate rustc_serialize;
extern crate regex;

//use

use config::*;
use word::*;
use std::io::prelude::*;
//use rustc_serialize::*;    //commented out as it is only used to emit test configs
use std::fs::File;
use rand::distributions::Weighted;

//static variables

static CONFIG_FILE_NAME: &'static str = "json_sample_config.json";

//functions

fn main() {
    let config: WordGeneratorConfig = load_config(&CONFIG_FILE_NAME);

    //comment other code and uncomment these to emit a new test config
    //let mut file = File::create(&CONFIG_FILE_NAME).unwrap();

    /*match file.write(json::encode(&config).unwrap().as_bytes()) {
        Ok(_) => (),
        Err(err) => panic!("Error writing test config: {}", err),
    }*/

    let mut first_syllable_list : Vec<Weighted<String>> = Vec::new();
    let mut last_syllable_list : Vec<Weighted<String>> = Vec::new();
    let mut normal_syllable_list : Vec<Weighted<String>> = Vec::new();

    //generate weighted lists of syllable configurations in each syllable position

    for syllable in config.syllables {
        let syllable_weighted : Weighted<String> = Weighted{
            weight: syllable.weight as u32,
            item: syllable.string.clone()
        };
        //if the syllable is eligible for all positions
        if !(syllable.only_first_syllable || syllable.only_last_syllable) {
            normal_syllable_list.push(syllable_weighted.clone());
            first_syllable_list.push(syllable_weighted.clone());
            last_syllable_list.push(syllable_weighted.clone());
        }
        //these need to be separate ifs, as marking only first + only last syllable means it can be used for either
        if syllable.only_first_syllable {
            first_syllable_list.push(syllable_weighted.clone());
        }
        if syllable.only_last_syllable {
            last_syllable_list.push(syllable_weighted.clone());
        }

    }

    let mut grapheme_groups : Vec<(String, Vec<Weighted<String>>)> = Vec::new();

    //convert list of graphemes in config into lists of Weighted
    for group in config.graphemes {
        let mut graphemes_converted : Vec<Weighted<String>> = Vec::new();

        graphemes_converted.extend(
            group.graphemes.iter().map(
                |g: &Grapheme| Weighted{weight: g.weight as u32, item: g.string.clone() }
                )
            );
        grapheme_groups.push((group.name.clone(), graphemes_converted));
    }

    let word_factory : WordFactory = WordFactory {
        first_syllable_list : first_syllable_list,
        normal_syllable_list : normal_syllable_list,
        last_syllable_list : last_syllable_list,
        graphemes : grapheme_groups,
        settings : config.settings.clone(),
        rewrites : config.rewrites.clone(),
        rejects : config.rejects.clone(),
    };

    let mut file = File::create(&config.output_settings.output_file).unwrap();

    for _ in 0..config.output_settings.word_count {

        //generate a string of syllables
        // - parse config list of syllables into 3 lists, one for first syllable, one for last, and one for rest

        //apply syllable level rewrites and rejects
        //generate graphemes for each syllable
        //apply grapheme level rewrites and rejects
        //write to file
        let mut word : Word = Word {
                syllables : String::new(),
                graphemes : String::new(),
                syllable_rewrite_history : Vec::new(),
                grapheme_rewrite_history : Vec::new(),
                syllable_rejects : Vec::new(),
                grapheme_rejects : Vec::new()
            };

        word_factory.generate_syllables(&mut word);
        word_factory.rewrite_syllables(&mut word);
        word_factory.mark_syllable_rejects(&mut word);

        word_factory.generate_graphemes(&mut word);
        word_factory.rewrite_graphemes(&mut word);
        word_factory.mark_grapheme_rejects(&mut word);

        match file.write(get_word_graphemes(&word).as_bytes()) {
                Err(error) => panic!("error {} writing to file", error),
                Ok(_) => (),
        };
        match file.write("\n".as_bytes()) {
                Err(error) => panic!("error {} writing to file", error),
                Ok(_) => (),
        };
    }
}
<commit_msg>implemented rewrites_before_rejects from settings<commit_after>#![feature(collections)]
#![feature(unicode)]

//modules

mod config;
mod word;

//extern crates

extern crate rand;
extern crate rustc_serialize;
extern crate regex;

//use

use config::*;
use word::*;
use std::io::prelude::*;
//use rustc_serialize::*;    //commented out as it is only used to emit test configs
use std::fs::File;
use rand::distributions::Weighted;

//static variables

static CONFIG_FILE_NAME: &'static str = "json_sample_config.json";

//functions

fn main() {
    let config: WordGeneratorConfig = load_config(&CONFIG_FILE_NAME);

    //comment other code and uncomment these to emit a new test config
    //let mut file = File::create(&CONFIG_FILE_NAME).unwrap();

    /*match file.write(json::encode(&config).unwrap().as_bytes()) {
        Ok(_) => (),
        Err(err) => panic!("Error writing test config: {}", err),
    }*/

    let mut first_syllable_list : Vec<Weighted<String>> = Vec::new();
    let mut last_syllable_list : Vec<Weighted<String>> = Vec::new();
    let mut normal_syllable_list : Vec<Weighted<String>> = Vec::new();

    //generate weighted lists of syllable configurations in each syllable position

    for syllable in config.syllables {
        let syllable_weighted : Weighted<String> = Weighted{
            weight: syllable.weight as u32,
            item: syllable.string.clone()
        };
        //if the syllable is eligible for all positions
        if !(syllable.only_first_syllable || syllable.only_last_syllable) {
            normal_syllable_list.push(syllable_weighted.clone());
            first_syllable_list.push(syllable_weighted.clone());
            last_syllable_list.push(syllable_weighted.clone());
        }
        //these need to be separate ifs, as marking only first + only last syllable means it can be used for either
        if syllable.only_first_syllable {
            first_syllable_list.push(syllable_weighted.clone());
        }
        if syllable.only_last_syllable {
            last_syllable_list.push(syllable_weighted.clone());
        }

    }

    let mut grapheme_groups : Vec<(String, Vec<Weighted<String>>)> = Vec::new();

    //convert list of graphemes in config into lists of Weighted
    for group in config.graphemes {
        let mut graphemes_converted : Vec<Weighted<String>> = Vec::new();

        graphemes_converted.extend(
            group.graphemes.iter().map(
                |g: &Grapheme| Weighted{weight: g.weight as u32, item: g.string.clone() }
                )
            );
        grapheme_groups.push((group.name.clone(), graphemes_converted));
    }

    let word_factory : WordFactory = WordFactory {
        first_syllable_list : first_syllable_list,
        normal_syllable_list : normal_syllable_list,
        last_syllable_list : last_syllable_list,
        graphemes : grapheme_groups,
        settings : config.settings.clone(),
        rewrites : config.rewrites.clone(),
        rejects : config.rejects.clone(),
    };

    let mut file = File::create(&config.output_settings.output_file).unwrap();

    for _ in 0..config.output_settings.word_count {

        let mut word : Word = Word {
                syllables : String::new(),
                graphemes : String::new(),
                syllable_rewrite_history : Vec::new(),
                grapheme_rewrite_history : Vec::new(),
                syllable_rejects : Vec::new(),
                grapheme_rejects : Vec::new()
            };

        word_factory.generate_syllables(&mut word);

        if (config.settings.rewrites_before_rejects) {
            word_factory.rewrite_syllables(&mut word);
            word_factory.mark_syllable_rejects(&mut word);
        } else {
            word_factory.mark_syllable_rejects(&mut word);
            word_factory.rewrite_syllables(&mut word);
        }

        word_factory.generate_graphemes(&mut word);

        if (config.settings.rewrites_before_rejects) {
            word_factory.rewrite_graphemes(&mut word);
            word_factory.mark_grapheme_rejects(&mut word);
        } else {
            word_factory.mark_grapheme_rejects(&mut word);
            word_factory.rewrite_graphemes(&mut word);
        }

        match file.write(get_word_graphemes(&word).as_bytes()) {
                Err(error) => panic!("error {} writing to file", error),
                Ok(_) => (),
        };
        match file.write("\n".as_bytes()) {
                Err(error) => panic!("error {} writing to file", error),
                Ok(_) => (),
        };
    }
}
<|endoftext|>


<commit_before>
mod graph;
mod dijkstra;
mod utils;
mod stats;
mod graph_state;
mod algorithms;
mod passive_routing_test;
mod importer;
mod exporter;
mod sim;
mod cmd;
mod progress;
mod tests;

extern crate rand;

use std::thread;
use std::sync::Arc;
use std::sync::Mutex;
use std::vec::Vec;
use std::env;

use crate::cmd::cmd_loop;
use crate::cmd::ext_loop;
use crate::graph::Graph;
use crate::tests::*;
use crate::sim::GlobalState;

pub const VERSION : &'static str = "1.0";
pub const CMD_SOCKET_ADDRESS : &'static str = "127.0.0.1:8011";

const HELP_TEXT: &'static str = concat!(
	"--help|-h         Display this help.\n",
	"--version|-v      Display version.\n",
	"--run <file>      Run commands from file.\n",
	"--bind <address>  Bind command socket to address. (Default: 127.0.0.1:8011)\n"
);

struct TZ {
	k: usize,
	graph: Graph,
	//a: Vec<Vec<u32>>
}

/*
impl TZ {
	fn new(k: usize) -> Self {
		Self {
			k: k,
			graph: Graph::new()
			//a: Vec::with_capacity(k)
		}
	}

	// return elements a - b
	fn intersect(a : &Vec<u32>, b: &Vec<u32>) -> Vec<u32> {
		let mut ret = Vec::new();
		for i in 0..a.len() {
			if let Err(_) = b.binary_search(&a[i]) {
				ret.push(a[i]);
			}
		}
		ret
	}

	fn dist(&self, source: u32, target: u32) -> f32 {
		let mut dijkstra = Dijkstra::new();
		dijkstra.find_shortest_distance(&self.graph, source, target)
	}

	fn nearest(&self, a: &Vec<u32>, v: u32) -> u32 {
		let mut v_dist = 10000.0;
		let mut v_min = 0;
		for e in a {
			let d = self.dist(*e, v);
			if d < v_dist {
				v_min = v;
				v_dist = d;
			}
		}
		//if v_dist > 0.0 {
			v_min
		//} else {
		//	0 //hm...
		//}
	}

	fn nearer(&self, a: &Vec<u32>, v: u32, d_max: f32) -> Vec<u32> {
		let mut ret = Vec::new();
		for e in a {
			let d = self.dist(*e, v);
			if d < d_max {
				ret.push(*e);
			}
		}
		ret
	}

	fn init(&mut self) {
		let mut a = Vec::new();
		let k = self.k;

//		let mut graph = Graph::new();
		self.graph.add_lattice4(3, 3);
		let vs : Vec<u32> = (0..self.graph.node_count() as u32).collect();

		use rand::seq::SliceRandom;

		// create A_i
		a.push(vs.clone());
		for i in 1..k {
			let len = a[i-1].len() as f32;
			let n = len.powf(1.0 - 1.0 / k as f32) as usize;
			// TODO: preserve order
			let mut next : Vec<u32> = a[i-1].as_slice().choose_multiple(&mut rand::thread_rng(), n).cloned().collect();
			next.sort_unstable();
			a.push(next);
		}

		for i in 0..a.len() {
			println!("i: {}, n: {}, {:?}", i, a[i].len(), a[i]);
		}

		for v in vs {
			let mut bv : Vec<u32> = Vec::new();
			for i in 0..(k-1) {
				// find nearest in next set
				let n = self.nearest(&a[i+1], v);
				let diff = Self::intersect(&a[i], &a[i+1]);

for i in 0..a[i].len() {
			if let Err(_) = a[i+1].binary_search(&a[i]) {
				if self.dist(v, n)
				diff.push(a[i]);
			}
		}

println!("diff.len(): {}, n: {}", diff.len(), n);
				let mut w = self.nearer(&diff, v, self.dist(v, n));
				bv.append(&mut w);
			}
			break;
		}
	}
}
*/

/*
fn get_p(tz: &TZ, i: usize, v: u32) {
}

fn zwick_dist(u: mut u32, v: mut u32) {
	let mut i = 0;
	let mut w = u;

	while get_b(v) {
		(u, v) = (v, u);
		w = get_p(i, u);
	}

	return delta(w, u) + delta(w, v);
}

fn zwick_prepro() {
	let mut graph = Graph::new();
	graph.add_lattice4(3, 3);
	let k = 3;
	let n = graph.nodes.len();
	let sets = Vector<Vector<u32>>::with_capacity(k);
	sets[0] = (0..n).collect();

	let dijkstra = Dijkstra::new();
	//let min = dijkstra.find_shortest_distance(&graph, source as ID, target as ID);

	// get element from set
	//fn smallest_dist(Vector<u32> &set, u32 v) -> u32

	//let mut dists = Vector<u32>::new();

	let prob = (n as f32).exp(-1.0 / k as f32);

	//let v = (0..n).collect();
	for v in 1..k {
		for i in 0..k {
			let delta = smallest_dist(sets[i], v);
		}
	}

	let bs = Vector<u32>::new();
	for i in 0..(k-1) {
		let v = sets[0][i]; // == V
		let smallest = smallest_dist(sets[i+1], v);
		for w in intersect(sets[i], sets[i+1]) {
			let d = distance(w, v);
			if d < smallest {
				bs.add(w);
			}
		}
	}
}*/

fn main() {
	/*
	let mut tz = TZ::new(4);
	tz.init();
	return;
	*/
	/*
	println!("Node: {}, Link: {}, Packet: {}, Value: {}, Vec<usize>: {}",
		size_of::<node::Node>(), size_of::<link::Link>(), size_of::<sim::Packet>(),
		size_of::<serde_json::Value>(), size_of::<Vec<usize>>()
	);*/

	let mut cmd_address = CMD_SOCKET_ADDRESS.to_string();
	let args: Vec<String> = env::args().skip(1).collect();
	let mut run_script = String::new();

	if let Some((cmd, args)) = args.split_first() {
		match cmd.as_ref() {
			"-h" | "--help" => {
				println!("{}", HELP_TEXT);
				std::process::exit(0);
			},
			"-v" | "--version" => {
				println!("{}", VERSION);
				std::process::exit(0);
			},
			"--bind" => {
				if args.len() == 1 {
					cmd_address = args[0].clone();
				} else if args.len() == 0 {
					println!("Address missing for \"--bind\".");
					std::process::exit(1);
				} else {
					println!("Too many arguments for \"--bind\".");
					std::process::exit(1);
				}
			},
			"--run" => {
				run_script = format!("run {}", args.join(" "));
			},
			// some hard coded scenarios - need to be removed
			"run1" => {
				run_test1();
				std::process::exit(0);
			},
			"run2" => {
				run_test2();
				std::process::exit(0);
			},
			"run3" => {
				run_test3();
				std::process::exit(0);
			},
			"run4" => {
				run_test4();
				std::process::exit(0);
			},
			_ => {
				println!("Unknown argument: {}", cmd);
				std::process::exit(1);
			}
		}
	}

	let sim = Arc::new(Mutex::new(GlobalState::new(&cmd_address)));

	// console
	let cmd_handle = sim.clone();
	let cmd_thread = thread::spawn(move || {
		cmd_loop(cmd_handle, &run_script);
	});

	// unix socket
	let ext_handle = sim.clone();
	let ext_thread = thread::spawn(move || {
		ext_loop(ext_handle, &cmd_address);
	});

	cmd_thread.join().unwrap();
	ext_thread.join().unwrap();

	// exit with error code
	if let Ok(sim) = sim.clone().lock() {
		if sim.abort_simulation {
			std::process::exit(1);
		}
	}
}
<commit_msg>decrease version<commit_after>
mod graph;
mod dijkstra;
mod utils;
mod stats;
mod graph_state;
mod algorithms;
mod passive_routing_test;
mod importer;
mod exporter;
mod sim;
mod cmd;
mod progress;
mod tests;

extern crate rand;

use std::thread;
use std::sync::Arc;
use std::sync::Mutex;
use std::vec::Vec;
use std::env;

use crate::cmd::cmd_loop;
use crate::cmd::ext_loop;
use crate::graph::Graph;
use crate::tests::*;
use crate::sim::GlobalState;

pub const VERSION : &'static str = "0.8";
pub const CMD_SOCKET_ADDRESS : &'static str = "127.0.0.1:8011";

const HELP_TEXT: &'static str = concat!(
	"--help|-h         Display this help.\n",
	"--version|-v      Display version.\n",
	"--run <file>      Run commands from file.\n",
	"--bind <address>  Bind command socket to address. (Default: 127.0.0.1:8011)\n"
);

struct TZ {
	k: usize,
	graph: Graph,
	//a: Vec<Vec<u32>>
}

/*
impl TZ {
	fn new(k: usize) -> Self {
		Self {
			k: k,
			graph: Graph::new()
			//a: Vec::with_capacity(k)
		}
	}

	// return elements a - b
	fn intersect(a : &Vec<u32>, b: &Vec<u32>) -> Vec<u32> {
		let mut ret = Vec::new();
		for i in 0..a.len() {
			if let Err(_) = b.binary_search(&a[i]) {
				ret.push(a[i]);
			}
		}
		ret
	}

	fn dist(&self, source: u32, target: u32) -> f32 {
		let mut dijkstra = Dijkstra::new();
		dijkstra.find_shortest_distance(&self.graph, source, target)
	}

	fn nearest(&self, a: &Vec<u32>, v: u32) -> u32 {
		let mut v_dist = 10000.0;
		let mut v_min = 0;
		for e in a {
			let d = self.dist(*e, v);
			if d < v_dist {
				v_min = v;
				v_dist = d;
			}
		}
		//if v_dist > 0.0 {
			v_min
		//} else {
		//	0 //hm...
		//}
	}

	fn nearer(&self, a: &Vec<u32>, v: u32, d_max: f32) -> Vec<u32> {
		let mut ret = Vec::new();
		for e in a {
			let d = self.dist(*e, v);
			if d < d_max {
				ret.push(*e);
			}
		}
		ret
	}

	fn init(&mut self) {
		let mut a = Vec::new();
		let k = self.k;

//		let mut graph = Graph::new();
		self.graph.add_lattice4(3, 3);
		let vs : Vec<u32> = (0..self.graph.node_count() as u32).collect();

		use rand::seq::SliceRandom;

		// create A_i
		a.push(vs.clone());
		for i in 1..k {
			let len = a[i-1].len() as f32;
			let n = len.powf(1.0 - 1.0 / k as f32) as usize;
			// TODO: preserve order
			let mut next : Vec<u32> = a[i-1].as_slice().choose_multiple(&mut rand::thread_rng(), n).cloned().collect();
			next.sort_unstable();
			a.push(next);
		}

		for i in 0..a.len() {
			println!("i: {}, n: {}, {:?}", i, a[i].len(), a[i]);
		}

		for v in vs {
			let mut bv : Vec<u32> = Vec::new();
			for i in 0..(k-1) {
				// find nearest in next set
				let n = self.nearest(&a[i+1], v);
				let diff = Self::intersect(&a[i], &a[i+1]);

for i in 0..a[i].len() {
			if let Err(_) = a[i+1].binary_search(&a[i]) {
				if self.dist(v, n)
				diff.push(a[i]);
			}
		}

println!("diff.len(): {}, n: {}", diff.len(), n);
				let mut w = self.nearer(&diff, v, self.dist(v, n));
				bv.append(&mut w);
			}
			break;
		}
	}
}
*/

/*
fn get_p(tz: &TZ, i: usize, v: u32) {
}

fn zwick_dist(u: mut u32, v: mut u32) {
	let mut i = 0;
	let mut w = u;

	while get_b(v) {
		(u, v) = (v, u);
		w = get_p(i, u);
	}

	return delta(w, u) + delta(w, v);
}

fn zwick_prepro() {
	let mut graph = Graph::new();
	graph.add_lattice4(3, 3);
	let k = 3;
	let n = graph.nodes.len();
	let sets = Vector<Vector<u32>>::with_capacity(k);
	sets[0] = (0..n).collect();

	let dijkstra = Dijkstra::new();
	//let min = dijkstra.find_shortest_distance(&graph, source as ID, target as ID);

	// get element from set
	//fn smallest_dist(Vector<u32> &set, u32 v) -> u32

	//let mut dists = Vector<u32>::new();

	let prob = (n as f32).exp(-1.0 / k as f32);

	//let v = (0..n).collect();
	for v in 1..k {
		for i in 0..k {
			let delta = smallest_dist(sets[i], v);
		}
	}

	let bs = Vector<u32>::new();
	for i in 0..(k-1) {
		let v = sets[0][i]; // == V
		let smallest = smallest_dist(sets[i+1], v);
		for w in intersect(sets[i], sets[i+1]) {
			let d = distance(w, v);
			if d < smallest {
				bs.add(w);
			}
		}
	}
}*/

fn main() {
	/*
	let mut tz = TZ::new(4);
	tz.init();
	return;
	*/
	/*
	println!("Node: {}, Link: {}, Packet: {}, Value: {}, Vec<usize>: {}",
		size_of::<node::Node>(), size_of::<link::Link>(), size_of::<sim::Packet>(),
		size_of::<serde_json::Value>(), size_of::<Vec<usize>>()
	);*/

	let mut cmd_address = CMD_SOCKET_ADDRESS.to_string();
	let args: Vec<String> = env::args().skip(1).collect();
	let mut run_script = String::new();

	if let Some((cmd, args)) = args.split_first() {
		match cmd.as_ref() {
			"-h" | "--help" => {
				println!("{}", HELP_TEXT);
				std::process::exit(0);
			},
			"-v" | "--version" => {
				println!("{}", VERSION);
				std::process::exit(0);
			},
			"--bind" => {
				if args.len() == 1 {
					cmd_address = args[0].clone();
				} else if args.len() == 0 {
					println!("Address missing for \"--bind\".");
					std::process::exit(1);
				} else {
					println!("Too many arguments for \"--bind\".");
					std::process::exit(1);
				}
			},
			"--run" => {
				run_script = format!("run {}", args.join(" "));
			},
			// some hard coded scenarios - need to be removed
			"run1" => {
				run_test1();
				std::process::exit(0);
			},
			"run2" => {
				run_test2();
				std::process::exit(0);
			},
			"run3" => {
				run_test3();
				std::process::exit(0);
			},
			"run4" => {
				run_test4();
				std::process::exit(0);
			},
			_ => {
				println!("Unknown argument: {}", cmd);
				std::process::exit(1);
			}
		}
	}

	let sim = Arc::new(Mutex::new(GlobalState::new(&cmd_address)));

	// console
	let cmd_handle = sim.clone();
	let cmd_thread = thread::spawn(move || {
		cmd_loop(cmd_handle, &run_script);
	});

	// unix socket
	let ext_handle = sim.clone();
	let ext_thread = thread::spawn(move || {
		ext_loop(ext_handle, &cmd_address);
	});

	cmd_thread.join().unwrap();
	ext_thread.join().unwrap();

	// exit with error code
	if let Ok(sim) = sim.clone().lock() {
		if sim.abort_simulation {
			std::process::exit(1);
		}
	}
}
<|endoftext|>