<commit_before>// Copyright (c) 2020-2021, 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 num_derive::*;

use crate::partition::BlockSize;
use crate::serialize::{Deserialize, Serialize};

use std::fmt;

// NOTE: Add Structures at the end.
/// Contains the speed settings.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct SpeedSettings {
  /// Enables inter-frames to have multiple reference frames.
  ///
  /// Enabled is slower.
  pub multiref: bool,

  /// Enables fast deblocking filter.
  pub fast_deblock: bool,

  /// Enables reduced transform set.
  ///
  /// Enabled is faster.
  pub reduced_tx_set: bool,

  /// Enables using transform-domain distortion instead of pixel-domain.
  ///
  /// Enabled is faster.
  pub tx_domain_distortion: bool,

  /// Enables using transform-domain rate estimation.
  ///
  /// Enabled is faster.
  pub tx_domain_rate: bool,

  /// Enables bottom-up encoding, rather than top-down.
  ///
  /// Enabled is slower.
  pub encode_bottomup: bool,

  /// The number of lookahead frames to be used for temporal RDO.
  ///
  /// Higher is slower.
  pub rdo_lookahead_frames: usize,

  /// Enables searching transform size and type with RDO.
  ///
  /// Enabled is slower.
  pub rdo_tx_decision: bool,

  /// Prediction modes to search.
  ///
  /// Complex settings are slower.
  pub prediction_modes: PredictionModesSetting,

  /// Enables searching near motion vectors during RDO.
  ///
  /// Enabled is slower.
  pub include_near_mvs: bool,

  /// Which scene detection mode to use. Standard is slower, but best.
  pub scene_detection_mode: SceneDetectionSpeed,

  /// Enables CDEF.
  pub cdef: bool,

  /// Enables LRF.
  pub lrf: bool,

  /// The amount of search done for self guided restoration.
  pub sgr_complexity: SGRComplexityLevel,

  /// Use SATD instead of SAD for subpixel search.
  ///
  /// Enabled is slower.
  pub use_satd_subpel: bool,

  /// Use non-square partition type everywhere
  ///
  /// Enabled is slower.
  pub non_square_partition: bool,

  /// Search level for segmentation.
  ///
  /// Full search is at least twice as slow.
  pub segmentation: SegmentationLevel,

  /// Enable tx split for inter mode block.
  pub enable_inter_tx_split: bool,

  /// Use fine directional intra prediction
  pub fine_directional_intra: bool,

  /// Enable full search in some parts of motion estimation. Allowing full
  /// search is slower.
  pub me_allow_full_search: bool,

  // NOTE: put enums and basic type fields above
  /// Range of partition sizes that can be used. Larger ranges are slower.
  ///
  /// Must be based on square block sizes, so e.g. 8×4 isn't allowed here.
  pub partition_range: PartitionRange,
}

impl Default for SpeedSettings {
  /// The default settings are equivalent to speed 0
  fn default() -> Self {
    SpeedSettings {
      partition_range: PartitionRange::new(
        BlockSize::BLOCK_4X4,
        BlockSize::BLOCK_64X64,
      ),
      multiref: true,
      fast_deblock: false,
      reduced_tx_set: false,
      // TX domain distortion is always faster, with no significant quality change,
      // although it will be ignored when Tune == Psychovisual.
      tx_domain_distortion: true,
      tx_domain_rate: false,
      encode_bottomup: true,
      rdo_tx_decision: true,
      rdo_lookahead_frames: 40,
      prediction_modes: PredictionModesSetting::ComplexAll,
      include_near_mvs: true,
      scene_detection_mode: SceneDetectionSpeed::Standard,
      cdef: true,
      lrf: true,
      sgr_complexity: SGRComplexityLevel::Full,
      use_satd_subpel: true,
      non_square_partition: true,
      segmentation: SegmentationLevel::Full,
      enable_inter_tx_split: false,
      fine_directional_intra: true,
      me_allow_full_search: true,
    }
  }
}

impl SpeedSettings {
  /// Set the speed setting according to a numeric speed preset.
  pub fn from_preset(speed: usize) -> Self {
    // The default settings are equivalent to speed 0
    let mut settings = SpeedSettings::default();

    if speed >= 1 {
      settings.segmentation = SegmentationLevel::Simple;
    }

    if speed >= 2 {
      settings.non_square_partition = false;
    }

    if speed >= 3 {
      settings.partition_range =
        PartitionRange::new(BlockSize::BLOCK_8X8, BlockSize::BLOCK_64X64);
      settings.rdo_lookahead_frames = 30;
      settings.prediction_modes = PredictionModesSetting::ComplexKeyframes;
    }

    if speed >= 4 {
      settings.encode_bottomup = false;
    }

    if speed >= 5 {
      settings.include_near_mvs = false;
      settings.sgr_complexity = SGRComplexityLevel::Reduced;
    }

    if speed >= 6 {
      settings.reduced_tx_set = true;
      settings.rdo_lookahead_frames = 20;
      settings.rdo_tx_decision = false;
      settings.me_allow_full_search = false;
    }

    if speed >= 7 {
      settings.prediction_modes = PredictionModesSetting::Simple;
    }

    if speed >= 8 {
      // Multiref is enabled automatically if low_latency is false.
      //
      // If low_latency is true, enabling multiref allows using multiple
      // backwards references. low_latency false enables both forward and
      // backwards references.
      settings.multiref = false;
      settings.fast_deblock = true;
    }

    if speed >= 9 {
      // 8x8 is fast enough to use until very high speed levels,
      // because 8x8 with reduced TX set is faster but with equivalent
      // or better quality compared to 16x16 (to which reduced TX set does not apply).
      settings.partition_range =
        PartitionRange::new(BlockSize::BLOCK_32X32, BlockSize::BLOCK_64X64);
      settings.rdo_lookahead_frames = 10;
      // FIXME: With unknown reasons, inter_tx_split does not work if reduced_tx_set is false
      settings.enable_inter_tx_split = true;
    }

    if speed >= 10 {
      settings.partition_range =
        PartitionRange::new(BlockSize::BLOCK_32X32, BlockSize::BLOCK_32X32);
      settings.scene_detection_mode = SceneDetectionSpeed::Fast;
      settings.lrf = false;
      settings.use_satd_subpel = false;
    }

    settings
  }
}

/// Range of block sizes to use.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct PartitionRange {
  pub(crate) min: BlockSize,
  pub(crate) max: BlockSize,
}

impl PartitionRange {
  /// Creates a new partition range with min and max partition sizes.
  pub fn new(min: BlockSize, max: BlockSize) -> Self {
    assert!(max >= min);
    // Topdown search checks the min block size for PARTITION_SPLIT only, so
    // the min block size must be square.
    assert!(min.is_sqr());
    // Rectangular max partition sizes have not been tested.
    assert!(max.is_sqr());

    Self { min, max }
  }
}

/// Prediction modes to search.
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum SceneDetectionSpeed {
  /// Fastest scene detection using pixel-wise comparison
  Fast,
  /// Scene detection using motion vectors and cost estimates
  Standard,
  /// Completely disable scene detection and only place keyframes
  /// at fixed intervals.
  None,
}

impl fmt::Display for SceneDetectionSpeed {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        SceneDetectionSpeed::Fast => "Fast",
        SceneDetectionSpeed::Standard => "Standard",
        SceneDetectionSpeed::None => "None",
      }
    )
  }
}

/// Prediction modes to search.
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum PredictionModesSetting {
  /// Only simple prediction modes.
  Simple,
  /// Search all prediction modes on key frames and simple modes on other
  /// frames.
  ComplexKeyframes,
  /// Search all prediction modes on all frames.
  ComplexAll,
}

impl fmt::Display for PredictionModesSetting {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        PredictionModesSetting::Simple => "Simple",
        PredictionModesSetting::ComplexKeyframes => "Complex-KFs",
        PredictionModesSetting::ComplexAll => "Complex-All",
      }
    )
  }
}

/// Search level for self guided restoration
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum SGRComplexityLevel {
  /// Search all sgr parameters
  Full,
  /// Search a reduced set of sgr parameters
  Reduced,
}

impl fmt::Display for SGRComplexityLevel {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        SGRComplexityLevel::Full => "Full",
        SGRComplexityLevel::Reduced => "Reduced",
      }
    )
  }
}

/// Search level for segmentation
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum SegmentationLevel {
  /// No segmentation is signalled.
  Disabled,
  /// Segmentation index is derived from source statistics.
  Simple,
  /// Search all segmentation indices.
  Full,
}

impl fmt::Display for SegmentationLevel {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        SegmentationLevel::Disabled => "Disabled",
        SegmentationLevel::Simple => "Simple",
        SegmentationLevel::Full => "Full",
      }
    )
  }
}
<commit_msg>Mark SpeedSettings struct as non-exhaustive<commit_after>// Copyright (c) 2020-2021, 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 num_derive::*;

use crate::partition::BlockSize;
use crate::serialize::{Deserialize, Serialize};

use std::fmt;

// NOTE: Add Structures at the end.
/// Contains the speed settings.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SpeedSettings {
  /// Enables inter-frames to have multiple reference frames.
  ///
  /// Enabled is slower.
  pub multiref: bool,

  /// Enables fast deblocking filter.
  pub fast_deblock: bool,

  /// Enables reduced transform set.
  ///
  /// Enabled is faster.
  pub reduced_tx_set: bool,

  /// Enables using transform-domain distortion instead of pixel-domain.
  ///
  /// Enabled is faster.
  pub tx_domain_distortion: bool,

  /// Enables using transform-domain rate estimation.
  ///
  /// Enabled is faster.
  pub tx_domain_rate: bool,

  /// Enables bottom-up encoding, rather than top-down.
  ///
  /// Enabled is slower.
  pub encode_bottomup: bool,

  /// The number of lookahead frames to be used for temporal RDO.
  ///
  /// Higher is slower.
  pub rdo_lookahead_frames: usize,

  /// Enables searching transform size and type with RDO.
  ///
  /// Enabled is slower.
  pub rdo_tx_decision: bool,

  /// Prediction modes to search.
  ///
  /// Complex settings are slower.
  pub prediction_modes: PredictionModesSetting,

  /// Enables searching near motion vectors during RDO.
  ///
  /// Enabled is slower.
  pub include_near_mvs: bool,

  /// Which scene detection mode to use. Standard is slower, but best.
  pub scene_detection_mode: SceneDetectionSpeed,

  /// Enables CDEF.
  pub cdef: bool,

  /// Enables LRF.
  pub lrf: bool,

  /// The amount of search done for self guided restoration.
  pub sgr_complexity: SGRComplexityLevel,

  /// Use SATD instead of SAD for subpixel search.
  ///
  /// Enabled is slower.
  pub use_satd_subpel: bool,

  /// Use non-square partition type everywhere
  ///
  /// Enabled is slower.
  pub non_square_partition: bool,

  /// Search level for segmentation.
  ///
  /// Full search is at least twice as slow.
  pub segmentation: SegmentationLevel,

  /// Enable tx split for inter mode block.
  pub enable_inter_tx_split: bool,

  /// Use fine directional intra prediction
  pub fine_directional_intra: bool,

  /// Enable full search in some parts of motion estimation. Allowing full
  /// search is slower.
  pub me_allow_full_search: bool,

  // NOTE: put enums and basic type fields above
  /// Range of partition sizes that can be used. Larger ranges are slower.
  ///
  /// Must be based on square block sizes, so e.g. 8×4 isn't allowed here.
  pub partition_range: PartitionRange,
}

impl Default for SpeedSettings {
  /// The default settings are equivalent to speed 0
  fn default() -> Self {
    SpeedSettings {
      partition_range: PartitionRange::new(
        BlockSize::BLOCK_4X4,
        BlockSize::BLOCK_64X64,
      ),
      multiref: true,
      fast_deblock: false,
      reduced_tx_set: false,
      // TX domain distortion is always faster, with no significant quality change,
      // although it will be ignored when Tune == Psychovisual.
      tx_domain_distortion: true,
      tx_domain_rate: false,
      encode_bottomup: true,
      rdo_tx_decision: true,
      rdo_lookahead_frames: 40,
      prediction_modes: PredictionModesSetting::ComplexAll,
      include_near_mvs: true,
      scene_detection_mode: SceneDetectionSpeed::Standard,
      cdef: true,
      lrf: true,
      sgr_complexity: SGRComplexityLevel::Full,
      use_satd_subpel: true,
      non_square_partition: true,
      segmentation: SegmentationLevel::Full,
      enable_inter_tx_split: false,
      fine_directional_intra: true,
      me_allow_full_search: true,
    }
  }
}

impl SpeedSettings {
  /// Set the speed setting according to a numeric speed preset.
  pub fn from_preset(speed: usize) -> Self {
    // The default settings are equivalent to speed 0
    let mut settings = SpeedSettings::default();

    if speed >= 1 {
      settings.segmentation = SegmentationLevel::Simple;
    }

    if speed >= 2 {
      settings.non_square_partition = false;
    }

    if speed >= 3 {
      settings.partition_range =
        PartitionRange::new(BlockSize::BLOCK_8X8, BlockSize::BLOCK_64X64);
      settings.rdo_lookahead_frames = 30;
      settings.prediction_modes = PredictionModesSetting::ComplexKeyframes;
    }

    if speed >= 4 {
      settings.encode_bottomup = false;
    }

    if speed >= 5 {
      settings.include_near_mvs = false;
      settings.sgr_complexity = SGRComplexityLevel::Reduced;
    }

    if speed >= 6 {
      settings.reduced_tx_set = true;
      settings.rdo_lookahead_frames = 20;
      settings.rdo_tx_decision = false;
      settings.me_allow_full_search = false;
    }

    if speed >= 7 {
      settings.prediction_modes = PredictionModesSetting::Simple;
    }

    if speed >= 8 {
      // Multiref is enabled automatically if low_latency is false.
      //
      // If low_latency is true, enabling multiref allows using multiple
      // backwards references. low_latency false enables both forward and
      // backwards references.
      settings.multiref = false;
      settings.fast_deblock = true;
    }

    if speed >= 9 {
      // 8x8 is fast enough to use until very high speed levels,
      // because 8x8 with reduced TX set is faster but with equivalent
      // or better quality compared to 16x16 (to which reduced TX set does not apply).
      settings.partition_range =
        PartitionRange::new(BlockSize::BLOCK_32X32, BlockSize::BLOCK_64X64);
      settings.rdo_lookahead_frames = 10;
      // FIXME: With unknown reasons, inter_tx_split does not work if reduced_tx_set is false
      settings.enable_inter_tx_split = true;
    }

    if speed >= 10 {
      settings.partition_range =
        PartitionRange::new(BlockSize::BLOCK_32X32, BlockSize::BLOCK_32X32);
      settings.scene_detection_mode = SceneDetectionSpeed::Fast;
      settings.lrf = false;
      settings.use_satd_subpel = false;
    }

    settings
  }
}

/// Range of block sizes to use.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct PartitionRange {
  pub(crate) min: BlockSize,
  pub(crate) max: BlockSize,
}

impl PartitionRange {
  /// Creates a new partition range with min and max partition sizes.
  pub fn new(min: BlockSize, max: BlockSize) -> Self {
    assert!(max >= min);
    // Topdown search checks the min block size for PARTITION_SPLIT only, so
    // the min block size must be square.
    assert!(min.is_sqr());
    // Rectangular max partition sizes have not been tested.
    assert!(max.is_sqr());

    Self { min, max }
  }
}

/// Prediction modes to search.
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum SceneDetectionSpeed {
  /// Fastest scene detection using pixel-wise comparison
  Fast,
  /// Scene detection using motion vectors and cost estimates
  Standard,
  /// Completely disable scene detection and only place keyframes
  /// at fixed intervals.
  None,
}

impl fmt::Display for SceneDetectionSpeed {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        SceneDetectionSpeed::Fast => "Fast",
        SceneDetectionSpeed::Standard => "Standard",
        SceneDetectionSpeed::None => "None",
      }
    )
  }
}

/// Prediction modes to search.
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum PredictionModesSetting {
  /// Only simple prediction modes.
  Simple,
  /// Search all prediction modes on key frames and simple modes on other
  /// frames.
  ComplexKeyframes,
  /// Search all prediction modes on all frames.
  ComplexAll,
}

impl fmt::Display for PredictionModesSetting {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        PredictionModesSetting::Simple => "Simple",
        PredictionModesSetting::ComplexKeyframes => "Complex-KFs",
        PredictionModesSetting::ComplexAll => "Complex-All",
      }
    )
  }
}

/// Search level for self guided restoration
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum SGRComplexityLevel {
  /// Search all sgr parameters
  Full,
  /// Search a reduced set of sgr parameters
  Reduced,
}

impl fmt::Display for SGRComplexityLevel {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        SGRComplexityLevel::Full => "Full",
        SGRComplexityLevel::Reduced => "Reduced",
      }
    )
  }
}

/// Search level for segmentation
#[derive(
  Clone,
  Copy,
  Debug,
  PartialOrd,
  PartialEq,
  FromPrimitive,
  Serialize,
  Deserialize,
)]
pub enum SegmentationLevel {
  /// No segmentation is signalled.
  Disabled,
  /// Segmentation index is derived from source statistics.
  Simple,
  /// Search all segmentation indices.
  Full,
}

impl fmt::Display for SegmentationLevel {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "{}",
      match self {
        SegmentationLevel::Disabled => "Disabled",
        SegmentationLevel::Simple => "Simple",
        SegmentationLevel::Full => "Full",
      }
    )
  }
}
<|endoftext|>


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

use crate::cpu_leaf::*;
#[cfg(target_arch = "x86_64")]
use kvm_bindings::CpuId;
#[cfg(target_arch = "x86")]
use std::arch::x86::{CpuidResult, __cpuid_count, __get_cpuid_max};
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::{CpuidResult, __cpuid_count, __get_cpuid_max};

/// Intel brand string.
pub const VENDOR_ID_INTEL: &[u8; 12] = b"GenuineIntel";
/// AMD brand string.
pub const VENDOR_ID_AMD: &[u8; 12] = b"AuthenticAMD";

/// cpuid related error.
#[derive(Clone, Debug)]
pub enum Error {
    /// The function was called with invalid parameters.
    InvalidParameters(String),
    /// Function not supported on the current architecture.
    NotSupported,
}

/// Extract entry from the cpuid.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_cpuid(function: u32, count: u32) -> Result<CpuidResult, Error> {
    // TODO: replace with validation based on `has_cpuid()` when it becomes stable:
    //  https://doc.rust-lang.org/core/arch/x86/fn.has_cpuid.html
    #[cfg(target_env = "sgx")]
    {
        return Err(Error::NotSupported);
    }
    // For x86 the host supports the `cpuid` instruction if SSE is enabled. Otherwise it's hard to check.
    #[cfg(target_arch = "x86")]
    {
        #[cfg(not(target_feature = "sse"))]
        {
            return Err(Error::NotSupported);
        }
    }

    // this is safe because the host supports the `cpuid` instruction
    let max_function = unsafe { __get_cpuid_max(function & leaf_0x80000000::LEAF_NUM).0 };
    if function > max_function {
        return Err(Error::InvalidParameters(format!(
            "Function not supported: 0x{:x}",
            function
        )));
    }

    // this is safe because the host supports the `cpuid` instruction
    let entry = unsafe { __cpuid_count(function, count) };
    if entry.eax == 0 && entry.ebx == 0 && entry.ecx == 0 && entry.edx == 0 {
        return Err(Error::InvalidParameters(format!(
            "Invalid count: {}",
            count
        )));
    }

    Ok(entry)
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_host() -> Result<[u8; 12], Error> {
    match get_cpuid(0, 0) {
        Ok(vendor_entry) => {
            let bytes: [u8; 12] = unsafe {
                std::mem::transmute([vendor_entry.ebx, vendor_entry.edx, vendor_entry.ecx])
            };
            Ok(bytes)
        }
        Err(e) => Err(e),
    }
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_cpuid(cpuid: &CpuId) -> Result<[u8; 12], Error> {
    // Search for vendor id entry.
    for entry in cpuid.as_slice().iter() {
        if entry.function == 0 && entry.index == 0 {
            let cpu_vendor_id: [u8; 12] =
                unsafe { std::mem::transmute([entry.ebx, entry.edx, entry.ecx]) };
            return Ok(cpu_vendor_id);
        }
    }

    Err(Error::NotSupported)
}

#[cfg(test)]
pub mod tests {
    use crate::common::*;

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub fn get_topoext_fn() -> u32 {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        let function = match &vendor_id.ok().unwrap() {
            VENDOR_ID_INTEL => leaf_0x4::LEAF_NUM,
            VENDOR_ID_AMD => leaf_0x8000001d::LEAF_NUM,
            _ => 0,
        };
        assert!(function != 0);

        function
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_cpu_id() {
        // get_cpu_id should work correctly here
        let topoext_fn = get_topoext_fn();

        // check that get_cpuid works for valid parameters
        match get_cpuid(topoext_fn, 0) {
            Ok(topoext_entry) => {
                assert!(topoext_entry.eax != 0);
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `function`
        match get_cpuid(0x9000_0000, 0) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Function not supported: 0x90000000");
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `count`
        match get_cpuid(topoext_fn, 100) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Invalid count: 100");
            }
            _ => panic!("Wrong behavior"),
        }
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_vendor_id() {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        matches!(&vendor_id.ok().unwrap(), VENDOR_ID_INTEL | VENDOR_ID_AMD);
    }
}
<commit_msg>cpuid: add function for identifying same cpu model<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::bit_helper::BitHelper;
use crate::cpu_leaf::*;
#[cfg(target_arch = "x86_64")]
use kvm_bindings::CpuId;
#[cfg(target_arch = "x86")]
use std::arch::x86::{CpuidResult, __cpuid_count, __get_cpuid_max};
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::{CpuidResult, __cpuid_count, __get_cpuid_max};

/// Intel brand string.
pub const VENDOR_ID_INTEL: &[u8; 12] = b"GenuineIntel";
/// AMD brand string.
pub const VENDOR_ID_AMD: &[u8; 12] = b"AuthenticAMD";

/// cpuid related error.
#[derive(Clone, Debug)]
pub enum Error {
    /// The function was called with invalid parameters.
    InvalidParameters(String),
    /// Function not supported on the current architecture.
    NotSupported,
}

/// Extract entry from the cpuid.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_cpuid(function: u32, count: u32) -> Result<CpuidResult, Error> {
    // TODO: replace with validation based on `has_cpuid()` when it becomes stable:
    //  https://doc.rust-lang.org/core/arch/x86/fn.has_cpuid.html
    #[cfg(target_env = "sgx")]
    {
        return Err(Error::NotSupported);
    }
    // For x86 the host supports the `cpuid` instruction if SSE is enabled. Otherwise it's hard to check.
    #[cfg(target_arch = "x86")]
    {
        #[cfg(not(target_feature = "sse"))]
        {
            return Err(Error::NotSupported);
        }
    }

    // this is safe because the host supports the `cpuid` instruction
    let max_function = unsafe { __get_cpuid_max(function & leaf_0x80000000::LEAF_NUM).0 };
    if function > max_function {
        return Err(Error::InvalidParameters(format!(
            "Function not supported: 0x{:x}",
            function
        )));
    }

    // this is safe because the host supports the `cpuid` instruction
    let entry = unsafe { __cpuid_count(function, count) };
    if entry.eax == 0 && entry.ebx == 0 && entry.ecx == 0 && entry.edx == 0 {
        return Err(Error::InvalidParameters(format!(
            "Invalid count: {}",
            count
        )));
    }

    Ok(entry)
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_host() -> Result<[u8; 12], Error> {
    match get_cpuid(0, 0) {
        Ok(vendor_entry) => {
            let bytes: [u8; 12] = unsafe {
                std::mem::transmute([vendor_entry.ebx, vendor_entry.edx, vendor_entry.ecx])
            };
            Ok(bytes)
        }
        Err(e) => Err(e),
    }
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_cpuid(cpuid: &CpuId) -> Result<[u8; 12], Error> {
    // Search for vendor id entry.
    for entry in cpuid.as_slice().iter() {
        if entry.function == 0 && entry.index == 0 {
            let cpu_vendor_id: [u8; 12] =
                unsafe { std::mem::transmute([entry.ebx, entry.edx, entry.ecx]) };
            return Ok(cpu_vendor_id);
        }
    }

    Err(Error::NotSupported)
}

/// Validates that the provided CPUID belongs to a CPU of the same
/// model as the host's.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn is_same_model(cpuid: &CpuId) -> bool {
    // Try to get the vendor IDs from the host and the CPUID struct.
    if let (Ok(host_vendor_id), Ok(cpuid_vendor_id)) =
        (get_vendor_id_from_host(), get_vendor_id_from_cpuid(cpuid))
    {
        // If the vendor IDs aren't the same, the CPUs are not identical.
        if host_vendor_id != cpuid_vendor_id {
            return false;
        }
    } else {
        // This only fails when CPUID is not supported, in which case
        // we can't tell if the CPUs are identical.
        return false;
    }

    // Try to get the feature information leaf from the host CPUID.
    let host_feature_info_leaf = get_cpuid(leaf_0x1::LEAF_NUM, 0);

    // The relevant information for this comparison is in the EAX register.
    let host_feature_info_leaf_eax = match host_feature_info_leaf {
        Ok(leaf) => leaf.eax,
        Err(_) => {
            // If this fails, we can't tell if the CPUs are identical.
            return false;
        }
    };

    // Search for the entry for leaf0x1.
    let feature_info_leaf = cpuid
        .as_slice()
        .iter()
        .find(|entry| entry.function == leaf_0x1::LEAF_NUM);

    // The relevant information is in EAX.
    let feature_info_leaf_eax = match feature_info_leaf {
        Some(leaf) => leaf.eax,
        None => {
            // Fail fast if we can't retrieve the relevant
            // information from CPUID.
            return false;
        }
    };

    // Validate that all of these properties are the same.
    for elem in &[
        leaf_0x1::eax::EXTENDED_FAMILY_ID_BITRANGE,
        leaf_0x1::eax::EXTENDED_PROCESSOR_MODEL_BITRANGE,
        leaf_0x1::eax::PROCESSOR_FAMILY_BITRANGE,
        leaf_0x1::eax::PROCESSOR_MODEL_BITRANGE,
        leaf_0x1::eax::STEPPING_BITRANGE,
    ] {
        if feature_info_leaf_eax.read_bits_in_range(elem)
            != host_feature_info_leaf_eax.read_bits_in_range(elem)
        {
            return false;
        }
    }

    true
}

#[cfg(test)]
pub mod tests {
    use crate::common::*;

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub fn get_topoext_fn() -> u32 {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        let function = match &vendor_id.ok().unwrap() {
            VENDOR_ID_INTEL => leaf_0x4::LEAF_NUM,
            VENDOR_ID_AMD => leaf_0x8000001d::LEAF_NUM,
            _ => 0,
        };
        assert!(function != 0);

        function
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_cpu_id() {
        // get_cpu_id should work correctly here
        let topoext_fn = get_topoext_fn();

        // check that get_cpuid works for valid parameters
        match get_cpuid(topoext_fn, 0) {
            Ok(topoext_entry) => {
                assert!(topoext_entry.eax != 0);
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `function`
        match get_cpuid(0x9000_0000, 0) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Function not supported: 0x90000000");
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `count`
        match get_cpuid(topoext_fn, 100) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Invalid count: 100");
            }
            _ => panic!("Wrong behavior"),
        }
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_vendor_id() {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        matches!(&vendor_id.ok().unwrap(), VENDOR_ID_INTEL | VENDOR_ID_AMD);
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_is_same_model() {
        let mut curr_cpuid = CpuId::new(2).unwrap();

        // Add the vendor ID leaf.
        let vendor = get_cpuid(0x0, 0).unwrap();
        curr_cpuid.as_mut_slice()[0].function = 0x0;
        curr_cpuid.as_mut_slice()[0].index = 0;
        curr_cpuid.as_mut_slice()[0].ebx = vendor.ebx;
        curr_cpuid.as_mut_slice()[0].ecx = vendor.ecx;
        curr_cpuid.as_mut_slice()[0].edx = vendor.edx;

        // Add the feature info leaf.
        let feature_info = get_cpuid(0x1, 0).unwrap();
        curr_cpuid.as_mut_slice()[1].function = 0x1;
        curr_cpuid.as_mut_slice()[1].index = 0;
        curr_cpuid.as_mut_slice()[1].eax = feature_info.eax;

        assert!(is_same_model(&curr_cpuid));

        let mut diff_vendor_cpuid = curr_cpuid.clone();
        for mut entry in diff_vendor_cpuid.as_mut_slice() {
            if entry.function == 0x0 && entry.index == 0 {
                entry.ebx = 0xFFFF_FFFF;
            }
        }
        assert!(!is_same_model(&diff_vendor_cpuid));

        let mut diff_feature_cpuid = curr_cpuid;
        for mut entry in diff_feature_cpuid.as_mut_slice() {
            if entry.function == 0x1 && entry.index == 0 {
                entry.eax ^= 0x1;
            }
        }
        assert!(!is_same_model(&diff_feature_cpuid));
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

pub mod cpu_time;

#[cfg(target_os = "linux")]
mod cgroup;

// re-export some traits for ease of use
use crate::config::{ReadableSize, KIB};
use lazy_static::lazy_static;
use std::sync::atomic::{AtomicU64, Ordering};
use sysinfo::RefreshKind;
pub use sysinfo::{DiskExt, NetworkExt, ProcessExt, ProcessorExt, SystemExt};

pub const HIGH_PRI: i32 = -1;
const CPU_CORES_QUOTA_ENV_VAR_KEY: &str = "TIKV_CPU_CORES_QUOTA";

static GLOBAL_MEMORY_USAGE: AtomicU64 = AtomicU64::new(0);

lazy_static! {
    #[cfg(target_os = "linux")]
    static ref SELF_CGROUP: cgroup::CGroupSys = cgroup::CGroupSys::default();
}

pub struct SysQuota;
impl SysQuota {
    #[cfg(target_os = "linux")]
    pub fn cpu_cores_quota() -> f64 {
        let mut cpu_num = num_cpus::get() as f64;
        let cpuset_cores = SELF_CGROUP.cpuset_cores().len() as f64;
        let cpu_quota = SELF_CGROUP.cpu_cores_quota().unwrap_or(0.);

        if cpuset_cores != 0. {
            cpu_num = cpu_num.min(cpuset_cores);
        }

        if cpu_quota != 0. {
            cpu_num = cpu_num.min(cpu_quota);
        }

        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(not(target_os = "linux"))]
    pub fn cpu_cores_quota() -> f64 {
        let cpu_num = num_cpus::get() as f64;
        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(target_os = "linux")]
    pub fn memory_limit_in_bytes() -> u64 {
        let total_mem = Self::sysinfo_memory_limit_in_bytes();
        let cgroup_memory_limit = SELF_CGROUP.memory_limit_in_bytes();
        if cgroup_memory_limit <= 0 {
            total_mem
        } else {
            std::cmp::min(total_mem, cgroup_memory_limit as u64)
        }
    }

    #[cfg(not(target_os = "linux"))]
    pub fn memory_limit_in_bytes() -> u64 {
        Self::sysinfo_memory_limit_in_bytes()
    }

    pub fn log_quota() {
        info!(
            "memory limit in bytes: {}, cpu cores quota: {}",
            Self::memory_limit_in_bytes(),
            Self::cpu_cores_quota()
        );
    }

    fn sysinfo_memory_limit_in_bytes() -> u64 {
        let system = sysinfo::System::new_with_specifics(RefreshKind::new().with_memory());
        system.get_total_memory() * KIB
    }
}

pub fn get_global_memory_usage() -> u64 {
    GLOBAL_MEMORY_USAGE.load(Ordering::Acquire)
}

pub fn record_global_memory_usage() {
    let mut system = sysinfo::System::new();
    system.refresh_memory();
    let usage = system.get_used_memory() * KIB;
    GLOBAL_MEMORY_USAGE.store(usage, Ordering::Release);
}

fn limit_cpu_cores_quota_by_env_var(quota: f64) -> f64 {
    match std::env::var(CPU_CORES_QUOTA_ENV_VAR_KEY)
        .ok()
        .and_then(|value| value.parse().ok())
    {
        Some(env_var_quota) if quota.is_sign_positive() => f64::min(quota, env_var_quota),
        _ => quota,
    }
}

#[cfg(target_os = "linux")]
pub mod thread {
    use libc::{self, c_int};
    use std::io::Error;

    pub fn set_priority(pri: i32) -> Result<(), Error> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            if libc::setpriority(libc::PRIO_PROCESS as u32, tid as u32, pri) != 0 {
                let e = Error::last_os_error();
                return Err(e);
            }
            Ok(())
        }
    }

    pub fn get_priority() -> Result<i32, Error> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            clear_errno();
            let ret = libc::getpriority(libc::PRIO_PROCESS as u32, tid as u32);
            if ret == -1 {
                let e = Error::last_os_error();
                if let Some(errno) = e.raw_os_error() {
                    if errno != 0 {
                        return Err(e);
                    }
                }
            }
            Ok(ret)
        }
    }

    // Sadly the std lib does not have any support for setting `errno`, so we
    // have to implement this ourselves.
    extern "C" {
        #[link_name = "__errno_location"]
        fn errno_location() -> *mut c_int;
    }

    fn clear_errno() {
        // Unsafe due to FFI.
        unsafe {
            *errno_location() = 0;
        }
    }

    #[cfg(test)]
    mod tests {
        use super::super::HIGH_PRI;
        use super::*;
        use std::io::ErrorKind;

        #[test]
        fn test_set_priority() {
            // priority is a value in range -20 to 19, the default priority
            // is 0, lower priorities cause more favorable scheduling.
            assert_eq!(get_priority().unwrap(), 0);
            set_priority(10).unwrap();
            assert_eq!(get_priority().unwrap(), 10);

            // only users who have `SYS_NICE_CAP` capability can increase priority.
            let ret = set_priority(HIGH_PRI);
            if let Err(e) = ret {
                assert_eq!(e.kind(), ErrorKind::PermissionDenied);
            } else {
                assert_eq!(get_priority().unwrap(), HIGH_PRI);
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
pub mod thread {
    use std::io::Error;

    pub fn set_priority(_: i32) -> Result<(), Error> {
        Ok(())
    }

    pub fn get_priority() -> Result<i32, Error> {
        Ok(0)
    }
}

fn read_size_in_cache(level: usize, field: &str) -> Option<u64> {
    std::fs::read_to_string(format!(
        "/sys/devices/system/cpu/cpu0/cache/index{}/{}",
        level, field
    ))
    .ok()
    .and_then(|s| s.parse::<ReadableSize>().ok())
    .map(|s| s.0)
}

/// Gets the size of given level cache.
///
/// It will only return `Some` on Linux.
pub fn cache_size(level: usize) -> Option<u64> {
    read_size_in_cache(level, "size")
}

/// Gets the size of given level cache line.
///
/// It will only return `Some` on Linux.
pub fn cache_line_size(level: usize) -> Option<u64> {
    read_size_in_cache(level, "coherency_line_size")
}
<commit_msg>*: fix build when using MacOS (#10250)<commit_after>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

pub mod cpu_time;

#[cfg(target_os = "linux")]
mod cgroup;

// re-export some traits for ease of use
use crate::config::{ReadableSize, KIB};
#[cfg(target_os = "linux")]
use lazy_static::lazy_static;
use std::sync::atomic::{AtomicU64, Ordering};
use sysinfo::RefreshKind;
pub use sysinfo::{DiskExt, NetworkExt, ProcessExt, ProcessorExt, SystemExt};

pub const HIGH_PRI: i32 = -1;
const CPU_CORES_QUOTA_ENV_VAR_KEY: &str = "TIKV_CPU_CORES_QUOTA";

static GLOBAL_MEMORY_USAGE: AtomicU64 = AtomicU64::new(0);

#[cfg(target_os = "linux")]
lazy_static! {
    static ref SELF_CGROUP: cgroup::CGroupSys = cgroup::CGroupSys::default();
}

pub struct SysQuota;
impl SysQuota {
    #[cfg(target_os = "linux")]
    pub fn cpu_cores_quota() -> f64 {
        let mut cpu_num = num_cpus::get() as f64;
        let cpuset_cores = SELF_CGROUP.cpuset_cores().len() as f64;
        let cpu_quota = SELF_CGROUP.cpu_cores_quota().unwrap_or(0.);

        if cpuset_cores != 0. {
            cpu_num = cpu_num.min(cpuset_cores);
        }

        if cpu_quota != 0. {
            cpu_num = cpu_num.min(cpu_quota);
        }

        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(not(target_os = "linux"))]
    pub fn cpu_cores_quota() -> f64 {
        let cpu_num = num_cpus::get() as f64;
        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(target_os = "linux")]
    pub fn memory_limit_in_bytes() -> u64 {
        let total_mem = Self::sysinfo_memory_limit_in_bytes();
        let cgroup_memory_limit = SELF_CGROUP.memory_limit_in_bytes();
        if cgroup_memory_limit <= 0 {
            total_mem
        } else {
            std::cmp::min(total_mem, cgroup_memory_limit as u64)
        }
    }

    #[cfg(not(target_os = "linux"))]
    pub fn memory_limit_in_bytes() -> u64 {
        Self::sysinfo_memory_limit_in_bytes()
    }

    pub fn log_quota() {
        info!(
            "memory limit in bytes: {}, cpu cores quota: {}",
            Self::memory_limit_in_bytes(),
            Self::cpu_cores_quota()
        );
    }

    fn sysinfo_memory_limit_in_bytes() -> u64 {
        let system = sysinfo::System::new_with_specifics(RefreshKind::new().with_memory());
        system.get_total_memory() * KIB
    }
}

pub fn get_global_memory_usage() -> u64 {
    GLOBAL_MEMORY_USAGE.load(Ordering::Acquire)
}

pub fn record_global_memory_usage() {
    let mut system = sysinfo::System::new();
    system.refresh_memory();
    let usage = system.get_used_memory() * KIB;
    GLOBAL_MEMORY_USAGE.store(usage, Ordering::Release);
}

fn limit_cpu_cores_quota_by_env_var(quota: f64) -> f64 {
    match std::env::var(CPU_CORES_QUOTA_ENV_VAR_KEY)
        .ok()
        .and_then(|value| value.parse().ok())
    {
        Some(env_var_quota) if quota.is_sign_positive() => f64::min(quota, env_var_quota),
        _ => quota,
    }
}

#[cfg(target_os = "linux")]
pub mod thread {
    use libc::{self, c_int};
    use std::io::Error;

    pub fn set_priority(pri: i32) -> Result<(), Error> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            if libc::setpriority(libc::PRIO_PROCESS as u32, tid as u32, pri) != 0 {
                let e = Error::last_os_error();
                return Err(e);
            }
            Ok(())
        }
    }

    pub fn get_priority() -> Result<i32, Error> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            clear_errno();
            let ret = libc::getpriority(libc::PRIO_PROCESS as u32, tid as u32);
            if ret == -1 {
                let e = Error::last_os_error();
                if let Some(errno) = e.raw_os_error() {
                    if errno != 0 {
                        return Err(e);
                    }
                }
            }
            Ok(ret)
        }
    }

    // Sadly the std lib does not have any support for setting `errno`, so we
    // have to implement this ourselves.
    extern "C" {
        #[link_name = "__errno_location"]
        fn errno_location() -> *mut c_int;
    }

    fn clear_errno() {
        // Unsafe due to FFI.
        unsafe {
            *errno_location() = 0;
        }
    }

    #[cfg(test)]
    mod tests {
        use super::super::HIGH_PRI;
        use super::*;
        use std::io::ErrorKind;

        #[test]
        fn test_set_priority() {
            // priority is a value in range -20 to 19, the default priority
            // is 0, lower priorities cause more favorable scheduling.
            assert_eq!(get_priority().unwrap(), 0);
            set_priority(10).unwrap();
            assert_eq!(get_priority().unwrap(), 10);

            // only users who have `SYS_NICE_CAP` capability can increase priority.
            let ret = set_priority(HIGH_PRI);
            if let Err(e) = ret {
                assert_eq!(e.kind(), ErrorKind::PermissionDenied);
            } else {
                assert_eq!(get_priority().unwrap(), HIGH_PRI);
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
pub mod thread {
    use std::io::Error;

    pub fn set_priority(_: i32) -> Result<(), Error> {
        Ok(())
    }

    pub fn get_priority() -> Result<i32, Error> {
        Ok(0)
    }
}

fn read_size_in_cache(level: usize, field: &str) -> Option<u64> {
    std::fs::read_to_string(format!(
        "/sys/devices/system/cpu/cpu0/cache/index{}/{}",
        level, field
    ))
    .ok()
    .and_then(|s| s.parse::<ReadableSize>().ok())
    .map(|s| s.0)
}

/// Gets the size of given level cache.
///
/// It will only return `Some` on Linux.
pub fn cache_size(level: usize) -> Option<u64> {
    read_size_in_cache(level, "size")
}

/// Gets the size of given level cache line.
///
/// It will only return `Some` on Linux.
pub fn cache_line_size(level: usize) -> Option<u64> {
    read_size_in_cache(level, "coherency_line_size")
}
<|endoftext|>


<commit_before>use core::pin::Pin;
use futures_core::future::{Future, FusedFuture};
use futures_core::task::{Context, Poll};
use crate::future::{Either, FutureExt};

/// Future for the [`select()`] function.
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[derive(Debug)]
pub struct Select<A, B> {
    inner: Option<(A, B)>,
}

impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}

/// Waits for either one of two differently-typed futures to complete.
///
/// This function will return a new future which awaits for either one of both
/// futures to complete. The returned future will finish with both the value
/// resolved and a future representing the completion of the other work.
///
/// Note that this function consumes the receiving futures and returns a
/// wrapped version of them.
///
/// Also note that if both this and the second future have the same
/// output type you can use the `Either::factor_first` method to
/// conveniently extract out the value at the end.
///
/// # Examples
///
/// ```
/// use futures::future::{self, Either, Future, FutureExt};
///
/// // A poor-man's join implemented on top of select
///
/// fn join<A, B>(a: A, b: B) -> impl Future<Output=(A::Output, B::Output)>
///     where A: Future + Unpin,
///           B: Future + Unpin,
/// {
///     future::select(a, b).then(|either| {
///         match either {
///             Either::Left((x, b)) => b.map(move |y| (x, y)).left_future(),
///             Either::Right((y, a)) => a.map(move |x| (x, y)).right_future(),
///         }
///     })
/// }
/// ```
pub fn select<A, B>(future1: A, future2: B) -> Select<A, B>
    where A: Future + Unpin, B: Future + Unpin
{
    Select { inner: Some((future1, future2)) }
}

impl<A, B> Future for Select<A, B>
where
    A: Future + Unpin,
    B: Future + Unpin,
{
    type Output = Either<(A::Output, B), (B::Output, A)>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let (mut a, mut b) = self.inner.take().expect("cannot poll Select twice");
        match a.poll_unpin(cx) {
            Poll::Ready(x) => Poll::Ready(Either::Left((x, b))),
            Poll::Pending => match b.poll_unpin(cx) {
                Poll::Ready(x) => Poll::Ready(Either::Right((x, a))),
                Poll::Pending => {
                    self.inner = Some((a, b));
                    Poll::Pending
                }
            }
        }
    }
}

impl<A, B> FusedFuture for Select<A, B>
where
    A: Future + Unpin,
    B: Future + Unpin,
{
    fn is_terminated(&self) -> bool {
        self.inner.is_none()
    }
}
<commit_msg>Add a simple example for future::select() (#2182)<commit_after>use core::pin::Pin;
use futures_core::future::{Future, FusedFuture};
use futures_core::task::{Context, Poll};
use crate::future::{Either, FutureExt};

/// Future for the [`select()`] function.
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[derive(Debug)]
pub struct Select<A, B> {
    inner: Option<(A, B)>,
}

impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}

/// Waits for either one of two differently-typed futures to complete.
///
/// This function will return a new future which awaits for either one of both
/// futures to complete. The returned future will finish with both the value
/// resolved and a future representing the completion of the other work.
///
/// Note that this function consumes the receiving futures and returns a
/// wrapped version of them.
///
/// Also note that if both this and the second future have the same
/// output type you can use the `Either::factor_first` method to
/// conveniently extract out the value at the end.
///
/// # Examples
///
/// A simple example
///
/// ```
/// # futures::executor::block_on(async {
/// use futures::future::{self, Either};
/// use futures::pin_mut;
///
/// // These two futures have different types even though their outputs have the same type
/// let future1 = async { 1 };
/// let future2 = async { 2 };
///
/// // 'select' requires Future + Unpin bounds
/// pin_mut!(future1);
/// pin_mut!(future2);
///
/// let value = match future::select(future1, future2).await {
///     Either::Left((value1, _)) => value1, // `value1` is resolved from `future1`
///                                          // `_` represents `future2`
///     Either::Right((value2, _)) => value2, // `value2` is resolved from `future2`
///                                           // `_` represents `future1`
/// };
///
/// assert!(value == 1 || value == 2);
/// # });
/// ```
///
/// A more complex example
///
/// ```
/// use futures::future::{self, Either, Future, FutureExt};
///
/// // A poor-man's join implemented on top of select
///
/// fn join<A, B>(a: A, b: B) -> impl Future<Output=(A::Output, B::Output)>
///     where A: Future + Unpin,
///           B: Future + Unpin,
/// {
///     future::select(a, b).then(|either| {
///         match either {
///             Either::Left((x, b)) => b.map(move |y| (x, y)).left_future(),
///             Either::Right((y, a)) => a.map(move |x| (x, y)).right_future(),
///         }
///     })
/// }
/// ```
pub fn select<A, B>(future1: A, future2: B) -> Select<A, B>
    where A: Future + Unpin, B: Future + Unpin
{
    Select { inner: Some((future1, future2)) }
}

impl<A, B> Future for Select<A, B>
where
    A: Future + Unpin,
    B: Future + Unpin,
{
    type Output = Either<(A::Output, B), (B::Output, A)>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let (mut a, mut b) = self.inner.take().expect("cannot poll Select twice");
        match a.poll_unpin(cx) {
            Poll::Ready(x) => Poll::Ready(Either::Left((x, b))),
            Poll::Pending => match b.poll_unpin(cx) {
                Poll::Ready(x) => Poll::Ready(Either::Right((x, a))),
                Poll::Pending => {
                    self.inner = Some((a, b));
                    Poll::Pending
                }
            }
        }
    }
}

impl<A, B> FusedFuture for Select<A, B>
where
    A: Future + Unpin,
    B: Future + Unpin,
{
    fn is_terminated(&self) -> bool {
        self.inner.is_none()
    }
}
<|endoftext|>


<commit_before>use alloc::sync::Arc;

/// A way of waking up a specific task.
///
/// By implementing this trait, types that are expected to be wrapped in an `Arc`
/// can be converted into `Waker` objects.
/// Those Wakers can be used to signal executors that a task it owns
/// is ready to be `poll`ed again.
// Note: Send + Sync required because `Arc<T>` doesn't automatically imply
// those bounds, but `Waker` implements them.
pub trait ArcWake: Send + Sync {
    /// Indicates that the associated task is ready to make progress and should
    /// be `poll`ed.
    ///
    /// This function can be called from an arbitrary thread, including threads which
    /// did not create the `ArcWake` based `Waker`.
    ///
    /// Executors generally maintain a queue of "ready" tasks; `wake` should place
    /// the associated task onto this queue.
    fn wake(self: Arc<Self>) {
        Self::wake_by_ref(&self)
    }

    /// Indicates that the associated task is ready to make progress and should
    /// be `poll`ed.
    ///
    /// This function can be called from an arbitrary thread, including threads which
    /// did not create the `ArcWake` based `Waker`.
    ///
    /// Executors generally maintain a queue of "ready" tasks; `wake_by_ref` should place
    /// the associated task onto this queue.
    ///
    /// This function is similar to `wake`, but must not consume the provided data
    /// pointer.
    fn wake_by_ref(arc_self: &Arc<Self>);
}
<commit_msg>Add docs for ArcWake -> Waker conversion<commit_after>use alloc::sync::Arc;

/// A way of waking up a specific task.
///
/// By implementing this trait, types that are expected to be wrapped in an `Arc`
/// can be converted into [`Waker`] objects.
/// Those Wakers can be used to signal executors that a task it owns
/// is ready to be `poll`ed again.
///
/// Currently, there are two ways to convert `ArcWake` into [`Waker`]:
///
/// * [`waker`](crate::task::waker()) converts `Arc<impl ArcWake>` into [`Waker`].
/// * [`waker_ref`](crate::task::waker_ref()) converts `&Arc<impl ArcWake>` into [`WakerRef`] that
///   provides access to a [`&Waker`][`Waker`].
///
/// [`Waker`]: std::task::Waker
/// [`WakerRef`]: crate::task::WakerRef
// Note: Send + Sync required because `Arc<T>` doesn't automatically imply
// those bounds, but `Waker` implements them.
pub trait ArcWake: Send + Sync {
    /// Indicates that the associated task is ready to make progress and should
    /// be `poll`ed.
    ///
    /// This function can be called from an arbitrary thread, including threads which
    /// did not create the `ArcWake` based [`Waker`].
    ///
    /// Executors generally maintain a queue of "ready" tasks; `wake` should place
    /// the associated task onto this queue.
    ///
    /// [`Waker`]: std::task::Waker
    fn wake(self: Arc<Self>) {
        Self::wake_by_ref(&self)
    }

    /// Indicates that the associated task is ready to make progress and should
    /// be `poll`ed.
    ///
    /// This function can be called from an arbitrary thread, including threads which
    /// did not create the `ArcWake` based [`Waker`].
    ///
    /// Executors generally maintain a queue of "ready" tasks; `wake_by_ref` should place
    /// the associated task onto this queue.
    ///
    /// This function is similar to [`wake`](ArcWake::wake), but must not consume the provided data
    /// pointer.
    ///
    /// [`Waker`]: std::task::Waker
    fn wake_by_ref(arc_self: &Arc<Self>);
}
<|endoftext|>


<commit_before>use {Backend, QueueFamily};
use device::{Device, PhysicalDevice};
use internal::Channel;
use native;

use std::sync::Arc;

use hal::{self, format, image};
use hal::{Backbuffer, SwapchainConfig};
use hal::window::Extent2D;

use core_graphics::geometry::{CGRect, CGSize};
use foreign_types::{ForeignType, ForeignTypeRef};
use parking_lot::{Mutex, MutexGuard};
use metal;
use objc::rc::autoreleasepool;
use objc::runtime::Object;

pub type CAMetalLayer = *mut Object;

pub struct Surface {
    pub(crate) inner: Arc<SurfaceInner>,
    pub(crate) has_swapchain: bool
}

#[derive(Debug)]
pub(crate) struct SurfaceInner {
    pub(crate) view: *mut Object,
    pub(crate) render_layer: Mutex<CAMetalLayer>,
}

unsafe impl Send for SurfaceInner {}
unsafe impl Sync for SurfaceInner {}

impl Drop for SurfaceInner {
    fn drop(&mut self) {
        unsafe { msg_send![self.view, release]; }
    }
}

impl SurfaceInner {
    fn next_frame<'a>(&self, frames: &'a [Frame]) -> Result<(usize, MutexGuard<'a, FrameInner>), ()> {
        let layer_ref = self.render_layer.lock();
        autoreleasepool(|| { // for the drawable
            let (drawable, texture_temp): (&metal::DrawableRef, &metal::TextureRef) = unsafe {
                let drawable = msg_send![*layer_ref, nextDrawable];
                (drawable, msg_send![drawable, texture])
            };

            trace!("looking for {:?}", texture_temp);
            match frames.iter().position(|f| f.texture.as_ptr() == texture_temp.as_ptr()) {
                Some(index) => {
                    let mut frame = frames[index].inner.lock();
                    assert!(frame.drawable.is_none());
                    frame.drawable = Some(drawable.to_owned());

                    debug!("next is frame[{}]", index);
                    Ok((index, frame))
                }
                None => Err(()),
            }
        })
    }
}

#[derive(Debug)]
struct FrameInner {
    drawable: Option<metal::Drawable>,
    /// If there is a `drawable`, availability indicates if it's free for grabs.
    /// If there is `None`, `available == false` means that the frame has already
    /// been acquired and the `drawable` will appear at some point.
    available: bool,
    last_frame: usize,
}

#[derive(Debug)]
struct Frame {
    inner: Mutex<FrameInner>,
    texture: metal::Texture,
}

unsafe impl Send for Frame {}
unsafe impl Sync for Frame {}

impl Drop for Frame {
    fn drop(&mut self) {
        info!("dropping Frame");
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AcquireMode {
    Wait,
    Oldest,
}

pub struct Swapchain {
    frames: Arc<Vec<Frame>>,
    surface: Arc<SurfaceInner>,
    extent: Extent2D,
    last_frame: usize,
    image_ready_callbacks: Vec<Arc<Mutex<Option<SwapchainImage>>>>,
    pub acquire_mode: AcquireMode,
}

impl Drop for Swapchain {
    fn drop(&mut self) {
        info!("dropping Swapchain");
        for ir in self.image_ready_callbacks.drain(..) {
            if ir.lock().take().is_some() {
                debug!("\twith a callback");
            }
        }
    }
}

impl Swapchain {
    /// Returns the drawable for the specified swapchain image index,
    /// marks the index as free for future use.
    pub(crate) fn take_drawable(&self, index: hal::SwapImageIndex) -> metal::Drawable {
        let mut frame = self
            .frames[index as usize]
            .inner
            .lock();
        assert!(!frame.available);
        frame.available = true;
        frame.drawable
            .take()
            .expect("Drawable has not been acquired!")
    }

    fn signal_sync(&self, sync: hal::FrameSync<Backend>) {
        match sync {
            hal::FrameSync::Semaphore(semaphore) => {
                if let Some(ref system) = semaphore.system {
                    system.signal();
                }
            }
            hal::FrameSync::Fence(fence) => {
                *fence.0.borrow_mut() = native::FenceInner::Idle { signaled: true };
            }
        }
    }
}

#[derive(Debug)]
pub struct SwapchainImage {
    frames: Arc<Vec<Frame>>,
    surface: Arc<SurfaceInner>,
    index: hal::SwapImageIndex,
}

impl SwapchainImage {
    /// Waits until the specified swapchain index is available for rendering.
    /// Returns the number of frames it had to wait.
    pub fn wait_until_ready(&self) -> usize {
        // check the target frame first
        {
            let frame = self.frames[self.index as usize].inner.lock();
            assert!(!frame.available);
            if frame.drawable.is_some() {
                return 0;
            }
        }
        // wait for new frames to come until we meet the chosen one
        let mut count = 1;
        while self.surface.next_frame(&self.frames).unwrap().0 != self.index as usize {
            count += 1;
        }
        debug!("Swapchain image is ready after {} frames", count);
        count
    }
}

impl hal::Surface<Backend> for Surface {
    fn kind(&self) -> image::Kind {
        let ex = self.inner.dimensions();
        image::Kind::D2(ex.width, ex.height, 1, 1)
    }

    fn compatibility(
        &self, _: &PhysicalDevice,
    ) -> (hal::SurfaceCapabilities, Option<Vec<format::Format>>, Vec<hal::PresentMode>) {
        let current_extent = Some(self.inner.dimensions());

        let caps = hal::SurfaceCapabilities {
            //Note: this is hardcoded in `CAMetalLayer` documentation
            image_count: 2 .. 4,
            current_extent,
            extents: Extent2D { width: 4, height: 4} .. Extent2D { width: 4096, height: 4096 },
            max_image_layers: 1,
            usage: image::Usage::COLOR_ATTACHMENT | image::Usage::SAMPLED |
                image::Usage::TRANSFER_SRC | image::Usage::TRANSFER_DST,
        };

        let formats = vec![
            format::Format::Bgra8Unorm,
            format::Format::Bgra8Srgb,
            format::Format::Rgba16Float,
        ];
        let present_modes = vec![
            hal::PresentMode::Fifo,
            hal::PresentMode::Immediate,
        ];

        (caps, Some(formats), present_modes)
    }

    fn supports_queue_family(&self, _queue_family: &QueueFamily) -> bool {
        // we only expose one family atm, so it's compatible
        true
    }
}

impl SurfaceInner {
    fn dimensions(&self) -> Extent2D {
        unsafe {
            // NSView/UIView bounds are measured in DIPs
            let bounds: CGRect = msg_send![self.view, bounds];
            //let bounds_pixel: NSRect = msg_send![self.nsview, convertRectToBacking:bounds];
            Extent2D {
                width: bounds.size.width as _,
                height: bounds.size.height as _,
            }
        }
    }
}

impl Device {
    pub(crate) fn build_swapchain(
        &self,
        surface: &mut Surface,
        config: SwapchainConfig,
    ) -> (Swapchain, Backbuffer<Backend>) {
        info!("build_swapchain {:?}", config);

        let mtl_format = self.private_caps
            .map_format(config.format)
            .expect("unsupported backbuffer format");

        let render_layer_borrow = surface.inner.render_layer.lock();
        let render_layer = *render_layer_borrow;
        let format_desc = config.format.surface_desc();
        let framebuffer_only = config.image_usage == image::Usage::COLOR_ATTACHMENT;
        let display_sync = match config.present_mode {
            hal::PresentMode::Immediate => false,
            _ => true,
        };
        let device = self.shared.device.lock();
        let device_raw: &metal::DeviceRef = &*device;

        unsafe {
            msg_send![render_layer, setDevice: device_raw];
            msg_send![render_layer, setPixelFormat: mtl_format];
            msg_send![render_layer, setFramebufferOnly: framebuffer_only];
            msg_send![render_layer, setMaximumDrawableCount: config.image_count as u64];
            msg_send![render_layer, setDrawableSize: CGSize::new(config.extent.width as f64, config.extent.height as f64)];
            //TODO: only set it where supported
            msg_send![render_layer, setAllowsNextDrawableTimeout:false];
            msg_send![render_layer, setDisplaySyncEnabled: display_sync];
        };

        let frames = (0 .. config.image_count)
            .map(|index| autoreleasepool(|| { // for the drawable & texture
                let (drawable, texture) = unsafe {
                    let drawable: &metal::DrawableRef = msg_send![render_layer, nextDrawable];
                    assert!(!drawable.as_ptr().is_null());
                    let texture: &metal::TextureRef = msg_send![drawable, texture];
                    (drawable, texture)
                };
                trace!("\tframe[{}] = {:?}", index, texture);

                let drawable = if index == 0 && surface.has_swapchain {
                    // when resizing, this trick frees up the currently shown frame
                    // HACK: the has_swapchain is unfortunate, and might not be
                    //       correct in all cases.
                    drawable.present();
                    None
                } else {
                    Some(drawable.to_owned())
                };
                Frame {
                    inner: Mutex::new(FrameInner {
                        drawable,
                        available: true,
                        last_frame: 0,
                    }),
                    texture: texture.to_owned(),
                }
            }))
            .collect::<Vec<_>>();

        let images = frames
            .iter()
            .map(|frame| native::Image {
                raw: frame.texture.clone(),
                kind: image::Kind::D2(config.extent.width, config.extent.height, 1, 1),
                format_desc,
                shader_channel: Channel::Float,
                mtl_format,
                mtl_type: metal::MTLTextureType::D2,
            })
            .collect();

        surface.has_swapchain = true;

        let swapchain = Swapchain {
            frames: Arc::new(frames),
            surface: surface.inner.clone(),
            extent: config.extent,
            last_frame: 0,
            image_ready_callbacks: Vec::new(),
            acquire_mode: AcquireMode::Wait,
        };

        (swapchain, Backbuffer::Images(images))
    }
}

impl hal::Swapchain<Backend> for Swapchain {
    fn acquire_image(
        &mut self, _timeout_ns: u64, sync: hal::FrameSync<Backend>
    ) -> Result<hal::SwapImageIndex, hal::AcquireError> {
        self.last_frame += 1;

        //TODO: figure out a proper story of HiDPI
        if false && self.surface.dimensions() != self.extent {
            unimplemented!()
        }

        let mut oldest_index = 0;
        let mut oldest_frame = self.last_frame;

        for (index, frame_arc) in self.frames.iter().enumerate() {
            let mut frame = frame_arc.inner.lock();
            if !frame.available {
                continue
            }
            if frame.drawable.is_some() {
                frame.available = false;
                frame.last_frame = self.last_frame;
                self.signal_sync(sync);
                return Ok(index as _);
            }
            if frame.last_frame < oldest_frame {
                oldest_frame = frame.last_frame;
                oldest_index = index;
            }
        }

        let (index, mut frame) = match self.acquire_mode {
            AcquireMode::Wait => {
                self.surface.next_frame(&self.frames)
                    .map_err(|_| hal::AcquireError::OutOfDate)?
            }
            AcquireMode::Oldest => {
                self.image_ready_callbacks.retain(|ir| ir.lock().is_some());
                match sync {
                    hal::FrameSync::Semaphore(semaphore) => {
                        self.image_ready_callbacks.push(Arc::clone(&semaphore.image_ready));
                        let mut sw_image = semaphore.image_ready.lock();
                        assert!(sw_image.is_none());
                        *sw_image = Some(SwapchainImage {
                            frames: self.frames.clone(),
                            surface: self.surface.clone(),
                            index: oldest_index as _,
                        });
                    }
                    hal::FrameSync::Fence(_fence) => {
                        //TODO: need presentation handlers always created and setting a bool
                        unimplemented!()
                    }
                }

                let frame = self.frames[oldest_index].inner.lock();
                (oldest_index, frame)
            }
        };

        assert!(frame.available);
        frame.last_frame = self.last_frame;
        frame.available = false;

        Ok(index as _)
    }
}
<commit_msg>[mtl] Add version check to `displaySyncEnabled`<commit_after>use {Backend, QueueFamily};
use device::{Device, PhysicalDevice};
use internal::Channel;
use native;

use std::sync::Arc;

use hal::{self, format, image};
use hal::{Backbuffer, SwapchainConfig};
use hal::window::Extent2D;

use core_graphics::geometry::{CGRect, CGSize};
use foreign_types::{ForeignType, ForeignTypeRef};
use parking_lot::{Mutex, MutexGuard};
use metal;
use objc::rc::autoreleasepool;
use objc::runtime::Object;

pub type CAMetalLayer = *mut Object;

pub struct Surface {
    pub(crate) inner: Arc<SurfaceInner>,
    pub(crate) has_swapchain: bool
}

#[derive(Debug)]
pub(crate) struct SurfaceInner {
    pub(crate) view: *mut Object,
    pub(crate) render_layer: Mutex<CAMetalLayer>,
}

unsafe impl Send for SurfaceInner {}
unsafe impl Sync for SurfaceInner {}

impl Drop for SurfaceInner {
    fn drop(&mut self) {
        unsafe { msg_send![self.view, release]; }
    }
}

impl SurfaceInner {
    fn next_frame<'a>(&self, frames: &'a [Frame]) -> Result<(usize, MutexGuard<'a, FrameInner>), ()> {
        let layer_ref = self.render_layer.lock();
        autoreleasepool(|| { // for the drawable
            let (drawable, texture_temp): (&metal::DrawableRef, &metal::TextureRef) = unsafe {
                let drawable = msg_send![*layer_ref, nextDrawable];
                (drawable, msg_send![drawable, texture])
            };

            trace!("looking for {:?}", texture_temp);
            match frames.iter().position(|f| f.texture.as_ptr() == texture_temp.as_ptr()) {
                Some(index) => {
                    let mut frame = frames[index].inner.lock();
                    assert!(frame.drawable.is_none());
                    frame.drawable = Some(drawable.to_owned());

                    debug!("next is frame[{}]", index);
                    Ok((index, frame))
                }
                None => Err(()),
            }
        })
    }
}

#[derive(Debug)]
struct FrameInner {
    drawable: Option<metal::Drawable>,
    /// If there is a `drawable`, availability indicates if it's free for grabs.
    /// If there is `None`, `available == false` means that the frame has already
    /// been acquired and the `drawable` will appear at some point.
    available: bool,
    last_frame: usize,
}

#[derive(Debug)]
struct Frame {
    inner: Mutex<FrameInner>,
    texture: metal::Texture,
}

unsafe impl Send for Frame {}
unsafe impl Sync for Frame {}

impl Drop for Frame {
    fn drop(&mut self) {
        info!("dropping Frame");
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AcquireMode {
    Wait,
    Oldest,
}

pub struct Swapchain {
    frames: Arc<Vec<Frame>>,
    surface: Arc<SurfaceInner>,
    extent: Extent2D,
    last_frame: usize,
    image_ready_callbacks: Vec<Arc<Mutex<Option<SwapchainImage>>>>,
    pub acquire_mode: AcquireMode,
}

impl Drop for Swapchain {
    fn drop(&mut self) {
        info!("dropping Swapchain");
        for ir in self.image_ready_callbacks.drain(..) {
            if ir.lock().take().is_some() {
                debug!("\twith a callback");
            }
        }
    }
}

impl Swapchain {
    /// Returns the drawable for the specified swapchain image index,
    /// marks the index as free for future use.
    pub(crate) fn take_drawable(&self, index: hal::SwapImageIndex) -> metal::Drawable {
        let mut frame = self
            .frames[index as usize]
            .inner
            .lock();
        assert!(!frame.available);
        frame.available = true;
        frame.drawable
            .take()
            .expect("Drawable has not been acquired!")
    }

    fn signal_sync(&self, sync: hal::FrameSync<Backend>) {
        match sync {
            hal::FrameSync::Semaphore(semaphore) => {
                if let Some(ref system) = semaphore.system {
                    system.signal();
                }
            }
            hal::FrameSync::Fence(fence) => {
                *fence.0.borrow_mut() = native::FenceInner::Idle { signaled: true };
            }
        }
    }
}

#[derive(Debug)]
pub struct SwapchainImage {
    frames: Arc<Vec<Frame>>,
    surface: Arc<SurfaceInner>,
    index: hal::SwapImageIndex,
}

impl SwapchainImage {
    /// Waits until the specified swapchain index is available for rendering.
    /// Returns the number of frames it had to wait.
    pub fn wait_until_ready(&self) -> usize {
        // check the target frame first
        {
            let frame = self.frames[self.index as usize].inner.lock();
            assert!(!frame.available);
            if frame.drawable.is_some() {
                return 0;
            }
        }
        // wait for new frames to come until we meet the chosen one
        let mut count = 1;
        while self.surface.next_frame(&self.frames).unwrap().0 != self.index as usize {
            count += 1;
        }
        debug!("Swapchain image is ready after {} frames", count);
        count
    }
}

impl hal::Surface<Backend> for Surface {
    fn kind(&self) -> image::Kind {
        let ex = self.inner.dimensions();
        image::Kind::D2(ex.width, ex.height, 1, 1)
    }

    fn compatibility(
        &self, _: &PhysicalDevice,
    ) -> (hal::SurfaceCapabilities, Option<Vec<format::Format>>, Vec<hal::PresentMode>) {
        let current_extent = Some(self.inner.dimensions());

        let caps = hal::SurfaceCapabilities {
            //Note: this is hardcoded in `CAMetalLayer` documentation
            image_count: 2 .. 4,
            current_extent,
            extents: Extent2D { width: 4, height: 4} .. Extent2D { width: 4096, height: 4096 },
            max_image_layers: 1,
            usage: image::Usage::COLOR_ATTACHMENT | image::Usage::SAMPLED |
                image::Usage::TRANSFER_SRC | image::Usage::TRANSFER_DST,
        };

        let formats = vec![
            format::Format::Bgra8Unorm,
            format::Format::Bgra8Srgb,
            format::Format::Rgba16Float,
        ];
        let present_modes = vec![
            hal::PresentMode::Fifo,
            hal::PresentMode::Immediate,
        ];

        (caps, Some(formats), present_modes)
    }

    fn supports_queue_family(&self, _queue_family: &QueueFamily) -> bool {
        // we only expose one family atm, so it's compatible
        true
    }
}

impl SurfaceInner {
    fn dimensions(&self) -> Extent2D {
        unsafe {
            // NSView/UIView bounds are measured in DIPs
            let bounds: CGRect = msg_send![self.view, bounds];
            //let bounds_pixel: NSRect = msg_send![self.nsview, convertRectToBacking:bounds];
            Extent2D {
                width: bounds.size.width as _,
                height: bounds.size.height as _,
            }
        }
    }
}

impl Device {
    pub(crate) fn build_swapchain(
        &self,
        surface: &mut Surface,
        config: SwapchainConfig,
    ) -> (Swapchain, Backbuffer<Backend>) {
        info!("build_swapchain {:?}", config);

        let caps = &self.private_caps;
        let mtl_format = caps
            .map_format(config.format)
            .expect("unsupported backbuffer format");

        let render_layer_borrow = surface.inner.render_layer.lock();
        let render_layer = *render_layer_borrow;
        let format_desc = config.format.surface_desc();
        let framebuffer_only = config.image_usage == image::Usage::COLOR_ATTACHMENT;
        let display_sync = config.present_mode != hal::PresentMode::Immediate;
        let can_set_display_sync = caps.os_is_mac && caps.has_version_at_least(10, 13);
        let device = self.shared.device.lock();
        let device_raw: &metal::DeviceRef = &*device;

        unsafe {
            msg_send![render_layer, setDevice: device_raw];
            msg_send![render_layer, setPixelFormat: mtl_format];
            msg_send![render_layer, setFramebufferOnly: framebuffer_only];
            msg_send![render_layer, setMaximumDrawableCount: config.image_count as u64];
            msg_send![render_layer, setDrawableSize: CGSize::new(config.extent.width as f64, config.extent.height as f64)];
            msg_send![render_layer, setAllowsNextDrawableTimeout:false];
            if can_set_display_sync {
                msg_send![render_layer, setDisplaySyncEnabled: display_sync];
            }
        };

        let frames = (0 .. config.image_count)
            .map(|index| autoreleasepool(|| { // for the drawable & texture
                let (drawable, texture) = unsafe {
                    let drawable: &metal::DrawableRef = msg_send![render_layer, nextDrawable];
                    assert!(!drawable.as_ptr().is_null());
                    let texture: &metal::TextureRef = msg_send![drawable, texture];
                    (drawable, texture)
                };
                trace!("\tframe[{}] = {:?}", index, texture);

                let drawable = if index == 0 && surface.has_swapchain {
                    // when resizing, this trick frees up the currently shown frame
                    // HACK: the has_swapchain is unfortunate, and might not be
                    //       correct in all cases.
                    drawable.present();
                    None
                } else {
                    Some(drawable.to_owned())
                };
                Frame {
                    inner: Mutex::new(FrameInner {
                        drawable,
                        available: true,
                        last_frame: 0,
                    }),
                    texture: texture.to_owned(),
                }
            }))
            .collect::<Vec<_>>();

        let images = frames
            .iter()
            .map(|frame| native::Image {
                raw: frame.texture.clone(),
                kind: image::Kind::D2(config.extent.width, config.extent.height, 1, 1),
                format_desc,
                shader_channel: Channel::Float,
                mtl_format,
                mtl_type: metal::MTLTextureType::D2,
            })
            .collect();

        surface.has_swapchain = true;

        let swapchain = Swapchain {
            frames: Arc::new(frames),
            surface: surface.inner.clone(),
            extent: config.extent,
            last_frame: 0,
            image_ready_callbacks: Vec::new(),
            acquire_mode: AcquireMode::Wait,
        };

        (swapchain, Backbuffer::Images(images))
    }
}

impl hal::Swapchain<Backend> for Swapchain {
    fn acquire_image(
        &mut self, _timeout_ns: u64, sync: hal::FrameSync<Backend>
    ) -> Result<hal::SwapImageIndex, hal::AcquireError> {
        self.last_frame += 1;

        //TODO: figure out a proper story of HiDPI
        if false && self.surface.dimensions() != self.extent {
            unimplemented!()
        }

        let mut oldest_index = 0;
        let mut oldest_frame = self.last_frame;

        for (index, frame_arc) in self.frames.iter().enumerate() {
            let mut frame = frame_arc.inner.lock();
            if !frame.available {
                continue
            }
            if frame.drawable.is_some() {
                frame.available = false;
                frame.last_frame = self.last_frame;
                self.signal_sync(sync);
                return Ok(index as _);
            }
            if frame.last_frame < oldest_frame {
                oldest_frame = frame.last_frame;
                oldest_index = index;
            }
        }

        let (index, mut frame) = match self.acquire_mode {
            AcquireMode::Wait => {
                self.surface.next_frame(&self.frames)
                    .map_err(|_| hal::AcquireError::OutOfDate)?
            }
            AcquireMode::Oldest => {
                self.image_ready_callbacks.retain(|ir| ir.lock().is_some());
                match sync {
                    hal::FrameSync::Semaphore(semaphore) => {
                        self.image_ready_callbacks.push(Arc::clone(&semaphore.image_ready));
                        let mut sw_image = semaphore.image_ready.lock();
                        assert!(sw_image.is_none());
                        *sw_image = Some(SwapchainImage {
                            frames: self.frames.clone(),
                            surface: self.surface.clone(),
                            index: oldest_index as _,
                        });
                    }
                    hal::FrameSync::Fence(_fence) => {
                        //TODO: need presentation handlers always created and setting a bool
                        unimplemented!()
                    }
                }

                let frame = self.frames[oldest_index].inner.lock();
                (oldest_index, frame)
            }
        };

        assert!(frame.available);
        frame.last_frame = self.last_frame;
        frame.available = false;

        Ok(index as _)
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::time::Instant;

use futures::task;

use reactor::{Message, Handle, Remote};

/// A token that identifies an active timeout.
pub struct TimeoutToken {
    token: usize,
}

impl TimeoutToken {
    /// Adds a new timeout to get fired at the specified instant, notifying the
    /// specified task.
    pub fn new(at: Instant, handle: &Handle) -> io::Result<TimeoutToken> {
        match handle.inner.upgrade() {
            Some(inner) => {
                let token = inner.borrow_mut().add_timeout(at);
                Ok(TimeoutToken { token: token })
            }
            None => Err(io::Error::new(io::ErrorKind::Other, "event loop gone")),
        }
    }

    /// Updates a previously added timeout to notify a new task instead.
    ///
    /// # Panics
    ///
    /// This method will panic if the timeout specified was not created by this
    /// loop handle's `add_timeout` method.
    pub fn update_timeout(&self, handle: &Remote) {
        handle.send(Message::UpdateTimeout(self.token, task::current()))
    }

    /// Resets previously added (or fired) timeout to an new timeout
    ///
    /// # Panics
    ///
    /// This method will panic if the timeout specified was not created by this
    /// loop handle's `add_timeout` method.
    pub fn reset_timeout(&mut self, at: Instant, handle: &Remote) {
        handle.send(Message::ResetTimeout(self.token, at));
    }

    /// Cancel a previously added timeout.
    ///
    /// # Panics
    ///
    /// This method will panic if the timeout specified was not created by this
    /// loop handle's `add_timeout` method.
    pub fn cancel_timeout(&self, handle: &Remote) {
        debug!("cancel timeout {}", self.token);
        handle.send(Message::CancelTimeout(self.token))
    }
}
<commit_msg>TokenTimeout: add doc around new panic<commit_after>use std::io;
use std::time::Instant;

use futures::task;

use reactor::{Message, Handle, Remote};

/// A token that identifies an active timeout.
pub struct TimeoutToken {
    token: usize,
}

impl TimeoutToken {
    /// Adds a new timeout to get fired at the specified instant, notifying the
    /// specified task.
    pub fn new(at: Instant, handle: &Handle) -> io::Result<TimeoutToken> {
        match handle.inner.upgrade() {
            Some(inner) => {
                let token = inner.borrow_mut().add_timeout(at);
                Ok(TimeoutToken { token: token })
            }
            None => Err(io::Error::new(io::ErrorKind::Other, "event loop gone")),
        }
    }

    /// Updates a previously added timeout to notify a new task instead.
    ///
    /// # Panics
    ///
    /// This method will panic if the timeout specified was not created by this
    /// loop handle's `add_timeout` method.
    pub fn update_timeout(&self, handle: &Remote) {
        handle.send(Message::UpdateTimeout(self.token, task::current()))
    }

    /// Resets previously added (or fired) timeout to an new timeout
    ///
    /// # Panics
    ///
    /// This method will panic if the timeout specified was not created by this
    /// loop handle's `add_timeout` method.
    pub fn reset_timeout(&mut self, at: Instant, handle: &Remote) {
        handle.send(Message::ResetTimeout(self.token, at));
    }

    /// Cancel a previously added timeout.
    ///
    /// # Panics
    ///
    /// This method will panic if the timeout specified was not created by this
    /// loop handle's `add_timeout` method or if called multiple times.
    pub fn cancel_timeout(&self, handle: &Remote) {
        debug!("cancel timeout {}", self.token);
        handle.send(Message::CancelTimeout(self.token))
    }
}
<|endoftext|>


<commit_before>use Score;
use DocId;
use fastfield::U64FastFieldReader;
use postings::DocSet;
use query::Scorer;
use postings::Postings;
use fastfield::FastFieldReader;

pub struct TermScorer<TPostings>
where
    TPostings: Postings,
{
    pub idf: Score,
    pub fieldnorm_reader_opt: Option<U64FastFieldReader>,
    pub postings: TPostings,
}

impl<TPostings> TermScorer<TPostings>
where
    TPostings: Postings,
{
    pub fn postings(&self) -> &TPostings {
        &self.postings
    }
}

impl<TPostings> DocSet for TermScorer<TPostings>
where
    TPostings: Postings,
{
    fn advance(&mut self) -> bool {
        self.postings.advance()
    }

    fn doc(&self) -> DocId {
        self.postings.doc()
    }

    fn size_hint(&self) -> usize {
        self.postings.size_hint()
    }
}

impl<TPostings> Scorer for TermScorer<TPostings>
where
    TPostings: Postings,
{
    fn score(&self) -> Score {
        let doc = self.postings.doc();
        let tf = match self.fieldnorm_reader_opt {
            Some(ref fieldnorm_reader) => {
                let field_norm = fieldnorm_reader.get(doc);
                (self.postings.term_freq() as f32 / field_norm as f32)
            }
            None => self.postings.term_freq() as f32,
        };
        self.idf * tf.sqrt()
    }
}
<commit_msg>NOBUG Use the skipping logic of segment postings in<commit_after>use Score;
use DocId;
use postings::SkipResult;
use fastfield::U64FastFieldReader;
use postings::DocSet;
use query::Scorer;
use postings::Postings;
use fastfield::FastFieldReader;

pub struct TermScorer<TPostings>
where
    TPostings: Postings,
{
    pub idf: Score,
    pub fieldnorm_reader_opt: Option<U64FastFieldReader>,
    pub postings: TPostings,
}

impl<TPostings> TermScorer<TPostings>
where
    TPostings: Postings,
{
    pub fn postings(&self) -> &TPostings {
        &self.postings
    }
}

impl<TPostings> DocSet for TermScorer<TPostings>
where
    TPostings: Postings,
{
    fn advance(&mut self) -> bool {
        self.postings.advance()
    }

    fn doc(&self) -> DocId {
        self.postings.doc()
    }

    fn size_hint(&self) -> usize {
        self.postings.size_hint()
    }

    fn skip_next(&mut self, target: DocId) -> SkipResult {
        self.postings.skip_next(target)
    }
}

impl<TPostings> Scorer for TermScorer<TPostings>
where
    TPostings: Postings,
{
    fn score(&self) -> Score {
        let doc = self.postings.doc();
        let tf = match self.fieldnorm_reader_opt {
            Some(ref fieldnorm_reader) => {
                let field_norm = fieldnorm_reader.get(doc);
                (self.postings.term_freq() as f32 / field_norm as f32)
            }
            None => self.postings.term_freq() as f32,
        };
        self.idf * tf.sqrt()
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use http::Request;
use http::Response;

pub struct Context {
    pub request: Request,
    pub response: Response,
    pub contexts: HashMap<String, Value>,
    stop: bool
}

impl Context {
    pub fn new(request: Request) -> Context {
        let response = Response::empty(200);

        Context {
            request: request,
            response: response,
            contexts: HashMap::new(),
            stop: false
        }
    }

    pub fn stop(&mut self) {
        self.stop = true;
    }

    pub fn next(&self) -> bool {
        !self.stop
    }
}

pub enum Value {
    String(String),
    Int32(i32),
    Int64(i64),
    Usize(usize),
    Isize(isize),
    Double(f64),
    Array(Vec<Value>),
    Map(HashMap<Value, Value>),
    Boolean(bool),
    Binary(Vec<u8>),
}

impl Value {
    pub fn as_str(&self) -> Option<&str> {
        match *self {
            Value::String(ref s) => Some(s),
            _ => None,
        }
    }

    pub fn as_i32(&self) -> Option<i32> {
        match *self {
            Value::Int32(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_i64(&self) -> Option<i64> {
        match *self {
            Value::Int64(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_usize(&self) -> Option<usize> {
        match *self {
            Value::Usize(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_isize(&self) -> Option<isize> {
        match *self {
            Value::Isize(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_f64(&self) -> Option<f64> {
        match *self {
            Value::Double(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_vec(&self) -> Option<&Vec<Value>> {
        match *self {
            Value::Array(ref i) => Some(i),
            _ => None
        }
    }

    pub fn as_map(&self) -> Option<&HashMap<Value, Value>> {
        match *self {
            Value::Map(ref i) => Some(i),
            _ => None
        }
    }

    pub fn as_bool(&self) -> Option<bool> {
        match *self {
            Value::Boolean(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_binary(&self) -> Option<&Vec<u8>> {
        match *self {
            Value::Binary(ref i) => Some(i),
            _ => None
        }
    }
}
<commit_msg>add contexts type<commit_after>use std::collections::HashMap;
use std::time::Instant;

use http::Request;
use http::Response;

pub struct Context {
    pub request: Request,
    pub response: Response,
    pub contexts: HashMap<String, Value>,
    stop: bool
}

impl Context {
    pub fn new(request: Request) -> Context {
        let response = Response::empty(200);

        Context {
            request: request,
            response: response,
            contexts: HashMap::new(),
            stop: false
        }
    }

    pub fn stop(&mut self) {
        self.stop = true;
    }

    pub fn next(&self) -> bool {
        !self.stop
    }
}

pub enum Value {
    String(String),
    Int32(i32),
    Int64(i64),
    Usize(usize),
    Isize(isize),
    Double(f64),
    Array(Vec<Value>),
    Map(HashMap<Value, Value>),
    Boolean(bool),
    Binary(Vec<u8>),
    Instant(Instant)
}

impl Value {
    pub fn as_str(&self) -> Option<&str> {
        match *self {
            Value::String(ref s) => Some(s),
            _ => None,
        }
    }

    pub fn as_i32(&self) -> Option<i32> {
        match *self {
            Value::Int32(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_i64(&self) -> Option<i64> {
        match *self {
            Value::Int64(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_usize(&self) -> Option<usize> {
        match *self {
            Value::Usize(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_isize(&self) -> Option<isize> {
        match *self {
            Value::Isize(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_f64(&self) -> Option<f64> {
        match *self {
            Value::Double(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_vec(&self) -> Option<&Vec<Value>> {
        match *self {
            Value::Array(ref i) => Some(i),
            _ => None
        }
    }

    pub fn as_map(&self) -> Option<&HashMap<Value, Value>> {
        match *self {
            Value::Map(ref i) => Some(i),
            _ => None
        }
    }

    pub fn as_bool(&self) -> Option<bool> {
        match *self {
            Value::Boolean(ref i) => Some(*i),
            _ => None
        }
    }

    pub fn as_binary(&self) -> Option<&Vec<u8>> {
        match *self {
            Value::Binary(ref i) => Some(i),
            _ => None
        }
    }

    pub fn as_instant(&self) -> Option<&Instant> {
        match *self {
            Value::Instant(ref i) => Some(i),
            _ => None
        }
    }
}
<|endoftext|>