2 files changed, 38 insertions, 14 deletions

diff --git a/ui/base/models/tree_node_model.h b/ui/base/models/tree_node_model.h
index 9720b8a..ec6e2c3 100644
--- a/ui/base/models/tree_node_model.h
+++ b/ui/base/models/tree_node_model.h
@@ -81,18 +81,19 @@ class TreeNode : public TreeModelNode {
     // If the node has a parent, remove it from its parent.
     NodeType* parent = node->parent();
     if (parent)
-      parent->Remove(parent->GetIndexOf(node));
+      parent->Remove(node);
     node->parent_ = static_cast<NodeType*>(this);
     children_->insert(children_->begin() + index, node);
   }
 
-  // Removes the node by index. This does NOT delete the specified node, it is
-  // up to the caller to delete it when done.
-  virtual NodeType* Remove(int index) {
-    DCHECK(index >= 0 && index < child_count());
-    NodeType* node = GetChild(index);
+  // Remove |node| from this Node and returns it. It's up to the caller to
+  // delete it.
+  virtual NodeType* Remove(NodeType* node) {
+    typename std::vector<NodeType*>::iterator i =
+        std::find(children_->begin(), children_->end(), node);
+    DCHECK(i != children_.end());
     node->parent_ = NULL;
-    children_->erase(index + children_->begin());
+    children_->erase(i);
     return node;
   }
 
@@ -213,17 +214,18 @@ class TreeNodeModel : public TreeModel {
     return static_cast<NodeType*>(model_node);
   }
 
-  void Add(NodeType* parent, NodeType* child, int index) {
-    DCHECK(parent && child);
-    parent->Add(child, index);
+  void Add(NodeType* parent, NodeType* node, int index) {
+    DCHECK(parent && node);
+    parent->Add(node, index);
     NotifyObserverTreeNodesAdded(parent, index, 1);
   }
 
-  NodeType* Remove(NodeType* parent, int index) {
+  NodeType* Remove(NodeType* parent, NodeType* node) {
     DCHECK(parent);
-    NodeType* child = parent->Remove(index);
+    int index = parent->GetIndexOf(node);
+    NodeType* delete_node = parent->Remove(node);
     NotifyObserverTreeNodesRemoved(parent, index, 1);
-    return child;
+    return delete_node;
   }
 
   void NotifyObserverTreeNodesAdded(NodeType* parent, int start, int count) {
diff --git a/ui/base/models/tree_node_model_unittest.cc b/ui/base/models/tree_node_model_unittest.cc
index 43e92623..c1ae8b4 100644
--- a/ui/base/models/tree_node_model_unittest.cc
+++ b/ui/base/models/tree_node_model_unittest.cc
@@ -114,7 +114,7 @@ TEST_F(TreeNodeModelTest, RemoveNode) {
   ASSERT_EQ(1, model.GetChildCount(root));
 
   // Now remove the |child1| from root and release the memory.
-  delete model.Remove(root, 0);
+  delete model.Remove(root, child1);
 
   AssertObserverCount(0, 1, 0);
 
@@ -292,4 +292,26 @@ TEST_F(TreeNodeModelTest, SetTitle) {
   EXPECT_EQ(title, root->GetTitle());
 }
 
+TEST_F(TreeNodeModelTest, BasicOperations) {
+  TreeNodeWithValue<int>* root = new TreeNodeWithValue<int>(0);
+  EXPECT_EQ(0, root->child_count());
+
+  TreeNodeWithValue<int>* child1 = new TreeNodeWithValue<int>(1);
+  root->Add(child1, root->child_count());
+  EXPECT_EQ(1, root->child_count());
+  EXPECT_EQ(root, child1->parent());
+
+  TreeNodeWithValue<int>* child2 = new TreeNodeWithValue<int>(1);
+  root->Add(child2, root->child_count());
+  EXPECT_EQ(2, root->child_count());
+  EXPECT_EQ(child1->parent(), child2->parent());
+
+  root->Remove(child2);
+  EXPECT_EQ(1, root->child_count());
+  EXPECT_EQ(NULL, child2->parent());
+
+  delete root->Remove(child1);
+  EXPECT_EQ(0, root->child_count());
+}
+
 }  // namespace ui