<commit_before>///! Handle the selections of items
use crate::event::{Event, EventArg, EventHandler, UpdateScreen};
use crate::item::{Item, MatchedItem, MatchedRange};
use crate::orderedvec::OrderedVec;
use crate::theme::{ColorTheme, DEFAULT_THEME};
use crate::util::{reshape_string, LinePrinter};
use crate::SkimOptions;
use std::cmp::max;
use std::cmp::min;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tuikit::prelude::*;

pub struct Selection {
    items: OrderedVec<MatchedItem>, // all items
    selected: HashMap<(usize, usize), MatchedItem>,

    //
    // |>------ items[items.len()-1]
    // |
    // +======+ screen end
    // |
    // |>------ line_cursor, position from screen start
    // |
    // +======+ item_cursor, screen start
    // |
    // |>------ item[0]
    //
    item_cursor: usize, // the index of matched item currently highlighted.
    line_cursor: usize, // line No.
    hscroll_offset: usize,
    height: AtomicUsize,
    tabstop: usize,

    // Options
    multi_selection: bool,
    reverse: bool,
    no_hscroll: bool,
    theme: Arc<ColorTheme>,
}

impl Selection {
    pub fn new() -> Self {
        Selection {
            items: OrderedVec::new(),
            selected: HashMap::new(),
            item_cursor: 0,
            line_cursor: 0,
            hscroll_offset: 0,
            height: AtomicUsize::new(0),
            tabstop: 0,
            multi_selection: false,
            reverse: false,
            no_hscroll: false,
            theme: Arc::new(DEFAULT_THEME),
        }
    }

    pub fn with_options(options: &SkimOptions) -> Self {
        let mut selection = Self::new();
        selection.parse_options(options);
        selection
    }

    fn parse_options(&mut self, options: &SkimOptions) {
        if options.multi {
            self.multi_selection = true;
        }

        if options.layout.starts_with("reverse") {
            self.reverse = true;
        }

        if options.no_hscroll {
            self.no_hscroll = true;
        }

        if let Some(tabstop_str) = options.tabstop {
            let tabstop = tabstop_str.parse::<usize>().unwrap_or(8);
            self.tabstop = max(1, tabstop);
        }
    }

    pub fn theme(mut self, theme: Arc<ColorTheme>) -> Self {
        self.theme = theme;
        self
    }

    pub fn append_sorted_items(&mut self, items: Vec<MatchedItem>) {
        self.items.append_ordered(items);

        let height = self.height.load(Ordering::Relaxed);
        if self.items.len() <= self.line_cursor {
            // if not enough items, move cursor down
            self.line_cursor = max(min(self.items.len(), height), 1) - 1;
        }

        if self.items.len() <= self.line_cursor + self.item_cursor {
            // if not enough items, scroll the cursor a page down
            self.item_cursor = max(height, self.item_cursor) - height;
        }
    }

    pub fn clear(&mut self) {
        self.items.clear();
    }

    pub fn act_move_line_cursor(&mut self, diff: i32) {
        let diff = if self.reverse { -diff } else { diff };

        let mut line_cursor = self.line_cursor as i32;
        let mut item_cursor = self.item_cursor as i32;
        let item_len = self.items.len() as i32;

        let height = self.height.load(Ordering::Relaxed) as i32;

        line_cursor += diff;
        if line_cursor >= height {
            item_cursor += line_cursor - height + 1;
            item_cursor = max(0, min(item_cursor, item_len - height));
            line_cursor = min(height - 1, item_len - item_cursor);
        } else if line_cursor < 0 {
            item_cursor += line_cursor;
            item_cursor = max(item_cursor, 0);
            line_cursor = 0;
        } else {
            line_cursor = max(0, min(line_cursor, item_len - 1 - item_cursor));
        }

        self.item_cursor = item_cursor as usize;
        self.line_cursor = line_cursor as usize;
    }

    pub fn act_toggle(&mut self) {
        if !self.multi_selection || self.items.is_empty() {
            return;
        }

        let cursor = self.item_cursor + self.line_cursor;
        let current_item = self
            .items
            .get(cursor)
            .unwrap_or_else(|| panic!("model:act_toggle: failed to get item {}", cursor));
        let index = current_item.item.get_full_index();
        if !self.selected.contains_key(&index) {
            self.selected.insert(index, current_item.clone());
        } else {
            self.selected.remove(&index);
        }
    }

    pub fn act_toggle_all(&mut self) {
        for current_item in self.items.iter() {
            let index = current_item.item.get_full_index();
            if !self.selected.contains_key(&index) {
                self.selected.insert(index, current_item.clone());
            } else {
                self.selected.remove(&index);
            }
        }
    }

    pub fn act_select_all(&mut self) {
        for current_item in self.items.iter() {
            let index = current_item.item.get_full_index();
            self.selected.insert(index, current_item.clone());
        }
    }

    pub fn act_deselect_all(&mut self) {
        self.selected.clear();
    }

    pub fn act_scroll(&mut self, offset: i32) {
        let mut hscroll_offset = self.hscroll_offset as i32;
        hscroll_offset += offset;
        hscroll_offset = max(0, hscroll_offset);
        self.hscroll_offset = hscroll_offset as usize;
    }

    pub fn get_selected_items(&mut self) -> Vec<Arc<Item>> {
        // select the current one
        if !self.items.is_empty() {
            let cursor = self.item_cursor + self.line_cursor;
            let current_item = self
                .items
                .get(cursor)
                .unwrap_or_else(|| panic!("model:act_output: failed to get item {}", cursor));
            let index = current_item.item.get_full_index();
            self.selected.insert(index, current_item.clone());
        }

        let mut selected: Vec<Arc<Item>> = self.selected.values().map(|item| item.item.clone()).collect();

        selected.sort_by_key(|item| item.get_full_index());
        selected
    }

    pub fn get_current_item_idx(&self) -> usize {
        self.item_cursor + self.line_cursor
    }

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

    pub fn is_multi_selection(&self) -> bool {
        self.multi_selection
    }

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

    pub fn get_current_item(&self) -> Option<Arc<Item>> {
        let item_idx = self.get_current_item_idx();
        self.items.get(item_idx).map(|item| item.item.clone())
    }
}

impl EventHandler for Selection {
    fn accept_event(&self, event: Event) -> bool {
        use crate::event::Event::*;
        match event {
            EvActUp | EvActDown | EvActToggle | EvActToggleDown | EvActToggleUp | EvActToggleAll | EvActSelectAll
            | EvActDeselectAll | EvActPageDown | EvActPageUp | EvActScrollLeft | EvActScrollRight => true,
            _ => false,
        }
    }

    fn handle(&mut self, event: Event, arg: &EventArg) -> UpdateScreen {
        use crate::event::Event::*;
        match event {
            EvActUp => {
                self.act_move_line_cursor(1);
            }
            EvActDown => {
                self.act_move_line_cursor(-1);
            }
            EvActToggle => {
                self.act_toggle();
            }
            EvActToggleDown => {
                self.act_toggle();
                self.act_move_line_cursor(-1);
            }
            EvActToggleUp => {
                self.act_toggle();
                self.act_move_line_cursor(1);
            }
            EvActToggleAll => {
                self.act_toggle_all();
            }
            EvActSelectAll => {
                self.act_select_all();
            }
            EvActDeselectAll => {
                self.act_deselect_all();
            }
            EvActPageDown => {
                let height = 1 - (self.height.load(Ordering::Relaxed) as i32);
                self.act_move_line_cursor(height);
            }
            EvActPageUp => {
                let height = (self.height.load(Ordering::Relaxed) as i32) - 1;
                self.act_move_line_cursor(height);
            }
            EvActScrollLeft => {
                self.act_scroll(*arg.downcast_ref::<i32>().unwrap_or(&-1));
            }
            EvActScrollRight => {
                self.act_scroll(*arg.downcast_ref::<i32>().unwrap_or(&1));
            }
            _ => {}
        }
        UpdateScreen::Redraw
    }
}

impl Selection {
    fn draw_item(&self, canvas: &mut Canvas, row: usize, matched_item: &MatchedItem, is_current: bool) -> Result<()> {
        let (screen_width, screen_height) = canvas.size()?;

        // update item heights
        self.height.store(screen_height, Ordering::Relaxed);

        if screen_width < 3 {
            return Err("screen width is too small".into());
        }

        let index = matched_item.item.get_full_index();

        let default_attr = if is_current {
            self.theme.current()
        } else {
            self.theme.normal()
        };

        // print selection cursor
        if self.selected.contains_key(&index) {
            let _ = canvas.print_with_attr(row, 1, ">", default_attr.extend(self.theme.selected()));
        } else {
            let _ = canvas.print_with_attr(row, 1, " ", default_attr);
        }

        let item = &matched_item.item;
        let text = item.get_text();
        let (match_start_char, match_end_char) = match matched_item.matched_range {
            Some(MatchedRange::Chars(ref matched_indices)) => {
                if !matched_indices.is_empty() {
                    (matched_indices[0], matched_indices[matched_indices.len() - 1] + 1)
                } else {
                    (0, 0)
                }
            }
            Some(MatchedRange::ByteRange(match_start, match_end)) => {
                let match_start_char = text[..match_start].chars().count();
                let diff = text[match_start..match_end].chars().count();
                (match_start_char, match_start_char + diff)
            }
            None => (0, 0),
        };

        let container_width = screen_width - 2;
        let (shift, full_width) =
            reshape_string(&text, container_width, match_start_char, match_end_char, self.tabstop);

        let mut printer = LinePrinter::builder()
            .row(row)
            .col(2)
            .tabstop(self.tabstop)
            .container_width(container_width)
            .shift(if self.no_hscroll { 0 } else { shift })
            .text_width(full_width)
            .hscroll_offset(self.hscroll_offset)
            .build();

        // print out the original content
        if item.get_text_struct().is_some() && item.get_text_struct().as_ref().unwrap().has_attrs() {
            for (ch, attr) in item.get_text_struct().as_ref().unwrap().iter() {
                printer.print_char(canvas, ch, default_attr.extend(attr), false);
            }
        } else {
            for ch in item.get_text().chars() {
                printer.print_char(canvas, ch, default_attr, false);
            }
        }

        // print the highlighted content
        printer.reset();
        match matched_item.matched_range {
            Some(MatchedRange::Chars(ref matched_indices)) => {
                let mut matched_indices_iter = matched_indices.iter().peekable();

                for (ch_idx, ch) in text.chars().enumerate() {
                    match matched_indices_iter.peek() {
                        Some(&&match_idx) if ch_idx == match_idx => {
                            printer.print_char(canvas, ch, default_attr.extend(self.theme.matched()), false);
                            let _ = matched_indices_iter.next();
                        }
                        Some(_) | None => {
                            printer.print_char(canvas, ch, default_attr, true);
                        }
                    }
                }
            }

            Some(MatchedRange::ByteRange(start, end)) => {
                for (idx, ch) in text.char_indices() {
                    printer.print_char(
                        canvas,
                        ch,
                        default_attr.extend(self.theme.matched()),
                        !(idx >= start && idx < end),
                    );
                }
            }

            _ => {}
        }

        Ok(())
    }
}

impl Draw for Selection {
    fn draw(&self, canvas: &mut Canvas) -> Result<()> {
        let (_screen_width, screen_height) = canvas.size()?;
        canvas.clear()?;

        let item_idx_lower = self.item_cursor;
        let max_upper = self.item_cursor + screen_height;
        let item_idx_upper = min(max_upper, self.items.len());

        for item_idx in item_idx_lower..item_idx_upper {
            let line_cursor = item_idx - item_idx_lower;
            let line_no = if self.reverse {
                // top down
                line_cursor
            } else {
                // bottom up
                screen_height - 1 - line_cursor
            };

            // print the cursor label
            let label = if line_cursor == self.line_cursor { ">" } else { " " };
            let _next_col = canvas.print_with_attr(line_no, 0, label, self.theme.cursor()).unwrap();

            let item = self
                .items
                .get(item_idx)
                .unwrap_or_else(|| panic!("model:draw_items: failed to get item at {}", item_idx));

            let _ = self.draw_item(canvas, line_no, &item, line_cursor == self.line_cursor);
        }

        Ok(())
    }
}
<commit_msg>[selection] fix cursor after items changed<commit_after>///! Handle the selections of items
use crate::event::{Event, EventArg, EventHandler, UpdateScreen};
use crate::item::{Item, MatchedItem, MatchedRange};
use crate::orderedvec::OrderedVec;
use crate::theme::{ColorTheme, DEFAULT_THEME};
use crate::util::{reshape_string, LinePrinter};
use crate::SkimOptions;
use std::cmp::max;
use std::cmp::min;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tuikit::prelude::*;

pub struct Selection {
    items: OrderedVec<MatchedItem>, // all items
    selected: HashMap<(usize, usize), MatchedItem>,

    //
    // |>------ items[items.len()-1]
    // |
    // +======+ screen end
    // |
    // |>------ line_cursor, position from screen start
    // |
    // +======+ item_cursor, screen start
    // |
    // |>------ item[0]
    //
    item_cursor: usize, // the index of matched item currently highlighted.
    line_cursor: usize, // line No.
    hscroll_offset: usize,
    height: AtomicUsize,
    tabstop: usize,

    // Options
    multi_selection: bool,
    reverse: bool,
    no_hscroll: bool,
    theme: Arc<ColorTheme>,
}

impl Selection {
    pub fn new() -> Self {
        Selection {
            items: OrderedVec::new(),
            selected: HashMap::new(),
            item_cursor: 0,
            line_cursor: 0,
            hscroll_offset: 0,
            height: AtomicUsize::new(0),
            tabstop: 0,
            multi_selection: false,
            reverse: false,
            no_hscroll: false,
            theme: Arc::new(DEFAULT_THEME),
        }
    }

    pub fn with_options(options: &SkimOptions) -> Self {
        let mut selection = Self::new();
        selection.parse_options(options);
        selection
    }

    fn parse_options(&mut self, options: &SkimOptions) {
        if options.multi {
            self.multi_selection = true;
        }

        if options.layout.starts_with("reverse") {
            self.reverse = true;
        }

        if options.no_hscroll {
            self.no_hscroll = true;
        }

        if let Some(tabstop_str) = options.tabstop {
            let tabstop = tabstop_str.parse::<usize>().unwrap_or(8);
            self.tabstop = max(1, tabstop);
        }
    }

    pub fn theme(mut self, theme: Arc<ColorTheme>) -> Self {
        self.theme = theme;
        self
    }

    pub fn append_sorted_items(&mut self, items: Vec<MatchedItem>) {
        self.items.append_ordered(items);

        let height = self.height.load(Ordering::Relaxed);
        if self.items.len() <= self.line_cursor {
            // if not enough items, move cursor down
            self.line_cursor = max(min(self.items.len(), height), 1) - 1;
        }

        if self.items.len() <= self.line_cursor + self.item_cursor {
            // if not enough items, scroll the cursor a page down
            self.item_cursor = max(self.items.len(), height) - height;
        }
    }

    pub fn clear(&mut self) {
        self.items.clear();
    }

    pub fn act_move_line_cursor(&mut self, diff: i32) {
        let diff = if self.reverse { -diff } else { diff };

        let mut line_cursor = self.line_cursor as i32;
        let mut item_cursor = self.item_cursor as i32;
        let item_len = self.items.len() as i32;

        let height = self.height.load(Ordering::Relaxed) as i32;

        line_cursor += diff;
        if line_cursor >= height {
            item_cursor += line_cursor - height + 1;
            item_cursor = max(0, min(item_cursor, item_len - height));
            line_cursor = min(height - 1, item_len - item_cursor);
        } else if line_cursor < 0 {
            item_cursor += line_cursor;
            item_cursor = max(item_cursor, 0);
            line_cursor = 0;
        } else {
            line_cursor = max(0, min(line_cursor, item_len - 1 - item_cursor));
        }

        self.item_cursor = item_cursor as usize;
        self.line_cursor = line_cursor as usize;
    }

    pub fn act_toggle(&mut self) {
        if !self.multi_selection || self.items.is_empty() {
            return;
        }

        let cursor = self.item_cursor + self.line_cursor;
        let current_item = self
            .items
            .get(cursor)
            .unwrap_or_else(|| panic!("model:act_toggle: failed to get item {}", cursor));
        let index = current_item.item.get_full_index();
        if !self.selected.contains_key(&index) {
            self.selected.insert(index, current_item.clone());
        } else {
            self.selected.remove(&index);
        }
    }

    pub fn act_toggle_all(&mut self) {
        for current_item in self.items.iter() {
            let index = current_item.item.get_full_index();
            if !self.selected.contains_key(&index) {
                self.selected.insert(index, current_item.clone());
            } else {
                self.selected.remove(&index);
            }
        }
    }

    pub fn act_select_all(&mut self) {
        for current_item in self.items.iter() {
            let index = current_item.item.get_full_index();
            self.selected.insert(index, current_item.clone());
        }
    }

    pub fn act_deselect_all(&mut self) {
        self.selected.clear();
    }

    pub fn act_scroll(&mut self, offset: i32) {
        let mut hscroll_offset = self.hscroll_offset as i32;
        hscroll_offset += offset;
        hscroll_offset = max(0, hscroll_offset);
        self.hscroll_offset = hscroll_offset as usize;
    }

    pub fn get_selected_items(&mut self) -> Vec<Arc<Item>> {
        // select the current one
        if !self.items.is_empty() {
            let cursor = self.item_cursor + self.line_cursor;
            let current_item = self
                .items
                .get(cursor)
                .unwrap_or_else(|| panic!("model:act_output: failed to get item {}", cursor));
            let index = current_item.item.get_full_index();
            self.selected.insert(index, current_item.clone());
        }

        let mut selected: Vec<Arc<Item>> = self.selected.values().map(|item| item.item.clone()).collect();

        selected.sort_by_key(|item| item.get_full_index());
        selected
    }

    pub fn get_current_item_idx(&self) -> usize {
        self.item_cursor + self.line_cursor
    }

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

    pub fn is_multi_selection(&self) -> bool {
        self.multi_selection
    }

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

    pub fn get_current_item(&self) -> Option<Arc<Item>> {
        let item_idx = self.get_current_item_idx();
        self.items.get(item_idx).map(|item| item.item.clone())
    }
}

impl EventHandler for Selection {
    fn accept_event(&self, event: Event) -> bool {
        use crate::event::Event::*;
        match event {
            EvActUp | EvActDown | EvActToggle | EvActToggleDown | EvActToggleUp | EvActToggleAll | EvActSelectAll
            | EvActDeselectAll | EvActPageDown | EvActPageUp | EvActScrollLeft | EvActScrollRight => true,
            _ => false,
        }
    }

    fn handle(&mut self, event: Event, arg: &EventArg) -> UpdateScreen {
        use crate::event::Event::*;
        match event {
            EvActUp => {
                self.act_move_line_cursor(1);
            }
            EvActDown => {
                self.act_move_line_cursor(-1);
            }
            EvActToggle => {
                self.act_toggle();
            }
            EvActToggleDown => {
                self.act_toggle();
                self.act_move_line_cursor(-1);
            }
            EvActToggleUp => {
                self.act_toggle();
                self.act_move_line_cursor(1);
            }
            EvActToggleAll => {
                self.act_toggle_all();
            }
            EvActSelectAll => {
                self.act_select_all();
            }
            EvActDeselectAll => {
                self.act_deselect_all();
            }
            EvActPageDown => {
                let height = 1 - (self.height.load(Ordering::Relaxed) as i32);
                self.act_move_line_cursor(height);
            }
            EvActPageUp => {
                let height = (self.height.load(Ordering::Relaxed) as i32) - 1;
                self.act_move_line_cursor(height);
            }
            EvActScrollLeft => {
                self.act_scroll(*arg.downcast_ref::<i32>().unwrap_or(&-1));
            }
            EvActScrollRight => {
                self.act_scroll(*arg.downcast_ref::<i32>().unwrap_or(&1));
            }
            _ => {}
        }
        UpdateScreen::Redraw
    }
}

impl Selection {
    fn draw_item(&self, canvas: &mut Canvas, row: usize, matched_item: &MatchedItem, is_current: bool) -> Result<()> {
        let (screen_width, screen_height) = canvas.size()?;

        // update item heights
        self.height.store(screen_height, Ordering::Relaxed);

        if screen_width < 3 {
            return Err("screen width is too small".into());
        }

        let index = matched_item.item.get_full_index();

        let default_attr = if is_current {
            self.theme.current()
        } else {
            self.theme.normal()
        };

        // print selection cursor
        if self.selected.contains_key(&index) {
            let _ = canvas.print_with_attr(row, 1, ">", default_attr.extend(self.theme.selected()));
        } else {
            let _ = canvas.print_with_attr(row, 1, " ", default_attr);
        }

        let item = &matched_item.item;
        let text = item.get_text();
        let (match_start_char, match_end_char) = match matched_item.matched_range {
            Some(MatchedRange::Chars(ref matched_indices)) => {
                if !matched_indices.is_empty() {
                    (matched_indices[0], matched_indices[matched_indices.len() - 1] + 1)
                } else {
                    (0, 0)
                }
            }
            Some(MatchedRange::ByteRange(match_start, match_end)) => {
                let match_start_char = text[..match_start].chars().count();
                let diff = text[match_start..match_end].chars().count();
                (match_start_char, match_start_char + diff)
            }
            None => (0, 0),
        };

        let container_width = screen_width - 2;
        let (shift, full_width) =
            reshape_string(&text, container_width, match_start_char, match_end_char, self.tabstop);

        let mut printer = LinePrinter::builder()
            .row(row)
            .col(2)
            .tabstop(self.tabstop)
            .container_width(container_width)
            .shift(if self.no_hscroll { 0 } else { shift })
            .text_width(full_width)
            .hscroll_offset(self.hscroll_offset)
            .build();

        // print out the original content
        if item.get_text_struct().is_some() && item.get_text_struct().as_ref().unwrap().has_attrs() {
            for (ch, attr) in item.get_text_struct().as_ref().unwrap().iter() {
                printer.print_char(canvas, ch, default_attr.extend(attr), false);
            }
        } else {
            for ch in item.get_text().chars() {
                printer.print_char(canvas, ch, default_attr, false);
            }
        }

        // print the highlighted content
        printer.reset();
        match matched_item.matched_range {
            Some(MatchedRange::Chars(ref matched_indices)) => {
                let mut matched_indices_iter = matched_indices.iter().peekable();

                for (ch_idx, ch) in text.chars().enumerate() {
                    match matched_indices_iter.peek() {
                        Some(&&match_idx) if ch_idx == match_idx => {
                            printer.print_char(canvas, ch, default_attr.extend(self.theme.matched()), false);
                            let _ = matched_indices_iter.next();
                        }
                        Some(_) | None => {
                            printer.print_char(canvas, ch, default_attr, true);
                        }
                    }
                }
            }

            Some(MatchedRange::ByteRange(start, end)) => {
                for (idx, ch) in text.char_indices() {
                    printer.print_char(
                        canvas,
                        ch,
                        default_attr.extend(self.theme.matched()),
                        !(idx >= start && idx < end),
                    );
                }
            }

            _ => {}
        }

        Ok(())
    }
}

impl Draw for Selection {
    fn draw(&self, canvas: &mut Canvas) -> Result<()> {
        let (_screen_width, screen_height) = canvas.size()?;
        canvas.clear()?;

        let item_idx_lower = self.item_cursor;
        let max_upper = self.item_cursor + screen_height;
        let item_idx_upper = min(max_upper, self.items.len());

        for item_idx in item_idx_lower..item_idx_upper {
            let line_cursor = item_idx - item_idx_lower;
            let line_no = if self.reverse {
                // top down
                line_cursor
            } else {
                // bottom up
                screen_height - 1 - line_cursor
            };

            // print the cursor label
            let label = if line_cursor == self.line_cursor { ">" } else { " " };
            let _next_col = canvas.print_with_attr(line_no, 0, label, self.theme.cursor()).unwrap();

            let item = self
                .items
                .get(item_idx)
                .unwrap_or_else(|| panic!("model:draw_items: failed to get item at {}", item_idx));

            let _ = self.draw_item(canvas, line_no, &item, line_cursor == self.line_cursor);
        }

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


<commit_before>use std::collections::HashSet;
use std::fmt;
use std::ptr;
use std::sync::Arc;

use crate::cpu::flush_icache;
use crate::dseg::DSeg;
use crate::gc::Address;
use crate::os;
use crate::ty::SourceTypeArray;
use crate::utils::GrowableVec;
use crate::vm::FctDefinitionId;
use crate::vm::VM;

use dora_parser::Position;

#[derive(Debug, Clone)]
pub enum CodeDescriptor {
    DoraFct(FctDefinitionId),
    CompileStub,
    TrapStub,
    AllocStub,
    VerifyStub,
    NativeStub(FctDefinitionId),
    DoraStub,
    GuardCheckStub,
    SafepointStub,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct CodeId(usize);

impl CodeId {
    pub fn from(idx: usize) -> CodeId {
        CodeId(idx)
    }

    pub fn idx(self) -> usize {
        self.0
    }
}

impl From<usize> for CodeId {
    fn from(data: usize) -> CodeId {
        CodeId(data)
    }
}

impl GrowableVec<Code> {
    pub fn idx(&self, index: CodeId) -> Arc<Code> {
        self.idx_usize(index.0)
    }
}

pub struct Code {
    code_start: Address,
    code_end: Address,

    // pointer to beginning of function
    function_entry: Address,

    desc: CodeDescriptor,

    framesize: i32,
    lazy_compilation: LazyCompilationData,
    gcpoints: GcPoints,
    comments: Comments,
    positions: PositionTable,
}

impl Code {
    pub fn from_optimized_buffer(vm: &VM, buffer: &[u8], desc: CodeDescriptor) -> Code {
        let dseg = DSeg::new();

        Code::from_buffer(
            vm,
            &dseg,
            buffer,
            LazyCompilationData::new(),
            GcPoints::new(),
            0,
            Comments::new(),
            PositionTable::new(),
            desc,
        )
    }

    pub fn from_buffer(
        vm: &VM,
        dseg: &DSeg,
        buffer: &[u8],
        lazy_compilation: LazyCompilationData,
        gcpoints: GcPoints,
        framesize: i32,
        comments: Comments,
        positions: PositionTable,
        desc: CodeDescriptor,
    ) -> Code {
        let code_space_size = dseg.size() as usize + buffer.len();
        let code_start = vm.gc.alloc_code(code_space_size);
        let code_end = code_start.offset(code_space_size);

        if code_start.is_null() {
            panic!("out of memory: not enough executable memory left!");
        }

        os::jit_writable();

        dseg.finish(code_start.to_ptr());

        let function_entry = code_start.offset(dseg.size() as usize);

        unsafe {
            ptr::copy_nonoverlapping(buffer.as_ptr(), function_entry.to_mut_ptr(), buffer.len());
        }

        os::jit_executable();

        flush_icache(code_start.to_ptr(), code_space_size);

        Code {
            code_start,
            code_end,
            function_entry,
            desc,
            lazy_compilation,
            gcpoints,
            comments,
            framesize,
            positions,
        }
    }

    pub fn position_for_offset(&self, offset: u32) -> Option<Position> {
        self.positions.get(offset)
    }

    pub fn gcpoint_for_offset(&self, offset: u32) -> Option<&GcPoint> {
        self.gcpoints.get(offset)
    }

    pub fn ptr_start(&self) -> Address {
        self.code_start
    }

    pub fn ptr_end(&self) -> Address {
        self.code_end
    }

    pub fn fct_id(&self) -> FctDefinitionId {
        match self.desc {
            CodeDescriptor::NativeStub(fct_id) => fct_id,
            CodeDescriptor::DoraFct(fct_id) => fct_id,
            _ => panic!("no fctid found"),
        }
    }

    pub fn instruction_start(&self) -> Address {
        self.function_entry
    }

    pub fn instruction_end(&self) -> Address {
        self.code_end
    }

    pub fn framesize(&self) -> i32 {
        self.framesize
    }

    pub fn comment_for_offset(&self, offset: u32) -> Option<&String> {
        self.comments.get(offset)
    }

    pub fn lazy_for_offset(&self, offset: u32) -> Option<&LazyCompilationSite> {
        self.lazy_compilation.get(offset)
    }

    pub fn descriptor(&self) -> CodeDescriptor {
        self.desc.clone()
    }
}

impl fmt::Debug for Code {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Code {{ start: {:?}, end: {:?}, desc: {:?} }}",
            self.ptr_start(),
            self.ptr_end(),
            self.desc,
        )
    }
}

#[derive(Debug)]
pub struct GcPoints {
    entries: Vec<(u32, GcPoint)>,
}

impl GcPoints {
    pub fn new() -> GcPoints {
        GcPoints {
            entries: Vec::new(),
        }
    }

    pub fn get(&self, offset: u32) -> Option<&GcPoint> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(_) => None,
        }
    }

    pub fn insert(&mut self, offset: u32, gcpoint: GcPoint) {
        if let Some(last) = self.entries.last_mut() {
            debug_assert!(offset > last.0);
        }

        self.entries.push((offset, gcpoint));
    }
}

#[derive(Debug)]
pub struct GcPoint {
    pub offsets: Vec<i32>,
}

impl GcPoint {
    pub fn new() -> GcPoint {
        GcPoint {
            offsets: Vec::new(),
        }
    }

    pub fn merge(lhs: GcPoint, rhs: GcPoint) -> GcPoint {
        let mut offsets = HashSet::new();

        for offset in lhs.offsets {
            offsets.insert(offset);
        }

        for offset in rhs.offsets {
            offsets.insert(offset);
        }

        GcPoint::from_offsets(offsets.drain().collect())
    }

    pub fn from_offsets(offsets: Vec<i32>) -> GcPoint {
        GcPoint { offsets }
    }
}

pub struct Comments {
    entries: Vec<(u32, String)>,
}

impl Comments {
    pub fn new() -> Comments {
        Comments {
            entries: Vec::new(),
        }
    }

    pub fn get(&self, offset: u32) -> Option<&String> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(_) => None,
        }
    }

    pub fn insert(&mut self, offset: u32, comment: String) {
        if let Some(last) = self.entries.last_mut() {
            debug_assert!(offset >= last.0);
        }

        self.entries.push((offset, comment));
    }
}

#[derive(Debug)]
pub struct PositionTable {
    entries: Vec<(u32, Position)>,
}

impl PositionTable {
    pub fn new() -> PositionTable {
        PositionTable {
            entries: Vec::new(),
        }
    }

    pub fn insert(&mut self, offset: u32, position: Position) {
        if let Some(last) = self.entries.last() {
            debug_assert!(offset > last.0);
        }

        self.entries.push((offset, position));
    }

    pub fn get(&self, offset: u32) -> Option<Position> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(self.entries[idx].1),
            Err(_) => None,
        }
    }
}

#[derive(Debug)]
pub struct LazyCompilationData {
    entries: Vec<(u32, LazyCompilationSite)>,
}

impl LazyCompilationData {
    pub fn new() -> LazyCompilationData {
        LazyCompilationData {
            entries: Vec::new(),
        }
    }

    pub fn insert(&mut self, offset: u32, info: LazyCompilationSite) {
        if let Some(last) = self.entries.last() {
            debug_assert!(offset > last.0);
        }

        self.entries.push((offset, info));
    }

    pub fn get(&self, offset: u32) -> Option<&LazyCompilationSite> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(_) => None,
        }
    }
}

#[derive(Clone, Debug)]
pub enum LazyCompilationSite {
    Direct(FctDefinitionId, i32, SourceTypeArray),
    Virtual(bool, FctDefinitionId, u32, SourceTypeArray),
}
<commit_msg>code: save memory for duplicate GcPoints<commit_after>use std::collections::HashSet;
use std::fmt;
use std::ptr;
use std::sync::Arc;

use crate::cpu::flush_icache;
use crate::dseg::DSeg;
use crate::gc::Address;
use crate::os;
use crate::ty::SourceTypeArray;
use crate::utils::GrowableVec;
use crate::vm::FctDefinitionId;
use crate::vm::VM;

use dora_parser::Position;

#[derive(Debug, Clone)]
pub enum CodeDescriptor {
    DoraFct(FctDefinitionId),
    CompileStub,
    TrapStub,
    AllocStub,
    VerifyStub,
    NativeStub(FctDefinitionId),
    DoraStub,
    GuardCheckStub,
    SafepointStub,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct CodeId(usize);

impl CodeId {
    pub fn from(idx: usize) -> CodeId {
        CodeId(idx)
    }

    pub fn idx(self) -> usize {
        self.0
    }
}

impl From<usize> for CodeId {
    fn from(data: usize) -> CodeId {
        CodeId(data)
    }
}

impl GrowableVec<Code> {
    pub fn idx(&self, index: CodeId) -> Arc<Code> {
        self.idx_usize(index.0)
    }
}

pub struct Code {
    code_start: Address,
    code_end: Address,

    // pointer to beginning of function
    function_entry: Address,

    desc: CodeDescriptor,

    framesize: i32,
    lazy_compilation: LazyCompilationData,
    gcpoints: GcPoints,
    comments: Comments,
    positions: PositionTable,
}

impl Code {
    pub fn from_optimized_buffer(vm: &VM, buffer: &[u8], desc: CodeDescriptor) -> Code {
        let dseg = DSeg::new();

        Code::from_buffer(
            vm,
            &dseg,
            buffer,
            LazyCompilationData::new(),
            GcPoints::new(),
            0,
            Comments::new(),
            PositionTable::new(),
            desc,
        )
    }

    pub fn from_buffer(
        vm: &VM,
        dseg: &DSeg,
        buffer: &[u8],
        lazy_compilation: LazyCompilationData,
        gcpoints: GcPoints,
        framesize: i32,
        comments: Comments,
        positions: PositionTable,
        desc: CodeDescriptor,
    ) -> Code {
        let code_space_size = dseg.size() as usize + buffer.len();
        let code_start = vm.gc.alloc_code(code_space_size);
        let code_end = code_start.offset(code_space_size);

        if code_start.is_null() {
            panic!("out of memory: not enough executable memory left!");
        }

        os::jit_writable();

        dseg.finish(code_start.to_ptr());

        let function_entry = code_start.offset(dseg.size() as usize);

        unsafe {
            ptr::copy_nonoverlapping(buffer.as_ptr(), function_entry.to_mut_ptr(), buffer.len());
        }

        os::jit_executable();

        flush_icache(code_start.to_ptr(), code_space_size);

        Code {
            code_start,
            code_end,
            function_entry,
            desc,
            lazy_compilation,
            gcpoints,
            comments,
            framesize,
            positions,
        }
    }

    pub fn position_for_offset(&self, offset: u32) -> Option<Position> {
        self.positions.get(offset)
    }

    pub fn gcpoint_for_offset(&self, offset: u32) -> Option<&GcPoint> {
        self.gcpoints.get(offset)
    }

    pub fn ptr_start(&self) -> Address {
        self.code_start
    }

    pub fn ptr_end(&self) -> Address {
        self.code_end
    }

    pub fn fct_id(&self) -> FctDefinitionId {
        match self.desc {
            CodeDescriptor::NativeStub(fct_id) => fct_id,
            CodeDescriptor::DoraFct(fct_id) => fct_id,
            _ => panic!("no fctid found"),
        }
    }

    pub fn instruction_start(&self) -> Address {
        self.function_entry
    }

    pub fn instruction_end(&self) -> Address {
        self.code_end
    }

    pub fn framesize(&self) -> i32 {
        self.framesize
    }

    pub fn comment_for_offset(&self, offset: u32) -> Option<&String> {
        self.comments.get(offset)
    }

    pub fn lazy_for_offset(&self, offset: u32) -> Option<&LazyCompilationSite> {
        self.lazy_compilation.get(offset)
    }

    pub fn descriptor(&self) -> CodeDescriptor {
        self.desc.clone()
    }
}

impl fmt::Debug for Code {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Code {{ start: {:?}, end: {:?}, desc: {:?} }}",
            self.ptr_start(),
            self.ptr_end(),
            self.desc,
        )
    }
}

#[derive(Debug)]
pub struct GcPoints {
    entries: Vec<(u32, GcPoint)>,
}

impl GcPoints {
    pub fn new() -> GcPoints {
        GcPoints {
            entries: Vec::new(),
        }
    }

    pub fn get(&self, offset: u32) -> Option<&GcPoint> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(_) => None,
        }
    }

    pub fn insert(&mut self, offset: u32, gcpoint: GcPoint) {
        if let Some((last_offset, last_gcpoint)) = self.entries.last() {
            debug_assert!(offset > *last_offset);

            if last_gcpoint == &gcpoint {
                // The last gcpoint stores exactly the same offsets, so reuse that one.
                return;
            }
        }

        self.entries.push((offset, gcpoint));
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct GcPoint {
    pub offsets: Vec<i32>,
}

impl GcPoint {
    pub fn new() -> GcPoint {
        GcPoint {
            offsets: Vec::new(),
        }
    }

    pub fn merge(lhs: GcPoint, rhs: GcPoint) -> GcPoint {
        let mut offsets = HashSet::new();

        for offset in lhs.offsets {
            offsets.insert(offset);
        }

        for offset in rhs.offsets {
            offsets.insert(offset);
        }

        GcPoint::from_offsets(offsets.drain().collect())
    }

    pub fn from_offsets(offsets: Vec<i32>) -> GcPoint {
        GcPoint { offsets }
    }
}

pub struct Comments {
    entries: Vec<(u32, String)>,
}

impl Comments {
    pub fn new() -> Comments {
        Comments {
            entries: Vec::new(),
        }
    }

    pub fn get(&self, offset: u32) -> Option<&String> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(_) => None,
        }
    }

    pub fn insert(&mut self, offset: u32, comment: String) {
        if let Some(last) = self.entries.last_mut() {
            debug_assert!(offset >= last.0);
        }

        self.entries.push((offset, comment));
    }
}

#[derive(Debug)]
pub struct PositionTable {
    entries: Vec<(u32, Position)>,
}

impl PositionTable {
    pub fn new() -> PositionTable {
        PositionTable {
            entries: Vec::new(),
        }
    }

    pub fn insert(&mut self, offset: u32, position: Position) {
        if let Some(last) = self.entries.last() {
            debug_assert!(offset > last.0);
        }

        self.entries.push((offset, position));
    }

    pub fn get(&self, offset: u32) -> Option<Position> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(self.entries[idx].1),
            Err(_) => None,
        }
    }
}

#[derive(Debug)]
pub struct LazyCompilationData {
    entries: Vec<(u32, LazyCompilationSite)>,
}

impl LazyCompilationData {
    pub fn new() -> LazyCompilationData {
        LazyCompilationData {
            entries: Vec::new(),
        }
    }

    pub fn insert(&mut self, offset: u32, info: LazyCompilationSite) {
        if let Some(last) = self.entries.last() {
            debug_assert!(offset > last.0);
        }

        self.entries.push((offset, info));
    }

    pub fn get(&self, offset: u32) -> Option<&LazyCompilationSite> {
        let result = self
            .entries
            .binary_search_by_key(&offset, |&(offset, _)| offset);

        match result {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(_) => None,
        }
    }
}

#[derive(Clone, Debug)]
pub enum LazyCompilationSite {
    Direct(FctDefinitionId, i32, SourceTypeArray),
    Virtual(bool, FctDefinitionId, u32, SourceTypeArray),
}
<|endoftext|>


<commit_before>use gmp::mpq::Mpq;
use std::borrow::Borrow;
use std::collections::VecDeque;
use std::vec::Vec;

struct BallotState {
    prev: usize,
    edge_flow: Box<[Mpq]>,
    count: i32,
}

#[derive(Clone)]
struct CandidateState {
    prev: usize,
    count: i32,
}

pub fn strength<Ballot>(num_seats: usize, ballots: &[(Ballot, Mpq)]) -> Mpq
    where Ballot: Borrow<[usize]>
{
    let ballot_candidates = &mut vec![Vec::new(); ballots.len()][..];
    let candidate_ballots = &mut vec![Vec::new(); num_seats][..];
    for (b, &(ref cs, _)) in ballots.iter().enumerate() {
        ballot_candidates[b].push((!0, !0));
        for (i, &c) in cs.borrow().iter().enumerate() {
            ballot_candidates[b].push((c, candidate_ballots[c].len()));
            candidate_ballots[c].push((b, i + 1));
        }
    }
    let ballot_states = &mut ballots.iter()
        .map(|&(ref cs, ref w)| {
            let mut edge_flow = vec![Mpq::zero(); 1 + cs.borrow().len()].into_boxed_slice();
            edge_flow[0] = w.clone();
            BallotState {
                prev: !0,
                edge_flow: edge_flow,
                count: 0,
            }
        })
        .collect::<Vec<_>>()[..];
    let candidate_states = &mut vec![CandidateState {
        prev: !0,
        count: 0,
    }; num_seats][..];

    let mut total_flow = Mpq::zero();
    let mut queue = VecDeque::new();

    loop {
        for (b, bs) in ballot_states.iter_mut().enumerate() {
            if !bs.edge_flow[0].is_zero() {
                bs.prev = 0;
                queue.push_back(b);
            }
        }

        let mut found = Vec::new();
        'search: loop {
            match queue.pop_front() {
                None => {
                    return total_flow;
                }
                Some(b) => {
                    for &(c, j) in ballot_candidates[b].iter().skip(1) {
                        if candidate_states[c].prev != !0 {
                            continue;
                        }
                        candidate_states[c].prev = j;
                        found.push(c);
                        if found.len() == num_seats {
                            break 'search;
                        }
                        for &(b1, i1) in &candidate_ballots[c] {
                            if ballot_states[b1].edge_flow[i1].is_zero() ||
                               ballot_states[b1].prev != !0 {
                                continue;
                            }
                            ballot_states[b1].prev = i1;
                            queue.push_back(b1);
                        }
                    }
                }
            }
        }

        let mut sunk = 0;
        for &c in found.iter().rev() {
            let j = candidate_states[c].prev;
            let (b, _) = candidate_ballots[c][j];
            let count = candidate_states[c].count + 1;
            ballot_states[b].count += count;
            let i1 = ballot_states[b].prev;
            if i1 == 0 {
                sunk += count;
            } else {
                let (c1, _) = ballot_candidates[b][i1];
                candidate_states[c1].count += count;
            }
        }
        debug_assert_eq!(sunk, num_seats as i32);

        let flow = found.iter()
            .map(|&c| {
                let j = candidate_states[c].prev;
                let (b, _) = candidate_ballots[c][j];
                &ballot_states[b].edge_flow[ballot_states[b].prev] /
                Mpq::from(ballot_states[b].count as i64)
            })
            .min()
            .unwrap();
        debug_assert!(!flow.is_zero());
        total_flow = total_flow + &flow;

        for c in found {
            let j = candidate_states[c].prev;
            let (b, i) = candidate_ballots[c][j];
            {
                let edge_flow = &mut ballot_states[b].edge_flow[i];
                *edge_flow = &*edge_flow + &flow * Mpq::from(candidate_states[c].count as i64 + 1);
            }
            if ballot_states[b].count != 0 {
                let edge_flow = &mut ballot_states[b].edge_flow[ballot_states[b].prev];
                *edge_flow = &*edge_flow - &flow * Mpq::from(ballot_states[b].count as i64);
                ballot_states[b].count = 0;
            }

            candidate_states[c].prev = !0;
            candidate_states[c].count = 0;
        }

        for bs in &mut *ballot_states {
            bs.prev = !0;
            debug_assert_eq!(bs.count, 0);
        }
        queue.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::strength;
    use gmp::mpq::Mpq;

    const Q: fn(i64) -> Mpq = Mpq::from;

    #[test]
    fn test_strength_1() {
        // Wikipedia
        let ballots: &[(&[usize], Mpq)] =
            &[(&[0], Q(12)), (&[1], Q(0)), (&[0, 1], Q(51)), (&[], Q(27))];
        assert_eq!(strength(2, ballots), Q(63) / Q(2));
    }

    #[test]
    fn test_strength_2() {
        // Wikipedia
        let ballots: &[(&[usize], Mpq)] = &[(&[0], Q(38)), (&[1], Q(27)), (&[0, 1], Q(12))];
        assert_eq!(strength(2, ballots), Q(77) / Q(2));
    }

    #[test]
    fn test_strength_3() {
        // Schulze’s calcul02.pdf
        let ballots: &[(&[usize], Mpq)] = &[(&[0, 1, 2, 3], Q(36_597383) / Q(1_000000)),
                                            (&[0, 1, 2], Q(5_481150) / Q(1_000000)),
                                            (&[0, 1, 3], Q(13_279131) / Q(1_000000)),
                                            (&[0, 1], Q(4_859413) / Q(1_000000)),
                                            (&[0, 2, 3], Q(35_425375) / Q(1_000000)),
                                            (&[0, 2], Q(5_490934) / Q(1_000000)),
                                            (&[0, 3], Q(22_855333) / Q(1_000000)),
                                            (&[0], Q(19_835570) / Q(1_000000)),
                                            (&[1, 2, 3], Q(22_928716) / Q(1_000000)),
                                            (&[1, 2], Q(5_538309) / Q(1_000000)),
                                            (&[1, 3], Q(13_130227) / Q(1_000000)),
                                            (&[1], Q(6_056291) / Q(1_000000)),
                                            (&[2, 3], Q(23_992772) / Q(1_000000)),
                                            (&[2], Q(16_699207) / Q(1_000000)),
                                            (&[3], Q(98_165759) / Q(1_000000)),
                                            (&[], Q(129_664430) / Q(1_000000))];
        assert_eq!(strength(4, ballots), Q(77_389937) / Q(1_000000));
    }
}
<commit_msg>vote-management: Replace Edmonds–Karp algorithm with Dinic’s algorithm<commit_after>use gmp::mpq::Mpq;
use std::borrow::Borrow;
use std::collections::VecDeque;
use std::vec::Vec;

struct BallotState {
    level: usize,
    prev: usize,
    edge_flow: Box<[Mpq]>,
    count: i32,
}

#[derive(Clone)]
struct CandidateState {
    level: usize,
    prev: usize,
    count: i32,
}

pub fn strength<Ballot>(num_seats: usize, ballots: &[(Ballot, Mpq)]) -> Mpq
    where Ballot: Borrow<[usize]>
{
    let ballot_candidates = &mut vec![Vec::new(); ballots.len()][..];
    let candidate_ballots = &mut vec![Vec::new(); num_seats][..];
    for (b, &(ref cs, _)) in ballots.iter().enumerate() {
        ballot_candidates[b].push((!0, !0));
        for (i, &c) in cs.borrow().iter().enumerate() {
            ballot_candidates[b].push((c, candidate_ballots[c].len()));
            candidate_ballots[c].push((b, i + 1));
        }
    }
    let ballot_states = &mut ballots.iter()
        .map(|&(ref cs, ref w)| {
            let mut edge_flow = vec![Mpq::zero(); 1 + cs.borrow().len()].into_boxed_slice();
            edge_flow[0] = w.clone();
            BallotState {
                level: !0,
                prev: !0,
                edge_flow: edge_flow,
                count: 0,
            }
        })
        .collect::<Vec<_>>()[..];
    let candidate_states = &mut vec![CandidateState {
        level: !0,
        prev: !0,
        count: 0,
    }; num_seats][..];

    let mut total_flow = Mpq::zero();
    let mut queue = VecDeque::new();

    loop {
        for (b, bs) in ballot_states.iter_mut().enumerate() {
            if !bs.edge_flow[0].is_zero() {
                bs.level = 0;
                bs.prev = 0;
                queue.push_back(b);
            }
        }

        let mut found = Vec::new();
        'search: loop {
            match queue.pop_front() {
                None => {
                    return total_flow;
                }
                Some(b) => {
                    let level = ballot_states[b].level;
                    for &(c, j) in ballot_candidates[b].iter().skip(1) {
                        if candidate_states[c].prev != !0 {
                            continue;
                        }
                        candidate_states[c].level = level + 1;
                        candidate_states[c].prev = j;
                        found.push(c);
                        if found.len() == num_seats {
                            break 'search;
                        }
                        for &(b1, i1) in &candidate_ballots[c] {
                            if ballot_states[b1].edge_flow[i1].is_zero() ||
                               ballot_states[b1].prev != !0 {
                                continue;
                            }
                            ballot_states[b1].level = level + 2;
                            ballot_states[b1].prev = i1;
                            queue.push_back(b1);
                        }
                    }
                }
            }
        }

        'iterate: loop {
            let mut sunk = 0;
            for &c in found.iter().rev() {
                let j = candidate_states[c].prev;
                let (b, _) = candidate_ballots[c][j];
                let count = candidate_states[c].count + 1;
                ballot_states[b].count += count;
                let i1 = ballot_states[b].prev;
                if i1 == 0 {
                    sunk += count;
                } else {
                    let (c1, _) = ballot_candidates[b][i1];
                    candidate_states[c1].count += count;
                }
            }
            debug_assert_eq!(sunk, num_seats as i32);

            let flow = found.iter()
                .map(|&c| {
                    let j = candidate_states[c].prev;
                    let (b, _) = candidate_ballots[c][j];
                    &ballot_states[b].edge_flow[ballot_states[b].prev] /
                    Mpq::from(ballot_states[b].count as i64)
                })
                .min()
                .unwrap();
            debug_assert!(!flow.is_zero());
            total_flow = total_flow + &flow;

            for &c in &found {
                let j = candidate_states[c].prev;
                let (b, i) = candidate_ballots[c][j];
                {
                    let edge_flow = &mut ballot_states[b].edge_flow[i];
                    *edge_flow = &*edge_flow +
                                 &flow * Mpq::from(candidate_states[c].count as i64 + 1);
                }
                if ballot_states[b].count != 0 {
                    let edge_flow = &mut ballot_states[b].edge_flow[ballot_states[b].prev];
                    *edge_flow = &*edge_flow - &flow * Mpq::from(ballot_states[b].count as i64);
                    ballot_states[b].count = 0;
                }

                candidate_states[c].count = 0;
            }

            for &c in &found {
                candidate_states[c].prev = (candidate_states[c].prev..candidate_ballots[c].len())
                    .find(|&j| {
                        let (b, _) = candidate_ballots[c][j];
                        let i1 = ballot_states[b].prev;
                        if i1 != !0 {
                            ballot_states[b].prev = (i1..ballot_candidates[b].len())
                                .find(|&i1| {
                                    let (c1, _) = ballot_candidates[b][i1];
                                    !ballot_states[b].edge_flow[i1].is_zero() &&
                                    (i1 == 0 || candidate_states[c1].level < ballot_states[b].level)
                                })
                                .unwrap_or(!0);
                            ballot_states[b].prev != !0 &&
                            ballot_states[b].level < candidate_states[c].level
                        } else {
                            false
                        }
                    })
                    .unwrap_or(!0);

                if candidate_states[c].prev == !0 {
                    break 'iterate;
                }
            }
        }

        for cs in &mut *candidate_states {
            cs.level = !0;
            cs.prev = !0;
            debug_assert_eq!(cs.count, 0);
        }

        for bs in &mut *ballot_states {
            bs.level = !0;
            bs.prev = !0;
            debug_assert_eq!(bs.count, 0);
        }
        queue.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::strength;
    use gmp::mpq::Mpq;

    const Q: fn(i64) -> Mpq = Mpq::from;

    #[test]
    fn test_strength_1() {
        // Wikipedia
        let ballots: &[(&[usize], Mpq)] =
            &[(&[0], Q(12)), (&[1], Q(0)), (&[0, 1], Q(51)), (&[], Q(27))];
        assert_eq!(strength(2, ballots), Q(63) / Q(2));
    }

    #[test]
    fn test_strength_2() {
        // Wikipedia
        let ballots: &[(&[usize], Mpq)] = &[(&[0], Q(38)), (&[1], Q(27)), (&[0, 1], Q(12))];
        assert_eq!(strength(2, ballots), Q(77) / Q(2));
    }

    #[test]
    fn test_strength_3() {
        // Schulze’s calcul02.pdf
        let ballots: &[(&[usize], Mpq)] = &[(&[0, 1, 2, 3], Q(36_597383) / Q(1_000000)),
                                            (&[0, 1, 2], Q(5_481150) / Q(1_000000)),
                                            (&[0, 1, 3], Q(13_279131) / Q(1_000000)),
                                            (&[0, 1], Q(4_859413) / Q(1_000000)),
                                            (&[0, 2, 3], Q(35_425375) / Q(1_000000)),
                                            (&[0, 2], Q(5_490934) / Q(1_000000)),
                                            (&[0, 3], Q(22_855333) / Q(1_000000)),
                                            (&[0], Q(19_835570) / Q(1_000000)),
                                            (&[1, 2, 3], Q(22_928716) / Q(1_000000)),
                                            (&[1, 2], Q(5_538309) / Q(1_000000)),
                                            (&[1, 3], Q(13_130227) / Q(1_000000)),
                                            (&[1], Q(6_056291) / Q(1_000000)),
                                            (&[2, 3], Q(23_992772) / Q(1_000000)),
                                            (&[2], Q(16_699207) / Q(1_000000)),
                                            (&[3], Q(98_165759) / Q(1_000000)),
                                            (&[], Q(129_664430) / Q(1_000000))];
        assert_eq!(strength(4, ballots), Q(77_389937) / Q(1_000000));
    }
}
<|endoftext|>


<commit_before>//! Compiler plugin defining macros that create PHF data structures.
//!
//! # Example
//!
//! ```rust
//! #![feature(plugin, core)]
//! #![plugin(phf_macros)]
//!
//! extern crate phf;
//!
//! #[derive(Clone)]
//! pub enum Keyword {
//!     Loop,
//!     Continue,
//!     Break,
//!     Fn,
//!     Extern,
//! }
//!
//! static KEYWORDS: phf::Map<&'static str, Keyword> = phf_map! {
//!     "loop" => Keyword::Loop,
//!     "continue" => Keyword::Continue,
//!     "break" => Keyword::Break,
//!     "fn" => Keyword::Fn,
//!     "extern" => Keyword::Extern,
//! };
//!
//! pub fn parse_keyword(keyword: &str) -> Option<Keyword> {
//!     KEYWORDS.get(keyword).cloned()
//! }
//! # fn main() {}
//! ```
#![doc(html_root_url="https://docs.rs/phf_macros/0.7.20")]
#![feature(plugin_registrar, quote, rustc_private)]

#[macro_use]
extern crate syntax;
extern crate rustc_plugin;
extern crate phf_shared;
extern crate phf_generator;
#[cfg(feature = "unicase_support")]
extern crate unicase;

use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::time::Instant;
use syntax::ast::{self, Expr, ExprKind, Mutability, TyKind};
use syntax::tokenstream::TokenTree;
use syntax::codemap::Span;
use syntax::ext::base::{DummyResult, ExtCtxt, MacResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse;
use syntax::parse::token::{Comma, Eof, FatArrow};
use syntax::print::pprust;
use syntax::ptr::P;
use syntax::symbol::Symbol;
use rustc_plugin::Registry;
use phf_generator::HashState;
use std::env;
#[cfg(feature = "unicase_support")]
use unicase::UniCase;

use util::{Entry, Key};
use util::{create_map, create_set, create_ordered_map, create_ordered_set};

pub mod util;
mod macros;

mod errors {
    pub use syntax::errors::*;
}

#[plugin_registrar]
#[doc(hidden)]
pub fn macro_registrar(reg: &mut Registry) {
    reg.register_macro("phf_map", expand_phf_map);
    reg.register_macro("phf_set", expand_phf_set);
    reg.register_macro("phf_ordered_map", expand_phf_ordered_map);
    reg.register_macro("phf_ordered_set", expand_phf_ordered_set);
}

fn generate_hash(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> HashState {
    let start = Instant::now();
    let state = phf_generator::generate_hash(entries);
    let time = Instant::now() - start;

    if env::var_os("PHF_STATS").is_some() {
        let time = time.as_secs() as f64 + (time.subsec_nanos() as f64 / 1_000_000_000.);
        cx.parse_sess
          .span_diagnostic
          .span_note_without_error(sp, &format!("PHF generation took {} seconds", time));
    }

    state
}

fn expand_phf_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_map(cx, sp, entries, state)
}

fn expand_phf_set(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_set(cx, sp, entries, state)
}

fn expand_phf_ordered_map(cx: &mut ExtCtxt,
                          sp: Span,
                          tts: &[TokenTree])
                          -> Box<MacResult + 'static> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_ordered_map(cx, sp, entries, state)
}

fn expand_phf_ordered_set(cx: &mut ExtCtxt,
                          sp: Span,
                          tts: &[TokenTree])
                          -> Box<MacResult + 'static> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_ordered_set(cx, sp, entries, state)
}

fn parse_map(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), tts.to_vec());
    let mut entries = Vec::new();

    let mut bad = false;
    while parser.token != Eof {
        let key = cx.expander().fold_expr(panictry!(parser.parse_expr()));
        let key_contents = parse_key(cx, &*key).unwrap_or_else(|| {
            bad = true;
            Key::Str(Symbol::intern("").as_str())
        });
        let key = adjust_key(cx, key);

        if !parser.eat(&FatArrow) {
            cx.span_err(parser.span, "expected `=>`");
            return None;
        }

        let value = panictry!(parser.parse_expr());

        entries.push(Entry {
            key_contents: key_contents,
            key: key,
            value: value,
        });

        if !parser.eat(&Comma) && parser.token != Eof {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_set(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), tts.to_vec());
    let mut entries = Vec::new();
    let value = quote_expr!(&*cx, ());

    let mut bad = false;
    while parser.token != Eof {
        let key = cx.expander().fold_expr(panictry!(parser.parse_expr()));
        let key_contents = parse_key(cx, &*key).unwrap_or_else(|| {
            bad = true;
            Key::Str(Symbol::intern("").as_str())
        });
        let key = adjust_key(cx, key);

        entries.push(Entry {
            key_contents: key_contents,
            key: key,
            value: value.clone(),
        });

        if !parser.eat(&Comma) && parser.token != Eof {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_key(cx: &mut ExtCtxt, e: &Expr) -> Option<Key> {
    match e.node {
        ExprKind::Lit(ref lit) => {
            match lit.node {
                ast::LitKind::Str(ref s, _) => Some(Key::Str(s.as_str())),
                ast::LitKind::ByteStr(ref b) => Some(Key::Binary(b.clone())),
                ast::LitKind::Byte(b) => Some(Key::U8(b)),
                ast::LitKind::Char(c) => Some(Key::Char(c)),
                ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I8)) =>
                    Some(Key::I8(i as i8)),
                ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I16)) =>
                    Some(Key::I16(i as i16)),
                ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I32)) =>
                    Some(Key::I32(i as i32)),
                ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I64)) =>
                    Some(Key::I64(i as i64)),
                ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U8)) =>
                    Some(Key::U8(i as u8)),
                ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U16)) =>
                    Some(Key::U16(i as u16)),
                ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U32)) =>
                    Some(Key::U32(i as u32)),
                ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U64)) =>
                    Some(Key::U64(i as u64)),
                ast::LitKind::Bool(b) => Some(Key::Bool(b)),
                _ => {
                    cx.span_err(e.span, "unsupported literal type");
                    None
                }
            }
        }
        ExprKind::Array(ref v) => {
            let bytes: Vec<Option<u8>> = v.iter().map(|expr|
                if let ExprKind::Lit(ref p) = expr.node {
                    match p.node {
                        ast::LitKind::Int(val, ast::LitIntType::Unsigned(ast::UintTy::U8)) if val < 256 =>
                            Some(val as u8),
                        ast::LitKind::Int(val, ast::LitIntType::Unsuffixed) if val < 256 =>
                            Some(val as u8),
                        _ => None,
                    }
                } else {
                    None
            }).collect();
            if bytes.iter().all(|x| x.is_some()) {
                Some(Key::Binary(std::rc::Rc::new(bytes.iter().map(|x| x.unwrap()).collect())))
            } else {
                cx.span_err(e.span,
                            "not all elements of an expected u8 array literal were u8 literals");
                None
            }
        }
        #[cfg(feature = "unicase_support")]
        ExprKind::Call(ref f, ref args) => {
            if let ExprKind::Path(_, ref path) = f.node {
                if &*path.segments.last().unwrap().identifier.name.as_str() == "UniCase" {
                    if args.len() == 1 {
                        if let ExprKind::Lit(ref lit) = args.first().unwrap().node {
                            if let ast::LitKind::Str(ref s, _) = lit.node {
                                return Some(Key::UniCase(UniCase(s.to_string())));
                            } else {
                                cx.span_err(e.span, "only a str literal is allowed in UniCase");
                                return None;
                            }
                        }
                    } else {
                        cx.span_err(e.span, "only one str literal is allowed in UniCase");
                        return None;
                    }
                }
            }
            cx.span_err(e.span, "only UniCase is allowed besides literals");
            None
        },
        _ => {
            cx.span_err(e.span, "expected a literal (or a UniCase if the unicase_support feature is enabled)");
            None
        }
    }
}

fn adjust_key(cx: &mut ExtCtxt, e: P<Expr>) -> P<Expr> {
    let coerce_as_slice = match e.node {
        ExprKind::Lit(ref lit) => {
            match lit.node {
                ast::LitKind::ByteStr(_) => true,
                _ => false,
            }
        },
        _ => false,
    };
    if coerce_as_slice {
        let u8_type = cx.ty_path(cx.path_ident(e.span, cx.ident_of("u8")));
        let array_type = cx.ty(e.span, TyKind::Slice(u8_type));
        let slice_type = cx.ty_rptr(e.span, array_type, None, Mutability::Immutable);
        cx.expr_cast(e.span, e, slice_type)
    } else {
        e
    }
}

fn has_duplicates(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> bool {
    let mut dups = false;
    let mut strings = HashMap::new();
    for entry in entries.iter() {
        let &mut (ref mut spans, _) = match strings.entry(&entry.key_contents) {
            Occupied(e) => e.into_mut(),
            Vacant(e) => e.insert((vec![], &entry.key)),
        };
        spans.push(entry.key.span);
    }

    for &(ref spans, key) in strings.values() {
        if spans.len() == 1 {
            continue;
        }

        dups = true;
        let mut err = cx.struct_span_err(sp, &*format!("duplicate key {}",
                                                       pprust::expr_to_string(&**key)));
        for span in spans.iter() {
            err.span_note(*span, "one occurrence here");
        }
        err.emit();
    }

    dups
}
<commit_msg>Fix build<commit_after>//! Compiler plugin defining macros that create PHF data structures.
//!
//! # Example
//!
//! ```rust
//! #![feature(plugin, core)]
//! #![plugin(phf_macros)]
//!
//! extern crate phf;
//!
//! #[derive(Clone)]
//! pub enum Keyword {
//!     Loop,
//!     Continue,
//!     Break,
//!     Fn,
//!     Extern,
//! }
//!
//! static KEYWORDS: phf::Map<&'static str, Keyword> = phf_map! {
//!     "loop" => Keyword::Loop,
//!     "continue" => Keyword::Continue,
//!     "break" => Keyword::Break,
//!     "fn" => Keyword::Fn,
//!     "extern" => Keyword::Extern,
//! };
//!
//! pub fn parse_keyword(keyword: &str) -> Option<Keyword> {
//!     KEYWORDS.get(keyword).cloned()
//! }
//! # fn main() {}
//! ```
#![doc(html_root_url = "https://docs.rs/phf_macros/0.7.20")]
#![feature(plugin_registrar, quote, rustc_private)]

#[macro_use]
extern crate syntax;
extern crate phf_generator;
extern crate phf_shared;
extern crate rustc_plugin;
#[cfg(feature = "unicase_support")]
extern crate unicase;

use phf_generator::HashState;
use rustc_plugin::Registry;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::HashMap;
use std::env;
use std::time::Instant;
use syntax::ast::{self, Expr, ExprKind, Mutability, TyKind};
use syntax::codemap::Span;
use syntax::ext::base::{DummyResult, ExtCtxt, MacResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse;
use syntax::parse::token::{Comma, Eof, FatArrow};
use syntax::print::pprust;
use syntax::ptr::P;
use syntax::symbol::Symbol;
use syntax::tokenstream::TokenTree;
#[cfg(feature = "unicase_support")]
use unicase::UniCase;

use util::{create_map, create_ordered_map, create_ordered_set, create_set};
use util::{Entry, Key};

mod macros;
pub mod util;

mod errors {
    pub use syntax::errors::*;
}

#[plugin_registrar]
#[doc(hidden)]
pub fn macro_registrar(reg: &mut Registry) {
    reg.register_macro("phf_map", expand_phf_map);
    reg.register_macro("phf_set", expand_phf_set);
    reg.register_macro("phf_ordered_map", expand_phf_ordered_map);
    reg.register_macro("phf_ordered_set", expand_phf_ordered_set);
}

fn generate_hash(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> HashState {
    let start = Instant::now();
    let state = phf_generator::generate_hash(entries);
    let time = Instant::now() - start;

    if env::var_os("PHF_STATS").is_some() {
        let time = time.as_secs() as f64 + (time.subsec_nanos() as f64 / 1_000_000_000.);
        cx.parse_sess
            .span_diagnostic
            .span_note_without_error(sp, &format!("PHF generation took {} seconds", time));
    }

    state
}

fn expand_phf_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_map(cx, sp, entries, state)
}

fn expand_phf_set(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_set(cx, sp, entries, state)
}

fn expand_phf_ordered_map(
    cx: &mut ExtCtxt,
    sp: Span,
    tts: &[TokenTree],
) -> Box<MacResult + 'static> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_ordered_map(cx, sp, entries, state)
}

fn expand_phf_ordered_set(
    cx: &mut ExtCtxt,
    sp: Span,
    tts: &[TokenTree],
) -> Box<MacResult + 'static> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, &*entries) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, &*entries);

    create_ordered_set(cx, sp, entries, state)
}

fn parse_map(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), tts.to_vec());
    let mut entries = Vec::new();

    let mut bad = false;
    while parser.token != Eof {
        let key = cx.expander().fold_expr(panictry!(parser.parse_expr()));
        let key_contents = parse_key(cx, &*key).unwrap_or_else(|| {
            bad = true;
            Key::Str(Symbol::intern("").as_str())
        });
        let key = adjust_key(cx, key);

        if !parser.eat(&FatArrow) {
            cx.span_err(parser.span, "expected `=>`");
            return None;
        }

        let value = panictry!(parser.parse_expr());

        entries.push(Entry {
            key_contents: key_contents,
            key: key,
            value: value,
        });

        if !parser.eat(&Comma) && parser.token != Eof {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_set(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), tts.to_vec());
    let mut entries = Vec::new();
    let value = quote_expr!(&*cx, ());

    let mut bad = false;
    while parser.token != Eof {
        let key = cx.expander().fold_expr(panictry!(parser.parse_expr()));
        let key_contents = parse_key(cx, &*key).unwrap_or_else(|| {
            bad = true;
            Key::Str(Symbol::intern("").as_str())
        });
        let key = adjust_key(cx, key);

        entries.push(Entry {
            key_contents: key_contents,
            key: key,
            value: value.clone(),
        });

        if !parser.eat(&Comma) && parser.token != Eof {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_key(cx: &mut ExtCtxt, e: &Expr) -> Option<Key> {
    match e.node {
        ExprKind::Lit(ref lit) => match lit.node {
            ast::LitKind::Str(ref s, _) => Some(Key::Str(s.as_str())),
            ast::LitKind::ByteStr(ref b) => Some(Key::Binary(b.clone())),
            ast::LitKind::Byte(b) => Some(Key::U8(b)),
            ast::LitKind::Char(c) => Some(Key::Char(c)),
            ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I8)) => Some(Key::I8(i as i8)),
            ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I16)) => {
                Some(Key::I16(i as i16))
            }
            ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I32)) => {
                Some(Key::I32(i as i32))
            }
            ast::LitKind::Int(i, ast::LitIntType::Signed(ast::IntTy::I64)) => {
                Some(Key::I64(i as i64))
            }
            ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U8)) => {
                Some(Key::U8(i as u8))
            }
            ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U16)) => {
                Some(Key::U16(i as u16))
            }
            ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U32)) => {
                Some(Key::U32(i as u32))
            }
            ast::LitKind::Int(i, ast::LitIntType::Unsigned(ast::UintTy::U64)) => {
                Some(Key::U64(i as u64))
            }
            ast::LitKind::Bool(b) => Some(Key::Bool(b)),
            _ => {
                cx.span_err(e.span, "unsupported literal type");
                None
            }
        },
        ExprKind::Array(ref v) => {
            let bytes: Vec<Option<u8>> = v.iter()
                .map(|expr| {
                    if let ExprKind::Lit(ref p) = expr.node {
                        match p.node {
                            ast::LitKind::Int(val, ast::LitIntType::Unsigned(ast::UintTy::U8))
                                if val < 256 =>
                            {
                                Some(val as u8)
                            }
                            ast::LitKind::Int(val, ast::LitIntType::Unsuffixed) if val < 256 => {
                                Some(val as u8)
                            }
                            _ => None,
                        }
                    } else {
                        None
                    }
                })
                .collect();
            if bytes.iter().all(|x| x.is_some()) {
                Some(Key::Binary(std::rc::Rc::new(
                    bytes.iter().map(|x| x.unwrap()).collect(),
                )))
            } else {
                cx.span_err(
                    e.span,
                    "not all elements of an expected u8 array literal were u8 literals",
                );
                None
            }
        }
        #[cfg(feature = "unicase_support")]
        ExprKind::Call(ref f, ref args) => {
            if let ExprKind::Path(_, ref path) = f.node {
                if &*path.segments.last().unwrap().ident.name.as_str() == "UniCase" {
                    if args.len() == 1 {
                        if let ExprKind::Lit(ref lit) = args.first().unwrap().node {
                            if let ast::LitKind::Str(ref s, _) = lit.node {
                                return Some(Key::UniCase(UniCase(s.to_string())));
                            } else {
                                cx.span_err(e.span, "only a str literal is allowed in UniCase");
                                return None;
                            }
                        }
                    } else {
                        cx.span_err(e.span, "only one str literal is allowed in UniCase");
                        return None;
                    }
                }
            }
            cx.span_err(e.span, "only UniCase is allowed besides literals");
            None
        }
        _ => {
            cx.span_err(
                e.span,
                "expected a literal (or a UniCase if the unicase_support feature is enabled)",
            );
            None
        }
    }
}

fn adjust_key(cx: &mut ExtCtxt, e: P<Expr>) -> P<Expr> {
    let coerce_as_slice = match e.node {
        ExprKind::Lit(ref lit) => match lit.node {
            ast::LitKind::ByteStr(_) => true,
            _ => false,
        },
        _ => false,
    };
    if coerce_as_slice {
        let u8_type = cx.ty_path(cx.path_ident(e.span, cx.ident_of("u8")));
        let array_type = cx.ty(e.span, TyKind::Slice(u8_type));
        let slice_type = cx.ty_rptr(e.span, array_type, None, Mutability::Immutable);
        cx.expr_cast(e.span, e, slice_type)
    } else {
        e
    }
}

fn has_duplicates(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> bool {
    let mut dups = false;
    let mut strings = HashMap::new();
    for entry in entries.iter() {
        let &mut (ref mut spans, _) = match strings.entry(&entry.key_contents) {
            Occupied(e) => e.into_mut(),
            Vacant(e) => e.insert((vec![], &entry.key)),
        };
        spans.push(entry.key.span);
    }

    for &(ref spans, key) in strings.values() {
        if spans.len() == 1 {
            continue;
        }

        dups = true;
        let mut err = cx.struct_span_err(
            sp,
            &*format!("duplicate key {}", pprust::expr_to_string(&**key)),
        );
        for span in spans.iter() {
            err.span_note(*span, "one occurrence here");
        }
        err.emit();
    }

    dups
}
<|endoftext|>


<commit_before>//! To run a LOGO program, run:
//!     cargo run --example logo_interpreter -- my_logo_program.txt

extern crate turtle;

use turtle::Turtle;
use std::env;
use std::fs::File;
use std::io::{self, BufRead, BufReader};

fn main() {
    let mut turtle = Turtle::new();

    let args: Vec<String> = env::args().collect();

    if args.len() != 2 {
        println!(
            "Interactive LOGO Mode.\nType the commands in stdin.\nPress '^C' (cmd/ctrl+C) to exit"
        );
        let stdin = io::stdin();
        interpret(&mut turtle, stdin)
    }
    else {
        let path = &args[1];
        let file = File::open(path).expect("Could not open provided program");
        interpret(&mut turtle, file)
    }
}

fn interpret<R: io::Read>(turtle: &mut Turtle, input: R) {
    let mut reader = BufReader::new(input);
    loop {
        let mut buffer = String::new();
        let read_bytes = reader.read_line(&mut buffer)
            .expect("Unable read input");
        if read_bytes == 0 {
            // Reached EOF, break loop
            break;
        }
        let cmd_args = buffer.split_whitespace();
        handle_command(turtle, cmd_args);
    }
}

fn handle_command<'a, A: Iterator<Item = &'a str>>(turtle: &mut Turtle, mut args: A) {
    while let Some(command) = args.next() {
        match command
        {
            "fd" | "forward" => {
                let distance = parse_distance(args.next()
                    .expect("Expected a distance value after fd/forward command")
                );
                turtle.forward(distance);
            },
            "bk" | "back" => {
                let distance = parse_distance(args.next()
                    .expect("Expect a distance value after bk/back command")
                );
                turtle.backward(distance);
            },
            "lt" | "left" => {
                let distance = parse_distance(args.next()
                    .expect("Expect a distance value after lt/left command")
                );
                turtle.left(distance);
            },
            "rt" | "right" => {
                let distance = parse_distance(args.next()
                    .expect("Expect a distance value after rt/right command")
                );
                turtle.right(distance);
            },
            _ => unimplemented!("Use of invalid or unsupported LOGO command"),
        }
    }
}

fn parse_distance(s: &str) -> f64 {
    let dist: f64 = s.parse().expect("Not a valid number");
    dist
}
<commit_msg>minor asthetic changes<commit_after>//! To run a LOGO program, run:
//!     cargo run --example logo_interpreter -- my_logo_program.txt

extern crate turtle;

use turtle::Turtle;
use std::env;
use std::fs::File;
use std::io::{self, BufRead, BufReader};

fn main() {
    let mut turtle = Turtle::new();

    let args: Vec<String> = env::args().collect();

    if args.len() != 2 {
        println!("Interactive LOGO Mode.\nType the commands in stdin.\nPress '^C' (cmd/ctrl+C) to exit");
        let stdin = io::stdin();
        interpret(&mut turtle, stdin)
    }
    else {
        let path = &args[1];
        let file = File::open(path).expect("Could not open provided program");
        interpret(&mut turtle, file)
    }
}

fn interpret<R: io::Read>(turtle: &mut Turtle, input: R) {
    let mut reader = BufReader::new(input);
    loop {
        let mut buffer = String::new();
        let read_bytes = reader.read_line(&mut buffer)
            .expect("Unable read input");
        if read_bytes == 0 {
            // Reached EOF, break loop
            break;
        }
        let cmd_args = buffer.split_whitespace();
        handle_command(turtle, cmd_args);
    }
}

fn handle_command<'a, A: Iterator<Item = &'a str>>(turtle: &mut Turtle, mut args: A) {
    while let Some(command) = args.next() {
        match command {
            "fd" | "forward" => {
                let distance = parse_distance(args.next()
                    .expect("Expected a distance value after fd/forward command"));
                turtle.forward(distance);
            },
            "bk" | "back" => {
                let distance = parse_distance(args.next()
                    .expect("Expect a distance value after bk/back command"));
                turtle.backward(distance);
            },
            "lt" | "left" => {
                let distance = parse_distance(args.next()
                    .expect("Expect a distance value after lt/left command"));
                turtle.left(distance);
            },
            "rt" | "right" => {
                let distance = parse_distance(args.next()
                    .expect("Expect a distance value after rt/right command"));
                turtle.right(distance);
            },
            _ => unimplemented!("Use of invalid or unsupported LOGO command"),
        }
    }
}

fn parse_distance(s: &str) -> f64 {
    let dist: f64 = s.parse().expect("Not a valid number");
    dist
}
<|endoftext|>