[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,
-            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)
     }
 }

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