Status GetNextInternal()

in tensorflow_compression/cc/kernels/y4m_dataset_kernels.cc [117:213]
71 lines of code
10 McCabe index (conditional complexity)

      Status GetNextInternal(IteratorContext* ctx,
                             vector<Tensor>* out_tensors,
                             bool* end_of_sequence) override {
        mutex_lock l(mu_);

        do {
          if (file_) {
            const string_view frame_header = "FRAME\n";
            size_t frame_size = width_ * height_ * 3;
            int64_t cbcr_width = width_;
            int64_t cbcr_height = height_;
            string_view frame_buffer;

            if (chroma_format_ == ChromaFormat::I420) {
              frame_size /= 2;
              cbcr_width /= 2;
              cbcr_height /= 2;
            }
            const size_t cbcr_size = cbcr_width * cbcr_height;

            // This is a no-op for the second and subsequent frames.
            buffer_.resize(frame_header.size() + frame_size);

            // Try to read the next frame.
            Status status = file_->Read(file_pos_, buffer_.size(),
                                        &frame_buffer, &buffer_[0]);

            // Yield frame on successful read of a complete frame.
            if (status.ok()) {
              DCHECK_EQ(frame_buffer.size(), buffer_.size());

              if (!absl::ConsumePrefix(&frame_buffer, frame_header)) {
                return errors::InvalidArgument(
                    "Input file '", dataset()->filenames_[file_index_],
                    "' has a FRAME marker at byte ", file_pos_, " which is "
                    "either invalid or has unsupported frame parameters.");
              }

              Tensor y_tensor(ctx->allocator({}), DT_UINT8,
                              {height_, width_, 1});
              Tensor cbcr_tensor(ctx->allocator({}), DT_UINT8,
                                 {cbcr_height, cbcr_width, 2});
              auto flat_y = y_tensor.flat<uint8_t>();
              auto flat_cbcr = cbcr_tensor.flat<uint8_t>();
              std::memcpy(flat_y.data(), frame_buffer.data(), flat_y.size());
              frame_buffer.remove_prefix(flat_y.size());
              for (int i = 0; i < cbcr_size; i++) {
                flat_cbcr.data()[2*i] = frame_buffer[i];
                flat_cbcr.data()[2*i+1] = frame_buffer[cbcr_size+i];
              }
              out_tensors->push_back(std::move(y_tensor));
              out_tensors->push_back(std::move(cbcr_tensor));

              file_pos_ += buffer_.size();
              *end_of_sequence = false;
              return status;
            }

            // Catch any other errors than out of range, which needs special
            // treatment.
            if (!errors::IsOutOfRange(status)) {
              return status;
            }

            // If frame buffer is not empty, we just read an incomplete frame
            // (or one that has frame parameters that change its size).
            if (!frame_buffer.empty()) {
              return errors::InvalidArgument(
                  "Input file '", dataset()->filenames_[file_index_],
                  "' has an incomplete or unsupported frame at byte ",
                  file_pos_, ". Expected to read ", buffer_.size(),
                  " bytes, only ", frame_buffer.size(), " were available.");
            }

            // Out of range error with empty frame buffer means correct end of
            // file. Clean up and check for next file.
            file_.reset();
            ++file_index_;
          }

          // Exit if there are no more files to process.
          if (file_index_ >= dataset()->filenames_.size()) {
            *end_of_sequence = true;
            return Status::OK();
          }

          // Open next file.
          TF_RETURN_IF_ERROR(ctx->env()->NewRandomAccessFile(
              dataset()->filenames_[file_index_], &file_));

          // Read and parse header.
          TF_RETURN_IF_ERROR(ReadHeader(*file_, file_index_, buffer_));
          TF_RETURN_IF_ERROR(ParseHeader(buffer_, file_index_, width_, height_,
                                         chroma_format_));
          file_pos_ = buffer_.size();
        } while (true);
      }