From 6a855f528bd5431d90e984c0c110df9683bd3e05 Mon Sep 17 00:00:00 2001
From: John Keiser <john@johnkeiser.com>
Date: Mon, 31 Aug 2020 13:09:19 -0700
Subject: [PATCH] Ensure only leaf values can push the iterator

---
 src/generic/ondemand/json_iterator-inl.h | 45 ++++++++++++++++--------
 src/generic/ondemand/json_iterator.h     |  8 +++--
 2 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/src/generic/ondemand/json_iterator-inl.h b/src/generic/ondemand/json_iterator-inl.h
index 1c4a9a1e..b562f6cd 100644
--- 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,7 +263,10 @@ 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);
 }
 
 //
@@ -269,48 +274,58 @@ simdjson_really_inline json_iterator_ref json_iterator::borrow() noexcept {
 //
 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
index ae7b3361..b3b0629a 100644
--- a/src/generic/ondemand/json_iterator.h
+++ b/src/generic/ondemand/json_iterator.h
@@ -137,6 +137,7 @@ public:
 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 @@ protected:
   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 @@ public:
   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