<commit_before>//! TODO Documentation
use std::{fmt, panic};
use std::mem::ManuallyDrop;
use std::ops::Deref;
use std::rc::{Rc, Weak};
use std::sync::atomic::{AtomicBool, Ordering};

use errors::{UpgradeHandleErr, UpgradeHandleResult};
use wlroots_sys::{wlr_input_device, wlr_key_state, wlr_keyboard, wlr_keyboard_get_modifiers,
                  wlr_keyboard_led, wlr_keyboard_led_update, wlr_keyboard_modifier,
                  wlr_keyboard_set_keymap};

use xkbcommon::xkb::{self, Keycode, Keymap, LedIndex, ModIndex, ModMask};

use InputDevice;

pub type KeyState = wlr_key_state;

/// The masks for each group of modifiers.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct ModifierMasks {
    pub depressed: ModMask,
    pub latched: ModMask,
    pub locked: ModMask,
    pub group: ModMask
}

/// Information about repeated keypresses for a particular Keyboard.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct RepeatInfo {
    /// The rate at which extended keypresses will fire more events.
    rate: i32,
    /// How long it takes for a keypress to register on this device.
    delay: i32
}

pub struct Keyboard {
    /// The structure that ensures weak handles to this structure are still alive.
    ///
    /// They contain weak handles, and will safely not use dead memory when this
    /// is freed by wlroots.
    ///
    /// If this is `None`, then this is from an upgraded `KeyboardHandle`, and
    /// the operations are **unchecked**.
    /// This is means safe operations might fail, but only if you use the unsafe
    /// marked function `upgrade` on a `KeyboardHandle`.
    liveliness: Option<Rc<AtomicBool>>,
    /// The device that refers to this keyboard.
    device: InputDevice,
    /// The underlying keyboard data.
    keyboard: *mut wlr_keyboard,
    /// The keymap used by wlroots.
    /// Stored here so that we can ensure you can only borrow the keymap
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    keymap: Option<ManuallyDrop<Keymap>>,
    /// The keyboard state used by wlroots.
    /// Stored here so that we can ensure you can only borrow the state
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    #[allow(dead_code)]
    xkb_state: Option<ManuallyDrop<xkb::State>>
}

pub struct KeyboardHandle {
    /// The Rc that ensures that this handle is still alive.
    ///
    /// When wlroots deallocates the keyboard associated with this handle,
    handle: Weak<AtomicBool>,
    /// The device that refers to this keyboard.
    device: InputDevice,
    /// The underlying keyboard data.
    keyboard: *mut wlr_keyboard,
    /// The keymap used by wlroots.
    /// Stored here so that we can ensure you can only borrow the keymap
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    #[allow(dead_code)]
    keymap: Option<ManuallyDrop<Keymap>>,
    /// The keyboard state used by wlroots.
    /// Stored here so that we can ensure you can only borrow the state
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    #[allow(dead_code)]
    xkb_state: Option<ManuallyDrop<xkb::State>>
}

impl Keyboard {
    /// Tries to convert an input device to a Keyboard
    ///
    /// Returns None if it is of a different type of input variant.
    ///
    /// # Safety
    /// This creates a totally new Keyboard (e.g with its own reference count)
    /// so only do this once per `wlr_input_device`!
    pub(crate) unsafe fn new_from_input_device(device: *mut wlr_input_device) -> Option<Self> {
        use wlroots_sys::wlr_input_device_type::*;
        match (*device).type_ {
            WLR_INPUT_DEVICE_KEYBOARD => {
                let keyboard = (*device).__bindgen_anon_1.keyboard;
                Some(Keyboard { liveliness: Some(Rc::new(AtomicBool::new(false))),
                                device: InputDevice::from_ptr(device),
                                keyboard,
                                keymap: None,
                                xkb_state: None })
            }
            _ => None
        }
    }

    unsafe fn from_handle(handle: &KeyboardHandle) -> Self {
        Keyboard { liveliness: None,
                   device: handle.input_device().clone(),
                   keyboard: handle.as_ptr(),
                   keymap: None,
                   xkb_state: None }
    }

    /// Gets the wlr_keyboard associated with this KeyboardHandle.
    pub(crate) unsafe fn as_ptr(&self) -> *mut wlr_keyboard {
        self.keyboard
    }

    /// Gets the wlr_input_device associated with this KeyboardHandle
    pub fn input_device(&self) -> &InputDevice {
        &self.device
    }

    /// Set the keymap for this Keyboard.
    pub fn set_keymap(&mut self, keymap: &Keymap) {
        unsafe {
            wlr_keyboard_set_keymap(self.keyboard, keymap.get_raw_ptr() as _);
        }
    }

    /// Get the XKB keymap associated with this Keyboard.
    pub fn get_keymap(&mut self) -> Option<&Keymap> {
        unsafe {
            let keymap_ptr = (*self.keyboard).keymap as *mut _;
            if keymap_ptr.is_null() {
                None
            } else {
                self.keymap = Some(ManuallyDrop::new(Keymap::from_raw_ptr(keymap_ptr)));
                self.keymap.as_ref().map(Deref::deref)
            }
        }
    }

    /// Get the keycodes for this keyboard as reported by XKB.
    ///
    /// # Limitations
    /// wlroots limits this list to `WLR_KEYBOARD_KEYS_CAP` elements,
    /// which at the time of writing is `32`.
    pub fn keycodes(&self) -> Vec<Keycode> {
        unsafe {
            let mut result = (*self.keyboard).keycodes.to_vec();
            result.truncate((*self.keyboard).num_keycodes);
            result
        }
    }

    /// Get the list of LEDs for this keyboard as reported by XKB.
    pub fn led_list(&self) -> Vec<LedIndex> {
        unsafe { (*self.keyboard).led_indexes.to_vec() }
    }

    /// Get the list of modifiers for this keyboard as reported by XKB.
    pub fn modifier_list(&self) -> Vec<ModIndex> {
        unsafe { (*self.keyboard).mod_indexes.to_vec() }
    }

    /// Get the size of the keymap.
    pub fn keymap_size(&self) -> usize {
        unsafe { (*self.keyboard).keymap_size }
    }

    /// Get the XKB state associated with this `Keyboard`.
    pub fn get_xkb_state(&mut self) -> Option<&xkb::State> {
        unsafe {
            let xkb_state_ptr = (*self.keyboard).xkb_state as *mut _;
            if xkb_state_ptr.is_null() {
                None
            } else {
                self.xkb_state = Some(ManuallyDrop::new(xkb::State::from_raw_ptr(xkb_state_ptr)));
                self.xkb_state.as_ref().map(Deref::deref)
            }
        }
    }

    /// Get the repeat info for this keyboard.
    pub fn repeat_info(&self) -> RepeatInfo {
        unsafe {
            RepeatInfo { rate: (*self.keyboard).repeat_info.rate,
                         delay: (*self.keyboard).repeat_info.delay }
        }
    }

    /// Update the LED lights using the provided bitmap.
    ///
    /// 1 means one, 0 means off.
    pub fn update_led(&mut self, leds: KeyboardLed) {
        unsafe {
            wlr_keyboard_led_update(self.keyboard, leds.bits() as u32);
        }
    }

    /// Get the modifiers that are currently pressed on the keyboard.
    pub fn get_modifiers(&self) -> KeyboardModifier {
        unsafe {
            KeyboardModifier::from_bits_truncate(wlr_keyboard_get_modifiers(self.keyboard))
        }
    }

    /// Get the modifier masks for each group.
    pub fn get_modifier_masks(&self) -> ModifierMasks {
        unsafe {
            ModifierMasks { depressed: (*self.keyboard).modifiers.depressed,
                            latched: (*self.keyboard).modifiers.latched,
                            locked: (*self.keyboard).modifiers.locked,
                            group: (*self.keyboard).modifiers.group }
        }
    }

    /// Creates a weak reference to a `Keyboard`.
    ///
    /// # Panics
    /// If this `Keyboard` is a previously upgraded `KeyboardHandle`,
    /// then this function will panic.
    pub fn weak_reference(&self) -> KeyboardHandle {
        let arc = self.liveliness.as_ref()
                      .expect("Cannot downgrade previously upgraded KeyboardHandle!");
        KeyboardHandle { handle: Rc::downgrade(arc),
                         // NOTE Rationale for cloning:
                         // We can't use the keyboard handle unless the keyboard is alive,
                         // which means the device pointer is still alive.
                         device: unsafe { self.device.clone() },
                         keyboard: self.keyboard,
                         keymap: None,
                         xkb_state: None }
    }

    /// Manually set the lock used to determine if a double-borrow is
    /// occuring on this structure.
    ///
    /// # Panics
    /// Panics when trying to set the lock on an upgraded handle.
    pub(crate) unsafe fn set_lock(&self, val: bool) {
        self.liveliness.as_ref()
            .expect("Tried to set lock on borrowed Keyboard")
            .store(val, Ordering::Release);
    }
}

impl Drop for Keyboard {
    fn drop(&mut self) {
        match self.liveliness {
            None => {}
            Some(ref liveliness) => {
                if Rc::strong_count(liveliness) == 1 {
                    wlr_log!(L_DEBUG, "Dropped Keyboard {:p}", self.keyboard);
                    let weak_count = Rc::weak_count(liveliness);
                    if weak_count > 0 {
                        wlr_log!(L_DEBUG,
                                 "Still {} weak pointers to Keyboard {:p}",
                                 weak_count,
                                 self.keyboard);
                    }
                }
            }
        }
    }
}

impl fmt::Debug for Keyboard {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f,
               "Keyboard {{ device: {:?}, keyboard: {:?} }}",
               self.device, self.keyboard)
    }
}

impl KeyboardHandle {
    /// Upgrades the keyboard handle to a reference to the backing `Keyboard`.
    ///
    /// # Unsafety
    /// This function is unsafe, because it creates an unbounded `Keyboard`
    /// which may live forever..
    /// But no keyboard lives forever and might be disconnected at any time.
    ///
    /// # Panics
    /// This function will panic if multiple mutable borrows are detected.
    pub(crate) unsafe fn upgrade(&self) -> UpgradeHandleResult<Keyboard> {
        self.handle.upgrade()
            .ok_or(UpgradeHandleErr::DoubleUpgrade)
            // NOTE
            // We drop the Rc here because having two would allow a dangling
            // pointer to exist!
            .and_then(|check| {
                let keyboard = Keyboard::from_handle(self);
                if check.load(Ordering::Acquire) {
                    return Err(UpgradeHandleErr::AlreadyBorrowed)
                }
                check.store(true, Ordering::Release);
                Ok(keyboard)
            })
    }

    /// Run a function on the referenced Keyboard, if it still exists
    ///
    /// Returns the result of the function, if successful
    ///
    /// # Safety
    /// By enforcing a rather harsh limit on the lifetime of the output
    /// to a short lived scope of an anonymous function,
    /// this function ensures the Keyboard does not live longer
    /// than it exists.
    ///
    /// # Panics
    /// This function will panic if multiple mutable borrows are detected.
    /// This will happen if you call `upgrade` directly within this callback,
    /// or if you run this function within the another run to the same `Output`.
    ///
    /// So don't nest `run` calls and everything will be ok :).
    pub fn run<F, R>(&mut self, runner: F) -> UpgradeHandleResult<Option<R>>
        where F: FnOnce(&mut Keyboard) -> R
    {
        let mut keyboard = unsafe { self.upgrade()? };
        let res = panic::catch_unwind(panic::AssertUnwindSafe(|| Some(runner(&mut keyboard))));
        self.handle.upgrade().map(|check| {
                                      // Sanity check that it hasn't been tampered with.
                                      if !check.load(Ordering::Acquire) {
                                          wlr_log!(L_ERROR,
                                                   "After running keyboard callback, mutable \
                                                    lock was false for: {:?}",
                                                   keyboard);
                                          panic!("Lock in incorrect state!");
                                      }
                                      check.store(false, Ordering::Release);
                                  });
        match res {
            Ok(res) => Ok(res),
            Err(err) => panic::resume_unwind(err)
        }
    }

    /// Gets the wlr_input_device associated with this KeyboardHandle
    pub fn input_device(&self) -> &InputDevice {
        &self.device
    }

    /// Gets the wlr_keyboard associated with this KeyboardHandle.
    pub(crate) unsafe fn as_ptr(&self) -> *mut wlr_keyboard {
        self.keyboard
    }
}

impl fmt::Debug for KeyboardHandle {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f,
               "Keyboard {{ device: {:?}, keyboard: {:?} }}",
               self.device, self.keyboard)
    }
}

bitflags! {
    pub struct KeyboardLed: u32 {
        const WLR_LED_NUM_LOCK = wlr_keyboard_led::WLR_LED_NUM_LOCK as u32;
        const WLR_LED_CAPS_LOCK = wlr_keyboard_led::WLR_LED_CAPS_LOCK as u32;
        const WLR_LED_SCROLL_LOCK = wlr_keyboard_led::WLR_LED_SCROLL_LOCK as u32;
    }
}

bitflags! {
    pub struct KeyboardModifier: u32 {
        const WLR_MODIFIER_SHIFT = wlr_keyboard_modifier::WLR_MODIFIER_SHIFT as u32;
        const WLR_MODIFIER_CAPS = wlr_keyboard_modifier::WLR_MODIFIER_CAPS as u32;
        const WLR_MODIFIER_CTRL = wlr_keyboard_modifier::WLR_MODIFIER_CTRL as u32;
        const WLR_MODIFIER_ALT = wlr_keyboard_modifier::WLR_MODIFIER_ALT as u32;
        const WLR_MODIFIER_MOD2 = wlr_keyboard_modifier::WLR_MODIFIER_MOD2 as u32;
        const WLR_MODIFIER_MOD3 = wlr_keyboard_modifier::WLR_MODIFIER_MOD3 as u32;
        const WLR_MODIFIER_LOGO = wlr_keyboard_modifier::WLR_MODIFIER_LOGO as u32;
        const WLR_MODIFIER_MOD5 = wlr_keyboard_modifier::WLR_MODIFIER_MOD5 as u32;
    }
}

impl fmt::Display for KeyboardModifier {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let mod_vec = vec![("Shift", KeyboardModifier::WLR_MODIFIER_SHIFT),
                           ("Caps lock", KeyboardModifier::WLR_MODIFIER_CAPS),
                           ("Ctrl", KeyboardModifier::WLR_MODIFIER_CTRL),
                           ("Alt", KeyboardModifier::WLR_MODIFIER_ALT),
                           ("Mod2", KeyboardModifier::WLR_MODIFIER_MOD2),
                           ("Mod3", KeyboardModifier::WLR_MODIFIER_MOD3),
                           ("Logo", KeyboardModifier::WLR_MODIFIER_LOGO),
                           ("Mod5", KeyboardModifier::WLR_MODIFIER_MOD5)];

        let mods: Vec<&str> = mod_vec.into_iter()
                                     .filter(|&(_, flag)| self.contains(flag))
                                     .map(|(st, _)| st)
                                     .collect();

        write!(formatter, "{:?}", mods)
    }
}
<commit_msg>Added note on why we can pass a &Keymap<commit_after>//! TODO Documentation
use std::{fmt, panic};
use std::mem::ManuallyDrop;
use std::ops::Deref;
use std::rc::{Rc, Weak};
use std::sync::atomic::{AtomicBool, Ordering};

use errors::{UpgradeHandleErr, UpgradeHandleResult};
use wlroots_sys::{wlr_input_device, wlr_key_state, wlr_keyboard, wlr_keyboard_get_modifiers,
                  wlr_keyboard_led, wlr_keyboard_led_update, wlr_keyboard_modifier,
                  wlr_keyboard_set_keymap};

use xkbcommon::xkb::{self, Keycode, Keymap, LedIndex, ModIndex, ModMask};

use InputDevice;

pub type KeyState = wlr_key_state;

/// The masks for each group of modifiers.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct ModifierMasks {
    pub depressed: ModMask,
    pub latched: ModMask,
    pub locked: ModMask,
    pub group: ModMask
}

/// Information about repeated keypresses for a particular Keyboard.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct RepeatInfo {
    /// The rate at which extended keypresses will fire more events.
    rate: i32,
    /// How long it takes for a keypress to register on this device.
    delay: i32
}

pub struct Keyboard {
    /// The structure that ensures weak handles to this structure are still alive.
    ///
    /// They contain weak handles, and will safely not use dead memory when this
    /// is freed by wlroots.
    ///
    /// If this is `None`, then this is from an upgraded `KeyboardHandle`, and
    /// the operations are **unchecked**.
    /// This is means safe operations might fail, but only if you use the unsafe
    /// marked function `upgrade` on a `KeyboardHandle`.
    liveliness: Option<Rc<AtomicBool>>,
    /// The device that refers to this keyboard.
    device: InputDevice,
    /// The underlying keyboard data.
    keyboard: *mut wlr_keyboard,
    /// The keymap used by wlroots.
    /// Stored here so that we can ensure you can only borrow the keymap
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    keymap: Option<ManuallyDrop<Keymap>>,
    /// The keyboard state used by wlroots.
    /// Stored here so that we can ensure you can only borrow the state
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    #[allow(dead_code)]
    xkb_state: Option<ManuallyDrop<xkb::State>>
}

pub struct KeyboardHandle {
    /// The Rc that ensures that this handle is still alive.
    ///
    /// When wlroots deallocates the keyboard associated with this handle,
    handle: Weak<AtomicBool>,
    /// The device that refers to this keyboard.
    device: InputDevice,
    /// The underlying keyboard data.
    keyboard: *mut wlr_keyboard,
    /// The keymap used by wlroots.
    /// Stored here so that we can ensure you can only borrow the keymap
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    #[allow(dead_code)]
    keymap: Option<ManuallyDrop<Keymap>>,
    /// The keyboard state used by wlroots.
    /// Stored here so that we can ensure you can only borrow the state
    /// as long as the keyboard.
    ///
    /// NOTE We have it marked as manually drop because otherwise
    /// we'd unref it twice (once here and again in wlroots).
    #[allow(dead_code)]
    xkb_state: Option<ManuallyDrop<xkb::State>>
}

impl Keyboard {
    /// Tries to convert an input device to a Keyboard
    ///
    /// Returns None if it is of a different type of input variant.
    ///
    /// # Safety
    /// This creates a totally new Keyboard (e.g with its own reference count)
    /// so only do this once per `wlr_input_device`!
    pub(crate) unsafe fn new_from_input_device(device: *mut wlr_input_device) -> Option<Self> {
        use wlroots_sys::wlr_input_device_type::*;
        match (*device).type_ {
            WLR_INPUT_DEVICE_KEYBOARD => {
                let keyboard = (*device).__bindgen_anon_1.keyboard;
                Some(Keyboard { liveliness: Some(Rc::new(AtomicBool::new(false))),
                                device: InputDevice::from_ptr(device),
                                keyboard,
                                keymap: None,
                                xkb_state: None })
            }
            _ => None
        }
    }

    unsafe fn from_handle(handle: &KeyboardHandle) -> Self {
        Keyboard { liveliness: None,
                   device: handle.input_device().clone(),
                   keyboard: handle.as_ptr(),
                   keymap: None,
                   xkb_state: None }
    }

    /// Gets the wlr_keyboard associated with this KeyboardHandle.
    pub(crate) unsafe fn as_ptr(&self) -> *mut wlr_keyboard {
        self.keyboard
    }

    /// Gets the wlr_input_device associated with this KeyboardHandle
    pub fn input_device(&self) -> &InputDevice {
        &self.device
    }

    /// Set the keymap for this Keyboard.
    pub fn set_keymap(&mut self, keymap: &Keymap) {
        unsafe {
            // NOTE wlr_keyboard_set_keymap updates the reference count,
            // so we don't need to mem::forget the key map here
            // or take it by value.
            wlr_keyboard_set_keymap(self.keyboard, keymap.get_raw_ptr() as _);
        }
    }

    /// Get the XKB keymap associated with this Keyboard.
    pub fn get_keymap(&mut self) -> Option<&Keymap> {
        unsafe {
            let keymap_ptr = (*self.keyboard).keymap as *mut _;
            if keymap_ptr.is_null() {
                None
            } else {
                self.keymap = Some(ManuallyDrop::new(Keymap::from_raw_ptr(keymap_ptr)));
                self.keymap.as_ref().map(Deref::deref)
            }
        }
    }

    /// Get the keycodes for this keyboard as reported by XKB.
    ///
    /// # Limitations
    /// wlroots limits this list to `WLR_KEYBOARD_KEYS_CAP` elements,
    /// which at the time of writing is `32`.
    pub fn keycodes(&self) -> Vec<Keycode> {
        unsafe {
            let mut result = (*self.keyboard).keycodes.to_vec();
            result.truncate((*self.keyboard).num_keycodes);
            result
        }
    }

    /// Get the list of LEDs for this keyboard as reported by XKB.
    pub fn led_list(&self) -> Vec<LedIndex> {
        unsafe { (*self.keyboard).led_indexes.to_vec() }
    }

    /// Get the list of modifiers for this keyboard as reported by XKB.
    pub fn modifier_list(&self) -> Vec<ModIndex> {
        unsafe { (*self.keyboard).mod_indexes.to_vec() }
    }

    /// Get the size of the keymap.
    pub fn keymap_size(&self) -> usize {
        unsafe { (*self.keyboard).keymap_size }
    }

    /// Get the XKB state associated with this `Keyboard`.
    pub fn get_xkb_state(&mut self) -> Option<&xkb::State> {
        unsafe {
            let xkb_state_ptr = (*self.keyboard).xkb_state as *mut _;
            if xkb_state_ptr.is_null() {
                None
            } else {
                self.xkb_state = Some(ManuallyDrop::new(xkb::State::from_raw_ptr(xkb_state_ptr)));
                self.xkb_state.as_ref().map(Deref::deref)
            }
        }
    }

    /// Get the repeat info for this keyboard.
    pub fn repeat_info(&self) -> RepeatInfo {
        unsafe {
            RepeatInfo { rate: (*self.keyboard).repeat_info.rate,
                         delay: (*self.keyboard).repeat_info.delay }
        }
    }

    /// Update the LED lights using the provided bitmap.
    ///
    /// 1 means one, 0 means off.
    pub fn update_led(&mut self, leds: KeyboardLed) {
        unsafe {
            wlr_keyboard_led_update(self.keyboard, leds.bits() as u32);
        }
    }

    /// Get the modifiers that are currently pressed on the keyboard.
    pub fn get_modifiers(&self) -> KeyboardModifier {
        unsafe {
            KeyboardModifier::from_bits_truncate(wlr_keyboard_get_modifiers(self.keyboard))
        }
    }

    /// Get the modifier masks for each group.
    pub fn get_modifier_masks(&self) -> ModifierMasks {
        unsafe {
            ModifierMasks { depressed: (*self.keyboard).modifiers.depressed,
                            latched: (*self.keyboard).modifiers.latched,
                            locked: (*self.keyboard).modifiers.locked,
                            group: (*self.keyboard).modifiers.group }
        }
    }

    /// Creates a weak reference to a `Keyboard`.
    ///
    /// # Panics
    /// If this `Keyboard` is a previously upgraded `KeyboardHandle`,
    /// then this function will panic.
    pub fn weak_reference(&self) -> KeyboardHandle {
        let arc = self.liveliness.as_ref()
                      .expect("Cannot downgrade previously upgraded KeyboardHandle!");
        KeyboardHandle { handle: Rc::downgrade(arc),
                         // NOTE Rationale for cloning:
                         // We can't use the keyboard handle unless the keyboard is alive,
                         // which means the device pointer is still alive.
                         device: unsafe { self.device.clone() },
                         keyboard: self.keyboard,
                         keymap: None,
                         xkb_state: None }
    }

    /// Manually set the lock used to determine if a double-borrow is
    /// occuring on this structure.
    ///
    /// # Panics
    /// Panics when trying to set the lock on an upgraded handle.
    pub(crate) unsafe fn set_lock(&self, val: bool) {
        self.liveliness.as_ref()
            .expect("Tried to set lock on borrowed Keyboard")
            .store(val, Ordering::Release);
    }
}

impl Drop for Keyboard {
    fn drop(&mut self) {
        match self.liveliness {
            None => {}
            Some(ref liveliness) => {
                if Rc::strong_count(liveliness) == 1 {
                    wlr_log!(L_DEBUG, "Dropped Keyboard {:p}", self.keyboard);
                    let weak_count = Rc::weak_count(liveliness);
                    if weak_count > 0 {
                        wlr_log!(L_DEBUG,
                                 "Still {} weak pointers to Keyboard {:p}",
                                 weak_count,
                                 self.keyboard);
                    }
                }
            }
        }
    }
}

impl fmt::Debug for Keyboard {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f,
               "Keyboard {{ device: {:?}, keyboard: {:?} }}",
               self.device, self.keyboard)
    }
}

impl KeyboardHandle {
    /// Upgrades the keyboard handle to a reference to the backing `Keyboard`.
    ///
    /// # Unsafety
    /// This function is unsafe, because it creates an unbounded `Keyboard`
    /// which may live forever..
    /// But no keyboard lives forever and might be disconnected at any time.
    ///
    /// # Panics
    /// This function will panic if multiple mutable borrows are detected.
    pub(crate) unsafe fn upgrade(&self) -> UpgradeHandleResult<Keyboard> {
        self.handle.upgrade()
            .ok_or(UpgradeHandleErr::DoubleUpgrade)
            // NOTE
            // We drop the Rc here because having two would allow a dangling
            // pointer to exist!
            .and_then(|check| {
                let keyboard = Keyboard::from_handle(self);
                if check.load(Ordering::Acquire) {
                    return Err(UpgradeHandleErr::AlreadyBorrowed)
                }
                check.store(true, Ordering::Release);
                Ok(keyboard)
            })
    }

    /// Run a function on the referenced Keyboard, if it still exists
    ///
    /// Returns the result of the function, if successful
    ///
    /// # Safety
    /// By enforcing a rather harsh limit on the lifetime of the output
    /// to a short lived scope of an anonymous function,
    /// this function ensures the Keyboard does not live longer
    /// than it exists.
    ///
    /// # Panics
    /// This function will panic if multiple mutable borrows are detected.
    /// This will happen if you call `upgrade` directly within this callback,
    /// or if you run this function within the another run to the same `Output`.
    ///
    /// So don't nest `run` calls and everything will be ok :).
    pub fn run<F, R>(&mut self, runner: F) -> UpgradeHandleResult<Option<R>>
        where F: FnOnce(&mut Keyboard) -> R
    {
        let mut keyboard = unsafe { self.upgrade()? };
        let res = panic::catch_unwind(panic::AssertUnwindSafe(|| Some(runner(&mut keyboard))));
        self.handle.upgrade().map(|check| {
                                      // Sanity check that it hasn't been tampered with.
                                      if !check.load(Ordering::Acquire) {
                                          wlr_log!(L_ERROR,
                                                   "After running keyboard callback, mutable \
                                                    lock was false for: {:?}",
                                                   keyboard);
                                          panic!("Lock in incorrect state!");
                                      }
                                      check.store(false, Ordering::Release);
                                  });
        match res {
            Ok(res) => Ok(res),
            Err(err) => panic::resume_unwind(err)
        }
    }

    /// Gets the wlr_input_device associated with this KeyboardHandle
    pub fn input_device(&self) -> &InputDevice {
        &self.device
    }

    /// Gets the wlr_keyboard associated with this KeyboardHandle.
    pub(crate) unsafe fn as_ptr(&self) -> *mut wlr_keyboard {
        self.keyboard
    }
}

impl fmt::Debug for KeyboardHandle {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f,
               "Keyboard {{ device: {:?}, keyboard: {:?} }}",
               self.device, self.keyboard)
    }
}

bitflags! {
    pub struct KeyboardLed: u32 {
        const WLR_LED_NUM_LOCK = wlr_keyboard_led::WLR_LED_NUM_LOCK as u32;
        const WLR_LED_CAPS_LOCK = wlr_keyboard_led::WLR_LED_CAPS_LOCK as u32;
        const WLR_LED_SCROLL_LOCK = wlr_keyboard_led::WLR_LED_SCROLL_LOCK as u32;
    }
}

bitflags! {
    pub struct KeyboardModifier: u32 {
        const WLR_MODIFIER_SHIFT = wlr_keyboard_modifier::WLR_MODIFIER_SHIFT as u32;
        const WLR_MODIFIER_CAPS = wlr_keyboard_modifier::WLR_MODIFIER_CAPS as u32;
        const WLR_MODIFIER_CTRL = wlr_keyboard_modifier::WLR_MODIFIER_CTRL as u32;
        const WLR_MODIFIER_ALT = wlr_keyboard_modifier::WLR_MODIFIER_ALT as u32;
        const WLR_MODIFIER_MOD2 = wlr_keyboard_modifier::WLR_MODIFIER_MOD2 as u32;
        const WLR_MODIFIER_MOD3 = wlr_keyboard_modifier::WLR_MODIFIER_MOD3 as u32;
        const WLR_MODIFIER_LOGO = wlr_keyboard_modifier::WLR_MODIFIER_LOGO as u32;
        const WLR_MODIFIER_MOD5 = wlr_keyboard_modifier::WLR_MODIFIER_MOD5 as u32;
    }
}

impl fmt::Display for KeyboardModifier {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let mod_vec = vec![("Shift", KeyboardModifier::WLR_MODIFIER_SHIFT),
                           ("Caps lock", KeyboardModifier::WLR_MODIFIER_CAPS),
                           ("Ctrl", KeyboardModifier::WLR_MODIFIER_CTRL),
                           ("Alt", KeyboardModifier::WLR_MODIFIER_ALT),
                           ("Mod2", KeyboardModifier::WLR_MODIFIER_MOD2),
                           ("Mod3", KeyboardModifier::WLR_MODIFIER_MOD3),
                           ("Logo", KeyboardModifier::WLR_MODIFIER_LOGO),
                           ("Mod5", KeyboardModifier::WLR_MODIFIER_MOD5)];

        let mods: Vec<&str> = mod_vec.into_iter()
                                     .filter(|&(_, flag)| self.contains(flag))
                                     .map(|(st, _)| st)
                                     .collect();

        write!(formatter, "{:?}", mods)
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

#[macro_use]
mod macros;
pub(crate) mod actor;
mod download;
mod error;
mod idle;
pub(crate) mod install;
mod park;
mod poll;
mod prepare_download;
mod probe;
mod reboot;
mod transition;

use self::{
    download::Download,
    error::Error,
    idle::Idle,
    install::Install,
    park::Park,
    poll::Poll,
    prepare_download::PrepareDownload,
    probe::{Probe, ServerAddress},
    reboot::Reboot,
};
use crate::{firmware::Metadata, http_api, runtime_settings::RuntimeSettings, settings::Settings};
use actix::{Actor, System};
use futures::future::Future;
use slog_scope::error;

trait StateChangeImpl {
    fn handle(self, shared_state: &mut SharedState) -> Result<StateMachine, failure::Error>;
    fn name(&self) -> &'static str;

    fn handle_download_abort(&self) -> actor::download_abort::Response {
        actor::download_abort::Response::InvalidState
    }

    fn handle_trigger_probe(&self) -> actor::probe::Response {
        actor::probe::Response::InvalidState(self.name().to_owned())
    }
}

trait TransitionCallback: StateChangeImpl + Into<State<Idle>> {}

trait ProgressReporter: TransitionCallback {
    fn package_uid(&self) -> String;
    fn report_enter_state_name(&self) -> &'static str;
    fn report_leave_state_name(&self) -> &'static str;
}

#[derive(Debug, PartialEq)]
struct State<S>(S)
where
    State<S>: StateChangeImpl;

#[derive(Debug, PartialEq)]
struct SharedState {
    settings: Settings,
    runtime_settings: RuntimeSettings,
    firmware: Metadata,
}

#[derive(Debug, PartialEq)]
enum StateMachine {
    Park(State<Park>),
    Idle(State<Idle>),
    Poll(State<Poll>),
    Probe(State<Probe>),
    PrepareDownload(State<PrepareDownload>),
    Download(State<Download>),
    Install(State<Install>),
    Reboot(State<Reboot>),
    Error(State<Error>),
}

impl<S> State<S>
where
    State<S>: TransitionCallback + ProgressReporter,
{
    fn handle_with_callback_and_report_progress(
        self,
        shared_state: &mut SharedState,
    ) -> Result<StateMachine, failure::Error> {
        use transition::{state_change_callback, Transition};

        let transition =
            state_change_callback(&shared_state.settings.firmware.metadata_path, self.name())?;

        match transition {
            Transition::Continue => Ok(self.handle_and_report_progress(shared_state)?),
            Transition::Cancel => Ok(StateMachine::Idle(self.into())),
        }
    }
}

impl<S> State<S>
where
    State<S>: ProgressReporter,
{
    fn handle_and_report_progress(
        self,
        shared_state: &mut SharedState,
    ) -> Result<StateMachine, failure::Error> {
        let server = &shared_state.settings.network.server_address.clone();
        let firmware = &shared_state.firmware.clone();
        let package_uid = &self.package_uid();
        let enter_state = self.report_enter_state_name();
        let leave_state = self.report_leave_state_name();

        let report = |state, previous_state, error_message, current_log| {
            crate::client::Api::new(&server).report(
                state,
                firmware,
                package_uid,
                previous_state,
                error_message,
                current_log,
            )
        };

        report(enter_state, None, None, None)?;
        self.handle(shared_state)
            .and_then(|state| {
                report(leave_state, None, None, None)?;
                Ok(state)
            })
            .or_else(|e| {
                report(
                    "error",
                    Some(enter_state),
                    Some(e.to_string()),
                    Some(crate::logger::buffer().lock().unwrap().to_string()),
                )?;
                Err(e)
            })
    }
}

impl StateMachine {
    fn new() -> Self {
        StateMachine::Idle(State(Idle {}))
    }

    fn move_to_next_state(self, shared_state: &mut SharedState) -> Result<Self, failure::Error> {
        match self {
            StateMachine::Error(s) => Ok(s.handle(shared_state)?),
            StateMachine::Park(s) => Ok(s.handle(shared_state)?),
            StateMachine::Idle(s) => Ok(s.handle(shared_state)?),
            StateMachine::Poll(s) => Ok(s.handle(shared_state)?),
            StateMachine::Probe(s) => Ok(s.handle(shared_state)?),
            StateMachine::PrepareDownload(s) => Ok(s.handle(shared_state)?),
            StateMachine::Download(s) => {
                Ok(s.handle_with_callback_and_report_progress(shared_state)?)
            }
            StateMachine::Install(s) => {
                Ok(s.handle_with_callback_and_report_progress(shared_state)?)
            }
            StateMachine::Reboot(s) => Ok(s.handle_with_callback_and_report_progress(shared_state)?),
        }
    }

    fn for_any_state<F, A>(&self, f: F) -> A
    where
        F: Fn(&dyn StateChangeImpl) -> A,
    {
        match self {
            StateMachine::Error(s) => f(s),
            StateMachine::Park(s) => f(s),
            StateMachine::Idle(s) => f(s),
            StateMachine::Poll(s) => f(s),
            StateMachine::Probe(s) => f(s),
            StateMachine::PrepareDownload(s) => f(s),
            StateMachine::Download(s) => f(s),
            StateMachine::Install(s) => f(s),
            StateMachine::Reboot(s) => f(s),
        }
    }
}

/// Runs the state machine up to completion handling all procing
/// states without extra manual work.
///
/// It supports following states, and transitions, as shown in the
/// below diagram:
///
/// ```text
///           .--------------.
///           |              v
/// Park <- Idle -> Poll -> Probe -> Download -> Install -> Reboot
///           ^      ^        '          '          '
///           '      '        '          '          '
///           '      `--------'          '          '
///           `---------------'          '          '
///           `--------------------------'          '
///           `-------------------------------------'
/// ```
///
/// # Example
/// ```
/// # extern crate failure;
/// # extern crate updatehub;
/// # use failure;
/// # fn run() -> Result<(), failure::Error> {
/// use updatehub;
///
/// updatehub::logger::init(0);
/// let settings = updatehub::Settings::load()?;
/// updatehub::run(settings)?;
/// # Ok(())
/// # }
/// ```
pub fn run(settings: Settings) -> Result<(), failure::Error> {
    let listen_socket = settings.network.listen_socket.clone();
    let mut runtime_settings = RuntimeSettings::new().load(&settings.storage.runtime_settings)?;
    if !settings.storage.read_only {
        runtime_settings.enable_persistency();
    }

    let firmware = Metadata::from_path(&settings.firmware.metadata_path)?;
    let agent_machine = actor::Machine::new(
        StateMachine::new(),
        SharedState {
            settings,
            runtime_settings,
            firmware,
        },
    );

    System::run(move || {
        let machine = agent_machine.start();
        let api = machine.clone();
        actix_web::HttpServer::new(move || {
            actix_web::App::new().configure(|cfg| http_api::API::configure(cfg, api.clone()))
        })
        .bind(listen_socket.clone())
        .expect(&format!(
            "Failed to bind listen socket, {:?}, for HTTP API",
            listen_socket,
        ))
        .start();

        // Iterate over the state machine on a separated thread.
        actix::Arbiter::new().exec_fn(move || {
            while machine.connected() {
                if let Err(e) = machine.send(actor::Step).wait() {
                    error!("Communication to actor failed: {:?}", e);
                }
            }
        });
    })?;

    unreachable!("actix System has stopped");
}
<commit_msg>updatehub: Run StateMachine Actor on a different arbiter<commit_after>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

#[macro_use]
mod macros;
pub(crate) mod actor;
mod download;
mod error;
mod idle;
pub(crate) mod install;
mod park;
mod poll;
mod prepare_download;
mod probe;
mod reboot;
mod transition;

use self::{
    download::Download,
    error::Error,
    idle::Idle,
    install::Install,
    park::Park,
    poll::Poll,
    prepare_download::PrepareDownload,
    probe::{Probe, ServerAddress},
    reboot::Reboot,
};
use crate::{firmware::Metadata, http_api, runtime_settings::RuntimeSettings, settings::Settings};
use actix::{Actor, System};
use futures::future::Future;
use slog_scope::{error, info};

trait StateChangeImpl {
    fn handle(self, shared_state: &mut SharedState) -> Result<StateMachine, failure::Error>;
    fn name(&self) -> &'static str;

    fn handle_download_abort(&self) -> actor::download_abort::Response {
        actor::download_abort::Response::InvalidState
    }

    fn handle_trigger_probe(&self) -> actor::probe::Response {
        actor::probe::Response::InvalidState(self.name().to_owned())
    }
}

trait TransitionCallback: StateChangeImpl + Into<State<Idle>> {}

trait ProgressReporter: TransitionCallback {
    fn package_uid(&self) -> String;
    fn report_enter_state_name(&self) -> &'static str;
    fn report_leave_state_name(&self) -> &'static str;
}

#[derive(Debug, PartialEq)]
struct State<S>(S)
where
    State<S>: StateChangeImpl;

#[derive(Debug, PartialEq)]
struct SharedState {
    settings: Settings,
    runtime_settings: RuntimeSettings,
    firmware: Metadata,
}

#[derive(Debug, PartialEq)]
enum StateMachine {
    Park(State<Park>),
    Idle(State<Idle>),
    Poll(State<Poll>),
    Probe(State<Probe>),
    PrepareDownload(State<PrepareDownload>),
    Download(State<Download>),
    Install(State<Install>),
    Reboot(State<Reboot>),
    Error(State<Error>),
}

impl<S> State<S>
where
    State<S>: TransitionCallback + ProgressReporter,
{
    fn handle_with_callback_and_report_progress(
        self,
        shared_state: &mut SharedState,
    ) -> Result<StateMachine, failure::Error> {
        use transition::{state_change_callback, Transition};

        let transition =
            state_change_callback(&shared_state.settings.firmware.metadata_path, self.name())?;

        match transition {
            Transition::Continue => Ok(self.handle_and_report_progress(shared_state)?),
            Transition::Cancel => Ok(StateMachine::Idle(self.into())),
        }
    }
}

impl<S> State<S>
where
    State<S>: ProgressReporter,
{
    fn handle_and_report_progress(
        self,
        shared_state: &mut SharedState,
    ) -> Result<StateMachine, failure::Error> {
        let server = &shared_state.settings.network.server_address.clone();
        let firmware = &shared_state.firmware.clone();
        let package_uid = &self.package_uid();
        let enter_state = self.report_enter_state_name();
        let leave_state = self.report_leave_state_name();

        let report = |state, previous_state, error_message, current_log| {
            crate::client::Api::new(&server).report(
                state,
                firmware,
                package_uid,
                previous_state,
                error_message,
                current_log,
            )
        };

        report(enter_state, None, None, None)?;
        self.handle(shared_state)
            .and_then(|state| {
                report(leave_state, None, None, None)?;
                Ok(state)
            })
            .or_else(|e| {
                report(
                    "error",
                    Some(enter_state),
                    Some(e.to_string()),
                    Some(crate::logger::buffer().lock().unwrap().to_string()),
                )?;
                Err(e)
            })
    }
}

impl StateMachine {
    fn new() -> Self {
        StateMachine::Idle(State(Idle {}))
    }

    fn move_to_next_state(self, shared_state: &mut SharedState) -> Result<Self, failure::Error> {
        match self {
            StateMachine::Error(s) => Ok(s.handle(shared_state)?),
            StateMachine::Park(s) => Ok(s.handle(shared_state)?),
            StateMachine::Idle(s) => Ok(s.handle(shared_state)?),
            StateMachine::Poll(s) => Ok(s.handle(shared_state)?),
            StateMachine::Probe(s) => Ok(s.handle(shared_state)?),
            StateMachine::PrepareDownload(s) => Ok(s.handle(shared_state)?),
            StateMachine::Download(s) => {
                Ok(s.handle_with_callback_and_report_progress(shared_state)?)
            }
            StateMachine::Install(s) => {
                Ok(s.handle_with_callback_and_report_progress(shared_state)?)
            }
            StateMachine::Reboot(s) => Ok(s.handle_with_callback_and_report_progress(shared_state)?),
        }
    }

    fn for_any_state<F, A>(&self, f: F) -> A
    where
        F: Fn(&dyn StateChangeImpl) -> A,
    {
        match self {
            StateMachine::Error(s) => f(s),
            StateMachine::Park(s) => f(s),
            StateMachine::Idle(s) => f(s),
            StateMachine::Poll(s) => f(s),
            StateMachine::Probe(s) => f(s),
            StateMachine::PrepareDownload(s) => f(s),
            StateMachine::Download(s) => f(s),
            StateMachine::Install(s) => f(s),
            StateMachine::Reboot(s) => f(s),
        }
    }
}

/// Runs the state machine up to completion handling all procing
/// states without extra manual work.
///
/// It supports following states, and transitions, as shown in the
/// below diagram:
///
/// ```text
///           .--------------.
///           |              v
/// Park <- Idle -> Poll -> Probe -> Download -> Install -> Reboot
///           ^      ^        '          '          '
///           '      '        '          '          '
///           '      `--------'          '          '
///           `---------------'          '          '
///           `--------------------------'          '
///           `-------------------------------------'
/// ```
///
/// # Example
/// ```
/// # extern crate failure;
/// # extern crate updatehub;
/// # use failure;
/// # fn run() -> Result<(), failure::Error> {
/// use updatehub;
///
/// updatehub::logger::init(0);
/// let settings = updatehub::Settings::load()?;
/// updatehub::run(settings)?;
/// # Ok(())
/// # }
/// ```
pub fn run(settings: Settings) -> Result<(), failure::Error> {
    let listen_socket = settings.network.listen_socket.clone();
    let mut runtime_settings = RuntimeSettings::new().load(&settings.storage.runtime_settings)?;
    if !settings.storage.read_only {
        runtime_settings.enable_persistency();
    }
    let firmware = Metadata::from_path(&settings.firmware.metadata_path)?;

    System::run(move || {
        let machine_addr = actor::Machine::start_in_arbiter(&actix::Arbiter::new(), move |_| {
            actor::Machine::new(
                StateMachine::new(),
                SharedState {
                    settings,
                    runtime_settings,
                    firmware,
                },
            )
        });

        let addr_clone = machine_addr.clone();
        actix_web::HttpServer::new(move || {
            actix_web::App::new().configure(|cfg| http_api::API::configure(cfg, addr_clone.clone()))
        })
        .bind(listen_socket.clone())
        .expect(&format!(
            "Failed to bind listen socket, {:?}, for HTTP API",
            listen_socket,
        ))
        .start();

        // Iterate over the state machine on a separated thread.
        actix::Arbiter::new().exec_fn(move || {
            while machine_addr.connected() {
                if let Err(e) = machine_addr.send(actor::Step).wait() {
                    error!("Communication to actor failed: {:?}", e);
                }
            }
            System::current().stop();
        });
    })?;

    info!("actix System has stopped");
    Ok(())
}
<|endoftext|>


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

//! Handles operations using platform Time Stamp Counter (TSC).

// TODO(b/213149158): Remove after uses are added.
#![allow(dead_code)]

use anyhow::{anyhow, Result};
use base::{error, info};
use lazy_static::lazy_static;
use std::arch::x86_64::_rdtsc;

mod calibrate;
mod cpuid;
mod grouping;

pub use calibrate::*;
pub use cpuid::*;

fn rdtsc_safe() -> u64 {
    // Safe because _rdtsc takes no arguments
    unsafe { _rdtsc() }
}

// Singleton for getting the state of the host TSCs, to avoid calibrating multiple times.
lazy_static! {
    static ref TSC_STATE: Option<TscState> = match calibrate_tsc_state(rdtsc_safe) {
        Ok(tsc_state) => {
            info!("Using calibrated tsc frequency: {} Hz", tsc_state.frequency);
            for (core, offset) in tsc_state.offsets.iter().enumerate() {
                info!("Core {} has tsc offset of {:?} ns", core, offset);
            }
            Some(tsc_state)
        }
        Err(e) => {
            error!("Failed to calibrate tsc state: {:#}", e);
            None
        }
    };
}

/// Returns the frequency of the host TSC. Uses the TSC_STATE lazy_static ref, so the calibration
/// only happens once.
pub fn tsc_frequency() -> Result<u64> {
    let state = TSC_STATE
        .as_ref()
        .ok_or(anyhow!("TSC calibration failed"))?;
    Ok(state.frequency)
}

/// Returns the state of the host TSCs. Uses the TSC_STATE lazy_static ref, so the calibration
/// only happens once.
pub fn tsc_state() -> Result<TscState> {
    Ok(TSC_STATE
        .as_ref()
        .ok_or(anyhow!("TSC calibration failed"))?
        .clone())
}

#[derive(Default, Debug)]
pub struct TscSyncMitigations {
    /// Vec of per-vcpu affinities to apply to each vcpu thread. If None, no affinity should be
    /// applied.
    pub affinities: Vec<Option<Vec<usize>>>,
    /// Vec of TSC offsets to set on each vcpu. If None, no offset should be applied.
    pub offsets: Vec<Option<u64>>,
}

impl TscSyncMitigations {
    fn new(num_vcpus: usize) -> Self {
        TscSyncMitigations {
            affinities: vec![None; num_vcpus],
            offsets: vec![None; num_vcpus],
        }
    }

    pub fn get_vcpu_affinity(&self, cpu_id: usize) -> Option<Vec<usize>> {
        self.affinities.get(cpu_id).unwrap().clone()
    }

    pub fn get_vcpu_tsc_offset(&self, cpu_id: usize) -> Option<u64> {
        *self.offsets.get(cpu_id).unwrap()
    }
}

/// Given the state of the host TSCs in `tsc_state`, and the number of vcpus that are intended to
/// be run, return a set of affinities and TSC offsets to apply to those vcpus.
pub fn get_tsc_sync_mitigations(tsc_state: &TscState, num_vcpus: usize) -> TscSyncMitigations {
    tsc_sync_mitigations_inner(tsc_state, num_vcpus, rdtsc_safe)
}

fn tsc_sync_mitigations_inner(
    tsc_state: &TscState,
    num_vcpus: usize,
    rdtsc: fn() -> u64,
) -> TscSyncMitigations {
    let mut mitigations = TscSyncMitigations::new(num_vcpus);
    // If there's only one core grouping that means all the TSCs are in sync and no mitigations are
    // needed.
    if tsc_state.core_grouping.size() == 1 {
        return mitigations;
    }

    let largest_group = tsc_state.core_grouping.largest_group();
    let num_cores = tsc_state.offsets.len();

    // If the largest core group is larger than the number of vcpus, just pin all vcpus to that core
    // group, and no need to set offsets.
    if largest_group.cores.len() >= num_vcpus {
        let affinity: Vec<usize> = largest_group.cores.iter().map(|core| core.core).collect();
        for i in 0..num_vcpus {
            mitigations.affinities[i] = Some(affinity.clone());
        }
    } else {
        // Otherwise, we pin each vcpu to a core and set it's offset to compensate.
        let host_tsc_now = rdtsc();

        for i in 0..num_vcpus {
            // This handles the case where num_vcpus > num_cores, even though we try to avoid that in
            // practice.
            let pinned_core = i % num_cores;

            mitigations.affinities[i] = Some(vec![pinned_core]);
            // The guest TSC value is calculated like so:
            //   host_tsc + tsc_offset = guest_tsc
            // If we assume that each host core has it's own error (core_offset), then it's more
            // like this:
            //   host_tsc + core_offset + tsc_offset = guest_tsc
            // We want guest_tsc to be 0 at boot, so the formula is this:
            //   host_tsc + core_offset + tsc_offset = 0
            // and then you subtract host_tsc and core_offset from both sides and you get:
            //   tsc_offset = 0 - host_tsc - core_offset
            mitigations.offsets[i] = Some(
                0u64.wrapping_sub(host_tsc_now)
                    // Note: wrapping_add and casting tsc_state from an i64 to a u64 should be the
                    //  same as using the future wrapping_add_signed function, which is only in
                    //  nightly. This should be switched to using wrapping_add_signed once that is
                    //  in stable.
                    .wrapping_add(tsc_state.offsets[pinned_core].wrapping_neg() as i64 as u64),
            );
        }
    }

    mitigations
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tsc::grouping::{CoreGroup, CoreGrouping, CoreOffset};
    use std::time::Duration;

    #[test]
    fn test_sync_mitigation_set_offsets() {
        let offsets = vec![0, 1000, -1000, 2000];
        // frequency of 1GHz means 20 nanos is 20 ticks
        let state = TscState::new(1_000_000_000, offsets, Duration::from_nanos(20))
            .expect("TscState::new should not fail for this test");

        assert_eq!(
            state.core_grouping,
            CoreGrouping::new(vec![
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 2,
                        offset: -1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset { core: 0, offset: 0 }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 1,
                        offset: 1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 3,
                        offset: 2000
                    }]
                },
            ])
            .expect("CoreGrouping::new should not fail here")
        );

        fn fake_rdtsc() -> u64 {
            u64::MAX
        }

        let mitigations = tsc_sync_mitigations_inner(&state, 4, fake_rdtsc);

        // core offsets are:
        //  - core 0: has an offset of 0, so TSC offset = 0 - u64::MAX - 0 = 1
        //  - core 1: has an offset of 1000, so TSC offset = 0 - u64::MAX - 1000 = -999
        //  - core 2: has an offset of -1000, so TSC offset = 0 - u64::MAX + 1000 = 1001
        //  - core 3: has an offset of 2000, so TSC offset = 0 - u64::MAX - 2000 = -1999
        let expected = vec![1, 1u64.wrapping_sub(1000), 1001u64, 1u64.wrapping_sub(2000)];

        for (i, expect) in expected.iter().enumerate() {
            assert_eq!(
                mitigations
                    .get_vcpu_tsc_offset(i)
                    .unwrap_or_else(|| panic!("core {} should have an offset of {}", i, expect)),
                *expect
            );

            assert_eq!(
                mitigations
                    .get_vcpu_affinity(i)
                    .unwrap_or_else(|| panic!("core {} should have an affinity of [{}]", i, i)),
                vec![i]
            );
        }
    }

    #[test]
    fn test_sync_mitigation_large_group() {
        // 8 cores, and cores 1,3,5,7 are in-sync at offset -1000
        let offsets = vec![0, -1000, 1000, -1000, 2000, -1000, 3000, -1000];
        // frequency of 1GHz means 20 nanos is 20 ticks
        let state = TscState::new(1_000_000_000, offsets, Duration::from_nanos(20))
            .expect("TscState::new should not fail for this test");

        assert_eq!(
            state.core_grouping,
            CoreGrouping::new(vec![
                CoreGroup {
                    cores: vec![
                        CoreOffset {
                            core: 1,
                            offset: -1000
                        },
                        CoreOffset {
                            core: 3,
                            offset: -1000
                        },
                        CoreOffset {
                            core: 5,
                            offset: -1000
                        },
                        CoreOffset {
                            core: 7,
                            offset: -1000
                        }
                    ]
                },
                CoreGroup {
                    cores: vec![CoreOffset { core: 0, offset: 0 }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 2,
                        offset: 1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 4,
                        offset: 2000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 6,
                        offset: 3000
                    }]
                },
            ])
            .expect("CoreGrouping::new should not fail here")
        );

        fn fake_rdtsc() -> u64 {
            u64::MAX
        }

        let num_vcpus = 4;
        let mitigations = tsc_sync_mitigations_inner(&state, num_vcpus, fake_rdtsc);

        let expected_affinity = vec![1, 3, 5, 7];
        for i in 0..num_vcpus {
            assert_eq!(
                mitigations.get_vcpu_affinity(i).unwrap_or_else(|| panic!(
                    "core {} should have an affinity of {:?}",
                    i, expected_affinity
                )),
                expected_affinity
            );
            assert_eq!(mitigations.get_vcpu_tsc_offset(i), None);
        }
    }

    #[test]
    fn more_vcpus_than_cores() {
        // 4 cores, two can be grouped but it doesn't matter because we'll have more vcpus than
        // the largest group.
        let offsets = vec![0, 0, 1000, 2000];
        // frequency of 1GHz means 20 nanos is 20 ticks
        let state = TscState::new(1_000_000_000, offsets, Duration::from_nanos(20))
            .expect("TscState::new should not fail for this test");

        assert_eq!(
            state.core_grouping,
            CoreGrouping::new(vec![
                CoreGroup {
                    cores: vec![
                        CoreOffset { core: 0, offset: 0 },
                        CoreOffset { core: 1, offset: 0 }
                    ]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 2,
                        offset: 1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 3,
                        offset: 2000
                    }]
                },
            ])
            .expect("CoreGrouping::new should not fail here")
        );

        fn fake_rdtsc() -> u64 {
            u64::MAX
        }

        // 8 vcpus, more than we have cores
        let num_vcpus = 8;
        let mitigations = tsc_sync_mitigations_inner(&state, num_vcpus, fake_rdtsc);
        let expected_offsets = vec![1, 1, 1u64.wrapping_sub(1000), 1u64.wrapping_sub(2000)];

        for i in 0..num_vcpus {
            assert_eq!(
                mitigations.get_vcpu_affinity(i).unwrap_or_else(|| panic!(
                    "core {} should have an affinity of {:?}",
                    i,
                    i % 4
                )),
                // expected affinity is the vcpu modulo 4
                vec![i % 4]
            );
            assert_eq!(
                mitigations.get_vcpu_tsc_offset(i).unwrap_or_else(|| panic!(
                    "core {} should have an offset of {:?}",
                    i,
                    expected_offsets[i % 4]
                )),
                expected_offsets[i % 4]
            );
        }
    }
}
<commit_msg>devices: tsc: move logs to debug level<commit_after>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Handles operations using platform Time Stamp Counter (TSC).

// TODO(b/213149158): Remove after uses are added.
#![allow(dead_code)]

use anyhow::{anyhow, Result};
use base::{debug, error};
use lazy_static::lazy_static;
use std::arch::x86_64::_rdtsc;

mod calibrate;
mod cpuid;
mod grouping;

pub use calibrate::*;
pub use cpuid::*;

fn rdtsc_safe() -> u64 {
    // Safe because _rdtsc takes no arguments
    unsafe { _rdtsc() }
}

// Singleton for getting the state of the host TSCs, to avoid calibrating multiple times.
lazy_static! {
    static ref TSC_STATE: Option<TscState> = match calibrate_tsc_state(rdtsc_safe) {
        Ok(tsc_state) => {
            debug!("Using calibrated tsc frequency: {} Hz", tsc_state.frequency);
            for (core, offset) in tsc_state.offsets.iter().enumerate() {
                debug!("Core {} has tsc offset of {:?} ns", core, offset);
            }
            Some(tsc_state)
        }
        Err(e) => {
            error!("Failed to calibrate tsc state: {:#}", e);
            None
        }
    };
}

/// Returns the frequency of the host TSC. Uses the TSC_STATE lazy_static ref, so the calibration
/// only happens once.
pub fn tsc_frequency() -> Result<u64> {
    let state = TSC_STATE
        .as_ref()
        .ok_or(anyhow!("TSC calibration failed"))?;
    Ok(state.frequency)
}

/// Returns the state of the host TSCs. Uses the TSC_STATE lazy_static ref, so the calibration
/// only happens once.
pub fn tsc_state() -> Result<TscState> {
    Ok(TSC_STATE
        .as_ref()
        .ok_or(anyhow!("TSC calibration failed"))?
        .clone())
}

#[derive(Default, Debug)]
pub struct TscSyncMitigations {
    /// Vec of per-vcpu affinities to apply to each vcpu thread. If None, no affinity should be
    /// applied.
    pub affinities: Vec<Option<Vec<usize>>>,
    /// Vec of TSC offsets to set on each vcpu. If None, no offset should be applied.
    pub offsets: Vec<Option<u64>>,
}

impl TscSyncMitigations {
    fn new(num_vcpus: usize) -> Self {
        TscSyncMitigations {
            affinities: vec![None; num_vcpus],
            offsets: vec![None; num_vcpus],
        }
    }

    pub fn get_vcpu_affinity(&self, cpu_id: usize) -> Option<Vec<usize>> {
        self.affinities.get(cpu_id).unwrap().clone()
    }

    pub fn get_vcpu_tsc_offset(&self, cpu_id: usize) -> Option<u64> {
        *self.offsets.get(cpu_id).unwrap()
    }
}

/// Given the state of the host TSCs in `tsc_state`, and the number of vcpus that are intended to
/// be run, return a set of affinities and TSC offsets to apply to those vcpus.
pub fn get_tsc_sync_mitigations(tsc_state: &TscState, num_vcpus: usize) -> TscSyncMitigations {
    tsc_sync_mitigations_inner(tsc_state, num_vcpus, rdtsc_safe)
}

fn tsc_sync_mitigations_inner(
    tsc_state: &TscState,
    num_vcpus: usize,
    rdtsc: fn() -> u64,
) -> TscSyncMitigations {
    let mut mitigations = TscSyncMitigations::new(num_vcpus);
    // If there's only one core grouping that means all the TSCs are in sync and no mitigations are
    // needed.
    if tsc_state.core_grouping.size() == 1 {
        return mitigations;
    }

    let largest_group = tsc_state.core_grouping.largest_group();
    let num_cores = tsc_state.offsets.len();

    // If the largest core group is larger than the number of vcpus, just pin all vcpus to that core
    // group, and no need to set offsets.
    if largest_group.cores.len() >= num_vcpus {
        let affinity: Vec<usize> = largest_group.cores.iter().map(|core| core.core).collect();
        for i in 0..num_vcpus {
            mitigations.affinities[i] = Some(affinity.clone());
        }
    } else {
        // Otherwise, we pin each vcpu to a core and set it's offset to compensate.
        let host_tsc_now = rdtsc();

        for i in 0..num_vcpus {
            // This handles the case where num_vcpus > num_cores, even though we try to avoid that in
            // practice.
            let pinned_core = i % num_cores;

            mitigations.affinities[i] = Some(vec![pinned_core]);
            // The guest TSC value is calculated like so:
            //   host_tsc + tsc_offset = guest_tsc
            // If we assume that each host core has it's own error (core_offset), then it's more
            // like this:
            //   host_tsc + core_offset + tsc_offset = guest_tsc
            // We want guest_tsc to be 0 at boot, so the formula is this:
            //   host_tsc + core_offset + tsc_offset = 0
            // and then you subtract host_tsc and core_offset from both sides and you get:
            //   tsc_offset = 0 - host_tsc - core_offset
            mitigations.offsets[i] = Some(
                0u64.wrapping_sub(host_tsc_now)
                    // Note: wrapping_add and casting tsc_state from an i64 to a u64 should be the
                    //  same as using the future wrapping_add_signed function, which is only in
                    //  nightly. This should be switched to using wrapping_add_signed once that is
                    //  in stable.
                    .wrapping_add(tsc_state.offsets[pinned_core].wrapping_neg() as i64 as u64),
            );
        }
    }

    mitigations
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tsc::grouping::{CoreGroup, CoreGrouping, CoreOffset};
    use std::time::Duration;

    #[test]
    fn test_sync_mitigation_set_offsets() {
        let offsets = vec![0, 1000, -1000, 2000];
        // frequency of 1GHz means 20 nanos is 20 ticks
        let state = TscState::new(1_000_000_000, offsets, Duration::from_nanos(20))
            .expect("TscState::new should not fail for this test");

        assert_eq!(
            state.core_grouping,
            CoreGrouping::new(vec![
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 2,
                        offset: -1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset { core: 0, offset: 0 }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 1,
                        offset: 1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 3,
                        offset: 2000
                    }]
                },
            ])
            .expect("CoreGrouping::new should not fail here")
        );

        fn fake_rdtsc() -> u64 {
            u64::MAX
        }

        let mitigations = tsc_sync_mitigations_inner(&state, 4, fake_rdtsc);

        // core offsets are:
        //  - core 0: has an offset of 0, so TSC offset = 0 - u64::MAX - 0 = 1
        //  - core 1: has an offset of 1000, so TSC offset = 0 - u64::MAX - 1000 = -999
        //  - core 2: has an offset of -1000, so TSC offset = 0 - u64::MAX + 1000 = 1001
        //  - core 3: has an offset of 2000, so TSC offset = 0 - u64::MAX - 2000 = -1999
        let expected = vec![1, 1u64.wrapping_sub(1000), 1001u64, 1u64.wrapping_sub(2000)];

        for (i, expect) in expected.iter().enumerate() {
            assert_eq!(
                mitigations
                    .get_vcpu_tsc_offset(i)
                    .unwrap_or_else(|| panic!("core {} should have an offset of {}", i, expect)),
                *expect
            );

            assert_eq!(
                mitigations
                    .get_vcpu_affinity(i)
                    .unwrap_or_else(|| panic!("core {} should have an affinity of [{}]", i, i)),
                vec![i]
            );
        }
    }

    #[test]
    fn test_sync_mitigation_large_group() {
        // 8 cores, and cores 1,3,5,7 are in-sync at offset -1000
        let offsets = vec![0, -1000, 1000, -1000, 2000, -1000, 3000, -1000];
        // frequency of 1GHz means 20 nanos is 20 ticks
        let state = TscState::new(1_000_000_000, offsets, Duration::from_nanos(20))
            .expect("TscState::new should not fail for this test");

        assert_eq!(
            state.core_grouping,
            CoreGrouping::new(vec![
                CoreGroup {
                    cores: vec![
                        CoreOffset {
                            core: 1,
                            offset: -1000
                        },
                        CoreOffset {
                            core: 3,
                            offset: -1000
                        },
                        CoreOffset {
                            core: 5,
                            offset: -1000
                        },
                        CoreOffset {
                            core: 7,
                            offset: -1000
                        }
                    ]
                },
                CoreGroup {
                    cores: vec![CoreOffset { core: 0, offset: 0 }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 2,
                        offset: 1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 4,
                        offset: 2000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 6,
                        offset: 3000
                    }]
                },
            ])
            .expect("CoreGrouping::new should not fail here")
        );

        fn fake_rdtsc() -> u64 {
            u64::MAX
        }

        let num_vcpus = 4;
        let mitigations = tsc_sync_mitigations_inner(&state, num_vcpus, fake_rdtsc);

        let expected_affinity = vec![1, 3, 5, 7];
        for i in 0..num_vcpus {
            assert_eq!(
                mitigations.get_vcpu_affinity(i).unwrap_or_else(|| panic!(
                    "core {} should have an affinity of {:?}",
                    i, expected_affinity
                )),
                expected_affinity
            );
            assert_eq!(mitigations.get_vcpu_tsc_offset(i), None);
        }
    }

    #[test]
    fn more_vcpus_than_cores() {
        // 4 cores, two can be grouped but it doesn't matter because we'll have more vcpus than
        // the largest group.
        let offsets = vec![0, 0, 1000, 2000];
        // frequency of 1GHz means 20 nanos is 20 ticks
        let state = TscState::new(1_000_000_000, offsets, Duration::from_nanos(20))
            .expect("TscState::new should not fail for this test");

        assert_eq!(
            state.core_grouping,
            CoreGrouping::new(vec![
                CoreGroup {
                    cores: vec![
                        CoreOffset { core: 0, offset: 0 },
                        CoreOffset { core: 1, offset: 0 }
                    ]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 2,
                        offset: 1000
                    }]
                },
                CoreGroup {
                    cores: vec![CoreOffset {
                        core: 3,
                        offset: 2000
                    }]
                },
            ])
            .expect("CoreGrouping::new should not fail here")
        );

        fn fake_rdtsc() -> u64 {
            u64::MAX
        }

        // 8 vcpus, more than we have cores
        let num_vcpus = 8;
        let mitigations = tsc_sync_mitigations_inner(&state, num_vcpus, fake_rdtsc);
        let expected_offsets = vec![1, 1, 1u64.wrapping_sub(1000), 1u64.wrapping_sub(2000)];

        for i in 0..num_vcpus {
            assert_eq!(
                mitigations.get_vcpu_affinity(i).unwrap_or_else(|| panic!(
                    "core {} should have an affinity of {:?}",
                    i,
                    i % 4
                )),
                // expected affinity is the vcpu modulo 4
                vec![i % 4]
            );
            assert_eq!(
                mitigations.get_vcpu_tsc_offset(i).unwrap_or_else(|| panic!(
                    "core {} should have an offset of {:?}",
                    i,
                    expected_offsets[i % 4]
                )),
                expected_offsets[i % 4]
            );
        }
    }
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Jordi Boggiano <j.boggiano@seld.be>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.
//  *
//  * Synced with http://lingrok.org/xref/coreutils/src/tty.c

// spell-checker:ignore (ToDO) ttyname filedesc

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use std::ffi::CStr;
use uucore::InvalidEncodingHandling;

static ABOUT: &str = "Print the file name of the terminal connected to standard input.";

mod options {
    pub const SILENT: &str = "silent";
}

fn get_usage() -> String {
    format!("{0} [OPTION]...", executable!())
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::SILENT)
                .long(options::SILENT)
                .visible_alias("quiet")
                .short("s")
                .help("print nothing, only return an exit status")
                .required(false),
        )
        .get_matches_from(args);

    let silent = matches.is_present(options::SILENT);

    // Call libc function ttyname
    let tty = unsafe {
        let ptr = libc::ttyname(libc::STDIN_FILENO);
        if !ptr.is_null() {
            String::from_utf8_lossy(CStr::from_ptr(ptr).to_bytes()).to_string()
        } else {
            "".to_owned()
        }
    };

    if !silent {
        if !tty.chars().all(|c| c.is_whitespace()) {
            println!("{}", tty);
        } else {
            println!("not a tty");
        }
    }

    if atty::is(atty::Stream::Stdin) {
        libc::EXIT_SUCCESS
    } else {
        libc::EXIT_FAILURE
    }
}
<commit_msg>tty: correct exit code for wrong args<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Jordi Boggiano <j.boggiano@seld.be>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.
//  *
//  * Synced with http://lingrok.org/xref/coreutils/src/tty.c

// spell-checker:ignore (ToDO) ttyname filedesc

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use std::ffi::CStr;
use uucore::InvalidEncodingHandling;

static ABOUT: &str = "Print the file name of the terminal connected to standard input.";

mod options {
    pub const SILENT: &str = "silent";
}

fn get_usage() -> String {
    format!("{0} [OPTION]...", executable!())
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::SILENT)
                .long(options::SILENT)
                .visible_alias("quiet")
                .short("s")
                .help("print nothing, only return an exit status")
                .required(false),
        )
        .get_matches_from_safe(args);

    let matches = match matches {
        Ok(m) => m,
        Err(e) => {
            eprint!("{}", e);
            return 2;
        }
    };

    let silent = matches.is_present(options::SILENT);

    // Call libc function ttyname
    let tty = unsafe {
        let ptr = libc::ttyname(libc::STDIN_FILENO);
        if !ptr.is_null() {
            String::from_utf8_lossy(CStr::from_ptr(ptr).to_bytes()).to_string()
        } else {
            "".to_owned()
        }
    };

    if !silent {
        if !tty.chars().all(|c| c.is_whitespace()) {
            println!("{}", tty);
        } else {
            println!("not a tty");
        }
    }

    if atty::is(atty::Stream::Stdin) {
        libc::EXIT_SUCCESS
    } else {
        libc::EXIT_FAILURE
    }
}
<|endoftext|>


<commit_before>
use ::{
    Direction,
    Incoming,
};

use visit::{
    GraphBase,
    GraphRef,
    GraphEdgeRef,
    IntoNodeIdentifiers,
    IntoNeighbors,
    IntoNeighborsDirected,
    IntoEdgeReferences,
    IntoExternals,
    Visitable,
    EdgeRef,
};

/// Wrapper type for walking the graph as if all edges are reversed.
#[derive(Copy, Clone)]
pub struct Reversed<G>(pub G);

impl<G: GraphBase> GraphBase for Reversed<G> {
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;
}

impl<G: GraphRef> GraphRef for Reversed<G> { }

impl<G> IntoNodeIdentifiers for Reversed<G>
    where G: IntoNodeIdentifiers
{
    type NodeIdentifiers = G::NodeIdentifiers;
    fn node_identifiers(self) -> Self::NodeIdentifiers {
        self.0.node_identifiers()
    }

    fn node_count(&self) -> usize {
        self.0.node_count()
    }
}

impl<G> IntoNeighbors for Reversed<G>
    where G: IntoNeighborsDirected
{
    type Neighbors = G::NeighborsDirected;
    fn neighbors(self, n: G::NodeId) -> G::NeighborsDirected
    {
        self.0.neighbors_directed(n, Incoming)
    }
}

impl<G> IntoNeighborsDirected for Reversed<G>
    where G: IntoNeighborsDirected
{
    type NeighborsDirected = G::NeighborsDirected;
    fn neighbors_directed(self, n: G::NodeId, d: Direction)
        -> G::NeighborsDirected
    {
        self.0.neighbors_directed(n, d.opposite())
    }
}

impl<G> IntoExternals for Reversed<G>
    where G: IntoExternals,
{
    type Externals = G::Externals;
    fn externals(self, d: Direction) -> G::Externals {
        self.0.externals(d.opposite())
    }
}

impl<G: Visitable> Visitable for Reversed<G>
{
    type Map = G::Map;
    fn visit_map(&self) -> G::Map {
        self.0.visit_map()
    }
    fn reset_map(&self, map: &mut Self::Map) {
        self.0.reset_map(map);
    }
}

/// An edge reference for `Reversed`.
#[derive(Copy, Clone, Debug)]
pub struct ReversedEdgeRef<R>(R);

/// An edge reference
impl<R> EdgeRef for ReversedEdgeRef<R>
    where R: EdgeRef,
{
    type NodeId = R::NodeId;
    type EdgeId = R::EdgeId;
    type Weight = R::Weight;
    fn source(&self) -> Self::NodeId {
        self.0.target()
    }
    fn target(&self) -> Self::NodeId {
        self.0.source()
    }
    fn weight(&self) -> &Self::Weight {
        self.0.weight()
    }
    fn id(&self) -> Self::EdgeId {
        self.0.id()
    }
}

impl<G> GraphEdgeRef for Reversed<G>
    where G: GraphEdgeRef
{
    type EdgeRef = ReversedEdgeRef<G::EdgeRef>;
}

impl<G> IntoEdgeReferences for Reversed<G>
    where G: IntoEdgeReferences
{
    type EdgeReferences = ReversedEdgeReferences<G::EdgeReferences>;
    fn edge_references(self) -> Self::EdgeReferences {
        ReversedEdgeReferences {
            iter: self.0.edge_references(),
        }
    }
}

/// An iterator of edge references for `Reversed`.
pub struct ReversedEdgeReferences<I> {
    iter: I,
}

impl<I> Iterator for ReversedEdgeReferences<I>
    where I: Iterator,
          I::Item: EdgeRef,
{
    type Item = ReversedEdgeRef<I::Item>;
    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(ReversedEdgeRef)
    }
}

<commit_msg>reversed: impl Debug<commit_after>
use ::{
    Direction,
    Incoming,
};

use visit::{
    GraphBase,
    GraphRef,
    GraphEdgeRef,
    IntoNodeIdentifiers,
    IntoNeighbors,
    IntoNeighborsDirected,
    IntoEdgeReferences,
    IntoExternals,
    Visitable,
    EdgeRef,
};

/// Wrapper type for walking the graph as if all edges are reversed.
#[derive(Copy, Clone, Debug)]
pub struct Reversed<G>(pub G);

impl<G: GraphBase> GraphBase for Reversed<G> {
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;
}

impl<G: GraphRef> GraphRef for Reversed<G> { }

impl<G> IntoNodeIdentifiers for Reversed<G>
    where G: IntoNodeIdentifiers
{
    type NodeIdentifiers = G::NodeIdentifiers;
    fn node_identifiers(self) -> Self::NodeIdentifiers {
        self.0.node_identifiers()
    }

    fn node_count(&self) -> usize {
        self.0.node_count()
    }
}

impl<G> IntoNeighbors for Reversed<G>
    where G: IntoNeighborsDirected
{
    type Neighbors = G::NeighborsDirected;
    fn neighbors(self, n: G::NodeId) -> G::NeighborsDirected
    {
        self.0.neighbors_directed(n, Incoming)
    }
}

impl<G> IntoNeighborsDirected for Reversed<G>
    where G: IntoNeighborsDirected
{
    type NeighborsDirected = G::NeighborsDirected;
    fn neighbors_directed(self, n: G::NodeId, d: Direction)
        -> G::NeighborsDirected
    {
        self.0.neighbors_directed(n, d.opposite())
    }
}

impl<G> IntoExternals for Reversed<G>
    where G: IntoExternals,
{
    type Externals = G::Externals;
    fn externals(self, d: Direction) -> G::Externals {
        self.0.externals(d.opposite())
    }
}

impl<G: Visitable> Visitable for Reversed<G>
{
    type Map = G::Map;
    fn visit_map(&self) -> G::Map {
        self.0.visit_map()
    }
    fn reset_map(&self, map: &mut Self::Map) {
        self.0.reset_map(map);
    }
}

/// An edge reference for `Reversed`.
#[derive(Copy, Clone, Debug)]
pub struct ReversedEdgeRef<R>(R);

/// An edge reference
impl<R> EdgeRef for ReversedEdgeRef<R>
    where R: EdgeRef,
{
    type NodeId = R::NodeId;
    type EdgeId = R::EdgeId;
    type Weight = R::Weight;
    fn source(&self) -> Self::NodeId {
        self.0.target()
    }
    fn target(&self) -> Self::NodeId {
        self.0.source()
    }
    fn weight(&self) -> &Self::Weight {
        self.0.weight()
    }
    fn id(&self) -> Self::EdgeId {
        self.0.id()
    }
}

impl<G> GraphEdgeRef for Reversed<G>
    where G: GraphEdgeRef
{
    type EdgeRef = ReversedEdgeRef<G::EdgeRef>;
}

impl<G> IntoEdgeReferences for Reversed<G>
    where G: IntoEdgeReferences
{
    type EdgeReferences = ReversedEdgeReferences<G::EdgeReferences>;
    fn edge_references(self) -> Self::EdgeReferences {
        ReversedEdgeReferences {
            iter: self.0.edge_references(),
        }
    }
}

/// An iterator of edge references for `Reversed`.
pub struct ReversedEdgeReferences<I> {
    iter: I,
}

impl<I> Iterator for ReversedEdgeReferences<I>
    where I: Iterator,
          I::Item: EdgeRef,
{
    type Item = ReversedEdgeRef<I::Item>;
    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(ReversedEdgeRef)
    }
}

<|endoftext|>


<commit_before>use cairo::{self, MatrixTrait};
use cairo_sys;
use glib::translate::*;
use std::cell::Cell;

use attributes::Attribute;
use coord_units::CoordUnits;
use draw;
use drawing_ctx::{self, RsvgDrawingCtx};
use handle::RsvgHandle;
use length::{LengthDir, RsvgLength};
use node::{NodeResult, NodeTrait, RsvgNode};
use parsers::{parse, parse_and_validate, Parse};
use property_bag::PropertyBag;
use state::Opacity;

coord_units!(MaskUnits, CoordUnits::ObjectBoundingBox);
coord_units!(MaskContentUnits, CoordUnits::UserSpaceOnUse);

// remove this once cairo-rs has this mask_surface()
fn cairo_mask_surface(cr: &cairo::Context, surface: &cairo::Surface, x: f64, y: f64) {
    unsafe {
        let raw_cr = cr.to_glib_none().0;

        cairo_sys::cairo_mask_surface(raw_cr, surface.to_raw_none(), x, y);
    }
}

pub struct NodeMask {
    x: Cell<RsvgLength>,
    y: Cell<RsvgLength>,
    width: Cell<RsvgLength>,
    height: Cell<RsvgLength>,

    units: Cell<MaskUnits>,
    content_units: Cell<MaskContentUnits>,
}

impl NodeMask {
    pub fn new() -> NodeMask {
        NodeMask {
            x: Cell::new(NodeMask::get_default_pos(LengthDir::Horizontal)),
            y: Cell::new(NodeMask::get_default_pos(LengthDir::Vertical)),

            width: Cell::new(NodeMask::get_default_size(LengthDir::Horizontal)),
            height: Cell::new(NodeMask::get_default_size(LengthDir::Vertical)),

            units: Cell::new(MaskUnits::default()),
            content_units: Cell::new(MaskContentUnits::default()),
        }
    }

    fn get_default_pos(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse_str("-10%", dir).unwrap()
    }

    fn get_default_size(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse_str("120%", dir).unwrap()
    }

    pub fn generate_cairo_mask(
        &self,
        node: &RsvgNode,
        affine_before_mask: &cairo::Matrix,
        draw_ctx: *mut RsvgDrawingCtx,
    ) {
        let cascaded = node.get_cascaded_values();
        let values = cascaded.get();

        let width = drawing_ctx::get_width(draw_ctx) as i32;
        let height = drawing_ctx::get_height(draw_ctx) as i32;

        let mut surface = match cairo::ImageSurface::create(cairo::Format::ARgb32, width, height) {
            Ok(surface) => surface,
            Err(_) => return,
        };

        let mask_units = CoordUnits::from(self.units.get());
        let content_units = CoordUnits::from(self.content_units.get());

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
        }

        let x = self.x.get().normalize(&values, draw_ctx);
        let y = self.y.get().normalize(&values, draw_ctx);
        let w = self.width.get().normalize(&values, draw_ctx);
        let h = self.height.get().normalize(&values, draw_ctx);

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::pop_view_box(draw_ctx);
        }

        // Use a scope because mask_cr needs to release the
        // reference to the surface before we access the pixels
        {
            let save_cr = drawing_ctx::get_cairo_context(draw_ctx);

            let mask_cr = cairo::Context::new(&surface);
            mask_cr.set_matrix(*affine_before_mask);
            mask_cr.transform(node.get_transform());

            drawing_ctx::set_cairo_context(draw_ctx, &mask_cr);

            if mask_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();

                draw::add_clipping_rect(
                    draw_ctx,
                    x * rect.width + rect.x,
                    y * rect.height + rect.y,
                    w * rect.width,
                    h * rect.height,
                );
            } else {
                draw::add_clipping_rect(draw_ctx, x, y, w, h);
            }

            if content_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();
                let bbtransform =
                    cairo::Matrix::new(rect.width, 0.0, 0.0, rect.height, rect.x, rect.y);

                mask_cr.transform(bbtransform);

                drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
            }

            drawing_ctx::push_discrete_layer(draw_ctx, values, false);
            node.draw_children(&cascaded, draw_ctx, false);
            drawing_ctx::pop_discrete_layer(draw_ctx, node, values, false);

            if content_units == CoordUnits::ObjectBoundingBox {
                drawing_ctx::pop_view_box(draw_ctx);
            }

            drawing_ctx::set_cairo_context(draw_ctx, &save_cr);
        }

        {
            let rowstride = surface.get_stride() as usize;
            let mut pixels = surface.get_data().unwrap();
            let opacity = {
                let Opacity(o) = values.opacity;
                u8::from(o)
            };

            for row in pixels.chunks_mut(rowstride) {
                for p in row[..4 * width as usize].chunks_mut(4) {
                    //  Assuming, the pixel is linear RGB (not sRGB)
                    //  y = luminance
                    //  Y = 0.2126 R + 0.7152 G + 0.0722 B
                    //  1.0 opacity = 255
                    //
                    //  When Y = 1.0, pixel for mask should be 0xFFFFFFFF
                    //    (you get 1.0 luminance from 255 from R, G and B)
                    //
                    // r_mult = 0xFFFFFFFF / (255.0 * 255.0) * .2126 = 14042.45  ~= 14042
                    // g_mult = 0xFFFFFFFF / (255.0 * 255.0) * .7152 = 47239.69  ~= 47240
                    // b_mult = 0xFFFFFFFF / (255.0 * 255.0) * .0722 =  4768.88  ~= 4769
                    //
                    // This allows for the following expected behaviour:
                    //    (we only care about the most sig byte)
                    // if pixel = 0x00FFFFFF, pixel' = 0xFF......
                    // if pixel = 0x00020202, pixel' = 0x02......
                    // if pixel = 0x00000000, pixel' = 0x00......
                    //
                    // NOTE: the following assumes little-endian
                    let (r, g, b, o) = (p[2] as u32, p[1] as u32, p[0] as u32, opacity as u32);
                    p[3] = (((r * 14042 + g * 47240 + b * 4769) * o) >> 24) as u8;
                }
            }
        }

        let cr = drawing_ctx::get_cairo_context(draw_ctx);

        cr.identity_matrix();

        let (xofs, yofs) = drawing_ctx::get_offset(draw_ctx);
        cairo_mask_surface(&cr, &surface, xofs, yofs);
    }
}

impl NodeTrait for NodeMask {
    fn set_atts(&self, _: &RsvgNode, _: *const RsvgHandle, pbag: &PropertyBag) -> NodeResult {
        for (_key, attr, value) in pbag.iter() {
            match attr {
                Attribute::X => self.x.set(parse("x", value, LengthDir::Horizontal)?),
                Attribute::Y => self.y.set(parse("y", value, LengthDir::Vertical)?),
                Attribute::Width => self.width.set(parse_and_validate(
                    "width",
                    value,
                    LengthDir::Horizontal,
                    RsvgLength::check_nonnegative,
                )?),
                Attribute::Height => self.height.set(parse_and_validate(
                    "height",
                    value,
                    LengthDir::Vertical,
                    RsvgLength::check_nonnegative,
                )?),

                Attribute::MaskUnits => self.units.set(parse("maskUnits", value, ())?),

                Attribute::MaskContentUnits => {
                    self.content_units
                        .set(parse("maskContentUnits", value, ())?)
                }

                _ => (),
            }
        }

        Ok(())
    }
}
<commit_msg>mask.rs: Use with_discrete_layer()<commit_after>use cairo::{self, MatrixTrait};
use cairo_sys;
use glib::translate::*;
use std::cell::Cell;

use attributes::Attribute;
use coord_units::CoordUnits;
use draw;
use drawing_ctx::{self, RsvgDrawingCtx};
use handle::RsvgHandle;
use length::{LengthDir, RsvgLength};
use node::{NodeResult, NodeTrait, RsvgNode};
use parsers::{parse, parse_and_validate, Parse};
use property_bag::PropertyBag;
use state::Opacity;

coord_units!(MaskUnits, CoordUnits::ObjectBoundingBox);
coord_units!(MaskContentUnits, CoordUnits::UserSpaceOnUse);

// remove this once cairo-rs has this mask_surface()
fn cairo_mask_surface(cr: &cairo::Context, surface: &cairo::Surface, x: f64, y: f64) {
    unsafe {
        let raw_cr = cr.to_glib_none().0;

        cairo_sys::cairo_mask_surface(raw_cr, surface.to_raw_none(), x, y);
    }
}

pub struct NodeMask {
    x: Cell<RsvgLength>,
    y: Cell<RsvgLength>,
    width: Cell<RsvgLength>,
    height: Cell<RsvgLength>,

    units: Cell<MaskUnits>,
    content_units: Cell<MaskContentUnits>,
}

impl NodeMask {
    pub fn new() -> NodeMask {
        NodeMask {
            x: Cell::new(NodeMask::get_default_pos(LengthDir::Horizontal)),
            y: Cell::new(NodeMask::get_default_pos(LengthDir::Vertical)),

            width: Cell::new(NodeMask::get_default_size(LengthDir::Horizontal)),
            height: Cell::new(NodeMask::get_default_size(LengthDir::Vertical)),

            units: Cell::new(MaskUnits::default()),
            content_units: Cell::new(MaskContentUnits::default()),
        }
    }

    fn get_default_pos(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse_str("-10%", dir).unwrap()
    }

    fn get_default_size(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse_str("120%", dir).unwrap()
    }

    pub fn generate_cairo_mask(
        &self,
        node: &RsvgNode,
        affine_before_mask: &cairo::Matrix,
        draw_ctx: *mut RsvgDrawingCtx,
    ) {
        let cascaded = node.get_cascaded_values();
        let values = cascaded.get();

        let width = drawing_ctx::get_width(draw_ctx) as i32;
        let height = drawing_ctx::get_height(draw_ctx) as i32;

        let mut surface = match cairo::ImageSurface::create(cairo::Format::ARgb32, width, height) {
            Ok(surface) => surface,
            Err(_) => return,
        };

        let mask_units = CoordUnits::from(self.units.get());
        let content_units = CoordUnits::from(self.content_units.get());

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
        }

        let x = self.x.get().normalize(&values, draw_ctx);
        let y = self.y.get().normalize(&values, draw_ctx);
        let w = self.width.get().normalize(&values, draw_ctx);
        let h = self.height.get().normalize(&values, draw_ctx);

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::pop_view_box(draw_ctx);
        }

        // Use a scope because mask_cr needs to release the
        // reference to the surface before we access the pixels
        {
            let save_cr = drawing_ctx::get_cairo_context(draw_ctx);

            let mask_cr = cairo::Context::new(&surface);
            mask_cr.set_matrix(*affine_before_mask);
            mask_cr.transform(node.get_transform());

            drawing_ctx::set_cairo_context(draw_ctx, &mask_cr);

            if mask_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();

                draw::add_clipping_rect(
                    draw_ctx,
                    x * rect.width + rect.x,
                    y * rect.height + rect.y,
                    w * rect.width,
                    h * rect.height,
                );
            } else {
                draw::add_clipping_rect(draw_ctx, x, y, w, h);
            }

            if content_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();
                let bbtransform =
                    cairo::Matrix::new(rect.width, 0.0, 0.0, rect.height, rect.x, rect.y);

                mask_cr.transform(bbtransform);

                drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
            }

            drawing_ctx::with_discrete_layer(draw_ctx, node, values, false, |_cr| {
                node.draw_children(&cascaded, draw_ctx, false);
            });

            if content_units == CoordUnits::ObjectBoundingBox {
                drawing_ctx::pop_view_box(draw_ctx);
            }

            drawing_ctx::set_cairo_context(draw_ctx, &save_cr);
        }

        {
            let rowstride = surface.get_stride() as usize;
            let mut pixels = surface.get_data().unwrap();
            let opacity = {
                let Opacity(o) = values.opacity;
                u8::from(o)
            };

            for row in pixels.chunks_mut(rowstride) {
                for p in row[..4 * width as usize].chunks_mut(4) {
                    //  Assuming, the pixel is linear RGB (not sRGB)
                    //  y = luminance
                    //  Y = 0.2126 R + 0.7152 G + 0.0722 B
                    //  1.0 opacity = 255
                    //
                    //  When Y = 1.0, pixel for mask should be 0xFFFFFFFF
                    //    (you get 1.0 luminance from 255 from R, G and B)
                    //
                    // r_mult = 0xFFFFFFFF / (255.0 * 255.0) * .2126 = 14042.45  ~= 14042
                    // g_mult = 0xFFFFFFFF / (255.0 * 255.0) * .7152 = 47239.69  ~= 47240
                    // b_mult = 0xFFFFFFFF / (255.0 * 255.0) * .0722 =  4768.88  ~= 4769
                    //
                    // This allows for the following expected behaviour:
                    //    (we only care about the most sig byte)
                    // if pixel = 0x00FFFFFF, pixel' = 0xFF......
                    // if pixel = 0x00020202, pixel' = 0x02......
                    // if pixel = 0x00000000, pixel' = 0x00......
                    //
                    // NOTE: the following assumes little-endian
                    let (r, g, b, o) = (p[2] as u32, p[1] as u32, p[0] as u32, opacity as u32);
                    p[3] = (((r * 14042 + g * 47240 + b * 4769) * o) >> 24) as u8;
                }
            }
        }

        let cr = drawing_ctx::get_cairo_context(draw_ctx);

        cr.identity_matrix();

        let (xofs, yofs) = drawing_ctx::get_offset(draw_ctx);
        cairo_mask_surface(&cr, &surface, xofs, yofs);
    }
}

impl NodeTrait for NodeMask {
    fn set_atts(&self, _: &RsvgNode, _: *const RsvgHandle, pbag: &PropertyBag) -> NodeResult {
        for (_key, attr, value) in pbag.iter() {
            match attr {
                Attribute::X => self.x.set(parse("x", value, LengthDir::Horizontal)?),
                Attribute::Y => self.y.set(parse("y", value, LengthDir::Vertical)?),
                Attribute::Width => self.width.set(parse_and_validate(
                    "width",
                    value,
                    LengthDir::Horizontal,
                    RsvgLength::check_nonnegative,
                )?),
                Attribute::Height => self.height.set(parse_and_validate(
                    "height",
                    value,
                    LengthDir::Vertical,
                    RsvgLength::check_nonnegative,
                )?),

                Attribute::MaskUnits => self.units.set(parse("maskUnits", value, ())?),

                Attribute::MaskContentUnits => {
                    self.content_units
                        .set(parse("maskContentUnits", value, ())?)
                }

                _ => (),
            }
        }

        Ok(())
    }
}
<|endoftext|>