<commit_before>extern crate image;
extern crate img_hash;
extern crate postgres;
#[macro_use] extern crate log;
extern crate log4rs;
extern crate time;
#[macro_use] extern crate hyper;
extern crate url;
extern crate toml;

use std::env;
use postgres::{Connection, SslMode};
use img_hash::{ImageHash, HashType};
use std::thread;
use std::default::Default;
use time::{Timespec, Duration};
use hyper::Client;
use hyper::header::ContentType;
use image::ImageFormat;
use std::io::Read;
use hyper::status::StatusCode;
use postgres::stmt::Statement;
use std::fs::File;
use url::form_urlencoded;

header! { (XExtra, "X-Extra") => [String] }
header! { (XStream, "X-Stream") => [String] }

fn prepare_statement<'a>(conn: &'a Connection, stmt: &str) -> Statement<'a> {
    loop {
        let prep_stmt = conn.prepare_cached(stmt);
        if prep_stmt.is_err() {
            error!("Error in prepare for '{}': {}",
                   stmt,
                   prep_stmt.unwrap_err());
            thread::sleep(Duration::seconds(4).to_std().unwrap());
            continue;
        }
        return prep_stmt.unwrap();
    }
}

fn main() {
    let mut config_string = String::new();
    File::open("config.toml").unwrap().read_to_string(&mut config_string).unwrap();
    let mut parser = toml::Parser::new(&config_string);

    let config = parser.parse().unwrap();
    let config_table = config.get("config").unwrap();
    let interval = Duration::minutes(config_table.lookup("refresh_minutes")
                                                 .unwrap()
                                                 .as_integer()
                                                 .unwrap());
    let imager_url: &str = &(String::from(config_table.lookup("imager_host")
                                                      .unwrap()
                                                      .as_str()
                                                      .unwrap()) +
                             &String::from("/streams"));
    let check_ms = Duration::milliseconds(
        config_table.lookup("check_ms").unwrap().as_integer().unwrap())
        .to_std().unwrap();

    log4rs::init_file("log.toml", Default::default()).unwrap();
    let db_url: &str = &env::var("DATABASE_URL").unwrap();
    let conn = Connection::connect(db_url, SslMode::None).unwrap();

    let query_stmt = prepare_statement(&conn,
                                       "select * from videos_video order by \"lastRetrieved\" \
                                        asc limit 1");
    let update_stmt = prepare_statement(&conn,
                                        "update videos_video set working = true, hash = $1, \
                                         motion = $2, \"lastRetrieved\" = $3, extra = $4, \"streamURL\" = $5 where id = $6");
    let not_working_stmt = prepare_statement(&conn,
                                     "update videos_video set working = false, \"lastRetrieved\" = $1 where id = $2");

    info!("Connected to Postgres and ready to update video...");
    loop {
        let res = query_stmt.query(&[]);
        if res.is_err() {
            error!("Error in query: {}", res.unwrap_err());
            thread::sleep(check_ms);
            continue;
        }

        let items = res.unwrap();
        if items.is_empty() {
            warn!("No video feeds");
            thread::sleep(check_ms);
            continue;
        }
        let item = items.get(0);
        let url: String = item.get("url");
        let when: postgres::Result<Timespec> = item.get_opt("lastRetrieved").unwrap();
        let id: i32 = item.get("id");
        let old_hash: postgres::Result<String> = item.get_opt("hash").unwrap();

        if when.is_err() {
            debug!("Item: {}, When: never", url);
        } else {
            let now = time::now().to_timespec();
            let diff: Duration = now - when.unwrap();
            debug!("Item: {}, When: {}", url, diff);
            if diff < interval {
                debug!("oldest item is young: {}", diff);
                thread::sleep(check_ms);
                continue;
            }
        }
        info!("Updating {}", url);

        let data: String = form_urlencoded::Serializer::new(String::new())
            .append_pair("url", &url)
            .finish();

        let raw_resp = Client::new()
                           .post(imager_url)
                           .header(ContentType::form_url_encoded())
                           .body(&data)
                           .send();

        if raw_resp.is_err() {
            warn!("{:?}", raw_resp.unwrap_err());
            warn!("Can't connect to imager");
            thread::sleep(check_ms);
            continue;
        }

        let mut resp = raw_resp.unwrap();

        if resp.status != StatusCode::Ok {
            warn!("{:?}", resp.status_raw());
            let now = time::now().to_timespec();
            not_working_stmt.execute(&[&now, &id]).unwrap();
            thread::sleep(check_ms);
            continue;
        }

        let extra: String = match resp.headers.get::<XExtra>() {
            Some(val) => val.to_string(),
            None => "{}".to_string()
        };

        let stream: String = match resp.headers.get::<XStream>() {
            Some(val) => val.to_string(),
            None => "".to_string()
        };

        let mut buf = Vec::new();
        resp.read_to_end(&mut buf).unwrap();
        let image = image::load_from_memory_with_format(&buf, ImageFormat::PNG).unwrap();
        let hash = ImageHash::hash(&image, 16, HashType::DoubleGradient);

        debug!("Image hash: {}", hash.to_base64());
        let motion: f64 = match old_hash {
            Ok(val) => {
                match ImageHash::from_base64(&val) {
                    Ok(val) => val.dist_ratio(&hash) as f64,
                    Err(_) => -1 as f64
                }
            },
            Err(_) => -1 as f64
        };
        debug!("Difference {}", motion);
        let now = time::now().to_timespec();
        match update_stmt.execute(&[&hash.to_base64(), &motion, &now, &extra, &stream, &id]) {
            Ok(_) => info!("Updated {} with motion {}", url, motion),
            Err(err) => warn!("Error executing update: {:?}", err),
        }
    }
}
<commit_msg>Cope with dodgy images from Imager<commit_after>extern crate image;
extern crate img_hash;
extern crate postgres;
#[macro_use] extern crate log;
extern crate log4rs;
extern crate time;
#[macro_use] extern crate hyper;
extern crate url;
extern crate toml;

use std::env;
use postgres::{Connection, SslMode};
use img_hash::{ImageHash, HashType};
use std::thread;
use std::default::Default;
use time::{Timespec, Duration};
use hyper::Client;
use hyper::header::ContentType;
use image::ImageFormat;
use image::GenericImage;
use std::io::Read;
use hyper::status::StatusCode;
use postgres::stmt::Statement;
use std::fs::File;
use url::form_urlencoded;

header! { (XExtra, "X-Extra") => [String] }
header! { (XStream, "X-Stream") => [String] }

fn prepare_statement<'a>(conn: &'a Connection, stmt: &str) -> Statement<'a> {
    loop {
        let prep_stmt = conn.prepare_cached(stmt);
        if prep_stmt.is_err() {
            error!("Error in prepare for '{}': {}",
                   stmt,
                   prep_stmt.unwrap_err());
            thread::sleep(Duration::seconds(4).to_std().unwrap());
            continue;
        }
        return prep_stmt.unwrap();
    }
}

fn main() {
    let mut config_string = String::new();
    File::open("config.toml").unwrap().read_to_string(&mut config_string).unwrap();
    let mut parser = toml::Parser::new(&config_string);

    let config = parser.parse().unwrap();
    let config_table = config.get("config").unwrap();
    let interval = Duration::minutes(config_table.lookup("refresh_minutes")
                                                 .unwrap()
                                                 .as_integer()
                                                 .unwrap());
    let imager_url: &str = &(String::from(config_table.lookup("imager_host")
                                                      .unwrap()
                                                      .as_str()
                                                      .unwrap()) +
                             &String::from("/streams"));
    let check_ms = Duration::milliseconds(
        config_table.lookup("check_ms").unwrap().as_integer().unwrap())
        .to_std().unwrap();

    log4rs::init_file("log.toml", Default::default()).unwrap();
    let db_url: &str = &env::var("DATABASE_URL").unwrap();
    let conn = Connection::connect(db_url, SslMode::None).unwrap();

    let query_stmt = prepare_statement(&conn,
                                       "select * from videos_video order by \"lastRetrieved\" \
                                        asc limit 1");
    let update_stmt = prepare_statement(&conn,
                                        "update videos_video set working = true, hash = $1, \
                                         motion = $2, \"lastRetrieved\" = $3, extra = $4, \"streamURL\" = $5 where id = $6");
    let not_working_stmt = prepare_statement(&conn,
                                     "update videos_video set working = false, \"lastRetrieved\" = $1 where id = $2");

    info!("Connected to Postgres and ready to update video...");
    loop {
        let res = query_stmt.query(&[]);
        if res.is_err() {
            error!("Error in query: {}", res.unwrap_err());
            thread::sleep(check_ms);
            continue;
        }

        let items = res.unwrap();
        if items.is_empty() {
            warn!("No video feeds");
            thread::sleep(check_ms);
            continue;
        }
        let item = items.get(0);
        let url: String = item.get("url");
        let when: postgres::Result<Timespec> = item.get_opt("lastRetrieved").unwrap();
        let id: i32 = item.get("id");
        let old_hash: postgres::Result<String> = item.get_opt("hash").unwrap();

        if when.is_err() {
            debug!("Item: {}, When: never", url);
        } else {
            let now = time::now().to_timespec();
            let diff: Duration = now - when.unwrap();
            debug!("Item: {}, When: {}", url, diff);
            if diff < interval {
                debug!("oldest item is young: {}", diff);
                thread::sleep(check_ms);
                continue;
            }
        }
        info!("Updating {}", url);

        let data: String = form_urlencoded::Serializer::new(String::new())
            .append_pair("url", &url)
            .finish();

        let raw_resp = Client::new()
                           .post(imager_url)
                           .header(ContentType::form_url_encoded())
                           .body(&data)
                           .send();

        if raw_resp.is_err() {
            warn!("{:?}", raw_resp.unwrap_err());
            warn!("Can't connect to imager");
            thread::sleep(check_ms);
            continue;
        }

        let mut resp = raw_resp.unwrap();

        if resp.status != StatusCode::Ok {
            warn!("{:?}", resp.status_raw());
            let now = time::now().to_timespec();
            not_working_stmt.execute(&[&now, &id]).unwrap();
            thread::sleep(check_ms);
            continue;
        }

        let extra: String = match resp.headers.get::<XExtra>() {
            Some(val) => val.to_string(),
            None => "{}".to_string()
        };

        let stream: String = match resp.headers.get::<XStream>() {
            Some(val) => val.to_string(),
            None => "".to_string()
        };

        let mut buf = Vec::new();
        resp.read_to_end(&mut buf).unwrap();
        let image = image::load_from_memory_with_format(&buf, ImageFormat::PNG).unwrap();
        let (width, height) = image.dimensions();
        if width == 0 || height == 0 {
            warn!("Dodgy image dimensions, skipping");
            thread::sleep(check_ms);
            continue;
        }
        let hash = ImageHash::hash(&image, 16, HashType::DoubleGradient);

        debug!("Image hash: {}", hash.to_base64());
        let motion: f64 = match old_hash {
            Ok(val) => {
                match ImageHash::from_base64(&val) {
                    Ok(val) => val.dist_ratio(&hash) as f64,
                    Err(_) => -1 as f64
                }
            },
            Err(_) => -1 as f64
        };
        debug!("Difference {}", motion);
        let now = time::now().to_timespec();
        match update_stmt.execute(&[&hash.to_base64(), &motion, &now, &extra, &stream, &id]) {
            Ok(_) => info!("Updated {} with motion {}", url, motion),
            Err(err) => warn!("Error executing update: {:?}", err),
        }
    }
}
<|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 bluetooth_blacklist::{Blacklist, uuid_is_blacklisted};
use core::clone::Clone;
use dom::bindings::codegen::Bindings::BluetoothBinding;
use dom::bindings::codegen::Bindings::BluetoothBinding::RequestDeviceOptions;
use dom::bindings::codegen::Bindings::BluetoothBinding::{BluetoothScanFilter, BluetoothMethods};
use dom::bindings::error::Error::{Security, Type};
use dom::bindings::error::Fallible;
use dom::bindings::global::GlobalRef;
use dom::bindings::js::Root;
use dom::bindings::reflector::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::str::DOMString;
use dom::bluetoothadvertisingdata::BluetoothAdvertisingData;
use dom::bluetoothdevice::BluetoothDevice;
use dom::bluetoothuuid::BluetoothUUID;
use ipc_channel::ipc::{self, IpcSender};
use net_traits::bluetooth_scanfilter::{BluetoothScanfilter, BluetoothScanfilterSequence};
use net_traits::bluetooth_scanfilter::{RequestDeviceoptions, ServiceUUIDSequence};
use net_traits::bluetooth_thread::BluetoothMethodMsg;

const FILTER_EMPTY_ERROR: &'static str = "'filters' member must be non - empty to find any devices.";
const FILTER_ERROR: &'static str = "A filter must restrict the devices in some way.";
const FILTER_NAME_TOO_LONG_ERROR: &'static str = "A 'name' or 'namePrefix' can't be longer then 29 bytes.";
// 248 is the maximum number of UTF-8 code units in a Bluetooth Device Name.
const MAX_DEVICE_NAME_LENGTH: usize = 248;
// A device name can never be longer than 29 bytes.
// An advertising packet is at most 31 bytes long.
// The length and identifier of the length field take 2 bytes.
// That leaves 29 bytes for the name.
const MAX_FILTER_NAME_LENGTH: usize = 29;
const NAME_PREFIX_ERROR: &'static str = "'namePrefix', if present, must be non - empty.";
const NAME_TOO_LONG_ERROR: &'static str = "A device name can't be longer than 248 bytes.";
const SERVICE_ERROR: &'static str = "'services', if present, must contain at least one service.";

// https://webbluetoothcg.github.io/web-bluetooth/#bluetooth
#[dom_struct]
pub struct Bluetooth {
    reflector_: Reflector,
}

impl Bluetooth {
    pub fn new_inherited() -> Bluetooth {
        Bluetooth {
            reflector_: Reflector::new(),
        }
    }

    pub fn new(global: GlobalRef) -> Root<Bluetooth> {
        reflect_dom_object(box Bluetooth::new_inherited(),
                           global,
                           BluetoothBinding::Wrap)
    }

    fn get_bluetooth_thread(&self) -> IpcSender<BluetoothMethodMsg> {
        let global_root = self.global();
        let global_ref = global_root.r();
        global_ref.as_window().bluetooth_thread()
    }
}

fn canonicalize_filter(filter: &BluetoothScanFilter, global: GlobalRef) -> Fallible<BluetoothScanfilter> {
    if filter.services.is_none() && filter.name.is_none() && filter.namePrefix.is_none() {
        return Err(Type(FILTER_ERROR.to_owned()));
    }

    let mut services_vec = vec!();
    if let Some(ref services) = filter.services {
        if services.is_empty() {
            return Err(Type(SERVICE_ERROR.to_owned()));
        }
        for service in services {
            let uuid = try!(BluetoothUUID::GetService(global, service.clone())).to_string();
            if uuid_is_blacklisted(uuid.as_ref(), Blacklist::All) {
                return Err(Security)
            }
            services_vec.push(uuid);
        }
    }

    let mut name = String::new();
    if let Some(ref filter_name) = filter.name {
        //NOTE: DOMString::len() gives back the size in bytes
        if filter_name.len() > MAX_DEVICE_NAME_LENGTH {
            return Err(Type(NAME_TOO_LONG_ERROR.to_owned()));
        }
        if filter_name.len() > MAX_FILTER_NAME_LENGTH {
            return Err(Type(FILTER_NAME_TOO_LONG_ERROR.to_owned()));
        }
        name = filter_name.to_string();
    }

    let mut name_prefix = String::new();
    if let Some(ref filter_name_prefix) = filter.namePrefix {
        if filter_name_prefix.is_empty() {
            return Err(Type(NAME_PREFIX_ERROR.to_owned()));
        }
        if filter_name_prefix.len() > MAX_DEVICE_NAME_LENGTH {
            return Err(Type(NAME_TOO_LONG_ERROR.to_owned()));
        }
        if filter_name_prefix.len() > MAX_FILTER_NAME_LENGTH {
            return Err(Type(FILTER_NAME_TOO_LONG_ERROR.to_owned()));
        }
        name_prefix = filter_name_prefix.to_string();
    }

    Ok(BluetoothScanfilter::new(name, name_prefix, services_vec))
}

fn convert_request_device_options(options: &RequestDeviceOptions,
                                  global: GlobalRef)
                                  -> Fallible<RequestDeviceoptions> {
    if options.filters.is_empty() {
        return Err(Type(FILTER_EMPTY_ERROR.to_owned()));
    }

    let mut filters = vec!();
    for filter in &options.filters {
        filters.push(try!(canonicalize_filter(&filter, global)));
    }

    let mut optional_services = vec!();
    if let Some(ref opt_services) = options.optionalServices {
        for opt_service in opt_services {
            let uuid = try!(BluetoothUUID::GetService(global, opt_service.clone())).to_string();
            if uuid_is_blacklisted(uuid.as_ref(), Blacklist::All) {
                return Err(Security)
            }
            optional_services.push(uuid);
        }
    }

    Ok(RequestDeviceoptions::new(BluetoothScanfilterSequence::new(filters),
                                 ServiceUUIDSequence::new(optional_services)))
}

impl BluetoothMethods for Bluetooth {
    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetooth-requestdevice
    fn RequestDevice(&self, option: &RequestDeviceOptions) -> Fallible<Root<BluetoothDevice>> {
        let (sender, receiver) = ipc::channel().unwrap();
        let option = try!(convert_request_device_options(option, self.global().r()));
        self.get_bluetooth_thread().send(BluetoothMethodMsg::RequestDevice(option, sender)).unwrap();
        let device = receiver.recv().unwrap();
        match device {
            Ok(device) => {
                let ad_data = BluetoothAdvertisingData::new(self.global().r(),
                                                            device.appearance,
                                                            device.tx_power,
                                                            device.rssi);
                Ok(BluetoothDevice::new(self.global().r(),
                                        DOMString::from(device.id),
                                        device.name.map(DOMString::from),
                                        &ad_data))
            },
            Err(error) => {
                Err(Type(error))
            },
        }
    }
}
<commit_msg>Auto merge of #11636 - szeged:optservicesfix, r=jdm<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 bluetooth_blacklist::{Blacklist, uuid_is_blacklisted};
use core::clone::Clone;
use dom::bindings::codegen::Bindings::BluetoothBinding;
use dom::bindings::codegen::Bindings::BluetoothBinding::RequestDeviceOptions;
use dom::bindings::codegen::Bindings::BluetoothBinding::{BluetoothScanFilter, BluetoothMethods};
use dom::bindings::error::Error::{Security, Type};
use dom::bindings::error::Fallible;
use dom::bindings::global::GlobalRef;
use dom::bindings::js::Root;
use dom::bindings::reflector::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::str::DOMString;
use dom::bluetoothadvertisingdata::BluetoothAdvertisingData;
use dom::bluetoothdevice::BluetoothDevice;
use dom::bluetoothuuid::BluetoothUUID;
use ipc_channel::ipc::{self, IpcSender};
use net_traits::bluetooth_scanfilter::{BluetoothScanfilter, BluetoothScanfilterSequence};
use net_traits::bluetooth_scanfilter::{RequestDeviceoptions, ServiceUUIDSequence};
use net_traits::bluetooth_thread::BluetoothMethodMsg;

const FILTER_EMPTY_ERROR: &'static str = "'filters' member must be non - empty to find any devices.";
const FILTER_ERROR: &'static str = "A filter must restrict the devices in some way.";
const FILTER_NAME_TOO_LONG_ERROR: &'static str = "A 'name' or 'namePrefix' can't be longer then 29 bytes.";
// 248 is the maximum number of UTF-8 code units in a Bluetooth Device Name.
const MAX_DEVICE_NAME_LENGTH: usize = 248;
// A device name can never be longer than 29 bytes.
// An advertising packet is at most 31 bytes long.
// The length and identifier of the length field take 2 bytes.
// That leaves 29 bytes for the name.
const MAX_FILTER_NAME_LENGTH: usize = 29;
const NAME_PREFIX_ERROR: &'static str = "'namePrefix', if present, must be non - empty.";
const NAME_TOO_LONG_ERROR: &'static str = "A device name can't be longer than 248 bytes.";
const SERVICE_ERROR: &'static str = "'services', if present, must contain at least one service.";

// https://webbluetoothcg.github.io/web-bluetooth/#bluetooth
#[dom_struct]
pub struct Bluetooth {
    reflector_: Reflector,
}

impl Bluetooth {
    pub fn new_inherited() -> Bluetooth {
        Bluetooth {
            reflector_: Reflector::new(),
        }
    }

    pub fn new(global: GlobalRef) -> Root<Bluetooth> {
        reflect_dom_object(box Bluetooth::new_inherited(),
                           global,
                           BluetoothBinding::Wrap)
    }

    fn get_bluetooth_thread(&self) -> IpcSender<BluetoothMethodMsg> {
        let global_root = self.global();
        let global_ref = global_root.r();
        global_ref.as_window().bluetooth_thread()
    }
}

fn canonicalize_filter(filter: &BluetoothScanFilter, global: GlobalRef) -> Fallible<BluetoothScanfilter> {
    if filter.services.is_none() && filter.name.is_none() && filter.namePrefix.is_none() {
        return Err(Type(FILTER_ERROR.to_owned()));
    }

    let mut services_vec = vec!();
    if let Some(ref services) = filter.services {
        if services.is_empty() {
            return Err(Type(SERVICE_ERROR.to_owned()));
        }
        for service in services {
            let uuid = try!(BluetoothUUID::GetService(global, service.clone())).to_string();
            if uuid_is_blacklisted(uuid.as_ref(), Blacklist::All) {
                return Err(Security)
            }
            services_vec.push(uuid);
        }
    }

    let mut name = String::new();
    if let Some(ref filter_name) = filter.name {
        //NOTE: DOMString::len() gives back the size in bytes
        if filter_name.len() > MAX_DEVICE_NAME_LENGTH {
            return Err(Type(NAME_TOO_LONG_ERROR.to_owned()));
        }
        if filter_name.len() > MAX_FILTER_NAME_LENGTH {
            return Err(Type(FILTER_NAME_TOO_LONG_ERROR.to_owned()));
        }
        name = filter_name.to_string();
    }

    let mut name_prefix = String::new();
    if let Some(ref filter_name_prefix) = filter.namePrefix {
        if filter_name_prefix.is_empty() {
            return Err(Type(NAME_PREFIX_ERROR.to_owned()));
        }
        if filter_name_prefix.len() > MAX_DEVICE_NAME_LENGTH {
            return Err(Type(NAME_TOO_LONG_ERROR.to_owned()));
        }
        if filter_name_prefix.len() > MAX_FILTER_NAME_LENGTH {
            return Err(Type(FILTER_NAME_TOO_LONG_ERROR.to_owned()));
        }
        name_prefix = filter_name_prefix.to_string();
    }

    Ok(BluetoothScanfilter::new(name, name_prefix, services_vec))
}

fn convert_request_device_options(options: &RequestDeviceOptions,
                                  global: GlobalRef)
                                  -> Fallible<RequestDeviceoptions> {
    if options.filters.is_empty() {
        return Err(Type(FILTER_EMPTY_ERROR.to_owned()));
    }

    let mut filters = vec!();
    for filter in &options.filters {
        filters.push(try!(canonicalize_filter(&filter, global)));
    }

    let mut optional_services = vec!();
    if let Some(ref opt_services) = options.optionalServices {
        for opt_service in opt_services {
            let uuid = try!(BluetoothUUID::GetService(global, opt_service.clone())).to_string();
            if !uuid_is_blacklisted(uuid.as_ref(), Blacklist::All) {
                optional_services.push(uuid);
            }
        }
    }

    Ok(RequestDeviceoptions::new(BluetoothScanfilterSequence::new(filters),
                                 ServiceUUIDSequence::new(optional_services)))
}

impl BluetoothMethods for Bluetooth {
    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetooth-requestdevice
    fn RequestDevice(&self, option: &RequestDeviceOptions) -> Fallible<Root<BluetoothDevice>> {
        let (sender, receiver) = ipc::channel().unwrap();
        let option = try!(convert_request_device_options(option, self.global().r()));
        self.get_bluetooth_thread().send(BluetoothMethodMsg::RequestDevice(option, sender)).unwrap();
        let device = receiver.recv().unwrap();
        match device {
            Ok(device) => {
                let ad_data = BluetoothAdvertisingData::new(self.global().r(),
                                                            device.appearance,
                                                            device.tx_power,
                                                            device.rssi);
                Ok(BluetoothDevice::new(self.global().r(),
                                        DOMString::from(device.id),
                                        device.name.map(DOMString::from),
                                        &ad_data))
            },
            Err(error) => {
                Err(Type(error))
            },
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 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.
use std::collections::BTreeMap;
use std::ops::RangeInclusive;

use base::pagesize;

use crate::address_allocator::{AddressAllocator, AddressAllocatorSet};
use crate::{Alloc, Error, Result};

/// Manages allocating system resources such as address space and interrupt numbers.

/// MMIO address Type
///    Low: address allocated from low_address_space
///    High: address allocated from high_address_space
#[derive(Copy, Clone)]
pub enum MmioType {
    Low,
    High,
}

/// Region of memory.
#[derive(Debug)]
pub struct MemRegion {
    pub base: u64,
    pub size: u64,
}

pub struct SystemAllocatorConfig {
    /// IO ports. Only for x86_64.
    pub io: Option<MemRegion>,
    /// Low (<=4GB) MMIO region.
    ///
    /// Parts of this region may be reserved or otherwise excluded from the
    /// created SystemAllocator's MmioType::Low allocator. However, no new
    /// regions will be added.
    pub low_mmio: MemRegion,
    /// High (>4GB) MMIO region.
    ///
    /// Parts of this region may be reserved or otherwise excluded from the
    /// created SystemAllocator's MmioType::High allocator. However, no new
    /// regions will be added.
    pub high_mmio: MemRegion,
    /// Platform MMIO space. Only for ARM.
    pub platform_mmio: Option<MemRegion>,
    /// The first IRQ number to give out.
    pub first_irq: u32,
}

#[derive(Debug)]
pub struct SystemAllocator {
    io_address_space: Option<AddressAllocator>,

    // Indexed by MmioType::Low and MmioType::High.
    mmio_address_spaces: [AddressAllocator; 2],
    mmio_platform_address_spaces: Option<AddressAllocator>,

    reserved_region: Option<MemRegion>,

    // Each bus number has a AddressAllocator
    pci_allocator: BTreeMap<u8, AddressAllocator>,
    irq_allocator: AddressAllocator,
    next_anon_id: usize,
}

fn to_range_inclusive(base: u64, size: u64) -> Result<RangeInclusive<u64>> {
    let end = base
        .checked_add(size.checked_sub(1).ok_or(Error::PoolSizeZero)?)
        .ok_or(Error::PoolOverflow { base, size })?;
    Ok(RangeInclusive::new(base, end))
}

fn range_intersect(r1: &RangeInclusive<u64>, r2: &RangeInclusive<u64>) -> RangeInclusive<u64> {
    RangeInclusive::new(
        u64::max(*r1.start(), *r2.start()),
        u64::min(*r1.end(), *r2.end()),
    )
}

impl SystemAllocator {
    /// Creates a new `SystemAllocator` for managing addresses and irq numbers.
    /// Will return an error if `base` + `size` overflows u64 (or allowed
    /// maximum for the specific type), or if alignment isn't a power of two.
    ///
    /// If `reserve_region_size` is not None, then a region is reserved from
    /// the start of `config.high_mmio` before the mmio allocator is created.
    ///
    /// If `mmio_address_ranges` is not empty, then `config.low_mmio` and
    /// `config.high_mmio` are intersected with the ranges specified.
    pub fn new(
        config: SystemAllocatorConfig,
        reserve_region_size: Option<u64>,
        mmio_address_ranges: &[RangeInclusive<u64>],
    ) -> Result<Self> {
        let page_size = pagesize() as u64;

        let (high_mmio, reserved_region) = match reserve_region_size {
            Some(len) => {
                if len > config.high_mmio.size {
                    return Err(Error::PoolSizeZero);
                }
                (
                    MemRegion {
                        base: config.high_mmio.base + len,
                        size: config.high_mmio.size - len,
                    },
                    Some(MemRegion {
                        base: config.high_mmio.base,
                        size: len,
                    }),
                )
            }
            None => (config.high_mmio, None),
        };

        let intersect_mmio_range = |src: MemRegion| -> Result<Vec<RangeInclusive<u64>>> {
            let src_range = to_range_inclusive(src.base, src.size)?;
            Ok(if mmio_address_ranges.is_empty() {
                vec![src_range]
            } else {
                mmio_address_ranges
                    .iter()
                    .map(|r| range_intersect(r, &src_range))
                    .collect()
            })
        };

        Ok(SystemAllocator {
            io_address_space: if let Some(io) = config.io {
                // TODO make sure we don't overlap with existing well known
                // ports such as 0xcf8 (serial ports).
                if io.base > 0x1_0000 || io.size + io.base > 0x1_0000 {
                    return Err(Error::IOPortOutOfRange(io.base, io.size));
                }
                Some(AddressAllocator::new(
                    to_range_inclusive(io.base, io.size)?,
                    Some(0x400),
                    None,
                )?)
            } else {
                None
            },
            mmio_address_spaces: [
                // MmioType::Low
                AddressAllocator::new_from_list(
                    intersect_mmio_range(config.low_mmio)?,
                    Some(page_size),
                    None,
                )?,
                // MmioType::High
                AddressAllocator::new_from_list(
                    intersect_mmio_range(high_mmio)?,
                    Some(page_size),
                    None,
                )?,
            ],

            pci_allocator: BTreeMap::new(),

            mmio_platform_address_spaces: if let Some(platform) = config.platform_mmio {
                Some(AddressAllocator::new(
                    to_range_inclusive(platform.base, platform.size)?,
                    Some(page_size),
                    None,
                )?)
            } else {
                None
            },

            reserved_region,

            irq_allocator: AddressAllocator::new(
                RangeInclusive::new(config.first_irq as u64, 1023),
                Some(1),
                None,
            )?,
            next_anon_id: 0,
        })
    }

    /// Reserves the next available system irq number.
    pub fn allocate_irq(&mut self) -> Option<u32> {
        let id = self.get_anon_alloc();
        self.irq_allocator
            .allocate(1, id, "irq-auto".to_string())
            .map(|v| v as u32)
            .ok()
    }

    /// release irq to system irq number pool
    pub fn release_irq(&mut self, irq: u32) {
        let _ = self.irq_allocator.release_containing(irq.into());
    }

    /// Reserves the next available system irq number.
    pub fn reserve_irq(&mut self, irq: u32) -> bool {
        let id = self.get_anon_alloc();
        self.irq_allocator
            .allocate_at(irq as u64, 1, id, "irq-fixed".to_string())
            .is_ok()
    }

    fn get_pci_allocator_mut(&mut self, bus: u8) -> Option<&mut AddressAllocator> {
        // pci root is 00:00.0, Bus 0 next device is 00:01.0 with mandatory function
        // number zero.
        if self.pci_allocator.get(&bus).is_none() {
            let base = if bus == 0 { 8 } else { 0 };

            // Each bus supports up to 32 (devices) x 8 (functions).
            // Prefer allocating at device granularity (preferred_align = 8), but fall back to
            // allocating individual functions (min_align = 1) when we run out of devices.
            match AddressAllocator::new(RangeInclusive::new(base, (32 * 8) - 1), Some(1), Some(8)) {
                Ok(v) => self.pci_allocator.insert(bus, v),
                Err(_) => return None,
            };
        }
        self.pci_allocator.get_mut(&bus)
    }

    // Check whether devices exist or not on the specified bus
    pub fn pci_bus_empty(&self, bus: u8) -> bool {
        if self.pci_allocator.get(&bus).is_none() {
            true
        } else {
            false
        }
    }

    /// Allocate PCI slot location.
    pub fn allocate_pci(&mut self, bus: u8, tag: String) -> Option<Alloc> {
        let id = self.get_anon_alloc();
        let allocator = match self.get_pci_allocator_mut(bus) {
            Some(v) => v,
            None => return None,
        };
        allocator
            .allocate(1, id, tag)
            .map(|v| Alloc::PciBar {
                bus,
                dev: (v >> 3) as u8,
                func: (v & 7) as u8,
                bar: 0,
            })
            .ok()
    }

    /// Reserve PCI slot location.
    pub fn reserve_pci(&mut self, alloc: Alloc, tag: String) -> bool {
        let id = self.get_anon_alloc();
        match alloc {
            Alloc::PciBar {
                bus,
                dev,
                func,
                bar: _,
            } => {
                let allocator = match self.get_pci_allocator_mut(bus) {
                    Some(v) => v,
                    None => return false,
                };
                let df = ((dev as u64) << 3) | (func as u64);
                allocator.allocate_at(df, 1, id, tag).is_ok()
            }
            _ => false,
        }
    }

    /// release PCI slot location.
    pub fn release_pci(&mut self, bus: u8, dev: u8, func: u8) -> bool {
        let allocator = match self.get_pci_allocator_mut(bus) {
            Some(v) => v,
            None => return false,
        };
        let df = ((dev as u64) << 3) | (func as u64);
        allocator.release_containing(df).is_ok()
    }

    /// Gets an allocator to be used for platform device MMIO allocation.
    pub fn mmio_platform_allocator(&mut self) -> Option<&mut AddressAllocator> {
        self.mmio_platform_address_spaces.as_mut()
    }

    /// Gets an allocator to be used for IO memory.
    pub fn io_allocator(&mut self) -> Option<&mut AddressAllocator> {
        self.io_address_space.as_mut()
    }

    /// Gets an allocator to be used for MMIO allocation.
    ///    MmioType::Low: low mmio allocator
    ///    MmioType::High: high mmio allocator
    pub fn mmio_allocator(&mut self, mmio_type: MmioType) -> &mut AddressAllocator {
        &mut self.mmio_address_spaces[mmio_type as usize]
    }

    /// Gets a set of allocators to be used for MMIO allocation.
    /// The set of allocators will try the low and high MMIO allocators, in that order.
    pub fn mmio_allocator_any(&mut self) -> AddressAllocatorSet {
        AddressAllocatorSet::new(&mut self.mmio_address_spaces)
    }

    /// Gets the pools of all mmio allocators.
    pub fn mmio_pools(&self) -> Vec<&RangeInclusive<u64>> {
        self.mmio_address_spaces
            .iter()
            .flat_map(|mmio_as| mmio_as.pools())
            .collect()
    }

    /// Gets the reserved address space region.
    pub fn reserved_region(&self) -> Option<&MemRegion> {
        self.reserved_region.as_ref()
    }

    /// Gets a unique anonymous allocation
    pub fn get_anon_alloc(&mut self) -> Alloc {
        self.next_anon_id += 1;
        Alloc::Anon(self.next_anon_id)
    }
}

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

    #[test]
    fn example() {
        let mut a = SystemAllocator::new(
            SystemAllocatorConfig {
                io: Some(MemRegion {
                    base: 0x1000,
                    size: 0xf000,
                }),
                low_mmio: MemRegion {
                    base: 0x3000_0000,
                    size: 0x1_0000,
                },
                high_mmio: MemRegion {
                    base: 0x1000_0000,
                    size: 0x1000_0000,
                },
                platform_mmio: None,
                first_irq: 5,
            },
            None,
            &[],
        )
        .unwrap();

        assert_eq!(a.allocate_irq(), Some(5));
        assert_eq!(a.allocate_irq(), Some(6));
        assert_eq!(
            a.mmio_allocator(MmioType::High).allocate(
                0x100,
                Alloc::PciBar {
                    bus: 0,
                    dev: 0,
                    func: 0,
                    bar: 0
                },
                "bar0".to_string()
            ),
            Ok(0x10000000)
        );
        assert_eq!(
            a.mmio_allocator(MmioType::High).get(&Alloc::PciBar {
                bus: 0,
                dev: 0,
                func: 0,
                bar: 0
            }),
            Some(&(0x10000000, 0x100, "bar0".to_string()))
        );
    }
}
<commit_msg>resource: Add reserve_mmio() function<commit_after>// Copyright 2018 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.
use std::collections::BTreeMap;
use std::ops::RangeInclusive;

use base::pagesize;

use crate::address_allocator::{AddressAllocator, AddressAllocatorSet};
use crate::{Alloc, Error, Result};

/// Manages allocating system resources such as address space and interrupt numbers.

/// MMIO address Type
///    Low: address allocated from low_address_space
///    High: address allocated from high_address_space
#[derive(Copy, Clone)]
pub enum MmioType {
    Low,
    High,
}

/// Region of memory.
#[derive(Debug)]
pub struct MemRegion {
    pub base: u64,
    pub size: u64,
}

pub struct SystemAllocatorConfig {
    /// IO ports. Only for x86_64.
    pub io: Option<MemRegion>,
    /// Low (<=4GB) MMIO region.
    ///
    /// Parts of this region may be reserved or otherwise excluded from the
    /// created SystemAllocator's MmioType::Low allocator. However, no new
    /// regions will be added.
    pub low_mmio: MemRegion,
    /// High (>4GB) MMIO region.
    ///
    /// Parts of this region may be reserved or otherwise excluded from the
    /// created SystemAllocator's MmioType::High allocator. However, no new
    /// regions will be added.
    pub high_mmio: MemRegion,
    /// Platform MMIO space. Only for ARM.
    pub platform_mmio: Option<MemRegion>,
    /// The first IRQ number to give out.
    pub first_irq: u32,
}

#[derive(Debug)]
pub struct SystemAllocator {
    io_address_space: Option<AddressAllocator>,

    // Indexed by MmioType::Low and MmioType::High.
    mmio_address_spaces: [AddressAllocator; 2],
    mmio_platform_address_spaces: Option<AddressAllocator>,

    reserved_region: Option<MemRegion>,

    // Each bus number has a AddressAllocator
    pci_allocator: BTreeMap<u8, AddressAllocator>,
    irq_allocator: AddressAllocator,
    next_anon_id: usize,
}

fn to_range_inclusive(base: u64, size: u64) -> Result<RangeInclusive<u64>> {
    let end = base
        .checked_add(size.checked_sub(1).ok_or(Error::PoolSizeZero)?)
        .ok_or(Error::PoolOverflow { base, size })?;
    Ok(RangeInclusive::new(base, end))
}

fn range_intersect(r1: &RangeInclusive<u64>, r2: &RangeInclusive<u64>) -> RangeInclusive<u64> {
    RangeInclusive::new(
        u64::max(*r1.start(), *r2.start()),
        u64::min(*r1.end(), *r2.end()),
    )
}

impl SystemAllocator {
    /// Creates a new `SystemAllocator` for managing addresses and irq numbers.
    /// Will return an error if `base` + `size` overflows u64 (or allowed
    /// maximum for the specific type), or if alignment isn't a power of two.
    ///
    /// If `reserve_region_size` is not None, then a region is reserved from
    /// the start of `config.high_mmio` before the mmio allocator is created.
    ///
    /// If `mmio_address_ranges` is not empty, then `config.low_mmio` and
    /// `config.high_mmio` are intersected with the ranges specified.
    pub fn new(
        config: SystemAllocatorConfig,
        reserve_region_size: Option<u64>,
        mmio_address_ranges: &[RangeInclusive<u64>],
    ) -> Result<Self> {
        let page_size = pagesize() as u64;

        let (high_mmio, reserved_region) = match reserve_region_size {
            Some(len) => {
                if len > config.high_mmio.size {
                    return Err(Error::PoolSizeZero);
                }
                (
                    MemRegion {
                        base: config.high_mmio.base + len,
                        size: config.high_mmio.size - len,
                    },
                    Some(MemRegion {
                        base: config.high_mmio.base,
                        size: len,
                    }),
                )
            }
            None => (config.high_mmio, None),
        };

        let intersect_mmio_range = |src: MemRegion| -> Result<Vec<RangeInclusive<u64>>> {
            let src_range = to_range_inclusive(src.base, src.size)?;
            Ok(if mmio_address_ranges.is_empty() {
                vec![src_range]
            } else {
                mmio_address_ranges
                    .iter()
                    .map(|r| range_intersect(r, &src_range))
                    .collect()
            })
        };

        Ok(SystemAllocator {
            io_address_space: if let Some(io) = config.io {
                // TODO make sure we don't overlap with existing well known
                // ports such as 0xcf8 (serial ports).
                if io.base > 0x1_0000 || io.size + io.base > 0x1_0000 {
                    return Err(Error::IOPortOutOfRange(io.base, io.size));
                }
                Some(AddressAllocator::new(
                    to_range_inclusive(io.base, io.size)?,
                    Some(0x400),
                    None,
                )?)
            } else {
                None
            },
            mmio_address_spaces: [
                // MmioType::Low
                AddressAllocator::new_from_list(
                    intersect_mmio_range(config.low_mmio)?,
                    Some(page_size),
                    None,
                )?,
                // MmioType::High
                AddressAllocator::new_from_list(
                    intersect_mmio_range(high_mmio)?,
                    Some(page_size),
                    None,
                )?,
            ],

            pci_allocator: BTreeMap::new(),

            mmio_platform_address_spaces: if let Some(platform) = config.platform_mmio {
                Some(AddressAllocator::new(
                    to_range_inclusive(platform.base, platform.size)?,
                    Some(page_size),
                    None,
                )?)
            } else {
                None
            },

            reserved_region,

            irq_allocator: AddressAllocator::new(
                RangeInclusive::new(config.first_irq as u64, 1023),
                Some(1),
                None,
            )?,
            next_anon_id: 0,
        })
    }

    /// Reserves the next available system irq number.
    pub fn allocate_irq(&mut self) -> Option<u32> {
        let id = self.get_anon_alloc();
        self.irq_allocator
            .allocate(1, id, "irq-auto".to_string())
            .map(|v| v as u32)
            .ok()
    }

    /// release irq to system irq number pool
    pub fn release_irq(&mut self, irq: u32) {
        let _ = self.irq_allocator.release_containing(irq.into());
    }

    /// Reserves the next available system irq number.
    pub fn reserve_irq(&mut self, irq: u32) -> bool {
        let id = self.get_anon_alloc();
        self.irq_allocator
            .allocate_at(irq as u64, 1, id, "irq-fixed".to_string())
            .is_ok()
    }

    fn get_pci_allocator_mut(&mut self, bus: u8) -> Option<&mut AddressAllocator> {
        // pci root is 00:00.0, Bus 0 next device is 00:01.0 with mandatory function
        // number zero.
        if self.pci_allocator.get(&bus).is_none() {
            let base = if bus == 0 { 8 } else { 0 };

            // Each bus supports up to 32 (devices) x 8 (functions).
            // Prefer allocating at device granularity (preferred_align = 8), but fall back to
            // allocating individual functions (min_align = 1) when we run out of devices.
            match AddressAllocator::new(RangeInclusive::new(base, (32 * 8) - 1), Some(1), Some(8)) {
                Ok(v) => self.pci_allocator.insert(bus, v),
                Err(_) => return None,
            };
        }
        self.pci_allocator.get_mut(&bus)
    }

    // Check whether devices exist or not on the specified bus
    pub fn pci_bus_empty(&self, bus: u8) -> bool {
        if self.pci_allocator.get(&bus).is_none() {
            true
        } else {
            false
        }
    }

    /// Allocate PCI slot location.
    pub fn allocate_pci(&mut self, bus: u8, tag: String) -> Option<Alloc> {
        let id = self.get_anon_alloc();
        let allocator = match self.get_pci_allocator_mut(bus) {
            Some(v) => v,
            None => return None,
        };
        allocator
            .allocate(1, id, tag)
            .map(|v| Alloc::PciBar {
                bus,
                dev: (v >> 3) as u8,
                func: (v & 7) as u8,
                bar: 0,
            })
            .ok()
    }

    /// Reserve PCI slot location.
    pub fn reserve_pci(&mut self, alloc: Alloc, tag: String) -> bool {
        let id = self.get_anon_alloc();
        match alloc {
            Alloc::PciBar {
                bus,
                dev,
                func,
                bar: _,
            } => {
                let allocator = match self.get_pci_allocator_mut(bus) {
                    Some(v) => v,
                    None => return false,
                };
                let df = ((dev as u64) << 3) | (func as u64);
                allocator.allocate_at(df, 1, id, tag).is_ok()
            }
            _ => false,
        }
    }

    /// release PCI slot location.
    pub fn release_pci(&mut self, bus: u8, dev: u8, func: u8) -> bool {
        let allocator = match self.get_pci_allocator_mut(bus) {
            Some(v) => v,
            None => return false,
        };
        let df = ((dev as u64) << 3) | (func as u64);
        allocator.release_containing(df).is_ok()
    }

    /// Reserve specified range from pci mmio, get the overlap of specified
    /// range with mmio pools, exclude the overlap from mmio allocator.
    ///
    /// If any part of the specified range has been allocated, return Error.
    pub fn reserve_mmio(&mut self, start: u64, len: u64) -> Result<()> {
        let target = to_range_inclusive(start, len)?;
        let mut pools = Vec::new();
        for pool in self.mmio_pools() {
            pools.push(RangeInclusive::new(*pool.start(), *pool.end()));
        }
        pools.sort_by(|a, b| a.start().cmp(b.start()));
        for pool in &pools {
            if pool.start() > target.end() {
                break;
            }

            let overlap = range_intersect(pool, &target);
            if !overlap.is_empty() {
                let id = self.get_anon_alloc();
                self.mmio_allocator_any().allocate_at(
                    *overlap.start(),
                    *overlap.end() - *overlap.start() + 1,
                    id,
                    "pci mmio reserve".to_string(),
                )?;
            }
        }

        Ok(())
    }

    /// Gets an allocator to be used for platform device MMIO allocation.
    pub fn mmio_platform_allocator(&mut self) -> Option<&mut AddressAllocator> {
        self.mmio_platform_address_spaces.as_mut()
    }

    /// Gets an allocator to be used for IO memory.
    pub fn io_allocator(&mut self) -> Option<&mut AddressAllocator> {
        self.io_address_space.as_mut()
    }

    /// Gets an allocator to be used for MMIO allocation.
    ///    MmioType::Low: low mmio allocator
    ///    MmioType::High: high mmio allocator
    pub fn mmio_allocator(&mut self, mmio_type: MmioType) -> &mut AddressAllocator {
        &mut self.mmio_address_spaces[mmio_type as usize]
    }

    /// Gets a set of allocators to be used for MMIO allocation.
    /// The set of allocators will try the low and high MMIO allocators, in that order.
    pub fn mmio_allocator_any(&mut self) -> AddressAllocatorSet {
        AddressAllocatorSet::new(&mut self.mmio_address_spaces)
    }

    /// Gets the pools of all mmio allocators.
    pub fn mmio_pools(&self) -> Vec<&RangeInclusive<u64>> {
        self.mmio_address_spaces
            .iter()
            .flat_map(|mmio_as| mmio_as.pools())
            .collect()
    }

    /// Gets the reserved address space region.
    pub fn reserved_region(&self) -> Option<&MemRegion> {
        self.reserved_region.as_ref()
    }

    /// Gets a unique anonymous allocation
    pub fn get_anon_alloc(&mut self) -> Alloc {
        self.next_anon_id += 1;
        Alloc::Anon(self.next_anon_id)
    }
}

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

    #[test]
    fn example() {
        let mut a = SystemAllocator::new(
            SystemAllocatorConfig {
                io: Some(MemRegion {
                    base: 0x1000,
                    size: 0xf000,
                }),
                low_mmio: MemRegion {
                    base: 0x3000_0000,
                    size: 0x1_0000,
                },
                high_mmio: MemRegion {
                    base: 0x1000_0000,
                    size: 0x1000_0000,
                },
                platform_mmio: None,
                first_irq: 5,
            },
            None,
            &[],
        )
        .unwrap();

        assert_eq!(a.allocate_irq(), Some(5));
        assert_eq!(a.allocate_irq(), Some(6));
        assert_eq!(
            a.mmio_allocator(MmioType::High).allocate(
                0x100,
                Alloc::PciBar {
                    bus: 0,
                    dev: 0,
                    func: 0,
                    bar: 0
                },
                "bar0".to_string()
            ),
            Ok(0x10000000)
        );
        assert_eq!(
            a.mmio_allocator(MmioType::High).get(&Alloc::PciBar {
                bus: 0,
                dev: 0,
                func: 0,
                bar: 0
            }),
            Some(&(0x10000000, 0x100, "bar0".to_string()))
        );

        let id = a.get_anon_alloc();
        assert_eq!(
            a.mmio_allocator(MmioType::Low).allocate_at(
                0x3000_5000,
                0x5000,
                id,
                "Test".to_string()
            ),
            Ok(())
        );
        assert_eq!(a.mmio_allocator(MmioType::Low).release(id), Ok(()));
        assert_eq!(a.reserve_mmio(0x3000_2000, 0x4000), Ok(()));
        assert_eq!(
            a.mmio_allocator(MmioType::Low)
                .allocate_at(0x3000_5000, 0x5000, id, "Test".to_string())
                .is_err(),
            true
        );
    }
}
<|endoftext|>


<commit_before>use crate::ops::{Deref, DerefMut};

#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(align(64))]
pub(super) struct Aligner;

#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(super) struct CacheAligned<T>(pub T, pub Aligner);

impl<T> Deref for CacheAligned<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> DerefMut for CacheAligned<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<T> CacheAligned<T> {
    pub(super) fn new(t: T) -> Self {
        CacheAligned(t, Aligner)
    }
}
<commit_msg>Auto merge of #90284 - tonyyzy:patch-1, r=JohnTitor<commit_after>use crate::ops::{Deref, DerefMut};

#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(align(64))]
pub(super) struct CacheAligned<T>(pub T);

impl<T> Deref for CacheAligned<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> DerefMut for CacheAligned<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<T> CacheAligned<T> {
    pub(super) fn new(t: T) -> Self {
        CacheAligned(t)
    }
}
<|endoftext|>


<commit_before>mod code_actions;
mod extend_selection;
mod folding_ranges;
mod line_index;
mod symbols;
#[cfg(test)]
mod test_utils;
mod typing;

pub use self::{
    code_actions::{add_derive, add_impl, flip_comma, introduce_variable, LocalEdit},
    extend_selection::extend_selection,
    folding_ranges::{folding_ranges, Fold, FoldKind},
    line_index::{LineCol, LineIndex},
    symbols::{file_structure, file_symbols, FileSymbol, StructureNode},
    typing::{join_lines, on_enter, on_eq_typed},
};
use ra_text_edit::{TextEdit, TextEditBuilder};
use ra_syntax::{
    algo::find_leaf_at_offset,
    ast::{self, AstNode, NameOwner},
    SourceFileNode,
    Location,
    SyntaxKind::{self, *},
    SyntaxNodeRef, TextRange, TextUnit,
};

#[derive(Debug)]
pub struct HighlightedRange {
    pub range: TextRange,
    pub tag: &'static str,
}

#[derive(Debug)]
pub struct Diagnostic {
    pub range: TextRange,
    pub msg: String,
}

#[derive(Debug)]
pub struct Runnable {
    pub range: TextRange,
    pub kind: RunnableKind,
}

#[derive(Debug)]
pub enum RunnableKind {
    Test { name: String },
    Bin,
}

pub fn matching_brace(file: &SourceFileNode, offset: TextUnit) -> Option<TextUnit> {
    const BRACES: &[SyntaxKind] = &[
        L_CURLY, R_CURLY, L_BRACK, R_BRACK, L_PAREN, R_PAREN, L_ANGLE, R_ANGLE,
    ];
    let (brace_node, brace_idx) = find_leaf_at_offset(file.syntax(), offset)
        .filter_map(|node| {
            let idx = BRACES.iter().position(|&brace| brace == node.kind())?;
            Some((node, idx))
        })
        .next()?;
    let parent = brace_node.parent()?;
    let matching_kind = BRACES[brace_idx ^ 1];
    let matching_node = parent
        .children()
        .find(|node| node.kind() == matching_kind)?;
    Some(matching_node.range().start())
}

pub fn highlight(file: &SourceFileNode) -> Vec<HighlightedRange> {
    let mut res = Vec::new();
    for node in file.syntax().descendants() {
        let tag = match node.kind() {
            COMMENT => "comment",
            STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => "string",
            ATTR => "attribute",
            NAME_REF => "text",
            NAME => "function",
            INT_NUMBER | FLOAT_NUMBER | CHAR | BYTE => "literal",
            LIFETIME => "parameter",
            k if k.is_keyword() => "keyword",
            _ => continue,
        };
        res.push(HighlightedRange {
            range: node.range(),
            tag,
        })
    }
    res
}

pub fn diagnostics(file: &SourceFileNode) -> Vec<Diagnostic> {
    fn location_to_range(location: Location) -> TextRange {
        match location {
            Location::Offset(offset) => TextRange::offset_len(offset, 1.into()),
            Location::Range(range) => range,
        }
    }

    file.errors()
        .into_iter()
        .map(|err| Diagnostic {
            range: location_to_range(err.location()),
            msg: format!("Syntax Error: {}", err),
        })
        .collect()
}

pub fn syntax_tree(file: &SourceFileNode) -> String {
    ::ra_syntax::utils::dump_tree(file.syntax())
}

pub fn runnables(file: &SourceFileNode) -> Vec<Runnable> {
    file.syntax()
        .descendants()
        .filter_map(ast::FnDef::cast)
        .filter_map(|f| {
            let name = f.name()?.text();
            let kind = if name == "main" {
                RunnableKind::Bin
            } else if f.has_atom_attr("test") {
                RunnableKind::Test {
                    name: name.to_string(),
                }
            } else {
                return None;
            };
            Some(Runnable {
                range: f.syntax().range(),
                kind,
            })
        })
        .collect()
}

pub fn find_node_at_offset<'a, N: AstNode<'a>>(
    syntax: SyntaxNodeRef<'a>,
    offset: TextUnit,
) -> Option<N> {
    find_leaf_at_offset(syntax, offset).find_map(|leaf| leaf.ancestors().find_map(N::cast))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{add_cursor, assert_eq_dbg, extract_offset, assert_eq_text};

    #[test]
    fn test_highlighting() {
        let file = SourceFileNode::parse(
            r#"
// comment
fn main() {}
    println!("Hello, {}!", 92);
"#,
        );
        let hls = highlight(&file);
        assert_eq_dbg(
            r#"[HighlightedRange { range: [1; 11), tag: "comment" },
                HighlightedRange { range: [12; 14), tag: "keyword" },
                HighlightedRange { range: [15; 19), tag: "function" },
                HighlightedRange { range: [29; 36), tag: "text" },
                HighlightedRange { range: [38; 50), tag: "string" },
                HighlightedRange { range: [52; 54), tag: "literal" }]"#,
            &hls,
        );
    }

    #[test]
    fn test_runnables() {
        let file = SourceFileNode::parse(
            r#"
fn main() {}

#[test]
fn test_foo() {}

#[test]
#[ignore]
fn test_foo() {}
"#,
        );
        let runnables = runnables(&file);
        assert_eq_dbg(
            r#"[Runnable { range: [1; 13), kind: Bin },
                Runnable { range: [15; 39), kind: Test { name: "test_foo" } },
                Runnable { range: [41; 75), kind: Test { name: "test_foo" } }]"#,
            &runnables,
        )
    }

    #[test]
    fn test_matching_brace() {
        fn do_check(before: &str, after: &str) {
            let (pos, before) = extract_offset(before);
            let file = SourceFileNode::parse(&before);
            let new_pos = match matching_brace(&file, pos) {
                None => pos,
                Some(pos) => pos,
            };
            let actual = add_cursor(&before, new_pos);
            assert_eq_text!(after, &actual);
        }

        do_check("struct Foo { a: i32, }<|>", "struct Foo <|>{ a: i32, }");
    }
}
<commit_msg>resolved #324: remove unnecessary braces in use statement.<commit_after>mod code_actions;
mod extend_selection;
mod folding_ranges;
mod line_index;
mod symbols;
#[cfg(test)]
mod test_utils;
mod typing;

pub use self::{
    code_actions::{add_derive, add_impl, flip_comma, introduce_variable, LocalEdit},
    extend_selection::extend_selection,
    folding_ranges::{folding_ranges, Fold, FoldKind},
    line_index::{LineCol, LineIndex},
    symbols::{file_structure, file_symbols, FileSymbol, StructureNode},
    typing::{join_lines, on_enter, on_eq_typed},
};
use ra_text_edit::{TextEdit, TextEditBuilder};
use ra_syntax::{
    algo::find_leaf_at_offset,
    ast::{self, AstNode, NameOwner},
    SourceFileNode,
    Location,
    SyntaxKind::{self, *},
    SyntaxNodeRef, TextRange, TextUnit,
};

#[derive(Debug)]
pub struct HighlightedRange {
    pub range: TextRange,
    pub tag: &'static str,
}

#[derive(Debug, Copy, Clone)]
pub enum Severity {
    Error,
    Warning
}

#[derive(Debug)]
pub struct Diagnostic {
    pub range: TextRange,
    pub msg: String,
    pub severity: Severity,
}

#[derive(Debug)]
pub struct Runnable {
    pub range: TextRange,
    pub kind: RunnableKind,
}

#[derive(Debug)]
pub enum RunnableKind {
    Test { name: String },
    Bin,
}

pub fn matching_brace(file: &SourceFileNode, offset: TextUnit) -> Option<TextUnit> {
    const BRACES: &[SyntaxKind] = &[
        L_CURLY, R_CURLY, L_BRACK, R_BRACK, L_PAREN, R_PAREN, L_ANGLE, R_ANGLE,
    ];
    let (brace_node, brace_idx) = find_leaf_at_offset(file.syntax(), offset)
        .filter_map(|node| {
            let idx = BRACES.iter().position(|&brace| brace == node.kind())?;
            Some((node, idx))
        })
        .next()?;
    let parent = brace_node.parent()?;
    let matching_kind = BRACES[brace_idx ^ 1];
    let matching_node = parent
        .children()
        .find(|node| node.kind() == matching_kind)?;
    Some(matching_node.range().start())
}

pub fn highlight(file: &SourceFileNode) -> Vec<HighlightedRange> {
    let mut res = Vec::new();
    for node in file.syntax().descendants() {
        let tag = match node.kind() {
            COMMENT => "comment",
            STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => "string",
            ATTR => "attribute",
            NAME_REF => "text",
            NAME => "function",
            INT_NUMBER | FLOAT_NUMBER | CHAR | BYTE => "literal",
            LIFETIME => "parameter",
            k if k.is_keyword() => "keyword",
            _ => continue,
        };
        res.push(HighlightedRange {
            range: node.range(),
            tag,
        })
    }
    res
}

pub fn diagnostics(file: &SourceFileNode) -> Vec<Diagnostic> {
    fn location_to_range(location: Location) -> TextRange {
        match location {
            Location::Offset(offset) => TextRange::offset_len(offset, 1.into()),
            Location::Range(range) => range,
        }
    }

    let mut errors: Vec<Diagnostic> = file.errors()
        .into_iter()
        .map(|err| Diagnostic {
            range: location_to_range(err.location()),
            msg: format!("Syntax Error: {}", err),
            severity: Severity::Error,
        })
        .collect();

    let warnings = check_unnecessary_braces_in_use_statement(file);

    errors.extend(warnings);
    errors
}

fn check_unnecessary_braces_in_use_statement(file: &SourceFileNode) -> Vec<Diagnostic> {
    let mut diagnostics = Vec::new();
    for node in file.syntax().descendants() {
        if let Some(use_tree_list) = ast::UseTreeList::cast(node) {
            if use_tree_list.use_trees().count() <= 1 {
                diagnostics.push(Diagnostic {
                    range: use_tree_list.syntax().range(),
                    msg: format!("Unnecessary braces in use statement"),
                    severity: Severity::Warning,
                })
            }
        }
    }

    diagnostics
}

pub fn syntax_tree(file: &SourceFileNode) -> String {
    ::ra_syntax::utils::dump_tree(file.syntax())
}

pub fn runnables(file: &SourceFileNode) -> Vec<Runnable> {
    file.syntax()
        .descendants()
        .filter_map(ast::FnDef::cast)
        .filter_map(|f| {
            let name = f.name()?.text();
            let kind = if name == "main" {
                RunnableKind::Bin
            } else if f.has_atom_attr("test") {
                RunnableKind::Test {
                    name: name.to_string(),
                }
            } else {
                return None;
            };
            Some(Runnable {
                range: f.syntax().range(),
                kind,
            })
        })
        .collect()
}

pub fn find_node_at_offset<'a, N: AstNode<'a>>(
    syntax: SyntaxNodeRef<'a>,
    offset: TextUnit,
) -> Option<N> {
    find_leaf_at_offset(syntax, offset).find_map(|leaf| leaf.ancestors().find_map(N::cast))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{add_cursor, assert_eq_dbg, extract_offset, assert_eq_text};

    #[test]
    fn test_highlighting() {
        let file = SourceFileNode::parse(
            r#"
// comment
fn main() {}
    println!("Hello, {}!", 92);
"#,
        );
        let hls = highlight(&file);
        assert_eq_dbg(
            r#"[HighlightedRange { range: [1; 11), tag: "comment" },
                HighlightedRange { range: [12; 14), tag: "keyword" },
                HighlightedRange { range: [15; 19), tag: "function" },
                HighlightedRange { range: [29; 36), tag: "text" },
                HighlightedRange { range: [38; 50), tag: "string" },
                HighlightedRange { range: [52; 54), tag: "literal" }]"#,
            &hls,
        );
    }

    #[test]
    fn test_runnables() {
        let file = SourceFileNode::parse(
            r#"
fn main() {}

#[test]
fn test_foo() {}

#[test]
#[ignore]
fn test_foo() {}
"#,
        );
        let runnables = runnables(&file);
        assert_eq_dbg(
            r#"[Runnable { range: [1; 13), kind: Bin },
                Runnable { range: [15; 39), kind: Test { name: "test_foo" } },
                Runnable { range: [41; 75), kind: Test { name: "test_foo" } }]"#,
            &runnables,
        )
    }

    #[test]
    fn test_matching_brace() {
        fn do_check(before: &str, after: &str) {
            let (pos, before) = extract_offset(before);
            let file = SourceFileNode::parse(&before);
            let new_pos = match matching_brace(&file, pos) {
                None => pos,
                Some(pos) => pos,
            };
            let actual = add_cursor(&before, new_pos);
            assert_eq_text!(after, &actual);
        }

        do_check("struct Foo { a: i32, }<|>", "struct Foo <|>{ a: i32, }");
    }

    #[test]
    fn test_check_unnecessary_braces_in_use_statement() {
        let file = SourceFileNode::parse(
            r#"
use a;
use {b};
use a::{c};
use a::{c, d::e};
use a::{c, d::{e}};
fn main() {}
"#,
        );
        let diagnostics = check_unnecessary_braces_in_use_statement(&file);
        assert_eq_dbg(
            r#"[Diagnostic { range: [12; 15), msg: "Unnecessary braces in use statement", severity: Warning },
                Diagnostic { range: [24; 27), msg: "Unnecessary braces in use statement", severity: Warning },
                Diagnostic { range: [61; 64), msg: "Unnecessary braces in use statement", severity: Warning }]"#,
            &diagnostics,
        )
    }
}
<|endoftext|>


<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.

use std::cell::UnsafeCell;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{spin_loop_hint, AtomicBool, Ordering};

const UNLOCKED: bool = false;
const LOCKED: bool = true;

/// A primitive that provides safe, mutable access to a shared resource.
///
/// Unlike `Mutex`, a `SpinLock` will not voluntarily yield its CPU time until the resource is
/// available and will instead keep spinning until the resource is acquired. For the vast majority
/// of cases, `Mutex` is a better choice than `SpinLock`. If a `SpinLock` must be used then users
/// should try to do as little work as possible while holding the `SpinLock` and avoid any sort of
/// blocking at all costs as it can severely penalize performance.
///
/// # Poisoning
///
/// This `SpinLock` does not implement lock poisoning so it is possible for threads to access
/// poisoned data if a thread panics while holding the lock. If lock poisoning is needed, it can be
/// implemented by wrapping the `SpinLock` in a new type that implements poisoning. See the
/// implementation of `std::sync::Mutex` for an example of how to do this.
#[repr(align(128))]
pub struct SpinLock<T: ?Sized> {
    lock: AtomicBool,
    value: UnsafeCell<T>,
}

impl<T> SpinLock<T> {
    /// Creates a new, unlocked `SpinLock` that's ready for use.
    pub fn new(value: T) -> SpinLock<T> {
        SpinLock {
            lock: AtomicBool::new(UNLOCKED),
            value: UnsafeCell::new(value),
        }
    }

    /// Consumes the `SpinLock` and returns the value guarded by it. This method doesn't perform any
    /// locking as the compiler guarantees that there are no references to `self`.
    pub fn into_inner(self) -> T {
        // No need to take the lock because the compiler can statically guarantee
        // that there are no references to the SpinLock.
        self.value.into_inner()
    }
}

impl<T: ?Sized> SpinLock<T> {
    /// Acquires exclusive, mutable access to the resource protected by the `SpinLock`, blocking the
    /// current thread until it is able to do so. Upon returning, the current thread will be the
    /// only thread with access to the resource. The `SpinLock` will be released when the returned
    /// `SpinLockGuard` is dropped. Attempting to call `lock` while already holding the `SpinLock`
    /// will cause a deadlock.
    pub fn lock(&self) -> SpinLockGuard<T> {
        while self
            .lock
            .compare_exchange_weak(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            spin_loop_hint();
        }

        SpinLockGuard {
            lock: self,
            value: unsafe { &mut *self.value.get() },
        }
    }

    fn unlock(&self) {
        // Don't need to compare and swap because we exclusively hold the lock.
        self.lock.store(UNLOCKED, Ordering::Release);
    }

    /// Returns a mutable reference to the contained value. This method doesn't perform any locking
    /// as the compiler will statically guarantee that there are no other references to `self`.
    pub fn get_mut(&mut self) -> &mut T {
        // Safe because the compiler can statically guarantee that there are no other references to
        // `self`. This is also why we don't need to acquire the lock.
        unsafe { &mut *self.value.get() }
    }
}

unsafe impl<T: ?Sized + Send> Send for SpinLock<T> {}
unsafe impl<T: ?Sized + Send> Sync for SpinLock<T> {}

impl<T: ?Sized + Default> Default for SpinLock<T> {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

impl<T> From<T> for SpinLock<T> {
    fn from(source: T) -> Self {
        Self::new(source)
    }
}

/// An RAII implementation of a "scoped lock" for a `SpinLock`. When this structure is dropped, the
/// lock will be released. The resource protected by the `SpinLock` can be accessed via the `Deref`
/// and `DerefMut` implementations of this structure.
pub struct SpinLockGuard<'a, T: 'a + ?Sized> {
    lock: &'a SpinLock<T>,
    value: &'a mut T,
}

impl<'a, T: ?Sized> Deref for SpinLockGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        self.value
    }
}

impl<'a, T: ?Sized> DerefMut for SpinLockGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        self.value
    }
}

impl<'a, T: ?Sized> Drop for SpinLockGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.unlock();
    }
}

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

    use std::mem;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;
    use std::thread;

    #[derive(PartialEq, Eq, Debug)]
    struct NonCopy(u32);

    #[test]
    fn it_works() {
        let sl = SpinLock::new(NonCopy(13));

        assert_eq!(*sl.lock(), NonCopy(13));
    }

    #[test]
    fn smoke() {
        let sl = SpinLock::new(NonCopy(7));

        mem::drop(sl.lock());
        mem::drop(sl.lock());
    }

    #[test]
    fn send() {
        let sl = SpinLock::new(NonCopy(19));

        thread::spawn(move || {
            let value = sl.lock();
            assert_eq!(*value, NonCopy(19));
        })
        .join()
        .unwrap();
    }

    #[test]
    fn high_contention() {
        const THREADS: usize = 23;
        const ITERATIONS: usize = 101;

        let mut threads = Vec::with_capacity(THREADS);

        let sl = Arc::new(SpinLock::new(0usize));
        for _ in 0..THREADS {
            let sl2 = sl.clone();
            threads.push(thread::spawn(move || {
                for _ in 0..ITERATIONS {
                    *sl2.lock() += 1;
                }
            }));
        }

        for t in threads.into_iter() {
            t.join().unwrap();
        }

        assert_eq!(*sl.lock(), THREADS * ITERATIONS);
    }

    #[test]
    fn get_mut() {
        let mut sl = SpinLock::new(NonCopy(13));
        *sl.get_mut() = NonCopy(17);

        assert_eq!(sl.into_inner(), NonCopy(17));
    }

    #[test]
    fn into_inner() {
        let sl = SpinLock::new(NonCopy(29));
        assert_eq!(sl.into_inner(), NonCopy(29));
    }

    #[test]
    fn into_inner_drop() {
        struct NeedsDrop(Arc<AtomicUsize>);
        impl Drop for NeedsDrop {
            fn drop(&mut self) {
                self.0.fetch_add(1, Ordering::AcqRel);
            }
        }

        let value = Arc::new(AtomicUsize::new(0));
        let needs_drop = SpinLock::new(NeedsDrop(value.clone()));
        assert_eq!(value.load(Ordering::Acquire), 0);

        {
            let inner = needs_drop.into_inner();
            assert_eq!(inner.0.load(Ordering::Acquire), 0);
        }

        assert_eq!(value.load(Ordering::Acquire), 1);
    }

    #[test]
    fn arc_nested() {
        // Tests nested sltexes and access to underlying data.
        let sl = SpinLock::new(1);
        let arc = Arc::new(SpinLock::new(sl));
        thread::spawn(move || {
            let nested = arc.lock();
            let lock2 = nested.lock();
            assert_eq!(*lock2, 1);
        })
        .join()
        .unwrap();
    }

    #[test]
    fn arc_access_in_unwind() {
        let arc = Arc::new(SpinLock::new(1));
        let arc2 = arc.clone();
        thread::spawn(move || {
            struct Unwinder {
                i: Arc<SpinLock<i32>>,
            }
            impl Drop for Unwinder {
                fn drop(&mut self) {
                    *self.i.lock() += 1;
                }
            }
            let _u = Unwinder { i: arc2 };
            panic!();
        })
        .join()
        .expect_err("thread did not panic");
        let lock = arc.lock();
        assert_eq!(*lock, 2);
    }

    #[test]
    fn unsized_value() {
        let sltex: &SpinLock<[i32]> = &SpinLock::new([1, 2, 3]);
        {
            let b = &mut *sltex.lock();
            b[0] = 4;
            b[2] = 5;
        }
        let expected: &[i32] = &[4, 2, 5];
        assert_eq!(&*sltex.lock(), expected);
    }
}
<commit_msg>sync: spinlock: Check state before compare_exchange_weak<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.

use std::cell::UnsafeCell;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{spin_loop_hint, AtomicBool, Ordering};

const UNLOCKED: bool = false;
const LOCKED: bool = true;

/// A primitive that provides safe, mutable access to a shared resource.
///
/// Unlike `Mutex`, a `SpinLock` will not voluntarily yield its CPU time until the resource is
/// available and will instead keep spinning until the resource is acquired. For the vast majority
/// of cases, `Mutex` is a better choice than `SpinLock`. If a `SpinLock` must be used then users
/// should try to do as little work as possible while holding the `SpinLock` and avoid any sort of
/// blocking at all costs as it can severely penalize performance.
///
/// # Poisoning
///
/// This `SpinLock` does not implement lock poisoning so it is possible for threads to access
/// poisoned data if a thread panics while holding the lock. If lock poisoning is needed, it can be
/// implemented by wrapping the `SpinLock` in a new type that implements poisoning. See the
/// implementation of `std::sync::Mutex` for an example of how to do this.
#[repr(align(128))]
pub struct SpinLock<T: ?Sized> {
    lock: AtomicBool,
    value: UnsafeCell<T>,
}

impl<T> SpinLock<T> {
    /// Creates a new, unlocked `SpinLock` that's ready for use.
    pub fn new(value: T) -> SpinLock<T> {
        SpinLock {
            lock: AtomicBool::new(UNLOCKED),
            value: UnsafeCell::new(value),
        }
    }

    /// Consumes the `SpinLock` and returns the value guarded by it. This method doesn't perform any
    /// locking as the compiler guarantees that there are no references to `self`.
    pub fn into_inner(self) -> T {
        // No need to take the lock because the compiler can statically guarantee
        // that there are no references to the SpinLock.
        self.value.into_inner()
    }
}

impl<T: ?Sized> SpinLock<T> {
    /// Acquires exclusive, mutable access to the resource protected by the `SpinLock`, blocking the
    /// current thread until it is able to do so. Upon returning, the current thread will be the
    /// only thread with access to the resource. The `SpinLock` will be released when the returned
    /// `SpinLockGuard` is dropped. Attempting to call `lock` while already holding the `SpinLock`
    /// will cause a deadlock.
    pub fn lock(&self) -> SpinLockGuard<T> {
        loop {
            let state = self.lock.load(Ordering::Relaxed);
            if state == UNLOCKED
                && self
                    .lock
                    .compare_exchange_weak(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
                    .is_ok()
            {
                break;
            }
            spin_loop_hint();
        }

        SpinLockGuard {
            lock: self,
            value: unsafe { &mut *self.value.get() },
        }
    }

    fn unlock(&self) {
        // Don't need to compare and swap because we exclusively hold the lock.
        self.lock.store(UNLOCKED, Ordering::Release);
    }

    /// Returns a mutable reference to the contained value. This method doesn't perform any locking
    /// as the compiler will statically guarantee that there are no other references to `self`.
    pub fn get_mut(&mut self) -> &mut T {
        // Safe because the compiler can statically guarantee that there are no other references to
        // `self`. This is also why we don't need to acquire the lock.
        unsafe { &mut *self.value.get() }
    }
}

unsafe impl<T: ?Sized + Send> Send for SpinLock<T> {}
unsafe impl<T: ?Sized + Send> Sync for SpinLock<T> {}

impl<T: ?Sized + Default> Default for SpinLock<T> {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

impl<T> From<T> for SpinLock<T> {
    fn from(source: T) -> Self {
        Self::new(source)
    }
}

/// An RAII implementation of a "scoped lock" for a `SpinLock`. When this structure is dropped, the
/// lock will be released. The resource protected by the `SpinLock` can be accessed via the `Deref`
/// and `DerefMut` implementations of this structure.
pub struct SpinLockGuard<'a, T: 'a + ?Sized> {
    lock: &'a SpinLock<T>,
    value: &'a mut T,
}

impl<'a, T: ?Sized> Deref for SpinLockGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        self.value
    }
}

impl<'a, T: ?Sized> DerefMut for SpinLockGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        self.value
    }
}

impl<'a, T: ?Sized> Drop for SpinLockGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.unlock();
    }
}

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

    use std::mem;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;
    use std::thread;

    #[derive(PartialEq, Eq, Debug)]
    struct NonCopy(u32);

    #[test]
    fn it_works() {
        let sl = SpinLock::new(NonCopy(13));

        assert_eq!(*sl.lock(), NonCopy(13));
    }

    #[test]
    fn smoke() {
        let sl = SpinLock::new(NonCopy(7));

        mem::drop(sl.lock());
        mem::drop(sl.lock());
    }

    #[test]
    fn send() {
        let sl = SpinLock::new(NonCopy(19));

        thread::spawn(move || {
            let value = sl.lock();
            assert_eq!(*value, NonCopy(19));
        })
        .join()
        .unwrap();
    }

    #[test]
    fn high_contention() {
        const THREADS: usize = 23;
        const ITERATIONS: usize = 101;

        let mut threads = Vec::with_capacity(THREADS);

        let sl = Arc::new(SpinLock::new(0usize));
        for _ in 0..THREADS {
            let sl2 = sl.clone();
            threads.push(thread::spawn(move || {
                for _ in 0..ITERATIONS {
                    *sl2.lock() += 1;
                }
            }));
        }

        for t in threads.into_iter() {
            t.join().unwrap();
        }

        assert_eq!(*sl.lock(), THREADS * ITERATIONS);
    }

    #[test]
    fn get_mut() {
        let mut sl = SpinLock::new(NonCopy(13));
        *sl.get_mut() = NonCopy(17);

        assert_eq!(sl.into_inner(), NonCopy(17));
    }

    #[test]
    fn into_inner() {
        let sl = SpinLock::new(NonCopy(29));
        assert_eq!(sl.into_inner(), NonCopy(29));
    }

    #[test]
    fn into_inner_drop() {
        struct NeedsDrop(Arc<AtomicUsize>);
        impl Drop for NeedsDrop {
            fn drop(&mut self) {
                self.0.fetch_add(1, Ordering::AcqRel);
            }
        }

        let value = Arc::new(AtomicUsize::new(0));
        let needs_drop = SpinLock::new(NeedsDrop(value.clone()));
        assert_eq!(value.load(Ordering::Acquire), 0);

        {
            let inner = needs_drop.into_inner();
            assert_eq!(inner.0.load(Ordering::Acquire), 0);
        }

        assert_eq!(value.load(Ordering::Acquire), 1);
    }

    #[test]
    fn arc_nested() {
        // Tests nested sltexes and access to underlying data.
        let sl = SpinLock::new(1);
        let arc = Arc::new(SpinLock::new(sl));
        thread::spawn(move || {
            let nested = arc.lock();
            let lock2 = nested.lock();
            assert_eq!(*lock2, 1);
        })
        .join()
        .unwrap();
    }

    #[test]
    fn arc_access_in_unwind() {
        let arc = Arc::new(SpinLock::new(1));
        let arc2 = arc.clone();
        thread::spawn(move || {
            struct Unwinder {
                i: Arc<SpinLock<i32>>,
            }
            impl Drop for Unwinder {
                fn drop(&mut self) {
                    *self.i.lock() += 1;
                }
            }
            let _u = Unwinder { i: arc2 };
            panic!();
        })
        .join()
        .expect_err("thread did not panic");
        let lock = arc.lock();
        assert_eq!(*lock, 2);
    }

    #[test]
    fn unsized_value() {
        let sltex: &SpinLock<[i32]> = &SpinLock::new([1, 2, 3]);
        {
            let b = &mut *sltex.lock();
            b[0] = 4;
            b[2] = 5;
        }
        let expected: &[i32] = &[4, 2, 5];
        assert_eq!(&*sltex.lock(), expected);
    }
}
<|endoftext|>