5 files changed, 55 insertions, 25 deletions

diff --git a/sync/internal_api/base_node.cc b/sync/internal_api/base_node.cc
index f9bf79b..a0ae01e 100644
--- a/sync/internal_api/base_node.cc
+++ b/sync/internal_api/base_node.cc
@@ -220,31 +220,7 @@ void BaseNode::GetChildIds(std::vector<int64>* result) const {
 }
 
 int BaseNode::GetTotalNodeCount() const {
-  syncable::BaseTransaction* trans = GetTransaction()->GetWrappedTrans();
-
-  int count = 1;  // Start with one to include the node itself.
-
-  std::stack<int64> stack;
-  stack.push(GetFirstChildId());
-  while (!stack.empty()) {
-    int64 handle = stack.top();
-    stack.pop();
-    if (handle == kInvalidId)
-      continue;
-    count++;
-    syncable::Entry entry(trans, syncable::GET_BY_HANDLE, handle);
-    if (!entry.good())
-      continue;
-    syncable::Id successor_id = entry.GetSuccessorId();
-    if (!successor_id.IsRoot())
-      stack.push(IdToMetahandle(trans, successor_id));
-    if (!entry.Get(syncable::IS_DIR))
-      continue;
-    syncable::Id child_id = entry.GetFirstChildId();
-    if (!child_id.IsRoot())
-      stack.push(IdToMetahandle(trans, child_id));
-  }
-  return count;
+  return GetEntry()->GetTotalNodeCount();
 }
 
 DictionaryValue* BaseNode::GetSummaryAsValue() const {
diff --git a/sync/syncable/directory.cc b/sync/syncable/directory.cc
index a844e66..9cbb73c 100644
--- a/sync/syncable/directory.cc
+++ b/sync/syncable/directory.cc
@@ -321,6 +321,46 @@ bool Directory::GetChildHandlesByHandle(
   return true;
 }
 
+int Directory::GetTotalNodeCount(
+    BaseTransaction* trans,
+    EntryKernel* kernel) const {
+  if (!SyncAssert(this == trans->directory(), FROM_HERE,
+                  "Directories don't match", trans))
+    return false;
+
+  int count = 1;
+  std::deque<const OrderedChildSet*> child_sets;
+
+  GetChildSetForKernel(trans, kernel, &child_sets);
+  while (!child_sets.empty()) {
+    const OrderedChildSet* set = child_sets.front();
+    child_sets.pop_front();
+    for (OrderedChildSet::const_iterator it = set->begin();
+         it != set->end(); ++it) {
+      count++;
+      GetChildSetForKernel(trans, *it, &child_sets);
+    }
+  }
+
+  return count;
+}
+
+void Directory::GetChildSetForKernel(
+    BaseTransaction* trans,
+    EntryKernel* kernel,
+    std::deque<const OrderedChildSet*>* child_sets) const {
+  if (!kernel->ref(IS_DIR))
+    return;  // Not a directory => no children.
+
+  const OrderedChildSet* descendants =
+      kernel_->parent_child_index->GetChildren(kernel->ref(ID));
+  if (!descendants)
+    return;  // This directory has no children.
+
+  // Add our children to the list of items to be traversed.
+  child_sets->push_back(descendants);
+}
+
 EntryKernel* Directory::GetRootEntry() {
   return GetEntryById(Id());
 }
diff --git a/sync/syncable/directory.h b/sync/syncable/directory.h
index 61c5522..1f787f28 100644
--- a/sync/syncable/directory.h
+++ b/sync/syncable/directory.h
@@ -321,6 +321,9 @@ class SYNC_EXPORT Directory {
   bool GetChildHandlesByHandle(BaseTransaction*, int64 handle,
       ChildHandles* result);
 
+  // Counts all items under the given node, including the node itself.
+  int GetTotalNodeCount(BaseTransaction*, EntryKernel* kernel_) const;
+
   // Returns true iff |id| has children.
   bool HasChildren(BaseTransaction* trans, const Id& id);
 
@@ -454,6 +457,12 @@ class SYNC_EXPORT Directory {
   bool SafeToPurgeFromMemory(WriteTransaction* trans,
                              const EntryKernel* const entry) const;
 
+  // A helper used by GetTotalNodeCount.
+  void GetChildSetForKernel(
+      BaseTransaction*,
+      EntryKernel* kernel_,
+      std::deque<const OrderedChildSet*>* child_sets) const;
+
   Directory& operator = (const Directory&);
 
  public:
diff --git a/sync/syncable/entry.cc b/sync/syncable/entry.cc
index ac62448..0cf1f5f 100644
--- a/sync/syncable/entry.cc
+++ b/sync/syncable/entry.cc
@@ -108,6 +108,10 @@ void Entry::GetChildHandles(std::vector<int64>* result) const {
   dir()->GetChildHandlesById(basetrans_, Get(ID), result);
 }
 
+int Entry::GetTotalNodeCount() const {
+  return dir()->GetTotalNodeCount(basetrans_, kernel_);
+}
+
 bool Entry::ShouldMaintainPosition() const {
   return kernel_->ShouldMaintainPosition();
 }
diff --git a/sync/syncable/entry.h b/sync/syncable/entry.h
index 26ccc0a..097f8a8 100644
--- a/sync/syncable/entry.h
+++ b/sync/syncable/entry.h
@@ -108,6 +108,7 @@ class SYNC_EXPORT Entry {
   Id GetPredecessorId() const;
   Id GetSuccessorId() const;
   Id GetFirstChildId() const;
+  int GetTotalNodeCount() const;
 
   // Returns a vector of this node's children's handles.
   // Clears |result| if there are no children.  If this node is of a type that