cc: Fix LayerTree deserialization for commit.

When we receive a commit protobuf from the engine, we collect all the
existing Layers in a map and proceed with deserializing the new LayerTree
in the commit protobuf, re-using any existing Layers if possible.

When deserializing a Layer, Layer::FromLayerNodeProto creates a copy of
the existing children, in order to clear the parent pointers for these
layers, if they have now been added to the children of another layer.
However, since we only traverse the layers from the updated LayerTree
during deserialization, if a Layer has been destroyed on the engine, it will
not be traversed during the derserialization step, and will still hold a
reference to all its children. When this Layer is being destroyed on the
client, it will call Layer::RemoveFromParent for all its children, which
may now have been added to the children list of a different Layer.

This patch clears all the properties that are sent with a Layer Hierarchy
update in Layer::ClearLayerHierarchyPropertiesForDeserializationAndAddToMap,
which is run for the existing LayerTree before parsing the new LayerTree from
the commit protobuf.

BUG=590102
CQ_INCLUDE_TRYBOTS=tryserver.blink:linux_blink_rel

Review URL: https://codereview.chromium.org/1804703004

Cr-Commit-Position: refs/heads/master@{#382953}

8 files changed, 211 insertions, 115 deletions

diff --git a/cc/layers/layer.cc b/cc/layers/layer.cc
index 0567738..1bd263d 100644
--- a/cc/layers/layer.cc
+++ b/cc/layers/layer.cc
@@ -1309,16 +1309,51 @@ void Layer::ToLayerNodeProto(proto::LayerNode* proto) const {
     replica_layer_->ToLayerNodeProto(proto->mutable_replica_layer());
 }
 
+void Layer::ClearLayerTreePropertiesForDeserializationAndAddToMap(
+    LayerIdMap* layer_map) {
+  (*layer_map)[layer_id_] = this;
+
+  if (layer_tree_host_)
+    layer_tree_host_->UnregisterLayer(this);
+
+  layer_tree_host_ = nullptr;
+  parent_ = nullptr;
+
+  // Clear these properties for all the children and add them to the map.
+  for (auto& child : children_) {
+    child->ClearLayerTreePropertiesForDeserializationAndAddToMap(layer_map);
+  }
+
+  children_.clear();
+
+  if (mask_layer_) {
+    mask_layer_->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        layer_map);
+    mask_layer_ = nullptr;
+  }
+
+  if (replica_layer_) {
+    replica_layer_->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        layer_map);
+    replica_layer_ = nullptr;
+  }
+}
+
 void Layer::FromLayerNodeProto(const proto::LayerNode& proto,
-                               const LayerIdMap& layer_map) {
+                               const LayerIdMap& layer_map,
+                               LayerTreeHost* layer_tree_host) {
+  DCHECK(!layer_tree_host_);
+  DCHECK(children_.empty());
+  DCHECK(!mask_layer_);
+  DCHECK(!replica_layer_);
+  DCHECK(layer_tree_host);
   DCHECK(proto.has_id());
+
   layer_id_ = proto.id();
 
-  // To deserialize the new children, make a copy of the old list, before
-  // inserting each of the new children in a new list, reusing old Layer objects
-  // if the Layer already exists.
-  LayerList old_children = children_;
-  children_.clear();
+  layer_tree_host_ = layer_tree_host;
+  layer_tree_host_->RegisterLayer(this);
+
   for (int i = 0; i < proto.children_size(); ++i) {
     const proto::LayerNode& child_proto = proto.children(i);
     DCHECK(child_proto.has_type());
@@ -1326,51 +1361,26 @@ void Layer::FromLayerNodeProto(const proto::LayerNode& proto,
         LayerProtoConverter::FindOrAllocateAndConstruct(child_proto, layer_map);
     // The child must now refer to this layer as its parent, and must also have
     // the same LayerTreeHost. This must be done before deserializing children.
+    DCHECK(!child->parent_);
     child->parent_ = this;
-    child->layer_tree_host_ = layer_tree_host_;
-    child->FromLayerNodeProto(child_proto, layer_map);
+    child->FromLayerNodeProto(child_proto, layer_map, layer_tree_host_);
     children_.push_back(child);
   }
 
-  // Remove now-unused children from the tree.
-  for (auto& child : old_children) {
-    // A child might have been moved to a different parent.
-    if (child->parent_ != this)
-      continue;
-    // Our own child is not part of our new children, so remove it.
-    if (std::find(children_.begin(), children_.end(), child) ==
-        children_.end()) {
-      child->parent_ = nullptr;
-      child->layer_tree_host_ = nullptr;
-    }
-  }
-
-  if (mask_layer_) {
-    mask_layer_->parent_ = nullptr;
-    mask_layer_->layer_tree_host_ = nullptr;
-  }
   if (proto.has_mask_layer()) {
     mask_layer_ = LayerProtoConverter::FindOrAllocateAndConstruct(
         proto.mask_layer(), layer_map);
     mask_layer_->parent_ = this;
-    mask_layer_->layer_tree_host_ = layer_tree_host_;
-    mask_layer_->FromLayerNodeProto(proto.mask_layer(), layer_map);
-  } else {
-    mask_layer_ = nullptr;
+    mask_layer_->FromLayerNodeProto(proto.mask_layer(), layer_map,
+                                    layer_tree_host_);
   }
 
-  if (replica_layer_) {
-    replica_layer_->parent_ = nullptr;
-    replica_layer_->layer_tree_host_ = nullptr;
-  }
   if (proto.has_replica_layer()) {
     replica_layer_ = LayerProtoConverter::FindOrAllocateAndConstruct(
         proto.replica_layer(), layer_map);
     replica_layer_->parent_ = this;
-    replica_layer_->layer_tree_host_ = layer_tree_host_;
-    replica_layer_->FromLayerNodeProto(proto.replica_layer(), layer_map);
-  } else {
-    replica_layer_ = nullptr;
+    replica_layer_->FromLayerNodeProto(proto.replica_layer(), layer_map,
+                                       layer_tree_host_);
   }
 }
 
diff --git a/cc/layers/layer.h b/cc/layers/layer.h
index 409a556..aa0d9a9 100644
--- a/cc/layers/layer.h
+++ b/cc/layers/layer.h
@@ -375,12 +375,20 @@ class CC_EXPORT Layer : public base::RefCounted<Layer> {
   // construction of the correct layer on the client.
   void ToLayerNodeProto(proto::LayerNode* proto) const;
 
+  // Recursively iterate over this layer and all children and reset the
+  // properties sent with the hierarchical structure in the LayerNode protos.
+  // This must be done before deserializing the new LayerTree from the Layernode
+  // protos.
+  void ClearLayerTreePropertiesForDeserializationAndAddToMap(
+      LayerIdMap* layer_map);
+
   // Recursively iterate over the given LayerNode proto and read the structure
   // into this node and its children. The |layer_map| should be used to look
   // for previously existing Layers, since they should be re-used between each
   // hierarchy update.
   void FromLayerNodeProto(const proto::LayerNode& proto,
-                          const LayerIdMap& layer_map);
+                          const LayerIdMap& layer_map,
+                          LayerTreeHost* layer_tree_host);
 
   // This method is similar to PushPropertiesTo, but instead of pushing to
   // a LayerImpl, it pushes the properties to proto::LayerProperties. It is
diff --git a/cc/layers/layer_proto_converter.cc b/cc/layers/layer_proto_converter.cc
index 626380d..f4c097f 100644
--- a/cc/layers/layer_proto_converter.cc
+++ b/cc/layers/layer_proto_converter.cc
@@ -30,10 +30,13 @@ void LayerProtoConverter::SerializeLayerHierarchy(
 // static
 scoped_refptr<Layer> LayerProtoConverter::DeserializeLayerHierarchy(
     scoped_refptr<Layer> existing_root,
-    const proto::LayerNode& root_node) {
+    const proto::LayerNode& root_node,
+    LayerTreeHost* layer_tree_host) {
   LayerIdMap layer_id_map;
-  if (existing_root)
-    RecursivelyFindAllLayers(existing_root.get(), &layer_id_map);
+  if (existing_root) {
+    existing_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &layer_id_map);
+  }
 
   scoped_refptr<Layer> new_root = existing_root;
   if (!existing_root ||
@@ -42,7 +45,7 @@ scoped_refptr<Layer> LayerProtoConverter::DeserializeLayerHierarchy(
     // so find or create the new root.
     new_root = FindOrAllocateAndConstruct(root_node, layer_id_map);
   }
-  new_root->FromLayerNodeProto(root_node, layer_id_map);
+  new_root->FromLayerNodeProto(root_node, layer_id_map, layer_tree_host);
   return new_root;
 }
 
diff --git a/cc/layers/layer_proto_converter.h b/cc/layers/layer_proto_converter.h
index 92a4b54..b86ae91 100644
--- a/cc/layers/layer_proto_converter.h
+++ b/cc/layers/layer_proto_converter.h
@@ -33,7 +33,8 @@ class CC_EXPORT LayerProtoConverter {
   // the first deserialize.
   static scoped_refptr<Layer> DeserializeLayerHierarchy(
       const scoped_refptr<Layer> existing_root,
-      const proto::LayerNode& root_node);
+      const proto::LayerNode& root_node,
+      LayerTreeHost* layer_tree_host);
 
   // Serializes the properties of all the dirty nodes in the Layer hierarchy.
   // The proto::LayerUpdate will contain all nodes that are dirty. These nodes
diff --git a/cc/layers/layer_proto_converter_unittest.cc b/cc/layers/layer_proto_converter_unittest.cc
index 2de7f7d..33d12ac 100644
--- a/cc/layers/layer_proto_converter_unittest.cc
+++ b/cc/layers/layer_proto_converter_unittest.cc
@@ -68,7 +68,8 @@ TEST_F(LayerProtoConverterTest, TestKeepingRoot) {
   child_c_node->set_parent_id(child_b_node->id());
 
   scoped_refptr<Layer> new_root =
-      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node);
+      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node,
+                                                     layer_tree_host_.get());
 
   // The new root should not be the same as the old root.
   EXPECT_EQ(old_root->id(), new_root->id());
@@ -84,6 +85,8 @@ TEST_F(LayerProtoConverterTest, TestKeepingRoot) {
 
   scoped_refptr<Layer> child_c = child_b->children()[0];
   EXPECT_EQ(child_c_node->id(), child_c->id());
+
+  new_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerProtoConverterTest, TestNoExistingRoot) {
@@ -104,7 +107,8 @@ TEST_F(LayerProtoConverterTest, TestNoExistingRoot) {
   child_a_node->set_parent_id(new_root_id);  // root_node
 
   scoped_refptr<Layer> new_root =
-      LayerProtoConverter::DeserializeLayerHierarchy(nullptr, root_node);
+      LayerProtoConverter::DeserializeLayerHierarchy(nullptr, root_node,
+                                                     layer_tree_host_.get());
 
   // The new root should not be the same as the old root.
   EXPECT_EQ(new_root_id, new_root->id());
@@ -113,6 +117,8 @@ TEST_F(LayerProtoConverterTest, TestNoExistingRoot) {
 
   EXPECT_EQ(child_a_node->id(), child_a->id());
   EXPECT_EQ(0u, child_a->children().size());
+
+  new_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerProtoConverterTest, TestSwappingRoot) {
@@ -145,7 +151,8 @@ TEST_F(LayerProtoConverterTest, TestSwappingRoot) {
 
   scoped_refptr<Layer> old_root = Layer::Create();
   scoped_refptr<Layer> new_root =
-      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node);
+      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node,
+                                                     layer_tree_host_.get());
 
   // The new root should not be the same as the old root.
   EXPECT_EQ(root_node.id(), new_root->id());
@@ -161,6 +168,8 @@ TEST_F(LayerProtoConverterTest, TestSwappingRoot) {
 
   scoped_refptr<Layer> child_c = child_b->children()[0];
   EXPECT_EQ(child_c_node->id(), child_c->id());
+
+  new_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerProtoConverterTest, RecursivePropertiesSerialization) {
@@ -343,7 +352,8 @@ TEST_F(LayerProtoConverterTest, PictureLayerTypeDeserialization) {
   root_node.set_type(proto::LayerNode::PICTURE_LAYER);
 
   scoped_refptr<Layer> new_root =
-      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node);
+      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node,
+                                                     layer_tree_host_.get());
 
   // Validate that the ids are equal.
   EXPECT_EQ(old_root->id(), new_root->id());
@@ -353,6 +363,8 @@ TEST_F(LayerProtoConverterTest, PictureLayerTypeDeserialization) {
   proto::LayerNode layer_node;
   new_root->SetTypeForProtoSerialization(&layer_node);
   EXPECT_EQ(proto::LayerNode::PICTURE_LAYER, layer_node.type());
+
+  new_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerProtoConverterTest, HudLayerTypeSerialization) {
@@ -374,7 +386,8 @@ TEST_F(LayerProtoConverterTest, HudLayerTypeDeserialization) {
   root_node.set_type(proto::LayerNode::HEADS_UP_DISPLAY_LAYER);
 
   scoped_refptr<Layer> new_root =
-      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node);
+      LayerProtoConverter::DeserializeLayerHierarchy(old_root, root_node,
+                                                     layer_tree_host_.get());
 
   // Validate that the ids are equal.
   EXPECT_EQ(old_root->id(), new_root->id());
@@ -384,6 +397,8 @@ TEST_F(LayerProtoConverterTest, HudLayerTypeDeserialization) {
   proto::LayerNode layer_node;
   new_root->SetTypeForProtoSerialization(&layer_node);
   EXPECT_EQ(proto::LayerNode::HEADS_UP_DISPLAY_LAYER, layer_node.type());
+
+  new_root->SetLayerTreeHost(nullptr);
 }
 
 }  // namespace
diff --git a/cc/layers/layer_unittest.cc b/cc/layers/layer_unittest.cc
index da644f5..752757b 100644
--- a/cc/layers/layer_unittest.cc
+++ b/cc/layers/layer_unittest.cc
@@ -393,11 +393,8 @@ class LayerSerializationTest : public testing::Test {
     Layer::LayerIdMap empty_dest_layer_map;
     scoped_refptr<Layer> layer_dest_root = Layer::Create();
 
-    // Forcefully set the layer tree host for the root layer, which should cause
-    // it to propagate to all the children.
-    layer_dest_root->layer_tree_host_ = layer_tree_host_.get();
-
-    layer_dest_root->FromLayerNodeProto(proto, empty_dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(proto, empty_dest_layer_map,
+                                        layer_tree_host_.get());
 
     EXPECT_EQ(layer_src_root->id(), layer_dest_root->id());
     EXPECT_EQ(nullptr, layer_dest_root->parent());
@@ -451,15 +448,19 @@ class LayerSerializationTest : public testing::Test {
     layer_root->ToLayerNodeProto(&root_proto);
 
     Layer::LayerIdMap dest_layer_map;
-    dest_layer_map[layer_root->id()] = layer_root;
-    dest_layer_map[layer_src_a->id()] = layer_src_a;
-    layer_root->FromLayerNodeProto(root_proto, dest_layer_map);
+    layer_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &dest_layer_map);
+    layer_root->FromLayerNodeProto(root_proto, dest_layer_map,
+                                   layer_tree_host_.get());
 
     EXPECT_EQ(33, layer_root->transform_tree_index_);
     ASSERT_EQ(1u, layer_root->children().size());
     scoped_refptr<Layer> layer_dest_a = layer_root->children()[0];
     EXPECT_EQ(layer_src_a, layer_dest_a);
     EXPECT_EQ(42, layer_dest_a->transform_tree_index_);
+
+    // Clear the reference to the LTH for all the layers.
+    layer_root->SetLayerTreeHost(nullptr);
   }
 
   void RunNonDestructiveDeserializationReorderChildrenTest() {
@@ -488,7 +489,8 @@ class LayerSerializationTest : public testing::Test {
     layer_src_root->ToLayerNodeProto(&root_proto_1);
     Layer::LayerIdMap dest_layer_map;
     scoped_refptr<Layer> layer_dest_root = Layer::Create();
-    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure initial copy is correct.
     ASSERT_EQ(2u, layer_dest_root->children().size());
@@ -510,10 +512,10 @@ class LayerSerializationTest : public testing::Test {
     // Now serialize and deserialize again.
     proto::LayerNode root_proto_2;
     layer_src_root->ToLayerNodeProto(&root_proto_2);
-    dest_layer_map[layer_dest_root->id()] = layer_dest_root;
-    dest_layer_map[layer_dest_a->id()] = layer_dest_a;
-    dest_layer_map[layer_dest_b->id()] = layer_dest_b;
-    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map);
+    layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure second copy is correct.
     EXPECT_EQ(33, layer_dest_root->transform_tree_index_);
@@ -524,6 +526,8 @@ class LayerSerializationTest : public testing::Test {
     layer_dest_a = layer_dest_root->children()[1];
     EXPECT_EQ(layer_src_a->id(), layer_dest_a->id());
     EXPECT_EQ(42, layer_dest_a->transform_tree_index_);
+
+    layer_dest_root->SetLayerTreeHost(nullptr);
   }
 
   void RunNonDestructiveDeserializationAddChildTest() {
@@ -550,7 +554,8 @@ class LayerSerializationTest : public testing::Test {
     layer_src_root->ToLayerNodeProto(&root_proto_1);
     Layer::LayerIdMap dest_layer_map;
     scoped_refptr<Layer> layer_dest_root = Layer::Create();
-    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure initial copy is correct.
     ASSERT_EQ(1u, layer_dest_root->children().size());
@@ -568,9 +573,10 @@ class LayerSerializationTest : public testing::Test {
     // Now serialize and deserialize again.
     proto::LayerNode root_proto_2;
     layer_src_root->ToLayerNodeProto(&root_proto_2);
-    dest_layer_map[layer_dest_root->id()] = layer_dest_root;
-    dest_layer_map[layer_dest_a->id()] = layer_dest_a;
-    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map);
+    layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure second copy is correct.
     EXPECT_EQ(33, layer_dest_root->transform_tree_index_);
@@ -580,6 +586,8 @@ class LayerSerializationTest : public testing::Test {
     EXPECT_EQ(42, layer_dest_a->transform_tree_index_);
     scoped_refptr<Layer> layer_dest_b = layer_dest_root->children()[1];
     EXPECT_EQ(layer_src_b->id(), layer_dest_b->id());
+
+    layer_dest_root->SetLayerTreeHost(nullptr);
   }
 
   void RunNonDestructiveDeserializationRemoveChildTest() {
@@ -608,7 +616,8 @@ class LayerSerializationTest : public testing::Test {
     layer_src_root->ToLayerNodeProto(&root_proto_1);
     Layer::LayerIdMap dest_layer_map;
     scoped_refptr<Layer> layer_dest_root = Layer::Create();
-    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure initial copy is correct.
     ASSERT_EQ(2u, layer_dest_root->children().size());
@@ -628,10 +637,10 @@ class LayerSerializationTest : public testing::Test {
     // Now serialize and deserialize again.
     proto::LayerNode root_proto_2;
     layer_src_root->ToLayerNodeProto(&root_proto_2);
-    dest_layer_map[layer_dest_root->id()] = layer_dest_root;
-    dest_layer_map[layer_dest_a->id()] = layer_dest_a;
-    dest_layer_map[layer_dest_b->id()] = layer_dest_b;
-    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map);
+    layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure second copy is correct.
     EXPECT_EQ(33, layer_dest_root->transform_tree_index_);
@@ -639,6 +648,8 @@ class LayerSerializationTest : public testing::Test {
     layer_dest_a = layer_dest_root->children()[0];
     EXPECT_EQ(layer_src_a->id(), layer_dest_a->id());
     EXPECT_EQ(42, layer_dest_a->transform_tree_index_);
+
+    layer_dest_root->SetLayerTreeHost(nullptr);
   }
 
   void RunNonDestructiveDeserializationMoveChildEarlierTest() {
@@ -673,7 +684,8 @@ class LayerSerializationTest : public testing::Test {
     layer_src_root->ToLayerNodeProto(&root_proto_1);
     Layer::LayerIdMap dest_layer_map;
     scoped_refptr<Layer> layer_dest_root = Layer::Create();
-    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure initial copy is correct.
     ASSERT_EQ(2u, layer_dest_root->children().size());
@@ -699,11 +711,10 @@ class LayerSerializationTest : public testing::Test {
     // Now serialize and deserialize again.
     proto::LayerNode root_proto_2;
     layer_src_root->ToLayerNodeProto(&root_proto_2);
-    dest_layer_map[layer_dest_root->id()] = layer_dest_root;
-    dest_layer_map[layer_dest_a->id()] = layer_dest_a;
-    dest_layer_map[layer_dest_b->id()] = layer_dest_b;
-    dest_layer_map[layer_dest_c->id()] = layer_dest_c;
-    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map);
+    layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure second copy is correct.
     EXPECT_EQ(33, layer_dest_root->transform_tree_index_);
@@ -718,6 +729,8 @@ class LayerSerializationTest : public testing::Test {
     layer_dest_c = layer_dest_a->children()[0];
     EXPECT_EQ(layer_src_c->id(), layer_dest_c->id());
     EXPECT_EQ(99, layer_dest_c->transform_tree_index_);
+
+    layer_dest_root->SetLayerTreeHost(nullptr);
   }
 
   void RunNonDestructiveDeserializationMoveChildLaterTest() {
@@ -752,7 +765,8 @@ class LayerSerializationTest : public testing::Test {
     layer_src_root->ToLayerNodeProto(&root_proto_1);
     Layer::LayerIdMap dest_layer_map;
     scoped_refptr<Layer> layer_dest_root = Layer::Create();
-    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure initial copy is correct.
     ASSERT_EQ(2u, layer_dest_root->children().size());
@@ -778,11 +792,10 @@ class LayerSerializationTest : public testing::Test {
     // Now serialize and deserialize again.
     proto::LayerNode root_proto_2;
     layer_src_root->ToLayerNodeProto(&root_proto_2);
-    dest_layer_map[layer_dest_root->id()] = layer_dest_root;
-    dest_layer_map[layer_dest_a->id()] = layer_dest_a;
-    dest_layer_map[layer_dest_b->id()] = layer_dest_b;
-    dest_layer_map[layer_dest_c->id()] = layer_dest_c;
-    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map);
+    layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+        &dest_layer_map);
+    layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map,
+                                        layer_tree_host_.get());
 
     // Ensure second copy is correct.
     EXPECT_EQ(33, layer_dest_root->transform_tree_index_);
@@ -797,6 +810,8 @@ class LayerSerializationTest : public testing::Test {
     layer_dest_c = layer_dest_b->children()[0];
     EXPECT_EQ(layer_src_c->id(), layer_dest_c->id());
     EXPECT_EQ(99, layer_dest_c->transform_tree_index_);
+
+    layer_dest_root->SetLayerTreeHost(nullptr);
   }
 
   TestTaskGraphRunner task_graph_runner_;
@@ -2232,7 +2247,8 @@ TEST_F(LayerTest, RecursiveHierarchySerialization) {
 
   Layer::LayerIdMap empty_dest_layer_map;
   scoped_refptr<Layer> layer_dest_root = Layer::Create();
-  layer_dest_root->FromLayerNodeProto(proto, empty_dest_layer_map);
+  layer_dest_root->FromLayerNodeProto(proto, empty_dest_layer_map,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(layer_src_root->id(), layer_dest_root->id());
   EXPECT_EQ(nullptr, layer_dest_root->parent());
@@ -2254,6 +2270,8 @@ TEST_F(LayerTest, RecursiveHierarchySerialization) {
   EXPECT_EQ(0u, layer_dest_c->children().size());
   EXPECT_EQ(layer_src_c_mask->id(), layer_dest_c->mask_layer()->id());
   EXPECT_EQ(layer_src_c_replica->id(), layer_dest_c->replica_layer()->id());
+
+  layer_dest_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerTest, RecursiveHierarchySerializationWithNodeReuse) {
@@ -2278,7 +2296,8 @@ TEST_F(LayerTest, RecursiveHierarchySerializationWithNodeReuse) {
 
   Layer::LayerIdMap dest_layer_map_1;
   scoped_refptr<Layer> layer_dest_root = Layer::Create();
-  layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map_1);
+  layer_dest_root->FromLayerNodeProto(root_proto_1, dest_layer_map_1,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(layer_src_root->id(), layer_dest_root->id());
   ASSERT_EQ(1u, layer_dest_root->children().size());
@@ -2287,8 +2306,8 @@ TEST_F(LayerTest, RecursiveHierarchySerializationWithNodeReuse) {
 
   // Setup new destination layer map.
   Layer::LayerIdMap dest_layer_map_2;
-  dest_layer_map_2[layer_dest_root->id()] = layer_dest_root;
-  dest_layer_map_2[layer_dest_a_1->id()] = layer_dest_a_1;
+  layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+      &dest_layer_map_2);
 
   // Add Layer |b|.
   scoped_refptr<Layer> layer_src_b = Layer::Create();
@@ -2299,7 +2318,8 @@ TEST_F(LayerTest, RecursiveHierarchySerializationWithNodeReuse) {
   layer_src_root->ToLayerNodeProto(&root_proto_2);
 
   // Second deserialization.
-  layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map_2);
+  layer_dest_root->FromLayerNodeProto(root_proto_2, dest_layer_map_2,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(layer_src_root->id(), layer_dest_root->id());
   ASSERT_EQ(2u, layer_dest_root->children().size());
@@ -2316,6 +2336,8 @@ TEST_F(LayerTest, RecursiveHierarchySerializationWithNodeReuse) {
 
   // Layer |a| should be the same.
   EXPECT_EQ(layer_dest_a_1.get(), layer_dest_a_2.get());
+
+  layer_dest_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerTest, DeletingSubtreeDeletesLayers) {
@@ -2347,7 +2369,8 @@ TEST_F(LayerTest, DeletingSubtreeDeletesLayers) {
   // Deserialization 1.
   Layer::LayerIdMap empty_dest_layer_map;
   scoped_refptr<Layer> layer_dest_root = Layer::Create();
-  layer_dest_root->FromLayerNodeProto(proto1, empty_dest_layer_map);
+  layer_dest_root->FromLayerNodeProto(proto1, empty_dest_layer_map,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(layer_src_root->id(), layer_dest_root->id());
   ASSERT_EQ(2u, layer_dest_root->children().size());
@@ -2367,13 +2390,15 @@ TEST_F(LayerTest, DeletingSubtreeDeletesLayers) {
 
   // Deserialization 2.
   Layer::LayerIdMap dest_layer_map_2;
-  dest_layer_map_2[layer_dest_root->id()] = layer_dest_root;
-  dest_layer_map_2[layer_dest_a->id()] = layer_dest_a;
-  dest_layer_map_2[layer_dest_b->id()] = layer_dest_b;
-  layer_dest_root->FromLayerNodeProto(proto2, dest_layer_map_2);
+  layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+      &dest_layer_map_2);
+  layer_dest_root->FromLayerNodeProto(proto2, dest_layer_map_2,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(0u, layer_dest_a->children().size());
   EXPECT_EQ(0u, layer_dest_b->children().size());
+
+  layer_dest_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerTest, DeleteMaskAndReplicaLayer) {
@@ -2390,7 +2415,8 @@ TEST_F(LayerTest, DeleteMaskAndReplicaLayer) {
   // Deserialization 1.
   Layer::LayerIdMap dest_layer_map;
   scoped_refptr<Layer> layer_dest_root = Layer::Create();
-  layer_dest_root->FromLayerNodeProto(proto1, dest_layer_map);
+  layer_dest_root->FromLayerNodeProto(proto1, dest_layer_map,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(layer_src_root->id(), layer_dest_root->id());
   ASSERT_TRUE(layer_dest_root->mask_layer());
@@ -2401,13 +2427,6 @@ TEST_F(LayerTest, DeleteMaskAndReplicaLayer) {
   EXPECT_EQ(layer_src_root->replica_layer()->id(),
             layer_dest_root->replica_layer()->id());
 
-  // Store the newly constructed layer structure in the id map.
-  dest_layer_map[layer_dest_root->id()] = layer_dest_root;
-  dest_layer_map[layer_dest_root->mask_layer()->id()] =
-      layer_dest_root->mask_layer();
-  dest_layer_map[layer_dest_root->replica_layer()->id()] =
-      layer_dest_root->replica_layer();
-
   // Clear mask and replica layers.
   layer_src_root->mask_layer()->RemoveFromParent();
   layer_src_root->replica_layer()->RemoveFromParent();
@@ -2417,10 +2436,15 @@ TEST_F(LayerTest, DeleteMaskAndReplicaLayer) {
   layer_src_root->ToLayerNodeProto(&proto2);
 
   // Deserialization 2.
-  layer_dest_root->FromLayerNodeProto(proto2, dest_layer_map);
+  layer_dest_root->ClearLayerTreePropertiesForDeserializationAndAddToMap(
+      &dest_layer_map);
+  layer_dest_root->FromLayerNodeProto(proto2, dest_layer_map,
+                                      layer_tree_host_.get());
 
   EXPECT_EQ(nullptr, layer_dest_root->mask_layer());
   EXPECT_EQ(nullptr, layer_dest_root->replica_layer());
+
+  layer_dest_root->SetLayerTreeHost(nullptr);
 }
 
 TEST_F(LayerSerializationTest, HierarchyDeserializationWithLayerTreeHost) {
diff --git a/cc/trees/layer_tree_host.cc b/cc/trees/layer_tree_host.cc
index 16197c3..64622e4 100644
--- a/cc/trees/layer_tree_host.cc
+++ b/cc/trees/layer_tree_host.cc
@@ -1504,19 +1504,11 @@ void LayerTreeHost::FromProtobufForCommit(const proto::LayerTreeHost& proto) {
   // Layer hierarchy.
   scoped_refptr<Layer> new_root_layer =
       LayerProtoConverter::DeserializeLayerHierarchy(root_layer_,
-                                                     proto.root_layer());
+                                                     proto.root_layer(), this);
   if (root_layer_ != new_root_layer) {
-    if (root_layer_)
-      root_layer_->SetLayerTreeHost(nullptr);
     root_layer_ = new_root_layer;
-    root_layer_->SetLayerTreeHost(this);
   }
 
-  // Populate layer_id_map_ with the new layers.
-  layer_id_map_.clear();
-  LayerTreeHostCommon::CallFunctionForSubtree(
-      root_layer(),
-      [this](Layer* layer) { layer_id_map_[layer->id()] = layer; });
   for (auto layer_id : proto.layers_that_should_push_properties())
     layers_that_should_push_properties_.insert(layer_id_map_[layer_id]);
 
diff --git a/cc/trees/layer_tree_host_unittest_serialization.cc b/cc/trees/layer_tree_host_unittest_serialization.cc
index 3e3e0a1..cae21ca 100644
--- a/cc/trees/layer_tree_host_unittest_serialization.cc
+++ b/cc/trees/layer_tree_host_unittest_serialization.cc
@@ -40,13 +40,19 @@ class LayerTreeHostSerializationTest : public testing::Test {
     layer_tree_host_src_->in_paint_layer_contents_ = false;
     layer_tree_host_dst_->in_paint_layer_contents_ = false;
 
-    // Need to reset root layer and LayerTreeHost pointers before tear down.
-    layer_tree_host_src_->SetRootLayer(nullptr);
+    // Need to reset LayerTreeHost pointers before tear down.
     layer_tree_host_src_ = nullptr;
-    layer_tree_host_dst_->SetRootLayer(nullptr);
     layer_tree_host_dst_ = nullptr;
   }
 
+  void VerifyHostHasAllExpectedLayersInTree(Layer* root_layer) {
+    LayerTreeHostCommon::CallFunctionForSubtree(
+        root_layer, [root_layer](Layer* layer) {
+          DCHECK(layer->layer_tree_host());
+          EXPECT_EQ(layer, layer->layer_tree_host()->LayerById(layer->id()));
+        });
+  }
+
   void VerifySerializationAndDeserialization() {
     proto::LayerTreeHost proto;
 
@@ -269,7 +275,7 @@ class LayerTreeHostSerializationTest : public testing::Test {
     VerifySerializationAndDeserialization();
   }
 
-  void LayersChangedMultipleSerializations() {
+  void RunLayersChangedMultipleSerializations() {
     // Just fake setup a layer for both source and dest.
     scoped_refptr<Layer> root_layer_src = Layer::Create();
     layer_tree_host_src_->SetRootLayer(root_layer_src);
@@ -307,6 +313,39 @@ class LayerTreeHostSerializationTest : public testing::Test {
     VerifySerializationAndDeserialization();
   }
 
+  void RunAddAndRemoveNodeFromLayerTree() {
+    /* Testing serialization when the tree hierarchy changes like this:
+         root                  root
+        /    \                /    \
+       a      b        =>    a      c
+               \                     \
+                c                     d
+    */
+    scoped_refptr<Layer> layer_src_root = Layer::Create();
+    layer_tree_host_src_->SetRootLayer(layer_src_root);
+
+    scoped_refptr<Layer> layer_src_a = Layer::Create();
+    scoped_refptr<Layer> layer_src_b = Layer::Create();
+    scoped_refptr<Layer> layer_src_c = Layer::Create();
+    scoped_refptr<Layer> layer_src_d = Layer::Create();
+
+    layer_src_root->AddChild(layer_src_a);
+    layer_src_root->AddChild(layer_src_b);
+    layer_src_b->AddChild(layer_src_c);
+
+    VerifySerializationAndDeserialization();
+    VerifyHostHasAllExpectedLayersInTree(layer_tree_host_dst_->root_layer());
+
+    // Now change the Layer Hierarchy
+    layer_src_c->RemoveFromParent();
+    layer_src_b->RemoveFromParent();
+    layer_src_root->AddChild(layer_src_c);
+    layer_src_c->AddChild(layer_src_d);
+
+    VerifySerializationAndDeserialization();
+    VerifyHostHasAllExpectedLayersInTree(layer_tree_host_dst_->root_layer());
+  }
+
  private:
   TestTaskGraphRunner task_graph_runner_src_;
   FakeLayerTreeHostClient client_src_;
@@ -326,7 +365,11 @@ TEST_F(LayerTreeHostSerializationTest, LayersChanged) {
 }
 
 TEST_F(LayerTreeHostSerializationTest, LayersChangedMultipleSerializations) {
-  LayersChangedMultipleSerializations();
+  RunLayersChangedMultipleSerializations();
+}
+
+TEST_F(LayerTreeHostSerializationTest, AddAndRemoveNodeFromLayerTree) {
+  RunAddAndRemoveNodeFromLayerTree();
 }
 
 }  // namespace cc