src/codec/SkPngRustCodec.cpp

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

#include "src/codec/SkPngRustCodec.h"

#include <cstring>
#include <limits>
#include <memory>
#include <tuple>
#include <utility>

#include "include/core/SkColorSpace.h"
#include "include/core/SkStream.h"
#include "include/private/SkAssert.h"
#include "include/private/SkEncodedInfo.h"
#include "include/private/SkHdrMetadata.h"
#include "include/private/SkSafe32.h"
#include "include/private/SkTemplates.h"
#include "modules/skcms/skcms.h"
#include "rust/common/SpanUtils.h"
#include "rust/png/FFI.rs.h"
#include "src/codec/SkCodecPriv.h"
#include "src/codec/SkFrameHolder.h"
#include "src/codec/SkParseEncodedOrigin.h"
#include "src/codec/SkPngCompositeChunkReader.h"
#include "src/codec/SkPngPriv.h"
#include "src/codec/SkSwizzler.h"
#include "src/core/SkRasterPipeline.h"
#include "src/core/SkRasterPipelineOpList.h"
#include "src/core/SkSafeMath.h"
#include "third_party/rust/cxx/v1/cxx.h"

#ifdef __clang__
#pragma clang diagnostic error "-Wconversion"
#endif

namespace {

SkEncodedInfo::Color ToColor(rust_png::ColorType colorType, const rust_png::Reader& reader) {
    switch (colorType) {
        case rust_png::ColorType::Grayscale:
            return SkEncodedInfo::kGray_Color;
        case rust_png::ColorType::Rgb:
            if (reader.has_sbit_chunk()) {
                SkSpan<const uint8_t> sBit = ToSkSpan(reader.get_sbit_chunk());

                // Restricting to 8-bit images, because of `SkEncodedInfo`
                // asserts.  See also https://crbug.com/458852444.
                if (sBit.size() == 3 && reader.output_bits_per_component() == 8) {
                  if (sBit[0] == 5 && sBit[1] == 6 && sBit[2] == 5) {
                      return SkEncodedInfo::k565_Color;
                  }
                }
            }
            return SkEncodedInfo::kRGB_Color;
        case rust_png::ColorType::GrayscaleAlpha:
            if (reader.has_sbit_chunk()) {
                SkSpan<const uint8_t> sBit = ToSkSpan(reader.get_sbit_chunk());

                // Restricting to 8-bit images, because of `SkEncodedInfo`
                // asserts.  See also https://crbug.com/458852444.
                if (sBit.size() == 2 && reader.output_bits_per_component() == 8) {
                    if (sBit[0] == kGraySigBit_GrayAlphaIsJustAlpha && sBit[1] == 8) {
                        return SkEncodedInfo::kXAlpha_Color;
                    }
                }
            }
            return SkEncodedInfo::kGrayAlpha_Color;
        case rust_png::ColorType::Rgba:
            return SkEncodedInfo::kRGBA_Color;
        case rust_png::ColorType::Indexed:
            return SkEncodedInfo::kPalette_Color;
    }
    SK_ABORT("Unexpected `rust_png::ColorType`: %d", static_cast<int>(colorType));
}

SkEncodedInfo::Alpha ToAlpha(rust_png::ColorType colorType, const rust_png::Reader& reader) {
    switch (colorType) {
        case rust_png::ColorType::Grayscale:
        case rust_png::ColorType::Rgb:
            return SkEncodedInfo::kOpaque_Alpha;
        case rust_png::ColorType::GrayscaleAlpha:
        case rust_png::ColorType::Rgba:
            return SkEncodedInfo::kUnpremul_Alpha;
        case rust_png::ColorType::Indexed:
            if (reader.has_trns_chunk()) {
                return SkEncodedInfo::kUnpremul_Alpha;
            } else {
                return SkEncodedInfo::kOpaque_Alpha;
            }
    }
    SK_ABORT("Unexpected `rust_png::ColorType`: %d", static_cast<int>(colorType));
}

SkCodecAnimation::DisposalMethod ToDisposalMethod(rust_png::DisposeOp op) {
    switch (op) {
        case rust_png::DisposeOp::None:
            return SkCodecAnimation::DisposalMethod::kKeep;
        case rust_png::DisposeOp::Background:
            return SkCodecAnimation::DisposalMethod::kRestoreBGColor;
        case rust_png::DisposeOp::Previous:
            return SkCodecAnimation::DisposalMethod::kRestorePrevious;
    }
    SK_ABORT("Unexpected `rust_png::DisposeOp`: %d", static_cast<int>(op));
}

SkCodecAnimation::Blend ToBlend(rust_png::BlendOp op) {
    switch (op) {
        case rust_png::BlendOp::Source:
            return SkCodecAnimation::Blend::kSrc;
        case rust_png::BlendOp::Over:
            return SkCodecAnimation::Blend::kSrcOver;
    }
    SK_ABORT("Unexpected `rust_png::BlendOp`: %d", static_cast<int>(op));
}

SkColorSpacePrimaries ToSkColorSpacePrimaries(const rust_png::ColorSpacePrimaries& p) {
    return SkColorSpacePrimaries({p.fRX, p.fRY, p.fGX, p.fGY, p.fBX, p.fBY, p.fWX, p.fWY});
}

skhdr::MasteringDisplayColorVolume ToSkMDCV(const rust_png::MasteringDisplayColorVolume& mdcv) {
    return skhdr::MasteringDisplayColorVolume({
        ToSkColorSpacePrimaries(mdcv.fDisplayPrimaries),
        mdcv.fMaximumDisplayMasteringLuminance,
        mdcv.fMinimumDisplayMasteringLuminance});
}

skhdr::ContentLightLevelInformation ToSkCLLI(const rust_png::ContentLightLevelInfo& clli) {
    return skhdr::ContentLightLevelInformation({clli.fMaxCLL, clli.fMaxFALL});
}

std::unique_ptr<SkCodecs::ColorProfile> CreateColorProfile(const rust_png::Reader& reader) {
    // NOTE: This method is based on `read_color_profile` in
    // `src/codec/SkPngCodec.cpp` but has been refactored to use Rust inputs
    // instead of `libpng`.

    // Considering the `cICP` chunk first, because the spec at
    // https://www.w3.org/TR/png-3/#cICP-chunk says: "This chunk, if understood
    // by the decoder, is the highest-precedence color chunk."
    uint8_t cicpPrimariesId = 0;
    uint8_t cicpTransferId = 0;
    uint8_t cicpMatrixId = 0;
    bool cicpIsFullRange = false;
    if (reader.try_get_cicp_chunk(cicpPrimariesId, cicpTransferId, cicpMatrixId, cicpIsFullRange)) {
        // https://www.w3.org/TR/png-3/#cICP-chunk says "RGB is currently the
        // only supported color model in PNG, and as such Matrix Coefficients
        // shall be set to 0."
        //
        // According to SkColorSpace::MakeCICP narrow range images are rare and
        // therefore not supported.
        if (cicpMatrixId == 0 && cicpIsFullRange) {
            sk_sp<SkColorSpace> colorSpace =
                    SkColorSpace::MakeCICP(static_cast<SkNamedPrimaries::CicpId>(cicpPrimariesId),
                                           static_cast<SkNamedTransferFn::CicpId>(cicpTransferId));
            if (colorSpace) {
                return SkCodecs::ColorProfile::Make(colorSpace);
            }
        }
    }

    if (reader.has_iccp_chunk()) {
        // `SkData::MakeWithCopy` is resilient against 0-sized inputs, so
        // no need to check `rust_slice.empty()` here.
        rust::Slice<const uint8_t> rust_slice = reader.get_iccp_chunk();
        sk_sp<SkData> owned_data = SkData::MakeWithCopy(rust_slice.data(), rust_slice.size());
        std::unique_ptr<SkCodecs::ColorProfile> parsed_data =
                SkCodecs::ColorProfile::MakeICCProfile(std::move(owned_data));
        if (parsed_data) {
            return parsed_data;
        }
    }

    if (reader.is_srgb()) {
        // TODO(https://crbug.com/362304558): Consider the intent field from the
        // `sRGB` chunk.
        return nullptr;
    }

    // Next, check for presence of `gAMA` and `cHRM` chunks.
    float gamma = 0;
    // Unlike libpng, image-rs does not seem to validate the gamma value.
    // So we check explicitly for zero, which is an invalid value encountered in the wild.
    // TODO: upstream?
    if (!reader.try_get_gama(gamma) || gamma <= 0) {
        // We ignore whether `chRM` is present or not.
        //
        // This preserves the behavior decided in Chromium's 83587041dc5f1428c09
        // (https://codereview.chromium.org/2469473002).  The PNG spec states
        // that cHRM is valid even without gAMA but we cannot apply the cHRM
        // without guessing a gAMA.  Color correction is not a guessing game,
        // so we match the behavior of Safari and Firefox instead (compat).
        return nullptr;
    }
    rust_png::ColorSpacePrimaries chrm;
    const bool got_chrm = reader.try_get_chrm(chrm);
    if (!got_chrm) {
        // If there is no `cHRM` chunk then check if `gamma` is neutral (in PNG
        // / `SkNamedTransferFn::k2Dot2` sense).  `kPngGammaThreshold` mimics
        // `PNG_GAMMA_THRESHOLD_FIXED` from `libpng`.
        constexpr float kPngGammaThreshold = 0.05f;
        constexpr float kMinNeutralValue = 1.0f - kPngGammaThreshold;
        constexpr float kMaxNeutralValue = 1.0f + kPngGammaThreshold;
        float tmp = gamma * 2.2f;
        bool is_neutral = kMinNeutralValue < tmp && tmp < kMaxNeutralValue;
        if (is_neutral) {
            // Don't construct a custom color profile if the only encoded color
            // space information is a "neutral" gamma.  This is primarily needed
            // for correctness (see // https://crbug.com/388025081), but may
            // also help with performance (using a slightly more direct
            // `SkSwizzler` instead of `skcms_Transform`).
            return nullptr;
        }
    }

    // Construct a color profile based on `cHRM` and `gAMA` chunks.
    skcms_Matrix3x3 toXYZD50;
    if (got_chrm) {
        if (!ToSkColorSpacePrimaries(chrm).toXYZD50(&toXYZD50)) {
            return nullptr;
        }
    } else {
        // `blink::PNGImageDecoder` returns a null color profile when `gAMA` is
        // present without `cHRM`.  We fall back to the sRGB profile instead
        // because we do gamma correction via `skcms_Transform` (rather than
        // relying on `libpng` gamma correction as the legacy Blink decoder does
        // in this scenario).
        toXYZD50 = SkNamedGamut::kSRGB;
    }

    skcms_TransferFunction fn;
    fn.a = 1.0f;
    fn.b = fn.c = fn.d = fn.e = fn.f = 0.0f;
    fn.g = 1.0f / gamma;

    return SkCodecs::ColorProfile::Make(fn, toXYZD50);
}

// Returns `nullopt` when input errors are encountered.
std::optional<SkEncodedInfo> CreateEncodedInfo(const rust_png::Reader& reader) {
    rust_png::ColorType rustColor = reader.output_color_type();
    SkEncodedInfo::Color skColor = ToColor(rustColor, reader);

    std::unique_ptr<SkCodecs::ColorProfile> profile = CreateColorProfile(reader);
    if (!SkPngCodecBase::isCompatibleColorProfileAndType(profile.get(), skColor)) {
        profile = nullptr;
    }

    skhdr::Metadata hdrMetadata;
    {
        rust_png::MasteringDisplayColorVolume rust_mdcv;
        if (reader.try_get_mdcv_chunk(rust_mdcv)) {
            hdrMetadata.setMasteringDisplayColorVolume(ToSkMDCV(rust_mdcv));
        }
        rust_png::ContentLightLevelInfo rust_clli;
        if (reader.try_get_clli_chunk(rust_clli)) {
            hdrMetadata.setContentLightLevelInformation(ToSkCLLI(rust_clli));
        }
    }

    // Protect against large PNGs. See http://bugzil.la/251381 for more details.
    constexpr uint32_t kMaxPNGSize = 1000000;
    if ((reader.width() > kMaxPNGSize) || (reader.height() > kMaxPNGSize)) {
        return std::nullopt;
    }
    // We checked image dimensions above, so here we can just assert that casts
    // from `uint32_t` to `int` work ok.
    //
    // We don't use a saturating cast, because this could invalidate `fcTL`
    // checks done within the `png` crate.  For example, the new / truncated
    // image width could end up smaller than `fcTL.frameWidth`.
    SkSafeMath safe;
    int width = safe.castTo<int>(reader.width());
    int height = safe.castTo<int>(reader.height());
    SkASSERT_RELEASE(safe.ok());

    return SkEncodedInfo::Make(width,
                               height,
                               skColor,
                               ToAlpha(rustColor, reader),
                               reader.output_bits_per_component(), // bitsPerComponent
                               reader.output_bits_per_component(), // colorDepth
                               std::move(profile),
                               hdrMetadata);
}

SkCodec::Result ToSkCodecResult(rust_png::DecodingResult rustResult) {
    switch (rustResult) {
        case rust_png::DecodingResult::Success:
        case rust_png::DecodingResult::EndOfFrame:
            return SkCodec::kSuccess;
        case rust_png::DecodingResult::FormatError:
            return SkCodec::kErrorInInput;
        case rust_png::DecodingResult::ParameterError:
            return SkCodec::kInvalidParameters;
        case rust_png::DecodingResult::OtherIoError:
        case rust_png::DecodingResult::LimitsExceededError:
            return SkCodec::kInternalError;
        case rust_png::DecodingResult::IncompleteInput:
            return SkCodec::kIncompleteInput;
    }
    SK_ABORT("Unexpected `rust_png::DecodingResult`: %d", static_cast<int>(rustResult));
}

void blendRow(SkSpan<uint8_t> dstRow,
              SkSpan<const uint8_t> srcRow,
              SkColorType color,
              SkAlphaType alpha) {
    SkASSERT_RELEASE(dstRow.size() >= srcRow.size());
    SkRasterPipeline_<256> p;

    SkRasterPipelineContexts::MemoryCtx dstCtx = {dstRow.data(), 0};
    p.appendLoadDst(color, &dstCtx);
    if (kUnpremul_SkAlphaType == alpha) {
        p.append(SkRasterPipelineOp::premul_dst);
    }

    SkRasterPipelineContexts::MemoryCtx srcCtx = {
        const_cast<void*>(static_cast<const void*>(srcRow.data())),
        0,
    };
    p.appendLoad(color, &srcCtx);
    if (kUnpremul_SkAlphaType == alpha) {
        p.append(SkRasterPipelineOp::premul);
    }

    p.append(SkRasterPipelineOp::srcover);

    if (kUnpremul_SkAlphaType == alpha) {
        p.append(SkRasterPipelineOp::unpremul);
    }
    p.appendStore(color, &dstCtx);

    SkSafeMath safe;
    size_t bpp = safe.castTo<size_t>(SkColorTypeBytesPerPixel(color));
    SkASSERT_RELEASE(safe.ok());

    size_t width = srcRow.size() / bpp;
    p.run(0, 0, width, 1);
}

void blendAllRows(SkSpan<uint8_t> dstFrame,
                  SkSpan<const uint8_t> srcFrame,
                  size_t rowSize,
                  size_t rowStride,
                  SkColorType color,
                  SkAlphaType alpha) {
    while (srcFrame.size() >= rowSize) {
        blendRow(dstFrame, srcFrame.first(rowSize), color, alpha);

        srcFrame = srcFrame.subspan(std::min(rowStride, srcFrame.size()));
        dstFrame = dstFrame.subspan(std::min(rowStride, dstFrame.size()));
    }
}

SkEncodedOrigin GetEncodedOrigin(const rust_png::Reader& reader) {
    if (reader.has_exif_chunk()) {
        rust::Slice<const uint8_t> rust_slice = reader.get_exif_chunk();
        SkEncodedOrigin origin;
        if (SkParseEncodedOrigin(rust_slice.data(), rust_slice.size(), &origin)) {
            return origin;
        }
    }

    return kTopLeft_SkEncodedOrigin;
}

bool IsValidFctlIfAny(const SkEncodedInfo& imageInfo,
                      const rust_png::Reader& reader) {
    // Enforce that if an `fcTL` appears before an `IDAT` chunk, then it has the
    // same dimensions as the ones in the earlier `IHDR` chunk.  This
    // corresponds to the restrictions that the spec at
    // https://www.w3.org/TR/png-3/#fcTL-chunk places on "fcTL chunk
    // corresponding to the default image".
    //
    // Doing this check here is more robust than doing it inside
    // `setFrameInfoFromCurrentFctlChunk`, because the code elsewhere in
    // `SkPngRustCodec` may ignore the `fcTL` chunk (e.g. if
    // `idatIsNotPartOfAnimation` and/or if `acTL` chunk is missing).
    //
    // Reporting a hard error is more robust than trying to ignore the `fcTL`
    // chunk and falling back to decoding a static image, because such a
    // fallback would risk discrepancies between dimensions used in different
    // layers of the software stack (see https://crbug.com/428205250).
    //
    // This check is kind of a defense-in-depth - in the long-term the
    // dimensions should be checked in the Rust `png` crate itself
    // (see https://github.com/image-rs/image-png/pull/614 which hasn't yet been
    // released in a new crates.io version).
    if (reader.has_fctl_chunk()) {
        uint32_t width = 0;
        uint32_t height = 0;
        uint32_t xOffset = 0;
        uint32_t yOffset = 0;
        auto ignoredDisposeOp = rust_png::DisposeOp::None;
        auto ignoredBlendOp = rust_png::BlendOp::Source;
        uint32_t ignoredDurationMs = 0;
        reader.get_fctl_info(width, height, xOffset, yOffset,
                             ignoredDisposeOp, ignoredBlendOp, ignoredDurationMs);

        SkSafeMath safe;
        int frameWidth = safe.castTo<int>(width);
        int frameHeight = safe.castTo<int>(height);
        if (!safe.ok()) {
            return false;
        }

        if (xOffset != 0 || frameWidth != imageInfo.width() ||
                yOffset != 0 || frameHeight != imageInfo.height()) {
            return false;
        }
    }

    return true;
}

void ReadUnknownChunks(const rust_png::Reader& reader, SkPngCompositeChunkReader* chunkReader) {
    if (!chunkReader) {
        return;
    }
    size_t unknownChunksCount = reader.get_unknown_chunks_count();
    for (size_t i = 0; i < unknownChunksCount; ++i) {
        auto name = reader.get_unknown_chunk_name(i);
        auto data = reader.get_unknown_chunk_data(i);
        std::string tag(reinterpret_cast<const char*>(name.data()), 4);
        chunkReader->readChunk(tag.c_str(), data.data(), data.size());
    }
}

}  // namespace

// static
std::unique_ptr<SkPngRustCodec> SkPngRustCodec::MakeFromStream(std::unique_ptr<SkStream> stream,
                                                               Result* result,
                                                               SkPngChunkReader* chunkReader) {
    SkASSERT_RELEASE(stream);
    SkASSERT_RELEASE(result);

    auto inputAdapter = std::make_unique<rust::stream::SkStreamAdapter>(stream.get());
    rust::Box<rust_png::ResultOfReader> resultOfReader =
            rust_png::new_reader(std::move(inputAdapter));
    *result = ToSkCodecResult(resultOfReader->err());
    if (*result != kSuccess) {
        return nullptr;
    }
    rust::Box<rust_png::Reader> reader = resultOfReader->unwrap();

    std::optional<SkEncodedInfo> maybeImageInfo = CreateEncodedInfo(*reader);
    if (!maybeImageInfo.has_value()) {
        *result = kErrorInInput;
        return nullptr;
    }
    SkEncodedInfo& imageInfo = *maybeImageInfo;

    if (!IsValidFctlIfAny(imageInfo, *reader)) {
        *result = kErrorInInput;
        return nullptr;
    }

    auto compositeReader = sk_make_sp<SkPngCompositeChunkReader>(chunkReader);
    ReadUnknownChunks(*reader, compositeReader.get());
    std::unique_ptr<SkStream> gainmapStream = compositeReader->takeGainmapStream();
    std::optional<SkGainmapInfo> gainmapInfo = compositeReader->getGainmapInfo();

    return std::make_unique<SkPngRustCodec>(std::move(imageInfo),
                                            std::move(stream),
                                            std::move(reader),
                                            std::move(compositeReader),
                                            std::move(gainmapStream),
                                            std::move(gainmapInfo));
}

SkPngRustCodec::SkPngRustCodec(SkEncodedInfo&& encodedInfo,
                               std::unique_ptr<SkStream> stream,
                               rust::Box<rust_png::Reader> reader,
                               sk_sp<SkPngCompositeChunkReader> chunkReader,
                               std::unique_ptr<SkStream> gainmapStream,
                               std::optional<SkGainmapInfo> gainmapInfo)
        : SkPngCodecBase(std::move(encodedInfo),
                         // TODO(https://crbug.com/370522089): If/when `SkCodec` can
                         // avoid unnecessary rewinding, then stop "hiding" our stream
                         // from it.
                         /* stream = */ nullptr,
                         GetEncodedOrigin(*reader),
                         std::move(chunkReader),
                         std::move(gainmapStream),
                         gainmapInfo)
        , fReader(std::move(reader))
        , fPrivStream(std::move(stream))
        , fFrameHolder(encodedInfo.width(), encodedInfo.height()) {
    SkASSERT_RELEASE(fPrivStream);

    bool idatIsNotPartOfAnimation = fReader->has_actl_chunk() && !fReader->has_fctl_chunk();
    fFrameAtCurrentStreamPosition = idatIsNotPartOfAnimation ? -1 : 0;
    fStreamIsPositionedAtStartOfFrameData = true;
    if (!idatIsNotPartOfAnimation) {
        // This `appendNewFrame` call should always succeed because:
        // * `fFrameHolder.size()` is 0 at this point
        // * Width and height are already capped when calling `SkEncodedInfo::Make`
        // * `!fReader->has_fctl_chunk()` means that we don't need to worry
        //   about validating other frame metadata.
        Result result = fFrameHolder.appendNewFrame(*fReader, this->getEncodedInfo());
        SkASSERT_RELEASE(result == kSuccess);
    }
}

SkPngRustCodec::~SkPngRustCodec() = default;

SkCodec::Result SkPngRustCodec::readToStartOfNextFrame() {
    SkASSERT_RELEASE(fFrameAtCurrentStreamPosition < this->getRawFrameCount());
    Result result = ToSkCodecResult(fReader->next_frame_info());
    if (result != kSuccess) {
        fStreamIsPositionedAtStartOfFrameData = false;
        return result;
    }

    fStreamIsPositionedAtStartOfFrameData = true;
    fFrameAtCurrentStreamPosition++;
    if (fFrameAtCurrentStreamPosition == fFrameHolder.size()) {
        result = fFrameHolder.appendNewFrame(*fReader, this->getEncodedInfo());
    }

    return result;
}

SkCodec::Result SkPngRustCodec::seekToStartOfFrame(int index) {
    // Callers of this `private` method should provide a valid `index`.
    //
    // `index == fFrameHolder.size()` means that we are seeking to the next
    // frame (i.e. to the first frame for which an `fcTL` chunk wasn't parsed
    // yet).
    SkASSERT_RELEASE((0 <= index) && (index <= fFrameHolder.size()));

    // TODO(https://crbug.com/371060427): Improve runtime performance by seeking
    // directly to the right offset in the stream, rather than calling `rewind`
    // here and moving one-frame-at-a-time via `readToStartOfNextFrame` below.
    if ((index < fFrameAtCurrentStreamPosition) ||
        (index == fFrameAtCurrentStreamPosition && !fStreamIsPositionedAtStartOfFrameData)) {
        if (!fPrivStream->rewind()) {
            return kCouldNotRewind;
        }

        auto inputAdapter =
                std::make_unique<rust::stream::SkStreamAdapter>(fPrivStream.get());
        rust::Box<rust_png::ResultOfReader> resultOfReader =
                rust_png::new_reader(std::move(inputAdapter));

        // `SkPngRustCodec` constructor must have run before, and the
        // constructor got a successfully created reader - we therefore also
        // expect success here.
        SkASSERT_RELEASE(kSuccess == ToSkCodecResult(resultOfReader->err()));
        fReader = resultOfReader->unwrap();
        this->processUnknownChunks();

        bool idatIsNotPartOfAnimation = fReader->has_actl_chunk() && !fReader->has_fctl_chunk();
        fFrameAtCurrentStreamPosition = idatIsNotPartOfAnimation ? -1 : 0;
        fStreamIsPositionedAtStartOfFrameData = true;
    }
    while (fFrameAtCurrentStreamPosition < index) {
        Result result = this->readToStartOfNextFrame();
        if (result != kSuccess) {
            return result;
        }
    }

    return kSuccess;
}

int SkPngRustCodec::getRawFrameCount() const {
    if (!fReader->has_actl_chunk()) {
        return 1;
    }

    static_assert(sizeof(int) >= sizeof(int32_t), "Is it ok to use Sk64_pin_to_s32 below?");
    uint32_t num_frames = fReader->get_actl_num_frames();
    return Sk64_pin_to_s32(num_frames);
}

SkCodec::Result SkPngRustCodec::parseAdditionalFrameInfos() {
    while (fFrameHolder.size() < this->getRawFrameCount()) {
        int oldFrameCount = fFrameHolder.size();

        Result result = this->seekToStartOfFrame(fFrameHolder.size());
        if (result != kSuccess) {
            return result;
        }
        SkASSERT_RELEASE(fFrameHolder.size() == (oldFrameCount + 1));
    }
    return kSuccess;
}

void SkPngRustCodec::getSubsetOrSampleFromFullImage(SkSpan<const uint8_t> fullImageBuffer,
                                                    SkSpan<uint8_t> dst,
                                                    size_t dstRowStride,
                                                    size_t offset,
                                                    int srcHeight,
                                                    int dstHeight,
                                                    size_t dstRowSize) {
    // This only needs to be used in the case of interlaced images that need a subset
    // or sampling, otherwise we can decode row by row.
    SkASSERT_RELEASE(fReader->interlaced());
    SkASSERT_RELEASE(this->options().fSubset || this->isSampling());

    // Verify we have enough source rows
    SkASSERT_RELEASE(fullImageBuffer.size() >= offset);  // Preventing subtraction underflow.
    SkASSERT_RELEASE(srcHeight >= 0);  // Preventing `static_cast<size_t>` overflow
    const size_t encodedRowBytes = this->getEncodedRowBytes();
    size_t fullRowsInFullImageBuffer = (fullImageBuffer.size() - offset) / encodedRowBytes;
    SkASSERT_RELEASE(fullRowsInFullImageBuffer >= static_cast<size_t>(srcHeight));

    // We want the whole row and applyXformRow does the rest, so only offset to correct y value.
    fullImageBuffer = fullImageBuffer.subspan(offset);

    int rowsWritten = 0;
    for (int y = 0; y < srcHeight && rowsWritten < dstHeight; ++y) {
        SkSpan<const uint8_t> srcRow = fullImageBuffer.first(encodedRowBytes);
        fullImageBuffer = fullImageBuffer.subspan(encodedRowBytes);

        // Skip rows not needed by the vertical sampler
        if (this->swizzler() && !this->swizzler()->rowNeeded(y)) {
            continue;
        }

        SkSpan<uint8_t> dstRow = dst.first(dstRowSize);
        // Safe advancement: robust if fDst lacks stride padding on the last row.
        dst = dst.subspan(std::min(dstRowStride, dst.size()));

        // Copy the source row into the correct position in the destination.
        this->applyXformRow(dstRow, srcRow);
        rowsWritten++;
    }
}

SkCodec::Result SkPngRustCodec::startDecoding(const SkImageInfo& dstInfo,
                                              void* pixels,
                                              size_t rowBytes,
                                              const Options& options,
                                              std::optional<DecodingState>& decodingState) {
    if (options.fFrameIndex == 0 && fFrameHolder.size() == 0) {
        Result result = this->seekToStartOfFrame(0);
        if (result != kSuccess) {
            return result;
        }
    }
    if (options.fFrameIndex < 0 || options.fFrameIndex >= fFrameHolder.size()) {
        return kInvalidParameters;
    }
    const SkFrame* frame = fFrameHolder.getFrame(options.fFrameIndex);
    SkASSERT_RELEASE(frame);

    // TODO(https://crbug.com/362830091): Consider handling `fSubset` for APNG.
    if (options.fSubset && !this->supportsSamplingOrSubsetting()) {
        return kUnimplemented;
    }

    Result result = this->seekToStartOfFrame(options.fFrameIndex);
    if (result != kSuccess) {
        return result;
    }

    // https://www.w3.org/TR/png-3/#11PLTE says that for color type 3
    // (indexed-color), the PLTE chunk is required. OTOH, `Codec_InvalidImages`
    // expects that we will succeed in this case and produce *some* output.
    //
    // This check needs to happen after `seekToStartOfFrame`, because the act
    // of seeking-and/or-rewinding may reset whether `fReader` has encountered
    // an PLTE chunk yet or not.
    if (this->getEncodedInfo().color() == SkEncodedInfo::kPalette_Color &&
        !fReader->has_plte_chunk()) {
        return kInvalidInput;
    }

    result = this->initializeXforms(dstInfo, options, frame->width());
    if (result != kSuccess) {
        return result;
    }

    SkSafeMath safe;
    DecodingState state;
    state.fDst = UninitializedDstInfo{
            .fDstPtr = static_cast<uint8_t*>(pixels),
            .fDstRowStride = rowBytes,
    };
    if (options.fSubset) {
        state.fYByteOffset = safe.mul(safe.castTo<size_t>(this->getEncodedRowBytes()),
                                      safe.castTo<size_t>(options.fSubset->top()));
        state.fFirstRow = options.fSubset->top();
        state.fLastRow = options.fSubset->bottom() - 1;
    } else {
        state.fYByteOffset = 0;
        state.fFirstRow = 0;
        state.fLastRow = frame->height() - 1;
    }

    if (!safe.ok()) {
        return kErrorInInput;
    }

    decodingState = std::move(state);
    return kSuccess;
}

SkCodec::Result SkPngRustCodec::initializeDecodingDst(DecodingState& decodingState) {
    if (std::holds_alternative<DecodingDstInfo>(decodingState.fDst)) {
        return kSuccess;
    }

    const SkFrame* frame = fFrameHolder.getFrame(this->options().fFrameIndex);
    SkASSERT_RELEASE(frame);

    const UninitializedDstInfo uninit = std::get<UninitializedDstInfo>(decodingState.fDst);

    SkSafeMath safe;
    uint8_t dstBytesPerPixel = safe.castTo<uint8_t>(this->dstInfo().bytesPerPixel());
    if (dstBytesPerPixel >= 32u) {
        return kInvalidParameters;
    }
    size_t dstRowCount = safe.castTo<size_t>(frame->height());

    size_t dstRowSize = 0;
    if (!this->canReadRow()) {
        this->initializeXformParams();
        dstRowSize = this->getDstRowSize();
    } else {
        size_t frameWidth = safe.castTo<size_t>(frame->width());
        dstRowSize = safe.mul(dstBytesPerPixel, frameWidth);
        if (!safe.ok()) {
            return kErrorInInput;
        }
    }

    SkSpan<uint8_t> dstSpan;
    if (this->isSampling() || this->options().fSubset) {
        if (!this->supportsSamplingOrSubsetting()) {
            return kInvalidParameters;
        }
        if (uninit.fDstRowStride < dstRowSize) {
            return kInvalidParameters;
        }
        const int srcHeight = decodingState.fLastRow - decodingState.fFirstRow + 1;
        const int sampleY = this->swizzler() ? this->swizzler()->sampleY() : 1;
        const size_t dstHeight =
                safe.castTo<size_t>(SkCodecPriv::GetSampledDimension(srcHeight, sampleY));
        const size_t dstBufferSize = safe.mul(dstHeight, uninit.fDstRowStride);
        if (!safe.ok()) {
            return kErrorInInput;
        }
        dstSpan = SkSpan(uninit.fDstPtr, dstBufferSize);
    } else {
        if (uninit.fDstRowStride < this->dstInfo().minRowBytes()) {
            return kInvalidParameters;
        }
        size_t imageHeight = safe.castTo<size_t>(this->dstInfo().height());
        size_t imageSize = safe.mul(uninit.fDstRowStride, imageHeight);
        size_t xPixelOffset = safe.castTo<size_t>(frame->xOffset());
        size_t xByteOffsetFrame = safe.mul(dstBytesPerPixel, xPixelOffset);
        size_t yPixelOffset = safe.castTo<size_t>(frame->yOffset());
        size_t yByteOffsetFrame = safe.mul(uninit.fDstRowStride, yPixelOffset);
        size_t frameHeightTimesRowStride =
                safe.mul(dstRowCount, uninit.fDstRowStride);
        if (!safe.ok()) {
            return kErrorInInput;
        }

        dstSpan = SkSpan(uninit.fDstPtr, imageSize)
                          .subspan(xByteOffsetFrame)
                          .subspan(yByteOffsetFrame);
        if (frameHeightTimesRowStride < dstSpan.size()) {
            dstSpan = dstSpan.first(frameHeightTimesRowStride);
        }
    }

    size_t preblendBufferSize = 0;
    if (fReader->interlaced()) {
        // Use fPreblendBuffer to decode whole image untransformed, then truncate/xform later.
        if (this->options().fSubset || this->isSampling()) {
            preblendBufferSize =
                    safe.mul(this->getEncodedRowBytes(),
                             safe.castTo<size_t>(this->getEncodedInfo().height()));
            if (!safe.ok()) {
                return kErrorInInput;
            }
        } else if (frame->getBlend() == SkCodecAnimation::Blend::kSrcOver) {
            preblendBufferSize = dstSpan.size();
        }
    } else if (frame->getBlend() == SkCodecAnimation::Blend::kSrcOver) {
        preblendBufferSize = dstRowSize;
    }

    if (Result result = decodingState.fPreblendBuffer.allocateFromBudget(this, preblendBufferSize);
        result != kSuccess) {
        return result;
    }

    if (fReader->interlaced() && !this->options().fSubset && !this->isSampling()) {
        if (Result result = decodingState.fDecodedInterlacedFullWidthRow.allocateFromBudget(
                    this, this->getEncodedRowBytes());
            result != kSuccess) {
            return result;
        }
        if (!this->canReadRow()) {
            if (Result result =
                        decodingState.fXformedInterlacedRow.allocateFromBudget(this, dstRowSize);
                result != kSuccess) {
                return result;
            }
        }
    }

    decodingState.fDst = DecodingDstInfo{
            .fDstRowStride = uninit.fDstRowStride,
            .fDst = dstSpan,
            .fDstRowSize = dstRowSize,
            .fDstRowCount = dstRowCount,
            .fDstBytesPerPixel = dstBytesPerPixel,
    };
    return kSuccess;
}

void SkPngRustCodec::expandDecodedInterlacedRow(DecodingState& decodingState,
                                                SkSpan<uint8_t> dstFrame,
                                                SkSpan<const uint8_t> srcRow,
                                                const DecodingDstInfo& decodingDst,
                                                bool xFormNeeded) {
    SkASSERT_RELEASE(fReader->interlaced());
    const size_t dstRowStride = decodingDst.fDstRowStride;
    if (xFormNeeded) {
        // Copy (potentially shorter for initial Adam7 passes) `srcRow` into a
        // full-frame-width `fDecodedInterlacedFullWidthRow`. This is needed because
        // `applyXformRow` requires full-width rows as input (can't change
        // `SkSwizzler::fSrcWidth` after `initializeXforms`).
        SkSpan<uint8_t> decodedInterlacedFullWidthRow =
                decodingState.fDecodedInterlacedFullWidthRow.span();
        SkASSERT_RELEASE(decodedInterlacedFullWidthRow.size() == this->getEncodedRowBytes());
        SkASSERT_RELEASE(decodedInterlacedFullWidthRow.size() >= srcRow.size());
        memcpy(decodedInterlacedFullWidthRow.data(), srcRow.data(), srcRow.size());
        memset(decodedInterlacedFullWidthRow.data() + srcRow.size(),
               0,
               decodedInterlacedFullWidthRow.size() - srcRow.size());

        // With `kYes_ZeroInitialized`, `fSwizzler` skips writing leading transparent pixels on the
        // assumption that the destination is already zeroed, so `fXformedInterlacedRow` must be
        // re-zeroed before each row. Otherwise, every byte of the row gets overwritten.
        SkSpan<uint8_t> xformedInterlacedRow = decodingState.fXformedInterlacedRow.span();
        SkASSERT_RELEASE(xformedInterlacedRow.size() == decodingDst.fDstRowSize);
        if (this->options().fZeroInitialized == kYes_ZeroInitialized) {
            memset(xformedInterlacedRow.data(), 0, xformedInterlacedRow.size());
        }
        this->applyXformRow(xformedInterlacedRow, decodedInterlacedFullWidthRow);
    }

    const uint8_t dstBytesPerPixel = decodingDst.fDstBytesPerPixel;
    SkASSERT_RELEASE(dstBytesPerPixel < 32u);  // Checked in `initializeDecodingDst`.

    rust::Slice<const uint8_t> expandedRow =
            xFormNeeded ? rust::Slice<const uint8_t>(decodingState.fXformedInterlacedRow)
                        : rust::Slice<const uint8_t>(srcRow);
    const uint8_t bitsPerPixel =
            xFormNeeded ? dstBytesPerPixel * 8u : this->getEncodedInfo().bitsPerPixel();

    // `splat_last_interlaced_row` infers the number of rows of the frame from
    // `dstFrame.size() / dstRowStride`, so it can only be used when `dstFrame`
    // spans whole rows for the full height of the frame.  This is not the case
    // when an APNG frame is offset horizontally - the tail of `dstFrame` then
    // stops short of a whole row.
    SkSafeMath safe;
    size_t wholeFrameSize = safe.mul(dstRowStride, decodingDst.fDstRowCount);
    if (safe.ok() && dstFrame.size() >= wholeFrameSize) {
        SkSpan<uint8_t> wholeFrame = dstFrame.first(wholeFrameSize);
        // Replicating each decoded sample over the pixels that later Adam7
        // passes will overwrite means that `dstFrame` holds a complete (if
        // blocky) image after each pass, rather than a mostly-transparent one.
        // This matters because partially decoded images are presented to the
        // user.  The final image is not affected.
        fReader->splat_last_interlaced_row(
                rust::Slice<uint8_t>(wholeFrame), dstRowStride, expandedRow, bitsPerPixel);
    } else {
        fReader->expand_last_interlaced_row(
                rust::Slice<uint8_t>(dstFrame), dstRowStride, expandedRow, bitsPerPixel);
    }
}

bool SkPngRustCodec::isLastFrame() {
    return (this->isAnimated() == IsAnimated::kNo) ||
           (this->options().fFrameIndex + 1 == this->getRawFrameCount());
}

bool SkPngRustCodec::isSampling() const {
    return this->swizzler() && (this->swizzler()->sampleX() > 1 || this->swizzler()->sampleY() > 1);
}

bool SkPngRustCodec::supportsSamplingOrSubsetting() {
    if (this->isAnimated() != IsAnimated::kNo || fFrameHolder.size() == 0) {
        return false;
    }
    const SkFrame* frame = fFrameHolder.getFrame(0);
    return frame && frame->frameRect() == this->bounds();
}

void SkPngRustCodec::finishStreamAfterLastFrame() { std::ignore = fReader->finish_decoding(); }

// Given the dstInfo and the rust colortype/bits per component, determines if we
// can use rust_png::Reader::read_row to decode directly into dst.
bool SkPngRustCodec::canReadRow() {
    // Check alpha types
    if (this->dstInfo().alphaType() != kUnpremul_SkAlphaType) {
        return false;
    }
    // We cannot decode directly when subsetting or sub-sampling.
    if (this->options().fSubset || this->isSampling()) {
        return false;
    }

    // Check color types
    rust_png::ColorType color_type = fReader->output_color_type();
    uint8_t bits_per_component =  fReader->output_bits_per_component();
    switch (this->dstInfo().colorType()) {
        case kRGBA_8888_SkColorType:
            if (color_type != rust_png::ColorType::Rgba
                || bits_per_component != 8) {
                return false;
            }
            break;
        case kGray_8_SkColorType:
            if (color_type != rust_png::ColorType::Grayscale
                || bits_per_component != 8) {
                return false;
            }
            break;
        default:
            return false;
    }

    // Check profiles
    if (!!this->getEncodedInfo().profile() != !!this->dstInfo().colorSpace()) {
        return false;
    }
    if (this->getEncodedInfo().profile()) {
        SkASSERT_RELEASE(this->dstInfo().colorSpace());
        skcms_ICCProfile dstProfile;
        this->dstInfo().colorSpace()->toProfile(&dstProfile);
        if (!skcms_ApproximatelyEqualProfiles(this->getEncodedInfo().profile(), &dstProfile)) {
            return false;
        }
    }

    return true;
}

SkCodec::Result SkPngRustCodec::incrementalDecodeXForm(DecodingState& decodingState) {
    SkASSERT_RELEASE(!this->canReadRow());

    const bool isSampling = this->isSampling();
    const bool interlaced = fReader->interlaced();
    const bool subset = this->options().fSubset;
    DecodingDstInfo& decodingDst = decodingState.dst();

    const int srcHeight = decodingState.fLastRow - decodingState.fFirstRow + 1;
    const int sampleY = this->swizzler() ? this->swizzler()->sampleY() : 1;
    const int dstHeight = SkCodecPriv::GetSampledDimension(srcHeight, sampleY);

    // A non-interlaced partial decode stops once it has written the rows it needs, leaving the
    // rest of the frame (and everything after it) unread. Interlaced decodes always walk the whole
    // frame, because every Adam7 pass has to be read before any row is final.
    const bool partialDecode = subset || isSampling;
    const bool stopsBeforeEndOfFrame = partialDecode && !interlaced;

    while (true) {
        rust::Slice<const uint8_t> decodedRow;
        fStreamIsPositionedAtStartOfFrameData = false;
        if (!stopsBeforeEndOfFrame || decodingState.fRowsWrittenToOutput < dstHeight) {
            Result result = ToSkCodecResult(fReader->next_interlaced_row(decodedRow));
            if (result != kSuccess) {
                return result;
            }
        }

        // This is how FFI layer says "no more rows", or if we have decoded all the rows
        // we need for our subset or sampling.
        if (decodedRow.empty()) {
            if (interlaced && !decodingState.fPreblendBuffer.empty()) {
                if (partialDecode) {
                    this->getSubsetOrSampleFromFullImage(
                            SkSpan<uint8_t>(decodingState.fPreblendBuffer),
                            decodingDst.fDst,
                            decodingDst.fDstRowStride,
                            decodingState.fYByteOffset,
                            srcHeight,
                            dstHeight,
                            decodingDst.fDstRowSize);
                } else {
                    blendAllRows(decodingDst.fDst,
                                 decodingState.fPreblendBuffer,
                                 decodingDst.fDstRowSize,
                                 decodingDst.fDstRowStride,
                                 this->dstInfo().colorType(),
                                 this->dstInfo().alphaType());
                }
            }
            if (!interlaced && !partialDecode) {
                // All of the original `fDst` should be filled out at this point.
                SkASSERT_RELEASE(decodingDst.fDst.empty());
            }

            if (!stopsBeforeEndOfFrame) {
                // `static_cast` is ok, because `startDecoding` already validated `fFrameIndex`.
                fFrameHolder.markFrameAsFullyReceived(
                        static_cast<size_t>(this->options().fFrameIndex));
            }
            if (!partialDecode && this->isLastFrame()) {
                this->finishStreamAfterLastFrame();
            }
            return kSuccess;
        }

        if (interlaced) {
            if (decodingState.fPreblendBuffer.empty()) {
                this->expandDecodedInterlacedRow(decodingState,
                                                 decodingDst.fDst,
                                                 decodedRow,
                                                 decodingDst,
                                                 /*xFormNeeded=*/true);
            } else {
                if (partialDecode) {
                    SkSafeMath safe;
                    uint8_t encodedBytesPerPixel = safe.castTo<uint8_t>(this->getEncodedInfo()
                                                      .makeImageInfo()
                                                      .bytesPerPixel());
                    SkASSERT_RELEASE(safe.ok());               // Checked in `startDecoding`.
                    DecodingDstInfo fullImageDecodingDst = {
                            .fDstRowStride = this->getEncodedRowBytes(),
                            .fDst = decodingState.fPreblendBuffer,
                            .fDstRowSize = this->getEncodedRowBytes(),
                            .fDstRowCount = static_cast<size_t>(
                                    this->getEncodedInfo().height()),
                            .fDstBytesPerPixel = encodedBytesPerPixel};
                    this->expandDecodedInterlacedRow(decodingState,
                                                     decodingState.fPreblendBuffer,
                                                     decodedRow,
                                                     fullImageDecodingDst,
                                                     /*xFormNeeded=*/false);
                } else {
                    this->expandDecodedInterlacedRow(decodingState,
                                                     decodingState.fPreblendBuffer,
                                                     decodedRow,
                                                     decodingDst,
                                                     /*xFormNeeded=*/true);
                }
            }
            // `fRowsWrittenToOutput` is not incremented, because full, contiguous rows
            // are not decoded until pass 6 (or 7 depending on how you look) of
            // Adam7 interlacing scheme.
        } else {
            int y = decodingState.fCurrentSourceRow - decodingState.fFirstRow;
            decodingState.fCurrentSourceRow++;
            if (y < 0) {
                continue;
            }
            if (this->swizzler() && !this->swizzler()->rowNeeded(y)) {
                continue;
            }

            if (decodingState.fPreblendBuffer.empty()) {
                this->applyXformRow(decodingDst.fDst, decodedRow);
            } else {
                // Non-empty `fPreblendBuffer` implies 2nd or later frame of
                // a APNG.  But `SkPngRustCodec::getSampler` returns nullptr
                // for APNGs.  So we don't need to worry about sub-sampling
                // here (maybe it works, but when enabling it we should be
                // careful, add tests, etc.).  Note that `this->swizzler()`
                // may be non-null when no sub-sampling happens, because we
                // may still need to "swizzle" between color formats.
                SkASSERT_RELEASE(!isSampling);
                this->applyXformRow(decodingState.fPreblendBuffer, decodedRow);
                blendRow(decodingDst.fDst,
                         decodingState.fPreblendBuffer,
                         this->dstInfo().colorType(),
                         this->dstInfo().alphaType());
            }

            decodingDst.fDst = decodingDst.fDst.subspan(
                    std::min(decodingDst.fDstRowStride, decodingDst.fDst.size()));
            decodingState.fRowsWrittenToOutput++;
        }
    }
}

SkCodec::Result SkPngRustCodec::incrementalDecode(DecodingState& decodingState) {
    SkASSERT_RELEASE(this->canReadRow());
    const bool interlaced = fReader->interlaced();
    rust::Slice<uint8_t> dstSlice;
    if (interlaced) {
        SkSpan<uint8_t> fullWidthRow = decodingState.fDecodedInterlacedFullWidthRow.span();
        SkASSERT_RELEASE(fullWidthRow.size() == this->getEncodedRowBytes());
        memset(fullWidthRow.data(), 0, fullWidthRow.size());
        dstSlice = rust::Slice<uint8_t>(fullWidthRow);
    }
    DecodingDstInfo& decodingDst = decodingState.dst();

    while (true) {
        if (!interlaced) {
            dstSlice = decodingState.fPreblendBuffer.empty()
                          ? rust::Slice<uint8_t>(decodingDst.fDst)
                          : rust::Slice<uint8_t>(decodingState.fPreblendBuffer);
        }
        fStreamIsPositionedAtStartOfFrameData = false;

        rust_png::DecodingResult rustResult = fReader->read_row(dstSlice);
        Result result = ToSkCodecResult(rustResult);

        if (result != kSuccess) {
            return result;
        }

        // No more rows.
        if (rustResult == rust_png::DecodingResult::EndOfFrame) {
            if (interlaced && !decodingState.fPreblendBuffer.empty()) {
                blendAllRows(decodingDst.fDst,
                             decodingState.fPreblendBuffer,
                             decodingDst.fDstRowSize,
                             decodingDst.fDstRowStride,
                             this->dstInfo().colorType(),
                             this->dstInfo().alphaType());
                }

            if (!interlaced) {
                // All of the original `fDst` should be filled out at this point.
                SkASSERT_RELEASE(decodingDst.fDst.empty());
            }
            // `static_cast` is ok, because `startDecoding` already validated `fFrameIndex`.
            fFrameHolder.markFrameAsFullyReceived(static_cast<size_t>(this->options().fFrameIndex));
            if (this->isLastFrame()) {
                // No partial-decode check here: `canReadRow` rules out subsets and sampling.
                this->finishStreamAfterLastFrame();
            }
            return kSuccess;
        }

        // Expand interlaced rows or blend into previous frame if needed.
        if (interlaced) {
            if (decodingState.fPreblendBuffer.empty()) {
                this->expandDecodedInterlacedRow(decodingState,
                                                 decodingDst.fDst,
                                                 dstSlice,
                                                 decodingDst,
                                                 /*xFormNeeded=*/false);
            } else {
                this->expandDecodedInterlacedRow(decodingState,
                                                 decodingState.fPreblendBuffer,
                                                 dstSlice,
                                                 decodingDst,
                                                 /*xFormNeeded=*/false);
            }
            // `fRowsWrittenToOutput` is not incremented, because full, contiguous rows
            // are not decoded until pass 6 (or 7 depending on how you look) of
            // Adam7 interlacing scheme.
        } else {
            if (!decodingState.fPreblendBuffer.empty()) {
                blendRow(decodingDst.fDst,
                         decodingState.fPreblendBuffer,
                         this->dstInfo().colorType(),
                         this->dstInfo().alphaType());
            }
            // Increment our pointer to dst memory.
            decodingDst.fDst = decodingDst.fDst.subspan(
                std::min(decodingDst.fDstRowStride, decodingDst.fDst.size()));
            decodingState.fRowsWrittenToOutput++;
        }
    }
}

void SkPngRustCodec::reportRowsDecoded(Result result,
                                       const DecodingState& decodingState,
                                       int* rowsDecoded) const {
    if (result == kSuccess || !rowsDecoded) {
        return;
    }
    *rowsDecoded = decodingState.fRowsWrittenToOutput;
}

SkCodec::Result SkPngRustCodec::onGetPixels(const SkImageInfo& dstInfo,
                                            void* pixels,
                                            size_t rowBytes,
                                            const Options& options,
                                            int* rowsDecoded) {
    std::optional<DecodingState> decodingState;
    Result result = this->startDecoding(dstInfo, pixels, rowBytes, options, decodingState);
    if (result != kSuccess) {
        return result;
    }

    result = this->initializeDecodingDst(*decodingState);
    if (result == kSuccess) {
        if (this->canReadRow()) {
            result = this->incrementalDecode(*decodingState);
        } else {
            result = this->incrementalDecodeXForm(*decodingState);
        }
    }
    this->reportRowsDecoded(result, *decodingState, rowsDecoded);
    return result;
}

SkCodec::Result SkPngRustCodec::onStartIncrementalDecode(const SkImageInfo& dstInfo,
                                                         void* pixels,
                                                         size_t rowBytes,
                                                         const Options& options) {
    // It is okay if `fIncrementalDecodingState` contains state of another,
    // partially decoded frame - in this case we want to clobber it.
    fIncrementalDecodingState.reset();
    return this->startDecoding(dstInfo, pixels, rowBytes, options, fIncrementalDecodingState);
}

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

    Result result = this->initializeDecodingDst(*fIncrementalDecodingState);
    if (result == kSuccess) {
        if (this->canReadRow()) {
            result = this->incrementalDecode(*fIncrementalDecodingState);
        } else {
            result = this->incrementalDecodeXForm(*fIncrementalDecodingState);
        }
    }
    this->reportRowsDecoded(result, *fIncrementalDecodingState, rowsDecoded);
    if (result != kIncompleteInput) {
        // After successfully reading the whole row (`kSuccess`), and after a
        // fatal error (only recoverable error is `kIncompleteInput`) our client
        // should not call `onIncrementalDecode` again.  This means that the
        // incremental decoding state can be discarded at this point.
        fIncrementalDecodingState.reset();
    }
    return result;
}

int SkPngRustCodec::onGetFrameCount() {
    do {
        if (!fCanParseAdditionalFrameInfos || fIncrementalDecodingState.has_value()) {
            break;
        }

        if (fPrivStream->hasLength()) {
            size_t currentLength = fPrivStream->getLength();
            if (fMaxStreamLengthSeenWhenParsingAdditionalFrameInfos.has_value()) {
                size_t oldLength = *fMaxStreamLengthSeenWhenParsingAdditionalFrameInfos;
                // We use `>=` instead of `==`, because the underlying stream
                // can be "cleared" - see https://crbug.com/431273809#comment4.
                if (oldLength >= currentLength) {
                    // Don't retry `parseAdditionalFrameInfos` if the input
                    // didn't change (or is smaller than last time).
                    break;
                }
            }
            fMaxStreamLengthSeenWhenParsingAdditionalFrameInfos = currentLength;
        }

        switch (this->parseAdditionalFrameInfos()) {
            case kIncompleteInput:
                fCanParseAdditionalFrameInfos = true;
                break;
            case kSuccess:
                SkASSERT_RELEASE(fFrameHolder.size() == this->getRawFrameCount());
                fCanParseAdditionalFrameInfos = false;
                break;
            default:
                fCanParseAdditionalFrameInfos = false;
                break;
        }
    } while (false);

    return fFrameHolder.size();
}

bool SkPngRustCodec::onGetFrameInfo(int index, FrameInfo* info) const {
    return fFrameHolder.getFrameInfo(index, info);
}

int SkPngRustCodec::onGetRepetitionCount() {
    if (!fReader->has_actl_chunk()) {
        return 0;
    }

    uint32_t numFrames = fReader->get_actl_num_frames();
    if (numFrames <= 1) {
        return 0;
    }

    // APNG spec says that "`num_plays` indicates the number of times that this
    // animation should play; if it is 0, the animation should play
    // indefinitely."
    SkSafeMath safe;
    int numPlays = safe.castTo<int>(fReader->get_actl_num_plays());
    if ((numPlays == 0) || !safe.ok()) {
        return kRepetitionCountInfinite;
    }

    // Subtracting 1, because `SkCodec::onGetRepetitionCount` doc comment says
    // that "This number does not include the first play through of each frame.
    // For example, a repetition count of 4 means that each frame is played 5
    // times and then the animation stops."
    return numPlays - 1;
}

SkCodec::IsAnimated SkPngRustCodec::onIsAnimated() {
    if (fReader->has_actl_chunk() && fReader->get_actl_num_frames() > 1) {
        return IsAnimated::kYes;
    }
    return IsAnimated::kNo;
}

std::unique_ptr<SkCodec> SkPngRustCodec::onDecodeGainmap(std::unique_ptr<SkStream> stream,
                                                         SkCodec::Result* result) {
    return SkPngRustCodec::MakeFromStream(std::move(stream), result, fPngChunkReader.get());
}

std::optional<SkSpan<const SkPngCodecBase::PaletteColorEntry>> SkPngRustCodec::onTryGetPlteChunk() {
    if (fReader->output_color_type() != rust_png::ColorType::Indexed) {
        return std::nullopt;
    }

    SkASSERT_RELEASE(fReader->has_plte_chunk());  // Checked in `startDecoding`.
    SkSpan<const uint8_t> bytes = ToSkSpan(fReader->get_plte_chunk());

    // Make sure that `bytes.size()` is a multiple of
    // `sizeof(PaletteColorEntry)`.
    constexpr size_t kEntrySize = sizeof(PaletteColorEntry);
    bytes = bytes.first((bytes.size() / kEntrySize) * kEntrySize);

    // Alignment of `PaletteColorEntry` is 1, because its size is 3, and size
    // has to be a multiple of alignment (every element of an array has to be
    // aligned) + alignment is always a power of 2.  And this means that
    // `bytes.data()` is already aligned.
    static_assert(kEntrySize == 3, "");
    static_assert(std::alignment_of<PaletteColorEntry>::value == 1, "");
    static_assert(std::alignment_of<uint8_t>::value == 1, "");
    SkSpan<const PaletteColorEntry> palette = SkSpan(
            reinterpret_cast<const PaletteColorEntry*>(bytes.data()), bytes.size() / kEntrySize);

    return palette;
}

std::optional<SkSpan<const uint8_t>> SkPngRustCodec::onTryGetTrnsChunk() {
    if (fReader->output_color_type() != rust_png::ColorType::Indexed) {
        return std::nullopt;
    }

    if (!fReader->has_trns_chunk()) {
        return std::nullopt;
    }

    return ToSkSpan(fReader->get_trns_chunk());
}

class SkPngRustCodec::FrameHolder::PngFrame final : public SkFrame {
public:
    PngFrame(int id, SkEncodedInfo::Alpha alpha) : SkFrame(id), fReportedAlpha(alpha) {}

    bool isFullyReceived() const { return fFullyReceived; }
    void markAsFullyReceived() { fFullyReceived = true; }

private:
    SkEncodedInfo::Alpha onReportedAlpha() const override { return fReportedAlpha; }

    const SkEncodedInfo::Alpha fReportedAlpha;
    bool fFullyReceived = false;
};

SkPngRustCodec::FrameHolder::FrameHolder(int width, int height) : SkFrameHolder() {
    fScreenWidth = width;
    fScreenHeight = height;
}

const SkFrameHolder* SkPngRustCodec::getFrameHolder() const { return &fFrameHolder; }

// We cannot use the SkCodec implementation since we pass nullptr to the superclass out of
// an abundance of caution w/r to rewinding the stream.
//
// TODO(https://crbug.com/370522089): See if `SkCodec` can be tweaked to avoid
// the need to hide the stream from it.
sk_sp<const SkData> SkPngRustCodec::getEncodedData() const {
    SkASSERT_RELEASE(fPrivStream);
    sk_sp<const SkData> data = fPrivStream->getData();
    if (data) {
        return data;
    }
    auto dStream = fPrivStream->duplicate();
    if (!dStream->hasLength()) {
        return nullptr;
    }
    return SkData::MakeFromStream(dStream.get(), dStream->getLength());
}

SkPngRustCodec::FrameHolder::~FrameHolder() = default;

const SkFrame* SkPngRustCodec::FrameHolder::onGetFrame(int unverifiedIndex) const {
    SkSafeMath safe;
    size_t index = safe.castTo<size_t>(unverifiedIndex);
    if (safe.ok() && (index < fFrames.size())) {
        return &fFrames[index];
    }
    return nullptr;
}

int SkPngRustCodec::FrameHolder::size() const {
    // This invariant is maintained in `appendNewFrame`.
    SkASSERT_RELEASE(SkTFitsIn<int>(fFrames.size()));
    return static_cast<int>(fFrames.size());
}

void SkPngRustCodec::FrameHolder::markFrameAsFullyReceived(size_t index) {
    SkASSERT_RELEASE(index < fFrames.size());
    fFrames[index].markAsFullyReceived();
}

bool SkPngRustCodec::FrameHolder::getFrameInfo(int index, FrameInfo* info) const {
    const SkFrame* frame = this->getFrame(index);
    if (frame && info) {
        bool isFullyReceived = static_cast<const PngFrame*>(frame)->isFullyReceived();
        frame->fillIn(info, isFullyReceived);
    }
    return !!frame;
}

SkCodec::Result SkPngRustCodec::FrameHolder::appendNewFrame(const rust_png::Reader& reader,
                                                            const SkEncodedInfo& info) {
    // Ensure that `this->size()` fits into an `int`.  `+ 1u` is used to account
    // for `push_back` / `emplace_back` below.
    if (!SkTFitsIn<int>(fFrames.size() + 1u)) {
        return kErrorInInput;
    }
    int id = static_cast<int>(fFrames.size());

    if (reader.has_actl_chunk() && reader.has_fctl_chunk()) {
        if (!fFrames.empty()) {
            // Having `fcTL` for a new frame means that the previous frame has been
            // fully received (since all of the previous frame's `fdAT` / `IDAT`
            // chunks must have come before the new frame's `fcTL` chunk).
            fFrames.back().markAsFullyReceived();
        }

        PngFrame frame(id, info.alpha());
        SkCodec::Result result = this->setFrameInfoFromCurrentFctlChunk(reader, &frame);
        if (result == SkCodec::kSuccess) {
            fFrames.push_back(std::move(frame));
        }
        return result;
    }

    SkASSERT_RELEASE(!reader.has_actl_chunk() || !reader.has_fctl_chunk());
    SkASSERT_RELEASE(id == 0);
    fFrames.emplace_back(id, info.alpha());
    SkFrame& frame = fFrames.back();
    frame.setXYWH(0, 0, info.width(), info.height());
    frame.setBlend(SkCodecAnimation::Blend::kSrc);
    this->setAlphaAndRequiredFrame(&frame);
    return kSuccess;
}

SkCodec::Result SkPngRustCodec::FrameHolder::setFrameInfoFromCurrentFctlChunk(
        const rust_png::Reader& reader, PngFrame* frame) {
    SkASSERT_RELEASE(reader.has_fctl_chunk());  // Caller should guarantee this
    SkASSERT_RELEASE(frame);

    uint32_t width = 0;
    uint32_t height = 0;
    uint32_t xOffset = 0;
    uint32_t yOffset = 0;
    auto disposeOp = rust_png::DisposeOp::None;
    auto blendOp = rust_png::BlendOp::Source;
    uint32_t durationMs = 0;
    reader.get_fctl_info(width, height, xOffset, yOffset, disposeOp, blendOp, durationMs);

    {
        SkSafeMath safe;
        frame->setXYWH(safe.castTo<int>(xOffset),
                       safe.castTo<int>(yOffset),
                       safe.castTo<int>(width),
                       safe.castTo<int>(height));
        frame->setDuration(safe.castTo<int>(durationMs));
        if (!safe.ok()) {
            return kErrorInInput;
        }
    }

    frame->setDisposalMethod(ToDisposalMethod(disposeOp));

    // https://wiki.mozilla.org/APNG_Specification#.60fcTL.60:_The_Frame_Control_Chunk
    // points out that "for the first frame the two blend modes are functionally
    // equivalent" so we use `BlendOp::Source` because it has better performance
    // characteristics.
    if (frame->frameId() == 0) {
        blendOp = rust_png::BlendOp::Source;
    }
    frame->setBlend(ToBlend(blendOp));

    // Note: `setAlphaAndRequiredFrame` needs to be called last, because it
    // depends on the other properties set above.
    this->setAlphaAndRequiredFrame(frame);
    return kSuccess;
}

void SkPngRustCodec::processUnknownChunks() { ReadUnknownChunks(*fReader, fPngChunkReader.get()); }

SkSampler* SkPngRustCodec::getSampler(bool createIfNecessary) {
    if (!this->supportsSamplingOrSubsetting()) {
        if (createIfNecessary) {
            // Clean up tentative decoding state set up by startIncrementalDecode.
            fIncrementalDecodingState.reset();
        }
        return nullptr;
    }
    return SkPngCodecBase::getSampler(createIfNecessary);
}