From 53233c487e2bdeb0bc14501d9719040753ec5d08 Mon Sep 17 00:00:00 2001
From: Sergej Jaskiewicz <sergej.jaskiewicz@jetbrains.com>
Date: Wed, 25 May 2022 23:36:17 +0300
Subject: [PATCH] [JS] Factor out source map path rewriting

---
 .../jetbrains/kotlin/cli/js/dce/K2JSDce.kt    | 36 ++++++---------
 .../kotlin/js/parser/sourcemaps/SourceMap.kt  | 46 +++++++++++++++++--
 2 files changed, 56 insertions(+), 26 deletions(-)

diff --git a/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/dce/K2JSDce.kt b/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/dce/K2JSDce.kt
index 269ebeb322b..3ae11a416d6 100644
--- a/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/dce/K2JSDce.kt
+++ b/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/dce/K2JSDce.kt
@@ -124,37 +124,29 @@ class K2JSDce : CLITool<K2JSDceArguments>() {
     }
 
     private fun mapSourcePaths(inputFile: File, targetFile: File): Boolean {
-        val json = try {
-            parseJson(inputFile)
-        } catch (e: JsonSyntaxException) {
-            return false
-        }
-
-        val sourcesArray = (json as? JsonObject)?.properties?.get("sources") as? JsonArray ?: return false
-        val sources = sourcesArray.elements.map {
-            (it as? JsonString)?.value ?: return false
-        }
-
+        val jsonOut = StringWriter()
         val pathCalculator = RelativePathCalculator(targetFile.parentFile)
-        val mappedSources = sources.map {
-            val result = pathCalculator.calculateRelativePathTo(File(inputFile.parentFile, it))
-            if (result != null) {
-                if (File(targetFile.parentFile, result).exists()) {
-                    result
+        val sourceMapChanged = try {
+            SourceMap.mapSources(inputFile.readText(), jsonOut) {
+                val result = pathCalculator.calculateRelativePathTo(File(inputFile.parentFile, it))
+                if (result != null) {
+                    if (File(targetFile.parentFile, result).exists()) {
+                        result
+                    } else {
+                        it
+                    }
                 } else {
                     it
                 }
-            } else {
-                it
             }
+        } catch (e: SourceMapSourceReplacementException) {
+            false
         }
 
-        if (mappedSources == sources) return false
-
-        json.properties["sources"] = JsonArray(*mappedSources.map { JsonString(it) }.toTypedArray())
+        if (!sourceMapChanged) return false
 
         targetFile.parentFile.mkdirs()
-        OutputStreamWriter(FileOutputStream(targetFile), "UTF-8").use { it.write(json.toString()) }
+        OutputStreamWriter(FileOutputStream(targetFile), "UTF-8").use { it.write(jsonOut.toString()) }
 
         return true
     }
diff --git a/js/js.parser/src/org/jetbrains/kotlin/js/parser/sourcemaps/SourceMap.kt b/js/js.parser/src/org/jetbrains/kotlin/js/parser/sourcemaps/SourceMap.kt
index 95c06fd1dab..33ac094893c 100644
--- a/js/js.parser/src/org/jetbrains/kotlin/js/parser/sourcemaps/SourceMap.kt
+++ b/js/js.parser/src/org/jetbrains/kotlin/js/parser/sourcemaps/SourceMap.kt
@@ -16,9 +16,7 @@
 
 package org.jetbrains.kotlin.js.parser.sourcemaps
 
-import java.io.File
-import java.io.PrintStream
-import java.io.Reader
+import java.io.*
 
 class SourceMap(val sourceContentResolver: (String) -> Reader?) {
     val groups = mutableListOf<SourceMapGroup>()
@@ -49,8 +47,48 @@ class SourceMap(val sourceContentResolver: (String) -> Reader?) {
             writer.println()
         }
     }
+
+    companion object {
+        @Throws(IOException::class, SourceMapSourceReplacementException::class)
+        fun replaceSources(sourceMapFile: File, mapping: (String) -> String): Boolean {
+            val content = sourceMapFile.readText()
+            return sourceMapFile.writer().buffered().use {
+                mapSources(content, it, mapping)
+            }
+        }
+
+        @Throws(IOException::class, SourceMapSourceReplacementException::class)
+        fun mapSources(content: String, output: Writer, mapping: (String) -> String): Boolean {
+            val json = try {
+                parseJson(content)
+            } catch (e: JsonSyntaxException) {
+                throw SourceMapSourceReplacementException(cause = e)
+            }
+            val jsonObject = json as? JsonObject ?: throw SourceMapSourceReplacementException("Top-level object expected")
+            val sources = jsonObject.properties["sources"]
+            if (sources != null) {
+                val sourcesArray =
+                    sources as? JsonArray ?: throw SourceMapSourceReplacementException("'sources' property is not of array type")
+                var changed = false
+                val fixedSources = sourcesArray.elements.mapTo(mutableListOf<JsonNode>()) {
+                    val sourcePath = it as? JsonString ?: throw SourceMapSourceReplacementException("'sources' array must contain strings")
+                    val replacedPath = mapping(sourcePath.value)
+                    if (!changed && replacedPath != sourcePath.value) {
+                        changed = true
+                    }
+                    JsonString(replacedPath)
+                }
+                if (!changed) return false
+                jsonObject.properties["sources"] = JsonArray(fixedSources)
+            }
+            jsonObject.write(output)
+            return true
+        }
+    }
 }
 
+class SourceMapSourceReplacementException(message: String? = null, cause: Throwable? = null) : Exception(message, cause)
+
 data class SourceMapSegment(
     val generatedColumnNumber: Int,
     val sourceFileName: String?,
@@ -66,4 +104,4 @@ sealed class SourceMapParseResult
 
 class SourceMapSuccess(val value: SourceMap) : SourceMapParseResult()
 
-class SourceMapError(val message: String) : SourceMapParseResult()
\ No newline at end of file
+class SourceMapError(val message: String) : SourceMapParseResult()