Change structure of ParameterInfos

timotheeandres · timotheeandres · commit da9bdfb2e289 · 2021-12-09T20:09:09.000+01:00
- Based on suggestion by Nicolas Stucki, here: https://gist.github.com/nicolasstucki/84ebcd5c2cfc9aa14abba96ae1a0e996 - Add ability to pass multiple ParameterAnnotations
diff --git a/compiler/src/dotty/tools/dotc/ast/MainProxies.scala b/compiler/src/dotty/tools/dotc/ast/MainProxies.scala
@@ -47,7 +47,7 @@ import Annotations.Annotation
   */
 object MainProxies {
   private type DefaultValueSymbols = Map[Int, Symbol]
-  private type ParameterAnnotations = Vector[Option[Annotation]]
+  private type ParameterAnnotationss = Seq[Seq[Annotation]]
 
   def mainProxies(stats: List[tpd.Tree])(using Context): List[untpd.Tree] = {
     import tpd._
@@ -69,19 +69,15 @@ object MainProxies {
       }
 
     /** Computes the list of main methods present in the code. */
-    def mainMethods(scope: Tree, stats: List[Tree]): List[(Symbol, ParameterAnnotations, DefaultValueSymbols, Option[Comment])] = stats.flatMap {
+    def mainMethods(scope: Tree, stats: List[Tree]): List[(Symbol, ParameterAnnotationss, DefaultValueSymbols, Option[Comment])] = stats.flatMap {
       case stat: DefDef =>
         val sym = stat.symbol
         sym.annotations.filter(_.matches(defn.MainAnnot)) match {
           case Nil =>
             Nil
           case _ :: Nil =>
             val paramAnnotations = stat.paramss.flatMap(_.map(
-              valdef => valdef.symbol.annotations.filter(_.matches(defn.MainAnnotParameterAnnotation)) match {
-                case Nil => None
-                case paramAnnot :: Nil => Some(paramAnnot)
-                case paramAnnot :: others => report.error(s"parameters cannot have multiple annotations", paramAnnot.tree); None
-              }
+              valdef => valdef.symbol.annotations.filter(_.matches(defn.MainAnnotParameterAnnotation))
             ))
             (sym, paramAnnotations.toVector, defaultValueSymbols(scope, sym), stat.rawComment) :: Nil
           case mainAnnot :: others =>
@@ -99,7 +95,7 @@ object MainProxies {
   }
 
   import untpd._
-  def mainProxy(mainFun: Symbol, paramAnnotations: ParameterAnnotations, defaultValueSymbols: DefaultValueSymbols, docComment: Option[Comment])(using Context): List[TypeDef] = {
+  def mainProxy(mainFun: Symbol, paramAnnotations: ParameterAnnotationss, defaultValueSymbols: DefaultValueSymbols, docComment: Option[Comment])(using Context): List[TypeDef] = {
     val mainAnnot = mainFun.getAnnotation(defn.MainAnnot).get
     def pos = mainFun.sourcePos
     val mainArgsName: TermName = nme.args
@@ -111,6 +107,11 @@ object MainProxies {
 
     inline def some(value: Tree): Tree = Apply(ref(defn.SomeClass.companionModule.termRef), value)
 
+    def unitToValue(value: Tree): Tree =
+      val anonName = nme.ANON_FUN
+      val defdef = DefDef(anonName, List(Nil), TypeTree(), value)
+      Block(defdef, Closure(Nil, Ident(anonName), EmptyTree))
+
     /**
       * Creates a list of references and definitions of arguments, the first referencing the second.
       * The goal is to create the
@@ -150,27 +151,24 @@ object MainProxies {
             /*
              * Assignations to be made after the creation of the ParameterInfos.
              * For example:
-             *   args0paramInfos.documentation = Some("my param x")
+             *   args0paramInfos.withDocumentation = Some("my param x")
              * is represented by the pair
-             *   ("documentation", some(lit("my param x")))
+             *   (defn.MainAnnotationParameterInfos_withDocumentation, some(lit("my param x")))
              */
-            var assignations: List[(String, Tree)] = Nil
+            var assignations: List[(Symbol, List[Tree])] = Nil
             for (dvSym <- defaultValueSymbols.get(n))
-              assignations = ("defaultValue" -> some(ref(dvSym.termRef))) :: assignations
-            for (annot <- paramAnnotations(n))
-              assignations = ("annotation" -> some(instanciateAnnotation(annot))) :: assignations
+              assignations = (defn.MainAnnotationParameterInfos_withDefaultValue -> List(unitToValue(ref(dvSym.termRef)))) :: assignations
             for (doc <- documentation.argDocs.get(param))
-              assignations = ("documentation" -> some(lit(doc))) :: assignations
+              assignations = (defn.MainAnnotationParameterInfos_withDocumentation -> List(lit(doc))) :: assignations
 
-            val assignationsTrees = assignations.map{
-              case (name, value) => Apply(Select(paramInfosIdent, defn.MainAnnotParameterInfos.requiredMethod(name + "_=").name), value)
-            }
+            val instanciatedAnnots = paramAnnotations(n).map(instanciateAnnotation).toList
+            if instanciatedAnnots.nonEmpty then
+              assignations = (defn.MainAnnotationParameterInfos_withAnnotations -> instanciatedAnnots) :: assignations
 
             if assignations.isEmpty then
               paramInfosTree
             else
-              val paramInfosInstance = ValDef(paramInfosName, TypeTree(), paramInfosTree)
-              Block(paramInfosInstance :: assignationsTrees, paramInfosIdent)
+              assignations.foldLeft[Tree](paramInfosTree){ case (tree, (setterSym, values)) => Apply(Select(tree, setterSym.name), values) }
           }
 
           val argDef = ValDef(
diff --git a/compiler/src/dotty/tools/dotc/core/Definitions.scala b/compiler/src/dotty/tools/dotc/core/Definitions.scala
@@ -859,6 +859,9 @@ class Definitions {
   @tu lazy val MainAnnot: ClassSymbol = requiredClass("scala.annotation.MainAnnotation")
     @tu lazy val MainAnnot_command: Symbol = MainAnnot.requiredMethod("command")
   @tu lazy val MainAnnotParameterInfos: ClassSymbol = requiredClass("scala.annotation.MainAnnotation.ParameterInfos")
+    @tu lazy val MainAnnotationParameterInfos_withDefaultValue: Symbol = MainAnnotParameterInfos.requiredMethod("withDefaultValue")
+    @tu lazy val MainAnnotationParameterInfos_withDocumentation: Symbol = MainAnnotParameterInfos.requiredMethod("withDocumentation")
+    @tu lazy val MainAnnotationParameterInfos_withAnnotations: Symbol = MainAnnotParameterInfos.requiredMethod("withAnnotations")
   @tu lazy val MainAnnotParameterAnnotation: ClassSymbol = requiredClass("scala.annotation.MainAnnotation.ParameterAnnotation")
   @tu lazy val MainAnnotCommand: ClassSymbol = requiredClass("scala.annotation.MainAnnotation.Command")
     @tu lazy val MainAnnotCommand_argGetter: Symbol = MainAnnotCommand.requiredMethod("argGetter")
diff --git a/library/src/scala/annotation/MainAnnotation.scala b/library/src/scala/annotation/MainAnnotation.scala
@@ -25,19 +25,34 @@ trait MainAnnotation extends StaticAnnotation:
 end MainAnnotation
 
 object MainAnnotation:
-  /**
-    * The information related to one of the parameters of the annotated method.
-    * @param name the name of the parameter
-    * @param typeName the name of the parameter's type
-    * @tparam T the type of the parameter
-    */
-  class ParameterInfos[T](var name: String, var typeName: String):
+  // Inspired by https://github.com/scala-js/scala-js/blob/0708917912938714d52be1426364f78a3d1fd269/linker-interface/shared/src/main/scala/org/scalajs/linker/interface/StandardConfig.scala#L23-L218
+  final class ParameterInfos[T] private (
+    /** The name of the parameter */
+    val name: String,
+    /** The name of the parameter's type */
+    val typeName: String,
     /** The docstring of the parameter. Defaults to None. */
-    var documentation: Option[String] = None
+    val documentation: Option[String],
     /** The default value that the parameter has. Defaults to None. */
-    var defaultValue: Option[T] = None
-    /** If there is one, the ParameterAnnotation associated with the parameter. Defaults to None. */
-    var annotation: Option[ParameterAnnotation] = None
+    val defaultValueOpt: Option[() => T],
+    /** The ParameterAnnotations associated with the parameter. Defaults to Seq.empty. */
+    val annotations: Seq[ParameterAnnotation],
+  ) {
+    // Main public constructor
+    def this(name: String, typeName: String) =
+      this(name, typeName, None, None, Seq.empty)
+
+    def withDefaultValue(defaultValueGetter: () => T): ParameterInfos[T] =
+      new ParameterInfos(name, typeName, documentation, Some(defaultValueGetter), annotations)
+
+    def withDocumentation(doc: String): ParameterInfos[T] =
+      new ParameterInfos(name, typeName, Some(doc), defaultValueOpt, annotations)
+
+    def withAnnotations(annots: ParameterAnnotation*): ParameterInfos[T] =
+      new ParameterInfos(name, typeName, documentation, defaultValueOpt, annots)
+
+    override def toString: String = s"$name: $typeName"
+  }
 
   /** A class representing a command to run */
   trait Command[ArgumentParser[_], MainResultType]:
diff --git a/library/src/scala/main.scala b/library/src/scala/main.scala
@@ -174,17 +174,11 @@ final class main(maxLineLength: Int) extends MainAnnotation:
             error(s"more than one value for $argName: ${multValues.mkString(", ")}")
         }
 
-      private def getAnnotationData[T](paramInfos: ParameterInfos[_], extractor: Arg => T): Option[T] =
-        paramInfos.annotation match {
-          case Some(annot: Arg) => Some(extractor(annot))
-          case _ => None
-        }
-
       private inline def getEffectiveName(paramInfos: ParameterInfos[_]): String =
-        getAnnotationData(paramInfos, _.name).filter(_.length > 0).getOrElse(paramInfos.name)
+        paramInfos.annotations.collectFirst{ case arg: Arg if arg.name.length > 0 => arg.name }.getOrElse(paramInfos.name)
 
       private inline def getShortName(paramInfos: ParameterInfos[_]): Option[Char] =
-        getAnnotationData(paramInfos, _.shortName).filterNot(_ == 0)
+        paramInfos.annotations.collectFirst{ case arg: Arg if arg.shortName != 0 => arg.shortName }
 
       private def registerArg(paramInfos: ParameterInfos[_], argKind: ArgumentKind): Unit =
         argNames += getEffectiveName(paramInfos)
@@ -198,8 +192,8 @@ final class main(maxLineLength: Int) extends MainAnnotation:
 
       override def argGetter[T](paramInfos: ParameterInfos[T])(using p: ArgumentParser[T]): () => T =
         val name = getEffectiveName(paramInfos)
-        val (defaultGetter, argumentKind) = paramInfos.defaultValue match {
-          case Some(value) => (() => () => value, ArgumentKind.OptionalArgument)
+        val (defaultGetter, argumentKind) = paramInfos.defaultValueOpt match {
+          case Some(value) => (() => value, ArgumentKind.OptionalArgument)
           case None => (() => error(s"missing argument for $name"), ArgumentKind.SimpleArgument)
         }
         registerArg(paramInfos, argumentKind)
diff --git a/tests/run/main-annotation-homemade-annot-1.scala b/tests/run/main-annotation-homemade-annot-1.scala
@@ -21,7 +21,8 @@ object Test:
 end Test
 
 class mainAwait(timeout: Int = 2) extends MainAnnotation:
-  self =>
+  import MainAnnotation._
+  import main.{Arg}
 
   private val maxLineLength = 120
 
@@ -33,8 +34,12 @@ class mainAwait(timeout: Int = 2) extends MainAnnotation:
   }
 
   override def command(args: Array[String], commandName: String, docComment: String) =
-    new MainAnnotation.Command[ArgumentParser, MainResultType]:
+    new Command[ArgumentParser, MainResultType]:
+      private val argMarker = "--"
+      private val shortArgMarker = "-"
+
       private var argNames = new mutable.ArrayBuffer[String]
+      private var argShortNames = new mutable.ArrayBuffer[Option[Char]]
       private var argTypes = new mutable.ArrayBuffer[String]
       private var argDocs = new mutable.ArrayBuffer[String]
       private var argKinds = new mutable.ArrayBuffer[ArgumentKind]
@@ -53,23 +58,34 @@ class mainAwait(timeout: Int = 2) extends MainAnnotation:
       private def argAt(idx: Int): Option[String] =
         if idx < args.length then Some(args(idx)) else None
 
+      private def isArgNameAt(idx: Int): Boolean =
+        val arg = args(argIdx)
+        val isFullName = arg.startsWith(argMarker)
+        val isShortName = arg.startsWith(shortArgMarker) && arg.length == 2 && shortNameIsValid(arg(1))
+
+        isFullName || isShortName
+
       private def nextPositionalArg(): Option[String] =
-        while argIdx < args.length && args(argIdx).startsWith("--") do argIdx += 2
+        while argIdx < args.length && isArgNameAt(argIdx) do argIdx += 2
         val result = argAt(argIdx)
         argIdx += 1
         result
 
+      private def shortNameIsValid(shortName: Char): Boolean =
+        shortName == 0 || shortName.isLetter
+
       private def convert[T](argName: String, arg: String, p: ArgumentParser[T]): () => T =
         p.fromStringOption(arg) match
           case Some(t) => () => t
           case None => error(s"invalid argument for $argName: $arg")
 
       private def argUsage(pos: Int): String =
         val name = argNames(pos)
+        val namePrint = argShortNames(pos).map(short => s"[$shortArgMarker$short | $argMarker$name]").getOrElse(s"[$argMarker$name]")
 
         argKinds(pos) match {
-          case ArgumentKind.SimpleArgument => s"[--$name] <${argTypes(pos)}>"
-          case ArgumentKind.OptionalArgument => s"[[--$name] <${argTypes(pos)}>]"
+          case ArgumentKind.SimpleArgument => s"$namePrint <${argTypes(pos)}>"
+          case ArgumentKind.OptionalArgument => s"[$namePrint <${argTypes(pos)}>]"
           case ArgumentKind.VarArgument => s"[<${argTypes(pos)}> [<${argTypes(pos)}> [...]]]"
         }
 
@@ -137,19 +153,18 @@ class mainAwait(timeout: Int = 2) extends MainAnnotation:
             println(argDoc)
         }
 
-      private def indicesOfArg(argName: String): Seq[Int] =
-        def allIndicesOf(s: String): Seq[Int] =
-          def recurse(s: String, from: Int): Seq[Int] =
-            val i = args.indexOf(s, from)
-            if i < 0 then Seq() else i +: recurse(s, i + 1)
-
-          recurse(s, 0)
+      private def indicesOfArg(argName: String, shortArgName: Option[Char]): Seq[Int] =
+        def allIndicesOf(s: String, from: Int): Seq[Int] =
+          val i = args.indexOf(s, from)
+          if i < 0 then Seq() else i +: allIndicesOf(s, i + 1)
 
-        val indices = allIndicesOf(s"--$argName")
-        indices.filter(_ >= 0)
+        val indices = allIndicesOf(s"$argMarker$argName", 0)
+        val indicesShort = shortArgName.map(shortName => allIndicesOf(s"$shortArgMarker$shortName", 0)).getOrElse(Seq())
+        (indices ++: indicesShort).filter(_ >= 0)
 
-      private def getArgGetter[T](argName: String, getDefaultGetter: () => () => T)(using p: ArgumentParser[T]): () => T =
-        indicesOfArg(argName) match {
+      private def getArgGetter[T](paramInfos: ParameterInfos[_], getDefaultGetter: () => () => T)(using p: ArgumentParser[T]): () => T =
+        val argName = getEffectiveName(paramInfos)
+        indicesOfArg(argName, getShortName(paramInfos)) match {
           case s @ (Seq() | Seq(_)) =>
             val argOpt = s.headOption.map(idx => argAt(idx + 1)).getOrElse(nextPositionalArg())
             argOpt match {
@@ -161,48 +176,64 @@ class mainAwait(timeout: Int = 2) extends MainAnnotation:
             error(s"more than one value for $argName: ${multValues.mkString(", ")}")
         }
 
-      private def registerArg(paramInfos: MainAnnotation.ParameterInfos[_], argKind: ArgumentKind): Unit =
-        argNames += paramInfos.name
+      private inline def getEffectiveName(paramInfos: ParameterInfos[_]): String =
+        paramInfos.annotations.collectFirst{ case arg: Arg if arg.name.length > 0 => arg.name }.getOrElse(paramInfos.name)
+
+      private inline def getShortName(paramInfos: ParameterInfos[_]): Option[Char] =
+        paramInfos.annotations.collectFirst{ case arg: Arg if arg.shortName != 0 => arg.shortName }
+
+      private def registerArg(paramInfos: ParameterInfos[_], argKind: ArgumentKind): Unit =
+        argNames += getEffectiveName(paramInfos)
         argTypes += paramInfos.typeName
         argDocs += paramInfos.documentation.getOrElse("")
         argKinds += argKind
 
-      override def argGetter[T](paramInfos: MainAnnotation.ParameterInfos[T])(using p: ArgumentParser[T]): () => T =
-        val name = paramInfos.name
-        val (defaultGetter, argumentKind) = paramInfos.defaultValue match {
-          case Some(value) => (() => () => value, ArgumentKind.OptionalArgument)
+        val shortName = getShortName(paramInfos)
+        shortName.foreach(c => if !shortNameIsValid(c) then throw IllegalArgumentException(s"Invalid short name: $shortArgMarker$c"))
+        argShortNames += shortName
+
+      override def argGetter[T](paramInfos: ParameterInfos[T])(using p: ArgumentParser[T]): () => T =
+        val name = getEffectiveName(paramInfos)
+        val (defaultGetter, argumentKind) = paramInfos.defaultValueOpt match {
+          case Some(value) => (() => value, ArgumentKind.OptionalArgument)
           case None => (() => error(s"missing argument for $name"), ArgumentKind.SimpleArgument)
         }
         registerArg(paramInfos, argumentKind)
-        getArgGetter(name, defaultGetter)
+        getArgGetter(paramInfos, defaultGetter)
 
-      override def varargGetter[T](paramInfos: MainAnnotation.ParameterInfos[T])(using p: ArgumentParser[T]): () => Seq[T] =
+      override def varargGetter[T](paramInfos: ParameterInfos[T])(using p: ArgumentParser[T]): () => Seq[T] =
         registerArg(paramInfos, ArgumentKind.VarArgument)
         def remainingArgGetters(): List[() => T] = nextPositionalArg() match
-          case Some(arg) => convert(paramInfos.name, arg, p) :: remainingArgGetters()
+          case Some(arg) => convert(getEffectiveName(paramInfos), arg, p) :: remainingArgGetters()
           case None => Nil
         val getters = remainingArgGetters()
         () => getters.map(_())
 
       override def run(f: => MainResultType): Unit =
+        def checkShortNamesUnique(): Unit =
+          val shortNameToIndices = argShortNames.collect{ case Some(short) => short }.zipWithIndex.groupBy(_._1).view.mapValues(_.map(_._2))
+          for ((shortName, indices) <- shortNameToIndices if indices.length > 1)
+            error(s"$shortName is used as short name for multiple parameters: ${indices.map(idx => argNames(idx)).mkString(", ")}")
+
         def flagUnused(): Unit = nextPositionalArg() match
           case Some(arg) =>
             error(s"unused argument: $arg")
             flagUnused()
           case None =>
             for
               arg <- args
-              if arg.startsWith("--") && !argNames.contains(arg.drop(2))
+              if arg.startsWith(argMarker) && !argNames.contains(arg.drop(2))
             do
               error(s"unknown argument name: $arg")
         end flagUnused
 
-        if args.contains("--help") then
+        if args.contains(s"${argMarker}help") then
           usage()
           println()
           explain()
         else
           flagUnused()
+          checkShortNamesUnique()
           if errors.nonEmpty then
             for msg <- errors do println(s"Error: $msg")
             usage()
diff --git a/tests/run/main-annotation-homemade-annot-2.scala b/tests/run/main-annotation-homemade-annot-2.scala
