2 files changed, 74 insertions, 40 deletions

diff --git a/odb/transaction.cxx b/odb/transaction.cxx
index 197fc63..508d1aa 100644
--- a/odb/transaction.cxx
+++ b/odb/transaction.cxx
@@ -226,74 +226,99 @@ namespace odb
     s->state = state;
   }
 
-  void transaction::
-  unregister (void* key)
+  size_t transaction::
+  find (void* key)
   {
-    // Note that it is ok for this function not to find the key.
-    //
     if (callback_count_ == 0)
-      return;
+      return 0;
 
     size_t stack_count;
 
-    // See if this is the last slot registered. This will be a fast path
-    // if things are going to be unregistered from destructors.
+    // See if this is the last slot registered. This will be a fast path if,
+    // for example, things are going to be unregistered from destructors.
     //
     if (callback_count_ <= stack_callback_count)
     {
       if (stack_callbacks_[callback_count_ - 1].key == key)
-      {
-        callback_count_--;
-        return;
-      }
+        return callback_count_ - 1;
 
       stack_count = callback_count_;
     }
     else
     {
       if (dyn_callbacks_.back ().key == key)
-      {
-        dyn_callbacks_.pop_back ();
-        callback_count_--;
-        return;
-      }
+        return callback_count_ - 1;
 
       stack_count = stack_callback_count;
     }
 
-    size_t dyn_count (callback_count_ - stack_count);
-
     // Otherwise do a linear search.
     //
     for (size_t i (0); i < stack_count; ++i)
+      if (stack_callbacks_[i].key == key)
+        return i;
+
+    for (size_t i (0), dyn_count (callback_count_ - stack_count);
+         i < dyn_count; ++i)
+      if (dyn_callbacks_[i].key == key)
+        return i + stack_callback_count;
+
+    return callback_count_;
+  }
+
+  void transaction::
+  unregister (void* key)
+  {
+    size_t i (find (key));
+
+    // It is ok for this function not to find the key.
+    //
+    if (i == callback_count_)
+      return;
+
+    // See if this is the last slot registered.
+    //
+    if (i == callback_count_ - 1)
     {
-      callback_data& d (stack_callbacks_[i]);
-      if (d.key == key)
-      {
-        // Add to the free list.
-        //
-        d.event = 0;
-        d.key = reinterpret_cast<void*> (free_callback_);
-        free_callback_ = i;
-        return;
-      }
-    }
+      if (i >= stack_callback_count)
+        dyn_callbacks_.pop_back ();
 
-    for (size_t i (0); i < dyn_count; ++i)
+      callback_count_--;
+    }
+    else
     {
-      callback_data& d (dyn_callbacks_[i]);
-      if (d.key == key)
-      {
-        // Add to the free list.
-        //
-        d.event = 0;
-        d.key = reinterpret_cast<void*> (free_callback_);
-        free_callback_ = stack_callback_count + i;
-        return;
-      }
+      callback_data& d (
+        i < stack_callback_count
+        ? stack_callbacks_[i]
+        : dyn_callbacks_[i - stack_callback_count]);
+
+      // Add to the free list.
+      //
+      d.event = 0;
+      d.key = reinterpret_cast<void*> (free_callback_);
+      free_callback_ = i;
     }
   }
 
+  void transaction::
+  update (void* key, unsigned long long data, transaction** state)
+  {
+    size_t i (find (key));
+
+    // It is ok for this function not to find the key.
+    //
+    if (i == callback_count_)
+      return;
+
+    callback_data& d (
+      i < stack_callback_count
+      ? stack_callbacks_[i]
+      : dyn_callbacks_[i - stack_callback_count]);
+
+    d.data = data;
+    d.state = state;
+  }
+
   //
   // transaction_impl
   //
diff --git a/odb/transaction.hxx b/odb/transaction.hxx
index 1f7b866..e035661 100644
--- a/odb/transaction.hxx
+++ b/odb/transaction.hxx
@@ -136,6 +136,12 @@ namespace odb
     void
     unregister (void* key);
 
+    // Update the data and state values for a callback. Note that just
+    // like unregister(), this is a potentially slow operation.
+    //
+    void
+    update (void* key, unsigned long long data, transaction** state = 0);
+
   public:
     transaction_impl&
     implementation ();
@@ -149,6 +155,9 @@ namespace odb
   protected:
     friend struct rollback_guard;
 
+    std::size_t
+    find (void* key);
+
     void
     call (unsigned short event);