Reduce overhead of retargeting branches. (microsoft#327)

* Add back-links to branches from targets.

Makes branch targets track what branches target them.
This reduces the overhead of retargeting branches after
an insert because it means that the entire graph does
not need to be searched for branches. Instead the
branch target holds its own list of branches that target
it, and that list can be iterated over instead. This reduces
the cost of InsertBeforeAndRetargetOffsets from O(G) to O(B)
where B is the number of branches that target the Instruction
being inserted before. The amortized cost should therefore be
near O(1) because the average instruction is not a branch
target, and most branch targets will have only 1 branch tareting
them. I have not done the math to be certain of this, though.

The list of branches targeting a single instruction is stored
as a CDataItem so that we don't have to add additional pointers
to instructions where it isn't necessary. The list is a list
of raw pointers because every pointer in the list should be
kept alive by the graph, and using CComPtrs would increase overhead
while iterating through the list.

* Reduce overhead of GetInstructionSize

* Remove message logging from Instruction.cpp

Message logging is taking up a lot of time during iteration over
instructions. It doesn't add a lot of useful info to the logs, so
I'm removing it.

* Added debugwait to profiler helpers

* code review comments. Fix memory leak.

* fix merge problems

* Fix linux build

src/InstrumentationEngine.Lib/DataContainer.cpp

@@ -4,7 +4,7 @@
 #include "stdafx.h"
 #include "DataContainer.h"
 
-MicrosoftInstrumentationEngine::CDataContainer::CDataContainer()
+MicrosoftInstrumentationEngine::CDataContainer::CDataContainer() : m_dataContainerMap(1)
 {
     DEFINE_REFCOUNT_NAME(CDataContainer);
 
@@ -32,27 +32,51 @@ MicrosoftInstrumentationEngine::CDataContainer::~CDataContainer()
 HRESULT MicrosoftInstrumentationEngine::CDataContainer::SetDataItem(
     _In_ const GUID* pComponentId,
     _In_ const GUID* pObjectGuid,
-    _In_ IUnknown* pDataItem
+    _In_opt_ IUnknown* pDataItem
     )
 {
     HRESULT hr = S_OK;
 
     CCriticalSectionHolder lock(&m_cs);
 
-    CAtlMap<GUID, CComPtr<IUnknown>>* pMap;
+    CAtlMap<GUID, CComPtr<IUnknown>>* pMap = nullptr;
 
     if (m_dataContainerMap.Lookup(*pComponentId, pMap) == false)
     {
-        pMap = new CAtlMap<GUID, CComPtr<IUnknown>>;
-        if (!pMap)
+        // Only create the new map if we aren't setting the
+        // data item to null.
+        if (pDataItem != nullptr)
         {
-            return E_OUTOFMEMORY;
-        }
+            pMap = new CAtlMap<GUID, CComPtr<IUnknown>>(1);
+            if (!pMap)
+            {
+                return E_OUTOFMEMORY;
+            }
 
-        m_dataContainerMap.SetAt(*pComponentId, pMap);
+            m_dataContainerMap.SetAt(*pComponentId, pMap);
+        }
     }
 
-    pMap->SetAt(*pObjectGuid, CComPtr<IUnknown>(pDataItem));
+    if (pMap != nullptr)
+    {
+        // if the data item is null, free memory.
+        if (pDataItem == nullptr)
+        {
+            pMap->RemoveKey(*pObjectGuid);
+
+            if (pMap->GetCount() == 0)
+            {
+                m_dataContainerMap.RemoveKey(*pComponentId);
+
+                // need to manually delete the map.
+                delete pMap;
+            }
+        }
+        else
+        {
+            pMap->SetAt(*pObjectGuid, CComPtr<IUnknown>(pDataItem));
+        }
+    }
 
     return hr;
 }

src/InstrumentationEngine.Lib/DataContainer.h

@@ -23,7 +23,7 @@ namespace MicrosoftInstrumentationEngine
         virtual HRESULT __stdcall SetDataItem(
             _In_ const GUID* componentId,
             _In_ const GUID* objectGuid,
-            _In_ IUnknown* pDataItem
+            _In_opt_ IUnknown* pDataItem
             );
 
         virtual HRESULT __stdcall GetDataItem(

src/InstrumentationEngine/BranchTargetInfo.cpp

@@ -0,0 +1,148 @@
+#include "stdafx.h"
+#include "refcount.h"
+#include "Instruction.h"
+#include "BranchTargetInfo.h"
+
+namespace MicrosoftInstrumentationEngine
+{
+    // Note: we construct m_branches with only one initial bucket to save memory and because
+    // most instructions will not be targeted by that many branches.
+    CBranchTargetInfo::CBranchTargetInfo(_In_ CInstruction* pInstruction) : m_pInstruction(pInstruction), m_branches(1)
+    {
+    }
+
+    HRESULT CBranchTargetInfo::GetInstance(_In_ CInstruction* pInstruction, _Outptr_ CBranchTargetInfo** ppResult)
+    {
+        IfNullRet(pInstruction);
+        IfNullRet(ppResult);
+
+        *ppResult = nullptr;
+
+        const GUID* uuid = &__uuidof(CBranchTargetInfo);
+        CComPtr<IUnknown> pUnknown;
+        HRESULT hr;
+        // Don't assert, it is common for there to not be target info.
+        IfFailRetNoLog(pInstruction->GetDataItem(uuid, uuid, &pUnknown));
+        CComPtr<CBranchTargetInfo> pTargetInfo;
+        IfFailRet(pUnknown.QueryInterface(&pTargetInfo));
+
+        IfFalseRet(pTargetInfo != nullptr, E_UNEXPECTED);
+        *ppResult = pTargetInfo.Detach();
+
+        return S_OK;
+    }
+
+    HRESULT CBranchTargetInfo::GetOrCreateInstance(_In_ CInstruction* pInstruction, _Out_ CBranchTargetInfo** ppResult)
+    {
+        IfNullRet(pInstruction);
+        IfNullRet(ppResult);
+        if (!SUCCEEDED(GetInstance(pInstruction, ppResult)))
+        {
+            CComPtr<CBranchTargetInfo> pResult;
+            pResult.Attach(new CBranchTargetInfo(pInstruction));
+            const GUID* uuid = &__uuidof(CBranchTargetInfo);
+            pInstruction->SetDataItem(uuid, uuid, pResult);
+            *ppResult = pResult.Detach();
+        }
+
+        return S_OK;
+    }
+
+    HRESULT CBranchTargetInfo::SetBranchTarget(_In_ CInstruction* pBranch, _In_opt_ CInstruction* pTarget, _In_opt_ CInstruction* pOldTarget)
+    {
+        IfNullRet(pBranch);
+
+        HRESULT hr;
+
+        if (pOldTarget != nullptr)
+        {
+            CComPtr<CBranchTargetInfo> pOldInfo;
+
+            // There may not be an instance of CBranchTargetInfo if we are in the
+            // middle of retargetting a branch. That is OK, just ignore it.
+            if (SUCCEEDED(CBranchTargetInfo::GetInstance(pOldTarget, &pOldInfo)))
+            {
+                pOldInfo->Remove(pBranch);
+            }
+        }
+
+        if (pTarget != nullptr)
+        {
+            CComPtr<CBranchTargetInfo> pInfo;
+            IfFailRet(CBranchTargetInfo::GetOrCreateInstance(pTarget, &pInfo));
+            pInfo.p->m_branches.emplace(pBranch);
+        }
+
+        return S_OK;
+    }
+
+    HRESULT CBranchTargetInfo::RetargetBranches(_In_ CInstruction* pOriginalInstruction, _In_ CInstruction* pNewInstruction)
+    {
+        IfNullRet(pOriginalInstruction);
+        IfNullRet(pNewInstruction);
+
+        CComPtr<CBranchTargetInfo> pTargetInfo;
+        if (SUCCEEDED(CBranchTargetInfo::GetInstance(pOriginalInstruction, &pTargetInfo)))
+        {
+            return pTargetInfo->Retarget(pNewInstruction);
+        }
+        return S_OK;
+    }
+
+    HRESULT CBranchTargetInfo::Retarget(_In_ CInstruction* pNewInstruction)
+    {
+        HRESULT hr = S_OK;
+
+        // Remove this CBranchTargetInfo from the data container. If a new one
+        // is needed, it will be regenerated when branch targets are set.
+        Disconnect();
+        for (CInstruction* pInstr : m_branches)
+        {
+            if (pInstr->GetIsBranchInternal())
+            {
+                CComPtr<CBranchInstruction> pBranch;
+                IfFailRet(pInstr->QueryInterface(__uuidof(CBranchInstruction), (void**)&pBranch));
+
+                // This indicates that the new instruction targets this instruction.
+                // If we were to retarget it, we would cause an infinite loop.
+                if (pBranch != pNewInstruction)
+                {
+                    IfFailRet(pBranch->SetBranchTarget(pNewInstruction));
+                }
+                else
+                {
+                    // reset the branch target, since the disconnection
+                    // above will have removed it.
+                    IfFailRet(pBranch->SetBranchTarget(m_pInstruction));
+                }
+            }
+            else if (pInstr->GetIsSwitchInternal())
+            {
+                CComPtr<ISwitchInstruction> pSwitch;
+                IfFailRet(pInstr->QueryInterface(__uuidof(ISwitchInstruction), (void**)&pSwitch));
+                IfFalseRet(pSwitch != nullptr, E_UNEXPECTED);
+                IfFailRet(pSwitch->ReplaceBranchTarget(m_pInstruction, pNewInstruction));
+            }
+        }
+
+        return S_OK;
+    }
+
+    void CBranchTargetInfo::Remove(_In_ CInstruction* pOldBranch)
+    {
+        m_branches.erase(pOldBranch);
+        if (m_branches.empty())
+        {
+            Disconnect();
+        }
+    }
+
+    void CBranchTargetInfo::Disconnect()
+    {
+        const GUID* uuid = &__uuidof(CBranchTargetInfo);
+
+        // break cycles.
+        m_pInstruction->SetDataItem(uuid, uuid, nullptr);
+    }
+}
+

src/InstrumentationEngine/BranchTargetInfo.h

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+
+namespace MicrosoftInstrumentationEngine
+{
+    class __declspec(uuid("0E23A44A-6700-4810-889C-8067638C512E"))
+        CBranchTargetInfo : public IUnknown, virtual CModuleRefCount
+    {
+    private:
+
+        // This is a raw pointer to avoid circular references.
+        CInstruction* m_pInstruction;
+
+        // Using raw pointers for speed. The graph should keep these pointers
+        // alive.
+        std::unordered_set<CInstruction*> m_branches;
+
+    public:
+        DEFINE_DELEGATED_REFCOUNT_ADDREF(CBranchTargetInfo);
+        DEFINE_DELEGATED_REFCOUNT_RELEASE(CBranchTargetInfo);
+        STDMETHOD(QueryInterface)(_In_ REFIID riid, _Out_ void** ppvObject) override
+        {
+            return ImplQueryInterface(
+                static_cast<IUnknown*>(this),
+                this,
+                riid,
+                ppvObject
+            );
+        }
+
+    public:
+        static HRESULT GetInstance(_In_ CInstruction* pInstruction, _Outptr_ CBranchTargetInfo** ppResult);
+        static HRESULT GetOrCreateInstance(_In_ CInstruction* pInstruction, _Out_ CBranchTargetInfo** ppResult);
+        static HRESULT SetBranchTarget(_In_ CInstruction* pBranch, _In_opt_ CInstruction* pNewTarget, _In_opt_ CInstruction* pOldTarget);
+        static HRESULT RetargetBranches(_In_ CInstruction* pOriginalInstruction, _In_ CInstruction* pNewInstruction);
+        void Disconnect();
+
+        virtual ~CBranchTargetInfo() = default;
+
+        CBranchTargetInfo(_In_ CInstruction* pInstruction);
+        HRESULT Retarget(_In_ CInstruction* pNewInstruction);
+        void Remove(_In_ CInstruction* pOldBranch);
+    };
+}

src/InstrumentationEngine/CMakeLists.txt

@@ -20,6 +20,7 @@ set(src_files
     ./AssemblyInfo.cpp
     #./AssemblyInjector.cpp
     #./AtlModule.cpp
+    ./BranchTargetInfo.cpp
     ./CompositeType.cpp
     ./CorHeaders.cpp
     ./CorMethodMalloc.cpp