Include XOR ops in merge sequence

Test: th
Bug: 177104308

Change-Id: I0ad017c7570879b115381f73920a9097500145cb
diff --git a/payload_generator/merge_sequence_generator.cc b/payload_generator/merge_sequence_generator.cc
index 289e2f8..c5fd988 100644
--- a/payload_generator/merge_sequence_generator.cc
+++ b/payload_generator/merge_sequence_generator.cc
@@ -18,16 +18,22 @@
 
 #include <algorithm>
 
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/extent_utils.h"
+#include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
 
 CowMergeOperation CreateCowMergeOperation(const Extent& src_extent,
-                                          const Extent& dst_extent) {
+                                          const Extent& dst_extent,
+                                          CowMergeOperation::Type op_type,
+                                          uint32_t src_offset) {
   CowMergeOperation ret;
-  ret.set_type(CowMergeOperation::COW_COPY);
+  ret.set_type(op_type);
   *ret.mutable_src_extent() = src_extent;
   *ret.mutable_dst_extent() = dst_extent;
+  ret.set_src_offset(src_offset);
   return ret;
 }
 
@@ -36,6 +42,17 @@
   os << "CowMergeOperation src extent: "
      << ExtentsToString({merge_operation.src_extent()})
      << ", dst extent: " << ExtentsToString({merge_operation.dst_extent()});
+  if (merge_operation.has_src_offset()) {
+    os << ", src offset: " << merge_operation.src_offset();
+  }
+  os << " op_type: ";
+  if (merge_operation.type() == CowMergeOperation::COW_COPY) {
+    os << "COW_COPY";
+  } else if (merge_operation.type() == CowMergeOperation::COW_XOR) {
+    os << "COW_XOR";
+  } else {
+    os << merge_operation.type();
+  }
   return os;
 }
 
@@ -56,12 +73,27 @@
   return abs_diff;
 }
 
+CowMergeOperation::Type GetCowOpType(InstallOperation::Type install_op_type) {
+  switch (install_op_type) {
+    case InstallOperation::SOURCE_COPY:
+      return CowMergeOperation::COW_COPY;
+    case InstallOperation::SOURCE_BSDIFF:
+    case InstallOperation::BROTLI_BSDIFF:
+    case InstallOperation::PUFFDIFF:
+      return CowMergeOperation::COW_XOR;
+    default:
+      CHECK(false) << "Unknown install op type: " << install_op_type;
+      return CowMergeOperation::COW_REPLACE;
+  }
+}
+
 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));
+    sequence->emplace_back(CreateCowMergeOperation(
+        src_extent, dst_extent, CowMergeOperation::COW_COPY));
     return;
   }
 
@@ -71,51 +103,89 @@
     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)));
+        ExtentForRange(i + dst_extent.start_block(), num_blocks),
+        CowMergeOperation::COW_COPY));
   }
 }
 
+static bool ProcessXorOps(std::vector<CowMergeOperation>* sequence,
+                          const AnnotatedOperation& aop) {
+  const auto size_before = sequence->size();
+  sequence->insert(sequence->end(), aop.xor_ops.begin(), aop.xor_ops.end());
+  std::for_each(
+      sequence->begin() + size_before,
+      sequence->end(),
+      [](CowMergeOperation& op) {
+        CHECK_EQ(op.type(), CowMergeOperation::COW_XOR);
+        // If |src_offset| is greater than 0, then we are reading 1
+        // extra block at the end of src_extent. This dependency must
+        // be honored during merge sequence generation, or we can end
+        // up with a corrupted device after merge.
+        if (op.src_offset() > 0) {
+          if (op.src_extent().num_blocks() == op.dst_extent().num_blocks()) {
+            op.mutable_src_extent()->set_num_blocks(
+                op.src_extent().num_blocks() + 1);
+          }
+          CHECK_EQ(op.src_extent().num_blocks(),
+                   op.dst_extent().num_blocks() + 1);
+        }
+      });
+  return true;
+}
+
+static bool ProcessCopyOps(std::vector<CowMergeOperation>* sequence,
+                           const AnnotatedOperation& aop) {
+  CHECK_EQ(GetCowOpType(aop.op.type()), CowMergeOperation::COW_COPY);
+  if (aop.op.dst_extents().size() != 1) {
+    std::vector<Extent> out_extents;
+    ExtentsToVector(aop.op.dst_extents(), &out_extents);
+    LOG(ERROR)
+        << "The dst extents for source_copy are expected to be contiguous,"
+        << " dst extents: " << ExtentsToString(out_extents);
+    return false;
+  }
+  // Split the source extents.
+  size_t used_blocks = 0;
+  for (const auto& src_extent : aop.op.src_extents()) {
+    // The dst_extent in the merge sequence will be a subset of
+    // InstallOperation's dst_extent. This will simplify the OTA -> COW
+    // conversion when we install the payload.
+    Extent dst_extent =
+        ExtentForRange(aop.op.dst_extents(0).start_block() + used_blocks,
+                       src_extent.num_blocks());
+    // Self-overlapping operation, 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, CowMergeOperation::COW_COPY));
+    }
+    used_blocks += src_extent.num_blocks();
+  }
+
+  if (used_blocks != aop.op.dst_extents(0).num_blocks()) {
+    LOG(ERROR) << "Number of blocks in src extents doesn't equal to the"
+               << " ones in the dst extents, src blocks " << used_blocks
+               << ", dst blocks " << aop.op.dst_extents(0).num_blocks();
+    return false;
+  }
+  return true;
+}
+
 std::unique_ptr<MergeSequenceGenerator> MergeSequenceGenerator::Create(
     const std::vector<AnnotatedOperation>& aops) {
   std::vector<CowMergeOperation> sequence;
+
   for (const auto& aop : aops) {
-    // Only handle SOURCE_COPY now for the cow size optimization.
-    if (aop.op.type() != InstallOperation::SOURCE_COPY) {
-      continue;
-    }
-    if (aop.op.dst_extents().size() != 1) {
-      std::vector<Extent> out_extents;
-      ExtentsToVector(aop.op.dst_extents(), &out_extents);
-      LOG(ERROR) << "The dst extents for source_copy expects to be contiguous,"
-                 << " dst extents: " << ExtentsToString(out_extents);
-      return nullptr;
-    }
-
-    // Split the source extents.
-    size_t used_blocks = 0;
-    for (const auto& src_extent : aop.op.src_extents()) {
-      // The dst_extent in the merge sequence will be a subset of
-      // InstallOperation's dst_extent. This will simplify the OTA -> COW
-      // conversion when we install the payload.
-      Extent dst_extent =
-          ExtentForRange(aop.op.dst_extents(0).start_block() + used_blocks,
-                         src_extent.num_blocks());
-
-      // 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));
+    if (aop.op.type() == InstallOperation::SOURCE_COPY) {
+      if (!ProcessCopyOps(&sequence, aop)) {
+        return nullptr;
       }
-      used_blocks += src_extent.num_blocks();
-    }
-
-    if (used_blocks != aop.op.dst_extents(0).num_blocks()) {
-      LOG(ERROR) << "Number of blocks in src extents doesn't equal to the"
-                 << " ones in the dst extents, src blocks " << used_blocks
-                 << ", dst blocks " << aop.op.dst_extents(0).num_blocks();
-      return nullptr;
+    } else if (!aop.xor_ops.empty()) {
+      if (!ProcessXorOps(&sequence, aop)) {
+        return nullptr;
+      }
     }
   }
 
@@ -167,7 +237,7 @@
       }
       auto ret = merge_after.emplace(op, std::move(operations));
       // Check the insertion indeed happens.
-      CHECK(ret.second);
+      CHECK(ret.second) << op;
     }
   }
 
@@ -221,9 +291,9 @@
     for (const auto& op : free_operations) {
       incoming_edges.erase(op);
 
-      // Now that this particular operation is merged, other operations blocked
-      // by this one may be free. Decrement the count of blocking operations,
-      // and set up the free operations for the next iteration.
+      // Now that this particular operation is merged, other operations
+      // blocked by this one may be free. Decrement the count of blocking
+      // operations, and set up the free operations for the next iteration.
       for (const auto& blocked : merge_after[op]) {
         auto it = incoming_edges.find(blocked);
         if (it == incoming_edges.end()) {
@@ -271,6 +341,7 @@
   LOG(INFO) << "Blocks in merge sequence " << blocks_in_sequence
             << ", blocks in raw " << blocks_in_raw;
   if (!ValidateSequence(merge_sequence)) {
+    LOG(ERROR) << "Invalid Sequence";
     return false;
   }
 
@@ -283,6 +354,16 @@
   LOG(INFO) << "Validating merge sequence";
   ExtentRanges visited;
   for (const auto& op : sequence) {
+    // If |src_offset| is greater than zero, dependency should include 1 extra
+    // block at end of src_extent, as the OP actually references data past
+    // original src_extent.
+    if (op.src_offset() > 0) {
+      CHECK_EQ(op.src_extent().num_blocks(), op.dst_extent().num_blocks() + 1)
+          << op;
+    } else {
+      CHECK_EQ(op.src_extent().num_blocks(), op.dst_extent().num_blocks())
+          << op;
+    }
     if (visited.OverlapsWithExtent(op.src_extent())) {
       LOG(ERROR) << "Transfer violates the merge sequence " << op
                  << "Visited extent ranges: ";
diff --git a/payload_generator/merge_sequence_generator.h b/payload_generator/merge_sequence_generator.h
index 385fcc3..d3b5eb9 100644
--- a/payload_generator/merge_sequence_generator.h
+++ b/payload_generator/merge_sequence_generator.h
@@ -31,7 +31,9 @@
 namespace chromeos_update_engine {
 // Constructs CowMergeOperation from src & dst extents
 CowMergeOperation CreateCowMergeOperation(const Extent& src_extent,
-                                          const Extent& dst_extent);
+                                          const Extent& dst_extent,
+                                          CowMergeOperation::Type op_type,
+                                          uint32_t src_offset = 0);
 
 // Comparator for CowMergeOperation.
 bool operator<(const CowMergeOperation& op1, const CowMergeOperation& op2);
@@ -67,7 +69,7 @@
   bool FindDependency(std::map<CowMergeOperation, std::set<CowMergeOperation>>*
                           merge_after) const;
   // The list of CowMergeOperations to sort.
-  std::vector<CowMergeOperation> operations_;
+  const std::vector<CowMergeOperation> operations_;
 };
 
 void SplitSelfOverlapping(const Extent& src_extent,
diff --git a/payload_generator/merge_sequence_generator_unittest.cc b/payload_generator/merge_sequence_generator_unittest.cc
index 579cca8..86d4fdd 100644
--- a/payload_generator/merge_sequence_generator_unittest.cc
+++ b/payload_generator/merge_sequence_generator_unittest.cc
@@ -20,11 +20,18 @@
 #include <gtest/gtest.h>
 
 #include "update_engine/payload_consumer/payload_constants.h"
+#include "update_engine/payload_generator/delta_diff_generator.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"
+#include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
+CowMergeOperation CreateCowMergeOperation(const Extent& src_extent,
+                                          const Extent& dst_extent) {
+  return CreateCowMergeOperation(
+      src_extent, dst_extent, CowMergeOperation::COW_COPY);
+}
 class MergeSequenceGeneratorTest : public ::testing::Test {
  protected:
   void VerifyTransfers(MergeSequenceGenerator* generator,
@@ -183,8 +190,8 @@
       CreateCowMergeOperation(ExtentForRange(10, 10), ExtentForRange(15, 10)),
   };
 
-  // file 1,2,3 form a cycle. And file3, whose dst ext has smallest offset, will
-  // be converted to raw blocks
+  // file 1,2,3 form a cycle. And file3, whose dst ext has smallest offset,
+  // will be converted to raw blocks
   std::vector<CowMergeOperation> expected{
       transfers[3], transfers[1], transfers[0]};
   GenerateSequence(transfers, expected);
@@ -247,4 +254,174 @@
   ValidateSplitSequence(b, a);
 }
 
+TEST_F(MergeSequenceGeneratorTest, GenerateSequenceWithXor) {
+  std::vector<CowMergeOperation> transfers = {
+      // cycle 1
+      CreateCowMergeOperation(ExtentForRange(10, 10),
+                              ExtentForRange(25, 10),
+                              CowMergeOperation::COW_XOR),
+      CreateCowMergeOperation(ExtentForRange(24, 5), ExtentForRange(35, 5)),
+      CreateCowMergeOperation(ExtentForRange(30, 10),
+                              ExtentForRange(15, 10),
+                              CowMergeOperation::COW_XOR),
+      // cycle 2
+      CreateCowMergeOperation(ExtentForRange(55, 10), ExtentForRange(60, 10)),
+      CreateCowMergeOperation(ExtentForRange(60, 10),
+                              ExtentForRange(70, 10),
+                              CowMergeOperation::COW_XOR),
+      CreateCowMergeOperation(ExtentForRange(70, 10), ExtentForRange(55, 10)),
+  };
+
+  // file 3, 6 will be converted to raw.
+  std::vector<CowMergeOperation> expected{
+      transfers[1], transfers[0], transfers[4], transfers[3]};
+  GenerateSequence(transfers, expected);
+}
+
+TEST_F(MergeSequenceGeneratorTest, CreateGeneratorWithXor) {
+  std::vector<AnnotatedOperation> aops;
+  auto& aop = aops.emplace_back();
+  aop.op.set_type(InstallOperation::SOURCE_BSDIFF);
+  *aop.op.mutable_src_extents()->Add() = ExtentForRange(10, 5);
+  *aop.op.mutable_dst_extents()->Add() = ExtentForRange(20, 5);
+  auto& xor_map = aop.xor_ops;
+  {
+    // xor_map[i] = i * kBlockSize + 123;
+    auto& op = xor_map.emplace_back();
+    *op.mutable_src_extent() = ExtentForRange(10, 5);
+    *op.mutable_dst_extent() = ExtentForRange(20, 5);
+    op.set_src_offset(123);
+    op.set_type(CowMergeOperation::COW_XOR);
+  }
+  auto generator = MergeSequenceGenerator::Create(aops);
+  ASSERT_NE(generator, nullptr);
+  std::vector<CowMergeOperation> sequence;
+  ASSERT_TRUE(generator->Generate(&sequence));
+  ASSERT_EQ(sequence.size(), 1UL);
+  ASSERT_EQ(sequence[0].src_extent().start_block(), 10UL);
+  ASSERT_EQ(sequence[0].dst_extent().start_block(), 20UL);
+  ASSERT_EQ(sequence[0].src_extent().num_blocks(), 6UL);
+  ASSERT_EQ(sequence[0].dst_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[0].type(), CowMergeOperation::COW_XOR);
+  ASSERT_EQ(sequence[0].src_offset(), 123UL);
+
+  ASSERT_TRUE(generator->ValidateSequence(sequence));
+}
+
+TEST_F(MergeSequenceGeneratorTest, CreateGeneratorWithXorMultipleExtents) {
+  std::vector<AnnotatedOperation> aops;
+  auto& aop = aops.emplace_back();
+  aop.op.set_type(InstallOperation::SOURCE_BSDIFF);
+  *aop.op.mutable_src_extents()->Add() = ExtentForRange(10, 10);
+  *aop.op.mutable_dst_extents()->Add() = ExtentForRange(30, 5);
+  *aop.op.mutable_dst_extents()->Add() = ExtentForRange(45, 5);
+  auto& xor_map = aop.xor_ops;
+  {
+    // xor_map[i] = i * kBlockSize + 123;
+    auto& op = xor_map.emplace_back();
+    *op.mutable_src_extent() = ExtentForRange(10, 5);
+    *op.mutable_dst_extent() = ExtentForRange(30, 5);
+    op.set_src_offset(123);
+    op.set_type(CowMergeOperation::COW_XOR);
+  }
+  {
+    // xor_map[i] = i * kBlockSize + 123;
+    auto& op = xor_map.emplace_back();
+    *op.mutable_src_extent() = ExtentForRange(15, 5);
+    *op.mutable_dst_extent() = ExtentForRange(45, 5);
+    op.set_src_offset(123);
+    op.set_type(CowMergeOperation::COW_XOR);
+  }
+  auto generator = MergeSequenceGenerator::Create(aops);
+  ASSERT_NE(generator, nullptr);
+  std::vector<CowMergeOperation> sequence;
+  ASSERT_TRUE(generator->Generate(&sequence));
+  ASSERT_EQ(sequence.size(), 2UL);
+  ASSERT_EQ(sequence[0].src_extent().start_block(), 10UL);
+  ASSERT_EQ(sequence[0].dst_extent().start_block(), 30UL);
+  ASSERT_EQ(sequence[0].src_extent().num_blocks(), 6UL);
+  ASSERT_EQ(sequence[0].dst_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[0].type(), CowMergeOperation::COW_XOR);
+  ASSERT_EQ(sequence[0].src_offset(), 123UL);
+
+  ASSERT_EQ(sequence[1].src_extent().start_block(), 15UL);
+  ASSERT_EQ(sequence[1].dst_extent().start_block(), 45UL);
+  ASSERT_EQ(sequence[1].src_extent().num_blocks(), 6UL);
+  ASSERT_EQ(sequence[1].dst_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[1].type(), CowMergeOperation::COW_XOR);
+  ASSERT_EQ(sequence[1].src_offset(), 123UL);
+
+  ASSERT_TRUE(generator->ValidateSequence(sequence));
+}
+
+TEST_F(MergeSequenceGeneratorTest, CreateGeneratorXorAppendBlock) {
+  std::vector<AnnotatedOperation> aops;
+  auto& aop = aops.emplace_back();
+  aop.op.set_type(InstallOperation::SOURCE_BSDIFF);
+  *aop.op.mutable_src_extents()->Add() = ExtentForRange(10, 10);
+  *aop.op.mutable_dst_extents()->Add() = ExtentForRange(20, 10);
+  auto& xor_map = aop.xor_ops;
+  {
+    auto& op = xor_map.emplace_back();
+    *op.mutable_src_extent() = ExtentForRange(10, 5);
+    *op.mutable_dst_extent() = ExtentForRange(20, 5);
+    op.set_type(CowMergeOperation::COW_XOR);
+  }
+  {
+    auto& op = xor_map.emplace_back();
+    *op.mutable_src_extent() = ExtentForRange(15, 5);
+    *op.mutable_dst_extent() = ExtentForRange(25, 5);
+    op.set_src_offset(123);
+    op.set_type(CowMergeOperation::COW_XOR);
+  }
+  auto generator = MergeSequenceGenerator::Create(aops);
+  ASSERT_NE(generator, nullptr);
+  std::vector<CowMergeOperation> sequence;
+  ASSERT_TRUE(generator->Generate(&sequence));
+  ASSERT_EQ(sequence.size(), 2UL);
+  ASSERT_EQ(sequence[0].src_extent().start_block(), 15UL);
+  ASSERT_EQ(sequence[0].dst_extent().start_block(), 25UL);
+  ASSERT_EQ(sequence[0].src_extent().num_blocks(), 6UL);
+  ASSERT_EQ(sequence[0].dst_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[0].type(), CowMergeOperation::COW_XOR);
+  ASSERT_EQ(sequence[0].src_offset(), 123UL);
+
+  ASSERT_EQ(sequence[1].src_extent().start_block(), 10UL);
+  ASSERT_EQ(sequence[1].dst_extent().start_block(), 20UL);
+  ASSERT_EQ(sequence[1].src_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[1].dst_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[1].type(), CowMergeOperation::COW_XOR);
+
+  ASSERT_TRUE(generator->ValidateSequence(sequence));
+}
+
+TEST_F(MergeSequenceGeneratorTest, CreateGeneratorXorAlreadyPlusOne) {
+  std::vector<AnnotatedOperation> aops;
+  auto& aop = aops.emplace_back();
+  aop.op.set_type(InstallOperation::SOURCE_BSDIFF);
+  *aop.op.mutable_src_extents()->Add() = ExtentForRange(10, 10);
+  *aop.op.mutable_dst_extents()->Add() = ExtentForRange(20, 10);
+  auto& xor_map = aop.xor_ops;
+  {
+    auto& op = xor_map.emplace_back();
+    *op.mutable_src_extent() = ExtentForRange(15, 6);
+    *op.mutable_dst_extent() = ExtentForRange(25, 5);
+    op.set_src_offset(123);
+    op.set_type(CowMergeOperation::COW_XOR);
+  }
+  auto generator = MergeSequenceGenerator::Create(aops);
+  ASSERT_NE(generator, nullptr);
+  std::vector<CowMergeOperation> sequence;
+  ASSERT_TRUE(generator->Generate(&sequence));
+  ASSERT_EQ(sequence.size(), 1UL);
+  ASSERT_EQ(sequence[0].src_extent().start_block(), 15UL);
+  ASSERT_EQ(sequence[0].dst_extent().start_block(), 25UL);
+  ASSERT_EQ(sequence[0].src_extent().num_blocks(), 6UL);
+  ASSERT_EQ(sequence[0].dst_extent().num_blocks(), 5UL);
+  ASSERT_EQ(sequence[0].type(), CowMergeOperation::COW_XOR);
+  ASSERT_EQ(sequence[0].src_offset(), 123UL);
+
+  ASSERT_TRUE(generator->ValidateSequence(sequence));
+}
+
 }  // namespace chromeos_update_engine