<commit_before>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! This module implements the virtio video encoder and decoder devices.
//! The current implementation uses [v3 RFC] of the virtio-video protocol.
//!
//! [v3 RFC]: https://markmail.org/thread/wxdne5re7aaugbjg

use std::fmt::{self, Display};
use std::os::unix::io::{AsRawFd, RawFd};
use std::thread;

use base::{error, Error as SysError, Event};
use data_model::{DataInit, Le32};
use vm_memory::GuestMemory;

use crate::virtio::resource_bridge::ResourceRequestSocket;
use crate::virtio::virtio_device::VirtioDevice;
use crate::virtio::{self, copy_config, DescriptorError, Interrupt, VIRTIO_F_VERSION_1};

#[macro_use]
mod macros;
mod async_cmd_desc_map;
mod command;
mod control;
mod decoder;
mod device;
mod encoder;
mod error;
mod event;
mod format;
mod params;
mod protocol;
mod response;
mod worker;

use command::ReadCmdError;
use device::AsyncCmdTag;
use worker::Worker;

const QUEUE_SIZE: u16 = 256;
const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE, QUEUE_SIZE];

/// An error indicating something went wrong in virtio-video's worker.
#[derive(Debug)]
pub enum Error {
    /// Failed to create a libvda instance.
    LibvdaCreationFailed(libvda::Error),
    /// Creating PollContext failed.
    PollContextCreationFailed(SysError),
    /// A DescriptorChain contains invalid data.
    InvalidDescriptorChain(DescriptorError),
    /// No available descriptor in which an event is written to.
    DescriptorNotAvailable,
    /// Error while polling for events.
    PollError(SysError),
    /// Failed to read a virtio-video command.
    ReadFailure(ReadCmdError),
    /// Got response for an unexpected asynchronous command.
    UnexpectedResponse(AsyncCmdTag),
    /// Failed to write an event into the event queue.
    WriteEventFailure {
        event: event::VideoEvt,
        error: std::io::Error,
    },
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Error::*;
        match self {
            LibvdaCreationFailed(e) => write!(f, "failed to create a libvda instance: {}", e),
            PollContextCreationFailed(e) => write!(f, "failed to create PollContext: {}", e),
            InvalidDescriptorChain(e) => write!(f, "DescriptorChain contains invalid data: {}", e),
            DescriptorNotAvailable => {
                write!(f, "no available descriptor in which an event is written to")
            }
            PollError(err) => write!(f, "failed to poll events: {}", err),
            ReadFailure(e) => write!(f, "failed to read a command from the guest: {}", e),
            UnexpectedResponse(tag) => {
                write!(f, "got a response for an untracked command: {:?}", tag)
            }
            WriteEventFailure { event, error } => write!(
                f,
                "failed to write an event {:?} into event queue: {}",
                event, error
            ),
        }
    }
}

impl std::error::Error for Error {}
pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
pub enum VideoDeviceType {
    Decoder,
    Encoder,
}

pub struct VideoDevice {
    device_type: VideoDeviceType,
    kill_evt: Option<Event>,
    resource_bridge: Option<ResourceRequestSocket>,
}

impl VideoDevice {
    pub fn new(
        device_type: VideoDeviceType,
        resource_bridge: Option<ResourceRequestSocket>,
    ) -> VideoDevice {
        VideoDevice {
            device_type,
            kill_evt: None,
            resource_bridge,
        }
    }
}

impl Drop for VideoDevice {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            // Ignore the result because there is nothing we can do about it.
            let _ = kill_evt.write(1);
        }
    }
}

impl VirtioDevice for VideoDevice {
    fn keep_fds(&self) -> Vec<RawFd> {
        let mut keep_fds = Vec::new();
        if let Some(resource_bridge) = &self.resource_bridge {
            keep_fds.push(resource_bridge.as_raw_fd());
        }
        keep_fds
    }

    fn device_type(&self) -> u32 {
        match &self.device_type {
            VideoDeviceType::Decoder => virtio::TYPE_VIDEO_DEC,
            VideoDeviceType::Encoder => virtio::TYPE_VIDEO_ENC,
        }
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn features(&self) -> u64 {
        1u64 << VIRTIO_F_VERSION_1
            | 1u64 << protocol::VIRTIO_VIDEO_F_RESOURCE_NON_CONTIG
            | 1u64 << protocol::VIRTIO_VIDEO_F_RESOURCE_VIRTIO_OBJECT
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let mut cfg = protocol::virtio_video_config {
            version: Le32::from(0),
            max_caps_length: Le32::from(1024), // Set a big number
            max_resp_length: Le32::from(1024), // Set a big number
        };
        copy_config(data, 0, cfg.as_mut_slice(), offset);
    }

    fn activate(
        &mut self,
        mem: GuestMemory,
        interrupt: Interrupt,
        mut queues: Vec<virtio::queue::Queue>,
        mut queue_evts: Vec<Event>,
    ) {
        if queues.len() != QUEUE_SIZES.len() {
            error!(
                "wrong number of queues are passed: expected {}, actual {}",
                queues.len(),
                QUEUE_SIZES.len()
            );
            return;
        }
        if queue_evts.len() != QUEUE_SIZES.len() {
            error!(
                "wrong number of events are passed: expected {}, actual {}",
                queue_evts.len(),
                QUEUE_SIZES.len()
            );
        }

        let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e))) {
            Ok(v) => v,
            Err(e) => {
                error!("failed to create kill Event pair: {:?}", e);
                return;
            }
        };
        self.kill_evt = Some(self_kill_evt);

        let cmd_queue = queues.remove(0);
        let cmd_evt = queue_evts.remove(0);
        let event_queue = queues.remove(0);
        let event_evt = queue_evts.remove(0);
        let resource_bridge = match self.resource_bridge.take() {
            Some(r) => r,
            None => {
                error!("no resource bridge is passed");
                return;
            }
        };
        let mut worker = Worker {
            interrupt,
            mem,
            cmd_evt,
            event_evt,
            kill_evt,
            resource_bridge,
        };
        let worker_result = match &self.device_type {
            VideoDeviceType::Decoder => thread::Builder::new()
                .name("virtio video decoder".to_owned())
                .spawn(move || {
                    let vda = libvda::decode::VdaInstance::new(libvda::decode::VdaImplType::Gavda)
                        .map_err(Error::LibvdaCreationFailed)?;
                    let device = decoder::Decoder::new(&vda);
                    worker.run(cmd_queue, event_queue, device)
                }),
            VideoDeviceType::Encoder => thread::Builder::new()
                .name("virtio video encoder".to_owned())
                .spawn(move || {
                    let device = encoder::Encoder::new();
                    worker.run(cmd_queue, event_queue, device)
                }),
        };
        if let Err(e) = worker_result {
            error!(
                "failed to spawn virtio_video worker for {:?}: {}",
                &self.device_type, e
            );
            return;
        }
    }
}
<commit_msg>devices: virtio: video: Ensure failures are logged<commit_after>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! This module implements the virtio video encoder and decoder devices.
//! The current implementation uses [v3 RFC] of the virtio-video protocol.
//!
//! [v3 RFC]: https://markmail.org/thread/wxdne5re7aaugbjg

use std::fmt::{self, Display};
use std::os::unix::io::{AsRawFd, RawFd};
use std::thread;

use base::{error, Error as SysError, Event};
use data_model::{DataInit, Le32};
use vm_memory::GuestMemory;

use crate::virtio::resource_bridge::ResourceRequestSocket;
use crate::virtio::virtio_device::VirtioDevice;
use crate::virtio::{self, copy_config, DescriptorError, Interrupt, VIRTIO_F_VERSION_1};

#[macro_use]
mod macros;
mod async_cmd_desc_map;
mod command;
mod control;
mod decoder;
mod device;
mod encoder;
mod error;
mod event;
mod format;
mod params;
mod protocol;
mod response;
mod worker;

use command::ReadCmdError;
use device::AsyncCmdTag;
use worker::Worker;

const QUEUE_SIZE: u16 = 256;
const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE, QUEUE_SIZE];

/// An error indicating something went wrong in virtio-video's worker.
#[derive(Debug)]
pub enum Error {
    /// Creating PollContext failed.
    PollContextCreationFailed(SysError),
    /// A DescriptorChain contains invalid data.
    InvalidDescriptorChain(DescriptorError),
    /// No available descriptor in which an event is written to.
    DescriptorNotAvailable,
    /// Error while polling for events.
    PollError(SysError),
    /// Failed to read a virtio-video command.
    ReadFailure(ReadCmdError),
    /// Got response for an unexpected asynchronous command.
    UnexpectedResponse(AsyncCmdTag),
    /// Failed to write an event into the event queue.
    WriteEventFailure {
        event: event::VideoEvt,
        error: std::io::Error,
    },
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Error::*;
        match self {
            PollContextCreationFailed(e) => write!(f, "failed to create PollContext: {}", e),
            InvalidDescriptorChain(e) => write!(f, "DescriptorChain contains invalid data: {}", e),
            DescriptorNotAvailable => {
                write!(f, "no available descriptor in which an event is written to")
            }
            PollError(err) => write!(f, "failed to poll events: {}", err),
            ReadFailure(e) => write!(f, "failed to read a command from the guest: {}", e),
            UnexpectedResponse(tag) => {
                write!(f, "got a response for an untracked command: {:?}", tag)
            }
            WriteEventFailure { event, error } => write!(
                f,
                "failed to write an event {:?} into event queue: {}",
                event, error
            ),
        }
    }
}

impl std::error::Error for Error {}
pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
pub enum VideoDeviceType {
    Decoder,
    Encoder,
}

pub struct VideoDevice {
    device_type: VideoDeviceType,
    kill_evt: Option<Event>,
    resource_bridge: Option<ResourceRequestSocket>,
}

impl VideoDevice {
    pub fn new(
        device_type: VideoDeviceType,
        resource_bridge: Option<ResourceRequestSocket>,
    ) -> VideoDevice {
        VideoDevice {
            device_type,
            kill_evt: None,
            resource_bridge,
        }
    }
}

impl Drop for VideoDevice {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            // Ignore the result because there is nothing we can do about it.
            let _ = kill_evt.write(1);
        }
    }
}

impl VirtioDevice for VideoDevice {
    fn keep_fds(&self) -> Vec<RawFd> {
        let mut keep_fds = Vec::new();
        if let Some(resource_bridge) = &self.resource_bridge {
            keep_fds.push(resource_bridge.as_raw_fd());
        }
        keep_fds
    }

    fn device_type(&self) -> u32 {
        match &self.device_type {
            VideoDeviceType::Decoder => virtio::TYPE_VIDEO_DEC,
            VideoDeviceType::Encoder => virtio::TYPE_VIDEO_ENC,
        }
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn features(&self) -> u64 {
        1u64 << VIRTIO_F_VERSION_1
            | 1u64 << protocol::VIRTIO_VIDEO_F_RESOURCE_NON_CONTIG
            | 1u64 << protocol::VIRTIO_VIDEO_F_RESOURCE_VIRTIO_OBJECT
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let mut cfg = protocol::virtio_video_config {
            version: Le32::from(0),
            max_caps_length: Le32::from(1024), // Set a big number
            max_resp_length: Le32::from(1024), // Set a big number
        };
        copy_config(data, 0, cfg.as_mut_slice(), offset);
    }

    fn activate(
        &mut self,
        mem: GuestMemory,
        interrupt: Interrupt,
        mut queues: Vec<virtio::queue::Queue>,
        mut queue_evts: Vec<Event>,
    ) {
        if queues.len() != QUEUE_SIZES.len() {
            error!(
                "wrong number of queues are passed: expected {}, actual {}",
                queues.len(),
                QUEUE_SIZES.len()
            );
            return;
        }
        if queue_evts.len() != QUEUE_SIZES.len() {
            error!(
                "wrong number of events are passed: expected {}, actual {}",
                queue_evts.len(),
                QUEUE_SIZES.len()
            );
        }

        let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e))) {
            Ok(v) => v,
            Err(e) => {
                error!("failed to create kill Event pair: {:?}", e);
                return;
            }
        };
        self.kill_evt = Some(self_kill_evt);

        let cmd_queue = queues.remove(0);
        let cmd_evt = queue_evts.remove(0);
        let event_queue = queues.remove(0);
        let event_evt = queue_evts.remove(0);
        let resource_bridge = match self.resource_bridge.take() {
            Some(r) => r,
            None => {
                error!("no resource bridge is passed");
                return;
            }
        };
        let mut worker = Worker {
            interrupt,
            mem,
            cmd_evt,
            event_evt,
            kill_evt,
            resource_bridge,
        };
        let worker_result = match &self.device_type {
            VideoDeviceType::Decoder => thread::Builder::new()
                .name("virtio video decoder".to_owned())
                .spawn(move || {
                    let vda = match libvda::decode::VdaInstance::new(
                        libvda::decode::VdaImplType::Gavda,
                    ) {
                        Ok(vda) => vda,
                        Err(e) => {
                            error!("Failed to initialize vda: {}", e);
                            return;
                        }
                    };
                    let device = decoder::Decoder::new(&vda);
                    if let Err(e) = worker.run(cmd_queue, event_queue, device) {
                        error!("Failed to start decoder worker: {}", e);
                    };
                    // Don't return any information since the return value is never checked.
                }),
            VideoDeviceType::Encoder => thread::Builder::new()
                .name("virtio video encoder".to_owned())
                .spawn(move || {
                    let device = encoder::Encoder::new();
                    if let Err(e) = worker.run(cmd_queue, event_queue, device) {
                        error!("Failed to start encoder worker: {}", e);
                    }
                }),
        };
        if let Err(e) = worker_result {
            error!(
                "failed to spawn virtio_video worker for {:?}: {}",
                &self.device_type, e
            );
            return;
        }
    }
}
<|endoftext|>


<commit_before>use aster;
use syntax::ast::{
    self,
    Item,
    MetaItem,
};
use syntax::codemap::Span;
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use attr::Attr;

pub fn expand_derive_queryable(
    cx: &mut ExtCtxt,
    span: Span,
    meta_item: &MetaItem,
    annotatable: &Annotatable,
    push: &mut FnMut(Annotatable)
) {
    if let Annotatable::Item(ref item) = *annotatable {
        let (generics, attrs) = match Attr::from_item(cx, item) {
            Some((generics, attrs)) => (generics, attrs),
            None => {
                cx.span_err(span, "`#[derive(Queryable)]` can only be applied to structs or tuple structs");
                return;
            }
        };
        let builder = aster::AstBuilder::new().span(span);

        let ty = builder.ty().path()
            .segment(item.ident).with_generics(generics.clone()).build()
            .build();

        let row_type = builder.ty().tuple()
            .with_tys(attrs.iter().map(|f| f.ty.clone()))
            .build();

        let build_impl = struct_literal_with_fields_assigned_to_row_elements(
            &item,
            &builder,
            &attrs,
        );
        let display_generics = builder.from_generics(generics)
            .ty_param_id("__ST")
            .build();

        let impl_item = quote_item!(cx,
            impl$display_generics ::diesel::Queryable<__ST> for $ty where
                __ST: ::diesel::types::NativeSqlType,
                $row_type: ::diesel::types::FromSqlRow<__ST>,
            {
                type Row = $row_type;

                fn build(row: Self::Row) -> Self {
                    $build_impl
                }
            }
        ).unwrap();

        push(Annotatable::Item(impl_item));
    } else {
        cx.span_err(meta_item.span,
                    "`derive` may only be applied to enums and structs");
    };
}

fn struct_literal_with_fields_assigned_to_row_elements(
    item: &Item,
    builder: &aster::AstBuilder,
    fields: &[Attr],
) -> P<ast::Expr> {
    let mut build_impl_builder = builder.expr().struct_path(item.ident);
    for (i, field) in fields.iter().enumerate() {
        build_impl_builder = build_impl_builder
            .field(field.field_name.unwrap())
            .tup_field(i)
            .id("row");
    }
    build_impl_builder.build()
}
<commit_msg>Remove use of `aster` from `queryable` in `codegen`<commit_after>use syntax::ast::{
    self,
    Item,
    MetaItem,
};
use syntax::codemap::Span;
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::parse::token::*;
use syntax::ptr::P;

use attr::Attr;

pub fn expand_derive_queryable(
    cx: &mut ExtCtxt,
    span: Span,
    meta_item: &MetaItem,
    annotatable: &Annotatable,
    push: &mut FnMut(Annotatable)
) {
    if let Annotatable::Item(ref item) = *annotatable {
        let (mut generics, attrs) = match Attr::from_item(cx, item) {
            Some((generics, attrs)) => (generics, attrs),
            None => {
                cx.span_err(span, "`#[derive(Queryable)]` can only be applied to structs or tuple structs");
                return;
            }
        };

        let ty = struct_ty(cx, span, item.ident, &generics);

        let row_type = cx.ty(span, ast::TyTup(attrs.iter().map(|f| f.ty.clone()).collect()));

        let build_impl = struct_literal_with_fields_assigned_to_row_elements(
            span, &item, cx, &attrs);
        let mut params = generics.ty_params.into_vec();
        params.push(ty_param_with_name(cx, span, "__ST"));
        generics.ty_params = params.into();

        let impl_item = quote_item!(cx,
            impl$generics ::diesel::Queryable<__ST> for $ty where
                __ST: ::diesel::types::NativeSqlType,
                $row_type: ::diesel::types::FromSqlRow<__ST>,
            {
                type Row = $row_type;

                fn build(row: Self::Row) -> Self {
                    $build_impl
                }
            }
        ).unwrap();

        push(Annotatable::Item(impl_item));
    } else {
        cx.span_err(meta_item.span,
                    "`derive` may only be applied to enums and structs");
    };
}

fn ty_param_with_name(cx: &mut ExtCtxt, span: Span, name: &str) -> ast::TyParam {
    cx.typaram(span, str_to_ident(name), P::empty(), None)
}

fn struct_ty(
    cx: &mut ExtCtxt,
    span: Span,
    name: ast::Ident,
    generics: &ast::Generics,
) -> P<ast::Ty> {
    let lifetimes = generics.lifetimes.iter().map(|lt| lt.lifetime).collect();
    let ty_params = generics.ty_params.iter()
        .map(|param| cx.ty_ident(span, param.ident))
        .collect();
    cx.ty_path(cx.path_all(span, false, vec![name], lifetimes, ty_params, Vec::new()))
}

fn struct_literal_with_fields_assigned_to_row_elements(
    span: Span,
    item: &Item,
    cx: &mut ExtCtxt,
    fields: &[Attr],
) -> P<ast::Expr> {
    let tup = cx.expr_ident(span, str_to_ident("row"));
    let fields = fields.iter().enumerate().map(|(i, field)| {
        cx.field_imm(
            span,
            field.field_name.unwrap(),
            cx.expr_tup_field_access(span, tup.clone(), i),
        )
    }).collect();
    cx.expr_struct_ident(span, item.ident, fields)
}
<|endoftext|>


<commit_before>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! ARM PL031 Real Time Clock
//!
//! This module implements a PL031 Real Time Clock (RTC) that provides to provides long time base counter.
//! This is achieved by generating an interrupt signal after counting for a programmed number of cycles of
//! a real-time clock input.
//!
use std::fmt;
use std::time::Instant;
use std::{io, result};

use crate::BusDevice;
use logger::{warn, IncMetric, METRICS};
use utils::byte_order;
use utils::eventfd::EventFd;
//use bus::Error;

// As you can see in https://static.docs.arm.com/ddi0224/c/real_time_clock_pl031_r1p3_technical_reference_manual_DDI0224C.pdf
// at section 3.2 Summary of RTC registers, the total size occupied by this device is 0x000 -> 0xFFC + 4 = 0x1000.
// From 0x0 to 0x1C we have following registers:
const RTCDR: u64 = 0x0; // Data Register.
const RTCMR: u64 = 0x4; // Match Register.
const RTCLR: u64 = 0x8; // Load Regiser.
const RTCCR: u64 = 0xc; // Control Register.
const RTCIMSC: u64 = 0x10; // Interrupt Mask Set or Clear Register.
const RTCRIS: u64 = 0x14; // Raw Interrupt Status.
const RTCMIS: u64 = 0x18; // Masked Interrupt Status.
const RTCICR: u64 = 0x1c; // Interrupt Clear Register.
                          // From 0x020 to 0xFDC => reserved space.
                          // From 0xFE0 to 0x1000 => Peripheral and PrimeCell Identification Registers which are Read Only registers.
                          // AMBA standard devices have CIDs (Cell IDs) and PIDs (Peripheral IDs). The linux kernel will look for these in order to assert the identity
                          // of these devices (i.e look at the `amba_device_try_add` function).
                          // We are putting the expected values (look at 'Reset value' column from above mentioned document) in an array.
const PL031_ID: [u8; 8] = [0x31, 0x10, 0x04, 0x00, 0x0d, 0xf0, 0x05, 0xb1];
// We are only interested in the margins.
const AMBA_ID_LOW: u64 = 0xFE0;
const AMBA_ID_HIGH: u64 = 0x1000;

#[derive(Debug)]
pub enum Error {
    BadWriteOffset(u64),
    InterruptFailure(io::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::BadWriteOffset(offset) => write!(f, "Bad Write Offset: {}", offset),
            Error::InterruptFailure(e) => write!(f, "Failed to trigger interrupt: {}", e),
        }
    }
}
type Result<T> = result::Result<T, Error>;

/// A RTC device following the PL031 specification..
pub struct RTC {
    previous_now: Instant,
    tick_offset: i64,
    // This is used for implementing the RTC alarm. However, in Firecracker we do not need it.
    match_value: u32,
    // Writes to this register load an update value into the RTC.
    load: u32,
    imsc: u32,
    ris: u32,
    interrupt_evt: EventFd,
}

impl RTC {
    /// Constructs an AMBA PL031 RTC device.
    pub fn new(interrupt_evt: EventFd) -> RTC {
        RTC {
            // This is used only for duration measuring purposes.
            previous_now: Instant::now(),
            tick_offset: utils::time::get_time_ns(utils::time::ClockType::Real) as i64,
            match_value: 0,
            load: 0,
            imsc: 0,
            ris: 0,
            interrupt_evt,
        }
    }

    fn trigger_interrupt(&mut self) -> Result<()> {
        self.interrupt_evt.write(1).map_err(Error::InterruptFailure)
    }

    fn get_time(&self) -> u32 {
        let ts = (self.tick_offset as i128)
            + (Instant::now().duration_since(self.previous_now).as_nanos() as i128);
        (ts / utils::time::NANOS_PER_SECOND as i128) as u32
    }

    fn handle_write(&mut self, offset: u64, val: u32) -> Result<()> {
        match offset {
            RTCMR => {
                // The MR register is used for implementing the RTC alarm. A real time clock alarm is
                // a feature that can be used to allow a computer to 'wake up' after shut down to execute
                // tasks every day or on a certain day. It can sometimes be found in the 'Power Management'
                // section of a motherboard's BIOS setup. This is functionality that extends beyond
                // Firecracker intended use. However, we increment a metric just in case.
                self.match_value = val;
                METRICS.rtc.missed_write_count.inc();
            }
            RTCLR => {
                self.load = val;
                self.previous_now = Instant::now();
                // If the unwrap fails, then the internal value of the clock has been corrupted and
                // we want to terminate the execution of the process.
                self.tick_offset = utils::time::seconds_to_nanoseconds(i64::from(val))
                    .expect("Time conversion overflow");
            }
            RTCIMSC => {
                self.imsc = val & 1;
                self.trigger_interrupt()?;
            }
            RTCICR => {
                // As per above mentioned doc, the interrupt is cleared by writing any data value to
                // the Interrupt Clear Register.
                self.ris = 0;
                self.trigger_interrupt()?;
            }
            RTCCR => (), // ignore attempts to turn off the timer.
            o => {
                return Err(Error::BadWriteOffset(o));
            }
        }
        Ok(())
    }
    /// Provides a reference to the interrupt event fd.
    pub fn interrupt_evt(&self) -> &EventFd {
        &self.interrupt_evt
    }
}

impl BusDevice for RTC {
    fn read(&mut self, offset: u64, data: &mut [u8]) {
        let v;
        let mut read_ok = true;

        if offset < AMBA_ID_HIGH && offset >= AMBA_ID_LOW {
            let index = ((offset - AMBA_ID_LOW) >> 2) as usize;
            v = u32::from(PL031_ID[index]);
        } else {
            v = match offset {
                RTCDR => self.get_time(),
                RTCMR => {
                    METRICS.rtc.missed_read_count.inc();
                    // Even though we are not implementing RTC alarm we return the last value
                    self.match_value
                }
                RTCLR => self.load,
                RTCCR => 1, // RTC is always enabled.
                RTCIMSC => self.imsc,
                RTCRIS => self.ris,
                RTCMIS => self.ris & self.imsc,
                _ => {
                    read_ok = false;
                    0
                }
            };
        }
        if read_ok && data.len() <= 4 {
            byte_order::write_le_u32(data, v);
        } else {
            warn!(
                "Invalid RTC PL031 read: offset {}, data length {}",
                offset,
                data.len()
            );
            METRICS.rtc.error_count.inc();
        }
    }

    fn write(&mut self, offset: u64, data: &[u8]) {
        if data.len() <= 4 {
            let v = byte_order::read_le_u32(&data[..]);
            if let Err(e) = self.handle_write(offset, v) {
                warn!("Failed to write to RTC PL031 device: {}", e);
                METRICS.rtc.error_count.inc();
            }
        } else {
            warn!(
                "Invalid RTC PL031 write: offset {}, data length {}",
                offset,
                data.len()
            );
            METRICS.rtc.error_count.inc();
        }
    }
}

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

    #[test]
    fn test_rtc_read_write_and_event() {
        let mut rtc = RTC::new(EventFd::new(libc::EFD_NONBLOCK).unwrap());
        let mut data = [0; 4];

        // Read and write to the MR register.
        byte_order::write_le_u32(&mut data, 123);
        rtc.write(RTCMR, &data);
        rtc.read(RTCMR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 123);

        // Read and write to the LR register.
        let v = utils::time::get_time_ns(utils::time::ClockType::Real);
        byte_order::write_le_u32(&mut data, (v / utils::time::NANOS_PER_SECOND) as u32);
        let previous_now_before = rtc.previous_now;
        rtc.write(RTCLR, &data);

        assert!(rtc.previous_now > previous_now_before);

        rtc.read(RTCLR, &mut data);
        let v_read = byte_order::read_le_u32(&data[..]);
        assert_eq!((v / utils::time::NANOS_PER_SECOND) as u32, v_read);

        // Read and write to IMSC register.
        // Test with non zero value.
        let non_zero = 1;
        byte_order::write_le_u32(&mut data, non_zero);
        rtc.write(RTCIMSC, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() == 1);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(non_zero & 1, v);

        // Now test with 0.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCIMSC, &data);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(0, v);

        // Read and write to the ICR register.
        byte_order::write_le_u32(&mut data, 1);
        rtc.write(RTCICR, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() > 1);
        let v_before = byte_order::read_le_u32(&data[..]);
        let no_errors_before = METRICS.rtc.error_count.count();

        rtc.read(RTCICR, &mut data);
        let no_errors_after = METRICS.rtc.error_count.count();
        let v = byte_order::read_le_u32(&data[..]);
        // ICR is a  write only register. Data received should stay equal to data sent.
        assert_eq!(v, v_before);
        assert_eq!(no_errors_after - no_errors_before, 1);

        // Attempts to turn off the RTC should not go through.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCCR, &data);
        rtc.read(RTCCR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 1);

        // Attempts to write beyond the writable space. Using here the space used to read
        // the CID and PID from.
        byte_order::write_le_u32(&mut data, 0);
        let no_errors_before = METRICS.rtc.error_count.count();
        rtc.write(AMBA_ID_LOW, &data);
        let no_errors_after = METRICS.rtc.error_count.count();
        assert_eq!(no_errors_after - no_errors_before, 1);
        // However, reading from the AMBA_ID_LOW should succeed upon read.

        let mut data = [0; 4];
        rtc.read(AMBA_ID_LOW, &mut data);
        let index = AMBA_ID_LOW + 3;
        assert_eq!(data[0], PL031_ID[((index - AMBA_ID_LOW) >> 2) as usize]);
    }
}
<commit_msg>rtc: unify read and write functions<commit_after>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! ARM PL031 Real Time Clock
//!
//! This module implements a PL031 Real Time Clock (RTC) that provides to provides long time base counter.
//! This is achieved by generating an interrupt signal after counting for a programmed number of cycles of
//! a real-time clock input.
//!
use std::fmt;
use std::time::Instant;
use std::{io, result};

use crate::BusDevice;
use logger::{warn, IncMetric, METRICS};
use utils::byte_order;
use utils::eventfd::EventFd;
//use bus::Error;

// As you can see in https://static.docs.arm.com/ddi0224/c/real_time_clock_pl031_r1p3_technical_reference_manual_DDI0224C.pdf
// at section 3.2 Summary of RTC registers, the total size occupied by this device is 0x000 -> 0xFFC + 4 = 0x1000.
// From 0x0 to 0x1C we have following registers:
const RTCDR: u64 = 0x0; // Data Register.
const RTCMR: u64 = 0x4; // Match Register.
const RTCLR: u64 = 0x8; // Load Regiser.
const RTCCR: u64 = 0xc; // Control Register.
const RTCIMSC: u64 = 0x10; // Interrupt Mask Set or Clear Register.
const RTCRIS: u64 = 0x14; // Raw Interrupt Status.
const RTCMIS: u64 = 0x18; // Masked Interrupt Status.
const RTCICR: u64 = 0x1c; // Interrupt Clear Register.
                          // From 0x020 to 0xFDC => reserved space.
                          // From 0xFE0 to 0x1000 => Peripheral and PrimeCell Identification Registers which are Read Only registers.
                          // AMBA standard devices have CIDs (Cell IDs) and PIDs (Peripheral IDs). The linux kernel will look for these in order to assert the identity
                          // of these devices (i.e look at the `amba_device_try_add` function).
                          // We are putting the expected values (look at 'Reset value' column from above mentioned document) in an array.
const PL031_ID: [u8; 8] = [0x31, 0x10, 0x04, 0x00, 0x0d, 0xf0, 0x05, 0xb1];
// We are only interested in the margins.
const AMBA_ID_LOW: u64 = 0xFE0;
const AMBA_ID_HIGH: u64 = 0x1000;

pub enum Error {
    BadOffset(u64, &'static str),
    InterruptFailure(io::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::BadOffset(offset, ops) => write!(f, "Bad {} offset: {}", ops, offset),
            Error::InterruptFailure(e) => write!(f, "Failed to trigger interrupt: {}", e),
        }
    }
}
type Result<T> = result::Result<T, Error>;

/// A RTC device following the PL031 specification..
pub struct RTC {
    previous_now: Instant,
    tick_offset: i64,
    // This is used for implementing the RTC alarm. However, in Firecracker we do not need it.
    match_value: u32,
    // Writes to this register load an update value into the RTC.
    load: u32,
    imsc: u32,
    ris: u32,
    interrupt_evt: EventFd,
}

impl RTC {
    /// Constructs an AMBA PL031 RTC device.
    pub fn new(interrupt_evt: EventFd) -> RTC {
        RTC {
            // This is used only for duration measuring purposes.
            previous_now: Instant::now(),
            tick_offset: utils::time::get_time_ns(utils::time::ClockType::Real) as i64,
            match_value: 0,
            load: 0,
            imsc: 0,
            ris: 0,
            interrupt_evt,
        }
    }

    fn trigger_interrupt(&mut self) -> Result<()> {
        self.interrupt_evt.write(1).map_err(Error::InterruptFailure)
    }

    fn get_time(&self) -> u32 {
        let ts = (self.tick_offset as i128)
            + (Instant::now().duration_since(self.previous_now).as_nanos() as i128);
        (ts / utils::time::NANOS_PER_SECOND as i128) as u32
    }

    fn handle_read(&mut self, offset: u64) -> Result<u32> {
        let val;

        if offset < AMBA_ID_HIGH && offset >= AMBA_ID_LOW {
            let index = ((offset - AMBA_ID_LOW) >> 2) as usize;
            val = u32::from(PL031_ID[index]);
        } else {
            val = match offset {
                RTCDR => self.get_time(),
                RTCMR => {
                    METRICS.rtc.missed_read_count.inc();
                    // Even though we are not implementing RTC alarm we return the last value
                    self.match_value
                }
                RTCLR => self.load,
                RTCCR => 1, // RTC is always enabled.
                RTCIMSC => self.imsc,
                RTCRIS => self.ris,
                RTCMIS => self.ris & self.imsc,
                off => {
                    return Err(Error::BadOffset(off, "read"));
                }
            };
        }
        Ok(val)
    }

    fn handle_write(&mut self, offset: u64, val: u32) -> Result<()> {
        match offset {
            RTCMR => {
                // The MR register is used for implementing the RTC alarm. A real time clock alarm is
                // a feature that can be used to allow a computer to 'wake up' after shut down to execute
                // tasks every day or on a certain day. It can sometimes be found in the 'Power Management'
                // section of a motherboard's BIOS setup. This is functionality that extends beyond
                // Firecracker intended use. However, we increment a metric just in case.
                self.match_value = val;
                METRICS.rtc.missed_write_count.inc();
            }
            RTCLR => {
                self.load = val;
                self.previous_now = Instant::now();
                // If the unwrap fails, then the internal value of the clock has been corrupted and
                // we want to terminate the execution of the process.
                self.tick_offset = utils::time::seconds_to_nanoseconds(i64::from(val))
                    .expect("Time conversion overflow");
            }
            RTCIMSC => {
                self.imsc = val & 1;
                self.trigger_interrupt()?;
            }
            RTCICR => {
                // As per above mentioned doc, the interrupt is cleared by writing any data value to
                // the Interrupt Clear Register.
                self.ris = 0;
                self.trigger_interrupt()?;
            }
            RTCCR => (), // ignore attempts to turn off the timer.
            off => {
                return Err(Error::BadOffset(off, "write"));
            }
        }
        Ok(())
    }
    /// Provides a reference to the interrupt event fd.
    pub fn interrupt_evt(&self) -> &EventFd {
        &self.interrupt_evt
    }
}

impl BusDevice for RTC {
    fn read(&mut self, offset: u64, data: &mut [u8]) {
        if data.len() <= 4 {
            match self.handle_read(offset) {
                Ok(val) => byte_order::write_le_u32(data, val),
                Err(e) => {
                    warn!("Failed to read from the RTC PL031 device: {}", e);
                    METRICS.rtc.error_count.inc();
                }
            }
        } else {
            warn!("Invalid RTC PL031 data length {}", data.len());
            METRICS.rtc.error_count.inc();
        }
    }

    fn write(&mut self, offset: u64, data: &[u8]) {
        if data.len() <= 4 {
            let v = byte_order::read_le_u32(&data[..]);
            if let Err(e) = self.handle_write(offset, v) {
                warn!("Failed to write to the RTC PL031 device: {}", e);
                METRICS.rtc.error_count.inc();
            }
        } else {
            warn!("Invalid RTC PL031 data length {}", data.len());
            METRICS.rtc.error_count.inc();
        }
    }
}

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

    #[test]
    fn test_rtc_read_write_and_event() {
        let mut rtc = RTC::new(EventFd::new(libc::EFD_NONBLOCK).unwrap());
        let mut data = [0; 4];

        // Read and write to the MR register.
        byte_order::write_le_u32(&mut data, 123);
        rtc.write(RTCMR, &data);
        rtc.read(RTCMR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 123);

        // Read and write to the LR register.
        let v = utils::time::get_time_ns(utils::time::ClockType::Real);
        byte_order::write_le_u32(&mut data, (v / utils::time::NANOS_PER_SECOND) as u32);
        let previous_now_before = rtc.previous_now;
        rtc.write(RTCLR, &data);

        assert!(rtc.previous_now > previous_now_before);

        rtc.read(RTCLR, &mut data);
        let v_read = byte_order::read_le_u32(&data[..]);
        assert_eq!((v / utils::time::NANOS_PER_SECOND) as u32, v_read);

        // Read and write to IMSC register.
        // Test with non zero value.
        let non_zero = 1;
        byte_order::write_le_u32(&mut data, non_zero);
        rtc.write(RTCIMSC, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() == 1);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(non_zero & 1, v);

        // Now test with 0.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCIMSC, &data);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(0, v);

        // Read and write to the ICR register.
        byte_order::write_le_u32(&mut data, 1);
        rtc.write(RTCICR, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() > 1);
        let v_before = byte_order::read_le_u32(&data[..]);
        let no_errors_before = METRICS.rtc.error_count.count();

        rtc.read(RTCICR, &mut data);
        let no_errors_after = METRICS.rtc.error_count.count();
        let v = byte_order::read_le_u32(&data[..]);
        // ICR is a  write only register. Data received should stay equal to data sent.
        assert_eq!(v, v_before);
        assert_eq!(no_errors_after - no_errors_before, 1);

        // Attempts to turn off the RTC should not go through.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCCR, &data);
        rtc.read(RTCCR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 1);

        // Attempts to write beyond the writable space. Using here the space used to read
        // the CID and PID from.
        byte_order::write_le_u32(&mut data, 0);
        let no_errors_before = METRICS.rtc.error_count.count();
        rtc.write(AMBA_ID_LOW, &data);
        let no_errors_after = METRICS.rtc.error_count.count();
        assert_eq!(no_errors_after - no_errors_before, 1);
        // However, reading from the AMBA_ID_LOW should succeed upon read.

        let mut data = [0; 4];
        rtc.read(AMBA_ID_LOW, &mut data);
        let index = AMBA_ID_LOW + 3;
        assert_eq!(data[0], PL031_ID[((index - AMBA_ID_LOW) >> 2) as usize]);
    }
}
<|endoftext|>


<commit_before>#![deny(unsafe_op_in_unsafe_fn)]

use crate::alloc::{GlobalAlloc, Layout, System};
use crate::ffi::c_void;
use crate::ptr;
use crate::sync::atomic::{AtomicPtr, Ordering};
use crate::sys::c;
use crate::sys_common::alloc::{realloc_fallback, MIN_ALIGN};

#[cfg(test)]
mod tests;

// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
// See https://docs.microsoft.com/windows/win32/api/heapapi/

// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
const HEAP_ZERO_MEMORY: c::DWORD = 0x00000008;

extern "system" {
    // Get a handle to the default heap of the current process, or null if the operation fails.
    //
    // SAFETY: Successful calls to this function within the same process are assumed to
    // always return the same handle, which remains valid for the entire lifetime of the process.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
    fn GetProcessHeap() -> c::HANDLE;

    // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
    // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
    // set to `HEAP_ZERO_MEMORY`.
    //
    // Returns a pointer to the newly-allocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
    fn HeapAlloc(hHeap: c::HANDLE, dwFlags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID;

    // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
    // to a block of at least `dwBytes` bytes, either shrinking the block in place,
    // or allocating at a new location, copying memory, and freeing the original location.
    //
    // Returns a pointer to the reallocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    // If the operation fails the given block will never have been freed.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
    //     `HeapReAlloc`, that has not already been freed.
    // If the block was successfully reallocated at a new location, pointers pointing to
    // the freed memory, such as `lpMem`, must not be dereferenced ever again.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
    fn HeapReAlloc(
        hHeap: c::HANDLE,
        dwFlags: c::DWORD,
        lpMem: c::LPVOID,
        dwBytes: c::SIZE_T,
    ) -> c::LPVOID;

    // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
    // Returns a nonzero value if the operation is successful, and zero if the operation fails.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
    //     that has not already been freed.
    // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
    // must not be dereferenced ever again.
    //
    // Note that `lpMem` is allowed to be null, which will not cause the operation to fail.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
    fn HeapFree(hHeap: c::HANDLE, dwFlags: c::DWORD, lpMem: c::LPVOID) -> c::BOOL;
}

// Cached handle to the default heap of the current process.
// Either a non-null handle returned by `GetProcessHeap`, or null when not yet initialized or `GetProcessHeap` failed.
static HEAP: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());

// Get a handle to the default heap of the current process, or null if the operation fails.
// SAFETY: If this operation is successful, `HEAP` will be successfully initialized and contain
// a non-null handle returned by `GetProcessHeap`.
#[inline]
unsafe fn init_or_get_process_heap() -> c::HANDLE {
    let heap = HEAP.load(Ordering::Relaxed);
    if heap.is_null() {
        // `HEAP` has not yet been successfully initialized
        let heap = unsafe { GetProcessHeap() };
        if !heap.is_null() {
            // SAFETY: No locking is needed because within the same process,
            // successful calls to `GetProcessHeap` will always return the same value, even on different threads.
            HEAP.store(heap, Ordering::Relaxed);

            // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
            heap
        } else {
            // Could not get the current process heap.
            ptr::null_mut()
        }
    } else {
        // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
        heap
    }
}

// Header containing a pointer to the start of an allocated block.
// SAFETY: Size and alignment must be <= `MIN_ALIGN`.
#[repr(C)]
struct Header(*mut u8);

// Allocate a block of optionally zeroed memory for a given `layout`.
// SAFETY: Returns a pointer satisfying the guarantees of `System` about allocated pointers.
#[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
    let heap = unsafe { init_or_get_process_heap() };
    if heap.is_null() {
        // Allocation has failed, could not get the current process heap.
        return ptr::null_mut();
    }

    // Allocated memory will be either zeroed or uninitialized.
    let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

    if layout.align() <= MIN_ALIGN {
        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        // The returned pointer points to the start of an allocated block.
        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
    } else {
        // Allocate extra padding in order to be able to satisfy the alignment.
        let total = layout.align() + layout.size();

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
        if ptr.is_null() {
            // Allocation has failed.
            return ptr::null_mut();
        }

        // Create a correctly aligned pointer offset from the start of the allocated block,
        // and write a header before it.

        let offset = layout.align() - (ptr as usize & (layout.align() - 1));
        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
        // least `MIN_ALIGN` bytes of padding before it.
        let aligned = unsafe { ptr.add(offset) };
        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
        // it, it is safe to write a header directly before it.
        unsafe { ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };

        // SAFETY: The returned pointer does not point to the to the start of an allocated block,
        // but there is a header readable directly before it containing the location of the start
        // of the block.
        aligned
    }
}

// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
// following properties:
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
// Instead there will be a header readable directly before the returned pointer, containing the actual
// location of the start of the block.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = false;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = true;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let block = {
            if layout.align() <= MIN_ALIGN {
                ptr
            } else {
                // The location of the start of the block is stored in the padding before `ptr`.

                // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                // and have a header readable directly before it.
                unsafe { ptr::read((ptr as *mut Header).offset(-1)).0 }
            }
        };

        // SAFETY: because `ptr` has been successfully allocated with this allocator,
        // `HEAP` must have been successfully initialized and contain a non-null handle
        // returned by `GetProcessHeap`.
        let heap = HEAP.load(Ordering::Relaxed);

        // SAFETY: `block` is a pointer to the start of an allocated block.
        unsafe {
            let err = HeapFree(heap, 0, block as c::LPVOID);
            debug_assert!(err != 0, "Failed to free heap memory: {}", c::GetLastError());
        }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            // SAFETY: because `ptr` has been successfully allocated with this allocator,
            // `HEAP` must have been successfully initialized and contain a non-null handle
            // returned by `GetProcessHeap`.
            let heap = HEAP.load(Ordering::Relaxed);

            // SAFETY: `ptr` is a pointer to the start of an allocated block.
            // The returned pointer points to the start of an allocated block.
            unsafe { HeapReAlloc(heap, 0, ptr as c::LPVOID, new_size) as *mut u8 }
        } else {
            // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
            // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
            unsafe { realloc_fallback(self, ptr, layout, new_size) }
        }
    }
}
<commit_msg>Introduce `get_process_heap` and fix atomic ordering.<commit_after>#![deny(unsafe_op_in_unsafe_fn)]

use crate::alloc::{GlobalAlloc, Layout, System};
use crate::ffi::c_void;
use crate::ptr;
use crate::sync::atomic::{AtomicPtr, Ordering};
use crate::sys::c;
use crate::sys_common::alloc::{realloc_fallback, MIN_ALIGN};

#[cfg(test)]
mod tests;

// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
// See https://docs.microsoft.com/windows/win32/api/heapapi/

// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
const HEAP_ZERO_MEMORY: c::DWORD = 0x00000008;

extern "system" {
    // Get a handle to the default heap of the current process, or null if the operation fails.
    //
    // SAFETY: Successful calls to this function within the same process are assumed to
    // always return the same handle, which remains valid for the entire lifetime of the process.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
    fn GetProcessHeap() -> c::HANDLE;

    // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
    // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
    // set to `HEAP_ZERO_MEMORY`.
    //
    // Returns a pointer to the newly-allocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
    fn HeapAlloc(hHeap: c::HANDLE, dwFlags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID;

    // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
    // to a block of at least `dwBytes` bytes, either shrinking the block in place,
    // or allocating at a new location, copying memory, and freeing the original location.
    //
    // Returns a pointer to the reallocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    // If the operation fails the given block will never have been freed.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
    //     `HeapReAlloc`, that has not already been freed.
    // If the block was successfully reallocated at a new location, pointers pointing to
    // the freed memory, such as `lpMem`, must not be dereferenced ever again.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
    fn HeapReAlloc(
        hHeap: c::HANDLE,
        dwFlags: c::DWORD,
        lpMem: c::LPVOID,
        dwBytes: c::SIZE_T,
    ) -> c::LPVOID;

    // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
    // Returns a nonzero value if the operation is successful, and zero if the operation fails.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
    //     that has not already been freed.
    // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
    // must not be dereferenced ever again.
    //
    // Note that `lpMem` is allowed to be null, which will not cause the operation to fail.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
    fn HeapFree(hHeap: c::HANDLE, dwFlags: c::DWORD, lpMem: c::LPVOID) -> c::BOOL;
}

// Cached handle to the default heap of the current process.
// Either a non-null handle returned by `GetProcessHeap`, or null when not yet initialized or `GetProcessHeap` failed.
static HEAP: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());

// Get a handle to the default heap of the current process, or null if the operation fails.
// If this operation is successful, `HEAP` will be successfully initialized and contain
// a non-null handle returned by `GetProcessHeap`.
#[inline]
fn init_or_get_process_heap() -> c::HANDLE {
    let heap = HEAP.load(Ordering::Relaxed);
    if heap.is_null() {
        // `HEAP` has not yet been successfully initialized
        let heap = unsafe { GetProcessHeap() };
        if !heap.is_null() {
            // SAFETY: No locking is needed because within the same process,
            // successful calls to `GetProcessHeap` will always return the same value, even on different threads.
            HEAP.store(heap, Ordering::Release);

            // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
            heap
        } else {
            // Could not get the current process heap.
            ptr::null_mut()
        }
    } else {
        // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
        heap
    }
}

// Get a non-null handle to the default heap of the current process.
// SAFETY: `HEAP` must have been successfully initialized.
#[inline]
unsafe fn get_process_heap() -> c::HANDLE {
    HEAP.load(Ordering::Acquire)
}

// Header containing a pointer to the start of an allocated block.
// SAFETY: Size and alignment must be <= `MIN_ALIGN`.
#[repr(C)]
struct Header(*mut u8);

// Allocate a block of optionally zeroed memory for a given `layout`.
// SAFETY: Returns a pointer satisfying the guarantees of `System` about allocated pointers,
// or null if the operation fails. If this returns non-null `HEAP` will have been successfully
// initialized.
#[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
    let heap = init_or_get_process_heap();
    if heap.is_null() {
        // Allocation has failed, could not get the current process heap.
        return ptr::null_mut();
    }

    // Allocated memory will be either zeroed or uninitialized.
    let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

    if layout.align() <= MIN_ALIGN {
        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        // The returned pointer points to the start of an allocated block.
        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
    } else {
        // Allocate extra padding in order to be able to satisfy the alignment.
        let total = layout.align() + layout.size();

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
        if ptr.is_null() {
            // Allocation has failed.
            return ptr::null_mut();
        }

        // Create a correctly aligned pointer offset from the start of the allocated block,
        // and write a header before it.

        let offset = layout.align() - (ptr as usize & (layout.align() - 1));
        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
        // least `MIN_ALIGN` bytes of padding before it.
        let aligned = unsafe { ptr.add(offset) };
        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
        // it, it is safe to write a header directly before it.
        unsafe { ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };

        // SAFETY: The returned pointer does not point to the to the start of an allocated block,
        // but there is a header readable directly before it containing the location of the start
        // of the block.
        aligned
    }
}

// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
// following properties:
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
// Instead there will be a header readable directly before the returned pointer, containing the actual
// location of the start of the block.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = false;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = true;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let block = {
            if layout.align() <= MIN_ALIGN {
                ptr
            } else {
                // The location of the start of the block is stored in the padding before `ptr`.

                // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                // and have a header readable directly before it.
                unsafe { ptr::read((ptr as *mut Header).offset(-1)).0 }
            }
        };

        // SAFETY: because `ptr` has been successfully allocated with this allocator,
        // `HEAP` must have been successfully initialized.
        let heap = unsafe { get_process_heap() };

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
        // `block` is a pointer to the start of an allocated block.
        unsafe {
            let err = HeapFree(heap, 0, block as c::LPVOID);
            debug_assert!(err != 0, "Failed to free heap memory: {}", c::GetLastError());
        }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            // SAFETY: because `ptr` has been successfully allocated with this allocator,
            // `HEAP` must have been successfully initialized.
            let heap = unsafe { get_process_heap() };

            // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
            // `ptr` is a pointer to the start of an allocated block.
            // The returned pointer points to the start of an allocated block.
            unsafe { HeapReAlloc(heap, 0, ptr as c::LPVOID, new_size) as *mut u8 }
        } else {
            // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
            // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
            unsafe { realloc_fallback(self, ptr, layout, new_size) }
        }
    }
}
<|endoftext|>


<commit_before>//! Collection types.

#![stable(feature = "rust1", since = "1.0.0")]

#[cfg(not(no_global_oom_handling))]
pub mod binary_heap;
#[cfg(not(no_global_oom_handling))]
mod btree;
#[cfg(not(no_global_oom_handling))]
pub mod linked_list;
#[cfg(not(no_global_oom_handling))]
pub mod vec_deque;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
pub mod btree_map {
    //! A map based on a B-Tree.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub use super::btree::map::*;
}

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
pub mod btree_set {
    //! A set based on a B-Tree.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub use super::btree::set::*;
}

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use binary_heap::BinaryHeap;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use btree_map::BTreeMap;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use btree_set::BTreeSet;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use linked_list::LinkedList;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use vec_deque::VecDeque;

use crate::alloc::{Layout, LayoutError};
use core::fmt::Display;

/// The error type for `try_reserve` methods.
#[derive(Clone, PartialEq, Eq, Debug)]
#[unstable(feature = "try_reserve", reason = "new API", issue = "48043")]
pub struct TryReserveError {
    kind: TryReserveErrorKind,
}

impl TryReserveError {
    /// Details about the allocation that caused the error
    #[inline]
    #[unstable(
        feature = "try_reserve_kind",
        reason = "Uncertain how much info should be exposed",
        issue = "48043"
    )]
    pub fn kind(&self) -> TryReserveErrorKind {
        self.kind.clone()
    }
}

/// Details of the allocation that caused a `TryReserveError`
#[derive(Clone, PartialEq, Eq, Debug)]
#[unstable(
    feature = "try_reserve_kind",
    reason = "Uncertain how much info should be exposed",
    issue = "48043"
)]
pub enum TryReserveErrorKind {
    /// Error due to the computed capacity exceeding the collection's maximum
    /// (usually `isize::MAX` bytes).
    CapacityOverflow,

    /// The memory allocator returned an error
    AllocError {
        /// The layout of allocation request that failed
        layout: Layout,

        #[doc(hidden)]
        #[unstable(
            feature = "container_error_extra",
            issue = "none",
            reason = "\
            Enable exposing the allocator’s custom error value \
            if an associated type is added in the future: \
            https://github.com/rust-lang/wg-allocators/issues/23"
        )]
        non_exhaustive: (),
    },
}

#[unstable(
    feature = "try_reserve_kind",
    reason = "Uncertain how much info should be exposed",
    issue = "48043"
)]
impl From<TryReserveErrorKind> for TryReserveError {
    #[inline]
    fn from(kind: TryReserveErrorKind) -> Self {
        Self { kind }
    }
}

#[unstable(feature = "try_reserve", reason = "new API", issue = "48043")]
impl From<LayoutError> for TryReserveError {
    /// Always evaluates to [`TryReserveErrorKind::CapacityOverflow`].
    #[inline]
    fn from(_: LayoutError) -> Self {
        TryReserveErrorKind::CapacityOverflow.into()
    }
}

#[unstable(feature = "try_reserve", reason = "new API", issue = "48043")]
impl Display for TryReserveError {
    fn fmt(
        &self,
        fmt: &mut core::fmt::Formatter<'_>,
    ) -> core::result::Result<(), core::fmt::Error> {
        fmt.write_str("memory allocation failed")?;
        let reason = match self.kind {
            TryReserveErrorKind::CapacityOverflow => {
                " because the computed capacity exceeded the collection's maximum"
            }
            TryReserveErrorKind::AllocError { .. } => {
                " because the memory allocator returned a error"
            }
        };
        fmt.write_str(reason)
    }
}

/// An intermediate trait for specialization of `Extend`.
#[doc(hidden)]
trait SpecExtend<I: IntoIterator> {
    /// Extends `self` with the contents of the given iterator.
    fn spec_extend(&mut self, iter: I);
}
<commit_msg>Rollup merge of #88216 - kornelski:from_layout_err, r=kennytm<commit_after>//! Collection types.

#![stable(feature = "rust1", since = "1.0.0")]

#[cfg(not(no_global_oom_handling))]
pub mod binary_heap;
#[cfg(not(no_global_oom_handling))]
mod btree;
#[cfg(not(no_global_oom_handling))]
pub mod linked_list;
#[cfg(not(no_global_oom_handling))]
pub mod vec_deque;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
pub mod btree_map {
    //! A map based on a B-Tree.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub use super::btree::map::*;
}

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
pub mod btree_set {
    //! A set based on a B-Tree.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub use super::btree::set::*;
}

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use binary_heap::BinaryHeap;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use btree_map::BTreeMap;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use btree_set::BTreeSet;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use linked_list::LinkedList;

#[cfg(not(no_global_oom_handling))]
#[stable(feature = "rust1", since = "1.0.0")]
#[doc(no_inline)]
pub use vec_deque::VecDeque;

use crate::alloc::{Layout, LayoutError};
use core::fmt::Display;

/// The error type for `try_reserve` methods.
#[derive(Clone, PartialEq, Eq, Debug)]
#[unstable(feature = "try_reserve", reason = "new API", issue = "48043")]
pub struct TryReserveError {
    kind: TryReserveErrorKind,
}

impl TryReserveError {
    /// Details about the allocation that caused the error
    #[inline]
    #[unstable(
        feature = "try_reserve_kind",
        reason = "Uncertain how much info should be exposed",
        issue = "48043"
    )]
    pub fn kind(&self) -> TryReserveErrorKind {
        self.kind.clone()
    }
}

/// Details of the allocation that caused a `TryReserveError`
#[derive(Clone, PartialEq, Eq, Debug)]
#[unstable(
    feature = "try_reserve_kind",
    reason = "Uncertain how much info should be exposed",
    issue = "48043"
)]
pub enum TryReserveErrorKind {
    /// Error due to the computed capacity exceeding the collection's maximum
    /// (usually `isize::MAX` bytes).
    CapacityOverflow,

    /// The memory allocator returned an error
    AllocError {
        /// The layout of allocation request that failed
        layout: Layout,

        #[doc(hidden)]
        #[unstable(
            feature = "container_error_extra",
            issue = "none",
            reason = "\
            Enable exposing the allocator’s custom error value \
            if an associated type is added in the future: \
            https://github.com/rust-lang/wg-allocators/issues/23"
        )]
        non_exhaustive: (),
    },
}

#[unstable(
    feature = "try_reserve_kind",
    reason = "Uncertain how much info should be exposed",
    issue = "48043"
)]
impl From<TryReserveErrorKind> for TryReserveError {
    #[inline]
    fn from(kind: TryReserveErrorKind) -> Self {
        Self { kind }
    }
}

#[unstable(feature = "try_reserve_kind", reason = "new API", issue = "48043")]
impl From<LayoutError> for TryReserveErrorKind {
    /// Always evaluates to [`TryReserveErrorKind::CapacityOverflow`].
    #[inline]
    fn from(_: LayoutError) -> Self {
        TryReserveErrorKind::CapacityOverflow
    }
}

#[unstable(feature = "try_reserve", reason = "new API", issue = "48043")]
impl Display for TryReserveError {
    fn fmt(
        &self,
        fmt: &mut core::fmt::Formatter<'_>,
    ) -> core::result::Result<(), core::fmt::Error> {
        fmt.write_str("memory allocation failed")?;
        let reason = match self.kind {
            TryReserveErrorKind::CapacityOverflow => {
                " because the computed capacity exceeded the collection's maximum"
            }
            TryReserveErrorKind::AllocError { .. } => {
                " because the memory allocator returned a error"
            }
        };
        fmt.write_str(reason)
    }
}

/// An intermediate trait for specialization of `Extend`.
#[doc(hidden)]
trait SpecExtend<I: IntoIterator> {
    /// Extends `self` with the contents of the given iterator.
    fn spec_extend(&mut self, iter: I);
}
<|endoftext|>


<commit_before>#[derive(Debug, Clone)]
struct Tree {
    enter: Node,
    left: Option<Box<Tree>>,
    right: Option<Box<Tree>>,
}

//for red-black tree
#[derive(Debug, Clone)]
enum Color {
    Black,
    Red,
    White,
}

#[derive(Debug, Clone)]
struct Node {
    val: i32,
    color: Color,
}

impl Node {
    fn new(v: &i32) -> Node {
        Node {
            val: *v,
            color: Color::White,
        }
    }
}

type ReTr<'a> = Result<&'a mut Tree, String>;

impl Tree {
    fn new(v: &i32) -> Tree {
        Tree {
            enter: Node::new(v),
            left: None,
            right: None,
        }
    }

    fn insert(&mut self, v: &i32) {
        let target = if *v > self.enter.val {
            &mut self.right
        } else {
            &mut self.left
        };
        //println!("target is -> {:?}", target);
        match target {
            Some(next) => {
                //println!("next is -> {:?}", &next);
                next.insert(v)
            }
            None => {
                let new_node = Node::new(v);
                let new_tree = Tree {
                    enter: new_node,
                    left: None,
                    right: None,
                };
                *target = Some(Box::new(new_tree));
            }
        }
    }

    fn look_up_recursive(&mut self, v: &i32) -> Result<&mut Self, String> {
        if *v == self.enter.val {
            return Ok(self);
        } else if *v > self.enter.val {
            if let Some(_) = self.right {
                let next = self.right.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        } else {
            if let Some(_) = self.left {
                let next = self.left.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        }
    }

    fn look_up(&mut self, v: &i32) -> ReTr {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return Ok(this);
            } else if *v > this.enter.val {
                if let None = this.right {
                    return Err("not find".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return Err("not find".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn look_up_father(&mut self, v: &i32) -> (ReTr, String) {
        let mut this: &mut Self = self;
        loop {
            if *v == this.right.as_ref().unwrap().enter.val {
                return (Ok(this), "right".to_string());
            } else if *v == this.left.as_ref().unwrap().enter.val {
                return (Ok(this), "left".to_string());
            } else if *v > this.enter.val {
                if let None = this.right {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn cut_left_tree(&mut self) -> Option<Box<Self>> {
        self.left.clone()
    }

    fn cut_right_tree(&mut self) -> Option<Box<Self>> {
        self.right.clone()
    }

    fn left_shift(&mut self, v: &i32) {
        let father = self.look_up_father(v).0.unwrap();
        let mut this = father.cut_left_tree().unwrap();
        let mut son = this.cut_right_tree().unwrap();
        let grand_son = son.cut_left_tree().unwrap();

        this.right = Some(grand_son);
        son.left = Some(this);
        father.left = Some(son);
    }

    fn right_shift(&mut self, v: &i32) {
        let father = self.look_up_father(v).0.unwrap();
        let mut this = father.cut_right_tree().unwrap();
        let mut son = this.cut_left_tree().unwrap();
        let grand_son = son.cut_right_tree().unwrap();

        this.left = Some(grand_son);
        son.right = Some(this);
        father.right = Some(son);
    }
}

fn main() {
    //let mut a = Tree::new(&3);
    //a.insert(&2);
    //a.insert(&1);
    //a.insert(&4);
    //a.insert(&5);

    //println!("{:?}", a);
    //println!("{:?}", a.look_up(&1));
    //println!("{:?}", a.look_up_father(&4));
    //println!("{:?}", a);

    //let mut b = Tree::new(&10);
    //b.insert(&2);
    //b.insert(&1);
    //b.insert(&5);
    //b.insert(&3);
    //b.insert(&6);
    //b.insert(&11);

    //println!("{:?}", b);
    //println!("{:?}", b.look_up_father(&2));
    //b.left_shift(&2);
    //println!("{:?}", b);

    let mut c = Tree::new(&2);
    c.insert(&8);
    c.insert(&1);
    c.insert(&5);
    c.insert(&4);
    c.insert(&7);
    c.insert(&9);
    //println!("{:?}", c.look_up_father(&8));
    c.right_shift(&8);
    println!("{:?}", c);
}
<commit_msg>add color<commit_after>#[derive(Debug, Clone)]
struct Tree {
    enter: Node,
    left: Option<Box<Tree>>,
    right: Option<Box<Tree>>,
}

//for red-black tree
#[derive(Debug, Clone)]
enum Color {
    Black,
    Red,
    White,
}

#[derive(Debug, Clone)]
struct Node {
    val: i32,
    color: Color,
}

impl Node {
    fn new(v: &i32) -> Node {
        Node {
            val: *v,
            color: Color::White,
        }
    }

    fn make_black(&mut self) {
        self.color = Color::Black;
    }

    fn make_red(&mut self) {
        self.color = Color::Red;
    }
}

type ReTr<'a> = Result<&'a mut Tree, String>;

impl Tree {
    fn new(v: &i32) -> Tree {
        let mut root = Node::new(v);
        root.make_black();
        Tree {
            enter: root,
            left: None,
            right: None,
        }
    }

    fn insert(&mut self, v: &i32) {
        let target = if *v > self.enter.val {
            &mut self.right
        } else {
            &mut self.left
        };

        match target {
            Some(next) => next.insert(v),
            None => {
                let mut new_node = Node::new(v);
                new_node.make_red();
                let new_tree = Tree {
                    enter: new_node,
                    left: None,
                    right: None,
                };
                *target = Some(Box::new(new_tree));
            }
        }
    }

    fn look_up_recursive(&mut self, v: &i32) -> ReTr {
        if *v == self.enter.val {
            return Ok(self);
        } else if *v > self.enter.val {
            if let Some(_) = self.right {
                let next = self.right.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        } else {
            if let Some(_) = self.left {
                let next = self.left.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        }
    }

    fn look_up(&mut self, v: &i32) -> ReTr {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return Ok(this);
            } else if *v > this.enter.val {
                if let None = this.right {
                    return Err("not find".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return Err("not find".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn look_up_father(&mut self, v: &i32) -> (ReTr, String) {
        let mut this: &mut Self = self;
        loop {
            if *v == this.right.as_ref().unwrap().enter.val {
                return (Ok(this), "right".to_string());
            } else if *v == this.left.as_ref().unwrap().enter.val {
                return (Ok(this), "left".to_string());
            } else if *v > this.enter.val {
                if let None = this.right {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn cut_left_tree(&mut self) -> Option<Box<Self>> {
        self.left.clone()
    }

    fn cut_right_tree(&mut self) -> Option<Box<Self>> {
        self.right.clone()
    }

    fn left_shift(&mut self, v: &i32) {
        let father = self.look_up_father(v).0.unwrap();
        let mut this = father.cut_left_tree().unwrap();
        let mut son = this.cut_right_tree().unwrap();
        let grand_son = son.cut_left_tree().unwrap();

        this.right = Some(grand_son);
        son.left = Some(this);
        father.left = Some(son);
    }

    fn right_shift(&mut self, v: &i32) {
        let father = self.look_up_father(v).0.unwrap();
        let mut this = father.cut_right_tree().unwrap();
        let mut son = this.cut_left_tree().unwrap();
        let grand_son = son.cut_right_tree().unwrap();

        this.left = Some(grand_son);
        son.right = Some(this);
        father.right = Some(son);
    }
}

fn main() {
    //let mut a = Tree::new(&3);
    //a.insert(&2);
    //a.insert(&1);
    //a.insert(&4);
    //a.insert(&5);

    //println!("{:?}", a);
    //println!("{:?}", a.look_up(&1));
    //println!("{:?}", a.look_up_father(&4));
    //println!("{:?}", a);

    //let mut b = Tree::new(&10);
    //b.insert(&2);
    //b.insert(&1);
    //b.insert(&5);
    //b.insert(&3);
    //b.insert(&6);
    //b.insert(&11);

    //println!("{:?}", b);
    //println!("{:?}", b.look_up_father(&2));
    //b.left_shift(&2);
    //println!("{:?}", b);

    let mut c = Tree::new(&2);
    c.insert(&8);
    c.insert(&1);
    c.insert(&5);
    c.insert(&4);
    c.insert(&7);
    c.insert(&9);
    //println!("{:?}", c.look_up_father(&8));
    c.right_shift(&8);
    println!("{:?}", c);
}
<|endoftext|>


<commit_before>#![cfg(not(all(
    feature = "os-poll",
    feature = "os-util",
    feature = "tcp",
    feature = "udp",
    feature = "uds"
)))]
compile_error!("run main Mio tests with `--all-features`");
<commit_msg>Remove the requirement to enable all features for testing<commit_after><|endoftext|>


<commit_before>//! Library for parsing and converting bencoded data.
//!
//! # Examples
//!
//! Decoding bencoded data:
//!
//! ```rust
//!     extern crate bip_bencode;
//!
//!     use std::default::Default;
//!     use bip_bencode::{BencodeRef, BRefAccess, BDecodeOpt};
//!
//!     fn main() {
//!         let data = b"d12:lucky_numberi7ee";
//!         let bencode = BencodeRef::decode(data, BDecodeOpt::default()).unwrap();
//!
//!         assert_eq!(7, bencode.dict().unwrap().lookup("lucky_number".as_bytes())
//!             .unwrap().int().unwrap());
//!     }
//! ```
//!
//! Encoding bencoded data:
//!
//! ```rust
//!     #[macro_use]
//!     extern crate bip_bencode;
//!
//!     fn main() {
//!         let message = (ben_map!{
//!             "lucky_number" => ben_int!(7),
//!             "lucky_string" => ben_bytes!("7")
//!         }).encode();
//!
//!         assert_eq!(&b"d12:lucky_numberi7e12:lucky_string1:7e"[..], &message[..]);
//!     }
//! ```

#[macro_use]
extern crate error_chain;

mod access;
mod cow;
mod mutable;
mod reference;
mod error;

/// Traits for implementation functionality.
pub mod inner {
    pub use cow::BCowConvert;
}

/// Traits for extended functionality.
pub mod ext {
    pub use access::convert::{BConvertExt};
    pub use access::bencode::{BRefAccessExt};
}

pub use reference::bencode_ref::{BencodeRef};
pub use mutable::bencode_mut::{BencodeMut};
pub use access::bencode::{BRefAccess, BencodeRefKind, BMutAccess, BencodeMutKind};
pub use access::convert::{BConvert};
pub use access::dict::BDictAccess;
pub use access::list::BListAccess;
pub use reference::decode_opt::BDecodeOpt;
pub use error::{BencodeParseError, BencodeParseErrorKind, BencodeParseResult};
pub use error::{BencodeConvertError, BencodeConvertErrorKind, BencodeConvertResult};

const BEN_END: u8 = b'e';
const DICT_START: u8 = b'd';
const LIST_START: u8 = b'l';
const INT_START: u8 = b'i';

const BYTE_LEN_LOW: u8 = b'0';
const BYTE_LEN_HIGH: u8 = b'9';
const BYTE_LEN_END: u8 = b':';

/// Construct a `BencodeMut` map by supplying string references as keys and `BencodeMut` as values.
#[macro_export]
macro_rules! ben_map {
( $($key:expr => $val:expr),* ) => {
        {
            use bip_bencode::{BMutAccess, BencodeMut};
            use bip_bencode::inner::BCowConvert;

            let mut bencode_map = BencodeMut::new_dict();
            {
                let mut map = bencode_map.dict_mut().unwrap();
                $(
                    map.insert(BCowConvert::convert($key), $val);
                )*
            }

            bencode_map
        }
    }
}

/// Construct a `BencodeMut` list by supplying a list of `BencodeMut` values.
#[macro_export]
macro_rules! ben_list {
    ( $($ben:expr),* ) => {
        {
            use bip_bencode::{BencodeMut, BMutAccess};

            let mut bencode_list = BencodeMut::new_list();
            {
                let mut list = bencode_list.list_mut().unwrap();
                $(
                    list.push($ben);
                )*
            }

            bencode_list
        }
    }
}

/// Construct `BencodeMut` bytes by supplying a type convertible to `Vec<u8>`.
#[macro_export]
macro_rules! ben_bytes {
    ( $ben:expr ) => {
        {
            use bip_bencode::{BencodeMut};
            use bip_bencode::inner::BCowConvert;

            BencodeMut::new_bytes(BCowConvert::convert($ben))
        }
    }
}

/// Construct a `BencodeMut` integer by supplying an `i64`.
#[macro_export]
macro_rules! ben_int {
    ( $ben:expr ) => {
        {
            use bip_bencode::{BencodeMut};

            BencodeMut::new_int($ben)
        }
    }
}
<commit_msg>Remove unused mut for macro ben_map<commit_after>//! Library for parsing and converting bencoded data.
//!
//! # Examples
//!
//! Decoding bencoded data:
//!
//! ```rust
//!     extern crate bip_bencode;
//!
//!     use std::default::Default;
//!     use bip_bencode::{BencodeRef, BRefAccess, BDecodeOpt};
//!
//!     fn main() {
//!         let data = b"d12:lucky_numberi7ee";
//!         let bencode = BencodeRef::decode(data, BDecodeOpt::default()).unwrap();
//!
//!         assert_eq!(7, bencode.dict().unwrap().lookup("lucky_number".as_bytes())
//!             .unwrap().int().unwrap());
//!     }
//! ```
//!
//! Encoding bencoded data:
//!
//! ```rust
//!     #[macro_use]
//!     extern crate bip_bencode;
//!
//!     fn main() {
//!         let message = (ben_map!{
//!             "lucky_number" => ben_int!(7),
//!             "lucky_string" => ben_bytes!("7")
//!         }).encode();
//!
//!         assert_eq!(&b"d12:lucky_numberi7e12:lucky_string1:7e"[..], &message[..]);
//!     }
//! ```

#[macro_use]
extern crate error_chain;

mod access;
mod cow;
mod mutable;
mod reference;
mod error;

/// Traits for implementation functionality.
pub mod inner {
    pub use cow::BCowConvert;
}

/// Traits for extended functionality.
pub mod ext {
    pub use access::convert::{BConvertExt};
    pub use access::bencode::{BRefAccessExt};
}

pub use reference::bencode_ref::{BencodeRef};
pub use mutable::bencode_mut::{BencodeMut};
pub use access::bencode::{BRefAccess, BencodeRefKind, BMutAccess, BencodeMutKind};
pub use access::convert::{BConvert};
pub use access::dict::BDictAccess;
pub use access::list::BListAccess;
pub use reference::decode_opt::BDecodeOpt;
pub use error::{BencodeParseError, BencodeParseErrorKind, BencodeParseResult};
pub use error::{BencodeConvertError, BencodeConvertErrorKind, BencodeConvertResult};

const BEN_END: u8 = b'e';
const DICT_START: u8 = b'd';
const LIST_START: u8 = b'l';
const INT_START: u8 = b'i';

const BYTE_LEN_LOW: u8 = b'0';
const BYTE_LEN_HIGH: u8 = b'9';
const BYTE_LEN_END: u8 = b':';

/// Construct a `BencodeMut` map by supplying string references as keys and `BencodeMut` as values.
#[macro_export]
macro_rules! ben_map {
( $($key:expr => $val:expr),* ) => {
        {
            use bip_bencode::{BMutAccess, BencodeMut};
            use bip_bencode::inner::BCowConvert;

            let mut bencode_map = BencodeMut::new_dict();
            {
                let map = bencode_map.dict_mut().unwrap();
                $(
                    map.insert(BCowConvert::convert($key), $val);
                )*
            }

            bencode_map
        }
    }
}

/// Construct a `BencodeMut` list by supplying a list of `BencodeMut` values.
#[macro_export]
macro_rules! ben_list {
    ( $($ben:expr),* ) => {
        {
            use bip_bencode::{BencodeMut, BMutAccess};

            let mut bencode_list = BencodeMut::new_list();
            {
                let mut list = bencode_list.list_mut().unwrap();
                $(
                    list.push($ben);
                )*
            }

            bencode_list
        }
    }
}

/// Construct `BencodeMut` bytes by supplying a type convertible to `Vec<u8>`.
#[macro_export]
macro_rules! ben_bytes {
    ( $ben:expr ) => {
        {
            use bip_bencode::{BencodeMut};
            use bip_bencode::inner::BCowConvert;

            BencodeMut::new_bytes(BCowConvert::convert($ben))
        }
    }
}

/// Construct a `BencodeMut` integer by supplying an `i64`.
#[macro_export]
macro_rules! ben_int {
    ( $ben:expr ) => {
        {
            use bip_bencode::{BencodeMut};

            BencodeMut::new_int($ben)
        }
    }
}
<|endoftext|>