<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, ensure, Result};
use sha2::{Digest, Sha256};
use std::time::{Duration, Instant};
use thiserror::Error;
use zerocopy::AsBytes;

use crate::app::TransportWrapper;
use crate::bootstrap::{Bootstrap, BootstrapOptions, UpdateProtocol};
use crate::impl_serializable_error;
use crate::io::uart::Uart;
use crate::transport::Capability;

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 1024 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const MAGIC_HEADER: [u8; 4] = [0xfd, 0xff, 0xff, 0xff];
    const CRYPTOLIB_TELL: [u8; 4] = [0x53, 0x53, 0x53, 0x53];

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::digest(&frame[Frame::HASH_LEN..]);
        sha.into()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let mut digest = Sha256::digest(self.as_bytes());
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest.into()
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Result<Vec<Frame>> {
        // The given payload will contain up to four sections concatenated together:
        // RO_A, RW_A optionally follwed by RO_B, RW_B
        // Each section starts with a magic number on at least a 256 byte boundary.

        // This rescue protocol uses the RW_A section only, which will start at the second
        // occurrance of the magic value, and end at the third occurrence or at the end of the
        // file.

        ensure!(
            payload.starts_with(&Self::MAGIC_HEADER),
            RescueError::ImageFormatError
        );

        // Find second occurrence of magic value, not followed by signature of encrypted
        // cryptolib.
        let min_addr = match payload[256..]
            .chunks(256)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER && &c[4..8] != &Self::CRYPTOLIB_TELL)
        {
            Some(n) => (n + 1) * 256,
            None => bail!(RescueError::ImageFormatError),
        };

        // Find third occurrence of magic value.
        let max_addr = match payload[min_addr + 256..]
            .chunks(256)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER)
        {
            Some(n) => (n + 1) * 256 + min_addr,
            None => payload.len(),
        };

        // Trim trailing 0xff bytes.
        let max_addr = (payload[..max_addr]
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frames = Vec::new();
        let mut frame_num = 0;
        let mut addr = min_addr;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: addr as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
            .iter_mut()
            .for_each(|f| f.header.hash = f.header_hash());
        Ok(frames)
    }
}

#[derive(Debug, Error, serde::Serialize, serde::Deserialize)]
pub enum RescueError {
    #[error("Unrecognized image file format")]
    ImageFormatError,
    #[error("Synchronization error communicating with boot rom")]
    SyncError,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}
impl_serializable_error!(RescueError);

/// Implements the UART rescue protocol of Google Ti50 firmware.
pub struct Rescue {}

impl Rescue {
    /// Abort if a block has not been accepted after this number of retries.
    const MAX_CONSECUTIVE_ERRORS: u32 = 50;
    /// Take some measure to regain protocol synchronization, in case of this number of retries
    /// of the same block.
    const RESYNC_AFTER_CONSECUTIVE_ERRORS: u32 = 3;

    /// Creates a new `Rescue` protocol updater from `options`.
    pub fn new(_options: &BootstrapOptions) -> Self {
        Self {}
    }

    /// Waits for some time for a character, returns None on timeout.
    fn read_char(&self, uart: &dyn Uart) -> Option<char> {
        let mut buf = [0u8; 1];
        match uart.read_timeout(&mut buf, Duration::from_millis(100)) {
            Ok(1) => Some(buf[0] as char),
            Ok(_) => None,
            _ => None,
        }
    }

    /// Waits some time for data, returning true if the given string was seen in full, or false
    /// as soon as a non-matching character is received or on timeout.
    fn expect_string(&self, uart: &dyn Uart, s: &str) -> bool {
        for expected_ch in s.chars() {
            match self.read_char(uart) {
                Some(ch) if ch == expected_ch => (),
                _ => return false,
            }
        }
        true
    }

    /// Reads and discards any characters in the receive buffer, waiting a little while for any
    /// more which will also be discarded.
    fn flush_rx(&self, uart: &dyn Uart) {
        let mut response = [0u8; Frame::HASH_LEN];
        loop {
            match uart.read_timeout(&mut response, Duration::from_millis(500)) {
                Ok(0) | Err(_) => break,
                Ok(_) => continue,
            }
        }
    }

    /// As the 1024 byte blocks sent to the chip have no discernible header, the sender and
    /// receiver could be "out of sync".  This is resolved by sending one byte at a time, and
    /// observing when the chip sends a response (which will be a rejection due to checksum).
    fn synchronize(&self, uart: &dyn Uart) -> Result<()> {
        // Most likely, only a few "extra" bytes have been sent during initial negotiation.
        // Send almost a complete block in one go, and then send each of the last 16 bytes one
        // at a time, slowly enough to detect a response before sending the next byte.
        uart.write(&[0u8; 1008])?;
        let mut response = [0u8; 1];
        let limit = match uart.read_timeout(&mut response, Duration::from_millis(50)) {
            Ok(0) | Err(_) => 16,
            Ok(_) => {
                // A response at this point must mean that more than 16 bytes had already been
                // sent before entering this method.  This will be resolved by doing another
                // slower round of 1024 bytes with delay in between every one.
                self.flush_rx(uart);
                1024
            }
        };
        for _ in 0..limit {
            uart.write(&[0u8; 1])?;
            match uart.read_timeout(&mut response, Duration::from_millis(50)) {
                Ok(0) | Err(_) => (),
                Ok(_) => {
                    self.flush_rx(uart);
                    return Ok(());
                }
            }
        }
        Err(RescueError::SyncError.into())
    }

    /// Reset the chip and send the magic 'r' character at the opportune moment during boot in
    /// order to enter rescue more, repeat if necessary.
    fn enter_rescue_mode(&self, container: &Bootstrap, uart: &dyn Uart) -> Result<()> {
        // Attempt getting the attention of the bootloader.
        let timeout = Duration::from_millis(2000);
        for _ in 0..Self::MAX_CONSECUTIVE_ERRORS {
            eprint!("Resetting...");
            container.reset_pin.write(false)?; // Low active
            std::thread::sleep(container.reset_delay);
            container.reset_pin.write(true)?; // Release reset
            let stopwatch = Instant::now();
            while stopwatch.elapsed() < timeout {
                if !self.expect_string(uart, "Bldr |") {
                    continue;
                }
                uart.write(&['r' as u8])?;
                eprint!("a.");
                while stopwatch.elapsed() < timeout {
                    if !self.expect_string(uart, "oops?|") {
                        continue;
                    }
                    uart.write(&['r' as u8])?;
                    eprint!("b.");
                    if self.expect_string(uart, "escue") {
                        eprintln!("c: Entered rescue mode!");
                        self.synchronize(uart)?;
                        return Ok(());
                    }
                }
            }
            eprintln!(" Failed to enter rescue mode.");
        }
        Err(RescueError::RepeatedErrors.into())
    }
}

impl UpdateProtocol for Rescue {
    fn verify_capabilities(
        &self,
        _container: &Bootstrap,
        transport: &TransportWrapper,
    ) -> Result<()> {
        transport
            .capabilities()?
            .request(Capability::GPIO | Capability::UART)
            .ok()?;
        Ok(())
    }

    /// Returns false, in order to as the containing Bootstrap struct to not perform standard
    /// BOOTSTRAP/RESET sequence.
    fn uses_common_bootstrap_reset(&self) -> bool {
        false
    }

    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(
        &self,
        container: &Bootstrap,
        transport: &TransportWrapper,
        payload: &[u8],
        progress: &dyn Fn(u32, u32),
    ) -> Result<()> {
        let frames = Frame::from_payload(payload)?;
        let uart = container.uart_params.create(transport)?;

        self.enter_rescue_mode(container, &*uart)?;

        // Send frames one at a time.
        'next_block: for (idx, frame) in frames.iter().enumerate() {
            for consecutive_errors in 0..Self::MAX_CONSECUTIVE_ERRORS {
                eprint!("{}.", idx);
                progress(frame.header.flash_offset, frame.data.len() as u32);
                uart.write(frame.as_bytes())?;
                let mut response = [0u8; Frame::HASH_LEN];
                let mut index = 0;
                while index < Frame::HASH_LEN {
                    let timeout = if index == 0 {
                        Duration::from_millis(1000)
                    } else {
                        Duration::from_millis(10)
                    };
                    match uart.read_timeout(&mut response[index..], timeout) {
                        Ok(0) | Err(_) => break,
                        Ok(n) => index += n,
                    }
                }
                if index < Frame::HASH_LEN {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                    continue;
                }
                if response[4..].chunks(4).all(|x| x == &response[..4]) {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                } else if response == frame.frame_hash() {
                    continue 'next_block;
                } else {
                    self.flush_rx(&*uart);
                    if consecutive_errors >= Self::RESYNC_AFTER_CONSECUTIVE_ERRORS {
                        eprint!("sync.");
                        self.synchronize(&*uart)?;
                    }
                }
            }
            bail!(RescueError::RepeatedErrors);
        }

        // Reset, in order to leave rescue mode.
        container.reset_pin.write(false)?; // Low active
        std::thread::sleep(container.reset_delay);
        container.reset_pin.write(true)?; // Release reset
        eprintln!("Success!");
        Ok(())
    }
}
<commit_msg>[opentitantool]: Adapt heuristics for Ti50 images<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, ensure, Result};
use sha2::{Digest, Sha256};
use std::time::{Duration, Instant};
use thiserror::Error;
use zerocopy::AsBytes;

use crate::app::TransportWrapper;
use crate::bootstrap::{Bootstrap, BootstrapOptions, UpdateProtocol};
use crate::impl_serializable_error;
use crate::io::uart::Uart;
use crate::transport::Capability;

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 1024 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const HEADER_ALIGNMENT: usize = 0x1000;
    const MAGIC_HEADER: [u8; 4] = [0xfd, 0xff, 0xff, 0xff];
    const CRYPTOLIB_TELL: [u8; 4] = [0x53, 0x53, 0x53, 0x53];

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::digest(&frame[Frame::HASH_LEN..]);
        sha.into()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let mut digest = Sha256::digest(self.as_bytes());
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest.into()
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Result<Vec<Frame>> {
        // The given payload will contain up to four sections concatenated together:
        // RO_A, RW_A optionally follwed by RO_B, RW_B
        // Each section starts with a magic number on at least a 256 byte boundary.

        // This rescue protocol uses the RW_A section only, which will start at the second
        // occurrance of the magic value, and end at the third occurrence or at the end of the
        // file.

        ensure!(
            payload.starts_with(&Self::MAGIC_HEADER),
            RescueError::ImageFormatError
        );

        // Find second occurrence of magic value, not followed by signature of encrypted
        // cryptolib.
        let min_addr = match payload[Self::HEADER_ALIGNMENT..]
            .chunks(Self::HEADER_ALIGNMENT)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER && &c[4..8] != &Self::CRYPTOLIB_TELL)
        {
            Some(n) => (n + 1) * Self::HEADER_ALIGNMENT,
            None => bail!(RescueError::ImageFormatError),
        };

        // Find third occurrence of magic value.
        let max_addr = match payload[min_addr + Self::HEADER_ALIGNMENT..]
            .chunks(Self::HEADER_ALIGNMENT)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER)
        {
            Some(n) => (n + 1) * Self::HEADER_ALIGNMENT + min_addr,
            None => payload.len(),
        };

        // Trim trailing 0xff bytes.
        let max_addr = (payload[..max_addr]
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frames = Vec::new();
        let mut frame_num = 0;
        let mut addr = min_addr;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: addr as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
            .iter_mut()
            .for_each(|f| f.header.hash = f.header_hash());
        Ok(frames)
    }
}

#[derive(Debug, Error, serde::Serialize, serde::Deserialize)]
pub enum RescueError {
    #[error("Unrecognized image file format")]
    ImageFormatError,
    #[error("Synchronization error communicating with boot rom")]
    SyncError,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}
impl_serializable_error!(RescueError);

/// Implements the UART rescue protocol of Google Ti50 firmware.
pub struct Rescue {}

impl Rescue {
    /// Abort if a block has not been accepted after this number of retries.
    const MAX_CONSECUTIVE_ERRORS: u32 = 50;
    /// Take some measure to regain protocol synchronization, in case of this number of retries
    /// of the same block.
    const RESYNC_AFTER_CONSECUTIVE_ERRORS: u32 = 3;

    /// Creates a new `Rescue` protocol updater from `options`.
    pub fn new(_options: &BootstrapOptions) -> Self {
        Self {}
    }

    /// Waits for some time for a character, returns None on timeout.
    fn read_char(&self, uart: &dyn Uart) -> Option<char> {
        let mut buf = [0u8; 1];
        match uart.read_timeout(&mut buf, Duration::from_millis(100)) {
            Ok(1) => Some(buf[0] as char),
            Ok(_) => None,
            _ => None,
        }
    }

    /// Waits some time for data, returning true if the given string was seen in full, or false
    /// as soon as a non-matching character is received or on timeout.
    fn expect_string(&self, uart: &dyn Uart, s: &str) -> bool {
        for expected_ch in s.chars() {
            match self.read_char(uart) {
                Some(ch) if ch == expected_ch => (),
                _ => return false,
            }
        }
        true
    }

    /// Reads and discards any characters in the receive buffer, waiting a little while for any
    /// more which will also be discarded.
    fn flush_rx(&self, uart: &dyn Uart) {
        let mut response = [0u8; Frame::HASH_LEN];
        loop {
            match uart.read_timeout(&mut response, Duration::from_millis(500)) {
                Ok(0) | Err(_) => break,
                Ok(_) => continue,
            }
        }
    }

    /// As the 1024 byte blocks sent to the chip have no discernible header, the sender and
    /// receiver could be "out of sync".  This is resolved by sending one byte at a time, and
    /// observing when the chip sends a response (which will be a rejection due to checksum).
    fn synchronize(&self, uart: &dyn Uart) -> Result<()> {
        // Most likely, only a few "extra" bytes have been sent during initial negotiation.
        // Send almost a complete block in one go, and then send each of the last 16 bytes one
        // at a time, slowly enough to detect a response before sending the next byte.
        uart.write(&[0u8; 1008])?;
        let mut response = [0u8; 1];
        let limit = match uart.read_timeout(&mut response, Duration::from_millis(50)) {
            Ok(0) | Err(_) => 16,
            Ok(_) => {
                // A response at this point must mean that more than 16 bytes had already been
                // sent before entering this method.  This will be resolved by doing another
                // slower round of 1024 bytes with delay in between every one.
                self.flush_rx(uart);
                1024
            }
        };
        for _ in 0..limit {
            uart.write(&[0u8; 1])?;
            match uart.read_timeout(&mut response, Duration::from_millis(50)) {
                Ok(0) | Err(_) => (),
                Ok(_) => {
                    self.flush_rx(uart);
                    return Ok(());
                }
            }
        }
        Err(RescueError::SyncError.into())
    }

    /// Reset the chip and send the magic 'r' character at the opportune moment during boot in
    /// order to enter rescue more, repeat if necessary.
    fn enter_rescue_mode(&self, container: &Bootstrap, uart: &dyn Uart) -> Result<()> {
        // Attempt getting the attention of the bootloader.
        let timeout = Duration::from_millis(2000);
        for _ in 0..Self::MAX_CONSECUTIVE_ERRORS {
            eprint!("Resetting...");
            container.reset_pin.write(false)?; // Low active
            std::thread::sleep(container.reset_delay);
            container.reset_pin.write(true)?; // Release reset
            let stopwatch = Instant::now();
            while stopwatch.elapsed() < timeout {
                if !self.expect_string(uart, "Bldr |") {
                    continue;
                }
                uart.write(&['r' as u8])?;
                eprint!("a.");
                while stopwatch.elapsed() < timeout {
                    if !self.expect_string(uart, "oops?|") {
                        continue;
                    }
                    uart.write(&['r' as u8])?;
                    eprint!("b.");
                    if self.expect_string(uart, "escue") {
                        eprintln!("c: Entered rescue mode!");
                        self.synchronize(uart)?;
                        return Ok(());
                    }
                }
            }
            eprintln!(" Failed to enter rescue mode.");
        }
        Err(RescueError::RepeatedErrors.into())
    }
}

impl UpdateProtocol for Rescue {
    fn verify_capabilities(
        &self,
        _container: &Bootstrap,
        transport: &TransportWrapper,
    ) -> Result<()> {
        transport
            .capabilities()?
            .request(Capability::GPIO | Capability::UART)
            .ok()?;
        Ok(())
    }

    /// Returns false, in order to as the containing Bootstrap struct to not perform standard
    /// BOOTSTRAP/RESET sequence.
    fn uses_common_bootstrap_reset(&self) -> bool {
        false
    }

    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(
        &self,
        container: &Bootstrap,
        transport: &TransportWrapper,
        payload: &[u8],
        progress: &dyn Fn(u32, u32),
    ) -> Result<()> {
        let frames = Frame::from_payload(payload)?;
        let uart = container.uart_params.create(transport)?;

        self.enter_rescue_mode(container, &*uart)?;

        // Send frames one at a time.
        'next_block: for (idx, frame) in frames.iter().enumerate() {
            for consecutive_errors in 0..Self::MAX_CONSECUTIVE_ERRORS {
                eprint!("{}.", idx);
                progress(frame.header.flash_offset, frame.data.len() as u32);
                uart.write(frame.as_bytes())?;
                let mut response = [0u8; Frame::HASH_LEN];
                let mut index = 0;
                while index < Frame::HASH_LEN {
                    let timeout = if index == 0 {
                        Duration::from_millis(1000)
                    } else {
                        Duration::from_millis(10)
                    };
                    match uart.read_timeout(&mut response[index..], timeout) {
                        Ok(0) | Err(_) => break,
                        Ok(n) => index += n,
                    }
                }
                if index < Frame::HASH_LEN {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                    continue;
                }
                if response[4..].chunks(4).all(|x| x == &response[..4]) {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                } else if response == frame.frame_hash() {
                    continue 'next_block;
                } else {
                    self.flush_rx(&*uart);
                    if consecutive_errors >= Self::RESYNC_AFTER_CONSECUTIVE_ERRORS {
                        eprint!("sync.");
                        self.synchronize(&*uart)?;
                    }
                }
            }
            bail!(RescueError::RepeatedErrors);
        }

        // Reset, in order to leave rescue mode.
        container.reset_pin.write(false)?; // Low active
        std::thread::sleep(container.reset_delay);
        container.reset_pin.write(true)?; // Release reset
        eprintln!("Success!");
        Ok(())
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use document_loader::DocumentLoader;
use dom::bindings::codegen::Bindings::DOMImplementationBinding;
use dom::bindings::codegen::Bindings::DOMImplementationBinding::DOMImplementationMethods;
use dom::bindings::codegen::Bindings::DocumentBinding::{DocumentMethods, ElementCreationOptions};
use dom::bindings::codegen::Bindings::NodeBinding::NodeMethods;
use dom::bindings::error::Fallible;
use dom::bindings::inheritance::Castable;
use dom::bindings::reflector::{Reflector, reflect_dom_object};
use dom::bindings::root::{Dom, DomRoot};
use dom::bindings::str::DOMString;
use dom::bindings::xmlname::{namespace_from_domstring, validate_qualified_name};
use dom::document::{Document, HasBrowsingContext, IsHTMLDocument};
use dom::document::DocumentSource;
use dom::documenttype::DocumentType;
use dom::htmlbodyelement::HTMLBodyElement;
use dom::htmlheadelement::HTMLHeadElement;
use dom::htmlhtmlelement::HTMLHtmlElement;
use dom::htmltitleelement::HTMLTitleElement;
use dom::node::Node;
use dom::text::Text;
use dom::xmldocument::XMLDocument;
use dom_struct::dom_struct;
use script_traits::DocumentActivity;

// https://dom.spec.whatwg.org/#domimplementation
#[dom_struct]
pub struct DOMImplementation {
    reflector_: Reflector,
    document: Dom<Document>,
}

impl DOMImplementation {
    fn new_inherited(document: &Document) -> DOMImplementation {
        DOMImplementation {
            reflector_: Reflector::new(),
            document: Dom::from_ref(document),
        }
    }

    pub fn new(document: &Document) -> DomRoot<DOMImplementation> {
        let window = document.window();
        reflect_dom_object(Box::new(DOMImplementation::new_inherited(document)),
                           window,
                           DOMImplementationBinding::Wrap)
    }
}

// https://dom.spec.whatwg.org/#domimplementation
impl DOMImplementationMethods for DOMImplementation {
    // https://dom.spec.whatwg.org/#dom-domimplementation-createdocumenttype
    fn CreateDocumentType(&self,
                          qualified_name: DOMString,
                          pubid: DOMString,
                          sysid: DOMString)
                          -> Fallible<DomRoot<DocumentType>> {
        validate_qualified_name(&qualified_name)?;
        Ok(DocumentType::new(qualified_name, Some(pubid), Some(sysid), &self.document))
    }

    // https://dom.spec.whatwg.org/#dom-domimplementation-createdocument
    fn CreateDocument(&self,
                      maybe_namespace: Option<DOMString>,
                      qname: DOMString,
                      maybe_doctype: Option<&DocumentType>)
                      -> Fallible<DomRoot<XMLDocument>> {
        let win = self.document.window();
        let loader = DocumentLoader::new(&self.document.loader());
        let namespace = namespace_from_domstring(maybe_namespace.to_owned());

        let content_type = match namespace {
            ns!(html) => "application/xhtml+xml",
            ns!(svg) => "image/svg+xml",
            _ => "application/xml"
        };

        // Step 1.
        let doc = XMLDocument::new(win,
                                   HasBrowsingContext::No,
                                   None,
                                   self.document.origin().clone(),
                                   IsHTMLDocument::NonHTMLDocument,
                                   Some(DOMString::from(content_type)),
                                   None,
                                   DocumentActivity::Inactive,
                                   DocumentSource::NotFromParser,
                                   loader);
        // Step 2-3.
        let maybe_elem = if qname.is_empty() {
            None
        } else {
            let options = ElementCreationOptions { is: None };
            match doc.upcast::<Document>().CreateElementNS(maybe_namespace, qname, &options) {
                Err(error) => return Err(error),
                Ok(elem) => Some(elem),
            }
        };

        {
            let doc_node = doc.upcast::<Node>();

            // Step 4.
            if let Some(doc_type) = maybe_doctype {
                doc_node.AppendChild(doc_type.upcast()).unwrap();
            }

            // Step 5.
            if let Some(ref elem) = maybe_elem {
                doc_node.AppendChild(elem.upcast()).unwrap();
            }
        }

        // Step 6.
        // FIXME: https://github.com/mozilla/servo/issues/1522

        // Step 7.
        Ok(doc)
    }

    // https://dom.spec.whatwg.org/#dom-domimplementation-createhtmldocument
    fn CreateHTMLDocument(&self, title: Option<DOMString>) -> DomRoot<Document> {
        let win = self.document.window();
        let loader = DocumentLoader::new(&self.document.loader());

        // Step 1-2.
        let doc = Document::new(win,
                                HasBrowsingContext::No,
                                None,
                                self.document.origin().clone(),
                                IsHTMLDocument::HTMLDocument,
                                None,
                                None,
                                DocumentActivity::Inactive,
                                DocumentSource::NotFromParser,
                                loader,
                                None,
                                None,
                                Default::default());

        {
            // Step 3.
            let doc_node = doc.upcast::<Node>();
            let doc_type = DocumentType::new(DOMString::from("html"), None, None, &doc);
            doc_node.AppendChild(doc_type.upcast()).unwrap();
        }

        {
            // Step 4.
            let doc_node = doc.upcast::<Node>();
            let doc_html = DomRoot::upcast::<Node>(HTMLHtmlElement::new(local_name!("html"),
                                                                     None,
                                                                     &doc));
            doc_node.AppendChild(&doc_html).expect("Appending failed");

            {
                // Step 5.
                let doc_head = DomRoot::upcast::<Node>(HTMLHeadElement::new(local_name!("head"),
                                                                         None,
                                                                         &doc));
                doc_html.AppendChild(&doc_head).unwrap();

                // Step 6.
                if let Some(title_str) = title {
                    // Step 6.1.
                    let doc_title =
                        DomRoot::upcast::<Node>(HTMLTitleElement::new(local_name!("title"),
                                                                   None,
                                                                   &doc));
                    doc_head.AppendChild(&doc_title).unwrap();

                    // Step 6.2.
                    let title_text = Text::new(title_str, &doc);
                    doc_title.AppendChild(title_text.upcast()).unwrap();
                }
            }

            // Step 7.
            let doc_body = HTMLBodyElement::new(local_name!("body"), None, &doc);
            doc_html.AppendChild(doc_body.upcast()).unwrap();
        }

        // Step 8.
        // FIXME: https://github.com/mozilla/servo/issues/1522

        // Step 9.
        doc
    }

    // https://dom.spec.whatwg.org/#dom-domimplementation-hasfeature
    fn HasFeature(&self) -> bool {
        true
    }
}
<commit_msg>Remove out of date comments from domimplementation<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use document_loader::DocumentLoader;
use dom::bindings::codegen::Bindings::DOMImplementationBinding;
use dom::bindings::codegen::Bindings::DOMImplementationBinding::DOMImplementationMethods;
use dom::bindings::codegen::Bindings::DocumentBinding::{DocumentMethods, ElementCreationOptions};
use dom::bindings::codegen::Bindings::NodeBinding::NodeMethods;
use dom::bindings::error::Fallible;
use dom::bindings::inheritance::Castable;
use dom::bindings::reflector::{Reflector, reflect_dom_object};
use dom::bindings::root::{Dom, DomRoot};
use dom::bindings::str::DOMString;
use dom::bindings::xmlname::{namespace_from_domstring, validate_qualified_name};
use dom::document::{Document, HasBrowsingContext, IsHTMLDocument};
use dom::document::DocumentSource;
use dom::documenttype::DocumentType;
use dom::htmlbodyelement::HTMLBodyElement;
use dom::htmlheadelement::HTMLHeadElement;
use dom::htmlhtmlelement::HTMLHtmlElement;
use dom::htmltitleelement::HTMLTitleElement;
use dom::node::Node;
use dom::text::Text;
use dom::xmldocument::XMLDocument;
use dom_struct::dom_struct;
use script_traits::DocumentActivity;

// https://dom.spec.whatwg.org/#domimplementation
#[dom_struct]
pub struct DOMImplementation {
    reflector_: Reflector,
    document: Dom<Document>,
}

impl DOMImplementation {
    fn new_inherited(document: &Document) -> DOMImplementation {
        DOMImplementation {
            reflector_: Reflector::new(),
            document: Dom::from_ref(document),
        }
    }

    pub fn new(document: &Document) -> DomRoot<DOMImplementation> {
        let window = document.window();
        reflect_dom_object(Box::new(DOMImplementation::new_inherited(document)),
                           window,
                           DOMImplementationBinding::Wrap)
    }
}

// https://dom.spec.whatwg.org/#domimplementation
impl DOMImplementationMethods for DOMImplementation {
    // https://dom.spec.whatwg.org/#dom-domimplementation-createdocumenttype
    fn CreateDocumentType(&self,
                          qualified_name: DOMString,
                          pubid: DOMString,
                          sysid: DOMString)
                          -> Fallible<DomRoot<DocumentType>> {
        validate_qualified_name(&qualified_name)?;
        Ok(DocumentType::new(qualified_name, Some(pubid), Some(sysid), &self.document))
    }

    // https://dom.spec.whatwg.org/#dom-domimplementation-createdocument
    fn CreateDocument(&self,
                      maybe_namespace: Option<DOMString>,
                      qname: DOMString,
                      maybe_doctype: Option<&DocumentType>)
                      -> Fallible<DomRoot<XMLDocument>> {
        let win = self.document.window();
        let loader = DocumentLoader::new(&self.document.loader());
        let namespace = namespace_from_domstring(maybe_namespace.to_owned());

        let content_type = match namespace {
            ns!(html) => "application/xhtml+xml",
            ns!(svg) => "image/svg+xml",
            _ => "application/xml"
        };

        // Step 1.
        let doc = XMLDocument::new(win,
                                   HasBrowsingContext::No,
                                   None,
                                   self.document.origin().clone(),
                                   IsHTMLDocument::NonHTMLDocument,
                                   Some(DOMString::from(content_type)),
                                   None,
                                   DocumentActivity::Inactive,
                                   DocumentSource::NotFromParser,
                                   loader);
        // Step 2-3.
        let maybe_elem = if qname.is_empty() {
            None
        } else {
            let options = ElementCreationOptions { is: None };
            match doc.upcast::<Document>().CreateElementNS(maybe_namespace, qname, &options) {
                Err(error) => return Err(error),
                Ok(elem) => Some(elem),
            }
        };

        {
            let doc_node = doc.upcast::<Node>();

            // Step 4.
            if let Some(doc_type) = maybe_doctype {
                doc_node.AppendChild(doc_type.upcast()).unwrap();
            }

            // Step 5.
            if let Some(ref elem) = maybe_elem {
                doc_node.AppendChild(elem.upcast()).unwrap();
            }
        }

        // Step 6.
        // The origin is already set

        // Step 7.
        Ok(doc)
    }

    // https://dom.spec.whatwg.org/#dom-domimplementation-createhtmldocument
    fn CreateHTMLDocument(&self, title: Option<DOMString>) -> DomRoot<Document> {
        let win = self.document.window();
        let loader = DocumentLoader::new(&self.document.loader());

        // Step 1-2.
        let doc = Document::new(win,
                                HasBrowsingContext::No,
                                None,
                                self.document.origin().clone(),
                                IsHTMLDocument::HTMLDocument,
                                None,
                                None,
                                DocumentActivity::Inactive,
                                DocumentSource::NotFromParser,
                                loader,
                                None,
                                None,
                                Default::default());

        {
            // Step 3.
            let doc_node = doc.upcast::<Node>();
            let doc_type = DocumentType::new(DOMString::from("html"), None, None, &doc);
            doc_node.AppendChild(doc_type.upcast()).unwrap();
        }

        {
            // Step 4.
            let doc_node = doc.upcast::<Node>();
            let doc_html = DomRoot::upcast::<Node>(HTMLHtmlElement::new(local_name!("html"),
                                                                     None,
                                                                     &doc));
            doc_node.AppendChild(&doc_html).expect("Appending failed");

            {
                // Step 5.
                let doc_head = DomRoot::upcast::<Node>(HTMLHeadElement::new(local_name!("head"),
                                                                         None,
                                                                         &doc));
                doc_html.AppendChild(&doc_head).unwrap();

                // Step 6.
                if let Some(title_str) = title {
                    // Step 6.1.
                    let doc_title =
                        DomRoot::upcast::<Node>(HTMLTitleElement::new(local_name!("title"),
                                                                   None,
                                                                   &doc));
                    doc_head.AppendChild(&doc_title).unwrap();

                    // Step 6.2.
                    let title_text = Text::new(title_str, &doc);
                    doc_title.AppendChild(title_text.upcast()).unwrap();
                }
            }

            // Step 7.
            let doc_body = HTMLBodyElement::new(local_name!("body"), None, &doc);
            doc_html.AppendChild(doc_body.upcast()).unwrap();
        }

        // Step 8.
        // The origin is already set

        // Step 9.
        doc
    }

    // https://dom.spec.whatwg.org/#dom-domimplementation-hasfeature
    fn HasFeature(&self) -> bool {
        true
    }
}
<|endoftext|>


<commit_before>extern crate bencode;

use self::bencode::{FromBencode, Bencode, StringFromBencodeError};
use self::bencode::util::ByteString;
use std::fs::File;
use std::io::Read;
use std::{convert, io};

macro_rules! get_field_with_default {
    ($m:expr, $field:expr, $default:expr) => (
        match $m.get(&ByteString::from_str($field)) {
            Some(a) => try!(FromBencode::from_bencode(a)),
            None => $default
        };
    )
}

macro_rules! get_field {
    ($m:expr, $field:expr) => (
        get_field_with_default!($m, $field, return Err(MetainfoError::DoesntContain($field)))
    )
}

#[derive(PartialEq, Debug)]
struct Metainfo {
    announce: String,
    created_by: String,
}

impl FromBencode for Metainfo {
    type Err = MetainfoError;

    fn from_bencode(bencode: &bencode::Bencode) -> Result<Metainfo, MetainfoError> {
        match bencode {
            &Bencode::Dict(ref m) => {
                let metainfo = Metainfo{
                    announce: get_field!(m, "announce"),
                    created_by: get_field_with_default!(m, "created by", "".to_string()),
                };
                Ok(metainfo)
            }
            _ => Err(MetainfoError::NotADict)
        }
    }
}

#[derive(Debug)]
enum MetainfoError {
    IoError(io::Error),
    DecodingError(bencode::streaming::Error),
    NotADict,
    DoesntContain(&'static str),
    NotAString(StringFromBencodeError),
}

impl convert::From<io::Error> for MetainfoError {
    fn from(err: io::Error) -> MetainfoError {
        MetainfoError::IoError(err)
    }
}

impl convert::From<bencode::streaming::Error> for MetainfoError {
    fn from(err: bencode::streaming::Error) -> MetainfoError {
        MetainfoError::DecodingError(err)
    }
}

impl convert::From<StringFromBencodeError> for MetainfoError {
    fn from(err: StringFromBencodeError) -> MetainfoError {
        MetainfoError::NotAString(err)
    }
}

pub fn run() {
    let result = parse_torrent_file("test_data/flagfromserver.torrent");
    match result {
        Ok(s)  => println!("Yay, it worked: {:?}", s),
        Err(e) => println!("Oops, it failed: {:?}", e)
    }
}

fn parse_torrent_file(filename: &str) -> Result<Metainfo, MetainfoError> {
    // read the torrent file into a byte vector
    let mut f = try!(File::open(filename));
    let mut v = Vec::new();
    try!(f.read_to_end(&mut v));

    // decode the byte vector into a struct
    let bencode = try!(bencode::from_vec(v));
    let result = try!(FromBencode::from_bencode(&bencode));

    Ok(result)
}
<commit_msg>Added unpacking of info struct.<commit_after>extern crate bencode;

use self::bencode::{FromBencode, Bencode, NumFromBencodeError, StringFromBencodeError};
use self::bencode::util::ByteString;
use std::fs::File;
use std::io::Read;
use std::{convert, io};

macro_rules! get_field_with_default {
    ($m:expr, $field:expr, $default:expr) => (
        match $m.get(&ByteString::from_str($field)) {
            Some(a) => try!(FromBencode::from_bencode(a)),
            None => $default
        };
    )
}

macro_rules! get_field {
    ($m:expr, $field:expr) => (
        get_field_with_default!($m, $field, return Err(MetainfoError::DoesntContain($field)))
    )
}

#[derive(PartialEq, Debug)]
struct Metainfo {
    announce: String,
    info: Info,
    created_by: String,
}

impl FromBencode for Metainfo {
    type Err = MetainfoError;

    fn from_bencode(bencode: &bencode::Bencode) -> Result<Metainfo, MetainfoError> {
        match bencode {
            &Bencode::Dict(ref m) => {
                let metainfo = Metainfo{
                    announce: get_field!(m, "announce"),
                    info: get_field!(m, "info"),
                    created_by: get_field_with_default!(m, "created by", "".to_string()),
                };
                Ok(metainfo)
            }
            _ => Err(MetainfoError::NotADict)
        }
    }
}

#[derive(PartialEq, Debug)]
struct Info {
    piece_length: u32,
    name: String,
}

impl FromBencode for Info {
    type Err = MetainfoError;

    fn from_bencode(bencode: &bencode::Bencode) -> Result<Info, MetainfoError> {
        match bencode {
            &Bencode::Dict(ref m) => {
                let info = Info{
                    piece_length: get_field!(m, "piece length"),
                    name: get_field!(m, "name"),
                };
                Ok(info)
            }
            _ => Err(MetainfoError::NotADict)
        }
    }
}

#[derive(Debug)]
enum MetainfoError {
    IoError(io::Error),
    DecodingError(bencode::streaming::Error),
    NotADict,
    DoesntContain(&'static str),
    NotANumber(NumFromBencodeError),
    NotAString(StringFromBencodeError),
}

impl convert::From<io::Error> for MetainfoError {
    fn from(err: io::Error) -> MetainfoError {
        MetainfoError::IoError(err)
    }
}

impl convert::From<bencode::streaming::Error> for MetainfoError {
    fn from(err: bencode::streaming::Error) -> MetainfoError {
        MetainfoError::DecodingError(err)
    }
}

impl convert::From<NumFromBencodeError> for MetainfoError {
    fn from(err: NumFromBencodeError) -> MetainfoError {
        MetainfoError::NotANumber(err)
    }
}

impl convert::From<StringFromBencodeError> for MetainfoError {
    fn from(err: StringFromBencodeError) -> MetainfoError {
        MetainfoError::NotAString(err)
    }
}

pub fn run() {
    let result = parse_torrent_file("test_data/flagfromserver.torrent");
    match result {
        Ok(s)  => println!("Yay, it worked: {:?}", s),
        Err(e) => println!("Oops, it failed: {:?}", e)
    }
}

fn parse_torrent_file(filename: &str) -> Result<Metainfo, MetainfoError> {
    // read the torrent file into a byte vector
    let mut f = try!(File::open(filename));
    let mut v = Vec::new();
    try!(f.read_to_end(&mut v));

    // decode the byte vector into a struct
    let bencode = try!(bencode::from_vec(v));
    let result = try!(FromBencode::from_bencode(&bencode));

    Ok(result)
}
<|endoftext|>


<commit_before>//! A Netherrack buffer backed by a VecDeque::<u8>

use std::collections::VecDeque;

extern crate bit_utils;
use self::bit_utils::BitInformation;

/// Defines a DequeBuffer
///
/// At some stage in the future, this will be changed to inherit the trait NetherrackBuffer<(Something extending Read and Write), but for now, we'll do everything here
pub struct DequeBuffer {

    data: VecDeque<u8>

}

impl DequeBuffer {

    /// Creates a DequeBuffer from a VecDeque
    ///
    /// After this method, we own the VecDeque
    pub fn from(deque: VecDeque<u8>) -> DequeBuffer {
        DequeBuffer { data: deque }
    }

    /// Creates a new empty DequeBuffer
    pub fn new() -> DequeBuffer {
        DequeBuffer { data: VecDeque::<u8>::new() }
    }

    /// Creates a new empty DequeBuffer with the specified initial capacity
    pub fn with_capacity(capacity: usize) -> DequeBuffer {
        DequeBuffer { data: VecDeque::<u8>::with_capacity(capacity) }
    }

    /// Gets the remaining bytes in this DequeBuffer
    pub fn remaining(&self) -> usize {
        self.data.len()
    }

    /// Reads a UTF-8 string from the buffer
    pub fn read_utf8_string(&mut self) -> String {
        let size_result = self.read_unsigned_varint_32();

        if size_result.is_err() {
            return String::from("!!!BAD_SIZE!!!");
        } else {
            let mut size: u32 = size_result.unwrap();

            let mut bytes: Vec<u8> = Vec::<u8>::with_capacity(size as usize);

            while size > 0 {
                bytes.push(self.read_unsigned_byte());
                size -= 1;
            }

            let result = String::from_utf8(bytes);

            if result.is_ok() {
                return result.unwrap();
            } else {
                return String::from("!!!INVALID_UTF8!!!");
            }

        }
    }

    /// Reads an unsigned byte from the buffer, returning 0 if no data is available
    pub fn read_unsigned_byte(&mut self) -> u8 {
        match self.data.pop_front() {
            Some(value) => value,
            None => 0
        }
    }

    /// Reads an unsigned byte from the buffer, returning -127 if no data is available
    pub fn read_signed_byte(&mut self) -> i8 {
        match self.data.pop_front() {
            Some(value) => value as i8,
            None => -127
        }
    }

    /// Reads an unsigned short from the buffer
    pub fn read_unsigned_short(&mut self) -> u16 {

        let high: u8 = self.read_unsigned_byte();
        let low: u8 = self.read_unsigned_byte();

        ((high as u16) << 8) | low as u16

    }

    /// Writes an unsigned byte to the buffer
    pub fn write_unsigned_byte(&mut self, value: u8) {
        self.data.push_back(value);
    }

    /// Writes a signed byte to the buffer
    pub fn write_signed_byte(&mut self, value: i8) {
        self.data.push_back(value as u8);
    }

    /// Reads an unsigned Varint from the buffer
    pub fn read_unsigned_varint_32(&mut self) -> Result<u32, &'static str> {
        if self.remaining() == 0 {
            return Err("Cannot read a Varint with no data to read.");
        } else {
            // The number of bits to shift by. <<0, <<7, <<14, etc
            let mut shift_amount: u32 = 0;

            // The decoded value
            let mut decoded_value: u32 = 0;

            while self.remaining() >= 1 {

                let byte_value: u8 = self.read_unsigned_byte(); //Yay for this existing!

                decoded_value |= ((byte_value & 0b01111111) as u32) << shift_amount; // We're not shifting for the first byte ;)

                // See if we're still reading bytes
                if byte_value.has_most_signifigant_bit() {
                    shift_amount += 7;
                } else {
                    return Ok(decoded_value);
                }

            }

            return Err("Reached end of buffer while reading a varint32");

        }
    }

}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn new_contains_no_elements() {
        let new = DequeBuffer::new();

        assert_eq!(0, new.remaining());
    }

    #[test]
    fn new_with_capacity_contains_no_elements() {
        let new = DequeBuffer::with_capacity(25000); //WOW, MUCH BIG PACKET

        assert_eq!(0, new.remaining());
    }

    #[test]
    fn unsigned_byte() {
        let mut buffer = DequeBuffer::with_capacity(2);

        buffer.write_unsigned_byte(1);
        buffer.write_unsigned_byte(15);

        assert_eq!(2, buffer.remaining());

        assert_eq!(1, buffer.read_unsigned_byte());
        assert_eq!(15, buffer.read_unsigned_byte());
    }

    #[test]
    fn signed_byte() {
        let mut buffer = DequeBuffer::with_capacity(2);

        buffer.write_signed_byte(-50);
        buffer.write_signed_byte(0);

        assert_eq!(2, buffer.remaining());

        assert_eq!(-50, buffer.read_signed_byte());
        assert_eq!(-0, buffer.read_signed_byte());
    }

}<commit_msg>DequeBuffer: Add write_unsigned_varint_32, and test read/write varint32<commit_after>//! A Netherrack buffer backed by a VecDeque::<u8>

use std::collections::VecDeque;

extern crate bit_utils;
use self::bit_utils::BitInformation;

/// Defines a DequeBuffer
///
/// At some stage in the future, this will be changed to inherit the trait NetherrackBuffer<(Something extending Read and Write), but for now, we'll do everything here
pub struct DequeBuffer {

    data: VecDeque<u8>

}

impl DequeBuffer {

    /// Creates a DequeBuffer from a VecDeque
    ///
    /// After this method, we own the VecDeque
    pub fn from(deque: VecDeque<u8>) -> DequeBuffer {
        DequeBuffer { data: deque }
    }

    /// Creates a new empty DequeBuffer
    pub fn new() -> DequeBuffer {
        DequeBuffer { data: VecDeque::<u8>::new() }
    }

    /// Creates a new empty DequeBuffer with the specified initial capacity
    pub fn with_capacity(capacity: usize) -> DequeBuffer {
        DequeBuffer { data: VecDeque::<u8>::with_capacity(capacity) }
    }

    /// Gets the remaining bytes in this DequeBuffer
    pub fn remaining(&self) -> usize {
        self.data.len()
    }

    /// Reads a UTF-8 string from the buffer
    pub fn read_utf8_string(&mut self) -> String {
        let size_result = self.read_unsigned_varint_32();

        if size_result.is_err() {
            return String::from("!!!BAD_SIZE!!!");
        } else {
            let mut size: u32 = size_result.unwrap();

            let mut bytes: Vec<u8> = Vec::<u8>::with_capacity(size as usize);

            while size > 0 {
                bytes.push(self.read_unsigned_byte());
                size -= 1;
            }

            let result = String::from_utf8(bytes);

            if result.is_ok() {
                return result.unwrap();
            } else {
                return String::from("!!!INVALID_UTF8!!!");
            }

        }
    }

    /// Reads an unsigned byte from the buffer, returning 0 if no data is available
    pub fn read_unsigned_byte(&mut self) -> u8 {
        match self.data.pop_front() {
            Some(value) => value,
            None => 0
        }
    }

    /// Reads an unsigned byte from the buffer, returning -127 if no data is available
    pub fn read_signed_byte(&mut self) -> i8 {
        match self.data.pop_front() {
            Some(value) => value as i8,
            None => -127
        }
    }

    /// Reads an unsigned short from the buffer
    pub fn read_unsigned_short(&mut self) -> u16 {

        let high: u8 = self.read_unsigned_byte();
        let low: u8 = self.read_unsigned_byte();

        ((high as u16) << 8) | low as u16

    }

    /// Writes an unsigned byte to the buffer
    pub fn write_unsigned_byte(&mut self, value: u8) {
        self.data.push_back(value);
    }

    /// Writes a signed byte to the buffer
    pub fn write_signed_byte(&mut self, value: i8) {
        self.data.push_back(value as u8);
    }

    /// Writes an unsigned 32-bit Varint from the buffer
    pub fn write_unsigned_varint_32(&mut self, value: u32) {

        let mut _value: u32 = value;

        if value == 0 {
            self.write_unsigned_byte(0);
        } else {

            while _value >= 0b10000000 {

                let next_byte: u8 = ((_value & 0b01111111) as u8) | 0b10000000;

                _value = _value >> 7;

                self.write_unsigned_byte(next_byte);

            }

            self.write_unsigned_byte((_value & 0b01111111) as u8);

        }

    }

    /// Reads an unsigned Varint from the buffer
    pub fn read_unsigned_varint_32(&mut self) -> Result<u32, &'static str> {
        if self.remaining() == 0 {
            return Err("Cannot read a Varint with no data to read.");
        } else {
            // The number of bits to shift by. <<0, <<7, <<14, etc
            let mut shift_amount: u32 = 0;

            // The decoded value
            let mut decoded_value: u32 = 0;

            while self.remaining() >= 1 {

                let byte_value: u8 = self.read_unsigned_byte(); //Yay for this existing!

                decoded_value |= ((byte_value & 0b01111111) as u32) << shift_amount; // We're not shifting for the first byte ;)

                // See if we're still reading bytes
                if byte_value.has_most_signifigant_bit() {
                    shift_amount += 7;
                } else {
                    return Ok(decoded_value);
                }

            }

            return Err("Reached end of buffer while reading a varint32");

        }
    }

}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn new_contains_no_elements() {
        let new = DequeBuffer::new();

        assert_eq!(0, new.remaining());
    }

    #[test]
    fn new_with_capacity_contains_no_elements() {
        let new = DequeBuffer::with_capacity(25000); //WOW, MUCH BIG PACKET

        assert_eq!(0, new.remaining());
    }

    #[test]
    fn unsigned_byte() {
        let mut buffer = DequeBuffer::with_capacity(2);

        buffer.write_unsigned_byte(1);
        buffer.write_unsigned_byte(15);

        assert_eq!(2, buffer.remaining());

        assert_eq!(1, buffer.read_unsigned_byte());
        assert_eq!(15, buffer.read_unsigned_byte());
    }

    #[test]
    fn signed_byte() {
        let mut buffer = DequeBuffer::with_capacity(2);

        buffer.write_signed_byte(-50);
        buffer.write_signed_byte(0);

        assert_eq!(2, buffer.remaining());

        assert_eq!(-50, buffer.read_signed_byte());
        assert_eq!(-0, buffer.read_signed_byte());
    }

    #[test]
    fn unsigned_varint_32() {

        let mut buffer = DequeBuffer::new();

        let var1: u32 = 0;

        let var2: u32 = 4003254234;

        let var3: u32 = 200;

        buffer.write_unsigned_varint_32(var1);
        buffer.write_unsigned_varint_32(var2);
        buffer.write_unsigned_varint_32(var3);

        assert_eq!(8, buffer.remaining());

        // No need to check for okay here, we know it will be. If you're getting it from the client, always check!
        assert_eq!(var1, buffer.read_unsigned_varint_32().unwrap());
        assert_eq!(var2, buffer.read_unsigned_varint_32().unwrap());
        assert_eq!(var3, buffer.read_unsigned_varint_32().unwrap());

    }

}<|endoftext|>


<commit_before>#![feature(plugin_registrar)]
#![feature(rustc_private)]

//! basic syntax extension to calculate the factorial of 10 at compile time for the compile-time
//! calculation task (the task itself is in `src/compile_time_calculation.rs`)
extern crate syntax;
extern crate rustc_plugin;

use syntax::ast;
use syntax::codemap::Span;
use syntax::parse;
use syntax::ext::base::{ExtCtxt, MacResult, DummyResult, MacEager};
use syntax::ext::build::AstBuilder;  // trait for expr_usize
use syntax::tokenstream::TokenTree;
use rustc_plugin::Registry;

fn exp_factorial(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    // extract the argument and ensure there is only one and it's a usize
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), tts.to_vec());

    // Try to parse a literal (doesn't need to be a number)
    let literal = match parser.parse_lit() {
        Ok(l) => l,
        Err(_) => {
            // span_err shows a compile time error to the user
            cx.span_err(sp, "fatal error");
            return DummyResult::any(sp);
        }
    };

    // Take the number of the literal (if it is a number)
    let n: u64 = match literal.node {
        ast::LitKind::Int(x, _) => x,
        _ => {
            cx.span_err(sp, "Invalid literal (expected unsigned integer)");
            return DummyResult::any(sp);
        }
    };

    // calculating the factorial
    let result = (1..(n as usize) + 1).fold(1, |accum, elem| accum * elem);

    MacEager::expr(cx.expr_usize(sp, result))
}

/// here's where we register the macro with the name of `factorial` so that it can be invoked a
/// like every other macro with a ! at the end
#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("factorial", exp_factorial);
}
<commit_msg>remove explicit type in compiler plugin<commit_after>#![feature(plugin_registrar)]
#![feature(rustc_private)]

//! basic syntax extension to calculate the factorial of 10 at compile time for the compile-time
//! calculation task (the task itself is in `src/compile_time_calculation.rs`)
extern crate syntax;
extern crate rustc_plugin;

use syntax::ast;
use syntax::codemap::Span;
use syntax::parse;
use syntax::ext::base::{ExtCtxt, MacResult, DummyResult, MacEager};
use syntax::ext::build::AstBuilder;  // trait for expr_usize
use syntax::tokenstream::TokenTree;
use rustc_plugin::Registry;

fn exp_factorial(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    // extract the argument and ensure there is only one and it's a usize
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), tts.to_vec());

    // Try to parse a literal (doesn't need to be a number)
    let literal = match parser.parse_lit() {
        Ok(l) => l,
        Err(_) => {
            // span_err shows a compile time error to the user
            cx.span_err(sp, "fatal error");
            return DummyResult::any(sp);
        }
    };

    // Take the number of the literal (if it is a number)
    let n = match literal.node {
        ast::LitKind::Int(x, _) => x,
        _ => {
            cx.span_err(sp, "Invalid literal (expected unsigned integer)");
            return DummyResult::any(sp);
        }
    };

    // calculating the factorial
    let result = (1..(n as usize) + 1).fold(1, |accum, elem| accum * elem);

    MacEager::expr(cx.expr_usize(sp, result))
}

/// here's where we register the macro with the name of `factorial` so that it can be invoked a
/// like every other macro with a ! at the end
#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("factorial", exp_factorial);
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

#![cfg_attr(feature = "cargo-clippy", allow(cast_lossless))]

use partition::TxSize;
use std::mem;

extern {
  static dc_qlookup_Q3: [i16; 256];
  static dc_qlookup_10_Q3: [i16; 256];
  static dc_qlookup_12_Q3: [i16; 256];

  static ac_qlookup_Q3: [i16; 256];
  static ac_qlookup_10_Q3: [i16; 256];
  static ac_qlookup_12_Q3: [i16; 256];
}

fn get_tx_scale(tx_size: TxSize) -> u8 {
  match tx_size {
    TxSize::TX_64X64 => 4,
    TxSize::TX_32X32 => 2,
    _ => 1
  }
}

pub fn dc_q(qindex: usize, bit_depth: usize) -> i16 {
  let &table = match bit_depth {
    8 => &dc_qlookup_Q3,
    10 => &dc_qlookup_10_Q3,
    12 => &dc_qlookup_12_Q3,
    _ => unimplemented!()
  };

  table[qindex]
}

pub fn ac_q(qindex: usize, bit_depth: usize) -> i16 {
  let &table = match bit_depth {
    8 => &ac_qlookup_Q3,
    10 => &ac_qlookup_10_Q3,
    12 => &ac_qlookup_12_Q3,
    _ => unimplemented!()
  };

  table[qindex]
}

#[derive(Debug, Default, Clone, Copy)]
pub struct QuantizationContext {
  tx_scale: i32,
  dc_quant: u32,
  dc_offset: i32,
  dc_mul_add: (u32, u32, u32),

  ac_quant: u32,
  ac_offset: i32,
  ac_mul_add: (u32, u32, u32)
}

fn divu_gen(d: u32) -> (u32, u32, u32) {
  let nbits = (mem::size_of_val(&d) as u64) * 8;
  let m = nbits - d.leading_zeros() as u64 - 1;
  if (d & (d - 1)) == 0 {
    (0xFFFFFFFF, 0xFFFFFFFF, m as u32)
  } else {
    let d = d as u64;
    let t = (1u64 << (m + nbits)) / d;
    let r = (t * d + d) & ((1 << nbits) - 1);
    if r <= 1u64 << m {
      (t as u32 + 1, 0u32, m as u32)
    } else {
      (t as u32, t as u32, m as u32)
    }
  }
}

#[inline]
fn divu_pair(x: i32, d: (u32, u32, u32)) -> i32 {
  let y = if x < 0 { -x } else { x } as u64;
  let (a, b, shift) = d;
  let shift = shift as u64;
  let a = a as u64;
  let b = b as u64;

  let y = (((a * y + b) >> 32) >> shift) as i32;
  if x < 0 {
    -y
  } else {
    y
  }
}

#[cfg(test)]
mod test {
  use super::*;
  #[test]
  fn test_divu_pair() {
    for d in 1..1024 {
      for x in -1000..1000 {
        let ab = divu_gen(d as u32);
        assert_eq!(x / d, divu_pair(x, ab));
      }
    }
  }
  #[test]
  fn gen_divu_table() {
    let b: Vec<(u32, u32, u32)> =
      dc_qlookup_Q3.iter().map(|&v| divu_gen(v as u32)).collect();

    println!("{:?}", b);
  }
}

impl QuantizationContext {
  pub fn update(&mut self, qindex: usize, tx_size: TxSize, is_intra: bool, bit_depth: usize) {
    self.tx_scale = get_tx_scale(tx_size) as i32;

    self.dc_quant = dc_q(qindex, bit_depth) as u32;
    self.dc_mul_add = divu_gen(self.dc_quant);

    self.ac_quant = ac_q(qindex, bit_depth) as u32;
    self.ac_mul_add = divu_gen(self.ac_quant);

    self.dc_offset = self.dc_quant as i32 * (if is_intra {21} else {15}) / 64;
    self.ac_offset = self.ac_quant as i32 * (if is_intra {21} else {15}) / 64;
  }

  #[inline]
  pub fn quantize(&self, coeffs: &mut [i32]) {
    coeffs[0] *= self.tx_scale;
    coeffs[0] += coeffs[0].signum() * self.dc_offset;
    coeffs[0] = divu_pair(coeffs[0], self.dc_mul_add);

    for c in coeffs[1..].iter_mut() {
      *c *= self.tx_scale;
      *c += c.signum() * self.ac_offset;
      *c = divu_pair(*c, self.ac_mul_add);
    }
  }
}

pub fn quantize_in_place(qindex: usize, coeffs: &mut [i32], tx_size: TxSize, bit_depth: usize) {
  let tx_scale = get_tx_scale(tx_size) as i32;

  let dc_quant = dc_q(qindex, bit_depth) as i32;
  let ac_quant = ac_q(qindex, bit_depth) as i32;

  // using 21/64=0.328125 as rounding offset. To be tuned
  let dc_offset = dc_quant * 21 / 64 as i32;
  let ac_offset = ac_quant * 21 / 64 as i32;

  coeffs[0] *= tx_scale;
  coeffs[0] += coeffs[0].signum() * dc_offset;
  coeffs[0] /= dc_quant;

  for c in coeffs[1..].iter_mut() {
    *c *= tx_scale;
    *c += c.signum() * ac_offset;
    *c /= ac_quant;
  }
}

pub fn dequantize(
  qindex: usize, coeffs: &[i32], rcoeffs: &mut [i32], tx_size: TxSize, bit_depth: usize
) {
  let tx_scale = get_tx_scale(tx_size) as i32;

  rcoeffs[0] = (coeffs[0] * dc_q(qindex, bit_depth) as i32) / tx_scale;
  let ac_quant = ac_q(qindex, bit_depth) as i32;

  for (r, &c) in rcoeffs.iter_mut().zip(coeffs.iter()).skip(1) {
    *r = c * ac_quant / tx_scale;
  }
}
<commit_msg>Don't use division in inverse quantizer<commit_after>// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

#![cfg_attr(feature = "cargo-clippy", allow(cast_lossless))]

use partition::TxSize;
use std::mem;

extern {
  static dc_qlookup_Q3: [i16; 256];
  static dc_qlookup_10_Q3: [i16; 256];
  static dc_qlookup_12_Q3: [i16; 256];

  static ac_qlookup_Q3: [i16; 256];
  static ac_qlookup_10_Q3: [i16; 256];
  static ac_qlookup_12_Q3: [i16; 256];
}

fn get_log_tx_scale(tx_size: TxSize) -> i32 {
  match tx_size {
    TxSize::TX_64X64 => 2,
    TxSize::TX_32X32 => 1,
    _ => 0
  }
}

pub fn dc_q(qindex: usize, bit_depth: usize) -> i16 {
  let &table = match bit_depth {
    8 => &dc_qlookup_Q3,
    10 => &dc_qlookup_10_Q3,
    12 => &dc_qlookup_12_Q3,
    _ => unimplemented!()
  };

  table[qindex]
}

pub fn ac_q(qindex: usize, bit_depth: usize) -> i16 {
  let &table = match bit_depth {
    8 => &ac_qlookup_Q3,
    10 => &ac_qlookup_10_Q3,
    12 => &ac_qlookup_12_Q3,
    _ => unimplemented!()
  };

  table[qindex]
}

#[derive(Debug, Default, Clone, Copy)]
pub struct QuantizationContext {
  log_tx_scale: i32,
  dc_quant: u32,
  dc_offset: i32,
  dc_mul_add: (u32, u32, u32),

  ac_quant: u32,
  ac_offset: i32,
  ac_mul_add: (u32, u32, u32)
}

fn divu_gen(d: u32) -> (u32, u32, u32) {
  let nbits = (mem::size_of_val(&d) as u64) * 8;
  let m = nbits - d.leading_zeros() as u64 - 1;
  if (d & (d - 1)) == 0 {
    (0xFFFFFFFF, 0xFFFFFFFF, m as u32)
  } else {
    let d = d as u64;
    let t = (1u64 << (m + nbits)) / d;
    let r = (t * d + d) & ((1 << nbits) - 1);
    if r <= 1u64 << m {
      (t as u32 + 1, 0u32, m as u32)
    } else {
      (t as u32, t as u32, m as u32)
    }
  }
}

#[inline]
fn divu_pair(x: i32, d: (u32, u32, u32)) -> i32 {
  let y = if x < 0 { -x } else { x } as u64;
  let (a, b, shift) = d;
  let shift = shift as u64;
  let a = a as u64;
  let b = b as u64;

  let y = (((a * y + b) >> 32) >> shift) as i32;
  if x < 0 {
    -y
  } else {
    y
  }
}

#[cfg(test)]
mod test {
  use super::*;
  #[test]
  fn test_divu_pair() {
    for d in 1..1024 {
      for x in -1000..1000 {
        let ab = divu_gen(d as u32);
        assert_eq!(x / d, divu_pair(x, ab));
      }
    }
  }
  #[test]
  fn gen_divu_table() {
    let b: Vec<(u32, u32, u32)> =
      dc_qlookup_Q3.iter().map(|&v| divu_gen(v as u32)).collect();

    println!("{:?}", b);
  }
}

impl QuantizationContext {
  pub fn update(&mut self, qindex: usize, tx_size: TxSize, is_intra: bool, bit_depth: usize) {
    self.log_tx_scale = get_log_tx_scale(tx_size);

    self.dc_quant = dc_q(qindex, bit_depth) as u32;
    self.dc_mul_add = divu_gen(self.dc_quant);

    self.ac_quant = ac_q(qindex, bit_depth) as u32;
    self.ac_mul_add = divu_gen(self.ac_quant);

    self.dc_offset = self.dc_quant as i32 * (if is_intra {21} else {15}) / 64;
    self.ac_offset = self.ac_quant as i32 * (if is_intra {21} else {15}) / 64;
  }

  #[inline]
  pub fn quantize(&self, coeffs: &mut [i32]) {
    coeffs[0] <<= self.log_tx_scale;
    coeffs[0] += coeffs[0].signum() * self.dc_offset;
    coeffs[0] = divu_pair(coeffs[0], self.dc_mul_add);

    for c in coeffs[1..].iter_mut() {
      *c <<= self.log_tx_scale;
      *c += c.signum() * self.ac_offset;
      *c = divu_pair(*c, self.ac_mul_add);
    }
  }
}

pub fn quantize_in_place(qindex: usize, coeffs: &mut [i32], tx_size: TxSize, bit_depth: usize) {
  let log_tx_scale = get_log_tx_scale(tx_size);

  let dc_quant = dc_q(qindex, bit_depth) as i32;
  let ac_quant = ac_q(qindex, bit_depth) as i32;

  // using 21/64=0.328125 as rounding offset. To be tuned
  let dc_offset = dc_quant * 21 / 64 as i32;
  let ac_offset = ac_quant * 21 / 64 as i32;

  coeffs[0] <<= log_tx_scale;
  coeffs[0] += coeffs[0].signum() * dc_offset;
  coeffs[0] /= dc_quant;

  for c in coeffs[1..].iter_mut() {
    *c <<= log_tx_scale;
    *c += c.signum() * ac_offset;
    *c /= ac_quant;
  }
}

pub fn dequantize(
  qindex: usize, coeffs: &[i32], rcoeffs: &mut [i32], tx_size: TxSize, bit_depth: usize
) {
  let log_tx_scale = get_log_tx_scale(tx_size) as i32;

  let val = coeffs[0] * dc_q(qindex, bit_depth) as i32;
  rcoeffs[0] = (val + if val < 0 { (1 << log_tx_scale) - 1 } else { 0 }) >> log_tx_scale;

  let ac_quant = ac_q(qindex, bit_depth) as i32;

  for (r, &c) in rcoeffs.iter_mut().zip(coeffs.iter()).skip(1) {
    *r = (c * ac_quant + if c < 0 { (1 << log_tx_scale) - 1 } else { 0 }) >> log_tx_scale;
  }
}
<|endoftext|>


<commit_before>use crate::ram::chunk::Chunk;
use crate::ram::entry;
use crate::ram::entry::EntryMeta;
use crate::ram::tombstone::TOMBSTONE_SIZE_U32;
use libc::*;
use lightning::list::WordList;
use parking_lot;
use std::fs::{copy, create_dir_all, remove_file, File};
use std::io;
use std::io::prelude::*;
use std::io::BufWriter;
use std::path::Path;
use std::ptr;
use std::sync::atomic::{
    AtomicBool, AtomicI64, AtomicU32, AtomicUsize, Ordering, Ordering::Relaxed,
};

pub const SEGMENT_SIZE_U32: u32 = 8 * 1024 * 1024;
pub const SEGMENT_SIZE: usize = SEGMENT_SIZE_U32 as usize;
pub const SEGMENT_MASK: usize = !(SEGMENT_SIZE - 1);
pub const SEGMENT_BITS_SHIFT: u32 = SEGMENT_SIZE.trailing_zeros();

pub struct Segment {
    pub id: u64,
    pub addr: usize,
    pub bound: usize,
    pub append_header: AtomicUsize,
    pub dead_space: AtomicU32,
    pub tombstones: AtomicU32,
    pub dead_tombstones: AtomicU32,
    pub last_tombstones_scanned: AtomicI64,
    pub references: AtomicUsize,
    pub backup_file_name: Option<String>,
    pub wal_file: Option<parking_lot::Mutex<BufWriter<File>>>,
    pub wal_file_name: Option<String>,
    pub archived: AtomicBool,
    pub dropped: AtomicBool,
}

impl Segment {
    pub fn new(
        id: u64,
        buffer_ptr: usize,
        backup_storage: &Option<String>,
        wal_storage: &Option<String>,
    ) -> Segment {
        let mut wal_file_name = None;
        let mut wal_file = None;
        let size = SEGMENT_SIZE;
        if let Some(wal_storage) = wal_storage {
            create_dir_all(wal_storage).unwrap();
            let file_name = format!("{}/{}.log", wal_storage, id);
            let file = BufWriter::with_capacity(
                4096, // most common disk block size
                File::create(&file_name).unwrap(),
            ); // fast fail
            wal_file_name = Some(file_name);
            wal_file = Some(parking_lot::Mutex::new(file));
        }
        debug!(
            "Creating new segment with id {}, size {}, address {}",
            id, size, buffer_ptr
        );
        Segment {
            addr: buffer_ptr,
            id,
            bound: buffer_ptr + size,
            append_header: AtomicUsize::new(buffer_ptr),
            dead_space: AtomicU32::new(0),
            tombstones: AtomicU32::new(0),
            dead_tombstones: AtomicU32::new(0),
            last_tombstones_scanned: AtomicI64::new(0),
            references: AtomicUsize::new(0),
            backup_file_name: backup_storage
                .clone()
                .map(|path| format!("{}/{}.backup", path, id)),
            archived: AtomicBool::new(false),
            dropped: AtomicBool::new(false),
            wal_file,
            wal_file_name,
        }
    }

    pub fn try_acquire(&self, size: u32) -> Option<usize> {
        let size = size as usize;
        loop {
            let curr_last = self.append_header.load(Ordering::SeqCst);
            let exp_last = curr_last + size;
            if exp_last > self.bound {
                return None;
            } else {
                if self
                    .append_header
                    .compare_and_swap(curr_last, exp_last, Ordering::SeqCst)
                    != curr_last
                {
                    continue;
                } else {
                    return Some(curr_last);
                }
            }
        }
    }

    pub fn shrink(&self, size: usize) {
        debug_assert!(
            size < SEGMENT_SIZE,
            "Shrink to {} max {}",
            size,
            SEGMENT_SIZE
        );
        punch_hole(self.addr, size);
    }

    fn append_header(&self) -> usize {
        self.append_header.load(Ordering::Relaxed)
    }

    pub fn entry_iter(&self) -> SegmentEntryIter {
        SegmentEntryIter {
            bound: self.append_header(),
            cursor: self.addr,
        }
    }

    pub fn dead_space(&self) -> u32 {
        self.dead_space.load(Ordering::Relaxed)
    }

    // dead space plus tombstone spaces
    pub fn total_dead_space(&self) -> u32 {
        let dead_tombstones_space =
            self.dead_tombstones.load(Ordering::Relaxed) * TOMBSTONE_SIZE_U32;
        let dead_cells_space = self.dead_space();
        return dead_tombstones_space + dead_cells_space;
    }

    pub fn used_spaces(&self) -> u32 {
        let space = self.append_header.load(Ordering::Relaxed) as usize - self.addr;
        debug_assert!(space <= SEGMENT_SIZE);
        return space as u32;
    }

    pub fn living_space(&self) -> u32 {
        let total_dead_space = self.total_dead_space();
        let used_space = self.used_spaces();
        if total_dead_space <= used_space {
            used_space - total_dead_space
        } else {
            warn!(
                "living space check error for segment {}, used {}, dead {}",
                self.id, used_space, total_dead_space
            );
            0
        }
    }

    pub fn valid_space(&self) -> u32 {
        return self.used_spaces() - self.dead_space();
    }

    pub fn living_rate(&self) -> f32 {
        let used_space = self.used_spaces() as f32;
        if used_space == 0f32 {
            // empty segment
            return 1f32;
        }
        return self.living_space() as f32 / used_space;
    }

    // archive this segment and write the data to backup storage
    pub fn archive(&self) -> Result<bool, io::Error> {
        if let &Some(ref backup_file) = &self.backup_file_name {
            while !self.no_references() { /* wait until all references released */ }
            let backup_file_path = Path::new(backup_file);
            if backup_file_path.exists() {
                warn!(
                    "Segment backup {} exists and can't archive twice",
                    backup_file
                );
                return Ok(false);
            }
            if let Some(ref wal_file) = self.wal_file_name {
                // if there is a WAL file ready, copy this file to backup
                if let Some(ref file_mutex) = self.wal_file {
                    // this should be redundant but I don't want to take the chance
                    // obtain the writer lock before continue
                    let _ = file_mutex.lock();
                    copy(wal_file, backup_file)?;
                    remove_file(wal_file)?;
                } else {
                    panic!()
                }
            } else {
                let backup_file = File::create(backup_file_path)?;
                let seg_size = self.append_header.load(Ordering::Relaxed) - self.addr;
                let mut buffer = BufWriter::with_capacity(seg_size, backup_file);
                unsafe {
                    for offset in 0..seg_size {
                        let ptr = (self.addr + offset) as *const u8;
                        let byte = *ptr;
                        buffer.write(&[byte])?;
                    }
                }
                buffer.flush()?;
                return Ok(true);
            }
        }
        return Ok(false);
    }

    pub fn write_wal(&self, addr: usize, size: u32) -> io::Result<()> {
        if let Some(ref wal_file) = self.wal_file {
            let mut file = wal_file.lock();
            unsafe {
                for offset in 0..size as usize {
                    let ptr = (addr + offset) as *const u8;
                    let byte = *ptr;
                    file.write(&[byte])?;
                }
            }
            file.flush()?;
        }
        return Ok(());
    }

    pub fn no_references(&self) -> bool {
        self.references.load(Ordering::Relaxed) == 0
    }

    pub fn mem_drop(&self, chunk: &Chunk) {
        if !self
            .dropped
            .compare_and_swap(false, true, Ordering::Relaxed)
        {
            chunk.allocator.free(self.addr);
        }
    }

    // remove the backup if it have one
    pub fn dispense(&self) {
        debug!("dispense segment {}", self.id);
        if let &Some(ref backup_storage) = &self.backup_file_name {
            let path = Path::new(backup_storage);
            if path.exists() {
                if let Err(_e) = remove_file(path) {
                    error!("cannot reclaim segment file on dispense {}", backup_storage)
                }
            } else {
                error!("cannot find segment backup to dispense {}", backup_storage)
            }
        }
    }
}

pub struct SegmentEntryIter {
    bound: usize,
    cursor: usize,
}

impl Iterator for SegmentEntryIter {
    type Item = EntryMeta;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        let cursor = self.cursor;
        if cursor >= self.bound {
            return None;
        }
        let (_, entry_meta) = entry::Entry::decode_from(cursor, |body_pos, header| {
            let entry_header_size = body_pos - cursor;
            let entry_size = entry_header_size + header.content_length as usize;
            trace!("Found body pos {}, entry header {:?}. Header size: {}, entry size: {}, entry pos: {}, content length {}, bound {}",
                       body_pos, header, entry_header_size, entry_size, cursor, header.content_length, self.bound);
            return EntryMeta {
                body_pos,
                entry_header: header,
                entry_size,
                entry_pos: cursor,
            };
        });
        self.cursor += entry_meta.entry_size;
        Some(entry_meta)
    }
}

pub const PAGE_SHIFT: usize = 12; // 4K
pub const PAGE_SIZE: usize = 1 << PAGE_SHIFT;

pub struct SegmentAllocator {
    base: usize,
    offset: AtomicUsize,
    limit: usize,
    gc_threshold: usize,
    free: WordList,
}

impl SegmentAllocator {
    pub fn new(chunk_size: usize) -> Self {
        let overflow = SEGMENT_SIZE - PAGE_SIZE;
        let aligned_size = chunk_size + overflow;
        let ptr = unsafe {
            libc::mmap(
                ptr::null_mut(),
                aligned_size,
                PROT_READ | PROT_WRITE,
                MAP_ANONYMOUS | MAP_PRIVATE,
                -1,
                0,
            )
        };
        let addr = ptr as usize;
        let start = addr + overflow;
        let aligned_addr = start & SEGMENT_MASK;
        Self {
            base: aligned_addr,
            offset: AtomicUsize::new(aligned_addr),
            limit: aligned_addr + chunk_size,
            gc_threshold: aligned_addr + (chunk_size as f64 * 0.9) as usize - SEGMENT_SIZE,
            free: WordList::with_capacity(chunk_size / SEGMENT_SIZE / 2),
        }
    }

    pub fn meet_gc_threshold(&self) -> bool {
        self.offset.load(Relaxed) > self.gc_threshold
    }

    pub fn alloc_seg(
        &self,
        backup_storage: &Option<String>,
        wal_storage: &Option<String>,
    ) -> Option<Segment> {
        self.free
            .pop()
            .or_else(|| loop {
                let addr = self.offset.load(Relaxed);
                let new_addr = addr + SEGMENT_SIZE;
                if addr > self.limit {
                    return None;
                } else {
                    if self.offset.compare_and_swap(addr, new_addr, Relaxed) == addr {
                        return Some(addr);
                    }
                }
            })
            .map(|addr| {
                let id = self.id_by_addr(addr);
                Segment::new(id as u64, addr, backup_storage, wal_storage)
            })
    }

    pub fn free(&self, seg_addr: usize) {
        debug_assert!(seg_addr >= self.base);
        debug_assert!(seg_addr < self.limit);
        debug!("Segment {} freed", seg_addr);
        self.free.push(seg_addr);
    }

    pub fn id_by_addr(&self, addr: usize) -> usize {
        let offset = addr - self.base;
        let id = offset >> SEGMENT_BITS_SHIFT;
        id
    }
}

#[cfg(target_os = "linux")]
unsafe fn madvise_free(addr: usize, size: usize) {
    madvise(addr as *mut c_void, size, MADV_REMOVE);
}

#[cfg(not(target_os = "linux"))]
unsafe fn madvise_free(addr: usize, size: usize) {
    madvise(addr as *mut c_void, size, MADV_DONTNEED);
}

fn punch_hole(seg_addr: usize, seg_size: usize) {
    let right_boundary = seg_addr + seg_size;
    let aligned_addr = (((right_boundary - 1) >> PAGE_SHIFT) + 1) << PAGE_SHIFT;
    let hole_length = (seg_addr + SEGMENT_SIZE) - aligned_addr;
    if hole_length > PAGE_SIZE {
        // Have pages to release
        debug!(
            "Partially free the segment by puching hole with size {}",
            hole_length
        );
        unsafe {
            madvise_free(aligned_addr, hole_length);
        }
    }
}
<commit_msg>fix bump pointer for segment allocation<commit_after>use crate::ram::chunk::Chunk;
use crate::ram::entry;
use crate::ram::entry::EntryMeta;
use crate::ram::tombstone::TOMBSTONE_SIZE_U32;
use libc::*;
use lightning::list::WordList;
use parking_lot;
use std::fs::{copy, create_dir_all, remove_file, File};
use std::io;
use std::io::prelude::*;
use std::io::BufWriter;
use std::path::Path;
use std::ptr;
use std::sync::atomic::{
    AtomicBool, AtomicI64, AtomicU32, AtomicUsize, Ordering, Ordering::Relaxed,
};

pub const SEGMENT_SIZE_U32: u32 = 8 * 1024 * 1024;
pub const SEGMENT_SIZE: usize = SEGMENT_SIZE_U32 as usize;
pub const SEGMENT_MASK: usize = !(SEGMENT_SIZE - 1);
pub const SEGMENT_BITS_SHIFT: u32 = SEGMENT_SIZE.trailing_zeros();

pub struct Segment {
    pub id: u64,
    pub addr: usize,
    pub bound: usize,
    pub append_header: AtomicUsize,
    pub dead_space: AtomicU32,
    pub tombstones: AtomicU32,
    pub dead_tombstones: AtomicU32,
    pub last_tombstones_scanned: AtomicI64,
    pub references: AtomicUsize,
    pub backup_file_name: Option<String>,
    pub wal_file: Option<parking_lot::Mutex<BufWriter<File>>>,
    pub wal_file_name: Option<String>,
    pub archived: AtomicBool,
    pub dropped: AtomicBool,
}

impl Segment {
    pub fn new(
        id: u64,
        buffer_ptr: usize,
        backup_storage: &Option<String>,
        wal_storage: &Option<String>,
    ) -> Segment {
        let mut wal_file_name = None;
        let mut wal_file = None;
        let size = SEGMENT_SIZE;
        if let Some(wal_storage) = wal_storage {
            create_dir_all(wal_storage).unwrap();
            let file_name = format!("{}/{}.log", wal_storage, id);
            let file = BufWriter::with_capacity(
                4096, // most common disk block size
                File::create(&file_name).unwrap(),
            ); // fast fail
            wal_file_name = Some(file_name);
            wal_file = Some(parking_lot::Mutex::new(file));
        }
        debug!(
            "Creating new segment with id {}, size {}, address {}",
            id, size, buffer_ptr
        );
        Segment {
            addr: buffer_ptr,
            id,
            bound: buffer_ptr + size,
            append_header: AtomicUsize::new(buffer_ptr),
            dead_space: AtomicU32::new(0),
            tombstones: AtomicU32::new(0),
            dead_tombstones: AtomicU32::new(0),
            last_tombstones_scanned: AtomicI64::new(0),
            references: AtomicUsize::new(0),
            backup_file_name: backup_storage
                .clone()
                .map(|path| format!("{}/{}.backup", path, id)),
            archived: AtomicBool::new(false),
            dropped: AtomicBool::new(false),
            wal_file,
            wal_file_name,
        }
    }

    pub fn try_acquire(&self, size: u32) -> Option<usize> {
        let size = size as usize;
        loop {
            let curr_last = self.append_header.load(Ordering::SeqCst);
            let exp_last = curr_last + size;
            if exp_last > self.bound {
                return None;
            } else {
                if self
                    .append_header
                    .compare_and_swap(curr_last, exp_last, Ordering::SeqCst)
                    != curr_last
                {
                    continue;
                } else {
                    return Some(curr_last);
                }
            }
        }
    }

    pub fn shrink(&self, size: usize) {
        debug_assert!(
            size < SEGMENT_SIZE,
            "Shrink to {} max {}",
            size,
            SEGMENT_SIZE
        );
        punch_hole(self.addr, size);
    }

    fn append_header(&self) -> usize {
        self.append_header.load(Ordering::Relaxed)
    }

    pub fn entry_iter(&self) -> SegmentEntryIter {
        SegmentEntryIter {
            bound: self.append_header(),
            cursor: self.addr,
        }
    }

    pub fn dead_space(&self) -> u32 {
        self.dead_space.load(Ordering::Relaxed)
    }

    // dead space plus tombstone spaces
    pub fn total_dead_space(&self) -> u32 {
        let dead_tombstones_space =
            self.dead_tombstones.load(Ordering::Relaxed) * TOMBSTONE_SIZE_U32;
        let dead_cells_space = self.dead_space();
        return dead_tombstones_space + dead_cells_space;
    }

    pub fn used_spaces(&self) -> u32 {
        let space = self.append_header.load(Ordering::Relaxed) as usize - self.addr;
        debug_assert!(space <= SEGMENT_SIZE);
        return space as u32;
    }

    pub fn living_space(&self) -> u32 {
        let total_dead_space = self.total_dead_space();
        let used_space = self.used_spaces();
        if total_dead_space <= used_space {
            used_space - total_dead_space
        } else {
            warn!(
                "living space check error for segment {}, used {}, dead {}",
                self.id, used_space, total_dead_space
            );
            0
        }
    }

    pub fn valid_space(&self) -> u32 {
        return self.used_spaces() - self.dead_space();
    }

    pub fn living_rate(&self) -> f32 {
        let used_space = self.used_spaces() as f32;
        if used_space == 0f32 {
            // empty segment
            return 1f32;
        }
        return self.living_space() as f32 / used_space;
    }

    // archive this segment and write the data to backup storage
    pub fn archive(&self) -> Result<bool, io::Error> {
        if let &Some(ref backup_file) = &self.backup_file_name {
            while !self.no_references() { /* wait until all references released */ }
            let backup_file_path = Path::new(backup_file);
            if backup_file_path.exists() {
                warn!(
                    "Segment backup {} exists and can't archive twice",
                    backup_file
                );
                return Ok(false);
            }
            if let Some(ref wal_file) = self.wal_file_name {
                // if there is a WAL file ready, copy this file to backup
                if let Some(ref file_mutex) = self.wal_file {
                    // this should be redundant but I don't want to take the chance
                    // obtain the writer lock before continue
                    let _ = file_mutex.lock();
                    copy(wal_file, backup_file)?;
                    remove_file(wal_file)?;
                } else {
                    panic!()
                }
            } else {
                let backup_file = File::create(backup_file_path)?;
                let seg_size = self.append_header.load(Ordering::Relaxed) - self.addr;
                let mut buffer = BufWriter::with_capacity(seg_size, backup_file);
                unsafe {
                    for offset in 0..seg_size {
                        let ptr = (self.addr + offset) as *const u8;
                        let byte = *ptr;
                        buffer.write(&[byte])?;
                    }
                }
                buffer.flush()?;
                return Ok(true);
            }
        }
        return Ok(false);
    }

    pub fn write_wal(&self, addr: usize, size: u32) -> io::Result<()> {
        if let Some(ref wal_file) = self.wal_file {
            let mut file = wal_file.lock();
            unsafe {
                for offset in 0..size as usize {
                    let ptr = (addr + offset) as *const u8;
                    let byte = *ptr;
                    file.write(&[byte])?;
                }
            }
            file.flush()?;
        }
        return Ok(());
    }

    pub fn no_references(&self) -> bool {
        self.references.load(Ordering::Relaxed) == 0
    }

    pub fn mem_drop(&self, chunk: &Chunk) {
        if !self
            .dropped
            .compare_and_swap(false, true, Ordering::Relaxed)
        {
            chunk.allocator.free(self.addr);
        }
    }

    // remove the backup if it have one
    pub fn dispense(&self) {
        debug!("dispense segment {}", self.id);
        if let &Some(ref backup_storage) = &self.backup_file_name {
            let path = Path::new(backup_storage);
            if path.exists() {
                if let Err(_e) = remove_file(path) {
                    error!("cannot reclaim segment file on dispense {}", backup_storage)
                }
            } else {
                error!("cannot find segment backup to dispense {}", backup_storage)
            }
        }
    }
}

pub struct SegmentEntryIter {
    bound: usize,
    cursor: usize,
}

impl Iterator for SegmentEntryIter {
    type Item = EntryMeta;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        let cursor = self.cursor;
        if cursor >= self.bound {
            return None;
        }
        let (_, entry_meta) = entry::Entry::decode_from(cursor, |body_pos, header| {
            let entry_header_size = body_pos - cursor;
            let entry_size = entry_header_size + header.content_length as usize;
            trace!("Found body pos {}, entry header {:?}. Header size: {}, entry size: {}, entry pos: {}, content length {}, bound {}",
                       body_pos, header, entry_header_size, entry_size, cursor, header.content_length, self.bound);
            return EntryMeta {
                body_pos,
                entry_header: header,
                entry_size,
                entry_pos: cursor,
            };
        });
        self.cursor += entry_meta.entry_size;
        Some(entry_meta)
    }
}

pub const PAGE_SHIFT: usize = 12; // 4K
pub const PAGE_SIZE: usize = 1 << PAGE_SHIFT;

pub struct SegmentAllocator {
    base: usize,
    offset: AtomicUsize,
    limit: usize,
    gc_threshold: usize,
    free: WordList,
}

impl SegmentAllocator {
    pub fn new(chunk_size: usize) -> Self {
        let overflow = SEGMENT_SIZE - PAGE_SIZE;
        let aligned_size = chunk_size + overflow;
        let ptr = unsafe {
            libc::mmap(
                ptr::null_mut(),
                aligned_size,
                PROT_READ | PROT_WRITE,
                MAP_ANONYMOUS | MAP_PRIVATE,
                -1,
                0,
            )
        };
        let addr = ptr as usize;
        let start = addr + overflow;
        let aligned_addr = start & SEGMENT_MASK;
        Self {
            base: aligned_addr,
            offset: AtomicUsize::new(aligned_addr),
            limit: aligned_addr + chunk_size,
            gc_threshold: aligned_addr + (chunk_size as f64 * 0.9) as usize - SEGMENT_SIZE,
            free: WordList::with_capacity(chunk_size / SEGMENT_SIZE / 2),
        }
    }

    pub fn meet_gc_threshold(&self) -> bool {
        self.offset.load(Relaxed) > self.gc_threshold
    }

    pub fn alloc_seg(
        &self,
        backup_storage: &Option<String>,
        wal_storage: &Option<String>,
    ) -> Option<Segment> {
        self.free
            .pop()
            .or_else(|| loop {
                debug!("Allocate segment by bump pointer");
                let addr = self.offset.load(Relaxed);
                let new_addr = addr + SEGMENT_SIZE;
                if new_addr > self.limit {
                    return None;
                } else {
                    if self.offset.compare_and_swap(addr, new_addr, Relaxed) == addr {
                        return Some(addr);
                    }
                }
            })
            .map(|addr| {
                let id = self.id_by_addr(addr);
                Segment::new(id as u64, addr, backup_storage, wal_storage)
            })
    }

    pub fn free(&self, seg_addr: usize) {
        debug_assert!(seg_addr >= self.base);
        debug_assert!(seg_addr < self.limit);
        debug!("Segment {} freed", seg_addr);
        self.free.push(seg_addr);
    }

    pub fn id_by_addr(&self, addr: usize) -> usize {
        let offset = addr - self.base;
        let id = offset >> SEGMENT_BITS_SHIFT;
        id
    }
}

#[cfg(target_os = "linux")]
unsafe fn madvise_free(addr: usize, size: usize) {
    madvise(addr as *mut c_void, size, MADV_REMOVE);
}

#[cfg(not(target_os = "linux"))]
unsafe fn madvise_free(addr: usize, size: usize) {
    madvise(addr as *mut c_void, size, MADV_DONTNEED);
}

fn punch_hole(seg_addr: usize, seg_size: usize) {
    let right_boundary = seg_addr + seg_size;
    let aligned_addr = (((right_boundary - 1) >> PAGE_SHIFT) + 1) << PAGE_SHIFT;
    let hole_length = (seg_addr + SEGMENT_SIZE) - aligned_addr;
    if hole_length > PAGE_SIZE {
        // Have pages to release
        debug!(
            "Partially free the segment by puching hole with size {}",
            hole_length
        );
        unsafe {
            madvise_free(aligned_addr, hole_length);
        }
    }
}
<|endoftext|>