mirror of
https://github.com/Swiftgram/Telegram-iOS.git
synced 2025-06-16 05:55:20 +00:00
557 lines
23 KiB
Swift
557 lines
23 KiB
Swift
import Foundation
|
|
import UIKit
|
|
import Metal
|
|
import MetalKit
|
|
import simd
|
|
import DctHuffman
|
|
|
|
private struct Vertex {
|
|
var position: vector_float2
|
|
var textureCoordinate: vector_float2
|
|
}
|
|
|
|
public final class CompressedImageRenderer {
|
|
private final class Shared {
|
|
static let shared: Shared = {
|
|
return Shared(sharedContext: AnimationCompressor.SharedContext.shared)!
|
|
}()
|
|
|
|
let sharedContext: AnimationCompressor.SharedContext
|
|
|
|
let computeIdctPipelineState: MTLComputePipelineState
|
|
let renderIdctPipelineState: MTLRenderPipelineState
|
|
let renderRgbPipelineState: MTLRenderPipelineState
|
|
let renderYuvaPipelineState: MTLRenderPipelineState
|
|
|
|
init?(sharedContext: AnimationCompressor.SharedContext) {
|
|
self.sharedContext = sharedContext
|
|
|
|
guard let idctFunction = self.sharedContext.defaultLibrary.makeFunction(name: "idctKernel") else {
|
|
return nil
|
|
}
|
|
|
|
guard let computeIdctPipelineState = try? self.sharedContext.device.makeComputePipelineState(function: idctFunction) else {
|
|
return nil
|
|
}
|
|
self.computeIdctPipelineState = computeIdctPipelineState
|
|
|
|
guard let vertexShader = self.sharedContext.defaultLibrary.makeFunction(name: "vertexShader") else {
|
|
return nil
|
|
}
|
|
guard let samplingIdctShader = self.sharedContext.defaultLibrary.makeFunction(name: "samplingIdctShader") else {
|
|
return nil
|
|
}
|
|
guard let samplingRgbShader = self.sharedContext.defaultLibrary.makeFunction(name: "samplingRgbShader") else {
|
|
return nil
|
|
}
|
|
guard let samplingYuvaShader = self.sharedContext.defaultLibrary.makeFunction(name: "samplingYuvaShader") else {
|
|
return nil
|
|
}
|
|
|
|
let idctPipelineStateDescriptor = MTLRenderPipelineDescriptor()
|
|
idctPipelineStateDescriptor.label = "Render IDCT Pipeline"
|
|
idctPipelineStateDescriptor.vertexFunction = vertexShader
|
|
idctPipelineStateDescriptor.fragmentFunction = samplingIdctShader
|
|
idctPipelineStateDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
|
|
|
|
guard let renderIdctPipelineState = try? self.sharedContext.device.makeRenderPipelineState(descriptor: idctPipelineStateDescriptor) else {
|
|
return nil
|
|
}
|
|
self.renderIdctPipelineState = renderIdctPipelineState
|
|
|
|
let rgbPipelineStateDescriptor = MTLRenderPipelineDescriptor()
|
|
rgbPipelineStateDescriptor.label = "Render RGB Pipeline"
|
|
rgbPipelineStateDescriptor.vertexFunction = vertexShader
|
|
rgbPipelineStateDescriptor.fragmentFunction = samplingRgbShader
|
|
rgbPipelineStateDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
|
|
|
|
guard let renderRgbPipelineState = try? self.sharedContext.device.makeRenderPipelineState(descriptor: rgbPipelineStateDescriptor) else {
|
|
return nil
|
|
}
|
|
self.renderRgbPipelineState = renderRgbPipelineState
|
|
|
|
let yuvaPipelineStateDescriptor = MTLRenderPipelineDescriptor()
|
|
yuvaPipelineStateDescriptor.label = "Render YUVA Pipeline"
|
|
yuvaPipelineStateDescriptor.vertexFunction = vertexShader
|
|
yuvaPipelineStateDescriptor.fragmentFunction = samplingYuvaShader
|
|
yuvaPipelineStateDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
|
|
|
|
guard let renderYuvaPipelineState = try? self.sharedContext.device.makeRenderPipelineState(descriptor: yuvaPipelineStateDescriptor) else {
|
|
return nil
|
|
}
|
|
self.renderYuvaPipelineState = renderYuvaPipelineState
|
|
}
|
|
}
|
|
|
|
private let sharedContext: AnimationCompressor.SharedContext
|
|
private let shared: Shared
|
|
|
|
private var compressedTextures: TextureSet?
|
|
private var outputTextures: TextureSet?
|
|
|
|
private var rgbTexture: Texture?
|
|
|
|
private var yuvaTextures: TextureSet?
|
|
|
|
private let commandQueue: MTLCommandQueue
|
|
|
|
public init?(sharedContext: AnimationCompressor.SharedContext) {
|
|
self.sharedContext = sharedContext
|
|
self.shared = Shared.shared
|
|
|
|
guard let commandQueue = self.sharedContext.device.makeCommandQueue() else {
|
|
return nil
|
|
}
|
|
self.commandQueue = commandQueue
|
|
}
|
|
|
|
private var drawableRequestTimestamp: Double?
|
|
|
|
private func getNextDrawable(metalLayer: CALayer, drawableSize: CGSize) -> CAMetalDrawable? {
|
|
#if targetEnvironment(simulator)
|
|
if #available(iOS 13.0, *) {
|
|
if let metalLayer = metalLayer as? CAMetalLayer {
|
|
if metalLayer.drawableSize != drawableSize {
|
|
metalLayer.drawableSize = drawableSize
|
|
}
|
|
return metalLayer.nextDrawable()
|
|
} else {
|
|
return nil
|
|
}
|
|
} else {
|
|
return nil
|
|
}
|
|
#else
|
|
if let metalLayer = metalLayer as? CAMetalLayer {
|
|
if metalLayer.drawableSize != drawableSize {
|
|
metalLayer.drawableSize = drawableSize
|
|
}
|
|
let beginTime = CFAbsoluteTimeGetCurrent()
|
|
let drawableRequestDuration: Double
|
|
if let drawableRequestTimestamp = self.drawableRequestTimestamp {
|
|
drawableRequestDuration = beginTime - drawableRequestTimestamp
|
|
if drawableRequestDuration < 1.0 / 60.0 {
|
|
return nil
|
|
}
|
|
} else {
|
|
drawableRequestDuration = 0.0
|
|
}
|
|
self.drawableRequestTimestamp = beginTime
|
|
let result = metalLayer.nextDrawable()
|
|
let duration = CFAbsoluteTimeGetCurrent() - beginTime
|
|
if duration > 1.0 / 200.0 {
|
|
print("lag \(duration * 1000.0) ms (\(drawableRequestDuration * 1000.0) ms)")
|
|
}
|
|
return result
|
|
} else {
|
|
return nil
|
|
}
|
|
#endif
|
|
}
|
|
|
|
private func updateIdctTextures(compressedImage: AnimationCompressor.CompressedImageData) {
|
|
self.rgbTexture = nil
|
|
self.yuvaTextures = nil
|
|
|
|
let readBuffer = ReadBuffer(data: compressedImage.data)
|
|
if readBuffer.readInt32() != 0x543ee445 {
|
|
return
|
|
}
|
|
if readBuffer.readInt32() != 4 {
|
|
return
|
|
}
|
|
|
|
let width = Int(readBuffer.readInt32())
|
|
let height = Int(readBuffer.readInt32())
|
|
|
|
let compressedTextures: TextureSet
|
|
if let current = self.compressedTextures, current.width == width, current.height == height {
|
|
compressedTextures = current
|
|
} else {
|
|
guard let textures = TextureSet(
|
|
device: self.sharedContext.device,
|
|
width: width,
|
|
height: height,
|
|
descriptions: [
|
|
TextureSet.Description(
|
|
fractionWidth: 1, fractionHeight: 1,
|
|
pixelFormat: .r32Float
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 2, fractionHeight: 2,
|
|
pixelFormat: .r32Float
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 2, fractionHeight: 2,
|
|
pixelFormat: .r32Float
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 1, fractionHeight: 1,
|
|
pixelFormat: .r32Float
|
|
)
|
|
],
|
|
usage: .shaderRead,
|
|
isShared: true
|
|
) else {
|
|
return
|
|
}
|
|
self.compressedTextures = textures
|
|
compressedTextures = textures
|
|
}
|
|
|
|
for i in 0 ..< 4 {
|
|
let planeWidth = Int(readBuffer.readInt32())
|
|
let planeHeight = Int(readBuffer.readInt32())
|
|
let bytesPerRow = Int(readBuffer.readInt32())
|
|
|
|
let planeSize = Int(readBuffer.readInt32())
|
|
let planeData = readBuffer.readDataNoCopy(length: planeSize)
|
|
|
|
var tempData: Data?
|
|
compressedTextures.textures[i].readDirect(width: planeWidth, height: planeHeight, bytesPerRow: bytesPerRow, read: { destinationBytes in
|
|
if let destinationBytes = destinationBytes {
|
|
readDCTBlocks(Int32(planeWidth), Int32(planeHeight), planeData, destinationBytes.assumingMemoryBound(to: Float32.self), Int32(bytesPerRow / 4))
|
|
return UnsafeRawPointer(destinationBytes)
|
|
} else {
|
|
tempData = Data(count: bytesPerRow * planeHeight)
|
|
return tempData!.withUnsafeMutableBytes { bytes -> UnsafeRawPointer in
|
|
readDCTBlocks(Int32(planeWidth), Int32(planeHeight), planeData, bytes.baseAddress!.assumingMemoryBound(to: Float32.self), Int32(bytesPerRow / 4))
|
|
return UnsafeRawPointer(bytes.baseAddress!)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
public func renderIdct(metalLayer: CALayer, compressedImage: AnimationCompressor.CompressedImageData, completion: @escaping () -> Void) {
|
|
DispatchQueue.global().async {
|
|
self.updateIdctTextures(compressedImage: compressedImage)
|
|
|
|
DispatchQueue.main.async {
|
|
guard let compressedTextures = self.compressedTextures else {
|
|
return
|
|
}
|
|
|
|
guard let commandBuffer = self.commandQueue.makeCommandBuffer() else {
|
|
return
|
|
}
|
|
commandBuffer.label = "MyCommand"
|
|
|
|
guard let computeEncoder = commandBuffer.makeComputeCommandEncoder() else {
|
|
return
|
|
}
|
|
|
|
computeEncoder.setComputePipelineState(self.shared.computeIdctPipelineState)
|
|
|
|
let outputTextures: TextureSet
|
|
if let current = self.outputTextures, current.width == compressedTextures.textures[0].width, current.height == compressedTextures.textures[0].height {
|
|
outputTextures = current
|
|
} else {
|
|
guard let textures = TextureSet(
|
|
device: self.sharedContext.device,
|
|
width: compressedTextures.textures[0].width,
|
|
height: compressedTextures.textures[0].height,
|
|
descriptions: [
|
|
TextureSet.Description(
|
|
fractionWidth: 1, fractionHeight: 1,
|
|
pixelFormat: .r8Unorm
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 2, fractionHeight: 2,
|
|
pixelFormat: .r8Unorm
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 2, fractionHeight: 2,
|
|
pixelFormat: .r8Unorm
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 1, fractionHeight: 1,
|
|
pixelFormat: .r8Unorm
|
|
)
|
|
],
|
|
usage: [.shaderRead, .shaderWrite],
|
|
isShared: false
|
|
) else {
|
|
return
|
|
}
|
|
self.outputTextures = textures
|
|
outputTextures = textures
|
|
}
|
|
|
|
for i in 0 ..< 4 {
|
|
computeEncoder.setTexture(compressedTextures.textures[i].texture, index: 0)
|
|
computeEncoder.setTexture(outputTextures.textures[i].texture, index: 1)
|
|
|
|
var colorPlaneInt32 = Int32(i)
|
|
computeEncoder.setBytes(&colorPlaneInt32, length: 4, index: 2)
|
|
|
|
let threadgroupSize = MTLSize(width: 8, height: 8, depth: 1)
|
|
let threadgroupCount = MTLSize(width: (compressedTextures.textures[i].width + threadgroupSize.width - 1) / threadgroupSize.width, height: (compressedTextures.textures[i].height + threadgroupSize.height - 1) / threadgroupSize.height, depth: 1)
|
|
|
|
computeEncoder.dispatchThreadgroups(threadgroupCount, threadsPerThreadgroup: threadgroupSize)
|
|
}
|
|
|
|
computeEncoder.endEncoding()
|
|
|
|
let drawableSize = CGSize(width: CGFloat(outputTextures.textures[0].width), height: CGFloat(outputTextures.textures[0].height))
|
|
|
|
guard let drawable = self.getNextDrawable(metalLayer: metalLayer, drawableSize: drawableSize) else {
|
|
commandBuffer.commit()
|
|
completion()
|
|
return
|
|
}
|
|
|
|
let renderPassDescriptor = MTLRenderPassDescriptor()
|
|
renderPassDescriptor.colorAttachments[0].texture = drawable.texture
|
|
renderPassDescriptor.colorAttachments[0].loadAction = .clear
|
|
renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0)
|
|
|
|
guard let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else {
|
|
return
|
|
}
|
|
renderEncoder.label = "MyRenderEncoder"
|
|
|
|
renderEncoder.setRenderPipelineState(self.shared.renderIdctPipelineState)
|
|
|
|
for i in 0 ..< 4 {
|
|
renderEncoder.setFragmentTexture(outputTextures.textures[i].texture, index: i)
|
|
}
|
|
|
|
renderEncoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 6)
|
|
|
|
renderEncoder.endEncoding()
|
|
|
|
var storedDrawable: MTLDrawable? = drawable
|
|
commandBuffer.addScheduledHandler { _ in
|
|
storedDrawable?.present()
|
|
storedDrawable = nil
|
|
}
|
|
|
|
|
|
#if targetEnvironment(simulator)
|
|
commandBuffer.addCompletedHandler { _ in
|
|
DispatchQueue.main.async {
|
|
completion()
|
|
}
|
|
}
|
|
#else
|
|
if #available(iOS 10.3, *) {
|
|
drawable.addPresentedHandler { _ in
|
|
DispatchQueue.main.async {
|
|
completion()
|
|
}
|
|
}
|
|
} else {
|
|
commandBuffer.addCompletedHandler { _ in
|
|
DispatchQueue.main.async {
|
|
completion()
|
|
}
|
|
}
|
|
}
|
|
#endif
|
|
|
|
commandBuffer.commit()
|
|
}
|
|
}
|
|
}
|
|
|
|
private func updateRgbTexture(width: Int, height: Int, bytesPerRow: Int, data: Data) {
|
|
self.compressedTextures = nil
|
|
self.outputTextures = nil
|
|
self.yuvaTextures = nil
|
|
|
|
let rgbTexture: Texture
|
|
if let current = self.rgbTexture, current.width == width, current.height == height {
|
|
rgbTexture = current
|
|
} else {
|
|
guard let texture = Texture(device: self.sharedContext.device, width: width, height: height, pixelFormat: .bgra8Unorm, usage: .shaderRead, isShared: true) else {
|
|
return
|
|
}
|
|
self.rgbTexture = texture
|
|
rgbTexture = texture
|
|
}
|
|
|
|
rgbTexture.readDirect(width: width, height: height, bytesPerRow: bytesPerRow, read: { destinationBytes in
|
|
return data.withUnsafeBytes { bytes -> UnsafeRawPointer in
|
|
if let destinationBytes = destinationBytes {
|
|
memcpy(destinationBytes, bytes.baseAddress!, bytes.count)
|
|
return UnsafeRawPointer(destinationBytes)
|
|
} else {
|
|
return bytes.baseAddress!
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
public func renderRgb(metalLayer: CALayer, width: Int, height: Int, bytesPerRow: Int, data: Data, completion: @escaping () -> Void) {
|
|
self.updateRgbTexture(width: width, height: height, bytesPerRow: bytesPerRow, data: data)
|
|
|
|
guard let rgbTexture = self.rgbTexture else {
|
|
return
|
|
}
|
|
|
|
guard let commandBuffer = self.commandQueue.makeCommandBuffer() else {
|
|
return
|
|
}
|
|
commandBuffer.label = "MyCommand"
|
|
|
|
let drawableSize = CGSize(width: CGFloat(rgbTexture.width), height: CGFloat(rgbTexture.height))
|
|
|
|
guard let drawable = self.getNextDrawable(metalLayer: metalLayer, drawableSize: drawableSize) else {
|
|
commandBuffer.commit()
|
|
completion()
|
|
return
|
|
}
|
|
|
|
let renderPassDescriptor = MTLRenderPassDescriptor()
|
|
renderPassDescriptor.colorAttachments[0].texture = drawable.texture
|
|
renderPassDescriptor.colorAttachments[0].loadAction = .clear
|
|
renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0)
|
|
|
|
guard let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else {
|
|
return
|
|
}
|
|
renderEncoder.label = "MyRenderEncoder"
|
|
|
|
renderEncoder.setRenderPipelineState(self.shared.renderRgbPipelineState)
|
|
renderEncoder.setFragmentTexture(rgbTexture.texture, index: 0)
|
|
|
|
renderEncoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 6)
|
|
|
|
renderEncoder.endEncoding()
|
|
|
|
commandBuffer.present(drawable)
|
|
|
|
commandBuffer.addCompletedHandler { _ in
|
|
DispatchQueue.main.async {
|
|
completion()
|
|
}
|
|
}
|
|
|
|
commandBuffer.commit()
|
|
}
|
|
|
|
private func updateYuvaTextures(width: Int, height: Int, data: Data) {
|
|
self.compressedTextures = nil
|
|
self.outputTextures = nil
|
|
self.rgbTexture = nil
|
|
|
|
let yuvaTextures: TextureSet
|
|
if let current = self.yuvaTextures, current.width == width, current.height == height {
|
|
yuvaTextures = current
|
|
} else {
|
|
guard let textures = TextureSet(
|
|
device: self.sharedContext.device,
|
|
width: width,
|
|
height: height,
|
|
descriptions: [
|
|
TextureSet.Description(
|
|
fractionWidth: 1, fractionHeight: 1,
|
|
pixelFormat: .r8Unorm
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 2, fractionHeight: 2,
|
|
pixelFormat: .rg8Unorm
|
|
),
|
|
TextureSet.Description(
|
|
fractionWidth: 2, fractionHeight: 1,
|
|
pixelFormat: .r8Uint
|
|
)
|
|
],
|
|
usage: .shaderRead,
|
|
isShared: true
|
|
) else {
|
|
return
|
|
}
|
|
self.yuvaTextures = textures
|
|
yuvaTextures = textures
|
|
}
|
|
|
|
data.withUnsafeBytes { yuvaBuffer in
|
|
guard let yuva = yuvaBuffer.baseAddress?.assumingMemoryBound(to: UInt8.self) else {
|
|
return
|
|
}
|
|
|
|
yuvaTextures.textures[0].readDirect(width: width, height: height, bytesPerRow: width, read: { destinationBytes in
|
|
if let destinationBytes = destinationBytes {
|
|
memcpy(destinationBytes, yuva.advanced(by: 0), width * height)
|
|
return UnsafeRawPointer(destinationBytes)
|
|
} else {
|
|
return UnsafeRawPointer(yuva.advanced(by: 0))
|
|
}
|
|
})
|
|
|
|
yuvaTextures.textures[1].readDirect(width: width / 2, height: height / 2, bytesPerRow: width, read: { destinationBytes in
|
|
if let destinationBytes = destinationBytes {
|
|
memcpy(destinationBytes, yuva.advanced(by: width * height), width * height / 2)
|
|
return UnsafeRawPointer(destinationBytes)
|
|
} else {
|
|
return UnsafeRawPointer(yuva.advanced(by: width * height))
|
|
}
|
|
})
|
|
|
|
yuvaTextures.textures[2].readDirect(width: width / 2, height: height, bytesPerRow: width / 2, read: { destinationBytes in
|
|
if let destinationBytes = destinationBytes {
|
|
memcpy(destinationBytes, yuva.advanced(by: width * height * 2), width / 2 * height)
|
|
return UnsafeRawPointer(destinationBytes)
|
|
} else {
|
|
return UnsafeRawPointer(yuva.advanced(by: width * height * 2))
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
public func renderYuva(metalLayer: CALayer, width: Int, height: Int, data: Data, completion: @escaping () -> Void) {
|
|
self.updateYuvaTextures(width: width, height: height, data: data)
|
|
|
|
guard let yuvaTextures = self.yuvaTextures else {
|
|
return
|
|
}
|
|
|
|
guard let commandBuffer = self.commandQueue.makeCommandBuffer() else {
|
|
return
|
|
}
|
|
commandBuffer.label = "MyCommand"
|
|
|
|
let drawableSize = CGSize(width: CGFloat(yuvaTextures.width), height: CGFloat(yuvaTextures.height))
|
|
|
|
guard let drawable = self.getNextDrawable(metalLayer: metalLayer, drawableSize: drawableSize) else {
|
|
commandBuffer.commit()
|
|
completion()
|
|
return
|
|
}
|
|
|
|
let renderPassDescriptor = MTLRenderPassDescriptor()
|
|
renderPassDescriptor.colorAttachments[0].texture = drawable.texture
|
|
renderPassDescriptor.colorAttachments[0].loadAction = .clear
|
|
renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0)
|
|
|
|
guard let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else {
|
|
return
|
|
}
|
|
renderEncoder.label = "MyRenderEncoder"
|
|
|
|
renderEncoder.setRenderPipelineState(self.shared.renderYuvaPipelineState)
|
|
renderEncoder.setFragmentTexture(yuvaTextures.textures[0].texture, index: 0)
|
|
renderEncoder.setFragmentTexture(yuvaTextures.textures[1].texture, index: 1)
|
|
renderEncoder.setFragmentTexture(yuvaTextures.textures[2].texture, index: 2)
|
|
|
|
var alphaSize = simd_uint2(UInt32(yuvaTextures.textures[0].texture.width), UInt32(yuvaTextures.textures[0].texture.height))
|
|
renderEncoder.setFragmentBytes(&alphaSize, length: 8, index: 3)
|
|
|
|
renderEncoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 6)
|
|
|
|
renderEncoder.endEncoding()
|
|
|
|
commandBuffer.present(drawable)
|
|
|
|
commandBuffer.addCompletedHandler { _ in
|
|
DispatchQueue.main.async {
|
|
completion()
|
|
}
|
|
}
|
|
|
|
commandBuffer.commit()
|
|
}
|
|
}
|