///|
/// Regex engine micro-benchmarks over synthetic inputs (no file system, so it
/// runs on every backend). Each case reports the best of `repeat` runs and a
/// checksum of all match spans (it must not change between engine versions).
///
/// Usage: regex_bench [repeat] [case-filter]
///        regex_bench --dump   (inputs and results as JSON lines, checked
///                             against Python by scripts/regex_bench_oracle.py)
fn main {
  let args = @env.args()
  if args.length() > 1 && args[1] == "--dump" {
    dump()
    return
  }
  let repeat = if args.length() > 1 {
    @string.parse_int(args[1]) catch {
      _ => 5
    }
  } else {
    5
  }
  let filter = if args.length() > 2 { args[2] } else { "" }
  let cases : Array[(String, () -> Int32)] = [
    ("repr find_all", bench_repr(repr_lines(20000))),
    ("json find_all", bench_json(json_docs(800))),
    ("lexer match_at", bench_tokenize(python_text(4000))),
    ("lexer match_end", bench_tokenize_end(python_text(4000))),
    ("rare search", bench_rare(log_text(300000))),
  ]
  for case in cases {
    let (name, run) = case
    if filter != "" && !name.contains(filter) {
      continue
    }
    let mut best = 1_000_000_000UL
    let mut sum = 0
    for _ in 0.. Int32 {
  self.state = self.state * 1664525U + 1013904223U
  ((self.state >> 8).reinterpret_as_int() & 0x7FFFFF) % n
}

///|
/// A backslash followed by a double quote.
let esc_dq : String = "\\\""

///|
/// A backslash followed by a single quote.
let esc_sq : String = "\\'"

///|
fn hex(n : Int32, width : Int32) -> String {
  let digits = "0123456789abcdef"
  let sb = StringBuilder()
  for i = width - 1; i >= 0; i = i - 1 {
    sb.write_char(digits.get_char((n >> (i * 4)) & 15).unwrap())
  }
  sb.to_string()
}

///|
/// Lines like the ones Rich's `ReprHighlighter` sees when pretty-printing.
fn repr_lines(n : Int32) -> Array[String] {
  let rng : Lcg = { state: 42, }
  let out = []
  for i in 0..
        "{'name': 'user_\{k}', 'id': \{k}, 'score': \{k % 97}.5, 'tags': ['alpha', 'beta', \"say \\\"hi\\\"\"], 'active': True, 'parent': None}"
      1 => " "
      2 =>
        "Point(x=\{k}, y=-\{k % 13}.25e-3, z=3+4j, w=1e10) at 0x7f\{hex(k, 8)}"
      3 =>
        "[IPv4Address('192.168.\{k % 256}.1'), IPv6Address('2001:db8::\{hex(k, 4)}'), 'https://example.com/api/v1/items/\{k}?q=x&page=2']"
      4 =>
        "PosixPath('/usr/local/lib/python3.14/site-packages/pkg_\{k}/module.py')"
      5 =>
        "UUID('\{hex(k, 8)}-1234-5678-9abc-\{hex(k, 8)}\{hex(k, 4)}'), datetime.datetime(2024, 1, \{k % 28 + 1}, 12, 30, 45)"
      6 =>
        "Processing item \{k} of 100000: status=ok elapsed=0.\{k % 1000}s worker=thread-\{k % 8} queue_depth=\{k % 50}"
      7 =>
        "b'\\x00\\x01 bytes' ... False ... {'nested': {'deep': [1, 2, [3, 4, [5, ...]]]}}"
      8 =>
        "    'description': 'A fairly long string value that goes on for a while without any spécial tökens in it at all 😀',"
      _ =>
        "Mac('00:1a:2b:3c:4d:\{hex(k % 256, 2)}') eui64=00-1a-2b-ff-fe-3c-4d-5e, ratio=\{k % 7}.\{k % 100}"
    }
    out.push(line)
  }
  out
}

///|
let repr_patterns : Array[String] = [
  "(?P<)(?P[-\\w.:|]*)(?P[\\w\\W]*)(?P>)",
  "(?P[\\w_]{1,50})=(?P\"?[\\w_]+\"?)?",
  "(?P[][{}()])",
  [
    "(?P[0-9]{1,3}\\.[0-9]{1,3}\\.[0-9]{1,3}\\.[0-9]{1,3})", "(?P([A-Fa-f0-9]{1,4}::?){1,7}[A-Fa-f0-9]{1,4})",
    "(?P(?:[0-9A-Fa-f]{1,2}-){7}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{1,2}:){7}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{4}\\.){3}[0-9A-Fa-f]{4})",
    "(?P(?:[0-9A-Fa-f]{1,2}-){5}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{1,2}:){5}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{4}\\.){2}[0-9A-Fa-f]{4})",
    "(?P[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12})",
    "(?P[\\w.]*?)\\(", "\\b(?PTrue)\\b|\\b(?PFalse)\\b|\\b(?PNone)\\b",
    "(?P\\.\\.\\.)", "(?P(?(?\\B(/[-\\w._+]+)*\\/)(?P[-\\w._+]*)?",
    "(?b?'''.*?(?(file|https|http|ws|wss)://[-0-9a-zA-Z$_+!`(),.?/;:&=%#~@]*)",
  ].join("|"),
]

///|
/// Sum of all group spans of all matches: a cheap fingerprint.
fn spans_checksum(ms : Array[@regex.Match]) -> Int32 {
  let mut sum = 0
  for m in ms {
    for g in 0..<=m.group_count() {
      sum = sum * 31 + m.start(group=g) * 7 + m.end(group=g)
    }
  }
  sum
}

///|
fn compile_all(patterns : Array[String]) -> Array[@regex.Regex] {
  patterns.map(p => {
    @regex.compile(p) catch {
      e => abort("bad pattern \{p}: \{Repr(e)}")
    }
  })
}

///|
fn bench_repr(lines : Array[String]) -> () -> Int32 {
  let res = compile_all(repr_patterns)
  fn() {
    let mut sum = 0
    for line in lines {
      for re in res {
        let ms = re.find_all(line) catch { _ => abort("budget") }
        sum = sum * 17 + spans_checksum(ms)
      }
    }
    sum
  }
}

///|
/// Pretty-printed JSON documents (Rich's `JSONHighlighter` input).
fn json_docs(n : Int32) -> Array[String] {
  let rng : Lcg = { state: 7, }
  let out = []
  for _ in 0..[\\{\\[\\(\\)\\]\\}])", "\\b(?Ptrue)\\b|\\b(?Pfalse)\\b|\\b(?Pnull)\\b",
    "(?P(? () -> Int32 {
  let res = compile_all(json_patterns)
  fn() {
    let mut sum = 0
    for doc in docs {
      for re in res {
        let ms = re.find_all(doc) catch { _ => abort("budget") }
        sum = sum * 17 + spans_checksum(ms)
      }
    }
    sum
  }
}

///|
/// Python-like source text.
fn python_text(n : Int32) -> String {
  let rng : Lcg = { state: 99, }
  let sb = StringBuilder()
  for i in 0..[^\\W\\d]\\w*)(?=\\s*\\()", 0),
  ("[^\\W\\d]\\w*", 0),
  ("\\*\\*|//|<<|>>|[-+*/%&|^~<>=!]=?", 0),
  ("[()\\[\\]{}:;,.@]", 0),
]

///|
fn compile_tokens() -> Array[@regex.Regex] {
  token_patterns.map(p => {
    @regex.compile(p.0, flags=p.1) catch {
      e => abort("bad pattern \{p.0}: \{Repr(e)}")
    }
  })
}

///|
/// A `RegexLexer`-style loop: at each position, the first rule that matches
/// (via `match_at`, which allocates a `Match`) wins.
fn bench_tokenize(text : String) -> () -> Int32 {
  let res = compile_tokens()
  fn() {
    let budget = @regex.Budget::new(1_000_000_000)
    let mut sum = 0
    let mut pos = 0
    let n = text.length()
    while pos < n {
      let mut next = pos + 1
      for i, re in res {
        let r = re.match_at(text, pos, budget~) catch { _ => abort("budget") }
        if r is Some(m) {
          if m.end() > pos {
            next = m.end()
          }
          sum = sum * 31 + i + m.end()
          break
        }
      }
      pos = next
    }
    sum
  }
}

///|
/// Like `bench_tokenize`, with the dispatch the pygments lexer uses:
/// `may_start_with` first, then `match_end` (no `Match` allocation).
fn bench_tokenize_end(text : String) -> () -> Int32 {
  let res = compile_tokens()
  fn() {
    let budget = @regex.Budget::new(1_000_000_000)
    let mut sum = 0
    let mut pos = 0
    let n = text.length()
    while pos < n {
      let c = text.unsafe_get(pos).to_int()
      let mut next = pos + 1
      for i, re in res {
        if !re.may_start_with(c) {
          continue
        }
        let e = re.match_end(text, pos, budget~) catch { _ => abort("budget") }
        if e >= 0 {
          if e > pos {
            next = e
          }
          sum = sum * 31 + i + e
          break
        }
      }
      pos = next
    }
    sum
  }
}

///|
/// A long log where the searched-for lines are rare.
fn log_text(n : Int32) -> String {
  let rng : Lcg = { state: 3, }
  let sb = StringBuilder()
  for i in 0...*)", "https?://[^\\s]+", "(?i)quota\\s+exceeded",
]

///|
fn bench_rare(text : String) -> () -> Int32 {
  let res = compile_all(rare_patterns)
  fn() {
    let mut sum = 0
    for re in res {
      let ms = re.find_all(text) catch { _ => abort("budget") }
      sum = sum * 17 + spans_checksum(ms)
    }
    sum
  }
}

///|
fn match_spans(ms : Array[@regex.Match]) -> Array[Int32] {
  let out = []
  for m in ms {
    for g in 0..<=m.group_count() {
      out.push(m.start(group=g))
      out.push(m.end(group=g))
    }
  }
  out
}

///|
fn dump_find_all(
  name : String,
  patterns : Array[String],
  texts : Array[String],
) -> Unit {
  let res = compile_all(patterns)
  let spans : Array[Array[Int32]] = []
  for t in texts {
    for re in res {
      spans.push(match_spans(re.find_all(t) catch { _ => abort("budget") }))
    }
  }
  let obj : Json = {
    "kind": "find_all",
    "name": name,
    "patterns": Json(patterns),
    "texts": Json(texts),
    "spans": Json(spans),
  }
  println(obj.stringify())
}

///|
fn dump_tokenize(text : String) -> Unit {
  let res = compile_tokens()
  let tokens : Array[Array[Int32]] = []
  let mut pos = 0
  let n = text.length()
  while pos < n {
    let mut next = pos + 1
    for i, re in res {
      let r = re.match_at(text, pos) catch { _ => abort("budget") }
      if r is Some(m) {
        if m.end() > pos {
          next = m.end()
        }
        tokens.push([pos, i, ..match_spans([m])])
        break
      }
    }
    pos = next
  }
  let obj : Json = {
    "kind": "tokenize",
    "name": "lexer match_at",
    "patterns": Json(token_patterns.map(p => [Json(p.0), Json(p.1)])),
    "texts": Json([text]),
    "spans": Json(tokens),
  }
  println(obj.stringify())
}

///|
fn dump() -> Unit {
  dump_find_all("repr find_all", repr_patterns, repr_lines(2000))
  dump_find_all("json find_all", json_patterns, json_docs(40))
  dump_find_all("rare search", rare_patterns, [log_text(20000)])
  dump_tokenize(python_text(300))
}