<commit_before>use crate::config::Config;
use crate::errors::RunError;
use crate::statemachine::*;
use log::{debug, trace};
use nix::errno::Errno;
use nix::sys::signal::Signal;
use nix::sys::wait::*;
use nix::unistd::Pid;
use nix::Error as NixErr;
use std::collections::{HashMap, HashSet};

pub fn create_state_machine<'a>(
    test: Pid,
    traces: &'a mut TraceMap,
    config: &'a Config,
) -> (TestState, LinuxData<'a>) {
    let mut data = LinuxData::new(traces, config);
    data.parent = test;
    (TestState::start_state(), data)
}

pub type UpdateContext = (TestState, TracerAction<ProcessInfo>);

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct ProcessInfo {
    pid: Pid,
    signal: Option<Signal>,
}

impl ProcessInfo {
    fn new(pid: Pid, signal: Option<Signal>) -> Self {
        Self {
            pid,
            signal
        }
    }
}

impl From<Pid> for ProcessInfo {
    fn from(pid: Pid) -> Self {
        ProcessInfo::new(pid, None)
    }
}

impl From<&Pid> for ProcessInfo {
    fn from(pid: &Pid) -> Self {
        ProcessInfo::new(*pid, None)
    }
}

/// Handle to linux process state
pub struct LinuxData<'a> {
    /// Recent results from waitpid to be handled by statemachine
    wait_queue: Vec<WaitStatus>,
    /// Current Pid to process
    current: Pid,
    /// Parent PID of test process
    parent: Pid,
    /// Map of addresses to breakpoints
    breakpoints: HashMap<u64, Breakpoint>,
    /// Instrumentation points in code with associated coverage data
    traces: &'a mut TraceMap,
    /// Program config
    config: &'a Config,
    /// Thread count. Hopefully getting rid of in future
    thread_count: isize,
}

impl<'a> StateData for LinuxData<'a> {
    fn start(&mut self) -> Result<Option<TestState>, RunError> {
        match waitpid(self.current, Some(WaitPidFlag::WNOHANG)) {
            Ok(WaitStatus::StillAlive) => Ok(None),
            Ok(sig @ WaitStatus::Stopped(_, Signal::SIGTRAP)) => {
                if let WaitStatus::Stopped(child, _) = sig {
                    self.current = child;
                }
                trace!("Caught inferior transitioning to Initialise state");
                Ok(Some(TestState::Initialise))
            }
            Ok(_) => Err(RunError::TestRuntime(
                "Unexpected signal when starting test".to_string(),
            )),
            Err(e) => Err(RunError::TestRuntime(format!(
                "Error when starting test: {}",
                e
            ))),
        }
    }

    fn init(&mut self) -> Result<TestState, RunError> {
        trace_children(self.current)?;
        for trace in self.traces.all_traces() {
            if let Some(addr) = trace.address {
                match Breakpoint::new(self.current, addr) {
                    Ok(bp) => {
                        let _ = self.breakpoints.insert(addr, bp);
                    }
                    Err(e) if e == NixErr::Sys(Errno::EIO) => {
                        return Err(RunError::TestRuntime(
                            "ERROR: Tarpaulin cannot find code addresses \
                             check that pie is disabled for your linker. \
                             If linking with gcc try adding -C link-args=-no-pie \
                             to your rust flags"
                                .to_string(),
                        ));
                    }
                    Err(NixErr::UnsupportedOperation) => {
                        debug!("Instrumentation address clash, ignoring 0x{:x}", addr);
                    }
                    Err(_) => {
                        return Err(RunError::TestRuntime(
                            "Failed to instrument test executable".to_string(),
                        ));
                    }
                }
            }
        }

        if continue_exec(self.parent, None).is_ok() {
            trace!("Initialised inferior, transitioning to wait state");
            Ok(TestState::wait_state())
        } else {
            Err(RunError::TestRuntime(
                "Test didn't launch correctly".to_string(),
            ))
        }
    }

    fn wait(&mut self) -> Result<Option<TestState>, RunError> {
        let mut result = Ok(None);
        let mut running = true;
        while running {
            let wait = waitpid(
                Pid::from_raw(-1),
                Some(WaitPidFlag::WNOHANG | WaitPidFlag::__WALL),
            );
            match wait {
                Ok(WaitStatus::StillAlive) => {
                    running = false;
                }
                Ok(WaitStatus::Exited(_, _)) => {
                    self.wait_queue.push(wait.unwrap());
                    result = Ok(Some(TestState::Stopped));
                    running = false;
                },
                Ok(WaitStatus::PtraceEvent(_, _, _)) => {
                    self.wait_queue.push(wait.unwrap());
                    result = Ok(Some(TestState::Stopped));
                    running = false;
                },
                Ok(s) => {
                    self.wait_queue.push(s);
                    result = Ok(Some(TestState::Stopped));
                }
                Err(e) => {
                    running = false;
                    result = Err(RunError::TestRuntime(
                            format!("An error occurred while waiting for response from test: {}", e)))
                },
            }
        }
        if !self.wait_queue.is_empty() {
            trace!("Result queue is {:?}", self.wait_queue);
        }
        result
    }

    fn stop(&mut self) -> Result<TestState, RunError> {
        let mut actions = Vec::new();
        let mut pcs = HashSet::new();
        let mut result = Ok(TestState::wait_state());
        let pending = self.wait_queue.clone();
        self.wait_queue.clear();
        for status in &pending {
            let state = match status {
                WaitStatus::PtraceEvent(c, s, e) => match self.handle_ptrace_event(*c, *s, *e) {
                    Ok(s) => Ok(s),
                    Err(e) => Err(RunError::TestRuntime(format!(
                        "Error occurred when handling ptrace event: {}",
                        e
                    ))),
                },
                WaitStatus::Stopped(c, Signal::SIGTRAP) => {
                    self.current = *c;
                    match self.collect_coverage_data(&mut pcs) {
                        Ok(s) => Ok(s),
                        Err(e) => Err(RunError::TestRuntime(format!(
                            "Error when collecting coverage: {}",
                            e
                        ))),
                    }
                }
                WaitStatus::Stopped(child, Signal::SIGSTOP) => {
                    Ok((TestState::wait_state(), TracerAction::Continue(child.into())))
                },
                WaitStatus::Stopped(_, Signal::SIGSEGV) => Err(RunError::TestRuntime(
                    "A segfault occurred while executing tests".to_string(),
                )),
                WaitStatus::Stopped(child, Signal::SIGILL) => Err(RunError::TestRuntime(format!(
                    "Error running test - SIGILL raised in {}",
                    child
                ))),
                WaitStatus::Stopped(c, s) => {
                    let sig = if self.config.forward_signals {
                        Some(*s)
                    } else {
                        None
                    };
                    let info = ProcessInfo::new(*c, sig);
                    Ok((TestState::wait_state(), TracerAction::TryContinue(info)))
                }
                WaitStatus::Signaled(c, s, f) => {
                    if let Ok(s) = self.handle_signaled(c, s, *f) {
                        Ok(s)
                    } else {
                        Err(RunError::TestRuntime(
                            "Attempting to handle tarpaulin being signaled".to_string(),
                        ))
                    }
                }
                WaitStatus::Exited(child, ec) => {
                    for ref mut value in self.breakpoints.values_mut() {
                        value.thread_killed(*child);
                    }
                    trace!("Exited {:?} parent {:?}", child, self.parent);
                    if child == &self.parent {
                        Ok((TestState::End(*ec), TracerAction::Nothing))
                    } else {
                        // Process may have already been destroyed. This is just incase
                        Ok((TestState::wait_state(), TracerAction::TryContinue(self.parent.into())))
                    }
                }
                _ => Err(RunError::TestRuntime(
                    "An unexpected signal has been caught by tarpaulin!".to_string(),
                )),
            };
            match state {
                Ok((TestState::Waiting{..}, action)) => {
                    actions.push(action);
                },
                Ok((state, action)) => {
                    result = Ok(state);
                    actions.push(action);
                },
                Err(e) => result = Err(e),
            }
        }
        let mut continued = false;
        for a in &actions {
            match a {
                TracerAction::TryContinue(t) => {
                    continued = true;
                    let _ = continue_exec(t.pid, t.signal);
                },
                TracerAction::Continue(t) => {
                    continued = true;
                    continue_exec(t.pid, t.signal)?;
                },
                TracerAction::Step(t) => {
                    continued = true;
                    single_step(t.pid)?;
                },
                TracerAction::Detach(t) => {
                    continued = true;
                    detach_child(t.pid)?;
                },
                _ => {},
            }
        }
        if !continued {
            trace!("No action suggested to continue tracee. Attempting a continue");
            let _ = continue_exec(self.parent, None);
        }
        result
    }
}

impl<'a> LinuxData<'a> {
    pub fn new(traces: &'a mut TraceMap, config: &'a Config) -> LinuxData<'a> {
        LinuxData {
            wait_queue: Vec::new(),
            current: Pid::from_raw(0),
            parent: Pid::from_raw(0),
            breakpoints: HashMap::new(),
            traces,
            config,
            thread_count: 0,
        }
    }

    fn handle_ptrace_event(&mut self,
        child: Pid,
        sig: Signal,
        event: i32,
    ) -> Result<(TestState, TracerAction<ProcessInfo>), RunError> {
        use nix::libc::*;

        if sig == Signal::SIGTRAP {
            match event {
                PTRACE_EVENT_CLONE => match get_event_data(child) {
                    Ok(t) => {
                        trace!("New thread spawned {}", t);
                        self.thread_count += 1;
                        Ok((TestState::wait_state(), TracerAction::Continue(child.into())))
                    }
                    Err(e) => {
                        trace!("Error in clone event {:?}", e);
                        Err(RunError::TestRuntime(
                            "Error occurred upon test executable thread creation".to_string(),
                        ))
                    }
                },
                PTRACE_EVENT_FORK | PTRACE_EVENT_VFORK => {
                    trace!("Caught fork event");
                    Ok((TestState::wait_state(), TracerAction::Continue(child.into())))
                }
                PTRACE_EVENT_EXEC => {
                    trace!("Child execed other process - detaching ptrace");
                    Ok((TestState::wait_state(), TracerAction::Detach(child.into())))
                }
                PTRACE_EVENT_EXIT => {
                    trace!("Child exiting");
                    self.thread_count -= 1;
                    Ok((TestState::wait_state(), TracerAction::TryContinue(child.into())))
                }
                _ => Err(RunError::TestRuntime(format!(
                    "Unrecognised ptrace event {}",
                    event
                ))),
            }
        } else {
            trace!("Unexpected signal with ptrace event {}", event);
            trace!("Signal: {:?}", sig);
            Err(RunError::TestRuntime("Unexpected signal".to_string()))
        }
    }

    fn collect_coverage_data(&mut self,
                             visited_pcs: &mut HashSet<u64>) -> Result<UpdateContext, RunError> {
        let mut action = None;
        if let Ok(rip) = current_instruction_pointer(self.current) {
            let rip = (rip - 1) as u64;
            trace!("Hit address 0x{:x}", rip);
            if self.breakpoints.contains_key(&rip) {
                let bp = &mut self.breakpoints.get_mut(&rip).unwrap();
                let updated = if visited_pcs.contains(&rip) {
                    let _ = bp.jump_to(self.current);
                    (true, TracerAction::Continue(self.current.into()))
                } else {
                    let enable = self.config.count;
                    // Don't reenable if multithreaded as can't yet sort out segfault issue
                    if let Ok(x) = bp.process(self.current, enable) {
                        x
                    } else {
                        // So failed to process a breakpoint.. Still continue to avoid
                        // stalling
                        (false, TracerAction::Continue(self.current.into()))
                    }
                };
                if updated.0 {
                    if let Some(ref mut t) = self.traces.get_trace_mut(rip) {
                        if let CoverageStat::Line(ref mut x) = t.stats {
                            trace!("Incrementing hit count for trace");
                            *x += 1;
                        }
                    }
                }
                action = Some(updated.1);
            }
        }
        let action = action.unwrap_or_else(|| TracerAction::Continue(self.current.into()));
        Ok((TestState::wait_state(), action))
    }

    fn handle_signaled(&mut self, pid: &Pid, sig: &Signal, flag: bool) -> Result<UpdateContext, RunError> {
        match (sig, flag) {
            (Signal::SIGTRAP, true) => {
                Ok((TestState::wait_state(), TracerAction::Continue(pid.into())))
            }
            _ => Err(RunError::StateMachine("Unexpected stop".to_string())),
        }
    }
}
<commit_msg>Added PC printout for SIGILL<commit_after>use crate::config::Config;
use crate::errors::RunError;
use crate::statemachine::*;
use log::{debug, trace};
use nix::errno::Errno;
use nix::sys::signal::Signal;
use nix::sys::wait::*;
use nix::unistd::Pid;
use nix::Error as NixErr;
use std::collections::{HashMap, HashSet};

pub fn create_state_machine<'a>(
    test: Pid,
    traces: &'a mut TraceMap,
    config: &'a Config,
) -> (TestState, LinuxData<'a>) {
    let mut data = LinuxData::new(traces, config);
    data.parent = test;
    (TestState::start_state(), data)
}

pub type UpdateContext = (TestState, TracerAction<ProcessInfo>);

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct ProcessInfo {
    pid: Pid,
    signal: Option<Signal>,
}

impl ProcessInfo {
    fn new(pid: Pid, signal: Option<Signal>) -> Self {
        Self {
            pid,
            signal
        }
    }
}

impl From<Pid> for ProcessInfo {
    fn from(pid: Pid) -> Self {
        ProcessInfo::new(pid, None)
    }
}

impl From<&Pid> for ProcessInfo {
    fn from(pid: &Pid) -> Self {
        ProcessInfo::new(*pid, None)
    }
}

/// Handle to linux process state
pub struct LinuxData<'a> {
    /// Recent results from waitpid to be handled by statemachine
    wait_queue: Vec<WaitStatus>,
    /// Current Pid to process
    current: Pid,
    /// Parent PID of test process
    parent: Pid,
    /// Map of addresses to breakpoints
    breakpoints: HashMap<u64, Breakpoint>,
    /// Instrumentation points in code with associated coverage data
    traces: &'a mut TraceMap,
    /// Program config
    config: &'a Config,
    /// Thread count. Hopefully getting rid of in future
    thread_count: isize,
}

impl<'a> StateData for LinuxData<'a> {
    fn start(&mut self) -> Result<Option<TestState>, RunError> {
        match waitpid(self.current, Some(WaitPidFlag::WNOHANG)) {
            Ok(WaitStatus::StillAlive) => Ok(None),
            Ok(sig @ WaitStatus::Stopped(_, Signal::SIGTRAP)) => {
                if let WaitStatus::Stopped(child, _) = sig {
                    self.current = child;
                }
                trace!("Caught inferior transitioning to Initialise state");
                Ok(Some(TestState::Initialise))
            }
            Ok(_) => Err(RunError::TestRuntime(
                "Unexpected signal when starting test".to_string(),
            )),
            Err(e) => Err(RunError::TestRuntime(format!(
                "Error when starting test: {}",
                e
            ))),
        }
    }

    fn init(&mut self) -> Result<TestState, RunError> {
        trace_children(self.current)?;
        for trace in self.traces.all_traces() {
            if let Some(addr) = trace.address {
                match Breakpoint::new(self.current, addr) {
                    Ok(bp) => {
                        let _ = self.breakpoints.insert(addr, bp);
                    }
                    Err(e) if e == NixErr::Sys(Errno::EIO) => {
                        return Err(RunError::TestRuntime(
                            "ERROR: Tarpaulin cannot find code addresses \
                             check that pie is disabled for your linker. \
                             If linking with gcc try adding -C link-args=-no-pie \
                             to your rust flags"
                                .to_string(),
                        ));
                    }
                    Err(NixErr::UnsupportedOperation) => {
                        debug!("Instrumentation address clash, ignoring 0x{:x}", addr);
                    }
                    Err(_) => {
                        return Err(RunError::TestRuntime(
                            "Failed to instrument test executable".to_string(),
                        ));
                    }
                }
            }
        }

        if continue_exec(self.parent, None).is_ok() {
            trace!("Initialised inferior, transitioning to wait state");
            Ok(TestState::wait_state())
        } else {
            Err(RunError::TestRuntime(
                "Test didn't launch correctly".to_string(),
            ))
        }
    }

    fn wait(&mut self) -> Result<Option<TestState>, RunError> {
        let mut result = Ok(None);
        let mut running = true;
        while running {
            let wait = waitpid(
                Pid::from_raw(-1),
                Some(WaitPidFlag::WNOHANG | WaitPidFlag::__WALL),
            );
            match wait {
                Ok(WaitStatus::StillAlive) => {
                    running = false;
                }
                Ok(WaitStatus::Exited(_, _)) => {
                    self.wait_queue.push(wait.unwrap());
                    result = Ok(Some(TestState::Stopped));
                    running = false;
                },
                Ok(WaitStatus::PtraceEvent(_, _, _)) => {
                    self.wait_queue.push(wait.unwrap());
                    result = Ok(Some(TestState::Stopped));
                    running = false;
                },
                Ok(s) => {
                    self.wait_queue.push(s);
                    result = Ok(Some(TestState::Stopped));
                }
                Err(e) => {
                    running = false;
                    result = Err(RunError::TestRuntime(
                            format!("An error occurred while waiting for response from test: {}", e)))
                },
            }
        }
        if !self.wait_queue.is_empty() {
            trace!("Result queue is {:?}", self.wait_queue);
        }
        result
    }

    fn stop(&mut self) -> Result<TestState, RunError> {
        let mut actions = Vec::new();
        let mut pcs = HashSet::new();
        let mut result = Ok(TestState::wait_state());
        let pending = self.wait_queue.clone();
        self.wait_queue.clear();
        for status in &pending {
            let state = match status {
                WaitStatus::PtraceEvent(c, s, e) => match self.handle_ptrace_event(*c, *s, *e) {
                    Ok(s) => Ok(s),
                    Err(e) => Err(RunError::TestRuntime(format!(
                        "Error occurred when handling ptrace event: {}",
                        e
                    ))),
                },
                WaitStatus::Stopped(c, Signal::SIGTRAP) => {
                    self.current = *c;
                    match self.collect_coverage_data(&mut pcs) {
                        Ok(s) => Ok(s),
                        Err(e) => Err(RunError::TestRuntime(format!(
                            "Error when collecting coverage: {}",
                            e
                        ))),
                    }
                }
                WaitStatus::Stopped(child, Signal::SIGSTOP) => {
                    Ok((TestState::wait_state(), TracerAction::Continue(child.into())))
                },
                WaitStatus::Stopped(_, Signal::SIGSEGV) => Err(RunError::TestRuntime(
                    "A segfault occurred while executing tests".to_string(),
                )),
                WaitStatus::Stopped(child, Signal::SIGILL) => {
                    let pc = current_instruction_pointer(*child).unwrap_or_else(|_| 1) - 1;
                    trace!("SIGILL raised. Child program counter is: 0x{:x}", pc);
                    Err(RunError::TestRuntime(format!("Error running test - SIGILL raised in {}", child)))
                },
                WaitStatus::Stopped(c, s) => {
                    let sig = if self.config.forward_signals {
                        Some(*s)
                    } else {
                        None
                    };
                    let info = ProcessInfo::new(*c, sig);
                    Ok((TestState::wait_state(), TracerAction::TryContinue(info)))
                }
                WaitStatus::Signaled(c, s, f) => {
                    if let Ok(s) = self.handle_signaled(c, s, *f) {
                        Ok(s)
                    } else {
                        Err(RunError::TestRuntime(
                            "Attempting to handle tarpaulin being signaled".to_string(),
                        ))
                    }
                }
                WaitStatus::Exited(child, ec) => {
                    for ref mut value in self.breakpoints.values_mut() {
                        value.thread_killed(*child);
                    }
                    trace!("Exited {:?} parent {:?}", child, self.parent);
                    if child == &self.parent {
                        Ok((TestState::End(*ec), TracerAction::Nothing))
                    } else {
                        // Process may have already been destroyed. This is just incase
                        Ok((TestState::wait_state(), TracerAction::TryContinue(self.parent.into())))
                    }
                }
                _ => Err(RunError::TestRuntime(
                    "An unexpected signal has been caught by tarpaulin!".to_string(),
                )),
            };
            match state {
                Ok((TestState::Waiting{..}, action)) => {
                    actions.push(action);
                },
                Ok((state, action)) => {
                    result = Ok(state);
                    actions.push(action);
                },
                Err(e) => result = Err(e),
            }
        }
        let mut continued = false;
        for a in &actions {
            match a {
                TracerAction::TryContinue(t) => {
                    continued = true;
                    let _ = continue_exec(t.pid, t.signal);
                },
                TracerAction::Continue(t) => {
                    continued = true;
                    continue_exec(t.pid, t.signal)?;
                },
                TracerAction::Step(t) => {
                    continued = true;
                    single_step(t.pid)?;
                },
                TracerAction::Detach(t) => {
                    continued = true;
                    detach_child(t.pid)?;
                },
                _ => {},
            }
        }
        if !continued {
            trace!("No action suggested to continue tracee. Attempting a continue");
            let _ = continue_exec(self.parent, None);
        }
        result
    }
}

impl<'a> LinuxData<'a> {
    pub fn new(traces: &'a mut TraceMap, config: &'a Config) -> LinuxData<'a> {
        LinuxData {
            wait_queue: Vec::new(),
            current: Pid::from_raw(0),
            parent: Pid::from_raw(0),
            breakpoints: HashMap::new(),
            traces,
            config,
            thread_count: 0,
        }
    }

    fn handle_ptrace_event(&mut self,
        child: Pid,
        sig: Signal,
        event: i32,
    ) -> Result<(TestState, TracerAction<ProcessInfo>), RunError> {
        use nix::libc::*;

        if sig == Signal::SIGTRAP {
            match event {
                PTRACE_EVENT_CLONE => match get_event_data(child) {
                    Ok(t) => {
                        trace!("New thread spawned {}", t);
                        self.thread_count += 1;
                        Ok((TestState::wait_state(), TracerAction::Continue(child.into())))
                    }
                    Err(e) => {
                        trace!("Error in clone event {:?}", e);
                        Err(RunError::TestRuntime(
                            "Error occurred upon test executable thread creation".to_string(),
                        ))
                    }
                },
                PTRACE_EVENT_FORK | PTRACE_EVENT_VFORK => {
                    trace!("Caught fork event");
                    Ok((TestState::wait_state(), TracerAction::Continue(child.into())))
                }
                PTRACE_EVENT_EXEC => {
                    trace!("Child execed other process - detaching ptrace");
                    Ok((TestState::wait_state(), TracerAction::Detach(child.into())))
                }
                PTRACE_EVENT_EXIT => {
                    trace!("Child exiting");
                    self.thread_count -= 1;
                    Ok((TestState::wait_state(), TracerAction::TryContinue(child.into())))
                }
                _ => Err(RunError::TestRuntime(format!(
                    "Unrecognised ptrace event {}",
                    event
                ))),
            }
        } else {
            trace!("Unexpected signal with ptrace event {}", event);
            trace!("Signal: {:?}", sig);
            Err(RunError::TestRuntime("Unexpected signal".to_string()))
        }
    }

    fn collect_coverage_data(&mut self,
                             visited_pcs: &mut HashSet<u64>) -> Result<UpdateContext, RunError> {
        let mut action = None;
        if let Ok(rip) = current_instruction_pointer(self.current) {
            let rip = (rip - 1) as u64;
            trace!("Hit address 0x{:x}", rip);
            if self.breakpoints.contains_key(&rip) {
                let bp = &mut self.breakpoints.get_mut(&rip).unwrap();
                let updated = if visited_pcs.contains(&rip) {
                    let _ = bp.jump_to(self.current);
                    (true, TracerAction::Continue(self.current.into()))
                } else {
                    let enable = self.config.count;
                    // Don't reenable if multithreaded as can't yet sort out segfault issue
                    if let Ok(x) = bp.process(self.current, enable) {
                        x
                    } else {
                        // So failed to process a breakpoint.. Still continue to avoid
                        // stalling
                        (false, TracerAction::Continue(self.current.into()))
                    }
                };
                if updated.0 {
                    if let Some(ref mut t) = self.traces.get_trace_mut(rip) {
                        if let CoverageStat::Line(ref mut x) = t.stats {
                            trace!("Incrementing hit count for trace");
                            *x += 1;
                        }
                    }
                }
                action = Some(updated.1);
            }
        }
        let action = action.unwrap_or_else(|| TracerAction::Continue(self.current.into()));
        Ok((TestState::wait_state(), action))
    }

    fn handle_signaled(&mut self, pid: &Pid, sig: &Signal, flag: bool) -> Result<UpdateContext, RunError> {
        match (sig, flag) {
            (Signal::SIGTRAP, true) => {
                Ok((TestState::wait_state(), TracerAction::Continue(pid.into())))
            }
            _ => Err(RunError::StateMachine("Unexpected stop".to_string())),
        }
    }
}
<|endoftext|>


<commit_before>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, thread_rng, Rng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 1;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {} - {}",
            self.filename,
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,
                Speed::Rapid if rating >= 2000 => 1.0,
                Speed::Rapid if rating >= 1800 => 2.0 / 5.0,
                Speed::Rapid => 1.0 / 8.0,
                Speed::Blitz if rating >= 2000 => 1.0,
                Speed::Blitz if rating >= 1800 => 1.0 / 4.0,
                Speed::Blitz => 1.0 / 15.0,
                Speed::Bullet if rating >= 2300 => 1.0,
                Speed::Bullet if rating >= 2200 => 4.0 / 5.0,
                Speed::Bullet if rating >= 2000 => 1.0 / 4.0,
                Speed::Bullet if rating >= 1800 => 1.0 / 7.0,
                _ => 1.0 / 20.0,
            }
        } else {
            if rating >= 1600 {
                1.0
            } else {
                0.5
            } // variant games
        };

        let rnd = thread_rng().sample(OpenClosed01);
        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1500
            && probability >= rnd
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<commit_msg>remove debug<commit_after>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, thread_rng, Rng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {} - {}",
            self.filename,
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,
                Speed::Rapid if rating >= 2000 => 1.0,
                Speed::Rapid if rating >= 1800 => 2.0 / 5.0,
                Speed::Rapid => 1.0 / 8.0,
                Speed::Blitz if rating >= 2000 => 1.0,
                Speed::Blitz if rating >= 1800 => 1.0 / 4.0,
                Speed::Blitz => 1.0 / 15.0,
                Speed::Bullet if rating >= 2300 => 1.0,
                Speed::Bullet if rating >= 2200 => 4.0 / 5.0,
                Speed::Bullet if rating >= 2000 => 1.0 / 4.0,
                Speed::Bullet if rating >= 1800 => 1.0 / 7.0,
                _ => 1.0 / 20.0,
            }
        } else {
            if rating >= 1600 {
                1.0
            } else {
                0.5
            } // variant games
        };

        let rnd = thread_rng().sample(OpenClosed01);
        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1500
            && probability >= rnd
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>use die::Die;
use die::DieType;
use uuid::Uuid;

// Rolls all the arguments into a single struct
pub struct ComposedRoll {
    pub advantage: bool,
    pub comment: Option<String>,
    pub d: i16,
    pub die: DieType,
    pub disadvantage: bool,
    pub e: i16,
    pub gt: u16,
    pub gte: u16,
    pub h: i16,
    pub l: i16,
    pub lt: u16,
    pub lte: u16,
    pub max: i16,
    pub min: i16,
    pub modifiers: Vec<i16>,
    pub n: i16,
    pub ro: i16,
    pub rr: i16,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Roll {
    /// Unique identifier for the roll
    pub _id: String,

    /// Roll comment
    pub comment: Option<String>,

    /// The dice that compose this roll
    pub dice: Vec<Die>,

    /// Modifiers to apply to the combined value
    pub modifiers: Vec<i16>,

    /// The combined value of the die before modifiers
    pub raw_value: i32,

    /// The associated token (optional)
    pub token: Option<String>,

    /// The final combined value of the die after modifiers
    pub value: i32,
}

impl Roll {
    pub fn new(mut dice: Vec<Die>) -> Roll {
        // Roll each dice
        for die in &mut dice {
            die.roll();
        }

        let value = dice.iter().fold(0, |sum, d| sum + d.value as i32);

        Roll {
            _id: Uuid::new_v4().to_string(),
            comment: None,
            dice,
            modifiers: Vec::new(),
            raw_value: value,
            token: None,
            value,
        }
    }

    /// Associate this roll with a comment
    pub fn add_comment(&mut self, comment: String) {
        self.comment = Some(comment)
    }

    /// Associate this roll with a token
    pub fn add_token(&mut self, token: String) {
        self.token = Some(token)
    }

    /// Add a modifier to the roll
    pub fn apply_modifier(&mut self, modifier: i16) {
        self.modifiers.push(modifier);
        self.value += modifier as i32;
    }

    /// Keep the dice greater than a number
    pub fn keep_greater_than(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) <= keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the dice greater than or equal to a number
    pub fn keep_greater_than_or_equal_to(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) < keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the dice less than a number
    pub fn keep_less_than(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) >= keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the dice less than or equal to a number
    pub fn keep_less_than_or_equal_to(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) > keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the highest rolled dice
    pub fn keep_high(&mut self, keep: u16) {
        // Sort the dice by value, drop everything below the keep value
        let mut count = 0;
        self.dice.sort_by(|a, b| b.value.cmp(&a.value));
        for die in &mut self.dice {
            if count >= keep {
                die.drop();
                self.value -= die.value as i32;
            }
            count += 1;
        }
        // sort by timestamp again before finishing the method
        self.dice.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
    }

    /// Keep the lowest rolled dice
    pub fn keep_low(&mut self, keep: u16) {
        // Sort the dice by value, drop everything below the keep value
        let mut count = 0;
        self.dice.sort_by(|a, b| a.value.cmp(&b.value));
        for die in &mut self.dice {
            if count >= keep {
                die.drop();
                self.value -= die.value as i32;
            }
            count += 1;
        }
        // sort by timestamp again before finishing the method
        self.dice.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
    }

    /// Reroll dice one time that are above a certain threshold
    pub fn reroll_dice_once_above(&mut self, threshold: i16) {
        let mut new_dice = Vec::new();
        for die in &mut self.dice {
            if !die.is_rerolled && die.value >= threshold {
                let mut d = Die::new(die.die);
                d.roll();
                let value = d.value;
                self.value += value as i32;
                &die.rerolled(&d);
                new_dice.push(d);
            }
        }

        self.dice.append(&mut new_dice);
    }

    /// Reroll dice one time that are below a certain threshold
    pub fn reroll_dice_once_below(&mut self, threshold: i16) {
        let mut new_dice = Vec::new();
        for die in &mut self.dice {
            if !die.is_rerolled && die.value <= threshold {
                let mut d = Die::new(die.die);
                d.roll();
                let value = d.value;
                self.value += value as i32;
                &die.rerolled(&d);
                new_dice.push(d);

            }
        }

        self.dice.append(&mut new_dice);
    }

    /// Reroll dice forever that are above a certain threshold (e.g. Exploding Dice)
    pub fn reroll_dice_forever_above(&mut self, threshold: i16) {
        // Reroll any dice that need to be rerolled
        self.reroll_dice_once_above(threshold);

        let mut has_more = false;
        for die in self.dice.iter() {
            if !die.is_rerolled && die.value >= threshold {
                has_more = true
            }
        }
        if has_more {
            self.reroll_dice_forever_above(threshold);
        }
    }

    /// Reroll dice forever that are below a certain threshold
    pub fn reroll_dice_forever_below(&mut self, threshold: i16) {
        // Reroll any dice that need to be rerolled
        self.reroll_dice_once_below(threshold);

        let mut has_more = false;
        for die in self.dice.iter() {
            if !die.is_rerolled && die.value <= threshold {
                has_more = true
            }
        }
        if has_more {
            self.reroll_dice_forever_below(threshold);
        }
    }
}

pub fn roll_and_keep_high(count: u16, die: DieType, keep: u16) -> Roll {
    let mut roll = roll_type(count, die);
    roll.keep_high(keep);
    roll
}

pub fn roll_and_keep_low(count: u16, die: DieType, keep: u16) -> Roll {
    let mut roll = roll_type(count, die);
    roll.keep_low(keep);
    roll
}

pub fn roll_d4(count: u16) -> Roll {
    roll_type(count, DieType::D4)
}

pub fn roll_d6(count: u16) -> Roll {
    roll_type(count, DieType::D6)
}

pub fn roll_d8(count: u16) -> Roll {
    roll_type(count, DieType::D8)
}

pub fn roll_d10(count: u16) -> Roll {
    roll_type(count, DieType::D10)
}

pub fn roll_d12(count: u16) -> Roll {
    roll_type(count, DieType::D12)
}

pub fn roll_d20(count: u16) -> Roll {
    roll_type(count, DieType::D20)
}

pub fn roll_d100(count: u16) -> Roll {
    roll_type(count, DieType::D100)
}

fn roll_type(count: u16, die: DieType) -> Roll {
    let mut dice = Vec::new();
    for _ in 0..count {
        dice.push(Die::new(die));
    }
    Roll::new(dice)
}

pub fn roll_with_advantage() -> Roll {
    let mut roll = roll_d20(2);
    roll.keep_high(1);
    roll
}

pub fn roll_with_disadvantage() -> Roll {
    let mut roll = roll_d20(2);
    roll.keep_low(1);
    roll
}
<commit_msg>add timestamp to roll<commit_after>use chrono::DateTime;
use chrono::prelude::Utc;
use die::Die;
use die::DieType;
use uuid::Uuid;

// Rolls all the arguments into a single struct
pub struct ComposedRoll {
    pub advantage: bool,
    pub comment: Option<String>,
    pub d: i16,
    pub die: DieType,
    pub disadvantage: bool,
    pub e: i16,
    pub gt: u16,
    pub gte: u16,
    pub h: i16,
    pub l: i16,
    pub lt: u16,
    pub lte: u16,
    pub max: i16,
    pub min: i16,
    pub modifiers: Vec<i16>,
    pub n: i16,
    pub ro: i16,
    pub rr: i16,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Roll {
    /// Unique identifier for the roll
    pub _id: String,

    /// Roll comment
    pub comment: Option<String>,

    /// The dice that compose this roll
    pub dice: Vec<Die>,

    /// Modifiers to apply to the combined value
    pub modifiers: Vec<i16>,

    /// The combined value of the die before modifiers
    pub raw_value: i32,

    /// Timestamp
    pub timestamp: DateTime<Utc>,

    /// The associated token (optional)
    pub token: Option<String>,

    /// The final combined value of the die after modifiers
    pub value: i32,
}

impl Roll {
    pub fn new(mut dice: Vec<Die>) -> Roll {
        let timestamp = Utc::now();

        // Roll each dice
        for die in &mut dice {
            die.roll();
        }

        let value = dice.iter().fold(0, |sum, d| sum + d.value as i32);

        Roll {
            _id: Uuid::new_v4().to_string(),
            comment: None,
            dice,
            modifiers: Vec::new(),
            raw_value: value,
            timestamp,
            token: None,
            value,
        }
    }

    /// Associate this roll with a comment
    pub fn add_comment(&mut self, comment: String) {
        self.comment = Some(comment)
    }

    /// Associate this roll with a token
    pub fn add_token(&mut self, token: String) {
        self.token = Some(token)
    }

    /// Add a modifier to the roll
    pub fn apply_modifier(&mut self, modifier: i16) {
        self.modifiers.push(modifier);
        self.value += modifier as i32;
    }

    /// Keep the dice greater than a number
    pub fn keep_greater_than(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) <= keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the dice greater than or equal to a number
    pub fn keep_greater_than_or_equal_to(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) < keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the dice less than a number
    pub fn keep_less_than(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) >= keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the dice less than or equal to a number
    pub fn keep_less_than_or_equal_to(&mut self, keep: u16) {
        for die in &mut self.dice {
            if (die.value as u16) > keep {
                die.drop();
                self.value -= die.value as i32;
            } else {
                die.success();
            }
        }
    }

    /// Keep the highest rolled dice
    pub fn keep_high(&mut self, keep: u16) {
        // Sort the dice by value, drop everything below the keep value
        let mut count = 0;
        self.dice.sort_by(|a, b| b.value.cmp(&a.value));
        for die in &mut self.dice {
            if count >= keep {
                die.drop();
                self.value -= die.value as i32;
            }
            count += 1;
        }
        // sort by timestamp again before finishing the method
        self.dice.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
    }

    /// Keep the lowest rolled dice
    pub fn keep_low(&mut self, keep: u16) {
        // Sort the dice by value, drop everything below the keep value
        let mut count = 0;
        self.dice.sort_by(|a, b| a.value.cmp(&b.value));
        for die in &mut self.dice {
            if count >= keep {
                die.drop();
                self.value -= die.value as i32;
            }
            count += 1;
        }
        // sort by timestamp again before finishing the method
        self.dice.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
    }

    /// Reroll dice one time that are above a certain threshold
    pub fn reroll_dice_once_above(&mut self, threshold: i16) {
        let mut new_dice = Vec::new();
        for die in &mut self.dice {
            if !die.is_rerolled && die.value >= threshold {
                let mut d = Die::new(die.die);
                d.roll();
                let value = d.value;
                self.value += value as i32;
                &die.rerolled(&d);
                new_dice.push(d);
            }
        }

        self.dice.append(&mut new_dice);
    }

    /// Reroll dice one time that are below a certain threshold
    pub fn reroll_dice_once_below(&mut self, threshold: i16) {
        let mut new_dice = Vec::new();
        for die in &mut self.dice {
            if !die.is_rerolled && die.value <= threshold {
                let mut d = Die::new(die.die);
                d.roll();
                let value = d.value;
                self.value += value as i32;
                &die.rerolled(&d);
                new_dice.push(d);

            }
        }

        self.dice.append(&mut new_dice);
    }

    /// Reroll dice forever that are above a certain threshold (e.g. Exploding Dice)
    pub fn reroll_dice_forever_above(&mut self, threshold: i16) {
        // Reroll any dice that need to be rerolled
        self.reroll_dice_once_above(threshold);

        let mut has_more = false;
        for die in self.dice.iter() {
            if !die.is_rerolled && die.value >= threshold {
                has_more = true
            }
        }
        if has_more {
            self.reroll_dice_forever_above(threshold);
        }
    }

    /// Reroll dice forever that are below a certain threshold
    pub fn reroll_dice_forever_below(&mut self, threshold: i16) {
        // Reroll any dice that need to be rerolled
        self.reroll_dice_once_below(threshold);

        let mut has_more = false;
        for die in self.dice.iter() {
            if !die.is_rerolled && die.value <= threshold {
                has_more = true
            }
        }
        if has_more {
            self.reroll_dice_forever_below(threshold);
        }
    }
}

pub fn roll_and_keep_high(count: u16, die: DieType, keep: u16) -> Roll {
    let mut roll = roll_type(count, die);
    roll.keep_high(keep);
    roll
}

pub fn roll_and_keep_low(count: u16, die: DieType, keep: u16) -> Roll {
    let mut roll = roll_type(count, die);
    roll.keep_low(keep);
    roll
}

pub fn roll_d4(count: u16) -> Roll {
    roll_type(count, DieType::D4)
}

pub fn roll_d6(count: u16) -> Roll {
    roll_type(count, DieType::D6)
}

pub fn roll_d8(count: u16) -> Roll {
    roll_type(count, DieType::D8)
}

pub fn roll_d10(count: u16) -> Roll {
    roll_type(count, DieType::D10)
}

pub fn roll_d12(count: u16) -> Roll {
    roll_type(count, DieType::D12)
}

pub fn roll_d20(count: u16) -> Roll {
    roll_type(count, DieType::D20)
}

pub fn roll_d100(count: u16) -> Roll {
    roll_type(count, DieType::D100)
}

fn roll_type(count: u16, die: DieType) -> Roll {
    let mut dice = Vec::new();
    for _ in 0..count {
        dice.push(Die::new(die));
    }
    Roll::new(dice)
}

pub fn roll_with_advantage() -> Roll {
    let mut roll = roll_d20(2);
    roll.keep_high(1);
    roll
}

pub fn roll_with_disadvantage() -> Roll {
    let mut roll = roll_d20(2);
    roll.keep_low(1);
    roll
}
<|endoftext|>


<commit_before>// Libtock regression tests to be used with real hardware.
// Requires P0.03 and P0.04 to be connected (on a nRF52 DK).

#![no_std]
extern crate alloc;

use alloc::string::String;
use core::fmt::Write;
use core::future::Future;
use core::mem;
use core::pin::Pin;
use core::task::Context;
use core::task::Poll;
use libtock::console::Console;
use libtock::console::ConsoleDriver;
use libtock::gpio::GpioDriverFactory;
use libtock::gpio::GpioState;
use libtock::gpio::ResistorMode;
use libtock::result::TockResult;
use libtock::timer::DriverContext;
use libtock::timer::Duration;

#[libtock::main]
async fn main() -> TockResult<()> {
    let mut drivers = libtock::retrieve_drivers()?;

    let mut test = LibtockTest::initialize(drivers.console);

    let test_result = libtock_test(&mut test, &mut drivers.timer, &mut drivers.gpio).await;

    if test_result.is_ok() && test.is_success() {
        test.log_success("Test suite finished with state SUCCESS")
    } else {
        test.log_failure("Test suite finished with state FAILURE")
    }
}

#[cfg_attr(
    feature = "__internal_disable_gpio_in_integration_test",
    allow(unused_variables)
)]
async fn libtock_test(
    test: &mut LibtockTest,
    timer: &mut DriverContext,
    gpio: &mut GpioDriverFactory,
) -> TockResult<()> {
    test.console()?;
    test.static_mut()?;
    test.dynamic_dispatch()?;
    test.formatting()?;
    test.heap()?;
    test.drivers_only_instantiable_once()?;
    test.callbacks(timer).await?;
    #[cfg(not(feature = "__internal_disable_gpio_in_integration_test"))]
    test.gpio(gpio)?;
    Ok(())
}

struct LibtockTest {
    console: Console,
    success: bool,
}

impl LibtockTest {
    fn initialize(console: ConsoleDriver) -> Self {
        Self {
            console: console.create_console(),
            success: true,
        }
    }

    fn console(&mut self) -> TockResult<()> {
        self.log_success("Console")
    }

    fn static_mut(&mut self) -> TockResult<()> {
        self.check_if_true(increment_static_mut() == 1, "static mut")
    }

    fn dynamic_dispatch(&mut self) -> TockResult<()> {
        let (x, y) = if foo() == "foo" {
            (&'0' as &dyn MyTrait, &0usize as &dyn MyTrait)
        } else {
            (&0usize as &dyn MyTrait, &'0' as &dyn MyTrait)
        };

        self.check_if_true(
            (x.dispatch(), y.dispatch()) == ("str", "usize"),
            "Dynamic dispatch",
        )
    }

    fn formatting(&mut self) -> TockResult<()> {
        let mut string = String::new();
        write!(string, "{}bar", foo())?;

        self.check_if_true(string == "foobar", "Formatting")
    }

    fn heap(&mut self) -> TockResult<()> {
        let mut string = String::new();
        string.push_str(foo());
        string.push_str("bar");

        self.check_if_true(string == "foobar", "Heap")
    }

    fn drivers_only_instantiable_once(&mut self) -> TockResult<()> {
        self.check_if_true(
            libtock::retrieve_drivers().is_err(),
            "Drivers only instantiable once",
        )
    }

    async fn callbacks(&mut self, timer_context: &mut DriverContext) -> TockResult<()> {
        let mut callback_hit = false;
        let mut with_callback = timer_context.with_callback(|_, _| callback_hit = true);
        let mut timer = with_callback.init()?;

        timer.set_alarm(Duration::from_ms(1000))?;

        AlternatingFuture { yielded: false }.await;

        mem::drop(timer);

        self.check_if_true(callback_hit, "Callbacks")
    }

    #[cfg_attr(
        feature = "__internal_disable_gpio_in_integration_test",
        allow(dead_code)
    )]
    fn gpio(&mut self, gpio: &mut GpioDriverFactory) -> TockResult<()> {
        let mut gpio_driver = gpio.init_driver().ok().unwrap();

        self.log_success("GPIO initialization")?;

        let mut gpios = gpio_driver.gpios();
        let mut pin_in = gpios.next().unwrap();
        let pin_in = pin_in.enable_input(ResistorMode::PullDown).ok().unwrap();
        let mut pin_out = gpios.next().unwrap();
        let pin_out = pin_out.enable_output().ok().unwrap();

        self.log_success("GPIO activation")?;

        pin_out.set_high().ok().unwrap();

        self.check_if_true(
            pin_in.read().ok() == Some(GpioState::High),
            "GPIO read/write",
        )
    }

    fn is_success(&self) -> bool {
        self.success
    }

    fn check_if_true(&mut self, condition: bool, message: &str) -> TockResult<()> {
        if condition {
            self.log_success(message)
        } else {
            self.log_failure(message)
        }
    }

    fn log_success(&mut self, message: &str) -> TockResult<()> {
        writeln!(&mut self.console, "[      OK ] {}", message)?;
        Ok(())
    }

    fn log_failure(&mut self, message: &str) -> TockResult<()> {
        writeln!(&mut self.console, "[ FAILURE ] {}", message)?;
        self.success = false;
        Ok(())
    }
}

#[inline(never)]
// Do not inline this to prevent compiler optimizations
fn foo() -> &'static str {
    "foo"
}

#[inline(never)]
fn increment_static_mut() -> usize {
    static mut STATIC: usize = 0;

    unsafe {
        STATIC += 1;
        STATIC
    }
}

trait MyTrait {
    fn dispatch(&self) -> &'static str;
}

impl MyTrait for usize {
    fn dispatch(&self) -> &'static str {
        "usize"
    }
}

impl MyTrait for char {
    fn dispatch(&self) -> &'static str {
        "str"
    }
}

struct AlternatingFuture {
    yielded: bool,
}

impl Future for AlternatingFuture {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
        self.yielded = !self.yielded;
        if self.yielded {
            Poll::Pending
        } else {
            Poll::Ready(())
        }
    }
}
<commit_msg>Add note about testing nRF52840dk<commit_after>// Libtock regression tests to be used with real hardware.
// Requires P0.03 and P0.04 to be connected (on a nRF52 DK) and
// P0.01 and P0.03 to be connected (on a nRF52840dk).

#![no_std]
extern crate alloc;

use alloc::string::String;
use core::fmt::Write;
use core::future::Future;
use core::mem;
use core::pin::Pin;
use core::task::Context;
use core::task::Poll;
use libtock::console::Console;
use libtock::console::ConsoleDriver;
use libtock::gpio::GpioDriverFactory;
use libtock::gpio::GpioState;
use libtock::gpio::ResistorMode;
use libtock::result::TockResult;
use libtock::timer::DriverContext;
use libtock::timer::Duration;

#[libtock::main]
async fn main() -> TockResult<()> {
    let mut drivers = libtock::retrieve_drivers()?;

    let mut test = LibtockTest::initialize(drivers.console);

    let test_result = libtock_test(&mut test, &mut drivers.timer, &mut drivers.gpio).await;

    if test_result.is_ok() && test.is_success() {
        test.log_success("Test suite finished with state SUCCESS")
    } else {
        test.log_failure("Test suite finished with state FAILURE")
    }
}

#[cfg_attr(
    feature = "__internal_disable_gpio_in_integration_test",
    allow(unused_variables)
)]
async fn libtock_test(
    test: &mut LibtockTest,
    timer: &mut DriverContext,
    gpio: &mut GpioDriverFactory,
) -> TockResult<()> {
    test.console()?;
    test.static_mut()?;
    test.dynamic_dispatch()?;
    test.formatting()?;
    test.heap()?;
    test.drivers_only_instantiable_once()?;
    test.callbacks(timer).await?;
    #[cfg(not(feature = "__internal_disable_gpio_in_integration_test"))]
    test.gpio(gpio)?;
    Ok(())
}

struct LibtockTest {
    console: Console,
    success: bool,
}

impl LibtockTest {
    fn initialize(console: ConsoleDriver) -> Self {
        Self {
            console: console.create_console(),
            success: true,
        }
    }

    fn console(&mut self) -> TockResult<()> {
        self.log_success("Console")
    }

    fn static_mut(&mut self) -> TockResult<()> {
        self.check_if_true(increment_static_mut() == 1, "static mut")
    }

    fn dynamic_dispatch(&mut self) -> TockResult<()> {
        let (x, y) = if foo() == "foo" {
            (&'0' as &dyn MyTrait, &0usize as &dyn MyTrait)
        } else {
            (&0usize as &dyn MyTrait, &'0' as &dyn MyTrait)
        };

        self.check_if_true(
            (x.dispatch(), y.dispatch()) == ("str", "usize"),
            "Dynamic dispatch",
        )
    }

    fn formatting(&mut self) -> TockResult<()> {
        let mut string = String::new();
        write!(string, "{}bar", foo())?;

        self.check_if_true(string == "foobar", "Formatting")
    }

    fn heap(&mut self) -> TockResult<()> {
        let mut string = String::new();
        string.push_str(foo());
        string.push_str("bar");

        self.check_if_true(string == "foobar", "Heap")
    }

    fn drivers_only_instantiable_once(&mut self) -> TockResult<()> {
        self.check_if_true(
            libtock::retrieve_drivers().is_err(),
            "Drivers only instantiable once",
        )
    }

    async fn callbacks(&mut self, timer_context: &mut DriverContext) -> TockResult<()> {
        let mut callback_hit = false;
        let mut with_callback = timer_context.with_callback(|_, _| callback_hit = true);
        let mut timer = with_callback.init()?;

        timer.set_alarm(Duration::from_ms(1000))?;

        AlternatingFuture { yielded: false }.await;

        mem::drop(timer);

        self.check_if_true(callback_hit, "Callbacks")
    }

    #[cfg_attr(
        feature = "__internal_disable_gpio_in_integration_test",
        allow(dead_code)
    )]
    fn gpio(&mut self, gpio: &mut GpioDriverFactory) -> TockResult<()> {
        let mut gpio_driver = gpio.init_driver().ok().unwrap();

        self.log_success("GPIO initialization")?;

        let mut gpios = gpio_driver.gpios();
        let mut pin_in = gpios.next().unwrap();
        let pin_in = pin_in.enable_input(ResistorMode::PullDown).ok().unwrap();
        let mut pin_out = gpios.next().unwrap();
        let pin_out = pin_out.enable_output().ok().unwrap();

        self.log_success("GPIO activation")?;

        pin_out.set_high().ok().unwrap();

        self.check_if_true(
            pin_in.read().ok() == Some(GpioState::High),
            "GPIO read/write",
        )
    }

    fn is_success(&self) -> bool {
        self.success
    }

    fn check_if_true(&mut self, condition: bool, message: &str) -> TockResult<()> {
        if condition {
            self.log_success(message)
        } else {
            self.log_failure(message)
        }
    }

    fn log_success(&mut self, message: &str) -> TockResult<()> {
        writeln!(&mut self.console, "[      OK ] {}", message)?;
        Ok(())
    }

    fn log_failure(&mut self, message: &str) -> TockResult<()> {
        writeln!(&mut self.console, "[ FAILURE ] {}", message)?;
        self.success = false;
        Ok(())
    }
}

#[inline(never)]
// Do not inline this to prevent compiler optimizations
fn foo() -> &'static str {
    "foo"
}

#[inline(never)]
fn increment_static_mut() -> usize {
    static mut STATIC: usize = 0;

    unsafe {
        STATIC += 1;
        STATIC
    }
}

trait MyTrait {
    fn dispatch(&self) -> &'static str;
}

impl MyTrait for usize {
    fn dispatch(&self) -> &'static str {
        "usize"
    }
}

impl MyTrait for char {
    fn dispatch(&self) -> &'static str {
        "str"
    }
}

struct AlternatingFuture {
    yielded: bool,
}

impl Future for AlternatingFuture {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
        self.yielded = !self.yielded;
        if self.yielded {
            Poll::Pending
        } else {
            Poll::Ready(())
        }
    }
}
<|endoftext|>


<commit_before>// TODO: Implement `libtock_runtime` and `libtock_unittest`, which are
// referenced in the comment on `RawSyscalls`.

/// `RawSyscalls` allows a fake Tock kernel to be injected into components for
/// unit testing. It is implemented by `libtock_runtime::TockSyscalls` and
/// `libtock_unittest::FakeSyscalls`. Components should not use `RawSyscalls`
/// directly; instead, use the `Syscalls` trait, which provides higher-level
/// interfaces to the system calls.

// RawSyscalls is designed to minimize the amount of handwritten assembly code
// needed without generating unnecessary instructions. There are a few major
// factors affecting its design:
//     1. Most system calls only clobber r0-r4 (*), while yield has a far longer
//        clobber list. As such, yield must have its own assembly
//        implementation.
//     2. The compiler is unable to optimize away unused arguments. For example,
//        memop's "get process RAM start address" operation only needs r0 set,
//        while memop's "break" operation needs both r0 and r1 set. If our
//        inline assembly calls "get process RAM start address" but sets both r0
//        and r1, the compiler doesn't know that r1 will be
//        ignored so setting that register will not be optimized away. Therefore
//        we want to set the minimum number of argument registers possible.
//     3. The cost of specifying unused return registers is only that of
//        unnecessarily marking a register as clobbered. Explanation: After
//        inlining, an unused register is marked as "changed by the assembly"
//        but can immediately be re-used by the compiler, which is the same as a
//        clobbered register. System calls should generally be
//        inlined -- and even if they aren't, the unused return values will
//        probably be passed in caller-saved registers (this is true for the C
//        ABI, so probably true for the Rust ABI), which are treated as
//        clobbered regardless.
//
// (*) When this file refers to registers, it uses the same naming convention as
// the Tock 2.0 syscalls TRD. Registers r0-r4 correspond to ARM registers r0-r4
// and RISC-V registers a0-a4.
//
// Currently, yield takes exactly one argument, to specify what yield type to
// do. Therefore we only need one raw yield call.
//
// Subscribe, command, read-write allow, and read-only allow all take four
// argument types. Even when calling command IDs that have unused arguments, we
// still need to clear the argument registers so as to avoid passing
// confidential data to capsules (this is in line with Tock's threat model). As
// such, four_arg_syscall() is used for all subscribe, command, read-only allow,
// and read-write allow system calls.
//
// Memop takes 1 or 2 arguments (operation and an optional argument). Because it
// is part of the core kernel, it is okay for us to leave arbitrary data in the
// argument register for operations where the argument register is unused
// (again, in line with Tock's threat model). As such, for efficiency, we need
// two raw memop calls: one for operations without an argument and one for
// operations with an argument.
//
// The success type for Memop calls depends on the operation perform. However,
// all *currently defined* memop operations return either Success or Success
// with u32. Therefore, the memop implementations only need to mark r0 and r1 as
// clobbered, not r2 and r3. This choice of clobbers will need to be revisited
// if and when a memop operation that returns more data is added.
//
// The decision of where to use u32 and usize can be a bit tricky. The Tock
// syscall ABI is currently only specified for 32-bit systems, so on real Tock
// systems both types match the size of a register, but the unit test
// environment can be either 32 bit or 64 bit. This interface uses usize for
// values that can contain pointers, so that pointers are not truncated in the
// unit test environment. To keep types as consistent as possible, it uses u32
// for all values that cannot be pointers.
pub trait RawSyscalls {
    // raw_yield should:
    //     1. Call syscall class 0
    //     2. Use register r0 for input and output as an inlateout register,
    //        passing in r0_in and returning its value.
    //     3. Mark all caller-saved registers as lateout clobbers.
    //     4. NOT provide any of the following options:
    //            pure             (yield has side effects)
    //            nomem            (a callback can read + write globals)
    //            readonly         (a callback can write globals)
    //            preserves_flags  (a callback can change flags)
    //            noreturn         (yield is expected to return)
    //            nostack          (a callback needs the stack)
    //
    // Design note: This is safe because the yield types that currently exist
    // are safe. If an unsafe yield type is added, we will need to make
    // raw_yield unsafe. Although raw_yield shouldn't be called by code outside
    // this crate, it can be, so that is a backwards-incompatible change. We
    // pass YieldType rather than a usize because if we used usize directly then
    // this API becomes unsound if the kernel adds support for an unsafe yield
    // type (or even one that takes one more argument).
    fn raw_yield(r0_in: YieldType) -> u32;

    // four_arg_syscall is used to invoke the subscribe, command, read-write
    // allow, and read-only allow system calls.
    //
    // four_arg_syscall's inline assembly should have the following properties:
    //     1. Calls the syscall class specified by class
    //     2. Passes r0-r3 in the corresponding registers as inlateout
    //        registers. Returns r0-r3 in order.
    //     3. Does not mark any registers as clobbered.
    //     4. Has all of the following options:
    //            preserves_flags  (these system calls do not touch flags)
    //            nostack          (these system calls do not touch the stack)
    //     5. Does NOT have any of the following options:
    //            pure      (these system calls have side effects)
    //            nomem     (the compiler needs to write to globals before allow)
    //            readonly  (rw allow can modify memory)
    //            noreturn  (all these system calls are expected to return)
    //
    // Note that subscribe's application data argument can potentially contain a
    // pointer, so r3 can contain a pointer (in addition to r1 and r2, which
    // more obviously contain pointers for subscribe and memop).
    //
    // For subscribe(), the callback pointer should be either 0 (for the null
    // callback) or an `unsafe extern fn(u32, u32, u32, usize)`.
    /// # Safety
    /// `four_arg_syscall` must NOT be used to invoke yield. Otherwise, it has
    /// the same safety invariants as the underlying system call, which varies
    /// depending on the system call class.
    unsafe fn four_arg_syscall(
        r0: u32,
        r1: usize,
        r2: usize,
        r3: usize,
        class: u8,
    ) -> (u32, usize, usize, usize);

    // zero_arg_memop is used to invoke memop operations that do not accept an
    // argument register. Because there are no memop commands that set r2 or r3,
    // this only needs to return r0 and r1.
    //
    // Memop commands may panic in the unit test environment, as not all memop
    // calls can be sensibly implemented in that environment.
    //
    // zero_arg_memop's inline assembly should have the following properties:
    //     1. Calls syscall class 5
    //     2. Specifies r0 as an inlateout register, and r1 as a lateout
    //        register.
    //     3. Does not mark any registers as clobbered.
    //     4. Has all of the following options:
    //            preserves_flags
    //            nostack
    //            nomem            (it is okay for the compiler to cache globals
    //                              across memop calls)
    //     5. Does NOT have any of the following options:
    //            pure      (two invocations of the same memop can return
    //                       different values)
    //            readonly  (incompatible with nomem)
    //            noreturn
    //
    // Design note: like raw_yield, this is safe because memops that currently
    // exist are safe. zero_arg_memop takes a ZeroArgMemop rather than a u32 so
    // that if the kernel adds an unsafe memop -- or one that can clobber r2/r3
    // --  this API doesn't become unsound.
    fn zero_arg_memop(r0_in: ZeroArgMemop) -> (u32, usize);

    // one_arg_memop is used to invoke memop operations that take an argument.
    // Because there are no memop operations that set r2 or r3, this only needs
    // to return r0 and r1.
    //
    // one_arg_memop's inline assembly should:
    //     1. Call syscall class 5
    //     2. Specify r0 and r1 as inlateout registers, and return (r0, r1)
    //     3. Not mark any registers as clobbered.
    //     4. Have all of the following options:
    //            preserves_flags
    //            nostack
    //            nomem            (the compiler can cache globals across memop
    //                              calls)
    //     5. Does NOT have any of the following options:
    //            pure      Two invocations of sbrk can return different values
    //            readonly  Incompatible with nomem
    //            noreturn
    //
    // Design note: like raw_yield, this is safe because memops that currently
    // exist are safe. zero_arg_memop takes a ZeroArgMemop rather than a u32 so
    // that if the kernel adds an unsafe memop -- or one that can clobber r2/r3
    // -- this API doesn't become unsound.
    fn one_arg_memop(r0_in: OneArgMemop, r1: usize) -> (u32, usize);
}

#[non_exhaustive]
#[repr(u32)]
pub enum OneArgMemop {
    Brk = 0,
    Sbrk = 1,
    FlashRegionStart = 8,
    FlashRegionEnd = 9,
    SpecifyStackTop = 10,
    SpecifyHeapStart = 11,
    // Note: before adding new memop operations, make sure the assumptions in
    // the design notes on `one_arg_memop` are valid for the new operation type.
}

// TODO: When the numeric values (0 and 1) are assigned to the yield types,
// specify those values here.
#[non_exhaustive]
#[repr(u32)]
pub enum YieldType {
    Wait,
    NoWait,
}

#[non_exhaustive]
#[repr(u32)]
pub enum ZeroArgMemop {
    MemoryStart = 2,
    MemoryEnd = 3,
    FlashStart = 4,
    FlashEnd = 5,
    GrantStart = 6,
    FlashRegions = 7,
    // Note: before adding new memop operations, make sure the assumptions in
    // the design notes on `zero_arg_memop` are valid for the new operation
    // type.
}
<commit_msg>Change r1 in four_arg_syscall to a u32.<commit_after>// TODO: Implement `libtock_runtime` and `libtock_unittest`, which are
// referenced in the comment on `RawSyscalls`.

/// `RawSyscalls` allows a fake Tock kernel to be injected into components for
/// unit testing. It is implemented by `libtock_runtime::TockSyscalls` and
/// `libtock_unittest::FakeSyscalls`. Components should not use `RawSyscalls`
/// directly; instead, use the `Syscalls` trait, which provides higher-level
/// interfaces to the system calls.

// RawSyscalls is designed to minimize the amount of handwritten assembly code
// needed without generating unnecessary instructions. There are a few major
// factors affecting its design:
//     1. Most system calls only clobber r0-r4 (*), while yield has a far longer
//        clobber list. As such, yield must have its own assembly
//        implementation.
//     2. The compiler is unable to optimize away unused arguments. For example,
//        memop's "get process RAM start address" operation only needs r0 set,
//        while memop's "break" operation needs both r0 and r1 set. If our
//        inline assembly calls "get process RAM start address" but sets both r0
//        and r1, the compiler doesn't know that r1 will be
//        ignored so setting that register will not be optimized away. Therefore
//        we want to set the minimum number of argument registers possible.
//     3. The cost of specifying unused return registers is only that of
//        unnecessarily marking a register as clobbered. Explanation: After
//        inlining, an unused register is marked as "changed by the assembly"
//        but can immediately be re-used by the compiler, which is the same as a
//        clobbered register. System calls should generally be
//        inlined -- and even if they aren't, the unused return values will
//        probably be passed in caller-saved registers (this is true for the C
//        ABI, so probably true for the Rust ABI), which are treated as
//        clobbered regardless.
//
// (*) When this file refers to registers, it uses the same naming convention as
// the Tock 2.0 syscalls TRD. Registers r0-r4 correspond to ARM registers r0-r4
// and RISC-V registers a0-a4.
//
// Currently, yield takes exactly one argument, to specify what yield type to
// do. Therefore we only need one raw yield call.
//
// Subscribe, command, read-write allow, and read-only allow all take four
// argument types. Even when calling command IDs that have unused arguments, we
// still need to clear the argument registers so as to avoid passing
// confidential data to capsules (this is in line with Tock's threat model). As
// such, four_arg_syscall() is used for all subscribe, command, read-only allow,
// and read-write allow system calls.
//
// Memop takes 1 or 2 arguments (operation and an optional argument). Because it
// is part of the core kernel, it is okay for us to leave arbitrary data in the
// argument register for operations where the argument register is unused
// (again, in line with Tock's threat model). As such, for efficiency, we need
// two raw memop calls: one for operations without an argument and one for
// operations with an argument.
//
// The success type for Memop calls depends on the operation perform. However,
// all *currently defined* memop operations return either Success or Success
// with u32. Therefore, the memop implementations only need to mark r0 and r1 as
// clobbered, not r2 and r3. This choice of clobbers will need to be revisited
// if and when a memop operation that returns more data is added.
//
// The decision of where to use u32 and usize can be a bit tricky. The Tock
// syscall ABI is currently only specified for 32-bit systems, so on real Tock
// systems both types match the size of a register, but the unit test
// environment can be either 32 bit or 64 bit. This interface uses usize for
// values that can contain pointers, so that pointers are not truncated in the
// unit test environment. To keep types as consistent as possible, it uses u32
// for all values that cannot be pointers.
pub trait RawSyscalls {
    // raw_yield should:
    //     1. Call syscall class 0
    //     2. Use register r0 for input and output as an inlateout register,
    //        passing in r0_in and returning its value.
    //     3. Mark all caller-saved registers as lateout clobbers.
    //     4. NOT provide any of the following options:
    //            pure             (yield has side effects)
    //            nomem            (a callback can read + write globals)
    //            readonly         (a callback can write globals)
    //            preserves_flags  (a callback can change flags)
    //            noreturn         (yield is expected to return)
    //            nostack          (a callback needs the stack)
    //
    // Design note: This is safe because the yield types that currently exist
    // are safe. If an unsafe yield type is added, we will need to make
    // raw_yield unsafe. Although raw_yield shouldn't be called by code outside
    // this crate, it can be, so that is a backwards-incompatible change. We
    // pass YieldType rather than a usize because if we used usize directly then
    // this API becomes unsound if the kernel adds support for an unsafe yield
    // type (or even one that takes one more argument).
    fn raw_yield(r0_in: YieldType) -> u32;

    // four_arg_syscall is used to invoke the subscribe, command, read-write
    // allow, and read-only allow system calls.
    //
    // four_arg_syscall's inline assembly should have the following properties:
    //     1. Calls the syscall class specified by class
    //     2. Passes r0-r3 in the corresponding registers as inlateout
    //        registers. Returns r0-r3 in order.
    //     3. Does not mark any registers as clobbered.
    //     4. Has all of the following options:
    //            preserves_flags  (these system calls do not touch flags)
    //            nostack          (these system calls do not touch the stack)
    //     5. Does NOT have any of the following options:
    //            pure      (these system calls have side effects)
    //            nomem     (the compiler needs to write to globals before allow)
    //            readonly  (rw allow can modify memory)
    //            noreturn  (all these system calls are expected to return)
    //
    // Note that subscribe's application data argument can potentially contain a
    // pointer, so r3 can contain a pointer (in addition to r1 and r2, which
    // more obviously contain pointers for subscribe and memop).
    //
    // For subscribe(), the callback pointer should be either 0 (for the null
    // callback) or an `unsafe extern fn(u32, u32, u32, usize)`.
    /// # Safety
    /// `four_arg_syscall` must NOT be used to invoke yield. Otherwise, it has
    /// the same safety invariants as the underlying system call, which varies
    /// depending on the system call class.
    unsafe fn four_arg_syscall(
        r0: u32,
        r1: u32,
        r2: usize,
        r3: usize,
        class: u8,
    ) -> (u32, usize, usize, usize);

    // zero_arg_memop is used to invoke memop operations that do not accept an
    // argument register. Because there are no memop commands that set r2 or r3,
    // this only needs to return r0 and r1.
    //
    // Memop commands may panic in the unit test environment, as not all memop
    // calls can be sensibly implemented in that environment.
    //
    // zero_arg_memop's inline assembly should have the following properties:
    //     1. Calls syscall class 5
    //     2. Specifies r0 as an inlateout register, and r1 as a lateout
    //        register.
    //     3. Does not mark any registers as clobbered.
    //     4. Has all of the following options:
    //            preserves_flags
    //            nostack
    //            nomem            (it is okay for the compiler to cache globals
    //                              across memop calls)
    //     5. Does NOT have any of the following options:
    //            pure      (two invocations of the same memop can return
    //                       different values)
    //            readonly  (incompatible with nomem)
    //            noreturn
    //
    // Design note: like raw_yield, this is safe because memops that currently
    // exist are safe. zero_arg_memop takes a ZeroArgMemop rather than a u32 so
    // that if the kernel adds an unsafe memop -- or one that can clobber r2/r3
    // --  this API doesn't become unsound.
    fn zero_arg_memop(r0_in: ZeroArgMemop) -> (u32, usize);

    // one_arg_memop is used to invoke memop operations that take an argument.
    // Because there are no memop operations that set r2 or r3, this only needs
    // to return r0 and r1.
    //
    // one_arg_memop's inline assembly should:
    //     1. Call syscall class 5
    //     2. Specify r0 and r1 as inlateout registers, and return (r0, r1)
    //     3. Not mark any registers as clobbered.
    //     4. Have all of the following options:
    //            preserves_flags
    //            nostack
    //            nomem            (the compiler can cache globals across memop
    //                              calls)
    //     5. Does NOT have any of the following options:
    //            pure      Two invocations of sbrk can return different values
    //            readonly  Incompatible with nomem
    //            noreturn
    //
    // Design note: like raw_yield, this is safe because memops that currently
    // exist are safe. zero_arg_memop takes a ZeroArgMemop rather than a u32 so
    // that if the kernel adds an unsafe memop -- or one that can clobber r2/r3
    // -- this API doesn't become unsound.
    fn one_arg_memop(r0_in: OneArgMemop, r1: usize) -> (u32, usize);
}

#[non_exhaustive]
#[repr(u32)]
pub enum OneArgMemop {
    Brk = 0,
    Sbrk = 1,
    FlashRegionStart = 8,
    FlashRegionEnd = 9,
    SpecifyStackTop = 10,
    SpecifyHeapStart = 11,
    // Note: before adding new memop operations, make sure the assumptions in
    // the design notes on `one_arg_memop` are valid for the new operation type.
}

// TODO: When the numeric values (0 and 1) are assigned to the yield types,
// specify those values here.
#[non_exhaustive]
#[repr(u32)]
pub enum YieldType {
    Wait,
    NoWait,
}

#[non_exhaustive]
#[repr(u32)]
pub enum ZeroArgMemop {
    MemoryStart = 2,
    MemoryEnd = 3,
    FlashStart = 4,
    FlashEnd = 5,
    GrantStart = 6,
    FlashRegions = 7,
    // Note: before adding new memop operations, make sure the assumptions in
    // the design notes on `zero_arg_memop` are valid for the new operation
    // type.
}
<|endoftext|>