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.

6 files changed, 99 insertions, 14 deletions

diff --git a/Userland/Services/WindowServer/ClientConnection.cpp b/Userland/Services/WindowServer/ClientConnection.cpp
index e0716ede4e..b6a267d3b0 100644
--- a/Userland/Services/WindowServer/ClientConnection.cpp
+++ b/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 (rect.has_value()) {
-        auto bitmap = Compositor::the().front_bitmap_for_screenshot({}, screen).cropped(rect.value());
+    if (screen_index.has_value()) {
+        auto* screen = Screen::find_by_index(screen_index.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 bitmap.to_shareable_bitmap();
+    return { {} };
 }
 
 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
-    auto scale_factor = screen.scale_factor();
+    // TODO: Mixed scale setups at what scale? Lowest? Highest? Configurable?
     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);
-    return bitmap->to_shareable_bitmap();
+    // Check if we need to compose from multiple screens. If not we can take a fast path
+    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)
diff --git a/Userland/Services/WindowServer/ClientConnection.h b/Userland/Services/WindowServer/ClientConnection.h
index 4f171712de..944b490c9f 100644
--- a/Userland/Services/WindowServer/ClientConnection.h
+++ b/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;
diff --git a/Userland/Services/WindowServer/Compositor.cpp b/Userland/Services/WindowServer/Compositor.cpp
index c70cd6807c..962b0c3b49 100644
--- a/Userland/Services/WindowServer/Compositor.cpp
+++ b/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;
diff --git a/Userland/Services/WindowServer/Compositor.h b/Userland/Services/WindowServer/Compositor.h
index ea245dfdb8..4040c972aa 100644
--- a/Userland/Services/WindowServer/Compositor.h
+++ b/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:
diff --git a/Userland/Services/WindowServer/WindowServer.ipc b/Userland/Services/WindowServer/WindowServer.ipc
index 7ea2b7580c..8559b1464b 100644
--- a/Userland/Services/WindowServer/WindowServer.ipc
+++ b/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() =|
diff --git a/Userland/Utilities/shot.cpp b/Userland/Utilities/shot.cpp
index 7897ea2230..c25ff95573 100644
--- a/Userland/Utilities/shot.cpp
+++ b/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) {