bugfixes

crypto/vm/dict.cpp

@@ -1308,6 +1308,294 @@ Ref<Cell> Dictionary::extract_minmax_key_ref(td::BitPtr key_buffer, int key_len,
   return extract_value_ref(extract_minmax_key(key_buffer, key_len, fetch_max, invert_first));
 }
 
+/*
+ *
+ *   ITERATORS (to be moved into a separate file)
+ *
+ */
+
+bool DictIterator::prevalidate(int mode) {
+  if (key_bits_ <= 0 || key_bits_ > Dictionary::max_key_bits) {
+    reset();
+    flags_ &= ~f_valid;
+    key_bits_ = 0;
+    return false;
+  } else {
+    if (mode >= 0) {
+      order_ = -(mode & 1) ^ ((mode >> 1) & 1);
+    }
+    flags_ |= f_valid;
+    return true;
+  }
+}
+
+bool DictIterator::bind(const DictionaryFixed& dict, int do_rewind) {
+  if (!is_valid() || !is_bound_to(dict)) {
+    return false;
+  }
+  dict_ = &dict;
+  label_mode_ = dict.label_mode();
+  return !do_rewind || rewind(do_rewind < 0);
+}
+
+bool DictIterator::rebind_to(const DictionaryFixed& dict, int do_rewind) {
+  reset();
+  dict_ = &dict;
+  label_mode_ = dict.label_mode();
+  root_ = dict.get_root_cell();
+  key_bits_ = dict.get_key_bits();
+  flags_ &= 3;
+  return prevalidate() && (!do_rewind || rewind(do_rewind < 0));
+}
+
+int DictIterator::compare_keys(td::ConstBitPtr a, td::ConstBitPtr b) const {
+  if (!key_bits_) {
+    return 0;
+  }
+  int c = (int)*a - (int)*b;
+  if (c) {
+    return (order_ & 1) ? -c : c;
+  }
+  c = a.compare(b, key_bits_);
+  return order_ >= 0 ? c : -c;
+}
+
+bool DictIterator::dive(int mode) {
+  int n = key_bits_, m = 0;
+  Ref<Cell> node = path_.empty() ? root_ : path_.back().next;
+  if (!path_.empty()) {
+    m = path_.back().pos + 1;
+    n -= m;
+    mode >>= 1;
+  }
+  // similar to dict_lookup_minmax: create new path down until the leaf
+  while (1) {
+    LabelParser label{std::move(node), n, label_mode_};
+    int l = label.extract_label_to(key(m));
+    assert(l >= 0 && l <= n);
+    m += l;
+    n -= l;
+    if (!n) {
+      leaf_ = std::move(label.remainder);
+      return true;
+    }
+    if (l) {
+      mode >>= 1;
+    }
+    int bit = mode & 1;
+    node = label.remainder->prefetch_ref(bit);
+    auto alt = label.remainder->prefetch_ref(1 - bit);
+    path_.emplace_back(node, std::move(alt), m, bit);
+    *key(m++) = (bool)bit;
+    --n;
+    mode >>= 1;
+  }
+}
+
+bool DictIterator::rewind(bool to_end) {
+  if (!is_valid()) {
+    return false;
+  }
+  if (root_.is_null()) {
+    return true;
+  }
+  auto node = root_;
+  int k = 0, mode = order_ ^ -(int)to_end;
+  // NB: can optimize by reusing several first entries of current path_
+  while (k < (int)path_.size()) {
+    auto& pe = path_[k++];
+    assert(pe.pos >= 0 && pe.pos < key_bits_);
+    if (pe.pos) {
+      mode >>= 1;
+    }
+    if (pe.v != (bool)(mode & 1)) {
+      // went different way at this node before, rotate and stop going down
+      pe.rotate(key());
+      leaf_.clear();
+      path_.resize(k);  // drop the remainder of the original path after first incorrect branch
+      return dive(mode);
+    }
+    mode >>= 1;
+  }
+  return !eof() || dive(mode);
+}
+
+bool DictIterator::next(bool go_back) {
+  if (!is_valid() || root_.is_null() || eof()) {
+    return false;
+  }
+  leaf_.clear();
+  int mode = order_ ^ -(int)go_back;
+  while (!path_.empty()) {
+    auto& pe = path_.back();
+    int bit = (mode >> (pe.pos > 0)) & 1;
+    if (pe.v == bit) {
+      pe.rotate(key());
+      return dive(mode);
+    }
+    path_.pop_back();
+  }
+  return false;
+}
+
+bool DictIterator::lookup(td::ConstBitPtr pos, int pos_bits, bool strict_after, bool backw) {
+  if (!is_valid() || root_.is_null() || pos_bits < 0 || pos_bits > key_bits_) {
+    return false;
+  }
+  int fill_mode = -(strict_after ^ backw) ^ order_;
+  if (!eof()) {
+    // reuse part of current path
+    std::size_t bp0 = 0;
+    int bp;
+    if (!key().compare(pos, pos_bits, &bp0)) {
+      bp = pos_bits;
+      if (bp >= key_bits_) {
+        // already at the desired element
+        return !strict_after || next(backw);
+      }
+    } else {
+      bp = (int)bp0;
+    }
+    int k = 0;
+    while (k < (int)path_.size() && path_[k].pos <= bp) {
+      auto& pe = path_[k];
+      if (pe.pos == bp) {
+        if (bp < pos_bits || pe.v != ((fill_mode >> (bp > 0)) & 1)) {
+          // rotate the last path element if it branched in incorrect direction
+          path_[k++].rotate(key());
+        }
+        break;
+      }
+      ++k;
+    }
+    path_.resize(k);  // drop the remainder of the path
+  }
+  int m = 0, n = key_bits_;
+  auto node = path_.empty() ? root_ : path_.back().next;
+  if (!path_.empty()) {
+    m = path_.back().pos + 1;  // m <= pos_bits + 1
+    n -= m;
+  }
+  int mode = -backw ^ order_, action = 0;
+  while (m < pos_bits && !action) {
+    LabelParser label{std::move(node), n, label_mode_};
+    int pfx_len = label.common_prefix_len(pos + m, pos_bits - m);
+    int l = label.extract_label_to(key() + m);
+    assert(pfx_len >= 0 && pfx_len <= label.l_bits && label.l_bits <= n);
+    if (pfx_len == pos_bits - m) {
+      // end of given position prefix
+      if (strict_after) {
+        // have to backtrace
+        action = 2;
+        break;
+      } else {
+        // label has correct prefix, register and dive down
+        action = 1;
+      }
+    } else if (pfx_len < l) {
+      // all subtree is either smaller or larger than required
+      if (pos[m + pfx_len] != ((mode >> (int)(m + pfx_len > 0)) & 1)) {
+        // label smaller than required, have to backtrace
+        action = 2;
+        break;
+      } else {
+        // label larger than required, register node and dive down
+        action = 1;
+      }
+    }
+    // if we are here, then either action=1 (dive down activated)
+    // ... or l = pfx_len < pos_bits - m
+    // continue going down
+    m += l;
+    n -= l;
+    if (!n) {
+      // key found in a leaf
+      leaf_ = std::move(label.remainder);
+      return true;
+    }
+    bool bit = action ? ((mode >> (m > 0)) & 1) : pos[m];
+    node = label.remainder->prefetch_ref(bit);
+    auto alt = label.remainder->prefetch_ref(1 - bit);
+    path_.emplace_back(node, std::move(alt), m, bit);
+    *key(m++) = (bool)bit;
+    --n;
+  }
+  if (!action) {
+    action = (strict_after ? 2 : 1);
+  }
+  if (action == 2) {
+    // have to backtrace to the "next" larger branch
+    // similar to next()
+    leaf_.clear();
+    while (!path_.empty()) {
+      auto& pe = path_.back();
+      int bit = (mode >> (pe.pos > 0)) & 1;
+      if (pe.v == bit) {
+        pe.rotate(key());
+        return dive(mode);
+      }
+      path_.pop_back();
+    }
+    return false;  // eof: no suitable element
+  }
+  // action=1, dive down
+  return dive(mode);
+}
+
+DictIterator DictionaryFixed::null_iterator() {
+  force_validate();
+  return DictIterator{*this};
+}
+
+DictIterator DictionaryFixed::make_iterator(int mode) {
+  force_validate();
+  DictIterator it{*this, mode};
+  it.rewind();
+  return it;
+}
+
+DictIterator DictionaryFixed::init_iterator(bool backw, bool invert_first) {
+  return make_iterator((int)backw + 2 * (int)invert_first);
+}
+
+DictIterator DictionaryFixed::begin() {
+  return init_iterator();
+}
+
+DictIterator DictionaryFixed::end() {
+  return null_iterator();
+}
+
+DictIterator DictionaryFixed::cbegin() {
+  return begin();
+}
+
+DictIterator DictionaryFixed::cend() {
+  return end();
+}
+
+DictIterator DictionaryFixed::rbegin() {
+  return init_iterator(true);
+}
+
+DictIterator DictionaryFixed::rend() {
+  return null_iterator();
+}
+
+DictIterator DictionaryFixed::crbegin() {
+  return rbegin();
+}
+
+DictIterator DictionaryFixed::crend() {
+  return rend();
+}
+
+/*
+ *
+ *   END (ITERATORS)
+ *
+ */
+
 std::pair<Ref<Cell>, bool> DictionaryFixed::extract_prefix_subdict_internal(Ref<Cell> dict, td::ConstBitPtr prefix,
                                                                             int prefix_len, bool remove_prefix) const {
   if (is_empty() || prefix_len <= 0) {
@@ -2309,6 +2597,14 @@ Ref<CellSlice> AugmentedDictionary::get_node_extra(Ref<Cell> cell_ref, int n) co
   return {};
 }
 
+Ref<CellSlice> AugmentedDictionary::extract_leaf_value(Ref<CellSlice> leaf) const {
+  if (leaf.not_null() && aug.skip_extra(leaf.write())) {
+    return std::move(leaf);
+  } else {
+    return Ref<CellSlice>{};
+  }
+}
+
 Ref<CellSlice> AugmentedDictionary::get_root_extra() const {
   return get_node_extra(root_cell, key_bits);
 }

crypto/vm/dict.h

@@ -14,7 +14,7 @@
     You should have received a copy of the GNU Lesser General Public License
     along with TON Blockchain Library.  If not, see <http://www.gnu.org/licenses/>.
 
-    Copyright 2017-2019 Telegram Systems LLP
+    Copyright 2017-2020 Telegram Systems LLP
 */
 #pragma once
 #include "common/bitstring.h"
@@ -171,6 +171,13 @@ class DictionaryBase {
   }
 };
 
+class DictIterator;
+
+template <typename T>
+std::pair<T, int> dict_range(T&& dict, bool rev = false, bool sgnd = false) {
+  return std::pair<T, int>{std::forward<T>(dict), (int)rev + 2 * (int)sgnd};
+}
+
 class DictionaryFixed : public DictionaryBase {
  public:
   typedef std::function<int(vm::CellSlice&, td::ConstBitPtr, int)> filter_func_t;
@@ -199,6 +206,9 @@ class DictionaryFixed : public DictionaryBase {
   static BitSlice integer_key(td::RefInt256 x, unsigned n, bool sgnd = true, unsigned char buffer[128] = 0,
                               bool quiet = false);
   static bool integer_key_simple(td::RefInt256 x, unsigned n, bool sgnd, td::BitPtr buffer, bool quiet = false);
+  td::RefInt256 key_as_integer(td::ConstBitPtr key, bool sgnd = false) const {
+    return td::bits_to_refint(key, key_bits, sgnd);
+  }
   bool key_exists(td::ConstBitPtr key, int key_len);
   bool int_key_exists(long long key);
   bool uint_key_exists(unsigned long long key);
@@ -221,6 +231,17 @@ class DictionaryFixed : public DictionaryBase {
   bool scan_diff(DictionaryFixed& dict2, const scan_diff_func_t& diff_func, int check_augm = 0);
   bool validate_check(const foreach_func_t& foreach_func, bool invert_first = false);
   bool validate_all();
+  DictIterator null_iterator();
+  DictIterator init_iterator(bool backw = false, bool invert_first = false);
+  DictIterator make_iterator(int mode);
+  DictIterator begin();
+  DictIterator end();
+  DictIterator cbegin();
+  DictIterator cend();
+  DictIterator rbegin();
+  DictIterator rend();
+  DictIterator crbegin();
+  DictIterator crend();
   template <typename T>
   bool key_exists(const T& key) {
     return key_exists(key.bits(), key.size());
@@ -247,6 +268,9 @@ class DictionaryFixed : public DictionaryBase {
   virtual int label_mode() const {
     return dict::LabelParser::chk_all;
   }
+  virtual Ref<CellSlice> extract_leaf_value(Ref<CellSlice> leaf) const {
+    return leaf;
+  }
   virtual Ref<Cell> finish_create_leaf(CellBuilder& cb, const CellSlice& value) const;
   virtual Ref<Cell> finish_create_fork(CellBuilder& cb, Ref<Cell> c1, Ref<Cell> c2, int n) const;
   virtual bool check_fork(CellSlice& cs, Ref<Cell> c1, Ref<Cell> c2, int n) const {
@@ -259,6 +283,7 @@ class DictionaryFixed : public DictionaryBase {
     return check_leaf(cs_ref.write(), key, key_len);
   }
   bool check_fork_raw(Ref<CellSlice> cs_ref, int n) const;
+  friend class DictIterator;
 
  private:
   std::pair<Ref<CellSlice>, Ref<Cell>> dict_lookup_delete(Ref<Cell> dict, td::ConstBitPtr key, int n) const;
@@ -277,6 +302,148 @@ class DictionaryFixed : public DictionaryBase {
                            const foreach_func_t& foreach_func, bool invert_first = false) const;
 };
 
+class DictIterator {
+  const DictionaryFixed* dict_{nullptr};
+  Ref<Cell> root_;
+  int label_mode_{dict::LabelParser::chk_size};
+  int key_bits_;
+  int flags_;
+  int order_;
+  unsigned char key_buffer[DictionaryBase::max_key_bytes];
+  bool prevalidate(int mode = -1);
+  enum { f_valid = 4 };
+
+ protected:
+  struct Fork {
+    Ref<Cell> next, alt;
+    int pos;
+    bool v;
+    Fork() : pos(-1) {
+    }
+    Fork(Ref<Cell> _next, Ref<Cell> _alt, int _pos, bool _v)
+        : next(std::move(_next)), alt(std::move(_alt)), pos(_pos), v(_v) {
+    }
+    void rotate(td::BitPtr key) {
+      std::swap(next, alt);
+      key[pos] = (v ^= true);
+    }
+  };
+  std::vector<Fork> path_;
+  Ref<CellSlice> leaf_;
+
+  td::BitPtr key(int offs = 0) {
+    return td::BitPtr{key_buffer, offs};
+  }
+  td::ConstBitPtr key(int offs = 0) const {
+    return td::ConstBitPtr{key_buffer, offs};
+  }
+  td::ConstBitPtr ckey(int offs = 0) const {
+    return td::ConstBitPtr{key_buffer, offs};
+  }
+
+ public:
+  DictIterator() : key_bits_(0), flags_(0), order_(0) {
+  }
+  // mode: 0 = bidir, +4 = fwd only, +8 = back only; +1 = reverse directions, +2 = signed int keys
+  enum { it_reverse = 1, it_signed = 2 };
+  DictIterator(Ref<Cell> root_cell, int key_bits, int mode = 0)
+      : root_(std::move(root_cell)), key_bits_(key_bits), flags_(mode >> 2) {
+    prevalidate(mode & 3);
+  }
+  DictIterator(const DictionaryFixed& dict, int mode = 0)
+      : DictIterator(dict.get_root_cell(), dict.get_key_bits(), mode) {
+    dict_ = &dict;
+    label_mode_ = dict.label_mode();
+  }
+  bool is_valid() const {
+    return flags_ & f_valid;
+  }
+  bool eof() const {
+    return leaf_.is_null();
+  }
+  bool reset() {
+    dict_ = nullptr;
+    root_.clear();
+    path_.clear();
+    leaf_.clear();
+    return true;
+  }
+  td::ConstBitPtr cur_pos() const {
+    return eof() ? td::ConstBitPtr{nullptr} : key();
+  }
+  Ref<Cell> get_root_cell() const {
+    return root_;
+  }
+  int get_key_bits() const {
+    return key_bits_;
+  }
+  bool is_bound() const {
+    return dict_;
+  }
+  bool is_bound_to(const DictionaryFixed& dict) const {
+    return root_.not_null() == dict.get_root_cell().not_null() &&
+           (root_.not_null() ? root_.get() == dict.get_root_cell().get() : key_bits_ == dict.get_key_bits());
+  }
+  bool bind(const DictionaryFixed& dict, int do_rewind = 0);
+  bool rebind_to(const DictionaryFixed& dict, int do_rewind = 0);
+  bool rewind(bool to_end = false);
+  bool next(bool backw = false);
+  bool prev() {
+    return next(true);
+  }
+  bool lookup(td::ConstBitPtr pos, int pos_bits, bool strict_after = false, bool backw = false);
+  template <typename T>
+  bool lookup(const T& key, bool strict_after = false, bool backw = false) {
+    return lookup(key.bits(), key.size(), strict_after, backw);
+  }
+  Ref<CellSlice> cur_value() const {
+    return dict_ ? dict_->extract_leaf_value(leaf_) : Ref<CellSlice>{};
+  }
+  Ref<CellSlice> cur_value_raw() const {
+    return leaf_;
+  }
+  std::pair<td::ConstBitPtr, Ref<CellSlice>> operator*() const {
+    return std::make_pair(cur_pos(), cur_value());
+  }
+  bool bound_to_same(const DictIterator& other) const {
+    return dict_ && dict_ == other.dict_;
+  }
+  bool operator==(const DictIterator& other) const {
+    return bound_to_same(other) && eof() == other.eof() && (eof() || key().equals(other.key(), key_bits_));
+  }
+  int compare_keys(td::ConstBitPtr a, td::ConstBitPtr b) const;
+  bool operator<(const DictIterator& other) const {
+    return bound_to_same(other) && !eof() && (other.eof() || compare_keys(key(), other.key()) < 0);
+  }
+  bool operator!=(const DictIterator& other) const {
+    return !(operator==(other));
+  }
+  bool operator>(const DictIterator& other) const {
+    return other < *this;
+  }
+  DictIterator& operator++() {
+    next();
+    return *this;
+  }
+  DictIterator& operator--() {
+    next(true);
+    return *this;
+  }
+
+ private:
+  bool dive(int mode);
+};
+
+template <typename T>
+DictIterator begin(std::pair<T, int> dictm) {
+  return dictm.first.make_iterator(dictm.second);
+}
+
+template <typename T>
+DictIterator end(std::pair<T, int> dictm) {
+  return dictm.first.null_iterator();
+}
+
 class Dictionary final : public DictionaryFixed {
  public:
   typedef std::function<bool(CellBuilder&, Ref<CellSlice>)> simple_map_func_t;
@@ -356,6 +523,9 @@ class Dictionary final : public DictionaryFixed {
   Ref<CellSlice> lookup_set_builder(const T& key, Ref<vm::CellBuilder> val_ref, SetMode mode = SetMode::Set) {
     return lookup_set_builder(key.bits(), key.size(), std::move(val_ref), mode);
   }
+  auto range(bool rev = false, bool sgnd = false) {
+    return dict_range(*this, rev, sgnd);
+  }
 
  private:
   bool check_fork(CellSlice& cs, Ref<Cell> c1, Ref<Cell> c2, int n) const override {
@@ -454,6 +624,9 @@ class AugmentedDictionary final : public DictionaryFixed {
   Ref<Cell> lookup_delete_ref(const T& key) {
     return lookup_delete_ref(key.bits(), key.size());
   }
+  auto range(bool rev = false, bool sgnd = false) {
+    return dict_range(*this, rev, sgnd);
+  }
 
   Ref<CellSlice> extract_value(Ref<CellSlice> value_extra) const;
   Ref<Cell> extract_value_ref(Ref<CellSlice> value_extra) const;
@@ -463,6 +636,7 @@ class AugmentedDictionary final : public DictionaryFixed {
  private:
   bool compute_root() const;
   Ref<CellSlice> get_node_extra(Ref<Cell> cell_ref, int n) const;
+  Ref<CellSlice> extract_leaf_value(Ref<CellSlice> leaf) const override;
   bool check_leaf(CellSlice& cs, td::ConstBitPtr key, int key_len) const override;
   bool check_fork(CellSlice& cs, Ref<Cell> c1, Ref<Cell> c2, int n) const override;
   Ref<Cell> finish_create_leaf(CellBuilder& cb, const CellSlice& value) const override;