1 files changed, 32 insertions, 18 deletions

diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index e02715015..21f60c090 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -1639,25 +1639,7 @@ void FrontEndCompileRequest::parseTranslationUnit(
             }
         }
 
-        if (shouldDocument)
-        {
-            RefPtr<DocMarkup> markup(new DocMarkup);
-            markup->extract(translationUnit->getModuleDecl(), getSourceManager(), getSink());
-
-            // Extract to a file
 
-            const String& path = sourceFile->getPathInfo().foundPath;
-            if (path.getLength())
-            {
-                String fileName = Path::getFileNameWithoutExt(path);
-                fileName.append(".md");
-
-                StringBuilder buf;
-                DocumentationUtil::writeMarkdown(markup, buf);
-
-                File::writeAllText(fileName, buf);
-            }
-        }
 
 #if 0
         // Test serialization
@@ -1816,6 +1798,38 @@ SlangResult FrontEndCompileRequest::executeActionsInner()
     if (getSink()->getErrorCount() != 0)
         return SLANG_FAIL;
 
+    // After semantic checking is performed we can try and output doc information for this
+    if (shouldDocument)
+    {
+        // Not 100% clear where best to get the ASTBuilder from, but from the linkage shouldn't
+        // cause any problems with scoping
+        ASTBuilder* astBuilder = getLinkage()->getASTBuilder();
+
+        for (TranslationUnitRequest* translationUnit : translationUnits)
+        {
+            RefPtr<DocMarkup> markup(new DocMarkup);
+
+            markup->extract(translationUnit->getModuleDecl(), getSourceManager(), getSink());
+
+            // Hmm.. we can have multiple sourcefiles. So fir now we just pick the first, so as to come up with
+            // a reasonable name
+            SourceFile* sourceFile = translationUnit->getSourceFiles()[0];
+
+            // Extract to a file
+            const String& path = sourceFile->getPathInfo().foundPath;
+            if (path.getLength())
+            {
+                String fileName = Path::getFileNameWithoutExt(path);
+                fileName.append(".md");
+
+                StringBuilder buf;
+                DocumentationUtil::writeMarkdown(markup, astBuilder, buf);
+
+                File::writeAllText(fileName, buf);
+            }
+        }
+    }
+
     // Look up all the entry points that are expected,
     // and use them to populate the `program` member.
     //