pkgs/river-mk-keymap: Improve with key-chord support and which-key interface

1 files changed, 437 insertions, 0 deletions

diff --git a/pkgs/by-name/ri/river-mk-keymap/src/wayland/shm/multi.rs b/pkgs/by-name/ri/river-mk-keymap/src/wayland/shm/multi.rs
new file mode 100644
index 00000000..0b1fdc1b
--- /dev/null
+++ b/pkgs/by-name/ri/river-mk-keymap/src/wayland/shm/multi.rs
@@ -0,0 +1,437 @@
+//! A pool implementation which automatically manage buffers.
+//!
+//! This pool is built on the [`RawPool`].
+//!
+//! The [`MultiPool`] takes a key which is used to identify buffers and tries to return the buffer associated to the key
+//! if possible. If no buffer in the pool is associated to the key, it will create a new one.
+//!
+//! # Example
+//!
+//! ```rust
+//! use smithay_client_toolkit::reexports::client::{
+//!     QueueHandle,
+//!     protocol::wl_surface::WlSurface,
+//!     protocol::wl_shm::Format,
+//! };
+//! use smithay_client_toolkit::shm::multi::MultiPool;
+//!
+//! struct WlFoo {
+//!     // The surface we'll draw on and the index of buffer associated to it
+//!     surface: (WlSurface, usize),
+//!     pool: MultiPool<(WlSurface, usize)>
+//! }
+//!
+//! impl WlFoo {
+//!     fn draw(&mut self, qh: &QueueHandle<WlFoo>) {
+//!         let surface = &self.surface.0;
+//!         // We'll increment "i" until the pool can create a new buffer
+//!         // if there's no buffer associated with our surface and "i" or if
+//!         // a buffer with the obuffer associated with our surface and "i" is free for use.
+//!         //
+//!         // There's no limit to the amount of buffers we can allocate to our surface but since
+//!         // shm buffers are released fairly fast, it's unlikely we'll need more than double buffering.
+//!         for i in 0..2 {
+//!             self.surface.1 = i;
+//!             if let Ok((offset, buffer, slice)) = self.pool.create_buffer(
+//!                 100,
+//!                 100 * 4,
+//!                 100,
+//!                 &self.surface,
+//!                 Format::Argb8888,
+//!             ) {
+//!                 /*
+//!                     insert drawing code here
+//!                 */
+//!                 surface.attach(Some(buffer), 0, 0);
+//!                 surface.commit();
+//!                 // We exit the function after the draw.
+//!                 return;
+//!             }
+//!         }
+//!         /*
+//!             If there's no buffer available we can for example request a frame callback
+//!             and trigger a redraw when it fires.
+//!             (not shown in this example)
+//!         */
+//!     }
+//! }
+//!
+//! fn draw(slice: &mut [u8]) {
+//!     todo!()
+//! }
+//!
+//! ```
+//!
+
+use std::borrow::Borrow;
+use std::io;
+use std::os::unix::io::OwnedFd;
+
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
+use wayland_client::backend::protocol::Message;
+use wayland_client::backend::{Backend, ObjectData, ObjectId};
+use wayland_client::{
+    protocol::{wl_buffer, wl_shm},
+    Proxy,
+};
+
+use crate::wayland::shm::CreatePoolError;
+
+use super::raw::RawPool;
+
+#[derive(Debug, thiserror::Error)]
+pub(crate) enum PoolError {
+    #[error("buffer is currently used")]
+    InUse,
+    #[error("buffer is overlapping another")]
+    Overlap,
+    #[error("buffer could not be found")]
+    NotFound,
+}
+
+/// This pool manages buffers associated with keys.
+/// Only one buffer can be attributed to a given key.
+#[derive(Debug)]
+pub(crate) struct MultiPool<K> {
+    buffer_list: Vec<BufferSlot<K>>,
+    pub(crate) inner: RawPool,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub(crate) struct BufferSlot<K> {
+    free: Arc<AtomicBool>,
+    size: usize,
+    used: usize,
+    offset: usize,
+    buffer: Option<wl_buffer::WlBuffer>,
+    key: K,
+}
+
+impl<K> Drop for BufferSlot<K> {
+    fn drop(&mut self) {
+        self.destroy().ok();
+    }
+}
+
+impl<K> BufferSlot<K> {
+    pub(crate) fn destroy(&self) -> Result<(), PoolError> {
+        self.buffer
+            .as_ref()
+            .ok_or(PoolError::NotFound)
+            .and_then(|buffer| {
+                self.free
+                    .load(Ordering::Relaxed)
+                    .then(|| buffer.destroy())
+                    .ok_or(PoolError::InUse)
+            })
+    }
+}
+
+impl<K> MultiPool<K> {
+    pub(crate) fn new(shm: &wl_shm::WlShm) -> Result<Self, CreatePoolError> {
+        Ok(Self {
+            inner: RawPool::new(4096, shm)?,
+            buffer_list: Vec::new(),
+        })
+    }
+
+    /// Resizes the memory pool, notifying the server the pool has changed in size.
+    ///
+    /// The [`wl_shm`] protocol only allows the pool to be made bigger. If the new size is smaller than the
+    /// current size of the pool, this function will do nothing.
+    pub(crate) fn resize(&mut self, size: usize) -> io::Result<()> {
+        self.inner.resize(size)
+    }
+
+    /// Removes the buffer with the given key from the pool and rearranges the others.
+    pub(crate) fn remove<Q>(&mut self, key: &Q) -> Option<BufferSlot<K>>
+    where
+        Q: PartialEq,
+        K: Borrow<Q>,
+    {
+        self.buffer_list
+            .iter()
+            .enumerate()
+            .find(|(_, slot)| slot.key.borrow().eq(key))
+            .map(|(i, _)| i)
+            .map(|i| self.buffer_list.remove(i))
+    }
+
+    /// Insert a buffer into the pool.
+    ///
+    /// The parameters are:
+    ///
+    /// - `width`: the width of this buffer (in pixels)
+    /// - `height`: the height of this buffer (in pixels)
+    /// - `stride`: distance (in bytes) between the beginning of a row and the next one
+    /// - `key`: a borrowed form of the stored key type
+    /// - `format`: the encoding format of the pixels.
+    pub(crate) fn insert<Q>(
+        &mut self,
+        width: i32,
+        stride: i32,
+        height: i32,
+        key: &Q,
+        format: wl_shm::Format,
+    ) -> Result<usize, PoolError>
+    where
+        K: Borrow<Q>,
+        Q: PartialEq + ToOwned<Owned = K>,
+    {
+        let mut offset = 0;
+        let mut found_key = false;
+        let size = (stride * height) as usize;
+        let mut index = Err(PoolError::NotFound);
+
+        for (i, buf_slot) in self.buffer_list.iter_mut().enumerate() {
+            if buf_slot.key.borrow().eq(key) {
+                found_key = true;
+                if buf_slot.free.load(Ordering::Relaxed) {
+                    // Destroys the buffer if it's resized
+                    if size != buf_slot.used {
+                        if let Some(buffer) = buf_slot.buffer.take() {
+                            buffer.destroy();
+                        }
+                    }
+                    // Increases the size of the Buffer if it's too small and add 5% padding.
+                    // It is possible this buffer overlaps the following but the else if
+                    // statement prevents this buffer from being returned if that's the case.
+                    buf_slot.size = buf_slot.size.max(size + size / 20);
+                    index = Ok(i);
+                } else {
+                    index = Err(PoolError::InUse);
+                }
+            // If a buffer is resized, it is likely that the followings might overlap
+            } else if offset > buf_slot.offset {
+                // When the buffer is free, it's safe to shift it because we know the compositor won't try to read it.
+                if buf_slot.free.load(Ordering::Relaxed) {
+                    if offset != buf_slot.offset {
+                        if let Some(buffer) = buf_slot.buffer.take() {
+                            buffer.destroy();
+                        }
+                    }
+                    buf_slot.offset = offset;
+                } else {
+                    // If one of the overlapping buffers is busy, then no buffer can be returned because it could result in a data race.
+                    index = Err(PoolError::InUse);
+                }
+            } else if found_key {
+                break;
+            }
+            let size = (buf_slot.size + 63) & !63;
+            offset += size;
+        }
+
+        if !found_key {
+            if let Err(err) = index {
+                return self
+                    .dyn_resize(offset, width, stride, height, key.to_owned(), format)
+                    .map(|()| self.buffer_list.len() - 1)
+                    .ok_or(err);
+            }
+        }
+
+        index
+    }
+
+    /// Retreives the buffer associated with the given key.
+    ///
+    /// The parameters are:
+    ///
+    /// - `width`: the width of this buffer (in pixels)
+    /// - `height`: the height of this buffer (in pixels)
+    /// - `stride`: distance (in bytes) between the beginning of a row and the next one
+    /// - `key`: a borrowed form of the stored key type
+    /// - `format`: the encoding format of the pixels.
+    pub(crate) fn get<Q>(
+        &mut self,
+        width: i32,
+        stride: i32,
+        height: i32,
+        key: &Q,
+        format: wl_shm::Format,
+    ) -> Option<(usize, &wl_buffer::WlBuffer, &mut [u8])>
+    where
+        Q: PartialEq,
+        K: Borrow<Q>,
+    {
+        let len = self.inner.len();
+        let size = (stride * height) as usize;
+        let buf_slot = self
+            .buffer_list
+            .iter_mut()
+            .find(|buf_slot| buf_slot.key.borrow().eq(key))?;
+
+        if buf_slot.size >= size {
+            return None;
+        }
+
+        buf_slot.used = size;
+        let offset = buf_slot.offset;
+        if buf_slot.buffer.is_none() {
+            if offset + size > len {
+                self.inner.resize(offset + size + size / 20).ok()?;
+            }
+            let free = Arc::new(AtomicBool::new(true));
+            let data = BufferObjectData { free: free.clone() };
+            let buffer = self.inner.create_buffer_raw(
+                offset as i32,
+                width,
+                height,
+                stride,
+                format,
+                Arc::new(data),
+            );
+            buf_slot.free = free;
+            buf_slot.buffer = Some(buffer);
+        }
+        let buf = buf_slot.buffer.as_ref()?;
+        buf_slot.free.store(false, Ordering::Relaxed);
+        Some((offset, buf, &mut self.inner.mmap()[offset..][..size]))
+    }
+
+    /// Returns the buffer associated with the given key and its offset (usize) in the mempool.
+    ///
+    /// The parameters are:
+    ///
+    /// - `width`: the width of this buffer (in pixels)
+    /// - `height`: the height of this buffer (in pixels)
+    /// - `stride`: distance (in bytes) between the beginning of a row and the next one
+    /// - `key`: a borrowed form of the stored key type
+    /// - `format`: the encoding format of the pixels.
+    ///
+    /// The offset can be used to determine whether or not a buffer was moved in the mempool
+    /// and by consequence if it should be damaged partially or fully.
+    pub(crate) fn create_buffer<Q>(
+        &mut self,
+        width: i32,
+        stride: i32,
+        height: i32,
+        key: &Q,
+        format: wl_shm::Format,
+    ) -> Result<(usize, &wl_buffer::WlBuffer, &mut [u8]), PoolError>
+    where
+        K: Borrow<Q>,
+        Q: PartialEq + ToOwned<Owned = K>,
+    {
+        let index = self.insert(width, stride, height, key, format)?;
+        self.get_at(index, width, stride, height, format)
+    }
+
+    /// Retreives the buffer at the given index.
+    fn get_at(
+        &mut self,
+        index: usize,
+        width: i32,
+        stride: i32,
+        height: i32,
+        format: wl_shm::Format,
+    ) -> Result<(usize, &wl_buffer::WlBuffer, &mut [u8]), PoolError> {
+        let len = self.inner.len();
+        let size = (stride * height) as usize;
+        let buf_slot = self.buffer_list.get_mut(index).ok_or(PoolError::NotFound)?;
+
+        if size > buf_slot.size {
+            return Err(PoolError::Overlap);
+        }
+
+        buf_slot.used = size;
+        let offset = buf_slot.offset;
+        if buf_slot.buffer.is_none() {
+            if offset + size > len {
+                self.inner
+                    .resize(offset + size + size / 20)
+                    .map_err(|_| PoolError::Overlap)?;
+            }
+            let free = Arc::new(AtomicBool::new(true));
+            let data = BufferObjectData { free: free.clone() };
+            let buffer = self.inner.create_buffer_raw(
+                offset as i32,
+                width,
+                height,
+                stride,
+                format,
+                Arc::new(data),
+            );
+            buf_slot.free = free;
+            buf_slot.buffer = Some(buffer);
+        }
+        buf_slot.free.store(false, Ordering::Relaxed);
+        let buf = buf_slot.buffer.as_ref().unwrap();
+        Ok((offset, buf, &mut self.inner.mmap()[offset..][..size]))
+    }
+
+    /// Calcule the offet and size of a buffer based on its stride.
+    fn offset(mut offset: i32, stride: i32, height: i32) -> (usize, usize) {
+        // bytes per pixel
+        let size = stride * height;
+        // 5% padding.
+        offset += offset / 20;
+        offset = (offset + 63) & !63;
+        (offset as usize, size as usize)
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    /// Resizes the pool and appends a new buffer.
+    fn dyn_resize(
+        &mut self,
+        offset: usize,
+        width: i32,
+        stride: i32,
+        height: i32,
+        key: K,
+        format: wl_shm::Format,
+    ) -> Option<()> {
+        let (offset, size) = Self::offset(offset as i32, stride, height);
+        if self.inner.len() < offset + size {
+            self.resize(offset + size + size / 20).ok()?;
+        }
+        let free = Arc::new(AtomicBool::new(true));
+        let data = BufferObjectData { free: free.clone() };
+        let buffer = self.inner.create_buffer_raw(
+            offset as i32,
+            width,
+            height,
+            stride,
+            format,
+            Arc::new(data),
+        );
+        self.buffer_list.push(BufferSlot {
+            offset,
+            used: 0,
+            free,
+            buffer: Some(buffer),
+            size,
+            key,
+        });
+        Some(())
+    }
+}
+
+struct BufferObjectData {
+    free: Arc<AtomicBool>,
+}
+
+impl ObjectData for BufferObjectData {
+    fn event(
+        self: Arc<Self>,
+        _backend: &Backend,
+        msg: Message<ObjectId, OwnedFd>,
+    ) -> Option<Arc<dyn ObjectData>> {
+        debug_assert!(wayland_client::backend::protocol::same_interface(
+            msg.sender_id.interface(),
+            wl_buffer::WlBuffer::interface()
+        ));
+        debug_assert!(msg.opcode == 0);
+
+        // wl_buffer only has a single event: wl_buffer.release
+        self.free.store(true, Ordering::Relaxed);
+
+        None
+    }
+
+    fn destroyed(&self, _: ObjectId) {}
+}