<commit_before>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use layout::{PointerBuilder, PointerReader};
use common::Word;

// TODO maybe we can simplify these traits. It seems that the only
// difference in the implementations is the FieldSize.

pub trait FromPointerReader<'a> {
    fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Self;
}

pub trait FromPointerBuilder<'a> {
    fn init_pointer(PointerBuilder<'a>, uint) -> Self;
    fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Self;
}

pub mod PrimitiveList {
    use super::{FromPointerReader, FromPointerBuilder};
    use layout::{ListReader, ListBuilder, PointerReader, PointerBuilder,
                 PrimitiveElement, element_size_for_type};
    use common::Word;

    pub struct Reader<'a, T> {
        // I want this field to be private, but then I can't access it in set_list()
        reader : ListReader<'a>
    }

    impl <'a, T : PrimitiveElement> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a, T : PrimitiveElement> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(element_size_for_type::<T>(), default_value) }
        }
    }

    impl <'a, T : PrimitiveElement> Index<uint, T> for Reader<'a, T> {
        fn index(&self, index : &uint) -> T {
            PrimitiveElement::get(&self.reader, *index)
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : PrimitiveElement> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : T) {
            PrimitiveElement::set(&self.builder, index, value);
        }
    }

    impl <'a, T : PrimitiveElement> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder { builder : builder.init_list(element_size_for_type::<T>(), size) }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder { builder : builder.get_list(element_size_for_type::<T>(), default_value) }
        }
    }

    impl <'a, T : PrimitiveElement> Index<uint, T> for Builder<'a, T> {
        fn index(&self, index : &uint) -> T {
            PrimitiveElement::get_from_builder(&self.builder, *index)
        }
    }
}

pub trait ToU16 {
    fn to_u16(self) -> u16;
}

pub mod EnumList {
    use layout::*;
    use list::*;
    use common::Word;

    pub struct Reader<'a, T> {
        reader : ListReader<'a>
    }

    impl <'a, T : FromPrimitive> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }

    }

    impl <'a, T : FromPrimitive> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(TwoBytes, default_value) }
        }
    }

    impl <'a, T : FromPrimitive> Index<uint, Option<T>> for Reader<'a, T> {
        fn index(&self, index : &uint) -> Option<T> {
            let result : u16 = PrimitiveElement::get(&self.reader, *index);
            FromPrimitive::from_u16(result)
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : ToU16 + FromPrimitive> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : T) {
            PrimitiveElement::set(&self.builder, index, value.to_u16());
        }
    }

    impl <'a, T : FromPrimitive> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder { builder : builder.init_list(TwoBytes, size) }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder { builder : builder.get_list(TwoBytes, default_value) }
        }
    }

    impl <'a, T : ToU16 + FromPrimitive> Index<uint, Option<T>> for Builder<'a, T> {
        fn index(&self, index : &uint) -> Option<T> {
            let result : u16 = PrimitiveElement::get_from_builder(&self.builder, *index);
            FromPrimitive::from_u16(result)
        }
    }
}

pub mod StructList {
    use super::{FromPointerReader, FromPointerBuilder};
    use common::Word;
    use layout::*;

    pub struct Reader<'a, T> {
        reader : ListReader<'a>
    }

    impl <'a, T : FromStructReader<'a>> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a, T : FromStructReader<'a>> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(InlineComposite, default_value) }
        }
    }

    impl <'a, T : FromStructReader<'a>> Index<uint, T> for Reader<'a, T> {
        fn index(&self, index : &uint) -> T {
            let result : T = FromStructReader::new(self.reader.get_struct_element(*index));
            result
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : FromStructBuilder<'a>> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

//        pub fn set(&self, index : uint, value : T) {
//        }
    }

    impl <'a, T : FromStructBuilder<'a> + HasStructSize> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder {
                builder : builder.init_struct_list(size, HasStructSize::struct_size(None::<T>))
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder {
                builder : builder.get_struct_list(HasStructSize::struct_size(None::<T>), default_value)
            }
        }
    }

    impl <'a, T : FromStructBuilder<'a>> Index<uint, T> for Builder<'a, T> {
        fn index(&self, index : &uint) -> T {
            let result : T =
                FromStructBuilder::new(self.builder.get_struct_element(*index));
            result
        }
    }
}

pub mod ListList {
    use super::{FromPointerReader, FromPointerBuilder};
    use std;
    use common::Word;
    use layout::*;

    pub struct Reader<'a, T> {
        reader : ListReader<'a>
    }

    impl <'a, T> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a, T : FromPointerReader<'a>> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(Pointer, default_value) }
        }
    }

    impl <'a, T : FromPointerReader<'a>> Index<uint, T> for Reader<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index <  self.size());
            FromPointerReader::get_from_pointer(
                &self.reader.get_pointer_element(*index), std::ptr::null())
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : FromPointerBuilder<'a>> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn init(&self, index : uint, size : uint) -> T {
            let result : T =
                FromPointerBuilder::init_pointer(self.builder.get_pointer_element(index), size);
            result
        }
    }

    impl <'a, T : FromPointerBuilder<'a>> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder {
                builder : builder.init_list(Pointer, size)
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder {
                builder : builder.get_list(Pointer, default_value)
            }
        }
    }

    impl <'a, T : FromPointerBuilder<'a>> Index<uint, T> for Builder<'a, T> {
        fn index(&self, index : &uint) -> T {
            let result : T =
                FromPointerBuilder::get_from_pointer(
                self.builder.get_pointer_element(*index),
                std::ptr::null());
            result
        }
    }

}

pub mod TextList {
    use super::{FromPointerReader, FromPointerBuilder};
    use std;
    use common::Word;
    use blob::Text;
    use layout::*;

    pub struct Reader<'a> {
        reader : ListReader<'a>
    }

    impl <'a> Reader<'a> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b> {
            Reader::<'b> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a> FromPointerReader<'a> for Reader<'a> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a> {
            Reader { reader : reader.get_list(Pointer, default_value) }
        }
    }

    impl <'a> Index<uint, Text::Reader<'a>> for Reader<'a> {
        fn index(&self, index : &uint) -> Text::Reader<'a> {
            assert!(*index <  self.size());
            self.reader.get_pointer_element(*index).get_text(std::ptr::null(), 0)
        }
    }

    pub struct Builder<'a> {
        builder : ListBuilder<'a>
    }

    impl <'a> Builder<'a> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : Text::Reader) {
            assert!(index < self.size());
            self.builder.get_pointer_element(index).set_text(value);
        }
    }

    impl <'a> FromPointerBuilder<'a> for Builder<'a> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a> {
            Builder {
                builder : builder.init_list(Pointer, size)
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a> {
            Builder {
                builder : builder.get_list(Pointer, default_value)
            }
        }
    }

    impl <'a> Index<uint, Text::Builder<'a>> for Builder<'a> {
        fn index(&self, index : &uint) -> Text::Builder<'a> {
            self.builder.get_pointer_element(*index).get_text(std::ptr::null(), 0)
        }
    }

}

pub mod DataList {
    use super::{FromPointerReader, FromPointerBuilder};
    use std;
    use common::Word;
    use blob::Data;
    use layout::*;

    pub struct Reader<'a> {
        reader : ListReader<'a>
    }

    impl <'a> Reader<'a> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b> {
            Reader::<'b> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a> FromPointerReader<'a> for Reader<'a> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a> {
            Reader { reader : reader.get_list(Pointer, default_value) }
        }
    }

    impl <'a> Index<uint, Data::Reader<'a>> for Reader<'a> {
        fn index(&self, index : &uint) -> Data::Reader<'a> {
            assert!(*index <  self.size());
            self.reader.get_pointer_element(*index).get_data(std::ptr::null(), 0)
        }
    }

    pub struct Builder<'a> {
        builder : ListBuilder<'a>
    }

    impl <'a> Builder<'a> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : Data::Reader) {
            assert!(index < self.size());
            self.builder.get_pointer_element(index).set_data(value);
        }
    }

    impl <'a> FromPointerBuilder<'a> for Builder<'a> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a> {
            Builder {
                builder : builder.init_list(Pointer, size)
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a> {
            Builder {
                builder : builder.get_list(Pointer, default_value)
            }
        }
    }

    impl <'a> Index<uint, Data::Builder<'a>> for Builder<'a> {
        fn index(&self, index : &uint) -> Data::Builder<'a> {
            self.builder.get_pointer_element(*index).get_data(std::ptr::null(), 0)
        }
    }

}
<commit_msg>bounds check list accessors<commit_after>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use layout::{PointerBuilder, PointerReader};
use common::Word;

// TODO maybe we can simplify these traits. It seems that the only
// difference in the implementations is the FieldSize.

pub trait FromPointerReader<'a> {
    fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Self;
}

pub trait FromPointerBuilder<'a> {
    fn init_pointer(PointerBuilder<'a>, uint) -> Self;
    fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Self;
}

pub mod PrimitiveList {
    use super::{FromPointerReader, FromPointerBuilder};
    use layout::{ListReader, ListBuilder, PointerReader, PointerBuilder,
                 PrimitiveElement, element_size_for_type};
    use common::Word;

    pub struct Reader<'a, T> {
        // I want this field to be private, but then I can't access it in set_list()
        reader : ListReader<'a>
    }

    impl <'a, T : PrimitiveElement> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a, T : PrimitiveElement> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(element_size_for_type::<T>(), default_value) }
        }
    }

    impl <'a, T : PrimitiveElement> Index<uint, T> for Reader<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index < self.size());
            PrimitiveElement::get(&self.reader, *index)
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : PrimitiveElement> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : T) {
            PrimitiveElement::set(&self.builder, index, value);
        }
    }

    impl <'a, T : PrimitiveElement> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder { builder : builder.init_list(element_size_for_type::<T>(), size) }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder { builder : builder.get_list(element_size_for_type::<T>(), default_value) }
        }
    }

    impl <'a, T : PrimitiveElement> Index<uint, T> for Builder<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index < self.size());
            PrimitiveElement::get_from_builder(&self.builder, *index)
        }
    }
}

pub trait ToU16 {
    fn to_u16(self) -> u16;
}

pub mod EnumList {
    use layout::*;
    use list::*;
    use common::Word;

    pub struct Reader<'a, T> {
        reader : ListReader<'a>
    }

    impl <'a, T : FromPrimitive> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }

    }

    impl <'a, T : FromPrimitive> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(TwoBytes, default_value) }
        }
    }

    impl <'a, T : FromPrimitive> Index<uint, Option<T>> for Reader<'a, T> {
        fn index(&self, index : &uint) -> Option<T> {
            assert!(*index < self.size());
            let result : u16 = PrimitiveElement::get(&self.reader, *index);
            FromPrimitive::from_u16(result)
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : ToU16 + FromPrimitive> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : T) {
            assert!(index < self.size());
            PrimitiveElement::set(&self.builder, index, value.to_u16());
        }
    }

    impl <'a, T : FromPrimitive> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder { builder : builder.init_list(TwoBytes, size) }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder { builder : builder.get_list(TwoBytes, default_value) }
        }
    }

    impl <'a, T : ToU16 + FromPrimitive> Index<uint, Option<T>> for Builder<'a, T> {
        fn index(&self, index : &uint) -> Option<T> {
            assert!(*index < self.size());
            let result : u16 = PrimitiveElement::get_from_builder(&self.builder, *index);
            FromPrimitive::from_u16(result)
        }
    }
}

pub mod StructList {
    use super::{FromPointerReader, FromPointerBuilder};
    use common::Word;
    use layout::*;

    pub struct Reader<'a, T> {
        reader : ListReader<'a>
    }

    impl <'a, T : FromStructReader<'a>> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a, T : FromStructReader<'a>> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(InlineComposite, default_value) }
        }
    }

    impl <'a, T : FromStructReader<'a>> Index<uint, T> for Reader<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index < self.size());
            let result : T = FromStructReader::new(self.reader.get_struct_element(*index));
            result
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : FromStructBuilder<'a>> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

//        pub fn set(&self, index : uint, value : T) {
//        }
    }

    impl <'a, T : FromStructBuilder<'a> + HasStructSize> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder {
                builder : builder.init_struct_list(size, HasStructSize::struct_size(None::<T>))
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder {
                builder : builder.get_struct_list(HasStructSize::struct_size(None::<T>), default_value)
            }
        }
    }

    impl <'a, T : FromStructBuilder<'a>> Index<uint, T> for Builder<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index < self.size());
            let result : T =
                FromStructBuilder::new(self.builder.get_struct_element(*index));
            result
        }
    }
}

pub mod ListList {
    use super::{FromPointerReader, FromPointerBuilder};
    use std;
    use common::Word;
    use layout::*;

    pub struct Reader<'a, T> {
        reader : ListReader<'a>
    }

    impl <'a, T> Reader<'a, T> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b, T> {
            Reader::<'b, T> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a, T : FromPointerReader<'a>> FromPointerReader<'a> for Reader<'a, T> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a, T> {
            Reader { reader : reader.get_list(Pointer, default_value) }
        }
    }

    impl <'a, T : FromPointerReader<'a>> Index<uint, T> for Reader<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index <  self.size());
            FromPointerReader::get_from_pointer(
                &self.reader.get_pointer_element(*index), std::ptr::null())
        }
    }

    pub struct Builder<'a, T> {
        builder : ListBuilder<'a>
    }

    impl <'a, T : FromPointerBuilder<'a>> Builder<'a, T> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a, T> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn init(&self, index : uint, size : uint) -> T {
            let result : T =
                FromPointerBuilder::init_pointer(self.builder.get_pointer_element(index), size);
            result
        }
    }

    impl <'a, T : FromPointerBuilder<'a>> FromPointerBuilder<'a> for Builder<'a, T> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a, T> {
            Builder {
                builder : builder.init_list(Pointer, size)
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a, T> {
            Builder {
                builder : builder.get_list(Pointer, default_value)
            }
        }
    }

    impl <'a, T : FromPointerBuilder<'a>> Index<uint, T> for Builder<'a, T> {
        fn index(&self, index : &uint) -> T {
            assert!(*index < self.size());
            let result : T =
                FromPointerBuilder::get_from_pointer(
                self.builder.get_pointer_element(*index),
                std::ptr::null());
            result
        }
    }

}

pub mod TextList {
    use super::{FromPointerReader, FromPointerBuilder};
    use std;
    use common::Word;
    use blob::Text;
    use layout::*;

    pub struct Reader<'a> {
        reader : ListReader<'a>
    }

    impl <'a> Reader<'a> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b> {
            Reader::<'b> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a> FromPointerReader<'a> for Reader<'a> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a> {
            Reader { reader : reader.get_list(Pointer, default_value) }
        }
    }

    impl <'a> Index<uint, Text::Reader<'a>> for Reader<'a> {
        fn index(&self, index : &uint) -> Text::Reader<'a> {
            assert!(*index <  self.size());
            self.reader.get_pointer_element(*index).get_text(std::ptr::null(), 0)
        }
    }

    pub struct Builder<'a> {
        builder : ListBuilder<'a>
    }

    impl <'a> Builder<'a> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : Text::Reader) {
            assert!(index < self.size());
            self.builder.get_pointer_element(index).set_text(value);
        }
    }

    impl <'a> FromPointerBuilder<'a> for Builder<'a> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a> {
            Builder {
                builder : builder.init_list(Pointer, size)
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a> {
            Builder {
                builder : builder.get_list(Pointer, default_value)
            }
        }
    }

    impl <'a> Index<uint, Text::Builder<'a>> for Builder<'a> {
        fn index(&self, index : &uint) -> Text::Builder<'a> {
            self.builder.get_pointer_element(*index).get_text(std::ptr::null(), 0)
        }
    }

}

pub mod DataList {
    use super::{FromPointerReader, FromPointerBuilder};
    use std;
    use common::Word;
    use blob::Data;
    use layout::*;

    pub struct Reader<'a> {
        reader : ListReader<'a>
    }

    impl <'a> Reader<'a> {
        pub fn new<'b>(reader : ListReader<'b>) -> Reader<'b> {
            Reader::<'b> { reader : reader }
        }

        pub fn size(&self) -> uint { self.reader.size() }
    }

    impl <'a> FromPointerReader<'a> for Reader<'a> {
        fn get_from_pointer(reader : &PointerReader<'a>, default_value : *Word) -> Reader<'a> {
            Reader { reader : reader.get_list(Pointer, default_value) }
        }
    }

    impl <'a> Index<uint, Data::Reader<'a>> for Reader<'a> {
        fn index(&self, index : &uint) -> Data::Reader<'a> {
            assert!(*index <  self.size());
            self.reader.get_pointer_element(*index).get_data(std::ptr::null(), 0)
        }
    }

    pub struct Builder<'a> {
        builder : ListBuilder<'a>
    }

    impl <'a> Builder<'a> {
        pub fn new(builder : ListBuilder<'a>) -> Builder<'a> {
            Builder { builder : builder }
        }

        pub fn size(&self) -> uint { self.builder.size() }

        pub fn set(&self, index : uint, value : Data::Reader) {
            assert!(index < self.size());
            self.builder.get_pointer_element(index).set_data(value);
        }
    }

    impl <'a> FromPointerBuilder<'a> for Builder<'a> {
        fn init_pointer(builder : PointerBuilder<'a>, size : uint) -> Builder<'a> {
            Builder {
                builder : builder.init_list(Pointer, size)
            }
        }
        fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *Word) -> Builder<'a> {
            Builder {
                builder : builder.get_list(Pointer, default_value)
            }
        }
    }

    impl <'a> Index<uint, Data::Builder<'a>> for Builder<'a> {
        fn index(&self, index : &uint) -> Data::Builder<'a> {
            assert!(*index < self.size());
            self.builder.get_pointer_element(*index).get_data(std::ptr::null(), 0)
        }
    }

}
<|endoftext|>


<commit_before>use std::{
    collections::VecDeque,
    fmt,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc,
        Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        loadbalancing::server::LoadBalancer,
        socks5::Address,
        tcprelay::client::ServerClient as TcpServerClient,
        udprelay::client::ServerClient as UdpServerClient,
    },
};

use log::{debug, info};
use tokio::{
    self,
    io::{AsyncReadExt, AsyncWriteExt},
    sync::Barrier,
    time,
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

#[derive(Debug, Copy, Clone)]
enum Score {
    Latency(u64),
    Errored,
}

struct ServerLatencyInner {
    latency_queue: VecDeque<Score>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: Score) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            // Never checked, assume it is the worst of all
            return 2 * 1000;
        }

        // 1. Mid Latency
        // 2. Proportion of Errors
        let mut vec_lat = Vec::with_capacity(self.latency_queue.len());
        let mut acc_err = 0;
        for lat in &self.latency_queue {
            match lat {
                Score::Latency(l) => vec_lat.push(l),
                Score::Errored => acc_err += 1,
            }
        }

        let max_lat = DEFAULT_CHECK_TIMEOUT_SEC * 1000;

        // Find the mid of latencies
        let mid_lat = if vec_lat.is_empty() {
            // The whole array are errors
            max_lat
        } else {
            vec_lat.sort();
            let mid = vec_lat.len() / 2;
            if vec_lat.len() % 2 == 0 {
                (vec_lat[mid] + vec_lat[mid - 1]) / 2
            } else {
                *vec_lat[mid]
            }
        };

        // Score = norm_lat + prop_err
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        let norm_lat = mid_lat as f64 / max_lat as f64;
        let prop_err = acc_err as f64 / self.latency_queue.len() as f64;

        ((norm_lat + prop_err) * 1000.0) as u64
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: Score) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 2; // Latency shouldn't greater than 2 secs, that's too long

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    async fn new(context: SharedContext, server_type: ServerType) -> Inner {
        let config = context.config();
        assert!(!config.server.is_empty());

        // Load balancer is only required in multi-server configuration
        let check_required = config.server.len() > 1;

        // Wait for all ping tasks to be started
        let barrier = Arc::new(Barrier::new(config.server.len() + 1));

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();
            let barrier = barrier.clone();

            tokio::spawn(
                // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                async move {
                    let start_time = Instant::now();

                    debug!(
                        "{:?} server {} latency ping task initializing",
                        server_type,
                        sc.config.addr()
                    );

                    // Quickly collect some latency data
                    //
                    // Maximum wait duration: DEFAULT_CHECK_TIMEOUT_SEC
                    while Instant::now() - start_time < Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC) {
                        let mut succ = false;

                        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
                            Ok(d) => {
                                succ = true;
                                latency.push(Score::Latency(d))
                            }
                            Err(..) => {
                                latency.push(Score::Errored) // Penalty
                            }
                        };
                        debug!(
                            "updated remote {:?} server {} (score: {})",
                            server_type,
                            sc.config.addr(),
                            score
                        );
                        sc.score.store(score, Ordering::Release);

                        if succ {
                            break;
                        }
                    }

                    // Wait until all the other tasks are finished initializing
                    barrier.wait().await;
                    drop(barrier);

                    debug!(
                        "{:?} server {} latency ping task started",
                        server_type,
                        sc.config.addr()
                    );

                    while context.server_running() {
                        // First round may be failed, plugins are started asynchronously

                        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
                            Ok(d) => latency.push(Score::Latency(d)),
                            Err(..) => latency.push(Score::Errored), // Penalty
                        };
                        debug!(
                            "updated remote {:?} server {} (score: {})",
                            server_type,
                            sc.config.addr(),
                            score
                        );
                        sc.score.store(score, Ordering::Release);

                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }

                    debug!(
                        "{:?} server {} latency ping task stopped",
                        server_type,
                        sc.config.addr()
                    );
                },
            );
        }

        if check_required {
            barrier.wait().await;
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    async fn check_request_tcp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static GET_BODY: &[u8] =
            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);

        let TcpServerClient { mut stream } = TcpServerClient::connect(context, &addr, sc).await?;
        stream.write_all(GET_BODY).await?;
        stream.flush().await?;
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await?;

        Ok(())
    }

    async fn check_request_udp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static DNS_QUERY: &[u8] = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01";

        let addr = Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53));

        let mut client = UdpServerClient::new(sc).await?;
        client.send_to(context, &addr, DNS_QUERY).await?;
        let _ = client.recv_from(context).await?;

        Ok(())
    }

    async fn check_request(sc: Arc<ServerConfig>, context: &Context, server_type: ServerType) -> io::Result<()> {
        match server_type {
            ServerType::Tcp => Inner::check_request_tcp(sc, context).await,
            ServerType::Udp => Inner::check_request_udp(sc, context).await,
        }
    }

    async fn check_delay(sc: &Server, context: &Context, server_type: ServerType) -> io::Result<u64> {
        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
        let res = time::timeout(timeout, Inner::check_request(sc.config.clone(), context, server_type)).await;

        let elapsed = Instant::now() - start;
        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
        match res {
            Ok(Ok(..)) => {
                // Got the result ... record its time
                debug!(
                    "checked remote {:?} server {} latency with {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );
                Ok(elapsed)
            }
            Ok(Err(err)) => {
                debug!(
                    "failed to check {:?} server {}, error: {}",
                    server_type,
                    sc.config.addr(),
                    err
                );

                // NOTE: connection / handshake error, server is down
                Err(err)
            }
            Err(..) => {
                // Timeout
                debug!(
                    "checked remote {:?} server {} latency timeout, elapsed {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );

                // NOTE: timeout is still available, but server is too slow
                Ok(elapsed)
            }
        }
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ServerType {
    Tcp,
    Udp,
}

/// Load balancer based on pinging latencies of all servers
#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    /// Create a PingBalancer
    pub async fn new(context: SharedContext, server_type: ServerType) -> PingBalancer {
        // Wait until all tasks are started
        let inner = Arc::new(Inner::new(context.clone(), server_type).await);

        if inner.servers.len() > 1 {
            let barrier = Arc::new(Barrier::new(2));

            let cloned_inner = inner.clone();
            let cloned_barrier = barrier.clone();
            tokio::spawn(async move {
                let inner = cloned_inner;
                assert!(!inner.servers.is_empty());

                let mut opt_barrier = Some(cloned_barrier);

                while context.server_running() {
                    // Choose the best one
                    let mut svr_idx = 0;

                    for (idx, svr) in inner.servers.iter().enumerate() {
                        let choosen_svr = &inner.servers[svr_idx];
                        if svr.score() < choosen_svr.score() {
                            svr_idx = idx;
                        }
                    }

                    let choosen_svr = &inner.servers[svr_idx];
                    if choosen_svr.score() != inner.servers[inner.best_idx()].score() {
                        if opt_barrier.is_none() {
                            info!(
                                "switched {:?} server from {} (score: {}) to {} (score: {})",
                                server_type,
                                inner.best_server().config.addr(),
                                inner.best_server().score(),
                                choosen_svr.config.addr(),
                                choosen_svr.score(),
                            );
                        }

                        inner.set_best_idx(svr_idx);
                    }

                    if let Some(barrier) = opt_barrier.take() {
                        barrier.wait().await;
                        debug!("ping {:?} server choosing task started", server_type);
                    }

                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                }
            });

            barrier.wait().await;
        }

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<commit_msg>Just check once before actual start<commit_after>use std::{
    collections::VecDeque,
    fmt,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc,
        Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        loadbalancing::server::LoadBalancer,
        socks5::Address,
        tcprelay::client::ServerClient as TcpServerClient,
        udprelay::client::ServerClient as UdpServerClient,
    },
};

use log::{debug, info};
use tokio::{
    self,
    io::{AsyncReadExt, AsyncWriteExt},
    sync::Barrier,
    time,
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

#[derive(Debug, Copy, Clone)]
enum Score {
    Latency(u64),
    Errored,
}

struct ServerLatencyInner {
    latency_queue: VecDeque<Score>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: Score) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            // Never checked, assume it is the worst of all
            return 2 * 1000;
        }

        // 1. Mid Latency
        // 2. Proportion of Errors
        let mut vec_lat = Vec::with_capacity(self.latency_queue.len());
        let mut acc_err = 0;
        for lat in &self.latency_queue {
            match lat {
                Score::Latency(l) => vec_lat.push(l),
                Score::Errored => acc_err += 1,
            }
        }

        let max_lat = DEFAULT_CHECK_TIMEOUT_SEC * 1000;

        // Find the mid of latencies
        let mid_lat = if vec_lat.is_empty() {
            // The whole array are errors
            max_lat
        } else {
            vec_lat.sort();
            let mid = vec_lat.len() / 2;
            if vec_lat.len() % 2 == 0 {
                (vec_lat[mid] + vec_lat[mid - 1]) / 2
            } else {
                *vec_lat[mid]
            }
        };

        // Score = norm_lat + prop_err
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        let norm_lat = mid_lat as f64 / max_lat as f64;
        let prop_err = acc_err as f64 / self.latency_queue.len() as f64;

        ((norm_lat + prop_err) * 1000.0) as u64
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: Score) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 2; // Latency shouldn't greater than 2 secs, that's too long

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    async fn new(context: SharedContext, server_type: ServerType) -> Inner {
        let config = context.config();
        assert!(!config.server.is_empty());

        // Load balancer is only required in multi-server configuration
        let check_required = config.server.len() > 1;

        // Wait for all ping tasks to be started
        let barrier = Arc::new(Barrier::new(config.server.len() + 1));

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();
            let barrier = barrier.clone();

            // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
            tokio::spawn(async move {
                debug!(
                    "{:?} server {} latency ping task initializing",
                    server_type,
                    sc.config.addr()
                );

                // Quickly collect some latency data
                //
                // Maximum wait duration: DEFAULT_CHECK_TIMEOUT_SEC
                Inner::check_update_score(&latency, &*sc, &*context, server_type).await;

                // Wait until all the other tasks are finished initializing
                barrier.wait().await;
                drop(barrier);

                debug!(
                    "{:?} server {} latency ping task started",
                    server_type,
                    sc.config.addr()
                );

                while context.server_running() {
                    // First round may be failed, plugins are started asynchronously
                    Inner::check_update_score(&latency, &*sc, &*context, server_type).await;

                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                }

                debug!(
                    "{:?} server {} latency ping task stopped",
                    server_type,
                    sc.config.addr()
                );
            });
        }

        if check_required {
            barrier.wait().await;
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    async fn check_update_score(latency: &ServerLatency, sc: &Server, context: &Context, server_type: ServerType) {
        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
            Ok(d) => latency.push(Score::Latency(d)),
            Err(..) => latency.push(Score::Errored), // Penalty
        };
        debug!(
            "updated remote {:?} server {} (score: {})",
            server_type,
            sc.config.addr(),
            score
        );
        sc.score.store(score, Ordering::Release);
    }

    async fn check_request_tcp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static GET_BODY: &[u8] =
            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);

        let TcpServerClient { mut stream } = TcpServerClient::connect(context, &addr, sc).await?;
        stream.write_all(GET_BODY).await?;
        stream.flush().await?;
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await?;

        Ok(())
    }

    async fn check_request_udp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static DNS_QUERY: &[u8] = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01";

        let addr = Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53));

        let mut client = UdpServerClient::new(sc).await?;
        client.send_to(context, &addr, DNS_QUERY).await?;
        let _ = client.recv_from(context).await?;

        Ok(())
    }

    async fn check_request(sc: Arc<ServerConfig>, context: &Context, server_type: ServerType) -> io::Result<()> {
        match server_type {
            ServerType::Tcp => Inner::check_request_tcp(sc, context).await,
            ServerType::Udp => Inner::check_request_udp(sc, context).await,
        }
    }

    async fn check_delay(sc: &Server, context: &Context, server_type: ServerType) -> io::Result<u64> {
        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
        let res = time::timeout(timeout, Inner::check_request(sc.config.clone(), context, server_type)).await;

        let elapsed = Instant::now() - start;
        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
        match res {
            Ok(Ok(..)) => {
                // Got the result ... record its time
                debug!(
                    "checked remote {:?} server {} latency with {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );
                Ok(elapsed)
            }
            Ok(Err(err)) => {
                debug!(
                    "failed to check {:?} server {}, error: {}",
                    server_type,
                    sc.config.addr(),
                    err
                );

                // NOTE: connection / handshake error, server is down
                Err(err)
            }
            Err(..) => {
                // Timeout
                debug!(
                    "checked remote {:?} server {} latency timeout, elapsed {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );

                // NOTE: timeout is still available, but server is too slow
                Ok(elapsed)
            }
        }
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ServerType {
    Tcp,
    Udp,
}

/// Load balancer based on pinging latencies of all servers
#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    /// Create a PingBalancer
    pub async fn new(context: SharedContext, server_type: ServerType) -> PingBalancer {
        // Wait until all tasks are started
        let inner = Arc::new(Inner::new(context.clone(), server_type).await);

        if inner.servers.len() > 1 {
            let barrier = Arc::new(Barrier::new(2));

            let cloned_inner = inner.clone();
            let cloned_barrier = barrier.clone();
            tokio::spawn(async move {
                let inner = cloned_inner;
                assert!(!inner.servers.is_empty());

                let mut opt_barrier = Some(cloned_barrier);

                while context.server_running() {
                    // Choose the best one
                    let mut svr_idx = 0;

                    for (idx, svr) in inner.servers.iter().enumerate() {
                        let choosen_svr = &inner.servers[svr_idx];
                        if svr.score() < choosen_svr.score() {
                            svr_idx = idx;
                        }
                    }

                    let choosen_svr = &inner.servers[svr_idx];
                    if choosen_svr.score() != inner.servers[inner.best_idx()].score() {
                        if opt_barrier.is_none() {
                            info!(
                                "switched {:?} server from {} (score: {}) to {} (score: {})",
                                server_type,
                                inner.best_server().config.addr(),
                                inner.best_server().score(),
                                choosen_svr.config.addr(),
                                choosen_svr.score(),
                            );
                        }

                        inner.set_best_idx(svr_idx);
                    }

                    if let Some(barrier) = opt_barrier.take() {
                        barrier.wait().await;
                        debug!("ping {:?} server choosing task started", server_type);
                    }

                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                }
            });

            barrier.wait().await;
        }

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<|endoftext|>


<commit_before>//! Renderer for function calls.

use hir::{AsAssocItem, HasSource, HirDisplay};
use ide_db::SymbolKind;
use itertools::Itertools;
use syntax::ast::Fn;

use crate::{
    item::{CompletionItem, CompletionItemKind, CompletionKind, CompletionRelevance, ImportEdit},
    render::{
        builder_ext::Params, compute_exact_name_match, compute_ref_match, compute_type_match,
        RenderContext,
    },
};

pub(crate) fn render_fn(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_fn");
    Some(FunctionRender::new(ctx, None, local_name, fn_, false)?.render(import_to_add))
}

pub(crate) fn render_method(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    receiver: Option<hir::Name>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_method");
    Some(FunctionRender::new(ctx, receiver, local_name, fn_, true)?.render(import_to_add))
}

#[derive(Debug)]
struct FunctionRender<'a> {
    ctx: RenderContext<'a>,
    name: String,
    receiver: Option<hir::Name>,
    func: hir::Function,
    ast_node: Fn,
    is_method: bool,
}

impl<'a> FunctionRender<'a> {
    fn new(
        ctx: RenderContext<'a>,
        receiver: Option<hir::Name>,
        local_name: Option<hir::Name>,
        fn_: hir::Function,
        is_method: bool,
    ) -> Option<FunctionRender<'a>> {
        let name = local_name.unwrap_or_else(|| fn_.name(ctx.db())).to_string();
        let ast_node = fn_.source(ctx.db())?.value;

        Some(FunctionRender { ctx, name, receiver, func: fn_, ast_node, is_method })
    }

    fn render(self, import_to_add: Option<ImportEdit>) -> CompletionItem {
        let params = self.params();
        let call = if let Some(receiver) = &self.receiver {
            format!("{}.{}", receiver, &self.name)
        } else {
            self.name.clone()
        };
        let mut item =
            CompletionItem::new(CompletionKind::Reference, self.ctx.source_range(), call.clone());
        item.kind(self.kind())
            .set_documentation(self.ctx.docs(self.func))
            .set_deprecated(
                self.ctx.is_deprecated(self.func) || self.ctx.is_deprecated_assoc_item(self.func),
            )
            .detail(self.detail())
            .add_call_parens(self.ctx.completion, call.clone(), params);

        if import_to_add.is_none() {
            let db = self.ctx.db();
            if let Some(actm) = self.func.as_assoc_item(db) {
                if let Some(trt) = actm.containing_trait_or_trait_impl(db) {
                    item.trait_name(trt.name(db).to_string());
                }
            }
        }

        item.add_import(import_to_add).lookup_by(self.name);

        let ret_type = self.func.ret_type(self.ctx.db());
        item.set_relevance(CompletionRelevance {
            type_match: compute_type_match(self.ctx.completion, &ret_type),
            exact_name_match: compute_exact_name_match(self.ctx.completion, &call),
            ..CompletionRelevance::default()
        });

        if let Some(ref_match) = compute_ref_match(self.ctx.completion, &ret_type) {
            // FIXME
            // For now we don't properly calculate the edits for ref match
            // completions on methods, so we've disabled them. See #8058.
            if !self.is_method {
                item.ref_match(ref_match);
            }
        }

        item.build()
    }

    fn detail(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());
        let ret = if ret_ty.is_unit() {
            // Omit the return type if it is the unit type
            String::new()
        } else {
            format!(" {}", self.ty_display())
        };

        format!("fn({}){}", self.params_display(), ret)
    }

    fn params_display(&self) -> String {
        if let Some(self_param) = self.func.self_param(self.ctx.db()) {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .skip(1) // skip the self param because we are manually handling that
                .map(|p| p.ty().display(self.ctx.db()).to_string());

            std::iter::once(self_param.display(self.ctx.db()).to_owned()).chain(params).join(", ")
        } else {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .map(|p| p.ty().display(self.ctx.db()).to_string())
                .join(", ");
            params
        }
    }

    fn ty_display(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());

        format!("-> {}", ret_ty.display(self.ctx.db()))
    }

    fn add_arg(&self, arg: &str, ty: &hir::Type) -> String {
        if let Some(derefed_ty) = ty.remove_ref() {
            for (name, local) in self.ctx.completion.locals.iter() {
                if name == arg && local.ty(self.ctx.db()) == derefed_ty {
                    let mutability = if ty.is_mutable_reference() { "&mut " } else { "&" };
                    return format!("{}{}", mutability, arg);
                }
            }
        }
        arg.to_string()
    }

    fn params(&self) -> Params {
        let ast_params = match self.ast_node.param_list() {
            Some(it) => it,
            None => return Params::Named(Vec::new()),
        };

        let mut params_pats = Vec::new();
        let params_ty = if self.ctx.completion.has_dot_receiver() || self.receiver.is_some() {
            self.func.method_params(self.ctx.db()).unwrap_or_default()
        } else {
            if let Some(s) = ast_params.self_param() {
                cov_mark::hit!(parens_for_method_call_as_assoc_fn);
                params_pats.push(Some(s.to_string()));
            }
            self.func.assoc_fn_params(self.ctx.db())
        };
        params_pats
            .extend(ast_params.params().into_iter().map(|it| it.pat().map(|it| it.to_string())));

        let params = params_pats
            .into_iter()
            .zip(params_ty)
            .flat_map(|(pat, param_ty)| {
                let pat = pat?;
                let name = pat;
                let arg = name.trim_start_matches("mut ").trim_start_matches('_');
                Some(self.add_arg(arg, param_ty.ty()))
            })
            .collect();
        Params::Named(params)
    }

    fn kind(&self) -> CompletionItemKind {
        if self.func.self_param(self.ctx.db()).is_some() {
            CompletionItemKind::Method
        } else {
            SymbolKind::Function.into()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_edit, check_edit_with_config, TEST_CONFIG},
        CompletionConfig,
    };

    #[test]
    fn inserts_parens_for_function_calls() {
        cov_mark::check!(inserts_parens_for_function_calls);
        check_edit(
            "no_args",
            r#"
fn no_args() {}
fn main() { no_$0 }
"#,
            r#"
fn no_args() {}
fn main() { no_args()$0 }
"#,
        );

        check_edit(
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args(${1:x}, ${2:y})$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.foo()$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.f$0
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.foo(${1:x})$0
}
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        $0
    }
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        self.foo(${1:x})$0
    }
}
"#,
        );
    }

    #[test]
    fn parens_for_method_call_as_assoc_fn() {
        cov_mark::check!(parens_for_method_call_as_assoc_fn);
        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::foo(${1:&self})$0 }
"#,
        );
    }

    #[test]
    fn suppress_arg_snippets() {
        cov_mark::check!(suppress_arg_snippets);
        check_edit_with_config(
            CompletionConfig { add_call_argument_snippets: false, ..TEST_CONFIG },
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args($0) }
"#,
        );
    }

    #[test]
    fn strips_underscores_from_args() {
        check_edit(
            "foo",
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { f$0 }
"#,
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { foo(${1:foo}, ${2:bar}, ${3:ho_ge_})$0 }
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn insert_mut_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&mut x})$0
}
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope_for_method() {
        check_edit(
            "apply_foo",
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.$0
}
"#,
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.apply_foo(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn trim_mut_keyword_in_func_completion() {
        check_edit(
            "take_mutably",
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_m$0
}
"#,
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_mutably(${1:x})$0
}
"#,
        );
    }
}
<commit_msg>minor: follow code style<commit_after>//! Renderer for function calls.

use hir::{AsAssocItem, HasSource, HirDisplay};
use ide_db::SymbolKind;
use itertools::Itertools;
use syntax::ast;

use crate::{
    item::{CompletionItem, CompletionItemKind, CompletionKind, CompletionRelevance, ImportEdit},
    render::{
        builder_ext::Params, compute_exact_name_match, compute_ref_match, compute_type_match,
        RenderContext,
    },
};

pub(crate) fn render_fn(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_fn");
    Some(FunctionRender::new(ctx, None, local_name, fn_, false)?.render(import_to_add))
}

pub(crate) fn render_method(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    receiver: Option<hir::Name>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_method");
    Some(FunctionRender::new(ctx, receiver, local_name, fn_, true)?.render(import_to_add))
}

#[derive(Debug)]
struct FunctionRender<'a> {
    ctx: RenderContext<'a>,
    name: String,
    receiver: Option<hir::Name>,
    func: hir::Function,
    ast_node: ast::Fn,
    is_method: bool,
}

impl<'a> FunctionRender<'a> {
    fn new(
        ctx: RenderContext<'a>,
        receiver: Option<hir::Name>,
        local_name: Option<hir::Name>,
        fn_: hir::Function,
        is_method: bool,
    ) -> Option<FunctionRender<'a>> {
        let name = local_name.unwrap_or_else(|| fn_.name(ctx.db())).to_string();
        let ast_node = fn_.source(ctx.db())?.value;

        Some(FunctionRender { ctx, name, receiver, func: fn_, ast_node, is_method })
    }

    fn render(self, import_to_add: Option<ImportEdit>) -> CompletionItem {
        let params = self.params();
        let call = if let Some(receiver) = &self.receiver {
            format!("{}.{}", receiver, &self.name)
        } else {
            self.name.clone()
        };
        let mut item =
            CompletionItem::new(CompletionKind::Reference, self.ctx.source_range(), call.clone());
        item.kind(self.kind())
            .set_documentation(self.ctx.docs(self.func))
            .set_deprecated(
                self.ctx.is_deprecated(self.func) || self.ctx.is_deprecated_assoc_item(self.func),
            )
            .detail(self.detail())
            .add_call_parens(self.ctx.completion, call.clone(), params);

        if import_to_add.is_none() {
            let db = self.ctx.db();
            if let Some(actm) = self.func.as_assoc_item(db) {
                if let Some(trt) = actm.containing_trait_or_trait_impl(db) {
                    item.trait_name(trt.name(db).to_string());
                }
            }
        }

        item.add_import(import_to_add).lookup_by(self.name);

        let ret_type = self.func.ret_type(self.ctx.db());
        item.set_relevance(CompletionRelevance {
            type_match: compute_type_match(self.ctx.completion, &ret_type),
            exact_name_match: compute_exact_name_match(self.ctx.completion, &call),
            ..CompletionRelevance::default()
        });

        if let Some(ref_match) = compute_ref_match(self.ctx.completion, &ret_type) {
            // FIXME
            // For now we don't properly calculate the edits for ref match
            // completions on methods, so we've disabled them. See #8058.
            if !self.is_method {
                item.ref_match(ref_match);
            }
        }

        item.build()
    }

    fn detail(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());
        let ret = if ret_ty.is_unit() {
            // Omit the return type if it is the unit type
            String::new()
        } else {
            format!(" {}", self.ty_display())
        };

        format!("fn({}){}", self.params_display(), ret)
    }

    fn params_display(&self) -> String {
        if let Some(self_param) = self.func.self_param(self.ctx.db()) {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .skip(1) // skip the self param because we are manually handling that
                .map(|p| p.ty().display(self.ctx.db()).to_string());

            std::iter::once(self_param.display(self.ctx.db()).to_owned()).chain(params).join(", ")
        } else {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .map(|p| p.ty().display(self.ctx.db()).to_string())
                .join(", ");
            params
        }
    }

    fn ty_display(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());

        format!("-> {}", ret_ty.display(self.ctx.db()))
    }

    fn add_arg(&self, arg: &str, ty: &hir::Type) -> String {
        if let Some(derefed_ty) = ty.remove_ref() {
            for (name, local) in self.ctx.completion.locals.iter() {
                if name == arg && local.ty(self.ctx.db()) == derefed_ty {
                    let mutability = if ty.is_mutable_reference() { "&mut " } else { "&" };
                    return format!("{}{}", mutability, arg);
                }
            }
        }
        arg.to_string()
    }

    fn params(&self) -> Params {
        let ast_params = match self.ast_node.param_list() {
            Some(it) => it,
            None => return Params::Named(Vec::new()),
        };

        let mut params_pats = Vec::new();
        let params_ty = if self.ctx.completion.has_dot_receiver() || self.receiver.is_some() {
            self.func.method_params(self.ctx.db()).unwrap_or_default()
        } else {
            if let Some(s) = ast_params.self_param() {
                cov_mark::hit!(parens_for_method_call_as_assoc_fn);
                params_pats.push(Some(s.to_string()));
            }
            self.func.assoc_fn_params(self.ctx.db())
        };
        params_pats
            .extend(ast_params.params().into_iter().map(|it| it.pat().map(|it| it.to_string())));

        let params = params_pats
            .into_iter()
            .zip(params_ty)
            .flat_map(|(pat, param_ty)| {
                let pat = pat?;
                let name = pat;
                let arg = name.trim_start_matches("mut ").trim_start_matches('_');
                Some(self.add_arg(arg, param_ty.ty()))
            })
            .collect();
        Params::Named(params)
    }

    fn kind(&self) -> CompletionItemKind {
        if self.func.self_param(self.ctx.db()).is_some() {
            CompletionItemKind::Method
        } else {
            SymbolKind::Function.into()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_edit, check_edit_with_config, TEST_CONFIG},
        CompletionConfig,
    };

    #[test]
    fn inserts_parens_for_function_calls() {
        cov_mark::check!(inserts_parens_for_function_calls);
        check_edit(
            "no_args",
            r#"
fn no_args() {}
fn main() { no_$0 }
"#,
            r#"
fn no_args() {}
fn main() { no_args()$0 }
"#,
        );

        check_edit(
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args(${1:x}, ${2:y})$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.foo()$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.f$0
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.foo(${1:x})$0
}
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        $0
    }
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        self.foo(${1:x})$0
    }
}
"#,
        );
    }

    #[test]
    fn parens_for_method_call_as_assoc_fn() {
        cov_mark::check!(parens_for_method_call_as_assoc_fn);
        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::foo(${1:&self})$0 }
"#,
        );
    }

    #[test]
    fn suppress_arg_snippets() {
        cov_mark::check!(suppress_arg_snippets);
        check_edit_with_config(
            CompletionConfig { add_call_argument_snippets: false, ..TEST_CONFIG },
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args($0) }
"#,
        );
    }

    #[test]
    fn strips_underscores_from_args() {
        check_edit(
            "foo",
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { f$0 }
"#,
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { foo(${1:foo}, ${2:bar}, ${3:ho_ge_})$0 }
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn insert_mut_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&mut x})$0
}
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope_for_method() {
        check_edit(
            "apply_foo",
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.$0
}
"#,
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.apply_foo(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn trim_mut_keyword_in_func_completion() {
        check_edit(
            "take_mutably",
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_m$0
}
"#,
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_mutably(${1:x})$0
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate regex;

include!("../clap-test.rs");

use clap::{App, Arg, ErrorKind, SubCommand};

static VISIBLE_ALIAS_HELP: &'static str = "clap-test 2.6

USAGE:
    clap-test [SUBCOMMAND]

FLAGS:
    -h, --help       Prints help information
    -V, --version    Prints version information

SUBCOMMANDS:
    help    Prints this message or the help of the given subcommand(s)
    test    Some help [aliases: dongle, done]";

static INVISIBLE_ALIAS_HELP: &'static str = "clap-test 2.6

USAGE:
    clap-test [SUBCOMMAND]

FLAGS:
    -h, --help       Prints help information
    -V, --version    Prints version information

SUBCOMMANDS:
    help    Prints this message or the help of the given subcommand(s)
    test    Some help";

#[cfg(feature = "suggestions")]
static DYM_SUBCMD: &'static str = "error: The subcommand 'subcm' wasn't recognized
	Did you mean 'subcmd'?

If you believe you received this message in error, try re-running with 'dym -- subcm'

USAGE:
    dym [SUBCOMMAND]

For more information try --help";

#[cfg(feature = "suggestions")]
static DYM_ARG: &'static str =
    "error: Found argument '--subcm' which wasn't expected, or isn't valid in this context
\tDid you mean to put '--subcmdarg' after the subcommand 'subcmd'?

USAGE:
    dym [SUBCOMMAND]

For more information try --help";

#[test]
fn subcommand() {
    let m = App::new("test")
        .subcommand(
            SubCommand::with_name("some").arg(
                Arg::with_name("test")
                    .short("t")
                    .long("test")
                    .takes_value(true)
                    .help("testing testing"),
            ),
        )
        .arg(Arg::with_name("other").long("other"))
        .get_matches_from(vec!["myprog", "some", "--test", "testing"]);

    assert_eq!(m.subcommand_name().unwrap(), "some");
    let sub_m = m.subcommand_matches("some").unwrap();
    assert!(sub_m.is_present("test"));
    assert_eq!(sub_m.value_of("test").unwrap(), "testing");
}

#[test]
fn subcommand_none_given() {
    let m = App::new("test")
        .subcommand(
            SubCommand::with_name("some").arg(
                Arg::with_name("test")
                    .short("t")
                    .long("test")
                    .takes_value(true)
                    .help("testing testing"),
            ),
        )
        .arg(Arg::with_name("other").long("other"))
        .get_matches_from(vec![""]);

    assert!(m.subcommand_name().is_none());
}

#[test]
fn subcommand_multiple() {
    let m = App::new("test")
        .subcommands(vec![
            SubCommand::with_name("some").arg(
                Arg::with_name("test")
                    .short("t")
                    .long("test")
                    .takes_value(true)
                    .help("testing testing"),
            ),
            SubCommand::with_name("add").arg(Arg::with_name("roster").short("r")),
        ])
        .arg(Arg::with_name("other").long("other"))
        .get_matches_from(vec!["myprog", "some", "--test", "testing"]);

    assert!(m.subcommand_matches("some").is_some());
    assert!(m.subcommand_matches("add").is_none());
    assert_eq!(m.subcommand_name().unwrap(), "some");
    let sub_m = m.subcommand_matches("some").unwrap();
    assert!(sub_m.is_present("test"));
    assert_eq!(sub_m.value_of("test").unwrap(), "testing");
}

#[test]
fn single_alias() {
    let m = App::new("myprog")
        .subcommand(SubCommand::with_name("test").alias("do-stuff"))
        .get_matches_from(vec!["myprog", "do-stuff"]);
    assert_eq!(m.subcommand_name(), Some("test"));
}

#[test]
fn multiple_aliases() {
    let m = App::new("myprog")
        .subcommand(SubCommand::with_name("test").aliases(&["do-stuff", "test-stuff"]))
        .get_matches_from(vec!["myprog", "test-stuff"]);
    assert_eq!(m.subcommand_name(), Some("test"));
}

#[test]
#[cfg(feature = "suggestions")]
fn subcmd_did_you_mean_output() {
    let app = App::new("dym").subcommand(SubCommand::with_name("subcmd"));
    assert!(test::compare_output(app, "dym subcm", DYM_SUBCMD, true));
}

#[test]
#[cfg(feature = "suggestions")]
fn subcmd_did_you_mean_output_arg() {
    let app = App::new("dym").subcommand(
        SubCommand::with_name("subcmd").arg_from_usage("-s --subcmdarg [subcmdarg] 'tests'"),
    );
    assert!(test::compare_output(app, "dym --subcm foo", DYM_ARG, true));
}

#[test]
fn alias_help() {
    let m = App::new("myprog")
        .subcommand(SubCommand::with_name("test").alias("do-stuff"))
        .get_matches_from_safe(vec!["myprog", "help", "do-stuff"]);
    assert!(m.is_err());
    assert_eq!(m.unwrap_err().kind, ErrorKind::HelpDisplayed);
}

#[test]
fn visible_aliases_help_output() {
    let app = App::new("clap-test").version("2.6").subcommand(
        SubCommand::with_name("test")
            .about("Some help")
            .alias("invisible")
            .visible_alias("dongle")
            .visible_alias("done"),
    );
    assert!(test::compare_output(
        app,
        "clap-test --help",
        VISIBLE_ALIAS_HELP,
        false
    ));
}

#[test]
fn invisible_aliases_help_output() {
    let app = App::new("clap-test").version("2.6").subcommand(
        SubCommand::with_name("test")
            .about("Some help")
            .alias("invisible"),
    );
    assert!(test::compare_output(
        app,
        "clap-test --help",
        INVISIBLE_ALIAS_HELP,
        false
    ));
}

#[test]
fn issue_1031_args_with_same_name() {
    let res = App::new("prog")
        .arg(Arg::from_usage("--ui-path=<PATH>"))
        .subcommand(SubCommand::with_name("signer"))
        .get_matches_from_safe(vec!["prog", "--ui-path", "signer"]);

    assert!(res.is_ok(), "{:?}", res.unwrap_err().kind);
    let m = res.unwrap();
    assert_eq!(m.value_of("ui-path"), Some("signer"));
}

#[test]
fn issue_1031_args_with_same_name_no_more_vals() {
    let res = App::new("prog")
        .arg(Arg::from_usage("--ui-path=<PATH>"))
        .subcommand(SubCommand::with_name("signer"))
        .get_matches_from_safe(vec!["prog", "--ui-path", "value", "signer"]);

    assert!(res.is_ok(), "{:?}", res.unwrap_err().kind);
    let m = res.unwrap();
    assert_eq!(m.value_of("ui-path"), Some("value"));
    assert_eq!(m.subcommand_name(), Some("signer"));
}

#[test]
fn issue_1161_multiple_hyphen_hyphen() {
    // from example 22
    let res = App::new("myprog")
        .arg(Arg::with_name("eff").short("f"))
        .arg(Arg::with_name("pea").short("p").takes_value(true))
        .arg(Arg::with_name("slop").multiple(true).last(true))
        .get_matches_from_safe(vec!["-f", "-p=bob", "--", "sloppy", "slop", "-a", "--", "subprogram", "position", "args"]);

    assert!(res.is_ok(), "{:?}", res.unwrap_err().kind);
    let m = res.unwrap();

    let expected = Some(vec!["sloppy", "slop", "-a", "--", "subprogram", "position", "args"]);
    let actual = m
            .values_of("slop")
            .map(|vals| vals.collect::<Vec<_>>());

    assert_eq!(expected, actual);

}<commit_msg>clean up whitespace<commit_after>extern crate clap;
extern crate regex;

include!("../clap-test.rs");

use clap::{App, Arg, ErrorKind, SubCommand};

static VISIBLE_ALIAS_HELP: &'static str = "clap-test 2.6

USAGE:
    clap-test [SUBCOMMAND]

FLAGS:
    -h, --help       Prints help information
    -V, --version    Prints version information

SUBCOMMANDS:
    help    Prints this message or the help of the given subcommand(s)
    test    Some help [aliases: dongle, done]";

static INVISIBLE_ALIAS_HELP: &'static str = "clap-test 2.6

USAGE:
    clap-test [SUBCOMMAND]

FLAGS:
    -h, --help       Prints help information
    -V, --version    Prints version information

SUBCOMMANDS:
    help    Prints this message or the help of the given subcommand(s)
    test    Some help";

#[cfg(feature = "suggestions")]
static DYM_SUBCMD: &'static str = "error: The subcommand 'subcm' wasn't recognized
	Did you mean 'subcmd'?

If you believe you received this message in error, try re-running with 'dym -- subcm'

USAGE:
    dym [SUBCOMMAND]

For more information try --help";

#[cfg(feature = "suggestions")]
static DYM_ARG: &'static str =
    "error: Found argument '--subcm' which wasn't expected, or isn't valid in this context
\tDid you mean to put '--subcmdarg' after the subcommand 'subcmd'?

USAGE:
    dym [SUBCOMMAND]

For more information try --help";

#[test]
fn subcommand() {
    let m = App::new("test")
        .subcommand(
            SubCommand::with_name("some").arg(
                Arg::with_name("test")
                    .short("t")
                    .long("test")
                    .takes_value(true)
                    .help("testing testing"),
            ),
        )
        .arg(Arg::with_name("other").long("other"))
        .get_matches_from(vec!["myprog", "some", "--test", "testing"]);

    assert_eq!(m.subcommand_name().unwrap(), "some");
    let sub_m = m.subcommand_matches("some").unwrap();
    assert!(sub_m.is_present("test"));
    assert_eq!(sub_m.value_of("test").unwrap(), "testing");
}

#[test]
fn subcommand_none_given() {
    let m = App::new("test")
        .subcommand(
            SubCommand::with_name("some").arg(
                Arg::with_name("test")
                    .short("t")
                    .long("test")
                    .takes_value(true)
                    .help("testing testing"),
            ),
        )
        .arg(Arg::with_name("other").long("other"))
        .get_matches_from(vec![""]);

    assert!(m.subcommand_name().is_none());
}

#[test]
fn subcommand_multiple() {
    let m = App::new("test")
        .subcommands(vec![
            SubCommand::with_name("some").arg(
                Arg::with_name("test")
                    .short("t")
                    .long("test")
                    .takes_value(true)
                    .help("testing testing"),
            ),
            SubCommand::with_name("add").arg(Arg::with_name("roster").short("r")),
        ])
        .arg(Arg::with_name("other").long("other"))
        .get_matches_from(vec!["myprog", "some", "--test", "testing"]);

    assert!(m.subcommand_matches("some").is_some());
    assert!(m.subcommand_matches("add").is_none());
    assert_eq!(m.subcommand_name().unwrap(), "some");
    let sub_m = m.subcommand_matches("some").unwrap();
    assert!(sub_m.is_present("test"));
    assert_eq!(sub_m.value_of("test").unwrap(), "testing");
}

#[test]
fn single_alias() {
    let m = App::new("myprog")
        .subcommand(SubCommand::with_name("test").alias("do-stuff"))
        .get_matches_from(vec!["myprog", "do-stuff"]);
    assert_eq!(m.subcommand_name(), Some("test"));
}

#[test]
fn multiple_aliases() {
    let m = App::new("myprog")
        .subcommand(SubCommand::with_name("test").aliases(&["do-stuff", "test-stuff"]))
        .get_matches_from(vec!["myprog", "test-stuff"]);
    assert_eq!(m.subcommand_name(), Some("test"));
}

#[test]
#[cfg(feature = "suggestions")]
fn subcmd_did_you_mean_output() {
    let app = App::new("dym").subcommand(SubCommand::with_name("subcmd"));
    assert!(test::compare_output(app, "dym subcm", DYM_SUBCMD, true));
}

#[test]
#[cfg(feature = "suggestions")]
fn subcmd_did_you_mean_output_arg() {
    let app = App::new("dym").subcommand(
        SubCommand::with_name("subcmd").arg_from_usage("-s --subcmdarg [subcmdarg] 'tests'"),
    );
    assert!(test::compare_output(app, "dym --subcm foo", DYM_ARG, true));
}

#[test]
fn alias_help() {
    let m = App::new("myprog")
        .subcommand(SubCommand::with_name("test").alias("do-stuff"))
        .get_matches_from_safe(vec!["myprog", "help", "do-stuff"]);
    assert!(m.is_err());
    assert_eq!(m.unwrap_err().kind, ErrorKind::HelpDisplayed);
}

#[test]
fn visible_aliases_help_output() {
    let app = App::new("clap-test").version("2.6").subcommand(
        SubCommand::with_name("test")
            .about("Some help")
            .alias("invisible")
            .visible_alias("dongle")
            .visible_alias("done"),
    );
    assert!(test::compare_output(
        app,
        "clap-test --help",
        VISIBLE_ALIAS_HELP,
        false
    ));
}

#[test]
fn invisible_aliases_help_output() {
    let app = App::new("clap-test").version("2.6").subcommand(
        SubCommand::with_name("test")
            .about("Some help")
            .alias("invisible"),
    );
    assert!(test::compare_output(
        app,
        "clap-test --help",
        INVISIBLE_ALIAS_HELP,
        false
    ));
}

#[test]
fn issue_1031_args_with_same_name() {
    let res = App::new("prog")
        .arg(Arg::from_usage("--ui-path=<PATH>"))
        .subcommand(SubCommand::with_name("signer"))
        .get_matches_from_safe(vec!["prog", "--ui-path", "signer"]);

    assert!(res.is_ok(), "{:?}", res.unwrap_err().kind);
    let m = res.unwrap();
    assert_eq!(m.value_of("ui-path"), Some("signer"));
}

#[test]
fn issue_1031_args_with_same_name_no_more_vals() {
    let res = App::new("prog")
        .arg(Arg::from_usage("--ui-path=<PATH>"))
        .subcommand(SubCommand::with_name("signer"))
        .get_matches_from_safe(vec!["prog", "--ui-path", "value", "signer"]);

    assert!(res.is_ok(), "{:?}", res.unwrap_err().kind);
    let m = res.unwrap();
    assert_eq!(m.value_of("ui-path"), Some("value"));
    assert_eq!(m.subcommand_name(), Some("signer"));
}

#[test]
fn issue_1161_multiple_hyphen_hyphen() {
    // from example 22
    let res = App::new("myprog")
        .arg(Arg::with_name("eff").short("f"))
        .arg(Arg::with_name("pea").short("p").takes_value(true))
        .arg(Arg::with_name("slop").multiple(true).last(true))
        .get_matches_from_safe(vec!["-f", "-p=bob", "--", "sloppy", "slop", "-a", "--", "subprogram", "position", "args"]);

    assert!(res.is_ok(), "{:?}", res.unwrap_err().kind);
    let m = res.unwrap();

    let expected = Some(vec!["sloppy", "slop", "-a", "--", "subprogram", "position", "args"]);
    let actual = m
            .values_of("slop")
            .map(|vals| vals.collect::<Vec<_>>());

    assert_eq!(expected, actual);
}<|endoftext|>