CHGenerator/Sources/Generator/Files/ImageGenerator.swift
2022-11-30 15:29:51 +01:00

457 lines
16 KiB
Swift

import Foundation
import AppKit
import CryptoKit
import Darwin.C
private struct ImageJob {
let destination: String
let width: Int
let path: String
let quality: Float
let alwaysGenerate: Bool
}
final class ImageGenerator {
private let imageOptimSupportedFileExtensions: Set<String> = ["jpg", "png", "svg"]
private let imageOptimizationBatchSize = 50
/**
The path to the input folder.
*/
private let input: URL
/**
The path to the output folder
*/
private let output: URL
/**
The images to generate.
The key is the image source path relative to the input folder, and the values are the destination path (relative to the output folder) and the required image width.
*/
private var imageJobs: [String : [ImageJob]] = [:]
/**
The images for which to generate multiple versions
The key is the source file, the value is the path of the requiring page.
*/
private var multiImageJobs: [String : String] = [:]
/**
The images which could not be found, but are required for the site.
The key is the image path, and the value is the page that requires it.
*/
private var missingImages: [String : String] = [:]
/**
All warnings produced for images during generation
*/
private var imageWarnings: Set<String> = []
/**
All images required by the site.
The values are the destination paths of the images, relative to the output folder
*/
private var requiredImages: Set<String> = []
/**
All images modified or created during this generator run.
*/
private var generatedImages: Set<String> = []
/**
The images optimized by ImageOptim
*/
private var optimizedImages: Set<String> = []
/**
A cache to get the size of source images, so that files don't have to be loaded multiple times.
The key is the absolute source path, and the value is the image size
*/
private var imageSizeCache: [String : NSSize] = [:]
private var fileUpdates: FileUpdateChecker
init(input: URL, output: URL) {
self.fileUpdates = FileUpdateChecker(input: input)
self.input = input
self.output = output
}
func writeDetectedFileChangesToDisk() {
fileUpdates.writeDetectedFileChangesToDisk()
}
private func getImageSize(atPath path: String) -> NSSize? {
if let size = imageSizeCache[path] {
return size
}
guard let image = getImage(atPath: path) else {
return nil
}
let size = image.size
imageSizeCache[path] = size
return size
}
private func getImage(atPath path: String) -> NSImage? {
guard let data = getData(atPath: path) else {
log.add(error: "Failed to load file", source: path)
return nil
}
guard let image = NSImage(data: data) else {
log.add(error: "Failed to read image", source: path)
return nil
}
return image
}
private func getData(atPath path: String) -> Data? {
let url = input.appendingPathComponent(path)
guard url.exists else {
return nil
}
do {
let data = try Data(contentsOf: url)
fileUpdates.didLoad(data, at: path)
return data
} catch {
log.add(error: "Failed to read data", source: path, error: error)
return nil
}
}
func requireImage(at destination: String, generatedFrom source: String, requiredBy path: String, quality: Float, width: Int, height: Int?, alwaysGenerate: Bool) -> NSSize {
requiredImages.insert(destination)
let height = height.unwrapped(CGFloat.init)
let sourceUrl = input.appendingPathComponent(source)
guard sourceUrl.exists else {
missingImages[source] = path
return .zero
}
guard let imageSize = getImageSize(atPath: source) else {
missingImages[source] = path
return .zero
}
let scaledSize = imageSize.scaledDown(to: CGFloat(width))
// Check desired height, then we can forget about it
if let height = height {
let expectedHeight = scaledSize.width / CGFloat(width) * height
if abs(expectedHeight - scaledSize.height) > 2 {
addWarning("Invalid height (\(scaledSize.height) instead of \(expectedHeight))", destination: destination, path: path)
}
}
let job = ImageJob(
destination: destination,
width: width,
path: path,
quality: quality,
alwaysGenerate: alwaysGenerate)
insert(job: job, source: source)
return scaledSize
}
private func insert(job: ImageJob, source: String) {
guard let existingSource = imageJobs[source] else {
imageJobs[source] = [job]
return
}
guard let existingJob = existingSource.first(where: { $0.destination == job.destination }) else {
imageJobs[source] = existingSource + [job]
return
}
if existingJob.width != job.width {
addWarning("Multiple image widths (\(existingJob.width) and \(job.width))", destination: job.destination, path: "\(existingJob.path) and \(job.path)")
}
}
func createImages() {
var count = 0
for (source, jobs) in imageJobs.sorted(by: { $0.key < $1.key }) {
print(String(format: "Creating images: %4d / %d\r", count, imageJobs.count), terminator: "")
fflush(stdout)
create(images: jobs, from: source)
count += 1
}
print(" \r", terminator: "")
createMultiImages()
optimizeImages()
printMissingImages()
printImageWarnings()
printGeneratedImages()
printTotalImageCount()
}
private func printMissingImages() {
guard !missingImages.isEmpty else {
return
}
print("\(missingImages.count) missing images:")
let sort = missingImages.sorted { (a, b) in
a.value < b.value && a.key < b.key
}
for (source, path) in sort {
print(" \(source) (required by \(path))")
}
}
private func printImageWarnings() {
guard !imageWarnings.isEmpty else {
return
}
print("\(imageWarnings.count) image warnings:")
for imageWarning in imageWarnings {
print(imageWarning)
}
}
private func printGeneratedImages() {
guard !generatedImages.isEmpty else {
return
}
print("\(generatedImages.count) images generated:")
for image in generatedImages {
print(" " + image)
}
}
private func printTotalImageCount() {
print("\(requiredImages.count) images")
}
private func addWarning(_ message: String, destination: String, path: String) {
let warning = " \(destination): \(message) required by \(path)"
imageWarnings.insert(warning)
}
private func addWarning(_ message: String, job: ImageJob) {
addWarning(message, destination: job.destination, path: job.path)
}
private func isMissing(_ job: ImageJob) -> Bool {
job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
}
private func create(images: [ImageJob], from source: String) {
// Only load image if required
let imageHasChanged = fileUpdates.fileHasChanged(at: source)
guard imageHasChanged || images.contains(where: isMissing) else {
return
}
guard let image = getImage(atPath: source) else {
missingImages[source] = images.first?.path
return
}
let jobs = imageHasChanged ? images : images.filter(isMissing)
// Update all images
jobs.forEach { job in
// Prevent memory overflow due to repeated NSImage operations
autoreleasepool {
create(job: job, from: image, source: source)
}
}
}
private func create(job: ImageJob, from image: NSImage, source: String) {
let destinationUrl = output.appendingPathComponent(job.destination)
create(job: job, from: image, source: source, at: destinationUrl)
}
private func create(job: ImageJob, from image: NSImage, source: String, at destinationUrl: URL) {
// Ensure that image file is supported
let ext = destinationUrl.pathExtension.lowercased()
guard ImageType(fileExtension: ext) != nil else {
fatalError()
}
let destinationExtension = destinationUrl.pathExtension.lowercased()
guard let type = ImageType(fileExtension: destinationExtension)?.fileType else {
addWarning("Invalid image extension \(destinationExtension)", job: job)
return
}
let desiredWidth = CGFloat(job.width)
let sourceRep = image.representations[0]
let destinationSize = NSSize(width: sourceRep.pixelsWide, height: sourceRep.pixelsHigh)
.scaledDown(to: desiredWidth)
// create NSBitmapRep manually, if using cgImage, the resulting size is wrong
let rep = NSBitmapImageRep(bitmapDataPlanes: nil,
pixelsWide: Int(destinationSize.width),
pixelsHigh: Int(destinationSize.height),
bitsPerSample: 8,
samplesPerPixel: 4,
hasAlpha: true,
isPlanar: false,
colorSpaceName: NSColorSpaceName.deviceRGB,
bytesPerRow: Int(destinationSize.width) * 4,
bitsPerPixel: 32)!
let ctx = NSGraphicsContext(bitmapImageRep: rep)
NSGraphicsContext.saveGraphicsState()
NSGraphicsContext.current = ctx
image.draw(in: NSMakeRect(0, 0, destinationSize.width, destinationSize.height))
ctx?.flushGraphics()
NSGraphicsContext.restoreGraphicsState()
// Get NSData, and save it
guard let data = rep.representation(using: type, properties: [.compressionFactor: NSNumber(value: job.quality)]) else {
addWarning("Failed to get data", job: job)
return
}
do {
try data.createFolderAndWrite(to: destinationUrl)
} catch {
addWarning("Failed to write image (\(error))", job: job)
return
}
generatedImages.insert(job.destination)
}
/**
Create images of different types.
This function generates versions for the given image, including png/jpg, avif, and webp. Different pixel density versions (1x and 2x) are also generated.
- Parameter destination: The path to the destination file
*/
@discardableResult
func requireMultiVersionImage(source: String, destination: String, requiredBy path: String, width: Int, desiredHeight: Int?) -> NSSize {
// Add @1x version
_ = requireScaledMultiImage(source: source, destination: destination, requiredBy: path, width: width, desiredHeight: desiredHeight)
// Add @2x version
return requireScaledMultiImage(
source: source,
destination: destination.insert("@2x", beforeLast: "."),
requiredBy: path,
width: width * 2,
desiredHeight: desiredHeight.unwrapped { $0 * 2 })
}
@discardableResult
private func requireScaledMultiImage(source: String, destination: String, requiredBy path: String, width: Int, desiredHeight: Int?) -> NSSize {
let rawDestinationPath = destination.dropAfterLast(".")
let avifPath = rawDestinationPath + ".avif"
let webpPath = rawDestinationPath + ".webp"
let needsGeneration = !output.appendingPathComponent(avifPath).exists || !output.appendingPathComponent(webpPath).exists
let size = requireImage(at: destination, generatedFrom: source, requiredBy: path, quality: 1.0, width: width, height: desiredHeight, alwaysGenerate: needsGeneration)
multiImageJobs[destination] = path
return size
}
private func createMultiImages() {
let sort = multiImageJobs.sorted { $0.value < $1.value && $0.key < $1.key }
var count = 1
for (baseImage, path) in sort {
print(String(format: "Creating image versions: %4d / %d\r", count, sort.count), terminator: "")
fflush(stdout)
createMultiImages(from: baseImage, path: path)
count += 1
}
print(" \r", terminator: "")
}
private func createMultiImages(from source: String, path: String) {
guard generatedImages.contains(source) else {
return
}
let sourceUrl = output.appendingPathComponent(source)
let sourcePath = sourceUrl.path
guard sourceUrl.exists else {
addWarning("No image at path \(sourcePath)", destination: source, path: path)
missingImages[source] = path
return
}
let avifPath = source.dropAfterLast(".") + ".avif"
createAVIF(at: output.appendingPathComponent(avifPath).path, from: sourcePath)
generatedImages.insert(avifPath)
let webpPath = source.dropAfterLast(".") + ".webp"
createWEBP(at: output.appendingPathComponent(webpPath).path, from: sourcePath)
generatedImages.insert(webpPath)
compress(at: source)
}
private func createAVIF(at destination: String, from source: String, quality: Int = 55, effort: Int = 5) {
let folder = destination.dropAfterLast("/")
let command = "npx avif --input=\(source) --quality=\(quality) --effort=\(effort) --output=\(folder) --overwrite"
do {
_ = try FileSystem.safeShell(command)
} catch {
addWarning("Failed to create AVIF image", destination: destination, path: destination)
}
}
private func createWEBP(at destination: String, from source: String, quality: Int = 75) {
let command = "cwebp \(source) -q \(quality) -o \(destination)"
do {
_ = try FileSystem.safeShell(command)
} catch {
addWarning("Failed to create WEBP image", destination: destination, path: destination)
}
}
private func compress(at destination: String, quality: Int = 70) {
let command = "magick convert \(destination) -quality \(quality)% \(destination)"
do {
_ = try FileSystem.safeShell(command)
} catch {
addWarning("Failed to compress image", destination: destination, path: destination)
}
}
private func optimizeImages() {
let all = generatedImages
.filter { imageOptimSupportedFileExtensions.contains($0.lastComponentAfter(".")) }
.map { output.appendingPathComponent($0).path }
for i in stride(from: 0, to: all.count, by: imageOptimizationBatchSize) {
let endIndex = min(i+imageOptimizationBatchSize, all.count)
let batch = all[i..<endIndex]
print(String(format: "Optimizing images: %4d / %d\r", endIndex, all.count), terminator: "")
fflush(stdout)
if optimizeImageBatch(batch) {
optimizedImages.formUnion(batch)
}
}
print(" \r", terminator: "")
fflush(stdout)
print("\(optimizedImages.count) images optimized")
}
private func optimizeImageBatch(_ batch: ArraySlice<String>) -> Bool {
let command = "imageoptim " + batch.joined(separator: " ")
do {
_ = try FileSystem.safeShell(command)
return true
} catch {
addWarning("Failed to optimize images", destination: "", path: "")
return false
}
}
}