/*
 * Copyright 2026 Google LLC.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "experimental/rust_jpeg/decoder/impl/SkJpegRustCodec.h"

#include <algorithm>

#include "include/codec/SkCodecAnimation.h"
#include "include/core/SkColorType.h"
#include "include/core/SkData.h"
#include "include/core/SkImageInfo.h"
#include "include/core/SkSpan.h"
#include "include/core/SkStream.h"
#include "include/private/SkAssert.h"
#include "include/private/SkEncodedInfo.h"
#include "include/private/SkGainmapInfo.h"
#include "include/private/SkJpegMetadataDecoder.h"
#include "include/private/SkTemplates.h"
#include "rust/common/SkStreamAdapter.h"
#include "rust/common/SpanUtils.h"
#include "src/codec/SkSwizzler.h"
#include "src/core/SkSafeMath.h"
#include "src/core/SkStreamPriv.h"

#include "include/private/SkExif.h"

#ifdef SK_CODEC_DECODES_JPEG_GAINMAPS
#include "include/core/SkColorSpace.h"
#include "include/private/SkXmp.h"
#include "modules/skcms/skcms.h"
#include "src/codec/SkJpegConstants.h"
#include "src/codec/SkJpegXmp.h"
#endif  // SK_CODEC_DECODES_JPEG_GAINMAPS

static constexpr SkColorType kXformSrcColorType = kRGBA_8888_SkColorType;

// Static assertions to validate that Rust JpegColor/JpegAlpha enum values
// match the corresponding SkEncodedInfo enum values.
static_assert(static_cast<int>(rust_jpeg::JpegColor::RGB) ==
              static_cast<int>(SkEncodedInfo::kRGB_Color),
              "JpegColor::RGB must match SkEncodedInfo::kRGB_Color");
static_assert(static_cast<int>(rust_jpeg::JpegColor::Grayscale) ==
              static_cast<int>(SkEncodedInfo::kGray_Color),
              "JpegColor::Grayscale must match SkEncodedInfo::kGray_Color");
static_assert(static_cast<int>(rust_jpeg::JpegColor::InvertedCMYK) ==
              static_cast<int>(SkEncodedInfo::kInvertedCMYK_Color),
              "JpegColor::InvertedCMYK must match SkEncodedInfo::kInvertedCMYK_Color");

static_assert(static_cast<int>(rust_jpeg::JpegAlpha::Opaque) ==
              static_cast<int>(SkEncodedInfo::kOpaque_Alpha),
              "JpegAlpha::Opaque must match SkEncodedInfo::kOpaque_Alpha");

namespace {

SkCodec::Result MapDecodingResult(rust_jpeg::DecodingResult rustResult) {
    switch (rustResult) {
        case rust_jpeg::DecodingResult::Success:
            return SkCodec::kSuccess;
        case rust_jpeg::DecodingResult::FormatError:
            return SkCodec::kErrorInInput;
        case rust_jpeg::DecodingResult::ParameterError:
            return SkCodec::kInvalidParameters;
        case rust_jpeg::DecodingResult::UnsupportedFeature:
            return SkCodec::kUnimplemented;
        case rust_jpeg::DecodingResult::IncompleteInput:
            return SkCodec::kIncompleteInput;
        case rust_jpeg::DecodingResult::MemoryError:
            return SkCodec::kInternalError;
        case rust_jpeg::DecodingResult::OtherError:
            return SkCodec::kErrorInInput;
    }
    SK_ABORT("Unexpected `rust_jpeg::DecodingResult`: %d", static_cast<int>(rustResult));
}

}  // namespace

std::unique_ptr<SkJpegRustCodec> SkJpegRustCodec::MakeFromStream(std::unique_ptr<SkStream> stream,
                                                                 Result* result) {
    Result resultStorage;
    if (!result) {
        result = &resultStorage;
    }

    if (!stream) {
        *result = kInvalidInput;
        return nullptr;
    }

    auto inputAdapter = std::make_unique<rust::stream::SkStreamAdapter>(stream.get());
    rust::Box<rust_jpeg::Reader> reader = rust_jpeg::new_reader(std::move(inputAdapter));

    if (!reader->metadata_loaded()) {
        rust_jpeg::DecodingResult metadataResult = reader->read_metadata();

        if (metadataResult != rust_jpeg::DecodingResult::Success) {
            *result = MapDecodingResult(metadataResult);
            return nullptr;
        }

        SkASSERT_RELEASE(reader->metadata_loaded());
    }

    uint32_t width = reader->width();
    uint32_t height = reader->height();
    rust_jpeg::JpegColor rustColor = reader->color();
    rust_jpeg::JpegAlpha rustAlpha = reader->alpha();

    SkEncodedInfo::Color color = static_cast<SkEncodedInfo::Color>(rustColor);
    SkEncodedInfo::Alpha alpha = static_cast<SkEncodedInfo::Alpha>(rustAlpha);
    constexpr int kBitsPerComponent = 8;

    std::unique_ptr<SkCodecs::ColorProfile> colorProfile;
    rust::Vec<uint8_t> profileData = reader->icc_profile();
    if (!profileData.empty()) {
        sk_sp<SkData> iccData = SkData::MakeWithCopy(profileData.data(), profileData.size());
        colorProfile = SkCodecs::ColorProfile::MakeICCProfile(std::move(iccData));
    }
    // Validate ICC profile data space against the image color space,
    // matching the existing C++ JPEG codec's behavior.  Drop profiles
    // whose data space is incompatible.
    if (colorProfile) {
        auto dataSpace = colorProfile->dataSpace();
        switch (color) {
            case SkEncodedInfo::kInvertedCMYK_Color:
                if (dataSpace != SkCodecs::ColorProfile::DataSpace::kCMYK) {
                    colorProfile = nullptr;
                }
                break;
            case SkEncodedInfo::kGray_Color:
                if (dataSpace != SkCodecs::ColorProfile::DataSpace::kGray &&
                    dataSpace != SkCodecs::ColorProfile::DataSpace::kRGB) {
                    colorProfile = nullptr;
                }
                break;
            default:
                if (dataSpace != SkCodecs::ColorProfile::DataSpace::kRGB) {
                    colorProfile = nullptr;
                }
                break;
        }
    }

    // Parse orientation from EXIF data using SkExif::Parse (shared with C++ codec).
    SkEncodedOrigin orientation = kDefault_SkEncodedOrigin;
    {
        rust::Vec<uint8_t> exifBytes = reader->exif_data();
        if (!exifBytes.empty()) {
            auto exifData = SkData::MakeWithCopy(exifBytes.data(), exifBytes.size());
            SkExif::Metadata exif;
            SkExif::Parse(exif, exifData.get());
            if (exif.fOrigin.has_value()) {
                orientation = exif.fOrigin.value();
            }
        }
    }

    SkEncodedInfo encodedInfo = SkEncodedInfo::Make(
        width,
        height,
        color,
        alpha,
        kBitsPerComponent,
        std::move(colorProfile)
    );

    // Guard against overflow: reject images whose raw pixel buffer would
    // exceed size_t range.
    size_t bytesPerPixel;
    switch (color) {
        case SkEncodedInfo::kRGB_Color:
            bytesPerPixel = 3;
            break;
        case SkEncodedInfo::kGray_Color:
            bytesPerPixel = 1;
            break;
        case SkEncodedInfo::kInvertedCMYK_Color:
            bytesPerPixel = 4;
            break;
        default:
            *result = kInvalidInput;
            return nullptr;
    }
    SkSafeMath safe;
    size_t srcRowBytes = safe.mul(safe.castTo<size_t>(width), bytesPerPixel);
    (void)safe.mul(safe.castTo<size_t>(height), srcRowBytes);
    if (!safe.ok()) {
        *result = kInternalError;
        return nullptr;
    }

    *result = kSuccess;
    return std::unique_ptr<SkJpegRustCodec>(new SkJpegRustCodec(
        std::move(encodedInfo),
        std::move(stream),
        std::move(reader),
        orientation
    ));
}

SkJpegRustCodec::SkJpegRustCodec(SkEncodedInfo&& encodedInfo,
                                 std::unique_ptr<SkStream> stream,
                                 rust::Box<rust_jpeg::Reader> reader,
                                 SkEncodedOrigin origin)
    : SkCodec(std::move(encodedInfo), skcms_PixelFormat_RGBA_8888,
              /* stream = */ nullptr, origin)
    , fPrivStream(std::move(stream))
    , fReader(std::move(reader)) {
    SkASSERT_RELEASE(fPrivStream);
}

SkJpegRustCodec::~SkJpegRustCodec() = default;

sk_sp<const SkData> SkJpegRustCodec::getEncodedData() const {
    SkASSERT_RELEASE(fPrivStream);
    sk_sp<const SkData> data = fPrivStream->getData();
    if (data) {
        return data;
    }
    auto dStream = fPrivStream->duplicate();
    if (!dStream || !dStream->hasLength()) {
        return nullptr;
    }
    return SkData::MakeFromStream(dStream.get(), dStream->getLength());
}

bool SkJpegRustCodec::onGetFrameInfo(int index, FrameInfo* info) const {
    if (index != 0) {
        return false;
    }
    if (info) {
        info->fRequiredFrame = SkCodec::kNoFrame;
        info->fDuration = 0;
        info->fFullyReceived = fReader->image_data_loaded();
        info->fAlphaType = this->getInfo().alphaType();
        info->fHasAlphaWithinBounds = false;  // JPEG never has alpha
        info->fDisposalMethod = SkCodecAnimation::DisposalMethod::kKeep;
        info->fBlend = SkCodecAnimation::Blend::kSrc;
        info->fFrameRect = SkIRect::MakeSize(this->dimensions());
    }
    return true;
}

bool SkJpegRustCodec::onGetGainmapInfo(SkGainmapInfo* info,
                                       std::unique_ptr<SkStream>* gainmapImageStream) {
#ifndef SK_CODEC_DECODES_JPEG_GAINMAPS
    return false;
#else
    // Segment scanning and MPF TIFF/IFD parsing happen in Rust. Higher-level
    // gainmap detection (XMP, ISO 21496-1) uses C++ SkJpegMetadataDecoder.
    if (!fReader->segments_scanned()) {
        return false;
    }
    if (!fReader->has_mpf()) {
        return false;
    }

    // Two lambdas needed: CXX exposes Reader and EmbeddedJpegScanner methods
    // with different prefixes in a flat namespace.
    auto buildSegmentsFromReader = [](const rust_jpeg::Reader& reader)
            -> std::vector<SkJpegMetadataDecoder::Segment> {
        std::vector<SkJpegMetadataDecoder::Segment> segs;
        const uint32_t count = reader.segment_count();
        segs.reserve(count);
        for (uint32_t i = 0; i < count; ++i) {
            rust_jpeg::JpegSegmentInfo si = reader.segment_info(i);
            if (si.parameter_length <= 2) continue;
            rust::Vec<uint8_t> pd = reader.segment_data(i);
            if (pd.empty()) continue;
            segs.emplace_back(si.marker, SkData::MakeWithCopy(pd.data(), pd.size()));
        }
        return segs;
    };

    auto buildSegmentsFromEmbedded = [](const rust_jpeg::EmbeddedJpegScanner& scanner)
            -> std::vector<SkJpegMetadataDecoder::Segment> {
        std::vector<SkJpegMetadataDecoder::Segment> segs;
        const uint32_t count = scanner.embedded_segment_count();
        segs.reserve(count);
        for (uint32_t i = 0; i < count; ++i) {
            rust_jpeg::JpegSegmentInfo si = scanner.embedded_segment_info(i);
            if (si.parameter_length <= 2) continue;
            rust::Vec<uint8_t> pd = scanner.embedded_segment_data(i);
            if (pd.empty()) continue;
            segs.emplace_back(si.marker, SkData::MakeWithCopy(pd.data(), pd.size()));
        }
        return segs;
    };

    auto baseMetadata = SkJpegMetadataDecoder::Make(buildSegmentsFromReader(*fReader));
    if (!baseMetadata) {
        return false;
    }

    SkExif::Metadata baseExif;
    SkExif::Parse(baseExif, baseMetadata->getExifMetadata(/*copyData=*/false).get());

    bool isoGainmapPresent = SkGainmapInfo::ParseVersion(
            baseMetadata->getISOGainmapMetadata(/*copyData=*/false).get());

    std::vector<sk_sp<const SkData>> baseApp1Params;
    {
        const uint32_t segCount = fReader->segment_count();
        for (uint32_t i = 0; i < segCount; ++i) {
            rust_jpeg::JpegSegmentInfo si = fReader->segment_info(i);
            if (si.marker == kXMPMarker && si.parameter_length > 2) {
                rust::Vec<uint8_t> pd = fReader->segment_data(i);
                if (!pd.empty()) {
                    baseApp1Params.push_back(SkData::MakeWithCopy(pd.data(), pd.size()));
                }
            }
        }
    }
    auto baseXmp = SkJpegMakeXmp(baseApp1Params);

    bool adobeGainmapPresent = baseXmp && baseXmp->getGainmapInfoAdobe(nullptr);

    sk_sp<const SkData> encodedData = this->getEncodedData();
    if (!encodedData) {
        return false;
    }

    const uint32_t mpfSegOff = fReader->mpf_segment_offset();
    const uint32_t mpfCount = fReader->mpf_image_count();

    // MPF entries are zero-based. Entry 0 is the primary image, so inspect
    // secondary images in [1, mpfCount).
    for (uint32_t mpIdx = 1; mpIdx < mpfCount; ++mpIdx) {
        rust_jpeg::MpfImageEntry entry = fReader->mpf_image_entry(mpIdx);

        // 2 (marker) + 2 (param length) + 4 (MPF signature).
        constexpr size_t kMpfHeaderSize = 8;
        SkSafeMath safe;
        size_t imageOffset = safe.add(safe.add(static_cast<size_t>(mpfSegOff), kMpfHeaderSize),
                                      static_cast<size_t>(entry.data_offset));
        size_t imageSize = entry.size;
        size_t imageEnd = safe.add(imageOffset, imageSize);
        if (!safe.ok() || imageEnd > encodedData->size()) {
            continue;
        }

        rust::Slice<const uint8_t> embSlice(
                encodedData->bytes() + imageOffset, imageSize);
        auto embScanner = rust_jpeg::scan_embedded_jpeg(embSlice);
        if (embScanner->embedded_had_error()) {
            continue;
        }

        auto embMetadata = SkJpegMetadataDecoder::Make(buildSegmentsFromEmbedded(*embScanner));
        if (!embMetadata) {
            continue;
        }

        SkGainmapInfo gainmapInfo;
        bool foundGainmap = false;

        if (isoGainmapPresent) {
            foundGainmap = SkGainmapInfo::Parse(
                    embMetadata->getISOGainmapMetadata(/*copyData=*/false).get(), gainmapInfo);
            if (foundGainmap && gainmapInfo.fGainmapMathColorSpace) {
                // Use the gain map image's ICC color space for gainmap math.
                sk_sp<SkColorSpace> imageColorSpace;
                auto iccData = embMetadata->getICCProfileData(/*copyData=*/false);
                skcms_ICCProfile iccProfile;
                if (iccData && skcms_Parse(iccData->data(), iccData->size(), &iccProfile)) {
                    imageColorSpace = SkColorSpace::Make(iccProfile);
                }
                gainmapInfo.fGainmapMathColorSpace = std::move(imageColorSpace);
            }
        }

        if (!foundGainmap) {
            std::vector<sk_sp<const SkData>> embApp1Params;
            const uint32_t embSegCount = embScanner->embedded_segment_count();
            for (uint32_t j = 0; j < embSegCount; ++j) {
                rust_jpeg::JpegSegmentInfo si = embScanner->embedded_segment_info(j);
                if (si.marker == kXMPMarker && si.parameter_length > 2) {
                    rust::Vec<uint8_t> pd = embScanner->embedded_segment_data(j);
                    if (!pd.empty()) {
                        embApp1Params.push_back(
                                SkData::MakeWithCopy(pd.data(), pd.size()));
                    }
                }
            }
            auto embXmp = SkJpegMakeXmp(embApp1Params);

            if (embXmp) {
                if (adobeGainmapPresent) {
                    foundGainmap = embXmp->getGainmapInfoAdobe(&gainmapInfo);
                }
                if (!foundGainmap && baseExif.fHdrHeadroom.has_value()) {
                    foundGainmap = embXmp->getGainmapInfoApple(
                            baseExif.fHdrHeadroom.value(), &gainmapInfo);
                }
            }
        }

        if (foundGainmap) {
            *info = gainmapInfo;
            *gainmapImageStream = SkMemoryStream::Make(
                    SkData::MakeWithCopy(encodedData->bytes() + imageOffset, imageSize));
            return true;
        }
    }

    return false;
#endif  // SK_CODEC_DECODES_JPEG_GAINMAPS
}

bool SkJpegRustCodec::onGetGainmapCodec(SkGainmapInfo* info,
                                        std::unique_ptr<SkCodec>* gainmapCodec) {
    std::unique_ptr<SkStream> stream;
    if (!this->onGetGainmapInfo(info, &stream)) {
        return false;
    }
    if (gainmapCodec) {
        Result result;
        *gainmapCodec = SkCodec::MakeFromStream(std::move(stream), &result);
        if (!*gainmapCodec) {
            return false;
        }
    }
    return true;
}

bool SkJpegRustCodec::onRewind() {
    fIncrementalDecodingState.reset();

    if (!fPrivStream->rewind()) {
        return false;
    }

    auto inputAdapter = std::make_unique<rust::stream::SkStreamAdapter>(fPrivStream.get());
    fReader = rust_jpeg::new_reader(std::move(inputAdapter));

    if (!fReader->metadata_loaded()) {
        rust_jpeg::DecodingResult metadataResult = fReader->read_metadata();
        if (metadataResult != rust_jpeg::DecodingResult::Success) {
            return false;
        }
    }

    return true;
}

SkCodec::Result SkJpegRustCodec::initializeSwizzler(const SkImageInfo& dstInfo,
                                                    const Options& opts) {
    SkImageInfo swizzlerInfo = dstInfo;
    SkCodec::Options swizzlerOptions = opts;

    if (this->xformOnDecode()) {
        fXformBuffer.reset(new uint32_t[dstInfo.width()]);
        swizzlerInfo = swizzlerInfo.makeColorType(kXformSrcColorType);
        if (kPremul_SkAlphaType == dstInfo.alphaType()) {
            swizzlerInfo = swizzlerInfo.makeAlphaType(kUnpremul_SkAlphaType);
        }
        swizzlerOptions.fZeroInitialized = kNo_ZeroInitialized;
    }

    fSwizzler = SkSwizzler::Make(this->getEncodedInfo(), nullptr, swizzlerInfo, swizzlerOptions);
    if (!fSwizzler) {
        return kInvalidConversion;
    }
    return kSuccess;
}

void SkJpegRustCodec::swizzleRow(const uint8_t* srcRow, void* dstRow) {
    if (this->xformOnDecode()) {
        const int32_t width = fSwizzler->swizzleWidth();
        fSwizzler->swizzle(fXformBuffer.get(), srcRow);
        this->applyColorXform(dstRow, fXformBuffer.get(), width);
    } else {
        fSwizzler->swizzle(dstRow, srcRow);
    }
}

SkCodec::Result SkJpegRustCodec::onGetPixels(const SkImageInfo& info,
                                             void* dst,
                                             size_t dstRowStride,
                                             const Options& options,
                                             int* rowsDecoded) {
    fReader->reset_decode_state();

    Result result = this->initializeSwizzler(info, options);
    if (result != kSuccess) {
        return result;
    }

    int decodedRows = 0;
    result = this->performFullDecode(info, dst, dstRowStride, &decodedRows);
    if (rowsDecoded) {
        *rowsDecoded = decodedRows;
    }
    return result;
}

SkCodec::Result SkJpegRustCodec::performFullDecode(const SkImageInfo& dstInfo,
                                                   void* dst,
                                                   size_t dstRowStride,
                                                   int* rowsDecoded) {
    rust_jpeg::DecodingResult readResult = fReader->read_image_data();
    if (readResult != rust_jpeg::DecodingResult::Success) {
        return MapDecodingResult(readResult);
    }

    const size_t srcRowBytes = fReader->row_bytes();

    rust::Slice<const uint8_t> imageData;
    rust_jpeg::DecodedRowsInfo rowsInfo = fReader->get_next_rows(imageData);

    if (imageData.empty()) {
        return kErrorInInput;
    }

    SkSpan<const uint8_t> srcImage(imageData.data(), imageData.size());
    SkSafeMath safe;
    size_t totalDstSize = safe.mul(safe.castTo<size_t>(dstInfo.height()), dstRowStride);
    if (!safe.ok()) {
        return kInternalError;
    }
    SkSpan<uint8_t> dstImage(static_cast<uint8_t*>(dst), totalDstSize);
    const uint32_t rowCount = rowsInfo.row_count;

    for (uint32_t y = 0; y < rowCount; ++y) {
        SkSpan<const uint8_t> srcRow = srcImage.subspan(y * srcRowBytes, srcRowBytes);
        SkSpan<uint8_t> dstRow = dstImage.subspan(y * dstRowStride, dstRowStride);
        this->swizzleRow(srcRow.data(), dstRow.data());
    }

    if (rowsDecoded) {
        *rowsDecoded = static_cast<int>(rowCount);
    }

    return (rowCount == static_cast<uint32_t>(this->dimensions().height())) ? kSuccess
                                                                            : kIncompleteInput;
}

SkCodec::Result SkJpegRustCodec::onStartIncrementalDecode(const SkImageInfo& dstInfo,
                                                          void* dst,
                                                          size_t dstRowBytes,
                                                          const Options& options) {
    if (dstRowBytes < dstInfo.minRowBytes()) {
        return kInvalidParameters;
    }

    if (options.fSubset) {
        return kInvalidParameters;
    }

    Result result = this->initializeSwizzler(dstInfo, options);
    if (result != kSuccess) {
        return result;
    }

    fReader->reset_decode_state();

    DecodingState state;

    const int32_t height = dstInfo.height();
    SkSafeMath safe;
    size_t totalDstSize = safe.mul(safe.castTo<size_t>(height), dstRowBytes);
    if (!safe.ok()) {
        return kInternalError;
    }

    state.fDst = SkSpan(static_cast<uint8_t*>(dst), totalDstSize);
    state.fDstRowStride = dstRowBytes;

    fIncrementalDecodingState = std::move(state);

    return kSuccess;
}

SkCodec::Result SkJpegRustCodec::onIncrementalDecode(int* rowsDecoded) {
    if (!fIncrementalDecodingState.has_value()) {
        return kInvalidParameters;
    }

    Result result = this->incrementalDecode(*fIncrementalDecodingState, rowsDecoded);

    if (result != kIncompleteInput) {
        fIncrementalDecodingState.reset();
    }

    return result;
}

SkCodec::Result SkJpegRustCodec::incrementalDecode(DecodingState& state, int* rowsDecodedPtr) {
    const size_t srcRowBytes = fReader->row_bytes();

    rust_jpeg::DecodingResult readResult = fReader->read_incremental_image_data();

    if (readResult != rust_jpeg::DecodingResult::Success &&
        readResult != rust_jpeg::DecodingResult::IncompleteInput) {
        return MapDecodingResult(readResult);
    }

    rust::Slice<const uint8_t> imageData;
    rust_jpeg::DecodedRowsInfo rowsInfo = fReader->get_next_rows(imageData);
    const uint32_t dstRowStart = rowsInfo.dst_row_start;
    const uint32_t rowCount = rowsInfo.row_count;
    if (rowsInfo.is_preview) {
        // Progressive previews replace the initialized frame from row zero.
        // SkCodec permits rowsDecoded to include initialized but unfinished rows.
        if (dstRowStart != 0) {
            return kErrorInInput;
        }
    }

    uint32_t copiedRows = 0;
    if (rowCount > 0) {
        const uint32_t height = static_cast<uint32_t>(this->dimensions().height());
        if (srcRowBytes == 0 || imageData.empty() || dstRowStart >= height ||
            rowCount > height - dstRowStart) {
            return kErrorInInput;
        }

        SkSafeMath safe;
        const size_t requiredSrcBytes = safe.mul(safe.castTo<size_t>(rowCount), srcRowBytes);
        if (!safe.ok() || imageData.size() < requiredSrcBytes) {
            return kErrorInInput;
        }

        SkSpan<const uint8_t> srcImage(imageData.data(), imageData.size());

        for (uint32_t i = 0; i < rowCount; ++i) {
            const uint32_t dstY = dstRowStart + i;
            const size_t srcOffset = safe.mul(safe.castTo<size_t>(i), srcRowBytes);
            const size_t dstOffset = safe.mul(safe.castTo<size_t>(dstY), state.fDstRowStride);
            if (!safe.ok()) {
                return kInternalError;
            }
            SkSpan<const uint8_t> srcRow = srcImage.subspan(srcOffset, srcRowBytes);
            SkSpan<uint8_t> dstRow = state.fDst.subspan(dstOffset, state.fDstRowStride);
            this->swizzleRow(srcRow.data(), dstRow.data());
            copiedRows++;
        }
        state.fTotalRowsInitialized = std::max(
            state.fTotalRowsInitialized, static_cast<int>(dstRowStart + copiedRows));
    }

    if (rowsDecodedPtr) {
        *rowsDecodedPtr = state.fTotalRowsInitialized;
    }

    return readResult == rust_jpeg::DecodingResult::Success &&
               state.fTotalRowsInitialized >= this->dimensions().height()
           ? kSuccess
           : kIncompleteInput;
}