WindowServer: Enable screen capture to span multiple screens

This enables the shot utility to capture all screens or just one, and
enables the Magnifier application to track the mouse cursor across
multiple screens.

Userland/Services/WindowServer/ClientConnection.cpp

@@ -921,30 +921,100 @@ void ClientConnection::did_become_responsive()
     set_unresponsive(false);
 }
 
-Messages::WindowServer::GetScreenBitmapResponse ClientConnection::get_screen_bitmap(Optional<Gfx::IntRect> const& rect)
+Messages::WindowServer::GetScreenBitmapResponse ClientConnection::get_screen_bitmap(Optional<Gfx::IntRect> const& rect, Optional<u32> const& screen_index)
 {
-    auto& screen = Screen::main(); // TODO: implement screenshots from other screens or areas spanning multiple screens
+    if (screen_index.has_value()) {
-    if (rect.has_value()) {
+        auto* screen = Screen::find_by_index(screen_index.value());
-        auto bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen).cropped(rect.value());
+        if (!screen) {
+            dbgln("get_screen_bitmap: Screen {} does not exist!", screen_index.value());
+            return { {} };
+        }
+        if (rect.has_value()) {
+            auto bitmap = Compositor::the().front_bitmap_for_screenshot({}, *screen).cropped(rect.value());
+            return bitmap->to_shareable_bitmap();
+        }
+        auto& bitmap = Compositor::the().front_bitmap_for_screenshot({}, *screen);
+        return bitmap.to_shareable_bitmap();
+    }
+    // TODO: Mixed scale setups at what scale? Lowest? Highest? Configurable?
+    if (auto bitmap = Gfx::Bitmap::create(Gfx::BitmapFormat::BGRx8888, Screen::bounding_rect().size(), 1)) {
+        Gfx::Painter painter(*bitmap);
+        Screen::for_each([&](auto& screen) {
+            auto screen_rect = screen.rect();
+            if (rect.has_value() && !rect.value().intersects(screen_rect))
+                return IterationDecision::Continue;
+            auto src_rect = rect.has_value() ? rect.value().intersected(screen_rect) : screen_rect;
+            VERIFY(Screen::bounding_rect().contains(src_rect));
+            auto& screen_bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen);
+            // TODO: painter does *not* support down-sampling!!!
+            painter.blit(screen_rect.location(), screen_bitmap, src_rect.translated(-screen_rect.location()), 1.0f, false);
+            return IterationDecision::Continue;
+        });
         return bitmap->to_shareable_bitmap();
     }
-    auto& bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen);
+    return { {} };
-    return bitmap.to_shareable_bitmap();
 }
 
 Messages::WindowServer::GetScreenBitmapAroundCursorResponse ClientConnection::get_screen_bitmap_around_cursor(Gfx::IntSize const& size)
 {
-    auto& screen = Screen::main(); // TODO: implement getting screen bitmaps from other screens or areas spanning multiple screens
+    // TODO: Mixed scale setups at what scale? Lowest? Highest? Configurable?
-    auto scale_factor = screen.scale_factor();
     auto cursor_location = ScreenInput::the().cursor_location();
-    Gfx::Rect rect { (cursor_location.x() * scale_factor) - (size.width() / 2), (cursor_location.y() * scale_factor) - (size.height() / 2), size.width(), size.height() };
+    Gfx::Rect rect { cursor_location.x() - (size.width() / 2), cursor_location.y() - (size.height() / 2), size.width(), size.height() };
 
     // Recompose the screen to make sure the cursor is painted in the location we think it is.
     // FIXME: This is rather wasteful. We can probably think of a way to avoid this.
     Compositor::the().compose();
 
-    auto bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen).cropped(rect);
+    // Check if we need to compose from multiple screens. If not we can take a fast path
-    return bitmap->to_shareable_bitmap();
+    size_t intersecting_with_screens = 0;
+    Screen::for_each([&](auto& screen) {
+        if (rect.intersects(screen.rect()))
+            intersecting_with_screens++;
+        return IterationDecision::Continue;
+    });
+
+    if (intersecting_with_screens == 1) {
+        auto& screen = Screen::closest_to_rect(rect);
+        auto bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen).cropped(rect.translated(-screen.rect().location()));
+        return bitmap->to_shareable_bitmap();
+    }
+
+    if (auto bitmap = Gfx::Bitmap::create(Gfx::BitmapFormat::BGRx8888, rect.size(), 1)) {
+        auto bounding_screen_src_rect = Screen::bounding_rect().intersected(rect);
+        Gfx::Painter painter(*bitmap);
+        Gfx::IntRect last_cursor_rect;
+        auto& screen_with_cursor = ScreenInput::the().cursor_location_screen();
+        auto cursor_rect = Compositor::the().current_cursor_rect();
+        Screen::for_each([&](auto& screen) {
+            auto screen_rect = screen.rect();
+            auto src_rect = screen_rect.intersected(bounding_screen_src_rect);
+            if (src_rect.is_empty())
+                return IterationDecision ::Continue;
+            auto& screen_bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen);
+            auto from_rect = src_rect.translated(-screen_rect.location());
+            auto target_location = rect.intersected(screen_rect).location().translated(-rect.location());
+            // TODO: painter does *not* support down-sampling!!!
+            painter.blit(target_location, screen_bitmap, from_rect, 1.0f, false);
+            // Check if we are a screen that doesn't have the cursor but the cursor would
+            // have normally been cut off (we don't draw portions of the cursor on a screen
+            // that doesn't actually have the cursor). In that case we need to render the remaining
+            // portion of the cursor on that screen's capture manually
+            if (&screen != &screen_with_cursor) {
+                auto screen_cursor_rect = cursor_rect.intersected(screen_rect);
+                if (!screen_cursor_rect.is_empty()) {
+                    if (auto const* cursor_bitmap = Compositor::the().cursor_bitmap_for_screenshot({}, screen)) {
+                        auto src_rect = screen_cursor_rect.translated(-cursor_rect.location());
+                        auto cursor_target = cursor_rect.intersected(screen_rect).location().translated(-rect.location());
+                        // TODO: painter does *not* support down-sampling!!!
+                        painter.blit(cursor_target, *cursor_bitmap, src_rect);
+                    }
+                }
+            }
+            return IterationDecision::Continue;
+        });
+        return bitmap->to_shareable_bitmap();
+    }
+    return { {} };
 }
 
 Messages::WindowServer::IsWindowModifiedResponse ClientConnection::is_window_modified(i32 window_id)

Userland/Services/WindowServer/ClientConnection.h

@@ -151,7 +151,7 @@ private:
     virtual Messages::WindowServer::GetMouseAccelerationResponse get_mouse_acceleration() override;
     virtual void set_scroll_step_size(u32) override;
     virtual Messages::WindowServer::GetScrollStepSizeResponse get_scroll_step_size() override;
-    virtual Messages::WindowServer::GetScreenBitmapResponse get_screen_bitmap(Optional<Gfx::IntRect> const&) override;
+    virtual Messages::WindowServer::GetScreenBitmapResponse get_screen_bitmap(Optional<Gfx::IntRect> const&, Optional<u32> const&) override;
     virtual Messages::WindowServer::GetScreenBitmapAroundCursorResponse get_screen_bitmap_around_cursor(Gfx::IntSize const&) override;
     virtual void set_double_click_speed(i32) override;
     virtual Messages::WindowServer::GetDoubleClickSpeedResponse get_double_click_speed() override;

Userland/Services/WindowServer/Compositor.cpp

@@ -64,6 +64,13 @@ Compositor::Compositor()
     init_bitmaps();
 }
 
+const Gfx::Bitmap* Compositor::cursor_bitmap_for_screenshot(Badge<ClientConnection>, Screen& screen) const
+{
+    if (!m_current_cursor)
+        return nullptr;
+    return &m_current_cursor->bitmap(screen.scale_factor());
+}
+
 const Gfx::Bitmap& Compositor::front_bitmap_for_screenshot(Badge<ClientConnection>, Screen& screen) const
 {
     return *m_screen_data[screen.index()].m_front_bitmap;

Userland/Services/WindowServer/Compositor.h

@@ -56,6 +56,7 @@ public:
 
     void did_construct_window_manager(Badge<WindowManager>);
 
+    const Gfx::Bitmap* cursor_bitmap_for_screenshot(Badge<ClientConnection>, Screen&) const;
     const Gfx::Bitmap& front_bitmap_for_screenshot(Badge<ClientConnection>, Screen&) const;
 
 private:

Userland/Services/WindowServer/WindowServer.ipc

@@ -125,7 +125,7 @@ endpoint WindowServer
     set_scroll_step_size(u32 step_size) => ()
     get_scroll_step_size() => (u32 step_size)
 
-    get_screen_bitmap(Optional<Gfx::IntRect> rect) => (Gfx::ShareableBitmap bitmap)
+    get_screen_bitmap(Optional<Gfx::IntRect> rect, Optional<u32> screen_index) => (Gfx::ShareableBitmap bitmap)
     get_screen_bitmap_around_cursor(Gfx::IntSize size) => (Gfx::ShareableBitmap bitmap)
 
     pong() =|

Userland/Utilities/shot.cpp

@@ -21,10 +21,12 @@ int main(int argc, char** argv)
     String output_path;
     bool output_to_clipboard = false;
     int delay = 0;
+    int screen = -1;
 
     args_parser.add_positional_argument(output_path, "Output filename", "output", Core::ArgsParser::Required::No);
     args_parser.add_option(output_to_clipboard, "Output to clipboard", "clipboard", 'c');
     args_parser.add_option(delay, "Seconds to wait before taking a screenshot", "delay", 'd', "seconds");
+    args_parser.add_option(screen, "The index of the screen (default: -1 for all screens)", "screen", 's', "index");
 
     args_parser.parse(argc, argv);
 
@@ -34,7 +36,12 @@ int main(int argc, char** argv)
 
     auto app = GUI::Application::construct(argc, argv);
     sleep(delay);
-    auto shared_bitmap = GUI::WindowServerConnection::the().get_screen_bitmap({});
+    Optional<u32> screen_index;
+    if (screen >= 0)
+        screen_index = (u32)screen;
+    dbgln("getting screenshot...");
+    auto shared_bitmap = GUI::WindowServerConnection::the().get_screen_bitmap({}, screen_index);
+    dbgln("got screenshot");
 
     auto* bitmap = shared_bitmap.bitmap();
     if (!bitmap) {