document_stream-inl.h

#ifndef SIMDJSON_GENERIC_ONDEMAND_DOCUMENT_STREAM_INL_H
 
#ifndef SIMDJSON_CONDITIONAL_INCLUDE
#define SIMDJSON_GENERIC_ONDEMAND_DOCUMENT_STREAM_INL_H
#include "simdjson/generic/ondemand/base.h"
#include "simdjson/generic/ondemand/document_stream.h"
#include "simdjson/generic/ondemand/document-inl.h"
#include "simdjson/generic/implementation_simdjson_result_base-inl.h"
#endif // SIMDJSON_CONDITIONAL_INCLUDE
 
#include <algorithm>
#include <cstring>
#include <stdexcept>
 
namespace simdjson {
namespace SIMDJSON_IMPLEMENTATION {
namespace ondemand {
 
#ifdef SIMDJSON_THREADS_ENABLED
 
inline void stage1_worker::finish() {
  // After calling "run" someone would call finish() to wait
  // for the end of the processing.
  // This function will wait until either the thread has done
  // the processing or, else, the destructor has been called.
  std::unique_lock<std::mutex> lock(locking_mutex);
  cond_var.wait(lock, [this]{return has_work == false;});
}
 
inline stage1_worker::~stage1_worker() {
  // The thread may never outlive the stage1_worker instance
  // and will always be stopped/joined before the stage1_worker
  // instance is gone.
  stop_thread();
}
 
inline void stage1_worker::start_thread() {
  std::unique_lock<std::mutex> lock(locking_mutex);
  if(thread.joinable()) {
    return; // This should never happen but we never want to create more than one thread.
  }
  thread = std::thread([this]{
      while(true) {
        std::unique_lock<std::mutex> thread_lock(locking_mutex);
        // We wait for either "run" or "stop_thread" to be called.
        cond_var.wait(thread_lock, [this]{return has_work || !can_work;});
        // If, for some reason, the stop_thread() method was called (i.e., the
        // destructor of stage1_worker is called, then we want to immediately destroy
        // the thread (and not do any more processing).
        if(!can_work) {
          break;
        }
        this->owner->stage1_thread_error = this->owner->run_stage1(*this->stage1_thread_parser,
              this->_next_batch_start);
        this->has_work = false;
        // The condition variable call should be moved after thread_lock.unlock() for performance
        // reasons but thread sanitizers may report it as a data race if we do.
        // See https://stackoverflow.com/questions/35775501/c-should-condition-variable-be-notified-under-lock
        cond_var.notify_one(); // will notify "finish"
        thread_lock.unlock();
      }
    }
  );
}
 
 
inline void stage1_worker::stop_thread() {
  std::unique_lock<std::mutex> lock(locking_mutex);
  // We have to make sure that all locks can be released.
  can_work = false;
  has_work = false;
  cond_var.notify_all();
  lock.unlock();
  if(thread.joinable()) {
    thread.join();
  }
}
 
inline void stage1_worker::run(document_stream * ds, parser * stage1, size_t next_batch_start) {
  std::unique_lock<std::mutex> lock(locking_mutex);
  owner = ds;
  _next_batch_start = next_batch_start;
  stage1_thread_parser = stage1;
  has_work = true;
  // The condition variable call should be moved after thread_lock.unlock() for performance
  // reasons but thread sanitizers may report it as a data race if we do.
  // See https://stackoverflow.com/questions/35775501/c-should-condition-variable-be-notified-under-lock
  cond_var.notify_one(); // will notify the thread lock that we have work
  lock.unlock();
}
 
#endif  // SIMDJSON_THREADS_ENABLED
 
simdjson_inline document_stream::document_stream(
  ondemand::parser &_parser,
  const uint8_t *_buf,
  size_t _len,
  size_t _batch_size,
  bool _allow_comma_separated,
  stream_format _format
) noexcept
  : parser{&_parser},
    buf{_buf},
    len{_len},
    batch_size{_batch_size <= MINIMAL_BATCH_SIZE ? MINIMAL_BATCH_SIZE : _batch_size},
    allow_comma_separated{_allow_comma_separated},
    format{_format},
    error{SUCCESS}
    #ifdef SIMDJSON_THREADS_ENABLED
    , use_thread(_parser.threaded) // we need to make a copy because _parser.threaded can change
    #endif
{
}
 
simdjson_inline document_stream::document_stream() noexcept
  : parser{nullptr},
    buf{nullptr},
    len{0},
    batch_size{0},
    allow_comma_separated{false},
    format{stream_format::whitespace_delimited},
    error{UNINITIALIZED}
    #ifdef SIMDJSON_THREADS_ENABLED
    , use_thread(false)
    #endif
{
}
 
simdjson_inline document_stream::~document_stream() noexcept
{
  #ifdef SIMDJSON_THREADS_ENABLED
  worker.reset();
  #endif
}
 
inline size_t document_stream::size_in_bytes() const noexcept {
  return len;
}
 
inline size_t document_stream::truncated_bytes() const noexcept {
  if(error == CAPACITY) { return len - batch_start; }
  return parser->implementation->structural_indexes[parser->implementation->n_structural_indexes] - parser->implementation->structural_indexes[parser->implementation->n_structural_indexes + 1];
}
 
simdjson_inline document_stream::iterator::iterator() noexcept
  : stream{nullptr}, finished{true} {
}
 
simdjson_inline document_stream::iterator::iterator(document_stream* _stream, bool is_end) noexcept
  : stream{_stream}, finished{is_end} {
}
 
simdjson_inline simdjson_result<ondemand::document_reference> document_stream::iterator::operator*() noexcept {
  return simdjson_result<ondemand::document_reference>(stream->doc, stream->error);
}
 
simdjson_inline document_stream::iterator& document_stream::iterator::operator++() noexcept {
  // If there is an error, then we want the iterator
  // to be finished, no matter what. (E.g., we do not
  // keep generating documents with errors, or go beyond
  // a document with errors.)
  //
  // Users do not have to call "operator*()" when they use operator++,
  // so we need to end the stream in the operator++ function.
  //
  // Note that setting finished = true is essential otherwise
  // we would enter an infinite loop.
  if (stream->error) { finished = true; }
  // Note that stream->error() is guarded against error conditions
  // (it will immediately return if stream->error casts to false).
  // In effect, this next function does nothing when (stream->error)
  // is true (hence the risk of an infinite loop).
  stream->next();
  // If that was the last document, we're finished.
  // It is the only type of error we do not want to appear
  // in operator*.
  if (stream->error == EMPTY) { finished = true; }
  // If we had any other kind of error (not EMPTY) then we want
  // to pass it along to the operator* and we cannot mark the result
  // as "finished" just yet.
  return *this;
}
 
simdjson_inline bool document_stream::iterator::at_end() const noexcept {
  return finished;
}
 
 
simdjson_inline bool document_stream::iterator::operator!=(const document_stream::iterator &other) const noexcept {
  return finished != other.finished;
}
 
simdjson_inline bool document_stream::iterator::operator==(const document_stream::iterator &other) const noexcept {
  return finished == other.finished;
}
 
simdjson_inline document_stream::iterator document_stream::begin() noexcept {
  start();
  // If there are no documents, we're finished.
  return iterator(this, error == EMPTY);
}
 
simdjson_inline document_stream::iterator document_stream::end() noexcept {
  return iterator(this, true);
}
 
inline void document_stream::start() noexcept {
  if (error) { return; }
  error = parser->allocate(batch_size);
  if (error) { return; }
  // Always run the first stage 1 parse immediately
  batch_start = 0;
  error = run_stage1(*parser, batch_start);
  while(error == EMPTY) {
    // In exceptional cases, we may start with an empty block
    batch_start = next_batch_start();
    if (batch_start >= len) { return; }
    error = run_stage1(*parser, batch_start);
  }
  if (error) { return; }
  // For json_sequence mode, structural_indexes[0] points to the actual JSON value
  // after the RS delimiter and any following whitespace. For regular mode, it is
  // the offset from batch_start to the first document in the batch.
  doc_index = batch_start + parser->implementation->structural_indexes[0];
  doc = document(json_iterator(&buf[batch_start], parser));
  doc.iter._streaming = true;
 
  #ifdef SIMDJSON_THREADS_ENABLED
  if (use_thread && next_batch_start() < len) {
    // Kick off the first thread on next batch if needed
    if (worker.get() == nullptr) {
      worker.reset(new(std::nothrow) stage1_worker());
      if (worker.get() == nullptr) { error = MEMALLOC; return; }
    }
    error = stage1_thread_parser.allocate(batch_size);
    if (error) { return; }
    worker->start_thread();
    start_stage1_thread();
    if (error) { return; }
  }
  #endif // SIMDJSON_THREADS_ENABLED
}
 
inline void document_stream::next() noexcept {
  // We always enter at once once in an error condition.
  if (error) { return; }
  next_document();
  if (error) { return; }
  auto cur_struct_index = doc.iter._root - parser->implementation->structural_indexes.get();
  doc_index = batch_start + parser->implementation->structural_indexes[cur_struct_index];
 
  // Check if at end of structural indexes (i.e. at end of batch)
  if(cur_struct_index >= static_cast<int64_t>(parser->implementation->n_structural_indexes)) {
    error = EMPTY;
    // Load another batch (if available)
    while (error == EMPTY) {
      batch_start = next_batch_start();
      if (batch_start >= len) { break; }
      #ifdef SIMDJSON_THREADS_ENABLED
      if(use_thread) {
        load_from_stage1_thread();
      } else {
        error = run_stage1(*parser, batch_start);
      }
      #else
      error = run_stage1(*parser, batch_start);
      #endif
      doc.iter = json_iterator(&buf[batch_start], parser);
      doc.iter._streaming = true;
      if (error) { continue; } // If the error was EMPTY, we may want to load another batch.
      doc_index = batch_start + parser->implementation->structural_indexes[0];
    }
  }
}
 
simdjson_inline uint8_t document_stream::document_delimiter() const noexcept {
  switch (format) {
    case stream_format::newline_delimited: return '\n';
    case stream_format::json_sequence: return 0x1E;
    default: return 0;
  }
}
 
simdjson_inline bool document_stream::skip_to_delimiter(uint8_t delimiter) noexcept {
  const uint8_t *const base = &buf[batch_start];
  const token_position pos = doc.iter.position();
  const token_position end = doc.iter.end_position();
  if (pos >= end) { return false; }
  const size_t here = size_t(doc.iter.token.peek(pos) - base);
  const size_t batch_len =
      (len - batch_start < batch_size) ? len - batch_start : batch_size;
  if (here >= batch_len) { return false; }
  const uint8_t *const found = static_cast<const uint8_t *>(
      std::memchr(base + here, delimiter, batch_len - here));
  if (found == nullptr) { return false; }
 
  const uint32_t boundary = uint32_t(found - base);
  // The answer is near `pos`: the delimiter ends the current document, while
  // `end` spans the whole batch. Gallop first so the cost follows the distance
  // rather than the size of the batch.
  token_position lo = pos;
  size_t hop = 1;
  while (lo + hop < end && lo[hop] < boundary) { lo += hop; hop <<= 1; }
  token_position hi = (lo + hop < end) ? lo + hop : end;
  while (lo < hi) {
    const token_position mid = lo + ((hi - lo) >> 1);
    if (*mid < boundary) { lo = mid + 1; } else { hi = mid; }
  }
  doc.iter.token.set_position(lo);
  return true;
}
 
inline void document_stream::next_document() noexcept {
  // A delimiter that cannot occur inside a document tells us where the current
  // one ends, so we can jump there instead of walking every structural. Only
  // valid while the iterator is still inside the document: a consumed document
  // already sits on the next one's first token, and skip_child() returns at
  // once for it.
  //
  // The jump does not structure-validate the unread remainder of the current
  // document: under newline_delimited / json_sequence the next delimiter is
  // assumed to be the true document boundary. Callers that leave depth() > 0
  // while violating that contract (e.g. pretty multi-line JSON under
  // newline_delimited) can mis-align following documents; use
  // whitespace_delimited if unsure.
  const uint8_t delimiter = document_delimiter();
  if (delimiter != 0 && !error && doc.iter.depth() > 0 &&
      skip_to_delimiter(delimiter)) {
    doc.iter._depth = 1;
    doc.iter._string_buf_loc = parser->string_buf.get();
    doc.iter._root = doc.iter.position();
    return;
  }
  // Go to next place where depth=0 (document depth)
  error = doc.iter.skip_child(0);
  if (error) { return; }
  // Always set depth=1 at the start of document
  doc.iter._depth = 1;
  // consume comma if comma separated is allowed
  if (allow_comma_separated) {
    error_code ignored = doc.iter.consume_character(',');
    static_cast<void>(ignored); // ignored on purpose
  }
  // Resets the string buffer at the beginning, thus invalidating the strings.
  doc.iter._string_buf_loc = parser->string_buf.get();
  doc.iter._root = doc.iter.position();
}
 
inline size_t document_stream::next_batch_start() const noexcept {
  return batch_start + parser->implementation->structural_indexes[parser->implementation->n_structural_indexes];
}
 
inline error_code document_stream::run_stage1(ondemand::parser &p, size_t _batch_start) noexcept {
  // This code only updates the structural index in the parser, it does not update any json_iterator
  // instance.
  size_t remaining = len - _batch_start;
  stage1_mode mode;
  if (remaining <= batch_size) {
    // Final batch
    switch (format) {
      case stream_format::json_sequence:
        mode = stage1_mode::json_sequence_final;
        break;
      case stream_format::comma_delimited:
        mode = stage1_mode::comma_delimited_final;
        break;
      default:
        mode = stage1_mode::streaming_final;
        break;
    }
    return p.implementation->stage1(&buf[_batch_start], remaining, mode);
  } else {
    // Partial batch
    switch (format) {
      case stream_format::json_sequence:
        mode = stage1_mode::json_sequence_partial;
        break;
      case stream_format::comma_delimited:
        mode = stage1_mode::comma_delimited_partial;
        break;
      default:
        mode = stage1_mode::streaming_partial;
        break;
    }
    return p.implementation->stage1(&buf[_batch_start], batch_size, mode);
  }
}
 
simdjson_inline size_t document_stream::iterator::current_index() const noexcept {
  return stream->doc_index;
}
 
simdjson_inline std::string_view document_stream::iterator::source() const noexcept {
  auto depth = stream->doc.iter.depth();
  auto cur_struct_index = stream->doc.iter._root - stream->parser->implementation->structural_indexes.get();
 
  // If at root, process the first token to determine if scalar value
  if (stream->doc.iter.at_root()) {
    switch (stream->buf[stream->batch_start + stream->parser->implementation->structural_indexes[cur_struct_index]]) {
      case '{': case '[':   // Depth=1 already at start of document
        break;
      case '}': case ']':
        depth--;
        break;
      default:    // Scalar value document
        // This returns a string spanning from start of value to the beginning of the next document (excluded)
        {
          auto next_index = stream->batch_start + stream->parser->implementation->structural_indexes[++cur_struct_index];
          // normally the length would be next_index - current_index() - 1, except for the last document
          size_t svlen = next_index - current_index();
          const char *start = reinterpret_cast<const char*>(stream->buf) + current_index();
          // Trim trailing whitespace, NUL, and RS (0x1E). In RFC 7464
          // json_sequence mode the scanner classifies RS as a scalar
          // character, so an RS-prefixed scalar document (number / true /
          // false / null / string) has no closing structural index and the
          // slice runs all the way up to the next document's RS. RS cannot
          // legally appear in a JSON value at the source level (control
          // characters in strings must be escaped as \u001E), so stripping
          // it is safe in every stream_format.
          while(svlen > 1 && (std::isspace(static_cast<unsigned char>(start[svlen-1])) || start[svlen-1] == '\0' || static_cast<uint8_t>(start[svlen-1]) == 0x1E || (stream->format == stream_format::comma_delimited && start[svlen-1] == ','))) {
            svlen--;
          }
          return std::string_view(start, svlen);
        }
    }
    cur_struct_index++;
  }
 
  while (cur_struct_index <= static_cast<int64_t>(stream->parser->implementation->n_structural_indexes)) {
    switch (stream->buf[stream->batch_start + stream->parser->implementation->structural_indexes[cur_struct_index]]) {
      case '{': case '[':
        depth++;
        break;
      case '}': case ']':
        depth--;
        break;
    }
    if (depth == 0) { break; }
    cur_struct_index++;
  }
 
  return std::string_view(reinterpret_cast<const char*>(stream->buf) + current_index(), stream->parser->implementation->structural_indexes[cur_struct_index] - current_index() + stream->batch_start + 1);;
}
 
inline error_code document_stream::iterator::error() const noexcept {
  return stream->error;
}
 
#ifdef SIMDJSON_THREADS_ENABLED
 
inline void document_stream::load_from_stage1_thread() noexcept {
  worker->finish();
  // Swap to the parser that was loaded up in the thread. Make sure the parser has
  // enough memory to swap to, as well.
  std::swap(stage1_thread_parser,*parser);
  error = stage1_thread_error;
  if (error) { return; }
 
  // If there's anything left, start the stage 1 thread!
  if (next_batch_start() < len) {
    start_stage1_thread();
  }
}
 
inline void document_stream::start_stage1_thread() noexcept {
  // we call the thread on a lambda that will update
  // this->stage1_thread_error
  // there is only one thread that may write to this value
  // TODO this is NOT exception-safe.
  this->stage1_thread_error = UNINITIALIZED; // In case something goes wrong, make sure it's an error
  size_t _next_batch_start = this->next_batch_start();
 
  worker->run(this, & this->stage1_thread_parser, _next_batch_start);
}
 
#endif // SIMDJSON_THREADS_ENABLED
 
} // namespace ondemand
} // namespace SIMDJSON_IMPLEMENTATION
} // namespace simdjson
 
namespace simdjson {
 
simdjson_inline simdjson_result<SIMDJSON_IMPLEMENTATION::ondemand::document_stream>::simdjson_result(
  error_code error
) noexcept :
    implementation_simdjson_result_base<SIMDJSON_IMPLEMENTATION::ondemand::document_stream>(error)
{
}
simdjson_inline simdjson_result<SIMDJSON_IMPLEMENTATION::ondemand::document_stream>::simdjson_result(
  SIMDJSON_IMPLEMENTATION::ondemand::document_stream &&value
) noexcept :
    implementation_simdjson_result_base<SIMDJSON_IMPLEMENTATION::ondemand::document_stream>(
      std::forward<SIMDJSON_IMPLEMENTATION::ondemand::document_stream>(value)
    )
{
}
 
}
 
#endif // SIMDJSON_GENERIC_ONDEMAND_DOCUMENT_STREAM_INL_H