[Options] Use a prefix trie for argument matching

Sources/SwiftDriver/Driver/Driver.swift

@@ -581,7 +581,11 @@ extension Driver {
         })
       }
 
-      parsedOption = .init(option: parsedOption.option, argument: translatedArgument)
+      parsedOption = .init(
+        option: parsedOption.option,
+        argument: translatedArgument,
+        index: parsedOption.index
+      )
     }
   }
 

Sources/SwiftDriver/Jobs/CompileJob.swift

@@ -251,7 +251,7 @@ extension Array where Element == Job.ArgTemplate {
   /// Append the last parsed option that matches one of the given options
   /// to this command line.
   mutating func appendLast(_ options: Option..., from parsedOptions: inout ParsedOptions) throws {
-    guard let parsedOption = parsedOptions.last(where: { options.contains($0.option) }) else {
+    guard let parsedOption = parsedOptions.last(for: options) else {
       return
     }
 
@@ -276,17 +276,15 @@ extension Array where Element == Job.ArgTemplate {
   /// Append all parsed options that match one of the given options
   /// to this command line.
   mutating func appendAll(_ options: Option..., from parsedOptions: inout ParsedOptions) throws {
-    let matchingOptions = parsedOptions.filter { options.contains($0.option) }
-    for matching in matchingOptions {
+    for matching in parsedOptions.arguments(for: options) {
       try append(matching)
     }
   }
 
   /// Append just the arguments from all parsed options that match one of the given options
   /// to this command line.
   mutating func appendAllArguments(_ options: Option..., from parsedOptions: inout ParsedOptions) throws {
-    let matchingOptions = parsedOptions.filter { options.contains($0.option) }
-    for matching in matchingOptions {
+    for matching in parsedOptions.arguments(for: options) {
       try self.appendSingleArgument(option: matching.option, argument: matching.argument.asSingle)
     }
   }
@@ -295,16 +293,12 @@ extension Array where Element == Job.ArgTemplate {
   /// or the flag that corresponds to the default value if neither
   /// appears.
   mutating func appendFlag(true trueFlag: Option, false falseFlag: Option, default defaultValue: Bool, from parsedOptions: inout ParsedOptions) {
-    guard let parsedOption = parsedOptions.last(where: { $0.option == trueFlag || $0.option == falseFlag }) else {
-      if defaultValue {
-        appendFlag(trueFlag)
-      } else {
-        appendFlag(falseFlag)
-      }
-      return
-    }
-
-    appendFlag(parsedOption.option)
+    let isTrue = parsedOptions.hasFlag(
+      positive: trueFlag,
+      negative: falseFlag,
+      default: defaultValue
+    )
+    appendFlag(isTrue ? trueFlag : falseFlag)
   }
 }
 

Sources/SwiftDriver/Jobs/DarwinToolchain+LinkerSupport.swift

@@ -208,9 +208,7 @@ extension DarwinToolchain {
       fallthrough
     case .executable:
       linkerTool = .dynamicLinker
-      let fSystemArgs = parsedOptions.filter {
-        $0.option == .F || $0.option == .Fsystem
-      }
+      let fSystemArgs = parsedOptions.arguments(for: .F, .Fsystem)
       for opt in fSystemArgs {
         commandLine.appendFlag(.F)
         commandLine.appendPath(try VirtualPath(path: opt.argument.asSingle))
@@ -245,7 +243,7 @@ extension DarwinToolchain {
       // These custom arguments should be right before the object file at the
       // end.
       try commandLine.append(
-        contentsOf: parsedOptions.filter { $0.option.group == .linkerOption }
+        contentsOf: parsedOptions.arguments(in: .linkerOption)
       )
       try commandLine.appendAllArguments(.Xlinker, from: &parsedOptions)
 

Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift

@@ -64,7 +64,7 @@ extension Driver {
     try commandLine.appendLast(.AssertConfig, from: &parsedOptions)
     try commandLine.appendLast(.autolinkForceLoad, from: &parsedOptions)
 
-    if let colorOption = parsedOptions.last(where: { $0.option == .colorDiagnostics || $0.option == .noColorDiagnostics }) {
+    if let colorOption = parsedOptions.last(for: .colorDiagnostics, .noColorDiagnostics) {
       commandLine.appendFlag(colorOption.option)
     } else if shouldColorDiagnostics() {
       commandLine.appendFlag(.colorDiagnostics)

Sources/SwiftDriver/Jobs/GenericUnixToolchain+LinkerSupport.swift

@@ -166,9 +166,7 @@ extension GenericUnixToolchain {
       }
       commandLine.append(contentsOf: inputFiles)
 
-      let fSystemArgs = parsedOptions.filter {
-        $0.option == .F || $0.option == .Fsystem
-      }
+      let fSystemArgs = parsedOptions.arguments(for: .F, .Fsystem)
       for opt in fSystemArgs {
         if opt.option == .Fsystem {
           commandLine.appendFlag("-iframework")
@@ -250,7 +248,7 @@ extension GenericUnixToolchain {
       // These custom arguments should be right before the object file at the
       // end.
       try commandLine.append(
-        contentsOf: parsedOptions.filter { $0.option.group == .linkerOption }
+        contentsOf: parsedOptions.arguments(in: .linkerOption)
       )
       try commandLine.appendAllArguments(.Xlinker, from: &parsedOptions)
       try commandLine.appendAllArguments(.XclangLinker, from: &parsedOptions)

Sources/SwiftDriver/Options/OptionParsing.swift

@@ -15,32 +15,15 @@ public enum OptionParseError : Error {
 }
 
 extension OptionTable {
-  private func matchOption(_ argument: String) -> Option? {
-    var bestOption: Option? = nil
-
-    // FIXME: Use a binary search or trie or similar.
-    for option in options {
-      // If there isn't a prefix match, keep going.
-      if !argument.starts(with: option.spelling) { continue }
-
-      // If this is the first option we've seen, or if it's a longer
-      // match than the best option so far, then we have a new best
-      // option.
-      if let bestOption = bestOption,
-        bestOption.spelling.count >= option.spelling.count {
-        continue
-      }
-
-      bestOption = option
-    }
-
-    return bestOption
-  }
-
   /// Parse the given command-line arguments into a set of options.
   ///
   /// Throws an error if the command line contains any errors.
   public func parse(_ arguments: [String]) throws -> ParsedOptions {
+    var trie = PrefixTrie<String.UTF8View, Option>()
+    for opt in Option.allOptions {
+      trie[opt.spelling.utf8] = opt
+    }
+
     var parsedOptions = ParsedOptions()
     var index = arguments.startIndex
     while index < arguments.endIndex {
@@ -59,8 +42,11 @@ extension OptionTable {
         continue
       }
 
-      // Match to an option, identified by key.
-      guard let option = matchOption(argument) else {
+      // Match to an option, identified by key. Note that this is a prefix
+      // match -- if the option is a `.flag`, we'll explicitly check to see if
+      // there's an unmatched suffix at the end, and pop an error. Otherwise,
+      // we'll treat the unmatched suffix as the argument to the option.
+      guard let option = trie[argument.utf8] else {
         throw OptionParseError.unknownOption(
           index: index - 1, argument: argument)
       }
@@ -123,7 +109,7 @@ extension OptionTable {
         index += 1
       }
     }
-
+    parsedOptions.buildIndex()
     return parsedOptions
   }
 }

Sources/SwiftDriver/Options/OptionTable.swift

@@ -14,6 +14,14 @@ public struct OptionTable {
 
   /// Retrieve the options.
   public var options: [Option] = Option.allOptions + Option.extraOptions
+  public lazy var groupMap: [Option.Group: [Option]] = {
+    var map = [Option.Group: [Option]]()
+    for opt in options {
+      guard let group = opt.group else { continue }
+      map[group, default: []].append(opt)
+    }
+    return map
+  }()
 }
 
 extension String {

Sources/SwiftDriver/Options/ParsedOptions.swift

@@ -45,6 +45,9 @@ public struct ParsedOption {
 
   /// The argument bound to this option.
   public let argument: Argument
+
+  /// The index in the command line where this argument appeared.
+  public let index: Int
 }
 
 extension ParsedOption: CustomStringConvertible {
@@ -84,15 +87,40 @@ public struct ParsedOptions {
   /// The parsed options, which match up an option with its argument(s).
   private var parsedOptions: [ParsedOption] = []
 
+  /// Maps the canonical spelling of an option to all the instances of
+  /// that option that we've seen in the map, and their index in the
+  /// parsedOptions array. Prefer to use this for lookup
+  /// whenever you can.
+  private var optionIndex = [String: [ParsedOption]]()
+
+  /// Maps option groups to the set of parsed options that are present for
+  /// them.
+  private var groupIndex = [Option.Group: [ParsedOption]]()
+
   /// Indication of which of the parsed options have been "consumed" by the
   /// driver. Any unconsumed options could have been omitted from the command
   /// line.
   private var consumed: [Bool] = []
 }
 
 extension ParsedOptions {
+  mutating func buildIndex() {
+    optionIndex.removeAll()
+    for parsed in parsedOptions {
+      optionIndex[parsed.option.canonical.spelling, default: []].append(parsed)
+      if let group = parsed.option.group {
+        groupIndex[group, default: []].append(parsed)
+      }
+    }
+  }
+
   mutating func addOption(_ option: Option, argument: Argument) {
-    parsedOptions.append(.init(option: option, argument: argument))
+    let parsed = ParsedOption(
+      option: option,
+      argument: argument,
+      index: parsedOptions.count
+    )
+    parsedOptions.append(parsed)
     consumed.append(false)
   }
 
@@ -146,16 +174,36 @@ extension ParsedOptions {
   /// Any options that match the `isIncluded` predicate will be marked "consumed".
   public mutating func filter(where isIncluded: (ParsedOption) throws -> Bool) rethrows -> [ParsedOption] {
     var result: [ParsedOption] = []
-    for index in parsedOptions.indices {
-      if try isIncluded(parsedOptions[index]) {
-        consumed[index] = true
-        result.append(parsedOptions[index])
+    for option in parsedOptions {
+      if try isIncluded(option) {
+        consumed[option.index] = true
+        result.append(option)
       }
     }
 
     return result
   }
 
+  public mutating func arguments(for options: Option...) -> [ParsedOption] {
+    return options.flatMap { lookup($0) }
+  }
+
+  public mutating func arguments(for options: [Option]) -> [ParsedOption] {
+    return options.flatMap { lookup($0) }
+  }
+
+  public mutating func arguments(in group: Option.Group) -> [ParsedOption] {
+    return groupIndex[group, default: []]
+  }
+
+  public mutating func last(for options: Option...) -> ParsedOption? {
+    return arguments(for: options).max { $0.index < $1.index }
+  }
+
+  public mutating func last(for options: [Option]) -> ParsedOption? {
+    return arguments(for: options).max { $0.index < $1.index }
+  }
+
   /// Return the last parsed options that matches the given predicate.
   ///
   /// Any options that match the `isIncluded` predicate will be marked "consumed".
@@ -166,7 +214,7 @@ extension ParsedOptions {
   /// Does this contain a particular option.
   public mutating func contains(_ option: Option) -> Bool {
     assert(option.alias == nil, "Don't check for aliased options")
-    return last { parsed in parsed.option.canonical == option } != nil
+    return !lookup(option).isEmpty
   }
 
   /// Determine whether the parsed options contains an option in the given
@@ -179,7 +227,7 @@ extension ParsedOptions {
   ///
   /// This operation does not consume any inputs.
   public var hasAnyInput: Bool {
-    return parsedOptions.contains { $0.option == .INPUT }
+    return !lookupWithoutConsuming(.INPUT).isEmpty
   }
 
   /// Walk through all of the parsed options, modifying each one.
@@ -189,56 +237,68 @@ extension ParsedOptions {
     for index in parsedOptions.indices {
       try body(&parsedOptions[index])
     }
+    buildIndex()
+  }
+
+  internal func lookupWithoutConsuming(_ option: Option) -> [ParsedOption] {
+    optionIndex[option.canonical.spelling, default: []]
+  }
+
+  internal mutating func lookup(_ option: Option, consume: Bool = true) -> [ParsedOption] {
+    let opts = lookupWithoutConsuming(option)
+    for opt in opts {
+      consumed[opt.index] = true
+    }
+    return opts
   }
 
   /// Find all of the inputs.
   public var allInputs: [String] {
     mutating get {
-      filter { $0.option == .INPUT }.map { $0.argument.asSingle }
+      lookup(.INPUT).map { $0.argument.asSingle }
     }
   }
 
   /// Determine whether the parsed options contain an argument with one of
   /// the given options
-  public mutating func hasArgument(_ options: Option...) -> Bool {
-    return last { parsed in
-      return options.contains(parsed.option)
-    } != nil
+  public func hasArgument(_ options: Option...) -> Bool {
+    return options.contains { !lookupWithoutConsuming($0).isEmpty }
   }
+
   /// Given an option and its negative form, return
   /// true if the option is present, false if the negation is present, and
   /// `default` if neither option is given. If both the option and its
   /// negation are present, the last one wins.
   public mutating func hasFlag(positive: Option,
                                negative: Option,
                                default: Bool) -> Bool {
-    let positiveIndexOpt = parsedOptions.lastIndex { $0.option == positive }
-    let negativeIndexOpt = parsedOptions.lastIndex { $0.option == negative }
+    let positiveOpt = lookup(positive).last
+    let negativeOpt = lookup(positive).last
 
     // If neither are present, return the default
-    guard positiveIndexOpt != nil || negativeIndexOpt != nil else {
+    guard positiveOpt != nil || negativeOpt != nil else {
       return `default`
     }
 
     // If the positive isn't provided, then the negative will be
-    guard let positiveIndex = positiveIndexOpt else { return false }
+    guard let positive = positiveOpt else { return false }
 
     // If the negative isn't provided, then the positive will be
-    guard let negativeIndex = negativeIndexOpt else { return true }
+    guard let negative = negativeOpt else { return true }
 
     // Otherwise, return true if the positive index is greater than the negative,
     // false otherwise
-    return positiveIndex > negativeIndex
+    return positive.index > negative.index
   }
 
   /// Get the last argument matching the given option.
   public mutating func getLastArgument(_ option: Option) -> Argument? {
     assert(option.alias == nil, "Don't check for aliased options")
-    return last { parsed in parsed.option.canonical == option }?.argument
+    return lookup(option).last?.argument
   }
 
   /// Get the last parsed option within the given option group.
   public mutating func getLast(in group: Option.Group) -> ParsedOption? {
-    return last { parsed in parsed.option.group == group }
+    return groupIndex[group]?.last
   }
 }