[K/Wasm] Generate source-map for WAT-files

2024-01-05 14:16:32 +00:00
parent aea2e8052a
commit e4c244d5db
12 changed files with 265 additions and 75 deletions
@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmCompiledModuleFragment
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmModuleFragmentGenerator
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.toJsStringLiteral
 import org.jetbrains.kotlin.backend.wasm.lower.markExportedDeclarations
+import org.jetbrains.kotlin.backend.wasm.utils.SourceMapGenerator
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.ir.backend.js.MainModule
 import org.jetbrains.kotlin.ir.backend.js.ModulesStructure
@@ -27,6 +28,7 @@ import org.jetbrains.kotlin.js.config.WasmTarget
 import org.jetbrains.kotlin.js.sourceMap.SourceFilePathResolver
 import org.jetbrains.kotlin.js.sourceMap.SourceMap3Builder
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.wasm.ir.convertors.WasmIrToBinary
 import org.jetbrains.kotlin.wasm.ir.convertors.WasmIrToText
 import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocation
@@ -39,7 +41,12 @@ class WasmCompilerResult(
    val jsUninstantiatedWrapper: String?,
    val jsWrapper: String,
    val wasm: ByteArray,
-    val sourceMap: String?
+    val debugInformation: DebugInformation?
+)
+
+class DebugInformation(
+    val sourceMapForBinary: String?,
+    val sourceMapForText: String?,
 )

 fun compileToLoweredIr(
@@ -110,9 +117,16 @@ fun compileWasm(
    allModules.forEach { codeGenerator.collectInterfaceTables(it) }
    allModules.forEach { codeGenerator.generateModule(it) }

+    val sourceMapGeneratorForBinary = runIf(generateSourceMaps) {
+        SourceMapGenerator("$baseFileName.wasm", backendContext.configuration)
+    }
+    val sourceMapGeneratorForText = runIf(generateWat && generateSourceMaps) {
+        SourceMapGenerator("$baseFileName.wat", backendContext.configuration)
+    }
+
    val linkedModule = compiledWasmModule.linkWasmCompiledFragments()
    val wat = if (generateWat) {
-        val watGenerator = WasmIrToText()
+        val watGenerator = WasmIrToText(sourceMapGeneratorForText)
        watGenerator.appendWasmModule(linkedModule)
        watGenerator.toString()
    } else {
@@ -121,18 +135,13 @@ fun compileWasm(

    val os = ByteArrayOutputStream()

-    val sourceMapFileName = "$baseFileName.wasm.map".takeIf { generateSourceMaps }
-    val sourceLocationMappings =
-        if (generateSourceMaps) mutableListOf<SourceLocationMapping>() else null
-
    val wasmIrToBinary =
        WasmIrToBinary(
            os,
            linkedModule,
            allModules.last().descriptor.name.asString(),
            emitNameSection,
-            sourceMapFileName,
-            sourceLocationMappings
+            sourceMapGeneratorForBinary
        )

    wasmIrToBinary.appendWasmModule()
@@ -156,43 +165,13 @@ fun compileWasm(
        jsUninstantiatedWrapper = jsUninstantiatedWrapper,
        jsWrapper = jsWrapper,
        wasm = byteArray,
-        sourceMap = generateSourceMap(backendContext.configuration, sourceLocationMappings)
+        debugInformation = DebugInformation(
+            sourceMapGeneratorForBinary?.generate(),
+            sourceMapGeneratorForText?.generate(),
+        ),
    )
 }

-private fun generateSourceMap(
-    configuration: CompilerConfiguration,
-    sourceLocationMappings: MutableList<SourceLocationMapping>?
-): String? {
-    if (sourceLocationMappings == null) return null
-
-    val sourceMapsInfo = SourceMapsInfo.from(configuration) ?: return null
-
-    val sourceMapBuilder =
-        SourceMap3Builder(null, { error("This should not be called for Kotlin/Wasm") }, sourceMapsInfo.sourceMapPrefix)
-
-    val pathResolver =
-        SourceFilePathResolver.create(sourceMapsInfo.sourceRoots, sourceMapsInfo.sourceMapPrefix, sourceMapsInfo.outputDir)
-
-    var prev: SourceLocation.Location? = null
-
-    for (mapping in sourceLocationMappings) {
-        when (val location = mapping.sourceLocation.takeIf { it != prev } ?: continue) {
-            is SourceLocation.NoLocation -> sourceMapBuilder.addEmptyMapping(mapping.offset)
-            is SourceLocation.Location -> {
-                location.apply {
-                    // TODO resulting path goes too deep since temporary directory we compiled first is deeper than final destination.
-                    val relativePath = pathResolver.getPathRelativeToSourceRoots(File(file)).replace(Regex("^\\.\\./"), "")
-                    sourceMapBuilder.addMapping(relativePath, null, { null }, line, column, null, mapping.offset)
-                    prev = this
-                }
-            }
-        }
-    }
-
-    return sourceMapBuilder.build()
-}
-
 fun WasmCompiledModuleFragment.generateAsyncWasiWrapper(wasmFilePath: String): String = """
 import { WASI } from 'wasi';
 import { argv, env } from 'node:process';
@@ -370,7 +349,10 @@ fun writeCompilationResult(
    }
    File(dir, "$fileNameBase.mjs").writeText(result.jsWrapper)

-    if (result.sourceMap != null) {
-        File(dir, "$fileNameBase.wasm.map").writeText(result.sourceMap)
+    result.debugInformation?.sourceMapForBinary?.let {
+        File(dir, "$fileNameBase.wasm.map").writeText(it)
+    }
+    result.debugInformation?.sourceMapForText?.let {
+        File(dir, "$fileNameBase.wat.map").writeText(it)
    }
 }
@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.wasm.utils
+
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.ir.backend.js.SourceMapsInfo
+import org.jetbrains.kotlin.js.sourceMap.SourceFilePathResolver
+import org.jetbrains.kotlin.js.sourceMap.SourceMap3Builder
+import org.jetbrains.kotlin.wasm.ir.debug.DebugData
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformation
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformationGenerator
+import org.jetbrains.kotlin.wasm.ir.debug.DebugSection
+import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocation
+import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocationMapping
+import java.io.File
+
+class SourceMapGenerator(
+    baseFileName: String,
+    private val configuration: CompilerConfiguration
+) : DebugInformationGenerator {
+    // TODO: eliminate duplication for the [org.jetbrains.kotlin.backend.wasm.writeCompilationResult] logic
+    private val sourceMapFileName = "$baseFileName.map"
+    private val sourceLocationMappings = mutableListOf<SourceLocationMapping>()
+
+    override fun addSourceLocation(location: SourceLocationMapping) {
+        sourceLocationMappings.add(location)
+    }
+
+    override fun generateDebugInformation(): DebugInformation {
+        return listOf(DebugSection("sourceMappingURL", DebugData.StringData(sourceMapFileName)))
+    }
+
+    fun generate(): String? {
+        val sourceMapsInfo = SourceMapsInfo.from(configuration) ?: return null
+
+        val sourceMapBuilder =
+            SourceMap3Builder(null, { error("This should not be called for Kotlin/Wasm") }, sourceMapsInfo.sourceMapPrefix)
+
+        val pathResolver =
+            SourceFilePathResolver.create(sourceMapsInfo.sourceRoots, sourceMapsInfo.sourceMapPrefix, sourceMapsInfo.outputDir)
+
+        var prev: SourceLocation? = null
+        var prevGeneratedLine = 0
+
+        for (mapping in sourceLocationMappings) {
+            val generatedLocation = mapping.generatedLocation
+            val sourceLocation = mapping.sourceLocation.takeIf { it != prev || prevGeneratedLine != generatedLocation.line } ?: continue
+
+            require(generatedLocation.line >= prevGeneratedLine) { "The order of the mapping is wrong" }
+
+            if (prevGeneratedLine != generatedLocation.line) {
+                repeat(generatedLocation.line - prevGeneratedLine) {
+                    sourceMapBuilder.newLine()
+                }
+                prevGeneratedLine = generatedLocation.line
+            }
+
+            when (sourceLocation) {
+                is SourceLocation.NoLocation -> sourceMapBuilder.addEmptyMapping(generatedLocation.column)
+                is SourceLocation.Location -> {
+                    sourceLocation.apply {
+                        // TODO resulting path goes too deep since temporary directory we compiled first is deeper than final destination.
+                        val relativePath = pathResolver.getPathRelativeToSourceRoots(File(file)).replace(Regex("^\\.\\./"), "")
+                        sourceMapBuilder.addMapping(relativePath, null, { null }, line, column, null, generatedLocation.column)
+                        prev = this
+                    }
+                }
+            }
+
+        }
+
+        return sourceMapBuilder.build()
+    }
+}
@@ -9,9 +9,11 @@ import org.jetbrains.kotlin.wasm.ir.*
 import java.io.ByteArrayOutputStream
 import java.io.OutputStream
 import kotlinx.collections.immutable.*
-import org.jetbrains.kotlin.wasm.ir.source.location.Box
-import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocation
-import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocationMapping
+import org.jetbrains.kotlin.wasm.ir.debug.DebugData
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformation
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformationConsumer
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformationGenerator
+import org.jetbrains.kotlin.wasm.ir.source.location.*

 private object WasmBinary {
    const val MAGIC = 0x6d736100u
@@ -56,9 +58,8 @@ class WasmIrToBinary(
    val module: WasmModule,
    val moduleName: String,
    val emitNameSection: Boolean,
-    private val sourceMapFileName: String? = null,
-    private val sourceLocationMappings: MutableList<SourceLocationMapping>? = null
-) {
+    private val debugInformationGenerator: DebugInformationGenerator? = null
+) : DebugInformationConsumer {
    private var b: ByteWriter = ByteWriter.OutputStream(outputStream)

    // "Stack" of offsets waiting initialization. 
@@ -66,6 +67,17 @@ class WasmIrToBinary(
    // until we generate the whole block and generate size. So, we put them into "stack" and initialize as soon as we have all required data.
    private var offsets = persistentListOf<Box>()

+    override fun consumeDebugInformation(debugInformation: DebugInformation) {
+        debugInformation.forEach {
+            appendSection(WasmBinary.Section.CUSTOM) {
+                b.writeString(it.name)
+                when (it.data) {
+                    is DebugData.StringData -> b.writeString(it.data.value)
+                }
+            }
+        }
+    }
+
    fun appendWasmModule() {
        b.writeUInt32(WasmBinary.MAGIC)
        b.writeUInt32(WasmBinary.VERSION)
@@ -174,13 +186,7 @@ class WasmIrToBinary(
                appendTextSection(definedFunctions)
            }

-            if (sourceMapFileName != null) {
-                // Custom section with URL to sourcemap
-                appendSection(WasmBinary.Section.CUSTOM) {
-                    b.writeString("sourceMappingURL")
-                    b.writeString(sourceMapFileName)
-                }
-            }
+            debugInformationGenerator?.let { consumeDebugInformation(it.generateDebugInformation()) }
        }
    }

@@ -247,7 +253,7 @@ class WasmIrToBinary(

    private fun appendInstr(instr: WasmInstr) {
        instr.location?.let {
-            sourceLocationMappings?.add(SourceLocationMapping(offsets + Box(b.written), it))
+            debugInformationGenerator?.addSourceLocation(SourceLocationMappingToBinary(it, offsets + Box(b.written)))
        }

        val opcode = instr.operator.opcode
@@ -717,3 +723,21 @@ abstract class ByteWriter {
        override fun createTemp() = OutputStream(ByteArrayOutputStream())
    }
 }
+
+private class SourceLocationMappingToBinary(
+    override val sourceLocation: SourceLocation,
+    // Offsets in generating binary, initialized lazily. Since blocks has as a prefix variable length number encoding its size
+    // we can't calculate absolute offsets inside those blocks until we generate whole block and generate size.
+    private val offsets: List<Box>,
+) : SourceLocationMapping() {
+    override val generatedLocation: SourceLocation.Location by lazy {
+        SourceLocation.Location(
+            file = "",
+            line = 0,
+            column = offsets.sumOf {
+                assert(it.value >= 0) { "Offset must be >=0 but ${it.value}" }
+                it.value
+            }
+        )
+    }
+}
@@ -6,9 +6,15 @@
 package org.jetbrains.kotlin.wasm.ir.convertors

 import org.jetbrains.kotlin.wasm.ir.*
+import org.jetbrains.kotlin.wasm.ir.debug.DebugData
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformation
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformationConsumer
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformationGenerator
+import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocation
+import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocationMappingToText

 open class SExpressionBuilder {
-    protected val stringBuilder = StringBuilder()
+    protected val stringBuilder = StringBuilderWithLocations()
    protected var indent = 0

    protected inline fun indented(body: () -> Unit) {
@@ -45,7 +51,21 @@ open class SExpressionBuilder {
 }


-class WasmIrToText : SExpressionBuilder() {
+class WasmIrToText(
+    private val debugInformationGenerator: DebugInformationGenerator? = null
+) : SExpressionBuilder(), DebugInformationConsumer {
+    override fun consumeDebugInformation(debugInformation: DebugInformation) {
+        debugInformation.forEach {
+            newLine()
+            stringBuilder.append("(; @custom ")
+            stringBuilder.append(it.name)
+            when (it.data) {
+                is DebugData.StringData -> stringBuilder.append(" \"${it.data.value}\"")
+            }
+            stringBuilder.append(" ;)")
+        }
+    }
+
    fun appendOffset(value: UInt) {
        if (value != 0u)
            appendElement("offset=$value")
@@ -59,6 +79,15 @@ class WasmIrToText : SExpressionBuilder() {
    }

    private fun appendInstr(wasmInstr: WasmInstr) {
+        wasmInstr.location?.let {
+            debugInformationGenerator?.addSourceLocation(
+                SourceLocationMappingToText(
+                    it,
+                    SourceLocation.Location("", stringBuilder.lineNumber, stringBuilder.columnNumber),
+                )
+            )
+        }
+
        val op = wasmInstr.operator

        if (op.opcode == WASM_OP_PSEUDO_OPCODE) {
@@ -264,6 +293,7 @@ class WasmIrToText : SExpressionBuilder() {
                startFunction?.let { appendStartFunction(it) }
                data.forEach { appendData(it) }
                tags.forEach { appendTag(it) }
+                debugInformationGenerator?.let { consumeDebugInformation(it.generateDebugInformation()) }
            }
        }
    }
@@ -571,6 +601,42 @@ class WasmIrToText : SExpressionBuilder() {
    }
 }

+class StringBuilderWithLocations {
+    private val builder = StringBuilder()
+
+    var lineNumber: Int = 0
+        private set
+
+    var columnNumber: Int = -1
+        private set
+
+    fun append(char: Char) {
+        if (char == '\n') {
+            appendLine()
+        } else {
+            builder.append(char)
+        }
+    }
+
+    fun append(text: String) {
+        builder.append(text)
+
+        val lines = text.split('\n').also {
+            if (it.size > 1) columnNumber = -1
+        }
+        lineNumber += lines.size - 1
+        columnNumber += lines.last().length
+    }
+
+    fun appendLine() {
+        builder.appendLine()
+        lineNumber += 1
+        columnNumber = -1
+    }
+
+    override fun toString() = builder.toString()
+}
+

 fun Byte.toWatData() = "\\" + this.toUByte().toString(16).padStart(2, '0')
 fun ByteArray.toWatData(): String = "\"" + joinToString("") { it.toWatData() } + "\""
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.wasm.ir.debug
+
+typealias DebugInformation = List<DebugSection>
+
+class DebugSection(val name: String, val data: DebugData)
+
+sealed interface DebugData {
+    @JvmInline
+    value class StringData(val value: String) : DebugData
+}
@@ -0,0 +1,10 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.wasm.ir.debug
+
+interface DebugInformationConsumer {
+    fun consumeDebugInformation(debugInformation: DebugInformation)
+}
@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.wasm.ir.debug
+
+import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocationMapping
+
+interface DebugInformationGenerator {
+    fun addSourceLocation(location: SourceLocationMapping)
+    fun generateDebugInformation(): DebugInformation
+}
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.wasm.ir.source.location
 sealed class SourceLocation {
    object NoLocation : SourceLocation()

+    // Both line and column are zero-based
    data class Location(val file: String, val line: Int, val column: Int) : SourceLocation()

    companion object {
@@ -5,16 +5,7 @@

 package org.jetbrains.kotlin.wasm.ir.source.location

-class SourceLocationMapping(
-    // Offsets in generating binary, initialized lazily. Since blocks has as a prefix variable length number encoding its size 
-    // we can't calculate absolute offsets inside those blocks until we generate whole block and generate size.
-    private val offsets: List<Box>,
-    val sourceLocation: SourceLocation
-) {
-    val offset by lazy {
-        offsets.sumOf {
-            assert(it.value >= 0) { "Offset must be >=0 but ${it.value}" }
-            it.value
-        }
-    }
+abstract class SourceLocationMapping {
+    abstract val sourceLocation: SourceLocation
+    abstract val generatedLocation: SourceLocation.Location
 }
@@ -0,0 +1,11 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.wasm.ir.source.location
+
+class SourceLocationMappingToText(
+    override val sourceLocation: SourceLocation,
+    override val generatedLocation: SourceLocation.Location
+) : SourceLocationMapping()
@@ -147,8 +147,8 @@ class WasmBackendFacade(
            wat = newWat,
            jsUninstantiatedWrapper = jsUninstantiatedWrapper,
            jsWrapper = jsWrapper,
-            sourceMap = null,
-            wasm = newWasm
+            wasm = newWasm,
+            debugInformation = null
        )
    }
 }
@@ -178,7 +178,7 @@ class WasmDebugRunner(testServices: TestServices) : AbstractWasmArtifactsCollect
    }

    private val WasmCompilerResult.parsedSourceMaps: SourceMap
-        get() = when (val parseResult = SourceMapParser.parse(sourceMap ?: error("Expect to have source maps for stepping test"))) {
+        get() = when (val parseResult = SourceMapParser.parse(debugInformation?.sourceMapForBinary ?: error("Expect to have source maps for stepping test"))) {
            is SourceMapSuccess -> parseResult.value
            is SourceMapError -> error(parseResult.message)
        }