Give mill the ability to re-write ESModule imports at link time

build.sc

@@ -78,6 +78,7 @@ object Deps {
     val scalajsEnvSelenium = ivy"org.scala-js::scalajs-env-selenium:1.1.1"
     val scalajsSbtTestAdapter = ivy"org.scala-js::scalajs-sbt-test-adapter:${scalaJsVersion}"
     val scalajsLinker = ivy"org.scala-js::scalajs-linker:${scalaJsVersion}"
+    val scalajsImportMap = ivy"com.armanbilge::scalajs-importmap:0.1.1"
   }
 
   object Scalanative_0_4 {
@@ -191,6 +192,7 @@ object Deps {
   val jarjarabrams = ivy"com.eed3si9n.jarjarabrams::jarjar-abrams-core:1.14.0"
   val requests = ivy"com.lihaoyi::requests:0.8.2"
 
+
   /** Used to manage transitive versions. */
   val transitiveDeps = Seq(
     ivy"org.apache.ant:ant:1.10.14",
@@ -796,7 +798,8 @@ object scalajslib extends MillStableScalaModule with BuildInfo {
         formatDep(Deps.Scalajs_1.scalajsEnvExoegoJsdomNodejs)
       ),
       BuildInfo.Value("scalajsEnvPhantomJs", formatDep(Deps.Scalajs_1.scalajsEnvPhantomjs)),
-      BuildInfo.Value("scalajsEnvSelenium", formatDep(Deps.Scalajs_1.scalajsEnvSelenium))
+      BuildInfo.Value("scalajsEnvSelenium", formatDep(Deps.Scalajs_1.scalajsEnvSelenium)),
+      BuildInfo.Value("scalajsImportMap", formatDep(Deps.Scalajs_1.scalajsImportMap))
     )
   }
 
@@ -817,7 +820,8 @@ object scalajslib extends MillStableScalaModule with BuildInfo {
       Deps.Scalajs_1.scalajsEnvJsdomNodejs,
       Deps.Scalajs_1.scalajsEnvExoegoJsdomNodejs,
       Deps.Scalajs_1.scalajsEnvPhantomjs,
-      Deps.Scalajs_1.scalajsEnvSelenium
+      Deps.Scalajs_1.scalajsEnvSelenium,
+      Deps.Scalajs_1.scalajsImportMap
     )
   }
 }

scalajslib/src/mill/scalajslib/ScalaJSModule.scala

@@ -75,6 +75,12 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
     val commonDeps = Seq(
       ivy"org.scala-js::scalajs-sbt-test-adapter:${scalaJSVersion()}"
     )
+    val scalajsImportMapDeps = scalaJSVersion() match {
+      case s"1.$n.$_" if n.toIntOption.exists(_ >= 16) && scalaJSImportMap().nonEmpty =>
+        Seq(ivy"${ScalaJSBuildInfo.scalajsImportMap}")
+      case _ => Seq.empty[Dep]
+    }
+
     val envDeps = scalaJSBinaryVersion() match {
       case "0.6" =>
         Seq(
@@ -89,7 +95,7 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
     // we need to use the scala-library of the currently running mill
     resolveDependencies(
       repositoriesTask(),
-      (commonDeps.iterator ++ envDeps)
+      (commonDeps.iterator ++ envDeps ++ scalajsImportMapDeps)
         .map(Lib.depToBoundDep(_, mill.main.BuildInfo.scalaVersion, "")),
       ctx = Some(T.log)
     )
@@ -130,7 +136,8 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
       esFeatures = esFeatures(),
       moduleSplitStyle = moduleSplitStyle(),
       outputPatterns = scalaJSOutputPatterns(),
-      minify = scalaJSMinify()
+      minify = scalaJSMinify(),
+      importMap = scalaJSImportMap()
     )
   }
 
@@ -172,7 +179,8 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
       esFeatures: ESFeatures,
       moduleSplitStyle: ModuleSplitStyle,
       outputPatterns: OutputPatterns,
-      minify: Boolean
+      minify: Boolean,
+      importMap: Seq[ESModuleImportMapping]
   )(implicit ctx: mill.api.Ctx): Result[Report] = {
     val outputPath = ctx.dest
 
@@ -192,7 +200,8 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
       esFeatures = esFeatures,
       moduleSplitStyle = moduleSplitStyle,
       outputPatterns = outputPatterns,
-      minify = minify
+      minify = minify,
+      importMap = importMap
     )
   }
 
@@ -266,6 +275,10 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
 
   def scalaJSOptimizer: Target[Boolean] = T { true }
 
+  def scalaJSImportMap: Target[Seq[ESModuleImportMapping]] = T {
+    Seq.empty[ESModuleImportMapping]
+  }
+
   /** Whether to emit a source map. */
   def scalaJSSourceMap: Target[Boolean] = T { true }
 
@@ -346,7 +359,8 @@ trait TestScalaJSModule extends ScalaJSModule with TestModule {
       esFeatures = esFeatures(),
       moduleSplitStyle = moduleSplitStyle(),
       outputPatterns = scalaJSOutputPatterns(),
-      minify = scalaJSMinify()
+      minify = scalaJSMinify(),
+      importMap = scalaJSImportMap()
     )
   }
 

scalajslib/src/mill/scalajslib/api/ScalaJSApi.scala

@@ -270,3 +270,11 @@ object OutputPatterns {
 
   implicit val rw: RW[OutputPatterns] = macroRW[OutputPatterns]
 }
+
+sealed trait ESModuleImportMapping
+object ESModuleImportMapping {
+  case class Prefix(prefix: String, replacement: String) extends ESModuleImportMapping
+
+  implicit def rwPrefix: RW[Prefix] = macroRW
+  implicit def rw: RW[ESModuleImportMapping] = macroRW
+}

scalajslib/src/mill/scalajslib/worker/ScalaJSWorker.scala

@@ -147,6 +147,13 @@ private[scalajslib] class ScalaJSWorker extends AutoCloseable {
     )
   }
 
+  private def toWorkerApi(importMap: api.ESModuleImportMapping): workerApi.ESModuleImportMapping = {
+    importMap match {
+      case api.ESModuleImportMapping.Prefix(prefix, replacement) =>
+        workerApi.ESModuleImportMapping.Prefix(prefix, replacement)
+    }
+  }
+
   def link(
       toolsClasspath: Agg[mill.PathRef],
       runClasspath: Agg[mill.PathRef],
@@ -161,7 +168,8 @@ private[scalajslib] class ScalaJSWorker extends AutoCloseable {
       esFeatures: api.ESFeatures,
       moduleSplitStyle: api.ModuleSplitStyle,
       outputPatterns: api.OutputPatterns,
-      minify: Boolean
+      minify: Boolean,
+      importMap: Seq[api.ESModuleImportMapping]
   )(implicit ctx: Ctx.Home): Result[api.Report] = {
     bridge(toolsClasspath).link(
       runClasspath = runClasspath.iterator.map(_.path.toNIO).toSeq,
@@ -176,7 +184,8 @@ private[scalajslib] class ScalaJSWorker extends AutoCloseable {
       esFeatures = toWorkerApi(esFeatures),
       moduleSplitStyle = toWorkerApi(moduleSplitStyle),
       outputPatterns = toWorkerApi(outputPatterns),
-      minify = minify
+      minify = minify,
+      importMap = importMap.map(toWorkerApi)
     ) match {
       case Right(report) => Result.Success(fromWorkerApi(report))
       case Left(message) => Result.Failure(message)

scalajslib/test/resources/esModuleRemap/src/app/App.scala

@@ -0,0 +1,16 @@
+package app
+
+import scala.scalajs.js
+import scala.scalajs.js.annotation._ 
+
+object App {  
+  def main(args: Array[String]): Unit = {
+    println(linspace(-10.0, 10.0, 10))
+  }
+}
+
+@js.native
+@JSImport("@stdlib/linspace", JSImport.Default)
+object linspace extends js.Object {
+  def apply(start: Double, stop: Double, num: Int): Any = js.native
+}

scalajslib/test/src/mill/scalajslib/RemapEsModuleTests.scala

@@ -0,0 +1,78 @@
+package mill.scalajslib
+
+import mill.api.Result
+import mill.define.Discover
+import mill.util.{TestEvaluator, TestUtil}
+import utest._
+import mill.define.Target
+import mill.scalajslib.api._
+
+object EsModuleRemapTests extends TestSuite {
+  val workspacePath = TestUtil.getOutPathStatic() / "esModuleRemap"
+
+  val remapTo = "https://cdn.jsdelivr.net/gh/stdlib-js/array-base-linspace@esm/index.mjs"
+
+  object EsModuleRemap extends TestUtil.BaseModule {
+
+    object sourceMapModule extends ScalaJSModule {
+      override def millSourcePath = workspacePath
+      override def scalaVersion = sys.props.getOrElse("TEST_SCALA_2_13_VERSION", ???)
+      override def scalaJSVersion = "1.16.0"
+      override def scalaJSSourceMap = false
+      override def moduleKind = ModuleKind.ESModule
+
+      override def scalaJSImportMap: Target[Seq[ESModuleImportMapping]] = Seq(
+        ESModuleImportMapping.Prefix("@stdlib/linspace", remapTo)
+      )
+    }
+
+    object OldJsModule extends ScalaJSModule {
+      override def millSourcePath = workspacePath
+      override def scalaVersion = sys.props.getOrElse("TEST_SCALA_2_13_VERSION", ???)
+      override def scalaJSVersion = "1.15.0"
+      override def scalaJSSourceMap = false
+      override def moduleKind = ModuleKind.ESModule
+
+      override def scalaJSImportMap: Target[Seq[ESModuleImportMapping]] = Seq(
+        ESModuleImportMapping.Prefix("@stdlib/linspace", remapTo)
+      )
+    }
+
+    override lazy val millDiscover = Discover[this.type]
+  }
+
+  val millSourcePath = os.pwd / "scalajslib" / "test" / "resources" / "esModuleRemap"
+
+  val evaluator = TestEvaluator.static(EsModuleRemap)
+
+  val tests: Tests = Tests {
+    prepareWorkspace()
+
+    test("should remap the esmodule") {
+      val Right((report, _)) =
+        evaluator(EsModuleRemap.sourceMapModule.fastLinkJS)
+      val publicModules = report.publicModules.toSeq
+      assert(publicModules.length == 1)
+      val main = publicModules.head
+      assert(main.jsFileName == "main.js")
+      val mainPath = report.dest.path / "main.js"
+      assert(os.exists(mainPath))
+      val rawJs = os.read.lines(mainPath)
+      assert(rawJs(1).contains(remapTo))
+    }
+
+    test("should throw for older scalaJS versions") {
+      val Left(Result.Exception(ex, _)) = evaluator(EsModuleRemap.OldJsModule.fastLinkJS)
+      val error = ex.getMessage
+      assert(error == "scalaJSImportMap is not supported with Scala.js < 1.16.")
+    }
+
+  }
+
+  def prepareWorkspace(): Unit = {
+    os.remove.all(workspacePath)
+    os.makeDir.all(workspacePath / os.up)
+    os.copy(millSourcePath, workspacePath)
+  }
+
+}

scalajslib/worker-api/src/mill/scalajslib/worker/api/ScalaJSWorkerApi.scala

@@ -17,7 +17,8 @@ private[scalajslib] trait ScalaJSWorkerApi {
       esFeatures: ESFeatures,
       moduleSplitStyle: ModuleSplitStyle,
       outputPatterns: OutputPatterns,
-      minify: Boolean
+      minify: Boolean,
+      importMap: Seq[ESModuleImportMapping]
   ): Either[String, Report]
 
   def run(config: JsEnvConfig, report: Report): Unit
@@ -122,3 +123,8 @@ private[scalajslib] final case class OutputPatterns(
     jsFileURI: String,
     sourceMapURI: String
 )
+
+private[scalajslib] sealed trait ESModuleImportMapping
+private[scalajslib] object ESModuleImportMapping {
+  case class Prefix(prefix: String, replacement: String) extends ESModuleImportMapping
+}

scalajslib/worker/1/src/mill/scalajslib/worker/ScalaJSWorkerImpl.scala

@@ -25,6 +25,8 @@ import org.scalajs.testing.adapter.{TestAdapterInitializer => TAI}
 import scala.collection.mutable
 import scala.ref.SoftReference
 
+import com.armanbilge.sjsimportmap.ImportMappedIRFile
+
 class ScalaJSWorkerImpl extends ScalaJSWorkerApi {
   private case class LinkerInput(
       isFullLinkJS: Boolean,
@@ -169,7 +171,8 @@ class ScalaJSWorkerImpl extends ScalaJSWorkerApi {
       esFeatures: ESFeatures,
       moduleSplitStyle: ModuleSplitStyle,
       outputPatterns: OutputPatterns,
-      minify: Boolean
+      minify: Boolean,
+      importMap: Seq[ESModuleImportMapping]
   ): Either[String, Report] = {
     // On Scala.js 1.2- we want to use the legacy mode either way since
     // the new mode is not supported and in tests we always use legacy = false
@@ -202,7 +205,24 @@ class ScalaJSWorkerImpl extends ScalaJSWorkerApi {
 
     val resultFuture = (for {
       (irContainers, _) <- irContainersAndPathsFuture
-      irFiles <- irFileCacheCache.cached(irContainers)
+      irFiles0 <- irFileCacheCache.cached(irContainers)
+      irFiles = if (importMap.isEmpty) {
+        irFiles0
+      } else {
+        if (!minorIsGreaterThanOrEqual(16)) {
+          throw new Exception("scalaJSImportMap is not supported with Scala.js < 1.16.")
+        }
+        val remapFunction = (rawImport: String) => {
+          importMap
+            .collectFirst {
+              case ESModuleImportMapping.Prefix(prefix, replacement)
+                  if rawImport.startsWith(prefix) =>
+                s"$replacement${rawImport.stripPrefix(prefix)}"
+            }
+            .getOrElse(rawImport)
+        }
+        irFiles0.map { ImportMappedIRFile.fromIRFile(_)(remapFunction) }
+      }
       report <-
         if (useLegacy) {
           val jsFileName = "out.js"