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

//! minigbm: implements swapchain allocation using ChromeOS's minigbm library.
//!
//! External code found at <https://chromium.googlesource.com/chromiumos/platform/minigbm>.

#![cfg(feature = "minigbm")]

use std::ffi::CStr;
use std::fs::File;
use std::io::Seek;
use std::io::SeekFrom;
use std::os::raw::c_char;
use std::sync::Arc;

use base::AsRawDescriptor;
use base::Error as BaseError;
use base::FromRawDescriptor;
use base::RawDescriptor;

use crate::rutabaga_gralloc::formats::DrmFormat;
use crate::rutabaga_gralloc::gralloc::Gralloc;
use crate::rutabaga_gralloc::gralloc::ImageAllocationInfo;
use crate::rutabaga_gralloc::gralloc::ImageMemoryRequirements;
use crate::rutabaga_gralloc::minigbm_bindings::*;
use crate::rutabaga_gralloc::rendernode;
use crate::rutabaga_utils::*;

struct MinigbmDeviceInner {
    fd: File,
    gbm: *mut gbm_device,
}

// Safe because minigbm handles synchronization internally.
unsafe impl Send for MinigbmDeviceInner {}
unsafe impl Sync for MinigbmDeviceInner {}

impl Drop for MinigbmDeviceInner {
    fn drop(&mut self) {
        // Safe because MinigbmDeviceInner is only constructed with a valid minigbm_device.
        unsafe {
            gbm_device_destroy(self.gbm);
        }
    }
}

/// A device capable of allocating `MinigbmBuffer`.
#[derive(Clone)]
pub struct MinigbmDevice {
    minigbm_device: Arc<MinigbmDeviceInner>,
    last_buffer: Option<Arc<MinigbmBuffer>>,
    device_name: &'static str,
}

impl MinigbmDevice {
    /// Returns a new `MinigbmDevice` if there is a rendernode in `/dev/dri/` that is accepted by
    /// the minigbm library.
    pub fn init() -> RutabagaResult<Box<dyn Gralloc>> {
        let undesired: &[&str] = &["vgem", "pvr"];
        let fd = rendernode::open_device(undesired)?;

        // gbm_create_device is safe to call with a valid fd, and we check that a valid one is
        // returned.  If the fd does not refer to a DRM device, gbm_create_device will reject it.
        let gbm = unsafe { gbm_create_device(fd.as_raw_descriptor()) };
        if gbm.is_null() {
            return Err(RutabagaError::BaseError(BaseError::last()));
        }

        // Safe because a valid minigbm device has a statically allocated string associated with
        // it, which is valid for the lifetime of the process.
        let backend_name: *const c_char = unsafe { gbm_device_get_backend_name(gbm) };
        let c_str: &CStr = unsafe { CStr::from_ptr(backend_name) };
        let device_name: &str = c_str.to_str()?;

        Ok(Box::new(MinigbmDevice {
            minigbm_device: Arc::new(MinigbmDeviceInner { fd, gbm }),
            last_buffer: None,
            device_name,
        }))
    }
}

impl Gralloc for MinigbmDevice {
    fn supports_external_gpu_memory(&self) -> bool {
        true
    }

    fn supports_dmabuf(&self) -> bool {
        true
    }

    fn get_image_memory_requirements(
        &mut self,
        info: ImageAllocationInfo,
    ) -> RutabagaResult<ImageMemoryRequirements> {
        let bo = unsafe {
            gbm_bo_create(
                self.minigbm_device.gbm,
                info.width,
                info.height,
                info.drm_format.0,
                info.flags.0,
            )
        };
        if bo.is_null() {
            return Err(RutabagaError::BaseError(BaseError::last()));
        }

        let mut reqs: ImageMemoryRequirements = Default::default();
        let gbm_buffer = MinigbmBuffer(bo, self.clone());

        // Intel GPUs typically only use cached memory buffers.  This will change with dGPUs, but
        // perhaps minigbm will be deprecated by then.  Other display drivers (rockchip, mediatek,
        // amdgpu) typically use write combine memory.  We can also consider use flags too if this
        // heuristic proves insufficient.
        //
        // Existing qcom devices use a mix of cached and WC buffers.  We don't *yet* have a good
        // way to differentiate, but https://patchwork.freedesktop.org/series/106847/ is a proposal
        // to fix that.  But existing devices without the FWB feature, the resulting mapping attrs
        // are the more restrictive of the combination for S1 and S2 mappings.  So if we map as
        // cached in S2 pgtables, but S1 has WC, then the result will be WC.
        if self.device_name == "i915" || self.device_name == "msm" {
            reqs.map_info = RUTABAGA_MAP_CACHE_CACHED;
        } else {
            reqs.map_info = RUTABAGA_MAP_CACHE_WC;
        }

        reqs.modifier = gbm_buffer.format_modifier();
        for plane in 0..gbm_buffer.num_planes() {
            reqs.strides[plane] = gbm_buffer.plane_stride(plane);
            reqs.offsets[plane] = gbm_buffer.plane_offset(plane);
        }

        let mut fd = gbm_buffer.export()?;
        let size = fd.seek(SeekFrom::End(0))?;

        // minigbm does have the ability to query image requirements without allocating memory
        // via the TEST_ALLOC flag.  However, support has only been added in i915.  Until this
        // flag is supported everywhere, do the actual allocation here and stash it away.
        if self.last_buffer.is_some() {
            return Err(RutabagaError::AlreadyInUse);
        }

        self.last_buffer = Some(Arc::new(gbm_buffer));
        reqs.info = info;
        reqs.size = size;
        Ok(reqs)
    }

    fn allocate_memory(&mut self, reqs: ImageMemoryRequirements) -> RutabagaResult<RutabagaHandle> {
        let last_buffer = self.last_buffer.take();
        if let Some(gbm_buffer) = last_buffer {
            if gbm_buffer.width() != reqs.info.width
                || gbm_buffer.height() != reqs.info.height
                || gbm_buffer.format() != reqs.info.drm_format
            {
                return Err(RutabagaError::InvalidGrallocDimensions);
            }

            let dmabuf = gbm_buffer.export()?.into();
            return Ok(RutabagaHandle {
                os_handle: dmabuf,
                handle_type: RUTABAGA_MEM_HANDLE_TYPE_DMABUF,
            });
        }

        let bo = unsafe {
            gbm_bo_create(
                self.minigbm_device.gbm,
                reqs.info.width,
                reqs.info.height,
                reqs.info.drm_format.0,
                reqs.info.flags.0,
            )
        };

        if bo.is_null() {
            return Err(RutabagaError::BaseError(BaseError::last()));
        }

        let gbm_buffer = MinigbmBuffer(bo, self.clone());
        let dmabuf = gbm_buffer.export()?.into();
        Ok(RutabagaHandle {
            os_handle: dmabuf,
            handle_type: RUTABAGA_MEM_HANDLE_TYPE_DMABUF,
        })
    }
}

/// An allocation from a `MinigbmDevice`.
pub struct MinigbmBuffer(*mut gbm_bo, MinigbmDevice);

// Safe because minigbm handles synchronization internally.
unsafe impl Send for MinigbmBuffer {}
unsafe impl Sync for MinigbmBuffer {}

impl MinigbmBuffer {
    /// Width in pixels.
    pub fn width(&self) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_width(self.0) }
    }

    /// Height in pixels.
    pub fn height(&self) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_height(self.0) }
    }

    /// `DrmFormat` of the buffer.
    pub fn format(&self) -> DrmFormat {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { DrmFormat(gbm_bo_get_format(self.0)) }
    }

    /// DrmFormat modifier flags for the buffer.
    pub fn format_modifier(&self) -> u64 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_modifier(self.0) }
    }

    /// Number of planes present in this buffer.
    pub fn num_planes(&self) -> usize {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_plane_count(self.0) as usize }
    }

    /// Offset in bytes for the given plane.
    pub fn plane_offset(&self, plane: usize) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_offset(self.0, plane) }
    }

    /// Length in bytes of one row for the given plane.
    pub fn plane_stride(&self, plane: usize) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_stride_for_plane(self.0, plane) }
    }

    /// Exports a new dmabuf/prime file descriptor.
    pub fn export(&self) -> RutabagaResult<File> {
        // This is always safe to call with a valid gbm_bo pointer.
        match unsafe { gbm_bo_get_fd(self.0) } {
            fd if fd >= 0 => {
                let dmabuf = unsafe { File::from_raw_descriptor(fd) };
                Ok(dmabuf)
            }
            ret => Err(RutabagaError::ComponentError(ret)),
        }
    }
}

impl Drop for MinigbmBuffer {
    fn drop(&mut self) {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_destroy(self.0) }
    }
}
<commit_msg>rutabaga_gfx: silence some compiler warnings<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.

//! minigbm: implements swapchain allocation using ChromeOS's minigbm library.
//!
//! External code found at <https://chromium.googlesource.com/chromiumos/platform/minigbm>.

#![cfg(feature = "minigbm")]

use std::ffi::CStr;
use std::fs::File;
use std::io::Seek;
use std::io::SeekFrom;
use std::os::raw::c_char;
use std::sync::Arc;

use base::AsRawDescriptor;
use base::Error as BaseError;
use base::FromRawDescriptor;

use crate::rutabaga_gralloc::formats::DrmFormat;
use crate::rutabaga_gralloc::gralloc::Gralloc;
use crate::rutabaga_gralloc::gralloc::ImageAllocationInfo;
use crate::rutabaga_gralloc::gralloc::ImageMemoryRequirements;
use crate::rutabaga_gralloc::minigbm_bindings::*;
use crate::rutabaga_gralloc::rendernode;
use crate::rutabaga_utils::*;

struct MinigbmDeviceInner {
    _fd: File,
    gbm: *mut gbm_device,
}

// Safe because minigbm handles synchronization internally.
unsafe impl Send for MinigbmDeviceInner {}
unsafe impl Sync for MinigbmDeviceInner {}

impl Drop for MinigbmDeviceInner {
    fn drop(&mut self) {
        // Safe because MinigbmDeviceInner is only constructed with a valid minigbm_device.
        unsafe {
            gbm_device_destroy(self.gbm);
        }
    }
}

/// A device capable of allocating `MinigbmBuffer`.
#[derive(Clone)]
pub struct MinigbmDevice {
    minigbm_device: Arc<MinigbmDeviceInner>,
    last_buffer: Option<Arc<MinigbmBuffer>>,
    device_name: &'static str,
}

impl MinigbmDevice {
    /// Returns a new `MinigbmDevice` if there is a rendernode in `/dev/dri/` that is accepted by
    /// the minigbm library.
    pub fn init() -> RutabagaResult<Box<dyn Gralloc>> {
        let undesired: &[&str] = &["vgem", "pvr"];
        let fd = rendernode::open_device(undesired)?;

        // gbm_create_device is safe to call with a valid fd, and we check that a valid one is
        // returned.  If the fd does not refer to a DRM device, gbm_create_device will reject it.
        let gbm = unsafe { gbm_create_device(fd.as_raw_descriptor()) };
        if gbm.is_null() {
            return Err(RutabagaError::BaseError(BaseError::last()));
        }

        // Safe because a valid minigbm device has a statically allocated string associated with
        // it, which is valid for the lifetime of the process.
        let backend_name: *const c_char = unsafe { gbm_device_get_backend_name(gbm) };
        let c_str: &CStr = unsafe { CStr::from_ptr(backend_name) };
        let device_name: &str = c_str.to_str()?;

        Ok(Box::new(MinigbmDevice {
            minigbm_device: Arc::new(MinigbmDeviceInner { _fd: fd, gbm }),
            last_buffer: None,
            device_name,
        }))
    }
}

impl Gralloc for MinigbmDevice {
    fn supports_external_gpu_memory(&self) -> bool {
        true
    }

    fn supports_dmabuf(&self) -> bool {
        true
    }

    fn get_image_memory_requirements(
        &mut self,
        info: ImageAllocationInfo,
    ) -> RutabagaResult<ImageMemoryRequirements> {
        let bo = unsafe {
            gbm_bo_create(
                self.minigbm_device.gbm,
                info.width,
                info.height,
                info.drm_format.0,
                info.flags.0,
            )
        };
        if bo.is_null() {
            return Err(RutabagaError::BaseError(BaseError::last()));
        }

        let mut reqs: ImageMemoryRequirements = Default::default();
        let gbm_buffer = MinigbmBuffer(bo, self.clone());

        // Intel GPUs typically only use cached memory buffers.  This will change with dGPUs, but
        // perhaps minigbm will be deprecated by then.  Other display drivers (rockchip, mediatek,
        // amdgpu) typically use write combine memory.  We can also consider use flags too if this
        // heuristic proves insufficient.
        //
        // Existing qcom devices use a mix of cached and WC buffers.  We don't *yet* have a good
        // way to differentiate, but https://patchwork.freedesktop.org/series/106847/ is a proposal
        // to fix that.  But existing devices without the FWB feature, the resulting mapping attrs
        // are the more restrictive of the combination for S1 and S2 mappings.  So if we map as
        // cached in S2 pgtables, but S1 has WC, then the result will be WC.
        if self.device_name == "i915" || self.device_name == "msm" {
            reqs.map_info = RUTABAGA_MAP_CACHE_CACHED;
        } else {
            reqs.map_info = RUTABAGA_MAP_CACHE_WC;
        }

        reqs.modifier = gbm_buffer.format_modifier();
        for plane in 0..gbm_buffer.num_planes() {
            reqs.strides[plane] = gbm_buffer.plane_stride(plane);
            reqs.offsets[plane] = gbm_buffer.plane_offset(plane);
        }

        let mut fd = gbm_buffer.export()?;
        let size = fd.seek(SeekFrom::End(0))?;

        // minigbm does have the ability to query image requirements without allocating memory
        // via the TEST_ALLOC flag.  However, support has only been added in i915.  Until this
        // flag is supported everywhere, do the actual allocation here and stash it away.
        if self.last_buffer.is_some() {
            return Err(RutabagaError::AlreadyInUse);
        }

        self.last_buffer = Some(Arc::new(gbm_buffer));
        reqs.info = info;
        reqs.size = size;
        Ok(reqs)
    }

    fn allocate_memory(&mut self, reqs: ImageMemoryRequirements) -> RutabagaResult<RutabagaHandle> {
        let last_buffer = self.last_buffer.take();
        if let Some(gbm_buffer) = last_buffer {
            if gbm_buffer.width() != reqs.info.width
                || gbm_buffer.height() != reqs.info.height
                || gbm_buffer.format() != reqs.info.drm_format
            {
                return Err(RutabagaError::InvalidGrallocDimensions);
            }

            let dmabuf = gbm_buffer.export()?.into();
            return Ok(RutabagaHandle {
                os_handle: dmabuf,
                handle_type: RUTABAGA_MEM_HANDLE_TYPE_DMABUF,
            });
        }

        let bo = unsafe {
            gbm_bo_create(
                self.minigbm_device.gbm,
                reqs.info.width,
                reqs.info.height,
                reqs.info.drm_format.0,
                reqs.info.flags.0,
            )
        };

        if bo.is_null() {
            return Err(RutabagaError::BaseError(BaseError::last()));
        }

        let gbm_buffer = MinigbmBuffer(bo, self.clone());
        let dmabuf = gbm_buffer.export()?.into();
        Ok(RutabagaHandle {
            os_handle: dmabuf,
            handle_type: RUTABAGA_MEM_HANDLE_TYPE_DMABUF,
        })
    }
}

/// An allocation from a `MinigbmDevice`.
pub struct MinigbmBuffer(*mut gbm_bo, MinigbmDevice);

// Safe because minigbm handles synchronization internally.
unsafe impl Send for MinigbmBuffer {}
unsafe impl Sync for MinigbmBuffer {}

impl MinigbmBuffer {
    /// Width in pixels.
    pub fn width(&self) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_width(self.0) }
    }

    /// Height in pixels.
    pub fn height(&self) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_height(self.0) }
    }

    /// `DrmFormat` of the buffer.
    pub fn format(&self) -> DrmFormat {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { DrmFormat(gbm_bo_get_format(self.0)) }
    }

    /// DrmFormat modifier flags for the buffer.
    pub fn format_modifier(&self) -> u64 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_modifier(self.0) }
    }

    /// Number of planes present in this buffer.
    pub fn num_planes(&self) -> usize {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_plane_count(self.0) as usize }
    }

    /// Offset in bytes for the given plane.
    pub fn plane_offset(&self, plane: usize) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_offset(self.0, plane) }
    }

    /// Length in bytes of one row for the given plane.
    pub fn plane_stride(&self, plane: usize) -> u32 {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_get_stride_for_plane(self.0, plane) }
    }

    /// Exports a new dmabuf/prime file descriptor.
    pub fn export(&self) -> RutabagaResult<File> {
        // This is always safe to call with a valid gbm_bo pointer.
        match unsafe { gbm_bo_get_fd(self.0) } {
            fd if fd >= 0 => {
                let dmabuf = unsafe { File::from_raw_descriptor(fd) };
                Ok(dmabuf)
            }
            ret => Err(RutabagaError::ComponentError(ret)),
        }
    }
}

impl Drop for MinigbmBuffer {
    fn drop(&mut self) {
        // This is always safe to call with a valid gbm_bo pointer.
        unsafe { gbm_bo_destroy(self.0) }
    }
}
<|endoftext|>


<commit_before>//! A server builder helper for the integration tests.

use std::io::{Read, Write};
use std::net;
use std::time::Duration;
use std::thread;

pub struct Server {
    addr: net::SocketAddr,
}

impl Server {
    pub fn addr(&self) -> net::SocketAddr {
        self.addr
    }
}

#[derive(Default)]
pub struct Txn {
    pub request: Vec<u8>,
    pub response: Vec<u8>,

    pub read_timeout: Option<Duration>,
    pub response_timeout: Option<Duration>,
    pub write_timeout: Option<Duration>,
    pub chunk_size: Option<usize>,
}

static DEFAULT_USER_AGENT: &'static str =
    concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));

pub fn spawn(txns: Vec<Txn>) -> Server {
    let listener = net::TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    thread::spawn(
        move || for txn in txns {
            let mut expected = txn.request;
            let reply = txn.response;
            let (mut socket, _addr) = listener.accept().unwrap();

            socket.set_read_timeout(Some(Duration::from_secs(5))).unwrap();

            replace_expected_vars(&mut expected, addr.to_string().as_ref(), DEFAULT_USER_AGENT.as_ref());

            if let Some(dur) = txn.read_timeout {
                thread::park_timeout(dur);
            }

            let mut buf = [0; 4096];
            assert!(buf.len() >= expected.len());

            let mut n = 0;
            while n < expected.len() {
                n += socket.read(&mut buf).unwrap();
            }

            match (::std::str::from_utf8(&expected), ::std::str::from_utf8(&buf[..n])) {
                (Ok(expected), Ok(received)) => assert_eq!(expected, received),
                _ => assert_eq!(expected, &buf[..n]),
            }

            if let Some(dur) = txn.response_timeout {
                thread::park_timeout(dur);
            }

            if let Some(dur) = txn.write_timeout {
                let headers_end = b"\r\n\r\n";
                let headers_end = reply.windows(headers_end.len()).position(|w| w == headers_end).unwrap() + 4;
                socket.write_all(&reply[..headers_end]).unwrap();

                let body = &reply[headers_end..];

                if let Some(chunk_size) = txn.chunk_size {
                    for content in body.chunks(chunk_size) {
                        thread::park_timeout(dur);
                        socket.write_all(&content).unwrap();
                    }
                } else {
                    thread::park_timeout(dur);
                    socket.write_all(&body).unwrap();
                }
            } else {
                socket.write_all(&reply).unwrap();
            }
        }
    );

    Server {
        addr: addr,
    }
}

fn replace_expected_vars(bytes: &mut Vec<u8>, host: &[u8], ua: &[u8]) {
    // plenty horrible, but these are just tests, and gets the job done
    let mut index = 0;
    loop {
        if index == bytes.len() {
            return;
        }

        for b in (&bytes[index..]).iter() {
            index += 1;
            if *b == b'$' {
                break;
            }
        }

        let has_host = (&bytes[index..]).starts_with(b"HOST");
        if has_host {
            bytes.drain(index - 1..index + 4);
            for (i, b) in host.iter().enumerate() {
                bytes.insert(index - 1 + i, *b);
            }
        } else {
            let has_ua = (&bytes[index..]).starts_with(b"USERAGENT");
            if has_ua {
                bytes.drain(index - 1..index + 9);
                for (i, b) in ua.iter().enumerate() {
                    bytes.insert(index - 1 + i, *b);
                }
            }
        }
    }
}

#[macro_export]
macro_rules! server {
    ($(request: $req:expr, response: $res:expr),*) => ({
        server!($(request: $req, response: $res;)*)
    });
    ($($($f:ident: $v:expr),*);*) => ({
        let txns = vec![
            $(__internal__txn! {
                $($f: $v,)*
            }),*
        ];
        ::support::server::spawn(txns)
    })
}

#[macro_export]
macro_rules! __internal__txn {
    ($($field:ident: $val:expr,)*) => (
        ::support::server::Txn {
            $( $field: __internal__prop!($field: $val), )*
            .. Default::default()
        }
    )
}

#[macro_export]
macro_rules! __internal__prop {
    (request: $val:expr) => (
        From::from(&$val[..])
    );
    (response: $val:expr) => (
        From::from(&$val[..])
    );
    ($field:ident: $val:expr) => (
        From::from($val)
    )
}
<commit_msg>fix: tests::support::server<commit_after>//! A server builder helper for the integration tests.

use std::io::{Read, Write};
use std::net;
use std::time::Duration;
use std::thread;

pub struct Server {
    addr: net::SocketAddr,
}

impl Server {
    pub fn addr(&self) -> net::SocketAddr {
        self.addr
    }
}

#[derive(Default)]
pub struct Txn {
    pub request: Vec<u8>,
    pub response: Vec<u8>,

    pub read_timeout: Option<Duration>,
    pub response_timeout: Option<Duration>,
    pub write_timeout: Option<Duration>,
    pub chunk_size: Option<usize>,
}

static DEFAULT_USER_AGENT: &'static str =
    concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));

pub fn spawn(txns: Vec<Txn>) -> Server {
    let listener = net::TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    thread::spawn(
        move || for txn in txns {
            let mut expected = txn.request;
            let reply = txn.response;
            let (mut socket, _addr) = listener.accept().unwrap();

            socket.set_read_timeout(Some(Duration::from_secs(5))).unwrap();

            replace_expected_vars(&mut expected, addr.to_string().as_ref(), DEFAULT_USER_AGENT.as_ref());

            if let Some(dur) = txn.read_timeout {
                thread::park_timeout(dur);
            }

            let mut buf = [0; 4096];
            assert!(buf.len() >= expected.len());

            let mut n = 0;
            while n < expected.len() {
                match socket.read(&mut buf[n..]) {
                    Ok(0) | Err(_) => break,
                    Ok(nread) => n += nread,
                }
            }

            match (::std::str::from_utf8(&expected), ::std::str::from_utf8(&buf[..n])) {
                (Ok(expected), Ok(received)) => assert_eq!(expected, received),
                _ => assert_eq!(expected, &buf[..n]),
            }

            if let Some(dur) = txn.response_timeout {
                thread::park_timeout(dur);
            }

            if let Some(dur) = txn.write_timeout {
                let headers_end = b"\r\n\r\n";
                let headers_end = reply.windows(headers_end.len()).position(|w| w == headers_end).unwrap() + 4;
                socket.write_all(&reply[..headers_end]).unwrap();

                let body = &reply[headers_end..];

                if let Some(chunk_size) = txn.chunk_size {
                    for content in body.chunks(chunk_size) {
                        thread::park_timeout(dur);
                        socket.write_all(&content).unwrap();
                    }
                } else {
                    thread::park_timeout(dur);
                    socket.write_all(&body).unwrap();
                }
            } else {
                socket.write_all(&reply).unwrap();
            }
        }
    );

    Server {
        addr: addr,
    }
}

fn replace_expected_vars(bytes: &mut Vec<u8>, host: &[u8], ua: &[u8]) {
    // plenty horrible, but these are just tests, and gets the job done
    let mut index = 0;
    loop {
        if index == bytes.len() {
            return;
        }

        for b in (&bytes[index..]).iter() {
            index += 1;
            if *b == b'$' {
                break;
            }
        }

        let has_host = (&bytes[index..]).starts_with(b"HOST");
        if has_host {
            bytes.drain(index - 1..index + 4);
            for (i, b) in host.iter().enumerate() {
                bytes.insert(index - 1 + i, *b);
            }
        } else {
            let has_ua = (&bytes[index..]).starts_with(b"USERAGENT");
            if has_ua {
                bytes.drain(index - 1..index + 9);
                for (i, b) in ua.iter().enumerate() {
                    bytes.insert(index - 1 + i, *b);
                }
            }
        }
    }
}

#[macro_export]
macro_rules! server {
    ($(request: $req:expr, response: $res:expr),*) => ({
        server!($(request: $req, response: $res;)*)
    });
    ($($($f:ident: $v:expr),*);*) => ({
        let txns = vec![
            $(__internal__txn! {
                $($f: $v,)*
            }),*
        ];
        ::support::server::spawn(txns)
    })
}

#[macro_export]
macro_rules! __internal__txn {
    ($($field:ident: $val:expr,)*) => (
        ::support::server::Txn {
            $( $field: __internal__prop!($field: $val), )*
            .. Default::default()
        }
    )
}

#[macro_export]
macro_rules! __internal__prop {
    (request: $val:expr) => (
        From::from(&$val[..])
    );
    (response: $val:expr) => (
        From::from(&$val[..])
    );
    ($field:ident: $val:expr) => (
        From::from($val)
    )
}
<|endoftext|>


<commit_before>// Allow unstable items until Rust hits 1.0
#![feature(core)]
#![feature(io)]
#![feature(env)]
#![feature(path)]

extern crate yaml;

use yaml::constructor::YamlStandardData;
use yaml::ffi::YamlEncoding;
use std::env;
use std::path;
use std::io::BufReader;
use std::fs::File;

fn match_utf8(filename: &str, expected: YamlStandardData) {
    match_file(YamlEncoding::YamlUtf8Encoding, filename, expected);
}

fn match_file(encoding: yaml::ffi::YamlEncoding, filename: &str, expected: YamlStandardData) {
    let mut args = env::args();
    let this_path = args.next().unwrap();
    let dir_path = path::Path::new(this_path.as_slice()).parent().unwrap();
    let file_path = dir_path.join("../tests/source").join(filename);
    println!("{}", file_path.display());
    let mut file = match File::open(&file_path) {
        Ok(f) => f,
        Err(e) => panic!("Failed to open file {}: {}", file_path.to_str().unwrap(), e)
    };
    let mut reader = BufReader::new(file);
    match yaml::parse_io(&mut reader, encoding) {
        Ok(docs) => if docs.len() == 1 {
            assert_eq!(docs.as_slice().first().unwrap(), &expected)
        } else {
            panic!("too many number of documents: {:?}", docs)
        },
        Err(e) => panic!("parse failure: {:?}", e)
    }
}

macro_rules! ystr{
    ($e: expr) => (
        YamlStandardData::YamlString($e.to_string())
    )
}

macro_rules! yint{
    ($e: expr) => (
        YamlStandardData::YamlInteger($e)
    )
}

macro_rules! yfloat{
    ($e: expr) => (
        YamlStandardData::YamlFloat($e)
    )
}

macro_rules! yseq{
    ($($e:expr),*) => (
        YamlStandardData::YamlSequence(vec![$(($e),)*])
    )
}

macro_rules! ymap{
    ($($k:expr => $v:expr),*) => (
        YamlStandardData::YamlMapping(vec![$((ystr!($k), $v),)*])
    )
}

macro_rules! y_cmp_map{
    ($($k:expr => $v:expr),*) => (
        YamlStandardData::YamlMapping(vec![$(($k, $v),)*])
    )
}

#[test]
fn sequence_of_scalars() {
    match_utf8("ball_players.yml", yseq![ystr!("Mark McGwire"), ystr!("Sammy Sosa"), ystr!("Ken Griffey")]);
}

#[test]
fn scalar_mappings() {
    match_utf8("player_stat.yml", ymap!{
                                    "hr" => yint!(65),
                                    "avg" => yfloat!(0.278),
                                    "rbi" => yint!(147)
                                })
}

#[test]
fn maps_of_sequences() {
    match_utf8("ball_clubs.yml", ymap!{
                                    "american" => yseq![ystr!("Boston Red Sox"), ystr!("Detroit Tigers"), ystr!("New York Yankees")],
                                    "national" => yseq![ystr!("New York Mets"), ystr!("Chicago Cubs"), ystr!("Atlanta Braves")]
                                })
}

#[test]
fn sequence_of_maps() {
    match_utf8("multiple_player_stat.yml",
    yseq![
        ymap!{
            "name" => ystr!("Mark McGwire"),
            "hr" => yint!(65)
        },
        ymap!{
            "name" => ystr!("Sammy Sosa"),
            "hr" => yint!(63)
        }
    ])
}

#[test]
fn sequence_of_sequences() {
    match_utf8("csv.yml",
    yseq![
        yseq![ystr!("name"), ystr!("hr")],
        yseq![ystr!("Mark McGwire"), yint!(65)],
        yseq![ystr!("Sammy Sosa"), yint!(63)]
    ])
}

#[test]
fn mapping_of_mappings() {
    match_utf8("map_map.yml",
    ymap!{
        "Mark McGwire" => ymap!{ "hr" => yint!(65) },
        "Sammy Sosa" => ymap!{ "hr" => yint!(63) }
    })
}

#[test]
fn alias() {
    match_utf8("alias.yml",
    ymap!{
        "hr" => yseq![ystr!("Mark McGwire"), ystr!("Sammy Sosa")],
        "rbi" => yseq![ystr!("Sammy Sosa"), ystr!("Ken Griffey")]
    })
}

#[test]
fn complex_keys() {
    match_utf8("complex_key.yml",
    y_cmp_map!{
        yseq![ystr!("Detroit Tigers"), ystr!("Chicago Cubs")] => yseq![ystr!("2001-07-23")],
        yseq![ystr!("New York Yankees"), ystr!("Atlanta Braves")] => yseq![ystr!("2001-07-02"), ystr!("2001-08-12"), ystr!("2001-08-14")]
    })
}

#[test]
fn block_literal() {
    match_utf8("block_literal.yml", ystr!("\\//||\\/||\n// ||  ||__\n"))
}

#[test]
fn plain_scalar() {
    match_utf8("plain_scalar.yml", ystr!("Mark McGwire's year was crippled by a knee injury."))
}

#[test]
fn quoted_scalar() {
    match_utf8("quoted_scalar.yml",
    ymap!{
        "unicode" => ystr!("Sosa did fine.\u{263A}"),
        "control" => ystr!("\x081998\t1999\t2000\n"),
        "hexesc" =>  ystr!("\x13\x10 is \r\n"),
        "single" => ystr!(r#""Howdy!" he cried."#),
        "quoted" => ystr!(" # not a 'comment'."),
        "tie-fighter" => ystr!(r"|\-*-/|")
    })
}

#[test]
fn multi_line_scalar() {
    match_utf8("multi_line_scalar.yml",
    ymap!{
        "plain" => ystr!("This unquoted scalar spans many lines."),
        "quoted" => ystr!("So does this quoted scalar.\n")
    })
}

#[test]
fn utf16le() {
    match_file(YamlEncoding::YamlUtf16LeEncoding, "utf16le.yml",
               yseq![ystr!("Hello"), ystr!("世界")]
    )
}

#[test]
fn utf16be() {
    match_file(YamlEncoding::YamlUtf16BeEncoding, "utf16be.yml",
               yseq![ystr!("Hello"), ystr!("世界")]
    )
}
<commit_msg>Add feature gating to test code<commit_after>// Allow unstable items until Rust hits 1.0
#![feature(core)]
#![feature(io)]
#![feature(env)]
#![feature(path)]
#![feature(fs)]

extern crate yaml;

use yaml::constructor::YamlStandardData;
use yaml::ffi::YamlEncoding;
use std::env;
use std::path;
use std::io::BufReader;
use std::fs::File;

fn match_utf8(filename: &str, expected: YamlStandardData) {
    match_file(YamlEncoding::YamlUtf8Encoding, filename, expected);
}

fn match_file(encoding: yaml::ffi::YamlEncoding, filename: &str, expected: YamlStandardData) {
    let mut args = env::args();
    let this_path = args.next().unwrap();
    let dir_path = path::Path::new(this_path.as_slice()).parent().unwrap();
    let file_path = dir_path.join("../tests/source").join(filename);
    println!("{}", file_path.display());
    let file = match File::open(&file_path) {
        Ok(f) => f,
        Err(e) => panic!("Failed to open file {}: {}", file_path.to_str().unwrap(), e)
    };
    let mut reader = BufReader::new(file);
    match yaml::parse_io(&mut reader, encoding) {
        Ok(docs) => if docs.len() == 1 {
            assert_eq!(docs.as_slice().first().unwrap(), &expected)
        } else {
            panic!("too many number of documents: {:?}", docs)
        },
        Err(e) => panic!("parse failure: {:?}", e)
    }
}

macro_rules! ystr{
    ($e: expr) => (
        YamlStandardData::YamlString($e.to_string())
    )
}

macro_rules! yint{
    ($e: expr) => (
        YamlStandardData::YamlInteger($e)
    )
}

macro_rules! yfloat{
    ($e: expr) => (
        YamlStandardData::YamlFloat($e)
    )
}

macro_rules! yseq{
    ($($e:expr),*) => (
        YamlStandardData::YamlSequence(vec![$(($e),)*])
    )
}

macro_rules! ymap{
    ($($k:expr => $v:expr),*) => (
        YamlStandardData::YamlMapping(vec![$((ystr!($k), $v),)*])
    )
}

macro_rules! y_cmp_map{
    ($($k:expr => $v:expr),*) => (
        YamlStandardData::YamlMapping(vec![$(($k, $v),)*])
    )
}

#[test]
fn sequence_of_scalars() {
    match_utf8("ball_players.yml", yseq![ystr!("Mark McGwire"), ystr!("Sammy Sosa"), ystr!("Ken Griffey")]);
}

#[test]
fn scalar_mappings() {
    match_utf8("player_stat.yml", ymap!{
                                    "hr" => yint!(65),
                                    "avg" => yfloat!(0.278),
                                    "rbi" => yint!(147)
                                })
}

#[test]
fn maps_of_sequences() {
    match_utf8("ball_clubs.yml", ymap!{
                                    "american" => yseq![ystr!("Boston Red Sox"), ystr!("Detroit Tigers"), ystr!("New York Yankees")],
                                    "national" => yseq![ystr!("New York Mets"), ystr!("Chicago Cubs"), ystr!("Atlanta Braves")]
                                })
}

#[test]
fn sequence_of_maps() {
    match_utf8("multiple_player_stat.yml",
    yseq![
        ymap!{
            "name" => ystr!("Mark McGwire"),
            "hr" => yint!(65)
        },
        ymap!{
            "name" => ystr!("Sammy Sosa"),
            "hr" => yint!(63)
        }
    ])
}

#[test]
fn sequence_of_sequences() {
    match_utf8("csv.yml",
    yseq![
        yseq![ystr!("name"), ystr!("hr")],
        yseq![ystr!("Mark McGwire"), yint!(65)],
        yseq![ystr!("Sammy Sosa"), yint!(63)]
    ])
}

#[test]
fn mapping_of_mappings() {
    match_utf8("map_map.yml",
    ymap!{
        "Mark McGwire" => ymap!{ "hr" => yint!(65) },
        "Sammy Sosa" => ymap!{ "hr" => yint!(63) }
    })
}

#[test]
fn alias() {
    match_utf8("alias.yml",
    ymap!{
        "hr" => yseq![ystr!("Mark McGwire"), ystr!("Sammy Sosa")],
        "rbi" => yseq![ystr!("Sammy Sosa"), ystr!("Ken Griffey")]
    })
}

#[test]
fn complex_keys() {
    match_utf8("complex_key.yml",
    y_cmp_map!{
        yseq![ystr!("Detroit Tigers"), ystr!("Chicago Cubs")] => yseq![ystr!("2001-07-23")],
        yseq![ystr!("New York Yankees"), ystr!("Atlanta Braves")] => yseq![ystr!("2001-07-02"), ystr!("2001-08-12"), ystr!("2001-08-14")]
    })
}

#[test]
fn block_literal() {
    match_utf8("block_literal.yml", ystr!("\\//||\\/||\n// ||  ||__\n"))
}

#[test]
fn plain_scalar() {
    match_utf8("plain_scalar.yml", ystr!("Mark McGwire's year was crippled by a knee injury."))
}

#[test]
fn quoted_scalar() {
    match_utf8("quoted_scalar.yml",
    ymap!{
        "unicode" => ystr!("Sosa did fine.\u{263A}"),
        "control" => ystr!("\x081998\t1999\t2000\n"),
        "hexesc" =>  ystr!("\x13\x10 is \r\n"),
        "single" => ystr!(r#""Howdy!" he cried."#),
        "quoted" => ystr!(" # not a 'comment'."),
        "tie-fighter" => ystr!(r"|\-*-/|")
    })
}

#[test]
fn multi_line_scalar() {
    match_utf8("multi_line_scalar.yml",
    ymap!{
        "plain" => ystr!("This unquoted scalar spans many lines."),
        "quoted" => ystr!("So does this quoted scalar.\n")
    })
}

#[test]
fn utf16le() {
    match_file(YamlEncoding::YamlUtf16LeEncoding, "utf16le.yml",
               yseq![ystr!("Hello"), ystr!("世界")]
    )
}

#[test]
fn utf16be() {
    match_file(YamlEncoding::YamlUtf16BeEncoding, "utf16be.yml",
               yseq![ystr!("Hello"), ystr!("世界")]
    )
}
<|endoftext|>


<commit_before>use crate::{
    transform::{MirPass, MirSource},
    util::patch::MirPatch,
};
use rustc_middle::mir::*;
use rustc_middle::ty::{Ty, TyCtxt};
use std::{borrow::Cow, fmt::Debug};

use super::simplify::simplify_cfg;

/// This pass optimizes something like
/// ```text
/// let x: Option<()>;
/// let y: Option<()>;
/// match (x,y) {
///     (Some(_), Some(_)) => {0},
///     _ => {1}
/// }
/// ```
/// into something like
/// ```text
/// let x: Option<()>;
/// let y: Option<()>;
/// let discriminant_x = // get discriminant of x
/// let discriminant_y = // get discriminant of y
/// if discriminant_x != discriminant_y || discriminant_x == None {1} else {0}
/// ```
pub struct EarlyOtherwiseBranch;

impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level < 1 {
            return;
        }
        trace!("running EarlyOtherwiseBranch on {:?}", source);
        // we are only interested in this bb if the terminator is a switchInt
        let bbs_with_switch =
            body.basic_blocks().iter_enumerated().filter(|(_, bb)| is_switch(bb.terminator()));

        let opts_to_apply: Vec<OptimizationToApply<'tcx>> = bbs_with_switch
            .flat_map(|(bb_idx, bb)| {
                let switch = bb.terminator();
                let helper = Helper { body, tcx };
                let infos = helper.go(bb, switch)?;
                Some(OptimizationToApply { infos, basic_block_first_switch: bb_idx })
            })
            .collect();

        let should_cleanup = !opts_to_apply.is_empty();

        for opt_to_apply in opts_to_apply {
            trace!("SUCCESS: found optimization possibility to apply: {:?}", &opt_to_apply);

            let statements_before =
                body.basic_blocks()[opt_to_apply.basic_block_first_switch].statements.len();
            let end_of_block_location = Location {
                block: opt_to_apply.basic_block_first_switch,
                statement_index: statements_before,
            };

            let mut patch = MirPatch::new(body);

            // create temp to store second discriminant in
            let discr_type = opt_to_apply.infos[0].second_switch_info.discr_ty;
            let discr_span = opt_to_apply.infos[0].second_switch_info.discr_source_info.span;
            let second_discriminant_temp = patch.new_temp(discr_type, discr_span);

            patch.add_statement(
                end_of_block_location,
                StatementKind::StorageLive(second_discriminant_temp),
            );

            // create assignment of discriminant
            let place_of_adt_to_get_discriminant_of =
                opt_to_apply.infos[0].second_switch_info.place_of_adt_discr_read;
            patch.add_assign(
                end_of_block_location,
                Place::from(second_discriminant_temp),
                Rvalue::Discriminant(place_of_adt_to_get_discriminant_of),
            );

            // create temp to store NotEqual comparison between the two discriminants
            let not_equal = BinOp::Ne;
            let not_equal_res_type = not_equal.ty(tcx, discr_type, discr_type);
            let not_equal_temp = patch.new_temp(not_equal_res_type, discr_span);
            patch.add_statement(end_of_block_location, StatementKind::StorageLive(not_equal_temp));

            // create NotEqual comparison between the two discriminants
            let first_descriminant_place =
                opt_to_apply.infos[0].first_switch_info.discr_used_in_switch;
            let not_equal_rvalue = Rvalue::BinaryOp(
                not_equal,
                Operand::Copy(Place::from(second_discriminant_temp)),
                Operand::Copy(Place::from(first_descriminant_place)),
            );
            patch.add_statement(
                end_of_block_location,
                StatementKind::Assign(box (Place::from(not_equal_temp), not_equal_rvalue)),
            );

            let (mut targets_to_jump_to, values_to_jump_to): (Vec<_>, Vec<_>) = opt_to_apply
                .infos
                .iter()
                .flat_map(|x| x.second_switch_info.targets_with_values.iter())
                .cloned()
                .unzip();

            // add otherwise case in the end
            targets_to_jump_to.push(opt_to_apply.infos[0].first_switch_info.otherwise_bb);
            // new block that jumps to the correct discriminant case. This block is switched to if the discriminants are equal
            let new_switch_data = BasicBlockData::new(Some(Terminator {
                source_info: opt_to_apply.infos[0].second_switch_info.discr_source_info,
                kind: TerminatorKind::SwitchInt {
                    // the first and second discriminants are equal, so just pick one
                    discr: Operand::Copy(first_descriminant_place),
                    switch_ty: discr_type,
                    values: Cow::from(values_to_jump_to),
                    targets: targets_to_jump_to,
                },
            }));

            let new_switch_bb = patch.new_block(new_switch_data);

            // switch on the NotEqual. If true, then jump to the `otherwise` case.
            // If false, then jump to a basic block that then jumps to the correct disciminant case
            let true_case = opt_to_apply.infos[0].first_switch_info.otherwise_bb;
            let false_case = new_switch_bb;
            patch.patch_terminator(
                opt_to_apply.basic_block_first_switch,
                TerminatorKind::if_(
                    tcx,
                    Operand::Move(Place::from(not_equal_temp)),
                    true_case,
                    false_case,
                ),
            );

            // generate StorageDead for the second_discriminant_temp not in use anymore
            patch.add_statement(
                end_of_block_location,
                StatementKind::StorageDead(second_discriminant_temp),
            );

            // Generate a StorageDead for not_equal_temp in each of the targets, since we moved it into the switch
            for bb in [false_case, true_case].iter() {
                patch.add_statement(
                    Location { block: *bb, statement_index: 0 },
                    StatementKind::StorageDead(not_equal_temp),
                );
            }

            patch.apply(body);
        }

        // Since this optimization adds new basic blocks and invalidates others,
        // clean up the cfg to make it nicer for other passes
        if should_cleanup {
            simplify_cfg(body);
        }
    }
}

fn is_switch<'tcx>(terminator: &Terminator<'tcx>) -> bool {
    match terminator.kind {
        TerminatorKind::SwitchInt { .. } => true,
        _ => false,
    }
}

struct Helper<'a, 'tcx> {
    body: &'a Body<'tcx>,
    tcx: TyCtxt<'tcx>,
}

#[derive(Debug, Clone)]
struct SwitchDiscriminantInfo<'tcx> {
    /// Type of the discriminant being switched on
    discr_ty: Ty<'tcx>,
    /// The basic block that the otherwise branch points to
    otherwise_bb: BasicBlock,
    /// Target along with the value being branched from. Otherwise is not included
    targets_with_values: Vec<(BasicBlock, u128)>,
    discr_source_info: SourceInfo,
    /// The place of the discriminant used in the switch
    discr_used_in_switch: Place<'tcx>,
    /// The place of the adt that has its discriminant read
    place_of_adt_discr_read: Place<'tcx>,
    /// The type of the adt that has its discriminant read
    type_adt_matched_on: Ty<'tcx>,
}

#[derive(Debug)]
struct OptimizationToApply<'tcx> {
    infos: Vec<OptimizationInfo<'tcx>>,
    /// Basic block of the original first switch
    basic_block_first_switch: BasicBlock,
}

#[derive(Debug)]
struct OptimizationInfo<'tcx> {
    /// Info about the first switch and discriminant
    first_switch_info: SwitchDiscriminantInfo<'tcx>,
    /// Info about the second switch and discriminant
    second_switch_info: SwitchDiscriminantInfo<'tcx>,
}

impl<'a, 'tcx> Helper<'a, 'tcx> {
    pub fn go(
        &self,
        bb: &BasicBlockData<'tcx>,
        switch: &Terminator<'tcx>,
    ) -> Option<Vec<OptimizationInfo<'tcx>>> {
        // try to find the statement that defines the discriminant that is used for the switch
        let discr = self.find_switch_discriminant_info(bb, switch)?;

        // go through each target, finding a discriminant read, and a switch
        let results = discr.targets_with_values.iter().map(|(target, value)| {
            self.find_discriminant_switch_pairing(&discr, target.clone(), value.clone())
        });

        // if the optimization did not apply for one of the targets, then abort
        if results.clone().any(|x| x.is_none()) || results.len() == 0 {
            trace!("NO: not all of the targets matched the pattern for optimization");
            return None;
        }

        Some(results.flatten().collect())
    }

    fn find_discriminant_switch_pairing(
        &self,
        discr_info: &SwitchDiscriminantInfo<'tcx>,
        target: BasicBlock,
        value: u128,
    ) -> Option<OptimizationInfo<'tcx>> {
        let bb = &self.body.basic_blocks()[target];
        // find switch
        let terminator = bb.terminator();
        if is_switch(terminator) {
            let this_bb_discr_info = self.find_switch_discriminant_info(bb, terminator)?;

            // the types of the two adts matched on have to be equalfor this optimization to apply
            if discr_info.type_adt_matched_on != this_bb_discr_info.type_adt_matched_on {
                trace!(
                    "NO: types do not match. LHS: {:?}, RHS: {:?}",
                    discr_info.type_adt_matched_on,
                    this_bb_discr_info.type_adt_matched_on
                );
                return None;
            }

            // the otherwise branch of the two switches have to point to the same bb
            if discr_info.otherwise_bb != this_bb_discr_info.otherwise_bb {
                trace!("NO: otherwise target is not the same");
                return None;
            }

            // check that the value being matched on is the same. The
            if this_bb_discr_info.targets_with_values.iter().find(|x| x.1 == value).is_none() {
                trace!("NO: values being matched on are not the same");
                return None;
            }

            // only allow optimization if the left and right of the tuple being matched are the same variants.
            // so the following should not optimize
            //  ```rust
            // let x: Option<()>;
            // let y: Option<()>;
            // match (x,y) {
            //     (Some(_), None) => {},
            //     _ => {}
            // }
            //  ```
            // We check this by seeing that the value of the first discriminant is the only other discriminant value being used as a target in the second switch
            if !(this_bb_discr_info.targets_with_values.len() == 1
                && this_bb_discr_info.targets_with_values[0].1 == value)
            {
                trace!(
                    "NO: The second switch did not have only 1 target (besides otherwise) that had the same value as the value from the first switch that got us here"
                );
                return None;
            }

            // if we reach this point, the optimization applies, and we should be able to optimize this case
            // store the info that is needed to apply the optimization

            Some(OptimizationInfo {
                first_switch_info: discr_info.clone(),
                second_switch_info: this_bb_discr_info,
            })
        } else {
            None
        }
    }

    fn find_switch_discriminant_info(
        &self,
        bb: &BasicBlockData<'tcx>,
        switch: &Terminator<'tcx>,
    ) -> Option<SwitchDiscriminantInfo<'tcx>> {
        match &switch.kind {
            TerminatorKind::SwitchInt { discr, targets, values, .. } => {
                let discr_local = discr.place()?.as_local()?;
                // the declaration of the discriminant read. Place of this read is being used in the switch
                let discr_decl = &self.body.local_decls()[discr_local];
                let discr_ty = discr_decl.ty;
                // the otherwise target lies as the last element
                let otherwise_bb = targets.get(values.len())?.clone();
                let targets_with_values = targets
                    .iter()
                    .zip(values.iter())
                    .map(|(t, v)| (t.clone(), v.clone()))
                    .collect();

                // find the place of the adt where the discriminant is being read from
                // assume this is the last statement of the block
                let place_of_adt_discr_read = match bb.statements.last()?.kind {
                    StatementKind::Assign(box (_, Rvalue::Discriminant(adt_place))) => {
                        Some(adt_place)
                    }
                    _ => None,
                }?;

                let type_adt_matched_on = place_of_adt_discr_read.ty(self.body, self.tcx).ty;

                Some(SwitchDiscriminantInfo {
                    discr_used_in_switch: discr.place()?,
                    discr_ty,
                    otherwise_bb,
                    targets_with_values,
                    discr_source_info: discr_decl.source_info,
                    place_of_adt_discr_read,
                    type_adt_matched_on,
                })
            }
            _ => unreachable!("must only be passed terminator that is a switch"),
        }
    }
}
<commit_msg>EarlyOtherwiseBranch::run_pass(): don't convert Place to Place (clippy::useless_conversion)<commit_after>use crate::{
    transform::{MirPass, MirSource},
    util::patch::MirPatch,
};
use rustc_middle::mir::*;
use rustc_middle::ty::{Ty, TyCtxt};
use std::{borrow::Cow, fmt::Debug};

use super::simplify::simplify_cfg;

/// This pass optimizes something like
/// ```text
/// let x: Option<()>;
/// let y: Option<()>;
/// match (x,y) {
///     (Some(_), Some(_)) => {0},
///     _ => {1}
/// }
/// ```
/// into something like
/// ```text
/// let x: Option<()>;
/// let y: Option<()>;
/// let discriminant_x = // get discriminant of x
/// let discriminant_y = // get discriminant of y
/// if discriminant_x != discriminant_y || discriminant_x == None {1} else {0}
/// ```
pub struct EarlyOtherwiseBranch;

impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level < 1 {
            return;
        }
        trace!("running EarlyOtherwiseBranch on {:?}", source);
        // we are only interested in this bb if the terminator is a switchInt
        let bbs_with_switch =
            body.basic_blocks().iter_enumerated().filter(|(_, bb)| is_switch(bb.terminator()));

        let opts_to_apply: Vec<OptimizationToApply<'tcx>> = bbs_with_switch
            .flat_map(|(bb_idx, bb)| {
                let switch = bb.terminator();
                let helper = Helper { body, tcx };
                let infos = helper.go(bb, switch)?;
                Some(OptimizationToApply { infos, basic_block_first_switch: bb_idx })
            })
            .collect();

        let should_cleanup = !opts_to_apply.is_empty();

        for opt_to_apply in opts_to_apply {
            trace!("SUCCESS: found optimization possibility to apply: {:?}", &opt_to_apply);

            let statements_before =
                body.basic_blocks()[opt_to_apply.basic_block_first_switch].statements.len();
            let end_of_block_location = Location {
                block: opt_to_apply.basic_block_first_switch,
                statement_index: statements_before,
            };

            let mut patch = MirPatch::new(body);

            // create temp to store second discriminant in
            let discr_type = opt_to_apply.infos[0].second_switch_info.discr_ty;
            let discr_span = opt_to_apply.infos[0].second_switch_info.discr_source_info.span;
            let second_discriminant_temp = patch.new_temp(discr_type, discr_span);

            patch.add_statement(
                end_of_block_location,
                StatementKind::StorageLive(second_discriminant_temp),
            );

            // create assignment of discriminant
            let place_of_adt_to_get_discriminant_of =
                opt_to_apply.infos[0].second_switch_info.place_of_adt_discr_read;
            patch.add_assign(
                end_of_block_location,
                Place::from(second_discriminant_temp),
                Rvalue::Discriminant(place_of_adt_to_get_discriminant_of),
            );

            // create temp to store NotEqual comparison between the two discriminants
            let not_equal = BinOp::Ne;
            let not_equal_res_type = not_equal.ty(tcx, discr_type, discr_type);
            let not_equal_temp = patch.new_temp(not_equal_res_type, discr_span);
            patch.add_statement(end_of_block_location, StatementKind::StorageLive(not_equal_temp));

            // create NotEqual comparison between the two discriminants
            let first_descriminant_place =
                opt_to_apply.infos[0].first_switch_info.discr_used_in_switch;
            let not_equal_rvalue = Rvalue::BinaryOp(
                not_equal,
                Operand::Copy(Place::from(second_discriminant_temp)),
                Operand::Copy(first_descriminant_place),
            );
            patch.add_statement(
                end_of_block_location,
                StatementKind::Assign(box (Place::from(not_equal_temp), not_equal_rvalue)),
            );

            let (mut targets_to_jump_to, values_to_jump_to): (Vec<_>, Vec<_>) = opt_to_apply
                .infos
                .iter()
                .flat_map(|x| x.second_switch_info.targets_with_values.iter())
                .cloned()
                .unzip();

            // add otherwise case in the end
            targets_to_jump_to.push(opt_to_apply.infos[0].first_switch_info.otherwise_bb);
            // new block that jumps to the correct discriminant case. This block is switched to if the discriminants are equal
            let new_switch_data = BasicBlockData::new(Some(Terminator {
                source_info: opt_to_apply.infos[0].second_switch_info.discr_source_info,
                kind: TerminatorKind::SwitchInt {
                    // the first and second discriminants are equal, so just pick one
                    discr: Operand::Copy(first_descriminant_place),
                    switch_ty: discr_type,
                    values: Cow::from(values_to_jump_to),
                    targets: targets_to_jump_to,
                },
            }));

            let new_switch_bb = patch.new_block(new_switch_data);

            // switch on the NotEqual. If true, then jump to the `otherwise` case.
            // If false, then jump to a basic block that then jumps to the correct disciminant case
            let true_case = opt_to_apply.infos[0].first_switch_info.otherwise_bb;
            let false_case = new_switch_bb;
            patch.patch_terminator(
                opt_to_apply.basic_block_first_switch,
                TerminatorKind::if_(
                    tcx,
                    Operand::Move(Place::from(not_equal_temp)),
                    true_case,
                    false_case,
                ),
            );

            // generate StorageDead for the second_discriminant_temp not in use anymore
            patch.add_statement(
                end_of_block_location,
                StatementKind::StorageDead(second_discriminant_temp),
            );

            // Generate a StorageDead for not_equal_temp in each of the targets, since we moved it into the switch
            for bb in [false_case, true_case].iter() {
                patch.add_statement(
                    Location { block: *bb, statement_index: 0 },
                    StatementKind::StorageDead(not_equal_temp),
                );
            }

            patch.apply(body);
        }

        // Since this optimization adds new basic blocks and invalidates others,
        // clean up the cfg to make it nicer for other passes
        if should_cleanup {
            simplify_cfg(body);
        }
    }
}

fn is_switch<'tcx>(terminator: &Terminator<'tcx>) -> bool {
    match terminator.kind {
        TerminatorKind::SwitchInt { .. } => true,
        _ => false,
    }
}

struct Helper<'a, 'tcx> {
    body: &'a Body<'tcx>,
    tcx: TyCtxt<'tcx>,
}

#[derive(Debug, Clone)]
struct SwitchDiscriminantInfo<'tcx> {
    /// Type of the discriminant being switched on
    discr_ty: Ty<'tcx>,
    /// The basic block that the otherwise branch points to
    otherwise_bb: BasicBlock,
    /// Target along with the value being branched from. Otherwise is not included
    targets_with_values: Vec<(BasicBlock, u128)>,
    discr_source_info: SourceInfo,
    /// The place of the discriminant used in the switch
    discr_used_in_switch: Place<'tcx>,
    /// The place of the adt that has its discriminant read
    place_of_adt_discr_read: Place<'tcx>,
    /// The type of the adt that has its discriminant read
    type_adt_matched_on: Ty<'tcx>,
}

#[derive(Debug)]
struct OptimizationToApply<'tcx> {
    infos: Vec<OptimizationInfo<'tcx>>,
    /// Basic block of the original first switch
    basic_block_first_switch: BasicBlock,
}

#[derive(Debug)]
struct OptimizationInfo<'tcx> {
    /// Info about the first switch and discriminant
    first_switch_info: SwitchDiscriminantInfo<'tcx>,
    /// Info about the second switch and discriminant
    second_switch_info: SwitchDiscriminantInfo<'tcx>,
}

impl<'a, 'tcx> Helper<'a, 'tcx> {
    pub fn go(
        &self,
        bb: &BasicBlockData<'tcx>,
        switch: &Terminator<'tcx>,
    ) -> Option<Vec<OptimizationInfo<'tcx>>> {
        // try to find the statement that defines the discriminant that is used for the switch
        let discr = self.find_switch_discriminant_info(bb, switch)?;

        // go through each target, finding a discriminant read, and a switch
        let results = discr.targets_with_values.iter().map(|(target, value)| {
            self.find_discriminant_switch_pairing(&discr, target.clone(), value.clone())
        });

        // if the optimization did not apply for one of the targets, then abort
        if results.clone().any(|x| x.is_none()) || results.len() == 0 {
            trace!("NO: not all of the targets matched the pattern for optimization");
            return None;
        }

        Some(results.flatten().collect())
    }

    fn find_discriminant_switch_pairing(
        &self,
        discr_info: &SwitchDiscriminantInfo<'tcx>,
        target: BasicBlock,
        value: u128,
    ) -> Option<OptimizationInfo<'tcx>> {
        let bb = &self.body.basic_blocks()[target];
        // find switch
        let terminator = bb.terminator();
        if is_switch(terminator) {
            let this_bb_discr_info = self.find_switch_discriminant_info(bb, terminator)?;

            // the types of the two adts matched on have to be equalfor this optimization to apply
            if discr_info.type_adt_matched_on != this_bb_discr_info.type_adt_matched_on {
                trace!(
                    "NO: types do not match. LHS: {:?}, RHS: {:?}",
                    discr_info.type_adt_matched_on,
                    this_bb_discr_info.type_adt_matched_on
                );
                return None;
            }

            // the otherwise branch of the two switches have to point to the same bb
            if discr_info.otherwise_bb != this_bb_discr_info.otherwise_bb {
                trace!("NO: otherwise target is not the same");
                return None;
            }

            // check that the value being matched on is the same. The
            if this_bb_discr_info.targets_with_values.iter().find(|x| x.1 == value).is_none() {
                trace!("NO: values being matched on are not the same");
                return None;
            }

            // only allow optimization if the left and right of the tuple being matched are the same variants.
            // so the following should not optimize
            //  ```rust
            // let x: Option<()>;
            // let y: Option<()>;
            // match (x,y) {
            //     (Some(_), None) => {},
            //     _ => {}
            // }
            //  ```
            // We check this by seeing that the value of the first discriminant is the only other discriminant value being used as a target in the second switch
            if !(this_bb_discr_info.targets_with_values.len() == 1
                && this_bb_discr_info.targets_with_values[0].1 == value)
            {
                trace!(
                    "NO: The second switch did not have only 1 target (besides otherwise) that had the same value as the value from the first switch that got us here"
                );
                return None;
            }

            // if we reach this point, the optimization applies, and we should be able to optimize this case
            // store the info that is needed to apply the optimization

            Some(OptimizationInfo {
                first_switch_info: discr_info.clone(),
                second_switch_info: this_bb_discr_info,
            })
        } else {
            None
        }
    }

    fn find_switch_discriminant_info(
        &self,
        bb: &BasicBlockData<'tcx>,
        switch: &Terminator<'tcx>,
    ) -> Option<SwitchDiscriminantInfo<'tcx>> {
        match &switch.kind {
            TerminatorKind::SwitchInt { discr, targets, values, .. } => {
                let discr_local = discr.place()?.as_local()?;
                // the declaration of the discriminant read. Place of this read is being used in the switch
                let discr_decl = &self.body.local_decls()[discr_local];
                let discr_ty = discr_decl.ty;
                // the otherwise target lies as the last element
                let otherwise_bb = targets.get(values.len())?.clone();
                let targets_with_values = targets
                    .iter()
                    .zip(values.iter())
                    .map(|(t, v)| (t.clone(), v.clone()))
                    .collect();

                // find the place of the adt where the discriminant is being read from
                // assume this is the last statement of the block
                let place_of_adt_discr_read = match bb.statements.last()?.kind {
                    StatementKind::Assign(box (_, Rvalue::Discriminant(adt_place))) => {
                        Some(adt_place)
                    }
                    _ => None,
                }?;

                let type_adt_matched_on = place_of_adt_discr_read.ty(self.body, self.tcx).ty;

                Some(SwitchDiscriminantInfo {
                    discr_used_in_switch: discr.place()?,
                    discr_ty,
                    otherwise_bb,
                    targets_with_values,
                    discr_source_info: discr_decl.source_info,
                    place_of_adt_discr_read,
                    type_adt_matched_on,
                })
            }
            _ => unreachable!("must only be passed terminator that is a switch"),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Location in memory that contains data.
//!
//! All buffers are guaranteed to be accessible from the GPU.
//!
//! # High-level wrappers
//!
//! The low level implementation of a buffer is `UnsafeBuffer`. However, the vulkano library
//! provides high-level wrappers around that type that are specialized depending on the way you
//! are going to use it:
//!
//! - `CpuAccessBuffer` designates a buffer located in RAM and whose content can be directly
//!   written by your application.
//! - `DeviceLocalBuffer` designates a buffer located in video memory and whose content can't be
//!   written by your application. Accessing this buffer from the GPU is usually faster than the
//!   `CpuAccessBuffer`.
//! - `ImmutableBuffer` designates a buffer in video memory and whose content can only be
//!   written once. Compared to `DeviceLocalBuffer`, this buffer requires less processing on the
//!   CPU because we don't need to keep track of the reads and writes.
//!
//! If you have data that is modified at every single frame, you are encouraged to use a
//! `CpuAccessibleBuffer`. If you have data that is very rarely modified, you are encouraged to
//! use an `ImmutableBuffer` or a `DeviceLocalBuffer` instead.
//!
//! If you just want to get started, you can use the `CpuAccessibleBuffer` everywhere, as it is
//! the most flexible type of buffer.
//!
//! # Buffers usage
//!
//! When you create a buffer object, you have to specify its *usage*. In other words, you have to
//! specify the way it is going to be used. Trying to use a buffer in a way that wasn't specified
//! when you created it will result in an error.
//!
//! You can use buffers for the following purposes:
//!
//! - Can contain arbitrary data that can be transferred from/to other buffers and images.
//! - Can be read and modified from a shader.
//! - Can be used as a source of vertices and indices.
//! - Can be used as a source of list of models for draw indirect commands.
//!
//! Accessing a buffer from a shader can be done in the following ways:
//!
//! - As a uniform buffer. Uniform buffers are read-only.
//! - As a storage buffer. Storage buffers can be read and written.
//! - As a uniform texel buffer. Contrary to a uniform buffer, the data is interpreted by the
//!   GPU and can be for example normalized.
//! - As a storage texel buffer. Additionnally, some data formats can be modified with atomic
//!   operations.
//!
//! Using uniform/storage texel buffers requires creating a *buffer view*. See the `view` module
//! for how to create a buffer view.
//!

pub use self::cpu_access::CpuAccessibleBuffer;
pub use self::device_local::DeviceLocalBuffer;
pub use self::immutable::ImmutableBuffer;
pub use self::slice::BufferSlice;
pub use self::sys::BufferCreationError;
pub use self::sys::Usage as BufferUsage;
pub use self::traits::Buffer;
pub use self::traits::BufferInner;
pub use self::traits::TypedBuffer;
pub use self::traits::TrackedBuffer;
pub use self::traits::TrackedBufferPipelineBarrierRequest;
pub use self::traits::TrackedBufferPipelineMemoryBarrierRequest;
pub use self::traits::TrackedBufferSubmitInfos;
pub use self::view::BufferView;

pub mod cpu_access;
pub mod device_local;
pub mod immutable;
mod slice;
pub mod sys;
mod traits;
pub mod view;
<commit_msg>Minor documentation change<commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Location in memory that contains data.
//!
//! A Vulkan buffer is very similar to a buffer that you would use in programming languages in
//! general, in the sense that it is a location in memory that contains data. The difference
//! between a Vulkan buffer and a regular buffer is that the content of a Vulkan buffer is
//! accessible from the GPU.
//!
//! # High-level wrappers
//!
//! The low level implementation of a buffer is `UnsafeBuffer`. However, the vulkano library
//! provides high-level wrappers around that type that are specialized depending on the way you
//! are going to use it:
//!
//! - `CpuAccessBuffer` designates a buffer located in RAM and whose content can be directly
//!   written by your application.
//! - `DeviceLocalBuffer` designates a buffer located in video memory and whose content can't be
//!   written by your application. Accessing this buffer from the GPU is usually faster than the
//!   `CpuAccessBuffer`.
//! - `ImmutableBuffer` designates a buffer in video memory and whose content can only be
//!   written once. Compared to `DeviceLocalBuffer`, this buffer requires less processing on the
//!   CPU because we don't need to keep track of the reads and writes.
//!
//! If you have data that is modified at every single frame, you are encouraged to use a
//! `CpuAccessibleBuffer`. If you have data that is very rarely modified, you are encouraged to
//! use an `ImmutableBuffer` or a `DeviceLocalBuffer` instead.
//!
//! If you just want to get started, you can use the `CpuAccessibleBuffer` everywhere, as it is
//! the most flexible type of buffer.
//!
//! # Buffers usage
//!
//! When you create a buffer object, you have to specify its *usage*. In other words, you have to
//! specify the way it is going to be used. Trying to use a buffer in a way that wasn't specified
//! when you created it will result in an error.
//!
//! You can use buffers for the following purposes:
//!
//! - Can contain arbitrary data that can be transferred from/to other buffers and images.
//! - Can be read and modified from a shader.
//! - Can be used as a source of vertices and indices.
//! - Can be used as a source of list of models for draw indirect commands.
//!
//! Accessing a buffer from a shader can be done in the following ways:
//!
//! - As a uniform buffer. Uniform buffers are read-only.
//! - As a storage buffer. Storage buffers can be read and written.
//! - As a uniform texel buffer. Contrary to a uniform buffer, the data is interpreted by the
//!   GPU and can be for example normalized.
//! - As a storage texel buffer. Additionnally, some data formats can be modified with atomic
//!   operations.
//!
//! Using uniform/storage texel buffers requires creating a *buffer view*. See the `view` module
//! for how to create a buffer view.
//!

pub use self::cpu_access::CpuAccessibleBuffer;
pub use self::device_local::DeviceLocalBuffer;
pub use self::immutable::ImmutableBuffer;
pub use self::slice::BufferSlice;
pub use self::sys::BufferCreationError;
pub use self::sys::Usage as BufferUsage;
pub use self::traits::Buffer;
pub use self::traits::BufferInner;
pub use self::traits::TypedBuffer;
pub use self::traits::TrackedBuffer;
pub use self::traits::TrackedBufferPipelineBarrierRequest;
pub use self::traits::TrackedBufferPipelineMemoryBarrierRequest;
pub use self::traits::TrackedBufferSubmitInfos;
pub use self::view::BufferView;

pub mod cpu_access;
pub mod device_local;
pub mod immutable;
mod slice;
pub mod sys;
mod traits;
pub mod view;
<|endoftext|>


<commit_before>
// extern crate mmap;
extern crate time;
extern crate timely;
extern crate dataflow_join;

use std::rc::Rc;

use dataflow_join::*;
use dataflow_join::graph::{GraphTrait, GraphMMap, GraphExtenderExt};

use timely::dataflow::operators::*;
use timely::dataflow::operators::aggregation::Aggregate;

fn main () {

    let filename = std::env::args().nth(1).unwrap();
    let inspect = std::env::args().nth(2).unwrap() == "inspect";
    let step_size = std::env::args().nth(3).unwrap().parse::<usize>().unwrap();

    timely::execute_from_args(std::env::args().skip(4), move |root| {

        let index = root.index();
        let peers = root.peers();

        let graph = Rc::new(GraphMMap::<u32>::new(&filename));

        let (mut input, probe) = root.dataflow(|builder| {

            let (input, cliques) = builder.new_input::<u32>();

            // pairs
            let cliques = cliques.extend(vec![&graph.extend_using(|&a| a as u64)]);

            // triangles
            let cliques = cliques.flat_map(|(p, es)| es.into_iter().map(move |e| (p, e)))
                                 .extend(vec![&graph.extend_using(|&(a,_)| a as u64)]);

            // quadrangles?
            let cliques = cliques.flat_map(|(p,es)| es.into_iter().map(move |e| (p, e)))
                                 .filter(|&((_,b),c)| b < c)
                                 .extend(vec![&graph.extend_using(|&((_,b),_)| b as u64),
                                              &graph.extend_using(|&((_,_),c)| c as u64)]);

            // // 5 cliques?
            // let cliques = cliques.flat_map(|(p,es)| es.into_iter().map(move |e| (p, e)))
            //                      .extend(vec![&graph.extend_using(|&(((a,_),_),_)| a as u64),
            //                                   &graph.extend_using(|&(((_,b),_),_)| b as u64),
            //                                   &graph.extend_using(|&(((_,_),c),_)| c as u64),
            //                                   &graph.extend_using(|&(((_,_),_),d)| d as u64)]);

            let mut count = 0;
            if inspect {
                cliques
                    .map(|x| ((), x.1.len()))
                    .aggregate(|_k, v: usize, a| *a += v, |_k, a: usize| a, |_k| 0)
                    .inspect_batch(move |_t, b| {
                        for x in b { count += *x; }
                        println!("count: {}", count);
                    });
            }

            (input, cliques.probe())
        });

        let nodes = graph.nodes() - 1;
        let limit = (nodes / step_size) + 1;
        for round in 0..limit {
            for source in 0..step_size {
                let candidate = source + round * step_size;
                if candidate % peers == index && candidate < nodes {
                    input.send(candidate as u32);
                }
            }

            input.advance_to(round as u64 + 1);
            root.step();
            while probe.less_than(input.time()) {
                root.step();
            }
        }
    }).unwrap();
}
<commit_msg>core pinning<commit_after>extern crate core_affinity;
extern crate time;
extern crate timely;
extern crate dataflow_join;

use std::rc::Rc;

use dataflow_join::*;
use dataflow_join::graph::{GraphTrait, GraphMMap, GraphExtenderExt};

use timely::dataflow::operators::*;
use timely::dataflow::operators::aggregation::Aggregate;

fn main () {

    let filename = std::env::args().nth(1).unwrap();
    let inspect = std::env::args().nth(2).unwrap() == "inspect";
    let step_size = std::env::args().nth(3).unwrap().parse::<usize>().unwrap();

    timely::execute_from_args(std::env::args().skip(4), move |root| {

        let index = root.index();
        let peers = root.peers();

        let core_ids = core_affinity::get_core_ids().unwrap();
        core_affinity::set_for_current(core_ids[index % core_ids.len()]);

        let graph = Rc::new(GraphMMap::<u32>::new(&filename));

        let (mut input, probe) = root.dataflow(|builder| {

            let (input, cliques) = builder.new_input::<u32>();

            // pairs
            let cliques = cliques.extend(vec![&graph.extend_using(|&a| a as u64)]);

            // triangles
            let cliques = cliques.flat_map(|(p, es)| es.into_iter().map(move |e| (p, e)))
                                 .extend(vec![&graph.extend_using(|&(a,_)| a as u64)]);

            // quadrangles?
            let cliques = cliques.flat_map(|(p,es)| es.into_iter().map(move |e| (p, e)))
                                 .filter(|&((_,b),c)| b < c)
                                 .extend(vec![&graph.extend_using(|&((_,b),_)| b as u64),
                                              &graph.extend_using(|&((_,_),c)| c as u64)]);

            // // 5 cliques?
            // let cliques = cliques.flat_map(|(p,es)| es.into_iter().map(move |e| (p, e)))
            //                      .extend(vec![&graph.extend_using(|&(((a,_),_),_)| a as u64),
            //                                   &graph.extend_using(|&(((_,b),_),_)| b as u64),
            //                                   &graph.extend_using(|&(((_,_),c),_)| c as u64),
            //                                   &graph.extend_using(|&(((_,_),_),d)| d as u64)]);

            let mut count = 0;
            if inspect {
                cliques
                    .map(|x| ((), x.1.len()))
                    .aggregate(|_k, v: usize, a| *a += v, |_k, a: usize| a, |_k| 0)
                    .inspect_batch(move |_t, b| {
                        for x in b { count += *x; }
                        println!("count: {}", count);
                    });
            }

            (input, cliques.probe())
        });

        let nodes = graph.nodes() - 1;
        let limit = (nodes / step_size) + 1;
        for round in 0..limit {
            for source in 0..step_size {
                let candidate = source + round * step_size;
                if candidate % peers == index && candidate < nodes {
                    input.send(candidate as u32);
                }
            }

            input.advance_to(round as u64 + 1);
            root.step();
            while probe.less_than(input.time()) {
                root.step();
            }
        }
    }).unwrap();
}
<|endoftext|>


<commit_before>//! Loads a Cargo project into a static instance of analysis, without support
//! for incorporating changes.
use std::{path::Path, sync::Arc};

use anyhow::Result;
use crossbeam_channel::{unbounded, Receiver};
use ra_db::{AbsPathBuf, CrateGraph};
use ra_ide::{AnalysisChange, AnalysisHost};
use ra_project_model::{CargoConfig, ProcMacroClient, ProjectManifest, ProjectWorkspace};
use vfs::{loader::Handle, AbsPath};

use crate::reload::{ProjectFolders, SourceRootConfig};

pub fn load_cargo(
    root: &Path,
    load_out_dirs_from_check: bool,
    with_proc_macro: bool,
) -> Result<(AnalysisHost, vfs::Vfs)> {
    let root = AbsPathBuf::assert(std::env::current_dir()?.join(root));
    let root = ProjectManifest::discover_single(&root)?;
    let ws = ProjectWorkspace::load(
        root,
        &CargoConfig { load_out_dirs_from_check, ..Default::default() },
        true,
    )?;

    let (sender, receiver) = unbounded();
    let mut vfs = vfs::Vfs::default();
    let mut loader = {
        let loader =
            vfs_notify::NotifyHandle::spawn(Box::new(move |msg| sender.send(msg).unwrap()));
        Box::new(loader)
    };

    let proc_macro_client = if with_proc_macro {
        let path = std::env::current_exe()?;
        ProcMacroClient::extern_process(path, &["proc-macro"]).unwrap()
    } else {
        ProcMacroClient::dummy()
    };

    let crate_graph = ws.to_crate_graph(None, &proc_macro_client, &mut |path: &AbsPath| {
        let contents = loader.load_sync(path);
        let path = vfs::VfsPath::from(path.to_path_buf());
        vfs.set_file_contents(path.clone(), contents);
        vfs.file_id(&path)
    });

    let project_folders = ProjectFolders::new(&[ws]);
    loader.set_config(vfs::loader::Config { load: project_folders.load, watch: vec![] });

    log::debug!("crate graph: {:?}", crate_graph);
    let host = load(crate_graph, project_folders.source_root_config, &mut vfs, &receiver);
    Ok((host, vfs))
}

pub(crate) fn load(
    crate_graph: CrateGraph,
    source_root_config: SourceRootConfig,
    vfs: &mut vfs::Vfs,
    receiver: &Receiver<vfs::loader::Message>,
) -> AnalysisHost {
    let lru_cap = std::env::var("RA_LRU_CAP").ok().and_then(|it| it.parse::<usize>().ok());
    let mut host = AnalysisHost::new(lru_cap);
    let mut analysis_change = AnalysisChange::new();

    // wait until Vfs has loaded all roots
    for task in receiver {
        match task {
            vfs::loader::Message::Progress { n_done, n_total } => {
                if n_done == n_total {
                    break;
                }
            }
            vfs::loader::Message::Loaded { files } => {
                for (path, contents) in files {
                    vfs.set_file_contents(path.into(), contents)
                }
            }
        }
    }
    let changes = vfs.take_changes();
    for file in changes {
        if file.exists() {
            let contents = vfs.file_contents(file.file_id).to_vec();
            if let Ok(text) = String::from_utf8(contents) {
                analysis_change.change_file(file.file_id, Some(Arc::new(text)))
            }
        }
    }
    let source_roots = source_root_config.partition(&vfs);
    analysis_change.set_roots(source_roots);

    analysis_change.set_crate_graph(crate_graph);

    host.apply_change(analysis_change);
    host
}

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

    use hir::Crate;

    #[test]
    fn test_loading_rust_analyzer() {
        let path = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap().parent().unwrap();
        let (host, _vfs) = load_cargo(path, false, false).unwrap();
        let n_crates = Crate::all(host.raw_database()).len();
        // RA has quite a few crates, but the exact count doesn't matter
        assert!(n_crates > 20);
    }
}
<commit_msg>reduce visibility<commit_after>//! Loads a Cargo project into a static instance of analysis, without support
//! for incorporating changes.
use std::{path::Path, sync::Arc};

use anyhow::Result;
use crossbeam_channel::{unbounded, Receiver};
use ra_db::{AbsPathBuf, CrateGraph};
use ra_ide::{AnalysisChange, AnalysisHost};
use ra_project_model::{CargoConfig, ProcMacroClient, ProjectManifest, ProjectWorkspace};
use vfs::{loader::Handle, AbsPath};

use crate::reload::{ProjectFolders, SourceRootConfig};

pub fn load_cargo(
    root: &Path,
    load_out_dirs_from_check: bool,
    with_proc_macro: bool,
) -> Result<(AnalysisHost, vfs::Vfs)> {
    let root = AbsPathBuf::assert(std::env::current_dir()?.join(root));
    let root = ProjectManifest::discover_single(&root)?;
    let ws = ProjectWorkspace::load(
        root,
        &CargoConfig { load_out_dirs_from_check, ..Default::default() },
        true,
    )?;

    let (sender, receiver) = unbounded();
    let mut vfs = vfs::Vfs::default();
    let mut loader = {
        let loader =
            vfs_notify::NotifyHandle::spawn(Box::new(move |msg| sender.send(msg).unwrap()));
        Box::new(loader)
    };

    let proc_macro_client = if with_proc_macro {
        let path = std::env::current_exe()?;
        ProcMacroClient::extern_process(path, &["proc-macro"]).unwrap()
    } else {
        ProcMacroClient::dummy()
    };

    let crate_graph = ws.to_crate_graph(None, &proc_macro_client, &mut |path: &AbsPath| {
        let contents = loader.load_sync(path);
        let path = vfs::VfsPath::from(path.to_path_buf());
        vfs.set_file_contents(path.clone(), contents);
        vfs.file_id(&path)
    });

    let project_folders = ProjectFolders::new(&[ws]);
    loader.set_config(vfs::loader::Config { load: project_folders.load, watch: vec![] });

    log::debug!("crate graph: {:?}", crate_graph);
    let host = load(crate_graph, project_folders.source_root_config, &mut vfs, &receiver);
    Ok((host, vfs))
}

fn load(
    crate_graph: CrateGraph,
    source_root_config: SourceRootConfig,
    vfs: &mut vfs::Vfs,
    receiver: &Receiver<vfs::loader::Message>,
) -> AnalysisHost {
    let lru_cap = std::env::var("RA_LRU_CAP").ok().and_then(|it| it.parse::<usize>().ok());
    let mut host = AnalysisHost::new(lru_cap);
    let mut analysis_change = AnalysisChange::new();

    // wait until Vfs has loaded all roots
    for task in receiver {
        match task {
            vfs::loader::Message::Progress { n_done, n_total } => {
                if n_done == n_total {
                    break;
                }
            }
            vfs::loader::Message::Loaded { files } => {
                for (path, contents) in files {
                    vfs.set_file_contents(path.into(), contents)
                }
            }
        }
    }
    let changes = vfs.take_changes();
    for file in changes {
        if file.exists() {
            let contents = vfs.file_contents(file.file_id).to_vec();
            if let Ok(text) = String::from_utf8(contents) {
                analysis_change.change_file(file.file_id, Some(Arc::new(text)))
            }
        }
    }
    let source_roots = source_root_config.partition(&vfs);
    analysis_change.set_roots(source_roots);

    analysis_change.set_crate_graph(crate_graph);

    host.apply_change(analysis_change);
    host
}

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

    use hir::Crate;

    #[test]
    fn test_loading_rust_analyzer() {
        let path = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap().parent().unwrap();
        let (host, _vfs) = load_cargo(path, false, false).unwrap();
        let n_crates = Crate::all(host.raw_database()).len();
        // RA has quite a few crates, but the exact count doesn't matter
        assert!(n_crates > 20);
    }
}
<|endoftext|>