diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index 2f76796..1592917 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -149,10 +149,12 @@
   diff_utils::FilterNoopOperations(&kernel_ops);
 
   // Write payload file to disk.
-  TEST_AND_RETURN_FALSE(payload.AddPartitionOperations(
-      kLegacyPartitionNameRoot, rootfs_ops));
-  TEST_AND_RETURN_FALSE(payload.AddPartitionOperations(
-      kLegacyPartitionNameKernel, kernel_ops));
+  TEST_AND_RETURN_FALSE(payload.AddPartition(config.source.rootfs,
+                                             config.target.rootfs,
+                                             rootfs_ops));
+  TEST_AND_RETURN_FALSE(payload.AddPartition(config.source.kernel,
+                                             config.target.kernel,
+                                             kernel_ops));
   TEST_AND_RETURN_FALSE(payload.WritePayload(output_path, temp_file_path,
                                              private_key_path, metadata_size));
   temp_file_unlinker.reset();
diff --git a/payload_generator/payload_file.cc b/payload_generator/payload_file.cc
index 2e0b22c..6888aaf 100644
--- a/payload_generator/payload_file.cc
+++ b/payload_generator/payload_file.cc
@@ -81,38 +81,29 @@
     *(manifest_.mutable_new_image_info()) = config.target.image_info;
 
   manifest_.set_block_size(config.block_size);
-
-  // Initialize the PartitionInfo objects if present.
-  if (!config.source.kernel.path.empty()) {
-    TEST_AND_RETURN_FALSE(diff_utils::InitializePartitionInfo(
-        config.source.kernel,
-        manifest_.mutable_old_kernel_info()));
-  }
-  TEST_AND_RETURN_FALSE(diff_utils::InitializePartitionInfo(
-      config.target.kernel,
-      manifest_.mutable_new_kernel_info()));
-  if (!config.source.rootfs.path.empty()) {
-    TEST_AND_RETURN_FALSE(diff_utils::InitializePartitionInfo(
-        config.source.rootfs,
-        manifest_.mutable_old_rootfs_info()));
-  }
-  TEST_AND_RETURN_FALSE(diff_utils::InitializePartitionInfo(
-      config.target.rootfs,
-      manifest_.mutable_new_rootfs_info()));
   return true;
 }
 
-bool PayloadFile::AddPartitionOperations(
-    const std::string& partition_name,
-    const vector<AnnotatedOperation>& aops) {
+bool PayloadFile::AddPartition(const PartitionConfig& old_conf,
+                               const PartitionConfig& new_conf,
+                               const vector<AnnotatedOperation>& aops) {
   // Check partitions order for Chrome OS
   if (major_version_ == kChromeOSMajorPayloadVersion) {
     const vector<const char*> part_order = { kLegacyPartitionNameRoot,
                                              kLegacyPartitionNameKernel };
-    TEST_AND_RETURN_FALSE(aops_vec_.size() < part_order.size());
-    TEST_AND_RETURN_FALSE(partition_name == part_order[aops_vec_.size()]);
+    TEST_AND_RETURN_FALSE(part_vec_.size() < part_order.size());
+    TEST_AND_RETURN_FALSE(new_conf.name == part_order[part_vec_.size()]);
   }
-  aops_vec_.emplace_back(partition_name, aops);
+  Partition part;
+  part.name = new_conf.name;
+  part.aops = aops;
+  // Initialize the PartitionInfo objects if present.
+  if (!old_conf.path.empty())
+    TEST_AND_RETURN_FALSE(diff_utils::InitializePartitionInfo(old_conf,
+                                                              &part.old_info));
+  TEST_AND_RETURN_FALSE(diff_utils::InitializePartitionInfo(new_conf,
+                                                            &part.new_info));
+  part_vec_.push_back(std::move(part));
   return true;
 }
 
@@ -131,8 +122,8 @@
 
   // Check that install op blobs are in order.
   uint64_t next_blob_offset = 0;
-  for (const auto& part_name_aops_pair : aops_vec_) {
-    for (const auto& aop : part_name_aops_pair.second) {
+  for (const auto& part : part_vec_) {
+    for (const auto& aop : part.aops) {
       if (!aop.op.has_data_offset())
         continue;
       if (aop.op.data_offset() != next_blob_offset) {
@@ -143,24 +134,37 @@
     }
   }
 
-  // Copy the operations from the aops_vec_ to the manifest.
+  // Copy the operations and partition info from the part_vec_ to the manifest.
   manifest_.clear_install_operations();
   manifest_.clear_kernel_install_operations();
   manifest_.clear_partitions();
-  for (const auto& part_name_aops_pair : aops_vec_) {
+  for (const auto& part : part_vec_) {
     if (major_version_ == kBrilloMajorPayloadVersion) {
       PartitionUpdate* partition = manifest_.add_partitions();
-      partition->set_partition_name(part_name_aops_pair.first);
-      for (const AnnotatedOperation& aop : part_name_aops_pair.second) {
+      partition->set_partition_name(part.name);
+      for (const AnnotatedOperation& aop : part.aops) {
         *partition->add_operations() = aop.op;
       }
+      if (part.old_info.has_size() || part.old_info.has_hash())
+        *(partition->mutable_old_partition_info()) = part.old_info;
+      if (part.new_info.has_size() || part.new_info.has_hash())
+        *(partition->mutable_new_partition_info()) = part.new_info;
     } else {
-      for (const AnnotatedOperation& aop : part_name_aops_pair.second) {
-        if (part_name_aops_pair.first == kLegacyPartitionNameKernel) {
+      // major_version_ == kChromeOSMajorPayloadVersion
+      if (part.name == kLegacyPartitionNameKernel) {
+        for (const AnnotatedOperation& aop : part.aops)
           *manifest_.add_kernel_install_operations() = aop.op;
-        } else {
+        if (part.old_info.has_size() || part.old_info.has_hash())
+          *manifest_.mutable_old_kernel_info() = part.old_info;
+        if (part.new_info.has_size() || part.new_info.has_hash())
+          *manifest_.mutable_new_kernel_info() = part.new_info;
+      } else {
+        for (const AnnotatedOperation& aop : part.aops)
           *manifest_.add_install_operations() = aop.op;
-        }
+        if (part.old_info.has_size() || part.old_info.has_hash())
+          *manifest_.mutable_old_rootfs_info() = part.old_info;
+        if (part.new_info.has_size() || part.new_info.has_hash())
+          *manifest_.mutable_new_rootfs_info() = part.new_info;
       }
     }
   }
@@ -251,8 +255,8 @@
   ScopedFileWriterCloser writer_closer(&writer);
   uint64_t out_file_size = 0;
 
-  for (auto& part_name_aops_pair : aops_vec_) {
-    for (AnnotatedOperation& aop : part_name_aops_pair.second) {
+  for (auto& part: part_vec_) {
+    for (AnnotatedOperation& aop : part.aops) {
       if (!aop.op.has_data_offset())
         continue;
       CHECK(aop.op.has_data_length());
@@ -285,8 +289,8 @@
   vector<DeltaObject> objects;
   off_t total_size = 0;
 
-  for (const auto& part_name_aops_pair : aops_vec_) {
-    for (const AnnotatedOperation& aop : part_name_aops_pair.second) {
+  for (const auto& part : part_vec_) {
+    for (const AnnotatedOperation& aop : part.aops) {
       objects.push_back(DeltaObject(aop.name,
                                     aop.op.type(),
                                     aop.op.data_length()));
diff --git a/payload_generator/payload_file.h b/payload_generator/payload_file.h
index 7cf9160..d877ef7 100644
--- a/payload_generator/payload_file.h
+++ b/payload_generator/payload_file.h
@@ -19,7 +19,6 @@
 
 #include <string>
 #include <vector>
-#include <utility>
 
 #include <chromeos/secure_blob.h>
 #include <gtest/gtest_prod.h>  // for FRIEND_TEST
@@ -39,10 +38,12 @@
   // required hashes of the requested partitions.
   bool Init(const PayloadGenerationConfig& config);
 
-  // Sets the list of operations to the payload manifest. The operations
+  // Add a partition to the payload manifest. Including partition name, list of
+  // operations and partition info. The operations in |aops|
   // reference a blob stored in the file provided to WritePayload().
-  bool AddPartitionOperations(const std::string& partition_name,
-                              const std::vector<AnnotatedOperation>& aops);
+  bool AddPartition(const PartitionConfig& old_conf,
+                    const PartitionConfig& new_conf,
+                    const std::vector<AnnotatedOperation>& aops);
 
   // Write the payload to the |payload_file| file. The operations reference
   // blobs in the |data_blobs_path| file and the blobs will be reordered in the
@@ -82,8 +83,19 @@
 
   DeltaArchiveManifest manifest_;
 
-  std::vector<std::pair<std::string,
-                        std::vector<AnnotatedOperation>>> aops_vec_;
+  // Struct has necessary information to write PartitionUpdate in protobuf.
+  struct Partition {
+    // The name of the partition.
+    std::string name;
+
+    // The operations to be performed to this partition.
+    std::vector<AnnotatedOperation> aops;
+
+    PartitionInfo old_info;
+    PartitionInfo new_info;
+  };
+
+  std::vector<Partition> part_vec_;
 };
 
 // Adds a dummy operation that points to a signature blob located at the
diff --git a/payload_generator/payload_file_unittest.cc b/payload_generator/payload_file_unittest.cc
index ff8c8a4..244d22b 100644
--- a/payload_generator/payload_file_unittest.cc
+++ b/payload_generator/payload_file_unittest.cc
@@ -54,6 +54,8 @@
       utils::MakeTempFile("ReorderBlobsTest.new.XXXXXX", &new_blobs, nullptr));
   ScopedPathUnlinker new_blobs_unlinker(new_blobs);
 
+  payload_.part_vec_.resize(2);
+
   vector<AnnotatedOperation> aops;
   AnnotatedOperation aop;
   aop.op.set_data_offset(8);
@@ -63,17 +65,16 @@
   aop.op.set_data_offset(7);
   aop.op.set_data_length(1);
   aops.push_back(aop);
-  payload_.AddPartitionOperations("part0", aops);
+  payload_.part_vec_[0].aops = aops;
 
   aop.op.set_data_offset(0);
   aop.op.set_data_length(6);
-  aops = {aop};
-  payload_.AddPartitionOperations("part1", aops);
+  payload_.part_vec_[1].aops = {aop};
 
   EXPECT_TRUE(payload_.ReorderDataBlobs(orig_blobs, new_blobs));
 
-  const vector<AnnotatedOperation>& part0_aops = payload_.aops_vec_[0].second;
-  const vector<AnnotatedOperation>& part1_aops = payload_.aops_vec_[1].second;
+  const vector<AnnotatedOperation>& part0_aops = payload_.part_vec_[0].aops;
+  const vector<AnnotatedOperation>& part1_aops = payload_.part_vec_[1].aops;
   string new_data;
   EXPECT_TRUE(utils::ReadFile(new_blobs, &new_data));
   // Kernel blobs should appear at the end.