Implement cow estimator with XOR enabled

Test: th
Bug: 177104308

Change-Id: Ifa59d750f9ec80de532044e534eeb1bb32b3aa7d
diff --git a/aosp/cow_converter.cc b/aosp/cow_converter.cc
index 8c641b8..3d314c4 100644
--- a/aosp/cow_converter.cc
+++ b/aosp/cow_converter.cc
@@ -60,11 +60,14 @@
   android::snapshot::CowWriter cow_writer{
       {.block_size = static_cast<uint32_t>(block_size), .compression = "gz"}};
   TEST_AND_RETURN_FALSE(cow_writer.Initialize(output_fd));
-  TEST_AND_RETURN_FALSE(CowDryRun(target_img_fd,
+  TEST_AND_RETURN_FALSE(CowDryRun(nullptr,
+                                  target_img_fd,
                                   partition.operations(),
                                   partition.merge_operations(),
                                   block_size,
-                                  &cow_writer));
+                                  &cow_writer,
+                                  partition.new_partition_info().size(),
+                                  false));
   TEST_AND_RETURN_FALSE(cow_writer.Finalize());
   return true;
 }
diff --git a/payload_generator/cow_size_estimator.cc b/payload_generator/cow_size_estimator.cc
index 36d6526..8e2c66d 100644
--- a/payload_generator/cow_size_estimator.cc
+++ b/payload_generator/cow_size_estimator.cc
@@ -16,6 +16,7 @@
 
 #include "update_engine/payload_generator/cow_size_estimator.h"
 
+#include <functional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -25,146 +26,139 @@
 
 #include "update_engine/common/cow_operation_convert.h"
 #include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/extent_utils.h"
 #include "update_engine/update_metadata.pb.h"
 #include "update_engine/payload_consumer/vabc_partition_writer.h"
 
 namespace chromeos_update_engine {
 using android::snapshot::CowWriter;
 
-namespace {
-bool PerformReplaceOp(const InstallOperation& op,
-                      CowWriter* writer,
-                      FileDescriptorPtr target_fd,
-                      size_t block_size) {
-  std::vector<unsigned char> buffer;
-  for (const auto& extent : op.dst_extents()) {
-    buffer.resize(extent.num_blocks() * block_size);
-    // No need to read from payload.bin then decompress, just read from target
-    // directly.
-    ssize_t bytes_read = 0;
-    auto success = utils::ReadAll(target_fd,
-                                  buffer.data(),
-                                  buffer.size(),
-                                  extent.start_block() * block_size,
-                                  &bytes_read);
-    TEST_AND_RETURN_FALSE(success);
-    CHECK_EQ(static_cast<size_t>(bytes_read), buffer.size());
-    TEST_AND_RETURN_FALSE(writer->AddRawBlocks(
-        extent.start_block(), buffer.data(), buffer.size()));
-  }
-  return true;
-}
-
-bool PerformZeroOp(const InstallOperation& op,
-                   CowWriter* writer,
-                   size_t block_size) {
-  for (const auto& extent : op.dst_extents()) {
-    TEST_AND_RETURN_FALSE(
-        writer->AddZeroBlocks(extent.start_block(), extent.num_blocks()));
-  }
-  return true;
-}
-
-bool WriteAllCowOps(size_t block_size,
-                    const std::vector<CowOperation>& converted,
-                    android::snapshot::ICowWriter* cow_writer,
-                    FileDescriptorPtr target_fd) {
-  std::vector<uint8_t> buffer(block_size);
-
-  for (const auto& cow_op : converted) {
-    switch (cow_op.op) {
-      case CowOperation::CowCopy:
-        if (cow_op.src_block == cow_op.dst_block) {
-          continue;
-        }
-        TEST_AND_RETURN_FALSE(
-            cow_writer->AddCopy(cow_op.dst_block, cow_op.src_block));
-        break;
-      case CowOperation::CowReplace:
-        ssize_t bytes_read = 0;
-        TEST_AND_RETURN_FALSE(chromeos_update_engine::utils::ReadAll(
-            target_fd,
-            buffer.data(),
-            block_size,
-            cow_op.dst_block * block_size,
-            &bytes_read));
-        if (bytes_read <= 0 || static_cast<size_t>(bytes_read) != block_size) {
-          LOG(ERROR) << "source_fd->Read failed: " << bytes_read;
-          return false;
-        }
-        TEST_AND_RETURN_FALSE(cow_writer->AddRawBlocks(
-            cow_op.dst_block, buffer.data(), block_size));
-        break;
-    }
-  }
-
-  return true;
-}
-}  // namespace
-
-size_t EstimateCowSize(
+bool CowDryRun(
+    FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
         merge_operations,
-    size_t block_size,
-    std::string compression) {
+    const size_t block_size,
+    android::snapshot::CowWriter* cow_writer,
+    const size_t partition_size,
+    const bool xor_enabled) {
+  CHECK_NE(target_fd, nullptr);
+  CHECK(target_fd->IsOpen());
+  VABCPartitionWriter::WriteMergeSequence(merge_operations, cow_writer);
+  ExtentRanges visited;
+  for (const auto& op : merge_operations) {
+    if (op.type() == CowMergeOperation::COW_COPY) {
+      visited.AddExtent(op.dst_extent());
+      for (size_t i = 0; i < op.dst_extent().num_blocks(); i++) {
+        cow_writer->AddCopy(op.dst_extent().start_block() + i,
+                            op.src_extent().start_block() + i);
+      }
+    } else if (op.type() == CowMergeOperation::COW_XOR && xor_enabled) {
+      CHECK_NE(source_fd, nullptr) << "Source fd is required to enable XOR ops";
+      CHECK(source_fd->IsOpen());
+      visited.AddExtent(op.dst_extent());
+      // dst block count is used, because
+      // src block count is probably(if src_offset > 0) 1 block
+      // larger than dst extent. Using it might lead to intreseting out of bound
+      // disk reads.
+      std::vector<unsigned char> old_data(op.dst_extent().num_blocks() *
+                                          block_size);
+      ssize_t bytes_read = 0;
+      if (!utils::PReadAll(
+              source_fd,
+              old_data.data(),
+              old_data.size(),
+              op.src_extent().start_block() * block_size + op.src_offset(),
+              &bytes_read)) {
+        PLOG(ERROR) << "Failed to read source data at " << op.src_extent();
+        return false;
+      }
+      std::vector<unsigned char> new_data(op.dst_extent().num_blocks() *
+                                          block_size);
+      if (!utils::PReadAll(target_fd,
+                           new_data.data(),
+                           new_data.size(),
+                           op.dst_extent().start_block() * block_size,
+                           &bytes_read)) {
+        PLOG(ERROR) << "Failed to read target data at " << op.dst_extent();
+        return false;
+      }
+      CHECK_GT(old_data.size(), 0UL);
+      CHECK_GT(new_data.size(), 0UL);
+      std::transform(new_data.begin(),
+                     new_data.end(),
+                     old_data.begin(),
+                     new_data.begin(),
+                     std::bit_xor<unsigned char>{});
+      CHECK(cow_writer->AddXorBlocks(op.dst_extent().start_block(),
+                                     new_data.data(),
+                                     new_data.size(),
+                                     op.src_extent().start_block(),
+                                     op.src_offset()));
+    }
+    // The value of label doesn't really matter, we just want to write some
+    // labels to simulate bahvior of update_engine. As update_engine writes
+    // labels every once a while when installing OTA, it's important that we do
+    // the same to get accurate size estimation.
+    cow_writer->AddLabel(0);
+  }
+  for (const auto& op : operations) {
+    if (op.type() == InstallOperation::ZERO) {
+      for (const auto& ext : op.dst_extents()) {
+        visited.AddExtent(ext);
+        cow_writer->AddZeroBlocks(ext.start_block(), ext.num_blocks());
+      }
+      cow_writer->AddLabel(0);
+    }
+  }
+  const size_t last_block = partition_size / block_size;
+  const auto unvisited_extents =
+      FilterExtentRanges({ExtentForRange(0, last_block)}, visited);
+  for (const auto& ext : unvisited_extents) {
+    std::vector<unsigned char> data(ext.num_blocks() * block_size);
+    ssize_t bytes_read = 0;
+    if (!utils::PReadAll(target_fd,
+                         data.data(),
+                         data.size(),
+                         ext.start_block() * block_size,
+                         &bytes_read)) {
+      PLOG(ERROR) << "Failed to read new block data at " << ext;
+      return false;
+    }
+    cow_writer->AddRawBlocks(ext.start_block(), data.data(), data.size());
+    cow_writer->AddLabel(0);
+  }
+
+  return cow_writer->Finalize();
+}
+
+size_t EstimateCowSize(
+    FileDescriptorPtr source_fd,
+    FileDescriptorPtr target_fd,
+    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
+    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations,
+    const size_t block_size,
+    std::string compression,
+    const size_t partition_size,
+    const bool xor_enabled) {
   android::snapshot::CowWriter cow_writer{
       {.block_size = static_cast<uint32_t>(block_size),
        .compression = std::move(compression)}};
   // CowWriter treats -1 as special value, will discard all the data but still
   // reports Cow size. Good for estimation purposes
   cow_writer.Initialize(android::base::borrowed_fd{-1});
-  CHECK(CowDryRun(
-      target_fd, operations, merge_operations, block_size, &cow_writer));
-  CHECK(cow_writer.Finalize());
+  CHECK(CowDryRun(source_fd,
+                  target_fd,
+                  operations,
+                  merge_operations,
+                  block_size,
+                  &cow_writer,
+                  partition_size,
+                  xor_enabled));
   return cow_writer.GetCowSize();
 }
 
-bool CowDryRun(
-    FileDescriptorPtr target_fd,
-    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
-    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
-        merge_operations,
-    size_t block_size,
-    android::snapshot::CowWriter* cow_writer) {
-  const auto converted = ConvertToCowOperations(operations, merge_operations);
-  WriteAllCowOps(block_size, converted, cow_writer, target_fd);
-  cow_writer->AddLabel(0);
-  for (const auto& op : operations) {
-    switch (op.type()) {
-      case InstallOperation::REPLACE:
-      case InstallOperation::REPLACE_BZ:
-      case InstallOperation::REPLACE_XZ:
-        TEST_AND_RETURN_FALSE(
-            PerformReplaceOp(op, cow_writer, target_fd, block_size));
-        break;
-      case InstallOperation::ZERO:
-      case InstallOperation::DISCARD:
-        TEST_AND_RETURN_FALSE(PerformZeroOp(op, cow_writer, block_size));
-        break;
-      case InstallOperation::SOURCE_COPY:
-      case InstallOperation::MOVE:
-        // Already handeled by WriteAllCowOps,
-        break;
-      case InstallOperation::SOURCE_BSDIFF:
-      case InstallOperation::BROTLI_BSDIFF:
-      case InstallOperation::PUFFDIFF:
-      case InstallOperation::BSDIFF:
-      case InstallOperation::ZUCCHINI:
-        // We might do something special by adding CowBsdiff to CowWriter.
-        // For now proceed the same way as normal REPLACE operation.
-        TEST_AND_RETURN_FALSE(
-            PerformReplaceOp(op, cow_writer, target_fd, block_size));
-        break;
-    }
-    // Arbitrary label number, we won't be resuming use these labels here.
-    // They are emitted just to keep size estimates accurate. As update_engine
-    // emits 1 label for every op.
-    cow_writer->AddLabel(2);
-  }
-  VABCPartitionWriter::WriteMergeSequence(merge_operations, cow_writer);
-  // TODO(zhangkelvin) Take FEC extents into account once VABC stabilizes
-  return true;
-}
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/cow_size_estimator.h b/payload_generator/cow_size_estimator.h
index 850c890..c0209d9 100644
--- a/payload_generator/cow_size_estimator.h
+++ b/payload_generator/cow_size_estimator.h
@@ -27,21 +27,28 @@
 // Virtual AB Compression enabled device. This is intended to be used by update
 // generators to put an estimate cow size in OTA payload. When installing an OTA
 // update, libsnapshot will take this estimate as a hint to allocate spaces.
+// If |xor_enabled| is true, then |source_fd| must be non-null.
 size_t EstimateCowSize(
+    FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
         merge_operations,
-    size_t block_size,
-    std::string compression);
+    const size_t block_size,
+    std::string compression,
+    const size_t partition_size,
+    bool xor_enabled);
 
 // Convert InstallOps to CowOps and apply the converted cow op to |cow_writer|
 bool CowDryRun(
+    FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
         merge_operations,
     size_t block_size,
-    android::snapshot::CowWriter* cow_writer);
+    android::snapshot::CowWriter* cow_writer,
+    size_t partition_size,
+    bool xor_enabled);
 
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index ab0f036..c6c70ef 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -120,9 +120,6 @@
     LOG(INFO) << "Estimating COW size for partition: " << new_part_.name;
     // Need the contents of source/target image bytes when doing
     // dry run.
-    FileDescriptorPtr source_fd{new EintrSafeFileDescriptor()};
-    source_fd->Open(old_part_.path.c_str(), O_RDONLY);
-
     auto target_fd = std::make_unique<EintrSafeFileDescriptor>();
     target_fd->Open(new_part_.path.c_str(), O_RDONLY);
 
@@ -131,18 +128,22 @@
     for (const AnnotatedOperation& aop : *aops_) {
       *operations.Add() = aop.op;
     }
+
+    FileDescriptorPtr source_fd = nullptr;
+    if (config_.enable_vabc_xor) {
+      source_fd = std::make_shared<EintrSafeFileDescriptor>();
+      source_fd->Open(old_part_.path.c_str(), O_RDONLY);
+    }
+
     *cow_size_ = EstimateCowSize(
+        std::move(source_fd),
         std::move(target_fd),
         std::move(operations),
         {cow_merge_sequence_->begin(), cow_merge_sequence_->end()},
         config_.block_size,
-        config_.target.dynamic_partition_metadata->vabc_compression_param());
-    if (!new_part_.disable_fec_computation) {
-      *cow_size_ +=
-          new_part_.verity.fec_extent.num_blocks() * config_.block_size;
-    }
-    *cow_size_ +=
-        new_part_.verity.hash_tree_extent.num_blocks() * config_.block_size;
+        config_.target.dynamic_partition_metadata->vabc_compression_param(),
+        new_part_.size,
+        config_.enable_vabc_xor);
     LOG(INFO) << "Estimated COW size for partition: " << new_part_.name << " "
               << *cow_size_;
   }