<commit_before>extern crate time;

use std::thread;
use std::sync::mpsc;
use std::sync::mpsc::{Sender, Receiver};

use game;

const BOARD_AREA: u8 = 64;

const STARTING_DEPTH: u8 = 5;
const TIME_LIMIT: f64 = 1.0;

const LIGHT_STARTING_SCORE: i16 = -10_000;
const DARK_STARTING_SCORE:  i16 =  10_000;

const RANDOMNESS: i16 = 3;

const BONUS_TURN: i16 = 3;

const MOBILITY: i16 = 1;

pub fn make_move(game: &game::Game) -> (usize, usize) {

    let mut start_time = time::precise_time_s();
    let mut end_time;
    let mut depth = STARTING_DEPTH;

    let mut best_move: (usize, usize)  = find_best_move(game, depth);

    end_time = time::precise_time_s();

    while ( end_time - start_time < TIME_LIMIT ) && ( depth + 1 + game.get_turn() <= BOARD_AREA ) {
        depth += 1;
        start_time = time::precise_time_s();
        best_move = find_best_move(game, depth);
        end_time = time::precise_time_s();
    }

    best_move
}

fn find_best_move(game: &game::Game, depth: u8) -> (usize, usize) {

    if let game::Status::Running { next_player } = game.get_status() {

        if depth > 0 {

            let mut best_move: (usize, usize) = (game::BOARD_SIZE, game::BOARD_SIZE);
            let mut best_score: i16;

            let mut best_end_move: (usize, usize) = (game::BOARD_SIZE, game::BOARD_SIZE);
            let mut best_end_score: i16;

            let mut num_moves: u8 = 0;
            let mut end_game: bool = true;

            match next_player {
                game::Player::Light => {
                    best_score = LIGHT_STARTING_SCORE;
                    best_end_score = LIGHT_STARTING_SCORE;
                }
                game::Player::Dark  => {
                    best_score = DARK_STARTING_SCORE;
                    best_end_score = DARK_STARTING_SCORE;
                }
            }

            let (tx, rx): (Sender<((usize, usize), (i16, bool))>, Receiver<((usize, usize), (i16, bool))>) = mpsc::channel();

            for row in 0..game::BOARD_SIZE {
                for col in 0..game::BOARD_SIZE {

                    if game.check_move((row, col)) {
                        num_moves +=1;

                        let thread_tx = tx.clone();
                        let mut game = game.clone();

                        thread::spawn(move || {
                            thread_tx.send(( (row, col), eval(game.make_move((row, col)), depth - 1) )).unwrap();
                        });
                    }
                }
            }

            for _ in 0..num_moves {
                let (current_move, (current_score, current_end_game)) = rx.recv().ok().expect("Could not receive answer");

                match next_player {
                    game::Player::Light => {
                        if current_end_game {
                            if current_score > best_end_score {
                                best_end_score = current_score;
                                best_end_move = current_move;
                            }
                        } else {
                            if current_score > best_score {
                                best_score = current_score;
                                best_move = current_move;
                                end_game = false;
                            }
                        }
                    }
                    game::Player::Dark  => {
                        if current_end_game {
                            if current_score < best_end_score {
                                best_end_score = current_score;
                                best_end_move = current_move;
                            }
                        } else {
                            if current_score < best_score {
                                best_score = current_score;
                                best_move = current_move;
                                end_game = false;
                            }
                        }
                    }
                }
            }

            match next_player {
                game::Player::Light  => {
                    if best_end_score > 0 || (best_end_score == 0 && best_score < 0) || end_game {
                        return best_end_move;
                    } else {
                        return best_move;
                    }
                }
                game::Player::Dark  => {
                    if best_end_score < 0 || (best_end_score == 0 && best_score > 0) || end_game {
                        return best_end_move;
                    } else {
                        return best_move;
                    }
                }
            }

        } else {
            panic!("Depth cannot be zero");
        }

    } else {
        panic!{"Game ended, cannot make a move!"};
    }
}

fn eval(game: &game::Game, depth: u8) -> (i16, bool) {

    match game.get_status() {
        game::Status::Ended => {
            let (score_light, score_dark) = game.get_score();
            let score: i16 = (score_light as i16) - (score_dark as i16);
            return (score, true);
        }
        game::Status::Running { next_player } => {
            if depth == 0 {
                match next_player {
                    game::Player::Light => return (heavy_eval(game) + BONUS_TURN, false),
                    game::Player::Dark  => return (heavy_eval(game) - BONUS_TURN, false),
                }
            } else {
                let mut end_game: bool = true;
                let mut num_moves: i16 = 0;
                let mut best_end_score: i16;
                let mut best_score: i16;

                match next_player {

                    game::Player::Light => {
                        best_score = LIGHT_STARTING_SCORE;
                        best_end_score = LIGHT_STARTING_SCORE;

                        for row in 0..game::BOARD_SIZE {
                            for col in 0..game::BOARD_SIZE {
                                if game.check_move((row, col)) {

                                    let (current_score, current_end_game) = eval(game.clone().make_move((row, col)), depth - 1);

                                    if current_end_game {
                                        if current_score > best_end_score {
                                            best_end_score = current_score;
                                        }
                                    } else {
                                        num_moves += 1;
                                        if current_score > best_score {
                                            best_score = current_score;
                                            end_game = false;
                                        }
                                    }

                                }
                            }
                        }

                        if best_end_score > 0 || (best_end_score == 0 && best_score < 0) || end_game {
                            return (best_end_score, true);
                        } else {
                            return (best_score + MOBILITY*num_moves, false);
                        }
                    }

                    game::Player::Dark  => {
                        best_score = DARK_STARTING_SCORE;
                        best_end_score = DARK_STARTING_SCORE;

                        for row in 0..game::BOARD_SIZE {
                            for col in 0..game::BOARD_SIZE {
                                if game.check_move((row, col)) {

                                    let (current_score, current_end_game) = eval(game.clone().make_move((row, col)), depth - 1);

                                    if current_end_game {
                                        if current_score < best_end_score {
                                            best_end_score = current_score;
                                        }
                                    } else {
                                        num_moves += 1;
                                        if current_score < best_score {
                                            best_score = current_score;
                                            end_game = false;
                                        }
                                    }

                                }
                            }
                        }

                        if best_end_score < 0 || (best_end_score == 0 && best_score > 0) || end_game {
                            return (best_end_score, true);
                        } else {
                            return (best_score - MOBILITY*num_moves, false);
                        }
                    }

                }

            }
        }
    }
}

fn heavy_eval(game: &game::Game) -> i16 {
    const CORNER_BONUS: i16 = 12;
    const ODD_MALUS: i16 = 5;
    const EVEN_BONUS: i16 = 3;
    const ODD_CORNER_MALUS: i16 = 8;
    const EVEN_CORNER_BONUS: i16 = 5;
    const FIXED_BONUS: i16 = 2;

    const SIDES: [( (usize, usize), (usize, usize), (usize, usize), (usize, usize), (usize, usize), (usize, usize), (usize, usize) ); 4] = [
        ( (0,0), (0,1), (1,1), (0,2), (2,2), (1,0), (2,0) ), // NW corner
        ( (0,7), (1,7), (1,6), (2,7), (2,5), (0,6), (0,5) ), // NE corner
        ( (7,0), (6,0), (6,1), (5,0), (5,2), (7,1), (7,2) ), // SW corner
        ( (7,7), (6,7), (6,6), (5,7), (5,5), (7,6), (7,6) ), // SE corner
        ];

    let (score_light, score_dark) = game.get_score();
    let mut score: i16 = (score_light as i16) - (score_dark as i16);

    for special_cells in SIDES.iter() {

        let (corner, odd, odd_corner, even, even_corner, counter_odd, counter_even) = *special_cells;

        if let game::Cell::Taken { player } = game.get_cell(corner) {
            match player {
                game::Player::Light => {
                    score += CORNER_BONUS;
                    if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(odd) {
                        score += FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(even) {
                            score += FIXED_BONUS;
                        }
                    }
                    if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(counter_odd) {
                        score += FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(counter_even) {
                            score += FIXED_BONUS;
                        }
                    }
                }
                game::Player::Dark => {
                    score -= CORNER_BONUS;
                    if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(odd) {
                        score -= FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(even) {
                            score -= FIXED_BONUS;
                        }
                    }
                    if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(counter_odd) {
                        score -= FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(counter_even) {
                            score -= FIXED_BONUS;
                        }
                    }
                }
            }

        } else {

            if let game::Cell::Taken { player } = game.get_cell(odd) {
                match player {
                    game::Player::Light => score -= ODD_MALUS,
                    game::Player::Dark  => score += ODD_MALUS,
                }
            } else if let game::Cell::Taken { player } = game.get_cell(even) {
                match player {
                    game::Player::Light => score += EVEN_BONUS,
                    game::Player::Dark  => score -= EVEN_BONUS,
                }
            }

            if let game::Cell::Taken { player } = game.get_cell(counter_odd) {
                match player {
                    game::Player::Light => score -= ODD_MALUS,
                    game::Player::Dark  => score += ODD_MALUS,
                }
            } else if let game::Cell::Taken { player } = game.get_cell(counter_even) {
                match player {
                    game::Player::Light => score += EVEN_BONUS,
                    game::Player::Dark  => score -= EVEN_BONUS,
                }
            }

            if let game::Cell::Taken { player } = game.get_cell(odd_corner) {
                match player {
                    game::Player::Light => score -= ODD_CORNER_MALUS,
                    game::Player::Dark  => score += ODD_CORNER_MALUS,
                }

            } else if let game::Cell::Taken { player } = game.get_cell(even_corner) {
                match player {
                    game::Player::Light => score += EVEN_CORNER_BONUS,
                    game::Player::Dark  => score -= EVEN_CORNER_BONUS,
                }
            }
        }
    }

    score
}
<commit_msg>Added RANDOMNESS back<commit_after>extern crate time;

use std::thread;
use std::sync::mpsc;
use std::sync::mpsc::{Sender, Receiver};

use game;

const BOARD_AREA: u8 = 64;

const STARTING_DEPTH: u8 = 5;
const TIME_LIMIT: f64 = 1.0;

const LIGHT_STARTING_SCORE: i16 = -10_000;
const DARK_STARTING_SCORE:  i16 =  10_000;

const RANDOMNESS: i16 = 3;

const BONUS_TURN: i16 = 3;

const MOBILITY: i16 = 1;

pub fn make_move(game: &game::Game) -> (usize, usize) {

    let mut start_time = time::precise_time_s();
    let mut end_time;
    let mut depth = STARTING_DEPTH;

    let mut best_move: (usize, usize)  = find_best_move(game, depth);

    end_time = time::precise_time_s();

    while ( end_time - start_time < TIME_LIMIT ) && ( depth + 1 + game.get_turn() <= BOARD_AREA ) {
        depth += 1;
        start_time = time::precise_time_s();
        best_move = find_best_move(game, depth);
        end_time = time::precise_time_s();
    }

    best_move
}

fn find_best_move(game: &game::Game, depth: u8) -> (usize, usize) {

    if let game::Status::Running { next_player } = game.get_status() {

        if depth > 0 {

            let mut best_move: (usize, usize) = (game::BOARD_SIZE, game::BOARD_SIZE);
            let mut best_score: i16;

            let mut best_end_move: (usize, usize) = (game::BOARD_SIZE, game::BOARD_SIZE);
            let mut best_end_score: i16;

            let mut num_moves: u8 = 0;
            let mut end_game: bool = true;

            match next_player {
                game::Player::Light => {
                    best_score = LIGHT_STARTING_SCORE;
                    best_end_score = LIGHT_STARTING_SCORE;
                }
                game::Player::Dark  => {
                    best_score = DARK_STARTING_SCORE;
                    best_end_score = DARK_STARTING_SCORE;
                }
            }

            let (tx, rx): (Sender<((usize, usize), (i16, bool))>, Receiver<((usize, usize), (i16, bool))>) = mpsc::channel();

            for row in 0..game::BOARD_SIZE {
                for col in 0..game::BOARD_SIZE {

                    if game.check_move((row, col)) {
                        num_moves +=1;

                        let thread_tx = tx.clone();
                        let mut game = game.clone();

                        thread::spawn(move || {
                            thread_tx.send(( (row, col), eval(game.make_move((row, col)), depth - 1) )).unwrap();
                        });
                    }
                }
            }

            for _ in 0..num_moves {
                let (current_move, (current_score, current_end_game)) = rx.recv().ok().expect("Could not receive answer");

                match next_player {
                    game::Player::Light => {
                        if current_end_game {
                            if current_score > best_end_score {
                                best_end_score = current_score;
                                best_end_move = current_move;
                            }
                        } else {
                            if current_score + RANDOMNESS > best_score {
                                best_score = current_score;
                                best_move = current_move;
                                end_game = false;
                            }
                        }
                    }
                    game::Player::Dark  => {
                        if current_end_game {
                            if current_score < best_end_score {
                                best_end_score = current_score;
                                best_end_move = current_move;
                            }
                        } else {
                            if current_score - RANDOMNESS < best_score {
                                best_score = current_score;
                                best_move = current_move;
                                end_game = false;
                            }
                        }
                    }
                }
            }

            match next_player {
                game::Player::Light  => {
                    if best_end_score > 0 || (best_end_score == 0 && best_score < 0) || end_game {
                        return best_end_move;
                    } else {
                        return best_move;
                    }
                }
                game::Player::Dark  => {
                    if best_end_score < 0 || (best_end_score == 0 && best_score > 0) || end_game {
                        return best_end_move;
                    } else {
                        return best_move;
                    }
                }
            }

        } else {
            panic!("Depth cannot be zero");
        }

    } else {
        panic!{"Game ended, cannot make a move!"};
    }
}

fn eval(game: &game::Game, depth: u8) -> (i16, bool) {

    match game.get_status() {
        game::Status::Ended => {
            let (score_light, score_dark) = game.get_score();
            let score: i16 = (score_light as i16) - (score_dark as i16);
            return (score, true);
        }
        game::Status::Running { next_player } => {
            if depth == 0 {
                match next_player {
                    game::Player::Light => return (heavy_eval(game) + BONUS_TURN, false),
                    game::Player::Dark  => return (heavy_eval(game) - BONUS_TURN, false),
                }
            } else {
                let mut end_game: bool = true;
                let mut num_moves: i16 = 0;
                let mut best_end_score: i16;
                let mut best_score: i16;

                match next_player {

                    game::Player::Light => {
                        best_score = LIGHT_STARTING_SCORE;
                        best_end_score = LIGHT_STARTING_SCORE;

                        for row in 0..game::BOARD_SIZE {
                            for col in 0..game::BOARD_SIZE {
                                if game.check_move((row, col)) {

                                    let (current_score, current_end_game) = eval(game.clone().make_move((row, col)), depth - 1);

                                    if current_end_game {
                                        if current_score > best_end_score {
                                            best_end_score = current_score;
                                        }
                                    } else {
                                        num_moves += 1;
                                        if current_score > best_score {
                                            best_score = current_score;
                                            end_game = false;
                                        }
                                    }

                                }
                            }
                        }

                        if best_end_score > 0 || (best_end_score == 0 && best_score < 0) || end_game {
                            return (best_end_score, true);
                        } else {
                            return (best_score + MOBILITY*num_moves, false);
                        }
                    }

                    game::Player::Dark  => {
                        best_score = DARK_STARTING_SCORE;
                        best_end_score = DARK_STARTING_SCORE;

                        for row in 0..game::BOARD_SIZE {
                            for col in 0..game::BOARD_SIZE {
                                if game.check_move((row, col)) {

                                    let (current_score, current_end_game) = eval(game.clone().make_move((row, col)), depth - 1);

                                    if current_end_game {
                                        if current_score < best_end_score {
                                            best_end_score = current_score;
                                        }
                                    } else {
                                        num_moves += 1;
                                        if current_score < best_score {
                                            best_score = current_score;
                                            end_game = false;
                                        }
                                    }

                                }
                            }
                        }

                        if best_end_score < 0 || (best_end_score == 0 && best_score > 0) || end_game {
                            return (best_end_score, true);
                        } else {
                            return (best_score - MOBILITY*num_moves, false);
                        }
                    }

                }

            }
        }
    }
}

fn heavy_eval(game: &game::Game) -> i16 {
    const CORNER_BONUS: i16 = 12;
    const ODD_MALUS: i16 = 5;
    const EVEN_BONUS: i16 = 3;
    const ODD_CORNER_MALUS: i16 = 8;
    const EVEN_CORNER_BONUS: i16 = 5;
    const FIXED_BONUS: i16 = 2;

    const SIDES: [( (usize, usize), (usize, usize), (usize, usize), (usize, usize), (usize, usize), (usize, usize), (usize, usize) ); 4] = [
        ( (0,0), (0,1), (1,1), (0,2), (2,2), (1,0), (2,0) ), // NW corner
        ( (0,7), (1,7), (1,6), (2,7), (2,5), (0,6), (0,5) ), // NE corner
        ( (7,0), (6,0), (6,1), (5,0), (5,2), (7,1), (7,2) ), // SW corner
        ( (7,7), (6,7), (6,6), (5,7), (5,5), (7,6), (7,6) ), // SE corner
        ];

    let (score_light, score_dark) = game.get_score();
    let mut score: i16 = (score_light as i16) - (score_dark as i16);

    for special_cells in SIDES.iter() {

        let (corner, odd, odd_corner, even, even_corner, counter_odd, counter_even) = *special_cells;

        if let game::Cell::Taken { player } = game.get_cell(corner) {
            match player {
                game::Player::Light => {
                    score += CORNER_BONUS;
                    if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(odd) {
                        score += FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(even) {
                            score += FIXED_BONUS;
                        }
                    }
                    if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(counter_odd) {
                        score += FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Light } = game.get_cell(counter_even) {
                            score += FIXED_BONUS;
                        }
                    }
                }
                game::Player::Dark => {
                    score -= CORNER_BONUS;
                    if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(odd) {
                        score -= FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(even) {
                            score -= FIXED_BONUS;
                        }
                    }
                    if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(counter_odd) {
                        score -= FIXED_BONUS;
                        if let game::Cell::Taken { player: game::Player::Dark } = game.get_cell(counter_even) {
                            score -= FIXED_BONUS;
                        }
                    }
                }
            }

        } else {

            if let game::Cell::Taken { player } = game.get_cell(odd) {
                match player {
                    game::Player::Light => score -= ODD_MALUS,
                    game::Player::Dark  => score += ODD_MALUS,
                }
            } else if let game::Cell::Taken { player } = game.get_cell(even) {
                match player {
                    game::Player::Light => score += EVEN_BONUS,
                    game::Player::Dark  => score -= EVEN_BONUS,
                }
            }

            if let game::Cell::Taken { player } = game.get_cell(counter_odd) {
                match player {
                    game::Player::Light => score -= ODD_MALUS,
                    game::Player::Dark  => score += ODD_MALUS,
                }
            } else if let game::Cell::Taken { player } = game.get_cell(counter_even) {
                match player {
                    game::Player::Light => score += EVEN_BONUS,
                    game::Player::Dark  => score -= EVEN_BONUS,
                }
            }

            if let game::Cell::Taken { player } = game.get_cell(odd_corner) {
                match player {
                    game::Player::Light => score -= ODD_CORNER_MALUS,
                    game::Player::Dark  => score += ODD_CORNER_MALUS,
                }

            } else if let game::Cell::Taken { player } = game.get_cell(even_corner) {
                match player {
                    game::Player::Light => score += EVEN_CORNER_BONUS,
                    game::Player::Dark  => score -= EVEN_CORNER_BONUS,
                }
            }
        }
    }

    score
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashSet};
use std::ops::Not;

use cranelift_codegen::cursor::{Cursor, FuncCursor};
use cranelift_codegen::entity::EntitySet;
use cranelift_codegen::ir::{InstructionData, Opcode, ProgramOrder, ValueDef};
use cranelift_codegen::ir::immediates::Offset32;

use crate::prelude::*;

/// Workaround for `StackSlot` not implementing `Ord`.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct OrdStackSlot(StackSlot);

impl PartialOrd for OrdStackSlot {
    fn partial_cmp(&self, rhs: &Self) -> Option<std::cmp::Ordering> {
        self.0.as_u32().partial_cmp(&rhs.0.as_u32())
    }
}

impl Ord for OrdStackSlot {
    fn cmp(&self, rhs: &Self) -> std::cmp::Ordering {
        self.0.as_u32().cmp(&rhs.0.as_u32())
    }
}

#[derive(Debug, Default)]
struct StackSlotUsage {
    stack_addr: HashSet<Inst>,
    stack_load: HashSet<Inst>,
    stack_store: HashSet<Inst>,
}

pub(super) fn optimize_function(
    ctx: &mut Context,
    clif_comments: &mut crate::pretty_clif::CommentWriter,
    name: String, // FIXME remove
) {
    ctx.flowgraph(); // Compute cfg and domtree.

    combine_stack_addr_with_load_store(&mut ctx.func);

    // Record all stack_addr, stack_load and stack_store instructions.
    let mut stack_slot_usage_map = BTreeMap::<OrdStackSlot, StackSlotUsage>::new();

    let mut cursor = FuncCursor::new(&mut ctx.func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::StackLoad {
                    opcode: Opcode::StackAddr,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_addr.insert(inst);
                }
                InstructionData::StackLoad {
                    opcode: Opcode::StackLoad,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_load.insert(inst);
                }
                InstructionData::StackStore {
                    opcode: Opcode::StackStore,
                    arg: _,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_store.insert(inst);
                }
                _ => {}
            }
        }
    }

    // FIXME Repeat following instructions until fixpoint.

    remove_unused_stack_addr_and_stack_load(&mut ctx.func, &mut stack_slot_usage_map);

    println!("stack slot usage: {:?}", stack_slot_usage_map);

    for (stack_slot, users) in stack_slot_usage_map.iter_mut() {
        if users.stack_addr.is_empty().not() {
            // Stack addr leaked; there may be unknown loads and stores.
            // FIXME use stacked borrows to optimize
            continue;
        }

        for load in users.stack_load.clone().drain() {
            let load_ebb = ctx.func.layout.inst_ebb(load).unwrap();
            let loaded_value = ctx.func.dfg.inst_results(load)[0];
            let loaded_type = ctx.func.dfg.value_type(loaded_value);

            let potential_stores = users.stack_store.iter().cloned().filter(|&store| {
                match spatial_overlap(&ctx.func, load, store) {
                    SpatialOverlap::No => false, // Can never be the source of the loaded value.
                    SpatialOverlap::Partial | SpatialOverlap::Full => true,
                }
            }).filter(|&store| {
                match temporal_order(&*ctx, load, store) {
                    TemporalOrder::NeverBefore => false, // Can never be the source of the loaded value.
                    TemporalOrder::MaybeBefore | TemporalOrder::DefinitivelyBefore => true,
                }
            }).collect::<Vec<Inst>>();

            for &store in &potential_stores {
                println!(
                    "Potential store -> load forwarding {} -> {} ({:?}, {:?})",
                    ctx.func.dfg.display_inst(store, None),
                    ctx.func.dfg.display_inst(load, None),
                    spatial_overlap(&ctx.func, store, load),
                    temporal_order(&*ctx, store, load),
                );
            }

            match *potential_stores {
                [] => println!("[{}] [BUG?] Reading uninitialized memory", name),
                [store] if spatial_overlap(&ctx.func, store, load) == SpatialOverlap::Full && temporal_order(&ctx, store, load) == TemporalOrder::DefinitivelyBefore => {
                    // Only one store could have been the origin of the value.
                    let store_ebb = ctx.func.layout.inst_ebb(store).unwrap();
                    let stored_value = ctx.func.dfg.inst_args(store)[0];
                    let stored_type = ctx.func.dfg.value_type(stored_value);
                    if stored_type == loaded_type && store_ebb == load_ebb {
                        println!("Store to load forward {} -> {}", store, load);
                        ctx.func.dfg.detach_results(load);
                        ctx.func.dfg.replace(load).nop();
                        ctx.func.dfg.change_to_alias(loaded_value, stored_value);
                        users.stack_load.remove(&load);
                    }
                }
                _ => {} // FIXME implement this
            }
        }

        for store in users.stack_store.clone().drain() {
            let potential_loads = users.stack_load.iter().cloned().filter(|&load| {
                match spatial_overlap(&ctx.func, store, load) {
                    SpatialOverlap::No => false, // Can never be the source of the loaded value.
                    SpatialOverlap::Partial | SpatialOverlap::Full => true,
                }
            }).filter(|&load| {
                match temporal_order(&*ctx, store, load) {
                    TemporalOrder::NeverBefore => false, // Can never be the source of the loaded value.
                    TemporalOrder::MaybeBefore | TemporalOrder::DefinitivelyBefore => true,
                }
            }).collect::<Vec<Inst>>();

            for &load in &potential_loads {
                println!(
                    "Potential load from store {} <- {} ({:?}, {:?})",
                    ctx.func.dfg.display_inst(load, None),
                    ctx.func.dfg.display_inst(store, None),
                    spatial_overlap(&ctx.func, store, load),
                    temporal_order(&*ctx, store, load),
                );
            }

            if potential_loads.is_empty() {
                // Never loaded; can safely remove all stores and the stack slot.
                // FIXME also remove stores when there is always a next store before a load.
                println!("[{}] Remove dead stack store {} of {}", name, ctx.func.dfg.display_inst(store, None), stack_slot.0);
                ctx.func.dfg.replace(store).nop();
                users.stack_store.remove(&store);
            }
        }

        if users.stack_store.is_empty() && users.stack_load.is_empty() {
            // FIXME make stack_slot zero sized.
        }
    }

    println!();
}

fn combine_stack_addr_with_load_store(func: &mut Function) {
    // Turn load and store into stack_load and stack_store when possible.
    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::Load { opcode: Opcode::Load, arg: addr, flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUD: stack_load.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            let ty = cursor.func.dfg.ctrl_typevar(inst);
                            cursor.func.dfg.replace(inst).stack_load(ty, stack_slot, combined_offset);
                        }
                    }
                }
                InstructionData::Store { opcode: Opcode::Store, args: [value, addr], flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUND: stack_store.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            cursor.func.dfg.replace(inst).stack_store(value, stack_slot, combined_offset);
                        }
                    }
                }
                _ => {}
            }
        }
    }
}

fn remove_unused_stack_addr_and_stack_load(func: &mut Function, stack_slot_usage_map: &mut BTreeMap<OrdStackSlot, StackSlotUsage>) {
    // FIXME incrementally rebuild on each call?
    let mut stack_addr_load_insts_users = HashMap::<Inst, HashSet<Inst>>::new();

    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            for &arg in cursor.func.dfg.inst_args(inst) {
                if let ValueDef::Result(arg_origin, 0) = cursor.func.dfg.value_def(arg) {
                    match cursor.func.dfg[arg_origin].opcode() {
                        Opcode::StackAddr | Opcode::StackLoad => {
                            stack_addr_load_insts_users.entry(arg_origin).or_insert_with(HashSet::new).insert(inst);
                        }
                        _ => {}
                    }
                }
            }
        }
    }

    for inst in stack_addr_load_insts_users.keys() {
        let mut is_recorded_stack_addr_or_stack_load = false;
        for stack_slot_users in stack_slot_usage_map.values() {
            is_recorded_stack_addr_or_stack_load |= stack_slot_users.stack_addr.contains(inst) || stack_slot_users.stack_load.contains(inst);
        }
        assert!(is_recorded_stack_addr_or_stack_load);
    }

    // Replace all unused stack_addr and stack_load instructions with nop.
    for stack_slot_users in stack_slot_usage_map.values_mut() {
        // FIXME remove clone
        for &inst in stack_slot_users.stack_addr.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_addr.remove(&inst);
            }
        }

        for &inst in stack_slot_users.stack_load.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_load.remove(&inst);
            }
        }
    }
}

fn try_get_stack_slot_and_offset_for_addr(func: &Function, addr: Value) -> Option<(StackSlot, Offset32)> {
    if let ValueDef::Result(addr_inst, 0) = func.dfg.value_def(addr) {
        if let InstructionData::StackLoad {
            opcode: Opcode::StackAddr,
            stack_slot,
            offset,
        } = func.dfg[addr_inst] {
            return Some((stack_slot, offset));
        }
    }
    None
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum SpatialOverlap {
    No,
    Partial,
    Full,
}

fn spatial_overlap(func: &Function, src: Inst, dest: Inst) -> SpatialOverlap {
    fn inst_info(func: &Function, inst: Inst) -> (StackSlot, Offset32, u32) {
        match func.dfg[inst] {
            InstructionData::StackLoad {
                opcode: Opcode::StackAddr,
                stack_slot,
                offset,
            }
            | InstructionData::StackLoad {
                opcode: Opcode::StackLoad,
                stack_slot,
                offset,
            }
            | InstructionData::StackStore {
                opcode: Opcode::StackStore,
                stack_slot,
                offset,
                arg: _,
            } => (stack_slot, offset, func.dfg.ctrl_typevar(inst).bytes()),
            _ => unreachable!("{:?}", func.dfg[inst]),
        }
    }

    debug_assert_ne!(src, dest);

    let (src_ss, src_offset, src_size) = inst_info(func, src);
    let (dest_ss, dest_offset, dest_size) = inst_info(func, dest);

    if src_ss != dest_ss {
        return SpatialOverlap::No;
    }

    if src_offset == dest_offset && src_size == dest_size {
        return SpatialOverlap::Full;
    }

    let src_end: i64 = src_offset.try_add_i64(i64::from(src_size)).unwrap().into();
    let dest_end: i64 = dest_offset.try_add_i64(i64::from(dest_size)).unwrap().into();
    if src_end <= dest_offset.into() || dest_end <= src_offset.into() {
        return SpatialOverlap::No;
    }

    SpatialOverlap::Partial
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum TemporalOrder {
    /// `src` will never be executed before `dest`.
    NeverBefore,

    /// `src` may be executed before `dest`.
    MaybeBefore,

    /// `src` will always be executed before `dest`.
    /// There may still be other instructions in between.
    DefinitivelyBefore,
}

fn temporal_order(ctx: &Context, src: Inst, dest: Inst) -> TemporalOrder {
    debug_assert_ne!(src, dest);

    let src_ebb = ctx.func.layout.inst_ebb(src).unwrap();
    let dest_ebb = ctx.func.layout.inst_ebb(dest).unwrap();
    if src_ebb == dest_ebb {
        use std::cmp::Ordering::*;
        match ctx.func.layout.cmp(src, dest) {
            Less => TemporalOrder::DefinitivelyBefore,
            Equal => unreachable!(),
            Greater => TemporalOrder::MaybeBefore, // FIXME use dominator to check for loops
        }
    } else {
        // FIXME O(stack_load count * ebb count)
        // FIXME reuse memory allocations
        // FIXME return DefinitivelyBefore is src dominates dest
        let mut visited = EntitySet::new();
        let mut todo = EntitySet::new();
        todo.insert(dest_ebb);
        while let Some(ebb) = todo.pop() {
            if visited.contains(ebb) {
                continue;
            }
            visited.insert(ebb);
            if ebb == src_ebb {
                return TemporalOrder::MaybeBefore;
            }
            for bb in ctx.cfg.pred_iter(ebb) {
                todo.insert(bb.ebb);
            }
        }
        TemporalOrder::NeverBefore
    }
}
<commit_msg>Add OptimizeContext<commit_after>use std::collections::{BTreeMap, HashSet};
use std::ops::Not;

use cranelift_codegen::cursor::{Cursor, FuncCursor};
use cranelift_codegen::entity::EntitySet;
use cranelift_codegen::ir::{InstructionData, Opcode, ProgramOrder, ValueDef};
use cranelift_codegen::ir::immediates::Offset32;

use crate::prelude::*;

/// Workaround for `StackSlot` not implementing `Ord`.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct OrdStackSlot(StackSlot);

impl PartialOrd for OrdStackSlot {
    fn partial_cmp(&self, rhs: &Self) -> Option<std::cmp::Ordering> {
        self.0.as_u32().partial_cmp(&rhs.0.as_u32())
    }
}

impl Ord for OrdStackSlot {
    fn cmp(&self, rhs: &Self) -> std::cmp::Ordering {
        self.0.as_u32().cmp(&rhs.0.as_u32())
    }
}

#[derive(Debug, Default)]
struct StackSlotUsage {
    stack_addr: HashSet<Inst>,
    stack_load: HashSet<Inst>,
    stack_store: HashSet<Inst>,
}

struct OptimizeContext<'a> {
    ctx: &'a mut Context,
    stack_slot_usage_map: BTreeMap<OrdStackSlot, StackSlotUsage>,
}

impl<'a> OptimizeContext<'a> {
    fn for_context(ctx: &'a mut Context) -> Self {
        ctx.flowgraph(); // Compute cfg and domtree.

        // Record all stack_addr, stack_load and stack_store instructions.
        let mut stack_slot_usage_map = BTreeMap::<OrdStackSlot, StackSlotUsage>::new();

        let mut cursor = FuncCursor::new(&mut ctx.func);
        while let Some(_ebb) = cursor.next_ebb() {
            while let Some(inst) = cursor.next_inst() {
                match cursor.func.dfg[inst] {
                    InstructionData::StackLoad {
                        opcode: Opcode::StackAddr,
                        stack_slot,
                        offset: _,
                    } => {
                        stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_addr.insert(inst);
                    }
                    InstructionData::StackLoad {
                        opcode: Opcode::StackLoad,
                        stack_slot,
                        offset: _,
                    } => {
                        stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_load.insert(inst);
                    }
                    InstructionData::StackStore {
                        opcode: Opcode::StackStore,
                        arg: _,
                        stack_slot,
                        offset: _,
                    } => {
                        stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_store.insert(inst);
                    }
                    _ => {}
                }
            }
        }

        OptimizeContext {
            ctx,
            stack_slot_usage_map,
        }
    }
}

pub(super) fn optimize_function(
    ctx: &mut Context,
    clif_comments: &mut crate::pretty_clif::CommentWriter,
    name: String, // FIXME remove
) {
    combine_stack_addr_with_load_store(&mut ctx.func);

    let mut opt_ctx = OptimizeContext::for_context(ctx);

    // FIXME Repeat following instructions until fixpoint.

    remove_unused_stack_addr_and_stack_load(&mut opt_ctx.ctx.func, &mut opt_ctx.stack_slot_usage_map);

    println!("stack slot usage: {:?}", opt_ctx.stack_slot_usage_map);

    for (stack_slot, users) in opt_ctx.stack_slot_usage_map.iter_mut() {
        if users.stack_addr.is_empty().not() {
            // Stack addr leaked; there may be unknown loads and stores.
            // FIXME use stacked borrows to optimize
            continue;
        }

        for load in users.stack_load.clone().drain() {
            let load_ebb = opt_ctx.ctx.func.layout.inst_ebb(load).unwrap();
            let loaded_value = opt_ctx.ctx.func.dfg.inst_results(load)[0];
            let loaded_type = opt_ctx.ctx.func.dfg.value_type(loaded_value);

            let ctx = &*opt_ctx.ctx;
            let potential_stores = users.stack_store.iter().cloned().filter(|&store| {
                match spatial_overlap(&ctx.func, load, store) {
                    SpatialOverlap::No => false, // Can never be the source of the loaded value.
                    SpatialOverlap::Partial | SpatialOverlap::Full => true,
                }
            }).filter(|&store| {
                match temporal_order(ctx, load, store) {
                    TemporalOrder::NeverBefore => false, // Can never be the source of the loaded value.
                    TemporalOrder::MaybeBefore | TemporalOrder::DefinitivelyBefore => true,
                }
            }).collect::<Vec<Inst>>();

            for &store in &potential_stores {
                println!(
                    "Potential store -> load forwarding {} -> {} ({:?}, {:?})",
                    opt_ctx.ctx.func.dfg.display_inst(store, None),
                    opt_ctx.ctx.func.dfg.display_inst(load, None),
                    spatial_overlap(&opt_ctx.ctx.func, store, load),
                    temporal_order(&*opt_ctx.ctx, store, load),
                );
            }

            match *potential_stores {
                [] => println!("[{}] [BUG?] Reading uninitialized memory", name),
                [store] if spatial_overlap(&opt_ctx.ctx.func, store, load) == SpatialOverlap::Full && temporal_order(&opt_ctx.ctx, store, load) == TemporalOrder::DefinitivelyBefore => {
                    // Only one store could have been the origin of the value.
                    let store_ebb = opt_ctx.ctx.func.layout.inst_ebb(store).unwrap();
                    let stored_value = opt_ctx.ctx.func.dfg.inst_args(store)[0];
                    let stored_type = opt_ctx.ctx.func.dfg.value_type(stored_value);
                    if stored_type == loaded_type && store_ebb == load_ebb {
                        println!("Store to load forward {} -> {}", store, load);
                        opt_ctx.ctx.func.dfg.detach_results(load);
                        opt_ctx.ctx.func.dfg.replace(load).nop();
                        opt_ctx.ctx.func.dfg.change_to_alias(loaded_value, stored_value);
                        users.stack_load.remove(&load);
                    }
                }
                _ => {} // FIXME implement this
            }
        }

        for store in users.stack_store.clone().drain() {
            let ctx = &*opt_ctx.ctx;
            let potential_loads = users.stack_load.iter().cloned().filter(|&load| {
                match spatial_overlap(&ctx.func, store, load) {
                    SpatialOverlap::No => false, // Can never be the source of the loaded value.
                    SpatialOverlap::Partial | SpatialOverlap::Full => true,
                }
            }).filter(|&load| {
                match temporal_order(ctx, store, load) {
                    TemporalOrder::NeverBefore => false, // Can never be the source of the loaded value.
                    TemporalOrder::MaybeBefore | TemporalOrder::DefinitivelyBefore => true,
                }
            }).collect::<Vec<Inst>>();

            for &load in &potential_loads {
                println!(
                    "Potential load from store {} <- {} ({:?}, {:?})",
                    opt_ctx.ctx.func.dfg.display_inst(load, None),
                    opt_ctx.ctx.func.dfg.display_inst(store, None),
                    spatial_overlap(&ctx.func, store, load),
                    temporal_order(&*opt_ctx.ctx, store, load),
                );
            }

            if potential_loads.is_empty() {
                // Never loaded; can safely remove all stores and the stack slot.
                // FIXME also remove stores when there is always a next store before a load.
                println!("[{}] Remove dead stack store {} of {}", name, opt_ctx.ctx.func.dfg.display_inst(store, None), stack_slot.0);
                opt_ctx.ctx.func.dfg.replace(store).nop();
                users.stack_store.remove(&store);
            }
        }

        if users.stack_store.is_empty() && users.stack_load.is_empty() {
            // FIXME make stack_slot zero sized.
        }
    }

    println!();
}

fn combine_stack_addr_with_load_store(func: &mut Function) {
    // Turn load and store into stack_load and stack_store when possible.
    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::Load { opcode: Opcode::Load, arg: addr, flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUD: stack_load.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            let ty = cursor.func.dfg.ctrl_typevar(inst);
                            cursor.func.dfg.replace(inst).stack_load(ty, stack_slot, combined_offset);
                        }
                    }
                }
                InstructionData::Store { opcode: Opcode::Store, args: [value, addr], flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUND: stack_store.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            cursor.func.dfg.replace(inst).stack_store(value, stack_slot, combined_offset);
                        }
                    }
                }
                _ => {}
            }
        }
    }
}

fn remove_unused_stack_addr_and_stack_load(func: &mut Function, stack_slot_usage_map: &mut BTreeMap<OrdStackSlot, StackSlotUsage>) {
    // FIXME incrementally rebuild on each call?
    let mut stack_addr_load_insts_users = HashMap::<Inst, HashSet<Inst>>::new();

    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            for &arg in cursor.func.dfg.inst_args(inst) {
                if let ValueDef::Result(arg_origin, 0) = cursor.func.dfg.value_def(arg) {
                    match cursor.func.dfg[arg_origin].opcode() {
                        Opcode::StackAddr | Opcode::StackLoad => {
                            stack_addr_load_insts_users.entry(arg_origin).or_insert_with(HashSet::new).insert(inst);
                        }
                        _ => {}
                    }
                }
            }
        }
    }

    for inst in stack_addr_load_insts_users.keys() {
        let mut is_recorded_stack_addr_or_stack_load = false;
        for stack_slot_users in stack_slot_usage_map.values() {
            is_recorded_stack_addr_or_stack_load |= stack_slot_users.stack_addr.contains(inst) || stack_slot_users.stack_load.contains(inst);
        }
        assert!(is_recorded_stack_addr_or_stack_load);
    }

    // Replace all unused stack_addr and stack_load instructions with nop.
    for stack_slot_users in stack_slot_usage_map.values_mut() {
        // FIXME remove clone
        for &inst in stack_slot_users.stack_addr.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_addr.remove(&inst);
            }
        }

        for &inst in stack_slot_users.stack_load.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_load.remove(&inst);
            }
        }
    }
}

fn try_get_stack_slot_and_offset_for_addr(func: &Function, addr: Value) -> Option<(StackSlot, Offset32)> {
    if let ValueDef::Result(addr_inst, 0) = func.dfg.value_def(addr) {
        if let InstructionData::StackLoad {
            opcode: Opcode::StackAddr,
            stack_slot,
            offset,
        } = func.dfg[addr_inst] {
            return Some((stack_slot, offset));
        }
    }
    None
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum SpatialOverlap {
    No,
    Partial,
    Full,
}

fn spatial_overlap(func: &Function, src: Inst, dest: Inst) -> SpatialOverlap {
    fn inst_info(func: &Function, inst: Inst) -> (StackSlot, Offset32, u32) {
        match func.dfg[inst] {
            InstructionData::StackLoad {
                opcode: Opcode::StackAddr,
                stack_slot,
                offset,
            }
            | InstructionData::StackLoad {
                opcode: Opcode::StackLoad,
                stack_slot,
                offset,
            }
            | InstructionData::StackStore {
                opcode: Opcode::StackStore,
                stack_slot,
                offset,
                arg: _,
            } => (stack_slot, offset, func.dfg.ctrl_typevar(inst).bytes()),
            _ => unreachable!("{:?}", func.dfg[inst]),
        }
    }

    debug_assert_ne!(src, dest);

    let (src_ss, src_offset, src_size) = inst_info(func, src);
    let (dest_ss, dest_offset, dest_size) = inst_info(func, dest);

    if src_ss != dest_ss {
        return SpatialOverlap::No;
    }

    if src_offset == dest_offset && src_size == dest_size {
        return SpatialOverlap::Full;
    }

    let src_end: i64 = src_offset.try_add_i64(i64::from(src_size)).unwrap().into();
    let dest_end: i64 = dest_offset.try_add_i64(i64::from(dest_size)).unwrap().into();
    if src_end <= dest_offset.into() || dest_end <= src_offset.into() {
        return SpatialOverlap::No;
    }

    SpatialOverlap::Partial
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum TemporalOrder {
    /// `src` will never be executed before `dest`.
    NeverBefore,

    /// `src` may be executed before `dest`.
    MaybeBefore,

    /// `src` will always be executed before `dest`.
    /// There may still be other instructions in between.
    DefinitivelyBefore,
}

fn temporal_order(ctx: &Context, src: Inst, dest: Inst) -> TemporalOrder {
    debug_assert_ne!(src, dest);

    let src_ebb = ctx.func.layout.inst_ebb(src).unwrap();
    let dest_ebb = ctx.func.layout.inst_ebb(dest).unwrap();
    if src_ebb == dest_ebb {
        use std::cmp::Ordering::*;
        match ctx.func.layout.cmp(src, dest) {
            Less => TemporalOrder::DefinitivelyBefore,
            Equal => unreachable!(),
            Greater => TemporalOrder::MaybeBefore, // FIXME use dominator to check for loops
        }
    } else {
        // FIXME O(stack_load count * ebb count)
        // FIXME reuse memory allocations
        // FIXME return DefinitivelyBefore is src dominates dest
        let mut visited = EntitySet::new();
        let mut todo = EntitySet::new();
        todo.insert(dest_ebb);
        while let Some(ebb) = todo.pop() {
            if visited.contains(ebb) {
                continue;
            }
            visited.insert(ebb);
            if ebb == src_ebb {
                return TemporalOrder::MaybeBefore;
            }
            for bb in ctx.cfg.pred_iter(ebb) {
                todo.insert(bb.ebb);
            }
        }
        TemporalOrder::NeverBefore
    }
}
<|endoftext|>


<commit_before>/// RouteGraph defines a Directed Acyclic Graph of Effects.
/// The edges connecting each Effect have a source and destination slot, tag, and channel.
/// Edges are also allowed to go to null, in which case they only have a destination slot and
/// channel. These are outputs.
/// Edges can also COME from null, in which case the source has the format (slot, channel)
extern crate online_dag;
use self::online_dag::iodag;
use self::online_dag::iodag::IODag;

use super::effect::Effect;

#[derive(Copy, Clone, Hash, Eq, PartialEq)]
struct EdgeWeight {
    from_slot: u32,
    from_ch: u8,
    to_slot: u32,
    to_ch: u8,
}

struct NodeData {
    effect: Effect,
}

type DagImpl=IODag<NodeData, EdgeWeight>;
pub type NodeHandle=iodag::NodeHandle;
pub type Edge=iodag::Edge<EdgeWeight>;
pub struct RouteGraph {
    dag: DagImpl,
}

impl RouteGraph {
    pub fn new() -> Self {
        RouteGraph {
            dag: DagImpl::new(),
        }
    }
    /*pub fn iter_topo_rev(&self) -> impl Iterator<Item=poscostdag::NodeHandle<RouteNode, RouteEdge>> {
        self.dag.iter_topo_rev(&self.root)
    }*/
    /*pub fn children_of(&self, of: &RouteNodeHandle) -> impl Iterator<Item=RouteHalfEdge> {
        self.dag.children(of)
    }
    pub fn add_node(&mut self, data: RouteNode) -> RouteNodeHandle {
        self.dag.add_node(data)
    }
    pub fn add_edge(&mut self, from: &RouteNodeHandle, to: &RouteNodeHandle, data: RouteEdge) -> Result<(), ()> {
        self.dag.add_edge(from, to, data)
    }
    pub fn rm_edge(&mut self, from: &RouteNodeHandle, to: &RouteNodeHandle, data: RouteEdge) {
        self.dag.rm_edge(from, to, data);
    }
    /// Return only the inputs into the left (i.e. non-delayed) channel of `of`
    pub fn left_children_of(&self, of: &RouteNodeHandle) -> impl
      Iterator<Item=RouteHalfEdge> {
        self.dag.children(of).filter(|edge| edge.weight().is_left())
    }
    /// Return only the inputs into the right (i.e. non-delayed) channel of `of`
    pub fn right_children_of(&self, of: &RouteNodeHandle) -> impl
      Iterator<Item=RouteHalfEdge> {
        self.dag.children(of).filter(|edge| edge.weight().is_right())
    }
    /// Returns 1 + the index of the highest channel number.
    /// e.g. if we have ch0 and ch2 (no ch1), this returns 3.
    pub fn n_channels(&self) -> u32 {
        self.right_children_of(&self.root())
            .map(|edge| edge.weight().slot_idx())
            .max()
            .unwrap_or(0u32) // default no channels
    }
    pub fn make_channel_output(&mut self, node: &RouteNodeHandle, ch: u32) {
        let root = &self.root.clone();
        self.add_edge(&root, &node, RouteEdge::new_right(ch));
    }*/
}
<commit_msg>Impl node/edge addition and removal<commit_after>/// RouteGraph defines a Directed Acyclic Graph of Effects.
/// The edges connecting each Effect have a source and destination slot, tag, and channel.
/// Edges are also allowed to go to null, in which case they only have a destination slot and
/// channel. These are outputs.
/// Edges can also COME from null, in which case the source has the format (slot, channel)
extern crate online_dag;
use self::online_dag::iodag;
use self::online_dag::iodag::IODag;

use super::effect::Effect;

type DagImpl=IODag<NodeData, EdgeWeight>;
pub type NodeHandle=iodag::NodeHandle;
pub type Edge=iodag::Edge<EdgeWeight>;

#[derive(Copy, Clone, Hash, Eq, PartialEq)]
struct EdgeWeight {
    from_slot: u32,
    from_ch: u8,
    to_slot: u32,
    to_ch: u8,
}

struct NodeData {
    effect: Effect,
}

pub struct RouteGraph {
    dag: DagImpl,
}

impl RouteGraph {
    pub fn new() -> Self {
        RouteGraph {
            dag: DagImpl::new(),
        }
    }
    pub fn add_node(&mut self, node: NodeData) -> NodeHandle {
        self.dag.add_node(node)
    }
    pub fn add_edge(&mut self, edge: Edge) -> Result<(), ()> {
        self.dag.add_edge(edge)
    }
    pub fn del_edge(&mut self, edge: Edge) {
        self.dag.del_edge(edge)
    }
}
<|endoftext|>


<commit_before>use crate::schema::FacetOptions;
use crate::schema::TextOptions;
use crate::schema::{is_valid_field_name, IntOptions};

use crate::schema::bytes_options::BytesOptions;
use crate::schema::FieldType;
use serde::de::{self, MapAccess, Visitor};
use serde::ser::SerializeStruct;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::fmt;

/// A `FieldEntry` represents a field and its configuration.
/// `Schema` are a collection of `FieldEntry`
///
/// It consists of
/// - a field name
/// - a field type, itself wrapping up options describing
/// how the field should be indexed.
#[derive(Clone, Debug, PartialEq)]
pub struct FieldEntry {
    name: String,
    field_type: FieldType,
}

impl FieldEntry {
    /// Creates a new u64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_text(field_name: String, text_options: TextOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::Str(text_options),
        }
    }

    /// Creates a new u64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_u64(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::U64(field_type),
        }
    }

    /// Creates a new i64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_i64(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::I64(field_type),
        }
    }

    /// Creates a new f64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_f64(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::F64(field_type),
        }
    }

    /// Creates a new date field entry in the schema, given
    /// a name, and some options.
    pub fn new_date(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::Date(field_type),
        }
    }

    /// Creates a field entry for a facet.
    pub fn new_facet(field_name: String, field_type: FacetOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::HierarchicalFacet(field_type),
        }
    }

    /// Creates a field entry for a bytes field
    pub fn new_bytes(field_name: String, bytes_type: BytesOptions) -> FieldEntry {
        FieldEntry {
            name: field_name,
            field_type: FieldType::Bytes(bytes_type),
        }
    }

    /// Returns the name of the field
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the field type
    pub fn field_type(&self) -> &FieldType {
        &self.field_type
    }

    /// Returns true iff the field is indexed
    pub fn is_indexed(&self) -> bool {
        match self.field_type {
            FieldType::Str(ref options) => options.get_indexing_options().is_some(),
            FieldType::U64(ref options)
            | FieldType::I64(ref options)
            | FieldType::F64(ref options)
            | FieldType::Date(ref options) => options.is_indexed(),
            FieldType::HierarchicalFacet(ref options) => options.is_indexed(),
            FieldType::Bytes(ref options) => options.is_indexed(),
        }
    }

    /// Returns true iff the field is a int (signed or unsigned) fast field
    pub fn is_fast(&self) -> bool {
        match self.field_type {
            FieldType::U64(ref options)
            | FieldType::I64(ref options)
            | FieldType::Date(ref options)
            | FieldType::F64(ref options) => options.is_fast(),
            _ => false,
        }
    }

    /// Returns true iff the field is stored
    pub fn is_stored(&self) -> bool {
        match self.field_type {
            FieldType::U64(ref options)
            | FieldType::I64(ref options)
            | FieldType::F64(ref options)
            | FieldType::Date(ref options) => options.is_stored(),
            FieldType::Str(ref options) => options.is_stored(),
            FieldType::HierarchicalFacet(ref options) => options.is_stored(),
            FieldType::Bytes(ref options) => options.is_stored(),
        }
    }
}

impl Serialize for FieldEntry {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut s = serializer.serialize_struct("field_entry", 3)?;
        s.serialize_field("name", &self.name)?;

        match self.field_type {
            FieldType::Str(ref options) => {
                s.serialize_field("type", "text")?;
                s.serialize_field("options", options)?;
            }
            FieldType::U64(ref options) => {
                s.serialize_field("type", "u64")?;
                s.serialize_field("options", options)?;
            }
            FieldType::I64(ref options) => {
                s.serialize_field("type", "i64")?;
                s.serialize_field("options", options)?;
            }
            FieldType::F64(ref options) => {
                s.serialize_field("type", "f64")?;
                s.serialize_field("options", options)?;
            }
            FieldType::Date(ref options) => {
                s.serialize_field("type", "date")?;
                s.serialize_field("options", options)?;
            }
            FieldType::HierarchicalFacet(ref options) => {
                s.serialize_field("type", "hierarchical_facet")?;
                s.serialize_field("options", options)?;
            }
            FieldType::Bytes(ref options) => {
                s.serialize_field("type", "bytes")?;
                s.serialize_field("options", options)?;
            }
        }

        s.end()
    }
}

impl<'de> Deserialize<'de> for FieldEntry {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(field_identifier, rename_all = "lowercase")]
        enum Field {
            Name,
            Type,
            Options,
        }

        const FIELDS: &[&str] = &["name", "type", "options"];

        struct FieldEntryVisitor;

        impl<'de> Visitor<'de> for FieldEntryVisitor {
            type Value = FieldEntry;

            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                formatter.write_str("struct FieldEntry")
            }

            fn visit_map<V>(self, mut map: V) -> Result<FieldEntry, V::Error>
            where
                V: MapAccess<'de>,
            {
                let mut name = None;
                let mut ty = None;
                let mut field_type = None;
                while let Some(key) = map.next_key()? {
                    match key {
                        Field::Name => {
                            if name.is_some() {
                                return Err(de::Error::duplicate_field("name"));
                            }
                            name = Some(map.next_value()?);
                        }
                        Field::Type => {
                            if ty.is_some() {
                                return Err(de::Error::duplicate_field("type"));
                            }
                            let type_string = map.next_value::<String>()?;
                            match type_string.as_str() {
                                "text" | "u64" | "i64" | "f64" | "date" | "bytes"
                                | "hierarchical_facet" => {
                                    // These types require additional options to create a field_type
                                }
                                _ => panic!("unhandled type"),
                            }
                            ty = Some(type_string);
                        }
                        Field::Options => match ty {
                            None => {
                                let msg = "The `type` field must be \
                                           specified before `options`";
                                return Err(de::Error::custom(msg));
                            }
                            Some(ref ty) => match ty.as_str() {
                                "text" => field_type = Some(FieldType::Str(map.next_value()?)),
                                "u64" => field_type = Some(FieldType::U64(map.next_value()?)),
                                "i64" => field_type = Some(FieldType::I64(map.next_value()?)),
                                "f64" => field_type = Some(FieldType::F64(map.next_value()?)),
                                "date" => field_type = Some(FieldType::Date(map.next_value()?)),
                                "bytes" => field_type = Some(FieldType::Bytes(map.next_value()?)),
                                "hierarchical_facet" => {
                                    field_type =
                                        Some(FieldType::HierarchicalFacet(map.next_value()?))
                                }
                                _ => {
                                    let msg = format!("Unrecognised type {}", ty);
                                    return Err(de::Error::custom(msg));
                                }
                            },
                        },
                    }
                }

                let name = name.ok_or_else(|| de::Error::missing_field("name"))?;
                ty.ok_or_else(|| de::Error::missing_field("ty"))?;
                let field_type = field_type.ok_or_else(|| de::Error::missing_field("options"))?;

                Ok(FieldEntry { name, field_type })
            }
        }

        deserializer.deserialize_struct("field_entry", FIELDS, FieldEntryVisitor)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::TEXT;
    use serde_json;

    #[test]
    #[should_panic]
    fn test_invalid_field_name_should_panic() {
        FieldEntry::new_text("-hello".to_string(), TEXT);
    }

    #[test]
    fn test_json_serialization() {
        let field_value = FieldEntry::new_text(String::from("title"), TEXT);

        let expected = r#"{
  "name": "title",
  "type": "text",
  "options": {
    "indexing": {
      "record": "position",
      "tokenizer": "default"
    },
    "stored": false
  }
}"#;
        let field_value_json = serde_json::to_string_pretty(&field_value).unwrap();

        assert_eq!(expected, &field_value_json);

        let field_value: FieldEntry = serde_json::from_str(expected).unwrap();

        assert_eq!("title", field_value.name);

        match field_value.field_type {
            FieldType::Str(_) => assert!(true),
            _ => panic!("expected FieldType::Str"),
        }
    }
}
<commit_msg>add FieldEntry constructor, closes #1086 (#1090)<commit_after>use crate::schema::FacetOptions;
use crate::schema::TextOptions;
use crate::schema::{is_valid_field_name, IntOptions};

use crate::schema::bytes_options::BytesOptions;
use crate::schema::FieldType;
use serde::de::{self, MapAccess, Visitor};
use serde::ser::SerializeStruct;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::fmt;

/// A `FieldEntry` represents a field and its configuration.
/// `Schema` are a collection of `FieldEntry`
///
/// It consists of
/// - a field name
/// - a field type, itself wrapping up options describing
/// how the field should be indexed.
#[derive(Clone, Debug, PartialEq)]
pub struct FieldEntry {
    name: String,
    field_type: FieldType,
}

impl FieldEntry {
    /// Creates a new field entry given a name and a field type
    pub fn new(field_name: String, field_type: FieldType) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type,
        }
    }

    /// Creates a new u64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_text(field_name: String, text_options: TextOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::Str(text_options),
        }
    }

    /// Creates a new u64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_u64(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::U64(field_type),
        }
    }

    /// Creates a new i64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_i64(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::I64(field_type),
        }
    }

    /// Creates a new f64 field entry in the schema, given
    /// a name, and some options.
    pub fn new_f64(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::F64(field_type),
        }
    }

    /// Creates a new date field entry in the schema, given
    /// a name, and some options.
    pub fn new_date(field_name: String, field_type: IntOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::Date(field_type),
        }
    }

    /// Creates a field entry for a facet.
    pub fn new_facet(field_name: String, field_type: FacetOptions) -> FieldEntry {
        assert!(is_valid_field_name(&field_name));
        FieldEntry {
            name: field_name,
            field_type: FieldType::HierarchicalFacet(field_type),
        }
    }

    /// Creates a field entry for a bytes field
    pub fn new_bytes(field_name: String, bytes_type: BytesOptions) -> FieldEntry {
        FieldEntry {
            name: field_name,
            field_type: FieldType::Bytes(bytes_type),
        }
    }

    /// Returns the name of the field
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the field type
    pub fn field_type(&self) -> &FieldType {
        &self.field_type
    }

    /// Returns true iff the field is indexed
    pub fn is_indexed(&self) -> bool {
        match self.field_type {
            FieldType::Str(ref options) => options.get_indexing_options().is_some(),
            FieldType::U64(ref options)
            | FieldType::I64(ref options)
            | FieldType::F64(ref options)
            | FieldType::Date(ref options) => options.is_indexed(),
            FieldType::HierarchicalFacet(ref options) => options.is_indexed(),
            FieldType::Bytes(ref options) => options.is_indexed(),
        }
    }

    /// Returns true iff the field is a int (signed or unsigned) fast field
    pub fn is_fast(&self) -> bool {
        match self.field_type {
            FieldType::U64(ref options)
            | FieldType::I64(ref options)
            | FieldType::Date(ref options)
            | FieldType::F64(ref options) => options.is_fast(),
            _ => false,
        }
    }

    /// Returns true iff the field is stored
    pub fn is_stored(&self) -> bool {
        match self.field_type {
            FieldType::U64(ref options)
            | FieldType::I64(ref options)
            | FieldType::F64(ref options)
            | FieldType::Date(ref options) => options.is_stored(),
            FieldType::Str(ref options) => options.is_stored(),
            FieldType::HierarchicalFacet(ref options) => options.is_stored(),
            FieldType::Bytes(ref options) => options.is_stored(),
        }
    }
}

impl Serialize for FieldEntry {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut s = serializer.serialize_struct("field_entry", 3)?;
        s.serialize_field("name", &self.name)?;

        match self.field_type {
            FieldType::Str(ref options) => {
                s.serialize_field("type", "text")?;
                s.serialize_field("options", options)?;
            }
            FieldType::U64(ref options) => {
                s.serialize_field("type", "u64")?;
                s.serialize_field("options", options)?;
            }
            FieldType::I64(ref options) => {
                s.serialize_field("type", "i64")?;
                s.serialize_field("options", options)?;
            }
            FieldType::F64(ref options) => {
                s.serialize_field("type", "f64")?;
                s.serialize_field("options", options)?;
            }
            FieldType::Date(ref options) => {
                s.serialize_field("type", "date")?;
                s.serialize_field("options", options)?;
            }
            FieldType::HierarchicalFacet(ref options) => {
                s.serialize_field("type", "hierarchical_facet")?;
                s.serialize_field("options", options)?;
            }
            FieldType::Bytes(ref options) => {
                s.serialize_field("type", "bytes")?;
                s.serialize_field("options", options)?;
            }
        }

        s.end()
    }
}

impl<'de> Deserialize<'de> for FieldEntry {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(field_identifier, rename_all = "lowercase")]
        enum Field {
            Name,
            Type,
            Options,
        }

        const FIELDS: &[&str] = &["name", "type", "options"];

        struct FieldEntryVisitor;

        impl<'de> Visitor<'de> for FieldEntryVisitor {
            type Value = FieldEntry;

            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                formatter.write_str("struct FieldEntry")
            }

            fn visit_map<V>(self, mut map: V) -> Result<FieldEntry, V::Error>
            where
                V: MapAccess<'de>,
            {
                let mut name = None;
                let mut ty = None;
                let mut field_type = None;
                while let Some(key) = map.next_key()? {
                    match key {
                        Field::Name => {
                            if name.is_some() {
                                return Err(de::Error::duplicate_field("name"));
                            }
                            name = Some(map.next_value()?);
                        }
                        Field::Type => {
                            if ty.is_some() {
                                return Err(de::Error::duplicate_field("type"));
                            }
                            let type_string = map.next_value::<String>()?;
                            match type_string.as_str() {
                                "text" | "u64" | "i64" | "f64" | "date" | "bytes"
                                | "hierarchical_facet" => {
                                    // These types require additional options to create a field_type
                                }
                                _ => panic!("unhandled type"),
                            }
                            ty = Some(type_string);
                        }
                        Field::Options => match ty {
                            None => {
                                let msg = "The `type` field must be \
                                           specified before `options`";
                                return Err(de::Error::custom(msg));
                            }
                            Some(ref ty) => match ty.as_str() {
                                "text" => field_type = Some(FieldType::Str(map.next_value()?)),
                                "u64" => field_type = Some(FieldType::U64(map.next_value()?)),
                                "i64" => field_type = Some(FieldType::I64(map.next_value()?)),
                                "f64" => field_type = Some(FieldType::F64(map.next_value()?)),
                                "date" => field_type = Some(FieldType::Date(map.next_value()?)),
                                "bytes" => field_type = Some(FieldType::Bytes(map.next_value()?)),
                                "hierarchical_facet" => {
                                    field_type =
                                        Some(FieldType::HierarchicalFacet(map.next_value()?))
                                }
                                _ => {
                                    let msg = format!("Unrecognised type {}", ty);
                                    return Err(de::Error::custom(msg));
                                }
                            },
                        },
                    }
                }

                let name = name.ok_or_else(|| de::Error::missing_field("name"))?;
                ty.ok_or_else(|| de::Error::missing_field("ty"))?;
                let field_type = field_type.ok_or_else(|| de::Error::missing_field("options"))?;

                Ok(FieldEntry { name, field_type })
            }
        }

        deserializer.deserialize_struct("field_entry", FIELDS, FieldEntryVisitor)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::TEXT;
    use serde_json;

    #[test]
    #[should_panic]
    fn test_invalid_field_name_should_panic() {
        FieldEntry::new_text("-hello".to_string(), TEXT);
    }

    #[test]
    fn test_json_serialization() {
        let field_value = FieldEntry::new_text(String::from("title"), TEXT);

        let expected = r#"{
  "name": "title",
  "type": "text",
  "options": {
    "indexing": {
      "record": "position",
      "tokenizer": "default"
    },
    "stored": false
  }
}"#;
        let field_value_json = serde_json::to_string_pretty(&field_value).unwrap();

        assert_eq!(expected, &field_value_json);

        let field_value: FieldEntry = serde_json::from_str(expected).unwrap();

        assert_eq!("title", field_value.name);

        match field_value.field_type {
            FieldType::Str(_) => assert!(true),
            _ => panic!("expected FieldType::Str"),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::net::SocketAddr;

use futures::sync::mpsc::{self, UnboundedSender};
use futures::sync::oneshot::{self, Sender};
use futures::{Future, Sink, Stream, stream};
use grpc::{Environment, ChannelBuilder, WriteFlags};
use kvproto::raft_serverpb::RaftMessage;
use kvproto::tikvpb_grpc::TikvClient;

const MAX_GRPC_RECV_MSG_LEN: usize = 10 * 1024 * 1024;
const MAX_GRPC_SEND_MSG_LEN: usize = 10 * 1024 * 1024;
const INITIAL_BUFFER_CAP: usize = 1024;

use util::collections::HashMap;
use super::{Error, Result, Config};

struct Conn {
    stream: UnboundedSender<Vec<(RaftMessage, WriteFlags)>>,
    buffer: Option<Vec<(RaftMessage, WriteFlags)>>,
    store_id: u64,
    alive: Arc<AtomicBool>,

    _client: TikvClient,
    _close: Sender<()>,
}

impl Conn {
    fn new(env: Arc<Environment>, addr: SocketAddr, cfg: &Config, store_id: u64) -> Conn {
        info!("server: new connection with tikv endpoint: {}", addr);

        let alive = Arc::new(AtomicBool::new(true));
        let alive1 = alive.clone();
        let channel = ChannelBuilder::new(env)
            .stream_initial_window_size(cfg.grpc_stream_initial_window_size)
            .max_receive_message_len(MAX_GRPC_RECV_MSG_LEN)
            .max_send_message_len(MAX_GRPC_SEND_MSG_LEN)
            .connect(&format!("{}", addr));
        let client = TikvClient::new(channel);
        let (tx, rx) = mpsc::unbounded();
        let (tx_close, rx_close) = oneshot::channel();
        let (sink, _) = client.raft();
        client.spawn(rx_close.map_err(|_| ())
            .select(sink.sink_map_err(Error::from)
                .send_all(rx.map(|msgs: Vec<(RaftMessage, WriteFlags)>| {
                        stream::iter::<_, _, ()>(msgs.into_iter().map(Ok))
                    })
                    .flatten()
                    .map_err(|_| Error::Sink))
                .then(move |r| {
                    alive.store(false, Ordering::SeqCst);
                    r
                })
                .map(|_| ())
                .map_err(move |e| warn!("send raftmessage to {} failed: {:?}", addr, e)))
            .map(|_| ())
            .map_err(|_| ()));
        Conn {
            stream: tx,
            buffer: Some(Vec::with_capacity(INITIAL_BUFFER_CAP)),
            store_id: store_id,
            alive: alive1,

            _client: client,
            _close: tx_close,
        }
    }
}

/// `RaftClient` is used for sending raft messages to other stores.
pub struct RaftClient {
    env: Arc<Environment>,
    conns: HashMap<(SocketAddr, usize), Conn>,
    pub addrs: HashMap<u64, SocketAddr>,
    cfg: Config,
}

impl RaftClient {
    pub fn new(env: Arc<Environment>, cfg: Config) -> RaftClient {
        RaftClient {
            env: env,
            conns: HashMap::default(),
            addrs: HashMap::default(),
            cfg: cfg,
        }
    }

    fn get_conn(&mut self, addr: SocketAddr, region_id: u64, store_id: u64) -> &mut Conn {
        let index = region_id as usize % self.cfg.grpc_raft_conn_num;
        let cfg = &self.cfg;
        let env = &self.env;
        self.conns
            .entry((addr, index))
            .or_insert_with(|| Conn::new(env.clone(), addr, cfg, store_id))
    }

    pub fn send(&mut self, store_id: u64, addr: SocketAddr, msg: RaftMessage) -> Result<()> {
        let mut conn = self.get_conn(addr, msg.region_id, store_id);
        conn.buffer.as_mut().unwrap().push((msg, WriteFlags::default().buffer_hint(true)));
        Ok(())
    }

    pub fn flush(&mut self) {
        let addrs = &mut self.addrs;
        self.conns.retain(|&mut (addr, _), conn| {
            let store_id = conn.store_id;
            if !conn.alive.load(Ordering::SeqCst) {
                if let Some(addr_current) = addrs.remove(&store_id) {
                    if addr_current != addr {
                        addrs.insert(store_id, addr_current);
                    }
                }
                return false;
            }

            if conn.buffer.as_ref().unwrap().is_empty() {
                return true;
            }

            let mut msgs = conn.buffer.take().unwrap();
            msgs.last_mut().unwrap().1 = WriteFlags::default();
            if let Err(e) = UnboundedSender::send(&conn.stream, msgs) {
                error!("server: drop conn with tikv endpoint {} flush conn error: {:?}",
                       addr,
                       e);

                if let Some(addr_current) = addrs.remove(&store_id) {
                    if addr_current != addr {
                        addrs.insert(store_id, addr_current);
                    }
                }
                return false;
            }

            conn.buffer = Some(Vec::with_capacity(INITIAL_BUFFER_CAP));
            true
        });
    }
}

impl Drop for RaftClient {
    fn drop(&mut self) {
        // Drop conns here to make sure all streams are dropped before Environment.
        self.conns.clear();
    }
}
<commit_msg>server: inc unreachable metrics when raft_client has error. (#2114)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::net::SocketAddr;

use futures::sync::mpsc::{self, UnboundedSender};
use futures::sync::oneshot::{self, Sender};
use futures::{Future, Sink, Stream, stream};
use grpc::{Environment, ChannelBuilder, WriteFlags};
use kvproto::raft_serverpb::RaftMessage;
use kvproto::tikvpb_grpc::TikvClient;

const MAX_GRPC_RECV_MSG_LEN: usize = 10 * 1024 * 1024;
const MAX_GRPC_SEND_MSG_LEN: usize = 10 * 1024 * 1024;
const INITIAL_BUFFER_CAP: usize = 1024;

use util::collections::HashMap;
use super::{Error, Result, Config};
use super::metrics::*;

struct Conn {
    stream: UnboundedSender<Vec<(RaftMessage, WriteFlags)>>,
    buffer: Option<Vec<(RaftMessage, WriteFlags)>>,
    store_id: u64,
    alive: Arc<AtomicBool>,

    _client: TikvClient,
    _close: Sender<()>,
}

impl Conn {
    fn new(env: Arc<Environment>, addr: SocketAddr, cfg: &Config, store_id: u64) -> Conn {
        info!("server: new connection with tikv endpoint: {}", addr);

        let alive = Arc::new(AtomicBool::new(true));
        let alive1 = alive.clone();
        let channel = ChannelBuilder::new(env)
            .stream_initial_window_size(cfg.grpc_stream_initial_window_size)
            .max_receive_message_len(MAX_GRPC_RECV_MSG_LEN)
            .max_send_message_len(MAX_GRPC_SEND_MSG_LEN)
            .connect(&format!("{}", addr));
        let client = TikvClient::new(channel);
        let (tx, rx) = mpsc::unbounded();
        let (tx_close, rx_close) = oneshot::channel();
        let (sink, _) = client.raft();
        client.spawn(rx_close.map_err(|_| ())
            .select(sink.sink_map_err(Error::from)
                .send_all(rx.map(|msgs: Vec<(RaftMessage, WriteFlags)>| {
                        stream::iter::<_, _, ()>(msgs.into_iter().map(Ok))
                    })
                    .flatten()
                    .map_err(|_| Error::Sink))
                .then(move |r| {
                    alive.store(false, Ordering::SeqCst);
                    r
                })
                .map(|_| ())
                .map_err(move |e| {
                    let store = store_id.to_string();
                    REPORT_FAILURE_MSG_COUNTER.with_label_values(&["unreachable", &*store]).inc();
                    warn!("send raftmessage to {} failed: {:?}", addr, e);
                }))
            .map(|_| ())
            .map_err(|_| ()));
        Conn {
            stream: tx,
            buffer: Some(Vec::with_capacity(INITIAL_BUFFER_CAP)),
            store_id: store_id,
            alive: alive1,

            _client: client,
            _close: tx_close,
        }
    }
}

/// `RaftClient` is used for sending raft messages to other stores.
pub struct RaftClient {
    env: Arc<Environment>,
    conns: HashMap<(SocketAddr, usize), Conn>,
    pub addrs: HashMap<u64, SocketAddr>,
    cfg: Config,
}

impl RaftClient {
    pub fn new(env: Arc<Environment>, cfg: Config) -> RaftClient {
        RaftClient {
            env: env,
            conns: HashMap::default(),
            addrs: HashMap::default(),
            cfg: cfg,
        }
    }

    fn get_conn(&mut self, addr: SocketAddr, region_id: u64, store_id: u64) -> &mut Conn {
        let index = region_id as usize % self.cfg.grpc_raft_conn_num;
        let cfg = &self.cfg;
        let env = &self.env;
        self.conns
            .entry((addr, index))
            .or_insert_with(|| Conn::new(env.clone(), addr, cfg, store_id))
    }

    pub fn send(&mut self, store_id: u64, addr: SocketAddr, msg: RaftMessage) -> Result<()> {
        let mut conn = self.get_conn(addr, msg.region_id, store_id);
        conn.buffer.as_mut().unwrap().push((msg, WriteFlags::default().buffer_hint(true)));
        Ok(())
    }

    pub fn flush(&mut self) {
        let addrs = &mut self.addrs;
        self.conns.retain(|&mut (addr, _), conn| {
            let store_id = conn.store_id;
            if !conn.alive.load(Ordering::SeqCst) {
                if let Some(addr_current) = addrs.remove(&store_id) {
                    if addr_current != addr {
                        addrs.insert(store_id, addr_current);
                    }
                }
                return false;
            }

            if conn.buffer.as_ref().unwrap().is_empty() {
                return true;
            }

            let mut msgs = conn.buffer.take().unwrap();
            msgs.last_mut().unwrap().1 = WriteFlags::default();
            if let Err(e) = UnboundedSender::send(&conn.stream, msgs) {
                error!("server: drop conn with tikv endpoint {} flush conn error: {:?}",
                       addr,
                       e);

                if let Some(addr_current) = addrs.remove(&store_id) {
                    if addr_current != addr {
                        addrs.insert(store_id, addr_current);
                    }
                }
                return false;
            }

            conn.buffer = Some(Vec::with_capacity(INITIAL_BUFFER_CAP));
            true
        });
    }
}

impl Drop for RaftClient {
    fn drop(&mut self) {
        // Drop conns here to make sure all streams are dropped before Environment.
        self.conns.clear();
    }
}
<|endoftext|>