[WIP] Stickers editor

submodules/TelegramUI/Components/MediaEditor/Sources/UniversalTextureSource.swift

@@ -3,11 +3,13 @@ import AVFoundation
 import Metal
 import MetalKit
 import ImageTransparency
+import SwiftSignalKit
 
 final class UniversalTextureSource: TextureSource {
     enum Input {
         case image(UIImage)
         case video(AVPlayerItem)
+        case entity(MediaEditorComposerEntity)
         
         fileprivate func createContext(renderTarget: RenderTarget, queue: DispatchQueue, additional: Bool) -> InputContext {
             switch self {
@@ -15,6 +17,8 @@ final class UniversalTextureSource: TextureSource {
                 return ImageInputContext(input: self, renderTarget: renderTarget, queue: queue)
             case .video:
                 return VideoInputContext(input: self, renderTarget: renderTarget, queue: queue, additional: additional)
+            case .entity:
+                return EntityInputContext(input: self, renderTarget: renderTarget, queue: queue)
             }
         }
     }
@@ -76,9 +80,15 @@ final class UniversalTextureSource: TextureSource {
     }
     
     private var previousAdditionalOutput: MediaEditorRenderer.Input?
+    private var readyForMoreData = Atomic<Bool>(value: true)
     private func update(forced: Bool) {
         let time = CACurrentMediaTime()
         
+        var fps: Int = 60
+        if self.mainInputContext?.useAsyncOutput == true {
+            fps = 30
+        }
+        
         let needsDisplayLink = (self.mainInputContext?.needsDisplayLink ?? false) || (self.additionalInputContext?.needsDisplayLink ?? false)
         if needsDisplayLink {
             if self.displayLink == nil {
@@ -87,7 +97,7 @@ final class UniversalTextureSource: TextureSource {
                         self.update(forced: self.forceUpdates)
                     }
                 }), selector: #selector(DisplayLinkTarget.handleDisplayLinkUpdate(sender:)))
-                displayLink.preferredFramesPerSecond = 60
+                displayLink.preferredFramesPerSecond = fps
                 displayLink.add(to: .main, forMode: .common)
                 self.displayLink = displayLink
             }
@@ -102,19 +112,33 @@ final class UniversalTextureSource: TextureSource {
             return
         }
                 
-        let main = self.mainInputContext?.output(time: time)
-        var additional = self.additionalInputContext?.output(time: time)
-        if let additional {
-            self.previousAdditionalOutput = additional
-        } else if self.additionalInputContext != nil {
-            additional = self.previousAdditionalOutput
+        if let mainInputContext = self.mainInputContext, mainInputContext.useAsyncOutput {
+            guard self.readyForMoreData.with({ $0 }) else {
+                return
+            }
+            let _ = self.readyForMoreData.swap(false)
+            mainInputContext.asyncOutput(time: time, completion: { [weak self] main in
+                guard let self else {
+                    return
+                }
+                if let main {
+                    self.output?.consume(main: main, additional: nil, render: true)
+                }
+                let _ = self.readyForMoreData.swap(true)
+            })
+        } else {
+            let main = self.mainInputContext?.output(time: time)
+            var additional = self.additionalInputContext?.output(time: time)
+            if let additional {
+                self.previousAdditionalOutput = additional
+            } else if self.additionalInputContext != nil {
+                additional = self.previousAdditionalOutput
+            }
+            guard let main else {
+                return
+            }
+            self.output?.consume(main: main, additional: additional, render: true)
         }
-        
-        guard let main else {
-            return
-        }
-        
-        self.output?.consume(main: main, additional: additional, render: true)
     }
     
     func connect(to consumer: MediaEditorRenderer) {
@@ -138,18 +162,31 @@ final class UniversalTextureSource: TextureSource {
     }
 }
 
-private protocol InputContext {
+protocol InputContext {
     typealias Input = UniversalTextureSource.Input
     typealias Output = MediaEditorRenderer.Input
     
     var input: Input { get }
+    
+    var useAsyncOutput: Bool { get }
     func output(time: Double) -> Output?
+    func asyncOutput(time: Double, completion: @escaping (Output?) -> Void)
     
     var needsDisplayLink: Bool { get }
     
     func invalidate()
 }
 
+extension InputContext {
+    var useAsyncOutput: Bool {
+        return false
+    }
+    
+    func asyncOutput(time: Double, completion: @escaping (Output?) -> Void) {
+        completion(self.output(time: time))
+    }
+}
+
 private class ImageInputContext: InputContext {
     fileprivate var input: Input
     private var texture: MTLTexture?
@@ -248,3 +285,59 @@ private class VideoInputContext: NSObject, InputContext, AVPlayerItemOutputPullD
         return true
     }
 }
+
+final class EntityInputContext: NSObject, InputContext, AVPlayerItemOutputPullDelegate {
+    internal var input: Input
+    private var textureRotation: TextureRotation = .rotate0Degrees
+    
+    var entity: MediaEditorComposerEntity {
+        guard case let .entity(entity) = self.input else {
+            fatalError()
+        }
+        return entity
+    }
+    
+    private let ciContext: CIContext
+    private let startTime: Double
+    
+    init(input: Input, renderTarget: RenderTarget, queue: DispatchQueue) {
+        guard case .entity = input else {
+            fatalError()
+        }
+        self.input = input
+        self.ciContext = CIContext(options: [.workingColorSpace : CGColorSpaceCreateDeviceRGB()])
+        self.startTime = CACurrentMediaTime()
+        super.init()
+        
+        self.textureRotation = .rotate0Degrees
+    }
+    
+    func output(time: Double) -> Output? {
+        return nil
+    }
+    
+    func asyncOutput(time: Double, completion: @escaping (Output?) -> Void) {
+        let deltaTime = max(0.0, time - self.startTime)
+        let timestamp = CMTime(seconds: deltaTime, preferredTimescale: CMTimeScale(NSEC_PER_SEC))
+        self.entity.image(for: timestamp, frameRate: 30, context: self.ciContext, completion: { image in
+            Queue.mainQueue().async {
+                completion(image.flatMap { .ciImage($0, timestamp) })
+            }
+        })
+    }
+    
+    func invalidate() {
+
+    }
+    
+    var needsDisplayLink: Bool {
+        if let entity = self.entity as? MediaEditorComposerStickerEntity, entity.isAnimated {
+            return true
+        }
+        return false
+    }
+    
+    var useAsyncOutput: Bool {
+        return true
+    }
+}