<commit_before>use Cursive;
use Printer;
use direction;
use event::*;
use menu::{MenuItem, MenuTree};
use std::rc::Rc;
use theme::ColorStyle;
use unicode_width::UnicodeWidthStr;
use vec::Vec2;
use view::{Position, View};
use views::{MenuPopup, OnEventView};

/// Current state of the menubar
#[derive(PartialEq, Debug)]
enum State {
    /// The menubar is inactive.
    Inactive,
    /// The menubar is actively selected.
    ///
    /// It will receive input.
    Selected,
    /// The menubar is still visible, but a submenu is open.
    ///
    /// It will not receive input.
    Submenu,
}

/// Shows a single-line list of items, with pop-up menus when one is selected.
///
/// The [`Cursive`] root already includes a menubar
/// that you just need to configure.
///
/// [`Cursive`]: ../struct.Cursive.html#method.menubar
pub struct Menubar {
    /// Menu items in this menubar.
    root: MenuTree,

    /// TODO: move this out of this view.
    pub autohide: bool,
    focus: usize,

    // TODO: make Menubar impl View and take out the State management
    state: State,
}

new_default!(Menubar);

impl Menubar {
    /// Creates a new, empty menubar.
    pub fn new() -> Self {
        Menubar {
            root: MenuTree::new(),
            autohide: true,
            state: State::Inactive,
            focus: 0,
        }
    }

    /// Hides the menubar.
    fn hide(&mut self) {
        self.state = State::Inactive;
    }

    /// True if we should be receiving events.
    pub fn receive_events(&self) -> bool {
        self.state == State::Selected
    }

    /// True if some submenus are visible.
    pub fn has_submenu(&self) -> bool {
        self.state == State::Submenu
    }

    /// Returns `true` if we should be drawn.
    pub fn visible(&self) -> bool {
        !self.autohide || self.state != State::Inactive
    }

    /// Adds a new item to the menubar.
    ///
    /// The item will use the given title, and on selection, will open a
    /// popup-menu with the given menu tree.
    pub fn add_subtree<S>(&mut self, title: S, menu: MenuTree) -> &mut Self
    where
        S: Into<String>,
    {
        let i = self.root.len();
        self.insert_subtree(i, title, menu)
    }

    /// Adds a delimiter to the menubar.
    pub fn add_delimiter(&mut self) -> &mut Self {
        let i = self.root.len();
        self.insert_delimiter(i)
    }

    /// Adds a leaf node to the menubar.
    pub fn add_leaf<S, F>(&mut self, title: S, cb: F) -> &mut Self
    where
        S: Into<String>,
        F: 'static + Fn(&mut Cursive),
    {
        let i = self.root.len();
        self.insert_leaf(i, title, cb)
    }

    /// Insert a new item at the given position.
    pub fn insert_subtree<S>(
        &mut self, i: usize, title: S, menu: MenuTree
    ) -> &mut Self
    where
        S: Into<String>,
    {
        self.root.insert_subtree(i, title, menu);
        self
    }

    /// Inserts a new delimiter at the given position.
    ///
    /// It will show up as `|`.
    pub fn insert_delimiter(&mut self, i: usize) -> &mut Self {
        self.root.insert_delimiter(i);
        self
    }

    /// Inserts a new leaf node at the given position.
    ///
    /// It will be directly actionable.
    pub fn insert_leaf<S, F>(&mut self, i: usize, title: S, cb: F) -> &mut Self
    where
        S: Into<String>,
        F: 'static + Fn(&mut Cursive),
    {
        self.root.insert_leaf(i, title, cb);
        self
    }

    /// Removes all menu items from this menubar.
    pub fn clear(&mut self) {
        self.root.clear();
        self.focus = 0;
    }

    /// Returns the number of items in this menubar.
    pub fn len(&self) -> usize {
        self.root.len()
    }

    /// Returns `true` if this menubar is empty.
    pub fn is_empty(&self) -> bool {
        self.root.is_empty()
    }

    /// Returns the item at the given position.
    ///
    /// Returns `None` if `i > self.len()`
    pub fn get_subtree(&mut self, i: usize) -> Option<&mut MenuTree> {
        self.root.get_subtree(i)
    }

    /// Looks for an item with the given label.
    pub fn find_subtree(&mut self, label: &str) -> Option<&mut MenuTree> {
        self.root.find_subtree(label)
    }

    /// Returns the position of the item with the given label.
    ///
    /// Returns `None` if no such label was found.
    pub fn find_position(&mut self, label: &str) -> Option<usize> {
        self.root.find_position(label)
    }

    /// Remove the item at the given position.
    pub fn remove(&mut self, i: usize) {
        self.root.remove(i);
    }

    fn child_at(&self, x: usize) -> Option<usize> {
        if x == 0 {
            return None;
        }
        let mut offset = 1;

        for (i, child) in self.root.children.iter().enumerate() {
            offset += child.label().width() + 2;
            if x < offset {
                return Some(i);
            }
        }

        None
    }

    fn select_child(&mut self, open_only: bool) -> EventResult {
        match self.root.children[self.focus] {
            MenuItem::Leaf(_, ref cb) if !open_only => {
                EventResult::Consumed(Some(cb.clone()))
            }
            MenuItem::Subtree(_, ref tree) => {
                // First, we need a new Rc to send the callback,
                // since we don't know when it will be called.
                let menu = Rc::clone(tree);

                self.state = State::Submenu;
                let offset = Vec2::new(
                    self.root.children[..self.focus]
                        .iter()
                        .map(|child| child.label().width() + 2)
                        .sum(),
                    if self.autohide { 1 } else { 0 },
                );
                // Since the closure will be called multiple times,
                // we also need a new Rc on every call.
                EventResult::with_cb(move |s| {
                    show_child(s, offset, Rc::clone(&menu))
                })
            }
            _ => EventResult::Ignored,
        }
    }
}

fn show_child(s: &mut Cursive, offset: Vec2, menu: Rc<MenuTree>) {
    // Adds a new layer located near the item title with the menu popup.
    // Also adds two key callbacks on this new view, to handle `left` and
    // `right` key presses.
    // (If the view itself listens for a `left` or `right` press, it will
    // consume it before our OnEventView. This means sub-menus can properly
    // be entered.)
    s.screen_mut().add_layer_at(
        Position::absolute(offset),
        OnEventView::new(
            MenuPopup::new(menu)
                .on_dismiss(|s| s.select_menubar())
                .on_action(|s| s.menubar().state = State::Inactive),
        ).on_event(Key::Right, |s| {
            s.pop_layer();
            s.select_menubar();
            // Act as if we sent "Right" then "Down"
            s.menubar()
                .on_event(Event::Key(Key::Right))
                .process(s);
            if let EventResult::Consumed(Some(cb)) =
                s.menubar().on_event(Event::Key(Key::Down))
            {
                cb(s);
            }
        })
            .on_event(Key::Left, |s| {
                s.pop_layer();
                s.select_menubar();
                // Act as if we sent "Left" then "Down"
                s.menubar()
                    .on_event(Event::Key(Key::Left))
                    .process(s);
                if let EventResult::Consumed(Some(cb)) =
                    s.menubar().on_event(Event::Key(Key::Down))
                {
                    cb(s);
                }
            }),
    );
}

impl View for Menubar {
    fn draw(&self, printer: &Printer) {
        // Draw the bar at the top
        printer.with_color(ColorStyle::primary(), |printer| {
            printer.print_hline((0, 0), printer.size.x, " ");
        });

        // TODO: draw the rest
        let mut offset = 1;
        for (i, item) in self.root.children.iter().enumerate() {
            let title = item.label();

            // We don't want to show HighlightInactive when we're not selected,
            // because it's ugly on the menubar.
            let selected =
                (self.state != State::Inactive) && (i == self.focus);
            printer.with_selection(selected, |printer| {
                printer.print((offset, 0), &format!(" {} ", title));
            });
            offset += title.width() + 2;
        }
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        match event {
            Event::Key(Key::Esc) => {
                self.hide();
                return EventResult::with_cb(|s| s.clear());
            }
            Event::Key(Key::Left) => loop {
                if self.focus > 0 {
                    self.focus -= 1;
                } else {
                    self.focus = self.root.len() - 1;
                }
                if !self.root.children[self.focus].is_delimiter() {
                    break;
                }
            },
            Event::Key(Key::Right) => loop {
                if self.focus + 1 < self.root.len() {
                    self.focus += 1;
                } else {
                    self.focus = 0;
                }
                if !self.root.children[self.focus].is_delimiter() {
                    break;
                }
            },
            Event::Key(Key::Down) => {
                return self.select_child(true);
            }
            Event::Key(Key::Enter) => {
                return self.select_child(false);
            }
            Event::Mouse {
                event: MouseEvent::Press(btn),
                position,
                offset,
            } if position.fits(offset) && position.y == offset.y =>
            {
                if let Some(child) = position
                    .checked_sub(offset)
                    .and_then(|pos| self.child_at(pos.x))
                {
                    if !self.root.children[child].is_delimiter() {
                        self.focus = child;
                        if btn == MouseButton::Left {
                            return self.select_child(true);
                        }
                    }
                }
            }
            Event::Mouse {
                event: MouseEvent::Release(btn),
                position,
                offset,
            } if position.fits(offset) && position.y == offset.y =>
            {
                if let Some(child) = position
                    .checked_sub(offset)
                    .and_then(|pos| self.child_at(pos.x))
                {
                    if self.focus == child && btn == MouseButton::Left
                        && self.root.children[child].is_leaf()
                    {
                        return self.select_child(false);
                    }
                }
            }
            Event::Mouse {
                event: MouseEvent::Press(_),
                ..
            } => {
                self.hide();
                return EventResult::with_cb(|s| s.clear());
            }
            _ => return EventResult::Ignored,
        }
        EventResult::Consumed(None)
    }

    fn take_focus(&mut self, _: direction::Direction) -> bool {
        self.state = State::Selected;
        true
    }

    fn required_size(&mut self, _: Vec2) -> Vec2 {
        // TODO: scroll the options if the screen is too small?

        // We add 2 to the length of every label for marin.
        // Also, we add 1 at the beginning.
        // (See the `draw()` method)
        let width = self.root
            .children
            .iter()
            .map(|item| item.label().len() + 2)
            .sum();

        Vec2::new(width, 1)
    }
}
<commit_msg>Implement Menubar::important_area<commit_after>use Cursive;
use Printer;
use direction;
use event::*;
use menu::{MenuItem, MenuTree};
use rect::Rect;
use std::rc::Rc;
use theme::ColorStyle;
use unicode_width::UnicodeWidthStr;
use vec::Vec2;
use view::{Position, View};
use views::{MenuPopup, OnEventView};

/// Current state of the menubar
#[derive(PartialEq, Debug)]
enum State {
    /// The menubar is inactive.
    Inactive,
    /// The menubar is actively selected.
    ///
    /// It will receive input.
    Selected,
    /// The menubar is still visible, but a submenu is open.
    ///
    /// It will not receive input.
    Submenu,
}

/// Shows a single-line list of items, with pop-up menus when one is selected.
///
/// The [`Cursive`] root already includes a menubar
/// that you just need to configure.
///
/// [`Cursive`]: ../struct.Cursive.html#method.menubar
pub struct Menubar {
    /// Menu items in this menubar.
    root: MenuTree,

    /// TODO: move this out of this view.
    pub autohide: bool,
    focus: usize,

    // TODO: make Menubar impl View and take out the State management
    state: State,
}

new_default!(Menubar);

impl Menubar {
    /// Creates a new, empty menubar.
    pub fn new() -> Self {
        Menubar {
            root: MenuTree::new(),
            autohide: true,
            state: State::Inactive,
            focus: 0,
        }
    }

    /// Hides the menubar.
    fn hide(&mut self) {
        self.state = State::Inactive;
    }

    /// True if we should be receiving events.
    pub fn receive_events(&self) -> bool {
        self.state == State::Selected
    }

    /// True if some submenus are visible.
    pub fn has_submenu(&self) -> bool {
        self.state == State::Submenu
    }

    /// Returns `true` if we should be drawn.
    pub fn visible(&self) -> bool {
        !self.autohide || self.state != State::Inactive
    }

    /// Adds a new item to the menubar.
    ///
    /// The item will use the given title, and on selection, will open a
    /// popup-menu with the given menu tree.
    pub fn add_subtree<S>(&mut self, title: S, menu: MenuTree) -> &mut Self
    where
        S: Into<String>,
    {
        let i = self.root.len();
        self.insert_subtree(i, title, menu)
    }

    /// Adds a delimiter to the menubar.
    pub fn add_delimiter(&mut self) -> &mut Self {
        let i = self.root.len();
        self.insert_delimiter(i)
    }

    /// Adds a leaf node to the menubar.
    pub fn add_leaf<S, F>(&mut self, title: S, cb: F) -> &mut Self
    where
        S: Into<String>,
        F: 'static + Fn(&mut Cursive),
    {
        let i = self.root.len();
        self.insert_leaf(i, title, cb)
    }

    /// Insert a new item at the given position.
    pub fn insert_subtree<S>(
        &mut self, i: usize, title: S, menu: MenuTree
    ) -> &mut Self
    where
        S: Into<String>,
    {
        self.root.insert_subtree(i, title, menu);
        self
    }

    /// Inserts a new delimiter at the given position.
    ///
    /// It will show up as `|`.
    pub fn insert_delimiter(&mut self, i: usize) -> &mut Self {
        self.root.insert_delimiter(i);
        self
    }

    /// Inserts a new leaf node at the given position.
    ///
    /// It will be directly actionable.
    pub fn insert_leaf<S, F>(&mut self, i: usize, title: S, cb: F) -> &mut Self
    where
        S: Into<String>,
        F: 'static + Fn(&mut Cursive),
    {
        self.root.insert_leaf(i, title, cb);
        self
    }

    /// Removes all menu items from this menubar.
    pub fn clear(&mut self) {
        self.root.clear();
        self.focus = 0;
    }

    /// Returns the number of items in this menubar.
    pub fn len(&self) -> usize {
        self.root.len()
    }

    /// Returns `true` if this menubar is empty.
    pub fn is_empty(&self) -> bool {
        self.root.is_empty()
    }

    /// Returns the item at the given position.
    ///
    /// Returns `None` if `i > self.len()`
    pub fn get_subtree(&mut self, i: usize) -> Option<&mut MenuTree> {
        self.root.get_subtree(i)
    }

    /// Looks for an item with the given label.
    pub fn find_subtree(&mut self, label: &str) -> Option<&mut MenuTree> {
        self.root.find_subtree(label)
    }

    /// Returns the position of the item with the given label.
    ///
    /// Returns `None` if no such label was found.
    pub fn find_position(&mut self, label: &str) -> Option<usize> {
        self.root.find_position(label)
    }

    /// Remove the item at the given position.
    pub fn remove(&mut self, i: usize) {
        self.root.remove(i);
    }

    fn child_at(&self, x: usize) -> Option<usize> {
        if x == 0 {
            return None;
        }
        let mut offset = 1;

        for (i, child) in self.root.children.iter().enumerate() {
            offset += child.label().width() + 2;
            if x < offset {
                return Some(i);
            }
        }

        None
    }

    fn select_child(&mut self, open_only: bool) -> EventResult {
        match self.root.children[self.focus] {
            MenuItem::Leaf(_, ref cb) if !open_only => {
                EventResult::Consumed(Some(cb.clone()))
            }
            MenuItem::Subtree(_, ref tree) => {
                // First, we need a new Rc to send the callback,
                // since we don't know when it will be called.
                let menu = Rc::clone(tree);

                self.state = State::Submenu;
                let offset = Vec2::new(
                    self.root.children[..self.focus]
                        .iter()
                        .map(|child| child.label().width() + 2)
                        .sum(),
                    if self.autohide { 1 } else { 0 },
                );
                // Since the closure will be called multiple times,
                // we also need a new Rc on every call.
                EventResult::with_cb(move |s| {
                    show_child(s, offset, Rc::clone(&menu))
                })
            }
            _ => EventResult::Ignored,
        }
    }
}

fn show_child(s: &mut Cursive, offset: Vec2, menu: Rc<MenuTree>) {
    // Adds a new layer located near the item title with the menu popup.
    // Also adds two key callbacks on this new view, to handle `left` and
    // `right` key presses.
    // (If the view itself listens for a `left` or `right` press, it will
    // consume it before our OnEventView. This means sub-menus can properly
    // be entered.)
    s.screen_mut().add_layer_at(
        Position::absolute(offset),
        OnEventView::new(
            MenuPopup::new(menu)
                .on_dismiss(|s| s.select_menubar())
                .on_action(|s| s.menubar().state = State::Inactive),
        ).on_event(Key::Right, |s| {
            s.pop_layer();
            s.select_menubar();
            // Act as if we sent "Right" then "Down"
            s.menubar()
                .on_event(Event::Key(Key::Right))
                .process(s);
            if let EventResult::Consumed(Some(cb)) =
                s.menubar().on_event(Event::Key(Key::Down))
            {
                cb(s);
            }
        })
            .on_event(Key::Left, |s| {
                s.pop_layer();
                s.select_menubar();
                // Act as if we sent "Left" then "Down"
                s.menubar()
                    .on_event(Event::Key(Key::Left))
                    .process(s);
                if let EventResult::Consumed(Some(cb)) =
                    s.menubar().on_event(Event::Key(Key::Down))
                {
                    cb(s);
                }
            }),
    );
}

impl View for Menubar {
    fn draw(&self, printer: &Printer) {
        // Draw the bar at the top
        printer.with_color(ColorStyle::primary(), |printer| {
            printer.print_hline((0, 0), printer.size.x, " ");
        });

        // TODO: draw the rest
        let mut offset = 1;
        for (i, item) in self.root.children.iter().enumerate() {
            let title = item.label();

            // We don't want to show HighlightInactive when we're not selected,
            // because it's ugly on the menubar.
            let selected =
                (self.state != State::Inactive) && (i == self.focus);
            printer.with_selection(selected, |printer| {
                printer.print((offset, 0), &format!(" {} ", title));
            });
            offset += title.width() + 2;
        }
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        match event {
            Event::Key(Key::Esc) => {
                self.hide();
                return EventResult::with_cb(|s| s.clear());
            }
            Event::Key(Key::Left) => loop {
                if self.focus > 0 {
                    self.focus -= 1;
                } else {
                    self.focus = self.root.len() - 1;
                }
                if !self.root.children[self.focus].is_delimiter() {
                    break;
                }
            },
            Event::Key(Key::Right) => loop {
                if self.focus + 1 < self.root.len() {
                    self.focus += 1;
                } else {
                    self.focus = 0;
                }
                if !self.root.children[self.focus].is_delimiter() {
                    break;
                }
            },
            Event::Key(Key::Down) => {
                return self.select_child(true);
            }
            Event::Key(Key::Enter) => {
                return self.select_child(false);
            }
            Event::Mouse {
                event: MouseEvent::Press(btn),
                position,
                offset,
            } if position.fits(offset) && position.y == offset.y =>
            {
                if let Some(child) = position
                    .checked_sub(offset)
                    .and_then(|pos| self.child_at(pos.x))
                {
                    if !self.root.children[child].is_delimiter() {
                        self.focus = child;
                        if btn == MouseButton::Left {
                            return self.select_child(true);
                        }
                    }
                }
            }
            Event::Mouse {
                event: MouseEvent::Release(btn),
                position,
                offset,
            } if position.fits(offset) && position.y == offset.y =>
            {
                if let Some(child) = position
                    .checked_sub(offset)
                    .and_then(|pos| self.child_at(pos.x))
                {
                    if self.focus == child && btn == MouseButton::Left
                        && self.root.children[child].is_leaf()
                    {
                        return self.select_child(false);
                    }
                }
            }
            Event::Mouse {
                event: MouseEvent::Press(_),
                ..
            } => {
                self.hide();
                return EventResult::with_cb(|s| s.clear());
            }
            _ => return EventResult::Ignored,
        }
        EventResult::Consumed(None)
    }

    fn take_focus(&mut self, _: direction::Direction) -> bool {
        self.state = State::Selected;
        true
    }

    fn required_size(&mut self, _: Vec2) -> Vec2 {
        // TODO: scroll the options if the screen is too small?

        // We add 2 to the length of every label for marin.
        // Also, we add 1 at the beginning.
        // (See the `draw()` method)
        let width = self.root
            .children
            .iter()
            .map(|item| item.label().len() + 2)
            .sum();

        Vec2::new(width, 1)
    }

    fn important_area(&self, _: Vec2) -> Rect {
        if self.root.is_empty() {
            return Rect::from((0, 0));
        }

        // X position is 1 (margin before the first item) + sum of widths
        // And each item has a 2 cells padding.
        let x = 1 + self.root.children[..self.focus]
            .iter()
            .map(|child| child.label().width() + 2)
            .sum::<usize>();

        let width = self.root.children[self.focus].label().width();

        Rect::from_size((x, 0), (width, 1))
    }
}
<|endoftext|>


<commit_before>use std::iter::{Fuse,Peekable};

/// An iterator adaptor that wraps each element in an [`Position`](../enum.Position.html).
///
/// Iterator element type is `Position<I::Item>`.
///
/// See [`.with_position()`](../trait.Itertools.html#method.with_position) for more information.
pub struct WithPosition<I>
    where I: Iterator,
{
    handled_first: bool,
    peekable: Peekable<Fuse<I>>,
}

/// Create a new `WithPosition` iterator.
pub fn with_position<I>(iter: I) -> WithPosition<I>
    where I: Iterator,
{
    WithPosition {
        handled_first: false,
        peekable: iter.fuse().peekable(),
    }
}

/// A value yielded by `WithPosition`.
/// Indicates the position of this element in the iterator results.
///
/// See [`.with_position()`](trait.Itertools.html#method.with_position) for more information.
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Position<T> {
    /// This is the first element.
    First(T),
    /// This is neither the first nor the last element.
    Middle(T),
    /// This is the last element.
    Last(T),
    /// This is the only element.
    Only(T),
}

impl<I: Iterator> Iterator for WithPosition<I> {
    type Item = Position<I::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.peekable.next() {
            Some(item) => {
                if !self.handled_first {
                    // Haven't seen the first item yet, and there is one to give.
                    self.handled_first = true;
                    // Peek to see if this is also the last item,
                    // in which case tag it as `Only`.
                    match self.peekable.peek() {
                        Some(_) => Some(Position::First(item)),
                        None => Some(Position::Only(item)),
                    }
                } else {
                    // Have seen the first item, and there's something left.
                    // Peek to see if this is the last item.
                    match self.peekable.peek() {
                        Some(_) => Some(Position::Middle(item)),
                        None => Some(Position::Last(item)),
                    }
                }
            }
            // Iterator is finished.
            None => None,
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.peekable.size_hint()
    }
}
<commit_msg>Add Position::into_inner<commit_after>use std::iter::{Fuse,Peekable};

/// An iterator adaptor that wraps each element in an [`Position`](../enum.Position.html).
///
/// Iterator element type is `Position<I::Item>`.
///
/// See [`.with_position()`](../trait.Itertools.html#method.with_position) for more information.
pub struct WithPosition<I>
    where I: Iterator,
{
    handled_first: bool,
    peekable: Peekable<Fuse<I>>,
}

/// Create a new `WithPosition` iterator.
pub fn with_position<I>(iter: I) -> WithPosition<I>
    where I: Iterator,
{
    WithPosition {
        handled_first: false,
        peekable: iter.fuse().peekable(),
    }
}

/// A value yielded by `WithPosition`.
/// Indicates the position of this element in the iterator results.
///
/// See [`.with_position()`](trait.Itertools.html#method.with_position) for more information.
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Position<T> {
    /// This is the first element.
    First(T),
    /// This is neither the first nor the last element.
    Middle(T),
    /// This is the last element.
    Last(T),
    /// This is the only element.
    Only(T),
}

impl<T> Position<T> {
    /// Return the inner value.
    pub fn into_inner(self) -> T {
        match self {
            Position::First(x) |
            Position::Middle(x) |
            Position::Last(x) |
            Position::Only(x) => x,
        }
    }
}

impl<I: Iterator> Iterator for WithPosition<I> {
    type Item = Position<I::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.peekable.next() {
            Some(item) => {
                if !self.handled_first {
                    // Haven't seen the first item yet, and there is one to give.
                    self.handled_first = true;
                    // Peek to see if this is also the last item,
                    // in which case tag it as `Only`.
                    match self.peekable.peek() {
                        Some(_) => Some(Position::First(item)),
                        None => Some(Position::Only(item)),
                    }
                } else {
                    // Have seen the first item, and there's something left.
                    // Peek to see if this is the last item.
                    match self.peekable.peek() {
                        Some(_) => Some(Position::Middle(item)),
                        None => Some(Position::Last(item)),
                    }
                }
            }
            // Iterator is finished.
            None => None,
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.peekable.size_hint()
    }
}
<|endoftext|>


<commit_before>//! Helper module to enable serializing bytes more efficiently

use core::{ops, fmt, char, iter, slice};
use core::fmt::Write;

use ser;

#[cfg(any(feature = "std", feature = "collections"))]
pub use self::bytebuf::{ByteBuf, ByteBufVisitor};

#[cfg(feature = "collections")]
use collections::Vec;

///////////////////////////////////////////////////////////////////////////////

/// `Bytes` wraps a `&[u8]` in order to serialize into a byte array.
#[derive(Clone, Copy, Eq, Hash, PartialEq, PartialOrd, Ord)]
pub struct Bytes<'a> {
    bytes: &'a [u8],
}

impl<'a> Bytes<'a> {
    /// Wrap an existing `&[u8]`.
    pub fn new(bytes: &'a [u8]) -> Self {
        Bytes {
            bytes: bytes,
        }
    }
}

impl<'a> fmt::Debug for Bytes<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(f.write_str("b\""));
        for c in escape_bytestring(self.bytes) {
            try!(f.write_char(c));
        }
        f.write_char('"')
    }
}

impl<'a> From<&'a [u8]> for Bytes<'a> {
    fn from(bytes: &'a [u8]) -> Self {
        Bytes {
            bytes: bytes,
        }
    }
}

#[cfg(any(feature = "std", feature = "collections"))]
impl<'a> From<&'a Vec<u8>> for Bytes<'a> {
    fn from(bytes: &'a Vec<u8>) -> Self {
        Bytes {
            bytes: bytes,
        }
    }
}

impl<'a> Into<&'a [u8]> for Bytes<'a> {
    fn into(self) -> &'a [u8] {
        self.bytes
    }
}

impl<'a> ops::Deref for Bytes<'a> {
    type Target = [u8];

    fn deref(&self) -> &[u8] { self.bytes }
}

impl<'a> ser::Serialize for Bytes<'a> {
    #[inline]
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: ser::Serializer
    {
        serializer.serialize_bytes(self.bytes)
    }
}

///////////////////////////////////////////////////////////////////////////////

#[cfg(any(feature = "std", feature = "collections"))]
mod bytebuf {
    use core::ops;
    use core::fmt;
    use core::fmt::Write;

    use ser;
    use de;

    #[cfg(feature = "collections")]
    use collections::Vec;

    /// `ByteBuf` wraps a `Vec<u8>` and serializes as a byte array.
    #[derive(Clone, Default, Eq, Hash, PartialEq, PartialOrd, Ord)]
    pub struct ByteBuf {
        bytes: Vec<u8>,
    }

    impl ByteBuf {
        /// Construct a new, empty `ByteBuf`.
        pub fn new() -> Self {
            ByteBuf {
                bytes: Vec::new(),
            }
        }

        /// Construct a new, empty `ByteBuf` with the specified capacity.
        pub fn with_capacity(cap: usize) -> Self {
            ByteBuf {
                bytes: Vec::with_capacity(cap)
            }
        }

        /// Wrap existing bytes in a `ByteBuf`.
        pub fn from<T: Into<Vec<u8>>>(bytes: T) -> Self {
            ByteBuf {
                bytes: bytes.into(),
            }
        }
    }

    impl fmt::Debug for ByteBuf {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            try!(f.write_str("b\""));
            for c in super::escape_bytestring(self.bytes.as_ref()) {
                try!(f.write_char(c));
            }
            f.write_char('"')
        }
    }

    impl Into<Vec<u8>> for ByteBuf {
        fn into(self) -> Vec<u8> {
            self.bytes
        }
    }

    impl From<Vec<u8>> for ByteBuf {
        fn from(bytes: Vec<u8>) -> Self {
            ByteBuf {
                bytes: bytes,
            }
        }
    }

    impl AsRef<Vec<u8>> for ByteBuf {
        fn as_ref(&self) -> &Vec<u8> {
            &self.bytes
        }
    }

    impl AsRef<[u8]> for ByteBuf {
        fn as_ref(&self) -> &[u8] {
            &self.bytes
        }
    }

    impl AsMut<Vec<u8>> for ByteBuf {
        fn as_mut(&mut self) -> &mut Vec<u8> {
            &mut self.bytes
        }
    }

    impl AsMut<[u8]> for ByteBuf {
        fn as_mut(&mut self) -> &mut [u8] {
            &mut self.bytes
        }
    }

    impl ops::Deref for ByteBuf {
        type Target = [u8];

        fn deref(&self) -> &[u8] { &self.bytes[..] }
    }

    impl ops::DerefMut for ByteBuf {
        fn deref_mut(&mut self) -> &mut [u8] { &mut self.bytes[..] }
    }

    impl ser::Serialize for ByteBuf {
        fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
            where S: ser::Serializer
        {
            serializer.serialize_bytes(self)
        }
    }

    /// This type implements the `serde::de::Visitor` trait for a `ByteBuf`.
    pub struct ByteBufVisitor;

    impl de::Visitor for ByteBufVisitor {
        type Value = ByteBuf;

        #[inline]
        fn visit_unit<E>(&mut self) -> Result<ByteBuf, E>
            where E: de::Error,
        {
            Ok(ByteBuf {
                bytes: Vec::new(),
            })
        }

        #[inline]
        fn visit_seq<V>(&mut self, mut visitor: V) -> Result<ByteBuf, V::Error>
            where V: de::SeqVisitor,
        {
            let (len, _) = visitor.size_hint();
            let mut values = Vec::with_capacity(len);

            while let Some(value) = try!(visitor.visit()) {
                values.push(value);
            }

            try!(visitor.end());

            Ok(ByteBuf {
                bytes: values,
            })
        }

        #[inline]
        fn visit_bytes<E>(&mut self, v: &[u8]) -> Result<ByteBuf, E>
            where E: de::Error,
        {
            self.visit_byte_buf(v.to_vec())
        }

        #[inline]
        fn visit_byte_buf<E>(&mut self, v: Vec<u8>) -> Result<ByteBuf, E>
            where E: de::Error,
        {
            Ok(ByteBuf {
                bytes: v,
            })
        }
    }

    impl de::Deserialize for ByteBuf {
        #[inline]
        fn deserialize<D>(deserializer: &mut D) -> Result<ByteBuf, D::Error>
            where D: de::Deserializer
        {
            deserializer.deserialize_bytes(ByteBufVisitor)
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

#[inline]
fn escape_bytestring<'a>(bytes: &'a [u8]) -> iter::FlatMap<slice::Iter<'a, u8>, char::EscapeDefault, fn(&u8) -> char::EscapeDefault> {
    fn f(b: &u8) -> char::EscapeDefault {
        char::from_u32(*b as u32).unwrap().escape_default()
    }
    bytes.iter().flat_map(f as fn(&u8) -> char::EscapeDefault)
}
<commit_msg>Use constructors to create Bytes and ByteBuf<commit_after>//! Helper module to enable serializing bytes more efficiently

use core::{ops, fmt, char, iter, slice};
use core::fmt::Write;

use ser;

#[cfg(any(feature = "std", feature = "collections"))]
pub use self::bytebuf::{ByteBuf, ByteBufVisitor};

#[cfg(feature = "collections")]
use collections::Vec;

///////////////////////////////////////////////////////////////////////////////

/// `Bytes` wraps a `&[u8]` in order to serialize into a byte array.
#[derive(Clone, Copy, Eq, Hash, PartialEq, PartialOrd, Ord)]
pub struct Bytes<'a> {
    bytes: &'a [u8],
}

impl<'a> Bytes<'a> {
    /// Wrap an existing `&[u8]`.
    pub fn new(bytes: &'a [u8]) -> Self {
        Bytes {
            bytes: bytes,
        }
    }
}

impl<'a> fmt::Debug for Bytes<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(f.write_str("b\""));
        for c in escape_bytestring(self.bytes) {
            try!(f.write_char(c));
        }
        f.write_char('"')
    }
}

impl<'a> From<&'a [u8]> for Bytes<'a> {
    fn from(bytes: &'a [u8]) -> Self {
        Bytes::new(bytes)
    }
}

#[cfg(any(feature = "std", feature = "collections"))]
impl<'a> From<&'a Vec<u8>> for Bytes<'a> {
    fn from(bytes: &'a Vec<u8>) -> Self {
        Bytes::new(bytes)
    }
}

impl<'a> Into<&'a [u8]> for Bytes<'a> {
    fn into(self) -> &'a [u8] {
        self.bytes
    }
}

impl<'a> ops::Deref for Bytes<'a> {
    type Target = [u8];

    fn deref(&self) -> &[u8] { self.bytes }
}

impl<'a> ser::Serialize for Bytes<'a> {
    #[inline]
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: ser::Serializer
    {
        serializer.serialize_bytes(self.bytes)
    }
}

///////////////////////////////////////////////////////////////////////////////

#[cfg(any(feature = "std", feature = "collections"))]
mod bytebuf {
    use core::ops;
    use core::fmt;
    use core::fmt::Write;

    use ser;
    use de;

    #[cfg(feature = "collections")]
    use collections::Vec;

    /// `ByteBuf` wraps a `Vec<u8>` and serializes as a byte array.
    #[derive(Clone, Default, Eq, Hash, PartialEq, PartialOrd, Ord)]
    pub struct ByteBuf {
        bytes: Vec<u8>,
    }

    impl ByteBuf {
        /// Construct a new, empty `ByteBuf`.
        pub fn new() -> Self {
            ByteBuf::from(Vec::new())
        }

        /// Construct a new, empty `ByteBuf` with the specified capacity.
        pub fn with_capacity(cap: usize) -> Self {
            ByteBuf::from(Vec::with_capacity(cap))
        }

        /// Wrap existing bytes in a `ByteBuf`.
        pub fn from<T: Into<Vec<u8>>>(bytes: T) -> Self {
            ByteBuf {
                bytes: bytes.into(),
            }
        }
    }

    impl fmt::Debug for ByteBuf {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            try!(f.write_str("b\""));
            for c in super::escape_bytestring(self.bytes.as_ref()) {
                try!(f.write_char(c));
            }
            f.write_char('"')
        }
    }

    impl Into<Vec<u8>> for ByteBuf {
        fn into(self) -> Vec<u8> {
            self.bytes
        }
    }

    impl From<Vec<u8>> for ByteBuf {
        fn from(bytes: Vec<u8>) -> Self {
            ByteBuf::from(bytes)
        }
    }

    impl AsRef<Vec<u8>> for ByteBuf {
        fn as_ref(&self) -> &Vec<u8> {
            &self.bytes
        }
    }

    impl AsRef<[u8]> for ByteBuf {
        fn as_ref(&self) -> &[u8] {
            &self.bytes
        }
    }

    impl AsMut<Vec<u8>> for ByteBuf {
        fn as_mut(&mut self) -> &mut Vec<u8> {
            &mut self.bytes
        }
    }

    impl AsMut<[u8]> for ByteBuf {
        fn as_mut(&mut self) -> &mut [u8] {
            &mut self.bytes
        }
    }

    impl ops::Deref for ByteBuf {
        type Target = [u8];

        fn deref(&self) -> &[u8] { &self.bytes[..] }
    }

    impl ops::DerefMut for ByteBuf {
        fn deref_mut(&mut self) -> &mut [u8] { &mut self.bytes[..] }
    }

    impl ser::Serialize for ByteBuf {
        fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
            where S: ser::Serializer
        {
            serializer.serialize_bytes(self)
        }
    }

    /// This type implements the `serde::de::Visitor` trait for a `ByteBuf`.
    pub struct ByteBufVisitor;

    impl de::Visitor for ByteBufVisitor {
        type Value = ByteBuf;

        #[inline]
        fn visit_unit<E>(&mut self) -> Result<ByteBuf, E>
            where E: de::Error,
        {
            Ok(ByteBuf::new())
        }

        #[inline]
        fn visit_seq<V>(&mut self, mut visitor: V) -> Result<ByteBuf, V::Error>
            where V: de::SeqVisitor,
        {
            let (len, _) = visitor.size_hint();
            let mut values = Vec::with_capacity(len);

            while let Some(value) = try!(visitor.visit()) {
                values.push(value);
            }

            try!(visitor.end());

            Ok(ByteBuf::from(values))
        }

        #[inline]
        fn visit_bytes<E>(&mut self, v: &[u8]) -> Result<ByteBuf, E>
            where E: de::Error,
        {
            self.visit_byte_buf(v.to_vec())
        }

        #[inline]
        fn visit_byte_buf<E>(&mut self, v: Vec<u8>) -> Result<ByteBuf, E>
            where E: de::Error,
        {
            Ok(ByteBuf::from(v))
        }
    }

    impl de::Deserialize for ByteBuf {
        #[inline]
        fn deserialize<D>(deserializer: &mut D) -> Result<ByteBuf, D::Error>
            where D: de::Deserializer
        {
            deserializer.deserialize_bytes(ByteBufVisitor)
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

#[inline]
fn escape_bytestring<'a>(bytes: &'a [u8]) -> iter::FlatMap<slice::Iter<'a, u8>, char::EscapeDefault, fn(&u8) -> char::EscapeDefault> {
    fn f(b: &u8) -> char::EscapeDefault {
        char::from_u32(*b as u32).unwrap().escape_default()
    }
    bytes.iter().flat_map(f as fn(&u8) -> char::EscapeDefault)
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![allow(non_upper_case_globals)]

use seccomp::{
    allow_syscall, allow_syscall_if, Error, SeccompAction, SeccompCmpOp::Eq,
    SeccompCondition as Cond, SeccompFilter, SeccompRule,
};

// Currently these variables are missing from rust libc:
// https://github.com/rust-lang/libc/blob/master/src/unix/notbsd/linux/musl/b64/aarch64.rs
// even though they are defined in musl libc:
// https://git.musl-libc.org/cgit/musl/tree/arch/aarch64/bits/syscall.h.in.
// Submitted issue in rust-lang: https://github.com/rust-lang/libc/issues/1348.
const SYS_fcntl: ::std::os::raw::c_long = 25;
const SYS_lseek: ::std::os::raw::c_long = 62;
const SYS_newfstatat: ::std::os::raw::c_long = 79;
const SYS_fstat: ::std::os::raw::c_long = 80;
const SYS_mmap: ::std::os::raw::c_long = 222;

/// The default filter containing the white listed syscall rules required by `Firecracker` to
/// function.
pub fn default_filter() -> Result<SeccompFilter, Error> {
    Ok(SeccompFilter::new(
        vec![
            #[cfg(target_env = "musl")]
            allow_syscall(libc::SYS_accept),
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_accept4),
            allow_syscall(libc::SYS_brk),
            allow_syscall(libc::SYS_clock_gettime),
            allow_syscall(libc::SYS_close),
            allow_syscall(libc::SYS_dup),
            allow_syscall_if(
                libc::SYS_epoll_ctl,
                or![
                    and![Cond::new(1, Eq, super::EPOLL_CTL_ADD)?],
                    and![Cond::new(1, Eq, super::EPOLL_CTL_DEL)?],
                ],
            ),
            #[cfg(target_env = "musl")]
            allow_syscall(libc::SYS_epoll_pwait),
            allow_syscall(libc::SYS_exit),
            allow_syscall(libc::SYS_exit_group),
            allow_syscall_if(
                SYS_fcntl,
                or![and![
                    Cond::new(1, Eq, super::FCNTL_F_SETFD)?,
                    Cond::new(2, Eq, super::FCNTL_FD_CLOEXEC)?,
                ]],
            ),
            allow_syscall(SYS_fstat),
            allow_syscall(SYS_newfstatat),
            allow_syscall_if(
                libc::SYS_futex,
                or![
                    and![Cond::new(1, Eq, super::FUTEX_WAIT_PRIVATE)?],
                    and![Cond::new(1, Eq, super::FUTEX_WAKE_PRIVATE)?],
                    and![Cond::new(1, Eq, super::FUTEX_REQUEUE_PRIVATE)?],
                ],
            ),
            allow_syscall(libc::SYS_getrandom),
            allow_syscall_if(libc::SYS_ioctl, super::create_ioctl_seccomp_rule()?),
            allow_syscall(SYS_lseek),
            #[cfg(target_env = "musl")]
            allow_syscall_if(
                libc::SYS_madvise,
                or![and![Cond::new(2, Eq, libc::MADV_DONTNEED as u64)?],],
            ),
            allow_syscall(SYS_mmap),
            allow_syscall(libc::SYS_munmap),
            #[cfg(target_env = "musl")]
            allow_syscall(libc::SYS_openat),
            allow_syscall(libc::SYS_read),
            allow_syscall(libc::SYS_readv),
            // SYS_rt_sigreturn is needed in case a fault does occur, so that the signal handler
            // can return. Otherwise we get stuck in a fault loop.
            allow_syscall(libc::SYS_rt_sigreturn),
            allow_syscall(libc::SYS_timerfd_create),
            allow_syscall(libc::SYS_timerfd_settime),
            allow_syscall(libc::SYS_write),
            allow_syscall(libc::SYS_writev),
        ]
        .into_iter()
        .collect(),
        SeccompAction::Trap,
    )?)
}
<commit_msg>aarch64: missed gnu syscalls<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![allow(non_upper_case_globals)]

use seccomp::{
    allow_syscall, allow_syscall_if, Error, SeccompAction, SeccompCmpOp::Eq,
    SeccompCondition as Cond, SeccompFilter, SeccompRule,
};

// Currently these variables are missing from rust libc:
// https://github.com/rust-lang/libc/blob/master/src/unix/notbsd/linux/musl/b64/aarch64.rs
// even though they are defined in musl libc:
// https://git.musl-libc.org/cgit/musl/tree/arch/aarch64/bits/syscall.h.in.
// Submitted issue in rust-lang: https://github.com/rust-lang/libc/issues/1348.
const SYS_fcntl: ::std::os::raw::c_long = 25;
const SYS_lseek: ::std::os::raw::c_long = 62;
const SYS_newfstatat: ::std::os::raw::c_long = 79;
const SYS_fstat: ::std::os::raw::c_long = 80;
const SYS_mmap: ::std::os::raw::c_long = 222;

/// The default filter containing the white listed syscall rules required by `Firecracker` to
/// function.
pub fn default_filter() -> Result<SeccompFilter, Error> {
    Ok(SeccompFilter::new(
        vec![
            #[cfg(target_env = "musl")]
            allow_syscall(libc::SYS_accept),
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_accept4),
            allow_syscall(libc::SYS_brk),
            allow_syscall(libc::SYS_clock_gettime),
            allow_syscall(libc::SYS_close),
            allow_syscall(libc::SYS_dup),
            allow_syscall_if(
                libc::SYS_epoll_ctl,
                or![
                    and![Cond::new(1, Eq, super::EPOLL_CTL_ADD)?],
                    and![Cond::new(1, Eq, super::EPOLL_CTL_DEL)?],
                ],
            ),
            allow_syscall(libc::SYS_epoll_pwait),
            allow_syscall(libc::SYS_exit),
            allow_syscall(libc::SYS_exit_group),
            allow_syscall_if(
                SYS_fcntl,
                or![and![
                    Cond::new(1, Eq, super::FCNTL_F_SETFD)?,
                    Cond::new(2, Eq, super::FCNTL_FD_CLOEXEC)?,
                ]],
            ),
            allow_syscall(SYS_fstat),
            allow_syscall(SYS_newfstatat),
            allow_syscall_if(
                libc::SYS_futex,
                or![
                    and![Cond::new(1, Eq, super::FUTEX_WAIT_PRIVATE)?],
                    and![Cond::new(1, Eq, super::FUTEX_WAKE_PRIVATE)?],
                    and![Cond::new(1, Eq, super::FUTEX_REQUEUE_PRIVATE)?],
                ],
            ),
            allow_syscall(libc::SYS_getrandom),
            allow_syscall_if(libc::SYS_ioctl, super::create_ioctl_seccomp_rule()?),
            allow_syscall(SYS_lseek),
            #[cfg(target_env = "musl")]
            allow_syscall_if(
                libc::SYS_madvise,
                or![and![Cond::new(2, Eq, libc::MADV_DONTNEED as u64)?],],
            ),
            allow_syscall(SYS_mmap),
            allow_syscall(libc::SYS_munmap),
            allow_syscall(libc::SYS_openat),
            allow_syscall(libc::SYS_read),
            allow_syscall(libc::SYS_readv),
            // SYS_rt_sigreturn is needed in case a fault does occur, so that the signal handler
            // can return. Otherwise we get stuck in a fault loop.
            allow_syscall(libc::SYS_rt_sigreturn),
            allow_syscall(libc::SYS_timerfd_create),
            allow_syscall(libc::SYS_timerfd_settime),
            allow_syscall(libc::SYS_write),
            allow_syscall(libc::SYS_writev),
        ]
        .into_iter()
        .collect(),
        SeccompAction::Trap,
    )?)
}
<|endoftext|>


<commit_before>extern crate rusty_machine as rm;
extern crate rand;

use rm::linalg::{Matrix, BaseMatrix};
use rand::{Rng, ThreadRng};
use std::collections::HashMap;

fn generate_arm(content_id: &usize, arms: &HashMap<usize, Arm>, feature_dim: &usize) -> Arm {
    let arm: Arm;
    if !arms.contains_key(content_id) {
        arm = Arm::new(*feature_dim, *content_id, 0.0001, 0.0);
    }else{
        arm = arms.get(content_id).unwrap().clone();
    };
    arm
}

#[derive(Debug, Clone)]
struct Arm{
    content_id : usize,
    alpha : f64,
    norm_mean : Matrix<f64>,
    cov_matrix: Matrix<f64>,
    win_rate : f64,
    win : f64,
    lose : f64,
}

impl Arm {
    fn new(feature_dim: usize, content_id: usize, alpha: f64, win_rate: f64) -> Arm {
        Arm {
            content_id: content_id,
            norm_mean: Matrix::zeros(1, feature_dim),
            cov_matrix: Matrix::ones(1, feature_dim),
            alpha: alpha,
            win_rate: win_rate,
            win: 0.0,
            lose: 0.0,
        }
    }

    fn update(&mut self, features: Matrix<f64>, is_click: bool) {
        // error of diag
        self.cov_matrix += Matrix::new(self.cov_matrix.rows(),
                                        self.cov_matrix.cols(),
                                        (&features.transpose()*&features).diag());
        let cost_of_click = is_click as i8 as f64;

        let feat_mul_cost_vec = features.iter().map(|x| cost_of_click*x).collect::<Vec<f64>>();
        let feat_mul_cost = Matrix::new(features.rows(), features.cols(), feat_mul_cost_vec);
        self.norm_mean += feat_mul_cost;

        if is_click{
            self.win += 1.;
        } else{
            self.lose += 1.;
        }
        self.win_rate = &self.win/(&self.win + &self.lose);
    }

    fn predict(&mut self, features: &Matrix<f64>) -> f64 {
        let one_div_cov_vec = self.cov_matrix.iter().map(|x| 1.0/x).collect::<Vec<f64>>();
        let one_div_cov = Matrix::new(self.cov_matrix.rows(),
                                      self.cov_matrix.cols(), one_div_cov_vec);
        // Since the covariance matrix preserves only the diagonal components,
        // it suffices to take the inverse matrix
        let theta = &one_div_cov.elemul(&self.norm_mean);
        // Again, the inverse matrix of the covariance matrix
        // is computed by taking reciprocal
        let mut tmp: f64 = ((features.elemul(&one_div_cov))*features.transpose()).data()[0];
        tmp = &tmp.sqrt()*&self.alpha;
        (theta*features.transpose()).data()[0] + &tmp
    }

    fn print_result(&mut self) {
        println!("content_id:{}, total_num:{}, win_rate:{}",
                &self.content_id, &self.win+&self.lose, &self.win_rate);
    }
}

struct Viewer{
    gender: String,
    rng: ThreadRng,
}

impl Viewer {
    fn new(gender: String) -> Viewer{
        Viewer{
            gender: gender,
            rng: rand::thread_rng(),
        }
    }

    fn view(&mut self, content_id: &usize) -> bool{
        if &self.gender=="man" {
            // Men are easy to click on ads with id 5 or less
            if *content_id < 6 {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        } else {
            // Women are easy to click on ads with id 6 or higher
            if *content_id > 5{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        }
    }
}

struct Rulet{
    rng: ThreadRng,
}

impl Rulet{
    fn new() -> Rulet {
        Rulet{
            rng: rand::thread_rng(),
        }
    }

    fn generate_features(&mut self, viewer: &Viewer) -> Matrix<f64> {
        let features = Some(&viewer.gender)
            .and_then(|gender|
                if gender=="man" {Some(Matrix::new(1,2,vec![1.,0.]))}
                else {Some(Matrix::new(1,2,vec![0.,1.]))}
            ).unwrap();
        features
    }

    fn generate_content(&mut self) -> usize{
        self.rng.gen_range(0, 10)
    }

    fn generate_gender(&mut self) -> String {
        if self.rng.next_f32() > 0.5{
            return "man".to_string();
        } else{
            return "women".to_string();
        }
    }
}

fn main() {
    /*Context is for men and women only
    Men are easy to click on ads with id 5 or less
    Women are easy to click on ads with id 6 or higher
    */
    let feature_dim = 2;
    let num_of_views = 10000;
    let mut rulet = Rulet::new();

    let mut content_id: usize;
    let mut features: Matrix<f64>;
    let mut is_clicked: bool;
    let mut arms: HashMap<usize, Arm> = HashMap::new();

    for _ in 0..num_of_views {
        let mut viewer = Viewer::new(rulet.generate_gender());

        content_id = rulet.generate_content();
        features = rulet.generate_features(&viewer);
        is_clicked = viewer.view(&content_id);

        let mut arm = generate_arm(&content_id, &arms, &feature_dim);
        arm.update(features, is_clicked);
        arms.remove(&content_id);
        arms.insert(content_id, arm.clone());
    }

    let man_mat: Matrix<f64> = Matrix::new(1,2, vec![1.,0.]);
    let woman_mat: Matrix<f64> = Matrix::new(1,2, vec![0.,1.]);

    println!("print result======");
    for (_, arm) in arms.iter() {
        let mut arm = arm.clone();
        &arm.print_result();
        println!("Click rate when men browse: {}", &arm.predict(&man_mat) );
        println!("Click rate when women browse: {}", &arm.predict(&woman_mat) );
    }

}
<commit_msg>modify<commit_after>extern crate rusty_machine as rm;
extern crate rand;

use rm::linalg::{Matrix, BaseMatrix};
use rand::{Rng, ThreadRng};
use std::collections::HashMap;

fn generate_arm(content_id: &usize, arms: &HashMap<usize, Arm>, feature_dim: &usize) -> Arm {
    let arm: Arm;
    if !arms.contains_key(content_id) {
        arm = Arm::new(*feature_dim, *content_id);
    }else{
        arm = arms.get(content_id).unwrap().clone();
    };
    arm
}

#[derive(Debug, Clone)]
struct Arm{
    content_id : usize,
    alpha : f64,
    norm_mean : Matrix<f64>,
    cov_matrix: Matrix<f64>,
    win_rate : f64,
    win : f64,
    lose : f64,
}

impl Arm {
    fn new(feature_dim: usize, content_id: usize) -> Arm {
        Arm {
            content_id: content_id,
            norm_mean: Matrix::zeros(1, feature_dim),
            cov_matrix: Matrix::ones(1, feature_dim),
            alpha: 0.0001,
            win_rate: 0.,
            win: 0.,
            lose: 0.,
        }
    }

    fn update(&mut self, features: Matrix<f64>, is_click: bool) {
        // error of diag
        self.cov_matrix += Matrix::new(self.cov_matrix.rows(),
                                        self.cov_matrix.cols(),
                                        (&features.transpose()*&features).diag());
        let cost_of_click = is_click as i8 as f64;

        let feat_mul_cost_vec = features.iter().map(|x| cost_of_click*x).collect::<Vec<f64>>();
        let feat_mul_cost = Matrix::new(features.rows(), features.cols(), feat_mul_cost_vec);
        self.norm_mean += feat_mul_cost;

        if is_click{
            self.win += 1.;
        } else{
            self.lose += 1.;
        }
        self.win_rate = &self.win/(&self.win + &self.lose);
    }

    fn predict(&mut self, features: &Matrix<f64>) -> f64 {
        let one_div_cov_vec = self.cov_matrix.iter().map(|x| 1.0/x).collect::<Vec<f64>>();
        let one_div_cov = Matrix::new(self.cov_matrix.rows(),
                                      self.cov_matrix.cols(), one_div_cov_vec);
        // Since the covariance matrix preserves only the diagonal components,
        // it suffices to take the inverse matrix
        let theta = &one_div_cov.elemul(&self.norm_mean);
        // Again, the inverse matrix of the covariance matrix
        // is computed by taking reciprocal
        let mut tmp: f64 = ((features.elemul(&one_div_cov))*features.transpose()).data()[0];
        tmp = &tmp.sqrt()*&self.alpha;
        (theta*features.transpose()).data()[0] + &tmp
    }

    fn print_result(&mut self) {
        println!("content_id:{}, total_num:{}, win_rate:{}",
                &self.content_id, &self.win+&self.lose, &self.win_rate);
    }
}

struct Viewer{
    gender: String,
    rng: ThreadRng,
}

impl Viewer {
    fn new(gender: String) -> Viewer{
        Viewer{
            gender: gender,
            rng: rand::thread_rng(),
        }
    }

    fn view(&mut self, content_id: &usize) -> bool{
        if &self.gender=="man" {
            // Men are easy to click on ads with id 5 or less
            if *content_id < 6 {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        } else {
            // Women are easy to click on ads with id 6 or higher
            if *content_id > 5{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        }
    }
}

struct Rulet{
    rng: ThreadRng,
}

impl Rulet{
    fn new() -> Rulet {
        Rulet{
            rng: rand::thread_rng(),
        }
    }

    fn generate_features(&mut self, viewer: &Viewer) -> Matrix<f64> {
        let features = Some(&viewer.gender)
            .and_then(|gender|
                if gender=="man" {Some(Matrix::new(1,2,vec![1.,0.]))}
                else {Some(Matrix::new(1,2,vec![0.,1.]))}
            ).unwrap();
        features
    }

    fn generate_content(&mut self) -> usize{
        self.rng.gen_range(0, 10)
    }

    fn generate_gender(&mut self) -> String {
        if self.rng.next_f32() > 0.5{
            return "man".to_string();
        } else{
            return "women".to_string();
        }
    }
}

fn main() {
    /*Context is for men and women only
    Men are easy to click on ads with id 5 or less
    Women are easy to click on ads with id 6 or higher
    */
    let feature_dim = 2;
    let num_of_views = 10000;
    let mut rulet = Rulet::new();

    let mut content_id: usize;
    let mut features: Matrix<f64>;
    let mut is_clicked: bool;
    let mut arms: HashMap<usize, Arm> = HashMap::new();

    for _ in 0..num_of_views {
        let mut viewer = Viewer::new(rulet.generate_gender());

        content_id = rulet.generate_content();
        features = rulet.generate_features(&viewer);
        is_clicked = viewer.view(&content_id);

        let mut arm = generate_arm(&content_id, &arms, &feature_dim);
        arm.update(features, is_clicked);
        arms.remove(&content_id);
        arms.insert(content_id, arm.clone());
    }

    let man_mat: Matrix<f64> = Matrix::new(1,2, vec![1.,0.]);
    let woman_mat: Matrix<f64> = Matrix::new(1,2, vec![0.,1.]);

    println!("print result======");
    for (_, arm) in arms.iter() {
        let mut arm = arm.clone();
        &arm.print_result();
        println!("Click rate when men browse: {}", &arm.predict(&man_mat) );
        println!("Click rate when women browse: {}", &arm.predict(&woman_mat) );
    }

}
<|endoftext|>


<commit_before>
extern crate cratesfyi;
extern crate clap;
#[macro_use]
extern crate log;
extern crate time;

use std::env;
use std::fs;
use std::process::{Command, exit};
use std::path::PathBuf;

use cratesfyi::docbuilder::{DocBuilder, DocBuilderError, command_result};
use cratesfyi::docbuilder::crte::Crate;
use cratesfyi::db;
use clap::{Arg, App, SubCommand};
use log::{LogLevel, LogLevelFilter, LogRecord, LogMetadata};
use time::now;

// I am using really simple Logger in this program
struct Logger;

impl log::Log for Logger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} {} - {}", now().to_utc().rfc3339(), record.level(), record.args());
        }
    }
}

fn update_crates_io_index(path: &PathBuf) -> Result<String, String> {
    info!("Updating crates.io-index");

    if path.exists() {
        let cwd = env::current_dir().unwrap();
        env::set_current_dir(path).unwrap();
        let res = command_result(Command::new("git").arg("pull").output().unwrap());
        env::set_current_dir(cwd).unwrap();
        res
    } else {
        command_result(Command::new("git")
            .arg("clone")
            .arg("https://github.com/rust-lang/crates.io-index.git")
            .arg(path.to_str().unwrap())
            .output().unwrap())
    }
}

// This will remove everything in CWD!!!
fn clean_build_dir() -> Result<(), DocBuilderError> {

    for file in try!(env::current_dir().unwrap()
                     .read_dir().map_err(DocBuilderError::RemoveBuildDir)) {
        let file = try!(file.map_err(DocBuilderError::RemoveBuildDir));
        let path = file.path();

        if path.file_name().unwrap() == ".cargo" ||
            path.file_name().unwrap() == ".crates.io-index" {
                continue;
            }

        if path.is_dir() {
            try!(fs::remove_dir_all(path).map_err(DocBuilderError::RemoveBuildDir));
        } else {
            try!(fs::remove_file(path).map_err(DocBuilderError::RemoveBuildDir));
        }
    }

    Ok(())
}

fn main() {
    log::set_logger(|max_log_level| {
        max_log_level.set(LogLevelFilter::Info);
        Box::new(Logger)
    }).unwrap();

    let matches = App::new("cratesfyi")
                      .version(env!("CARGO_PKG_VERSION"))
                      .about("Crates for your info!")
                      .subcommand(SubCommand::with_name("build")
                                      .about("Builds documentation in a chroot environment")
                                      .arg(Arg::with_name("PREFIX")
                                               .short("P")
                                               .long("prefix")
                                               .takes_value(true))
                                      .arg(Arg::with_name("DESTINATION")
                                               .short("d")
                                               .long("destination")
                                               .help("Sets destination path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CHROOT_PATH")
                                               .short("c")
                                               .long("chroot")
                                               .help("Sets chroot path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CHROOT_USER")
                                               .short("u")
                                               .long("chroot-user")
                                               .help("Sets chroot user name")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CRATES_IO_INDEX_PATH")
                                               .long("crates-io-index-path")
                                               .help("Sets crates.io-index path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("LOGS_PATH")
                                               .long("logs-path")
                                               .help("Sets logs path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("SKIP_IF_EXISTS")
                                               .short("s")
                                               .long("skip")
                                               .help("Skips building documentation if \
                                                      documentation exists"))
                                      .arg(Arg::with_name("SKIP_IF_LOG_EXISTS")
                                               .long("skip-if-log-exists")
                                               .help("Skips building documentation if build \
                                                      log exists"))
                                      .arg(Arg::with_name("KEEP_BUILD_DIRECTORY")
                                               .short("-k")
                                               .long("keep-build-directory")
                                               .help("Keeps build directory after build."))
                                      .subcommand(SubCommand::with_name("download-sources")
                                                      .about("Downloads sources of all crates"))
                                      .subcommand(SubCommand::with_name("world")
                                                      .about("Builds documentation of every \
                                                              crate")
                                                      .arg(Arg::with_name("SKIP_OLDEST_VERSIONS")
                                                               .long("skip-oldest-versions")
                                                               .help("Skips trying to build \
                                                                      oldest versions \
                                                                      if build fails"))
                                                      .arg(Arg::with_name("BUILD_ONLY_LATEST_V\
                                                                           ERSION")
                                                               .long("build-only-latest-versio\
                                                                      n")
                                                               .help("Builds only latest \
                                                                      version of crate and \
                                                                      skips oldest versions")))
                                      .subcommand(SubCommand::with_name("crate")
                                                      .about("Builds documentation for a crate")
                                                      .arg(Arg::with_name("CRATE_NAME")
                                                               .index(1)
                                                               .required(true)
                                                               .help("Crate name"))
                                                      .arg(Arg::with_name("CRATE_VERSION")
                                                               .index(2)
                                                               .required(true)
                                                               .help("Version of crate"))))
                      .subcommand(SubCommand::with_name("build-doc")
                                      .about("Builds documentation in CWD")
                                      .arg(Arg::with_name("CRATES_IO_INDEX_PATH")
                                               .short("p")
                                               .long("crates-io-index-path")
                                               .help("Sets crates.io-index path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CLEAN")
                                               .short("c")
                                               .long("clean")
                                               .help("Clean build dir before building"))
                                      .arg(Arg::with_name("CRATE_NAME")
                                               .index(1)
                                               .required(true)
                                               .help("Crate name"))
                                      .arg(Arg::with_name("CRATE_VERSION")
                                               .index(2)
                                               .required(true)
                                               .help("Version of crate")))
                      .subcommand(SubCommand::with_name("database")
                                      .about("Database operations")
                                      .subcommand(SubCommand::with_name("init")
                                                      .about("Initialize database. Currently \
                                                             only creates tables in database.")))
                      .get_matches();

    // DocBuilder
    if let Some(matches) = matches.subcommand_matches("build") {
        let mut dbuilder = {
            if let Some(prefix) = matches.value_of("PREFIX") {
                DocBuilder::from_prefix(PathBuf::from(prefix))
            } else {
                DocBuilder::default()
            }
        };

        // set destination
        if let Some(destination) = matches.value_of("DESTINATION") {
            dbuilder.destination(PathBuf::from(destination));
        }

        // set chroot path
        if let Some(chroot_path) = matches.value_of("CHROOT_PATH") {
            dbuilder.destination(PathBuf::from(chroot_path));
        }

        // set chroot user name
        if let Some(chroot_user) = matches.value_of("CHROOT_USER") {
            dbuilder.chroot_user(chroot_user.to_string());
        }

        // set crates.io-index path
        if let Some(crates_io_index_path) = matches.value_of("CRATES_IO_INDEX_PATH") {
            dbuilder.crates_io_index_path(PathBuf::from(crates_io_index_path));
        }

        // set logs path
        if let Some(logs_path) = matches.value_of("LOGS_PATH") {
            dbuilder.logs_path(PathBuf::from(logs_path));
        }

        dbuilder.skip_if_exists(matches.is_present("SKIP_IF_EXISTS"));
        dbuilder.skip_if_log_exists(matches.is_present("SKIP_IF_LOG_EXISTS"));
        dbuilder.keep_build_directory(matches.is_present("KEEP_BUILD_DIRECTORY"));

        // check paths
        if let Err(e) = dbuilder.check_paths() {
            println!("{:?}\nUse --help to get more information", e);
            std::process::exit(1);
        }

        // build world
        if let Some(matches) = matches.subcommand_matches("world") {
            dbuilder.build_only_latest_version(matches.is_present("BUILD_ONLY_LATEST_VERSION"));
            dbuilder.skip_oldest_versions(matches.is_present("SKIP_OLDEST_VERSIONS"));
            if let Err(e) = dbuilder.build_doc_for_every_crate() {
                println!("Failed to build world: {:#?}", e);
            }
        } else if let Some(matches) = matches.subcommand_matches("crate") {
            // Safe to call unwrap here
            let crte_name = matches.value_of("CRATE_NAME").unwrap();
            let version = matches.value_of("CRATE_VERSION").unwrap();
            let crte = Crate::new(crte_name.to_string(), vec![version.to_string()]);

            if let Err(e) = dbuilder.build_doc_for_crate_version(&crte, 0) {
                match e {
                    DocBuilderError::SkipDocumentationExists => {
                        println!("Skipping {} documentation already exists",
                                 crte.canonical_name(0))
                    }
                    _ => {
                        println!("Failed to build documentation for {}: {:?}",
                                 crte.canonical_name(0),
                                 e)
                    }
                }
            }
        } else if let Some(_) = matches.subcommand_matches("download-sources") {
            if let Err(e) = dbuilder.download_sources() {
                println!("{:?}", e);
            }
        }
    }

    // build-doc
    else if let Some(matches) = matches.subcommand_matches("build-doc") {
        let mut docbuilder = DocBuilder::default();

        // set crates.io-index path
        let crates_io_index_path = {
            if let Some(crates_io_index_path_conf) = matches.value_of("CRATES_IO_INDEX_PATH") {
                PathBuf::from(crates_io_index_path_conf)
            } else {
                let mut path = env::home_dir().unwrap();
                path.push(".crates.io-index");
                path
            }
        };

        docbuilder.crates_io_index_path(PathBuf::from(&crates_io_index_path));

        // update crates.io-index path
        if let Err(e) = update_crates_io_index(&crates_io_index_path) {
            panic!("{}", e);
        }

        // crates.io-index required for single crate
        if let Err(e) = docbuilder.check_crates_io_index_path() {
            panic!("{:?}", e);
        }

        let crte_name = matches.value_of("CRATE_NAME").unwrap();
        let version = matches.value_of("CRATE_VERSION").unwrap();

        let crte = Crate::new(crte_name.to_string(), vec![version.to_string()]);

        if matches.is_present("CLEAN") {
            clean_build_dir().unwrap();
        }

        if let Err(e) = crte.build_crate_doc(0, &docbuilder) {
            error!("Failed to build crate\n{:?}", e);
            exit(1);
        } else {
            info!("Crate successfully built!");
        }
    }

    // database operations
    if let Some(matches) = matches.subcommand_matches("database") {
        if let Some(_) = matches.subcommand_matches("init") {
            let conn = db::connect_db().unwrap();
            db::create_tables(&conn);
        }
    }

}
<commit_msg>Use cratesfyi::BUILD_VERSION for version<commit_after>
extern crate cratesfyi;
extern crate clap;
#[macro_use]
extern crate log;
extern crate time;

use std::env;
use std::fs;
use std::process::{Command, exit};
use std::path::PathBuf;

use cratesfyi::docbuilder::{DocBuilder, DocBuilderError, command_result};
use cratesfyi::docbuilder::crte::Crate;
use cratesfyi::db;
use clap::{Arg, App, SubCommand};
use log::{LogLevel, LogLevelFilter, LogRecord, LogMetadata};
use time::now;

// I am using really simple Logger in this program
struct Logger;

impl log::Log for Logger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} {} - {}", now().to_utc().rfc3339(), record.level(), record.args());
        }
    }
}

fn update_crates_io_index(path: &PathBuf) -> Result<String, String> {
    info!("Updating crates.io-index");

    if path.exists() {
        let cwd = env::current_dir().unwrap();
        env::set_current_dir(path).unwrap();
        let res = command_result(Command::new("git").arg("pull").output().unwrap());
        env::set_current_dir(cwd).unwrap();
        res
    } else {
        command_result(Command::new("git")
            .arg("clone")
            .arg("https://github.com/rust-lang/crates.io-index.git")
            .arg(path.to_str().unwrap())
            .output().unwrap())
    }
}

// This will remove everything in CWD!!!
fn clean_build_dir() -> Result<(), DocBuilderError> {

    for file in try!(env::current_dir().unwrap()
                     .read_dir().map_err(DocBuilderError::RemoveBuildDir)) {
        let file = try!(file.map_err(DocBuilderError::RemoveBuildDir));
        let path = file.path();

        if path.file_name().unwrap() == ".cargo" ||
            path.file_name().unwrap() == ".crates.io-index" {
                continue;
            }

        if path.is_dir() {
            try!(fs::remove_dir_all(path).map_err(DocBuilderError::RemoveBuildDir));
        } else {
            try!(fs::remove_file(path).map_err(DocBuilderError::RemoveBuildDir));
        }
    }

    Ok(())
}

fn main() {
    log::set_logger(|max_log_level| {
        max_log_level.set(LogLevelFilter::Info);
        Box::new(Logger)
    }).unwrap();

    let matches = App::new("cratesfyi")
                      .version(cratesfyi::BUILD_VERSION)
                      .about("Crates for your info!")
                      .subcommand(SubCommand::with_name("build")
                                      .about("Builds documentation in a chroot environment")
                                      .arg(Arg::with_name("PREFIX")
                                               .short("P")
                                               .long("prefix")
                                               .takes_value(true))
                                      .arg(Arg::with_name("DESTINATION")
                                               .short("d")
                                               .long("destination")
                                               .help("Sets destination path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CHROOT_PATH")
                                               .short("c")
                                               .long("chroot")
                                               .help("Sets chroot path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CHROOT_USER")
                                               .short("u")
                                               .long("chroot-user")
                                               .help("Sets chroot user name")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CRATES_IO_INDEX_PATH")
                                               .long("crates-io-index-path")
                                               .help("Sets crates.io-index path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("LOGS_PATH")
                                               .long("logs-path")
                                               .help("Sets logs path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("SKIP_IF_EXISTS")
                                               .short("s")
                                               .long("skip")
                                               .help("Skips building documentation if \
                                                      documentation exists"))
                                      .arg(Arg::with_name("SKIP_IF_LOG_EXISTS")
                                               .long("skip-if-log-exists")
                                               .help("Skips building documentation if build \
                                                      log exists"))
                                      .arg(Arg::with_name("KEEP_BUILD_DIRECTORY")
                                               .short("-k")
                                               .long("keep-build-directory")
                                               .help("Keeps build directory after build."))
                                      .subcommand(SubCommand::with_name("download-sources")
                                                      .about("Downloads sources of all crates"))
                                      .subcommand(SubCommand::with_name("world")
                                                      .about("Builds documentation of every \
                                                              crate")
                                                      .arg(Arg::with_name("SKIP_OLDEST_VERSIONS")
                                                               .long("skip-oldest-versions")
                                                               .help("Skips trying to build \
                                                                      oldest versions \
                                                                      if build fails"))
                                                      .arg(Arg::with_name("BUILD_ONLY_LATEST_V\
                                                                           ERSION")
                                                               .long("build-only-latest-versio\
                                                                      n")
                                                               .help("Builds only latest \
                                                                      version of crate and \
                                                                      skips oldest versions")))
                                      .subcommand(SubCommand::with_name("crate")
                                                      .about("Builds documentation for a crate")
                                                      .arg(Arg::with_name("CRATE_NAME")
                                                               .index(1)
                                                               .required(true)
                                                               .help("Crate name"))
                                                      .arg(Arg::with_name("CRATE_VERSION")
                                                               .index(2)
                                                               .required(true)
                                                               .help("Version of crate"))))
                      .subcommand(SubCommand::with_name("build-doc")
                                      .about("Builds documentation in CWD")
                                      .arg(Arg::with_name("CRATES_IO_INDEX_PATH")
                                               .short("p")
                                               .long("crates-io-index-path")
                                               .help("Sets crates.io-index path")
                                               .takes_value(true))
                                      .arg(Arg::with_name("CLEAN")
                                               .short("c")
                                               .long("clean")
                                               .help("Clean build dir before building"))
                                      .arg(Arg::with_name("CRATE_NAME")
                                               .index(1)
                                               .required(true)
                                               .help("Crate name"))
                                      .arg(Arg::with_name("CRATE_VERSION")
                                               .index(2)
                                               .required(true)
                                               .help("Version of crate")))
                      .subcommand(SubCommand::with_name("database")
                                      .about("Database operations")
                                      .subcommand(SubCommand::with_name("init")
                                                      .about("Initialize database. Currently \
                                                             only creates tables in database.")))
                      .get_matches();

    // DocBuilder
    if let Some(matches) = matches.subcommand_matches("build") {
        let mut dbuilder = {
            if let Some(prefix) = matches.value_of("PREFIX") {
                DocBuilder::from_prefix(PathBuf::from(prefix))
            } else {
                DocBuilder::default()
            }
        };

        // set destination
        if let Some(destination) = matches.value_of("DESTINATION") {
            dbuilder.destination(PathBuf::from(destination));
        }

        // set chroot path
        if let Some(chroot_path) = matches.value_of("CHROOT_PATH") {
            dbuilder.destination(PathBuf::from(chroot_path));
        }

        // set chroot user name
        if let Some(chroot_user) = matches.value_of("CHROOT_USER") {
            dbuilder.chroot_user(chroot_user.to_string());
        }

        // set crates.io-index path
        if let Some(crates_io_index_path) = matches.value_of("CRATES_IO_INDEX_PATH") {
            dbuilder.crates_io_index_path(PathBuf::from(crates_io_index_path));
        }

        // set logs path
        if let Some(logs_path) = matches.value_of("LOGS_PATH") {
            dbuilder.logs_path(PathBuf::from(logs_path));
        }

        dbuilder.skip_if_exists(matches.is_present("SKIP_IF_EXISTS"));
        dbuilder.skip_if_log_exists(matches.is_present("SKIP_IF_LOG_EXISTS"));
        dbuilder.keep_build_directory(matches.is_present("KEEP_BUILD_DIRECTORY"));

        // check paths
        if let Err(e) = dbuilder.check_paths() {
            println!("{:?}\nUse --help to get more information", e);
            std::process::exit(1);
        }

        // build world
        if let Some(matches) = matches.subcommand_matches("world") {
            dbuilder.build_only_latest_version(matches.is_present("BUILD_ONLY_LATEST_VERSION"));
            dbuilder.skip_oldest_versions(matches.is_present("SKIP_OLDEST_VERSIONS"));
            if let Err(e) = dbuilder.build_doc_for_every_crate() {
                println!("Failed to build world: {:#?}", e);
            }
        } else if let Some(matches) = matches.subcommand_matches("crate") {
            // Safe to call unwrap here
            let crte_name = matches.value_of("CRATE_NAME").unwrap();
            let version = matches.value_of("CRATE_VERSION").unwrap();
            let crte = Crate::new(crte_name.to_string(), vec![version.to_string()]);

            if let Err(e) = dbuilder.build_doc_for_crate_version(&crte, 0) {
                match e {
                    DocBuilderError::SkipDocumentationExists => {
                        println!("Skipping {} documentation already exists",
                                 crte.canonical_name(0))
                    }
                    _ => {
                        println!("Failed to build documentation for {}: {:?}",
                                 crte.canonical_name(0),
                                 e)
                    }
                }
            }
        } else if let Some(_) = matches.subcommand_matches("download-sources") {
            if let Err(e) = dbuilder.download_sources() {
                println!("{:?}", e);
            }
        }
    }

    // build-doc
    else if let Some(matches) = matches.subcommand_matches("build-doc") {
        let mut docbuilder = DocBuilder::default();

        // set crates.io-index path
        let crates_io_index_path = {
            if let Some(crates_io_index_path_conf) = matches.value_of("CRATES_IO_INDEX_PATH") {
                PathBuf::from(crates_io_index_path_conf)
            } else {
                let mut path = env::home_dir().unwrap();
                path.push(".crates.io-index");
                path
            }
        };

        docbuilder.crates_io_index_path(PathBuf::from(&crates_io_index_path));

        // update crates.io-index path
        if let Err(e) = update_crates_io_index(&crates_io_index_path) {
            panic!("{}", e);
        }

        // crates.io-index required for single crate
        if let Err(e) = docbuilder.check_crates_io_index_path() {
            panic!("{:?}", e);
        }

        let crte_name = matches.value_of("CRATE_NAME").unwrap();
        let version = matches.value_of("CRATE_VERSION").unwrap();

        let crte = Crate::new(crte_name.to_string(), vec![version.to_string()]);

        if matches.is_present("CLEAN") {
            clean_build_dir().unwrap();
        }

        if let Err(e) = crte.build_crate_doc(0, &docbuilder) {
            error!("Failed to build crate\n{:?}", e);
            exit(1);
        } else {
            info!("Crate successfully built!");
        }
    }

    // database operations
    if let Some(matches) = matches.subcommand_matches("database") {
        if let Some(_) = matches.subcommand_matches("init") {
            let conn = db::connect_db().unwrap();
            db::create_tables(&conn);
        }
    }

}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate docopt;
#[macro_use]
extern crate log;
extern crate mimir;
extern crate mimirsbrunn;
extern crate rustc_serialize;

use mimir::rubber::Rubber;
use mimirsbrunn::osm_reader::admin::{administrative_regions, compute_admin_weight};
use mimirsbrunn::osm_reader::poi::{pois, compute_poi_weight, PoiConfig};
use mimirsbrunn::osm_reader::street::{streets, compute_street_weight};
use mimirsbrunn::osm_reader::parse_osm_pbf;
use mimirsbrunn::admin_geofinder::AdminGeoFinder;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_city_level: u32,
    flag_connection_string: String,
    flag_import_way: bool,
    flag_import_admin: bool,
    flag_import_poi: bool,
    flag_dataset: String,
    flag_poi_config: Option<String>,
}

static USAGE: &'static str = r#"Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] [--import-way] [--import-admin] [--import-poi] [--dataset=<dataset>] [--city-level=<level>] --level=<level> ...

Options:
    -h, --help                Show this message.
    -i, --input=<file>        OSM PBF file.
    -l, --level=<level>       Admin levels to keep.
    -C, --city-level=<level>  City level to  calculate weight, [default: 8]
    -w, --import-way          Import ways
    -a, --import-admin        Import admins
    -p, --import-poi          Import POIs
    -c, --connection-string=<connection-string>
                              Elasticsearch parameters, [default: http://localhost:9200/munin]
    -d, --dataset=<dataset>   Name of the dataset, [default: fr]
    -j, --poi-config=<json>   POI configuration
"#;

fn main() {
    mimir::logger_init().unwrap();
    let args: Args = docopt::Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let city_level = args.flag_city_level;
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    debug!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);

    info!("creating adminstrative regions");
    let admins = administrative_regions(&mut parsed_pbf, levels);
    let admins_geofinder = admins.iter().cloned().collect::<AdminGeoFinder>();
    {
        info!("Extracting streets from osm");
        let mut streets = streets(&mut parsed_pbf, &admins_geofinder, city_level);

        info!("computing city weight");
        compute_admin_weight(&mut streets);

        info!("computing street weight");
        compute_street_weight(&mut streets, city_level);

        if args.flag_import_way {
            info!("importing streets into Mimir");
            let nb_streets = rubber.index(&args.flag_dataset, streets.into_iter())
                .unwrap();
            info!("Nb of indexed street: {}", nb_streets);
        }
    }
    let nb_admins = rubber.index(&args.flag_dataset, admins.iter())
        .unwrap();
    info!("Nb of indexed admin: {}", nb_admins);

    if args.flag_import_poi {
        let matcher = match args.flag_poi_config {
            None => PoiConfig::default(),
            Some(filename) => {
                let path = std::path::Path::new(&filename);
                let r = std::fs::File::open(&path).unwrap();
                PoiConfig::from_reader(r).unwrap()
            }
        };
        info!("Extracting pois from osm");
        let mut pois = pois(&mut parsed_pbf, &matcher, &admins_geofinder, city_level);

        info!("computing poi weight");
        compute_poi_weight(&mut pois, city_level);

        info!("Importing pois into Mimir");
        let nb_pois = rubber.index(&args.flag_dataset, pois.iter())
            .unwrap();

        info!("Nb of indexed pois: {}", nb_pois);
    }
}
<commit_msg>simplify docopt<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate docopt;
#[macro_use]
extern crate log;
extern crate mimir;
extern crate mimirsbrunn;
extern crate rustc_serialize;

use mimir::rubber::Rubber;
use mimirsbrunn::osm_reader::admin::{administrative_regions, compute_admin_weight};
use mimirsbrunn::osm_reader::poi::{pois, compute_poi_weight, PoiConfig};
use mimirsbrunn::osm_reader::street::{streets, compute_street_weight};
use mimirsbrunn::osm_reader::parse_osm_pbf;
use mimirsbrunn::admin_geofinder::AdminGeoFinder;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_city_level: u32,
    flag_connection_string: String,
    flag_import_way: bool,
    flag_import_admin: bool,
    flag_import_poi: bool,
    flag_dataset: String,
    flag_poi_config: Option<String>,
}

static USAGE: &'static str = r#"Usage:
    osm2mimir --help
    osm2mimir [options] --input=<file> --level=<level> ...

Options:
    -h, --help                Show this message.
    -i, --input=<file>        OSM PBF file.
    -l, --level=<level>       Admin levels to keep.
    -C, --city-level=<level>  City level to  calculate weight, [default: 8]
    -w, --import-way          Import ways
    -a, --import-admin        Import admins
    -p, --import-poi          Import POIs
    -c, --connection-string=<connection-string>
                              Elasticsearch parameters, [default: http://localhost:9200/munin]
    -d, --dataset=<dataset>   Name of the dataset, [default: fr]
    -j, --poi-config=<json>   POI configuration
"#;

fn main() {
    mimir::logger_init().unwrap();
    let args: Args = docopt::Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let city_level = args.flag_city_level;
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    debug!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);

    info!("creating adminstrative regions");
    let admins = administrative_regions(&mut parsed_pbf, levels);
    let admins_geofinder = admins.iter().cloned().collect::<AdminGeoFinder>();
    {
        info!("Extracting streets from osm");
        let mut streets = streets(&mut parsed_pbf, &admins_geofinder, city_level);

        info!("computing city weight");
        compute_admin_weight(&mut streets);

        info!("computing street weight");
        compute_street_weight(&mut streets, city_level);

        if args.flag_import_way {
            info!("importing streets into Mimir");
            let nb_streets = rubber.index(&args.flag_dataset, streets.into_iter())
                .unwrap();
            info!("Nb of indexed street: {}", nb_streets);
        }
    }
    let nb_admins = rubber.index(&args.flag_dataset, admins.iter())
        .unwrap();
    info!("Nb of indexed admin: {}", nb_admins);

    if args.flag_import_poi {
        let matcher = match args.flag_poi_config {
            None => PoiConfig::default(),
            Some(filename) => {
                let path = std::path::Path::new(&filename);
                let r = std::fs::File::open(&path).unwrap();
                PoiConfig::from_reader(r).unwrap()
            }
        };
        info!("Extracting pois from osm");
        let mut pois = pois(&mut parsed_pbf, &matcher, &admins_geofinder, city_level);

        info!("computing poi weight");
        compute_poi_weight(&mut pois, city_level);

        info!("Importing pois into Mimir");
        let nb_pois = rubber.index(&args.flag_dataset, pois.iter())
            .unwrap();

        info!("Nb of indexed pois: {}", nb_pois);
    }
}
<|endoftext|>