Make FlatSet and FlatMap movable.

PiperOrigin-RevId: 196156010

5 files changed, 79 insertions, 10 deletions

diff --git a/tensorflow/core/lib/gtl/flatmap.h b/tensorflow/core/lib/gtl/flatmap.h
index 889d2ddaa6..9dc439c163 100644
--- a/tensorflow/core/lib/gtl/flatmap.h
+++ b/tensorflow/core/lib/gtl/flatmap.h
@@ -76,6 +76,10 @@ class FlatMap {
 
   FlatMap(const FlatMap& src) : rep_(src.rep_) {}
 
+  // Move constructor leaves src in a valid but unspecified state (same as
+  // std::unordered_map).
+  FlatMap(FlatMap&& src) : rep_(std::move(src.rep_)) {}
+
   template <typename InputIter>
   FlatMap(InputIter first, InputIter last, size_t N = 1,
           const Hash& hf = Hash(), const Eq& eq = Eq())
@@ -92,6 +96,13 @@ class FlatMap {
     return *this;
   }
 
+  // Move-assignment operator leaves src in a valid but unspecified state (same
+  // as std::unordered_map).
+  FlatMap& operator=(FlatMap&& src) {
+    rep_.MoveFrom(std::move(src.rep_));
+    return *this;
+  }
+
   ~FlatMap() {}
 
   void swap(FlatMap& x) { rep_.swap(x.rep_); }
diff --git a/tensorflow/core/lib/gtl/flatmap_test.cc b/tensorflow/core/lib/gtl/flatmap_test.cc
index 0901eba926..0fd22ab37b 100644
--- a/tensorflow/core/lib/gtl/flatmap_test.cc
+++ b/tensorflow/core/lib/gtl/flatmap_test.cc
@@ -656,19 +656,33 @@ TEST(FlatMap, UniqueMap) {
   }
   EXPECT_EQ(map.size(), N);
 
+  // move constructor
+  UniqMap map2(std::move(map));
+
   // Lookups
   for (int i = 0; i < N; i++) {
-    EXPECT_EQ(*map.at(MakeUniq(i)), i + 100);
+    EXPECT_EQ(*map2.at(MakeUniq(i)), i + 100);
   }
 
+  // move assignment
+  UniqMap map3;
+  map3 = std::move(map2);
+
   // find+erase
-  EXPECT_EQ(map.count(MakeUniq(2)), 1);
-  map.erase(MakeUniq(2));
-  EXPECT_EQ(map.count(MakeUniq(2)), 0);
+  EXPECT_EQ(map3.count(MakeUniq(2)), 1);
+  map3.erase(MakeUniq(2));
+  EXPECT_EQ(map3.count(MakeUniq(2)), 0);
 
   // clear
-  map.clear();
-  EXPECT_EQ(map.size(), 0);
+  map3.clear();
+  EXPECT_EQ(map3.size(), 0);
+
+  // Check that moved-from maps are in a valid (though unspecified) state.
+  EXPECT_GE(map.size(), 0);
+  EXPECT_GE(map2.size(), 0);
+  // This insert should succeed no matter what state `map` is in, because
+  // MakeUniq(-1) is never called above: This key can't possibly exist.
+  EXPECT_TRUE(map.emplace(MakeUniq(-1), MakeUniq(-1)).second);
 }
 
 TEST(FlatMap, UniqueMapIter) {
diff --git a/tensorflow/core/lib/gtl/flatrep.h b/tensorflow/core/lib/gtl/flatrep.h
index 0d7e7487fc..65a076b0f3 100644
--- a/tensorflow/core/lib/gtl/flatrep.h
+++ b/tensorflow/core/lib/gtl/flatrep.h
@@ -51,10 +51,23 @@ class FlatRep {
   FlatRep(size_t N, const Hash& hf, const Eq& eq) : hash_(hf), equal_(eq) {
     Init(N);
   }
-  explicit FlatRep(const FlatRep& src) : hash_(src.hash_), equal_(src.equal_) {
+  FlatRep(const FlatRep& src) : hash_(src.hash_), equal_(src.equal_) {
     Init(src.size());
     CopyEntries(src.array_, src.end_, CopyEntry());
   }
+
+  FlatRep(FlatRep&& src)
+      // Copy rather than move src.hash_ and src.equal_.  This is necessary to
+      // leave src in a valid state -- otherwise e.g. if hash_ is an
+      // std::function, moving it would null it out.
+      : hash_(src.hash_), equal_(src.equal_) {
+    // TODO(jlebar): Init(1) still allocates some memory, so this isn't as cheap
+    // as it could be.  The fundamental problem is that we need to leave src in
+    // a valid state, and FlatRep *always* owns a nonzero amount of memory.
+    Init(1);
+    swap(src);
+  }
+
   ~FlatRep() {
     clear_no_resize();
     delete[] array_;
@@ -78,6 +91,12 @@ class FlatRep {
     }
   }
 
+  void MoveFrom(FlatRep&& src) {
+    if (this != &src) {
+      swap(src);
+    }
+  }
+
   void clear_no_resize() {
     for (Bucket* b = array_; b != end_; b++) {
       for (uint32 i = 0; i < kWidth; i++) {
diff --git a/tensorflow/core/lib/gtl/flatset.h b/tensorflow/core/lib/gtl/flatset.h
index f31e3abe41..311b7abe4d 100644
--- a/tensorflow/core/lib/gtl/flatset.h
+++ b/tensorflow/core/lib/gtl/flatset.h
@@ -59,6 +59,10 @@ class FlatSet {
 
   FlatSet(const FlatSet& src) : rep_(src.rep_) {}
 
+  // Move constructor leaves src in a valid but unspecified state (same as
+  // std::unordered_set).
+  FlatSet(FlatSet&& src) : rep_(std::move(src.rep_)) {}
+
   template <typename InputIter>
   FlatSet(InputIter first, InputIter last, size_t N = 1,
           const Hash& hf = Hash(), const Eq& eq = Eq())
@@ -75,6 +79,13 @@ class FlatSet {
     return *this;
   }
 
+  // Move-assignment operator leaves src in a valid but unspecified state (same
+  // as std::unordered_set).
+  FlatSet& operator=(FlatSet&& src) {
+    rep_.MoveFrom(std::move(src.rep_));
+    return *this;
+  }
+
   ~FlatSet() {}
 
   void swap(FlatSet& x) { rep_.swap(x.rep_); }
diff --git a/tensorflow/core/lib/gtl/flatset_test.cc b/tensorflow/core/lib/gtl/flatset_test.cc
index 010b4bb5df..8f8a953568 100644
--- a/tensorflow/core/lib/gtl/flatset_test.cc
+++ b/tensorflow/core/lib/gtl/flatset_test.cc
@@ -552,18 +552,32 @@ TEST(FlatSet, UniqueSet) {
   }
   EXPECT_EQ(set.size(), N);
 
+  // Move constructor
+  UniqSet set2(std::move(set));
+
   // Lookups
   for (int i = 0; i < N; i++) {
-    EXPECT_EQ(set.count(MakeUniq(i)), 1);
+    EXPECT_EQ(set2.count(MakeUniq(i)), 1);
   }
 
+  // Move-assignment operator
+  UniqSet set3;
+  set3 = std::move(set2);
+
   // erase
-  set.erase(MakeUniq(2));
-  EXPECT_EQ(set.count(MakeUniq(2)), 0);
+  set3.erase(MakeUniq(2));
+  EXPECT_EQ(set3.count(MakeUniq(2)), 0);
 
   // clear
   set.clear();
   EXPECT_EQ(set.size(), 0);
+
+  // Check that moved-from sets are in a valid (though unspecified) state.
+  EXPECT_GE(set.size(), 0);
+  EXPECT_GE(set2.size(), 0);
+  // This insert should succeed no matter what state `set` is in, because
+  // MakeUniq(-1) is never called above: This key can't possibly exist.
+  EXPECT_TRUE(set.emplace(MakeUniq(-1)).second);
 }
 
 TEST(FlatSet, UniqueSetIter) {