use std::error::Error;
use std::os::fd::AsFd;
use std::time::{Duration, Instant};
use wayland_client::protocol::{
wl_buffer::{self, WlBuffer},
wl_callback, wl_output, wl_shm,
wl_subsurface::WlSubsurface,
wl_surface::WlSurface,
};
use wayland_client::{Connection, Dispatch, QueueHandle, WEnum};
use wayland_protocols::ext::image_capture_source::v1::client::ext_image_capture_source_v1::ExtImageCaptureSourceV1;
use wayland_protocols::ext::image_copy_capture::v1::client::{
ext_image_copy_capture_frame_v1::{self, ExtImageCopyCaptureFrameV1, FailureReason},
ext_image_copy_capture_manager_v1,
ext_image_copy_capture_session_v1::{self, ExtImageCopyCaptureSessionV1},
};
use wayland_protocols::wp::viewporter::client::wp_viewport::WpViewport;
use crate::app::App;
use crate::shm;
use crate::target::{Kind, Target};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Live {
All,
Current,
None,
}
impl Live {
pub fn parse(s: &str) -> Result<Self, String> {
match s.trim() {
"all" => Ok(Live::All),
"current" => Ok(Live::Current),
"none" => Ok(Live::None),
other => Err(format!("{other:?} is not all, current or none")),
}
}
}
pub struct Slot {
pub(crate) buffer: WlBuffer,
pub(crate) busy: bool,
}
fn choose_format(
offered: &[wl_shm::Format],
supported: &[wl_shm::Format],
) -> Option<(wl_shm::Format, i32)> {
let sized = |format: &wl_shm::Format| Some((*format, bytes_per_pixel(*format)?));
offered
.iter()
.filter(|f| matches!(f, wl_shm::Format::Xrgb8888 | wl_shm::Format::Argb8888))
.find_map(sized)
.or_else(|| {
offered
.iter()
.filter(|f| supported.contains(f))
.find_map(sized)
})
}
fn bytes_per_pixel(format: wl_shm::Format) -> Option<i32> {
use wl_shm::Format;
Some(match format {
Format::Rgb565 | Format::Bgr565 => 2,
Format::Xrgb1555 | Format::Argb1555 => 2,
Format::Rgbx4444 | Format::Rgba4444 | Format::Bgrx4444 | Format::Bgra4444 => 2,
Format::Rgbx5551 | Format::Rgba5551 | Format::Bgrx5551 | Format::Bgra5551 => 2,
Format::Rgb888 | Format::Bgr888 => 3,
Format::Xrgb8888 | Format::Argb8888 | Format::Xbgr8888 | Format::Abgr8888 => 4,
Format::Rgbx8888 | Format::Rgba8888 | Format::Bgrx8888 | Format::Bgra8888 => 4,
Format::Xrgb2101010 | Format::Argb2101010 => 4,
Format::Xbgr2101010 | Format::Abgr2101010 => 4,
Format::Xbgr16161616 | Format::Abgr16161616 => 8,
Format::Xbgr16161616f | Format::Abgr16161616f => 8,
_ => return None,
})
}
#[derive(Default)]
struct PoolPlan {
total: usize,
slots: Vec<(usize, Vec<usize>)>,
}
pub struct Tile {
pub(crate) target: Target,
pub(crate) session: Option<ExtImageCopyCaptureSessionV1>,
pub(crate) frame: Option<ExtImageCopyCaptureFrameV1>,
pub(crate) filling: Option<usize>,
pub(crate) slots: Vec<Slot>,
pub(crate) showing: Option<usize>,
pub(crate) attached: Option<usize>,
pub(crate) formats: Vec<wl_shm::Format>,
pub(crate) format: Option<wl_shm::Format>,
pub(crate) stride: i32,
pub(crate) size: (u32, u32),
pub(crate) stale: bool,
pub(crate) transform: wl_output::Transform,
pub(crate) session_done: bool,
pub(crate) ready: bool,
pub(crate) settled: bool,
pub(crate) asked: Option<Instant>,
pub(crate) frames: u32,
pub(crate) surface: Option<WlSurface>,
pub(crate) subsurface: Option<WlSubsurface>,
pub(crate) viewport: Option<WpViewport>,
}
impl Tile {
pub fn new(target: Target) -> Self {
Self {
target,
session: None,
frame: None,
filling: None,
slots: Vec::new(),
showing: None,
attached: None,
formats: Vec::new(),
format: None,
stride: 0,
size: (0, 0),
stale: false,
transform: wl_output::Transform::Normal,
session_done: false,
ready: false,
settled: false,
asked: None,
frames: 0,
surface: None,
subsurface: None,
viewport: None,
}
}
pub fn bytes(&self) -> usize {
self.stride as usize * self.size.1 as usize
}
pub fn rotated(&self) -> bool {
use wl_output::Transform;
matches!(
self.transform,
Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270
)
}
}
impl App {
pub fn open_sessions(&mut self, qh: &QueueHandle<Self>) {
for (i, tile) in self.tiles.iter_mut().enumerate() {
let source: Option<ExtImageCaptureSourceV1> = match tile.target.kind {
Kind::Window => self
.toplevels
.iter()
.find(|(_, id)| !id.is_empty() && *id == tile.target.ft_id)
.map(|(handle, _)| self.src_mgr.create_source(handle, qh, ())),
Kind::Output => self
.outputs
.iter()
.find(|(_, n)| *n == tile.target.id)
.and_then(|(output, _)| {
self.output_src_mgr
.as_ref()
.map(|mgr| mgr.create_source(output, qh, ()))
}),
};
let Some(source) = source else {
tile.settled = true;
continue;
};
tile.session = Some(self.copy_mgr.create_session(
&source,
ext_image_copy_capture_manager_v1::Options::empty(),
qh,
i,
));
source.destroy();
}
}
pub fn start_captures(&mut self, qh: &QueueHandle<Self>) -> Result<(), Box<dyn Error>> {
let shm_formats = self.shm_formats.clone();
for tile in &mut self.tiles {
if tile.session.is_none() {
continue;
}
if !tile.session_done || tile.size.0 == 0 || tile.size.1 == 0 {
tile.settled = true;
continue;
}
let Some((format, bytes)) = choose_format(&tile.formats, &shm_formats) else {
eprintln!(
"wl-pick: no usable buffer format for {:?} (offered {:?})",
tile.target.title, tile.formats
);
tile.settled = true;
continue;
};
tile.format = Some(format);
tile.stride = tile.size.0 as i32 * bytes;
}
let live = self.live != Live::None;
let mut plan = self.plan_pool(|t| {
if live && t.target.kind != Kind::Output {
2
} else {
1
}
});
if plan.total > i32::MAX as usize {
eprintln!(
"wl-pick: {} MB of capture buffers is more than wl_shm allows; \
using one per window",
plan.total >> 20
);
plan = self.plan_pool(|_| 1);
}
if plan.total > i32::MAX as usize {
return Err(format!(
"{} MB of capture buffers is more than wl_shm allows; try --no-outputs",
plan.total >> 20
)
.into());
}
if plan.total == 0 {
return Ok(());
}
self.stats.pool_bytes = plan.total;
let file = shm::memfd("wl-pick-capture", plan.total)?;
let pool = self
.shm
.create_pool(file.as_fd(), plan.total as i32, qh, ());
for (i, offsets) in plan.slots {
let t = &self.tiles[i];
let (w, h, stride) = (t.size.0 as i32, t.size.1 as i32, t.stride);
let format = t.format.expect("planned tiles have a format");
for offset in offsets {
let slot = self.tiles[i].slots.len();
let buffer = pool.create_buffer(offset as i32, w, h, stride, format, qh, (i, slot));
self.tiles[i].slots.push(Slot {
buffer,
busy: false,
});
}
self.request_capture(i, qh);
}
pool.destroy(); Ok(())
}
fn plan_pool(&self, slots_for: impl Fn(&Tile) -> usize) -> PoolPlan {
const PAGE: usize = 4096;
let mut plan = PoolPlan::default();
for (i, tile) in self.tiles.iter().enumerate() {
if tile.format.is_none() {
continue;
}
let offsets = (0..slots_for(tile))
.map(|_| {
let offset = plan.total;
plan.total += tile.bytes().div_ceil(PAGE) * PAGE;
offset
})
.collect();
plan.slots.push((i, offsets));
}
plan
}
fn request_capture(&mut self, i: usize, qh: &QueueHandle<Self>) {
let t = &mut self.tiles[i];
if t.frame.is_some() || t.session.is_none() {
return; }
let Some(slot) = t.slots.iter().position(|s| !s.busy) else {
self.stats.starved += 1;
return; };
let (w, h) = (t.size.0 as i32, t.size.1 as i32);
let frame = t
.session
.as_ref()
.expect("checked above")
.create_frame(qh, i);
frame.attach_buffer(&t.slots[slot].buffer);
frame.damage_buffer(0, 0, w, h);
frame.capture();
t.frame = Some(frame);
t.filling = Some(slot);
t.asked = Some(Instant::now());
}
fn frame_ready(&mut self, i: usize) -> bool {
let visible = self.layout.tile(i as i32, self.scroll).is_some();
let t = &mut self.tiles[i];
let Some(slot) = t.filling.take() else {
return false;
};
t.frames += 1;
t.ready = true;
t.settled = true;
t.slots[slot].busy = true; let previous = t.showing.replace(slot);
match t.surface.clone() {
Some(surface) if visible => {
let (w, h) = (t.size.0 as i32, t.size.1 as i32);
surface.attach(Some(&t.slots[slot].buffer), 0, 0);
surface.damage_buffer(0, 0, w, h);
surface.commit();
t.attached = Some(slot);
}
_ => {
if let Some(prev) = previous {
t.slots[prev].busy = false;
}
}
}
t.surface.is_none()
}
pub fn arm_frame_callback(&mut self, qh: &QueueHandle<Self>) {
if self.live == Live::None {
return;
}
if let Some(surface) = self.surface.clone() {
surface.frame(qh, ());
surface.commit();
}
}
pub fn tick(&mut self, qh: &QueueHandle<Self>) {
self.stats.ticks += 1;
if self.live == Live::None {
return;
}
let interval = Duration::from_secs_f64(1.0 / self.fps.max(1) as f64);
let now = Instant::now();
for i in 0..self.tiles.len() {
if self.live == Live::Current && i != self.sel {
continue;
}
if self.tiles[i].target.kind == Kind::Output {
continue;
}
if self.layout.tile(i as i32, self.scroll).is_none() {
continue;
}
let t = &self.tiles[i];
if t.session.is_none() || t.stale || t.slots.is_empty() || t.frame.is_some() {
continue;
}
if t.asked.is_some_and(|a| now.duration_since(a) < interval) {
continue;
}
self.request_capture(i, qh);
}
}
}
impl Dispatch<ExtImageCopyCaptureSessionV1, usize> for App {
fn event(
app: &mut Self,
_: &ExtImageCopyCaptureSessionV1,
event: ext_image_copy_capture_session_v1::Event,
&i: &usize,
_: &Connection,
_: &QueueHandle<Self>,
) {
let Some(tile) = app.tiles.get_mut(i) else {
return;
};
match event {
ext_image_copy_capture_session_v1::Event::BufferSize { width, height } => {
if tile.slots.is_empty() {
tile.size = (width, height);
} else if tile.size != (width, height) {
tile.stale = true;
}
}
ext_image_copy_capture_session_v1::Event::ShmFormat {
format: WEnum::Value(f),
} => {
if !tile.formats.contains(&f) {
tile.formats.push(f);
}
}
ext_image_copy_capture_session_v1::Event::Done => tile.session_done = true,
ext_image_copy_capture_session_v1::Event::Stopped => {
if let Some(session) = tile.session.take() {
session.destroy();
}
tile.settled = true;
}
_ => {}
}
}
}
impl Dispatch<ExtImageCopyCaptureFrameV1, usize> for App {
fn event(
app: &mut Self,
_: &ExtImageCopyCaptureFrameV1,
event: ext_image_copy_capture_frame_v1::Event,
&i: &usize,
_: &Connection,
qh: &QueueHandle<Self>,
) {
let Some(tile) = app.tiles.get_mut(i) else {
return;
};
match event {
ext_image_copy_capture_frame_v1::Event::Transform {
transform: WEnum::Value(t),
} => tile.transform = t,
ext_image_copy_capture_frame_v1::Event::Ready => {
if let Some(frame) = tile.frame.take() {
frame.destroy();
}
if app.frame_ready(i) {
app.sync_tiles(qh);
}
}
ext_image_copy_capture_frame_v1::Event::Failed { reason } => {
if matches!(reason, WEnum::Value(FailureReason::BufferConstraints)) {
tile.stale = true;
}
if tile.frames == 0 {
eprintln!(
"wl-pick: capture failed for {:?} ({reason:?})",
tile.target.title
);
}
tile.settled = true;
if let Some(slot) = tile.filling.take() {
tile.slots[slot].busy = false;
}
if let Some(frame) = tile.frame.take() {
frame.destroy();
}
}
_ => {}
}
}
}
impl Dispatch<WlBuffer, (usize, usize)> for App {
fn event(
app: &mut Self,
_: &WlBuffer,
event: wl_buffer::Event,
&(tile, slot): &(usize, usize),
_: &Connection,
_: &QueueHandle<Self>,
) {
if let wl_buffer::Event::Release = event
&& let Some(t) = app.tiles.get_mut(tile)
{
t.slots[slot].busy = false;
}
}
}
impl Dispatch<wl_callback::WlCallback, ()> for App {
fn event(
app: &mut Self,
_: &wl_callback::WlCallback,
event: wl_callback::Event,
_: &(),
_: &Connection,
qh: &QueueHandle<Self>,
) {
if let wl_callback::Event::Done { .. } = event {
app.tick(qh);
app.arm_frame_callback(qh);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use wl_shm::Format;
fn supported() -> Vec<Format> {
vec![
Format::Argb8888,
Format::Xrgb8888,
Format::Xbgr8888,
Format::Bgr888,
Format::Xrgb2101010,
Format::Abgr16161616f,
Format::Xbgr16161616f,
]
}
#[test]
fn the_guaranteed_formats_win() {
let offered = [Format::Abgr16161616f, Format::Bgr888, Format::Xrgb8888];
assert_eq!(
choose_format(&offered, &supported()),
Some((Format::Xrgb8888, 4))
);
}
fn wlroots_accepts(stride: i32, width: i32, bytes: i32) -> bool {
stride % bytes == 0 && stride >= width * bytes
}
#[test]
fn the_declared_stride_matches_the_format() {
let width = 3830;
for (format, bytes) in [
(Format::Bgr888, 3),
(Format::Xbgr8888, 4),
(Format::Abgr16161616f, 8),
] {
let (chosen, chosen_bytes) =
choose_format(&[format], &supported()).unwrap_or_else(|| panic!("{format:?}"));
assert_eq!((chosen, chosen_bytes), (format, bytes));
assert!(
wlroots_accepts(width * chosen_bytes, width, bytes),
"{format:?} stride {} rejected",
width * chosen_bytes
);
}
assert!(!wlroots_accepts(width * 4, width, 3), "15320 for BGR888");
assert!(!wlroots_accepts(width * 4, width, 8), "15320 for ABGR16F");
assert!(wlroots_accepts(width * 4, width, 4), "four bytes was fine");
}
#[test]
fn formats_we_cannot_measure_are_declined() {
assert_eq!(choose_format(&[Format::Nv12], &supported()), None);
assert_eq!(choose_format(&[Format::Yuv420], &supported()), None);
assert_eq!(choose_format(&[], &supported()), None);
assert_eq!(
choose_format(&[Format::Abgr16161616f], &[Format::Xrgb8888]),
None,
"offered but not displayable"
);
}
#[test]
fn pixel_sizes_match_the_names() {
assert_eq!(bytes_per_pixel(Format::Rgb565), Some(2));
assert_eq!(bytes_per_pixel(Format::Bgr888), Some(3));
assert_eq!(bytes_per_pixel(Format::Xrgb2101010), Some(4));
assert_eq!(bytes_per_pixel(Format::Xbgr16161616f), Some(8));
assert_eq!(bytes_per_pixel(Format::Yuyv), None);
assert_eq!(bytes_per_pixel(Format::Rgb332), None);
}
}