Add Extent map

When installing OTA, we stream operations as data becomes available. For
each operation, update_engine needs to look at the merge sequence and
figure out which blocks should be converted to COW_XOR. To do that, we
need a data structure for querying a disjoint set of extents.

Test: th
Bug: 177104308

Change-Id: I14ec9d72a32e859b65516b894ca94d0153aa7e17
diff --git a/Android.bp b/Android.bp
index 51105eb..b3fbccf 100644
--- a/Android.bp
+++ b/Android.bp
@@ -804,6 +804,7 @@
         "payload_consumer/partition_writer_unittest.cc",
         "payload_consumer/extent_reader_unittest.cc",
         "payload_consumer/extent_writer_unittest.cc",
+        "payload_consumer/extent_map_unittest.cc",
         "payload_consumer/snapshot_extent_writer_unittest.cc",
         "payload_consumer/fake_file_descriptor.cc",
         "payload_consumer/file_descriptor_utils_unittest.cc",
diff --git a/common/utils.h b/common/utils.h
index 59f236e..4e78097 100644
--- a/common/utils.h
+++ b/common/utils.h
@@ -465,6 +465,17 @@
   DISALLOW_COPY_AND_ASSIGN(ScopedActionCompleter);
 };
 
+// Simple wrapper for creating a slice of some container,
+// similar to string_view but for other containers.
+template <typename T>
+struct Range {
+  Range(T t1, T t2) : t1_(t1), t2_(t2) {}
+  constexpr auto begin() const noexcept { return t1_; }
+  constexpr auto end() const noexcept { return t2_; }
+  T t1_;
+  T t2_;
+};
+
 }  // namespace chromeos_update_engine
 
 #define TEST_AND_RETURN_FALSE_ERRNO(_x)                              \
diff --git a/payload_consumer/extent_map.h b/payload_consumer/extent_map.h
new file mode 100644
index 0000000..aeb7066
--- /dev/null
+++ b/payload_consumer/extent_map.h
@@ -0,0 +1,90 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
+#define UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
+
+#include <functional>
+#include <map>
+#include <utility>
+
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+// Data structure for storing a disjoint set of extents.
+// Currently the only usecase is for VABCPartitionWriter to keep track of which
+// block belongs to which merge operation. Therefore this class only contains
+// the minimal set of functions needed.
+template <typename T, typename Comparator = ExtentLess>
+class ExtentMap {
+ public:
+  bool AddExtent(const Extent& extent, T&& value) {
+    if (Get(extent)) {
+      return false;
+    }
+    const auto& [it, inserted] = map_.insert({extent, std::forward<T>(value)});
+    if (inserted) {
+      set_.AddExtent(extent);
+    }
+    return inserted;
+  }
+
+  size_t size() const { return map_.size(); }
+
+  // Return a pointer to entry which is intersecting |extent|. If T is already
+  // a pointer type, return T on success. This function always return
+  // |nullptr| on failure. Therefore you cannot store nullptr as an entry.
+  std::optional<T> Get(const Extent& extent) const {
+    const auto it = map_.find(extent);
+    if (it == map_.end()) {
+      LOG_IF(WARNING, set_.OverlapsWithExtent(extent))
+          << "Looking up a partially intersecting extent isn't supported by "
+             "this data structure.";
+      return {};
+    }
+    return {it->second};
+  }
+
+  // Return a set of extents that are contained in this extent map.
+  // If |extent| is completely covered by this extent map, a vector of itself
+  // will be returned.
+  // If only a subset of |extent| is covered by this extent map, a vector of
+  // parts in this map will be returned.
+  // If |extent| has no intersection with this map, an empty vector will be
+  // returned.
+  // E.g. extent map contains [0,5] and [10,15], GetIntersectingExtents([3, 12])
+  // would return [3,5] and [10,12]
+  std::vector<Extent> GetIntersectingExtents(const Extent& extent) const {
+    return set_.GetIntersectingExtents(extent);
+  }
+
+  // Complement of |GetIntersectingExtents|, return vector of extents which are
+  // part of |extent| but not covered by this map.
+  std::vector<Extent> GetNonIntersectingExtents(const Extent& extent) const {
+    return FilterExtentRanges({extent}, set_);
+  }
+
+ private:
+  // Get a range of exents that potentially intersect with parameter |extent|
+  std::map<Extent, T, Comparator> map_;
+  ExtentRanges set_;
+};
+}  // namespace chromeos_update_engine
+
+#endif  // UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_MAP_H_
diff --git a/payload_consumer/extent_map_unittest.cc b/payload_consumer/extent_map_unittest.cc
new file mode 100644
index 0000000..e7972cd
--- /dev/null
+++ b/payload_consumer/extent_map_unittest.cc
@@ -0,0 +1,122 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include <gtest/gtest.h>
+#include <optional>
+
+#include "update_engine/payload_consumer/extent_map.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/extent_utils.h"
+
+namespace chromeos_update_engine {
+
+class ExtentMapTest : public ::testing::Test {
+ public:
+  ExtentMap<int> map_;
+};
+
+TEST_F(ExtentMapTest, QueryExactExtent) {
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
+  auto ret = map_.Get(ExtentForRange(0, 5));
+  ASSERT_NE(ret, std::nullopt);
+  ASSERT_EQ(*ret, 7);
+}
+
+TEST_F(ExtentMapTest, QuerySubset) {
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
+  auto ret = map_.Get(ExtentForRange(1, 2));
+  ASSERT_EQ(ret, std::nullopt);
+}
+
+TEST_F(ExtentMapTest, QueryTouching) {
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
+  auto ret = map_.Get(ExtentForRange(3, 2));
+  ASSERT_EQ(ret, std::nullopt);
+  ret = map_.Get(ExtentForRange(4, 1));
+  ASSERT_EQ(ret, std::nullopt);
+  ret = map_.Get(ExtentForRange(5, 5));
+  ASSERT_EQ(ret, std::nullopt);
+  ret = map_.Get(ExtentForRange(5, 6));
+  ASSERT_EQ(ret, std::nullopt);
+}
+
+TEST_F(ExtentMapTest, GetIntersectingExtents) {
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 7));
+  auto ret = std::vector<Extent>{};
+  ret = map_.GetIntersectingExtents(ExtentForRange(2, 10));
+  ASSERT_EQ(ret.size(), 2U);
+  ASSERT_EQ(ret[0].start_block(), 2U);
+  ASSERT_EQ(ret[0].num_blocks(), 3U);
+
+  ASSERT_EQ(ret[1].start_block(), 10U);
+  ASSERT_EQ(ret[1].num_blocks(), 2U);
+
+  ret = map_.GetIntersectingExtents(ExtentForRange(2, 17));
+  ASSERT_EQ(ret.size(), 2U);
+  ASSERT_EQ(ret[0].start_block(), 2U);
+  ASSERT_EQ(ret[0].num_blocks(), 3U);
+
+  ASSERT_EQ(ret[1].start_block(), 10U);
+  ASSERT_EQ(ret[1].num_blocks(), 5U);
+
+  ret = map_.GetIntersectingExtents(ExtentForRange(2, 2));
+  ASSERT_EQ(ret, std::vector<Extent>{ExtentForRange(2, 2)});
+
+  ret = map_.GetIntersectingExtents(ExtentForRange(10, 5));
+  ASSERT_EQ(ret, std::vector<Extent>{ExtentForRange(10, 5)});
+
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(20, 5), 7));
+  ret = map_.GetIntersectingExtents(ExtentForRange(0, 30));
+  ASSERT_EQ(ret.size(), 3U);
+  ASSERT_EQ(ret[0].start_block(), 0U);
+  ASSERT_EQ(ret[0].num_blocks(), 5U);
+
+  ASSERT_EQ(ret[1].start_block(), 10U);
+  ASSERT_EQ(ret[1].num_blocks(), 5U);
+
+  ASSERT_EQ(ret[2].start_block(), 20U);
+  ASSERT_EQ(ret[2].num_blocks(), 5U);
+}
+
+TEST_F(ExtentMapTest, GetNonIntersectingExtents) {
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(20, 5), 7));
+
+  auto ret = std::vector<Extent>{};
+  ret = map_.GetNonIntersectingExtents(ExtentForRange(2, 13));
+
+  ASSERT_EQ(ret.size(), 1U);
+  ASSERT_EQ(ret[0].start_block(), 5U);
+  ASSERT_EQ(ret[0].num_blocks(), 5U);
+
+  ret = map_.GetNonIntersectingExtents(ExtentForRange(7, 20));
+  ASSERT_EQ(ret.size(), 3U);
+  ASSERT_EQ(ret[0].start_block(), 7U);
+  ASSERT_EQ(ret[0].num_blocks(), 3U);
+
+  ASSERT_EQ(ret[1].start_block(), 15U);
+  ASSERT_EQ(ret[1].num_blocks(), 5U);
+
+  ASSERT_EQ(ret[2].start_block(), 25U);
+  ASSERT_EQ(ret[2].num_blocks(), 2U);
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_ranges.cc b/payload_generator/extent_ranges.cc
index 2098639..4cc2d6f 100644
--- a/payload_generator/extent_ranges.cc
+++ b/payload_generator/extent_ranges.cc
@@ -45,7 +45,7 @@
 
 bool ExtentRanges::ExtentsOverlap(const Extent& a, const Extent& b) {
   if (a.start_block() == b.start_block())
-    return true;
+    return a.num_blocks() != 0;
   if (a.start_block() == kSparseHole || b.start_block() == kSparseHole)
     return false;
   if (a.start_block() < b.start_block()) {
@@ -281,6 +281,31 @@
   return out;
 }
 
+Range<ExtentRanges::ExtentSet::const_iterator> ExtentRanges::GetCandidateRange(
+    const Extent& extent) const {
+  auto lower_it = extent_set_.lower_bound(extent);
+  if (lower_it != extent_set_.begin()) {
+    --lower_it;
+  }
+
+  const auto upper_it = extent_set_.upper_bound(
+      ExtentForRange(extent.start_block() + extent.num_blocks(), 1));
+  return {lower_it, upper_it};
+}
+
+std::vector<Extent> ExtentRanges::GetIntersectingExtents(
+    const Extent& extent) const {
+  const auto candidates = GetCandidateRange(extent);
+  std::vector<Extent> result;
+  for (const auto& ext : candidates) {
+    if (auto intersection = GetOverlapExtent(ext, extent);
+        intersection.num_blocks() != 0) {
+      result.emplace_back(std::move(intersection));
+    }
+  }
+  return result;
+}
+
 vector<Extent> FilterExtentRanges(const vector<Extent>& extents,
                                   const ExtentRanges& ranges) {
   vector<Extent> result;
@@ -331,4 +356,15 @@
   return result;
 }
 
+Extent GetOverlapExtent(const Extent& extent1, const Extent& extent2) {
+  if (!ExtentRanges::ExtentsOverlap(extent1, extent2)) {
+    return {};
+  }
+  const auto start_block =
+      std::max(extent1.start_block(), extent2.start_block());
+  const auto end_block = std::min(extent1.start_block() + extent1.num_blocks(),
+                                  extent2.start_block() + extent2.num_blocks());
+  return ExtentForRange(start_block, end_block - start_block);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_ranges.h b/payload_generator/extent_ranges.h
index 68aa27f..4894831 100644
--- a/payload_generator/extent_ranges.h
+++ b/payload_generator/extent_ranges.h
@@ -23,6 +23,7 @@
 
 #include <base/macros.h>
 
+#include "update_engine/common/utils.h"
 #include "update_engine/update_metadata.pb.h"
 
 // An ExtentRanges object represents an unordered collection of extents (and
@@ -84,6 +85,18 @@
   // the number of blocks in this extent set.
   std::vector<Extent> GetExtentsForBlockCount(uint64_t count) const;
 
+  // Compute the intersection between this ExtentRange and the |extent|
+  // parameter. Return results in a vector. If there's no intersection, an empty
+  // vector is returned.
+  std::vector<Extent> GetIntersectingExtents(const Extent& extent) const;
+
+  // Get a range of extents that possibly intersect with |extent|. (Returned
+  // extents do not necessarily intersect!). It is perfectly acceptable to just
+  // return all extents in this set, though more efficient solution using binary
+  // search is preferred.
+  Range<ExtentSet::const_iterator> GetCandidateRange(
+      const Extent& extent) const;
+
  private:
   ExtentSet extent_set_;
   uint64_t blocks_;
@@ -95,6 +108,8 @@
 std::vector<Extent> FilterExtentRanges(const std::vector<Extent>& extents,
                                        const ExtentRanges& ranges);
 
+Extent GetOverlapExtent(const Extent& extent1, const Extent& extent2);
+
 }  // namespace chromeos_update_engine
 
 #endif  // UPDATE_ENGINE_PAYLOAD_GENERATOR_EXTENT_RANGES_H_
diff --git a/payload_generator/extent_ranges_unittest.cc b/payload_generator/extent_ranges_unittest.cc
index f55bb73..314ec54 100644
--- a/payload_generator/extent_ranges_unittest.cc
+++ b/payload_generator/extent_ranges_unittest.cc
@@ -315,4 +315,20 @@
                 ranges));
 }
 
+TEST(ExtentRangesTest, GetOverlapExtent) {
+  const auto ret1 =
+      GetOverlapExtent(ExtentForRange(5, 5), ExtentForRange(10, 10));
+  ASSERT_EQ(ret1.num_blocks(), 0UL) << ret1;
+  const auto ret2 =
+      GetOverlapExtent(ExtentForRange(5, 5), ExtentForRange(9, 10));
+  ASSERT_EQ(ret2, ExtentForRange(9, 1));
+
+  const auto ret3 =
+      GetOverlapExtent(ExtentForRange(7, 5), ExtentForRange(3, 10));
+  ASSERT_EQ(ret3, ExtentForRange(7, 5));
+  const auto ret4 =
+      GetOverlapExtent(ExtentForRange(7, 5), ExtentForRange(3, 3));
+  ASSERT_EQ(ret4.num_blocks(), 0UL);
+}
+
 }  // namespace chromeos_update_engine