From 9e3a152b302073c2fdd46c6de5ba328f40ba101c Mon Sep 17 00:00:00 2001
From: Hiroshi Horie <548776+hiroshihorie@users.noreply.github.com>
Date: Mon, 29 Jul 2024 17:10:21 +0800
Subject: [PATCH] Optimize SampleBufferVideoRenderer (#441)

---
 Sources/LiveKit/Protocols/Mirrorable.swift    |  2 +-
 .../Views/SampleBufferVideoRenderer.swift     | 71 +++++++++++--------
 Sources/LiveKit/Views/VideoView.swift         |  8 +--
 3 files changed, 48 insertions(+), 33 deletions(-)

diff --git a/Sources/LiveKit/Protocols/Mirrorable.swift b/Sources/LiveKit/Protocols/Mirrorable.swift
index f7136f25b..04f3d4ec0 100644
--- a/Sources/LiveKit/Protocols/Mirrorable.swift
+++ b/Sources/LiveKit/Protocols/Mirrorable.swift
@@ -18,5 +18,5 @@ import Foundation
 
 // Internal only
 protocol Mirrorable {
-    func set(mirrored: Bool)
+    func set(isMirrored: Bool)
 }
diff --git a/Sources/LiveKit/Views/SampleBufferVideoRenderer.swift b/Sources/LiveKit/Views/SampleBufferVideoRenderer.swift
index 735517269..797ed3784 100644
--- a/Sources/LiveKit/Views/SampleBufferVideoRenderer.swift
+++ b/Sources/LiveKit/Views/SampleBufferVideoRenderer.swift
@@ -22,14 +22,14 @@ internal import LiveKitWebRTC
 
 class SampleBufferVideoRenderer: NativeView, Loggable {
     public let sampleBufferDisplayLayer: AVSampleBufferDisplayLayer
-    
-    private var firstFrameReceived = false
-    private var bufferTransform = CATransform3DIdentity
-    private var mirroredTransform = CATransform3DIdentity
-    private var displayLayerTransform: CATransform3D {
-        return CATransform3DConcat(bufferTransform, mirroredTransform)
+
+    private struct State {
+        var isMirrored: Bool = false
+        var videoRotation: VideoRotation = ._0
     }
 
+    private let _state = StateSync(State())
+
     override init(frame: CGRect) {
         sampleBufferDisplayLayer = AVSampleBufferDisplayLayer()
         super.init(frame: frame)
@@ -53,14 +53,16 @@ class SampleBufferVideoRenderer: NativeView, Loggable {
     override func performLayout() {
         super.performLayout()
         sampleBufferDisplayLayer.frame = bounds
+
+        let (rotation, isMirrored) = _state.read { ($0.videoRotation, $0.isMirrored) }
+        sampleBufferDisplayLayer.transform = CATransform3D.from(rotation: rotation, isMirrored: isMirrored)
+
         sampleBufferDisplayLayer.removeAllAnimations()
     }
 }
 
 extension SampleBufferVideoRenderer: LKRTCVideoRenderer {
-    func setSize(_: CGSize) {
-        //
-    }
+    func setSize(_: CGSize) {}
 
     func renderFrame(_ frame: LKRTCVideoFrame?) {
         guard let frame else { return }
@@ -82,43 +84,56 @@ extension SampleBufferVideoRenderer: LKRTCVideoRenderer {
             log("Failed to convert CVPixelBuffer to CMSampleBuffer", .error)
             return
         }
-        
-        if !firstFrameReceived {
-            bufferTransform = .fromFrameRotation(frame)
-            updateSampleBufferTransform()
-            firstFrameReceived = true
+
+        let rotation = frame.rotation.toLKType()
+        let didUpdateRotation = _state.mutate {
+            let result = $0.videoRotation != rotation
+            $0.videoRotation = rotation
+            return result
         }
 
         Task.detached { @MainActor in
             self.sampleBufferDisplayLayer.enqueue(sampleBuffer)
+            if didUpdateRotation {
+                self.setNeedsLayout()
+            }
         }
     }
 }
 
 extension SampleBufferVideoRenderer: Mirrorable {
-    func set(mirrored: Bool) {
-        mirroredTransform = mirrored ? VideoView.mirrorTransform : CATransform3DIdentity
-        updateSampleBufferTransform()
-    }
-}
+    func set(isMirrored: Bool) {
+        let didUpdateIsMirrored = _state.mutate {
+            let result = $0.isMirrored != isMirrored
+            $0.isMirrored = isMirrored
+            return result
+        }
 
-private extension SampleBufferVideoRenderer {
-    private func updateSampleBufferTransform() {
-        sampleBufferDisplayLayer.transform = displayLayerTransform
+        if didUpdateIsMirrored {
+            setNeedsLayout()
+        }
     }
 }
 
 private extension CATransform3D {
-    static func fromFrameRotation(_ frame: LKRTCVideoFrame) -> CATransform3D {
-        switch frame.rotation {
+    static func from(rotation: VideoRotation, isMirrored: Bool) -> CATransform3D {
+        var transform: CATransform3D
+
+        switch rotation {
         case ._0:
-            return CATransform3DIdentity
+            transform = CATransform3DIdentity
         case ._90:
-            return CATransform3DMakeRotation(.pi / 2.0, 0, 0, 1)
+            transform = CATransform3DMakeRotation(.pi / 2.0, 0, 0, 1)
         case ._180:
-            return CATransform3DMakeRotation(.pi, 0, 0, 1)
+            transform = CATransform3DMakeRotation(.pi, 0, 0, 1)
         case ._270:
-            return CATransform3DMakeRotation(-.pi / 0, 0, 0, 1)
+            transform = CATransform3DMakeRotation(-.pi / 2.0, 0, 0, 1)
         }
+
+        if isMirrored {
+            transform = CATransform3DConcat(transform, VideoView.mirrorTransform)
+        }
+
+        return transform
     }
 }
diff --git a/Sources/LiveKit/Views/VideoView.swift b/Sources/LiveKit/Views/VideoView.swift
index d5d772be1..0d59c811a 100644
--- a/Sources/LiveKit/Views/VideoView.swift
+++ b/Sources/LiveKit/Views/VideoView.swift
@@ -500,9 +500,9 @@ public class VideoView: NativeView, Loggable {
 
             if let _secondaryRenderer {
                 _secondaryRenderer.frame = rendererFrame
-                _secondaryRenderer.set(mirrored: _shouldMirror())
+                _secondaryRenderer.set(isMirrored: _shouldMirror())
             } else {
-                _primaryRenderer.set(mirrored: _shouldMirror())
+                _primaryRenderer.set(isMirrored: _shouldMirror())
             }
         }
     }
@@ -810,8 +810,8 @@ extension NSView {
 #endif
 
 extension LKRTCMTLVideoView: Mirrorable {
-    func set(mirrored: Bool) {
-        if mirrored {
+    func set(isMirrored: Bool) {
+        if isMirrored {
             #if os(macOS)
             // This is required for macOS
             wantsLayer = true