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

use crate::context::*;

use std::marker::PhantomData;
use std::ops::{Index, IndexMut};
use std::slice;

/// Tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocks<'a> {
  data: *const Block,
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  stride: usize, // number of cols in the underlying FrameBlocks
  phantom: PhantomData<&'a Block>,
}

/// Mutable tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocksMut<'a> {
  data: *mut Block,
  // private to guarantee borrowing rules
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  stride: usize, // number of cols in the underlying FrameBlocks
  phantom: PhantomData<&'a mut Block>,
}

// common impl for TileBlocks and TileBlocksMut
macro_rules! tile_blocks_common {
  // $name: TileBlocks or TileBlocksMut
  // $opt_mut: nothing or mut
  ($name:ident $(,$opt_mut:tt)?) => {
    impl<'a> $name<'a> {

      #[inline(always)]
      pub fn new(
        frame_blocks: &'a $($opt_mut)? FrameBlocks,
        x: usize,
        y: usize,
        cols: usize,
        rows: usize,
      ) -> Self {
        Self {
          data: & $($opt_mut)? frame_blocks[y][x],
          x,
          y,
          cols,
          rows,
          stride: frame_blocks.cols,
          phantom: PhantomData,
        }
      }

      #[inline(always)]
      pub fn x(&self) -> usize {
        self.x
      }

      #[inline(always)]
      pub fn y(&self) -> usize {
        self.y
      }

      #[inline(always)]
      pub fn cols(&self) -> usize {
        self.cols
      }

      #[inline(always)]
      pub fn rows(&self) -> usize {
        self.rows
      }
    }

    unsafe impl Send for $name<'_> {}
    unsafe impl Sync for $name<'_> {}

    impl Index<usize> for $name<'_> {
      type Output = [Block];
      #[inline(always)]
      fn index(&self, index: usize) -> &Self::Output {
        assert!(index < self.rows);
        unsafe {
          let ptr = self.data.add(index * self.stride);
          slice::from_raw_parts(ptr, self.cols)
        }
      }
    }

    // for convenience, also index by BlockOffset
    impl Index<BlockOffset> for $name<'_> {
      type Output = Block;
      #[inline(always)]
      fn index(&self, bo: BlockOffset) -> &Self::Output {
        &self[bo.y][bo.x]
      }
    }
  }
}

tile_blocks_common!(TileBlocks);
tile_blocks_common!(TileBlocksMut, mut);

impl TileBlocksMut<'_> {
  #[inline(always)]
  pub fn as_const(&self) -> TileBlocks<'_> {
    TileBlocks {
      data: self.data,
      x: self.x,
      y: self.y,
      cols: self.cols,
      rows: self.rows,
      stride: self.stride,
      phantom: PhantomData,
    }
  }
}

impl IndexMut<usize> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
    assert!(index < self.rows);
    unsafe {
      let ptr = self.data.add(index * self.stride);
      slice::from_raw_parts_mut(ptr, self.cols)
    }
  }
}

impl IndexMut<BlockOffset> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, bo: BlockOffset) -> &mut Self::Output {
    &mut self[bo.y][bo.x]
  }
}
<commit_msg>Implement FrameBlocks methods on TileBlocks<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::context::*;
use crate::partition::*;

use std::cmp;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};
use std::slice;

/// Tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocks<'a> {
  data: *const Block,
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  stride: usize, // number of cols in the underlying FrameBlocks
  phantom: PhantomData<&'a Block>,
}

/// Mutable tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocksMut<'a> {
  data: *mut Block,
  // private to guarantee borrowing rules
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  stride: usize, // number of cols in the underlying FrameBlocks
  phantom: PhantomData<&'a mut Block>,
}

// common impl for TileBlocks and TileBlocksMut
macro_rules! tile_blocks_common {
  // $name: TileBlocks or TileBlocksMut
  // $opt_mut: nothing or mut
  ($name:ident $(,$opt_mut:tt)?) => {
    impl<'a> $name<'a> {

      #[inline(always)]
      pub fn new(
        frame_blocks: &'a $($opt_mut)? FrameBlocks,
        x: usize,
        y: usize,
        cols: usize,
        rows: usize,
      ) -> Self {
        Self {
          data: & $($opt_mut)? frame_blocks[y][x],
          x,
          y,
          cols,
          rows,
          stride: frame_blocks.cols,
          phantom: PhantomData,
        }
      }

      #[inline(always)]
      pub fn x(&self) -> usize {
        self.x
      }

      #[inline(always)]
      pub fn y(&self) -> usize {
        self.y
      }

      #[inline(always)]
      pub fn cols(&self) -> usize {
        self.cols
      }

      #[inline(always)]
      pub fn rows(&self) -> usize {
        self.rows
      }

      #[inline(always)]
      pub fn above_of(&self, bo: BlockOffset) -> &Block {
        &self[bo.y - 1][bo.x]
      }

      #[inline(always)]
      pub fn left_of(&self, bo: BlockOffset) -> &Block {
        &self[bo.y][bo.x - 1]
      }

      #[inline(always)]
      pub fn above_left_of(&self, bo: BlockOffset) -> &Block {
        &self[bo.y - 1][bo.x - 1]
      }

      pub fn get_cdef(&self, sbo: SuperBlockOffset) -> u8 {
        let bo = sbo.block_offset(0, 0);
        self[bo.y][bo.x].cdef_index
      }
    }

    unsafe impl Send for $name<'_> {}
    unsafe impl Sync for $name<'_> {}

    impl Index<usize> for $name<'_> {
      type Output = [Block];
      #[inline(always)]
      fn index(&self, index: usize) -> &Self::Output {
        assert!(index < self.rows);
        unsafe {
          let ptr = self.data.add(index * self.stride);
          slice::from_raw_parts(ptr, self.cols)
        }
      }
    }

    // for convenience, also index by BlockOffset
    impl Index<BlockOffset> for $name<'_> {
      type Output = Block;
      #[inline(always)]
      fn index(&self, bo: BlockOffset) -> &Self::Output {
        &self[bo.y][bo.x]
      }
    }
  }
}

tile_blocks_common!(TileBlocks);
tile_blocks_common!(TileBlocksMut, mut);

impl TileBlocksMut<'_> {
  #[inline(always)]
  pub fn as_const(&self) -> TileBlocks<'_> {
    TileBlocks {
      data: self.data,
      x: self.x,
      y: self.y,
      cols: self.cols,
      rows: self.rows,
      stride: self.stride,
      phantom: PhantomData,
    }
  }

  pub fn for_each<F>(&mut self, bo: BlockOffset, bsize: BlockSize, f: F)
  where
    F: Fn(&mut Block) -> (),
  {
    let bw = bsize.width_mi();
    let bh = bsize.height_mi();
    for y in 0..bh {
      for x in 0..bw {
        f(&mut self[bo.y + y as usize][bo.x + x as usize]);
      }
    }
  }

  pub fn set_mode(
    &mut self,
    bo: BlockOffset,
    bsize: BlockSize,
    mode: PredictionMode,
  ) {
    self.for_each(bo, bsize, |block| block.mode = mode);
  }

  pub fn set_block_size(&mut self, bo: BlockOffset, bsize: BlockSize) {
    let n4_w = bsize.width_mi();
    let n4_h = bsize.height_mi();
    self.for_each(bo, bsize, |block| {
      block.bsize = bsize;
      block.n4_w = n4_w;
      block.n4_h = n4_h
    });
  }

  pub fn set_tx_size(
    &mut self,
    bo: BlockOffset,
    bsize: BlockSize,
    tx_size: TxSize,
  ) {
    self.for_each(bo, bsize, |block| block.txsize = tx_size);
  }

  pub fn set_skip(&mut self, bo: BlockOffset, bsize: BlockSize, skip: bool) {
    self.for_each(bo, bsize, |block| block.skip = skip);
  }

  pub fn set_segmentation_idx(
    &mut self,
    bo: BlockOffset,
    bsize: BlockSize,
    idx: u8,
  ) {
    self.for_each(bo, bsize, |block| block.segmentation_idx = idx);
  }

  pub fn set_ref_frames(
    &mut self,
    bo: BlockOffset,
    bsize: BlockSize,
    r: [RefType; 2],
  ) {
    self.for_each(bo, bsize, |block| block.ref_frames = r);
  }

  pub fn set_motion_vectors(
    &mut self,
    bo: BlockOffset,
    bsize: BlockSize,
    mvs: [MotionVector; 2],
  ) {
    self.for_each(bo, bsize, |block| block.mv = mvs);
  }

  pub fn set_cdef(&mut self, sbo: SuperBlockOffset, cdef_index: u8) {
    let bo = sbo.block_offset(0, 0);
    // Checkme: Is 16 still the right block unit for 128x128 superblocks?
    let bw = cmp::min(bo.x + MAX_MIB_SIZE, self.cols);
    let bh = cmp::min(bo.y + MAX_MIB_SIZE, self.rows);
    for y in bo.y..bh {
      for x in bo.x..bw {
        self[y as usize][x as usize].cdef_index = cdef_index;
      }
    }
  }
}

impl IndexMut<usize> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
    assert!(index < self.rows);
    unsafe {
      let ptr = self.data.add(index * self.stride);
      slice::from_raw_parts_mut(ptr, self.cols)
    }
  }
}

impl IndexMut<BlockOffset> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, bo: BlockOffset) -> &mut Self::Output {
    &mut self[bo.y][bo.x]
  }
}
<|endoftext|>


<commit_before>// Copyright 2016 Joe Wilm, The Alacritty Project Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::{Pty, HANDLE};

use std::i16;
use std::mem;
use std::os::windows::io::IntoRawHandle;
use std::ptr;
use std::sync::Arc;

use dunce::canonicalize;
use mio_anonymous_pipes::{EventedAnonRead, EventedAnonWrite};
use miow;
use widestring::U16CString;
use winapi::shared::basetsd::{PSIZE_T, SIZE_T};
use winapi::shared::minwindef::{BYTE, DWORD};
use winapi::shared::ntdef::{HANDLE, HRESULT, LPWSTR};
use winapi::shared::winerror::S_OK;
use winapi::um::libloaderapi::{GetModuleHandleA, GetProcAddress};
use winapi::um::processthreadsapi::{
    CreateProcessW, InitializeProcThreadAttributeList, UpdateProcThreadAttribute,
    PROCESS_INFORMATION, STARTUPINFOW,
};
use winapi::um::winbase::{EXTENDED_STARTUPINFO_PRESENT, STARTF_USESTDHANDLES, STARTUPINFOEXW};
use winapi::um::wincontypes::{COORD, HPCON};

use crate::cli::Options;
use crate::config::{Config, Shell};
use crate::display::OnResize;
use crate::term::SizeInfo;

/// Dynamically-loaded Pseudoconsole API from kernel32.dll
///
/// The field names are deliberately PascalCase as this matches
/// the defined symbols in kernel32 and also is the convention
/// that the `winapi` crate follows.
#[allow(non_snake_case)]
struct ConptyApi {
    CreatePseudoConsole:
        unsafe extern "system" fn(COORD, HANDLE, HANDLE, DWORD, *mut HPCON) -> HRESULT,
    ResizePseudoConsole: unsafe extern "system" fn(HPCON, COORD) -> HRESULT,
    ClosePseudoConsole: unsafe extern "system" fn(HPCON),
}

impl ConptyApi {
    /// Load the API or None if it cannot be found.
    pub fn new() -> Option<Self> {
        // Unsafe because windows API calls
        unsafe {
            let hmodule = GetModuleHandleA("kernel32\0".as_ptr() as _);
            assert!(!hmodule.is_null());

            let cpc = GetProcAddress(hmodule, "CreatePseudoConsole\0".as_ptr() as _);
            let rpc = GetProcAddress(hmodule, "ResizePseudoConsole\0".as_ptr() as _);
            let clpc = GetProcAddress(hmodule, "ClosePseudoConsole\0".as_ptr() as _);

            if cpc.is_null() || rpc.is_null() || clpc.is_null() {
                None
            } else {
                Some(Self {
                    CreatePseudoConsole: mem::transmute(cpc),
                    ResizePseudoConsole: mem::transmute(rpc),
                    ClosePseudoConsole: mem::transmute(clpc),
                })
            }
        }
    }
}

/// RAII Pseudoconsole
pub struct Conpty {
    pub handle: HPCON,
    api: ConptyApi,
}

/// Handle can be cloned freely and moved between threads.
pub type ConptyHandle = Arc<Conpty>;

impl Drop for Conpty {
    fn drop(&mut self) {
        unsafe { (self.api.ClosePseudoConsole)(self.handle) }
    }
}

// The Conpty API can be accessed from multiple threads.
unsafe impl Send for Conpty {}
unsafe impl Sync for Conpty {}

pub fn new<'a>(
    config: &Config,
    options: &Options,
    size: &SizeInfo,
    _window_id: Option<usize>,
) -> Option<Pty<'a>> {
    if !config.enable_experimental_conpty_backend() {
        return None;
    }

    let api = ConptyApi::new()?;

    let mut pty_handle = 0 as HPCON;

    // Passing 0 as the size parameter allows the "system default" buffer
    // size to be used. There may be small performance and memory advantages
    // to be gained by tuning this in the future, but it's likely a reasonable
    // start point.
    let (conout, conout_pty_handle) = miow::pipe::anonymous(0).unwrap();
    let (conin_pty_handle, conin) = miow::pipe::anonymous(0).unwrap();

    let coord =
        coord_from_sizeinfo(size).expect("Overflow when creating initial size on pseudoconsole");

    // Create the Pseudo Console, using the pipes
    let result = unsafe {
        (api.CreatePseudoConsole)(
            coord,
            conin_pty_handle.into_raw_handle(),
            conout_pty_handle.into_raw_handle(),
            0,
            &mut pty_handle as *mut HPCON,
        )
    };

    assert!(result == S_OK);

    let mut success;

    // Prepare child process startup info

    let mut size: SIZE_T = 0;

    let mut startup_info_ex: STARTUPINFOEXW = Default::default();

    let title = options.title.as_ref().map(String::as_str).unwrap_or("Alacritty");
    let title = U16CString::from_str(title).unwrap();
    startup_info_ex.StartupInfo.lpTitle = title.as_ptr() as LPWSTR;

    startup_info_ex.StartupInfo.cb = mem::size_of::<STARTUPINFOEXW>() as u32;

    // Setting this flag but leaving all the handles as default (null) ensures the
    // pty process does not inherit any handles from this Alacritty process.
    startup_info_ex.StartupInfo.dwFlags |= STARTF_USESTDHANDLES;

    // Create the appropriately sized thread attribute list.
    unsafe {
        success =
            InitializeProcThreadAttributeList(ptr::null_mut(), 1, 0, &mut size as PSIZE_T) > 0;

        // This call was expected to return false.
        assert!(!success);
    }

    let mut attr_list: Box<[BYTE]> = vec![0; size].into_boxed_slice();

    // Set startup info's attribute list & initialize it
    //
    // Lint failure is spurious; it's because winapi's definition of PROC_THREAD_ATTRIBUTE_LIST
    // implies it is one pointer in size (32 or 64 bits) but really this is just a dummy value.
    // Casting a *mut u8 (pointer to 8 bit type) might therefore not be aligned correctly in
    // the compiler's eyes.
    #[allow(clippy::cast_ptr_alignment)]
    {
        startup_info_ex.lpAttributeList = attr_list.as_mut_ptr() as _;
    }

    unsafe {
        success = InitializeProcThreadAttributeList(
            startup_info_ex.lpAttributeList,
            1,
            0,
            &mut size as PSIZE_T,
        ) > 0;

        assert!(success);
    }

    // Set thread attribute list's Pseudo Console to the specified ConPTY
    unsafe {
        success = UpdateProcThreadAttribute(
            startup_info_ex.lpAttributeList,
            0,
            22 | 0x0002_0000, // PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE
            pty_handle,
            mem::size_of::<HPCON>(),
            ptr::null_mut(),
            ptr::null_mut(),
        ) > 0;

        assert!(success);
    }

    // Get process commandline
    let default_shell = &Shell::new("powershell");
    let shell = config.shell().unwrap_or(default_shell);
    let initial_command = options.command().unwrap_or(shell);
    let mut cmdline = initial_command.args().to_vec();
    cmdline.insert(0, initial_command.program().into());

    // Warning, here be borrow hell
    let cwd = options.working_dir.as_ref().map(|dir| canonicalize(dir).unwrap());
    let cwd = cwd.as_ref().map(|dir| dir.to_str().unwrap());

    // Create the client application, using startup info containing ConPTY info
    let cmdline = U16CString::from_str(&cmdline.join(" ")).unwrap();
    let cwd = cwd.map(|s| U16CString::from_str(&s).unwrap());

    let mut proc_info: PROCESS_INFORMATION = Default::default();
    unsafe {
        success = CreateProcessW(
            ptr::null(),
            cmdline.as_ptr() as LPWSTR,
            ptr::null_mut(),
            ptr::null_mut(),
            false as i32,
            EXTENDED_STARTUPINFO_PRESENT,
            ptr::null_mut(),
            cwd.as_ref().map_or_else(ptr::null, |s| s.as_ptr()),
            &mut startup_info_ex.StartupInfo as *mut STARTUPINFOW,
            &mut proc_info as *mut PROCESS_INFORMATION,
        ) > 0;

        assert!(success);
    }

    // Store handle to console
    unsafe {
        HANDLE = proc_info.hProcess;
    }

    let conin = EventedAnonWrite::new(conin);
    let conout = EventedAnonRead::new(conout);

    let agent = Conpty { handle: pty_handle, api };

    Some(Pty {
        handle: super::PtyHandle::Conpty(ConptyHandle::new(agent)),
        conout: super::EventedReadablePipe::Anonymous(conout),
        conin: super::EventedWritablePipe::Anonymous(conin),
        read_token: 0.into(),
        write_token: 0.into(),
    })
}

impl OnResize for ConptyHandle {
    fn on_resize(&mut self, sizeinfo: &SizeInfo) {
        if let Some(coord) = coord_from_sizeinfo(sizeinfo) {
            let result = unsafe { (self.api.ResizePseudoConsole)(self.handle, coord) };
            assert!(result == S_OK);
        }
    }
}

/// Helper to build a COORD from a SizeInfo, returing None in overflow cases.
fn coord_from_sizeinfo(sizeinfo: &SizeInfo) -> Option<COORD> {
    let cols = sizeinfo.cols().0;
    let lines = sizeinfo.lines().0;

    if cols <= i16::MAX as usize && lines <= i16::MAX as usize {
        Some(COORD { X: sizeinfo.cols().0 as i16, Y: sizeinfo.lines().0 as i16 })
    } else {
        None
    }
}
<commit_msg>Improve errors for missing shell file on Windows <commit_after>// Copyright 2016 Joe Wilm, The Alacritty Project Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::{Pty, HANDLE};

use std::i16;
use std::io::Error;
use std::mem;
use std::os::windows::io::IntoRawHandle;
use std::ptr;
use std::sync::Arc;

use dunce::canonicalize;
use mio_anonymous_pipes::{EventedAnonRead, EventedAnonWrite};
use miow;
use widestring::U16CString;
use winapi::shared::basetsd::{PSIZE_T, SIZE_T};
use winapi::shared::minwindef::{BYTE, DWORD};
use winapi::shared::ntdef::{HANDLE, HRESULT, LPWSTR};
use winapi::shared::winerror::S_OK;
use winapi::um::libloaderapi::{GetModuleHandleA, GetProcAddress};
use winapi::um::processthreadsapi::{
    CreateProcessW, InitializeProcThreadAttributeList, UpdateProcThreadAttribute,
    PROCESS_INFORMATION, STARTUPINFOW,
};
use winapi::um::winbase::{EXTENDED_STARTUPINFO_PRESENT, STARTF_USESTDHANDLES, STARTUPINFOEXW};
use winapi::um::wincontypes::{COORD, HPCON};

use crate::cli::Options;
use crate::config::{Config, Shell};
use crate::display::OnResize;
use crate::term::SizeInfo;

/// Dynamically-loaded Pseudoconsole API from kernel32.dll
///
/// The field names are deliberately PascalCase as this matches
/// the defined symbols in kernel32 and also is the convention
/// that the `winapi` crate follows.
#[allow(non_snake_case)]
struct ConptyApi {
    CreatePseudoConsole:
        unsafe extern "system" fn(COORD, HANDLE, HANDLE, DWORD, *mut HPCON) -> HRESULT,
    ResizePseudoConsole: unsafe extern "system" fn(HPCON, COORD) -> HRESULT,
    ClosePseudoConsole: unsafe extern "system" fn(HPCON),
}

impl ConptyApi {
    /// Load the API or None if it cannot be found.
    pub fn new() -> Option<Self> {
        // Unsafe because windows API calls
        unsafe {
            let hmodule = GetModuleHandleA("kernel32\0".as_ptr() as _);
            assert!(!hmodule.is_null());

            let cpc = GetProcAddress(hmodule, "CreatePseudoConsole\0".as_ptr() as _);
            let rpc = GetProcAddress(hmodule, "ResizePseudoConsole\0".as_ptr() as _);
            let clpc = GetProcAddress(hmodule, "ClosePseudoConsole\0".as_ptr() as _);

            if cpc.is_null() || rpc.is_null() || clpc.is_null() {
                None
            } else {
                Some(Self {
                    CreatePseudoConsole: mem::transmute(cpc),
                    ResizePseudoConsole: mem::transmute(rpc),
                    ClosePseudoConsole: mem::transmute(clpc),
                })
            }
        }
    }
}

/// RAII Pseudoconsole
pub struct Conpty {
    pub handle: HPCON,
    api: ConptyApi,
}

/// Handle can be cloned freely and moved between threads.
pub type ConptyHandle = Arc<Conpty>;

impl Drop for Conpty {
    fn drop(&mut self) {
        unsafe { (self.api.ClosePseudoConsole)(self.handle) }
    }
}

// The Conpty API can be accessed from multiple threads.
unsafe impl Send for Conpty {}
unsafe impl Sync for Conpty {}

pub fn new<'a>(
    config: &Config,
    options: &Options,
    size: &SizeInfo,
    _window_id: Option<usize>,
) -> Option<Pty<'a>> {
    if !config.enable_experimental_conpty_backend() {
        return None;
    }

    let api = ConptyApi::new()?;

    let mut pty_handle = 0 as HPCON;

    // Passing 0 as the size parameter allows the "system default" buffer
    // size to be used. There may be small performance and memory advantages
    // to be gained by tuning this in the future, but it's likely a reasonable
    // start point.
    let (conout, conout_pty_handle) = miow::pipe::anonymous(0).unwrap();
    let (conin_pty_handle, conin) = miow::pipe::anonymous(0).unwrap();

    let coord =
        coord_from_sizeinfo(size).expect("Overflow when creating initial size on pseudoconsole");

    // Create the Pseudo Console, using the pipes
    let result = unsafe {
        (api.CreatePseudoConsole)(
            coord,
            conin_pty_handle.into_raw_handle(),
            conout_pty_handle.into_raw_handle(),
            0,
            &mut pty_handle as *mut HPCON,
        )
    };

    assert!(result == S_OK);

    let mut success;

    // Prepare child process startup info

    let mut size: SIZE_T = 0;

    let mut startup_info_ex: STARTUPINFOEXW = Default::default();

    let title = options.title.as_ref().map(String::as_str).unwrap_or("Alacritty");
    let title = U16CString::from_str(title).unwrap();
    startup_info_ex.StartupInfo.lpTitle = title.as_ptr() as LPWSTR;

    startup_info_ex.StartupInfo.cb = mem::size_of::<STARTUPINFOEXW>() as u32;

    // Setting this flag but leaving all the handles as default (null) ensures the
    // pty process does not inherit any handles from this Alacritty process.
    startup_info_ex.StartupInfo.dwFlags |= STARTF_USESTDHANDLES;

    // Create the appropriately sized thread attribute list.
    unsafe {
        let failure =
            InitializeProcThreadAttributeList(ptr::null_mut(), 1, 0, &mut size as PSIZE_T) > 0;

        // This call was expected to return false.
        if (!failure) {
            panic_shell_spawn();
        }
    }

    let mut attr_list: Box<[BYTE]> = vec![0; size].into_boxed_slice();

    // Set startup info's attribute list & initialize it
    //
    // Lint failure is spurious; it's because winapi's definition of PROC_THREAD_ATTRIBUTE_LIST
    // implies it is one pointer in size (32 or 64 bits) but really this is just a dummy value.
    // Casting a *mut u8 (pointer to 8 bit type) might therefore not be aligned correctly in
    // the compiler's eyes.
    #[allow(clippy::cast_ptr_alignment)]
    {
        startup_info_ex.lpAttributeList = attr_list.as_mut_ptr() as _;
    }

    unsafe {
        success = InitializeProcThreadAttributeList(
            startup_info_ex.lpAttributeList,
            1,
            0,
            &mut size as PSIZE_T,
        ) > 0;

        if (!success) {
            panic_shell_spawn();
        }
    }

    // Set thread attribute list's Pseudo Console to the specified ConPTY
    unsafe {
        success = UpdateProcThreadAttribute(
            startup_info_ex.lpAttributeList,
            0,
            22 | 0x0002_0000, // PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE
            pty_handle,
            mem::size_of::<HPCON>(),
            ptr::null_mut(),
            ptr::null_mut(),
        ) > 0;

        if (!success) {
            panic_shell_spawn();
        }
    }

    // Get process commandline
    let default_shell = &Shell::new("powershell");
    let shell = config.shell().unwrap_or(default_shell);
    let initial_command = options.command().unwrap_or(shell);
    let mut cmdline = initial_command.args().to_vec();
    cmdline.insert(0, initial_command.program().into());

    // Warning, here be borrow hell
    let cwd = options.working_dir.as_ref().map(|dir| canonicalize(dir).unwrap());
    let cwd = cwd.as_ref().map(|dir| dir.to_str().unwrap());

    // Create the client application, using startup info containing ConPTY info
    let cmdline = U16CString::from_str(&cmdline.join(" ")).unwrap();
    let cwd = cwd.map(|s| U16CString::from_str(&s).unwrap());

    let mut proc_info: PROCESS_INFORMATION = Default::default();
    unsafe {
        success = CreateProcessW(
            ptr::null(),
            cmdline.as_ptr() as LPWSTR,
            ptr::null_mut(),
            ptr::null_mut(),
            false as i32,
            EXTENDED_STARTUPINFO_PRESENT,
            ptr::null_mut(),
            cwd.as_ref().map_or_else(ptr::null, |s| s.as_ptr()),
            &mut startup_info_ex.StartupInfo as *mut STARTUPINFOW,
            &mut proc_info as *mut PROCESS_INFORMATION,
        ) > 0;

        if (!success) {
            panic_shell_spawn();
        }
    }

    // Store handle to console
    unsafe {
        HANDLE = proc_info.hProcess;
    }

    let conin = EventedAnonWrite::new(conin);
    let conout = EventedAnonRead::new(conout);

    let agent = Conpty { handle: pty_handle, api };

    Some(Pty {
        handle: super::PtyHandle::Conpty(ConptyHandle::new(agent)),
        conout: super::EventedReadablePipe::Anonymous(conout),
        conin: super::EventedWritablePipe::Anonymous(conin),
        read_token: 0.into(),
        write_token: 0.into(),
    })
}

// Panic with the last os error as message
fn panic_shell_spawn() {
    panic!("Unable to spawn shell: {}", Error::last_os_error());
}

impl OnResize for ConptyHandle {
    fn on_resize(&mut self, sizeinfo: &SizeInfo) {
        if let Some(coord) = coord_from_sizeinfo(sizeinfo) {
            let result = unsafe { (self.api.ResizePseudoConsole)(self.handle, coord) };
            assert!(result == S_OK);
        }
    }
}

/// Helper to build a COORD from a SizeInfo, returing None in overflow cases.
fn coord_from_sizeinfo(sizeinfo: &SizeInfo) -> Option<COORD> {
    let cols = sizeinfo.cols().0;
    let lines = sizeinfo.lines().0;

    if cols <= i16::MAX as usize && lines <= i16::MAX as usize {
        Some(COORD { X: sizeinfo.cols().0 as i16, Y: sizeinfo.lines().0 as i16 })
    } else {
        None
    }
}
<|endoftext|>


<commit_before>//! The SubscriptionManager used to manage subscription based RPCs.
//!
//! The manager provides four main things in terms of functionality:
//!
//! 1. The ability to create unique subscription IDs through the
//! use of the `IdProvider` trait. Two implementations are availble
//! out of the box, a `NumericIdProvider` and a `RandomStringIdProvider`.
//!
//! 2. An executor with which to drive `Future`s to completion.
//!
//! 3. A way to add new subscriptions. Subscriptions should come in the form
//! of a `Stream`. These subscriptions will be transformed into notifications
//! by the manager, which can be consumed by the client.
//!
//! 4. A way to cancel any currently active subscription.

use std::collections::HashMap;
use std::iter;
use std::sync::{
	atomic::{AtomicUsize, Ordering},
	Arc,
};

use crate::core::futures::sync::oneshot;
use crate::core::futures::{future as future01, Future as Future01};
use crate::{
	typed::{Sink, Subscriber},
	SubscriptionId,
};

use log::{error, warn};
use parking_lot::Mutex;
use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};

/// Alias for an implementation of `futures::future::Executor`.
pub type TaskExecutor = Arc<dyn future01::Executor<Box<dyn Future01<Item = (), Error = ()> + Send>> + Send + Sync>;

type ActiveSubscriptions = Arc<Mutex<HashMap<SubscriptionId, oneshot::Sender<()>>>>;

/// Trait used to provide unique subscription IDs.
pub trait IdProvider {
	/// A unique ID used to identify a subscription.
	type Id: Default + Into<SubscriptionId>;

	/// Returns the next ID for the subscription.
	fn next_id(&self) -> Self::Id;
}

/// Provides a thread-safe incrementing integer which
/// can be used as a subscription ID.
#[derive(Debug)]
pub struct NumericIdProvider {
	current_id: AtomicUsize,
}

impl NumericIdProvider {
	/// Create a new NumericIdProvider.
	pub fn new() -> Self {
		Default::default()
	}

	/// Create a new NumericIdProvider starting from
	/// the given ID.
	pub fn with_id(id: AtomicUsize) -> Self {
		Self { current_id: id }
	}
}

impl IdProvider for NumericIdProvider {
	type Id = u64;

	fn next_id(&self) -> Self::Id {
		self.current_id.fetch_add(1, Ordering::AcqRel) as u64
	}
}

impl Default for NumericIdProvider {
	fn default() -> Self {
		NumericIdProvider {
			current_id: AtomicUsize::new(1),
		}
	}
}

/// Used to generate random strings for use as
/// subscription IDs.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct RandomStringIdProvider {
	len: usize,
}

impl RandomStringIdProvider {
	/// Create a new RandomStringIdProvider.
	pub fn new() -> Self {
		Default::default()
	}

	/// Create a new RandomStringIdProvider, which will generate
	/// random id strings of the given length.
	pub fn with_len(len: usize) -> Self {
		Self { len }
	}
}

impl IdProvider for RandomStringIdProvider {
	type Id = String;

	fn next_id(&self) -> Self::Id {
		let mut rng = thread_rng();
		let id: String = iter::repeat(())
			.map(|()| rng.sample(Alphanumeric))
			.take(self.len)
			.collect();
		id
	}
}

impl Default for RandomStringIdProvider {
	fn default() -> Self {
		Self { len: 16 }
	}
}

/// Subscriptions manager.
///
/// Takes care of assigning unique subscription ids and
/// driving the sinks into completion.
#[derive(Clone)]
pub struct SubscriptionManager<I: IdProvider = RandomStringIdProvider> {
	id_provider: I,
	active_subscriptions: ActiveSubscriptions,
	executor: TaskExecutor,
}

impl SubscriptionManager {
	/// Creates a new SubscriptionManager.
	///
	/// Uses `RandomStringIdProvider` as the ID provider.
	pub fn new(executor: TaskExecutor) -> Self {
		Self {
			id_provider: RandomStringIdProvider::default(),
			active_subscriptions: Default::default(),
			executor,
		}
	}
}

impl<I: IdProvider> SubscriptionManager<I> {
	/// Creates a new SubscriptionManager with the specified
	/// ID provider.
	pub fn with_id_provider(id_provider: I, executor: TaskExecutor) -> Self {
		Self {
			id_provider,
			active_subscriptions: Default::default(),
			executor,
		}
	}

	/// Borrows the internal task executor.
	///
	/// This can be used to spawn additional tasks on the underlying event loop.
	pub fn executor(&self) -> &TaskExecutor {
		&self.executor
	}

	/// Creates new subscription for given subscriber.
	///
	/// Second parameter is a function that converts Subscriber Sink into a Future.
	/// This future will be driven to completion by the underlying event loop
	pub fn add<T, E, G, R, F>(&self, subscriber: Subscriber<T, E>, into_future: G) -> SubscriptionId
	where
		G: FnOnce(Sink<T, E>) -> R,
		R: future01::IntoFuture<Future = F, Item = (), Error = ()>,
		F: future01::Future<Item = (), Error = ()> + Send + 'static,
	{
		let id = self.id_provider.next_id();
		let subscription_id: SubscriptionId = id.into();
		if let Ok(sink) = subscriber.assign_id(subscription_id.clone()) {
			let (tx, rx) = oneshot::channel();
			let future = into_future(sink)
				.into_future()
				.select(rx.map_err(|e| warn!("Error timing out: {:?}", e)))
				.then(|_| Ok(()));

			self.active_subscriptions.lock().insert(subscription_id.clone(), tx);
			if self.executor.execute(Box::new(future)).is_err() {
				error!("Failed to spawn RPC subscription task");
			}
		}

		subscription_id
	}

	/// Cancel subscription.
	///
	/// Returns true if subscription existed or false otherwise.
	pub fn cancel(&self, id: SubscriptionId) -> bool {
		if let Some(tx) = self.active_subscriptions.lock().remove(&id) {
			let _ = tx.send(());
			return true;
		}

		false
	}
}

impl<I: Default + IdProvider> SubscriptionManager<I> {
	/// Creates a new SubscriptionManager.
	pub fn with_executor(executor: TaskExecutor) -> Self {
		Self {
			id_provider: Default::default(),
			active_subscriptions: Default::default(),
			executor,
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::typed::Subscriber;
	use futures::{compat::Future01CompatExt, executor, FutureExt};
	use futures::{stream, StreamExt, TryStreamExt};

	use crate::core::futures::sink::Sink as Sink01;
	use crate::core::futures::stream::Stream as Stream01;

	// Executor shared by all tests.
	//
	// This shared executor is used to prevent `Too many open files` errors
	// on systems with a lot of cores.
	lazy_static::lazy_static! {
		static ref EXECUTOR: executor::ThreadPool = executor::ThreadPool::new()
			.expect("Failed to create thread pool executor for tests");
	}

	pub struct TestTaskExecutor;
	type Boxed01Future01 = Box<dyn future01::Future<Item = (), Error = ()> + Send + 'static>;

	impl future01::Executor<Boxed01Future01> for TestTaskExecutor {
		fn execute(&self, future: Boxed01Future01) -> std::result::Result<(), future01::ExecuteError<Boxed01Future01>> {
			EXECUTOR.spawn_ok(future.compat().map(drop));
			Ok(())
		}
	}

	#[test]
	fn making_a_numeric_id_provider_works() {
		let provider = NumericIdProvider::new();
		let expected_id = 1;
		let actual_id = provider.next_id();

		assert_eq!(actual_id, expected_id);
	}

	#[test]
	fn default_numeric_id_provider_works() {
		let provider: NumericIdProvider = Default::default();
		let expected_id = 1;
		let actual_id = provider.next_id();

		assert_eq!(actual_id, expected_id);
	}

	#[test]
	fn numeric_id_provider_with_id_works() {
		let provider = NumericIdProvider::with_id(AtomicUsize::new(5));
		let expected_id = 5;
		let actual_id = provider.next_id();

		assert_eq!(actual_id, expected_id);
	}

	#[test]
	fn random_string_provider_returns_id_with_correct_default_len() {
		let provider = RandomStringIdProvider::new();
		let expected_len = 16;
		let actual_len = provider.next_id().len();

		assert_eq!(actual_len, expected_len);
	}

	#[test]
	fn random_string_provider_returns_id_with_correct_user_given_len() {
		let expected_len = 10;
		let provider = RandomStringIdProvider::with_len(expected_len);
		let actual_len = provider.next_id().len();

		assert_eq!(actual_len, expected_len);
	}

	#[test]
	fn new_subscription_manager_defaults_to_random_string_provider() {
		let manager = SubscriptionManager::new(Arc::new(TestTaskExecutor));
		let subscriber = Subscriber::<u64>::new_test("test_subTest").0;
		let stream = stream::iter(vec![Ok(1)]).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		if let SubscriptionId::String(_) = id {
			assert!(true)
		} else {
			assert!(false, "Expected SubscriptionId::String");
		}
	}

	#[test]
	fn new_subscription_manager_works_with_numeric_id_provider() {
		let id_provider = NumericIdProvider::default();
		let manager = SubscriptionManager::with_id_provider(id_provider, Arc::new(TestTaskExecutor));

		let subscriber = Subscriber::<u64>::new_test("test_subTest").0;
		let stream = stream::iter(vec![Ok(1)]).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		if let SubscriptionId::Number(_) = id {
			assert!(true)
		} else {
			assert!(false, "Expected SubscriptionId::Number");
		}
	}

	#[test]
	fn new_subscription_manager_works_with_random_string_provider() {
		let id_provider = RandomStringIdProvider::default();
		let manager = SubscriptionManager::with_id_provider(id_provider, Arc::new(TestTaskExecutor));

		let subscriber = Subscriber::<u64>::new_test("test_subTest").0;
		let stream = stream::iter(vec![Ok(1)]).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		if let SubscriptionId::String(_) = id {
			assert!(true)
		} else {
			assert!(false, "Expected SubscriptionId::String");
		}
	}

	#[test]
	fn subscription_is_canceled_if_it_existed() {
		let manager = SubscriptionManager::<NumericIdProvider>::with_executor(Arc::new(TestTaskExecutor));
		// Need to bind receiver here (unlike the other tests) or else the subscriber
		// will think the client has disconnected and not update `active_subscriptions`
		let (subscriber, _recv, _) = Subscriber::<u64>::new_test("test_subTest");

		let (mut tx, rx) = futures::channel::mpsc::channel(8);
		tx.start_send(1).unwrap();
		let stream = rx.map(|v| Ok::<_, ()>(v)).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		let is_cancelled = manager.cancel(id);
		assert!(is_cancelled);
	}

	#[test]
	fn subscription_is_not_canceled_because_it_didnt_exist() {
		let manager = SubscriptionManager::new(Arc::new(TestTaskExecutor));

		let id: SubscriptionId = 23u32.into();
		let is_cancelled = manager.cancel(id);
		let is_not_cancelled = !is_cancelled;

		assert!(is_not_cancelled);
	}
}
<commit_msg>Place NumericIdProvider's Id behind an Arc (#554)<commit_after>//! The SubscriptionManager used to manage subscription based RPCs.
//!
//! The manager provides four main things in terms of functionality:
//!
//! 1. The ability to create unique subscription IDs through the
//! use of the `IdProvider` trait. Two implementations are availble
//! out of the box, a `NumericIdProvider` and a `RandomStringIdProvider`.
//!
//! 2. An executor with which to drive `Future`s to completion.
//!
//! 3. A way to add new subscriptions. Subscriptions should come in the form
//! of a `Stream`. These subscriptions will be transformed into notifications
//! by the manager, which can be consumed by the client.
//!
//! 4. A way to cancel any currently active subscription.

use std::collections::HashMap;
use std::iter;
use std::sync::{
	atomic::{AtomicUsize, Ordering},
	Arc,
};

use crate::core::futures::sync::oneshot;
use crate::core::futures::{future as future01, Future as Future01};
use crate::{
	typed::{Sink, Subscriber},
	SubscriptionId,
};

use log::{error, warn};
use parking_lot::Mutex;
use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};

/// Alias for an implementation of `futures::future::Executor`.
pub type TaskExecutor = Arc<dyn future01::Executor<Box<dyn Future01<Item = (), Error = ()> + Send>> + Send + Sync>;

type ActiveSubscriptions = Arc<Mutex<HashMap<SubscriptionId, oneshot::Sender<()>>>>;

/// Trait used to provide unique subscription IDs.
pub trait IdProvider {
	/// A unique ID used to identify a subscription.
	type Id: Default + Into<SubscriptionId>;

	/// Returns the next ID for the subscription.
	fn next_id(&self) -> Self::Id;
}

/// Provides a thread-safe incrementing integer which
/// can be used as a subscription ID.
#[derive(Clone, Debug)]
pub struct NumericIdProvider {
	current_id: Arc<AtomicUsize>,
}

impl NumericIdProvider {
	/// Create a new NumericIdProvider.
	pub fn new() -> Self {
		Default::default()
	}

	/// Create a new NumericIdProvider starting from
	/// the given ID.
	pub fn with_id(id: AtomicUsize) -> Self {
		Self {
			current_id: Arc::new(id),
		}
	}
}

impl IdProvider for NumericIdProvider {
	type Id = u64;

	fn next_id(&self) -> Self::Id {
		self.current_id.fetch_add(1, Ordering::AcqRel) as u64
	}
}

impl Default for NumericIdProvider {
	fn default() -> Self {
		NumericIdProvider {
			current_id: Arc::new(AtomicUsize::new(1)),
		}
	}
}

/// Used to generate random strings for use as
/// subscription IDs.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct RandomStringIdProvider {
	len: usize,
}

impl RandomStringIdProvider {
	/// Create a new RandomStringIdProvider.
	pub fn new() -> Self {
		Default::default()
	}

	/// Create a new RandomStringIdProvider, which will generate
	/// random id strings of the given length.
	pub fn with_len(len: usize) -> Self {
		Self { len }
	}
}

impl IdProvider for RandomStringIdProvider {
	type Id = String;

	fn next_id(&self) -> Self::Id {
		let mut rng = thread_rng();
		let id: String = iter::repeat(())
			.map(|()| rng.sample(Alphanumeric))
			.take(self.len)
			.collect();
		id
	}
}

impl Default for RandomStringIdProvider {
	fn default() -> Self {
		Self { len: 16 }
	}
}

/// Subscriptions manager.
///
/// Takes care of assigning unique subscription ids and
/// driving the sinks into completion.
#[derive(Clone)]
pub struct SubscriptionManager<I: IdProvider = RandomStringIdProvider> {
	id_provider: I,
	active_subscriptions: ActiveSubscriptions,
	executor: TaskExecutor,
}

impl SubscriptionManager {
	/// Creates a new SubscriptionManager.
	///
	/// Uses `RandomStringIdProvider` as the ID provider.
	pub fn new(executor: TaskExecutor) -> Self {
		Self {
			id_provider: RandomStringIdProvider::default(),
			active_subscriptions: Default::default(),
			executor,
		}
	}
}

impl<I: IdProvider> SubscriptionManager<I> {
	/// Creates a new SubscriptionManager with the specified
	/// ID provider.
	pub fn with_id_provider(id_provider: I, executor: TaskExecutor) -> Self {
		Self {
			id_provider,
			active_subscriptions: Default::default(),
			executor,
		}
	}

	/// Borrows the internal task executor.
	///
	/// This can be used to spawn additional tasks on the underlying event loop.
	pub fn executor(&self) -> &TaskExecutor {
		&self.executor
	}

	/// Creates new subscription for given subscriber.
	///
	/// Second parameter is a function that converts Subscriber Sink into a Future.
	/// This future will be driven to completion by the underlying event loop
	pub fn add<T, E, G, R, F>(&self, subscriber: Subscriber<T, E>, into_future: G) -> SubscriptionId
	where
		G: FnOnce(Sink<T, E>) -> R,
		R: future01::IntoFuture<Future = F, Item = (), Error = ()>,
		F: future01::Future<Item = (), Error = ()> + Send + 'static,
	{
		let id = self.id_provider.next_id();
		let subscription_id: SubscriptionId = id.into();
		if let Ok(sink) = subscriber.assign_id(subscription_id.clone()) {
			let (tx, rx) = oneshot::channel();
			let future = into_future(sink)
				.into_future()
				.select(rx.map_err(|e| warn!("Error timing out: {:?}", e)))
				.then(|_| Ok(()));

			self.active_subscriptions.lock().insert(subscription_id.clone(), tx);
			if self.executor.execute(Box::new(future)).is_err() {
				error!("Failed to spawn RPC subscription task");
			}
		}

		subscription_id
	}

	/// Cancel subscription.
	///
	/// Returns true if subscription existed or false otherwise.
	pub fn cancel(&self, id: SubscriptionId) -> bool {
		if let Some(tx) = self.active_subscriptions.lock().remove(&id) {
			let _ = tx.send(());
			return true;
		}

		false
	}
}

impl<I: Default + IdProvider> SubscriptionManager<I> {
	/// Creates a new SubscriptionManager.
	pub fn with_executor(executor: TaskExecutor) -> Self {
		Self {
			id_provider: Default::default(),
			active_subscriptions: Default::default(),
			executor,
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::typed::Subscriber;
	use futures::{compat::Future01CompatExt, executor, FutureExt};
	use futures::{stream, StreamExt, TryStreamExt};

	use crate::core::futures::sink::Sink as Sink01;
	use crate::core::futures::stream::Stream as Stream01;

	// Executor shared by all tests.
	//
	// This shared executor is used to prevent `Too many open files` errors
	// on systems with a lot of cores.
	lazy_static::lazy_static! {
		static ref EXECUTOR: executor::ThreadPool = executor::ThreadPool::new()
			.expect("Failed to create thread pool executor for tests");
	}

	pub struct TestTaskExecutor;
	type Boxed01Future01 = Box<dyn future01::Future<Item = (), Error = ()> + Send + 'static>;

	impl future01::Executor<Boxed01Future01> for TestTaskExecutor {
		fn execute(&self, future: Boxed01Future01) -> std::result::Result<(), future01::ExecuteError<Boxed01Future01>> {
			EXECUTOR.spawn_ok(future.compat().map(drop));
			Ok(())
		}
	}

	#[test]
	fn making_a_numeric_id_provider_works() {
		let provider = NumericIdProvider::new();
		let expected_id = 1;
		let actual_id = provider.next_id();

		assert_eq!(actual_id, expected_id);
	}

	#[test]
	fn default_numeric_id_provider_works() {
		let provider: NumericIdProvider = Default::default();
		let expected_id = 1;
		let actual_id = provider.next_id();

		assert_eq!(actual_id, expected_id);
	}

	#[test]
	fn numeric_id_provider_with_id_works() {
		let provider = NumericIdProvider::with_id(AtomicUsize::new(5));
		let expected_id = 5;
		let actual_id = provider.next_id();

		assert_eq!(actual_id, expected_id);
	}

	#[test]
	fn random_string_provider_returns_id_with_correct_default_len() {
		let provider = RandomStringIdProvider::new();
		let expected_len = 16;
		let actual_len = provider.next_id().len();

		assert_eq!(actual_len, expected_len);
	}

	#[test]
	fn random_string_provider_returns_id_with_correct_user_given_len() {
		let expected_len = 10;
		let provider = RandomStringIdProvider::with_len(expected_len);
		let actual_len = provider.next_id().len();

		assert_eq!(actual_len, expected_len);
	}

	#[test]
	fn new_subscription_manager_defaults_to_random_string_provider() {
		let manager = SubscriptionManager::new(Arc::new(TestTaskExecutor));
		let subscriber = Subscriber::<u64>::new_test("test_subTest").0;
		let stream = stream::iter(vec![Ok(1)]).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		if let SubscriptionId::String(_) = id {
			assert!(true)
		} else {
			assert!(false, "Expected SubscriptionId::String");
		}
	}

	#[test]
	fn new_subscription_manager_works_with_numeric_id_provider() {
		let id_provider = NumericIdProvider::default();
		let manager = SubscriptionManager::with_id_provider(id_provider, Arc::new(TestTaskExecutor));

		let subscriber = Subscriber::<u64>::new_test("test_subTest").0;
		let stream = stream::iter(vec![Ok(1)]).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		if let SubscriptionId::Number(_) = id {
			assert!(true)
		} else {
			assert!(false, "Expected SubscriptionId::Number");
		}
	}

	#[test]
	fn new_subscription_manager_works_with_random_string_provider() {
		let id_provider = RandomStringIdProvider::default();
		let manager = SubscriptionManager::with_id_provider(id_provider, Arc::new(TestTaskExecutor));

		let subscriber = Subscriber::<u64>::new_test("test_subTest").0;
		let stream = stream::iter(vec![Ok(1)]).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		if let SubscriptionId::String(_) = id {
			assert!(true)
		} else {
			assert!(false, "Expected SubscriptionId::String");
		}
	}

	#[test]
	fn subscription_is_canceled_if_it_existed() {
		let manager = SubscriptionManager::<NumericIdProvider>::with_executor(Arc::new(TestTaskExecutor));
		// Need to bind receiver here (unlike the other tests) or else the subscriber
		// will think the client has disconnected and not update `active_subscriptions`
		let (subscriber, _recv, _) = Subscriber::<u64>::new_test("test_subTest");

		let (mut tx, rx) = futures::channel::mpsc::channel(8);
		tx.start_send(1).unwrap();
		let stream = rx.map(|v| Ok::<_, ()>(v)).compat();

		let id = manager.add(subscriber, |sink| {
			let stream = stream.map(|res| Ok(res));

			sink.sink_map_err(|_| ()).send_all(stream).map(|_| ())
		});

		let is_cancelled = manager.cancel(id);
		assert!(is_cancelled);
	}

	#[test]
	fn subscription_is_not_canceled_because_it_didnt_exist() {
		let manager = SubscriptionManager::new(Arc::new(TestTaskExecutor));

		let id: SubscriptionId = 23u32.into();
		let is_cancelled = manager.cancel(id);
		let is_not_cancelled = !is_cancelled;

		assert!(is_not_cancelled);
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Implementation of ElasticSearch [aggregations](https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations.html)

use std::collections::{BTreeMap, HashMap};

use rustc_serialize::json::{Json, ToJson};

use error::EsError;
use units::JsonVal;

/// Script attributes for various attributes
#[derive(Debug)]
pub struct Script<'a> {
    field:  &'a str,
    script: &'a str,
    params: Option<Json>
}

impl<'a> Script<'a> {
    fn add_to_object(&self, obj: &mut BTreeMap<String, Json>) {
        obj.insert("field".to_owned(), self.field.to_json());
        obj.insert("script".to_owned(), self.script.to_json());
        match self.params {
            Some(ref json) => {
                obj.insert("params".to_owned(), json.clone());
            },
            None           => ()
        }
    }

    pub fn with_params(mut self, params: Json) -> Self {
        self.params = Some(params);
        self
    }
}

/// Min aggregation
#[derive(Debug)]
pub enum Min<'a> {
    /// Field
    Field(&'a str),

    /// By Script
    Script(Script<'a>)
}

impl<'a> From<Min<'a>> for Aggregation<'a> {
    fn from(from: Min<'a>) -> Aggregation<'a> {
        Aggregation::Metrics(MetricsAggregation::Min(from))
    }
}

impl<'a> ToJson for Min<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        match self {
            &Min::Field(field) => {
                d.insert("field".to_string(), field.to_json());
            },
            &Min::Script(ref script) => {
                script.add_to_object(&mut d);
            }
        }
        Json::Object(d)
    }
}

/// Individual aggregations and their options
#[derive(Debug)]
pub enum MetricsAggregation<'a> {
    Min(Min<'a>)
}

impl<'a> ToJson for MetricsAggregation<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        match self {
            &MetricsAggregation::Min(ref min_agg) => {
                d.insert("min".to_owned(), min_agg.to_json());
            }
        }
        Json::Object(d)
    }
}

/// Order - used for some bucketing aggregations to determine the order of
/// buckets
#[derive(Debug)]
pub enum OrderKey<'a> {
    Count,
    Term,
    Expr(&'a str)
}

impl<'a> From<&'a str> for OrderKey<'a> {
    fn from(from: &'a str) -> OrderKey<'a> {
        OrderKey::Expr(from)
    }
}

impl<'a> ToString for OrderKey<'a> {
    fn to_string(&self) -> String {
        match *self {
            OrderKey::Count   => "_count".to_owned(),
            OrderKey::Term    => "_term".to_owned(),
            OrderKey::Expr(e) => e.to_owned()
        }
    }
}

/// Used to define the ordering of buckets in a some bucketted aggregations
///
/// # Examples
///
/// ```
/// use rs_es::operations::search::aggregations::{Order, OrderKey};
///
/// let order1 = Order::asc(OrderKey::Count);
/// let order2 = Order::desc("field_name");
/// ```
///
/// The first will produce a JSON fragment: `{"_count": "asc"}`; the second will
/// produce a JSON fragment: `{"field_name", "desc"}`
#[derive(Debug)]
pub struct Order<'a>(OrderKey<'a>, super::Order);

impl<'a> Order<'a> {
    /// Create an `Order` ascending
    pub fn asc<O: Into<OrderKey<'a>>>(key: O) -> Order<'a> {
        Order(key.into(), super::Order::Asc)
    }

    /// Create an `Order` descending
    pub fn desc<O: Into<OrderKey<'a>>>(key: O) -> Order<'a> {
        Order(key.into(), super::Order::Desc)
    }
}

impl<'a> ToJson for Order<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        d.insert(self.0.to_string(), self.1.to_json());
        Json::Object(d)
    }
}

/// Add to JSON trait
trait AddToJson {
    fn add_to_json(&self, &mut BTreeMap<String, Json>);
}

/// Terms aggregation
#[derive(Debug)]
pub struct Terms<'a> {
    field:      &'a str,
    size:       Option<u64>,
    shard_size: Option<u64>,
    order:      Option<Order<'a>>
}

impl<'a> Terms<'a> {
    pub fn new(field: &'a str) -> Terms<'a> {
        Terms {
            field:      field,
            size:       None,
            shard_size: None,
            order:      None
        }
    }

    add_field!(with_size, size, u64);
    add_field!(with_shard_size, shard_size, u64);
    add_field!(with_order, order, Order<'a>);
}

impl<'a> From<(Terms<'a>, Aggregations<'a>)> for Aggregation<'a> {
    fn from(from: (Terms<'a>, Aggregations<'a>)) -> Aggregation<'a> {
        Aggregation::Bucket(BucketAggregation::Terms(from.0), Some(from.1))
    }
}

impl<'a> From<Terms<'a>> for Aggregation<'a> {
    fn from(from: Terms<'a>) -> Aggregation<'a> {
        Aggregation::Bucket(BucketAggregation::Terms(from), None)
    }
}

impl<'a> ToJson for Terms<'a> {
    fn to_json(&self) -> Json {
        let mut json = BTreeMap::new();
        json.insert("field".to_owned(), Json::String(self.field.to_owned()));
        optional_add!(json, self.size, "size");
        optional_add!(json, self.shard_size, "shard_size");
        optional_add!(json, self.order, "order");
        Json::Object(json)
    }
}

/// The set of bucket aggregations
#[derive(Debug)]
pub enum BucketAggregation<'a> {
    Terms(Terms<'a>)
}

impl<'a> AddToJson for BucketAggregation<'a> {
    fn add_to_json(&self, json: &mut BTreeMap<String, Json>) {
        match self {
            &BucketAggregation::Terms(ref terms) => {
                json.insert("terms".to_owned(), terms.to_json());
            }
        }
    }
}

/// Aggregations are either metrics or bucket-based aggregations
#[derive(Debug)]
pub enum Aggregation<'a> {
    /// A metric aggregation (e.g. min)
    Metrics(MetricsAggregation<'a>),

    /// A bucket aggregation, groups data into buckets and optionally applies
    /// sub-aggregations
    Bucket(BucketAggregation<'a>, Option<Aggregations<'a>>)
}

impl<'a> ToJson for Aggregation<'a> {
    fn to_json(&self) -> Json {
        match self {
            &Aggregation::Metrics(ref ma)          => {
                ma.to_json()
            },
            &Aggregation::Bucket(ref ba, ref aggs) => {
                let mut d = BTreeMap::new();
                ba.add_to_json(&mut d);
                match aggs {
                    &Some(ref a) => {
                        d.insert("aggs".to_owned(), a.to_json());
                    },
                    &None        => ()
                }
                Json::Object(d)
            }
        }
    }
}

/// The set of aggregations
///
/// There are many ways of creating aggregations, either standalone or via a
/// conversion trait
#[derive(Debug)]
pub struct Aggregations<'a>(HashMap<&'a str, Aggregation<'a>>);

impl<'a> Aggregations<'a> {
    /// Create an empty-set of aggregations, individual aggregations should be
    /// added via the `add` method
    ///
    /// # Examples
    ///
    /// ```
    /// use rs_es::operations::search::aggregations::{Aggregations, Min};
    ///
    /// let mut aggs = Aggregations::new();
    /// aggs.add("agg_name", Min::Field("field_name"));
    /// ```
    pub fn new() -> Aggregations<'a> {
        Aggregations(HashMap::new())
    }

    /// Add an aggregation to the set of aggregations
    pub fn add<A: Into<Aggregation<'a>>>(&mut self, key: &'a str, val: A) {
        self.0.insert(key, val.into());
    }
}

impl<'b> From<Vec<(&'b str, Aggregation<'b>)>> for Aggregations<'b> {
    fn from(from: Vec<(&'b str, Aggregation<'b>)>) -> Aggregations<'b> {
        let mut aggs = Aggregations::new();
        for (name, agg) in from {
            aggs.add(name, agg);
        }
        aggs
    }
}

impl <'a, A: Into<Aggregation<'a>>> From<(&'a str, A)> for Aggregations<'a> {
    fn from(from: (&'a str, A)) -> Aggregations<'a> {
        let mut aggs = Aggregations::new();
        aggs.add(from.0, from.1.into());
        aggs
    }
}

impl<'a> ToJson for Aggregations<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        for (k, ref v) in self.0.iter() {
            d.insert((*k).to_owned(), v.to_json());
        }
        Json::Object(d)
    }
}

// Result objects

// Metrics result

#[derive(Debug)]
pub struct MinResult {
    pub value: JsonVal
}

impl<'a> From<&'a Json> for MinResult {
    fn from(from: &'a Json) -> MinResult {
        MinResult {
            value: JsonVal::from(from.find("value").expect("No 'value' value"))
        }
    }
}

// Buckets result

/// Macros for buckets to return a reference to the sub-aggregations
macro_rules! add_aggs_ref {
    () => {
        pub fn aggs_ref<'a>(&'a self) -> Option<&'a AggregationsResult> {
            self.aggs.as_ref()
        }
    }
}

#[derive(Debug)]
pub struct TermsBucketResult {
    pub key: JsonVal,
    pub doc_count: u64,
    pub aggs: Option<AggregationsResult>
}

impl TermsBucketResult {
    fn from(from: &Json, aggs: &Option<Aggregations>) -> TermsBucketResult {
        info!("Creating TermsBucketResult from: {:?} with {:?}", from, aggs);

        TermsBucketResult {
            key: JsonVal::from(from.find("key").expect("No 'key' value")),
            doc_count: get_json_u64!(from, "doc_count"),
            aggs: match aggs {
                &Some(ref agg) => {
                    Some(object_to_result(agg, from.as_object().expect("Not an object")))
                },
                &None          => None
            }
        }
    }

    add_aggs_ref!();
}

#[derive(Debug)]
pub struct TermsResult {
    pub doc_count_error_upper_bound: u64,
    pub sum_other_doc_count: u64,
    pub buckets: Vec<TermsBucketResult>
}

impl TermsResult {
    fn from(from: &Json, aggs: &Option<Aggregations>) -> TermsResult {
        TermsResult {
            doc_count_error_upper_bound: get_json_u64!(from, "doc_count_error_upper_bound"),
            sum_other_doc_count: get_json_u64!(from, "sum_other_doc_count"),
            buckets: from.find("buckets").expect("No buckets")
                .as_array().expect("Not an array")
                .iter().map(|bucket| {
                    TermsBucketResult::from(bucket, aggs)
                }).collect()
        }
    }
}

/// The result of one specific aggregation
///
/// The data returned varies depending on aggregation type
#[derive(Debug)]
pub enum AggregationResult {
    // Metrics
    Min(MinResult),

    // Buckets
    Terms(TermsResult)
}

/// Macro to implement the various as... functions that return the details of an
/// aggregation for that particular type
macro_rules! agg_as {
    ($n:ident,$t:ident,$rt:ty) => {
        pub fn $n<'a>(&'a self) -> Result<&'a $rt, EsError> {
            match self {
                &AggregationResult::$t(ref res) => Ok(res),
                _                               => {
                    Err(EsError::EsError(format!("Wrong type: {:?}", self)))
                }
            }
        }
    }
}

impl AggregationResult {
    // Metrics
    agg_as!(as_min, Min, MinResult);

    // buckets
    agg_as!(as_terms, Terms, TermsResult);
}

#[derive(Debug)]
pub struct AggregationsResult(HashMap<String, AggregationResult>);

/// Loads a Json object of aggregation results into an `AggregationsResult`.
fn object_to_result(aggs: &Aggregations, object: &BTreeMap<String, Json>) -> AggregationsResult {
    let mut ar_map = HashMap::new();

    for (key, val) in aggs.0.iter() {
        let owned_key = (*key).to_owned();
        let json = object.get(&owned_key).expect(&format!("No key: {}", &owned_key));
        ar_map.insert(owned_key, match val {
            &Aggregation::Metrics(ref ma) => {
                match ma {
                    &MetricsAggregation::Min(_) => {
                        AggregationResult::Min(MinResult::from(json))
                    }
                }
            },
            &Aggregation::Bucket(ref ba, ref aggs) => {
                match ba {
                    &BucketAggregation::Terms(_) => {
                        AggregationResult::Terms(TermsResult::from(json, aggs))
                    }
                }
            }
        });
    }

    info!("Processed aggs - From: {:?}. To: {:?}", object, ar_map);

    AggregationsResult(ar_map)
}

impl AggregationsResult {
    pub fn get<'a>(&'a self, key: &str) -> Result<&'a AggregationResult, EsError> {
        match self.0.get(key) {
            Some(ref agg_res) => Ok(agg_res),
            None              => Err(EsError::EsError(format!("No agg for key: {}",
                                                              key)))
        }
    }

    pub fn from(aggs: &Aggregations, json: &Json) -> AggregationsResult {
        let object = json.find("aggregations")
            .expect("No aggregations")
            .as_object()
            .expect("No aggregations");

        object_to_result(aggs, object)
    }
}
<commit_msg>DRYing conversion functions<commit_after>/*
 * Copyright 2015 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Implementation of ElasticSearch [aggregations](https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations.html)

use std::collections::{BTreeMap, HashMap};

use rustc_serialize::json::{Json, ToJson};

use error::EsError;
use units::JsonVal;

/// Script attributes for various attributes
#[derive(Debug)]
pub struct Script<'a> {
    field:  &'a str,
    script: &'a str,
    params: Option<Json>
}

impl<'a> Script<'a> {
    fn add_to_object(&self, obj: &mut BTreeMap<String, Json>) {
        obj.insert("field".to_owned(), self.field.to_json());
        obj.insert("script".to_owned(), self.script.to_json());
        match self.params {
            Some(ref json) => {
                obj.insert("params".to_owned(), json.clone());
            },
            None           => ()
        }
    }

    pub fn with_params(mut self, params: Json) -> Self {
        self.params = Some(params);
        self
    }
}

/// Min aggregation
#[derive(Debug)]
pub enum Min<'a> {
    /// Field
    Field(&'a str),

    /// By Script
    Script(Script<'a>)
}

macro_rules! metric_agg {
    ($b:ident) => {
        impl<'a> From<$b<'a>> for Aggregation<'a> {
            fn from(from: $b<'a>) -> Aggregation<'a> {
                Aggregation::Metrics(MetricsAggregation::$b(from))
            }
        }
    }
}

metric_agg!(Min);

impl<'a> ToJson for Min<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        match self {
            &Min::Field(field) => {
                d.insert("field".to_string(), field.to_json());
            },
            &Min::Script(ref script) => {
                script.add_to_object(&mut d);
            }
        }
        Json::Object(d)
    }
}

/// Individual aggregations and their options
#[derive(Debug)]
pub enum MetricsAggregation<'a> {
    Min(Min<'a>)
}

impl<'a> ToJson for MetricsAggregation<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        match self {
            &MetricsAggregation::Min(ref min_agg) => {
                d.insert("min".to_owned(), min_agg.to_json());
            }
        }
        Json::Object(d)
    }
}

/// Order - used for some bucketing aggregations to determine the order of
/// buckets
#[derive(Debug)]
pub enum OrderKey<'a> {
    Count,
    Term,
    Expr(&'a str)
}

impl<'a> From<&'a str> for OrderKey<'a> {
    fn from(from: &'a str) -> OrderKey<'a> {
        OrderKey::Expr(from)
    }
}

impl<'a> ToString for OrderKey<'a> {
    fn to_string(&self) -> String {
        match *self {
            OrderKey::Count   => "_count".to_owned(),
            OrderKey::Term    => "_term".to_owned(),
            OrderKey::Expr(e) => e.to_owned()
        }
    }
}

/// Used to define the ordering of buckets in a some bucketted aggregations
///
/// # Examples
///
/// ```
/// use rs_es::operations::search::aggregations::{Order, OrderKey};
///
/// let order1 = Order::asc(OrderKey::Count);
/// let order2 = Order::desc("field_name");
/// ```
///
/// The first will produce a JSON fragment: `{"_count": "asc"}`; the second will
/// produce a JSON fragment: `{"field_name", "desc"}`
#[derive(Debug)]
pub struct Order<'a>(OrderKey<'a>, super::Order);

impl<'a> Order<'a> {
    /// Create an `Order` ascending
    pub fn asc<O: Into<OrderKey<'a>>>(key: O) -> Order<'a> {
        Order(key.into(), super::Order::Asc)
    }

    /// Create an `Order` descending
    pub fn desc<O: Into<OrderKey<'a>>>(key: O) -> Order<'a> {
        Order(key.into(), super::Order::Desc)
    }
}

impl<'a> ToJson for Order<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        d.insert(self.0.to_string(), self.1.to_json());
        Json::Object(d)
    }
}

/// Add to JSON trait
trait AddToJson {
    fn add_to_json(&self, &mut BTreeMap<String, Json>);
}

/// Terms aggregation
#[derive(Debug)]
pub struct Terms<'a> {
    field:      &'a str,
    size:       Option<u64>,
    shard_size: Option<u64>,
    order:      Option<Order<'a>>
}

impl<'a> Terms<'a> {
    pub fn new(field: &'a str) -> Terms<'a> {
        Terms {
            field:      field,
            size:       None,
            shard_size: None,
            order:      None
        }
    }

    add_field!(with_size, size, u64);
    add_field!(with_shard_size, shard_size, u64);
    add_field!(with_order, order, Order<'a>);
}

macro_rules! bucket_agg {
    ($b:ident) => {
        impl<'a> From<($b<'a>, Aggregations<'a>)> for Aggregation<'a> {
            fn from(from: ($b<'a>, Aggregations<'a>)) -> Aggregation<'a> {
                Aggregation::Bucket(BucketAggregation::$b(from.0), Some(from.1))
            }
        }

        impl<'a> From<$b<'a>> for Aggregation<'a> {
            fn from(from: $b<'a>) -> Aggregation<'a> {
                Aggregation::Bucket(BucketAggregation::$b(from), None)
            }
        }
    }
}

bucket_agg!(Terms);

impl<'a> ToJson for Terms<'a> {
    fn to_json(&self) -> Json {
        let mut json = BTreeMap::new();
        json.insert("field".to_owned(), Json::String(self.field.to_owned()));
        optional_add!(json, self.size, "size");
        optional_add!(json, self.shard_size, "shard_size");
        optional_add!(json, self.order, "order");
        Json::Object(json)
    }
}

/// The set of bucket aggregations
#[derive(Debug)]
pub enum BucketAggregation<'a> {
    Terms(Terms<'a>)
}

impl<'a> AddToJson for BucketAggregation<'a> {
    fn add_to_json(&self, json: &mut BTreeMap<String, Json>) {
        match self {
            &BucketAggregation::Terms(ref terms) => {
                json.insert("terms".to_owned(), terms.to_json());
            }
        }
    }
}

/// Aggregations are either metrics or bucket-based aggregations
#[derive(Debug)]
pub enum Aggregation<'a> {
    /// A metric aggregation (e.g. min)
    Metrics(MetricsAggregation<'a>),

    /// A bucket aggregation, groups data into buckets and optionally applies
    /// sub-aggregations
    Bucket(BucketAggregation<'a>, Option<Aggregations<'a>>)
}

impl<'a> ToJson for Aggregation<'a> {
    fn to_json(&self) -> Json {
        match self {
            &Aggregation::Metrics(ref ma)          => {
                ma.to_json()
            },
            &Aggregation::Bucket(ref ba, ref aggs) => {
                let mut d = BTreeMap::new();
                ba.add_to_json(&mut d);
                match aggs {
                    &Some(ref a) => {
                        d.insert("aggs".to_owned(), a.to_json());
                    },
                    &None        => ()
                }
                Json::Object(d)
            }
        }
    }
}

/// The set of aggregations
///
/// There are many ways of creating aggregations, either standalone or via a
/// conversion trait
#[derive(Debug)]
pub struct Aggregations<'a>(HashMap<&'a str, Aggregation<'a>>);

impl<'a> Aggregations<'a> {
    /// Create an empty-set of aggregations, individual aggregations should be
    /// added via the `add` method
    ///
    /// # Examples
    ///
    /// ```
    /// use rs_es::operations::search::aggregations::{Aggregations, Min};
    ///
    /// let mut aggs = Aggregations::new();
    /// aggs.add("agg_name", Min::Field("field_name"));
    /// ```
    pub fn new() -> Aggregations<'a> {
        Aggregations(HashMap::new())
    }

    /// Add an aggregation to the set of aggregations
    pub fn add<A: Into<Aggregation<'a>>>(&mut self, key: &'a str, val: A) {
        self.0.insert(key, val.into());
    }
}

impl<'b> From<Vec<(&'b str, Aggregation<'b>)>> for Aggregations<'b> {
    fn from(from: Vec<(&'b str, Aggregation<'b>)>) -> Aggregations<'b> {
        let mut aggs = Aggregations::new();
        for (name, agg) in from {
            aggs.add(name, agg);
        }
        aggs
    }
}

impl <'a, A: Into<Aggregation<'a>>> From<(&'a str, A)> for Aggregations<'a> {
    fn from(from: (&'a str, A)) -> Aggregations<'a> {
        let mut aggs = Aggregations::new();
        aggs.add(from.0, from.1.into());
        aggs
    }
}

impl<'a> ToJson for Aggregations<'a> {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        for (k, ref v) in self.0.iter() {
            d.insert((*k).to_owned(), v.to_json());
        }
        Json::Object(d)
    }
}

// Result objects

// Metrics result

#[derive(Debug)]
pub struct MinResult {
    pub value: JsonVal
}

impl<'a> From<&'a Json> for MinResult {
    fn from(from: &'a Json) -> MinResult {
        MinResult {
            value: JsonVal::from(from.find("value").expect("No 'value' value"))
        }
    }
}

// Buckets result

/// Macros for buckets to return a reference to the sub-aggregations
macro_rules! add_aggs_ref {
    () => {
        pub fn aggs_ref<'a>(&'a self) -> Option<&'a AggregationsResult> {
            self.aggs.as_ref()
        }
    }
}

#[derive(Debug)]
pub struct TermsBucketResult {
    pub key: JsonVal,
    pub doc_count: u64,
    pub aggs: Option<AggregationsResult>
}

impl TermsBucketResult {
    fn from(from: &Json, aggs: &Option<Aggregations>) -> TermsBucketResult {
        info!("Creating TermsBucketResult from: {:?} with {:?}", from, aggs);

        TermsBucketResult {
            key: JsonVal::from(from.find("key").expect("No 'key' value")),
            doc_count: get_json_u64!(from, "doc_count"),
            aggs: match aggs {
                &Some(ref agg) => {
                    Some(object_to_result(agg, from.as_object().expect("Not an object")))
                },
                &None          => None
            }
        }
    }

    add_aggs_ref!();
}

#[derive(Debug)]
pub struct TermsResult {
    pub doc_count_error_upper_bound: u64,
    pub sum_other_doc_count: u64,
    pub buckets: Vec<TermsBucketResult>
}

impl TermsResult {
    fn from(from: &Json, aggs: &Option<Aggregations>) -> TermsResult {
        TermsResult {
            doc_count_error_upper_bound: get_json_u64!(from, "doc_count_error_upper_bound"),
            sum_other_doc_count: get_json_u64!(from, "sum_other_doc_count"),
            buckets: from.find("buckets").expect("No buckets")
                .as_array().expect("Not an array")
                .iter().map(|bucket| {
                    TermsBucketResult::from(bucket, aggs)
                }).collect()
        }
    }
}

/// The result of one specific aggregation
///
/// The data returned varies depending on aggregation type
#[derive(Debug)]
pub enum AggregationResult {
    // Metrics
    Min(MinResult),

    // Buckets
    Terms(TermsResult)
}

/// Macro to implement the various as... functions that return the details of an
/// aggregation for that particular type
macro_rules! agg_as {
    ($n:ident,$t:ident,$rt:ty) => {
        pub fn $n<'a>(&'a self) -> Result<&'a $rt, EsError> {
            match self {
                &AggregationResult::$t(ref res) => Ok(res),
                _                               => {
                    Err(EsError::EsError(format!("Wrong type: {:?}", self)))
                }
            }
        }
    }
}

impl AggregationResult {
    // Metrics
    agg_as!(as_min, Min, MinResult);

    // buckets
    agg_as!(as_terms, Terms, TermsResult);
}

#[derive(Debug)]
pub struct AggregationsResult(HashMap<String, AggregationResult>);

/// Loads a Json object of aggregation results into an `AggregationsResult`.
fn object_to_result(aggs: &Aggregations, object: &BTreeMap<String, Json>) -> AggregationsResult {
    let mut ar_map = HashMap::new();

    for (key, val) in aggs.0.iter() {
        let owned_key = (*key).to_owned();
        let json = object.get(&owned_key).expect(&format!("No key: {}", &owned_key));
        ar_map.insert(owned_key, match val {
            &Aggregation::Metrics(ref ma) => {
                match ma {
                    &MetricsAggregation::Min(_) => {
                        AggregationResult::Min(MinResult::from(json))
                    }
                }
            },
            &Aggregation::Bucket(ref ba, ref aggs) => {
                match ba {
                    &BucketAggregation::Terms(_) => {
                        AggregationResult::Terms(TermsResult::from(json, aggs))
                    }
                }
            }
        });
    }

    info!("Processed aggs - From: {:?}. To: {:?}", object, ar_map);

    AggregationsResult(ar_map)
}

impl AggregationsResult {
    pub fn get<'a>(&'a self, key: &str) -> Result<&'a AggregationResult, EsError> {
        match self.0.get(key) {
            Some(ref agg_res) => Ok(agg_res),
            None              => Err(EsError::EsError(format!("No agg for key: {}",
                                                              key)))
        }
    }

    pub fn from(aggs: &Aggregations, json: &Json) -> AggregationsResult {
        let object = json.find("aggregations")
            .expect("No aggregations")
            .as_object()
            .expect("No aggregations");

        object_to_result(aggs, object)
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use rustyline;
use rustyline::error::ReadlineError as RLError;
use combine;
use error::*;

pub struct Readline {
    rl: rustyline::Editor<()>,
    prompt: String,
    buffer: Vec<u8>,
}

impl Readline {
    fn new<S>(prompt: &str) -> Self {
        Readline {
            rl: rustyline::Editor::new(),
            prompt: prompt.to_owned(),
            buffer: Vec::with_capacity(128),
        }
    }
}

impl Read for Readline {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.write(line.as_bytes());
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Interrupted) => return Err(io::ErrorKind::Interrupted.into()),
                Err(RLError::Eof) => return Ok(0),
                Err(RLError::Errno(num)) => return Err(io::Error::from_raw_os_error(num)),
                Err(err) => return Err(io::ErrorKind::Other.into()),
            }
        }

        // if buf > buffer, drain entire buffer
        if buf.len() >= self.buffer.len() {
            let result = buf.write(&self.buffer);
            self.buffer.clear();
            result
        } else {
            let drained = self.buffer.drain(0..buf.len()).collect::<Vec<_>>();
            buf.write(&drained)
        }
    }
}

impl io::BufRead for Readline {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        unimplemented!()
    }

    fn consume(&mut self, amt: usize) {
        unimplemented!()
    }
}

#[derive(Clone, Debug)]
pub struct StringStream<R: io::Read> {
    src: R,
    buffer: String,
}

impl<R: io::Read> StringStream<R> {
    fn new(src: R) -> Self {
        StringStream {
            src: src,
            buffer: String::with_capacity(128),
        }
    }
}

impl<R: io::Read> combine::StreamOnce for StringStream<R> {
    type Item = char;
    type Range = str;
    type Position usize;
    fn uncons(&mut self) -> Result<Self::Item, Error<Self::Item, Self::Range>>;
    fn position(&self) -> Self::Position;
}
<commit_msg>Create LineStream to implement StreamOnce for BufStream sources<commit_after>use std::io;
use std::io::prelude::*;
use std::collections::VecDeque;
use rustyline;
use rustyline::error::ReadlineError as RLError;
use combine;
use combine::primitives::SourcePosition;

pub struct Readline {
    rl: rustyline::Editor<()>,
    prompt: String,
    buffer: Vec<u8>,
}

impl Readline {
    fn new<S>(prompt: &str) -> Self {
        Readline {
            rl: rustyline::Editor::new(),
            prompt: prompt.to_owned(),
            buffer: Vec::with_capacity(128),
        }
    }
}

impl Read for Readline {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.write(line.as_bytes());
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Interrupted) => return Err(io::ErrorKind::Interrupted.into()),
                Err(RLError::Eof) => return Ok(0),
                #[cfg(unix)]
                Err(RLError::Errno(num)) => return Err(io::Error::from_raw_os_error(num)),
                Err(err) => return Err(io::ErrorKind::Other.into()),
            }
        }

        // if buf > buffer, drain entire buffer
        if buf.len() >= self.buffer.len() {
            let result = buf.write(&self.buffer);
            self.buffer.clear();
            result
        } else {
            let drained = self.buffer.drain(0..buf.len()).collect::<Vec<_>>();
            buf.write(&drained)
        }
    }
}

impl io::BufRead for Readline {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        unimplemented!()
    }

    fn consume(&mut self, amt: usize) {
        unimplemented!()
    }
}

#[derive(Debug)]
pub struct LineStream<B: io::BufRead> {
    lines: io::Lines<B>,
    buffer: VecDeque<char>,
    position: SourcePosition,
}

impl<B: io::BufRead> LineStream<B> {
    fn new(src: B) -> Self {
        LineStream {
            lines: src.lines(),
            buffer: VecDeque::with_capacity(128),
            position: SourcePosition { line: 0, column: 0 },
        }
    }
}

impl<B: io::BufRead> combine::StreamOnce for LineStream<B> {
    type Item = char;
    type Range = String;
    type Position = combine::primitives::SourcePosition;

    fn uncons(&mut self) -> Result<Self::Item, combine::primitives::Error<Self::Item, Self::Range>> {
        if self.buffer.is_empty() {
            match self.lines.next() {
                Some(result) => {
                    match result {
                        Ok(line) => self.buffer.extend(line.chars()),
                        Err(err) => return Err(combine::primitives::Error::Other(Box::new(err)))
                    };
                },
                None => return Err(combine::primitives::Error::end_of_input()),
            }
            self.position.line += 1;
            self.position.column = 0;
        }
        let item = self.buffer.pop_front().unwrap();
        self.position.column += 1;
        Ok(item)
    }

    fn position(&self) -> Self::Position {
        self.position
    }
}
<|endoftext|>


<commit_before>extern crate typedb;
extern crate serde;
#[macro_use]
extern crate serde_derive;

use typedb::KV;

#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
enum MyKey {
    String(String),
    Int(i32),
}

#[derive(Clone, Serialize, Deserialize, Debug)]
enum MyValue {
    String(String),
    Int(i32),
}

fn main() {
    let cab_path = "./db.cab";
    let mut test_store = KV::<MyKey, MyValue>::new(cab_path).unwrap();

    let _ = test_store.insert(MyKey::Int(1i32), MyValue::String("value".to_string()));
    println!("{:?}", test_store.get(MyKey::Int(1i32)));
    let _ = test_store.remove(MyKey::Int(1i32));

    // clean up the cab
    let _ = std::fs::remove_file(cab_path);
}
<commit_msg>Update the custom_types example to use the key git diff and value git diff macros<commit_after>#[macro_use]
extern crate typedb;
#[macro_use]
extern crate serde_derive;

use typedb::KV;

key!(MyKey:
    String(String),
    Int(i32),
);

value!(MyValue:
    String(String),
    Int(i32),
);

fn main() {
    let cab_path = "./db.cab";
    let mut test_store = KV::<MyKey, MyValue>::new(cab_path).unwrap();

    let _ = test_store.insert(MyKey::Int(1i32), MyValue::String("value".to_string()));
    println!("{:?}", test_store.get(MyKey::Int(1i32)));
    let _ = test_store.remove(MyKey::Int(1i32));

    // clean up the cab
    let _ = std::fs::remove_file(cab_path);
}
<|endoftext|>


<commit_before>//! ARMv8 intrinsics.
//!
//! The reference is [ARMv8-A Reference Manual](http://infocenter.arm.com/help/index.jsp?topic=/com.arm.doc.ddi0487a.k_10775/index.html).

pub use super::v7::*;

#[cfg(test)]
use stdsimd_test::assert_instr;

/// Reverse the order of the bytes.
#[inline(always)]
#[cfg_attr(test, assert_instr(rev))]
pub unsafe fn _rev_u64(x: u64) -> u64 {
    x.swap_bytes() as u64
}

/// Count Leading Zeros.
#[inline(always)]
#[cfg_attr(test, assert_instr(clz))]
pub unsafe fn _clz_u64(x: u64) -> u64 {
    x.leading_zeros() as u64
}

#[allow(dead_code)]
extern "C" {
    #[link_name="llvm.bitreverse.i64"]
    fn rbit_u64(i: i64) -> i64;
}

/// Reverse the bit order.
#[inline(always)]
#[cfg_attr(test, assert_instr(rbit))]
pub unsafe fn _rbit_u64(x: u64) -> u64 {
    unsafe { rbit_u64(x as i64) as u64 }
}

/// Counts the leading most significant bits set.
///
/// When all bits of the operand are set it returns the size of the operand in
/// bits.
#[inline(always)]
#[cfg_attr(test, assert_instr(cls))]
pub unsafe fn _cls_u32(x: u32) -> u32 {
    u32::leading_zeros(((((((x as i32) >> 31) as u32) ^ x) << 1) | 1)) as u32
}

/// Counts the leading most significant bits set.
///
/// When all bits of the operand are set it returns the size of the operand in
/// bits.
#[inline(always)]
#[cfg_attr(test, assert_instr(cls))]
pub unsafe fn _cls_u64(x: u64) -> u64 {
    u64::leading_zeros(((((((x as i64) >> 63) as u64) ^ x) << 1) | 1)) as u64
}
<commit_msg>[arm] fix unused unsafe warning<commit_after>//! ARMv8 intrinsics.
//!
//! The reference is [ARMv8-A Reference Manual](http://infocenter.arm.com/help/index.jsp?topic=/com.arm.doc.ddi0487a.k_10775/index.html).

pub use super::v7::*;

#[cfg(test)]
use stdsimd_test::assert_instr;

/// Reverse the order of the bytes.
#[inline(always)]
#[cfg_attr(test, assert_instr(rev))]
pub unsafe fn _rev_u64(x: u64) -> u64 {
    x.swap_bytes() as u64
}

/// Count Leading Zeros.
#[inline(always)]
#[cfg_attr(test, assert_instr(clz))]
pub unsafe fn _clz_u64(x: u64) -> u64 {
    x.leading_zeros() as u64
}

#[allow(dead_code)]
extern "C" {
    #[link_name="llvm.bitreverse.i64"]
    fn rbit_u64(i: i64) -> i64;
}

/// Reverse the bit order.
#[inline(always)]
#[cfg_attr(test, assert_instr(rbit))]
pub unsafe fn _rbit_u64(x: u64) -> u64 {
    rbit_u64(x as i64) as u64
}

/// Counts the leading most significant bits set.
///
/// When all bits of the operand are set it returns the size of the operand in
/// bits.
#[inline(always)]
#[cfg_attr(test, assert_instr(cls))]
pub unsafe fn _cls_u32(x: u32) -> u32 {
    u32::leading_zeros(((((((x as i32) >> 31) as u32) ^ x) << 1) | 1)) as u32
}

/// Counts the leading most significant bits set.
///
/// When all bits of the operand are set it returns the size of the operand in
/// bits.
#[inline(always)]
#[cfg_attr(test, assert_instr(cls))]
pub unsafe fn _cls_u64(x: u64) -> u64 {
    u64::leading_zeros(((((((x as i64) >> 63) as u64) ^ x) << 1) | 1)) as u64
}
<|endoftext|>