diff --git a/aosp/cow_converter.cc b/aosp/cow_converter.cc
index 3e8e5fd..32aa12f 100644
--- a/aosp/cow_converter.cc
+++ b/aosp/cow_converter.cc
@@ -95,6 +95,7 @@
                                   manifest.block_size(),
                                   cow_writer.get(),
                                   partition.new_partition_info().size(),
+                                  partition.old_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 4be7ea2..cc1034d 100644
--- a/payload_generator/cow_size_estimator.cc
+++ b/payload_generator/cow_size_estimator.cc
@@ -57,7 +57,8 @@
         merge_operations,
     const size_t block_size,
     android::snapshot::ICowWriter* cow_writer,
-    const size_t partition_size,
+    const size_t new_partition_size,
+    const size_t old_partition_size,
     const bool xor_enabled) {
   CHECK_NE(target_fd, nullptr);
   CHECK(target_fd->IsOpen());
@@ -84,7 +85,7 @@
         if (xor_enabled) {
           std::unique_ptr<XORExtentWriter> writer =
               std::make_unique<XORExtentWriter>(
-                  op, source_fd, cow_writer, xor_map, partition_size);
+                  op, source_fd, cow_writer, xor_map, old_partition_size);
           TEST_AND_RETURN_FALSE(writer->Init(op.dst_extents(), block_size));
           for (const auto& ext : op.dst_extents()) {
             visited.AddExtent(ext);
@@ -153,7 +154,7 @@
     }
   }
 
-  const size_t last_block = partition_size / block_size;
+  const size_t last_block = new_partition_size / block_size;
   const auto unvisited_extents =
       FilterExtentRanges({ExtentForRange(0, last_block)}, visited);
   for (const auto& ext : unvisited_extents) {
@@ -194,14 +195,15 @@
         merge_operations,
     const size_t block_size,
     std::string compression,
-    const size_t partition_size,
+    const size_t new_partition_size,
+    const size_t old_partition_size,
     const bool xor_enabled,
     uint32_t cow_version,
     uint64_t compression_factor) {
   android::snapshot::CowOptions options{
       .block_size = static_cast<uint32_t>(block_size),
       .compression = std::move(compression),
-      .max_blocks = (partition_size / block_size),
+      .max_blocks = (new_partition_size / block_size),
       .compression_factor = compression_factor};
   auto cow_writer = CreateCowEstimator(cow_version, options);
   CHECK_NE(cow_writer, nullptr) << "Could not create cow estimator";
@@ -211,7 +213,8 @@
                   merge_operations,
                   block_size,
                   cow_writer.get(),
-                  partition_size,
+                  new_partition_size,
+                  old_partition_size,
                   xor_enabled));
   return cow_writer->GetCowSizeInfo();
 }
diff --git a/payload_generator/cow_size_estimator.h b/payload_generator/cow_size_estimator.h
index 060da43..9af3938 100644
--- a/payload_generator/cow_size_estimator.h
+++ b/payload_generator/cow_size_estimator.h
@@ -36,7 +36,8 @@
         merge_operations,
     const size_t block_size,
     std::string compression,
-    const size_t partition_size,
+    const size_t new_partition_size,
+    const size_t old_partition_size,
     bool xor_enabled,
     uint32_t cow_version,
     uint64_t compression_factor);
@@ -50,7 +51,8 @@
         merge_operations,
     size_t block_size,
     android::snapshot::ICowWriter* cow_writer,
-    size_t partition_size,
+    const size_t new_partition_size,
+    const size_t old_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 44d97df..426efbc 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -139,6 +139,7 @@
         config_.block_size,
         config_.target.dynamic_partition_metadata->vabc_compression_param(),
         new_part_.size,
+        old_part_.size,
         config_.enable_vabc_xor,
         config_.target.dynamic_partition_metadata->cow_version(),
         config_.target.dynamic_partition_metadata->compression_factor());