<commit_before>extern mod std;
extern mod extra;

use basics::*;
use std::fmt;
use sim::Environment;
use std::hashmap::HashSet;

static ELECTION_TIMEOUT : (uint, uint) = (150000, 299999);
static HEARTBEAT_TIMEOUT : (uint, uint) = (75000, 75000);

#[deriving(Clone)]
pub struct Message {
    from: ServerID,
    to: ServerID,
    body: MessageBody,
}

impl fmt::Default for Message {
    fn fmt(msg: &Message, f: &mut fmt::Formatter) {
        write!(f.buf, "{} to {}: {}",
               *msg.from,
               *msg.to,
               msg.body);
    }
}

#[deriving(Eq, Clone)]
enum VoteGranted {
    GRANTED,
    TERM_STALE,
    LOG_STALE,
    VOTED,
}

impl fmt::Default for VoteGranted {
    fn fmt(g: &VoteGranted, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", match *g {
            GRANTED    => "GRANTED",
            TERM_STALE => "TERM_STALE",
            LOG_STALE  => "LOG_STALE",
            VOTED      => "VOTED",
        });
    }
}

#[deriving(Clone)]
pub enum MessageBody {
    RequestVoteRequest {
        term: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        term: Term,
        granted: VoteGranted,
        logOk: bool,
    },
    AppendEntriesRequest {
        term: Term,
        seqno: uint,
    },
    AppendEntriesResponse {
        term: Term,
        seqno: uint,
    },
}

impl fmt::Default for MessageBody {
    fn fmt(b: &MessageBody, f: &mut fmt::Formatter) {
        match *b {
            RequestVoteRequest{term, lastLogIndex} =>
                write!(f.buf, "RequestVoteRequest(term: {}, lastLogIndex: {})",
                       term, lastLogIndex),
            RequestVoteResponse{term, granted, logOk} =>
                write!(f.buf, "RequestVoteResponse(term: {}, granted: {}, logOk: {})",
                       term, granted, logOk),
            AppendEntriesRequest{term, seqno} =>
                write!(f.buf, "AppendEntriesRequest(term: {}, seqno: {})",
                       term, seqno),
            AppendEntriesResponse{term, seqno} =>
                write!(f.buf, "AppendEntriesResponse(term: {}, seqno: {})",
                       term, seqno),
        }
    }
}

enum ServerState {
    Follower  {
        timer: Time,
    },
    Candidate {
        timer: Time,
        votes: HashSet<ServerID>,
        /**
         * If true, when the timer fires, the candidate should start a new election.
         * If false, when the timer fires, set a new timer.
         * Always true for submission algo, but sometimes false in hesitant, hesitant2.
         */
        should_retry: bool,
        pre: bool,
        max_term: Term,
        endorsements: HashSet<ServerID>,
    },
    Leader {
        timer: Time,
        heartbeat_seqno: uint,
        acks: HashSet<ServerID>,
        start_time: Time,
    },
}

impl ServerState {
    pub fn to_char(&self) -> char {
        match *self {
            Follower {..}  => 'F',
            Candidate {..} => 'C',
            Leader {..}    => 'L',
        }
    }
}

impl fmt::Default for HashSet<ServerID> {
    fn fmt(set: &HashSet<ServerID>, f: &mut fmt::Formatter) {
        write!(f.buf, "\\{");
        let mut sorted = set.iter().to_owned_vec();
        extra::sort::tim_sort(sorted);
        for s in sorted.iter() {
            write!(f.buf, "{}", **s);
        }
        write!(f.buf, "\\}");
    }
}

impl fmt::Default for Configuration {
    fn fmt(config: &Configuration, f: &mut fmt::Formatter) {
        write!(f.buf, "[");
        for c in config.iter() {
            write!(f.buf, "{}", *c);
        }
        write!(f.buf, "]");
    }
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => {
                write!(f.buf, "Follower(timer: {})",
                       timer)
            },
            Candidate{timer, votes: ref votes, should_retry, pre, max_term, endorsements: ref endorsements} => {
                write!(f.buf, "Candidate(timer: {}, votes: {}, should_retry: {}, pre: {}, max_term: {}, endorsements: {})",
                       timer, *votes, should_retry, pre, max_term, *endorsements)
            },
            Leader{timer, heartbeat_seqno, acks: ref acks, start_time} => {
                write!(f.buf, "Leader(timer: {}, seqno: {}, acks: {}, start_time: {})",
                       timer, heartbeat_seqno, *acks, start_time)
            },
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Configuration(~[HashSet<ServerID>]);

impl Configuration {
    fn is_quorum(&self, servers: &HashSet<ServerID>) -> bool {
        for c in self.iter() {
            if c.intersection(servers).len() <= c.len() / 2 {
                return false;
            }
        }
        return true;
    }
}

struct Server {
    id: ServerID,
    config: Configuration,
    algorithm: ~str,
    term: Term,
    state: ServerState,
    vote: Option<ServerID>,
    lastLogIndex: Index,
}

impl fmt::Default for Server {
    fn fmt(server: &Server, f: &mut fmt::Formatter) {
        write!(f.buf, "Server(id: {}, term: {}, log: {}, vote: {}, state: {})",
               server.id,
               server.term,
               server.lastLogIndex,
               server.vote,
               server.state)
    }
}

impl Server {
    pub fn new(id : ServerID, config: Configuration, env: &Environment,
               algorithm: &str) -> Server {
        match algorithm {
            "nograntnobump" |
            "stalelognobump" |
            "hesitant"      |
            "hesitant2"     |
            "zookeeper"     |
            "zookeeper2"    |
            "submission"    => {},
            _ => fail!("Unknown algorithm: {}", algorithm)
        };
        Server {
            id: id,
            config: config,
            algorithm: algorithm.into_owned(),
            term: Term(0),
            state: Follower { timer: env.make_time(ELECTION_TIMEOUT) },
            lastLogIndex: Index(0),
            vote: None,
        }
    }
    fn start_new_election(&mut self, env: &mut Environment, forceTerm: Option<Term>) {
        let pre = forceTerm.is_none() && match self.algorithm {
            ~"zookeeper"  => true,
            ~"zookeeper2" => true,
            _ => false,
        };
        if !pre {
            self.term = match forceTerm {
                Some(t) if t > self.term => Term(*t + 1),
                _                        => Term(*self.term + 1),
            };
            self.vote = None;
        }
        self.state = Candidate {
            timer: env.make_time(ELECTION_TIMEOUT),
            votes: newHashSet([self.id]),
            should_retry: match self.algorithm {
                ~"hesitant"  => false,
                ~"hesitant2" => false,
                _ => true,
            },
            pre: pre,
            max_term: Term(0),
            endorsements: newHashSet([self.id]),
        };
        self.vote = Some(self.id);
        env.multicast(self.id, &self.config, &RequestVoteRequest {
            term: self.term,
            lastLogIndex: self.lastLogIndex,
        });
        self.try_become_leader(env);
    }
    pub fn tick(&mut self, env: &mut Environment) {
        let mut need_step_down = false;

        match self.state {
            Follower  {timer, ..} => {
                if timer <= env.clock {
                    self.start_new_election(env, None)
                }
            }
            Candidate {timer, should_retry, ..}
                if timer <= env.clock && should_retry => {
                    self.start_new_election(env, None);
            },
            Candidate {timer: ref mut timer, should_retry, ..}
                if *timer <= env.clock && !should_retry => {
                    *timer = env.make_time(ELECTION_TIMEOUT);
            },
            Candidate {..} => {},
            Leader {timer: ref mut timer, heartbeat_seqno: ref mut heartbeat_seqno, acks: ref mut acks, ..} => {
                if *timer <= env.clock {
                    if (self.config.is_quorum(acks)) {
                        *heartbeat_seqno += 1;
                        acks.clear();
                        acks.insert(self.id);
                        *timer = env.make_time(HEARTBEAT_TIMEOUT);
                        env.multicast(self.id, &self.config, &AppendEntriesRequest {
                            term: self.term,
                            seqno: *heartbeat_seqno,
                        });
                    } else {
                        // TODO: this is probably too agressive, should only
                        // step down after election timeout
                        need_step_down = true;
                    }
                }
            },
        }
        if need_step_down {
            println!("Failed to maintain leadership: {}", *self);
            self.step_down(env, self.term)
        }
    }

    pub fn next_tick(&self) -> Time {
        let mut r = NEVER;
        match self.state {
            Follower { timer, .. }  => r = std::cmp::min(r, timer),
            Candidate { timer, .. } => r = std::cmp::min(r, timer),
            Leader { timer, .. }    => r = std::cmp::min(r, timer),
        }
        return r;
    }

    fn step_down(&mut self, env: &Environment, term: Term) {
        let t = match self.state {
            Follower { timer, .. } => timer,
            Candidate { timer, .. } => timer,
            Leader{..} => env.make_time(ELECTION_TIMEOUT),
        };
        self.term = term;
        self.vote = None;
        self.state = Follower {
            timer: t,
        };
    }

    fn try_become_leader(&mut self, env: &mut Environment) {
        let become_leader = match self.state {
            Follower {..} => false,
            Candidate {votes: ref votes, ..} => self.config.is_quorum(votes),
            Leader {..} => false,
        };
        if become_leader {
            self.state = Leader {
                timer: env.make_time(HEARTBEAT_TIMEOUT),
                heartbeat_seqno: 0,
                acks: newHashSet([self.id]),
                start_time: env.clock,
            };
            env.multicast(self.id, &self.config, &AppendEntriesRequest {
                term: self.term,
                seqno: 0,
            });
        }
    }

    pub fn stable_leader_start_time(&self, stable_req: uint) -> Option<Time> {
        match self.state {
            Leader {heartbeat_seqno, start_time, ..} => {
                if heartbeat_seqno > stable_req {
                    Some(start_time)
                } else {
                    None
                }
            },
            _ => None,
        }
    }

    pub fn handle(&mut self, env: &mut Environment, msg: &Message) {
        match msg.body {
            RequestVoteRequest {term, lastLogIndex} => {
                let reply = |granted| {
                    env.reply(msg, &RequestVoteResponse {
                        // use max here for nograntnobump algorithm
                        term: std::cmp::max(term, self.term),
                        granted: granted,
                        logOk: lastLogIndex >= self.lastLogIndex,
                    });
                };
                if term < self.term {
                    reply(TERM_STALE);
                } else {
                    if term > self.term {
                        match self.algorithm {
                            ~"nograntnobump"  => {},
                            ~"stalelognobump" => {},
                            _ => self.step_down(env, term)
                        }
                    }
                    // careful ordering to support algorithm hesitant:
                    // reply with LOG_STALE first
                    if lastLogIndex < self.lastLogIndex {
                        reply(LOG_STALE);
                    } else {
                        if self.algorithm == ~"stalelognobump" && term > self.term {
                            self.step_down(env, term);
                        }
                        match self.vote {
                            None => {
                                if self.algorithm == ~"nograntnobump" && term > self.term {
                                    self.step_down(env, term);
                                }
                                self.vote = Some(msg.from);
                                match self.state {
                                    Follower { timer: ref mut timer, .. } =>
                                        *timer = env.make_time(ELECTION_TIMEOUT),
                                    _ => {},
                                }
                                reply(GRANTED);
                            },
                            Some(c) => {
                                reply(if c == msg.from { GRANTED } else { VOTED });
                            },
                        }
                    }
                }
            },
            RequestVoteResponse {term, granted, logOk} => {
                let msg_term = term;
                let term = match self.state {
                    Candidate {pre, ..} => {
                        if (term > self.term &&
                            pre &&
                            self.algorithm == ~"zookeeper2") {
                            self.term
                        } else {
                            term
                        }
                    },
                    _ => term,
                };
                if term == self.term {
                    let mut forceNewElectionTerm = None;
                    match self.state {
                        Follower {..} => {},
                        Candidate {votes: ref mut votes,
                                   should_retry: ref mut should_retry,
                                   pre: ref mut pre,
                                   max_term: ref mut max_term,
                                   endorsements: ref mut endorsements,
                                   ..} => {
                            if *pre {
                                *max_term = std::cmp::max(*max_term, msg_term);
                                if logOk {
                                    endorsements.insert(msg.from);
                                    if self.config.is_quorum(endorsements) {
                                        if self.algorithm == ~"zookeeper2" {
                                            forceNewElectionTerm = Some(*max_term);
                                        } else {
                                            forceNewElectionTerm = Some(Term(0));
                                        }
                                    }
                                }
                            } else {
                                match granted {
                                    GRANTED => {
                                        votes.insert(msg.from);
                                    },
                                    TERM_STALE | VOTED => {
                                        endorsements.insert(msg.from);
                                        if self.algorithm == ~"hesitant" {
                                            *should_retry = true;
                                        }
                                        if self.algorithm == ~"hesitant2" {
                                            if self.config.is_quorum(endorsements) {
                                                *should_retry = true;
                                            }
                                        }
                                    },
                                    LOG_STALE => {},
                                }
                            }
                        },
                        Leader {..} => {},
                    }
                    match forceNewElectionTerm {
                        Some(t) => self.start_new_election(env, Some(t)),
                        None => {},
                    };
                    self.try_become_leader(env)
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
            AppendEntriesRequest {term, seqno} => {
                if term == self.term {
                    self.state = Follower {
                        timer: env.make_time(ELECTION_TIMEOUT),
                    };
                } else if term > self.term {
                    self.step_down(env, term);
                }
                env.reply(msg, &AppendEntriesResponse {
                    term: self.term,
                    seqno: seqno,
                });
            },
            AppendEntriesResponse {term, seqno} => {
                if term == self.term {
                    match self.state {
                        Follower {..} | Candidate {..} => {},
                        Leader { heartbeat_seqno, acks: ref mut acks, .. } => {
                            if seqno == heartbeat_seqno {
                                acks.insert(msg.from);
                            }
                        },
                    }
                } else if term > self.term {
                    self.step_down(env, term)
                }
            },
        }
    }

} // Server impl

<commit_msg>Change election timeout to 100-200ms<commit_after>extern mod std;
extern mod extra;

use basics::*;
use std::fmt;
use sim::Environment;
use std::hashmap::HashSet;

static ELECTION_TIMEOUT : (uint, uint) = (100000, 199999);
static HEARTBEAT_TIMEOUT : (uint, uint) = (50000, 50000);

#[deriving(Clone)]
pub struct Message {
    from: ServerID,
    to: ServerID,
    body: MessageBody,
}

impl fmt::Default for Message {
    fn fmt(msg: &Message, f: &mut fmt::Formatter) {
        write!(f.buf, "{} to {}: {}",
               *msg.from,
               *msg.to,
               msg.body);
    }
}

#[deriving(Eq, Clone)]
enum VoteGranted {
    GRANTED,
    TERM_STALE,
    LOG_STALE,
    VOTED,
}

impl fmt::Default for VoteGranted {
    fn fmt(g: &VoteGranted, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", match *g {
            GRANTED    => "GRANTED",
            TERM_STALE => "TERM_STALE",
            LOG_STALE  => "LOG_STALE",
            VOTED      => "VOTED",
        });
    }
}

#[deriving(Clone)]
pub enum MessageBody {
    RequestVoteRequest {
        term: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        term: Term,
        granted: VoteGranted,
        logOk: bool,
    },
    AppendEntriesRequest {
        term: Term,
        seqno: uint,
    },
    AppendEntriesResponse {
        term: Term,
        seqno: uint,
    },
}

impl fmt::Default for MessageBody {
    fn fmt(b: &MessageBody, f: &mut fmt::Formatter) {
        match *b {
            RequestVoteRequest{term, lastLogIndex} =>
                write!(f.buf, "RequestVoteRequest(term: {}, lastLogIndex: {})",
                       term, lastLogIndex),
            RequestVoteResponse{term, granted, logOk} =>
                write!(f.buf, "RequestVoteResponse(term: {}, granted: {}, logOk: {})",
                       term, granted, logOk),
            AppendEntriesRequest{term, seqno} =>
                write!(f.buf, "AppendEntriesRequest(term: {}, seqno: {})",
                       term, seqno),
            AppendEntriesResponse{term, seqno} =>
                write!(f.buf, "AppendEntriesResponse(term: {}, seqno: {})",
                       term, seqno),
        }
    }
}

enum ServerState {
    Follower  {
        timer: Time,
    },
    Candidate {
        timer: Time,
        votes: HashSet<ServerID>,
        /**
         * If true, when the timer fires, the candidate should start a new election.
         * If false, when the timer fires, set a new timer.
         * Always true for submission algo, but sometimes false in hesitant, hesitant2.
         */
        should_retry: bool,
        pre: bool,
        max_term: Term,
        endorsements: HashSet<ServerID>,
    },
    Leader {
        timer: Time,
        heartbeat_seqno: uint,
        acks: HashSet<ServerID>,
        start_time: Time,
    },
}

impl ServerState {
    pub fn to_char(&self) -> char {
        match *self {
            Follower {..}  => 'F',
            Candidate {..} => 'C',
            Leader {..}    => 'L',
        }
    }
}

impl fmt::Default for HashSet<ServerID> {
    fn fmt(set: &HashSet<ServerID>, f: &mut fmt::Formatter) {
        write!(f.buf, "\\{");
        let mut sorted = set.iter().to_owned_vec();
        extra::sort::tim_sort(sorted);
        for s in sorted.iter() {
            write!(f.buf, "{}", **s);
        }
        write!(f.buf, "\\}");
    }
}

impl fmt::Default for Configuration {
    fn fmt(config: &Configuration, f: &mut fmt::Formatter) {
        write!(f.buf, "[");
        for c in config.iter() {
            write!(f.buf, "{}", *c);
        }
        write!(f.buf, "]");
    }
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => {
                write!(f.buf, "Follower(timer: {})",
                       timer)
            },
            Candidate{timer, votes: ref votes, should_retry, pre, max_term, endorsements: ref endorsements} => {
                write!(f.buf, "Candidate(timer: {}, votes: {}, should_retry: {}, pre: {}, max_term: {}, endorsements: {})",
                       timer, *votes, should_retry, pre, max_term, *endorsements)
            },
            Leader{timer, heartbeat_seqno, acks: ref acks, start_time} => {
                write!(f.buf, "Leader(timer: {}, seqno: {}, acks: {}, start_time: {})",
                       timer, heartbeat_seqno, *acks, start_time)
            },
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Configuration(~[HashSet<ServerID>]);

impl Configuration {
    fn is_quorum(&self, servers: &HashSet<ServerID>) -> bool {
        for c in self.iter() {
            if c.intersection(servers).len() <= c.len() / 2 {
                return false;
            }
        }
        return true;
    }
}

struct Server {
    id: ServerID,
    config: Configuration,
    algorithm: ~str,
    term: Term,
    state: ServerState,
    vote: Option<ServerID>,
    lastLogIndex: Index,
}

impl fmt::Default for Server {
    fn fmt(server: &Server, f: &mut fmt::Formatter) {
        write!(f.buf, "Server(id: {}, term: {}, log: {}, vote: {}, state: {})",
               server.id,
               server.term,
               server.lastLogIndex,
               server.vote,
               server.state)
    }
}

impl Server {
    pub fn new(id : ServerID, config: Configuration, env: &Environment,
               algorithm: &str) -> Server {
        match algorithm {
            "nograntnobump" |
            "stalelognobump" |
            "hesitant"      |
            "hesitant2"     |
            "zookeeper"     |
            "zookeeper2"    |
            "submission"    => {},
            _ => fail!("Unknown algorithm: {}", algorithm)
        };
        Server {
            id: id,
            config: config,
            algorithm: algorithm.into_owned(),
            term: Term(0),
            state: Follower { timer: env.make_time(ELECTION_TIMEOUT) },
            lastLogIndex: Index(0),
            vote: None,
        }
    }
    fn start_new_election(&mut self, env: &mut Environment, forceTerm: Option<Term>) {
        let pre = forceTerm.is_none() && match self.algorithm {
            ~"zookeeper"  => true,
            ~"zookeeper2" => true,
            _ => false,
        };
        if !pre {
            self.term = match forceTerm {
                Some(t) if t > self.term => Term(*t + 1),
                _                        => Term(*self.term + 1),
            };
            self.vote = None;
        }
        self.state = Candidate {
            timer: env.make_time(ELECTION_TIMEOUT),
            votes: newHashSet([self.id]),
            should_retry: match self.algorithm {
                ~"hesitant"  => false,
                ~"hesitant2" => false,
                _ => true,
            },
            pre: pre,
            max_term: Term(0),
            endorsements: newHashSet([self.id]),
        };
        self.vote = Some(self.id);
        env.multicast(self.id, &self.config, &RequestVoteRequest {
            term: self.term,
            lastLogIndex: self.lastLogIndex,
        });
        self.try_become_leader(env);
    }
    pub fn tick(&mut self, env: &mut Environment) {
        let mut need_step_down = false;

        match self.state {
            Follower  {timer, ..} => {
                if timer <= env.clock {
                    self.start_new_election(env, None)
                }
            }
            Candidate {timer, should_retry, ..}
                if timer <= env.clock && should_retry => {
                    self.start_new_election(env, None);
            },
            Candidate {timer: ref mut timer, should_retry, ..}
                if *timer <= env.clock && !should_retry => {
                    *timer = env.make_time(ELECTION_TIMEOUT);
            },
            Candidate {..} => {},
            Leader {timer: ref mut timer, heartbeat_seqno: ref mut heartbeat_seqno, acks: ref mut acks, ..} => {
                if *timer <= env.clock {
                    if (self.config.is_quorum(acks)) {
                        *heartbeat_seqno += 1;
                        acks.clear();
                        acks.insert(self.id);
                        *timer = env.make_time(HEARTBEAT_TIMEOUT);
                        env.multicast(self.id, &self.config, &AppendEntriesRequest {
                            term: self.term,
                            seqno: *heartbeat_seqno,
                        });
                    } else {
                        // TODO: this is probably too agressive, should only
                        // step down after election timeout
                        need_step_down = true;
                    }
                }
            },
        }
        if need_step_down {
            println!("Failed to maintain leadership: {}", *self);
            self.step_down(env, self.term)
        }
    }

    pub fn next_tick(&self) -> Time {
        let mut r = NEVER;
        match self.state {
            Follower { timer, .. }  => r = std::cmp::min(r, timer),
            Candidate { timer, .. } => r = std::cmp::min(r, timer),
            Leader { timer, .. }    => r = std::cmp::min(r, timer),
        }
        return r;
    }

    fn step_down(&mut self, env: &Environment, term: Term) {
        let t = match self.state {
            Follower { timer, .. } => timer,
            Candidate { timer, .. } => timer,
            Leader{..} => env.make_time(ELECTION_TIMEOUT),
        };
        self.term = term;
        self.vote = None;
        self.state = Follower {
            timer: t,
        };
    }

    fn try_become_leader(&mut self, env: &mut Environment) {
        let become_leader = match self.state {
            Follower {..} => false,
            Candidate {votes: ref votes, ..} => self.config.is_quorum(votes),
            Leader {..} => false,
        };
        if become_leader {
            self.state = Leader {
                timer: env.make_time(HEARTBEAT_TIMEOUT),
                heartbeat_seqno: 0,
                acks: newHashSet([self.id]),
                start_time: env.clock,
            };
            env.multicast(self.id, &self.config, &AppendEntriesRequest {
                term: self.term,
                seqno: 0,
            });
        }
    }

    pub fn stable_leader_start_time(&self, stable_req: uint) -> Option<Time> {
        match self.state {
            Leader {heartbeat_seqno, start_time, ..} => {
                if heartbeat_seqno > stable_req {
                    Some(start_time)
                } else {
                    None
                }
            },
            _ => None,
        }
    }

    pub fn handle(&mut self, env: &mut Environment, msg: &Message) {
        match msg.body {
            RequestVoteRequest {term, lastLogIndex} => {
                let reply = |granted| {
                    env.reply(msg, &RequestVoteResponse {
                        // use max here for nograntnobump algorithm
                        term: std::cmp::max(term, self.term),
                        granted: granted,
                        logOk: lastLogIndex >= self.lastLogIndex,
                    });
                };
                if term < self.term {
                    reply(TERM_STALE);
                } else {
                    if term > self.term {
                        match self.algorithm {
                            ~"nograntnobump"  => {},
                            ~"stalelognobump" => {},
                            _ => self.step_down(env, term)
                        }
                    }
                    // careful ordering to support algorithm hesitant:
                    // reply with LOG_STALE first
                    if lastLogIndex < self.lastLogIndex {
                        reply(LOG_STALE);
                    } else {
                        if self.algorithm == ~"stalelognobump" && term > self.term {
                            self.step_down(env, term);
                        }
                        match self.vote {
                            None => {
                                if self.algorithm == ~"nograntnobump" && term > self.term {
                                    self.step_down(env, term);
                                }
                                self.vote = Some(msg.from);
                                match self.state {
                                    Follower { timer: ref mut timer, .. } =>
                                        *timer = env.make_time(ELECTION_TIMEOUT),
                                    _ => {},
                                }
                                reply(GRANTED);
                            },
                            Some(c) => {
                                reply(if c == msg.from { GRANTED } else { VOTED });
                            },
                        }
                    }
                }
            },
            RequestVoteResponse {term, granted, logOk} => {
                let msg_term = term;
                let term = match self.state {
                    Candidate {pre, ..} => {
                        if (term > self.term &&
                            pre &&
                            self.algorithm == ~"zookeeper2") {
                            self.term
                        } else {
                            term
                        }
                    },
                    _ => term,
                };
                if term == self.term {
                    let mut forceNewElectionTerm = None;
                    match self.state {
                        Follower {..} => {},
                        Candidate {votes: ref mut votes,
                                   should_retry: ref mut should_retry,
                                   pre: ref mut pre,
                                   max_term: ref mut max_term,
                                   endorsements: ref mut endorsements,
                                   ..} => {
                            if *pre {
                                *max_term = std::cmp::max(*max_term, msg_term);
                                if logOk {
                                    endorsements.insert(msg.from);
                                    if self.config.is_quorum(endorsements) {
                                        if self.algorithm == ~"zookeeper2" {
                                            forceNewElectionTerm = Some(*max_term);
                                        } else {
                                            forceNewElectionTerm = Some(Term(0));
                                        }
                                    }
                                }
                            } else {
                                match granted {
                                    GRANTED => {
                                        votes.insert(msg.from);
                                    },
                                    TERM_STALE | VOTED => {
                                        endorsements.insert(msg.from);
                                        if self.algorithm == ~"hesitant" {
                                            *should_retry = true;
                                        }
                                        if self.algorithm == ~"hesitant2" {
                                            if self.config.is_quorum(endorsements) {
                                                *should_retry = true;
                                            }
                                        }
                                    },
                                    LOG_STALE => {},
                                }
                            }
                        },
                        Leader {..} => {},
                    }
                    match forceNewElectionTerm {
                        Some(t) => self.start_new_election(env, Some(t)),
                        None => {},
                    };
                    self.try_become_leader(env)
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
            AppendEntriesRequest {term, seqno} => {
                if term == self.term {
                    self.state = Follower {
                        timer: env.make_time(ELECTION_TIMEOUT),
                    };
                } else if term > self.term {
                    self.step_down(env, term);
                }
                env.reply(msg, &AppendEntriesResponse {
                    term: self.term,
                    seqno: seqno,
                });
            },
            AppendEntriesResponse {term, seqno} => {
                if term == self.term {
                    match self.state {
                        Follower {..} | Candidate {..} => {},
                        Leader { heartbeat_seqno, acks: ref mut acks, .. } => {
                            if seqno == heartbeat_seqno {
                                acks.insert(msg.from);
                            }
                        },
                    }
                } else if term > self.term {
                    self.step_down(env, term)
                }
            },
        }
    }

} // Server impl

<|endoftext|>


<commit_before>//! A very simple wrapper for LLVM that can JIT functions written as IR strings.

extern crate llvm_sys as llvm;

use std::error::Error;
use std::ffi::{CStr, CString, NulError};
use std::fmt;
use std::result::Result;
use std::ops::Drop;
use std::os::raw::c_char;
use std::process::Command;
use std::sync::{Once, ONCE_INIT};

use llvm::prelude::{LLVMContextRef, LLVMModuleRef, LLVMMemoryBufferRef};
use llvm::execution_engine::{LLVMExecutionEngineRef, LLVMMCJITCompilerOptions};
use llvm::analysis::LLVMVerifierFailureAction;
use llvm::transforms::pass_manager_builder as pmb;

#[cfg(test)]
mod tests;

// Helper objects to make sure we only initialize once
static ONCE: Once = ONCE_INIT;
static mut initialize_failed: bool = false;

/// Error type returned by easy_ll.
#[derive(Debug)]
pub struct LlvmError(String);

impl LlvmError {
    pub fn new(description: &str) -> LlvmError {
        LlvmError(description.to_string())
    }

    pub fn to_string(self) -> String {
        self.0
    }
}

impl fmt::Display for LlvmError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl Error for LlvmError {
    fn description(&self) -> &str {
        &self.0
    }

    fn cause(&self) -> Option<&Error> {
        None
    }
}

impl From<NulError> for LlvmError {
    fn from(_: NulError) -> LlvmError {
        LlvmError::new("Null byte in string")
    }
}

/// The type of our "run" function pointer.
type RunFunc = extern "C" fn(i64) -> i64;

/// A compiled module returned by `compile_module`, wrapping a `run` function that takes `i64`
/// and returns `i64`. This structure includes (and manages) an LLVM execution engine, which is
/// freed when this structure is dropped.
#[derive(Debug)]
pub struct CompiledModule {
    context: LLVMContextRef,
    engine: Option<LLVMExecutionEngineRef>,
    function: Option<RunFunc>,
}

impl CompiledModule {
    /// Call the module's `run` function.
    pub fn run(&self, arg: i64) -> i64 {
        (self.function.unwrap())(arg)
    }
}

impl Drop for CompiledModule {
    /// Disposes of the LLVM execution engine and compiled function.
    fn drop(&mut self) {
        unsafe {
            self.engine.map(|e| llvm::execution_engine::LLVMDisposeExecutionEngine(e));
            llvm::core::LLVMContextDispose(self.context);
        }
    }
}

/// Compile a string of LLVM IR (in human readable format) into a `CompiledModule` that can then
/// be executed. The LLVM IR should contain an entry point function called `run` that takes `i64`
/// and returns `i64`, which will be called by `CompiledModule::run`.
pub fn compile_module(code: &str, num_threads: i32) -> Result<CompiledModule, LlvmError> {
    unsafe {
        // Initialize LLVM
        ONCE.call_once(|| initialize());
        if initialize_failed {
            return Err(LlvmError::new("LLVM initialization failed"));
        }

        // Create an LLVM context
        let context = llvm::core::LLVMContextCreate();
        if context.is_null() {
            return Err(LlvmError::new("LLVMContextCreate returned null"));
        }

        // Create a CompiledModule to wrap the context and our result (will clean it on Drop).
        let mut result = CompiledModule {
            context: context,
            engine: None,
            function: None,
        };

        // Parse the IR to get an LLVMModuleRef
        let module = try!(parse_module_str(context, code));

        // Validate and optimize the module
        try!(verify_module(module));
        try!(check_run_function(module));
        try!(optimize_module(module));

        // Create an execution engine for the module and find its run function
        let engine = try!(create_exec_engine(module));

        // TODO maybe make the set of libraries a parameter
        #[allow(unused_must_use)]
        {
            Command::new("make").arg("-C").arg("cpp").arg(format!("W={}", num_threads)).output();
        }
        let pl_module = try!(parse_module_file(context, "cpp/parlib.bc"));
        llvm::execution_engine::LLVMAddModule(engine, pl_module);

        result.engine = Some(engine);
        result.function = Some(try!(find_run_function(engine)));

        Ok(result)
    }
}

/// Initialize LLVM or save an error message in `initialize_failed` if this does not work.
/// We call this function only once in cases some steps are expensive.
fn initialize() {
    unsafe {
        if llvm::target::LLVM_InitializeNativeTarget() != 0 {
            initialize_failed = true;
            return;
        }
        if llvm::target::LLVM_InitializeNativeAsmPrinter() != 0 {
            initialize_failed = true;
            return;
        }
        if llvm::target::LLVM_InitializeNativeAsmParser() != 0 {
            initialize_failed = true;
            return;
        }
        llvm::execution_engine::LLVMLinkInMCJIT();
    }
}

unsafe fn parse_module_helper(context: LLVMContextRef,
                              buffer: LLVMMemoryBufferRef)
                              -> Result<LLVMModuleRef, LlvmError> {
    // Parse IR into a module
    let mut module = 0 as LLVMModuleRef;
    let mut error_str = 0 as *mut c_char;
    let result_code =
        llvm::ir_reader::LLVMParseIRInContext(context, buffer, &mut module, &mut error_str);
    if result_code != 0 {
        let msg = format!("Compile error: {}",
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }

    Ok(module)
}

/// Parse a string of IR code into an `LLVMModuleRef` for the given context.
unsafe fn parse_module_str(context: LLVMContextRef,
                           code: &str)
                           -> Result<LLVMModuleRef, LlvmError> {
    // Create an LLVM memory buffer around the code
    let code_len = code.len();
    let name = try!(CString::new("module"));
    let code = try!(CString::new(code));
    let buffer = llvm::core::LLVMCreateMemoryBufferWithMemoryRange(code.as_ptr(),
                                                                   code_len,
                                                                   name.as_ptr(),
                                                                   0);
    if buffer.is_null() {
        return Err(LlvmError::new("LLVMCreateMemoryBufferWithMemoryRange failed"));
    }

    parse_module_helper(context, buffer)
}

/// Parse a file of IR code into an `LLVMModuleRef` for the given context.
unsafe fn parse_module_file(context: LLVMContextRef,
                            file: &str)
                            -> Result<LLVMModuleRef, LlvmError> {
    let mut buffer = 0 as LLVMMemoryBufferRef;
    let mut error_str = 0 as *mut c_char;
    let file_name = CString::new(file).unwrap();
    let result_code = llvm::core::LLVMCreateMemoryBufferWithContentsOfFile(file_name.as_ptr(),
                                                                           &mut buffer,
                                                                           &mut error_str);
    if result_code != 0 {
        let msg = format!("Error reading module file {}: {}",
                          file,
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }
    if buffer.is_null() {
        return Err(LlvmError::new("LLVMCreateMemoryBufferWithContentsOfFile failed"));
    }
    parse_module_helper(context, buffer)
}

/// Verify a module using LLVM's verifier (for basic block structure, etc).
unsafe fn verify_module(module: LLVMModuleRef) -> Result<(), LlvmError> {
    let mut error_str = 0 as *mut c_char;
    let result_code =
        llvm::analysis::LLVMVerifyModule(module,
                                         LLVMVerifierFailureAction::LLVMReturnStatusAction,
                                         &mut error_str);
    if result_code != 0 {
        let msg = format!("Module verification failed: {}",
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }
    Ok(())
}

/// Check that a module has a "run" function of type i64 -> i64.
unsafe fn check_run_function(module: LLVMModuleRef) -> Result<(), LlvmError> {
    let run = CString::new("run").unwrap();
    let func = llvm::core::LLVMGetNamedFunction(module, run.as_ptr());
    if func.is_null() {
        println!("EEEK");
        return Err(LlvmError::new("No run function in module"));
    }
    let c_str = llvm::core::LLVMPrintTypeToString(llvm::core::LLVMTypeOf(func));
    let func_type = CStr::from_ptr(c_str).to_str().unwrap();
    if func_type != "i64 (i64)*" {
        return Err(LlvmError(format!("Run function has wrong type: {}", func_type)));
    }
    Ok(())
}

/// Optimize an LLVM module using our chosen passes (currently uses standard passes for -O2).
unsafe fn optimize_module(module: LLVMModuleRef) -> Result<(), LlvmError> {
    let manager = llvm::core::LLVMCreatePassManager();
    if manager.is_null() {
        return Err(LlvmError::new("LLVMCreatePassManager returned null"));
    }
    let builder = pmb::LLVMPassManagerBuilderCreate();
    if builder.is_null() {
        return Err(LlvmError::new("LLVMPassManagerBuilderCreate returned null"));
    }
    // TODO: not clear we need both Module and LTO calls here; just LTO might work
    pmb::LLVMPassManagerBuilderSetOptLevel(builder, 2);
    pmb::LLVMPassManagerBuilderPopulateModulePassManager(builder, manager);
    pmb::LLVMPassManagerBuilderPopulateLTOPassManager(builder, manager, 1, 1);
    pmb::LLVMPassManagerBuilderDispose(builder);
    llvm::core::LLVMRunPassManager(manager, module);
    llvm::core::LLVMDisposePassManager(manager);
    Ok(())
}

/// Create an MCJIT execution engine for a given module.
unsafe fn create_exec_engine(module: LLVMModuleRef) -> Result<LLVMExecutionEngineRef, LlvmError> {
    let mut engine = 0 as LLVMExecutionEngineRef;
    let mut error_str = 0 as *mut c_char;
    let mut options: LLVMMCJITCompilerOptions = std::mem::uninitialized();
    let options_size = std::mem::size_of::<LLVMMCJITCompilerOptions>();
    llvm::execution_engine::LLVMInitializeMCJITCompilerOptions(&mut options, options_size);
    options.OptLevel = 2;
    let result_code = llvm::execution_engine::LLVMCreateMCJITCompilerForModule(&mut engine,
                                                                               module,
                                                                               &mut options,
                                                                               options_size,
                                                                               &mut error_str);
    if result_code != 0 {
        let msg = format!("Creating execution engine failed: {}",
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }
    Ok(engine)
}

/// Get a pointer to the "run" function in an execution engine.
unsafe fn find_run_function(engine: LLVMExecutionEngineRef) -> Result<RunFunc, LlvmError> {
    let run = CString::new("run").unwrap();
    let func_addr = llvm::execution_engine::LLVMGetFunctionAddress(engine, run.as_ptr());
    if func_addr == 0 {
        return Err(LlvmError::new("No run function in module"));
    }
    let function: RunFunc = std::mem::transmute(func_addr);
    Ok(function)
}
<commit_msg>added make clean to llvm build process<commit_after>//! A very simple wrapper for LLVM that can JIT functions written as IR strings.

extern crate llvm_sys as llvm;

use std::error::Error;
use std::ffi::{CStr, CString, NulError};
use std::fmt;
use std::result::Result;
use std::ops::Drop;
use std::os::raw::c_char;
use std::process::Command;
use std::sync::{Once, ONCE_INIT};

use llvm::prelude::{LLVMContextRef, LLVMModuleRef, LLVMMemoryBufferRef};
use llvm::execution_engine::{LLVMExecutionEngineRef, LLVMMCJITCompilerOptions};
use llvm::analysis::LLVMVerifierFailureAction;
use llvm::transforms::pass_manager_builder as pmb;

#[cfg(test)]
mod tests;

// Helper objects to make sure we only initialize once
static ONCE: Once = ONCE_INIT;
static mut initialize_failed: bool = false;

/// Error type returned by easy_ll.
#[derive(Debug)]
pub struct LlvmError(String);

impl LlvmError {
    pub fn new(description: &str) -> LlvmError {
        LlvmError(description.to_string())
    }

    pub fn to_string(self) -> String {
        self.0
    }
}

impl fmt::Display for LlvmError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl Error for LlvmError {
    fn description(&self) -> &str {
        &self.0
    }

    fn cause(&self) -> Option<&Error> {
        None
    }
}

impl From<NulError> for LlvmError {
    fn from(_: NulError) -> LlvmError {
        LlvmError::new("Null byte in string")
    }
}

/// The type of our "run" function pointer.
type RunFunc = extern "C" fn(i64) -> i64;

/// A compiled module returned by `compile_module`, wrapping a `run` function that takes `i64`
/// and returns `i64`. This structure includes (and manages) an LLVM execution engine, which is
/// freed when this structure is dropped.
#[derive(Debug)]
pub struct CompiledModule {
    context: LLVMContextRef,
    engine: Option<LLVMExecutionEngineRef>,
    function: Option<RunFunc>,
}

impl CompiledModule {
    /// Call the module's `run` function.
    pub fn run(&self, arg: i64) -> i64 {
        (self.function.unwrap())(arg)
    }
}

impl Drop for CompiledModule {
    /// Disposes of the LLVM execution engine and compiled function.
    fn drop(&mut self) {
        unsafe {
            self.engine.map(|e| llvm::execution_engine::LLVMDisposeExecutionEngine(e));
            llvm::core::LLVMContextDispose(self.context);
        }
    }
}

/// Compile a string of LLVM IR (in human readable format) into a `CompiledModule` that can then
/// be executed. The LLVM IR should contain an entry point function called `run` that takes `i64`
/// and returns `i64`, which will be called by `CompiledModule::run`.
pub fn compile_module(code: &str, num_threads: i32) -> Result<CompiledModule, LlvmError> {
    unsafe {
        // Initialize LLVM
        ONCE.call_once(|| initialize());
        if initialize_failed {
            return Err(LlvmError::new("LLVM initialization failed"));
        }

        // Create an LLVM context
        let context = llvm::core::LLVMContextCreate();
        if context.is_null() {
            return Err(LlvmError::new("LLVMContextCreate returned null"));
        }

        // Create a CompiledModule to wrap the context and our result (will clean it on Drop).
        let mut result = CompiledModule {
            context: context,
            engine: None,
            function: None,
        };

        // Parse the IR to get an LLVMModuleRef
        let module = try!(parse_module_str(context, code));

        // Validate and optimize the module
        try!(verify_module(module));
        try!(check_run_function(module));
        try!(optimize_module(module));

        // Create an execution engine for the module and find its run function
        let engine = try!(create_exec_engine(module));

        // TODO maybe make the set of libraries a parameter
        #[allow(unused_must_use)]
        {
            Command::new("make").arg("-C").arg("cpp").arg("clean").output();
            Command::new("make").arg("-C").arg("cpp").arg(format!("W={}", num_threads)).output();
        }
        let pl_module = try!(parse_module_file(context, "cpp/parlib.bc"));
        llvm::execution_engine::LLVMAddModule(engine, pl_module);

        result.engine = Some(engine);
        result.function = Some(try!(find_run_function(engine)));

        Ok(result)
    }
}

/// Initialize LLVM or save an error message in `initialize_failed` if this does not work.
/// We call this function only once in cases some steps are expensive.
fn initialize() {
    unsafe {
        if llvm::target::LLVM_InitializeNativeTarget() != 0 {
            initialize_failed = true;
            return;
        }
        if llvm::target::LLVM_InitializeNativeAsmPrinter() != 0 {
            initialize_failed = true;
            return;
        }
        if llvm::target::LLVM_InitializeNativeAsmParser() != 0 {
            initialize_failed = true;
            return;
        }
        llvm::execution_engine::LLVMLinkInMCJIT();
    }
}

unsafe fn parse_module_helper(context: LLVMContextRef,
                              buffer: LLVMMemoryBufferRef)
                              -> Result<LLVMModuleRef, LlvmError> {
    // Parse IR into a module
    let mut module = 0 as LLVMModuleRef;
    let mut error_str = 0 as *mut c_char;
    let result_code =
        llvm::ir_reader::LLVMParseIRInContext(context, buffer, &mut module, &mut error_str);
    if result_code != 0 {
        let msg = format!("Compile error: {}",
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }

    Ok(module)
}

/// Parse a string of IR code into an `LLVMModuleRef` for the given context.
unsafe fn parse_module_str(context: LLVMContextRef,
                           code: &str)
                           -> Result<LLVMModuleRef, LlvmError> {
    // Create an LLVM memory buffer around the code
    let code_len = code.len();
    let name = try!(CString::new("module"));
    let code = try!(CString::new(code));
    let buffer = llvm::core::LLVMCreateMemoryBufferWithMemoryRange(code.as_ptr(),
                                                                   code_len,
                                                                   name.as_ptr(),
                                                                   0);
    if buffer.is_null() {
        return Err(LlvmError::new("LLVMCreateMemoryBufferWithMemoryRange failed"));
    }

    parse_module_helper(context, buffer)
}

/// Parse a file of IR code into an `LLVMModuleRef` for the given context.
unsafe fn parse_module_file(context: LLVMContextRef,
                            file: &str)
                            -> Result<LLVMModuleRef, LlvmError> {
    let mut buffer = 0 as LLVMMemoryBufferRef;
    let mut error_str = 0 as *mut c_char;
    let file_name = CString::new(file).unwrap();
    let result_code = llvm::core::LLVMCreateMemoryBufferWithContentsOfFile(file_name.as_ptr(),
                                                                           &mut buffer,
                                                                           &mut error_str);
    if result_code != 0 {
        let msg = format!("Error reading module file {}: {}",
                          file,
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }
    if buffer.is_null() {
        return Err(LlvmError::new("LLVMCreateMemoryBufferWithContentsOfFile failed"));
    }
    parse_module_helper(context, buffer)
}

/// Verify a module using LLVM's verifier (for basic block structure, etc).
unsafe fn verify_module(module: LLVMModuleRef) -> Result<(), LlvmError> {
    let mut error_str = 0 as *mut c_char;
    let result_code =
        llvm::analysis::LLVMVerifyModule(module,
                                         LLVMVerifierFailureAction::LLVMReturnStatusAction,
                                         &mut error_str);
    if result_code != 0 {
        let msg = format!("Module verification failed: {}",
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }
    Ok(())
}

/// Check that a module has a "run" function of type i64 -> i64.
unsafe fn check_run_function(module: LLVMModuleRef) -> Result<(), LlvmError> {
    let run = CString::new("run").unwrap();
    let func = llvm::core::LLVMGetNamedFunction(module, run.as_ptr());
    if func.is_null() {
        println!("EEEK");
        return Err(LlvmError::new("No run function in module"));
    }
    let c_str = llvm::core::LLVMPrintTypeToString(llvm::core::LLVMTypeOf(func));
    let func_type = CStr::from_ptr(c_str).to_str().unwrap();
    if func_type != "i64 (i64)*" {
        return Err(LlvmError(format!("Run function has wrong type: {}", func_type)));
    }
    Ok(())
}

/// Optimize an LLVM module using our chosen passes (currently uses standard passes for -O2).
unsafe fn optimize_module(module: LLVMModuleRef) -> Result<(), LlvmError> {
    let manager = llvm::core::LLVMCreatePassManager();
    if manager.is_null() {
        return Err(LlvmError::new("LLVMCreatePassManager returned null"));
    }
    let builder = pmb::LLVMPassManagerBuilderCreate();
    if builder.is_null() {
        return Err(LlvmError::new("LLVMPassManagerBuilderCreate returned null"));
    }
    // TODO: not clear we need both Module and LTO calls here; just LTO might work
    pmb::LLVMPassManagerBuilderSetOptLevel(builder, 2);
    pmb::LLVMPassManagerBuilderPopulateModulePassManager(builder, manager);
    pmb::LLVMPassManagerBuilderPopulateLTOPassManager(builder, manager, 1, 1);
    pmb::LLVMPassManagerBuilderDispose(builder);
    llvm::core::LLVMRunPassManager(manager, module);
    llvm::core::LLVMDisposePassManager(manager);
    Ok(())
}

/// Create an MCJIT execution engine for a given module.
unsafe fn create_exec_engine(module: LLVMModuleRef) -> Result<LLVMExecutionEngineRef, LlvmError> {
    let mut engine = 0 as LLVMExecutionEngineRef;
    let mut error_str = 0 as *mut c_char;
    let mut options: LLVMMCJITCompilerOptions = std::mem::uninitialized();
    let options_size = std::mem::size_of::<LLVMMCJITCompilerOptions>();
    llvm::execution_engine::LLVMInitializeMCJITCompilerOptions(&mut options, options_size);
    options.OptLevel = 2;
    let result_code = llvm::execution_engine::LLVMCreateMCJITCompilerForModule(&mut engine,
                                                                               module,
                                                                               &mut options,
                                                                               options_size,
                                                                               &mut error_str);
    if result_code != 0 {
        let msg = format!("Creating execution engine failed: {}",
                          CStr::from_ptr(error_str).to_str().unwrap());
        return Err(LlvmError(msg));
    }
    Ok(engine)
}

/// Get a pointer to the "run" function in an execution engine.
unsafe fn find_run_function(engine: LLVMExecutionEngineRef) -> Result<RunFunc, LlvmError> {
    let run = CString::new("run").unwrap();
    let func_addr = llvm::execution_engine::LLVMGetFunctionAddress(engine, run.as_ptr());
    if func_addr == 0 {
        return Err(LlvmError::new("No run function in module"));
    }
    let function: RunFunc = std::mem::transmute(func_addr);
    Ok(function)
}
<|endoftext|>


<commit_before>//
// Copyright 2022 The Project Oak Authors
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! Main 'kernel' for baremetal Oak Functions.
//!
//! This code takes care of initializing the x86-64 machine properly and
//! handing the reins over to the oak_baremetal_runtime. It is in a separate crate so that we
//! could support different boot protocols (eg Linux boot protocol or PVH) that
//! require different targets, linker scripts and/or provide machine
//! information in different data structures.
//!
//! Bootloaders (and VMMs under them) have to adhere to the following protocol:
//!   * Enter 64-bit long mode, and set up basic paging -- enough to load the code, as we will set
//!     up a full page table in `start_kernel`.
//!   * Implement a `#[panic_handler]` that delegates to `panic()` in this crate.
//!   * Call `start_kernel` from the entry point of the bootloader.

#![cfg_attr(not(test), no_std)]
#![feature(abi_x86_interrupt)]
#![feature(allocator_api)]
#![feature(asm_sym)]
#![feature(naked_functions)]
#![feature(once_cell)]

mod args;
mod avx;
mod boot;
mod descriptors;
mod elf;
mod ghcb;
#[cfg(feature = "c_interface")]
mod interface;
mod interrupts;
mod libm;
mod logging;
mod memory;
mod mm;
#[cfg(feature = "serial_channel")]
mod serial;
pub mod shutdown;
#[cfg(feature = "simple_io_channel")]
mod simpleio;
mod snp;
#[cfg(any(feature = "virtio_console_channel", feature = "vsock_channel"))]
mod virtio;

extern crate alloc;

use crate::{
    mm::Translator,
    snp::{get_snp_page_addresses, init_snp_pages},
};
use alloc::{alloc::Allocator, boxed::Box};
use core::{cell::OnceCell, marker::Sync, panic::PanicInfo, str::FromStr};
use linked_list_allocator::LockedHeap;
use log::{error, info};
use oak_channel::Channel;
use oak_linux_boot_params::BootParams;
use sev_guest::msr::{change_snp_state_for_frame, get_sev_status, PageAssignment, SevStatus};
use strum::{EnumIter, EnumString, IntoEnumIterator};
use x86_64::{
    structures::paging::{Page, Size2MiB},
    VirtAddr,
};

static mut GUEST_HOST_HEAP: OnceCell<LockedHeap> = OnceCell::new();

/// Main entry point for the kernel, to be called from bootloader.
pub fn start_kernel(info: &BootParams) -> Box<dyn Channel> {
    avx::enable_avx();
    descriptors::init_gdt();
    interrupts::init_idt();
    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    logging::init_logging();

    // We need to be done with the boot info struct before intializing memory. For example, the
    // multiboot protocol explicitly states data can be placed anywhere in memory; therefore, it's
    // highly likely we will overwrite some data after we initialize the heap. args::init_args()
    // caches the arguments (as long they are of reasonable length) in a static variable, allowing
    // us to refer to the args in the future.
    let kernel_args = args::init_args(info.args()).unwrap();
    info!("Kernel boot args: {}", kernel_args.args());

    let protocol = info.protocol();
    info!("Boot protocol:  {}", protocol);
    let snp_pages = if sev_status.contains(SevStatus::SNP_ACTIVE) {
        // We have to get the physical addresses of the CPUID pages now while the identity mapping
        // is still in place, but we can only initialize the instances after the new page
        // mappings have been set up.
        Some(get_snp_page_addresses(info))
    } else {
        None
    };

    // Safety: in the linker script we specify that the ELF header should be placed at 0x200000.
    let program_headers = unsafe { elf::get_phdrs(VirtAddr::new(0x20_0000)) };

    // Physical frame allocator: support up to 128 GiB of memory, for now.
    let mut frame_allocator = mm::init::<1024>(info.e820_table(), program_headers);

    let mut mapper = mm::init_paging(&mut frame_allocator, program_headers).unwrap();

    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        // Now that the page tables have been updated, we have to re-share the GHCB with the
        // hypervisor.
        ghcb::reshare_ghcb(&mut mapper);
        if sev_status.contains(SevStatus::SNP_ACTIVE) {
            // We must also initialise the CPUID and secrets pages when SEV-SNP is active. Panicking
            // is OK at this point, because these pages are required to support the full features
            // and we don't want to run without them.
            init_snp_pages(
                snp_pages.expect("Missing SNP CPUID and secrets pages."),
                &mapper,
            );
        }
    }

    // Allocate a section for guest-host communication (without the `ENCRYPTED` bit set)
    // We'll allocate 2*2MiB, as virtio needs more than 2 MiB for its data structures.
    let guest_host_frames = frame_allocator.allocate_contiguous(2).unwrap();
    let guest_host_pages = Page::range(
        mapper
            .translate_physical_frame(guest_host_frames.start)
            .unwrap(),
        mapper
            .translate_physical_frame(guest_host_frames.end)
            .unwrap(),
    );

    // If we are running on SNP we have to mark the frames as shared in the RMP. It is OK to crash
    // if we cannot mark the pages as shared in the RMP.
    if sev_status.contains(SevStatus::SNP_ACTIVE) {
        // TODO(#3414): Use the GHCB protocol when it is available.
        for frame in guest_host_frames {
            change_snp_state_for_frame(&frame, PageAssignment::Shared)
                .expect("Could not change SNP state for frame.");
        }
    }

    // Safety: initializing the new heap is safe as the frame allocator guarantees we're not
    // overwriting any other memory; writing to the static mut is safe as we're in the
    // initialization code and thus there can be no concurrent access.
    let guest_host_heap = unsafe {
        GUEST_HOST_HEAP
            .get_or_init(|| memory::init_guest_host_heap(guest_host_pages, &mut mapper).unwrap())
    };

    // If we don't find memory for heap, it's ok to panic.
    let heap_phys_frames = frame_allocator.largest_available().unwrap();
    memory::init_kernel_heap::<Size2MiB>(Page::range(
        mapper
            .translate_physical_frame(heap_phys_frames.start)
            .unwrap(),
        mapper
            .translate_physical_frame(heap_phys_frames.end + 1)
            .unwrap(),
    ))
    .unwrap();

    get_channel(&kernel_args, &mapper, guest_host_heap, sev_status)
}

#[derive(EnumIter, EnumString)]
#[strum(ascii_case_insensitive, serialize_all = "snake_case")]
enum ChannelType {
    #[cfg(feature = "virtio_console_channel")]
    VirtioConsole,
    #[cfg(feature = "vsock_channel")]
    VirtioVsock,
    #[cfg(feature = "serial_channel")]
    Serial,
    #[cfg(feature = "simple_io_channel")]
    SimpleIo,
}

/// Create a channel for communicating with the Untrusted Launcher.
fn get_channel<'a, X: Translator, A: Allocator + Sync>(
    kernel_args: &args::Args,
    mapper: &X,
    alloc: &'a A,
    sev_status: SevStatus,
) -> Box<dyn Channel + 'a> {
    // If we weren't told which channel to use, arbitrarily pick the first one in the `ChannelType`
    // enum. Depending on features that are enabled, this means that the enum acts as kind of a
    // reverse priority list for defaults.
    let chan_type = kernel_args
        .get("channel")
        .map(|chan_type| ChannelType::from_str(chan_type).unwrap())
        .unwrap_or_else(|| ChannelType::iter().next().unwrap());

    match chan_type {
        #[cfg(feature = "virtio_console_channel")]
        ChannelType::VirtioConsole => Box::new(virtio::get_console_channel(mapper, alloc)),
        #[cfg(feature = "vsock_channel")]
        ChannelType::VirtioVsock => Box::new(virtio::get_vsock_channel(mapper, alloc)),
        #[cfg(feature = "serial_channel")]
        ChannelType::Serial => Box::new(serial::Serial::new()),
        #[cfg(feature = "simple_io_channel")]
        ChannelType::SimpleIo => {
            Box::new(simpleio::SimpleIoChannel::new(mapper, alloc, sev_status))
        }
    }
}

/// Common panic routine for the kernel. This needs to be wrapped in a
/// panic_handler function in individual bootloader crates.
pub fn panic(info: &PanicInfo) -> ! {
    error!("PANIC: {}", info);
    shutdown::shutdown();
}
<commit_msg>Make the guest-host heap and mapper public (#3459)<commit_after>//
// Copyright 2022 The Project Oak Authors
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! Main 'kernel' for baremetal Oak Functions.
//!
//! This code takes care of initializing the x86-64 machine properly and
//! handing the reins over to the oak_baremetal_runtime. It is in a separate crate so that we
//! could support different boot protocols (eg Linux boot protocol or PVH) that
//! require different targets, linker scripts and/or provide machine
//! information in different data structures.
//!
//! Bootloaders (and VMMs under them) have to adhere to the following protocol:
//!   * Enter 64-bit long mode, and set up basic paging -- enough to load the code, as we will set
//!     up a full page table in `start_kernel`.
//!   * Implement a `#[panic_handler]` that delegates to `panic()` in this crate.
//!   * Call `start_kernel` from the entry point of the bootloader.

#![cfg_attr(not(test), no_std)]
#![feature(abi_x86_interrupt)]
#![feature(allocator_api)]
#![feature(asm_sym)]
#![feature(naked_functions)]
#![feature(once_cell)]

mod args;
mod avx;
mod boot;
mod descriptors;
mod elf;
mod ghcb;
#[cfg(feature = "c_interface")]
mod interface;
mod interrupts;
mod libm;
mod logging;
mod memory;
mod mm;
#[cfg(feature = "serial_channel")]
mod serial;
pub mod shutdown;
#[cfg(feature = "simple_io_channel")]
mod simpleio;
mod snp;
#[cfg(any(feature = "virtio_console_channel", feature = "vsock_channel"))]
mod virtio;

extern crate alloc;

use crate::{
    mm::Translator,
    snp::{get_snp_page_addresses, init_snp_pages},
};
use alloc::{alloc::Allocator, boxed::Box};
use core::{marker::Sync, panic::PanicInfo, str::FromStr};
use linked_list_allocator::LockedHeap;
use log::{error, info};
use mm::encrypted_mapper::{EncryptedPageTable, PhysOffset};
use oak_channel::Channel;
use oak_core::sync::OnceCell;
use oak_linux_boot_params::BootParams;
use sev_guest::msr::{change_snp_state_for_frame, get_sev_status, PageAssignment, SevStatus};
use strum::{EnumIter, EnumString, IntoEnumIterator};
use x86_64::{
    structures::paging::{MappedPageTable, Page, Size2MiB},
    VirtAddr,
};

/// The allocator for allocating space in the memory area that is shared with the hypervisor.
pub static GUEST_HOST_HEAP: OnceCell<LockedHeap> = OnceCell::new();

/// Translator that can convert between physical and virtual addresses.
pub static ADDRESS_TRANSLATOR: OnceCell<EncryptedPageTable<MappedPageTable<'static, PhysOffset>>> =
    OnceCell::new();

/// Main entry point for the kernel, to be called from bootloader.
pub fn start_kernel(info: &BootParams) -> Box<dyn Channel> {
    avx::enable_avx();
    descriptors::init_gdt();
    interrupts::init_idt();
    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    logging::init_logging();

    // We need to be done with the boot info struct before intializing memory. For example, the
    // multiboot protocol explicitly states data can be placed anywhere in memory; therefore, it's
    // highly likely we will overwrite some data after we initialize the heap. args::init_args()
    // caches the arguments (as long they are of reasonable length) in a static variable, allowing
    // us to refer to the args in the future.
    let kernel_args = args::init_args(info.args()).unwrap();
    info!("Kernel boot args: {}", kernel_args.args());

    let protocol = info.protocol();
    info!("Boot protocol:  {}", protocol);
    let snp_pages = if sev_status.contains(SevStatus::SNP_ACTIVE) {
        // We have to get the physical addresses of the CPUID pages now while the identity mapping
        // is still in place, but we can only initialize the instances after the new page
        // mappings have been set up.
        Some(get_snp_page_addresses(info))
    } else {
        None
    };

    // Safety: in the linker script we specify that the ELF header should be placed at 0x200000.
    let program_headers = unsafe { elf::get_phdrs(VirtAddr::new(0x20_0000)) };

    // Physical frame allocator: support up to 128 GiB of memory, for now.
    let mut frame_allocator = mm::init::<1024>(info.e820_table(), program_headers);

    let mut mapper = mm::init_paging(&mut frame_allocator, program_headers).unwrap();

    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        // Now that the page tables have been updated, we have to re-share the GHCB with the
        // hypervisor.
        ghcb::reshare_ghcb(&mut mapper);
        if sev_status.contains(SevStatus::SNP_ACTIVE) {
            // We must also initialise the CPUID and secrets pages when SEV-SNP is active. Panicking
            // is OK at this point, because these pages are required to support the full features
            // and we don't want to run without them.
            init_snp_pages(
                snp_pages.expect("Missing SNP CPUID and secrets pages."),
                &mapper,
            );
        }
    }

    // Allocate a section for guest-host communication (without the `ENCRYPTED` bit set)
    // We'll allocate 2*2MiB, as virtio needs more than 2 MiB for its data structures.
    let guest_host_frames = frame_allocator.allocate_contiguous(2).unwrap();
    let guest_host_pages = Page::range(
        mapper
            .translate_physical_frame(guest_host_frames.start)
            .unwrap(),
        mapper
            .translate_physical_frame(guest_host_frames.end)
            .unwrap(),
    );

    // If we are running on SNP we have to mark the frames as shared in the RMP. It is OK to crash
    // if we cannot mark the pages as shared in the RMP.
    if sev_status.contains(SevStatus::SNP_ACTIVE) {
        // TODO(#3414): Use the GHCB protocol when it is available.
        for frame in guest_host_frames {
            change_snp_state_for_frame(&frame, PageAssignment::Shared)
                .expect("Could not change SNP state for frame.");
        }
    }

    // Safety: initializing the new heap is safe as the frame allocator guarantees we're not
    // overwriting any other memory; writing to the static mut is safe as we're in the
    // initialization code and thus there can be no concurrent access.
    if GUEST_HOST_HEAP
        .set(unsafe { memory::init_guest_host_heap(guest_host_pages, &mut mapper) }.unwrap())
        .is_err()
    {
        panic!("Could not initialize the guest-host heap.");
    }

    if ADDRESS_TRANSLATOR.set(mapper).is_err() {
        panic!("Could not initialize the address translator.");
    }
    let mapper = ADDRESS_TRANSLATOR.get().unwrap();

    // If we don't find memory for heap, it's ok to panic.
    let heap_phys_frames = frame_allocator.largest_available().unwrap();
    memory::init_kernel_heap::<Size2MiB>(Page::range(
        mapper
            .translate_physical_frame(heap_phys_frames.start)
            .unwrap(),
        mapper
            .translate_physical_frame(heap_phys_frames.end + 1)
            .unwrap(),
    ))
    .unwrap();

    get_channel(
        &kernel_args,
        mapper,
        GUEST_HOST_HEAP.get().unwrap(),
        sev_status,
    )
}

#[derive(EnumIter, EnumString)]
#[strum(ascii_case_insensitive, serialize_all = "snake_case")]
enum ChannelType {
    #[cfg(feature = "virtio_console_channel")]
    VirtioConsole,
    #[cfg(feature = "vsock_channel")]
    VirtioVsock,
    #[cfg(feature = "serial_channel")]
    Serial,
    #[cfg(feature = "simple_io_channel")]
    SimpleIo,
}

/// Create a channel for communicating with the Untrusted Launcher.
fn get_channel<'a, X: Translator, A: Allocator + Sync>(
    kernel_args: &args::Args,
    mapper: &X,
    alloc: &'a A,
    sev_status: SevStatus,
) -> Box<dyn Channel + 'a> {
    // If we weren't told which channel to use, arbitrarily pick the first one in the `ChannelType`
    // enum. Depending on features that are enabled, this means that the enum acts as kind of a
    // reverse priority list for defaults.
    let chan_type = kernel_args
        .get("channel")
        .map(|chan_type| ChannelType::from_str(chan_type).unwrap())
        .unwrap_or_else(|| ChannelType::iter().next().unwrap());

    match chan_type {
        #[cfg(feature = "virtio_console_channel")]
        ChannelType::VirtioConsole => Box::new(virtio::get_console_channel(mapper, alloc)),
        #[cfg(feature = "vsock_channel")]
        ChannelType::VirtioVsock => Box::new(virtio::get_vsock_channel(mapper, alloc)),
        #[cfg(feature = "serial_channel")]
        ChannelType::Serial => Box::new(serial::Serial::new()),
        #[cfg(feature = "simple_io_channel")]
        ChannelType::SimpleIo => {
            Box::new(simpleio::SimpleIoChannel::new(mapper, alloc, sev_status))
        }
    }
}

/// Common panic routine for the kernel. This needs to be wrapped in a
/// panic_handler function in individual bootloader crates.
pub fn panic(info: &PanicInfo) -> ! {
    error!("PANIC: {}", info);
    shutdown::shutdown();
}
<|endoftext|>


<commit_before>use std::io::Cursor;

use msgpack::decode::*;

#[test]
fn pass_read_nil() {
    let buf = [0xc0];
    let mut cur = Cursor::new(&buf[..]);

    assert_eq!((), read_nil(&mut cur).unwrap());
    assert_eq!(1, cur.position());
}

#[test]
fn fail_read_nil_invalid_marker() {
    let buf = [0xc1];
    let mut cur = Cursor::new(&buf[..]);

    match read_nil(&mut cur) {
        Err(FixedValueReadError::TypeMismatch(..)) => (),
        other => panic!("unexpected result: {:?}", other)
    }
    assert_eq!(1, cur.position());
}

#[test]
fn fail_read_nil_invalid_marker_read() {
    let buf = [];
    let mut cur = Cursor::new(&buf[..]);

    match read_nil(&mut cur) {
        Err(FixedValueReadError::UnexpectedEOF) => (),
        other => panic!("unexpected result: {:?}", other)
    }
    assert_eq!(0, cur.position());
}
<commit_msg>[UT] Naming.<commit_after>use std::io::Cursor;

use msgpack::decode::*;

#[test]
fn pass() {
    let buf = [0xc0];
    let mut cur = Cursor::new(&buf[..]);

    assert_eq!((), read_nil(&mut cur).unwrap());
    assert_eq!(1, cur.position());
}

#[test]
fn fail_invalid_marker() {
    let buf = [0xc1];
    let mut cur = Cursor::new(&buf[..]);

    match read_nil(&mut cur) {
        Err(FixedValueReadError::TypeMismatch(..)) => (),
        other => panic!("unexpected result: {:?}", other)
    }
    assert_eq!(1, cur.position());
}

#[test]
fn fail_unexpected_eof() {
    let buf = [];
    let mut cur = Cursor::new(&buf[..]);

    match read_nil(&mut cur) {
        Err(FixedValueReadError::UnexpectedEOF) => (),
        other => panic!("unexpected result: {:?}", other)
    }
    assert_eq!(0, cur.position());
}
<|endoftext|>


<commit_before>extern crate test;
extern crate wheel_timer;

use test::Bencher;

use wheel_timer::WheelTimer;
use wheel_timer::Node;
use wheel_timer::Cons;
use wheel_timer::Nil;

#[test]
fn wheel_timer_schedule_test() {
  let mut timer = WheelTimer::<&'static str>::new(10);
  timer.schedule(3, "tick");

  timer.tick();
  timer.tick();
  timer.tick();

  let node = timer.tick();
  assert!(node != Nil);

  let val = match node {
    Cons(val, _) => val,
    Nil => ""
  };
  assert!(val == "tick");
}

#[test]
fn wheel_timer_tick_test() {
  let mut timer = WheelTimer::<uint>::new(10);

  for i in range(0, 10) {
    timer.schedule(i, i)
  }

  for i in range(0, 10) {
    let node = timer.tick();
    assert!(node != Nil);

    let val = match node {
      Cons(val, _) => val,
      Nil => -1
    };
    assert!(val == i);
  }
}

#[test]
fn wheel_timer_size_test() {
  let mut timer = WheelTimer::<uint>::new(10);

  for i in range(0, 10) {
    timer.schedule(i, i)
  }

  assert!(timer.size() == 10);

  for _ in timer {
  }

  assert!(timer.size() == 0);
}

#[bench]
fn bench_wheel_timer_drain(b: &mut Bencher) {
  let maxInterval = 20;
  let mut timer = WheelTimer::<uint>::new(maxInterval);

  b.iter(|| {
    // Fill
    for j in range(0u, 100u) {
      timer.schedule(j%maxInterval, j%maxInterval);
    }

    // Drain
    for k in timer {
      continue;
    }
  });
}

#[bench]
fn bench_wheel_timer_fill(b: &mut Bencher) {
  let maxInterval = 20;
  let mut timer = WheelTimer::<uint>::new(maxInterval);
  let mut i = 0;

  b.iter(|| {
    timer.schedule(i%maxInterval, i%maxInterval);
    i = i + 1;
  });
}

#[bench]
fn bench_wheel_timer_fast(b: &mut Bencher) {
  let maxInterval = 2;
  let mut timer = WheelTimer::<uint>::new(maxInterval);
  let mut i = 0;

  b.iter(|| {
    timer.schedule(i%maxInterval, i%maxInterval);
    timer.tick();
    i = i + 1;
  });
}
<commit_msg>Simplifying instantiation<commit_after>extern crate test;
extern crate wheel_timer;

use test::Bencher;

use wheel_timer::WheelTimer;
use wheel_timer::Cons;
use wheel_timer::Nil;

#[test]
fn wheel_timer_schedule_test() {
  let mut timer = WheelTimer::new(10);
  timer.schedule(3, "tick");

  timer.tick();
  timer.tick();
  timer.tick();

  let node = timer.tick();
  assert!(node != Nil);

  let val = match node {
    Cons(val, _) => val,
    Nil => ""
  };
  assert!(val == "tick");
}

#[test]
fn wheel_timer_tick_test() {
  let mut timer = WheelTimer::new(10);

  for i in range(0, 10) {
    timer.schedule(i, i)
  }

  for i in range(0, 10) {
    let node = timer.tick();
    assert!(node != Nil);

    let val = match node {
      Cons(val, _) => val,
      Nil => -1
    };
    assert!(val == i);
  }
}

#[test]
fn wheel_timer_size_test() {
  let mut timer = WheelTimer::new(10);

  for i in range(0, 10) {
    timer.schedule(i, i)
  }

  assert!(timer.size() == 10);

  for _ in timer {
  }

  assert!(timer.size() == 0);
}

#[bench]
fn bench_wheel_timer_drain(b: &mut Bencher) {
  let maxInterval = 20;
  let mut timer = WheelTimer::new(maxInterval);

  b.iter(|| {
    // Fill
    for j in range(0u, 100u) {
      timer.schedule(j%maxInterval, j%maxInterval);
    }

    // Drain
    for k in timer {
      continue;
    }
  });
}

#[bench]
fn bench_wheel_timer_fill(b: &mut Bencher) {
  let maxInterval = 20;
  let mut timer = WheelTimer::new(maxInterval);
  let mut i = 0;

  b.iter(|| {
    timer.schedule(i%maxInterval, i%maxInterval);
    i = i + 1;
  });
}

#[bench]
fn bench_wheel_timer_fast(b: &mut Bencher) {
  let maxInterval = 2;
  let mut timer = WheelTimer::new(maxInterval);
  let mut i = 0;

  b.iter(|| {
    timer.schedule(i%maxInterval, i%maxInterval);
    timer.tick();
    i = i + 1;
  });
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate maplit;
extern crate dot_dsl;

use dot_dsl::graph::graph_items::edge::Edge;
use dot_dsl::graph::graph_items::node::Node;
use dot_dsl::graph::Graph;

#[test]
fn test_empty_graph() {
    let graph = Graph::new();

    assert!(graph.nodes.is_empty());

    assert!(graph.edges.is_empty());

    assert!(graph.attrs.is_empty());
}

#[test]
#[ignore]
fn test_graph_with_one_node() {
    let nodes = vec![Node::new("a")];

    let graph = Graph::new().with_nodes(&nodes);

    assert!(graph.edges.is_empty());

    assert!(graph.attrs.is_empty());

    assert_eq!(graph.nodes, vec![Node::new("a")]);
}

#[test]
#[ignore]
fn test_graph_with_one_node_with_keywords() {
    let nodes = vec![Node::new("a").with_attrs(&[("color", "green")])];

    let graph = Graph::new().with_nodes(&nodes);

    assert!(graph.edges.is_empty());

    assert!(graph.attrs.is_empty());

    assert_eq!(
        graph.nodes,
        vec![Node::new("a").with_attrs(&[("color", "green")])]
    );
}

#[test]
#[ignore]
fn test_graph_with_one_edge() {
    let edges = vec![Edge::new("a", "b")];

    let graph = Graph::new().with_edges(&edges);

    assert!(graph.nodes.is_empty());

    assert!(graph.attrs.is_empty());

    assert_eq!(graph.edges, vec![Edge::new("a", "b")]);
}

#[test]
#[ignore]
fn test_graph_with_one_attribute() {
    let graph = Graph::new().with_attrs(&[("foo", "1")]);

    let expected_attrs = hashmap! {
        "foo".to_string() => "1".to_string(),
    };

    assert!(graph.nodes.is_empty());

    assert!(graph.edges.is_empty());

    assert_eq!(graph.attrs, expected_attrs);
}

#[test]
#[ignore]
fn test_graph_with_attributes() {
    let nodes = vec![
        Node::new("a").with_attrs(&[("color", "green")]),
        Node::new("c"),
        Node::new("b").with_attrs(&[("label", "Beta!")]),
    ];

    let edges = vec![
        Edge::new("b", "c"),
        Edge::new("a", "b").with_attrs(&[("color", "blue")]),
    ];

    let attrs = vec![("foo", "1"), ("title", "Testing Attrs"), ("bar", "true")];

    let expected_attrs = hashmap! {
        "foo".to_string() => "1".to_string(),
        "title".to_string() => "Testing Attrs".to_string(),
        "bar".to_string() => "true".to_string(),
    };

    let graph = Graph::new()
        .with_nodes(&nodes)
        .with_edges(&edges)
        .with_attrs(&attrs);

    assert_eq!(
        graph.nodes,
        vec![
            Node::new("a").with_attrs(&[("color", "green")]),
            Node::new("c"),
            Node::new("b").with_attrs(&[("label", "Beta!")]),
        ]
    );

    assert_eq!(
        graph.edges,
        vec![
            Edge::new("b", "c"),
            Edge::new("a", "b").with_attrs(&[("color", "blue")]),
        ]
    );

    assert_eq!(graph.attrs, expected_attrs);
}

#[test]
#[ignore]
fn test_graph_stores_attributes() {
    let attributes = [("foo", "bar"), ("bat", "baz"), ("bim", "bef")];
    let graph = Graph::new().with_nodes(
        &['a', 'b', 'c']
            .iter()
            .enumerate()
            .map(|(i, n)| Node::new(&n.to_string()).with_attrs(&attributes[i..i + 1]))
            .collect::<Vec<_>>(),
    );

    assert_eq!(
        graph
            .get_node("c")
            .expect("node must be stored")
            .get_attr("bim"),
        Some("bef")
    );
}
<commit_msg>dot-dsl: Use zip instead of enumerate + map-with-indexing (#817)<commit_after>#[macro_use]
extern crate maplit;
extern crate dot_dsl;

use dot_dsl::graph::graph_items::edge::Edge;
use dot_dsl::graph::graph_items::node::Node;
use dot_dsl::graph::Graph;

#[test]
fn test_empty_graph() {
    let graph = Graph::new();

    assert!(graph.nodes.is_empty());

    assert!(graph.edges.is_empty());

    assert!(graph.attrs.is_empty());
}

#[test]
#[ignore]
fn test_graph_with_one_node() {
    let nodes = vec![Node::new("a")];

    let graph = Graph::new().with_nodes(&nodes);

    assert!(graph.edges.is_empty());

    assert!(graph.attrs.is_empty());

    assert_eq!(graph.nodes, vec![Node::new("a")]);
}

#[test]
#[ignore]
fn test_graph_with_one_node_with_keywords() {
    let nodes = vec![Node::new("a").with_attrs(&[("color", "green")])];

    let graph = Graph::new().with_nodes(&nodes);

    assert!(graph.edges.is_empty());

    assert!(graph.attrs.is_empty());

    assert_eq!(
        graph.nodes,
        vec![Node::new("a").with_attrs(&[("color", "green")])]
    );
}

#[test]
#[ignore]
fn test_graph_with_one_edge() {
    let edges = vec![Edge::new("a", "b")];

    let graph = Graph::new().with_edges(&edges);

    assert!(graph.nodes.is_empty());

    assert!(graph.attrs.is_empty());

    assert_eq!(graph.edges, vec![Edge::new("a", "b")]);
}

#[test]
#[ignore]
fn test_graph_with_one_attribute() {
    let graph = Graph::new().with_attrs(&[("foo", "1")]);

    let expected_attrs = hashmap! {
        "foo".to_string() => "1".to_string(),
    };

    assert!(graph.nodes.is_empty());

    assert!(graph.edges.is_empty());

    assert_eq!(graph.attrs, expected_attrs);
}

#[test]
#[ignore]
fn test_graph_with_attributes() {
    let nodes = vec![
        Node::new("a").with_attrs(&[("color", "green")]),
        Node::new("c"),
        Node::new("b").with_attrs(&[("label", "Beta!")]),
    ];

    let edges = vec![
        Edge::new("b", "c"),
        Edge::new("a", "b").with_attrs(&[("color", "blue")]),
    ];

    let attrs = vec![("foo", "1"), ("title", "Testing Attrs"), ("bar", "true")];

    let expected_attrs = hashmap! {
        "foo".to_string() => "1".to_string(),
        "title".to_string() => "Testing Attrs".to_string(),
        "bar".to_string() => "true".to_string(),
    };

    let graph = Graph::new()
        .with_nodes(&nodes)
        .with_edges(&edges)
        .with_attrs(&attrs);

    assert_eq!(
        graph.nodes,
        vec![
            Node::new("a").with_attrs(&[("color", "green")]),
            Node::new("c"),
            Node::new("b").with_attrs(&[("label", "Beta!")]),
        ]
    );

    assert_eq!(
        graph.edges,
        vec![
            Edge::new("b", "c"),
            Edge::new("a", "b").with_attrs(&[("color", "blue")]),
        ]
    );

    assert_eq!(graph.attrs, expected_attrs);
}

#[test]
#[ignore]
fn test_graph_stores_attributes() {
    let attributes = [("foo", "bar"), ("bat", "baz"), ("bim", "bef")];
    let graph = Graph::new().with_nodes(
        &['a', 'b', 'c']
            .iter()
            .zip(attributes.iter())
            .map(|(name, &attr)| Node::new(&name.to_string()).with_attrs(&[attr]))
            .collect::<Vec<_>>(),
    );

    assert_eq!(
        graph
            .get_node("c")
            .expect("node must be stored")
            .get_attr("bim"),
        Some("bef")
    );
}
<|endoftext|>


<commit_before>use crate::faker::filesystem::raw::*;
use crate::impls::std::path::PathFaker;
use crate::locales::Data;
use crate::{Dummy, Fake};
use rand::seq::SliceRandom;
use rand::Rng;
use std::path::PathBuf;

impl<L: Data> Dummy<FilePath<L>> for PathBuf {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FilePath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, L::PATH_EXTENSIONS, 4);
        faker.fake_with_rng(rng)
    }
}

impl<L: Data> Dummy<FilePath<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FilePath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, L::PATH_EXTENSIONS, 4);
        let p: PathBuf = faker.fake_with_rng(rng);
        p.to_string_lossy().into()
    }
}

impl<L: Data> Dummy<FileName<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FileName<L>, rng: &mut R) -> Self {
        let name = L::PATH_SEGMENTS.choose(rng).unwrap();
        let ext = L::PATH_EXTENSIONS.choose(rng).unwrap();
        format!("{}.{}", name, ext)
    }
}

impl<L: Data> Dummy<FileExtension<L>> for &str {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FileExtension<L>, rng: &mut R) -> Self {
        L::PATH_EXTENSIONS.choose(rng).unwrap()
    }
}

impl<L: Data> Dummy<FileExtension<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FileExtension<L>, rng: &mut R) -> Self {
        let ext = L::PATH_EXTENSIONS.choose(rng).unwrap();
        ext.to_string()
    }
}

impl<L: Data> Dummy<DirPath<L>> for PathBuf {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &DirPath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, &[], 4);
        faker.fake_with_rng(rng)
    }
}

impl<L: Data> Dummy<DirPath<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &DirPath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, &[], 4);
        let p: PathBuf = faker.fake_with_rng(rng);
        p.to_string_lossy().into()
    }
}
<commit_msg>fix clippy::inefficient_to_string<commit_after>use crate::faker::filesystem::raw::*;
use crate::impls::std::path::PathFaker;
use crate::locales::Data;
use crate::{Dummy, Fake};
use rand::seq::SliceRandom;
use rand::Rng;
use std::path::PathBuf;

impl<L: Data> Dummy<FilePath<L>> for PathBuf {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FilePath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, L::PATH_EXTENSIONS, 4);
        faker.fake_with_rng(rng)
    }
}

impl<L: Data> Dummy<FilePath<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FilePath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, L::PATH_EXTENSIONS, 4);
        let p: PathBuf = faker.fake_with_rng(rng);
        p.to_string_lossy().into()
    }
}

impl<L: Data> Dummy<FileName<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FileName<L>, rng: &mut R) -> Self {
        let name = L::PATH_SEGMENTS.choose(rng).unwrap();
        let ext = L::PATH_EXTENSIONS.choose(rng).unwrap();
        format!("{}.{}", name, ext)
    }
}

impl<L: Data> Dummy<FileExtension<L>> for &str {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FileExtension<L>, rng: &mut R) -> Self {
        L::PATH_EXTENSIONS.choose(rng).unwrap()
    }
}

impl<L: Data> Dummy<FileExtension<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &FileExtension<L>, rng: &mut R) -> Self {
        let ext = L::PATH_EXTENSIONS.choose(rng).unwrap();
        (*ext).to_string()
    }
}

impl<L: Data> Dummy<DirPath<L>> for PathBuf {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &DirPath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, &[], 4);
        faker.fake_with_rng(rng)
    }
}

impl<L: Data> Dummy<DirPath<L>> for String {
    fn dummy_with_rng<R: Rng + ?Sized>(_: &DirPath<L>, rng: &mut R) -> Self {
        let faker = PathFaker::new(L::PATH_ROOT_DIRS, L::PATH_SEGMENTS, &[], 4);
        let p: PathBuf = faker.fake_with_rng(rng);
        p.to_string_lossy().into()
    }
}
<|endoftext|>