Ensure only leaf values can push the iterator

jkeiser · jkeiser · commit 4e822b51be06 · 2020-09-05T02:33:11.000-07:00
diff --git a/src/generic/ondemand/json_iterator-inl.h b/src/generic/ondemand/json_iterator-inl.h
@@ -19,11 +19,13 @@ simdjson_really_inline json_iterator &json_iterator::operator=(json_iterator &&o
   return *this;
 }
 simdjson_really_inline json_iterator::json_iterator(ondemand::parser *_parser) noexcept
-  : token_iterator(_parser->dom_parser.buf, _parser->dom_parser.structural_indexes.get()), parser{_parser}
+  : token_iterator(_parser->dom_parser.buf, _parser->dom_parser.structural_indexes.get()),
+    parser{_parser},
+    current_string_buf_loc{parser->string_buf.get()},
+    active_lease_depth{0}
 {
   // Release the string buf so it can be reused by the next document
   logger::log_headers();
-  current_string_buf_loc = parser->string_buf.get();
 }
 simdjson_really_inline json_iterator::~json_iterator() noexcept = default;
 
@@ -261,56 +263,69 @@ simdjson_really_inline bool json_iterator::is_alive() const noexcept {
 }
 
 simdjson_really_inline json_iterator_ref json_iterator::borrow() noexcept {
-  return json_iterator_ref(this);
+  SIMDJSON_ASSUME(active_lease_depth == 0);
+  const uint32_t child_depth = 1;
+  active_lease_depth = child_depth;
+  return json_iterator_ref(this, child_depth);
 }
 
 //
 // json_iterator_ref
 //
 simdjson_really_inline json_iterator_ref::json_iterator_ref() noexcept = default;
 simdjson_really_inline json_iterator_ref::json_iterator_ref(json_iterator_ref &&other) noexcept
-  : iter{std::forward<json_iterator_ref>(other).iter}
+  : iter{std::exchange(other.iter, nullptr)},
+    lease_depth{other.lease_depth}
 {
-  other.iter = nullptr;
 }
 simdjson_really_inline json_iterator_ref &json_iterator_ref::operator=(json_iterator_ref &&other) noexcept {
-  iter = std::forward<json_iterator_ref>(other).iter;
-  other.iter = nullptr;
+  SIMDJSON_ASSUME(!is_active());
+  iter = std::exchange(other.iter, nullptr);
+  lease_depth = other.lease_depth;
   return *this;
 }
-simdjson_really_inline json_iterator_ref::json_iterator_ref(json_iterator *_iter) noexcept
-  : iter{_iter}
+simdjson_really_inline json_iterator_ref::json_iterator_ref(json_iterator *_iter, uint32_t _lease_depth) noexcept
+  : iter{_iter},
+    lease_depth{_lease_depth}
 {
+  SIMDJSON_ASSUME(is_active());
 }
 simdjson_really_inline json_iterator_ref::~json_iterator_ref() noexcept {
   // The caller MUST consume their value and release the iterator before they die
-  SIMDJSON_ASSUME(!iter);
+  SIMDJSON_ASSUME(!is_alive());
 }
 
 simdjson_really_inline json_iterator_ref json_iterator_ref::borrow() noexcept {
-  return json_iterator_ref(iter);
+  SIMDJSON_ASSUME(is_active());
+  const uint32_t child_depth = lease_depth + 1;
+  iter->active_lease_depth = child_depth;
+  return json_iterator_ref(iter, child_depth);
 }
 simdjson_really_inline void json_iterator_ref::release() noexcept {
-  SIMDJSON_ASSUME(is_alive());
+  SIMDJSON_ASSUME(is_active());
+  iter->active_lease_depth = lease_depth - 1;
   iter = nullptr;
 }
 
 simdjson_really_inline json_iterator *json_iterator_ref::operator->() noexcept {
-  SIMDJSON_ASSUME(is_alive());
+  SIMDJSON_ASSUME(is_active());
   return iter;
 }
 simdjson_really_inline json_iterator &json_iterator_ref::operator*() noexcept {
-  SIMDJSON_ASSUME(is_alive());
+  SIMDJSON_ASSUME(is_active());
   return *iter;
 }
 simdjson_really_inline const json_iterator &json_iterator_ref::operator*() const noexcept {
-  SIMDJSON_ASSUME(is_alive());
+  SIMDJSON_ASSUME(is_active());
   return *iter;
 }
 
 simdjson_really_inline bool json_iterator_ref::is_alive() const noexcept {
   return iter != nullptr;
 }
+simdjson_really_inline bool json_iterator_ref::is_active() const noexcept {
+  return is_alive() && lease_depth == iter->active_lease_depth;
+}
 
 
 } // namespace ondemand
diff --git a/src/generic/ondemand/json_iterator.h b/src/generic/ondemand/json_iterator.h
@@ -137,6 +137,7 @@ class json_iterator : public token_iterator {
 protected:
   ondemand::parser *parser{};
   uint8_t *current_string_buf_loc{};
+  uint32_t active_lease_depth{};
 
   simdjson_really_inline json_iterator(ondemand::parser *parser) noexcept;
   template<int N>
@@ -150,6 +151,7 @@ class json_iterator : public token_iterator {
   friend class value;
   friend class raw_json_string;
   friend class parser;
+  friend class json_iterator_ref;
   friend simdjson_really_inline void logger::log_line(const json_iterator &iter, const char *title_prefix, const char *title, std::string_view detail, int delta, int depth_delta) noexcept;
 }; // json_iterator
 
@@ -170,10 +172,12 @@ class json_iterator_ref {
   simdjson_really_inline const json_iterator &operator*() const noexcept;
 
   simdjson_really_inline bool is_alive() const noexcept;
+  simdjson_really_inline bool is_active() const noexcept;
 
 private:
-  simdjson_really_inline json_iterator_ref(json_iterator *iter) noexcept;
-  json_iterator *iter;
+  simdjson_really_inline json_iterator_ref(json_iterator *iter, uint32_t lease_depth) noexcept;
+  json_iterator *iter{};
+  uint32_t lease_depth{};
 
   friend class json_iterator;
 }; // class json_iterator_ref
