feat(twoslash-vue): support queries that matching with the source pos…

…ition

packages/twoslash-vue/src/index.ts

@@ -5,10 +5,13 @@ import type {
   Range,
   TwoSlashExecuteOptions,
   TwoSlashInstance,
+  TwoSlashReturnMeta,
 } from 'twoslash'
 import {
+  createPositionConverter,
   createTwoSlasher as createTwoSlasherBase,
   defaultCompilerOptions,
+  findQueryMarkers,
   removeCodeRanges,
   resolveNodePositions,
 } from 'twoslash'
@@ -32,7 +35,27 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}, flag
       },
     )
 
-    const fileSource = lang.createVirtualFile('index.vue', ts.ScriptSnapshot.fromString(code), 'vue')!
+    const sourceMeta = {
+      removals: [] as Range[],
+      positionCompletions: [] as number[],
+      positionQueries: [] as number[],
+      positionHighlights: [] as Range[],
+    } satisfies Partial<TwoSlashReturnMeta>
+
+    const pc = createPositionConverter(code)
+    // we get the markers with the original code so the position is correct
+    findQueryMarkers(code, sourceMeta, pc.getIndexOfLineAbove)
+
+    // replace non-whitespace in the already extracted markers
+    let strippedCode = code
+    for (const [start, end] of sourceMeta.removals) {
+      strippedCode
+       = strippedCode.slice(0, start)
+       + strippedCode.slice(start, end).replace(/\S/g, ' ')
+       + strippedCode.slice(end)
+    }
+
+    const fileSource = lang.createVirtualFile('index.vue', ts.ScriptSnapshot.fromString(strippedCode), 'vue')!
     const fileCompiled = fileSource.getEmbeddedFiles()[0]
     const typeHelpers = sharedTypes.getTypesCode(fileSource.vueCompilerOptions)
     const compiled = [
@@ -41,6 +64,8 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}, flag
       typeHelpers,
     ].join('\n')
 
+    const map = new SourceMap(fileCompiled.mappings)
+
     // Pass compiled to TS file to twoslash
     const result = twoslasherBase(compiled, 'tsx', {
       ...options,
@@ -59,13 +84,17 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}, flag
         // ignore internal types
         return !id.startsWith('__VLS')
       },
+      positionCompletions: sourceMeta.positionCompletions
+        .map(p => map.toGeneratedOffset(p)![0]),
+      positionQueries: sourceMeta.positionQueries
+        .map(p => map.toGeneratedOffset(p)![0]),
+      positionHighlights: sourceMeta.positionHighlights
+        .map(([start, end]) => [map.toGeneratedOffset(start)![0], map.toGeneratedOffset(end)![0]] as Range),
     })
 
     if (!flag)
       return result
 
-    const map = new SourceMap(fileCompiled.mappings)
-
     // Map the tokens
     const mappedNodes = result.nodes
       .map((q) => {
@@ -83,19 +112,28 @@ export function createTwoSlasher(createOptions: CreateTwoSlashOptions = {}, flag
       })
       .filter(isNotNull)
 
-    const mappedRemovals = result.meta.removals
-      .map((r) => {
-        const start = map.toSourceOffset(r[0])?.[0]
-        const end = map.toSourceOffset(r[1])?.[0]
-        if (start == null || end == null || start < 0 || end < 0 || start >= end)
-          return undefined
-        return [start, end] as Range
-      })
-      .filter(isNotNull)
+    const mappedRemovals = [
+      ...sourceMeta.removals,
+      ...result.meta.removals
+        .map((r) => {
+          const start = map.toSourceOffset(r[0])?.[0]
+          const end = map.toSourceOffset(r[1])?.[0]
+          if (start == null || end == null || start < 0 || end < 0 || start >= end)
+            return undefined
+          return [start, end] as Range
+        })
+        .filter(isNotNull),
+    ]
 
-    const removed = removeCodeRanges(code, mappedRemovals, mappedNodes)
-    result.code = removed.code
-    result.nodes = resolveNodePositions(removed.nodes, result.code)
+    if (!options.handbookOptions?.keepNotations) {
+      const removed = removeCodeRanges(code, mappedRemovals, mappedNodes)
+      result.code = removed.code
+      result.meta.removals = removed.removals
+      result.nodes = resolveNodePositions(removed.nodes, result.code)
+    }
+    else {
+      result.meta.removals = mappedRemovals
+    }
 
     return result
   }

packages/twoslash-vue/test/query.test.ts

@@ -1,9 +1,10 @@
-import { expect, it } from 'vitest'
+import { describe, expect, it } from 'vitest'
 import { createTwoSlasher } from '../src/index'
 
 const code = `
 <script setup lang="ts">
 import { ref, computed } from 'vue'
+//             ^?
 
 const count = ref(0)
 const double = computed(() => count.value * 2)
@@ -12,20 +13,57 @@ const double = computed(() => count.value * 2)
 
 <template>
   <button @click="count++">count is: {{ count }}</button>
+//           ^?
 </template>
 `
 
 const twoslasher = createTwoSlasher()
 
-it('basic query', () => {
+describe('basic', () => {
   const result = twoslasher(code, 'vue')
 
-  expect(result.nodes.find(n => n.type === 'hover' && n.target === 'button'))
-    .toHaveProperty('text', '(property) button: ButtonHTMLAttributes & ReservedProps')
+  it('has correct hover types', () => {
+    expect(result.nodes.find(n => n.type === 'hover' && n.target === 'button'))
+      .toHaveProperty('text', '(property) button: ButtonHTMLAttributes & ReservedProps')
+    expect(result.nodes.find(n => n.type === 'hover' && n.target === 'click'))
+      .toHaveProperty('text', '(property) \'click\': ((payload: MouseEvent) => void) | undefined')
+  })
 
-  expect(result.nodes.find(n => n.type === 'hover' && n.target === 'click'))
-    .toHaveProperty('text', '(property) \'click\': ((payload: MouseEvent) => void) | undefined')
+  it('has correct query', () => {
+    expect(result.meta.positionQueries)
+      .toMatchInlineSnapshot(`
+        [
+          38,
+          235,
+          1970,
+        ]
+      `)
 
-  expect(result.nodes.find(n => n.type === 'query' && n.target === 'double'))
-    .toHaveProperty('text', 'const double: ComputedRef<number>')
+    expect(result.nodes.find(n => n.type === 'query' && n.target === 'double'))
+      .toHaveProperty('text', 'const double: ComputedRef<number>')
+
+    expect(result.nodes.find(n => n.type === 'query' && n.target === 'computed'))
+      .toMatchInlineSnapshot(`
+        {
+          "character": 15,
+          "docs": undefined,
+          "length": 8,
+          "line": 2,
+          "start": 41,
+          "target": "computed",
+          "text": "(alias) const computed: {
+            <T>(getter: ComputedGetter<T>, debugOptions?: DebuggerOptions | undefined): ComputedRef<T>;
+            <T>(options: WritableComputedOptions<T>, debugOptions?: DebuggerOptions | undefined): WritableComputedRef<...>;
+        }
+        import computed",
+          "type": "query",
+        }
+      `)
+
+    // TODO: support this, and also it should throw an error if it's not found
+    // expect(result.nodes.find(n => n.type === 'query' && n.target === 'click'))
+    //   .toMatchInlineSnapshot(`undefined`)
+    // expect(result.nodes.filter(n => n.type === 'query'))
+    //   .toHaveLength(3)
+  })
 })