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

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt

@@ -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,
+            sourceMapGeneratorForBinary
-            sourceLocationMappings
         )
 
     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) {
+    result.debugInformation?.sourceMapForBinary?.let {
-        File(dir, "$fileNameBase.wasm.map").writeText(result.sourceMap)
+        File(dir, "$fileNameBase.wasm.map").writeText(it)
+    }
+    result.debugInformation?.sourceMapForText?.let {
+        File(dir, "$fileNameBase.wat.map").writeText(it)
     }
 }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/utils/SourceMapGenerator.kt

@@ -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()
+    }
+}

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/convertors/WasmIrToBinary.kt

@@ -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.debug.DebugData
-import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocation
+import org.jetbrains.kotlin.wasm.ir.debug.DebugInformation
-import org.jetbrains.kotlin.wasm.ir.source.location.SourceLocationMapping
+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 debugInformationGenerator: DebugInformationGenerator? = null
-    private val sourceLocationMappings: MutableList<SourceLocationMapping>? = 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) {
+            debugInformationGenerator?.let { consumeDebugInformation(it.generateDebugInformation()) }
-                // Custom section with URL to sourcemap
-                appendSection(WasmBinary.Section.CUSTOM) {
-                    b.writeString("sourceMappingURL")
-                    b.writeString(sourceMapFileName)
-                }
-            }
         }
     }
 
@@ -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
+            }
+        )
+    }
+}

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/convertors/WasmIrToText.kt

@@ -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() } + "\""

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/debug/DebugInfo.kt

@@ -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
+}

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/debug/DebugInformationConsumer.kt

@@ -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)
+}

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/debug/DebugInformationGenerator.kt

@@ -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
+}

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/source/location/SourceLocation.kt

@@ -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 {

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/source/location/SourceLocationMapping.kt

@@ -5,16 +5,7 @@
 
 package org.jetbrains.kotlin.wasm.ir.source.location
 
-class SourceLocationMapping(
+abstract class SourceLocationMapping {
-    // Offsets in generating binary, initialized lazily. Since blocks has as a prefix variable length number encoding its size 
+    abstract val sourceLocation: SourceLocation
-    // we can't calculate absolute offsets inside those blocks until we generate whole block and generate size.
+    abstract val generatedLocation: SourceLocation.Location
-    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
-        }
-    }
 }

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/source/location/SourceLocationMappingToText.kt

@@ -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()

wasm/wasm.tests/test/org/jetbrains/kotlin/wasm/test/converters/WasmBackendFacade.kt

@@ -147,8 +147,8 @@ class WasmBackendFacade(
             wat = newWat,
             jsUninstantiatedWrapper = jsUninstantiatedWrapper,
             jsWrapper = jsWrapper,
-            sourceMap = null,
+            wasm = newWasm,
-            wasm = newWasm
+            debugInformation = null
         )
     }
 }

wasm/wasm.tests/test/org/jetbrains/kotlin/wasm/test/handlers/WasmDebugRunner.kt

@@ -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)
         }