[vector] Allow to iterate with mutables (google#7586)

Co-authored-by: ArnaudD-FR <arnaud.desmier@gmail.com> Co-authored-by: Derek Bailey <derekbailey@google.com>
zdrealityhub · Oct 21, 2022 · 5b3fadc · 5b3fadc
1 parent 872a497
commit 5b3fadc
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 6 deletions.
diff --git a/include/flatbuffers/array.h b/include/flatbuffers/array.h
@@ -35,7 +35,7 @@ template<typename T, uint16_t length> class Array {
  public:
   typedef uint16_t size_type;
   typedef typename IndirectHelper<IndirectHelperType>::return_type return_type;
-  typedef VectorIterator<T, return_type> const_iterator;
+  typedef VectorConstIterator<T, return_type> const_iterator;
   typedef VectorReverseIterator<const_iterator> const_reverse_iterator;
 
   // If T is a LE-scalar or a struct (!scalar_tag::value).

diff --git a/include/flatbuffers/buffer.h b/include/flatbuffers/buffer.h
@@ -76,6 +76,9 @@ template<typename T> struct IndirectHelper {
   static return_type Read(const uint8_t *p, uoffset_t i) {
     return EndianScalar((reinterpret_cast<const T *>(p))[i]);
   }
+  static return_type Read(uint8_t *p, uoffset_t i) {
+    return Read(const_cast<const uint8_t *>(p), i);
+  }
 };
 template<typename T> struct IndirectHelper<Offset<T>> {
   typedef const T *return_type;
@@ -85,13 +88,20 @@ template<typename T> struct IndirectHelper<Offset<T>> {
     p += i * sizeof(uoffset_t);
     return reinterpret_cast<return_type>(p + ReadScalar<uoffset_t>(p));
   }
+  static mutable_return_type Read(uint8_t *p, uoffset_t i) {
+    p += i * sizeof(uoffset_t);
+    return reinterpret_cast<mutable_return_type>(p + ReadScalar<uoffset_t>(p));
+  }
 };
 template<typename T> struct IndirectHelper<const T *> {
   typedef const T *return_type;
   typedef T *mutable_return_type;
   static const size_t element_stride = sizeof(T);
   static return_type Read(const uint8_t *p, uoffset_t i) {
-    return reinterpret_cast<const T *>(p + i * sizeof(T));
+    return reinterpret_cast<return_type>(p + i * sizeof(T));
+  }
+  static mutable_return_type Read(uint8_t *p, uoffset_t i) {
+    return reinterpret_cast<mutable_return_type>(p + i * sizeof(T));
   }
 };
 

diff --git a/include/flatbuffers/vector.h b/include/flatbuffers/vector.h
@@ -27,14 +27,15 @@ struct String;
 
 // An STL compatible iterator implementation for Vector below, effectively
 // calling Get() for every element.
-template<typename T, typename IT> struct VectorIterator {
+template<typename T, typename IT, typename Data = uint8_t *>
+struct VectorIterator {
   typedef std::random_access_iterator_tag iterator_category;
   typedef IT value_type;
   typedef ptrdiff_t difference_type;
   typedef IT *pointer;
   typedef IT &reference;
 
-  VectorIterator(const uint8_t *data, uoffset_t i)
+  VectorIterator(Data data, uoffset_t i)
       : data_(data + IndirectHelper<T>::element_stride * i) {}
   VectorIterator(const VectorIterator &other) : data_(other.data_) {}
   VectorIterator() : data_(nullptr) {}
@@ -116,9 +117,12 @@ template<typename T, typename IT> struct VectorIterator {
   }
 
  private:
-  const uint8_t *data_;
+  Data data_;
 };
 
+template<typename T, typename IT>
+using VectorConstIterator = VectorIterator<T, IT, const uint8_t *>;
+
 template<typename Iterator>
 struct VectorReverseIterator : public std::reverse_iterator<Iterator> {
   explicit VectorReverseIterator(Iterator iter)
@@ -145,7 +149,7 @@ template<typename T> class Vector {
  public:
   typedef VectorIterator<T, typename IndirectHelper<T>::mutable_return_type>
       iterator;
-  typedef VectorIterator<T, typename IndirectHelper<T>::return_type>
+  typedef VectorConstIterator<T, typename IndirectHelper<T>::return_type>
       const_iterator;
   typedef VectorReverseIterator<iterator> reverse_iterator;
   typedef VectorReverseIterator<const_iterator> const_reverse_iterator;

diff --git a/tests/monster_test.cpp b/tests/monster_test.cpp
@@ -439,6 +439,16 @@ void MutateFlatBuffersTest(uint8_t *flatbuf, std::size_t length) {
   TEST_EQ(first->hp(), 0);
   first->mutate_hp(1000);
 
+  // Test for each loop over mutable entries
+  for (auto item: *tables)
+  {
+    TEST_EQ(item->hp(), 1000);
+    item->mutate_hp(0);
+    TEST_EQ(item->hp(), 0);
+    item->mutate_hp(1000);
+    break; // one iteration is enough, just testing compilation
+  }
+
   // Mutate via LookupByKey
   TEST_NOTNULL(tables->MutableLookupByKey("Barney"));
   TEST_EQ(static_cast<Monster *>(nullptr),