import Foundation

struct Parser<A> {
  let run: (inout Substring) -> A?
}

extension Parser {
  func map<B>(_ f: @escaping (A) -> B) -> Parser<B> {
    return Parser<B> { str -> B? in
      self.run(&str).map(f)
    }
  }
  func run(_ str: String) -> (match: A?, rest: Substring) {
    var str = str[...]
    let match = self.run(&str)
    return (match, str)
  }
}

func zip<A, B>(_ a: Parser<A>, _ b: Parser<B>) -> Parser<(A, B)> {
  return Parser<(A, B)> { str -> (A, B)? in
    let original = str
    guard let matchA = a.run(&str) else { return nil }
    guard let matchB = b.run(&str) else {
      str = original
      return nil
    }
    return (matchA, matchB)
  }
}

func zip<A, B, C>(
  _ a: Parser<A>,
  _ b: Parser<B>,
  _ c: Parser<C>
  ) -> Parser<(A, B, C)> {
  return zip(a, zip(b, c))
    .map { a, bc in (a, bc.0, bc.1) }
}

func prefix(while p: @escaping (Character) -> Bool) -> Parser<Substring> {
  return Parser<Substring> { str in
    let prefix = str.prefix(while: p)
    str.removeFirst(prefix.count)
    return prefix
  }
}

func literal(_ p: String) -> Parser<Void> {
  return Parser<Void> { str in
    guard str.hasPrefix(p) else { return nil }
    str.removeFirst(p.count)
    return ()
  }
}

let word = prefix(while: { $0 != " " && $0 != "-" })

func optional<A>(_ parser: Parser<A>) -> Parser<A?> {
  return Parser<A?> { str in
    .some(parser.run(&str))
  }
}

let testTwoWords: Parser<[Substring?]> = zip(
  optional(word),
  literal(" "),
  word
  ).map { word1, _, word2 in
    return [word1, word2]
}


testTwoWords.run("OptionalWord RequiredWord")
// Output: (["Hello", "There"], rest "")

testTwoWords.run(" RequiredWord")
// Output: (["", "There"], rest "")
// However, ideally I would like ([nil, "There"], rest "")