Generate update operations per partition

In order to support multiple partitions, OperationsGenerator should take one
partition at a time, and be called multiple times. This also reduce code
repetition.

TEST=FEATURES=test emerge-peppy update_engine; cros_generate_update_payload
generate same payload as before.

BUG: 23037923
Change-Id: Iefbe87401217873e6dbdc616791667cfd1ec76a3
diff --git a/payload_generator/ab_generator.cc b/payload_generator/ab_generator.cc
index 23256d8..459b3b9 100644
--- a/payload_generator/ab_generator.cc
+++ b/payload_generator/ab_generator.cc
@@ -34,44 +34,30 @@
 
 bool ABGenerator::GenerateOperations(
     const PayloadGenerationConfig& config,
+    const PartitionConfig& old_part,
+    const PartitionConfig& new_part,
     BlobFileWriter* blob_file,
-    vector<AnnotatedOperation>* rootfs_ops,
-    vector<AnnotatedOperation>* kernel_ops) {
+    vector<AnnotatedOperation>* aops) {
 
   ssize_t hard_chunk_blocks = (config.hard_chunk_size == -1 ? -1 :
                                config.hard_chunk_size / config.block_size);
   size_t soft_chunk_blocks = config.soft_chunk_size / config.block_size;
 
-  rootfs_ops->clear();
+  aops->clear();
   TEST_AND_RETURN_FALSE(diff_utils::DeltaReadPartition(
-      rootfs_ops,
-      config.source.rootfs,
-      config.target.rootfs,
+      aops,
+      old_part,
+      new_part,
       hard_chunk_blocks,
       soft_chunk_blocks,
       blob_file,
       true));  // src_ops_allowed
-  LOG(INFO) << "done reading normal files";
+  LOG(INFO) << "done reading " << PartitionNameString(new_part.name);
 
-  // Read kernel partition
-  TEST_AND_RETURN_FALSE(diff_utils::DeltaReadPartition(
-      kernel_ops,
-      config.source.kernel,
-      config.target.kernel,
-      hard_chunk_blocks,
-      soft_chunk_blocks,
-      blob_file,
-      true));  // src_ops_allowed
-  LOG(INFO) << "done reading kernel";
-
-  TEST_AND_RETURN_FALSE(FragmentOperations(rootfs_ops,
-                                           config.target.rootfs.path,
+  TEST_AND_RETURN_FALSE(FragmentOperations(aops,
+                                           new_part.path,
                                            blob_file));
-  TEST_AND_RETURN_FALSE(FragmentOperations(kernel_ops,
-                                           config.target.kernel.path,
-                                           blob_file));
-  SortOperationsByDestination(rootfs_ops);
-  SortOperationsByDestination(kernel_ops);
+  SortOperationsByDestination(aops);
 
   // Use the soft_chunk_size when merging operations to prevent merging all
   // the operations into a huge one if there's no hard limit.
@@ -81,13 +67,9 @@
     merge_chunk_blocks = hard_chunk_blocks;
   }
 
-  TEST_AND_RETURN_FALSE(MergeOperations(rootfs_ops,
+  TEST_AND_RETURN_FALSE(MergeOperations(aops,
                                         merge_chunk_blocks,
-                                        config.target.rootfs.path,
-                                        blob_file));
-  TEST_AND_RETURN_FALSE(MergeOperations(kernel_ops,
-                                        merge_chunk_blocks,
-                                        config.target.kernel.path,
+                                        new_part.path,
                                         blob_file));
   return true;
 }
diff --git a/payload_generator/ab_generator.h b/payload_generator/ab_generator.h
index c4cb7af..7c89ed6 100644
--- a/payload_generator/ab_generator.h
+++ b/payload_generator/ab_generator.h
@@ -40,29 +40,30 @@
  public:
   ABGenerator() = default;
 
-  // Generate the update payload operations for the kernel and rootfs using
+  // Generate the update payload operations for the given partition using
   // SOURCE_* operations, used for generating deltas for the minor version
   // kSourceMinorPayloadVersion. This function will generate operations in the
-  // rootfs that will read blocks from the source partition in random order and
-  // write the new image on the target partition, also possibly in random order.
-  // The rootfs operations are stored in |rootfs_ops| and should be executed in
-  // that order. The kernel operations are stored in |kernel_ops|. All
-  // the offsets in the operations reference the data written to |blob_file|.
+  // partition that will read blocks from the source partition in random order
+  // and write the new image on the target partition, also possibly in random
+  // order. The operations are stored in |aops| and should be executed in that
+  // order. All the offsets in the operations reference the data written to
+  // |blob_file|.
   bool GenerateOperations(
       const PayloadGenerationConfig& config,
+      const PartitionConfig& old_part,
+      const PartitionConfig& new_part,
       BlobFileWriter* blob_file,
-      std::vector<AnnotatedOperation>* rootfs_ops,
-      std::vector<AnnotatedOperation>* kernel_ops) override;
+      std::vector<AnnotatedOperation>* aops) override;
 
   // Split the operations in the vector of AnnotatedOperations |aops|
   // such that for every operation there is only one dst extent and updates
   // |aops| with the new list of operations. All kinds of operations are
   // fragmented except BSDIFF and SOURCE_BSDIFF operations.
-  // The |target_rootfs_part| is the filename of the new image, where the
+  // The |target_part_path| is the filename of the new image, where the
   // destination extents refer to. The blobs of the operations in |aops| should
   // reference |blob_file|. |blob_file| are updated if needed.
   static bool FragmentOperations(std::vector<AnnotatedOperation>* aops,
-                                 const std::string& target_rootfs_part,
+                                 const std::string& target_part_path,
                                  BlobFileWriter* blob_file);
 
   // Takes a vector of AnnotatedOperations |aops| and sorts them by the first
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index d88e20f..c653433 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -131,8 +131,14 @@
     BlobFileWriter blob_file(data_file_fd, &data_file_size);
     // Generate the operations using the strategy we selected above.
     TEST_AND_RETURN_FALSE(strategy->GenerateOperations(config,
+                                                       config.source.rootfs,
+                                                       config.target.rootfs,
                                                        &blob_file,
-                                                       &rootfs_ops,
+                                                       &rootfs_ops));
+    TEST_AND_RETURN_FALSE(strategy->GenerateOperations(config,
+                                                       config.source.kernel,
+                                                       config.target.kernel,
+                                                       &blob_file,
                                                        &kernel_ops));
   }
 
diff --git a/payload_generator/full_update_generator.cc b/payload_generator/full_update_generator.cc
index a627823..101f28a 100644
--- a/payload_generator/full_update_generator.cc
+++ b/payload_generator/full_update_generator.cc
@@ -121,10 +121,11 @@
 
 bool FullUpdateGenerator::GenerateOperations(
     const PayloadGenerationConfig& config,
+    const PartitionConfig& old_part,
+    const PartitionConfig& new_part,
     BlobFileWriter* blob_file,
-    vector<AnnotatedOperation>* rootfs_ops,
-    vector<AnnotatedOperation>* kernel_ops) {
-  TEST_AND_RETURN_FALSE(config.Validate());
+    vector<AnnotatedOperation>* aops) {
+  TEST_AND_RETURN_FALSE(new_part.ValidateExists());
 
   // FullUpdateGenerator requires a positive chunk_size, otherwise there will
   // be only one operation with the whole partition which should not be allowed.
@@ -143,31 +144,11 @@
   TEST_AND_RETURN_FALSE(full_chunk_size > 0);
   TEST_AND_RETURN_FALSE(full_chunk_size % config.block_size == 0);
 
-  TEST_AND_RETURN_FALSE(GenerateOperationsForPartition(
-      config.target.rootfs,
-      config.block_size,
-      full_chunk_size / config.block_size,
-      blob_file,
-      rootfs_ops));
-  TEST_AND_RETURN_FALSE(GenerateOperationsForPartition(
-      config.target.kernel,
-      config.block_size,
-      full_chunk_size / config.block_size,
-      blob_file,
-      kernel_ops));
-  return true;
-}
-
-bool FullUpdateGenerator::GenerateOperationsForPartition(
-    const PartitionConfig& new_part,
-    size_t block_size,
-    size_t chunk_blocks,
-    BlobFileWriter* blob_file,
-    vector<AnnotatedOperation>* aops) {
+  size_t chunk_blocks = full_chunk_size / config.block_size;
   size_t max_threads = std::max(sysconf(_SC_NPROCESSORS_ONLN), 4L);
   LOG(INFO) << "Compressing partition " << PartitionNameString(new_part.name)
             << " from " << new_part.path << " splitting in chunks of "
-            << chunk_blocks << " blocks (" << block_size
+            << chunk_blocks << " blocks (" << config.block_size
             << " bytes each) using " << max_threads << " threads";
 
   int in_fd = open(new_part.path.c_str(), O_RDONLY, 0);
@@ -176,7 +157,7 @@
 
   // We potentially have all the ChunkProcessors in memory but only
   // |max_threads| will actually hold a block in memory while we process.
-  size_t partition_blocks = new_part.size / block_size;
+  size_t partition_blocks = new_part.size / config.block_size;
   size_t num_chunks = (partition_blocks + chunk_blocks - 1) / chunk_blocks;
   aops->resize(num_chunks);
   vector<ChunkProcessor> chunk_processors;
@@ -202,8 +183,8 @@
 
     chunk_processors.emplace_back(
         in_fd,
-        static_cast<off_t>(start_block) * block_size,
-        num_blocks * block_size,
+        static_cast<off_t>(start_block) * config.block_size,
+        num_blocks * config.block_size,
         blob_file,
         aop);
   }
diff --git a/payload_generator/full_update_generator.h b/payload_generator/full_update_generator.h
index 47829d8..d722028 100644
--- a/payload_generator/full_update_generator.h
+++ b/payload_generator/full_update_generator.h
@@ -35,28 +35,14 @@
   // OperationsGenerator override.
   // Creates a full update for the target image defined in |config|. |config|
   // must be a valid payload generation configuration for a full payload.
-  // Populates |rootfs_ops| and |kernel_ops|, with data about the update
-  // operations, and writes relevant data to |data_file_fd|, updating
-  // |data_file_size| as it does.
+  // Populates |aops|, with data about the update operations, and writes
+  // relevant data to |blob_file|.
   bool GenerateOperations(
       const PayloadGenerationConfig& config,
-      BlobFileWriter* blob_file,
-      std::vector<AnnotatedOperation>* rootfs_ops,
-      std::vector<AnnotatedOperation>* kernel_ops) override;
-
-  // Generates the list of operations to update inplace from the partition
-  // |old_part| to |new_part|. The |partition_size| should be at least
-  // |new_part.size| and any extra space there could be used as scratch space.
-  // The operations generated will not write more than |chunk_blocks| blocks.
-  // The new operations will create blobs in |data_file_fd| and update
-  // the file size pointed by |data_file_size| if needed.
-  // On success, stores the new operations in |aops| and returns true.
-  static bool GenerateOperationsForPartition(
+      const PartitionConfig& old_part,
       const PartitionConfig& new_part,
-      size_t block_size,
-      size_t chunk_blocks,
       BlobFileWriter* blob_file,
-      std::vector<AnnotatedOperation>* aops);
+      std::vector<AnnotatedOperation>* aops) override;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(FullUpdateGenerator);
diff --git a/payload_generator/full_update_generator_unittest.cc b/payload_generator/full_update_generator_unittest.cc
index 4c3c8a4..ac37158 100644
--- a/payload_generator/full_update_generator_unittest.cc
+++ b/payload_generator/full_update_generator_unittest.cc
@@ -39,25 +39,22 @@
     config_.hard_chunk_size = 128 * 1024;
     config_.block_size = 4096;
 
-    EXPECT_TRUE(utils::MakeTempFile("FullUpdateTest_rootfs.XXXXXX",
-                                    &config_.target.rootfs.path,
-                                    nullptr));
-    EXPECT_TRUE(utils::MakeTempFile("FullUpdateTest_kernel.XXXXXX",
-                                    &config_.target.kernel.path,
+    EXPECT_TRUE(utils::MakeTempFile("FullUpdateTest_partition.XXXXXX",
+                                    &new_part_conf.path,
                                     nullptr));
     EXPECT_TRUE(utils::MakeTempFile("FullUpdateTest_blobs.XXXXXX",
                                     &out_blobs_path_,
                                     &out_blobs_fd_));
 
     blob_file_.reset(new BlobFileWriter(out_blobs_fd_, &out_blobs_length_));
-    rootfs_part_unlinker_.reset(
-        new ScopedPathUnlinker(config_.target.rootfs.path));
-    kernel_part_unlinker_.reset(
-        new ScopedPathUnlinker(config_.target.kernel.path));
+    part_path_unlinker_.reset(new ScopedPathUnlinker(new_part_conf.path));
     out_blobs_unlinker_.reset(new ScopedPathUnlinker(out_blobs_path_));
   }
 
   PayloadGenerationConfig config_;
+  PartitionConfig new_part_conf{PartitionName::kRootfs};
+
+  vector<AnnotatedOperation> aops;
 
   // Output file holding the payload blobs.
   string out_blobs_path_;
@@ -66,8 +63,7 @@
   ScopedFdCloser out_blobs_fd_closer_{&out_blobs_fd_};
 
   std::unique_ptr<BlobFileWriter> blob_file_;
-  std::unique_ptr<ScopedPathUnlinker> rootfs_part_unlinker_;
-  std::unique_ptr<ScopedPathUnlinker> kernel_part_unlinker_;
+  std::unique_ptr<ScopedPathUnlinker> part_path_unlinker_;
   std::unique_ptr<ScopedPathUnlinker> out_blobs_unlinker_;
 
   // FullUpdateGenerator under test.
@@ -75,78 +71,72 @@
 };
 
 TEST_F(FullUpdateGeneratorTest, RunTest) {
-  chromeos::Blob new_root(9 * 1024 * 1024);
-  chromeos::Blob new_kern(3 * 1024 * 1024);
-  FillWithData(&new_root);
-  FillWithData(&new_kern);
+  chromeos::Blob new_part(9 * 1024 * 1024);
+  FillWithData(&new_part);
+  new_part_conf.size = new_part.size();
 
-  // Assume hashes take 2 MiB beyond the rootfs.
-  config_.rootfs_partition_size = new_root.size();
-  config_.target.rootfs.size = new_root.size() - 2 * 1024 * 1024;
-  config_.target.kernel.size = new_kern.size();
-
-  EXPECT_TRUE(test_utils::WriteFileVector(config_.target.rootfs.path,
-                                          new_root));
-  EXPECT_TRUE(test_utils::WriteFileVector(config_.target.kernel.path,
-                                          new_kern));
-
-  vector<AnnotatedOperation> rootfs_ops;
-  vector<AnnotatedOperation> kernel_ops;
+  EXPECT_TRUE(test_utils::WriteFileVector(new_part_conf.path, new_part));
 
   EXPECT_TRUE(generator_.GenerateOperations(config_,
+                                            new_part_conf,  // this is ignored
+                                            new_part_conf,
                                             blob_file_.get(),
-                                            &rootfs_ops,
-                                            &kernel_ops));
-  int64_t target_rootfs_chunks =
-      config_.target.rootfs.size / config_.hard_chunk_size;
-  EXPECT_EQ(target_rootfs_chunks, rootfs_ops.size());
-  EXPECT_EQ(new_kern.size() / config_.hard_chunk_size, kernel_ops.size());
-  for (off_t i = 0; i < target_rootfs_chunks; ++i) {
-    EXPECT_EQ(1, rootfs_ops[i].op.dst_extents_size());
+                                            &aops));
+  int64_t new_part_chunks = new_part_conf.size / config_.hard_chunk_size;
+  EXPECT_EQ(new_part_chunks, aops.size());
+  for (off_t i = 0; i < new_part_chunks; ++i) {
+    EXPECT_EQ(1, aops[i].op.dst_extents_size());
     EXPECT_EQ(i * config_.hard_chunk_size / config_.block_size,
-              rootfs_ops[i].op.dst_extents(0).start_block()) << "i = " << i;
+              aops[i].op.dst_extents(0).start_block()) << "i = " << i;
     EXPECT_EQ(config_.hard_chunk_size / config_.block_size,
-              rootfs_ops[i].op.dst_extents(0).num_blocks());
-    if (rootfs_ops[i].op.type() != InstallOperation::REPLACE) {
-      EXPECT_EQ(InstallOperation::REPLACE_BZ, rootfs_ops[i].op.type());
+              aops[i].op.dst_extents(0).num_blocks());
+    if (aops[i].op.type() != InstallOperation::REPLACE) {
+      EXPECT_EQ(InstallOperation::REPLACE_BZ, aops[i].op.type());
     }
   }
 }
 
 // Test that if the chunk size is not a divisor of the image size, it handles
-// correctly the last chunk of each partition.
+// correctly the last chunk of the partition.
 TEST_F(FullUpdateGeneratorTest, ChunkSizeTooBig) {
   config_.hard_chunk_size = 1024 * 1024;
   config_.soft_chunk_size = config_.hard_chunk_size;
-  chromeos::Blob new_root(1536 * 1024);  // 1.5 MiB
-  chromeos::Blob new_kern(128 * 1024);
-  config_.rootfs_partition_size = new_root.size();
-  config_.target.rootfs.size = new_root.size();
-  config_.target.kernel.size = new_kern.size();
+  chromeos::Blob new_part(1536 * 1024);  // 1.5 MiB
+  new_part_conf.size = new_part.size();
 
-  EXPECT_TRUE(test_utils::WriteFileVector(config_.target.rootfs.path,
-                                          new_root));
-  EXPECT_TRUE(test_utils::WriteFileVector(config_.target.kernel.path,
-                                          new_kern));
-
-  vector<AnnotatedOperation> rootfs_ops;
-  vector<AnnotatedOperation> kernel_ops;
+  EXPECT_TRUE(test_utils::WriteFileVector(new_part_conf.path, new_part));
 
   EXPECT_TRUE(generator_.GenerateOperations(config_,
+                                            new_part_conf,  // this is ignored
+                                            new_part_conf,
                                             blob_file_.get(),
-                                            &rootfs_ops,
-                                            &kernel_ops));
-  // rootfs has one chunk and a half.
-  EXPECT_EQ(2, rootfs_ops.size());
+                                            &aops));
+  // new_part has one chunk and a half.
+  EXPECT_EQ(2, aops.size());
   EXPECT_EQ(config_.hard_chunk_size / config_.block_size,
-            BlocksInExtents(rootfs_ops[0].op.dst_extents()));
-  EXPECT_EQ((new_root.size() - config_.hard_chunk_size) / config_.block_size,
-            BlocksInExtents(rootfs_ops[1].op.dst_extents()));
+            BlocksInExtents(aops[0].op.dst_extents()));
+  EXPECT_EQ((new_part.size() - config_.hard_chunk_size) / config_.block_size,
+            BlocksInExtents(aops[1].op.dst_extents()));
+}
 
-  // kernel has less than one chunk.
-  EXPECT_EQ(1, kernel_ops.size());
-  EXPECT_EQ(new_kern.size() / config_.block_size,
-            BlocksInExtents(kernel_ops[0].op.dst_extents()));
+// Test that if the image size is much smaller than the chunk size, it handles
+// correctly the only chunk of the partition.
+TEST_F(FullUpdateGeneratorTest, ImageSizeTooSmall) {
+  chromeos::Blob new_part(16 * 1024);
+  new_part_conf.size = new_part.size();
+
+  EXPECT_TRUE(test_utils::WriteFileVector(new_part_conf.path, new_part));
+
+  EXPECT_TRUE(generator_.GenerateOperations(config_,
+                                            new_part_conf,  // this is ignored
+                                            new_part_conf,
+                                            blob_file_.get(),
+                                            &aops));
+
+  // new_part has less than one chunk.
+  EXPECT_EQ(1, aops.size());
+  EXPECT_EQ(new_part.size() / config_.block_size,
+            BlocksInExtents(aops[0].op.dst_extents()));
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/inplace_generator.cc b/payload_generator/inplace_generator.cc
index 9534230..9d108bd 100644
--- a/payload_generator/inplace_generator.cc
+++ b/payload_generator/inplace_generator.cc
@@ -784,66 +784,38 @@
   return true;
 }
 
-bool InplaceGenerator::GenerateOperationsForPartition(
-    const PartitionConfig& old_part,
-    const PartitionConfig& new_part,
-    uint64_t partition_size,
-    size_t block_size,
-    ssize_t hard_chunk_blocks,
-    size_t soft_chunk_blocks,
-    BlobFileWriter* blob_file,
-    vector<AnnotatedOperation>* aops) {
-  const string part_name = PartitionNameString(new_part.name);
-  LOG(INFO) << "Delta compressing " << part_name << " partition...";
-  TEST_AND_RETURN_FALSE(
-      diff_utils::DeltaReadPartition(aops,
-                                     old_part,
-                                     new_part,
-                                     hard_chunk_blocks,
-                                     soft_chunk_blocks,
-                                     blob_file,
-                                     false));  // src_ops_allowed
-  LOG(INFO) << "Done reading " << part_name;
-
-  TEST_AND_RETURN_FALSE(
-      ResolveReadAfterWriteDependencies(new_part,
-                                        partition_size,
-                                        block_size,
-                                        blob_file,
-                                        aops));
-  LOG(INFO) << "Done reordering " << part_name;
-  return true;
-}
-
 bool InplaceGenerator::GenerateOperations(
     const PayloadGenerationConfig& config,
+    const PartitionConfig& old_part,
+    const PartitionConfig& new_part,
     BlobFileWriter* blob_file,
-    vector<AnnotatedOperation>* rootfs_ops,
-    vector<AnnotatedOperation>* kernel_ops) {
+    vector<AnnotatedOperation>* aops) {
   ssize_t hard_chunk_blocks = (config.hard_chunk_size == -1 ? -1 :
                                config.hard_chunk_size / config.block_size);
   size_t soft_chunk_blocks = config.soft_chunk_size / config.block_size;
+  uint64_t partition_size = new_part.size;
+  if (new_part.name == PartitionName::kRootfs)
+    partition_size = config.rootfs_partition_size;
 
-  TEST_AND_RETURN_FALSE(GenerateOperationsForPartition(
-      config.source.rootfs,
-      config.target.rootfs,
-      config.rootfs_partition_size,
-      config.block_size,
-      hard_chunk_blocks,
-      soft_chunk_blocks,
-      blob_file,
-      rootfs_ops));
+  const string part_name = PartitionNameString(new_part.name);
+  LOG(INFO) << "Delta compressing " << part_name << " partition...";
+  TEST_AND_RETURN_FALSE(
+    diff_utils::DeltaReadPartition(aops,
+                                   old_part,
+                                   new_part,
+                                   hard_chunk_blocks,
+                                   soft_chunk_blocks,
+                                   blob_file,
+                                   false));  // src_ops_allowed
+  LOG(INFO) << "Done reading " << part_name;
 
-  TEST_AND_RETURN_FALSE(GenerateOperationsForPartition(
-      config.source.kernel,
-      config.target.kernel,
-      config.target.kernel.size,  // kernel "filesystem" is the whole partition.
-      config.block_size,
-      hard_chunk_blocks,
-      soft_chunk_blocks,
-      blob_file,
-      kernel_ops));
-
+  TEST_AND_RETURN_FALSE(
+    ResolveReadAfterWriteDependencies(new_part,
+                                      partition_size,
+                                      config.block_size,
+                                      blob_file,
+                                      aops));
+  LOG(INFO) << "Done reordering " << part_name;
   return true;
 }
 
diff --git a/payload_generator/inplace_generator.h b/payload_generator/inplace_generator.h
index bca983c..4839824 100644
--- a/payload_generator/inplace_generator.h
+++ b/payload_generator/inplace_generator.h
@@ -218,39 +218,20 @@
       BlobFileWriter* blob_file,
       std::vector<AnnotatedOperation>* aops);
 
-  // Generates the list of operations to update inplace from the partition
-  // |old_part| to |new_part|. The |partition_size| should be at least
-  // |new_part.size| and any extra space there could be used as scratch space.
-  // The operations generated will not write more than |chunk_blocks| blocks.
-  // The new operations will create blobs in |data_file_fd| and update
-  // the file size pointed by |data_file_size| if needed.
-  // On success, stores the new operations in |aops| and returns true.
-  static bool GenerateOperationsForPartition(
-      const PartitionConfig& old_part,
-      const PartitionConfig& new_part,
-      uint64_t partition_size,
-      size_t block_size,
-      ssize_t hard_chunk_blocks,
-      size_t soft_chunk_blocks,
-      BlobFileWriter* blob_file,
-      std::vector<AnnotatedOperation>* aops);
-
-  // Generate the update payload operations for the kernel and rootfs using
+  // Generate the update payload operations for the given partition using
   // only operations that read from the target and/or write to the target,
   // hence, applying the payload "in-place" in the target partition. This method
   // assumes that the contents of the source image are pre-copied to the target
   // partition, up to the size of the source image. Use this method to generate
   // a delta update with the minor version kInPlaceMinorPayloadVersion.
-  // The rootfs operations are stored in |graph| and should be executed in the
-  // |final_order| order. The kernel operations are stored in |kernel_ops|. All
-  // the offsets in the operations reference the data written to |data_file_fd|.
-  // The total amount of data written to that file is stored in
-  // |data_file_size|.
+  // The operations are stored in |aops|. All the offsets in the operations
+  // reference the data written to |blob_file|.
   bool GenerateOperations(
       const PayloadGenerationConfig& config,
+      const PartitionConfig& old_part,
+      const PartitionConfig& new_part,
       BlobFileWriter* blob_file,
-      std::vector<AnnotatedOperation>* rootfs_ops,
-      std::vector<AnnotatedOperation>* kernel_ops) override;
+      std::vector<AnnotatedOperation>* aops) override;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(InplaceGenerator);
diff --git a/payload_generator/operations_generator.h b/payload_generator/operations_generator.h
index 7d34fd0..9127d7b 100644
--- a/payload_generator/operations_generator.h
+++ b/payload_generator/operations_generator.h
@@ -31,22 +31,21 @@
  public:
   virtual ~OperationsGenerator() = default;
 
-  // This method generates two lists of operations, one for the rootfs and one
-  // for the kernel and stores the generated operations in |rootfs_ops| and
-  // |kernel_ops| respectivelly. These operations are generated based on the
-  // given |config|. The operations should be applied in the order specified in
-  // the list, and they respect the payload version and type (delta or full)
-  // specified in |config|.
-  // The operations generated will refer to offsets in the file |data_file_fd|,
+  // This method generates a list of operations to update from the partition
+  // |old_part| to |new_part| and stores the generated operations in |aops|.
+  // These operations are generated based on the given |config|.
+  // The operations should be applied in the order specified in the list, and
+  // they respect the payload version and type (delta or full) specified in
+  // |config|.
+  // The operations generated will refer to offsets in the file |blob_file|,
   // where this function stores the output, but not necessarily in the same
-  // order as they appear in the |rootfs_ops| and |kernel_ops|.
-  // This function stores the amount of data written to |data_file_fd| in
-  // |data_file_size|.
+  // order as they appear in the |aops|.
   virtual bool GenerateOperations(
       const PayloadGenerationConfig& config,
+      const PartitionConfig& old_part,
+      const PartitionConfig& new_part,
       BlobFileWriter* blob_file,
-      std::vector<AnnotatedOperation>* rootfs_ops,
-      std::vector<AnnotatedOperation>* kernel_ops) = 0;
+      std::vector<AnnotatedOperation>* aops) = 0;
 
  protected:
   OperationsGenerator() = default;