<commit_before>// This is an implementation of simple cooperative processes
//
// TODO: kill
// TODO: check killed

use alloc::boxed::Box;

use core::option::Option::{self, Some, None};
use core::fmt::{Debug, Formatter, Result};
use core::cmp::PartialEq;

use core::sync::atomic::{AtomicUsize, Ordering};

use core::ops::Drop;

use super::data_structures::ProcessQueue;

use super::interrupts::{on, off};
use super::interrupts::pit::JIFFIES;

use super::machine::{self, context_switch};

use super::memory::AddressSpace;

use self::context::{KContext};

use self::idle::IDLE_PROCESS;

pub use self::current::{CURRENT_PROCESS};

pub mod context;
pub mod current;

pub mod ready_queue;

mod init;
mod idle;
mod reaper;
mod user;

// constants
const STACK_SIZE: usize = 2048;

// next pid
static NEXT_ID: AtomicUsize = AtomicUsize::new(0);

// Process state
#[repr(C)]
#[derive(Clone, Copy)]
enum State {
    INIT,
    READY,
    RUNNING,
    BLOCKED,
    TERMINATED,
}

// Process struct
pub struct Process {
    // Name of the process; some string
    name: &'static str,
    pid: usize,

    // This is the code of the process
    run: fn(&Process) -> usize,

    state: State,

    // The kheap-allocated stack
    // - this is a pointer to the stack, but *mut usize is not Send/Sync for sone reason
    // - this pointer is to the bottom of the stack, not the head
    stack: usize,

    // The saved kernel context for context switching
    kcontext: KContext,

    // Address space
    pub addr_space: AddressSpace,

    // Number of calls to interrupts::on() while this process was running
    // Interrupts are on if disable_cnt == 0
    pub disable_cnt: usize,

    // A link to the next process to be scheduled
    pub next_proc: *mut Process,
}

impl Process {
    pub fn new(name: &'static str, run: fn(&Process) -> usize) -> *mut Process {

        let mut p = Process {
            name: name,
            pid: NEXT_ID.fetch_add(1, Ordering::Relaxed),
            run: run,
            state: State::INIT,
            stack: 0,
            kcontext: KContext::new(),
            addr_space: AddressSpace::new(),
            disable_cnt: 0,
            next_proc: 0 as *mut Process,
        };

        p.get_stack();

        Box::into_raw(box p)
    }

    fn get_stack(&mut self) {
        // TODO: fudge
        // Allocate a stack
        let stack: Box<[usize; STACK_SIZE]> = box [0; STACK_SIZE];

        let stack_ptr = Box::into_raw(stack) as *mut usize;

        // set the pointer
        self.stack = stack_ptr as usize;

        // smash the stack
        unsafe {
            // put RA on stack to return to start_proc
            *stack_ptr.offset(STACK_SIZE as isize - 2) = start_proc as usize;
        }

        self.kcontext.esp = unsafe { stack_ptr.offset(STACK_SIZE as isize - 2) } as usize;

        // printf!("stack for {:?} is at 0x{:x}\n", self, self.stack);
    }

    fn set_state(&mut self, s: State) {
        self.state = s;
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        unsafe {
            Box::from_raw(self.stack as *mut [usize; 2048]);
        }
        // let the box go out of scope to dealloc
    }
}

impl Debug for Process {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "Process #{} @ 0x{:x}: {}",
               self.pid, self as *const Process as usize, self.name)
    }
}

impl PartialEq for Process {
    fn eq(&self, other: &Process) -> bool {
        self.pid == other.pid
    }
}

pub fn init() {
    // Init ready q and current proc
    current::init(); // THIS MUST BE FIRST
    ready_queue::init();

    // Add the init process to the ready q
    ready_queue::make_ready(Process::new("init", self::init::run));

    // Create the idle process
    idle::init();

    printf!("processes inited\n");
}

// The entry point of all processes
#[allow(private_no_mangle_fns)]
#[no_mangle]
fn start_proc() {
    // TODO check killed

    // run current process
    unsafe {
        let code = if !CURRENT_PROCESS.is_null() {
            printf!("{:?} [Starting]\n", *CURRENT_PROCESS);
            (*CURRENT_PROCESS).set_state(State::RUNNING);
            ((*CURRENT_PROCESS).run)(&*CURRENT_PROCESS)
        } else {
            panic!("No current process to start")
        };

        exit(code);
    }
}

// A safe wrapper around machine::proc_yield
pub fn proc_yield<'a>(q: Option<&'a mut ProcessQueue>) {
    unsafe {
        off();
        machine::proc_yield(q);
        on();
    }
}

// Yield to the next process waiting
//
// This function is called by the current process to
// yield to the next process on the ready q.
//
// To yield normally, call process::proc_yield()
//
// Interrupts should already be disabled here
#[no_mangle]
#[inline(never)]
pub unsafe fn _proc_yield<'a>(q: Option<&'a mut ProcessQueue>) {
    // yield onto the right queue
    if !CURRENT_PROCESS.is_null() {
        if let Some(queue) = q {
            (*CURRENT_PROCESS).set_state(State::BLOCKED);
            queue.push_tail(CURRENT_PROCESS);
        } else {
            ready_queue::make_ready(CURRENT_PROCESS);
        }

        CURRENT_PROCESS = 0 as *mut Process;
    }

    // get next process from ready q
    // spawn a reaper when there are a few processes that have died
    let mut next = if self::reaper::DEAD_COUNT >= 5 &&
        JIFFIES % 75 == 0 && 1 == 3 {
        Process::new("reaper", self::reaper::run)
    } else {
        ready_queue::get_next()
    };

    // run the idle process when everyone is done
    if next.is_null() {
        next = IDLE_PROCESS;
    }

    // switch address spaces
    (*next).addr_space.activate();

    // set the CURRENT_PROCESS
    CURRENT_PROCESS = next;

    //bootlog!("{:?} [Switching]\n", *CURRENT_PROCESS);

    // context switch
    context_switch((*CURRENT_PROCESS).kcontext,
                   if (*CURRENT_PROCESS).disable_cnt == 0 {
                       1 << 9
                   } else {
                       0
                   });

    panic!("The impossible has happened!");
}

// Called by the current process to exit
pub fn exit(code: usize) {
    unsafe {
        if CURRENT_PROCESS.is_null() {
            panic!("Exiting with no current process!\n");
        }

        (*CURRENT_PROCESS).set_state(State::TERMINATED);

        // clean up address space
        (*CURRENT_PROCESS).addr_space.clear();

        // NOTE: need to print *before* adding to reaper q
        printf!("{:?} [Exit 0x{:X}]\n",
                *CURRENT_PROCESS, code);

        // Disable interrupts
        off();

        self::reaper::reaper_add(CURRENT_PROCESS);

        // set current to None, so we will never run this again
        CURRENT_PROCESS = 0 as *mut Process;

        // switch to next ready process
        _proc_yield(None);
    }

    panic!("The impossible has happened!");
}
<commit_msg>comment<commit_after>// This is an implementation of simple cooperative processes
//
// TODO: kill
// TODO: check killed

use alloc::boxed::Box;

use core::option::Option::{self, Some, None};
use core::fmt::{Debug, Formatter, Result};
use core::cmp::PartialEq;

use core::sync::atomic::{AtomicUsize, Ordering};

use core::ops::Drop;

use super::data_structures::ProcessQueue;

use super::interrupts::{on, off};
use super::interrupts::pit::JIFFIES;

use super::machine::{self, context_switch};

use super::memory::AddressSpace;

use self::context::{KContext};

use self::idle::IDLE_PROCESS;

pub use self::current::{CURRENT_PROCESS};

pub mod context;
pub mod current;

pub mod ready_queue;

mod init;
mod idle;
mod reaper;
mod user;

// constants
const STACK_SIZE: usize = 2048;

// next pid
static NEXT_ID: AtomicUsize = AtomicUsize::new(0);

// Process state
#[repr(C)]
#[derive(Clone, Copy)]
enum State {
    INIT,
    READY,
    RUNNING,
    BLOCKED,
    TERMINATED,
}

// Process struct
pub struct Process {
    // Name of the process; some string
    name: &'static str,
    pid: usize,

    // This is the code of the process
    run: fn(&Process) -> usize,

    state: State,

    // The kheap-allocated stack
    // - this is a pointer to the stack, but *mut usize is not Send/Sync for sone reason
    // - this pointer is to the bottom of the stack, not the head
    stack: usize,

    // The saved kernel context for context switching
    kcontext: KContext,

    // Address space
    pub addr_space: AddressSpace,

    // Number of calls to interrupts::on() while this process was running
    // Interrupts are on if disable_cnt == 0
    pub disable_cnt: usize,

    // A link to the next process to be scheduled
    pub next_proc: *mut Process,
}

impl Process {
    pub fn new(name: &'static str, run: fn(&Process) -> usize) -> *mut Process {

        let mut p = Process {
            name: name,
            pid: NEXT_ID.fetch_add(1, Ordering::Relaxed),
            run: run,
            state: State::INIT,
            stack: 0,
            kcontext: KContext::new(),
            addr_space: AddressSpace::new(),
            disable_cnt: 0,
            next_proc: 0 as *mut Process,
        };

        p.get_stack();

        Box::into_raw(box p)
    }

    fn get_stack(&mut self) {
        // TODO: fudge
        // Allocate a stack
        let stack: Box<[usize; STACK_SIZE]> = box [0; STACK_SIZE];

        let stack_ptr = Box::into_raw(stack) as *mut usize;

        // set the pointer
        self.stack = stack_ptr as usize;

        // smash the stack
        unsafe {
            // put RA on stack to return to start_proc
            *stack_ptr.offset(STACK_SIZE as isize - 2) = start_proc as usize;
        }

        self.kcontext.esp = unsafe { stack_ptr.offset(STACK_SIZE as isize - 2) } as usize;

        // printf!("stack for {:?} is at 0x{:x}\n", self, self.stack);
    }

    fn set_state(&mut self, s: State) {
        self.state = s;
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        unsafe {
            Box::from_raw(self.stack as *mut [usize; 2048]);
        }
        // let the box go out of scope to dealloc
    }
}

impl Debug for Process {
    fn fmt(&self, f: &mut Formatter) -> Result {
        write!(f, "Process #{} @ 0x{:x}: {}",
               self.pid, self as *const Process as usize, self.name)
    }
}

impl PartialEq for Process {
    fn eq(&self, other: &Process) -> bool {
        self.pid == other.pid
    }
}

pub fn init() {
    // Init ready q and current proc
    current::init(); // THIS MUST BE FIRST
    ready_queue::init();

    // Add the init process to the ready q
    ready_queue::make_ready(Process::new("init", self::init::run));

    // Create the idle process
    idle::init();

    printf!("processes inited\n");
}

// The entry point of all processes
#[allow(private_no_mangle_fns)]
#[no_mangle]
fn start_proc() {
    // TODO check killed

    // run current process
    unsafe {
        let code = if !CURRENT_PROCESS.is_null() {
            printf!("{:?} [Starting]\n", *CURRENT_PROCESS);
            (*CURRENT_PROCESS).set_state(State::RUNNING);
            ((*CURRENT_PROCESS).run)(&*CURRENT_PROCESS)
        } else {
            panic!("No current process to start")
        };

        exit(code);
    }
}

// A safe wrapper around machine::proc_yield
pub fn proc_yield<'a>(q: Option<&'a mut ProcessQueue>) {
    unsafe {
        off();
        machine::proc_yield(q);
        on();
    }
}

// Yield to the next process waiting
//
// This function is called by the current process to
// yield to the next process on the ready q.
//
// To yield normally, call process::proc_yield()
//
// Interrupts should already be disabled here
#[no_mangle]
#[inline(never)]
pub unsafe fn _proc_yield<'a>(q: Option<&'a mut ProcessQueue>) {
    // yield onto the right queue
    if !CURRENT_PROCESS.is_null() {
        if let Some(queue) = q {
            (*CURRENT_PROCESS).set_state(State::BLOCKED);
            queue.push_tail(CURRENT_PROCESS);
        } else {
            ready_queue::make_ready(CURRENT_PROCESS);
        }

        CURRENT_PROCESS = 0 as *mut Process;
    }

    // get next process from ready q
    // spawn a reaper when there are a few processes that have died
    let mut next = if self::reaper::DEAD_COUNT >= 5 &&
        JIFFIES % 75 == 0 && 1 == 3 { // TODO: allow reaper
        Process::new("reaper", self::reaper::run)
    } else {
        ready_queue::get_next()
    };

    // run the idle process when everyone is done
    if next.is_null() {
        next = IDLE_PROCESS;
    }

    // switch address spaces
    (*next).addr_space.activate();

    // set the CURRENT_PROCESS
    CURRENT_PROCESS = next;

    //bootlog!("{:?} [Switching]\n", *CURRENT_PROCESS);

    // context switch
    context_switch((*CURRENT_PROCESS).kcontext,
                   if (*CURRENT_PROCESS).disable_cnt == 0 {
                       1 << 9
                   } else {
                       0
                   });

    panic!("The impossible has happened!");
}

// Called by the current process to exit
pub fn exit(code: usize) {
    unsafe {
        if CURRENT_PROCESS.is_null() {
            panic!("Exiting with no current process!\n");
        }

        (*CURRENT_PROCESS).set_state(State::TERMINATED);

        // clean up address space
        (*CURRENT_PROCESS).addr_space.clear();

        // NOTE: need to print *before* adding to reaper q
        printf!("{:?} [Exit 0x{:X}]\n",
                *CURRENT_PROCESS, code);

        // Disable interrupts
        off();

        self::reaper::reaper_add(CURRENT_PROCESS);

        // set current to None, so we will never run this again
        CURRENT_PROCESS = 0 as *mut Process;

        // switch to next ready process
        _proc_yield(None);
    }

    panic!("The impossible has happened!");
}
<|endoftext|>


<commit_before>// Copyright (C) 2020, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

//! Delivery rate estimation.
//!
//! This implements the algorithm for estimating delivery rate as described in
//! https://tools.ietf.org/html/draft-cheng-iccrg-delivery-rate-estimation-00

use std::cmp;

use std::time::Duration;
use std::time::Instant;

use crate::recovery::Sent;

#[derive(Default)]
pub struct Rate {
    delivered: usize,

    delivered_time: Option<Instant>,

    recent_delivered_packet_sent_time: Option<Instant>,

    app_limited_at_pkt: usize,

    rate_sample: RateSample,
}

impl Rate {
    pub fn on_packet_sent(&mut self, pkt: &mut Sent, now: Instant) {
        if self.delivered_time.is_none() {
            self.delivered_time = Some(now);
        }

        if self.recent_delivered_packet_sent_time.is_none() {
            self.recent_delivered_packet_sent_time = Some(now);
        }

        pkt.delivered = self.delivered;
        pkt.delivered_time = self.delivered_time.unwrap();

        pkt.recent_delivered_packet_sent_time =
            self.recent_delivered_packet_sent_time.unwrap();

        pkt.is_app_limited = self.app_limited_at_pkt > 0;
    }

    pub fn on_ack_received(&mut self, pkt: Sent, now: Instant) {
        self.rate_sample.prior_time = Some(pkt.delivered_time);

        self.delivered += pkt.size;
        self.delivered_time = Some(now);

        if pkt.delivered > self.rate_sample.prior_delivered {
            self.rate_sample.prior_delivered = pkt.delivered;
            self.rate_sample.is_app_limited = pkt.is_app_limited;

            self.rate_sample.send_elapsed =
                pkt.time - pkt.recent_delivered_packet_sent_time;

            self.rate_sample.ack_elapsed = self
                .delivered_time
                .unwrap()
                .duration_since(pkt.delivered_time);

            self.recent_delivered_packet_sent_time = Some(pkt.time);
        }
    }

    pub fn estimate(&mut self) {
        if (self.app_limited_at_pkt > 0) &&
            (self.delivered > self.app_limited_at_pkt)
        {
            self.app_limited_at_pkt = 0;
        }

        match self.rate_sample.prior_time {
            Some(_) => {
                self.rate_sample.delivered =
                    self.delivered - self.rate_sample.prior_delivered;

                self.rate_sample.interval = cmp::max(
                    self.rate_sample.send_elapsed,
                    self.rate_sample.ack_elapsed,
                );
            },
            None => return,
        }

        if self.rate_sample.interval.as_secs_f64() > 0.0 {
            self.rate_sample.delivery_rate = (self.rate_sample.delivered as f64 /
                self.rate_sample.interval.as_secs_f64())
                as u64;
        }
    }

    pub fn check_app_limited(&mut self, bytes_in_flight: usize) {
        let limited = self.delivered + bytes_in_flight;
        self.app_limited_at_pkt = if limited > 0 { limited } else { 1 };
    }

    pub fn delivery_rate(&self) -> u64 {
        self.rate_sample.delivery_rate
    }
}

impl std::fmt::Debug for Rate {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "delivered={:?} ", self.delivered)?;

        if let Some(t) = self.delivered_time {
            write!(f, "delivered_time={:?} ", t.elapsed())?;
        }

        if let Some(t) = self.recent_delivered_packet_sent_time {
            write!(f, "recent_delivered_packet_sent_time={:?} ", t.elapsed())?;
        }

        write!(f, "app_limited_at_pkt={:?} ", self.app_limited_at_pkt)?;

        Ok(())
    }
}

#[derive(Default)]
struct RateSample {
    delivery_rate: u64,

    is_app_limited: bool,

    interval: Duration,

    delivered: usize,

    prior_delivered: usize,

    prior_time: Option<Instant>,

    send_elapsed: Duration,

    ack_elapsed: Duration,
}

impl std::fmt::Debug for RateSample {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "delivery_rate={:?} ", self.delivery_rate)?;
        write!(f, "interval={:?} ", self.interval)?;
        write!(f, "delivered={:?} ", self.delivered)?;
        write!(f, "prior_delivered={:?} ", self.prior_delivered)?;
        write!(f, "send_elapsed={:?} ", self.send_elapsed)?;
        if let Some(t) = self.prior_time {
            write!(f, "prior_time={:?} ", t.elapsed())?;
        }
        write!(f, "ack_elapsed={:?}", self.ack_elapsed)?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::recovery::*;

    #[test]
    fn rate_check() {
        let config = Config::new(0xbabababa).unwrap();
        let mut recovery = Recovery::new(&config);

        let mut pkt_1 = Sent {
            pkt_num: 0,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recovery.delivery_rate.on_packet_sent(&mut pkt_1, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recovery.delivery_rate.on_ack_received(pkt_1, Instant::now());

        let mut pkt_2 = Sent {
            pkt_num: 1,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recovery.delivery_rate.on_packet_sent(&mut pkt_2, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recovery.delivery_rate.on_ack_received(pkt_2, Instant::now());
        recovery.delivery_rate.estimate();

        assert!(recovery.delivery_rate() > 0);
    }

    #[test]
    fn app_limited_check() {
        let config = Config::new(0xbabababa).unwrap();
        let mut recvry = Recovery::new(&config);

        let mut pkt_1 = Sent {
            pkt_num: 0,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recvry.delivery_rate.on_packet_sent(&mut pkt_1, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recvry.delivery_rate.on_ack_received(pkt_1, Instant::now());

        let mut pkt_2 = Sent {
            pkt_num: 1,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recvry.app_limited = true;
        recvry.delivery_rate.check_app_limited(recvry.bytes_in_flight);
        recvry.delivery_rate.on_packet_sent(&mut pkt_2, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recvry.delivery_rate.on_ack_received(pkt_2, Instant::now());
        recvry.delivery_rate.estimate();

        assert_eq!(recvry.delivery_rate.app_limited_at_pkt, 0);
    }
}
<commit_msg>delivery_rate: fmt<commit_after>// Copyright (C) 2020, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

//! Delivery rate estimation.
//!
//! This implements the algorithm for estimating delivery rate as described in
//! https://tools.ietf.org/html/draft-cheng-iccrg-delivery-rate-estimation-00

use std::cmp;

use std::time::Duration;
use std::time::Instant;

use crate::recovery::Sent;

#[derive(Default)]
pub struct Rate {
    delivered: usize,

    delivered_time: Option<Instant>,

    recent_delivered_packet_sent_time: Option<Instant>,

    app_limited_at_pkt: usize,

    rate_sample: RateSample,
}

impl Rate {
    pub fn on_packet_sent(&mut self, pkt: &mut Sent, now: Instant) {
        if self.delivered_time.is_none() {
            self.delivered_time = Some(now);
        }

        if self.recent_delivered_packet_sent_time.is_none() {
            self.recent_delivered_packet_sent_time = Some(now);
        }

        pkt.delivered = self.delivered;
        pkt.delivered_time = self.delivered_time.unwrap();

        pkt.recent_delivered_packet_sent_time =
            self.recent_delivered_packet_sent_time.unwrap();

        pkt.is_app_limited = self.app_limited_at_pkt > 0;
    }

    pub fn on_ack_received(&mut self, pkt: Sent, now: Instant) {
        self.rate_sample.prior_time = Some(pkt.delivered_time);

        self.delivered += pkt.size;
        self.delivered_time = Some(now);

        if pkt.delivered > self.rate_sample.prior_delivered {
            self.rate_sample.prior_delivered = pkt.delivered;
            self.rate_sample.is_app_limited = pkt.is_app_limited;

            self.rate_sample.send_elapsed =
                pkt.time - pkt.recent_delivered_packet_sent_time;

            self.rate_sample.ack_elapsed = self
                .delivered_time
                .unwrap()
                .duration_since(pkt.delivered_time);

            self.recent_delivered_packet_sent_time = Some(pkt.time);
        }
    }

    pub fn estimate(&mut self) {
        if (self.app_limited_at_pkt > 0) &&
            (self.delivered > self.app_limited_at_pkt)
        {
            self.app_limited_at_pkt = 0;
        }

        match self.rate_sample.prior_time {
            Some(_) => {
                self.rate_sample.delivered =
                    self.delivered - self.rate_sample.prior_delivered;

                self.rate_sample.interval = cmp::max(
                    self.rate_sample.send_elapsed,
                    self.rate_sample.ack_elapsed,
                );
            },
            None => return,
        }

        if self.rate_sample.interval.as_secs_f64() > 0.0 {
            self.rate_sample.delivery_rate = (self.rate_sample.delivered as f64 /
                self.rate_sample.interval.as_secs_f64())
                as u64;
        }
    }

    pub fn check_app_limited(&mut self, bytes_in_flight: usize) {
        let limited = self.delivered + bytes_in_flight;
        self.app_limited_at_pkt = if limited > 0 { limited } else { 1 };
    }

    pub fn delivery_rate(&self) -> u64 {
        self.rate_sample.delivery_rate
    }
}

impl std::fmt::Debug for Rate {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "delivered={:?} ", self.delivered)?;

        if let Some(t) = self.delivered_time {
            write!(f, "delivered_time={:?} ", t.elapsed())?;
        }

        if let Some(t) = self.recent_delivered_packet_sent_time {
            write!(f, "recent_delivered_packet_sent_time={:?} ", t.elapsed())?;
        }

        write!(f, "app_limited_at_pkt={:?} ", self.app_limited_at_pkt)?;

        Ok(())
    }
}

#[derive(Default)]
struct RateSample {
    delivery_rate: u64,

    is_app_limited: bool,

    interval: Duration,

    delivered: usize,

    prior_delivered: usize,

    prior_time: Option<Instant>,

    send_elapsed: Duration,

    ack_elapsed: Duration,
}

impl std::fmt::Debug for RateSample {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "delivery_rate={:?} ", self.delivery_rate)?;
        write!(f, "interval={:?} ", self.interval)?;
        write!(f, "delivered={:?} ", self.delivered)?;
        write!(f, "prior_delivered={:?} ", self.prior_delivered)?;
        write!(f, "send_elapsed={:?} ", self.send_elapsed)?;
        if let Some(t) = self.prior_time {
            write!(f, "prior_time={:?} ", t.elapsed())?;
        }
        write!(f, "ack_elapsed={:?}", self.ack_elapsed)?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::recovery::*;

    #[test]
    fn rate_check() {
        let config = Config::new(0xbabababa).unwrap();
        let mut recovery = Recovery::new(&config);

        let mut pkt_1 = Sent {
            pkt_num: 0,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recovery
            .delivery_rate
            .on_packet_sent(&mut pkt_1, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recovery
            .delivery_rate
            .on_ack_received(pkt_1, Instant::now());

        let mut pkt_2 = Sent {
            pkt_num: 1,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recovery
            .delivery_rate
            .on_packet_sent(&mut pkt_2, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recovery
            .delivery_rate
            .on_ack_received(pkt_2, Instant::now());
        recovery.delivery_rate.estimate();

        assert!(recovery.delivery_rate() > 0);
    }

    #[test]
    fn app_limited_check() {
        let config = Config::new(0xbabababa).unwrap();
        let mut recvry = Recovery::new(&config);

        let mut pkt_1 = Sent {
            pkt_num: 0,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recvry
            .delivery_rate
            .on_packet_sent(&mut pkt_1, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recvry.delivery_rate.on_ack_received(pkt_1, Instant::now());

        let mut pkt_2 = Sent {
            pkt_num: 1,
            frames: vec![],
            time: Instant::now(),
            size: 1200,
            ack_eliciting: true,
            in_flight: true,
            delivered: 0,
            delivered_time: Instant::now(),
            recent_delivered_packet_sent_time: Instant::now(),
            is_app_limited: false,
        };

        recvry.app_limited = true;
        recvry
            .delivery_rate
            .check_app_limited(recvry.bytes_in_flight);
        recvry
            .delivery_rate
            .on_packet_sent(&mut pkt_2, Instant::now());
        std::thread::sleep(Duration::from_millis(50));
        recvry.delivery_rate.on_ack_received(pkt_2, Instant::now());
        recvry.delivery_rate.estimate();

        assert_eq!(recvry.delivery_rate.app_limited_at_pkt, 0);
    }
}
<|endoftext|>


<commit_before>//! Thin wrappers around `std::path`, distinguishing between absolute and
//! relative paths.
use std::{
    borrow::Borrow,
    convert::{TryFrom, TryInto},
    ffi::OsStr,
    ops,
    path::{Component, Path, PathBuf},
};

/// Wrapper around an absolute [`PathBuf`].
#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct AbsPathBuf(PathBuf);

impl From<AbsPathBuf> for PathBuf {
    fn from(AbsPathBuf(path_buf): AbsPathBuf) -> PathBuf {
        path_buf
    }
}

impl ops::Deref for AbsPathBuf {
    type Target = AbsPath;
    fn deref(&self) -> &AbsPath {
        self.as_path()
    }
}

impl AsRef<Path> for AbsPathBuf {
    fn as_ref(&self) -> &Path {
        self.0.as_path()
    }
}

impl AsRef<AbsPath> for AbsPathBuf {
    fn as_ref(&self) -> &AbsPath {
        self.as_path()
    }
}

impl Borrow<AbsPath> for AbsPathBuf {
    fn borrow(&self) -> &AbsPath {
        self.as_path()
    }
}

impl TryFrom<PathBuf> for AbsPathBuf {
    type Error = PathBuf;
    fn try_from(path_buf: PathBuf) -> Result<AbsPathBuf, PathBuf> {
        if !path_buf.is_absolute() {
            return Err(path_buf);
        }
        Ok(AbsPathBuf(path_buf))
    }
}

impl TryFrom<&str> for AbsPathBuf {
    type Error = PathBuf;
    fn try_from(path: &str) -> Result<AbsPathBuf, PathBuf> {
        AbsPathBuf::try_from(PathBuf::from(path))
    }
}

impl PartialEq<AbsPath> for AbsPathBuf {
    fn eq(&self, other: &AbsPath) -> bool {
        self.as_path() == other
    }
}

impl serde::Serialize for AbsPathBuf {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.0.serialize(serializer)
    }
}

impl<'de> serde::Deserialize<'de> for AbsPathBuf {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let path = PathBuf::deserialize(deserializer)?;
        AbsPathBuf::try_from(path).map_err(|path| {
            serde::de::Error::custom(format!("expected absolute path, got {}", path.display()))
        })
    }
}

impl AbsPathBuf {
    /// Wrap the given absolute path in `AbsPathBuf`
    ///
    /// # Panics
    ///
    /// Panics if `path` is not absolute.
    pub fn assert(path: PathBuf) -> AbsPathBuf {
        AbsPathBuf::try_from(path)
            .unwrap_or_else(|path| panic!("expected absolute path, got {}", path.display()))
    }

    /// Coerces to a `AbsPath` slice.
    ///
    /// Equivalent of [`PathBuf::as_path`] for `AbsPathBuf`.
    pub fn as_path(&self) -> &AbsPath {
        AbsPath::assert(self.0.as_path())
    }

    /// Equivalent of [`PathBuf::pop`] for `AbsPathBuf`.
    ///
    /// Note that this won't remove the root component, so `self` will still be
    /// absolute.
    pub fn pop(&mut self) -> bool {
        self.0.pop()
    }
}

/// Wrapper around an absolute [`Path`].
#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
#[repr(transparent)]
pub struct AbsPath(Path);

impl AsRef<Path> for AbsPath {
    fn as_ref(&self) -> &Path {
        &self.0
    }
}

impl<'a> TryFrom<&'a Path> for &'a AbsPath {
    type Error = &'a Path;
    fn try_from(path: &'a Path) -> Result<&'a AbsPath, &'a Path> {
        if !path.is_absolute() {
            return Err(path);
        }
        Ok(AbsPath::assert(path))
    }
}

impl AbsPath {
    /// Wrap the given absolute path in `AbsPath`
    ///
    /// # Panics
    ///
    /// Panics if `path` is not absolute.
    pub fn assert(path: &Path) -> &AbsPath {
        assert!(path.is_absolute());
        unsafe { &*(path as *const Path as *const AbsPath) }
    }

    /// Equivalent of [`Path::parent`] for `AbsPath`.
    pub fn parent(&self) -> Option<&AbsPath> {
        self.0.parent().map(AbsPath::assert)
    }

    /// Equivalent of [`Path::join`] for `AbsPath`.
    pub fn join(&self, path: impl AsRef<Path>) -> AbsPathBuf {
        self.as_ref().join(path).try_into().unwrap()
    }

    /// Normalize the given path:
    /// - Removes repeated separators: `/a//b` becomes `/a/b`
    /// - Removes occurrences of `.` and resolves `..`.
    /// - Removes trailing slashes: `/a/b/` becomes `/a/b`.
    ///
    /// # Example
    /// ```
    /// # use paths::AbsPathBuf;
    /// let abs_path_buf = AbsPathBuf::assert("/a/../../b/.//c//".into());
    /// let normalized = abs_path_buf.normalize();
    /// assert_eq!(normalized, AbsPathBuf::assert("/b/c".into()));
    /// ```
    pub fn normalize(&self) -> AbsPathBuf {
        AbsPathBuf(normalize_path(&self.0))
    }

    /// Equivalent of [`Path::to_path_buf`] for `AbsPath`.
    pub fn to_path_buf(&self) -> AbsPathBuf {
        AbsPathBuf::try_from(self.0.to_path_buf()).unwrap()
    }

    /// Equivalent of [`Path::strip_prefix`] for `AbsPath`.
    ///
    /// Returns a relative path.
    pub fn strip_prefix(&self, base: &AbsPath) -> Option<&RelPath> {
        self.0.strip_prefix(base).ok().map(RelPath::new_unchecked)
    }
    pub fn starts_with(&self, base: &AbsPath) -> bool {
        self.0.starts_with(&base.0)
    }
    pub fn ends_with(&self, suffix: &RelPath) -> bool {
        self.0.starts_with(&suffix.0)
    }

    // region:delegate-methods

    // Note that we deliberately don't implement `Deref<Target = Path>` here.
    //
    // The problem with `Path` is that it directly exposes convenience IO-ing
    // methods. For example, `Path::exists` delegates to `fs::metadata`.
    //
    // For `AbsPath`, we want to make sure that this is a POD type, and that all
    // IO goes via `fs`. That way, it becomes easier to mock IO when we need it.

    pub fn file_name(&self) -> Option<&OsStr> {
        self.0.file_name()
    }
    pub fn extension(&self) -> Option<&OsStr> {
        self.0.extension()
    }
    pub fn file_stem(&self) -> Option<&OsStr> {
        self.0.file_stem()
    }
    pub fn as_os_str(&self) -> &OsStr {
        self.0.as_os_str()
    }
    pub fn display(&self) -> std::path::Display<'_> {
        self.0.display()
    }
    #[deprecated(note = "use std::fs::metadata().is_ok() instead")]
    pub fn exists(&self) -> bool {
        self.0.exists()
    }
    // endregion:delegate-methods
}

/// Wrapper around a relative [`PathBuf`].
#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct RelPathBuf(PathBuf);

impl From<RelPathBuf> for PathBuf {
    fn from(RelPathBuf(path_buf): RelPathBuf) -> PathBuf {
        path_buf
    }
}

impl ops::Deref for RelPathBuf {
    type Target = RelPath;
    fn deref(&self) -> &RelPath {
        self.as_path()
    }
}

impl AsRef<Path> for RelPathBuf {
    fn as_ref(&self) -> &Path {
        self.0.as_path()
    }
}

impl TryFrom<PathBuf> for RelPathBuf {
    type Error = PathBuf;
    fn try_from(path_buf: PathBuf) -> Result<RelPathBuf, PathBuf> {
        if !path_buf.is_relative() {
            return Err(path_buf);
        }
        Ok(RelPathBuf(path_buf))
    }
}

impl TryFrom<&str> for RelPathBuf {
    type Error = PathBuf;
    fn try_from(path: &str) -> Result<RelPathBuf, PathBuf> {
        RelPathBuf::try_from(PathBuf::from(path))
    }
}

impl RelPathBuf {
    /// Coerces to a `RelPath` slice.
    ///
    /// Equivalent of [`PathBuf::as_path`] for `RelPathBuf`.
    pub fn as_path(&self) -> &RelPath {
        RelPath::new_unchecked(self.0.as_path())
    }
}

/// Wrapper around a relative [`Path`].
#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
#[repr(transparent)]
pub struct RelPath(Path);

impl AsRef<Path> for RelPath {
    fn as_ref(&self) -> &Path {
        &self.0
    }
}

impl RelPath {
    /// Creates a new `RelPath` from `path`, without checking if it is relative.
    pub fn new_unchecked(path: &Path) -> &RelPath {
        unsafe { &*(path as *const Path as *const RelPath) }
    }
}

/// Taken from <https://github.com/rust-lang/cargo/blob/79c769c3d7b4c2cf6a93781575b7f592ef974255/src/cargo/util/paths.rs#L60-L85>
fn normalize_path(path: &Path) -> PathBuf {
    let mut components = path.components().peekable();
    let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().cloned() {
        components.next();
        PathBuf::from(c.as_os_str())
    } else {
        PathBuf::new()
    };

    for component in components {
        match component {
            Component::Prefix(..) => unreachable!(),
            Component::RootDir => {
                ret.push(component.as_os_str());
            }
            Component::CurDir => {}
            Component::ParentDir => {
                ret.pop();
            }
            Component::Normal(c) => {
                ret.push(c);
            }
        }
    }
    ret
}
<commit_msg>Fix AbsPath::ends_with<commit_after>//! Thin wrappers around `std::path`, distinguishing between absolute and
//! relative paths.
use std::{
    borrow::Borrow,
    convert::{TryFrom, TryInto},
    ffi::OsStr,
    ops,
    path::{Component, Path, PathBuf},
};

/// Wrapper around an absolute [`PathBuf`].
#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct AbsPathBuf(PathBuf);

impl From<AbsPathBuf> for PathBuf {
    fn from(AbsPathBuf(path_buf): AbsPathBuf) -> PathBuf {
        path_buf
    }
}

impl ops::Deref for AbsPathBuf {
    type Target = AbsPath;
    fn deref(&self) -> &AbsPath {
        self.as_path()
    }
}

impl AsRef<Path> for AbsPathBuf {
    fn as_ref(&self) -> &Path {
        self.0.as_path()
    }
}

impl AsRef<AbsPath> for AbsPathBuf {
    fn as_ref(&self) -> &AbsPath {
        self.as_path()
    }
}

impl Borrow<AbsPath> for AbsPathBuf {
    fn borrow(&self) -> &AbsPath {
        self.as_path()
    }
}

impl TryFrom<PathBuf> for AbsPathBuf {
    type Error = PathBuf;
    fn try_from(path_buf: PathBuf) -> Result<AbsPathBuf, PathBuf> {
        if !path_buf.is_absolute() {
            return Err(path_buf);
        }
        Ok(AbsPathBuf(path_buf))
    }
}

impl TryFrom<&str> for AbsPathBuf {
    type Error = PathBuf;
    fn try_from(path: &str) -> Result<AbsPathBuf, PathBuf> {
        AbsPathBuf::try_from(PathBuf::from(path))
    }
}

impl PartialEq<AbsPath> for AbsPathBuf {
    fn eq(&self, other: &AbsPath) -> bool {
        self.as_path() == other
    }
}

impl serde::Serialize for AbsPathBuf {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.0.serialize(serializer)
    }
}

impl<'de> serde::Deserialize<'de> for AbsPathBuf {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let path = PathBuf::deserialize(deserializer)?;
        AbsPathBuf::try_from(path).map_err(|path| {
            serde::de::Error::custom(format!("expected absolute path, got {}", path.display()))
        })
    }
}

impl AbsPathBuf {
    /// Wrap the given absolute path in `AbsPathBuf`
    ///
    /// # Panics
    ///
    /// Panics if `path` is not absolute.
    pub fn assert(path: PathBuf) -> AbsPathBuf {
        AbsPathBuf::try_from(path)
            .unwrap_or_else(|path| panic!("expected absolute path, got {}", path.display()))
    }

    /// Coerces to a `AbsPath` slice.
    ///
    /// Equivalent of [`PathBuf::as_path`] for `AbsPathBuf`.
    pub fn as_path(&self) -> &AbsPath {
        AbsPath::assert(self.0.as_path())
    }

    /// Equivalent of [`PathBuf::pop`] for `AbsPathBuf`.
    ///
    /// Note that this won't remove the root component, so `self` will still be
    /// absolute.
    pub fn pop(&mut self) -> bool {
        self.0.pop()
    }
}

/// Wrapper around an absolute [`Path`].
#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
#[repr(transparent)]
pub struct AbsPath(Path);

impl AsRef<Path> for AbsPath {
    fn as_ref(&self) -> &Path {
        &self.0
    }
}

impl<'a> TryFrom<&'a Path> for &'a AbsPath {
    type Error = &'a Path;
    fn try_from(path: &'a Path) -> Result<&'a AbsPath, &'a Path> {
        if !path.is_absolute() {
            return Err(path);
        }
        Ok(AbsPath::assert(path))
    }
}

impl AbsPath {
    /// Wrap the given absolute path in `AbsPath`
    ///
    /// # Panics
    ///
    /// Panics if `path` is not absolute.
    pub fn assert(path: &Path) -> &AbsPath {
        assert!(path.is_absolute());
        unsafe { &*(path as *const Path as *const AbsPath) }
    }

    /// Equivalent of [`Path::parent`] for `AbsPath`.
    pub fn parent(&self) -> Option<&AbsPath> {
        self.0.parent().map(AbsPath::assert)
    }

    /// Equivalent of [`Path::join`] for `AbsPath`.
    pub fn join(&self, path: impl AsRef<Path>) -> AbsPathBuf {
        self.as_ref().join(path).try_into().unwrap()
    }

    /// Normalize the given path:
    /// - Removes repeated separators: `/a//b` becomes `/a/b`
    /// - Removes occurrences of `.` and resolves `..`.
    /// - Removes trailing slashes: `/a/b/` becomes `/a/b`.
    ///
    /// # Example
    /// ```
    /// # use paths::AbsPathBuf;
    /// let abs_path_buf = AbsPathBuf::assert("/a/../../b/.//c//".into());
    /// let normalized = abs_path_buf.normalize();
    /// assert_eq!(normalized, AbsPathBuf::assert("/b/c".into()));
    /// ```
    pub fn normalize(&self) -> AbsPathBuf {
        AbsPathBuf(normalize_path(&self.0))
    }

    /// Equivalent of [`Path::to_path_buf`] for `AbsPath`.
    pub fn to_path_buf(&self) -> AbsPathBuf {
        AbsPathBuf::try_from(self.0.to_path_buf()).unwrap()
    }

    /// Equivalent of [`Path::strip_prefix`] for `AbsPath`.
    ///
    /// Returns a relative path.
    pub fn strip_prefix(&self, base: &AbsPath) -> Option<&RelPath> {
        self.0.strip_prefix(base).ok().map(RelPath::new_unchecked)
    }
    pub fn starts_with(&self, base: &AbsPath) -> bool {
        self.0.starts_with(&base.0)
    }
    pub fn ends_with(&self, suffix: &RelPath) -> bool {
        self.0.ends_with(&suffix.0)
    }

    // region:delegate-methods

    // Note that we deliberately don't implement `Deref<Target = Path>` here.
    //
    // The problem with `Path` is that it directly exposes convenience IO-ing
    // methods. For example, `Path::exists` delegates to `fs::metadata`.
    //
    // For `AbsPath`, we want to make sure that this is a POD type, and that all
    // IO goes via `fs`. That way, it becomes easier to mock IO when we need it.

    pub fn file_name(&self) -> Option<&OsStr> {
        self.0.file_name()
    }
    pub fn extension(&self) -> Option<&OsStr> {
        self.0.extension()
    }
    pub fn file_stem(&self) -> Option<&OsStr> {
        self.0.file_stem()
    }
    pub fn as_os_str(&self) -> &OsStr {
        self.0.as_os_str()
    }
    pub fn display(&self) -> std::path::Display<'_> {
        self.0.display()
    }
    #[deprecated(note = "use std::fs::metadata().is_ok() instead")]
    pub fn exists(&self) -> bool {
        self.0.exists()
    }
    // endregion:delegate-methods
}

/// Wrapper around a relative [`PathBuf`].
#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct RelPathBuf(PathBuf);

impl From<RelPathBuf> for PathBuf {
    fn from(RelPathBuf(path_buf): RelPathBuf) -> PathBuf {
        path_buf
    }
}

impl ops::Deref for RelPathBuf {
    type Target = RelPath;
    fn deref(&self) -> &RelPath {
        self.as_path()
    }
}

impl AsRef<Path> for RelPathBuf {
    fn as_ref(&self) -> &Path {
        self.0.as_path()
    }
}

impl TryFrom<PathBuf> for RelPathBuf {
    type Error = PathBuf;
    fn try_from(path_buf: PathBuf) -> Result<RelPathBuf, PathBuf> {
        if !path_buf.is_relative() {
            return Err(path_buf);
        }
        Ok(RelPathBuf(path_buf))
    }
}

impl TryFrom<&str> for RelPathBuf {
    type Error = PathBuf;
    fn try_from(path: &str) -> Result<RelPathBuf, PathBuf> {
        RelPathBuf::try_from(PathBuf::from(path))
    }
}

impl RelPathBuf {
    /// Coerces to a `RelPath` slice.
    ///
    /// Equivalent of [`PathBuf::as_path`] for `RelPathBuf`.
    pub fn as_path(&self) -> &RelPath {
        RelPath::new_unchecked(self.0.as_path())
    }
}

/// Wrapper around a relative [`Path`].
#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
#[repr(transparent)]
pub struct RelPath(Path);

impl AsRef<Path> for RelPath {
    fn as_ref(&self) -> &Path {
        &self.0
    }
}

impl RelPath {
    /// Creates a new `RelPath` from `path`, without checking if it is relative.
    pub fn new_unchecked(path: &Path) -> &RelPath {
        unsafe { &*(path as *const Path as *const RelPath) }
    }
}

/// Taken from <https://github.com/rust-lang/cargo/blob/79c769c3d7b4c2cf6a93781575b7f592ef974255/src/cargo/util/paths.rs#L60-L85>
fn normalize_path(path: &Path) -> PathBuf {
    let mut components = path.components().peekable();
    let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().cloned() {
        components.next();
        PathBuf::from(c.as_os_str())
    } else {
        PathBuf::new()
    };

    for component in components {
        match component {
            Component::Prefix(..) => unreachable!(),
            Component::RootDir => {
                ret.push(component.as_os_str());
            }
            Component::CurDir => {}
            Component::ParentDir => {
                ret.pop();
            }
            Component::Normal(c) => {
                ret.push(c);
            }
        }
    }
    ret
}
<|endoftext|>


<commit_before>use super::*;

use std::{collections::BTreeSet, ops::Bound};

#[derive(Clone)]
pub(crate) struct View<'a> {
    pub(crate) pid: PageId,
    pub(crate) lo: &'a IVec,
    pub(crate) hi: &'a IVec,
    pub(crate) is_index: bool,
    pub(crate) next: Option<PageId>,
    pub(crate) ptr: TreePtr<'a>,
    frags: Vec<&'a Frag>,
    base_offset: usize,
    pub(crate) base_data: &'a Data,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
}

impl<'a> View<'a> {
    pub(crate) fn new(
        pid: PageId,
        ptr: TreePtr<'a>,
        frags: Vec<&'a Frag>,
    ) -> View<'a> {
        let mut view = View {
            pid,
            ptr,
            frags,
            lo: unsafe { std::mem::uninitialized() },
            hi: unsafe { std::mem::uninitialized() },
            is_index: false,
            base_offset: usize::max_value(),
            base_data: unsafe { std::mem::uninitialized() },
            next: None,
            merging_child: None,
            merging: false,
        };

        let mut merge_confirmed = false;

        for (offset, frag) in view.frags.iter().enumerate() {
            match frag {
                Frag::Base(node) => {
                    if view.base_offset == usize::max_value() {
                        // hi and next may be changed via a
                        // parent split if we re-add that frag
                        view.hi = &node.hi;
                        view.next = node.next;
                    }
                    view.lo = &node.lo;
                    view.is_index = node.data.index_ref().is_some();
                    view.base_offset = offset;
                    view.base_data = &node.data;
                    break;
                }
                Frag::ParentMergeIntention(pid) => {
                    if merge_confirmed {
                        merge_confirmed = false;
                    } else {
                        assert!(view.merging_child.is_none());
                        view.merging_child = Some(*pid);
                    }
                }
                Frag::ParentMergeConfirm => {
                    assert!(!merge_confirmed);
                    merge_confirmed = true;
                }
                Frag::ChildMergeCap => {
                    assert!(!view.merging);
                    assert_eq!(offset, 0);
                    view.merging = true;
                }
                _ => {}
            }
        }

        assert_ne!(
            view.base_offset,
            usize::max_value(),
            "view was never initialized with a base"
        );

        view
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.hi.is_empty(),
            Bound::Included(bound) => self.hi > bound,
            Bound::Excluded(bound) => self.hi >= bound,
        }
    }

    pub(crate) fn contains_lower_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.lo.is_empty(),
            Bound::Included(bound) => self.lo <= bound,
            Bound::Excluded(bound) => self.lo < bound,
        }
    }

    fn keys(&self) -> BTreeSet<&IVec> {
        let mut keys: BTreeSet<&IVec> = self
            .base_data
            .leaf_ref()
            .unwrap()
            .iter()
            .map(|(k, _v)| k)
            .collect();

        for offset in (0..self.base_offset).rev() {
            match self.frags[offset] {
                Frag::Set(k, _) => {
                    keys.insert(k);
                }
                Frag::Del(k) => {
                    keys.remove(k);
                }
                Frag::Merge(k, _) => {
                    keys.insert(k);
                }
                Frag::Base(_) => {
                    panic!(
                        "somehow hit 2 base nodes while \
                         searching for a successor"
                    );
                }
                Frag::ChildMergeCap => {}
                Frag::ParentMergeIntention(_) | Frag::ParentMergeConfirm => {
                    panic!(
                        "somehow hit parent merge \
                         frags while searching for a \
                         successor"
                    )
                }
            }
        }

        keys
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(self.lo, self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
        };

        let keys = self.keys();

        let successor_keys = keys.range(predecessor_key..);

        for encoded_key in successor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    prefix_encode(self.lo, &[255; 1024 * 1024])
                } else {
                    prefix_encode(self.lo, self.hi)
                }
            }
            Bound::Included(b) => {
                let min = std::cmp::min(b, self.hi);
                prefix_encode(self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = std::cmp::min(b, self.hi);
                prefix_encode(self.lo, min)
            }
        };

        let keys = self.keys();

        let predecessor_keys = keys.range(..=successor_key).rev();

        for encoded_key in predecessor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn is_free(&self) -> bool {
        self.frags.is_empty()
    }

    pub(crate) fn leaf_value_for_key(
        &self,
        key: &[u8],
        config: &Config,
    ) -> Option<IVec> {
        assert!(!self.is_index);

        let mut merge_base = None;
        let mut merges = vec![];

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(k, val) if self.key_eq(k, key) => {
                    if merges.is_empty() {
                        return Some(val.clone());
                    } else {
                        merge_base = Some(val);
                        break;
                    }
                }
                Frag::Del(k) if self.key_eq(k, key) => return None,
                Frag::Merge(k, val) if self.key_eq(k, key) => merges.push(val),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.leaf_ref().expect("last_node should be a leaf");
                    let search = items
                        .binary_search_by(|&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        })
                        .ok();

                    let val = search.map(|idx| &items[idx].1);
                    if merges.is_empty() {
                        return val.cloned();
                    } else {
                        merge_base = val;
                    }
                }
                _ => {}
            }
        }

        if merges.is_empty() {
            None
        } else {
            let merge_fn_ptr = config
                .merge_operator
                .expect("must have a merge operator set");

            unsafe {
                let merge_fn: MergeOperator = std::mem::transmute(merge_fn_ptr);

                let mut ret = merge_fn(
                    key,
                    merge_base.map(|iv| &**iv),
                    &merges.pop().unwrap(),
                );
                                       ;
                for merge in merges.into_iter().rev() {
                    if let Some(v) = ret {
                        ret = merge_fn(key, Some(&*v), merge);
                    } else {
                        ret = merge_fn(key, None, merge);
                    }
                }

                ret.map(IVec::from)
            }
        }
    }

    #[inline]
    fn key_eq(&self, encoded: &[u8], not_encoded: &[u8]) -> bool {
        prefix_cmp_encoded(encoded, not_encoded, self.lo)
            == std::cmp::Ordering::Equal
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> PageId {
        assert!(self.is_index);

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(..) => unimplemented!(),
                Frag::Del(..) => unimplemented!(),
                Frag::Merge(..) => unimplemented!(),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.index_ref().expect("last_node should be a leaf");
                    let search =
                        binary_search_lub(items, |&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        });

                    // This might be none if ord is Less and we're
                    // searching for the empty key
                    let index = search.expect("failed to traverse index");

                    return items[index].1;
                }
                Frag::ParentMergeIntention(_)
                | Frag::ParentMergeConfirm
                | Frag::ChildMergeCap => {
                    // nothing to do for these frags
                }
            }
        }
        panic!("no index found")
    }

    pub(crate) fn should_split(&self, max_sz: u64) -> bool {
        let children = self.base_data.len();
        children > 2
            && self.size_in_bytes() > max_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn should_merge(&self, min_sz: u64) -> bool {
        self.size_in_bytes() < min_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn compact(&self, config: &Config) -> Node {
        let mut lhs = self.frags[self.base_offset].unwrap_base().clone();
        for offset in (0..self.base_offset).rev() {
            let frag = self.frags[offset];
            lhs.apply(frag, config.merge_operator);
        }
        lhs
    }

    pub(crate) fn split(&self, config: &Config) -> (Node, Node) {
        let mut lhs = self.compact(config);
        let rhs = lhs.split();

        lhs.data.drop_gte(&rhs.lo, &lhs.lo);
        lhs.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        lhs.next = None;

        (lhs, rhs)
    }

    #[inline]
    pub(crate) fn size_in_bytes(&self) -> u64 {
        // TODO needs to better account for
        // sizes that don't actually fall under
        // a merge threshold once we support one.
        self.frags[..self.base_offset + 1]
            .iter()
            .map(|f| f.size_in_bytes())
            .sum()
    }
}
<commit_msg>Fix bug with hi bounds being incorrectly handled while iterating<commit_after>use super::*;

use std::{collections::BTreeSet, ops::Bound};

#[derive(Debug, Clone)]
pub(crate) struct View<'a> {
    pub(crate) pid: PageId,
    pub(crate) lo: &'a IVec,
    pub(crate) hi: &'a IVec,
    pub(crate) is_index: bool,
    pub(crate) next: Option<PageId>,
    pub(crate) ptr: TreePtr<'a>,
    frags: Vec<&'a Frag>,
    base_offset: usize,
    pub(crate) base_data: &'a Data,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
}

impl<'a> View<'a> {
    pub(crate) fn new(
        pid: PageId,
        ptr: TreePtr<'a>,
        frags: Vec<&'a Frag>,
    ) -> View<'a> {
        let mut view = View {
            pid,
            ptr,
            frags,
            lo: unsafe { std::mem::uninitialized() },
            hi: unsafe { std::mem::uninitialized() },
            is_index: false,
            base_offset: usize::max_value(),
            base_data: unsafe { std::mem::uninitialized() },
            next: None,
            merging_child: None,
            merging: false,
        };

        let mut merge_confirmed = false;

        for (offset, frag) in view.frags.iter().enumerate() {
            match frag {
                Frag::Base(node) => {
                    if view.base_offset == usize::max_value() {
                        // hi and next may be changed via a
                        // parent split if we re-add that frag
                        view.hi = &node.hi;
                        view.next = node.next;
                    }
                    view.lo = &node.lo;
                    view.is_index = node.data.index_ref().is_some();
                    view.base_offset = offset;
                    view.base_data = &node.data;
                    break;
                }
                Frag::ParentMergeIntention(pid) => {
                    if merge_confirmed {
                        merge_confirmed = false;
                    } else {
                        assert!(view.merging_child.is_none());
                        view.merging_child = Some(*pid);
                    }
                }
                Frag::ParentMergeConfirm => {
                    assert!(!merge_confirmed);
                    merge_confirmed = true;
                }
                Frag::ChildMergeCap => {
                    assert!(!view.merging);
                    assert_eq!(offset, 0);
                    view.merging = true;
                }
                _ => {}
            }
        }

        assert_ne!(
            view.base_offset,
            usize::max_value(),
            "view was never initialized with a base"
        );

        view
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.hi.is_empty(),
            Bound::Included(bound) => self.hi > bound || self.hi.is_empty(),
            Bound::Excluded(bound) => self.hi >= bound || self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.lo.is_empty(),
            Bound::Included(bound) => self.lo <= bound,
            Bound::Excluded(bound) => self.lo < bound,
        }
    }

    fn keys(&self) -> BTreeSet<&IVec> {
        let mut keys: BTreeSet<&IVec> = self
            .base_data
            .leaf_ref()
            .unwrap()
            .iter()
            .map(|(k, _v)| k)
            .collect();

        for offset in (0..self.base_offset).rev() {
            match self.frags[offset] {
                Frag::Set(k, _) => {
                    keys.insert(k);
                }
                Frag::Del(k) => {
                    keys.remove(k);
                }
                Frag::Merge(k, _) => {
                    keys.insert(k);
                }
                Frag::Base(_) => {
                    panic!(
                        "somehow hit 2 base nodes while \
                         searching for a successor"
                    );
                }
                Frag::ChildMergeCap => {}
                Frag::ParentMergeIntention(_) | Frag::ParentMergeConfirm => {
                    panic!(
                        "somehow hit parent merge \
                         frags while searching for a \
                         successor"
                    )
                }
            }
        }

        keys
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(self.lo, self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
        };

        let keys = self.keys();

        let successor_keys = keys.range(predecessor_key..);

        for encoded_key in successor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    prefix_encode(self.lo, &[255; 1024 * 1024])
                } else {
                    prefix_encode(self.lo, self.hi)
                }
            }
            Bound::Included(b) => {
                let min = std::cmp::min(b, self.hi);
                prefix_encode(self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = std::cmp::min(b, self.hi);
                prefix_encode(self.lo, min)
            }
        };

        let keys = self.keys();

        let predecessor_keys = keys.range(..=successor_key).rev();

        for encoded_key in predecessor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn is_free(&self) -> bool {
        self.frags.is_empty()
    }

    pub(crate) fn leaf_value_for_key(
        &self,
        key: &[u8],
        config: &Config,
    ) -> Option<IVec> {
        assert!(!self.is_index);

        let mut merge_base = None;
        let mut merges = vec![];

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(k, val) if self.key_eq(k, key) => {
                    if merges.is_empty() {
                        return Some(val.clone());
                    } else {
                        merge_base = Some(val);
                        break;
                    }
                }
                Frag::Del(k) if self.key_eq(k, key) => return None,
                Frag::Merge(k, val) if self.key_eq(k, key) => merges.push(val),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.leaf_ref().expect("last_node should be a leaf");
                    let search = items
                        .binary_search_by(|&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        })
                        .ok();

                    let val = search.map(|idx| &items[idx].1);
                    if merges.is_empty() {
                        return val.cloned();
                    } else {
                        merge_base = val;
                    }
                }
                _ => {}
            }
        }

        if merges.is_empty() {
            None
        } else {
            let merge_fn_ptr = config
                .merge_operator
                .expect("must have a merge operator set");

            unsafe {
                let merge_fn: MergeOperator = std::mem::transmute(merge_fn_ptr);

                let mut ret = merge_fn(
                    key,
                    merge_base.map(|iv| &**iv),
                    &merges.pop().unwrap(),
                );
                                       ;
                for merge in merges.into_iter().rev() {
                    if let Some(v) = ret {
                        ret = merge_fn(key, Some(&*v), merge);
                    } else {
                        ret = merge_fn(key, None, merge);
                    }
                }

                ret.map(IVec::from)
            }
        }
    }

    #[inline]
    fn key_eq(&self, encoded: &[u8], not_encoded: &[u8]) -> bool {
        prefix_cmp_encoded(encoded, not_encoded, self.lo)
            == std::cmp::Ordering::Equal
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> PageId {
        assert!(self.is_index);

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(..) => unimplemented!(),
                Frag::Del(..) => unimplemented!(),
                Frag::Merge(..) => unimplemented!(),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.index_ref().expect("last_node should be a leaf");
                    let search =
                        binary_search_lub(items, |&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        });

                    // This might be none if ord is Less and we're
                    // searching for the empty key
                    let index = search.expect("failed to traverse index");

                    return items[index].1;
                }
                Frag::ParentMergeIntention(_)
                | Frag::ParentMergeConfirm
                | Frag::ChildMergeCap => {
                    // nothing to do for these frags
                }
            }
        }
        panic!("no index found")
    }

    pub(crate) fn should_split(&self, max_sz: u64) -> bool {
        let children = self.base_data.len();
        children > 2
            && self.size_in_bytes() > max_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn should_merge(&self, min_sz: u64) -> bool {
        self.size_in_bytes() < min_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn compact(&self, config: &Config) -> Node {
        let mut lhs = self.frags[self.base_offset].unwrap_base().clone();
        for offset in (0..self.base_offset).rev() {
            let frag = self.frags[offset];
            lhs.apply(frag, config.merge_operator);
        }
        lhs
    }

    pub(crate) fn split(&self, config: &Config) -> (Node, Node) {
        let mut lhs = self.compact(config);
        let rhs = lhs.split();

        lhs.data.drop_gte(&rhs.lo, &lhs.lo);
        lhs.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        lhs.next = None;

        (lhs, rhs)
    }

    #[inline]
    pub(crate) fn size_in_bytes(&self) -> u64 {
        // TODO needs to better account for
        // sizes that don't actually fall under
        // a merge threshold once we support one.
        self.frags[..self.base_offset + 1]
            .iter()
            .map(|f| f.size_in_bytes())
            .sum()
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::io::{Cursor, Read};
use std::sync::{Arc, Mutex};

use zip::ZipArchive;
use image::{self, DynamicImage, GenericImage, ImageLuma8, ImageFormat};

use cache::{AssetLoadContext, GeneratedAsset, WebAsset};
use terrain::raster::{GlobalRaster, Raster, RasterSource};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum LandCoverKind {
    TreeCover,
    WaterMask,
}
impl RasterSource for LandCoverKind {
    type Type = u8;
    fn load(
        &self,
        context: &mut AssetLoadContext,
        latitude: i16,
        longitude: i16,
    ) -> Option<Raster<u8>> {
        LandCoverParams {
            latitude,
            longitude,
            kind: *self,
            raw: None,
        }.load(context)
            .ok()
    }
}

/// Coordinates are of the lower left corner of the 10x10 degree cell.
#[derive(Debug, Eq, PartialEq)]
struct RawLandCoverParams {
    pub latitude: i16,
    pub longitude: i16,
    pub kind: LandCoverKind,
}
impl WebAsset for RawLandCoverParams {
    type Type = Arc<Mutex<DynamicImage>>;

    fn url(&self) -> String {
        let (latitude, longitude) = (self.latitude + 10, self.longitude);
        assert_eq!(latitude % 10, 0);
        assert_eq!(longitude % 10, 0);
        let n_or_s = if latitude >= 0 { 'N' } else { 'S' };
        let e_or_w = if longitude >= 0 { 'E' } else { 'W' };

        match self.kind {
            LandCoverKind::TreeCover => {
                format!("https://edcintl.cr.usgs.gov/downloads/sciweb1/shared/gtc/downloads/\
                         treecover2010_v3_individual/{:02}{}_{:03}{}_treecover2010_v3.tif.zip",
                latitude.abs(),
                n_or_s,
                longitude.abs(),
                e_or_w,
            )
            }
            LandCoverKind::WaterMask => {
                format!("https://edcintl.cr.usgs.gov/downloads/sciweb1/shared/gtc/downloads/\
                         WaterMask2010_UMD_individual/Hansen_GFC2013_datamask_{:02}{}_{:03}{}\
                         .tif.zip",
                latitude.abs(),
                n_or_s,
                longitude.abs(),
                e_or_w,
            )
            }
        }
    }
    fn filename(&self) -> String {
        let (latitude, longitude) = (self.latitude + 10, self.longitude);
        assert_eq!(latitude % 10, 0);
        assert_eq!(longitude % 10, 0);
        let n_or_s = if latitude >= 0 { 'N' } else { 'S' };
        let e_or_w = if longitude >= 0 { 'E' } else { 'W' };

        match self.kind {
            LandCoverKind::TreeCover => {
                format!("treecover/raw/{:02}{}_{:03}{}_treecover2010_v3.tif.zip",
                        latitude.abs(),
                        n_or_s,
                        longitude.abs(),
                        e_or_w,
                )
            }
            LandCoverKind::WaterMask => {
                format!("watermask/raw/Hansen_GFC2013_datamask_{:02}{}_{:03}{}.tif.zip",
                        latitude.abs(),
                        n_or_s,
                        longitude.abs(),
                        e_or_w,
                )
            }
        }
    }
    fn parse(
        &self,
        context: &mut AssetLoadContext,
        data: Vec<u8>,
    ) -> Result<Self::Type, Box<::std::error::Error>> {
        let mut zip = ZipArchive::new(Cursor::new(data))?;
        assert_eq!(zip.len(), 1);

        let mut data = Vec::new();
        zip.by_index(0)?.read_to_end(&mut data)?;

        let image = image::load_from_memory_with_format(&data[..], ImageFormat::TIFF)?;
        let image = Arc::new(Mutex::new(image));

        for latitude in self.latitude..(self.latitude + 9) {
            for longitude in self.longitude..(self.longitude + 9) {
                let params = LandCoverParams {
                    latitude,
                    longitude,
                    kind: self.kind,
                    raw: Some(image.clone()),
                };
                assert_eq!(params.raw_params(), *self);
                params.load(context)?;
            }
        }

        Ok(image)
    }
}

pub struct LandCoverParams {
    pub latitude: i16,
    pub longitude: i16,
    pub kind: LandCoverKind,
    pub raw: Option<Arc<Mutex<DynamicImage>>>,
}
impl LandCoverParams {
    fn raw_params(&self) -> RawLandCoverParams {
        RawLandCoverParams {
            latitude: self.latitude - ((self.latitude % 10) + 10) % 10,
            longitude: self.longitude - ((self.longitude % 10) + 10) % 10,
            kind: self.kind,
        }
    }

    fn generate_from_raw(&self) -> Result<Raster<u8>, Box<Error>> {
        let (w, h, image) = {
            let mut image = self.raw.as_ref().unwrap().lock().unwrap();
            let (w, h) = (image.width(), image.height());
            assert_eq!(w, h);
            let image = image
                .crop(
                    (w / 10) * ((self.longitude % 10 + 10) % 10) as u32,
                    (h / 10) * (9 - ((self.latitude % 10 + 10) % 10) as u32),
                    w / 10 + 1,
                    h / 10 + 1,
                )
                .clone();
            (w, h, image)
        };
        let image = image.rotate90().flipv();
        let values = if let ImageLuma8(image) = image {
            image.into_raw().into_iter()
        } else {
            unreachable!()
        };

        let values = match self.kind {
            LandCoverKind::TreeCover => values.collect(),
            LandCoverKind::WaterMask => {
                values
                    .map(|v| if v == 1 {
                        0
                    } else if v == 0 || v == 2 {
                        255
                    } else {
                        unreachable!()
                    })
                    .collect()
            }
        };

        Ok(Raster {
            width: w as usize / 10 + 1,
            height: h as usize / 10 + 1,
            cell_size: 10.0 / (w - 1) as f64,
            xllcorner: self.latitude as f64,
            yllcorner: self.longitude as f64,
            values,
        })
    }
}
impl GeneratedAsset for LandCoverParams {
    type Type = Raster<u8>;

    fn filename(&self) -> String {
        let n_or_s = if self.latitude >= 0 { 'N' } else { 'S' };
        let e_or_w = if self.longitude >= 0 { 'E' } else { 'W' };

        let directory = match self.kind {
            LandCoverKind::TreeCover => "treecover/processed",
            LandCoverKind::WaterMask => "watermask/processed",
        };

        format!("{}/{:02}{}_{:03}{}.raster",
                directory,
                self.latitude.abs(),
                n_or_s,
                self.longitude.abs(),
                e_or_w,
        )
    }

    fn generate(&self, context: &mut AssetLoadContext) -> Result<Self::Type, Box<Error>> {
        if self.raw.is_some() {
            return self.generate_from_raw();
        }

        Self {
            latitude: self.latitude,
            longitude: self.longitude,
            kind: self.kind,
            raw: Some(self.raw_params().load(context)?),
        }.generate_from_raw()
    }
}

pub struct BlueMarble;
impl WebAsset for BlueMarble {
    type Type = GlobalRaster<u8>;

    fn url(&self) -> String {
        "https://eoimages.gsfc.nasa.gov/images/imagerecords/74000/74218/\
         world.200412.3x21600x10800.png"
            .to_owned()
    }
    fn filename(&self) -> String {
        "bluemarble/world.200412.3x21600x10800.png".to_owned()
    }
    fn parse(
        &self,
        _context: &mut AssetLoadContext,
        data: Vec<u8>,
    ) -> Result<Self::Type, Box<::std::error::Error>> {
        let image = image::load_from_memory_with_format(&data, ImageFormat::PNG)?
            .to_rgb();
        Ok(GlobalRaster {
            width: image.width() as usize,
            height: image.height() as usize,
            bands: 3,
            values: image.into_raw(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn to_raw_params() {
        assert_eq!(
            LandCoverParams {
                latitude: 165,
                longitude: 31,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: 160,
                longitude: 30,
                kind: LandCoverKind::TreeCover,
            }
        );

        assert_eq!(
            LandCoverParams {
                latitude: 20,
                longitude: 20,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: 20,
                longitude: 20,
                kind: LandCoverKind::TreeCover,
            }
        );

        assert_eq!(
            LandCoverParams {
                latitude: -18,
                longitude: -18,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: -20,
                longitude: -20,
                kind: LandCoverKind::TreeCover,
            }
        );

        assert_eq!(
            LandCoverParams {
                latitude: -30,
                longitude: -30,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: -30,
                longitude: -30,
                kind: LandCoverKind::TreeCover,
            }
        );
    }
}
<commit_msg>Fix off by one in landcover file loading<commit_after>use std::error::Error;
use std::io::{Cursor, Read};
use std::sync::{Arc, Mutex};

use zip::ZipArchive;
use image::{self, DynamicImage, GenericImage, ImageLuma8, ImageFormat};

use cache::{AssetLoadContext, GeneratedAsset, WebAsset};
use terrain::raster::{GlobalRaster, Raster, RasterSource};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum LandCoverKind {
    TreeCover,
    WaterMask,
}
impl RasterSource for LandCoverKind {
    type Type = u8;
    fn load(
        &self,
        context: &mut AssetLoadContext,
        latitude: i16,
        longitude: i16,
    ) -> Option<Raster<u8>> {
        LandCoverParams {
            latitude,
            longitude,
            kind: *self,
            raw: None,
        }.load(context)
            .ok()
    }
}

/// Coordinates are of the lower left corner of the 10x10 degree cell.
#[derive(Debug, Eq, PartialEq)]
struct RawLandCoverParams {
    pub latitude: i16,
    pub longitude: i16,
    pub kind: LandCoverKind,
}
impl WebAsset for RawLandCoverParams {
    type Type = Arc<Mutex<DynamicImage>>;

    fn url(&self) -> String {
        let (latitude, longitude) = (self.latitude + 10, self.longitude);
        assert_eq!(latitude % 10, 0);
        assert_eq!(longitude % 10, 0);
        let n_or_s = if latitude >= 0 { 'N' } else { 'S' };
        let e_or_w = if longitude >= 0 { 'E' } else { 'W' };

        match self.kind {
            LandCoverKind::TreeCover => {
                format!("https://edcintl.cr.usgs.gov/downloads/sciweb1/shared/gtc/downloads/\
                         treecover2010_v3_individual/{:02}{}_{:03}{}_treecover2010_v3.tif.zip",
                latitude.abs(),
                n_or_s,
                longitude.abs(),
                e_or_w,
            )
            }
            LandCoverKind::WaterMask => {
                format!("https://edcintl.cr.usgs.gov/downloads/sciweb1/shared/gtc/downloads/\
                         WaterMask2010_UMD_individual/Hansen_GFC2013_datamask_{:02}{}_{:03}{}\
                         .tif.zip",
                latitude.abs(),
                n_or_s,
                longitude.abs(),
                e_or_w,
            )
            }
        }
    }
    fn filename(&self) -> String {
        let (latitude, longitude) = (self.latitude + 10, self.longitude);
        assert_eq!(latitude % 10, 0);
        assert_eq!(longitude % 10, 0);
        let n_or_s = if latitude >= 0 { 'N' } else { 'S' };
        let e_or_w = if longitude >= 0 { 'E' } else { 'W' };

        match self.kind {
            LandCoverKind::TreeCover => {
                format!("treecover/raw/{:02}{}_{:03}{}_treecover2010_v3.tif.zip",
                        latitude.abs(),
                        n_or_s,
                        longitude.abs(),
                        e_or_w,
                )
            }
            LandCoverKind::WaterMask => {
                format!("watermask/raw/Hansen_GFC2013_datamask_{:02}{}_{:03}{}.tif.zip",
                        latitude.abs(),
                        n_or_s,
                        longitude.abs(),
                        e_or_w,
                )
            }
        }
    }
    fn parse(
        &self,
        context: &mut AssetLoadContext,
        data: Vec<u8>,
    ) -> Result<Self::Type, Box<::std::error::Error>> {
        let mut zip = ZipArchive::new(Cursor::new(data))?;
        assert_eq!(zip.len(), 1);

        let mut data = Vec::new();
        zip.by_index(0)?.read_to_end(&mut data)?;

        let image = image::load_from_memory_with_format(&data[..], ImageFormat::TIFF)?;
        let image = Arc::new(Mutex::new(image));

        for latitude in self.latitude..(self.latitude + 10) {
            for longitude in self.longitude..(self.longitude + 10) {
                let params = LandCoverParams {
                    latitude,
                    longitude,
                    kind: self.kind,
                    raw: Some(image.clone()),
                };
                assert_eq!(params.raw_params(), *self);
                params.load(context)?;
            }
        }

        Ok(image)
    }
}

pub struct LandCoverParams {
    pub latitude: i16,
    pub longitude: i16,
    pub kind: LandCoverKind,
    pub raw: Option<Arc<Mutex<DynamicImage>>>,
}
impl LandCoverParams {
    fn raw_params(&self) -> RawLandCoverParams {
        RawLandCoverParams {
            latitude: self.latitude - ((self.latitude % 10) + 10) % 10,
            longitude: self.longitude - ((self.longitude % 10) + 10) % 10,
            kind: self.kind,
        }
    }

    fn generate_from_raw(&self) -> Result<Raster<u8>, Box<Error>> {
        let (w, h, image) = {
            let mut image = self.raw.as_ref().unwrap().lock().unwrap();
            let (w, h) = (image.width(), image.height());
            assert_eq!(w, h);
            let image = image
                .crop(
                    (w / 10) * ((self.longitude % 10 + 10) % 10) as u32,
                    (h / 10) * (9 - ((self.latitude % 10 + 10) % 10) as u32),
                    w / 10 + 1,
                    h / 10 + 1,
                )
                .clone();
            (w, h, image)
        };
        let image = image.rotate90().flipv();
        let values = if let ImageLuma8(image) = image {
            image.into_raw().into_iter()
        } else {
            unreachable!()
        };

        let values = match self.kind {
            LandCoverKind::TreeCover => values.collect(),
            LandCoverKind::WaterMask => {
                values
                    .map(|v| if v == 1 {
                        0
                    } else if v == 0 || v == 2 {
                        255
                    } else {
                        unreachable!()
                    })
                    .collect()
            }
        };

        Ok(Raster {
            width: w as usize / 10 + 1,
            height: h as usize / 10 + 1,
            cell_size: 10.0 / (w - 1) as f64,
            xllcorner: self.latitude as f64,
            yllcorner: self.longitude as f64,
            values,
        })
    }
}
impl GeneratedAsset for LandCoverParams {
    type Type = Raster<u8>;

    fn filename(&self) -> String {
        let n_or_s = if self.latitude >= 0 { 'N' } else { 'S' };
        let e_or_w = if self.longitude >= 0 { 'E' } else { 'W' };

        let directory = match self.kind {
            LandCoverKind::TreeCover => "treecover/processed",
            LandCoverKind::WaterMask => "watermask/processed",
        };

        format!("{}/{:02}{}_{:03}{}.raster",
                directory,
                self.latitude.abs(),
                n_or_s,
                self.longitude.abs(),
                e_or_w,
        )
    }

    fn generate(&self, context: &mut AssetLoadContext) -> Result<Self::Type, Box<Error>> {
        if self.raw.is_some() {
            return self.generate_from_raw();
        }

        Self {
            latitude: self.latitude,
            longitude: self.longitude,
            kind: self.kind,
            raw: Some(self.raw_params().load(context)?),
        }.generate_from_raw()
    }
}

pub struct BlueMarble;
impl WebAsset for BlueMarble {
    type Type = GlobalRaster<u8>;

    fn url(&self) -> String {
        "https://eoimages.gsfc.nasa.gov/images/imagerecords/74000/74218/\
         world.200412.3x21600x10800.png"
            .to_owned()
    }
    fn filename(&self) -> String {
        "bluemarble/world.200412.3x21600x10800.png".to_owned()
    }
    fn parse(
        &self,
        _context: &mut AssetLoadContext,
        data: Vec<u8>,
    ) -> Result<Self::Type, Box<::std::error::Error>> {
        let image = image::load_from_memory_with_format(&data, ImageFormat::PNG)?
            .to_rgb();
        Ok(GlobalRaster {
            width: image.width() as usize,
            height: image.height() as usize,
            bands: 3,
            values: image.into_raw(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn to_raw_params() {
        assert_eq!(
            LandCoverParams {
                latitude: 165,
                longitude: 31,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: 160,
                longitude: 30,
                kind: LandCoverKind::TreeCover,
            }
        );

        assert_eq!(
            LandCoverParams {
                latitude: 20,
                longitude: 20,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: 20,
                longitude: 20,
                kind: LandCoverKind::TreeCover,
            }
        );

        assert_eq!(
            LandCoverParams {
                latitude: -18,
                longitude: -18,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: -20,
                longitude: -20,
                kind: LandCoverKind::TreeCover,
            }
        );

        assert_eq!(
            LandCoverParams {
                latitude: -30,
                longitude: -30,
                kind: LandCoverKind::TreeCover,
            }.raw_params(),
            RawLandCoverParams {
                latitude: -30,
                longitude: -30,
                kind: LandCoverKind::TreeCover,
            }
        );
    }
}
<|endoftext|>


<commit_before>use std::{ env, fs };
use std::path::PathBuf;

#[derive(Clone)]
pub struct LighthouseConfig {
    pub data_dir: PathBuf,
    pub p2p_listen_port: u16,
}

const DEFAULT_LIGHTHOUSE_DIR: &str = ".lighthouse";

impl LighthouseConfig {
    /// Build a new lighthouse configuration from defaults.
    pub fn default() -> Self{
        let data_dir = {
            let home = env::home_dir()
                .expect("Unable to determine home dir.");
            home.join(DEFAULT_LIGHTHOUSE_DIR)
        };
        fs::create_dir_all(&data_dir).expect(&format!("Unable to create {:?}", &data_dir));
        let p2p_listen_port = 0;
        Self {
            data_dir,
            p2p_listen_port,
        }
    }
}
<commit_msg>Tidy config struct<commit_after>use std::{ env, fs };
use std::path::PathBuf;

/// Stores the core configuration for this Lighthouse instance.
/// This struct is general, other components may implement more
/// specialized config structs.
#[derive(Clone)]
pub struct LighthouseConfig {
    pub data_dir: PathBuf,
    pub p2p_listen_port: u16,
}

const DEFAULT_LIGHTHOUSE_DIR: &str = ".lighthouse";

impl LighthouseConfig {
    /// Build a new lighthouse configuration from defaults.
    pub fn default() -> Self{
        let data_dir = {
            let home = env::home_dir()
                .expect("Unable to determine home dir.");
            home.join(DEFAULT_LIGHTHOUSE_DIR)
        };
        fs::create_dir_all(&data_dir)
            .expect(&format!("Unable to create {:?}", &data_dir));
        let p2p_listen_port = 0;
        Self {
            data_dir,
            p2p_listen_port,
        }
    }
}
<|endoftext|>


<commit_before>extern crate gltf;

use gltf::json;
use std::{fs, io};

use std::boxed::Box;
use std::error::Error as StdError;

fn run(path: &str) -> Result<(), Box<StdError>> {
    use io::Read;
    let file = fs::File::open(&path)?;
    let mut data = Vec::with_capacity(file.metadata()?.len() as usize);
    let mut reader = io::BufReader::new(file);
    let _ = reader.read_to_end(&mut data)?;
    let json: json::Root = if gltf::is_binary(&data) {
        let glb = gltf::Glb::from_slice(&data)?;
        json::from_slice(glb.json)
    } else {
        json::from_slice(&data)
    }?;
    println!("{:#?}", json);
    Ok(())
}

fn main() {
    if let Some(path) = std::env::args().nth(1) {
        run(&path).expect("runtime error");
    } else {
        println!("usage: gltf-display <FILE>");
    }
}
<commit_msg>Fix example<commit_after>extern crate gltf;

use gltf::json;
use std::{fs, io};

use std::boxed::Box;
use std::error::Error as StdError;

fn run(path: &str) -> Result<(), Box<StdError>> {
    use io::Read;
    let file = fs::File::open(&path)?;
    let mut data = Vec::with_capacity(file.metadata()?.len() as usize);
    let mut reader = io::BufReader::new(file);
    let _ = reader.read_to_end(&mut data)?;
    let json: json::Root = if gltf::is_binary(&data) {
        let glb = gltf::Glb::from_slice(&data)?;
        json::from_slice(&glb.json)
    } else {
        json::from_slice(&data)
    }?;
    println!("{:#?}", json);
    Ok(())
}

fn main() {
    if let Some(path) = std::env::args().nth(1) {
        run(&path).expect("runtime error");
    } else {
        println!("usage: gltf-display <FILE>");
    }
}
<|endoftext|>