1/**
2 * @since 4.0.0
3 */
4
5import { RolldownMagicString } from "rolldown"
6import { type ESTree, parseSync } from "rolldown/utils"
7
8const markerPattern = /^ => (.+)$/
9const helper = "__effect_doctest_assert_"
10const statementTypes = new Set([
11 "ClassDeclaration",
12 "ExportAllDeclaration",
13 "ExportDefaultDeclaration",
14 "ExportNamedDeclaration",
15 "FunctionDeclaration",
16 "ImportDeclaration",
17 "TSDeclareFunction",
18 "TSEnumDeclaration",
19 "TSExportAssignment",
20 "TSImportEqualsDeclaration",
21 "TSInterfaceDeclaration",
22 "TSModuleDeclaration",
23 "TSNamespaceExportDeclaration",
24 "TSTypeAliasDeclaration",
25 "VariableDeclaration"
26])
27
28interface Node {
29 readonly type: string
30 readonly start: number
31 readonly end: number
32 readonly [key: string]: unknown
33}
34
35interface Candidate {
36 readonly node: Node
37 readonly parent: Node | undefined
38 readonly depth: number
39}
40
41const isNode = (value: unknown): value is Node =>
42 typeof value === "object" && value !== null &&
43 "type" in value && typeof value.type === "string" &&
44 "start" in value && typeof value.start === "number" &&
45 "end" in value && typeof value.end === "number"
46
47const isStatement = (node: Node): boolean => node.type.endsWith("Statement") || statementTypes.has(node.type)
48
49const collect = (program: ESTree.Program): { readonly candidates: Array<Candidate>; readonly names: Set<string> } => {
50 const candidates: Array<Candidate> = []
51 const names = new Set<string>()
52
53 const visit = (node: Node, parent: Node | undefined, depth: number): void => {
54 if (isStatement(node)) candidates.push({ node, parent, depth })
55 if (node.type === "Identifier" && typeof node.name === "string") names.add(node.name)
56
57 for (const value of Object.values(node)) {
58 if (isNode(value)) {
59 visit(value, node, depth + 1)
60 } else if (Array.isArray(value)) {
61 for (const child of value) {
62 if (isNode(child)) visit(child, node, depth + 1)
63 }
64 }
65 }
66 }
67
68 visit(program as unknown as Node, undefined, 0)
69 return { candidates, names }
70}
71
72const location = (source: string, offset: number, file: string, line: number): string => {
73 let localLine = 1
74 let column = 1
75 for (let index = 0; index < offset; index++) {
76 if (source.charCodeAt(index) === 10) {
77 localLine++
78 column = 1
79 } else {
80 column++
81 }
82 }
83 return `${file}:${line + localLine}:${column}`
84}
85
86const fail = (source: string, offset: number, file: string, line: number, message: string): never => {
87 throw new Error(`${location(source, offset, file, line)} ${message}`)
88}
89
90const isUnbracedControlFlow = (parent: Node | undefined): boolean =>
91 parent !== undefined && new Set([
92 "DoWhileStatement",
93 "ForInStatement",
94 "ForOfStatement",
95 "ForStatement",
96 "IfStatement",
97 "LabeledStatement",
98 "WhileStatement",
99 "WithStatement"
100 ]).has(parent.type)
101
102const assertionTarget = (
103 source: string,
104 candidates: ReadonlyArray<Candidate>,
105 comment: { readonly start: number },
106 file: string,
107 line: number
108): Candidate => {
109 const target = candidates
110 .filter(({ node }) => node.end <= comment.start && /^[ \t]*$/.test(source.slice(node.end, comment.start)))
111 .sort((left, right) => right.node.end - left.node.end || right.depth - left.depth)[0]
112
113 if (target === undefined) {
114 return fail(source, comment.start, file, line, "doctest assertion must trail a complete statement on the same line")
115 }
116 if (isUnbracedControlFlow(target.parent)) {
117 return fail(source, comment.start, file, line, "doctest assertions in control flow require an explicit block")
118 }
119 return target
120}
121
122const parseExpected = (expected: string, source: string, offset: number, file: string, line: number): void => {
123 const parsed = parseSync(`${file}?doctest-expected`, `const __expected = (${expected})`, { lang: "ts" })
124 const error = parsed.errors[0]
125 if (error !== undefined) {
126 fail(source, offset, file, line, `invalid doctest expected expression: ${error.message}`)
127 }
128}
129
130/**
131 * Transforms trailing doctest assertion comments into executable assertions.
132 *
133 * **Details**
134 *
135 * An assertion may trail an expression statement or a single initialized
136 * `const` identifier. Expected values are TypeScript expressions evaluated in
137 * the same lexical scope.
138 *
139 * @category testing
140 * @since 4.0.0
141 */
142export const transform = (source: string, file: string, line: number): string => {
143 if (!source.includes("// =>")) return source
144
145 const parsed = parseSync(file, source, { lang: "ts" })
146 const parseError = parsed.errors[0]
147 if (parseError !== undefined) {
148 fail(source, parseError.labels[0]?.start ?? 0, file, line, `invalid doctest source: ${parseError.message}`)
149 }
150
151 const markers = parsed.comments.flatMap((comment) => {
152 if (comment.type !== "Line") return []
153 const match = markerPattern.exec(comment.value)
154 return match === null ? [] : [{ comment, expected: match[1].trim() }]
155 })
156 if (markers.length === 0) return source
157
158 const { candidates, names } = collect(parsed.program)
159 let index = 0
160 while (names.has(`${helper}${index}`)) index++
161 const alias = `${helper}${index}`
162 const output = new RolldownMagicString(source)
163
164 for (const { comment, expected } of markers) {
165 parseExpected(expected, source, comment.start, file, line)
166 const { node, parent } = assertionTarget(source, candidates, comment, file, line)
167
168 if (node.type === "ExpressionStatement") {
169 if (node.directive !== null) {
170 fail(source, comment.start, file, line, "doctest assertions cannot target directives")
171 }
172 const expression = isNode(node.expression)
173 ? node.expression
174 : fail(source, comment.start, file, line, "doctest assertion has no expression")
175 output.overwrite(
176 expression.start,
177 comment.end,
178 `${alias}(${source.slice(expression.start, expression.end)}, ${expected});`
179 )
180 continue
181 }
182
183 if (node.type === "VariableDeclaration" && parent?.type !== "ExportNamedDeclaration") {
184 const declarations = node.declarations
185 const declaration = Array.isArray(declarations) && declarations.length === 1 ? declarations[0] : undefined
186 const id = isNode(declaration) && isNode(declaration.id) ? declaration.id : undefined
187 if (
188 node.kind === "const" && node.declare === false && isNode(declaration?.init) &&
189 id?.type === "Identifier" && typeof id.name === "string"
190 ) {
191 const indentation = source.slice(source.lastIndexOf("\n", node.start - 1) + 1, node.start)
192 output.overwrite(node.end, comment.end, `\n${indentation}${alias}(${id.name}, ${expected});`)
193 continue
194 }
195 }
196
197 fail(
198 source,
199 comment.start,
200 file,
201 line,
202 "doctest assertions can only target expression statements or a single initialized const identifier"
203 )
204 }
205
206 output.append(`\n\nimport { assertEquals as ${alias} } from "@effect/doctest/Runtime"`)
207 return output.toString()
208}