AU: Ranges class to managing unordered collection of blocks/extents.

Sometimes it's useful to have an unordered collection of ranges and
manipulate them. For example, we may want to do set subtraction on two
collections of extents. This class helps facilitate that.

BUG=7294
TEST=attached unittest

Review URL: http://codereview.chromium.org/3604005
diff --git a/SConstruct b/SConstruct
index 70803e7..585380a 100644
--- a/SConstruct
+++ b/SConstruct
@@ -1,4 +1,4 @@
-# Copyright (c) 2009 The Chromium OS Authors. All rights reserved.
+# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
@@ -203,6 +203,7 @@
                    delta_performer.cc
                    download_action.cc
                    extent_mapper.cc
+                   extent_ranges.cc
                    extent_writer.cc
                    filesystem_copier_action.cc
                    filesystem_iterator.cc
@@ -241,6 +242,7 @@
                             delta_performer_unittest.cc
                             download_action_unittest.cc
                             extent_mapper_unittest.cc
+                            extent_ranges_unittest.cc
                             extent_writer_unittest.cc
                             file_writer_unittest.cc
                             filesystem_copier_action_unittest.cc
diff --git a/extent_ranges.cc b/extent_ranges.cc
new file mode 100755
index 0000000..a8836c9
--- /dev/null
+++ b/extent_ranges.cc
@@ -0,0 +1,219 @@
+// Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "update_engine/extent_ranges.h"
+
+#include <set>
+#include <utility>
+#include <vector>
+
+#include <base/logging.h>
+
+using std::max;
+using std::min;
+using std::set;
+using std::vector;
+
+namespace chromeos_update_engine {
+
+bool ExtentRanges::ExtentsOverlapOrTouch(const Extent& a, const Extent& b) {
+  if (a.start_block() == b.start_block())
+    return true;
+  if (a.start_block() < b.start_block()) {
+    return a.start_block() + a.num_blocks() >= b.start_block();
+  } else {
+    return b.start_block() + b.num_blocks() >= a.start_block();
+  }
+}
+
+bool ExtentRanges::ExtentsOverlap(const Extent& a, const Extent& b) {
+  if (a.start_block() == b.start_block())
+    return true;
+  if (a.start_block() < b.start_block()) {
+    return a.start_block() + a.num_blocks() > b.start_block();
+  } else {
+    return b.start_block() + b.num_blocks() > a.start_block();
+  }
+}
+
+void ExtentRanges::AddBlock(uint64_t block) {
+  AddExtent(ExtentForRange(block, 1));
+}
+
+void ExtentRanges::SubtractBlock(uint64_t block) {
+  SubtractExtent(ExtentForRange(block, 1));
+}
+
+namespace {
+
+Extent UnionOverlappingExtents(const Extent& first, const Extent& second) {
+  uint64_t start = min(first.start_block(), second.start_block());
+  uint64_t end = max(first.start_block() + first.num_blocks(),
+                     second.start_block() + second.num_blocks());
+  return ExtentForRange(start, end - start);
+}
+  
+}  // namespace {}
+
+void ExtentRanges::AddExtent(Extent extent) {
+  if (extent.num_blocks() == 0)
+    return;
+
+  ExtentSet::iterator begin_del = extent_set_.end();
+  ExtentSet::iterator end_del = extent_set_.end();
+  uint64_t del_blocks = 0;
+  for (ExtentSet::iterator it = extent_set_.begin(), e = extent_set_.end();
+       it != e; ++it) {
+    if (ExtentsOverlapOrTouch(*it, extent)) {
+      end_del = it;
+      ++end_del;
+      del_blocks += it->num_blocks();
+      if (begin_del == extent_set_.end())
+        begin_del = it;
+
+      extent = UnionOverlappingExtents(extent, *it);
+    }
+  }
+  extent_set_.erase(begin_del, end_del);
+  extent_set_.insert(extent);
+  blocks_ -= del_blocks;
+  blocks_ += extent.num_blocks();
+}
+
+namespace {
+// Returns base - subtractee (set subtraction).
+ExtentRanges::ExtentSet SubtractOverlappingExtents(const Extent& base,
+                                                   const Extent& subtractee) {
+  ExtentRanges::ExtentSet ret;
+  if (subtractee.start_block() > base.start_block()) {
+    ret.insert(ExtentForRange(base.start_block(),
+                              subtractee.start_block() - base.start_block()));
+  }
+  uint64_t base_end = base.start_block() + base.num_blocks();
+  uint64_t subtractee_end = subtractee.start_block() + subtractee.num_blocks();
+  if (base_end > subtractee_end) {
+    ret.insert(ExtentForRange(subtractee_end, base_end - subtractee_end));
+  }
+  return ret;
+}
+}  // namespace {}
+
+void ExtentRanges::SubtractExtent(const Extent& extent) {
+  if (extent.num_blocks() == 0)
+    return;
+
+  ExtentSet::iterator begin_del = extent_set_.end();
+  ExtentSet::iterator end_del = extent_set_.end();
+  uint64_t del_blocks = 0;
+  ExtentSet new_extents;
+  for (ExtentSet::iterator it = extent_set_.begin(), e = extent_set_.end();
+       it != e; ++it) {
+    if (!ExtentsOverlap(*it, extent))
+      continue;
+    
+    if (begin_del == extent_set_.end())
+      begin_del = it;
+    end_del = it;
+    ++end_del;
+    
+    del_blocks += it->num_blocks();
+
+    ExtentSet subtraction = SubtractOverlappingExtents(*it, extent);
+    for (ExtentSet::iterator jt = subtraction.begin(), je = subtraction.end();
+         jt != je; ++jt) {
+      new_extents.insert(*jt);
+      del_blocks -= jt->num_blocks();
+    }
+  }
+  extent_set_.erase(begin_del, end_del);
+  extent_set_.insert(new_extents.begin(), new_extents.end());
+  blocks_ -= del_blocks;
+}
+
+void ExtentRanges::AddRanges(const ExtentRanges& ranges) {
+  for (ExtentSet::const_iterator it = ranges.extent_set_.begin(),
+           e = ranges.extent_set_.end(); it != e; ++it) {
+    AddExtent(*it);
+  }
+}
+
+void ExtentRanges::SubtractRanges(const ExtentRanges& ranges) {
+  for (ExtentSet::const_iterator it = ranges.extent_set_.begin(),
+           e = ranges.extent_set_.end(); it != e; ++it) {
+    SubtractExtent(*it);
+  }
+}
+
+void ExtentRanges::AddExtents(const vector<Extent>& extents) {
+  for (vector<Extent>::const_iterator it = extents.begin(), e = extents.end();
+       it != e; ++it) {
+    AddExtent(*it);
+  }
+}
+
+void ExtentRanges::SubtractExtents(const vector<Extent>& extents) {
+  for (vector<Extent>::const_iterator it = extents.begin(), e = extents.end();
+       it != e; ++it) {
+    SubtractExtent(*it);
+  }
+}
+
+void ExtentRanges::AddRepeatedExtents(
+    const ::google::protobuf::RepeatedPtrField<Extent> &exts) {
+  for (int i = 0, e = exts.size(); i != e; ++i) {
+    AddExtent(exts.Get(i));
+  }
+}
+
+void ExtentRanges::SubtractRepeatedExtents(
+    const ::google::protobuf::RepeatedPtrField<Extent> &exts) {
+  for (int i = 0, e = exts.size(); i != e; ++i) {
+    SubtractExtent(exts.Get(i));
+  }
+}
+
+void ExtentRanges::Dump() const {
+  LOG(INFO) << "ExtentRanges Dump. blocks: " << blocks_;
+  for (ExtentSet::const_iterator it = extent_set_.begin(),
+           e = extent_set_.end();
+       it != e; ++it) {
+    LOG(INFO) << "{" << it->start_block() << ", " << it->num_blocks() << "}";
+  }
+}
+
+Extent ExtentForRange(uint64_t start_block, uint64_t num_blocks) {
+  Extent ret;
+  ret.set_start_block(start_block);
+  ret.set_num_blocks(num_blocks);
+  return ret;
+}
+
+std::vector<Extent> ExtentRanges::GetExtentsForBlockCount(
+    uint64_t count) const {
+  vector<Extent> out;
+  if (count == 0)
+    return out;
+  uint64_t out_blocks = 0;
+  CHECK(count <= blocks_);
+  for (ExtentSet::const_iterator it = extent_set_.begin(),
+           e = extent_set_.end();
+       it != e; ++it) {
+    const uint64_t blocks_needed = count - out_blocks;
+    const Extent& extent = *it;
+    out.push_back(extent);
+    out_blocks += extent.num_blocks();
+    if (extent.num_blocks() < blocks_needed)
+      continue;
+    if (extent.num_blocks() == blocks_needed)
+      break;
+    // If we get here, we just added the last extent needed, but it's too big
+    out_blocks -= extent.num_blocks();
+    out_blocks += blocks_needed;
+    out.back().set_num_blocks(blocks_needed);
+    break;
+  }
+  return out;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/extent_ranges.h b/extent_ranges.h
new file mode 100644
index 0000000..8fdc302
--- /dev/null
+++ b/extent_ranges.h
@@ -0,0 +1,72 @@
+// Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROMEOS_PLATFORM_UPDATE_ENGINE_EXTENT_RANGES_H__
+#define CHROMEOS_PLATFORM_UPDATE_ENGINE_EXTENT_RANGES_H__
+
+#include <map>
+#include <set>
+#include <vector>
+
+#include <base/basictypes.h>
+
+#include "update_engine/delta_diff_generator.h"
+#include "update_engine/graph_types.h"
+#include "update_engine/update_metadata.pb.h"
+
+// An ExtentRanges object represents an unordered collection of extents
+// (and therefore blocks). Such an object may be modified by adding or
+// subtracting blocks (think: set addition or set subtraction).
+
+namespace chromeos_update_engine {
+
+struct ExtentLess {
+  bool operator()(const Extent& x, const Extent& y) const {
+    return x.start_block() < y.start_block();
+  }
+};
+
+Extent ExtentForRange(uint64_t start_block, uint64_t num_blocks);
+
+class ExtentRanges {
+ public:
+  typedef std::set<Extent, ExtentLess> ExtentSet;
+
+  ExtentRanges() : blocks_(0) {}
+  void AddBlock(uint64_t block);
+  void SubtractBlock(uint64_t block);
+  void AddExtent(Extent extent);
+  void SubtractExtent(const Extent& extent);
+  void AddExtents(const std::vector<Extent>& extents);
+  void SubtractExtents(const std::vector<Extent>& extents);
+  void AddRepeatedExtents(
+      const ::google::protobuf::RepeatedPtrField<Extent> &exts);
+  void SubtractRepeatedExtents(
+      const ::google::protobuf::RepeatedPtrField<Extent> &exts);
+  void AddRanges(const ExtentRanges& ranges);
+  void SubtractRanges(const ExtentRanges& ranges);
+
+  static bool ExtentsOverlapOrTouch(const Extent& a, const Extent& b);
+  static bool ExtentsOverlap(const Extent& a, const Extent& b);
+
+  // Dumps contents to the log file. Useful for debugging.
+  void Dump() const;
+  
+  uint64_t blocks() const { return blocks_; }
+  const ExtentSet& extent_set() const { return extent_set_; }
+
+  // Returns an ordered vector of extents for |count| blocks,
+  // using extents in extent_set_. The returned extents are not
+  // removed from extent_set_. |count| must be less than or equal to
+  // the number of blocks in this extent set.
+  std::vector<Extent> GetExtentsForBlockCount(uint64_t count) const;
+
+ private:
+  ExtentSet extent_set_;
+  uint64_t blocks_;
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // CHROMEOS_PLATFORM_UPDATE_ENGINE_EXTENT_RANGES_H__
diff --git a/extent_ranges_unittest.cc b/extent_ranges_unittest.cc
new file mode 100755
index 0000000..b5a32d9
--- /dev/null
+++ b/extent_ranges_unittest.cc
@@ -0,0 +1,237 @@
+// Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include "update_engine/extent_ranges.h"
+#include "update_engine/test_utils.h"
+
+using std::vector;
+
+namespace chromeos_update_engine {
+
+class ExtentRangesTest : public ::testing::Test {};
+
+namespace {
+void ExpectRangeEq(const ExtentRanges& ranges,
+                   uint64_t* expected,
+                   size_t sz,
+                   int line) {
+  uint64_t blocks = 0;
+  for (size_t i = 1; i < sz; i += 2) {
+    blocks += expected[i];
+  }
+  EXPECT_EQ(blocks, ranges.blocks()) << "line: " << line;
+
+  const ExtentRanges::ExtentSet& result = ranges.extent_set();
+  ExtentRanges::ExtentSet::const_iterator it = result.begin();
+  for (size_t i = 0; i < sz; i += 2) {
+    EXPECT_FALSE(it == result.end()) << "line: " << line;
+    EXPECT_EQ(expected[i], it->start_block()) << "line: " << line;
+    EXPECT_EQ(expected[i + 1], it->num_blocks()) << "line: " << line;
+    ++it;
+  }
+}
+
+#define EXPECT_RANGE_EQ(ranges, var)                            \
+  do {                                                          \
+    ExpectRangeEq(ranges, var, arraysize(var), __LINE__);       \
+  } while (0)
+
+void ExpectRangesOverlapOrTouch(uint64_t a_start, uint64_t a_num,
+                                uint64_t b_start, uint64_t b_num) {
+  EXPECT_TRUE(ExtentRanges::ExtentsOverlapOrTouch(ExtentForRange(a_start,
+                                                                 a_num),
+                                                  ExtentForRange(b_start,
+                                                                 b_num)));
+  EXPECT_TRUE(ExtentRanges::ExtentsOverlapOrTouch(ExtentForRange(b_start,
+                                                                 b_num),
+                                                  ExtentForRange(a_start,
+                                                                 a_num)));
+}
+
+void ExpectFalseRangesOverlapOrTouch(uint64_t a_start, uint64_t a_num,
+                                     uint64_t b_start, uint64_t b_num) {
+  EXPECT_FALSE(ExtentRanges::ExtentsOverlapOrTouch(ExtentForRange(a_start,
+                                                                  a_num),
+                                                   ExtentForRange(b_start,
+                                                                  b_num)));
+  EXPECT_FALSE(ExtentRanges::ExtentsOverlapOrTouch(ExtentForRange(b_start,
+                                                                  b_num),
+                                                   ExtentForRange(a_start,
+                                                                  a_num)));
+  EXPECT_FALSE(ExtentRanges::ExtentsOverlap(ExtentForRange(a_start,
+                                                           a_num),
+                                            ExtentForRange(b_start,
+                                                           b_num)));
+  EXPECT_FALSE(ExtentRanges::ExtentsOverlap(ExtentForRange(b_start,
+                                                           b_num),
+                                            ExtentForRange(a_start,
+                                                           a_num)));
+}
+
+void ExpectRangesOverlap(uint64_t a_start, uint64_t a_num,
+                         uint64_t b_start, uint64_t b_num) {
+  EXPECT_TRUE(ExtentRanges::ExtentsOverlap(ExtentForRange(a_start,
+                                                          a_num),
+                                           ExtentForRange(b_start,
+                                                          b_num)));
+  EXPECT_TRUE(ExtentRanges::ExtentsOverlap(ExtentForRange(b_start,
+                                                          b_num),
+                                           ExtentForRange(a_start,
+                                                          a_num)));
+  EXPECT_TRUE(ExtentRanges::ExtentsOverlapOrTouch(ExtentForRange(a_start,
+                                                                 a_num),
+                                                  ExtentForRange(b_start,
+                                                                 b_num)));
+  EXPECT_TRUE(ExtentRanges::ExtentsOverlapOrTouch(ExtentForRange(b_start,
+                                                                 b_num),
+                                                  ExtentForRange(a_start,
+                                                                 a_num)));
+}
+
+void ExpectFalseRangesOverlap(uint64_t a_start, uint64_t a_num,
+                              uint64_t b_start, uint64_t b_num) {
+  EXPECT_FALSE(ExtentRanges::ExtentsOverlap(ExtentForRange(a_start,
+                                                           a_num),
+                                            ExtentForRange(b_start,
+                                                           b_num)));
+  EXPECT_FALSE(ExtentRanges::ExtentsOverlap(ExtentForRange(b_start,
+                                                           b_num),
+                                            ExtentForRange(a_start,
+                                                           a_num)));
+}
+
+}  // namespace {}
+
+TEST(ExtentRangesTest, ExtentsOverlapTest) {
+  ExpectRangesOverlapOrTouch(10, 20, 30, 10);
+  ExpectRangesOverlap(10, 20, 25, 10);
+  ExpectFalseRangesOverlapOrTouch(10, 20, 35, 10);
+  ExpectFalseRangesOverlap(10, 20, 30, 10);
+  ExpectRangesOverlap(12, 4, 12, 3);
+}
+
+TEST(ExtentRangesTest, SimpleTest) {
+  ExtentRanges ranges;
+  {
+    uint64_t expected[] = {};
+    // Can't use arraysize() on 0-length arrays:
+    ExpectRangeEq(ranges, expected, 0, __LINE__);
+  }
+  ranges.SubtractBlock(2);
+  {
+    uint64_t expected[] = {};
+    // Can't use arraysize() on 0-length arrays:
+    ExpectRangeEq(ranges, expected, 0, __LINE__);
+  }
+  
+  ranges.AddBlock(0);
+  ranges.Dump();
+  ranges.AddBlock(1);
+  ranges.AddBlock(3);
+  
+  {
+    uint64_t expected[] = {0, 2, 3, 1};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+  ranges.AddBlock(2);
+  {
+    uint64_t expected[] = {0, 4};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+  ranges.SubtractBlock(2);
+  {
+    uint64_t expected[] = {0, 2, 3, 1};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+
+  for (uint64_t i = 100; i < 1000; i += 100) {
+    ranges.AddExtent(ExtentForRange(i, 50));
+  }
+  {
+    uint64_t expected[] = {0, 2, 3, 1, 100, 50, 200, 50, 300, 50, 400, 50,
+                           500, 50, 600, 50, 700, 50, 800, 50, 900, 50};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+
+  ranges.SubtractExtent(ExtentForRange(210, 410 - 210));
+  {
+    uint64_t expected[] = {0, 2, 3, 1, 100, 50, 200, 10, 410, 40, 500, 50,
+                           600, 50, 700, 50, 800, 50, 900, 50};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+  ranges.AddExtent(ExtentForRange(100000, 0));
+  {
+    uint64_t expected[] = {0, 2, 3, 1, 100, 50, 200, 10, 410, 40, 500, 50,
+                           600, 50, 700, 50, 800, 50, 900, 50};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+  ranges.SubtractExtent(ExtentForRange(3, 0));
+  {
+    uint64_t expected[] = {0, 2, 3, 1, 100, 50, 200, 10, 410, 40, 500, 50,
+                           600, 50, 700, 50, 800, 50, 900, 50};
+    EXPECT_RANGE_EQ(ranges, expected);
+  }
+}
+
+TEST(ExtentRangesTest, MultipleRanges) {
+  ExtentRanges ranges_a, ranges_b;
+  ranges_a.AddBlock(0);
+  ranges_b.AddBlock(4);
+  ranges_b.AddBlock(3);
+  {
+    uint64_t expected[] = {3, 2};
+    EXPECT_RANGE_EQ(ranges_b, expected);
+  }
+  ranges_a.AddRanges(ranges_b);
+  {
+    uint64_t expected[] = {0, 1, 3, 2};
+    EXPECT_RANGE_EQ(ranges_a, expected);
+  }
+  ranges_a.SubtractRanges(ranges_b);
+  {
+    uint64_t expected[] = {0, 1};
+    EXPECT_RANGE_EQ(ranges_a, expected);
+  }
+  {
+    uint64_t expected[] = {3, 2};
+    EXPECT_RANGE_EQ(ranges_b, expected);
+  }
+}
+
+TEST(ExtentRangesTest, GetExtentsForBlockCountTest) {
+  ExtentRanges ranges;
+  ranges.AddExtents(vector<Extent>(1, ExtentForRange(10, 30)));
+  {
+    vector<Extent> zero_extents = ranges.GetExtentsForBlockCount(0);
+    EXPECT_TRUE(zero_extents.empty());
+  }
+  ::google::protobuf::RepeatedPtrField<Extent> rep_field;
+  *rep_field.Add() = ExtentForRange(30, 40);
+  ranges.AddRepeatedExtents(rep_field);
+  ranges.SubtractExtents(vector<Extent>(1, ExtentForRange(20, 10)));
+  *rep_field.Mutable(0) = ExtentForRange(50, 10);
+  ranges.SubtractRepeatedExtents(rep_field);
+  EXPECT_EQ(40, ranges.blocks());
+
+  for (int i = 0; i < 2; i++) {
+    vector<Extent> expected(2);
+    expected[0] = ExtentForRange(10, 10);
+    expected[1] = ExtentForRange(30, i == 0 ? 10 : 20);
+    vector<Extent> actual =
+        ranges.GetExtentsForBlockCount(10 + expected[1].num_blocks());
+    EXPECT_EQ(expected.size(), actual.size());
+    for (vector<Extent>::size_type j = 0, e = expected.size(); j != e; ++j) {
+      EXPECT_EQ(expected[j].start_block(), actual[j].start_block())
+          << "j = " << j;
+      EXPECT_EQ(expected[j].num_blocks(), actual[j].num_blocks())
+          << "j = " << j;
+    }
+  }
+}
+
+}  // namespace chromeos_update_engine