<commit_before>// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "lib/jxl/dec_patch_dictionary.h"

#include <stdint.h>
#include <stdlib.h>
#include <sys/types.h>

#include <algorithm>
#include <random>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "lib/jxl/ans_params.h"
#include "lib/jxl/base/compiler_specific.h"
#include "lib/jxl/base/override.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/blending.h"
#include "lib/jxl/chroma_from_luma.h"
#include "lib/jxl/color_management.h"
#include "lib/jxl/common.h"
#include "lib/jxl/dec_ans.h"
#include "lib/jxl/dec_frame.h"
#include "lib/jxl/entropy_coder.h"
#include "lib/jxl/frame_header.h"
#include "lib/jxl/image.h"
#include "lib/jxl/image_bundle.h"
#include "lib/jxl/image_ops.h"
#include "lib/jxl/patch_dictionary_internal.h"

namespace jxl {

constexpr int kMaxPatches = 1 << 24;

Status PatchDictionary::Decode(BitReader* br, size_t xsize, size_t ysize,
                               bool* uses_extra_channels) {
  positions_.clear();
  std::vector<uint8_t> context_map;
  ANSCode code;
  JXL_RETURN_IF_ERROR(
      DecodeHistograms(br, kNumPatchDictionaryContexts, &code, &context_map));
  ANSSymbolReader decoder(&code, br);

  auto read_num = [&](size_t context) {
    size_t r = decoder.ReadHybridUint(context, br, context_map);
    return r;
  };

  size_t num_ref_patch = read_num(kNumRefPatchContext);
  // TODO(veluca): does this make sense?
  if (num_ref_patch > kMaxPatches) {
    return JXL_FAILURE("Too many patches in dictionary");
  }

  for (size_t id = 0; id < num_ref_patch; id++) {
    PatchReferencePosition ref_pos;
    ref_pos.ref = read_num(kReferenceFrameContext);
    if (ref_pos.ref >= kMaxNumReferenceFrames ||
        shared_->reference_frames[ref_pos.ref].frame->xsize() == 0) {
      return JXL_FAILURE("Invalid reference frame ID");
    }
    if (!shared_->reference_frames[ref_pos.ref].ib_is_in_xyb) {
      return JXL_FAILURE(
          "Patches cannot use frames saved post color transforms");
    }
    const ImageBundle& ib = *shared_->reference_frames[ref_pos.ref].frame;
    ref_pos.x0 = read_num(kPatchReferencePositionContext);
    ref_pos.y0 = read_num(kPatchReferencePositionContext);
    ref_pos.xsize = read_num(kPatchSizeContext) + 1;
    ref_pos.ysize = read_num(kPatchSizeContext) + 1;
    if (ref_pos.x0 + ref_pos.xsize > ib.xsize()) {
      return JXL_FAILURE("Invalid position specified in reference frame");
    }
    if (ref_pos.y0 + ref_pos.ysize > ib.ysize()) {
      return JXL_FAILURE("Invalid position specified in reference frame");
    }
    size_t id_count = read_num(kPatchCountContext) + 1;
    if (id_count > kMaxPatches) {
      return JXL_FAILURE("Too many patches in dictionary");
    }
    positions_.reserve(positions_.size() + id_count);
    for (size_t i = 0; i < id_count; i++) {
      PatchPosition pos;
      pos.ref_pos = ref_pos;
      if (i == 0) {
        pos.x = read_num(kPatchPositionContext);
        pos.y = read_num(kPatchPositionContext);
      } else {
        pos.x =
            positions_.back().x + UnpackSigned(read_num(kPatchOffsetContext));
        pos.y =
            positions_.back().y + UnpackSigned(read_num(kPatchOffsetContext));
      }
      if (pos.x + ref_pos.xsize > xsize) {
        return JXL_FAILURE("Invalid patch x: at %zu + %zu > %zu", pos.x,
                           ref_pos.xsize, xsize);
      }
      if (pos.y + ref_pos.ysize > ysize) {
        return JXL_FAILURE("Invalid patch y: at %zu + %zu > %zu", pos.y,
                           ref_pos.ysize, ysize);
      }
      for (size_t i = 0; i < shared_->metadata->m.extra_channel_info.size() + 1;
           i++) {
        uint32_t blend_mode = read_num(kPatchBlendModeContext);
        if (blend_mode >= uint32_t(PatchBlendMode::kNumBlendModes)) {
          return JXL_FAILURE("Invalid patch blend mode: %u", blend_mode);
        }
        PatchBlending info;
        info.mode = static_cast<PatchBlendMode>(blend_mode);
        if (UsesAlpha(info.mode)) {
          *uses_extra_channels = true;
        }
        if (info.mode != PatchBlendMode::kNone && i > 0) {
          *uses_extra_channels = true;
        }
        if (UsesAlpha(info.mode) &&
            shared_->metadata->m.extra_channel_info.size() > 1) {
          info.alpha_channel = read_num(kPatchAlphaChannelContext);
          if (info.alpha_channel >=
              shared_->metadata->m.extra_channel_info.size()) {
            return JXL_FAILURE(
                "Invalid alpha channel for blending: %u out of %u\n",
                info.alpha_channel,
                (uint32_t)shared_->metadata->m.extra_channel_info.size());
          }
        } else {
          info.alpha_channel = 0;
        }
        if (UsesClamp(info.mode)) {
          info.clamp = read_num(kPatchClampContext);
        } else {
          info.clamp = false;
        }
        pos.blending.push_back(info);
      }
      positions_.push_back(std::move(pos));
    }
  }

  if (!decoder.CheckANSFinalState()) {
    return JXL_FAILURE("ANS checksum failure.");
  }
  if (!HasAny()) {
    return JXL_FAILURE("Decoded patch dictionary but got none");
  }

  ComputePatchCache();
  return true;
}

int PatchDictionary::GetReferences() const {
  int result = 0;
  for (size_t i = 0; i < positions_.size(); ++i) {
    result |= (1 << static_cast<int>(positions_[i].ref_pos.ref));
  }
  return result;
}

void PatchDictionary::ComputePatchCache() {
  patch_starts_.clear();
  sorted_patches_.clear();
  if (positions_.empty()) return;
  std::vector<std::pair<size_t, size_t>> sorted_patches_y;
  for (size_t i = 0; i < positions_.size(); i++) {
    const PatchPosition& pos = positions_[i];<commit_msg>Fix patch decoding fuzzer-found slow behaviour.<commit_after>// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "lib/jxl/dec_patch_dictionary.h"

#include <stdint.h>
#include <stdlib.h>
#include <sys/types.h>

#include <algorithm>
#include <random>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "lib/jxl/ans_params.h"
#include "lib/jxl/base/compiler_specific.h"
#include "lib/jxl/base/override.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/blending.h"
#include "lib/jxl/chroma_from_luma.h"
#include "lib/jxl/color_management.h"
#include "lib/jxl/common.h"
#include "lib/jxl/dec_ans.h"
#include "lib/jxl/dec_frame.h"
#include "lib/jxl/entropy_coder.h"
#include "lib/jxl/frame_header.h"
#include "lib/jxl/image.h"
#include "lib/jxl/image_bundle.h"
#include "lib/jxl/image_ops.h"
#include "lib/jxl/patch_dictionary_internal.h"

namespace jxl {

constexpr size_t kMaxPatches = 1 << 24;

Status PatchDictionary::Decode(BitReader* br, size_t xsize, size_t ysize,
                               bool* uses_extra_channels) {
  positions_.clear();
  std::vector<uint8_t> context_map;
  ANSCode code;
  JXL_RETURN_IF_ERROR(
      DecodeHistograms(br, kNumPatchDictionaryContexts, &code, &context_map));
  ANSSymbolReader decoder(&code, br);

  auto read_num = [&](size_t context) {
    size_t r = decoder.ReadHybridUint(context, br, context_map);
    return r;
  };

  size_t num_ref_patch = read_num(kNumRefPatchContext);
  // TODO(veluca): does this make sense?
  if (num_ref_patch > kMaxPatches) {
    return JXL_FAILURE("Too many patches in dictionary");
  }

  size_t total_patches = 0;
  size_t next_size = 1;

  for (size_t id = 0; id < num_ref_patch; id++) {
    PatchReferencePosition ref_pos;
    ref_pos.ref = read_num(kReferenceFrameContext);
    if (ref_pos.ref >= kMaxNumReferenceFrames ||
        shared_->reference_frames[ref_pos.ref].frame->xsize() == 0) {
      return JXL_FAILURE("Invalid reference frame ID");
    }
    if (!shared_->reference_frames[ref_pos.ref].ib_is_in_xyb) {
      return JXL_FAILURE(
          "Patches cannot use frames saved post color transforms");
    }
    const ImageBundle& ib = *shared_->reference_frames[ref_pos.ref].frame;
    ref_pos.x0 = read_num(kPatchReferencePositionContext);
    ref_pos.y0 = read_num(kPatchReferencePositionContext);
    ref_pos.xsize = read_num(kPatchSizeContext) + 1;
    ref_pos.ysize = read_num(kPatchSizeContext) + 1;
    if (ref_pos.x0 + ref_pos.xsize > ib.xsize()) {
      return JXL_FAILURE("Invalid position specified in reference frame");
    }
    if (ref_pos.y0 + ref_pos.ysize > ib.ysize()) {
      return JXL_FAILURE("Invalid position specified in reference frame");
    }
    size_t id_count = read_num(kPatchCountContext) + 1;
    total_patches += id_count;
    if (total_patches > kMaxPatches) {
      return JXL_FAILURE("Too many patches in dictionary");
    }
    if (next_size < total_patches) {
      next_size *= 2;
      next_size = std::min<size_t>(next_size, kMaxPatches);
    }
    positions_.reserve(next_size);
    for (size_t i = 0; i < id_count; i++) {
      PatchPosition pos;
      pos.ref_pos = ref_pos;
      if (i == 0) {
        pos.x = read_num(kPatchPositionContext);
        pos.y = read_num(kPatchPositionContext);
      } else {
        pos.x =
            positions_.back().x + UnpackSigned(read_num(kPatchOffsetContext));
        pos.y =
            positions_.back().y + UnpackSigned(read_num(kPatchOffsetContext));
      }
      if (pos.x + ref_pos.xsize > xsize) {
        return JXL_FAILURE("Invalid patch x: at %zu + %zu > %zu", pos.x,
                           ref_pos.xsize, xsize);
      }
      if (pos.y + ref_pos.ysize > ysize) {
        return JXL_FAILURE("Invalid patch y: at %zu + %zu > %zu", pos.y,
                           ref_pos.ysize, ysize);
      }
      for (size_t i = 0; i < shared_->metadata->m.extra_channel_info.size() + 1;
           i++) {
        uint32_t blend_mode = read_num(kPatchBlendModeContext);
        if (blend_mode >= uint32_t(PatchBlendMode::kNumBlendModes)) {
          return JXL_FAILURE("Invalid patch blend mode: %u", blend_mode);
        }
        PatchBlending info;
        info.mode = static_cast<PatchBlendMode>(blend_mode);
        if (UsesAlpha(info.mode)) {
          *uses_extra_channels = true;
        }
        if (info.mode != PatchBlendMode::kNone && i > 0) {
          *uses_extra_channels = true;
        }
        if (UsesAlpha(info.mode) &&
            shared_->metadata->m.extra_channel_info.size() > 1) {
          info.alpha_channel = read_num(kPatchAlphaChannelContext);
          if (info.alpha_channel >=
              shared_->metadata->m.extra_channel_info.size()) {
            return JXL_FAILURE(
                "Invalid alpha channel for blending: %u out of %u\n",
                info.alpha_channel,
                (uint32_t)shared_->metadata->m.extra_channel_info.size());
          }
        } else {
          info.alpha_channel = 0;
        }
        if (UsesClamp(info.mode)) {
          info.clamp = read_num(kPatchClampContext);
        } else {
          info.clamp = false;
        }
        pos.blending.push_back(info);
      }
      positions_.push_back(std::move(pos));
    }
  }
  positions_.shrink_to_fit();

  if (!decoder.CheckANSFinalState()) {
    return JXL_FAILURE("ANS checksum failure.");
  }
  if (!HasAny()) {
    return JXL_FAILURE("Decoded patch dictionary but got none");
  }

  ComputePatchCache();
  return true;
}

int PatchDictionary::GetReferences() const {
  int result = 0;
  for (size_t i = 0; i < positions_.size(); ++i) {
    result |= (1 << static_cast<int>(positions_[i].ref_pos.ref));
  }
  return result;
}

void PatchDictionary::ComputePatchCache() {
  patch_starts_.clear();
  sorted_patches_.clear();
  if (positions_.empty()) return;
  std::vector<std::pair<size_t, size_t>> sorted_patches_y;
  for (size_t i = 0; i < positions_.size(); i++) {
    const PatchPosition& pos = positions_[i];
{'fraction_non_alphanumeric': 0.1297940043543795, 'fraction_numerical': 0.005275498241500586, 'mean_word_length': 2.975699067909454, 'pattern_counts': {'":': 0, '<': 51, '<?xml version=': 0, '>': 82, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2623385', 'n_tokens_mistral': 4182, 'n_tokens_neox': 4066, 'n_words': 832}