From 77ac92ad4ae14b1ba0128e26fbe98a74fadad514 Mon Sep 17 00:00:00 2001
From: Jeremy Koritzinsky <jekoritz@microsoft.com>
Date: Wed, 19 Jun 2024 11:04:55 -0700
Subject: [PATCH] Remove IDNMDOwner indirections in the regular importing
 scenarios

---
 src/interfaces/dispenser.cpp  |  6 +++---
 src/interfaces/dnmdowner.hpp  | 16 ++++++++++------
 src/interfaces/threadsafe.hpp | 15 ++++++++++-----
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/interfaces/dispenser.cpp b/src/interfaces/dispenser.cpp
index 7ad4e918..d5fd2b47 100644
--- a/src/interfaces/dispenser.cpp
+++ b/src/interfaces/dispenser.cpp
@@ -34,9 +34,9 @@ namespace
             threadSafeUnknown.Attach(new ControllingIUnknown());
             
             // Define an IDNMDOwner* tear-off here so the thread-safe object can be identified as a DNMD object.
-            (void)threadSafeUnknown->CreateAndAddTearOff<DelegatingDNMDOwner>(owner);
-            (void)threadSafeUnknown->CreateAndAddTearOff<ThreadSafeImportEmit<MetadataImportRO, MetadataEmit>>(import, emit);
-            // DelegatingDNMDOwner took ownership of owner, so we can return our thread-safe object here.
+            (void)threadSafeUnknown->CreateAndAddTearOff<DelegatingDNMDOwner>(handle_view);
+            (void)threadSafeUnknown->CreateAndAddTearOff<ThreadSafeImportEmit<MetadataImportRO, MetadataEmit>>(std::move(unknown), import, emit);
+            // ThreadSafeImportEmit took ownership of owner through unknown.
             return threadSafeUnknown;
         }
 
diff --git a/src/interfaces/dnmdowner.hpp b/src/interfaces/dnmdowner.hpp
index a4b6a557..e79a272e 100644
--- a/src/interfaces/dnmdowner.hpp
+++ b/src/interfaces/dnmdowner.hpp
@@ -19,6 +19,8 @@ struct IDNMDOwner : IUnknown
     virtual mdhandle_t MetaData() = 0;
 };
 
+class DNMDOwner;
+
 // We use a reference wrapper around the handle to allow the handle to be swapped out.
 // We plan to use swapping to implement table sorting as DNMD itself does not support
 // sorting tables or remapping tokens.
@@ -27,9 +29,9 @@ struct IDNMDOwner : IUnknown
 class mdhandle_view final
 {
 private:
-    IDNMDOwner* _owner;
+    DNMDOwner* _owner;
 public:
-    explicit mdhandle_view(IDNMDOwner* owner)
+    explicit mdhandle_view(DNMDOwner* owner)
         : _owner{ owner }
     {
     }
@@ -42,10 +44,7 @@ class mdhandle_view final
 
     mdhandle_view& operator=(mdhandle_view&& other) = default;
 
-    mdhandle_t get() const
-    {
-        return _owner->MetaData();
-    }
+    mdhandle_t get() const;
 
     bool operator==(std::nullptr_t) const
     {
@@ -110,4 +109,9 @@ class DNMDOwner final : public TearOffBase<IDNMDOwner>
     }
 };
 
+inline mdhandle_t mdhandle_view::get() const
+{
+    return _owner->MetaData();
+}
+
 #endif // !_SRC_INTERFACES_DNMDOWNER_HPP_
\ No newline at end of file
diff --git a/src/interfaces/threadsafe.hpp b/src/interfaces/threadsafe.hpp
index a36638a4..cfa0ab5f 100644
--- a/src/interfaces/threadsafe.hpp
+++ b/src/interfaces/threadsafe.hpp
@@ -14,10 +14,10 @@
 #include <cstdint>
 #include <mutex>
 
-// An tear-off that provides ownership semantics of an inner IDNMDOwner object.
+// A tear-off that re-exposes an mdhandle_view as an IDNMDOwner*.
 class DelegatingDNMDOwner final : public TearOffBase<IDNMDOwner>
 {
-    dncp::com_ptr<IDNMDOwner> _inner;
+    mdhandle_view _inner;
 
 protected:
     virtual bool TryGetInterfaceOnThis(REFIID riid, void** ppvObject) override
@@ -31,7 +31,7 @@ class DelegatingDNMDOwner final : public TearOffBase<IDNMDOwner>
     }
 
 public:
-    DelegatingDNMDOwner(IUnknown* controllingUnknown, IDNMDOwner* inner)
+    DelegatingDNMDOwner(IUnknown* controllingUnknown, mdhandle_view inner)
         : TearOffBase(controllingUnknown)
         , _inner{ inner }
     { }
@@ -40,7 +40,7 @@ class DelegatingDNMDOwner final : public TearOffBase<IDNMDOwner>
 
     mdhandle_t MetaData() override
     {
-        return _inner->MetaData();
+        return _inner.get();
     }
 };
 
@@ -48,6 +48,8 @@ template<typename TImport, typename TEmit>
 class ThreadSafeImportEmit : public TearOffBase<IMetaDataImport2, IMetaDataEmit2, IMetaDataAssemblyImport, IMetaDataAssemblyEmit>
 {
     pal::ReadWriteLock _lock;
+    // owning reference to the thread-unsafe object that provides the underlying implementation.
+    dncp::com_ptr<ControllingIUnknown> _threadUnsafe;
     // non-owning reference to the concrete non-locking implementations
     TImport* _import;
     TEmit* _emit;
@@ -80,12 +82,15 @@ class ThreadSafeImportEmit : public TearOffBase<IMetaDataImport2, IMetaDataEmit2
     }
 
 public:
-    ThreadSafeImportEmit(IUnknown* controllingUnknown, TImport* import, TEmit* emit)
+    ThreadSafeImportEmit(IUnknown* controllingUnknown, dncp::com_ptr<ControllingIUnknown>&& threadUnsafe, TImport* import, TEmit* emit)
         : TearOffBase(controllingUnknown)
         , _lock { }
+        , _threadUnsafe{ std::move(threadUnsafe) }
         , _import{ import }
         , _emit{ emit }
     {
+        assert(_import != nullptr);
+        assert(_emit != nullptr);
     }
 
     virtual ~ThreadSafeImportEmit() = default;