Refactor for new libsnapshot_cow API.

- Don't use CowWriter directly, use ICowWriter.
- Instantiate based on manifest version.
- Replace Emit* calls with Add*.
- Remove manual initialization.

Bug: 280529365
Test: update_engine_unittests
      ota applies
Change-Id: I9c9ac5e45c260b2dca8c41c4718f9667adf90c54
diff --git a/aosp/cow_converter.cc b/aosp/cow_converter.cc
index b9ec00a..80d161f 100644
--- a/aosp/cow_converter.cc
+++ b/aosp/cow_converter.cc
@@ -35,7 +35,6 @@
 #include "update_engine/payload_generator/cow_size_estimator.h"
 #include "update_engine/update_metadata.pb.h"
 
-using android::snapshot::CowWriter;
 DEFINE_string(partitions,
               "",
               "Comma separated list of partitions to extract, leave empty for "
@@ -43,9 +42,10 @@
 
 namespace chromeos_update_engine {
 
-bool ProcessPartition(const chromeos_update_engine::PartitionUpdate& partition,
-                      const char* image_dir,
-                      size_t block_size) {
+bool ProcessPartition(
+    const chromeos_update_engine::DeltaArchiveManifest& manifest,
+    const chromeos_update_engine::PartitionUpdate& partition,
+    const char* image_dir) {
   base::FilePath img_dir{image_dir};
   auto target_img = img_dir.Append(partition.partition_name() + ".img");
   auto output_cow = img_dir.Append(partition.partition_name() + ".cow");
@@ -61,18 +61,23 @@
     return false;
   }
 
-  android::snapshot::CowWriter cow_writer{
-      {.block_size = static_cast<uint32_t>(block_size), .compression = "gz"}};
-  TEST_AND_RETURN_FALSE(cow_writer.Initialize(output_fd));
+  const auto& dap = manifest.dynamic_partition_metadata();
+
+  android::snapshot::CowOptions options{
+      .block_size = static_cast<uint32_t>(manifest.block_size()),
+      .compression = dap.vabc_compression_param()};
+  auto cow_writer = android::snapshot::CreateCowWriter(
+      dap.cow_version(), options, std::move(output_fd));
+  TEST_AND_RETURN_FALSE(cow_writer);
   TEST_AND_RETURN_FALSE(CowDryRun(nullptr,
                                   target_img_fd,
                                   partition.operations(),
                                   partition.merge_operations(),
-                                  block_size,
-                                  &cow_writer,
+                                  manifest.block_size(),
+                                  cow_writer.get(),
                                   partition.new_partition_info().size(),
                                   false));
-  TEST_AND_RETURN_FALSE(cow_writer.Finalize());
+  TEST_AND_RETURN_FALSE(cow_writer->Finalize());
   return true;
 }
 
@@ -148,7 +153,7 @@
       continue;
     }
     LOG(INFO) << partition.partition_name();
-    if (!ProcessPartition(partition, images_dir, manifest.block_size())) {
+    if (!ProcessPartition(manifest, partition, images_dir)) {
       return 6;
     }
     base::FilePath img_dir{images_dir};
diff --git a/payload_consumer/cow_writer_file_descriptor.cc b/payload_consumer/cow_writer_file_descriptor.cc
index 19fb545..2dca1b3 100644
--- a/payload_consumer/cow_writer_file_descriptor.cc
+++ b/payload_consumer/cow_writer_file_descriptor.cc
@@ -87,9 +87,9 @@
 
 ssize_t CowWriterFileDescriptor::Write(const void* buf, size_t count) {
   auto offset = cow_reader_->Seek(0, SEEK_CUR);
-  CHECK_EQ(offset % cow_writer_->options().block_size, 0);
+  CHECK_EQ(offset % cow_writer_->GetBlockSize(), 0);
   auto success = cow_writer_->AddRawBlocks(
-      offset / cow_writer_->options().block_size, buf, count);
+      offset / cow_writer_->GetBlockSize(), buf, count);
   if (success) {
     if (cow_reader_->Seek(count, SEEK_CUR) < 0) {
       return -1;
diff --git a/payload_consumer/snapshot_extent_writer_unittest.cc b/payload_consumer/snapshot_extent_writer_unittest.cc
index d43d3a1..01c3d2e 100644
--- a/payload_consumer/snapshot_extent_writer_unittest.cc
+++ b/payload_consumer/snapshot_extent_writer_unittest.cc
@@ -43,9 +43,9 @@
   using ICowWriter::ICowWriter;
   ~FakeCowWriter() = default;
 
-  bool EmitCopy(uint64_t new_block,
-                uint64_t old_block,
-                uint64_t num_blocks) override {
+  bool AddCopy(uint64_t new_block,
+               uint64_t old_block,
+               uint64_t num_blocks) override {
     for (size_t i = 0; i < num_blocks; i++) {
       operations_[new_block + i] = {
           .type = CowOp::COW_COPY,
@@ -53,23 +53,23 @@
     }
     return true;
   }
-  bool EmitRawBlocks(uint64_t new_block_start,
-                     const void* data,
-                     size_t size) override {
+  bool AddRawBlocks(uint64_t new_block_start,
+                    const void* data,
+                    size_t size) override {
     auto&& op = operations_[new_block_start];
     const auto uint8_ptr = static_cast<const unsigned char*>(data);
     op.data.insert(op.data.end(), uint8_ptr, uint8_ptr + size);
     return true;
   }
-  bool EmitZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) override {
+  bool AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) override {
     operations_[new_block_start] = {.type = CowOp::COW_ZERO};
     return true;
   }
-  bool EmitXorBlocks(uint32_t new_block_start,
-                     const void* data,
-                     size_t size,
-                     uint32_t old_block,
-                     uint16_t offset) override {
+  bool AddXorBlocks(uint32_t new_block_start,
+                    const void* data,
+                    size_t size,
+                    uint32_t old_block,
+                    uint16_t offset) override {
     return false;
   }
   bool Finalize() override {
@@ -77,15 +77,18 @@
     return true;
   }
 
-  bool EmitLabel(uint64_t label) {
+  bool AddLabel(uint64_t label) {
     label_count_++;
     return true;
   }
 
-  bool EmitSequenceData(size_t num_ops, const uint32_t* data) override {
+  bool AddSequenceData(size_t num_ops, const uint32_t* data) override {
     return false;
   }
 
+  uint32_t GetBlockSize() const override { return 4096; }
+  std::optional<uint32_t> GetMaxBlocks() const override { return {}; }
+
   // Return number of bytes the cow image occupies on disk.
   uint64_t GetCowSize() override {
     return std::accumulate(
@@ -108,7 +111,7 @@
  protected:
   android::snapshot::CowOptions options_ = {
       .block_size = static_cast<uint32_t>(kBlockSize)};
-  FakeCowWriter cow_writer_{options_};
+  FakeCowWriter cow_writer_;
   SnapshotExtentWriter writer_{&cow_writer_};
 };
 
diff --git a/payload_consumer/vabc_partition_writer_unittest.cc b/payload_consumer/vabc_partition_writer_unittest.cc
index 3cdf3bc..93c692b 100644
--- a/payload_consumer/vabc_partition_writer_unittest.cc
+++ b/payload_consumer/vabc_partition_writer_unittest.cc
@@ -57,7 +57,7 @@
   }
 
  protected:
-  void EmitBlockTest(bool xor_enabled);
+  void AddBlockTest(bool xor_enabled);
   CowMergeOperation* AddMergeOp(PartitionUpdate* partition,
                                 std::array<size_t, 2> src_extent,
                                 std::array<size_t, 2> dst_extent,
@@ -75,7 +75,7 @@
 
   android::snapshot::CowOptions options_ = {
       .block_size = static_cast<uint32_t>(kBlockSize)};
-  android::snapshot::MockSnapshotWriter cow_writer_{options_};
+  android::snapshot::MockSnapshotWriter cow_writer_;
   MockDynamicPartitionControl dynamic_control_;
   PartitionUpdate partition_update_;
   InstallPlan install_plan_;
@@ -98,15 +98,14 @@
                           const std::optional<std::string>&,
                           bool) {
         auto cow_writer =
-            std::make_unique<android::snapshot::MockSnapshotWriter>(
-                android::snapshot::CowOptions{});
+            std::make_unique<android::snapshot::MockSnapshotWriter>();
         auto expected_merge_sequence = {10, 14, 13, 20, 25, 40, 41, 42, 43, 44};
         EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
-        EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+        EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
             .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
             .WillOnce(Return(true));
-        ON_CALL(*cow_writer, EmitCopy(_, _, _)).WillByDefault(Return(true));
-        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+        ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
+        ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
         return cow_writer;
       }));
   EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
@@ -123,15 +122,14 @@
       .WillOnce(Invoke(
           [](const std::string&, const std::optional<std::string>&, bool) {
             auto cow_writer =
-                std::make_unique<android::snapshot::MockSnapshotWriter>(
-                    android::snapshot::CowOptions{});
+                std::make_unique<android::snapshot::MockSnapshotWriter>();
             auto expected_merge_sequence = {19, 20, 21};
             EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
-            EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+            EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
                 .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
                 .WillOnce(Return(true));
-            ON_CALL(*cow_writer, EmitCopy(_, _, _)).WillByDefault(Return(true));
-            ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+            ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
+            ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
             return cow_writer;
           }));
   EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
@@ -139,15 +137,15 @@
   ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
 }
 
-TEST_F(VABCPartitionWriterTest, EmitBlockTestXor) {
-  return EmitBlockTest(true);
+TEST_F(VABCPartitionWriterTest, AddBlockTestXor) {
+  return AddBlockTest(true);
 }
 
-TEST_F(VABCPartitionWriterTest, EmitBlockTestNoXor) {
-  return EmitBlockTest(false);
+TEST_F(VABCPartitionWriterTest, AddBlockTestNoXor) {
+  return AddBlockTest(false);
 }
 
-void VABCPartitionWriterTest::EmitBlockTest(bool xor_enabled) {
+void VABCPartitionWriterTest::AddBlockTest(bool xor_enabled) {
   if (xor_enabled) {
     ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
         .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
@@ -177,33 +175,32 @@
                                      const std::optional<std::string>&,
                                      bool) {
         auto cow_writer =
-            std::make_unique<android::snapshot::MockSnapshotWriter>(
-                android::snapshot::CowOptions{});
-        ON_CALL(*cow_writer, EmitCopy(_, _, _)).WillByDefault(Return(true));
-        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+            std::make_unique<android::snapshot::MockSnapshotWriter>();
+        ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
+        ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
         ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
         EXPECT_CALL(*cow_writer, Initialize());
         if (xor_enabled) {
-          EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+          EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
               .WillOnce(Return(true));
-          EXPECT_CALL(*cow_writer, EmitCopy(10, 5, 1));
-          EXPECT_CALL(*cow_writer, EmitCopy(15, 10, 1));
+          EXPECT_CALL(*cow_writer, AddCopy(10, 5, 1));
+          EXPECT_CALL(*cow_writer, AddCopy(15, 10, 1));
           // libsnapshot want blocks in reverser order, so 21 goes before 20
-          EXPECT_CALL(*cow_writer, EmitCopy(20, 15, 2));
+          EXPECT_CALL(*cow_writer, AddCopy(20, 15, 2));
 
-          EXPECT_CALL(*cow_writer, EmitCopy(25, 20, 1));
-          EXPECT_CALL(*cow_writer, EmitRawBlocks(26, _, 4096))
+          EXPECT_CALL(*cow_writer, AddCopy(25, 20, 1));
+          EXPECT_CALL(*cow_writer, AddRawBlocks(26, _, 4096))
               .WillOnce(Return(true));
           EXPECT_CALL(*cow_writer, Finalize());
         } else {
           Sequence s;
-          EXPECT_CALL(*cow_writer, EmitCopy(10, 5, 1)).InSequence(s);
-          EXPECT_CALL(*cow_writer, EmitCopy(15, 10, 1)).InSequence(s);
+          EXPECT_CALL(*cow_writer, AddCopy(10, 5, 1)).InSequence(s);
+          EXPECT_CALL(*cow_writer, AddCopy(15, 10, 1)).InSequence(s);
           // libsnapshot want blocks in reverser order, so 21 goes before 20
-          EXPECT_CALL(*cow_writer, EmitCopy(20, 15, 2)).InSequence(s);
+          EXPECT_CALL(*cow_writer, AddCopy(20, 15, 2)).InSequence(s);
 
-          EXPECT_CALL(*cow_writer, EmitCopy(25, 20, 1)).InSequence(s);
-          EXPECT_CALL(*cow_writer, EmitRawBlocks(26, _, 4096))
+          EXPECT_CALL(*cow_writer, AddCopy(25, 20, 1)).InSequence(s);
+          EXPECT_CALL(*cow_writer, AddRawBlocks(26, _, 4096))
               .InSequence(s)
               .WillOnce(Return(true));
         }
@@ -254,29 +251,28 @@
                           const std::optional<std::string>&,
                           bool) {
         auto cow_writer =
-            std::make_unique<android::snapshot::MockSnapshotWriter>(
-                android::snapshot::CowOptions{});
-        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+            std::make_unique<android::snapshot::MockSnapshotWriter>();
+        ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
         auto expected_merge_sequence = {10, 11, 13, 14};
         auto expected_merge_sequence_rev = {11, 10, 14, 13};
         const bool is_ascending = android::base::GetBoolProperty(
             "ro.virtual_ab.userspace.snapshots.enabled", false);
         ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
         if (!is_ascending) {
-          EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+          EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
               .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence_rev)))
               .WillOnce(Return(true));
         } else {
-          EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+          EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
               .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
               .WillOnce(Return(true));
         }
         EXPECT_CALL(*cow_writer, Initialize()).Times(1);
-        EXPECT_CALL(*cow_writer, EmitCopy(_, _, _)).Times(0);
-        EXPECT_CALL(*cow_writer, EmitRawBlocks(_, _, _)).WillOnce(Return(true));
-        EXPECT_CALL(*cow_writer, EmitXorBlocks(10, _, kBlockSize * 2, 5, 0))
+        EXPECT_CALL(*cow_writer, AddCopy(_, _, _)).Times(0);
+        EXPECT_CALL(*cow_writer, AddRawBlocks(_, _, _)).WillOnce(Return(true));
+        EXPECT_CALL(*cow_writer, AddXorBlocks(10, _, kBlockSize * 2, 5, 0))
             .WillOnce(Return(true));
-        EXPECT_CALL(*cow_writer, EmitXorBlocks(13, _, kBlockSize * 2, 8, 0))
+        EXPECT_CALL(*cow_writer, AddXorBlocks(13, _, kBlockSize * 2, 8, 0))
             .WillOnce(Return(true));
         return cow_writer;
       }));
diff --git a/payload_consumer/xor_extent_writer_unittest.cc b/payload_consumer/xor_extent_writer_unittest.cc
index 55c3c6c..e6cd89a 100644
--- a/payload_consumer/xor_extent_writer_unittest.cc
+++ b/payload_consumer/xor_extent_writer_unittest.cc
@@ -75,8 +75,7 @@
 
 TEST_F(XorExtentWriterTest, StreamTest) {
   constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
-  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
-      .WillByDefault(Return(true));
+  ON_CALL(cow_writer_, AddXorBlocks(_, _, _, _, _)).WillByDefault(Return(true));
   const auto op1 = CreateCowMergeOperation(
       ExtentForRange(5, 2), ExtentForRange(5, 2), COW_XOR);
   ASSERT_TRUE(xor_map_.AddExtent(op1.dst_extent(), &op1));
@@ -109,23 +108,23 @@
 
   auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 10);
   EXPECT_CALL(cow_writer_,
-              EmitRawBlocks(455, zeros->data() + 2 * kBlockSize, kBlockSize))
+              AddRawBlocks(455, zeros->data() + 2 * kBlockSize, kBlockSize))
       .With(Args<1, 2>(BytesEqual(zeros->data(), kBlockSize)))
       .WillOnce(Return(true));
   EXPECT_CALL(cow_writer_,
-              EmitRawBlocks(320, zeros->data() + 5 * kBlockSize, kBlockSize))
+              AddRawBlocks(320, zeros->data() + 5 * kBlockSize, kBlockSize))
       .With(Args<1, 2>(BytesEqual(zeros->data(), kBlockSize)))
       .WillOnce(Return(true));
   EXPECT_CALL(cow_writer_,
-              EmitRawBlocks(323, zeros->data() + 8 * kBlockSize, kBlockSize))
+              AddRawBlocks(323, zeros->data() + 8 * kBlockSize, kBlockSize))
       .With(Args<1, 2>(BytesEqual(zeros->data(), kBlockSize)))
       .WillOnce(Return(true));
 
-  EXPECT_CALL(cow_writer_, EmitXorBlocks(5, _, kBlockSize * 2, 5, 0))
+  EXPECT_CALL(cow_writer_, AddXorBlocks(5, _, kBlockSize * 2, 5, 0))
       .WillOnce(Return(true));
-  EXPECT_CALL(cow_writer_, EmitXorBlocks(456, _, kBlockSize * 2, 45, 0))
+  EXPECT_CALL(cow_writer_, AddXorBlocks(456, _, kBlockSize * 2, 45, 0))
       .WillOnce(Return(true));
-  EXPECT_CALL(cow_writer_, EmitXorBlocks(321, _, kBlockSize * 2, 12, 777))
+  EXPECT_CALL(cow_writer_, AddXorBlocks(321, _, kBlockSize * 2, 12, 777))
       .WillOnce(Return(true));
 
   ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
@@ -134,8 +133,7 @@
 
 TEST_F(XorExtentWriterTest, SubsetExtentTest) {
   constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
-  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
-      .WillByDefault(Return(true));
+  ON_CALL(cow_writer_, AddXorBlocks(_, _, _, _, _)).WillByDefault(Return(true));
 
   const auto op3 = CreateCowMergeOperation(
       ExtentForRange(12, 4), ExtentForRange(320, 4), COW_XOR, 777);
@@ -161,14 +159,13 @@
   // [420-422] should be regular replace blocks
 
   auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 7);
-  EXPECT_CALL(
-      cow_writer_,
-      EmitRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
+  EXPECT_CALL(cow_writer_,
+              AddRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
       .WillOnce(Return(true));
 
-  EXPECT_CALL(cow_writer_, EmitXorBlocks(320, _, kBlockSize * 3, 12, 777))
+  EXPECT_CALL(cow_writer_, AddXorBlocks(320, _, kBlockSize * 3, 12, 777))
       .WillOnce(Return(true));
-  EXPECT_CALL(cow_writer_, EmitXorBlocks(323, _, kBlockSize, 15, 777))
+  EXPECT_CALL(cow_writer_, AddXorBlocks(323, _, kBlockSize, 15, 777))
       .WillOnce(Return(true));
 
   ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
@@ -177,8 +174,7 @@
 
 TEST_F(XorExtentWriterTest, LastBlockTest) {
   constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
-  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
-      .WillByDefault(Return(true));
+  ON_CALL(cow_writer_, AddXorBlocks(_, _, _, _, _)).WillByDefault(Return(true));
 
   const auto op3 = CreateCowMergeOperation(
       ExtentForRange(NUM_BLOCKS - 1, 1), ExtentForRange(2, 1), COW_XOR, 777);
@@ -207,18 +203,17 @@
   // [2-4] should be REPLACE blocks
 
   auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 9);
-  EXPECT_CALL(cow_writer_, EmitRawBlocks(320, zeros->data(), kBlockSize * 3))
+  EXPECT_CALL(cow_writer_, AddRawBlocks(320, zeros->data(), kBlockSize * 3))
       .WillOnce(Return(true));
-  EXPECT_CALL(
-      cow_writer_,
-      EmitRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
+  EXPECT_CALL(cow_writer_,
+              AddRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
       .WillOnce(Return(true));
 
   EXPECT_CALL(cow_writer_,
-              EmitRawBlocks(2, zeros->data() + 6 * kBlockSize, kBlockSize))
+              AddRawBlocks(2, zeros->data() + 6 * kBlockSize, kBlockSize))
       .WillOnce(Return(true));
   EXPECT_CALL(cow_writer_,
-              EmitRawBlocks(3, zeros->data() + 7 * kBlockSize, kBlockSize * 2))
+              AddRawBlocks(3, zeros->data() + 7 * kBlockSize, kBlockSize * 2))
       .WillOnce(Return(true));
 
   ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
diff --git a/payload_generator/cow_size_estimator.cc b/payload_generator/cow_size_estimator.cc
index 0c918e7..704aeaa 100644
--- a/payload_generator/cow_size_estimator.cc
+++ b/payload_generator/cow_size_estimator.cc
@@ -33,7 +33,8 @@
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
-using android::snapshot::CowWriter;
+using android::snapshot::CreateCowEstimator;
+using android::snapshot::ICowWriter;
 
 bool CowDryRun(
     FileDescriptorPtr source_fd,
@@ -42,7 +43,7 @@
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
         merge_operations,
     const size_t block_size,
-    android::snapshot::CowWriter* cow_writer,
+    android::snapshot::ICowWriter* cow_writer,
     const size_t partition_size,
     const bool xor_enabled) {
   CHECK_NE(target_fd, nullptr);
@@ -145,21 +146,21 @@
     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});
+  android::snapshot::CowOptions options{
+      .block_size = static_cast<uint32_t>(block_size),
+      .compression = std::move(compression)};
+  auto cow_writer =
+      CreateCowEstimator(android::snapshot::kCowVersionManifest, options);
+  CHECK_NE(cow_writer, nullptr) << "Could not create cow estimator";
   CHECK(CowDryRun(source_fd,
                   target_fd,
                   operations,
                   merge_operations,
                   block_size,
-                  &cow_writer,
+                  cow_writer.get(),
                   partition_size,
                   xor_enabled));
-  return cow_writer.GetCowSize();
+  return cow_writer->GetCowSize();
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/cow_size_estimator.h b/payload_generator/cow_size_estimator.h
index c0209d9..3bf6880 100644
--- a/payload_generator/cow_size_estimator.h
+++ b/payload_generator/cow_size_estimator.h
@@ -47,7 +47,7 @@
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
         merge_operations,
     size_t block_size,
-    android::snapshot::CowWriter* cow_writer,
+    android::snapshot::ICowWriter* cow_writer,
     size_t partition_size,
     bool xor_enabled);