vm: get screenshot PNG data early

Fixes #4009
utmapp · Feb 25, 2024 · da9c5c4 · da9c5c4
1 parent 07650fa
commit da9c5c4
Show file tree

Hide file tree

Showing 8 changed files with 57 additions and 33 deletions.
diff --git a/Platform/VMData.swift b/Platform/VMData.swift
@@ -67,7 +67,7 @@ import SwiftUI
     @Published var state: UTMVirtualMachineState = .stopped
 
     /// Copy from wrapped VM
-    @Published var screenshot: PlatformImage?
+    @Published var screenshot: UTMVirtualMachineScreenshot?
 
     /// Allows changes in the config, registry, and VM to be reflected
     private var observers: [AnyCancellable] = []
@@ -426,7 +426,7 @@ extension VMData {
 
     /// If non-null, is the most recent screenshot image of the running VM
     var screenshotImage: PlatformImage? {
-        wrapped?.screenshot
+        wrapped?.screenshot?.image
     }
 }
 

diff --git a/Platform/macOS/Display/VMDisplayAppleWindowController.swift b/Platform/macOS/Display/VMDisplayAppleWindowController.swift
@@ -257,9 +257,9 @@ extension VMDisplayAppleWindowController {
 }
 
 extension VMDisplayAppleWindowController: UTMScreenshotProvider {
-    var screenshot: PlatformImage? {
+    var screenshot: UTMVirtualMachineScreenshot? {
         if let image = mainView?.image() {
-            return image
+            return UTMVirtualMachineScreenshot(wrapping: image)
         } else {
             return nil
         }

diff --git a/Platform/macOS/Display/VMDisplayQemuMetalWindowController.swift b/Platform/macOS/Display/VMDisplayQemuMetalWindowController.swift
@@ -149,7 +149,7 @@ class VMDisplayQemuMetalWindowController: VMDisplayQemuWindowController {
     override func enterSuspended(isBusy busy: Bool) {
         if !busy {
             metalView.isHidden = true
-            screenshotView.image = vm.screenshot
+            screenshotView.image = vm.screenshot?.image
             screenshotView.isHidden = false
         }
         if vm.state == .stopped {

diff --git a/Remote/UTMRemoteSpiceVirtualMachine.swift b/Remote/UTMRemoteSpiceVirtualMachine.swift
@@ -98,7 +98,7 @@ final class UTMRemoteSpiceVirtualMachine: UTMSpiceVirtualMachine {
         }
     }
 
-    var screenshot: PlatformImage? {
+    var screenshot: UTMVirtualMachineScreenshot? {
         willSet {
             onStateChange?()
         }
@@ -273,11 +273,11 @@ extension UTMRemoteSpiceVirtualMachine {
     }
 
     func loadScreenshot(from url: URL) {
-        screenshot = UIImage(contentsOfURL: url)
+        screenshot = UTMVirtualMachineScreenshot(contentsOfURL: url)
     }
 
     func saveScreenshot() async {
-        if let data = screenshot?.pngData() {
+        if let data = screenshot?.pngData {
             try? await server.sendPackageFile(for: id, relativePathComponents: [kUTMBundleScreenshotFilename], data: data)
         }
     }

diff --git a/Services/UTMAppleVirtualMachine.swift b/Services/UTMAppleVirtualMachine.swift
@@ -89,7 +89,7 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
         }
     }
 
-    private(set) var screenshot: PlatformImage? {
+    private(set) var screenshot: UTMVirtualMachineScreenshot? {
         willSet {
             onStateChange?()
         }
@@ -729,7 +729,7 @@ extension UTMAppleVirtualMachine: VZVirtualMachineDelegate {
 }
 
 protocol UTMScreenshotProvider: AnyObject {
-    var screenshot: PlatformImage? { get }
+    var screenshot: UTMVirtualMachineScreenshot? { get }
 }
 
 enum UTMAppleVirtualMachineError: Error {

diff --git a/Services/UTMQemuVirtualMachine.swift b/Services/UTMQemuVirtualMachine.swift
@@ -95,7 +95,7 @@ final class UTMQemuVirtualMachine: UTMSpiceVirtualMachine {
         }
     }
 
-    var screenshot: PlatformImage? {
+    var screenshot: UTMVirtualMachineScreenshot? {
         willSet {
             onStateChange?()
         }

diff --git a/Services/UTMSpiceVirtualMachine.swift b/Services/UTMSpiceVirtualMachine.swift
@@ -22,7 +22,7 @@ protocol UTMSpiceVirtualMachine: UTMVirtualMachine where Configuration == UTMQem
     var isRunningAsDisposible: Bool { get }
 
     /// Get and set screenshot
-    var screenshot: PlatformImage? { get set }
+    var screenshot: UTMVirtualMachineScreenshot? { get set }
 
     /// Handles IO
     var ioServiceDelegate: UTMSpiceIODelegate? { get set }
@@ -72,8 +72,9 @@ extension UTMSpiceVirtualMachine {
 extension UTMSpiceVirtualMachine {
     @MainActor @discardableResult
     func takeScreenshot() async -> Bool {
-        let screenshot = await ioService?.screenshot()
-        self.screenshot = screenshot?.image
+        if let screenshot = await ioService?.screenshot() {
+            self.screenshot = UTMVirtualMachineScreenshot(wrapping: screenshot.image)
+        }
         return true
     }
 

diff --git a/Services/UTMVirtualMachine.swift b/Services/UTMVirtualMachine.swift
@@ -66,8 +66,8 @@ protocol UTMVirtualMachine: AnyObject, Identifiable {
     var state: UTMVirtualMachineState { get }
 
     /// If non-null, is the most recent screenshot of the running VM
-    var screenshot: PlatformImage? { get }
-    
+    var screenshot: UTMVirtualMachineScreenshot? { get }
+
     /// If non-null, `saveSnapshot` and `restoreSnapshot` will not work due to the reason specified
     var snapshotUnsupportedError: Error? { get }
 
@@ -290,6 +290,43 @@ extension UTMVirtualMachine {
 
 // MARK: - Screenshot
 
+struct UTMVirtualMachineScreenshot {
+    let image: PlatformImage
+    let pngData: Data?
+
+    init?(contentsOfURL url: URL) {
+        #if canImport(AppKit)
+        guard let image = NSImage(contentsOf: url) else {
+            return nil
+        }
+        #elseif canImport(UIKit)
+        guard let image = UIImage(contentsOfURL: url) else {
+            return nil
+        }
+        #endif
+        self.image = image
+        self.pngData = Self.createData(from: image)
+    }
+
+    init(wrapping image: PlatformImage) {
+        self.image = image
+        self.pngData = Self.createData(from: image)
+    }
+
+    private static func createData(from image: PlatformImage) -> Data? {
+        #if canImport(AppKit)
+        guard let cgref = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
+            return nil
+        }
+        let newrep = NSBitmapImageRep(cgImage: cgref)
+        newrep.size = image.size
+        return newrep.representation(using: .png, properties: [:])
+        #elseif canImport(UIKit)
+        return image.pngData()
+        #endif
+    }
+}
+
 extension UTMVirtualMachine {
     private var isScreenshotSaveEnabled: Bool {
         !UserDefaults.standard.bool(forKey: "NoSaveScreenshot")
@@ -319,12 +356,8 @@ extension UTMVirtualMachine {
         return timer
     }
 
-    func loadScreenshot() -> PlatformImage? {
-        #if canImport(AppKit)
-        return NSImage(contentsOf: screenshotUrl)
-        #elseif canImport(UIKit)
-        return UIImage(contentsOfURL: screenshotUrl)
-        #endif
+    func loadScreenshot() -> UTMVirtualMachineScreenshot? {
+        UTMVirtualMachineScreenshot(contentsOfURL: screenshotUrl)
     }
 
     func saveScreenshot() throws {
@@ -334,17 +367,7 @@ extension UTMVirtualMachine {
         guard let screenshot = screenshot else {
             return
         }
-        #if canImport(AppKit)
-        guard let cgref = screenshot.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
-            return
-        }
-        let newrep = NSBitmapImageRep(cgImage: cgref)
-        newrep.size = screenshot.size
-        let pngdata = newrep.representation(using: .png, properties: [:])
-        try pngdata?.write(to: screenshotUrl)
-        #elseif canImport(UIKit)
-        try screenshot.pngData()?.write(to: screenshotUrl)
-        #endif
+        try screenshot.pngData?.write(to: screenshotUrl)
     }
 
     func deleteScreenshot() throws {