Screenshot difference engine with swift-assembly

difference.swift

import AsmMacro
import CoreGraphics
import UIKit

// arm64 C ABI:
//   x0 = before RGBA bytes
//   x1 = after RGBA bytes
//   x2 = pixel count
//   x3 = output mask bytes
//
// Each vector pass compares four complete 32-bit RGBA pixels. Changed pixels
// become 0xffffffff (opaque white) in the mask; unchanged pixels become zero.
@Asm(
    """
    mov x4, #0                         // x4 accumulates the changed-pixel count
    cmp x2, #4                         // the NEON loop needs four RGBA pixels
    b.lo tail

    vector_loop:
      ldr q0, [x0], #16                // load four before pixels; advance x0
      ldr q1, [x1], #16                // load four after pixels; advance x1
      eor v0.16b, v0.16b, v1.16b       // zero bytes match; nonzero bytes differ
      cmeq v0.4s, v0.4s, #0            // regroup as four 32-bit RGBA pixels
      mvn v0.16b, v0.16b               // changed = 0xffffffff, unchanged = 0
      str q0, [x3], #16                // write four pixels into the visual mask
      ushr v0.4s, v0.4s, #31           // turn each changed lane into the integer 1
      addv s0, v0.4s                   // horizontally add those four 0/1 lanes
      umov w5, v0.s[0]                 // move the vector sum into an integer register
      add x4, x4, x5                   // add this batch to the running count
      sub x2, x2, #4
      cmp x2, #4
      b.hs vector_loop

    tail:
      cbz x2, done

    tail_loop:
      ldr w5, [x0], #4                 // scalar path: load one complete RGBA pixel
      ldr w6, [x1], #4
      cmp w5, w6
      cset w5, ne                      // w5 becomes 1 when that pixel changed
      add x4, x4, x5
      neg w5, w5                       // 1 -> 0xffffffff for an opaque mask pixel
      str w5, [x3], #4
      subs x2, x2, #1
      b.ne tail_loop

    done:
      mov x0, x4                       // x0 is the ARM64 C return-value register
      ret
    """,
    arch: .arm64
)
func armPixelDiff(
    _ before: UnsafePointer<UInt8>,
    _ after: UnsafePointer<UInt8>,
    _ pixelCount: UInt64,
    _ mask: UnsafeMutablePointer<UInt8>
) -> UInt64

struct PixelDiffResult {
    let beforeImage: UIImage
    let afterImage: UIImage
    let maskImage: UIImage
    let width: Int
    let height: Int
    let changedPixelCount: UInt64
    let elapsedMilliseconds: Double

    var totalPixelCount: UInt64 {
        UInt64(width * height)
    }

    var changedPercentage: Double {
        guard totalPixelCount > 0 else { return 0 }
        return Double(changedPixelCount) / Double(totalPixelCount) * 100
    }

    var aspectRatio: CGFloat {
        CGFloat(width) / CGFloat(height)
    }
}

enum PixelDiffError: LocalizedError {
    case missingImage(String)
    case unreadableImage(String)
    case mismatchedDimensions(before: CGSize, after: CGSize)
    case bitmapCreationFailed

    var errorDescription: String? {
        switch self {
        case .missingImage(let name):
            return "Could not find \(name).png in the app bundle."
        case .unreadableImage(let name):
            return "\(name).png could not be decoded as a bitmap."
        case .mismatchedDimensions(let before, let after):
            return "The screenshots must have identical pixel dimensions. Before is \(Int(before.width)) × \(Int(before.height)); after is \(Int(after.width)) × \(Int(after.height))."
        case .bitmapCreationFailed:
            return "Core Graphics could not create the RGBA bitmap."
        }
    }
}

enum PixelDiffEngine {
    static func runBundledPair() throws -> PixelDiffResult {
        guard let beforeImage = UIImage(named: "ScreenshotBefore") else {
            throw PixelDiffError.missingImage("ScreenshotBefore")
        }
        guard let afterImage = UIImage(named: "ScreenshotAfter") else {
            throw PixelDiffError.missingImage("ScreenshotAfter")
        }

        let before = try RGBAImage(image: beforeImage, name: "ScreenshotBefore")
        let after = try RGBAImage(image: afterImage, name: "ScreenshotAfter")

        guard before.width == after.width, before.height == after.height else {
            throw PixelDiffError.mismatchedDimensions(
                before: CGSize(width: before.width, height: before.height),
                after: CGSize(width: after.width, height: after.height)
            )
        }

        let pixelCount = before.width * before.height
        var mask = [UInt8](repeating: 0, count: pixelCount * 4)

        let start = DispatchTime.now().uptimeNanoseconds
        let changedPixelCount = before.bytes.withUnsafeBufferPointer { beforeBuffer in
            after.bytes.withUnsafeBufferPointer { afterBuffer in
                mask.withUnsafeMutableBufferPointer { maskBuffer in
                    armPixelDiff(
                        beforeBuffer.baseAddress!,
                        afterBuffer.baseAddress!,
                        UInt64(pixelCount),
                        maskBuffer.baseAddress!
                    )
                }
            }
        }
        let end = DispatchTime.now().uptimeNanoseconds

        #if DEBUG
        validateWithSwift(
            before: before.bytes,
            after: after.bytes,
            mask: mask,
            assemblyCount: changedPixelCount
        )
        #endif

        return PixelDiffResult(
            beforeImage: beforeImage,
            afterImage: afterImage,
            maskImage: try makeMaskImage(bytes: mask, width: before.width, height: before.height),
            width: before.width,
            height: before.height,
            changedPixelCount: changedPixelCount,
            elapsedMilliseconds: Double(end - start) / 1_000_000
        )
    }

    private static func makeMaskImage(bytes: [UInt8], width: Int, height: Int) throws -> UIImage {
        guard let provider = CGDataProvider(data: Data(bytes) as CFData),
              let image = CGImage(
                width: width,
                height: height,
                bitsPerComponent: 8,
                bitsPerPixel: 32,
                bytesPerRow: width * 4,
                space: CGColorSpaceCreateDeviceRGB(),
                bitmapInfo: CGBitmapInfo(
                    rawValue: CGBitmapInfo.byteOrder32Big.rawValue
                        | CGImageAlphaInfo.premultipliedLast.rawValue
                ),
                provider: provider,
                decode: nil,
                shouldInterpolate: false,
                intent: .defaultIntent
              )
        else {
            throw PixelDiffError.bitmapCreationFailed
        }

        return UIImage(cgImage: image, scale: 1, orientation: .up)
    }

    #if DEBUG
    private static func validateWithSwift(
        before: [UInt8],
        after: [UInt8],
        mask: [UInt8],
        assemblyCount: UInt64
    ) {
        var expectedMask = [UInt8](repeating: 0, count: mask.count)
        var expectedCount: UInt64 = 0

        for offset in stride(from: 0, to: before.count, by: 4) {
            let changed =
                before[offset] != after[offset]
                || before[offset + 1] != after[offset + 1]
                || before[offset + 2] != after[offset + 2]
                || before[offset + 3] != after[offset + 3]

            if changed {
                expectedCount += 1
                expectedMask[offset] = 255
                expectedMask[offset + 1] = 255
                expectedMask[offset + 2] = 255
                expectedMask[offset + 3] = 255
            }
        }

        assert(assemblyCount == expectedCount, "ARM count differs from the Swift reference.")
        assert(mask == expectedMask, "ARM mask differs from the Swift reference.")
    }
    #endif
}

private struct RGBAImage {
    let width: Int
    let height: Int
    let bytes: [UInt8]

    init(image: UIImage, name: String) throws {
        guard let cgImage = image.cgImage else {
            throw PixelDiffError.unreadableImage(name)
        }

        let pixelWidth = cgImage.width
        let pixelHeight = cgImage.height
        var output = [UInt8](repeating: 0, count: pixelWidth * pixelHeight * 4)

        let rendered = output.withUnsafeMutableBytes { rawBuffer -> Bool in
            guard let context = CGContext(
                data: rawBuffer.baseAddress,
                width: pixelWidth,
                height: pixelHeight,
                bitsPerComponent: 8,
                bytesPerRow: pixelWidth * 4,
                space: CGColorSpaceCreateDeviceRGB(),
                bitmapInfo: CGBitmapInfo.byteOrder32Big.rawValue
                    | CGImageAlphaInfo.premultipliedLast.rawValue
            ) else {
                return false
            }

            context.interpolationQuality = .none
            context.draw(
                cgImage,
                in: CGRect(x: 0, y: 0, width: pixelWidth, height: pixelHeight)
            )
            return true
        }

        guard rendered else {
            throw PixelDiffError.bitmapCreationFailed
        }

        width = pixelWidth
        height = pixelHeight
        bytes = output
    }
}