Support findInjectedSource (#573)

tools/clang/tools/dxcompiler/dxcdia.cpp

@@ -351,7 +351,7 @@ public:
 
   __override STDMETHODIMP findInjectedSource(
     /* [in] */ LPCOLESTR srcFile,
-    /* [out] */ IDiaEnumInjectedSources **ppResult) { return E_NOTIMPL; }
+    /* [out] */ IDiaEnumInjectedSources **ppResult);
 
   __override STDMETHODIMP getEnumDebugStreams(
     /* [out] */ IDiaEnumDebugStreams **ppEnumDebugStreams) { return E_NOTIMPL; }
@@ -1856,12 +1856,30 @@ public:
   }
 
   __override HRESULT GetItem(DWORD index, IDiaInjectedSource **ppItem) {
-    *ppItem = new (std::nothrow)DxcDiaInjectedSource(m_pSession, index);
+    if (index >= m_count)
+      return E_INVALIDARG;
+    unsigned itemIndex = index;
+    if (m_count == m_indexList.size())
+      itemIndex = m_indexList[index];
+    *ppItem = new (std::nothrow)DxcDiaInjectedSource(m_pSession, itemIndex);
     if (*ppItem == nullptr)
       return E_OUTOFMEMORY;
     (*ppItem)->AddRef();
     return S_OK;
   }
+  void Init(StringRef filename) {
+    for (unsigned i = 0; i < m_pSession->Contents()->getNumOperands(); ++i) {
+      StringRef fn =
+          dyn_cast<MDString>(m_pSession->Contents()->getOperand(i)->getOperand(0))
+              ->getString();
+      if (fn.equals(filename)) {
+        m_indexList.emplace_back(i);
+      }
+    }
+    m_count = m_indexList.size();
+  }
+private:
+  std::vector<unsigned> m_indexList;
 };
 
 class DxcDiaTableFrameData : public DxcDiaTableBase<IDiaEnumFrameData, IDiaFrameData> {
@@ -1887,6 +1905,22 @@ public:
   // HLSL is not based on IL, so no data to return.
 };
 
+__override STDMETHODIMP DxcDiaSession::findInjectedSource(
+    /* [in] */ LPCOLESTR srcFile,
+    /* [out] */ IDiaEnumInjectedSources **ppResult) {
+  if (Contents() != nullptr) {
+    CW2A pUtf8FileName(srcFile);
+    IDiaTable *pTable;
+    IFT(CreateDxcDiaTable(this, DiaTableKind::InjectedSource, &pTable));
+    DxcDiaTableInjectedSource *pInjectedSource =
+        dynamic_cast<DxcDiaTableInjectedSource *>(pTable);
+    pInjectedSource->Init(pUtf8FileName.m_psz);
+    *ppResult = pInjectedSource;
+    return S_OK;
+  }
+  return S_FALSE;
+}
+
 static
 HRESULT CreateDxcDiaTable(DxcDiaSession *pSession, DiaTableKind kind, IDiaTable **ppTable) {
   *ppTable = nullptr;

tools/clang/unittests/HLSL/CompilerTest.cpp

@@ -1186,6 +1186,81 @@ public:
 
     return o.str();
   }
+  std::wstring GetDebugFileContent(_In_ IDiaDataSource *pDataSource) {
+    CComPtr<IDiaSession> pSession;
+    CComPtr<IDiaTable> pTable;
+
+    CComPtr<IDiaTable> pSourcesTable;
+
+    CComPtr<IDiaEnumTables> pEnumTables;
+    std::wstringstream o;
+
+    VERIFY_SUCCEEDED(pDataSource->openSession(&pSession));
+    VERIFY_SUCCEEDED(pSession->getEnumTables(&pEnumTables));
+
+    ULONG fetched = 0;
+    while (pEnumTables->Next(1, &pTable, &fetched) == S_OK && fetched == 1) {
+      CComBSTR name;
+      IFT(pTable->get_name(&name));
+
+      if (wcscmp(name, L"SourceFiles") == 0) {
+        pSourcesTable = pTable.Detach();
+        continue;
+      }
+
+      pTable.Release();
+    }
+
+    if (!pSourcesTable) {
+      return L"cannot find source";
+    }
+
+    // Get source file contents.
+    // NOTE: "SourceFiles" has the root file first while "InjectedSources" is in
+    // alphabetical order.
+    //       It is important to keep the root file first for recompilation, so
+    //       iterate "SourceFiles" and look up the corresponding injected
+    //       source.
+    LONG count;
+    IFT(pSourcesTable->get_Count(&count));
+
+    CComPtr<IDiaSourceFile> pSourceFile;
+    CComBSTR pName;
+    CComPtr<IUnknown> pSymbolUnk;
+    CComPtr<IDiaEnumInjectedSources> pEnumInjectedSources;
+    CComPtr<IDiaInjectedSource> pInjectedSource;
+    std::wstring sourceText, sourceFilename;
+
+    while (SUCCEEDED(pSourcesTable->Next(1, &pSymbolUnk, &fetched)) &&
+           fetched == 1) {
+      sourceText = sourceFilename = L"";
+
+      IFT(pSymbolUnk->QueryInterface(&pSourceFile));
+      IFT(pSourceFile->get_fileName(&pName));
+
+      IFT(pSession->findInjectedSource(pName, &pEnumInjectedSources));
+
+      if (SUCCEEDED(pEnumInjectedSources->get_Count(&count)) && count == 1) {
+        IFT(pEnumInjectedSources->Item(0, &pInjectedSource));
+
+        DWORD cbData = 0;
+        std::string tempString;
+        CComBSTR bstr;
+        IFT(pInjectedSource->get_filename(&bstr));
+        IFT(pInjectedSource->get_source(0, &cbData, nullptr));
+
+        tempString.resize(cbData);
+        IFT(pInjectedSource->get_source(
+            cbData, &cbData, reinterpret_cast<BYTE *>(&tempString[0])));
+
+        CA2W tempWString(tempString.data());
+        o << tempWString.m_psz;
+      }
+      pSymbolUnk.Release();
+    }
+
+    return o.str();
+  }
 
   std::string GetOption(std::string &cmd, char *opt) {
     std::string option = cmd.substr(cmd.find(opt));
@@ -1440,6 +1515,8 @@ TEST_F(CompilerTest, CompileWhenDebugThenDIPresent) {
   VERIFY_IS_NOT_NULL(wcsstr(diaDump.c_str(), L"symIndexId: 5, CompilandEnv, name: hlslTarget, value: ps_6_0"));
   VERIFY_IS_NOT_NULL(wcsstr(diaDump.c_str(), L"lineNumber: 2"));
   VERIFY_IS_NOT_NULL(wcsstr(diaDump.c_str(), L"length: 99, filename: source.hlsl"));
+  std::wstring diaFileContent = GetDebugFileContent(pDiaSource).c_str();
+  VERIFY_IS_NOT_NULL(wcsstr(diaFileContent.c_str(), L"loat4 main(float4 pos : SV_Position) : SV_Target"));
 
 #if SUPPORT_FXC_PDB
   // Now, fake it by loading from a .pdb!