#2092: location: get object pointer more directly

src/vt/topos/location/location.h

@@ -96,7 +96,7 @@ struct EntityLocationCoord : LocationCoord {
   using LocRecType = LocRecord<EntityID>;
   using LocCacheType = LocLookup<EntityID, LocRecType>;
   using LocEntityMsg = LocEntity<EntityID>;
-  using LocalRegisteredContType = std::unordered_set<EntityID>;
+  using LocalRegisteredContType = std::unordered_map<EntityID, void*>;
   using LocalRegisteredMsgContType = std::unordered_map<EntityID, LocEntityMsg>;
   using ActionListType = std::list<NodeActionType>;
   using PendingType = PendingLocationLookup<EntityID>;
@@ -138,10 +138,12 @@ struct EntityLocationCoord : LocationCoord {
    * \param[in] home the home node for this entity
    * \param[in] msg_action function to trigger when message arrives for it
    * \param[in] migrated whether it migrated in: \c entityEmigrated is preferred
+   * \param[in] obj pointer to the object associated with this entity
    */
   void registerEntity(
     EntityID const& id, NodeType const& home,
-    LocMsgActionType msg_action = nullptr, bool const& migrated = false
+    LocMsgActionType msg_action = nullptr, bool const& migrated = false,
+    void* obj = nullptr
   );
 
   /**
@@ -191,11 +193,13 @@ struct EntityLocationCoord : LocationCoord {
    * \param[in] id the entity ID
    * \param[in] home_node the home node for the entity
    * \param[in] msg_action function to trigger when message arrives for it
+   * \param[in] obj pointer to the object associated with this entity
    */
   void entityImmigrated(
     EntityID const& id, NodeType const& home_node,
     NodeType const& __attribute__((unused)) from_node,
-    LocMsgActionType msg_action = nullptr
+    LocMsgActionType msg_action = nullptr,
+    void* obj = nullptr
   );
 
   /**
@@ -263,9 +267,10 @@ struct EntityLocationCoord : LocationCoord {
    * \brief Route a message with a custom handler where the element is local
    *
    * \param[in] m message shared pointer
+   * \param[in] obj the object pointer
    */
   template <typename MessageT, ActiveTypedFnType<MessageT> *f>
-  void routeMsgHandlerLocal(MsgSharedPtr<MessageT> const& msg);
+  void routeMsgHandlerLocal(MsgSharedPtr<MessageT> const& msg, void* obj);
 
   /**
    * \brief Route a message to the default handler
@@ -441,6 +446,8 @@ struct EntityLocationCoord : LocationCoord {
    */
   LocInstType getInst() const;
 
+  void* getObjContext() const { return obj_context_; }
+
 private:
   LocInstType this_inst = no_loc_inst;
 
@@ -461,6 +468,9 @@ struct EntityLocationCoord : LocationCoord {
 
   // List of nodes that inquire about an entity that require an update
   LocAsksType loc_asks_;
+
+  /// Current object pointer context
+  void* obj_context_ = nullptr;
 };
 
 }}  // end namespace vt::location

src/vt/topos/location/location.impl.h

@@ -95,7 +95,7 @@ template <typename EntityID>
 template <typename EntityID>
 void EntityLocationCoord<EntityID>::registerEntity(
   EntityID const& id, NodeType const& home, LocMsgActionType msg_action,
-  bool const& migrated
+  bool const& migrated, void* obj
 ) {
   auto const& this_node = theContext()->getNode();
   auto reg_iter = local_registered_.find(id);
@@ -112,7 +112,7 @@ void EntityLocationCoord<EntityID>::registerEntity(
     this_inst, home, migrated, id
   );
 
-  local_registered_.insert(id);
+  local_registered_[id] = obj;
 
   recs_.insert(id, home, LocRecType{id, eLocState::Local, this_node});
 
@@ -260,11 +260,11 @@ void EntityLocationCoord<EntityID>::entityEmigrated(
 template <typename EntityID>
 void EntityLocationCoord<EntityID>::entityImmigrated(
   EntityID const& id, NodeType const& home_node, NodeType const& from,
-  LocMsgActionType msg_action
+  LocMsgActionType msg_action, void* obj
 ) {
   // @todo: currently `from' is unused, but is passed to this method in case we
   // need it in the future
-  return registerEntity(id, home_node, msg_action, true);
+  return registerEntity(id, home_node, msg_action, true, obj);
 }
 
 template <typename EntityID>
@@ -562,9 +562,11 @@ void EntityLocationCoord<EntityID>::routeMsgNode(
           hid, from, handler, envelopeGetRef(msg->env)
         );
 
+	obj_context_ = local_registered_.find(hid)->second;
         runnable::makeRunnable(msg, true, handler, from)
           .withTDEpochFromMsg()
           .run();
+	obj_context_ = nullptr;
       } else {
         auto reg_han_iter = local_registered_msg_han_.find(hid);
         vtAssert(
@@ -667,10 +669,11 @@ template <typename MessageT, ActiveTypedFnType<MessageT> *f>
 void EntityLocationCoord<EntityID>::routePreparedMsgHandler(
   MsgSharedPtr<MessageT> const& msg
 ) {
-  if (local_registered_.find(msg->getEntity()) == local_registered_.end()) {
+  auto iter = local_registered_.find(msg->getEntity());
+  if (iter == local_registered_.end()) {
     return routePreparedMsg(msg);
   } else {
-    return routeMsgHandlerLocal<MessageT, f>(msg);
+    return routeMsgHandlerLocal<MessageT, f>(msg, iter->second);
   }
 }
 
@@ -700,9 +703,11 @@ void EntityLocationCoord<EntityID>::setupMessageForRouting(
 template <typename EntityID>
 template <typename MessageT, ActiveTypedFnType<MessageT> *f>
 void EntityLocationCoord<EntityID>::routeMsgHandlerLocal(
-  MsgSharedPtr<MessageT> const& msg
+  MsgSharedPtr<MessageT> const& msg, void* obj
 ) {
+  obj_context_ = obj;
   f(msg.get());
+  obj_context_ = nullptr;
 }
 
 template <typename EntityID>

src/vt/vrt/collection/manager.impl.h

@@ -292,32 +292,16 @@ template <typename ColT, typename IndexT, typename MsgT>
   auto const& col = entity_proxy.getCollectionProxy();
   auto const& elm = entity_proxy.getElementProxy();
   auto const& idx = elm.getIndex();
-  auto elm_holder = theCollection()->findElmHolder<IndexT>(col);
-
-  bool const exists = elm_holder->exists(idx);
-
-  vt_debug_print(
-    terse, vrt_coll,
-    "collectionMsgTypedHandler: exists={}, idx={}, cur_epoch={:x}\n",
-    exists, idx, cur_epoch
-  );
-
-  vtAssertInfo(exists, "Proxy must exist", cur_epoch, idx);
-
-  auto& inner_holder = elm_holder->lookup(idx);
-
   auto const sub_handler = col_msg->getVrtHandler();
-  auto const col_ptr = inner_holder.getRawPtr();
 
   vt_debug_print(
-    verbose, vrt_coll,
-    "collectionMsgTypedHandler: sub_handler={}\n", sub_handler
+    terse, vrt_coll,
+    "collectionMsgTypedHandler: idx={}, cur_epoch={:x}, sub_handler={}\n",
+    idx, cur_epoch, sub_handler
   );
 
-  vtAssertInfo(
-    col_ptr != nullptr, "Must be valid pointer",
-    sub_handler, HandlerManager::isHandlerMember(sub_handler), cur_epoch, idx, exists
-  );
+  auto lm = theLocMan()->getCollectionLM<IndexT>(col);
+  auto obj = reinterpret_cast<Indexable<IndexT>*>(lm->getObjContext());
 
   // Dispatch the handler after pushing the contextual epoch
   theMsg()->pushEpoch(cur_epoch);
@@ -328,7 +312,7 @@ template <typename ColT, typename IndexT, typename MsgT>
     trace_event = col_msg->getFromTraceEvent();
   #endif
   collectionAutoMsgDeliver<ColT,IndexT,MsgT>(
-    msg, col_ptr, sub_handler, from, trace_event, false
+    msg, obj, sub_handler, from, trace_event, false
   );
   theMsg()->popEpoch(cur_epoch);
 }
@@ -1179,22 +1163,25 @@ bool CollectionManager::insertCollectionElement(
   );
 
   if (!destroyed) {
+    void* obj_ptr = vc.get();
     elm_holder->insert(idx, typename Holder<IndexT>::InnerHolder{
       std::move(vc)
     });
 
     if (is_migrated_in) {
       theLocMan()->getCollectionLM<IndexT>(proxy)->entityImmigrated(
         idx, home_node, migrated_from,
-        CollectionManager::collectionMsgHandler<ColT, IndexT>
+        CollectionManager::collectionMsgHandler<ColT, IndexT>,
+	obj_ptr
       );
       elm_holder->applyListeners(
         listener::ElementEventEnum::ElementMigratedIn, idx, home_node
       );
     } else {
       theLocMan()->getCollectionLM<IndexT>(proxy)->registerEntity(
         idx, home_node,
-        CollectionManager::collectionMsgHandler<ColT, IndexT>
+        CollectionManager::collectionMsgHandler<ColT, IndexT>,
+	false, obj_ptr
       );
       elm_holder->applyListeners(
         listener::ElementEventEnum::ElementCreated, idx, home_node