Refactor argument preprocessing logic

Author: nicolasstucki

library/src/scala/annotation/newMain.scala

@@ -96,89 +96,17 @@ final class newMain extends MainAnnotation[FromString, Any]:
   private def shortNameWithMarker(name: String): String = shortArgMarker + name
 
   def command(info: Info, args: Seq[String]): Option[Seq[String]] =
-    checkNames(info)
-
     val canonicalNames = CanonicalNames(info)
-
-    var hasError: Boolean = false
-    def error(msg: String): Unit = {
-      hasError = true
-      println(s"Error: $msg")
-    }
-
+    canonicalNames.checkNames()
     if Help.hasHelpArg(canonicalNames, args) then
-      val help = new Help(info)
-      help.printUsage()
-      println()
-      help.printExplain()
+      Help.printUsage(info)
+      Help.printExplain(info)
       None
     else
-      val (positionalArgs, byNameArgs) = {
-        def recurse(remainingArgs: Seq[String], pa: mutable.Queue[String], bna: Seq[(String, String)]): (mutable.Queue[String], Seq[(String, String)]) =
-          remainingArgs match {
-            case Seq() =>
-              (pa, bna)
-            case (argName @ longArgRegex()) +: argValue +: rest =>
-              canonicalNames.getName(argName) match
-                case Some(canonicalName) => recurse(rest, pa, bna :+ (canonicalName -> argValue))
-                case None =>
-                  error(s"unknown argument name: $argName")
-                  recurse(rest, pa, bna)
-            case (argName @ shortArgRegex()) +: argValue +: rest =>
-              canonicalNames.getShortName(argName) match
-                case Some(canonicalName) => recurse(rest, pa, bna :+ (canonicalName -> argValue))
-                case None =>
-                  error(s"unknown argument name: $argName")
-                  recurse(rest, pa, bna)
-            case arg +: rest =>
-              recurse(rest, pa :+ arg, bna)
-          }
-
-        val (pa, bna) = recurse(args.toSeq, mutable.Queue.empty, Vector())
-        val nameToArgValues: Map[String, Seq[String]] = if bna.isEmpty then Map.empty else bna.groupMapReduce(_._1)(p => List(p._2))(_ ++ _)
-        (pa, nameToArgValues)
-      }
-
-      val argStrings: Seq[Seq[String]] =
-        for paramInfo <- info.parameters yield {
-          if (paramInfo.isVarargs) {
-            val byNameGetters = byNameArgs.getOrElse(paramInfo.name, Seq())
-            val positionalGetters = positionalArgs.removeAll()
-            // First take arguments passed by name, then those passed by position
-            byNameGetters ++ positionalGetters
-          } else {
-            byNameArgs.get(paramInfo.name) match
-            case Some(Nil) =>
-              throw AssertionError(s"${paramInfo.name} present in byNameArgs, but it has no argument value")
-            case Some(argValues) =>
-              if argValues.length > 1 then
-                // Do not accept multiple values
-                // Remove this test to take last given argument
-                error(s"more than one value for ${paramInfo.name}: ${argValues.mkString(", ")}")
-                Nil
-              else
-                List(argValues.last)
-            case None =>
-              if positionalArgs.length > 0 then
-                List(positionalArgs.dequeue())
-              else if paramInfo.hasDefault then
-                List("")
-              else
-                error(s"missing argument for ${paramInfo.name}")
-                Nil
-          }
-        }
-
-      // Handle unused and invalid args
-      for (remainingArg <- positionalArgs) error(s"unused argument: $remainingArg")
-
-      if hasError then
-        val help = new Help(info)
-        help.printUsage()
+      preProcessArgs(info, canonicalNames, args).orElse {
+        Help.printUsage(info)
         None
-      else
-        Some(argStrings.flatten)
-    end if
+      }
   end command
 
   def argGetter[T](param: Parameter, arg: String, defaultArgument: Option[() => T])(using p: FromString[T]): () => T = {
@@ -197,6 +125,71 @@ final class newMain extends MainAnnotation[FromString, Any]:
   def run(execProgram: () => Any): Unit =
     if !hasParseErrors then execProgram()
 
+  private def preProcessArgs(info: Info, canonicalNames: CanonicalNames, args: Seq[String]): Option[Seq[String]] =
+    var hasError: Boolean = false
+    def error(msg: String): Unit = {
+      hasError = true
+      println(s"Error: $msg")
+    }
+
+    val (positionalArgs, byNameArgsMap) =
+      val positionalArgs = List.newBuilder[String]
+      val byNameArgs = List.newBuilder[(String, String)]
+      var i = 0
+      while i < args.length do
+        args(i) match
+          case name @ (longArgRegex() | shortArgRegex()) =>
+            if i == args.length - 1 then // last argument -x ot --xyz
+              error(s"missing argument for ${name}")
+            else args(i + 1) match
+              case longArgRegex() | shortArgRegex() =>
+                error(s"missing argument for ${name}")
+              case value =>
+                canonicalNames.getName(name) match
+                  case Some(canonicalName) =>
+                    byNameArgs += ((canonicalName, value))
+                  case None =>
+                    error(s"unknown argument name: $name")
+                i += 1 // consume `value`
+          case value =>
+            positionalArgs += value
+        i += 1
+      end while
+      (positionalArgs.result(), byNameArgs.result().groupMap(_._1)(_._2))
+
+    // List of arguments in the order they should be passed to the main function
+    val orderedArgs: List[String] =
+      def rec(params: List[Parameter], acc: List[String], remainingArgs: List[String]): List[String] =
+        params match
+          case Nil =>
+            for (remainingArg <- remainingArgs) error(s"unused argument: $remainingArg")
+            acc.reverse
+          case param :: tailParams =>
+            if param.isVarargs then // also last arguments
+              byNameArgsMap.get(param.name) match
+                case Some(byNameVarargs) => acc ::: byNameVarargs.toList ::: remainingArgs
+                case None => acc ::: remainingArgs
+            else byNameArgsMap.get(param.name) match
+              case Some(argValues) =>
+                assert(argValues.nonEmpty, s"${param.name} present in byNameArgsMap, but it has no argument value")
+                if argValues.length > 1 then
+                  error(s"more than one value for ${param.name}: ${argValues.mkString(", ")}")
+                rec(tailParams, argValues.last :: acc, remainingArgs)
+
+              case None =>
+                remainingArgs match
+                  case arg :: rest =>
+                    rec(tailParams, arg :: acc, rest)
+                  case Nil =>
+                    if !param.hasDefault then
+                      error(s"missing argument for ${param.name}")
+                    rec(tailParams, "" :: acc, Nil)
+      rec(info.parameters.toList, Nil, positionalArgs)
+
+    if hasError then None
+    else Some(orderedArgs)
+  end preProcessArgs
+
   private var hasParseErrors: Boolean = false
 
   /** Issue an error, and return an uncallable getter */
@@ -210,31 +203,22 @@ final class newMain extends MainAnnotation[FromString, Any]:
       case Some(t) => () => t
       case None => parseError(s"could not parse argument for `${param.name}` of type ${param.typeName.split('.').last}: $arg")
 
-  private def checkNames(info: Info): Unit =
-    def checkDuplicateNames() =
-      val nameAndCanonicalName = info.parameters.flatMap { paramInfo =>
-        (getNameWithMarker(paramInfo.name) +: paramInfo.longAliases ++: paramInfo.shortAliases).map(_ -> paramInfo.name)
-      }
-      val nameToCanonicalNames = nameAndCanonicalName.groupMap(_._1)(_._2)
-      for (name, canonicalNames) <- nameToCanonicalNames if canonicalNames.length > 1 do
-        throw IllegalArgumentException(s"$name is used for multiple parameters: ${canonicalNames.mkString(", ")}")
-    def checkValidNames() =
-      def isValidArgName(name: String): Boolean =
-        longArgNameRegex.matches(name) || shortArgNameRegex.matches(name)
-      for param <- info.parameters do
-        if !isValidArgName(param.name) then
-          throw IllegalArgumentException(s"The following argument name is invalid: ${param.name}")
-        for alias <- param.aliasNames if !isValidArgName(alias) do
-          throw IllegalArgumentException(s"The following alias is invalid: $alias")
-
-    checkValidNames()
-    checkDuplicateNames()
-
-  private class Help(info: Info):
+
+  private object Help:
+
+    /** The name of the special argument to display the method's help.
+     *  If one of the method's parameters is called the same, will be ignored.
+     */
+    private inline val helpArg = "help"
+
+    /** The short name of the special argument to display the method's help.
+     *  If one of the method's parameters uses the same short name, will be ignored.
+     */
+    private inline val shortHelpArg = 'h'
 
     private inline val maxUsageLineLength = 120
 
-    def printUsage(): Unit =
+    def printUsage(info: Info): Unit =
       def argsUsage: Seq[String] =
         for (param <- info.parameters)
         yield {
@@ -266,7 +250,7 @@ final class newMain extends MainAnnotation[FromString, Any]:
       println(printUsageBeginning + printUsages.mkString("\n" + " " * argsOffset))
     end printUsage
 
-    def printExplain(): Unit =
+    def printExplain(info: Info): Unit =
       def shiftLines(s: Seq[String], shift: Int): String = s.map(" " * shift + _).mkString("\n")
 
       def wrapLongLine(line: String, maxLength: Int): List[String] = {
@@ -282,6 +266,8 @@ final class newMain extends MainAnnotation[FromString, Any]:
         recurse(line, Vector()).toList
       }
 
+      println()
+
       if (info.documentation.nonEmpty)
         println(wrapLongLine(info.documentation, maxUsageLineLength).mkString("\n"))
       if (info.parameters.nonEmpty) {
@@ -312,25 +298,16 @@ final class newMain extends MainAnnotation[FromString, Any]:
       }
     end printExplain
 
-  private object Help:
-    /** The name of the special argument to display the method's help.
-     *  If one of the method's parameters is called the same, will be ignored.
-     */
-    private inline val helpArg = "help"
-
-    /** The short name of the special argument to display the method's help.
-     *  If one of the method's parameters uses the same short name, will be ignored.
-     */
-    private inline val shortHelpArg = 'h'
-
     def hasHelpArg(canonicalNames: CanonicalNames, args: Seq[String]): Boolean =
       val helpIsOverridden = canonicalNames.getName(argMarker + helpArg).isDefined
       val shortHelpIsOverridden = canonicalNames.getShortName(shortArgMarker + shortHelpArg).isDefined
       (!helpIsOverridden && args.contains(longNameWithMarker(helpArg))) ||
       (!shortHelpIsOverridden && args.contains(shortNameWithMarker(shortHelpArg.toString)))
+
   end Help
 
   private class CanonicalNames(info: Info):
+
     private val namesToCanonicalName: Map[String, String] = info.parameters.flatMap(
       param =>
         val names = param.longAliases.map(_.drop(2))
@@ -347,10 +324,32 @@ final class newMain extends MainAnnotation[FromString, Any]:
         else names.map(_ -> canonicalName)
     ).toMap
 
-    def getName(name: String): Option[String] = namesToCanonicalName.get(name.drop(2))
+    def getName(name: String): Option[String] = namesToCanonicalName.get(name.drop(2)).orElse(getShortName(name))
 
     def getShortName(name: String): Option[String] = shortNamesToCanonicalName.get(name.drop(1))
 
+    override def toString(): String =
+      s"CanonicalNames($namesToCanonicalName, $shortNamesToCanonicalName)"
+
+    def checkNames(): Unit =
+      def checkDuplicateNames() =
+        val nameAndCanonicalName = info.parameters.flatMap { paramInfo =>
+          (getNameWithMarker(paramInfo.name) +: paramInfo.longAliases ++: paramInfo.shortAliases).map(_ -> paramInfo.name)
+        }
+        val nameToCanonicalNames = nameAndCanonicalName.groupMap(_._1)(_._2)
+        for (name, canonicalNames) <- nameToCanonicalNames if canonicalNames.length > 1 do
+          throw IllegalArgumentException(s"$name is used for multiple parameters: ${canonicalNames.mkString(", ")}")
+      def checkValidNames() =
+        def isValidArgName(name: String): Boolean =
+          longArgNameRegex.matches(name) || shortArgNameRegex.matches(name)
+        for param <- info.parameters do
+          if !isValidArgName(param.name) then
+            throw IllegalArgumentException(s"The following argument name is invalid: ${param.name}")
+          for alias <- param.aliasNames if !isValidArgName(alias) do
+            throw IllegalArgumentException(s"The following alias is invalid: $alias")
+
+      checkValidNames()
+      checkDuplicateNames()
   end CanonicalNames
 
 end newMain

tests/run/main-annotation-help-override.check

@@ -4,23 +4,33 @@ Usage: helpOverride1 [--notHelp] <Int>
 A method that should let --help and -h display help.
 Arguments:
   --notHelp - Int
-Error: could not parse argument for `help` of type Int: --help
+Error: missing argument for --help
+Error: missing argument for help
+Usage: helpOverride2 [--help] <Int>
 Usage: helpOverride3 [-h] <Int>
 
 A method that should let --help display help, but not -h.
 Arguments:
   -h - Int
+Error: missing argument for --help
+Error: missing argument for help
 Error: missing argument for h
 Usage: helpOverride4 [--help] <Int> [-h] <Int>
-Error: could not parse argument for `notHelp` of type Int: --help
+Error: missing argument for --help
+Error: missing argument for notHelp
+Usage: helpOverride5 [--notHelp | --help] <Int>
 Usage: helpOverride6 [--notHelp | -h] <Int>
 
 A method that should let --help display help, but not -h.
 Arguments:
   --notHelp (-h) - Int
+Error: missing argument for --help
+Error: missing argument for notHelp
 Error: missing argument for notH
 Usage: helpOverride7 [--notHelp | --help] <Int> [--notH | -h] <Int>
-Error: could not parse argument for `notHelp` of type Int: --help
+Error: missing argument for --help
+Error: missing argument for notHelp
+Usage: helpOverride8 [--notHelp | --help | -h] <Int>
 ##### -h
 Usage: helpOverride1 [--notHelp] <Int>
 
@@ -32,15 +42,25 @@ Usage: helpOverride2 [--help] <Int>
 A method that should let -h display help, but not --help.
 Arguments:
   --help - Int
-Error: could not parse argument for `h` of type Int: -h
+Error: missing argument for -h
+Error: missing argument for h
+Usage: helpOverride3 [-h] <Int>
+Error: missing argument for -h
+Error: missing argument for help
 Error: missing argument for h
 Usage: helpOverride4 [--help] <Int> [-h] <Int>
 Usage: helpOverride5 [--notHelp | --help] <Int>
 
 A method that should let -h display help, but not --help.
 Arguments:
   --notHelp (--help) - Int
-Error: could not parse argument for `notHelp` of type Int: -h
+Error: missing argument for -h
+Error: missing argument for notHelp
+Usage: helpOverride6 [--notHelp | -h] <Int>
+Error: missing argument for -h
+Error: missing argument for notHelp
 Error: missing argument for notH
 Usage: helpOverride7 [--notHelp | --help] <Int> [--notH | -h] <Int>
-Error: could not parse argument for `notHelp` of type Int: -h
+Error: missing argument for -h
+Error: missing argument for notHelp
+Usage: helpOverride8 [--notHelp | --help | -h] <Int>