OpenGLES + SceneKit render context for planar frame buffers

NextLevelGLContext.swift

//
//  NextLevelGLContext.swift
//  NextLevel
//
//  Created by Patrick Piemonte on 9/26/16.
//
//  The MIT License (MIT)
//
//  Copyright (c) 2016-present patrick piemonte (http://patrickpiemonte.com/)
//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.

import Foundation
import AVFoundation
import SceneKit
import OpenGLES
import CoreVideo

// MARK: - NextLevelGLContext

internal class NextLevelGLContext: NSObject {
    
    // MARK: - constants
    
    let NextLevelGLContextAttributeVertex = "a_position"
    let NextLevelGLContextAttributeTextureCoord = "a_texture"
    let NextLevelGLContextUniformTextureSamplerY = "u_samplerY"
    let NextLevelGLContextUniformTextureSamplerUV = "u_samplerUV"

    let NextLevelGLContextYUVVertexShader = "\n" +
        "attribute vec4 a_position;\n" +
        "attribute vec2 a_texture;\n" +
        "\n" +
        "varying vec2 v_texture;\n" +
        "\n" +
        "void main()\n" +
        "{\n" +
        "    v_texture = a_texture;\n" +
        "    gl_Position = a_position;\n" +
        "}\n" +
        ""
    
    let NextLevelGLContextYUVFragmentShader = "\n" +
        "uniform sampler2D u_samplerY;\n" +
        "uniform sampler2D u_samplerUV;\n" +
        "\n" +
        "varying highp vec2 v_texture;\n" +
        "\n" +
        "void main()\n" +
        "{\n" +
        "    mediump vec3 yuv;\n" +
        "    lowp vec3 rgb;\n" +
        "\n" +
        "    yuv.x = texture2D(u_samplerY, v_texture).r;\n" +
        "    yuv.yz = texture2D(u_samplerUV, v_texture).rg - vec2(0.5, 0.5);\n" +
        "\n" +
        "    // BT.709, the standard for HDTV\n" +
        "    rgb = mat3(  1.164,   1.164,  1.164,\n" +
        "                     0,  -0.213,  2.112,\n" +
        "                 1.793,  -0.533,      0) * yuv;\n" +
        "\n" +
        "    gl_FragColor = vec4(rgb, 1);\n" +
        "}\n" +
        ""
    
    // MARK: - private ivars
    
    internal var _renderNode: SCNNode
    
    internal var _context: EAGLContext?
    internal var _program: SCNProgram?
    internal var _material: SCNMaterial?
    
    internal var _lumaTexture: CVOpenGLESTexture?
    internal var _chromaTexture: CVOpenGLESTexture?
    internal var _textureCache: CVOpenGLESTextureCache?
    
    internal var _bufferWidth: Int
    internal var _bufferHeight: Int
    internal var _bufferFormatType: OSType
    internal var _presentationFrame: CGRect
    
    internal var _pixelBuffer: CVPixelBuffer?
    
    // MARK: - object lifecycle
    
    convenience init(view: NextLevelView) {
        self.init()
        self._context = view.eaglContext
        self._presentationFrame = view.bounds
        self.setupContext()
    }
    
    override init() {
        self._renderNode = SCNNode()
        self._bufferWidth = 0
        self._bufferHeight = 0
        self._bufferFormatType = OSType(kCVPixelFormatType_32BGRA)
        self._presentationFrame = UIScreen.main.bounds
        super.init()
    }
    
    deinit {
        self.destroyContext()
        self._pixelBuffer = nil
    }
}

// MARK: - NextLevelViewRenderContext

extension NextLevelGLContext: NextLevelViewRenderContext {

    // MARK: - properties
    
    internal var renderNode: SCNNode {
        get {
            return self._renderNode
        }
    }

    // MARK: - resource lifecycle
    
    func setupContext() {
        
        // ensure we have a context
        if self._context == nil {
            return
        }
        
        // setup texture cache
        
        if let context = self._context {
            EAGLContext.setCurrent(context)
            
            let error = CVOpenGLESTextureCacheCreate(kCFAllocatorDefault,
                                                     nil,
                                                     context,
                                                     nil,
                                                     &self._textureCache)
            
            if error != kCVReturnSuccess {
                debugPrint("CVOpenGLESTextureCacheCreate error \(error)")
            }
            
        }
        
        self.clearTextureCache()

        // setup material
        if self._material == nil {
            self._material = SCNMaterial()
            self._material?.writesToDepthBuffer = false
            self._material?.readsFromDepthBuffer = true
        }
        
        // Note: i have no plans on adding BGRA shader support
        
        // setup program
        
        let shaderProgram = SCNProgram()
        shaderProgram.delegate = self
        shaderProgram.vertexShader = NextLevelGLContextYUVVertexShader
        shaderProgram.fragmentShader = NextLevelGLContextYUVFragmentShader
        
        // setup uniforms for program
        shaderProgram.setSemantic(SCNGeometrySource.Semantic.vertex.rawValue, forSymbol: NextLevelGLContextAttributeVertex, options: nil)
        shaderProgram.setSemantic(SCNGeometrySource.Semantic.texcoord.rawValue, forSymbol: NextLevelGLContextAttributeTextureCoord, options: nil)
        
        if let material = self._material {
            material.program = shaderProgram
            
            material.handleBinding(ofSymbol: NextLevelGLContextUniformTextureSamplerY, handler: { (programId: UInt32, location: UInt32, node: SCNNode?, renderer: SCNRenderer) in
                glUniform1i(GLint(location), 0);
            })
            
            material.handleBinding(ofSymbol: NextLevelGLContextUniformTextureSamplerUV, handler: { (programId: UInt32, location: UInt32, node: SCNNode?, renderer: SCNRenderer) in
                glUniform1i(GLint(location), 1)
            })
        }
    }
    
    func destroyContext() {
        EAGLContext.setCurrent(self._context)
        
        self._program = nil
        
        if EAGLContext.current() == self._context {
            EAGLContext.setCurrent(nil)
        }
        
    }
    
    func clearTextureCache() {
        if let textureCache = self._textureCache {
            CVOpenGLESTextureCacheFlush(textureCache, 0)
            self._lumaTexture = nil
            self._chromaTexture = nil
        }
    }
}

// MARK: - layout

extension NextLevelGLContext {

    func updateLayout() {
        let bufferSize = CGSize(width: self._bufferWidth, height: self._bufferHeight)
        let insetRect = AVMakeRect(aspectRatio: bufferSize, insideRect: self._presentationFrame)
        
        let widthScale = Float(self._presentationFrame.size.height / insetRect.size.height)
        let heightScale = Float(self._presentationFrame.size.width / insetRect.size.width)
        
        let vertices: [GLfloat] = [-widthScale, -heightScale,
                                   widthScale, -heightScale,
                                   -widthScale, heightScale,
                                   widthScale, heightScale]
        
        let texCoord: [GLfloat] = [0.0, 1.0,
                                   1.0, 1.0,
                                   0.0, 0.0,
                                   1.0, 0.0]
        
        let triangleStrip: [CInt] = [0, 1, 2, 3]
        
        let vertexData = Data(bytes: vertices, count: vertices.count * MemoryLayout<GLfloat>.stride)
        
        let vertexSource = SCNGeometrySource(data: vertexData,
                                             semantic: SCNGeometrySource.Semantic.vertex,
                                             vectorCount: 4,
                                             usesFloatComponents: true,
                                             componentsPerVector: 2,
                                             bytesPerComponent: MemoryLayout<GLfloat>.stride,
                                             dataOffset: 0,
                                             dataStride: MemoryLayout<GLfloat>.stride * 2)

        let texCoordData = Data(bytes: texCoord, count: texCoord.count * MemoryLayout<GLfloat>.stride)
        
        let texCoordSource = SCNGeometrySource(data: texCoordData,
                                               semantic: SCNGeometrySource.Semantic.texcoord,
                                               vectorCount: 4,
                                               usesFloatComponents: true,
                                               componentsPerVector: 2,
                                               bytesPerComponent: MemoryLayout<GLfloat>.stride,
                                               dataOffset: 0,
                                               dataStride: MemoryLayout<GLfloat>.stride * 2)
        
        let element = SCNGeometryElement(indices: triangleStrip, primitiveType: .triangleStrip)
        
        // update the vertices and texture coordinates in the render node
        if let material = self._material {
            self._renderNode.geometry = SCNGeometry(sources: [vertexSource, texCoordSource], elements: [element])
            self._renderNode.geometry?.materials = [material]
        }
    }
    
    // MARK: - render
    
    func willRender(time: TimeInterval) {
    }
    
    func render(withImageBuffer imageBuffer: CVImageBuffer, time: TimeInterval) {
        
        if self._context == nil {
            return
        }
        EAGLContext.setCurrent(self._context)
        
        self.clearTextureCache()
        
        if let cache = self._textureCache {

            let width = CVPixelBufferGetWidth(imageBuffer)
            let height = CVPixelBufferGetHeight(imageBuffer)
            
            if self._bufferWidth != width || self._bufferHeight != height {
                self._bufferWidth = width
                self._bufferHeight = height
                self.updateLayout()
            }
            
            self._bufferFormatType = CVPixelBufferGetPixelFormatType(imageBuffer)
            
            switch self._bufferFormatType {
            case kCVPixelFormatType_32BGRA:
                self.renderBGRA(withImageBuffer: imageBuffer, cache: cache)
                break
            case kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange:
                self.renderYpCbCr8(withImageBuffer: imageBuffer, cache: cache)
                break
            case kCVPixelFormatType_420YpCbCr8BiPlanarFullRange:
                self.renderYpCbCr8(withImageBuffer: imageBuffer, cache: cache)
                break
            default:
                debugPrint("unsupported buffer format")
                break
            }
        }
    }
    
    func didRender(scene: SCNScene, pointOfView: SCNNode, time: TimeInterval) {
    }
    
}

// MARK: - post

extension NextLevelGLContext {
    
    internal func renderedPixelBuffer() -> CVPixelBuffer? {
        return nil
    }
    
}

// MARK: - private

extension NextLevelGLContext {
    
    fileprivate func renderBGRA(withImageBuffer imageBuffer: CVImageBuffer, cache: CVOpenGLESTextureCache) {
        debugPrint("unsupported pixel buffer format for now")
    }
    
    fileprivate func renderYpCbCr8(withImageBuffer imageBuffer: CVImageBuffer, cache: CVOpenGLESTextureCache) {
    
        let width = CVPixelBufferGetWidth(imageBuffer)
        let height = CVPixelBufferGetHeight(imageBuffer)
        
        // always upload the textures since the input may be changing
        var error: CVReturn = 0
        
        // Y-plane
        glActiveTexture(GLenum(GL_TEXTURE0));
        error = CVOpenGLESTextureCacheCreateTextureFromImage(kCFAllocatorDefault,
                                                                cache,
                                                                imageBuffer,
                                                                nil,
                                                                GLenum(GL_TEXTURE_2D),
                                                                GL_RED_EXT, // GLES2, not sure if this also supports FullRange
                                                                GLsizei(width),
                                                                GLsizei(height),
                                                                GLenum(GL_RED_EXT),
                                                                GLenum(GL_UNSIGNED_BYTE),
                                                                0,
                                                                &self._lumaTexture);
        if error != kCVReturnSuccess {
            debugPrint("CVOpenGLESTextureCacheCreateTextureFromImage error \(error)")
        }
        
        if let luma = self._lumaTexture {
            glBindTexture(CVOpenGLESTextureGetTarget(luma), CVOpenGLESTextureGetName(luma));
            glTexParameterf(GLenum(GL_TEXTURE_2D), GLenum(GL_TEXTURE_WRAP_S), GLfloat(GL_CLAMP_TO_EDGE));
            glTexParameterf(GLenum(GL_TEXTURE_2D), GLenum(GL_TEXTURE_WRAP_T), GLfloat(GL_CLAMP_TO_EDGE));
        }
        
        // UV-plane
        glActiveTexture(GLenum(GL_TEXTURE1));
        error = CVOpenGLESTextureCacheCreateTextureFromImage(kCFAllocatorDefault,
                                                                cache,
                                                                imageBuffer,
                                                                nil,
                                                                GLenum(GL_TEXTURE_2D),
                                                                GL_RG_EXT,
                                                                GLsizei( Float(width) * 0.5 ),
                                                                GLsizei( Float(height) * 0.5 ),
                                                                GLenum(GL_RG_EXT),
                                                                GLenum(GL_UNSIGNED_BYTE),
                                                                1,
                                                                &self._chromaTexture);
                                                                
        if error != kCVReturnSuccess {
            debugPrint("CVOpenGLESTextureCacheCreateTextureFromImage error \(error)")
        }
        
        if let chroma = self._chromaTexture {
            glBindTexture(CVOpenGLESTextureGetTarget(chroma), CVOpenGLESTextureGetName(chroma))
            glTexParameterf(GLenum(GL_TEXTURE_2D), GLenum(GL_TEXTURE_WRAP_S), GLfloat(GL_CLAMP_TO_EDGE))
            glTexParameterf(GLenum(GL_TEXTURE_2D), GLenum(GL_TEXTURE_WRAP_T), GLfloat(GL_CLAMP_TO_EDGE))
        }
    }
}

// MARK: - SCNProgramDelegate

extension NextLevelGLContext: SCNProgramDelegate {
    
    internal func program(_ program: SCNProgram, handleError error: Error) {
        debugPrint("program \(program) error \(error)")
    }
    
}