Implement common library for converting InstallOps to Cow Ops

Test: generate && serve an OTA
Bug: 168554689
Change-Id: If9e87ca3e993372ebb6b24ed64e71b319630bb18
diff --git a/Android.bp b/Android.bp
index a8fdf51..178b7da 100644
--- a/Android.bp
+++ b/Android.bp
@@ -703,6 +703,7 @@
         "common/action_pipe_unittest.cc",
         "common/action_processor_unittest.cc",
         "common/action_unittest.cc",
+        "common/cow_operation_convert_unittest.cc",
         "common/cpu_limiter_unittest.cc",
         "common/fake_prefs.cc",
         "common/file_fetcher_unittest.cc",
diff --git a/common/cow_operation_convert.cc b/common/cow_operation_convert.cc
index a4eaba3..db17b5f 100644
--- a/common/cow_operation_convert.cc
+++ b/common/cow_operation_convert.cc
@@ -16,15 +16,58 @@
 
 #include "update_engine/common/cow_operation_convert.h"
 
+#include <base/logging.h>
+
 #include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/extent_utils.h"
 
 namespace chromeos_update_engine {
+
 std::vector<CowOperation> ConvertToCowOperations(
     const ::google::protobuf::RepeatedPtrField<
         ::chromeos_update_engine::InstallOperation>& operations,
     const ::google::protobuf::RepeatedPtrField<CowMergeOperation>&
         merge_operations) {
-  // TODO(zhangkelvin) Implement this.
-  return {};
+  ExtentRanges merge_extents;
+  std::vector<CowOperation> converted;
+
+  // We want all CowCopy ops to be done first, before any COW_REPLACE happen.
+  // Therefore we add these ops in 2 separate loops. This is because during
+  // merge, a CowReplace might modify a block needed by CowCopy, so we always
+  // perform CowCopy first.
+
+  // This loop handles CowCopy blocks within SOURCE_COPY, and the next loop
+  // converts the leftover blocks to CowReplace?
+  for (const auto& merge_op : merge_operations) {
+    merge_extents.AddExtent(merge_op.dst_extent());
+    const auto& src_extent = merge_op.src_extent();
+    const auto& dst_extent = merge_op.dst_extent();
+    for (uint64_t i = 0; i < src_extent.num_blocks(); i++) {
+      converted.push_back({CowOperation::CowCopy,
+                           src_extent.start_block() + i,
+                           dst_extent.start_block() + i});
+    }
+  }
+  // COW_REPLACE are added after COW_COPY, because replace might modify blocks
+  // needed by COW_COPY. Please don't merge this loop with the previous one.
+  for (const auto& operation : operations) {
+    if (operation.type() != InstallOperation::SOURCE_COPY) {
+      continue;
+    }
+    const auto& src_extents = operation.src_extents();
+    const auto& dst_extents = operation.dst_extents();
+    BlockIterator it1{src_extents};
+    BlockIterator it2{dst_extents};
+    while (!it1.is_end() && !it2.is_end()) {
+      auto src_block = *it1;
+      auto dst_block = *it2;
+      if (!merge_extents.ContainsBlock(dst_block)) {
+        converted.push_back({CowOperation::CowReplace, src_block, dst_block});
+      }
+      ++it1;
+      ++it2;
+    }
+  }
+  return converted;
 }
 }  // namespace chromeos_update_engine
diff --git a/common/cow_operation_convert.h b/common/cow_operation_convert.h
index bca10ac..c0543f7 100644
--- a/common/cow_operation_convert.h
+++ b/common/cow_operation_convert.h
@@ -13,6 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
+
 #ifndef __COW_OPERATION_CONVERT_H
 #define __COW_OPERATION_CONVERT_H
 
diff --git a/common/cow_operation_convert_unittest.cc b/common/cow_operation_convert_unittest.cc
new file mode 100644
index 0000000..b70dcdf
--- /dev/null
+++ b/common/cow_operation_convert_unittest.cc
@@ -0,0 +1,220 @@
+//
+// Copyright (C) 2020 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 <algorithm>
+#include <array>
+#include <initializer_list>
+
+#include <gtest/gtest.h>
+
+#include "update_engine/common/cow_operation_convert.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+using OperationList = ::google::protobuf::RepeatedPtrField<
+    ::chromeos_update_engine::InstallOperation>;
+using MergeOplist = ::google::protobuf::RepeatedPtrField<
+    ::chromeos_update_engine::CowMergeOperation>;
+
+std::ostream& operator<<(std::ostream& out, CowOperation::Type op) {
+  switch (op) {
+    case CowOperation::Type::CowCopy:
+      out << "CowCopy";
+      break;
+    case CowOperation::Type::CowReplace:
+      out << "CowReplace";
+      break;
+    default:
+      out << op;
+      break;
+  }
+  return out;
+}
+
+std::ostream& operator<<(std::ostream& out, const CowOperation& c) {
+  out << "{" << c.op << ", " << c.src_block << ", " << c.dst_block << "}";
+  return out;
+}
+
+class CowOperationConvertTest : public testing::Test {
+ public:
+  void VerifyCowMergeOp(const std::vector<CowOperation>& cow_ops) {
+    // Build a set of all extents covered by InstallOps.
+    ExtentRanges src_extent_set;
+    ExtentRanges dst_extent_set;
+    for (auto&& op : operations_) {
+      src_extent_set.AddRepeatedExtents(op.src_extents());
+      dst_extent_set.AddRepeatedExtents(op.dst_extents());
+    }
+    ExtentRanges modified_extents;
+    for (auto&& cow_op : cow_ops) {
+      if (cow_op.op == CowOperation::CowCopy) {
+        EXPECT_TRUE(src_extent_set.ContainsBlock(cow_op.src_block));
+        // converted operations should be conflict free.
+        EXPECT_FALSE(modified_extents.ContainsBlock(cow_op.src_block))
+            << "SOURCE_COPY operation " << cow_op
+            << " read from a modified block";
+        src_extent_set.SubtractExtent(ExtentForRange(cow_op.src_block, 1));
+      }
+      EXPECT_TRUE(dst_extent_set.ContainsBlock(cow_op.dst_block));
+      dst_extent_set.SubtractExtent(ExtentForRange(cow_op.dst_block, 1));
+      modified_extents.AddBlock(cow_op.dst_block);
+    }
+    // The generated CowOps should cover all extents in InstallOps.
+    EXPECT_EQ(dst_extent_set.blocks(), 0UL);
+    // It's possible that src_extent_set is non-empty, because some operations
+    // will be converted to CowReplace, and we don't count the source extent for
+    // those.
+  }
+  OperationList operations_;
+  MergeOplist merge_operations_;
+};
+
+void AddOperation(OperationList* operations,
+                  ::chromeos_update_engine::InstallOperation_Type op_type,
+                  std::initializer_list<std::array<int, 2>> src_extents,
+                  std::initializer_list<std::array<int, 2>> dst_extents) {
+  auto&& op = operations->Add();
+  op->set_type(op_type);
+  for (const auto& extent : src_extents) {
+    *op->add_src_extents() = ExtentForRange(extent[0], extent[1]);
+  }
+  for (const auto& extent : dst_extents) {
+    *op->add_dst_extents() = ExtentForRange(extent[0], extent[1]);
+  }
+}
+
+void AddMergeOperation(MergeOplist* operations,
+                       ::chromeos_update_engine::CowMergeOperation_Type op_type,
+                       std::array<int, 2> src_extent,
+                       std::array<int, 2> dst_extent) {
+  auto&& op = operations->Add();
+  op->set_type(op_type);
+  *op->mutable_src_extent() = ExtentForRange(src_extent[0], src_extent[1]);
+  *op->mutable_dst_extent() = ExtentForRange(dst_extent[0], dst_extent[1]);
+}
+
+TEST_F(CowOperationConvertTest, NoConflict) {
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{20, 1}}, {{30, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{10, 1}}, {{20, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{0, 1}}, {{10, 1}});
+
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {20, 1}, {30, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {10, 1}, {20, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {0, 1}, {10, 1});
+
+  auto cow_ops = ConvertToCowOperations(operations_, merge_operations_);
+  ASSERT_EQ(cow_ops.size(), 3UL);
+  ASSERT_TRUE(std::all_of(cow_ops.begin(), cow_ops.end(), [](auto&& cow_op) {
+    return cow_op.op == CowOperation::CowCopy;
+  }));
+  VerifyCowMergeOp(cow_ops);
+}
+
+TEST_F(CowOperationConvertTest, CowReplace) {
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{30, 1}}, {{0, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{20, 1}}, {{30, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{10, 1}}, {{20, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{0, 1}}, {{10, 1}});
+
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {20, 1}, {30, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {10, 1}, {20, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {0, 1}, {10, 1});
+
+  auto cow_ops = ConvertToCowOperations(operations_, merge_operations_);
+  ASSERT_EQ(cow_ops.size(), 4UL);
+  // Expect 3 COW_COPY and 1 COW_REPLACE
+  ASSERT_EQ(std::count_if(cow_ops.begin(),
+                          cow_ops.end(),
+                          [](auto&& cow_op) {
+                            return cow_op.op == CowOperation::CowCopy;
+                          }),
+            3);
+  ASSERT_EQ(std::count_if(cow_ops.begin(),
+                          cow_ops.end(),
+                          [](auto&& cow_op) {
+                            return cow_op.op == CowOperation::CowReplace;
+                          }),
+            1);
+  VerifyCowMergeOp(cow_ops);
+}
+
+TEST_F(CowOperationConvertTest, ReOrderSourceCopy) {
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{30, 1}}, {{20, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{20, 1}}, {{10, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{10, 1}}, {{0, 1}});
+
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {10, 1}, {0, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {20, 1}, {10, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {30, 1}, {20, 1});
+
+  auto cow_ops = ConvertToCowOperations(operations_, merge_operations_);
+  ASSERT_EQ(cow_ops.size(), 3UL);
+  // Expect 3 COW_COPY
+  ASSERT_TRUE(std::all_of(cow_ops.begin(), cow_ops.end(), [](auto&& cow_op) {
+    return cow_op.op == CowOperation::CowCopy;
+  }));
+  VerifyCowMergeOp(cow_ops);
+}
+
+TEST_F(CowOperationConvertTest, InterleavingSrcExtent) {
+  AddOperation(&operations_,
+               InstallOperation::SOURCE_COPY,
+               {{30, 5}, {35, 5}},
+               {{20, 10}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{20, 1}}, {{10, 1}});
+  AddOperation(
+      &operations_, InstallOperation::SOURCE_COPY, {{10, 1}}, {{0, 1}});
+
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {10, 1}, {0, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {20, 1}, {10, 1});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {30, 5}, {20, 5});
+  AddMergeOperation(
+      &merge_operations_, CowMergeOperation::COW_COPY, {35, 5}, {25, 5});
+
+  auto cow_ops = ConvertToCowOperations(operations_, merge_operations_);
+  // Expect 4 COW_COPY
+  ASSERT_EQ(cow_ops.size(), 12UL);
+  ASSERT_TRUE(std::all_of(cow_ops.begin(), cow_ops.end(), [](auto&& cow_op) {
+    return cow_op.op == CowOperation::CowCopy;
+  }));
+  VerifyCowMergeOp(cow_ops);
+}
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.h b/payload_generator/extent_utils.h
index 9763b1f..f870b29 100644
--- a/payload_generator/extent_utils.h
+++ b/payload_generator/extent_utils.h
@@ -20,6 +20,8 @@
 #include <string>
 #include <vector>
 
+#include <base/logging.h>
+
 #include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/update_metadata.pb.h"
 
@@ -83,6 +85,43 @@
 
 bool operator==(const Extent& a, const Extent& b);
 
+// TODO(zhangkelvin) This is ugly. Rewrite using C++20's coroutine once
+// that's available. Unfortunately with C++17 this is the best I could do.
+
+// An iterator that takes a sequence of extents, and iterate over blocks
+// inside this sequence of extents.
+// Example usage:
+
+// BlockIterator it1{src_extents};
+// while(!it1.is_end()) {
+//    auto block = *it1;
+//    Do stuff with |block|
+// }
+struct BlockIterator {
+  explicit BlockIterator(
+      const google::protobuf::RepeatedPtrField<Extent>& src_extents)
+      : src_extents_(src_extents) {}
+
+  BlockIterator& operator++() {
+    CHECK_LT(cur_extent_, src_extents_.size());
+    block_offset_++;
+    if (block_offset_ >= src_extents_[cur_extent_].num_blocks()) {
+      cur_extent_++;
+      block_offset_ = 0;
+    }
+    return *this;
+  }
+
+  [[nodiscard]] bool is_end() { return cur_extent_ >= src_extents_.size(); }
+  [[nodiscard]] uint64_t operator*() {
+    return src_extents_[cur_extent_].start_block() + block_offset_;
+  }
+
+  const google::protobuf::RepeatedPtrField<Extent>& src_extents_;
+  int cur_extent_ = 0;
+  size_t block_offset_ = 0;
+};
+
 }  // namespace chromeos_update_engine
 
 #endif  // UPDATE_ENGINE_PAYLOAD_GENERATOR_EXTENT_UTILS_H_