feat(server): keep a framebuffer and tile-diff against it

Signed-off-by: Marc-André Lureau <marcandre.lureau@redhat.com>

Author: elmarco

crates/ironrdp-server/src/display.rs

@@ -3,6 +3,7 @@ use core::num::NonZeroU16;
 use anyhow::Result;
 use bytes::{Bytes, BytesMut};
 use ironrdp_displaycontrol::pdu::DisplayControlMonitorLayout;
+use ironrdp_graphics::diff;
 use ironrdp_pdu::pointer::PointerPositionAttribute;
 
 #[rustfmt::skip]
@@ -94,6 +95,61 @@ impl TryInto<Framebuffer> for BitmapUpdate {
     }
 }
 
+impl Framebuffer {
+    pub fn update(&mut self, bitmap: &BitmapUpdate) {
+        if self.format != bitmap.format {
+            warn!("Bitmap format mismatch, unsupported");
+            return;
+        }
+        let bpp = usize::from(self.format.bytes_per_pixel());
+        let x = usize::from(bitmap.x);
+        let y = usize::from(bitmap.y);
+        let width = usize::from(bitmap.width.get());
+        let height = usize::from(bitmap.height.get());
+
+        let start = y * self.stride + x * bpp;
+        let end = start + (height - 1) * self.stride + width * bpp;
+        let dst = &mut self.data[start..end];
+
+        for y in 0..height {
+            let start = y * bitmap.stride;
+            let end = start + width * bpp;
+            let src = bitmap.data.slice(start..end);
+
+            let start = y * self.stride;
+            let end = start + width * bpp;
+            let dst = &mut dst[start..end];
+
+            dst.copy_from_slice(&src);
+        }
+    }
+
+    pub(crate) fn update_diffs(&mut self, bitmap: &BitmapUpdate, diffs: &[diff::Rect]) {
+        for diff in diffs {
+            let Some(x) = u16::try_from(diff.x).ok() else {
+                continue;
+            };
+            let Some(y) = u16::try_from(diff.y).ok() else {
+                continue;
+            };
+            let Some(width) = u16::try_from(diff.width).ok() else {
+                continue;
+            };
+            let Some(width) = NonZeroU16::try_from(width).ok() else {
+                continue;
+            };
+            let Some(height) = u16::try_from(diff.height).ok() else {
+                continue;
+            };
+            let Some(height) = NonZeroU16::try_from(height).ok() else {
+                continue;
+            };
+            let sub = bitmap.sub(x, y, width, height).unwrap();
+            self.update(&sub)
+        }
+    }
+}
+
 /// Bitmap Display Update
 ///
 /// Bitmap updates are encoded using RDP 6.0 compression, fragmented and sent using
@@ -268,3 +324,42 @@ pub trait RdpServerDisplay: Send {
         debug!(?layout, "Requesting layout")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{BitmapUpdate, Framebuffer};
+    use core::num::NonZeroU16;
+    use ironrdp_graphics::{diff::Rect, image_processing::PixelFormat};
+
+    #[test]
+    fn framebuffer_update() {
+        let width = NonZeroU16::new(800).unwrap();
+        let height = NonZeroU16::new(600).unwrap();
+        let fmt = PixelFormat::ABgr32;
+        let bpp = usize::from(fmt.bytes_per_pixel());
+        let mut fb = Framebuffer::new(width, height, fmt);
+
+        let width = 15;
+        let stride = width * bpp;
+        let height = 20;
+        let data = vec![1u8; height * stride];
+        let update = BitmapUpdate::new(
+            1,
+            2,
+            NonZeroU16::new(15).unwrap(),
+            NonZeroU16::new(20).unwrap(),
+            fmt,
+            data,
+            stride,
+        );
+        let diffs = vec![Rect::new(2, 3, 4, 5)];
+        fb.update_diffs(&update, &diffs);
+        for y in 5..10 {
+            for x in 3..7 {
+                for b in 0..bpp {
+                    assert_eq!(fb.data[y * fb.stride + x * bpp + b], 1);
+                }
+            }
+        }
+    }
+}

crates/ironrdp-server/src/encoder/mod.rs

@@ -1,7 +1,9 @@
 use core::fmt;
+use core::num::NonZeroU16;
 
 use anyhow::{Context, Result};
 use ironrdp_acceptor::DesktopSize;
+use ironrdp_graphics::diff::{find_different_rects_sub, Rect};
 use ironrdp_pdu::encode_vec;
 use ironrdp_pdu::fast_path::UpdateCode;
 use ironrdp_pdu::geometry::ExclusiveRectangle;
@@ -65,7 +67,7 @@ impl UpdateEncoder {
     pub(crate) fn update(&mut self, update: DisplayUpdate) -> EncoderIter<'_> {
         EncoderIter {
             encoder: self,
-            update: Some(update),
+            state: State::Start(update),
         }
     }
 
@@ -129,14 +131,40 @@ impl UpdateEncoder {
         Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, encode_vec(&pos)?))
     }
 
-    async fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter> {
-        // Clone to satisfy spawn_blocking 'static requirement
-        // this should be cheap, even if using bitmap, since vec![] will be empty
-        let mut updater = self.bitmap_updater.clone();
-        let (res, bitmap) =
-            tokio::task::spawn_blocking(move || time_warn!("Encoding bitmap", 10, (updater.handle(&bitmap), bitmap)))
-                .await
-                .unwrap();
+    fn bitmap_diffs(&mut self, bitmap: &BitmapUpdate) -> Vec<Rect> {
+        const USE_DIFFS: bool = true;
+
+        if let Some(Framebuffer {
+            data,
+            stride,
+            width,
+            height,
+            ..
+        }) = USE_DIFFS.then_some(self.framebuffer.as_ref()).flatten()
+        {
+            find_different_rects_sub::<4>(
+                data,
+                *stride,
+                width.get().into(),
+                height.get().into(),
+                &bitmap.data,
+                bitmap.stride,
+                bitmap.width.get().into(),
+                bitmap.height.get().into(),
+                bitmap.x.into(),
+                bitmap.y.into(),
+            )
+        } else {
+            vec![Rect {
+                x: 0,
+                y: 0,
+                width: bitmap.width.get().into(),
+                height: bitmap.height.get().into(),
+            }]
+        }
+    }
+
+    fn bitmap_update_framebuffer(&mut self, bitmap: BitmapUpdate, diffs: &[Rect]) {
         if bitmap.x == 0
             && bitmap.y == 0
             && bitmap.width.get() == self.desktop_size.width
@@ -146,34 +174,88 @@ impl UpdateEncoder {
                 Ok(framebuffer) => self.framebuffer = Some(framebuffer),
                 Err(err) => warn!("Failed to convert bitmap to framebuffer: {}", err),
             }
+        } else if let Some(fb) = self.framebuffer.as_mut() {
+            fb.update_diffs(&bitmap, diffs);
         }
-        res
     }
+
+    async fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter> {
+        // Clone to satisfy spawn_blocking 'static requirement
+        // this should be cheap, even if using bitmap, since vec![] will be empty
+        let mut updater = self.bitmap_updater.clone();
+        tokio::task::spawn_blocking(move || time_warn!("Encoding bitmap", 10, updater.handle(&bitmap)))
+            .await
+            .unwrap()
+    }
+}
+
+#[derive(Debug, Default)]
+enum State {
+    Start(DisplayUpdate),
+    BitmapDiffs {
+        diffs: Vec<Rect>,
+        bitmap: BitmapUpdate,
+        pos: usize,
+    },
+    #[default]
+    Ended,
 }
 
 #[cfg_attr(feature = "__bench", visibility::make(pub))]
 pub(crate) struct EncoderIter<'a> {
     encoder: &'a mut UpdateEncoder,
-    update: Option<DisplayUpdate>,
+    state: State,
 }
 
 impl EncoderIter<'_> {
     #[cfg_attr(feature = "__bench", visibility::make(pub))]
     pub(crate) async fn next(&mut self) -> Option<Result<UpdateFragmenter>> {
-        let update = self.update.take()?;
-        let encoder = &mut self.encoder;
-
-        let res = match update {
-            DisplayUpdate::Bitmap(bitmap) => encoder.bitmap(bitmap).await,
-            DisplayUpdate::PointerPosition(pos) => encoder.pointer_position(pos),
-            DisplayUpdate::RGBAPointer(ptr) => encoder.rgba_pointer(ptr),
-            DisplayUpdate::ColorPointer(ptr) => encoder.color_pointer(ptr),
-            DisplayUpdate::HidePointer => encoder.hide_pointer(),
-            DisplayUpdate::DefaultPointer => encoder.default_pointer(),
-            DisplayUpdate::Resize(_) => return None,
-        };
-
-        Some(res)
+        loop {
+            let state = core::mem::take(&mut self.state);
+            let encoder = &mut self.encoder;
+
+            let res = match state {
+                State::Start(update) => match update {
+                    DisplayUpdate::Bitmap(bitmap) => {
+                        let diffs = encoder.bitmap_diffs(&bitmap);
+                        self.state = State::BitmapDiffs { diffs, bitmap, pos: 0 };
+                        continue;
+                    }
+                    DisplayUpdate::PointerPosition(pos) => encoder.pointer_position(pos),
+                    DisplayUpdate::RGBAPointer(ptr) => encoder.rgba_pointer(ptr),
+                    DisplayUpdate::ColorPointer(ptr) => encoder.color_pointer(ptr),
+                    DisplayUpdate::HidePointer => encoder.hide_pointer(),
+                    DisplayUpdate::DefaultPointer => encoder.default_pointer(),
+                    DisplayUpdate::Resize(_) => return None,
+                },
+                State::BitmapDiffs { diffs, bitmap, pos } => {
+                    let Some(rect) = diffs.get(pos) else {
+                        encoder.bitmap_update_framebuffer(bitmap, &diffs);
+                        self.state = State::Ended;
+                        return None;
+                    };
+                    let Rect { x, y, width, height } = *rect;
+                    let Some(sub) = bitmap.sub(
+                        u16::try_from(x).unwrap(),
+                        u16::try_from(y).unwrap(),
+                        NonZeroU16::new(u16::try_from(width).unwrap()).unwrap(),
+                        NonZeroU16::new(u16::try_from(height).unwrap()).unwrap(),
+                    ) else {
+                        warn!("Failed to extract bitmap subregion");
+                        return None;
+                    };
+                    self.state = State::BitmapDiffs {
+                        diffs,
+                        bitmap,
+                        pos: pos + 1,
+                    };
+                    encoder.bitmap(sub).await
+                }
+                State::Ended => return None,
+            };
+
+            return Some(res);
+        }
     }
 }