Repository: apache/arrow
File: cpp/src/arrow/ipc/message.cc (C++)

Code before:

Result<std::unique_ptr<Message>> ReadMessage(const int64_t offset,
                                             const int32_t metadata_length,
                                             const int64_t body_length,
                                             io::RandomAccessFile* file) {
  std::unique_ptr<Message> result;
  auto listener = std::make_shared<AssignMessageDecoderListener>(&result);
  MessageDecoder decoder(listener);

  if (metadata_length < decoder.next_required_size()) {
    return Status::Invalid("metadata_length should be at least ",
                           decoder.next_required_size());
  }

  ARROW_ASSIGN_OR_RAISE(auto metadata,
                        file->ReadAt(offset, metadata_length + body_length));
  if (metadata->size() < metadata_length) {
    return Status::Invalid("Expected to read ", metadata_length,
                           " metadata bytes at offset ", offset, " but got ",
                           metadata->size());
  }

  ARROW_RETURN_NOT_OK(decoder.Consume(SliceBuffer(metadata, 0, metadata_length)));

  switch (decoder.state()) {
    case MessageDecoder::State::INITIAL:
      return result;
    case MessageDecoder::State::METADATA_LENGTH:
      return Status::Invalid("metadata length is missing. File offset: ", offset,
                             ", metadata length: ", metadata_length);
    case MessageDecoder::State::METADATA:
      return Status::Invalid("flatbuffer size ", decoder.next_required_size(),
                             " invalid. File offset: ", offset,
                             ", metadata length: ", metadata_length);
    case MessageDecoder::State::BODY: {
      auto body = SliceBuffer(metadata, metadata_length,
                              std::min(body_length, metadata->size() - metadata_length));
      if (body->size() < decoder.next_required_size()) {
        return Status::IOError("Expected to be able to read ",
                               decoder.next_required_size(),
                               " bytes for message body, got ", body->size());
      }
      RETURN_NOT_OK(decoder.Consume(body));
      return result;
    }
    case MessageDecoder::State::EOS:
      return Status::Invalid("Unexpected empty message in IPC file format");
    default:
      return Status::Invalid("Unexpected state: ", decoder.state());
  }
}

Review comment (refactor):
There's a lot of code duplication with the previous `ReadMessage` overload above (the one that takes a `FieldsLoaderFunction`). Can we refactor using a common helper so that most of the duplication goes away?

Code after:
            body, file->ReadAt(offset + metadata_length, decoder.next_required_size()));
      }

      if (body->size() < decoder.next_required_size()) {
        return Status::IOError("Expected to be able to read ",
                               decoder.next_required_size(),
                               " bytes for message body, got ", body->size());
      }
      RETURN_NOT_OK(decoder.Consume(body));
      return result;
    }
    case MessageDecoder::State::EOS:
      return Status::Invalid("Unexpected empty message in IPC file format");
    default:
      return Status::Invalid("Unexpected state: ", decoder.state());
  }
}

Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_length,
                                             io::RandomAccessFile* file,
                                             const FieldsLoaderFunction& fields_loader) {
  return ReadMessageInternal(offset, metadata_length, /*body_length=*/std::nullopt, file,
                             fields_loader);
}

Result<std::unique_ptr<Message>> ReadMessage(const int64_t offset,
                                             const int32_t metadata_length,
                                             const int64_t body_length,
                                             io::RandomAccessFile* file) {
  return ReadMessageInternal(offset, metadata_length, body_length, file,
                             /*fields_loader=*/{});
}

Future<std::shared_ptr<Message>> ReadMessageAsync(int64_t offset, int32_t metadata_length,
                                                  int64_t body_length,
                                                  io::RandomAccessFile* file,
                                                  const io::IOContext& context) {
  struct State {
    std::unique_ptr<Message> result;
    std::shared_ptr<MessageDecoderListener> listener;
    std::shared_ptr<MessageDecoder> decoder;
  };
  auto state = std::make_shared<State>();
  state->listener = std::make_shared<AssignMessageDecoderListener>(&state->result);
  state->decoder = std::make_shared<MessageDecoder>(state->listener);

  if (metadata_length < state->decoder->next_required_size()) {
    return Status::Invalid("metadata_length should be at least ",
                           state->decoder->next_required_size());
  }
  return file->ReadAsync(context, offset, metadata_length + body_length)