Convert sync_notifier::RegistrationManager to use object IDs.

This is part of the continued effort to refactor the sync code out of the
invalidation client.

BUG=124149
TEST=sync_unit_tests --gtest_filter=RegistrationManagerTest.*


Review URL: https://chromiumcodereview.appspot.com/10548063

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@144660 0039d316-1c4b-4281-b951-d872f2087c98

6 files changed, 361 insertions, 329 deletions

diff --git a/sync/notifier/chrome_invalidation_client.cc b/sync/notifier/chrome_invalidation_client.cc
index 0ec2c40..0c1376e 100644
--- a/sync/notifier/chrome_invalidation_client.cc
+++ b/sync/notifier/chrome_invalidation_client.cc
@@ -104,9 +104,21 @@ void ChromeInvalidationClient::UpdateCredentials(
 
 void ChromeInvalidationClient::RegisterTypes(syncable::ModelTypeSet types) {
   DCHECK(CalledOnValidThread());
+  // TODO(dcheng): Even though these fields are redundant, we keep both around
+  // for convenience while the conversion is still in process.
   registered_types_ = types;
+  registered_ids_.clear();
+  for (syncable::ModelTypeSet::Iterator it = types.First();
+       it.Good(); it.Inc()) {
+    invalidation::ObjectId id;
+    if (!RealModelTypeToObjectId(it.Get(), &id)) {
+      LOG(WARNING) << "Invalid model type: " << it.Get();
+      continue;
+    }
+    registered_ids_.insert(id);
+  }
   if (GetState() == NO_NOTIFICATION_ERROR && registration_manager_.get()) {
-    registration_manager_->SetRegisteredTypes(registered_types_);
+    registration_manager_->SetRegisteredIds(registered_ids_);
   }
   // TODO(akalin): Clear invalidation versions for unregistered types.
 }
@@ -115,7 +127,7 @@ void ChromeInvalidationClient::Ready(
     invalidation::InvalidationClient* client) {
   ticl_state_ = NO_NOTIFICATION_ERROR;
   EmitStateChange();
-  registration_manager_->SetRegisteredTypes(registered_types_);
+  registration_manager_->SetRegisteredIds(registered_ids_);
 }
 
 void ChromeInvalidationClient::Invalidate(
@@ -222,15 +234,9 @@ void ChromeInvalidationClient::InformRegistrationStatus(
   DVLOG(1) << "InformRegistrationStatus: "
            << ObjectIdToString(object_id) << " " << new_state;
 
-  syncable::ModelType model_type;
-  if (!ObjectIdToRealModelType(object_id, &model_type)) {
-    LOG(WARNING) << "Could not get object id model type; ignoring";
-    return;
-  }
-
   if (new_state != InvalidationListener::REGISTERED) {
     // Let |registration_manager_| handle the registration backoff policy.
-    registration_manager_->MarkRegistrationLost(model_type);
+    registration_manager_->MarkRegistrationLost(object_id);
   }
 }
 
@@ -245,22 +251,16 @@ void ChromeInvalidationClient::InformRegistrationFailure(
            << "is_transient=" << is_transient
            << ", message=" << error_message;
 
-  syncable::ModelType model_type;
-  if (!ObjectIdToRealModelType(object_id, &model_type)) {
-    LOG(WARNING) << "Could not get object id model type; ignoring";
-    return;
-  }
-
   if (is_transient) {
     // We don't care about |unknown_hint|; we let
     // |registration_manager_| handle the registration backoff policy.
-    registration_manager_->MarkRegistrationLost(model_type);
+    registration_manager_->MarkRegistrationLost(object_id);
   } else {
     // Non-transient failures are permanent, so block any future
     // registration requests for |model_type|.  (This happens if the
     // server doesn't recognize the data type, which could happen for
     // brand-new data types.)
-    registration_manager_->DisableType(model_type);
+    registration_manager_->DisableId(object_id);
   }
 }
 
diff --git a/sync/notifier/chrome_invalidation_client.h b/sync/notifier/chrome_invalidation_client.h
index 20733e9..8d998c6 100644
--- a/sync/notifier/chrome_invalidation_client.h
+++ b/sync/notifier/chrome_invalidation_client.h
@@ -23,6 +23,7 @@
 #include "sync/internal_api/public/util/weak_handle.h"
 #include "sync/notifier/chrome_system_resources.h"
 #include "sync/notifier/invalidation_state_tracker.h"
+#include "sync/notifier/invalidation_util.h"
 #include "sync/notifier/notifications_disabled_reason.h"
 #include "sync/notifier/state_writer.h"
 
@@ -146,6 +147,9 @@ class ChromeInvalidationClient
   scoped_ptr<invalidation::InvalidationClient> invalidation_client_;
   scoped_ptr<RegistrationManager> registration_manager_;
   // Stored to pass to |registration_manager_| on start.
+  ObjectIdSet registered_ids_;
+  // TODO(dcheng): This is a mirror of the data in registered_ids_. It
+  // temporarily remains for convenience.
   syncable::ModelTypeSet registered_types_;
 
   // The states of the ticl and the push client (with
diff --git a/sync/notifier/invalidation_util.h b/sync/notifier/invalidation_util.h
index 628ffed..8a7c931 100644
--- a/sync/notifier/invalidation_util.h
+++ b/sync/notifier/invalidation_util.h
@@ -8,6 +8,7 @@
 #define SYNC_NOTIFIER_INVALIDATION_UTIL_H_
 #pragma once
 
+#include <set>
 #include <string>
 
 #include "sync/internal_api/public/syncable/model_type.h"
@@ -26,6 +27,8 @@ struct ObjectIdLessThan {
                   const invalidation::ObjectId& rhs) const;
 };
 
+typedef std::set<invalidation::ObjectId, ObjectIdLessThan> ObjectIdSet;
+
 bool RealModelTypeToObjectId(syncable::ModelType model_type,
                              invalidation::ObjectId* object_id);
 
diff --git a/sync/notifier/registration_manager.cc b/sync/notifier/registration_manager.cc
index 5dfc39c..6d7a898 100644
--- a/sync/notifier/registration_manager.cc
+++ b/sync/notifier/registration_manager.cc
@@ -7,34 +7,36 @@
 #include <algorithm>
 #include <cstddef>
 #include <string>
+#include <utility>
 
 #include "base/rand_util.h"
+#include "base/stl_util.h"
 #include "google/cacheinvalidation/include/invalidation-client.h"
 #include "google/cacheinvalidation/include/types.h"
-#include "sync/internal_api/public/syncable/model_type.h"
 #include "sync/notifier/invalidation_util.h"
 
 namespace csync {
 
 RegistrationManager::PendingRegistrationInfo::PendingRegistrationInfo() {}
 
-RegistrationManager::RegistrationStatus::RegistrationStatus()
-    : model_type(syncable::UNSPECIFIED),
-      registration_manager(NULL),
+RegistrationManager::RegistrationStatus::RegistrationStatus(
+    const invalidation::ObjectId& id, RegistrationManager* manager)
+    : id(id),
+      registration_manager(manager),
       enabled(true),
-      state(invalidation::InvalidationListener::UNREGISTERED) {}
+      state(invalidation::InvalidationListener::UNREGISTERED) {
+  DCHECK(registration_manager);
+}
 
 RegistrationManager::RegistrationStatus::~RegistrationStatus() {}
 
 void RegistrationManager::RegistrationStatus::DoRegister() {
-  DCHECK_NE(model_type, syncable::UNSPECIFIED);
-  DCHECK(registration_manager);
   CHECK(enabled);
   // We might be called explicitly, so stop the timer manually and
   // reset the delay.
   registration_timer.Stop();
   delay = base::TimeDelta();
-  registration_manager->DoRegisterType(model_type);
+  registration_manager->DoRegisterId(id);
   DCHECK(!last_registration_request.is_null());
 }
 
@@ -56,113 +58,78 @@ RegistrationManager::RegistrationManager(
     invalidation::InvalidationClient* invalidation_client)
     : invalidation_client_(invalidation_client) {
   DCHECK(invalidation_client_);
-  // Initialize statuses.
-  for (int i = syncable::FIRST_REAL_MODEL_TYPE;
-       i < syncable::MODEL_TYPE_COUNT; ++i) {
-    syncable::ModelType model_type = syncable::ModelTypeFromInt(i);
-    RegistrationStatus* status = &registration_statuses_[model_type];
-    status->model_type = model_type;
-    status->registration_manager = this;
-  }
 }
 
 RegistrationManager::~RegistrationManager() {
   DCHECK(CalledOnValidThread());
+  STLDeleteValues(&registration_statuses_);
 }
 
-void RegistrationManager::SetRegisteredTypes(
-    syncable::ModelTypeSet types) {
+void RegistrationManager::SetRegisteredIds(const ObjectIdSet& ids) {
   DCHECK(CalledOnValidThread());
 
-  for (int i = syncable::FIRST_REAL_MODEL_TYPE;
-       i < syncable::MODEL_TYPE_COUNT; ++i) {
-    syncable::ModelType model_type = syncable::ModelTypeFromInt(i);
-    if (types.Has(model_type)) {
-      if (!IsTypeRegistered(model_type)) {
-        TryRegisterType(model_type, false /* is_retry */);
-      }
-    } else {
-      if (IsTypeRegistered(model_type)) {
-        UnregisterType(model_type);
-      }
-    }
-  }
-}
+  const ObjectIdSet& old_ids = GetRegisteredIds();
+  const ObjectIdSet& to_register = ids;
+  ObjectIdSet to_unregister;
+  std::set_difference(old_ids.begin(), old_ids.end(),
+                      ids.begin(), ids.end(),
+                      std::inserter(to_unregister, to_unregister.begin()),
+                      ObjectIdLessThan());
 
-void RegistrationManager::MarkRegistrationLost(
-    syncable::ModelType model_type) {
-  DCHECK(CalledOnValidThread());
-  RegistrationStatus* status = &registration_statuses_[model_type];
-  if (!status->enabled) {
-    return;
+  for (ObjectIdSet::const_iterator it = to_unregister.begin();
+       it != to_unregister.end(); ++it) {
+    UnregisterId(*it);
   }
-  status->state = invalidation::InvalidationListener::UNREGISTERED;
-  bool is_retry = !status->last_registration_request.is_null();
-  TryRegisterType(model_type, is_retry);
-}
 
-void RegistrationManager::MarkAllRegistrationsLost() {
-  DCHECK(CalledOnValidThread());
-  for (int i = syncable::FIRST_REAL_MODEL_TYPE;
-       i < syncable::MODEL_TYPE_COUNT; ++i) {
-    syncable::ModelType model_type = syncable::ModelTypeFromInt(i);
-    if (IsTypeRegistered(model_type)) {
-      MarkRegistrationLost(model_type);
+  for (ObjectIdSet::const_iterator it = to_register.begin();
+       it != to_register.end(); ++it) {
+    if (!ContainsKey(registration_statuses_, *it)) {
+      registration_statuses_.insert(
+          std::make_pair(*it, new RegistrationStatus(*it, this)));
+    }
+    if (!IsIdRegistered(*it)) {
+      TryRegisterId(*it, false /* is-retry */);
     }
   }
 }
 
-void RegistrationManager::DisableType(syncable::ModelType model_type) {
-  DCHECK(CalledOnValidThread());
-  RegistrationStatus* status = &registration_statuses_[model_type];
-  LOG(INFO) << "Disabling " << syncable::ModelTypeToString(model_type);
-  status->Disable();
-}
-
-syncable::ModelTypeSet RegistrationManager::GetRegisteredTypes() const {
+void RegistrationManager::MarkRegistrationLost(
+    const invalidation::ObjectId& id) {
   DCHECK(CalledOnValidThread());
-  syncable::ModelTypeSet registered_types;
-  for (int i = syncable::FIRST_REAL_MODEL_TYPE;
-       i < syncable::MODEL_TYPE_COUNT; ++i) {
-    syncable::ModelType model_type = syncable::ModelTypeFromInt(i);
-    if (IsTypeRegistered(model_type)) {
-      registered_types.Put(model_type);
-    }
+  RegistrationStatusMap::const_iterator it = registration_statuses_.find(id);
+  if (it == registration_statuses_.end()) {
+    DLOG(WARNING) << "Attempt to mark non-existent registration for "
+                  << ObjectIdToString(id) << " as lost";
+    return;
   }
-  return registered_types;
+  if (!it->second->enabled) {
+    return;
+  }
+  it->second->state = invalidation::InvalidationListener::UNREGISTERED;
+  bool is_retry = !it->second->last_registration_request.is_null();
+  TryRegisterId(id, is_retry);
 }
 
-RegistrationManager::PendingRegistrationMap
-    RegistrationManager::GetPendingRegistrations() const {
+void RegistrationManager::MarkAllRegistrationsLost() {
   DCHECK(CalledOnValidThread());
-  PendingRegistrationMap pending_registrations;
-  for (int i = syncable::FIRST_REAL_MODEL_TYPE;
-       i < syncable::MODEL_TYPE_COUNT; ++i) {
-    syncable::ModelType model_type = syncable::ModelTypeFromInt(i);
-    const RegistrationStatus& status = registration_statuses_[model_type];
-    if (status.registration_timer.IsRunning()) {
-      pending_registrations[model_type].last_registration_request =
-          status.last_registration_request;
-      pending_registrations[model_type].registration_attempt =
-          status.last_registration_attempt;
-      pending_registrations[model_type].delay = status.delay;
-      pending_registrations[model_type].actual_delay =
-          status.registration_timer.GetCurrentDelay();
+  for (RegistrationStatusMap::const_iterator it =
+           registration_statuses_.begin();
+       it != registration_statuses_.end(); ++it) {
+    if (IsIdRegistered(it->first)) {
+      MarkRegistrationLost(it->first);
     }
   }
-  return pending_registrations;
 }
 
-void RegistrationManager::FirePendingRegistrationsForTest() {
+void RegistrationManager::DisableId(const invalidation::ObjectId& id) {
   DCHECK(CalledOnValidThread());
-  for (int i = syncable::FIRST_REAL_MODEL_TYPE;
-       i < syncable::MODEL_TYPE_COUNT; ++i) {
-    syncable::ModelType model_type = syncable::ModelTypeFromInt(i);
-    RegistrationStatus* status = &registration_statuses_[model_type];
-    if (status->registration_timer.IsRunning()) {
-      status->DoRegister();
-    }
+  RegistrationStatusMap::const_iterator it = registration_statuses_.find(id);
+  if (it == registration_statuses_.end()) {
+    DLOG(WARNING) << "Attempt to disable non-existent registration for "
+                  << ObjectIdToString(id);
+    return;
   }
+  it->second->Disable();
 }
 
 // static
@@ -184,6 +151,43 @@ double RegistrationManager::CalculateBackoff(
                   std::min(max_retry_interval, new_retry_interval));
 }
 
+ObjectIdSet RegistrationManager::GetRegisteredIdsForTest() const {
+  return GetRegisteredIds();
+}
+
+RegistrationManager::PendingRegistrationMap
+    RegistrationManager::GetPendingRegistrationsForTest() const {
+  DCHECK(CalledOnValidThread());
+  PendingRegistrationMap pending_registrations;
+  for (RegistrationStatusMap::const_iterator it =
+           registration_statuses_.begin();
+       it != registration_statuses_.end(); ++it) {
+    const invalidation::ObjectId& id = it->first;
+    RegistrationStatus* status = it->second;
+    if (status->registration_timer.IsRunning()) {
+      pending_registrations[id].last_registration_request =
+          status->last_registration_request;
+      pending_registrations[id].registration_attempt =
+          status->last_registration_attempt;
+      pending_registrations[id].delay = status->delay;
+      pending_registrations[id].actual_delay =
+          status->registration_timer.GetCurrentDelay();
+    }
+  }
+  return pending_registrations;
+}
+
+void RegistrationManager::FirePendingRegistrationsForTest() {
+  DCHECK(CalledOnValidThread());
+  for (RegistrationStatusMap::const_iterator it =
+           registration_statuses_.begin();
+       it != registration_statuses_.end(); ++it) {
+    if (it->second->registration_timer.IsRunning()) {
+      it->second->DoRegister();
+    }
+  }
+}
+
 double RegistrationManager::GetJitter() {
   // |jitter| lies in [-1.0, 1.0), which is low-biased, but only
   // barely.
@@ -192,10 +196,16 @@ double RegistrationManager::GetJitter() {
   return 2.0 * base::RandDouble() - 1.0;
 }
 
-void RegistrationManager::TryRegisterType(syncable::ModelType model_type,
-                                          bool is_retry) {
+void RegistrationManager::TryRegisterId(const invalidation::ObjectId& id,
+                                        bool is_retry) {
   DCHECK(CalledOnValidThread());
-  RegistrationStatus* status = &registration_statuses_[model_type];
+  RegistrationStatusMap::const_iterator it = registration_statuses_.find(id);
+  if (it == registration_statuses_.end()) {
+    DLOG(FATAL) << "TryRegisterId called on " << ObjectIdToString(id)
+                << " which is not in the registration map";
+    return;
+  }
+  RegistrationStatus* status = it->second;
   if (!status->enabled) {
     // Disabled, so do nothing.
     return;
@@ -213,7 +223,7 @@ void RegistrationManager::TryRegisterType(syncable::ModelType model_type,
         (status->delay <= base::TimeDelta()) ?
         base::TimeDelta() : status->delay;
     DVLOG(2) << "Registering "
-             << syncable::ModelTypeToString(model_type) << " in "
+             << ObjectIdToString(id) << " in "
              << delay.InMilliseconds() << " ms";
     status->registration_timer.Stop();
     status->registration_timer.Start(FROM_HERE,
@@ -229,47 +239,64 @@ void RegistrationManager::TryRegisterType(syncable::ModelType model_type,
     status->next_delay =
         base::TimeDelta::FromSeconds(static_cast<int64>(next_delay_seconds));
     DVLOG(2) << "New next delay for "
-             << syncable::ModelTypeToString(model_type) << " is "
+             << ObjectIdToString(id) << " is "
              << status->next_delay.InSeconds() << " seconds";
   } else {
     DVLOG(2) << "Not a retry -- registering "
-             << syncable::ModelTypeToString(model_type) << " immediately";
+             << ObjectIdToString(id) << " immediately";
     status->delay = base::TimeDelta();
     status->next_delay = base::TimeDelta();
     status->DoRegister();
   }
 }
 
-void RegistrationManager::DoRegisterType(syncable::ModelType model_type) {
+void RegistrationManager::DoRegisterId(const invalidation::ObjectId& id) {
   DCHECK(CalledOnValidThread());
-  invalidation::ObjectId object_id;
-  if (!RealModelTypeToObjectId(model_type, &object_id)) {
-    LOG(DFATAL) << "Invalid model type: " << model_type;
+  invalidation_client_->Register(id);
+  RegistrationStatusMap::const_iterator it = registration_statuses_.find(id);
+  if (it == registration_statuses_.end()) {
+    DLOG(FATAL) << "DoRegisterId called on " << ObjectIdToString(id)
+                << " which is not in the registration map";
     return;
   }
-  invalidation_client_->Register(object_id);
-  RegistrationStatus* status = &registration_statuses_[model_type];
-  status->state = invalidation::InvalidationListener::REGISTERED;
-  status->last_registration_request = base::Time::Now();
+  it->second->state = invalidation::InvalidationListener::REGISTERED;
+  it->second->last_registration_request = base::Time::Now();
 }
 
-void RegistrationManager::UnregisterType(syncable::ModelType model_type) {
+void RegistrationManager::UnregisterId(const invalidation::ObjectId& id) {
   DCHECK(CalledOnValidThread());
-  invalidation::ObjectId object_id;
-  if (!RealModelTypeToObjectId(model_type, &object_id)) {
-    LOG(DFATAL) << "Invalid model type: " << model_type;
+  invalidation_client_->Unregister(id);
+  RegistrationStatusMap::iterator it = registration_statuses_.find(id);
+  if (it == registration_statuses_.end()) {
+    DLOG(FATAL) << "UnregisterId called on " << ObjectIdToString(id)
+                << " which is not in the registration map";
     return;
   }
-  invalidation_client_->Unregister(object_id);
-  RegistrationStatus* status = &registration_statuses_[model_type];
-  status->state = invalidation::InvalidationListener::UNREGISTERED;
+  delete it->second;
+  registration_statuses_.erase(it);
+}
+
+
+ObjectIdSet RegistrationManager::GetRegisteredIds() const {
+  DCHECK(CalledOnValidThread());
+  ObjectIdSet ids;
+  for (RegistrationStatusMap::const_iterator it =
+           registration_statuses_.begin();
+       it != registration_statuses_.end(); ++it) {
+    if (IsIdRegistered(it->first)) {
+      ids.insert(it->first);
+    }
+  }
+  return ids;
 }
 
-bool RegistrationManager::IsTypeRegistered(
-    syncable::ModelType model_type) const {
+bool RegistrationManager::IsIdRegistered(
+    const invalidation::ObjectId& id) const {
   DCHECK(CalledOnValidThread());
-  return registration_statuses_[model_type].state ==
-      invalidation::InvalidationListener::REGISTERED;
+  RegistrationStatusMap::const_iterator it =
+      registration_statuses_.find(id);
+  return it != registration_statuses_.end() &&
+      it->second->state == invalidation::InvalidationListener::REGISTERED;
 }
 
 }  // namespace csync
diff --git a/sync/notifier/registration_manager.h b/sync/notifier/registration_manager.h
index 45a066f..cf0144c 100644
--- a/sync/notifier/registration_manager.h
+++ b/sync/notifier/registration_manager.h
@@ -13,12 +13,12 @@
 
 #include "base/basictypes.h"
 #include "base/threading/non_thread_safe.h"
-#include "base/threading/non_thread_safe.h"
 #include "base/time.h"
 #include "base/timer.h"
-#include "sync/internal_api/public/syncable/model_type.h"
 // For invalidation::InvalidationListener::RegistrationState.
 #include "google/cacheinvalidation/include/invalidation-listener.h"
+#include "google/cacheinvalidation/include/types.h"
+#include "sync/notifier/invalidation_util.h"
 
 namespace csync {
 
@@ -55,9 +55,11 @@ class RegistrationManager : public base::NonThreadSafe {
     // The delay of the timer, which should be max(delay, 0).
     base::TimeDelta actual_delay;
   };
-  // Map from types with pending registrations to info about the
+  // Map of object IDs with pending registrations to info about the
   // pending registration.
-  typedef std::map<syncable::ModelType, PendingRegistrationInfo>
+  typedef std::map<invalidation::ObjectId,
+                   PendingRegistrationInfo,
+                   ObjectIdLessThan>
       PendingRegistrationMap;
 
   // Does not take ownership of |invalidation_client_|.
@@ -66,35 +68,22 @@ class RegistrationManager : public base::NonThreadSafe {
 
   virtual ~RegistrationManager();
 
-  // Registers all types included in the given set (that are not
-  // already disabled) and sets all other types to be unregistered.
-  void SetRegisteredTypes(syncable::ModelTypeSet types);
+  // Registers all object IDs included in the given set (that are not
+  // already disabled) and unregisters all other object IDs.
+  void SetRegisteredIds(const ObjectIdSet& ids);
 
-  // Marks the registration for the |model_type| lost and re-registers
+  // Marks the registration for the |id| lost and re-registers
   // it (unless it's disabled).
-  void MarkRegistrationLost(syncable::ModelType model_type);
+  void MarkRegistrationLost(const invalidation::ObjectId& id);
 
-  // Marks registrations lost for all enabled types and re-registers
-  // them.
+  // Marks registrations lost for all enabled object IDs and re-registers them.
   void MarkAllRegistrationsLost();
 
-  // Marks the registration for the |model_type| permanently lost and
-  // blocks any future registration attempts.
-  void DisableType(syncable::ModelType model_type);
-
-  // The functions below should only be used in tests.
-
-  // Gets all currently-registered types.
-  syncable::ModelTypeSet GetRegisteredTypes() const;
-
-  // Gets all pending registrations and their next min delays.
-  PendingRegistrationMap GetPendingRegistrations() const;
+  // Marks the registration for the |id| permanently lost and blocks any future
+  // registration attempts.
+  void DisableId(const invalidation::ObjectId& id);
 
-  // Run pending registrations immediately.
-  void FirePendingRegistrationsForTest();
-
-  // Calculate exponential backoff.  |jitter| must be Uniform[-1.0,
-  // 1.0].
+  // Calculate exponential backoff.  |jitter| must be Uniform[-1.0, 1.0].
   static double CalculateBackoff(double retry_interval,
                                  double initial_retry_interval,
                                  double min_retry_interval,
@@ -103,13 +92,25 @@ class RegistrationManager : public base::NonThreadSafe {
                                  double jitter,
                                  double max_jitter);
 
+  // The functions below should only be used in tests.
+
+  // Gets all currently registered ids.
+  ObjectIdSet GetRegisteredIdsForTest() const;
+
+  // Gets all pending registrations and their next min delays.
+  PendingRegistrationMap GetPendingRegistrationsForTest() const;
+
+  // Run pending registrations immediately.
+  void FirePendingRegistrationsForTest();
+
  protected:
   // Overrideable for testing purposes.
   virtual double GetJitter();
 
  private:
   struct RegistrationStatus {
-    RegistrationStatus();
+    RegistrationStatus(const invalidation::ObjectId& id,
+                       RegistrationManager* manager);
     ~RegistrationStatus();
 
     // Calls registration_manager->DoRegister(model_type). (needed by
@@ -120,10 +121,10 @@ class RegistrationManager : public base::NonThreadSafe {
     // Sets |enabled| to false and resets other variables.
     void Disable();
 
-    // The model type for which this is the status.
-    syncable::ModelType model_type;
+    // The object for which this is the status.
+    const invalidation::ObjectId id;
     // The parent registration manager.
-    RegistrationManager* registration_manager;
+    RegistrationManager* const registration_manager;
 
     // Whether this data type should be registered.  Set to false if
     // we get a non-transient registration failure.
@@ -142,29 +143,38 @@ class RegistrationManager : public base::NonThreadSafe {
     base::TimeDelta next_delay;
     // The actual timer for registration.
     base::OneShotTimer<RegistrationStatus> registration_timer;
+
+    DISALLOW_COPY_AND_ASSIGN(RegistrationStatus);
   };
+  typedef std::map<invalidation::ObjectId,
+                   RegistrationStatus*,
+                   ObjectIdLessThan>
+      RegistrationStatusMap;
 
-  // Does nothing if the given type is disabled.  Otherwise, if
+  // Does nothing if the given id is disabled.  Otherwise, if
   // |is_retry| is not set, registers the given type immediately and
   // resets all backoff parameters.  If |is_retry| is set, registers
   // the given type at some point in the future and increases the
   // delay until the next retry.
-  void TryRegisterType(syncable::ModelType model_type,
-                       bool is_retry);
+  void TryRegisterId(const invalidation::ObjectId& id,
+                     bool is_retry);
 
-  // Registers the given type, which must be valid, immediately.
+  // Registers the given id, which must be valid, immediately.
   // Updates |last_registration| in the appropriate
   // RegistrationStatus.  Should only be called by
   // RegistrationStatus::DoRegister().
-  void DoRegisterType(syncable::ModelType model_type);
+  void DoRegisterId(const invalidation::ObjectId& id);
+
+  // Unregisters the given object ID.
+  void UnregisterId(const invalidation::ObjectId& id);
 
-  // Unregisters the given type, which must be valid.
-  void UnregisterType(syncable::ModelType model_type);
+  // Gets all currently registered ids.
+  ObjectIdSet GetRegisteredIds() const;
 
-  // Returns true iff the given type, which must be valid, is registered.
-  bool IsTypeRegistered(syncable::ModelType model_type) const;
+  // Returns true iff the given object ID is registered.
+  bool IsIdRegistered(const invalidation::ObjectId& id) const;
 
-  RegistrationStatus registration_statuses_[syncable::MODEL_TYPE_COUNT];
+  RegistrationStatusMap registration_statuses_;
   // Weak pointer.
   invalidation::InvalidationClient* invalidation_client_;
 
diff --git a/sync/notifier/registration_manager_unittest.cc b/sync/notifier/registration_manager_unittest.cc
index 4b82b06..17f25fe 100644
--- a/sync/notifier/registration_manager_unittest.cc
+++ b/sync/notifier/registration_manager_unittest.cc
@@ -12,21 +12,14 @@
 
 #include "base/basictypes.h"
 #include "base/message_loop.h"
+#include "base/stl_util.h"
 #include "google/cacheinvalidation/include/invalidation-client.h"
-#include "sync/internal_api/public/syncable/model_type.h"
 #include "sync/notifier/invalidation_util.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace csync {
 namespace {
 
-syncable::ModelType ObjectIdToModelType(
-    const invalidation::ObjectId& object_id) {
-  syncable::ModelType model_type = syncable::UNSPECIFIED;
-  EXPECT_TRUE(ObjectIdToRealModelType(object_id, &model_type));
-  return model_type;
-}
-
 // Fake registration manager that lets you override jitter.
 class FakeRegistrationManager : public RegistrationManager {
  public:
@@ -53,20 +46,20 @@ class FakeRegistrationManager : public RegistrationManager {
 };
 
 // Fake invalidation client that just stores the currently-registered
-// model types.
+// object IDs.
 class FakeInvalidationClient : public invalidation::InvalidationClient {
  public:
   FakeInvalidationClient() {}
 
   virtual ~FakeInvalidationClient() {}
 
-  void LoseRegistration(syncable::ModelType model_type) {
-    EXPECT_TRUE(registered_types_.Has(model_type));
-    registered_types_.Remove(model_type);
+  void LoseRegistration(const invalidation::ObjectId& oid) {
+    EXPECT_TRUE(ContainsKey(registered_ids_, oid));
+    registered_ids_.erase(oid);
   }
 
   void LoseAllRegistrations() {
-    registered_types_.Clear();
+    registered_ids_.clear();
   }
 
   // invalidation::InvalidationClient implementation.
@@ -76,9 +69,8 @@ class FakeInvalidationClient : public invalidation::InvalidationClient {
   virtual void Acknowledge(const invalidation::AckHandle& handle) {}
 
   virtual void Register(const invalidation::ObjectId& oid) {
-    syncable::ModelType model_type = ObjectIdToModelType(oid);
-    EXPECT_FALSE(registered_types_.Has(model_type));
-    registered_types_.Put(model_type);
+    EXPECT_FALSE(ContainsKey(registered_ids_, oid));
+    registered_ids_.insert(oid);
   }
 
   virtual void Register(const std::vector<invalidation::ObjectId>& oids) {
@@ -86,53 +78,53 @@ class FakeInvalidationClient : public invalidation::InvalidationClient {
   }
 
   virtual void Unregister(const invalidation::ObjectId& oid) {
-    syncable::ModelType model_type = ObjectIdToModelType(oid);
-    EXPECT_TRUE(registered_types_.Has(model_type));
-    registered_types_.Remove(model_type);
+    EXPECT_TRUE(ContainsKey(registered_ids_, oid));
+    registered_ids_.erase(oid);
   }
 
   virtual void Unregister(const std::vector<invalidation::ObjectId>& oids) {
     // Unused for now.
   }
 
-  const syncable::ModelTypeSet GetRegisteredTypes() const {
-    return registered_types_;
+  const ObjectIdSet& GetRegisteredIdsForTest() const {
+    return registered_ids_;
   }
 
  private:
-  syncable::ModelTypeSet registered_types_;
+  ObjectIdSet registered_ids_;
 
   DISALLOW_COPY_AND_ASSIGN(FakeInvalidationClient);
 };
 
-const syncable::ModelType kModelTypes[] = {
-  syncable::BOOKMARKS,
-  syncable::PREFERENCES,
-  syncable::THEMES,
-  syncable::AUTOFILL,
-  syncable::EXTENSIONS,
-};
-const size_t kModelTypeCount = arraysize(kModelTypes);
+size_t kObjectIdsCount = 5;
 
-syncable::ModelTypeSet FromPtr(
-    const syncable::ModelType* types, size_t count) {
-  syncable::ModelTypeSet type_set;
-  for (size_t i = 0; i < count; ++i) {
-    type_set.Put(types[i]);
-  }
-  return type_set;
+invalidation::ObjectId GetIdForIndex(size_t index) {
+  char name[2] = "a";
+  name[0] += static_cast<char>(index);
+  return invalidation::ObjectId(1 + index, name);
+}
+
+ObjectIdSet GetSequenceOfIdsStartingAt(size_t start, size_t count) {
+  ObjectIdSet ids;
+  for (size_t i = start; i < start + count; ++i)
+    ids.insert(GetIdForIndex(i));
+  return ids;
+}
+
+ObjectIdSet GetSequenceOfIds(size_t count) {
+  return GetSequenceOfIdsStartingAt(0, count);
 }
 
 void ExpectPendingRegistrations(
-    syncable::ModelTypeSet expected_pending_types,
+    const ObjectIdSet& expected_pending_ids,
     double expected_delay_seconds,
     const RegistrationManager::PendingRegistrationMap& pending_registrations) {
-  syncable::ModelTypeSet pending_types;
+  ObjectIdSet pending_ids;
   for (RegistrationManager::PendingRegistrationMap::const_iterator it =
            pending_registrations.begin(); it != pending_registrations.end();
        ++it) {
-    SCOPED_TRACE(syncable::ModelTypeToString(it->first));
-    pending_types.Put(it->first);
+    SCOPED_TRACE(ObjectIdToString(it->first));
+    pending_ids.insert(it->first);
     base::TimeDelta offset =
         it->second.last_registration_request -
         it->second.registration_attempt;
@@ -141,16 +133,16 @@ void ExpectPendingRegistrations(
             static_cast<int64>(expected_delay_seconds)) + offset;
     // TODO(akalin): Add base::PrintTo() for base::Time and
     // base::TimeDeltas.
-    EXPECT_EQ(it->second.delay, expected_delay)
-        << it->second.delay.InMicroseconds()
-        << ", " << expected_delay.InMicroseconds();
+    EXPECT_EQ(expected_delay, it->second.delay)
+        << expected_delay.InMicroseconds()
+        << ", " << it->second.delay.InMicroseconds();
     if (it->second.delay <= base::TimeDelta()) {
-      EXPECT_EQ(it->second.actual_delay, base::TimeDelta());
+      EXPECT_EQ(base::TimeDelta(), it->second.actual_delay);
     } else {
-      EXPECT_EQ(it->second.delay, it->second.actual_delay);
+      EXPECT_EQ(it->second.actual_delay, it->second.delay);
     }
   }
-  EXPECT_TRUE(pending_types.Equals(expected_pending_types));
+  EXPECT_EQ(expected_pending_ids, pending_ids);
 }
 
 class RegistrationManagerTest : public testing::Test {
@@ -160,38 +152,38 @@ class RegistrationManagerTest : public testing::Test {
 
   virtual ~RegistrationManagerTest() {}
 
-  void LoseRegistrations(syncable::ModelTypeSet types) {
-    for (syncable::ModelTypeSet::Iterator it = types.First();
-         it.Good(); it.Inc()) {
-      fake_invalidation_client_.LoseRegistration(it.Get());
-      fake_registration_manager_.MarkRegistrationLost(it.Get());
+  void LoseRegistrations(const ObjectIdSet& oids) {
+    for (ObjectIdSet::const_iterator it = oids.begin(); it != oids.end();
+         ++it) {
+      fake_invalidation_client_.LoseRegistration(*it);
+      fake_registration_manager_.MarkRegistrationLost(*it);
     }
   }
 
-  void DisableTypes(syncable::ModelTypeSet types) {
-    for (syncable::ModelTypeSet::Iterator it = types.First();
-         it.Good(); it.Inc()) {
-      fake_invalidation_client_.LoseRegistration(it.Get());
-      fake_registration_manager_.DisableType(it.Get());
+  void DisableIds(const ObjectIdSet& oids) {
+    for (ObjectIdSet::const_iterator it = oids.begin(); it != oids.end();
+         ++it) {
+      fake_invalidation_client_.LoseRegistration(*it);
+      fake_registration_manager_.DisableId(*it);
     }
   }
 
   // Used by MarkRegistrationLostBackoff* tests.
   void RunBackoffTest(double jitter) {
     fake_registration_manager_.SetJitter(jitter);
-    syncable::ModelTypeSet types = FromPtr(kModelTypes, kModelTypeCount);
-    fake_registration_manager_.SetRegisteredTypes(types);
+    ObjectIdSet ids = GetSequenceOfIds(kObjectIdsCount);
+    fake_registration_manager_.SetRegisteredIds(ids);
 
-    // Lose some types.
-    syncable::ModelTypeSet lost_types = FromPtr(kModelTypes, 2);
-    LoseRegistrations(lost_types);
+    // Lose some ids.
+    ObjectIdSet lost_ids = GetSequenceOfIds(2);
+    LoseRegistrations(lost_ids);
     ExpectPendingRegistrations(
-        lost_types, 0.0,
-        fake_registration_manager_.GetPendingRegistrations());
+        lost_ids, 0.0,
+        fake_registration_manager_.GetPendingRegistrationsForTest());
 
     // Trigger another failure to start delaying.
     fake_registration_manager_.FirePendingRegistrationsForTest();
-    LoseRegistrations(lost_types);
+    LoseRegistrations(lost_ids);
 
     double scaled_jitter =
         jitter * RegistrationManager::kRegistrationDelayMaxJitter;
@@ -201,31 +193,31 @@ class RegistrationManagerTest : public testing::Test {
         (1.0 + scaled_jitter);
     expected_delay = std::floor(expected_delay);
     ExpectPendingRegistrations(
-        lost_types, expected_delay,
-        fake_registration_manager_.GetPendingRegistrations());
+        lost_ids, expected_delay,
+        fake_registration_manager_.GetPendingRegistrationsForTest());
 
     // Trigger another failure.
     fake_registration_manager_.FirePendingRegistrationsForTest();
-    LoseRegistrations(lost_types);
+    LoseRegistrations(lost_ids);
     expected_delay *=
         RegistrationManager::kRegistrationDelayExponent + scaled_jitter;
     expected_delay = std::floor(expected_delay);
     ExpectPendingRegistrations(
-        lost_types, expected_delay,
-        fake_registration_manager_.GetPendingRegistrations());
+        lost_ids, expected_delay,
+        fake_registration_manager_.GetPendingRegistrationsForTest());
 
     // Trigger enough failures to hit the ceiling.
     while (expected_delay < RegistrationManager::kMaxRegistrationDelaySeconds) {
       fake_registration_manager_.FirePendingRegistrationsForTest();
-      LoseRegistrations(lost_types);
+      LoseRegistrations(lost_ids);
       expected_delay *=
           RegistrationManager::kRegistrationDelayExponent + scaled_jitter;
       expected_delay = std::floor(expected_delay);
     }
     ExpectPendingRegistrations(
-        lost_types,
+        lost_ids,
         RegistrationManager::kMaxRegistrationDelaySeconds,
-        fake_registration_manager_.GetPendingRegistrations());
+        fake_registration_manager_.GetPendingRegistrationsForTest());
   }
 
   FakeInvalidationClient fake_invalidation_client_;
@@ -238,21 +230,23 @@ class RegistrationManagerTest : public testing::Test {
   DISALLOW_COPY_AND_ASSIGN(RegistrationManagerTest);
 };
 
-TEST_F(RegistrationManagerTest, SetRegisteredTypes) {
-  syncable::ModelTypeSet types = FromPtr(kModelTypes, kModelTypeCount);
+// Basic test of SetRegisteredIds to make sure we properly register new IDs and
+// unregister any IDs no longer in the set.
+TEST_F(RegistrationManagerTest, SetRegisteredIds) {
+  ObjectIdSet ids = GetSequenceOfIds(kObjectIdsCount - 1);
 
-  EXPECT_TRUE(fake_registration_manager_.GetRegisteredTypes().Empty());
-  EXPECT_TRUE(fake_invalidation_client_.GetRegisteredTypes().Empty());
+  EXPECT_TRUE(fake_registration_manager_.GetRegisteredIdsForTest().empty());
+  EXPECT_TRUE(fake_invalidation_client_.GetRegisteredIdsForTest().empty());
 
-  fake_registration_manager_.SetRegisteredTypes(types);
-  EXPECT_TRUE(fake_registration_manager_.GetRegisteredTypes().Equals(types));
-  EXPECT_TRUE(fake_invalidation_client_.GetRegisteredTypes().Equals(types));
+  fake_registration_manager_.SetRegisteredIds(ids);
+  EXPECT_EQ(ids, fake_registration_manager_.GetRegisteredIdsForTest());
+  EXPECT_EQ(ids, fake_invalidation_client_.GetRegisteredIdsForTest());
 
-  types.Put(syncable::APPS);
-  types.Remove(syncable::BOOKMARKS);
-  fake_registration_manager_.SetRegisteredTypes(types);
-  EXPECT_TRUE(fake_registration_manager_.GetRegisteredTypes().Equals(types));
-  EXPECT_TRUE(fake_invalidation_client_.GetRegisteredTypes().Equals(types));
+  ids.insert(GetIdForIndex(kObjectIdsCount - 1));
+  ids.erase(GetIdForIndex(kObjectIdsCount - 2));
+  fake_registration_manager_.SetRegisteredIds(ids);
+  EXPECT_EQ(ids, fake_registration_manager_.GetRegisteredIdsForTest());
+  EXPECT_EQ(ids, fake_invalidation_client_.GetRegisteredIdsForTest());
 }
 
 int GetRoundedBackoff(double retry_interval, double jitter) {
@@ -293,32 +287,28 @@ TEST_F(RegistrationManagerTest, CalculateBackoff) {
   EXPECT_EQ(20, GetRoundedBackoff(13.0, +1.0));
 }
 
+// Losing a registration should queue automatic re-registration.
 TEST_F(RegistrationManagerTest, MarkRegistrationLost) {
-  syncable::ModelTypeSet types = FromPtr(kModelTypes, kModelTypeCount);
+  ObjectIdSet ids = GetSequenceOfIds(kObjectIdsCount);
 
-  fake_registration_manager_.SetRegisteredTypes(types);
-  EXPECT_TRUE(fake_registration_manager_.GetPendingRegistrations().empty());
+  fake_registration_manager_.SetRegisteredIds(ids);
+  EXPECT_TRUE(
+      fake_registration_manager_.GetPendingRegistrationsForTest().empty());
 
-  // Lose some types.
-  syncable::ModelTypeSet lost_types = FromPtr(
-      kModelTypes, 3);
-  syncable::ModelTypeSet non_lost_types = FromPtr(
-      kModelTypes + 3, kModelTypeCount - 3);
-  LoseRegistrations(lost_types);
+  // Lose some ids.
+  ObjectIdSet lost_ids = GetSequenceOfIds(3);
+  ObjectIdSet non_lost_ids = GetSequenceOfIdsStartingAt(3, kObjectIdsCount - 3);
+  LoseRegistrations(lost_ids);
   ExpectPendingRegistrations(
-      lost_types, 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
-  EXPECT_TRUE(
-      fake_registration_manager_.GetRegisteredTypes().Equals(non_lost_types));
-  EXPECT_TRUE(
-      fake_invalidation_client_.GetRegisteredTypes().Equals(non_lost_types));
+      lost_ids, 0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
+  EXPECT_EQ(non_lost_ids, fake_registration_manager_.GetRegisteredIdsForTest());
+  EXPECT_EQ(non_lost_ids, fake_invalidation_client_.GetRegisteredIdsForTest());
 
   // Pretend we waited long enough to re-register.
   fake_registration_manager_.FirePendingRegistrationsForTest();
-  EXPECT_TRUE(
-      fake_registration_manager_.GetRegisteredTypes().Equals(types));
-  EXPECT_TRUE(
-      fake_invalidation_client_.GetRegisteredTypes().Equals(types));
+  EXPECT_EQ(ids, fake_registration_manager_.GetRegisteredIdsForTest());
+  EXPECT_EQ(ids, fake_invalidation_client_.GetRegisteredIdsForTest());
 }
 
 TEST_F(RegistrationManagerTest, MarkRegistrationLostBackoffLow) {
@@ -333,51 +323,51 @@ TEST_F(RegistrationManagerTest, MarkRegistrationLostBackoffHigh) {
   RunBackoffTest(+1.0);
 }
 
+// Exponential backoff on lost registrations should be reset to zero if
+// SetRegisteredIds is called.
 TEST_F(RegistrationManagerTest, MarkRegistrationLostBackoffReset) {
-  syncable::ModelTypeSet types = FromPtr(kModelTypes, kModelTypeCount);
+  ObjectIdSet ids = GetSequenceOfIds(kObjectIdsCount);
 
-  fake_registration_manager_.SetRegisteredTypes(types);
+  fake_registration_manager_.SetRegisteredIds(ids);
 
-  // Lose some types.
-  syncable::ModelTypeSet lost_types = FromPtr(kModelTypes, 2);
-  LoseRegistrations(lost_types);
+  // Lose some ids.
+  ObjectIdSet lost_ids = GetSequenceOfIds(2);
+  LoseRegistrations(lost_ids);
   ExpectPendingRegistrations(
-      lost_types, 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
+      lost_ids, 0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 
   // Trigger another failure to start delaying.
   fake_registration_manager_.FirePendingRegistrationsForTest();
-  LoseRegistrations(lost_types);
+  LoseRegistrations(lost_ids);
   double expected_delay =
       RegistrationManager::kInitialRegistrationDelaySeconds;
   ExpectPendingRegistrations(
-      lost_types, expected_delay,
-      fake_registration_manager_.GetPendingRegistrations());
+      lost_ids, expected_delay,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 
-  // Set types again.
-  fake_registration_manager_.SetRegisteredTypes(types);
+  // Set ids again.
+  fake_registration_manager_.SetRegisteredIds(ids);
   ExpectPendingRegistrations(
-      syncable::ModelTypeSet(), 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
+      ObjectIdSet(),
+      0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 }
 
 TEST_F(RegistrationManagerTest, MarkAllRegistrationsLost) {
-  syncable::ModelTypeSet types = FromPtr(kModelTypes, kModelTypeCount);
+  ObjectIdSet ids = GetSequenceOfIds(kObjectIdsCount);
 
-  fake_registration_manager_.SetRegisteredTypes(types);
+  fake_registration_manager_.SetRegisteredIds(ids);
 
   fake_invalidation_client_.LoseAllRegistrations();
   fake_registration_manager_.MarkAllRegistrationsLost();
 
-  syncable::ModelTypeSet expected_types;
-  EXPECT_TRUE(
-      fake_registration_manager_.GetRegisteredTypes().Equals(expected_types));
-  EXPECT_TRUE(
-      fake_invalidation_client_.GetRegisteredTypes().Equals(expected_types));
+  EXPECT_TRUE(fake_registration_manager_.GetRegisteredIdsForTest().empty());
+  EXPECT_TRUE(fake_invalidation_client_.GetRegisteredIdsForTest().empty());
 
   ExpectPendingRegistrations(
-      types, 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
+      ids, 0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 
   // Trigger another failure to start delaying.
   fake_registration_manager_.FirePendingRegistrationsForTest();
@@ -386,51 +376,49 @@ TEST_F(RegistrationManagerTest, MarkAllRegistrationsLost) {
   double expected_delay =
       RegistrationManager::kInitialRegistrationDelaySeconds;
   ExpectPendingRegistrations(
-      types, expected_delay,
-      fake_registration_manager_.GetPendingRegistrations());
+      ids, expected_delay,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 
   // Pretend we waited long enough to re-register.
   fake_registration_manager_.FirePendingRegistrationsForTest();
-  EXPECT_TRUE(
-      fake_registration_manager_.GetRegisteredTypes().Equals(types));
-  EXPECT_TRUE(
-      fake_invalidation_client_.GetRegisteredTypes().Equals(types));
+  EXPECT_EQ(ids, fake_registration_manager_.GetRegisteredIdsForTest());
+  EXPECT_EQ(ids, fake_invalidation_client_.GetRegisteredIdsForTest());
 }
 
-TEST_F(RegistrationManagerTest, DisableType) {
-  syncable::ModelTypeSet types = FromPtr(kModelTypes, kModelTypeCount);
+// IDs that are disabled should not be re-registered by SetRegisteredIds or
+// automatic re-registration if that registration is lost.
+TEST_F(RegistrationManagerTest, DisableId) {
+  ObjectIdSet ids = GetSequenceOfIds(kObjectIdsCount);
 
-  fake_registration_manager_.SetRegisteredTypes(types);
-  EXPECT_TRUE(fake_registration_manager_.GetPendingRegistrations().empty());
+  fake_registration_manager_.SetRegisteredIds(ids);
+  EXPECT_TRUE(
+      fake_registration_manager_.GetPendingRegistrationsForTest().empty());
 
-  // Disable some types.
-  syncable::ModelTypeSet disabled_types = FromPtr(
-      kModelTypes, 3);
-  syncable::ModelTypeSet enabled_types = FromPtr(
-      kModelTypes + 3, kModelTypeCount - 3);
-  DisableTypes(disabled_types);
+  // Disable some ids.
+  ObjectIdSet disabled_ids = GetSequenceOfIds(3);
+  ObjectIdSet enabled_ids = GetSequenceOfIdsStartingAt(3, kObjectIdsCount - 3);
+  DisableIds(disabled_ids);
   ExpectPendingRegistrations(
-      syncable::ModelTypeSet(), 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
-  EXPECT_TRUE(
-      fake_registration_manager_.GetRegisteredTypes().Equals(enabled_types));
-  EXPECT_TRUE(
-      fake_invalidation_client_.GetRegisteredTypes().Equals(enabled_types));
+      ObjectIdSet(),
+      0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
+  EXPECT_EQ(enabled_ids, fake_registration_manager_.GetRegisteredIdsForTest());
+  EXPECT_EQ(enabled_ids, fake_invalidation_client_.GetRegisteredIdsForTest());
 
-  fake_registration_manager_.SetRegisteredTypes(types);
-  EXPECT_TRUE(
-      fake_registration_manager_.GetRegisteredTypes().Equals(enabled_types));
+  fake_registration_manager_.SetRegisteredIds(ids);
+  EXPECT_EQ(enabled_ids, fake_registration_manager_.GetRegisteredIdsForTest());
 
   fake_registration_manager_.MarkRegistrationLost(
-      disabled_types.First().Get());
+      *disabled_ids.begin());
   ExpectPendingRegistrations(
-      syncable::ModelTypeSet(), 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
+      ObjectIdSet(),
+      0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 
   fake_registration_manager_.MarkAllRegistrationsLost();
   ExpectPendingRegistrations(
-      enabled_types, 0.0,
-      fake_registration_manager_.GetPendingRegistrations());
+      enabled_ids, 0.0,
+      fake_registration_manager_.GetPendingRegistrationsForTest());
 }
 
 }  // namespace