O(1) socket removal

3 files changed, 26 insertions, 7 deletions

diff --git a/src/app_thread.cpp b/src/app_thread.cpp
index e4e5b19..e108594 100644
--- a/src/app_thread.cpp
+++ b/src/app_thread.cpp
@@ -138,16 +138,16 @@ zmq::socket_base_t *zmq::app_thread_t::create_socket (int type_)
     //  TODO: type is ignored for the time being.
     socket_base_t *s = new socket_base_t (this);
     zmq_assert (s);
+    s->set_index (sockets.size ());
     sockets.push_back (s);
     return s;
 }
 
 void zmq::app_thread_t::remove_socket (socket_base_t *socket_)
 {
-    //  TODO: To speed this up we can possibly use the system where each socket
-    //        holds its index (see I/O scheduler implementation).
-    sockets_t::iterator it = std::find (sockets.begin (), sockets.end (),
-        socket_);
-    zmq_assert (it != sockets.end ());
-    sockets.erase (it);
+    int i = socket_->get_index ();
+    socket_->set_index (-1);
+    sockets [i] = sockets.back ();
+    sockets [i]->set_index (i);
+    sockets.pop_back ();
 }
diff --git a/src/socket_base.cpp b/src/socket_base.cpp
index 4e14c68..6ad1f55 100644
--- a/src/socket_base.cpp
+++ b/src/socket_base.cpp
@@ -42,7 +42,8 @@ zmq::socket_base_t::socket_base_t (app_thread_t *parent_) :
     pending_term_acks (0),
     ticks (0),
     app_thread (parent_),
-    shutting_down (false)
+    shutting_down (false),
+    index (-1)
 {
 }
 
@@ -379,6 +380,17 @@ void zmq::socket_base_t::detach_outpipe (class writer_t *pipe_)
     out_pipes.pop_back ();
 }
 
+void zmq::socket_base_t::set_index (int index_)
+{
+    index = index_;
+}
+
+int zmq::socket_base_t::get_index ()
+{
+    zmq_assert (index != -1);
+    return index;
+}
+
 void zmq::socket_base_t::process_own (owned_t *object_)
 {
     io_objects.insert (object_);
diff --git a/src/socket_base.hpp b/src/socket_base.hpp
index 284d2c4..3f5774f 100644
--- a/src/socket_base.hpp
+++ b/src/socket_base.hpp
@@ -66,6 +66,10 @@ namespace zmq
         void detach_inpipe (class reader_t *pipe_);
         void detach_outpipe (class writer_t *pipe_);
 
+        //  Manipulating index in the app_thread's list of sockets.
+        void set_index (int index);
+        int get_index ();
+
     private:
 
         //  Handlers for incoming commands.
@@ -131,6 +135,9 @@ namespace zmq
         sessions_t sessions;
         mutex_t sessions_sync;
 
+        //  Index of the socket in the app_thread's list of sockets.
+        int index;
+
         socket_base_t (const socket_base_t&);
         void operator = (const socket_base_t&);
     };