void RangeEncodeImpl()

in tensorflow_compression/cc/kernels/unbounded_index_range_coding_kernels.cc [187:251]
53 lines of code
8 McCabe index (conditional complexity)

  void RangeEncodeImpl(TTypes<int32_t>::ConstFlat data,
                       TTypes<int32_t>::ConstFlat index,
                       TTypes<int32_t>::ConstMatrix cdf,
                       TTypes<int32_t>::ConstVec cdf_size,
                       TTypes<int32_t>::ConstVec offset,
                       std::string* output) const {
    RangeEncoder encoder;

    DCHECK_GE(cdf.dimension(1), 2);
    DCHECK_LE(cdf.dimension(1), std::numeric_limits<int16_t>::max());
    DCHECK_EQ(cdf.dimension(0), cdf_size.size());

    const uint32_t max_overflow = (1 << overflow_width_) - 1;

    const int64_t data_size = data.size();
    for (int64_t i = 0; i < data_size; ++i) {
      const int32_t cdf_index = index(i);

      DCHECK_GE(cdf_index, 0);
      DCHECK_LT(cdf_index, cdf.dimension(0));

      const int32_t max_value = cdf_size(cdf_index) - 2;
      DCHECK_GE(max_value, 0);
      DCHECK_LT(max_value + 1, cdf.dimension(1));

      int32_t value = data(i);
      // Map values with tracked probabilities to 0..max_value range.
      value -= offset(cdf_index);
      // If outside of this range, map value to non-negative integer overflow.
      // NOTE: It might be a good idea to check overflow is within uint32 range.
      uint32_t overflow = 0;
      if (value < 0) {
        overflow = -2 * value - 1;
        value = max_value;
      } else if (value >= max_value) {
        overflow = 2 * (value - max_value);
        value = max_value;
      }

      const int32_t* cdf_slice = &cdf(cdf_index, 0);
      encoder.Encode(cdf_slice[value], cdf_slice[value + 1], precision_,
                     output);

      // Encode overflow using variable length code.
      if (value == max_value) {
        int32_t widths = 0;
        while (overflow >> (widths * overflow_width_) != 0) {
          ++widths;
        }
        uint32_t val = widths;
        while (val >= max_overflow) {
          encoder.Encode(max_overflow, max_overflow + 1, overflow_width_,
                         output);
          val -= max_overflow;
        }
        encoder.Encode(val, val + 1, overflow_width_, output);
        for (int32_t j = 0; j < widths; ++j) {
          const uint32_t val =
              (overflow >> (j * overflow_width_)) & max_overflow;
          encoder.Encode(val, val + 1, overflow_width_, output);
        }
      }
    }
    encoder.Finalize(output);
  }