silica-viewer/SilicaViewer/Document.swift

import Cocoa
import ZIPFoundation
import CoreFoundation
import Accelerate
import CoreMedia

func objectRefGetValue2(_ objectRef: CFTypeRef) -> UInt32 {
    let val = unsafeBitCast(objectRef, to: UInt64.self)
    
    return valueForKeyedArchiverUID(val)
}

/// NSDocument for reading Silica Document files
class Document: NSDocument {
    // TODO: don't keep the entire file in memory, especially since we also store the image data raw
    var data: Data?
    
    var dict: NSDictionary?
    
    var info = SilicaDocument()
    
    struct SilicaParsingError: Error, LocalizedError {
        enum Kind {
            case invalid
        }
        
        let kind: Kind
        let filename: URL?
        
        public var errorDescription: String? {
            switch self.kind {
            case .invalid:
                return filename!.lastPathComponent + " is an invalid Silica Document."
            }
        }
    }
    
    func throwError(_ error: SilicaParsingError.Kind) {
        DispatchQueue.main.sync {
            let _ = presentError(SilicaParsingError(kind: error, filename: fileURL))
        }
    }
    
    override func makeWindowControllers() {
        let storyboard = NSStoryboard(name: NSStoryboard.Name("Main"), bundle: nil)
        let windowController = storyboard.instantiateController(withIdentifier: NSStoryboard.SceneIdentifier("Document Window Controller")) as! NSWindowController
        self.addWindowController(windowController)
    }
    
    override class func canConcurrentlyReadDocuments(ofType: String) -> Bool {
        return ofType == "com.procreate"
    }
    
    /// Figures out the best filename for the document, depending on what information we know.
    /// It prefers the actual name defined in the document, unless it's missing - where we'll fall back to the filename.
    /// - Returns: The best filename.
    func getIdealFilename() -> String {
        if (!(info.name.isEmpty)) {
            return info.name
        } else {
            return fileURL!.deletingPathExtension().lastPathComponent
        }
    }
    
    /// Pass in an object from the `$object` array, which always contains a `$class` key.
    /// - Parameter dict: The NSDictionary to parse.
    /// - Returns: The class name.
    func getDocumentClassName(dict: NSDictionary) -> String? {
        let objectsArray = self.dict?["$objects"] as! NSArray
        
        if let value = dict["$class"] {
            let classObjectId = objectRefGetValue2(value as CFTypeRef)
            let classObject = objectsArray[Int(classObjectId)] as! NSDictionary
            
            return classObject["$classname"] as? String
        }
        
        return nil
    }
    
    /// Converts a `CFKeyedArchiveUID` from a `NSKeyedArchive` to a Int for indexing an array or `NSDictionary`.
    /// - Parameter id: The `CFKeyedArchiveUID` to parse.
    /// - Returns: An integer representation fo the `CFKeyedArchiveUID`.
    func getClassID(id: Any?) -> Int {
        return Int(objectRefGetValue2(id! as CFTypeRef))
    }
    
    /// Parses a chunk filename to integer coordinates.
    /// "1~1.chunk" -> (1, 2).
    /// - Parameter filename: The chunk filename.
    /// - Returns: A tuple containing the parsed coordinates.
    func parseChunkFilename(_ filename: String) -> (Int, Int)? {
        let pathURL = URL(fileURLWithPath: filename)
        let pathComponents = pathURL.lastPathComponent.replacingOccurrences(of: ".chunk", with: "").components(separatedBy: "~")
        
        if let x = Int(pathComponents[0]), let y = Int(pathComponents[1]) {
            return (x, y + 1)
        } else {
            return nil
        }
    }
    
    /// Parses a raw integer blend mode from the Silica Document, into a `BlendMode` enumeration.
    /// - Parameters:
    ///   - blendMode: The integer value for the `blendMode` property.
    ///   - extendedBlend: The integer value for the `extendedBlend` property.
    /// - Returns: If the provided integer representations are correct, a `BlendMode` enumeration.
    func parseRawBlendMode(blendMode: Int, extendedBlend: Int) -> BlendMode? {
        if blendMode == 0 && blendMode != extendedBlend {
            return BlendMode(rawValue: extendedBlend)
        } else {
            return BlendMode(rawValue: blendMode)
        }
    }
    
    func parseSilicaLayer(archive: Archive, dict: NSDictionary, isMask: Bool) -> SilicaLayer? {
        if getDocumentClassName(dict: dict) == LayerClassName {
            var layer = SilicaLayer()
            
            // TODO: when does a layer actually not have a name? i think i hit this case before
            if let val = dict["name"] {
                if let name = parseBasicObject(key: val) as? String {
                    layer.name = name
                }
            }
            
            guard let uuid = parseBasicObject(key: dict["UUID"]) as? String else {
                throwError(.invalid)
                return nil
            }
            
            // sometimes mask is missing
            if let maskDict = parseBasicObject(key: dict["mask"]) as? NSDictionary {
                layer.mask = parseSilicaLayer(archive: archive, dict: maskDict, isMask: true)?.data
            }
            
            guard let blendMode = (dict["blend"] as? NSNumber)?.intValue,
                  let extendedBlendMode = (dict["extendedBlend"] as? NSNumber)?.intValue,
                  let parsedBlendMode = parseRawBlendMode(blendMode: blendMode, extendedBlend: extendedBlendMode) else {
                throwError(.invalid)
                return nil
            }
            
            layer.data.blendMode = parsedBlendMode
            
            if let opacity = (dict["opacity"] as? NSNumber)?.doubleValue {
                layer.data.opacity = opacity
            }
            
            if let hidden = dict["hidden"] as? Bool {
                layer.data.hidden = hidden
            }
            
            if let clipped = dict["clipped"] as? Bool {
                layer.clipped = clipped
            }
            
            var chunkPaths: [String] = []
            
            archive.forEach { (entry: Entry) in
                if entry.path.contains(uuid) {
                    chunkPaths.append(entry.path)
                }
            }
            
            layer.data.chunks = Array(repeating: SilicaChunk(), count: chunkPaths.count)
            
            let dispatchGroup = DispatchGroup()
            let queue = DispatchQueue(label: "imageWork")
            
            DispatchQueue.concurrentPerform(iterations: chunkPaths.count) { (i: Int) in
                guard let threadArchive = Archive(data: self.data!, accessMode: .read) else {
                    return
                }
                
                let threadEntry = threadArchive[chunkPaths[i]]
                
                guard let (tileX, tileY) = parseChunkFilename(threadEntry!.path) else {
                    return
                }
                
                let (tileWidth, tileHeight) = info.getTileSize(tileX, tileY)
                
                let numChannels = isMask ? 1 : 4
                let byteSize = tileWidth * tileHeight * numChannels
                
                let uncompressedMemory = UnsafeMutablePointer<UInt8>.allocate(capacity: byteSize)
                
                guard let lzoData = readData(archive: threadArchive, entry: threadEntry!) else {
                    return
                }
                
                lzoData.withUnsafeBytes({ (bytes: UnsafeRawBufferPointer) -> Void in
                    var len = lzo_uint(byteSize)
                    
                    lzo1x_decompress_safe(bytes.baseAddress!.assumingMemoryBound(to: uint8.self), lzo_uint(lzoData.count), uncompressedMemory, &len, nil)
                })
                
                let imageData = Data(bytes: uncompressedMemory, count: byteSize)
                
                let render: CGColorRenderingIntent = .defaultIntent
                let rgbColorSpace = isMask ? CGColorSpaceCreateDeviceGray() : info.colorSpace
                
                let bitmapInfo = CGBitmapInfo(rawValue: (isMask ? CGImageAlphaInfo.none : CGImageAlphaInfo.premultipliedLast).rawValue).union(.byteOrder32Big)
                let providerRef: CGDataProvider? = CGDataProvider(data: imageData as CFData)
                
                guard let cgimage = CGImage(width: tileWidth,
                                            height: tileHeight,
                                            bitsPerComponent: 8,
                                            bitsPerPixel: 8 * numChannels,
                                            bytesPerRow: tileWidth * numChannels,
                                            space: rgbColorSpace,
                                            bitmapInfo: bitmapInfo,
                                            provider: providerRef!,
                                            decode: nil,
                                            shouldInterpolate: false,
                                            intent: render) else {
                    return
                }
                
                queue.async(group: dispatchGroup) {
                    layer.data.chunks[i].image = cgimage
                    layer.data.chunks[i].x = tileX
                    layer.data.chunks[i].y = tileY
                }
            }
            
            dispatchGroup.wait()
            
            return layer
        }
        
        return nil
    }
    
    /// Parses a pair string into a proper typed tuple.
    /// For example, {0, 0} -> (0, 0).
    /// - Parameter str: The pair string.
    /// - Returns: If it can be parsed correctly, a tuple containing the integer values found in the string.
    func parsePairString(_ str: String) -> (Int, Int)? {
        let sizeComponents = str.replacingOccurrences(of: "{", with: "").replacingOccurrences(of: "}", with: "").components(separatedBy: ", ")
        
        if sizeComponents.count == 2 {
            let width = Int(sizeComponents[0])
            let height = Int(sizeComponents[1])
            
            return (width!, height!)
        } else {
            return nil
        }
    }
    
    func parseBasicObject(key: Any?) -> Any? {
        if let objectsArray = self.dict?["$objects"] as? NSArray {
            let classID = getClassID(id: key)
            
            return objectsArray[classID]
        } else {
            return nil
        }
    }
    
    func parseSilicaDocument(archive: Archive, dict: NSDictionary) {
        if getDocumentClassName(dict: dict) == DocumentClassName {
            // failing to parse these isn't a failing case, thus we don't throw errors here.
            if let trackedTime = (dict[TrackedTimeKey] as? NSNumber)?.intValue {
                info.trackedTime = trackedTime
            }
            
            if let tileSize = (dict[TileSizeKey] as? NSNumber)?.intValue {
                info.tileSize = tileSize
            }
            
            if let orientation = (dict[OrientationKey] as? NSNumber)?.intValue {
                info.orientation = orientation
            }
            
            if let flippedHorizontally = (dict[FlippedHorizontallyKey] as? NSNumber)?.boolValue {
                info.flippedHorizontally = flippedHorizontally
            }
            
            if let flippedVertically = (dict[FlippedVerticallyKey] as? NSNumber)?.boolValue {
                info.flippedVertically = flippedVertically
            }
            
            if let videoResolution = parseBasicObject(key: dict["SilicaDocumentVideoSegmentInfoKey"]) as? NSDictionary {
                guard let frameSize = parseBasicObject(key: videoResolution["frameSize"]) as? String,
                      let videoFrame = parsePairString(frameSize) else {
                    throwError(.invalid)
                    return
                }
                
                info.videoFrame = videoFrame
            }
            
            if let colorProfileClassKey = dict["colorProfile"] {
                let objectsArray = self.dict?["$objects"] as! NSArray
                
                let colorProfileClassID = getClassID(id: colorProfileClassKey)
                let colorProfile = objectsArray[colorProfileClassID] as! NSDictionary
                
                let colorProfileNameClassKey = colorProfile["SiColorProfileArchiveICCNameKey"]
                let colorProfileNameClassID = getClassID(id: colorProfileNameClassKey)
                let colorProfileName = objectsArray[colorProfileNameClassID] as! NSString
                
                // we only support the basic "Display P3" color space... does Procreate actually store the ICC data??
                if colorProfileName == "Display P3" {
                    info.colorSpace = CGColorSpace(name: CGColorSpace.displayP3)!
                }
            } else {
                info.colorSpace = CGColorSpace(name: CGColorSpace.displayP3)!
            }
            
            if let background = parseBasicObject(key: dict["backgroundColor"]) as? NSData {
                var backgroundArray: [Float] = [0.0, 0.0, 0.0, 0.0]
                
                background.getBytes(&backgroundArray, length: 16)
                let backgroundCgArray: [CGFloat] = [CGFloat(backgroundArray[0]), CGFloat(backgroundArray[1]), CGFloat(backgroundArray[2]), CGFloat(backgroundArray[3])]
                
                info.backgroundColor = CGColor(colorSpace: info.colorSpace, components: backgroundCgArray)!
            }
            
            if let strokeCount = parseBasicObject(key: dict[StrokeCountKey]) as? NSNumber {
                info.strokeCount = Int(truncating: strokeCount)
            }
            
            if let name = parseBasicObject(key: dict[NameKey]) as? String,
               name != "$null" {
                info.name = name
            }
            
            if let authorName = parseBasicObject(key: dict[AuthorNameKey]) as? String,
               authorName != "$null" {
                info.authorName = authorName
            }
            
            guard let size = parseBasicObject(key: dict[SizeKey]) as? String,
                  let (parsedWidth, parsedHeight) = parsePairString(size) else {
                throwError(.invalid)
                return
            }
            
            info.width = parsedWidth
            info.height = parsedHeight
            
            info.columns = Int(ceil(Float(info.width) / Float(info.tileSize)))
            info.rows = Int(ceil(Float(info.height) / Float(info.tileSize))) + 1 // TODO: lol why
            
            if info.width % info.tileSize != 0 {
                info.remainderWidth = (info.columns * info.tileSize) - info.width
            }
            
            if info.height % info.tileSize != 0 {
                info.remainderHeight = (info.rows * info.tileSize) - info.height
            }
            
            guard let layersDict = parseBasicObject(key: dict[LayersKey]) as? NSDictionary,
                  let layersArray = layersDict["NS.objects"] as? NSArray else {
                throwError(.invalid)
                return
            }
            
            for object in layersArray {
                guard let layerDict = parseBasicObject(key: object) as? NSDictionary,
                      let layer = parseSilicaLayer(archive: archive, dict: layerDict, isMask: false) else {
                    throwError(.invalid)
                    return
                }
                
                info.layers.append(layer)
            }
        }
    }
    
    func parseDocument(archive: Archive, dict: NSDictionary) {
        // double check if this archive is really correct
        if let value = dict["$version"] {
            if (value as! Int) != NSKeyedArchiveVersion {
                throwError(.invalid)
                return
            }
            
            self.dict = dict
            
            guard let objectsArray = dict["$objects"] as? NSArray else {
                throwError(.invalid)
                return
            }
            
            // let's begin by reading $top, which is always going to be SilicaDocument type.
            
            guard let topObject = dict["$top"] as? NSDictionary else {
                throwError(.invalid)
                return
            }
            
            let topClassID = objectRefGetValue2(topObject["root"] as CFTypeRef)
            
            guard let topObjectClass = objectsArray[Int(topClassID)] as? NSDictionary else {
                throwError(.invalid)
                return
            }
            
            parseSilicaDocument(archive: archive, dict: topObjectClass)
        }
    }
    
    override func read(from data: Data, ofType typeName: String) throws {
        self.data = data
        
        guard let archive = Archive(data: data, accessMode: Archive.AccessMode.read) else {
            throwError(.invalid)
            return
        }
        
        guard let documentEntry = archive[DocumentArchivePath] else {
            throwError(.invalid)
            return
        }
        
        guard let documentData = readData(archive: archive, entry: documentEntry) else {
            throwError(.invalid)
            return
        }
        
        var plistFormat = PropertyListSerialization.PropertyListFormat.binary
        guard let propertyList = try? PropertyListSerialization.propertyList(from: documentData, options: [], format: &plistFormat) else {
            throwError(.invalid)
            return
        }
        
        parseDocument(archive: archive, dict: propertyList as! NSDictionary)
    }
    
    /// Renders a full, composite image of the Silica Document - emulating the Procreate drawing engine.
    /// - Returns: If the image suceeds in drawing, then a `NSImage` of the canvas.
    func makeComposite() -> NSImage? {
        let engine = SilicaEngine()
        
        var image = NSImage(cgImage: engine.draw(document: info)!, size: info.nsSize())
        
        if info.orientation == 3 {
            image = image.imageRotatedByDegreess(degrees: 90)
        } else if info.orientation == 4 {
            image = image.imageRotatedByDegreess(degrees: -90)
        }
        
        if info.flippedHorizontally && (info.orientation == 1 || info.orientation == 2) {
            image = image.flipHorizontally()
        } else if info.flippedHorizontally && (info.orientation == 3 || info.orientation == 4) {
            image = image.flipVertically()
        } else if info.flippedVertically && (info.orientation == 1 || info.orientation == 2) {
            image = image.flipVertically()
        } else if !info.flippedVertically && (info.orientation == 3 || info.orientation == 4) {
            image = image.flipHorizontally()
        }
        
        return image
    }
    
    /// Extracts the thumbnail image generated by Procreate itself.
    /// - Returns: If it's able to find the thumbnail, an `NSImage`.
    func makeThumbnail() -> NSImage? {
        guard let archive = Archive(data: data!, accessMode: Archive.AccessMode.read) else  {
            return nil
        }
        
        guard let entry = archive[ThumbnailPath] else {
            return nil
        }
        
        guard let thumbnailData = readData(archive: archive, entry: entry) else {
            return nil
        }
        
        return NSImage(data: thumbnailData)
    }
}