Avoid self-overpping SOURCE_COPY ops by split them

snapuserd doesn't like self-overlapping copy operations. So we need to
split these operations. See linked bugs for examples.

Bug: 175137108
Test: treehugger
Change-Id: I24dca63ae89849330561fc26d9f2038982ed7ef2
diff --git a/payload_generator/merge_sequence_generator.cc b/payload_generator/merge_sequence_generator.cc
index eaffeac..289e2f8 100644
--- a/payload_generator/merge_sequence_generator.cc
+++ b/payload_generator/merge_sequence_generator.cc
@@ -50,6 +50,31 @@
          op1.dst_extent() == op2.dst_extent();
 }
 
+template <typename T>
+constexpr T GetDifference(T first, T second) {
+  T abs_diff = (first > second) ? (first - second) : (second - first);
+  return abs_diff;
+}
+
+void SplitSelfOverlapping(const Extent& src_extent,
+                          const Extent& dst_extent,
+                          std::vector<CowMergeOperation>* sequence) {
+  CHECK_EQ(src_extent.num_blocks(), dst_extent.num_blocks());
+  if (src_extent.start_block() == dst_extent.start_block()) {
+    sequence->emplace_back(CreateCowMergeOperation(src_extent, dst_extent));
+    return;
+  }
+
+  const size_t diff =
+      GetDifference(src_extent.start_block(), dst_extent.start_block());
+  for (size_t i = 0; i < src_extent.num_blocks(); i += diff) {
+    auto num_blocks = std::min<size_t>(diff, src_extent.num_blocks() - i);
+    sequence->emplace_back(CreateCowMergeOperation(
+        ExtentForRange(i + src_extent.start_block(), num_blocks),
+        ExtentForRange(i + dst_extent.start_block(), num_blocks)));
+  }
+}
+
 std::unique_ptr<MergeSequenceGenerator> MergeSequenceGenerator::Create(
     const std::vector<AnnotatedOperation>& aops) {
   std::vector<CowMergeOperation> sequence;
@@ -75,7 +100,14 @@
       Extent dst_extent =
           ExtentForRange(aop.op.dst_extents(0).start_block() + used_blocks,
                          src_extent.num_blocks());
-      sequence.emplace_back(CreateCowMergeOperation(src_extent, dst_extent));
+
+      // Self-overlapping SOURCE_COPY, must split into multiple non
+      // self-overlapping ops
+      if (ExtentRanges::ExtentsOverlap(src_extent, dst_extent)) {
+        SplitSelfOverlapping(src_extent, dst_extent, &sequence);
+      } else {
+        sequence.emplace_back(CreateCowMergeOperation(src_extent, dst_extent));
+      }
       used_blocks += src_extent.num_blocks();
     }
 
diff --git a/payload_generator/merge_sequence_generator.h b/payload_generator/merge_sequence_generator.h
index bc0158e..385fcc3 100644
--- a/payload_generator/merge_sequence_generator.h
+++ b/payload_generator/merge_sequence_generator.h
@@ -70,5 +70,9 @@
   std::vector<CowMergeOperation> operations_;
 };
 
+void SplitSelfOverlapping(const Extent& src_extent,
+                          const Extent& dst_extent,
+                          std::vector<CowMergeOperation>* sequence);
+
 }  // namespace chromeos_update_engine
 #endif
diff --git a/payload_generator/merge_sequence_generator_unittest.cc b/payload_generator/merge_sequence_generator_unittest.cc
index 1f0c2ea..666adbe 100644
--- a/payload_generator/merge_sequence_generator_unittest.cc
+++ b/payload_generator/merge_sequence_generator_unittest.cc
@@ -20,6 +20,7 @@
 #include <gtest/gtest.h>
 
 #include "update_engine/payload_consumer/payload_constants.h"
+#include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/extent_utils.h"
 #include "update_engine/payload_generator/merge_sequence_generator.h"
 
@@ -207,4 +208,49 @@
   GenerateSequence(transfers, expected);
 }
 
+std::ostream& operator<<(std::ostream& out, const Extent& extent) {
+  out << "[" << extent.start_block() << " - "
+      << extent.start_block() + extent.num_blocks() - 1 << "]";
+  return out;
+}
+
+void ValidateSplitSequence(const Extent& src_extent, const Extent& dst_extent) {
+  std::vector<CowMergeOperation> sequence;
+  SplitSelfOverlapping(src_extent, dst_extent, &sequence);
+  ExtentRanges src_extent_set;
+  src_extent_set.AddExtent(src_extent);
+  ExtentRanges dst_extent_set;
+  dst_extent_set.AddExtent(dst_extent);
+
+  size_t src_block_count = 0;
+  size_t dst_block_count = 0;
+  std::cout << "src_extent: " << src_extent << " dst_extent: " << dst_extent
+            << '\n';
+  for (const auto& merge_op : sequence) {
+    src_extent_set.SubtractExtent(merge_op.src_extent());
+    dst_extent_set.SubtractExtent(merge_op.dst_extent());
+    src_block_count += merge_op.src_extent().num_blocks();
+    dst_block_count += merge_op.dst_extent().num_blocks();
+    std::cout << merge_op.src_extent() << " -> " << merge_op.dst_extent()
+              << '\n';
+    ASSERT_FALSE(ExtentRanges::ExtentsOverlap(merge_op.src_extent(),
+                                              merge_op.dst_extent()));
+  }
+  std::cout << '\n';
+  // Check that all blocks are covered
+  ASSERT_EQ(src_extent_set.extent_set().size(), 0UL);
+  ASSERT_EQ(dst_extent_set.extent_set().size(), 0UL);
+
+  // Check that the split didn't cover extra blocks
+  ASSERT_EQ(src_block_count, src_extent.num_blocks());
+  ASSERT_EQ(dst_block_count, dst_extent.num_blocks());
+}
+
+TEST_F(MergeSequenceGeneratorTest, SplitSelfOverlappingTest) {
+  auto a = ExtentForRange(25, 16);
+  auto b = ExtentForRange(30, 16);
+  ValidateSplitSequence(a, b);
+  ValidateSplitSequence(b, a);
+}
+
 }  // namespace chromeos_update_engine