feat: support provide custom queries positions

src/core.ts

@@ -1,4 +1,4 @@
-import type { CompilerOptions, JsxEmit } from 'typescript'
+import type { CompilerOptions, JsxEmit, SourceFile } from 'typescript'
 import { createFSBackedSystem, createSystem, createVirtualTypeScriptEnvironment } from '@typescript/vfs'
 import { objectHash } from 'ohash'
 import { TwoslashError } from './error'
@@ -71,6 +71,10 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
       },
       removals: [],
       flagNotations: [],
+      virtualFiles: [],
+      positionQueries: options.positionQueries || [],
+      positionCompletions: options.positionCompletions || [],
+      positionHighlights: options.positionHighlights || [],
     }
     const {
       customTags = createOptions.customTags || [],
@@ -134,10 +138,60 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
       meta.removals.push([code.indexOf(cutAfterString), code.length])
     // #endregion
 
+    // #region extract markers
+    if (code.includes('//')) {
+      Array.from(code.matchAll(reAnnonateMarkers)).forEach((match) => {
+        const type = match[1] as '?' | '|' | '^^'
+        const index = match.index!
+        meta.removals.push([index, index + match[0].length + 1])
+        const markerIndex = match[0].indexOf('^')
+        const targetIndex = pc.getIndexOfLineAbove(index + markerIndex)
+        if (type === '?') {
+          meta.positionQueries.push(targetIndex)
+        }
+        else if (type === '|') {
+          meta.positionCompletions.push(targetIndex)
+        }
+        else {
+          const markerLength = match[0].lastIndexOf('^') - markerIndex + 1
+          meta.positionHighlights.push([
+            targetIndex,
+            targetIndex + markerLength,
+          ])
+        }
+      })
+    }
+    // #endregion
+
     const supportedFileTyes = ['js', 'jsx', 'ts', 'tsx']
-    const files = splitFiles(code, defaultFilename)
+    meta.virtualFiles = splitFiles(code, defaultFilename, fsRoot)
+
+    function getFileAtPosition(pos: number) {
+      return meta.virtualFiles.find(i => isInRange(pos, [i.offset, i.offset + i.content.length]))
+    }
+
+    function getQuickInfo(start: number, target: string): NodeWithoutPosition | undefined {
+      const file = getFileAtPosition(start)!
+      const quickInfo = ls.getQuickInfoAtPosition(file.filepath, start - file.offset)
+
+      if (quickInfo && quickInfo.displayParts) {
+        const text = quickInfo.displayParts.map(dp => dp.text).join('')
 
-    for (const file of files) {
+        // TODO: get different type of docs
+        const docs = quickInfo.documentation?.map(d => d.text).join('\n') || undefined
+
+        return {
+          type: 'hover',
+          text,
+          docs,
+          start,
+          length: target.length,
+          target,
+        }
+      }
+    }
+
+    for (const file of meta.virtualFiles) {
       // Only run the LSP-y things on source files
       if (file.extension === 'json') {
         if (!meta.compilerOptions.resolveJsonModule)
@@ -148,60 +202,16 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
       }
 
       const filepath = fsRoot + file.filename
-
       env.createFile(filepath, file.content)
 
-      if (!meta.handbookOptions.showEmit) {
-        const targetsQuery: number[] = []
-        const targetsCompletions: number[] = []
-        const targetsHighlights: Range[] = []
-        const source = ls.getProgram()!.getSourceFile(filepath)!
-
-        // #region extract markers
-        if (file.content.includes('//')) {
-          Array.from(file.content.matchAll(reAnnonateMarkers)).forEach((match) => {
-            const type = match[1] as '?' | '|' | '^^'
-            const index = match.index! + file.offset
-            meta.removals.push([index, index + match[0].length + 1])
-            const markerIndex = match[0].indexOf('^')
-            const targetIndex = pc.getIndexOfLineAbove(index + markerIndex)
-            if (type === '?') {
-              targetsQuery.push(targetIndex)
-            }
-            else if (type === '|') {
-              targetsCompletions.push(targetIndex)
-            }
-            else {
-              const markerLength = match[0].lastIndexOf('^') - markerIndex + 1
-              targetsHighlights.push([
-                targetIndex,
-                targetIndex + markerLength,
-              ])
-            }
-          })
-        }
-        // #endregion
+      const fileEnd = file.offset + file.content.length
+      function isInFile(pos: number) {
+        return file.offset <= pos && pos < fileEnd
+      }
 
+      if (!meta.handbookOptions.showEmit) {
         // #region get ts info for quick info
-        function getQuickInfo(start: number, target: string): NodeWithoutPosition | undefined {
-          const quickInfo = ls.getQuickInfoAtPosition(filepath, start - file.offset)
-
-          if (quickInfo && quickInfo.displayParts) {
-            const text = quickInfo.displayParts.map(dp => dp.text).join('')
-
-            // TODO: get different type of docs
-            const docs = quickInfo.documentation?.map(d => d.text).join('\n') || undefined
-
-            return {
-              type: 'hover',
-              text,
-              docs,
-              start,
-              length: target.length,
-              target,
-            }
-          }
-        }
+        const source = env.getSourceFile(filepath)!
 
         let identifiers: ReturnType<typeof getIdentifierTextSpans> | undefined
         if (!meta.handbookOptions.noStaticSemanticInfo) {
@@ -220,7 +230,9 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
         // #endregion
 
         // #region get query
-        for (const query of targetsQuery) {
+        for (const query of meta.positionQueries) {
+          if (!isInFile(query))
+            continue
           if (!identifiers)
             identifiers = getIdentifierTextSpans(ts, source, file.offset)
 
@@ -244,7 +256,9 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
         // #endregion
 
         // #region get highlights
-        for (const highlight of targetsHighlights) {
+        for (const highlight of meta.positionHighlights) {
+          if (!isInFile(highlight[0]))
+            continue
           if (!identifiers)
             identifiers = getIdentifierTextSpans(ts, source, file.offset)
 
@@ -268,9 +282,11 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
         // #endregion
 
         // #region get completions
-        targetsCompletions.forEach((target) => {
+        for (const target of meta.positionCompletions) {
+          if (!isInFile(target))
+            continue
           if (isInRemoval(target))
-            return
+            continue
           const completions = ls.getCompletionsAtPosition(filepath, target - 1, {})
           if (!completions && !meta.handbookOptions.noErrorValidation) {
             const pos = pc.indexToPos(target)
@@ -291,15 +307,15 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
             completions: (completions?.entries ?? []).filter(i => i.name.startsWith(prefix)),
             completionsPrefix: prefix,
           })
-        })
+        }
         // #endregion
       }
     }
 
     let errorNodes: Omit<NodeError, keyof Position>[] = []
 
     // #region get diagnostics, after all files are mounted
-    for (const file of files) {
+    for (const file of meta.virtualFiles) {
       if (!supportedFileTyes.includes(file.extension))
         continue
 
@@ -342,28 +358,32 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
 
     let outputCode = code
     if (meta.handbookOptions.showEmit) {
+      if (meta.handbookOptions.keepNotations) {
+        throw new TwoslashError(
+          `Option 'showEmit' cannot be used with 'keepNotations'`,
+          'With `showEmit` enabled, the output will always be the emitted code',
+          'Remove either option to continue',
+        )
+      }
       if (!meta.handbookOptions.keepNotations) {
         const { code: removedCode } = removeCodeRanges(outputCode, meta.removals)
-        const files = splitFiles(removedCode, defaultFilename)
-        for (const file of files) {
-          const filepath = fsRoot + file.filename
-          env.updateFile(filepath, file.content)
-        }
+        const files = splitFiles(removedCode, defaultFilename, fsRoot)
+        for (const file of files)
+          env.updateFile(file.filepath, file.content)
       }
       function removeExt(filename: string) {
         return filename.replace(/\.[^/.]+$/, '').replace(/\.d$/, '')
       }
 
-      const filenames = files.map(i => i.filename)
       const emitFilename = meta.handbookOptions.showEmittedFile
         ? meta.handbookOptions.showEmittedFile
         : meta.compilerOptions.jsx === 1 satisfies JsxEmit.Preserve
           ? 'index.jsx'
           : 'index.js'
-      let emitSource = files.find(i => removeExt(i.filename) === removeExt(emitFilename))?.filename
+      let emitSource = meta.virtualFiles.find(i => removeExt(i.filename) === removeExt(emitFilename))?.filename
 
       if (!emitSource && !meta.compilerOptions.outFile) {
-        const allFiles = filenames.join(', ')
+        const allFiles = meta.virtualFiles.map(i => i.filename).join(', ')
         throw new TwoslashError(
           `Could not find source file to show the emit for`,
           `Cannot find the corresponding **source** file  ${emitFilename} for completions via ^| returned no quickinfo from the compiler.`,
@@ -373,7 +393,7 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}): Two
 
       // Allow outfile, in which case you need any file.
       if (meta.compilerOptions.outFile)
-        emitSource = filenames[0]
+        emitSource = meta.virtualFiles[0].filename
 
       const output = ls.getEmitOutput(fsRoot + emitSource)
       const outfile = output.outputFiles

src/types.ts

@@ -1,5 +1,6 @@
 import type { VirtualTypeScriptEnvironment } from '@typescript/vfs'
 import type { CompilerOptions, CompletionEntry, CustomTransformers } from 'typescript'
+import type { VirtualFile } from './utils'
 
 type TS = typeof import('typescript')
 
@@ -13,7 +14,7 @@ export interface TwoSlashOptions extends CreateTwoSlashOptions, TwoSlashExecuteO
 /**
  * Options for twoslash instance
  */
-export interface TwoSlashExecuteOptions {
+export interface TwoSlashExecuteOptions extends Partial<Pick<TwoSlashReturnMeta, 'positionQueries' | 'positionCompletions' | 'positionHighlights'>> {
   /** Allows setting any of the handbook options from outside the function, useful if you don't want LSP identifiers */
   handbookOptions?: Partial<HandbookOptions>
 
@@ -112,6 +113,22 @@ export interface TwoSlashReturnMeta {
    * Flags which were parsed from the code
    */
   flagNotations: ParsedFlagNotation[]
+  /**
+   * The virtual files which were created
+   */
+  virtualFiles: VirtualFile[]
+  /**
+   * Positions of queries in the code
+   */
+  positionQueries: number[]
+  /**
+   * Positions of completions in the code
+   */
+  positionCompletions: number[]
+  /**
+   * Positions of errors in the code
+   */
+  positionHighlights: Range[]
 }
 
 export interface CompilerOptionDeclaration {

src/utils.ts

@@ -3,9 +3,10 @@ import { TwoslashError } from './error'
 import type { CompilerOptionDeclaration, NodeStartLength, NodeWithoutPosition, ParsedFlagNotation, Position, Range, TwoSlashNode } from './types'
 import { defaultHandbookOptions } from './defaults'
 
-export interface TemporaryFile {
+export interface VirtualFile {
   offset: number
   filename: string
+  filepath: string
   content: string
   extension: string
 }
@@ -155,11 +156,11 @@ export function createPositionConverter(code: string) {
 
 const reFilenamesMakers = /^\/\/\s?@filename: (.+)$/mg
 
-export function splitFiles(code: string, defaultFileName: string) {
+export function splitFiles(code: string, defaultFileName: string, root: string) {
   const matches = Array.from(code.matchAll(reFilenamesMakers))
 
   let currentFileName = defaultFileName
-  const files: TemporaryFile[] = []
+  const files: VirtualFile[] = []
 
   let index = 0
   for (const match of matches) {
@@ -169,6 +170,7 @@ export function splitFiles(code: string, defaultFileName: string) {
       files.push({
         offset: index,
         filename: currentFileName,
+        filepath: root + currentFileName,
         content,
         extension: getExtension(currentFileName),
       })
@@ -182,6 +184,7 @@ export function splitFiles(code: string, defaultFileName: string) {
     files.push({
       offset: index,
       filename: currentFileName,
+      filepath: root + currentFileName,
       content,
       extension: getExtension(currentFileName),
     })

test/fixtures/throws/show_emit_keep_notations.ts

@@ -0,0 +1,5 @@
+// @showEmit
+// @keepNotations
+
+console.log("Hello world" as string)
+//       ^?

test/new.test.ts

@@ -51,7 +51,7 @@ export function absolute(num: number) {
 import {absolute} from "./maths"
 const value = absolute(-1)
 //    ^?
-`, 'test.ts')
+`, 'test.ts', '')
   expect(files).toMatchInlineSnapshot(`
     [
       {
@@ -60,6 +60,7 @@ const value = absolute(-1)
     ",
         "extension": "ts",
         "filename": "test.ts",
+        "filepath": "test.ts",
         "offset": 0,
       },
       {
@@ -71,6 +72,7 @@ const value = absolute(-1)
     ",
         "extension": "ts",
         "filename": "maths.ts",
+        "filepath": "maths.ts",
         "offset": 20,
       },
       {
@@ -82,6 +84,7 @@ const value = absolute(-1)
     ",
         "extension": "ts",
         "filename": "index.ts",
+        "filepath": "index.ts",
         "offset": 131,
       },
     ]

test/results/throws/show_emit_keep_notations.txt

@@ -0,0 +1,6 @@
+
+## Option 'showEmit' cannot be used with 'keepNotations'
+
+With `showEmit` enabled, the output will always be the emitted code
+
+Remove either option to continue