1BRC study: the Nikaia compiler vs. human vs. AI

// `fast.nika` with the claims a programmer who knows the 1BRC format vouches
// for, each in `unsafe { assume(…) }` (ADR-376): the table index in range, the
// name's word reads inside the padded input, the station counter below its
// bound, the temperature's shape, the sums below a bound. In a build that
// removes checks they are believed, and the checks they imply are gone; in
// one that removes none, each is checked. See `README.md`.
//
// Why it exists: to show that Nikaia reaches the hardware as far as the
// fastest hand-tuned Rust does - the same instructions, the same tricks, the
// same claims about the input - not that it is cleverer than the people who
// wrote those programs. It is the proof that the language can express such
// optimisations; `fast.nika` beside it shows how much of them the prover
// already makes safe.
//
// What follows is `fast.nika`'s own description.
//
// The 1BRC aggregation written for speed against today's `std`, the
// counterpart of `src/bin/ai2.rs`.
//
// * the file is mapped and read as `Bytes`: nothing checks it as UTF-8 as a
//   whole; a station's name is checked once, where it is printed;
// * `find_all` over a `ByteSet` of `;` alone finds the separators; where a
//   line ends follows from the temperature's shape, so the `\n` is never
//   searched for;
// * a name's first sixteen bytes, read as two words with `u64_le`, are its
//   hash and its key; its bytes are compared only where it is longer;
// * the temperature is read from one word, without a branch;
// * an open-addressing table of our own in a `Vec`;
// * the last rows are read from a copy with zeros behind them, so that no
//   word read goes past the end;
// * on more than one core, `par_pieces` cuts the file at line ends and
//   `Par::fold` scans the pieces into a table each and merges them.
use std::cli
use std::collections
use std::fs
use std::hash
use std::simd
use std::target
use std::task

struct Slot {
    w0: u64,
    w1: u64,
    len: i64,
    min: i64,
    max: i64,
    sum: i64,
    count: i64,
}

struct Table {
    slots: Vec[Slot],
    // Each slot's name, kept apart: the loop reads it only for a name longer
    // than sixteen bytes, so a slot is the stats and the key alone.
    names: Vec[Bytes],
    used: i64,
}

comptime BITS: i64 = 14
// Whether the CPU level this program is built for (`--cpu`) hashes a word
// in one instruction: decided while it is built, so the other hash is not
// in the binary.
comptime CRC_IN_ONE: bool = target::has(target::Instruction::Crc32c)
comptime SLOTS: i64 = 16384
// What of a name's first and second word belongs to it, by its length: a
// name of `len` bytes keeps the low `len` bytes of the first word and the
// low `len - 8` of the second, computed while the program is built.
comptime FIRST: Array[u64, 101] = word_masks(0)
comptime SECOND: Array[u64, 101] = word_masks(8)
// How far behind the last row start a word read may reach: a name of 100
// bytes, its temperature and a word beyond.
comptime MARGIN: i64 = 256

// The masks of `FIRST` (`skip` 0) and `SECOND` (`skip` 8): for each length
// up to 100, a byte of ones for each of the name's bytes from `skip` on, at
// most eight.
fn word_masks(skip: i64) -> Vec[u64] {
    let mut masks: Vec[u64] = []
    for len in 0..<101 {
        let mut mask: u64 = 0
        let mut k = skip
        while k < len && k < skip + 8 {
            mask = (mask << 8) | 0xFF
            k += 1
        }
        masks.push(mask)
    }
    return masks
}

fn table() -> Table {
    let mut slots: Vec[Slot] = []
    slots.resize(SLOTS, Slot { w0: 0, w1: 0, len: 0, min: 0, max: 0, sum: 0, count: 0 })
    let mut names: Vec[Bytes] = []
    names.resize(SLOTS, Bytes([]))
    return Table { slots: slots, names: names, used: 0 }
}

impl Table {
    // The slot of the name `buf[pos..<pos + len]`, whose first sixteen bytes
    // are `w0` and `w1` (zeros past its end), taken where the name is new.
    fn slot_of(ref mut self, w0: u64, w1: u64, len: i64, buf: ref Bytes, pos: i64) -> i64 {
        // Both words go through the CRC whole: a first word cut to its low
        // half would let names that differ only in bytes 4 to 7 collide.
        let mut at = if CRC_IN_ONE {
            (hash::crc32c(hash::crc32c(0, w0), w1) as i64) & (SLOTS - 1)
        } else {
            ((w0 ^ (w1 << 5)).wrapping_mul(0x9E3779B97F4A7C15) >> (64 - BITS)).bits_as_i64()
        }
        while true {
            // The table holds `SLOTS` slots and `at` is masked into them; a
            // name's bytes lie inside the input.
            unsafe {
                assume(at >= 0 && at < self.slots.len() && at < self.names.len())
                assume(pos >= 0 && len >= 0 && len <= 100 && pos <= buf.len() && pos + len <= buf.len())
            }
            if self.slots[at].count == 0 {
                self.claim(at, w0, w1, len, buf, pos)
                return at
            }
            if self.slots[at].w0 == w0 && self.slots[at].w1 == w1 && self.slots[at].len == len
                && (len <= 16 || self.names[at][0..<len] == buf[pos..<pos + len]) {
                return at
            }
            at = (at + 1) & (SLOTS - 1)
        }
        panic("unreachable")
    }

    // A station seen for the first time: once per name, so kept out of the
    // loop that runs once per row.
    @inline(never)
    fn claim(ref mut self, at: i64, w0: u64, w1: u64, len: i64, buf: ref Bytes, pos: i64) {
        unsafe {
            assume(self.used < 10000)
            assume(at >= 0 && at < self.slots.len() && at < self.names.len())
            assume(pos >= 0 && len >= 0 && len <= 100 && pos <= buf.len() && pos + len <= buf.len())
        }
        self.used += 1
        if self.used > 10000 { panic("more than 10 000 stations") }
        self.slots[at] = Slot { w0: w0, w1: w1, len: len, min: 999, max: -999, sum: 0, count: 0 }
        self.names[at] = buf[pos..<pos + len].clone()
    }

    // Every row of `buf` that starts before `stop`.
    fn run(ref mut self, buf: ref Bytes, stop: i64) {
        if stop == 0 { return }
        let mut pos = 0
        // The `;` after each row's start, 32 bytes a compare: a row is about
        // 14 bytes, so it is mostly in the first 32, and the test whether it
        // was found is predicted. A name has at most 100 bytes, so at most
        // four loads look for it, each inside the input.
        let semis: simd::Fixed[u8, 32] = simd::splat(';')
        while true {
            // The row starts before `stop`, which leaves `MARGIN` bytes
            // behind it, and its `;` lies within 100 bytes of its start. The
            // first 32 bytes almost always hold it, so they are read before
            // the loop, which only a long name enters.
            let mut sep = pos
            unsafe { assume(sep >= 0 && sep + 32 <= buf.len()) }
            let first: simd::Fixed[u8, 32] = simd::load(buf, sep)
            let mut hit = first.eq(semis).first() ?? -1
            while hit < 0 {
                sep += 32
                unsafe { assume(sep >= pos && sep - pos <= 100 && sep + 32 <= buf.len()) }
                let block: simd::Fixed[u8, 32] = simd::load(buf, sep)
                hit = block.eq(semis).first() ?? -1
            }
            unsafe { assume(hit >= 0 && hit < 32) }
            sep += hit
            // A row starts at `pos`, before `stop`, which leaves `MARGIN`
            // bytes behind it: its name of at most 100 bytes, the `;` and the
            // words read after it lie inside the input.
            unsafe {
                assume(pos >= 0 && pos < sep && sep - pos <= 100)
                assume(sep + 16 <= buf.len())
            }
            // The name's first two words, zeros past its end.
            let len = sep - pos
            // A name has 1 to 100 bytes: each word is cut to it by a mask
            // looked up by its length, without a branch on the length, which
            // names of varying length would mispredict.
            unsafe { assume(len >= 1 && len <= 100) }
            let w0 = buf.u64_le(pos) & FIRST[len]
            let w1 = buf.u64_le(pos + 8) & SECOND[len]

            let at = self.slot_of(w0, w1, len, buf, pos)
            unsafe {
                assume(at >= 0 && at < self.slots.len() && at < self.names.len())
                assume(self.slots[at].sum >= -1000000000000000 && self.slots[at].sum <= 1000000000000000)
                assume(self.slots[at].count >= 0 && self.slots[at].count < 1000000000000)
            }
            // The temperature, `-?d?d.d`, from one word: where the `.` is
            // says how long it is, and the sign becomes a mask.
            let w = buf.u64_le(sep + 1)
            let dot = (!w & 0x10101000).trailing_zeros()
            // A temperature is `-?d?d.d`: its `.` is its second, third or
            // fourth byte.
            unsafe { assume(dot == 12 || dot == 20 || dot == 28) }
            let signed = ((!w) << 59).bits_as_i64() >> 63
            let keep = !(signed.bits_as_u64() & 0xFF)
            let digits = ((w & keep) << (28 - dot)) & 0x0F000F0F00
            let abs = (digits.wrapping_mul(0x640A0001) >> 32) & 0x3FF
            unsafe { assume(abs <= 999 && (signed == 0 || signed == -1)) }
            let v = (abs.bits_as_i64() ^ signed) - signed
            unsafe { assume(v >= -999 && v <= 999) }

            if v < self.slots[at].min { self.slots[at].min = v }
            if v > self.slots[at].max { self.slots[at].max = v }
            self.slots[at].sum += v
            self.slots[at].count += 1

            pos = sep + 4 + ((dot >> 3) as i64)
            if pos >= stop { break }
        }
    }

    // Every row of `b`: those up to the last `MARGIN` bytes in place, the
    // rest from a copy with zeros behind it.
    fn scan(ref mut self, b: ref Bytes) {
        let n = b.len()
        let mut cut = 0
        if n > MARGIN {
            cut = n - MARGIN
            while cut > 0 && b[cut - 1] != '\n' { cut -= 1 }
        }
        self.run(b, cut)
        let mut tail: Vec[u8] = []
        for c in b[cut..<n] { tail.push(c) }
        for _ in 0..<MARGIN { tail.push(0) }
        self.run(Bytes(tail), n - cut)
    }

    fn merge(ref mut self, other: Table) {
        let mut i = 0
        while i < SLOTS {
            let s = other.slots[i]
            let name = other.names[i]
            i += 1
            if s.count > 0 {
                let at = self.slot_of(s.w0, s.w1, s.len, name, 0)
                if s.min < self.slots[at].min { self.slots[at].min = s.min }
                if s.max > self.slots[at].max { self.slots[at].max = s.max }
                self.slots[at].sum += s.sum
                self.slots[at].count += s.count
            }
        }
    }
}

fn main() throws {
    let path = cli::args().nth(1) ?? "measurements.txt"
    let b = fs::map(ref path, fs::Root::Anywhere) catch {
        eprintln(f"cannot open {path}: {error}")
        return
    }
    let mut all = table()
    if task::cores() == 1 {
        all.scan(b)
    } else {
        all = b.par_pieces('\n').fold(table, fn(t, piece) { t.scan(piece) }, Table::merge; order: par::Order::Any)
    }

    // Each station's name is checked as UTF-8 here, once.
    let mut index: collections::HashMap[ref String, i64] = collections::HashMap()
    let mut at = 0
    while at < SLOTS {
        if all.slots[at].count > 0 { index.insert(all.names[at].text(), at) }
        at += 1
    }
    let mut names = index.keys().collect()
    names.sort()
    let body = names
        .map fn(name) {
            let s = all.slots[index[name] ?? panic("a key")]
            f"{name}={tenths(s.min)}/{tenths(mean(s.sum, s.count))}/{tenths(s.max)}"
        }
        .join(", ")
    println(f"{{{body}}}")
}

// The official rounding: a tie goes toward +infinity (`naive.nika`).
fn mean(sum: i64, count: i64) -> i64 {
    // At most 10^12 rows of at most 99.9 degrees each.
    unsafe { assume(count >= 1 && count <= 1000000000000 && sum >= -1000000000000000 && sum <= 1000000000000000) }
    return (2 * sum + count).div_euclid(2 * count)
}

fn tenths(t: i64) -> String {
    unsafe { assume(t >= -999 && t <= 999) }
    if t < 0 {
        let a = -t
        return f"-{a / 10}.{a % 10}"
    }
    return f"{t / 10}.{t % 10}"
}

// The 1BRC aggregation written for speed against today's `std`, the
// counterpart of `src/bin/ai2.rs`. See `README.md`.
//
// * the file is mapped and read as `Bytes`: nothing checks it as UTF-8 as a
//   whole; a station's name is checked once, where it is printed;
// * `find_all` over a `ByteSet` of `;` alone finds the separators; where a
//   line ends follows from the temperature's shape, so the `\n` is never
//   searched for;
// * a name's first sixteen bytes, read as two words with `u64_le`, are its
//   hash and its key; its bytes are compared only where it is longer;
// * the temperature is read from one word, without a branch;
// * an open-addressing table of our own in a `Vec`;
// * the last rows are read from a copy with zeros behind them, so that no
//   word read goes past the end;
// * on more than one core, `par_pieces` cuts the file at line ends and
//   `Par::fold` scans the pieces into a table each and merges them.
use std::cli
use std::collections
use std::fs
use std::simd
use std::task

struct Slot {
    w0: u64,
    w1: u64,
    len: i64,
    name: Bytes,
    min: i64,
    max: i64,
    sum: i64,
    count: i64,
}

struct Table {
    slots: Vec[Slot],
    used: i64,
    assert(slots.len() == SLOTS; message: "the table holds `SLOTS` slots")
}

comptime BITS: i64 = 14
comptime SLOTS: i64 = 16384
// What of a name's first and second word belongs to it, by its length: a
// name of `len` bytes keeps the low `len` bytes of the first word and the
// low `len - 8` of the second, computed while the program is built.
comptime FIRST: Array[u64, 101] = word_masks(0)
comptime SECOND: Array[u64, 101] = word_masks(8)
// How far behind the last row start a word read may reach: a name of 100
// bytes, its temperature and a word beyond.
comptime MARGIN: i64 = 256

// The masks of `FIRST` (`skip` 0) and `SECOND` (`skip` 8): for each length
// up to 100, a byte of ones for each of the name's bytes from `skip` on, at
// most eight.
fn word_masks(skip: i64) -> Vec[u64] {
    let mut masks: Vec[u64] = []
    for len in 0..<101 {
        let mut mask: u64 = 0
        let mut k = skip
        while k < len && k < skip + 8 {
            mask = (mask << 8) | 0xFF
            k += 1
        }
        masks.push(mask)
    }
    return masks
}

fn table() -> Table {
    let mut slots: Vec[Slot] = []
    slots.resize(SLOTS, Slot { w0: 0, w1: 0, len: 0, name: Bytes([]), min: 0, max: 0, sum: 0, count: 0 })
    return Table { slots: slots, used: 0 }
}

impl Table {
    // The slot of the name `buf[pos..<pos + len]`, whose first sixteen bytes
    // are `w0` and `w1` (zeros past its end), taken where the name is new.
    // The `assert` before each `return` is its postcondition: a caller knows
    // the slot lies inside the table, and indexes it without a check. The
    // first is its precondition: the name lies inside `buf`, which a caller
    // proves, so the cut of a long name is taken without a check.
    fn slot_of(ref mut self, w0: u64, w1: u64, len: i64, buf: ref Bytes, pos: i64) -> i64 {
        assert(pos >= 0 && len >= 0 && pos + len <= buf.len())
        let h = (w0 ^ (w1 << 5)).wrapping_mul(0x9E3779B97F4A7C15)
        let mut at = (h >> (64 - BITS)).bits_as_i64()
        while true {
            if self.slots[at].count == 0 {
                self.claim(at, w0, w1, len, buf, pos)
                assert(at >= 0 && at < SLOTS)
                return at
            }
            if self.slots[at].w0 == w0 && self.slots[at].w1 == w1 && self.slots[at].len == len
                && (len <= 16 || self.slots[at].name[0..<len] == buf[pos..<pos + len]) {
                assert(at >= 0 && at < SLOTS)
                return at
            }
            at = (at + 1) & (SLOTS - 1)
        }
        panic("unreachable")
    }

    // A station seen for the first time: once per name, so kept out of the
    // loop that runs once per row.
    @inline(never)
    fn claim(ref mut self, at: i64, w0: u64, w1: u64, len: i64, buf: ref Bytes, pos: i64) {
        assert(pos >= 0 && len >= 0 && pos + len <= buf.len())
        self.used += 1
        if self.used > 10000 { panic("more than 10 000 stations") }
        self.slots[at] = Slot { w0: w0, w1: w1, len: len, name: buf[pos..<pos + len].clone(), min: 999, max: -999, sum: 0, count: 0 }
    }

    // Every row of `buf` that starts before `stop`.
    fn run(ref mut self, buf: ref Bytes, stop: i64) {
        if stop == 0 { return }
        // A row starts before `stop`, which leaves `MARGIN` bytes behind it.
        assert(stop > 0 && stop <= buf.len() - MARGIN; message: "a row starts `MARGIN` bytes before the end")
        let mut pos = 0
        // The `;` after each row's start, 32 bytes a compare: a row is about
        // 14 bytes, so it is mostly in the first 32, and the test whether it
        // was found is predicted. A name has at most 100 bytes, so at most
        // four loads look for it, each inside the input.
        let semis: simd::Fixed[u8, 32] = simd::splat(';')
        while true {
            let mut sep = pos
            while true {
                // The input's, so checked here, and a fact from here on.
                if !(sep >= pos && sep - pos <= 100) { panic("a station's name has at most 100 bytes") }
                let block: simd::Fixed[u8, 32] = simd::load(buf, sep)
                let hit = block.eq(semis).first() ?? -1
                if hit >= 0 {
                    sep += hit
                    break
                }
                sep += 32
            }
            // The name's first two words, zeros past its end: each cut to
            // the name by a mask looked up by its length, without a branch on
            // the length, which names of varying length would mispredict.
            let len = sep - pos
            if !(len >= 0 && len <= 100) { panic("a station's name has at most 100 bytes") }
            let w0 = buf.u64_le(pos) & FIRST[len]
            let w1 = buf.u64_le(pos + 8) & SECOND[len]

            let at = self.slot_of(w0, w1, len, buf, pos)
            // The temperature, `-?d?d.d`, from one word: where the `.` is
            // says how long it is, and the sign becomes a mask.
            let w = buf.u64_le(sep + 1)
            let dot = (!w & 0x10101000).trailing_zeros()
            if dot > 28 { panic("a row without a temperature") }
            let signed = ((!w) << 59).bits_as_i64() >> 63
            let keep = !(signed.bits_as_u64() & 0xFF)
            let digits = ((w & keep) << (28 - dot)) & 0x0F000F0F00
            let abs = (digits.wrapping_mul(0x640A0001) >> 32) & 0x3FF
            let v = (abs.bits_as_i64() ^ signed) - signed

            if v < self.slots[at].min { self.slots[at].min = v }
            if v > self.slots[at].max { self.slots[at].max = v }
            self.slots[at].sum += v
            self.slots[at].count += 1

            pos = sep + 4 + ((dot >> 3) as i64)
            if pos >= stop { break }
        }
    }

    // Every row of `b`: those up to the last `MARGIN` bytes in place, the
    // rest from a copy with zeros behind it.
    fn scan(ref mut self, b: ref Bytes) {
        let n = b.len()
        let mut cut = 0
        if n > MARGIN {
            cut = n - MARGIN
            while cut > 0 && b[cut - 1] != '\n' { cut -= 1 }
        }
        self.run(b, cut)
        let mut tail: Vec[u8] = []
        for c in b[cut..<n] { tail.push(c) }
        for _ in 0..<MARGIN { tail.push(0) }
        self.run(Bytes(tail), n - cut)
    }

    fn merge(ref mut self, other: Table) {
        for s in other.slots {
            if s.count > 0 {
                let at = self.slot_of(s.w0, s.w1, s.len, s.name, 0)
                if s.min < self.slots[at].min { self.slots[at].min = s.min }
                if s.max > self.slots[at].max { self.slots[at].max = s.max }
                self.slots[at].sum += s.sum
                self.slots[at].count += s.count
            }
        }
    }
}

fn main() throws {
    let path = cli::args().nth(1) ?? "measurements.txt"
    let b = fs::map(ref path, fs::Root::Anywhere) catch {
        eprintln(f"cannot open {path}: {error}")
        return
    }
    let mut all = table()
    if task::cores() == 1 {
        all.scan(b)
    } else {
        all = b.par_pieces('\n').fold(table, fn(t, piece) { t.scan(piece) }, Table::merge; order: par::Order::Any)
    }

    // Each station's name is checked as UTF-8 here, once.
    let mut index: collections::HashMap[ref String, i64] = collections::HashMap()
    let mut at = 0
    while at < SLOTS {
        if all.slots[at].count > 0 { index.insert(all.slots[at].name.text(), at) }
        at += 1
    }
    let mut names = index.keys().collect()
    names.sort()
    let body = names
        .map fn(name) {
            let s = all.slots[index[name] ?? panic("a key")]
            f"{name}={tenths(s.min)}/{tenths(mean(s.sum, s.count))}/{tenths(s.max)}"
        }
        .join(", ")
    println(f"{{{body}}}")
}

// The official rounding: a tie goes toward +infinity (`naive.nika`).
fn mean(sum: i64, count: i64) -> i64 {
    return (2 * sum + count).div_euclid(2 * count)
}

fn tenths(t: i64) -> String {
    if t < 0 {
        let a = -t
        return f"-{a / 10}.{a % 10}"
    }
    return f"{t / 10}.{t % 10}"
}

// The One Billion Row Challenge: the lowest, mean and highest temperature
// of each weather station, from one reading per line.
//
//     printf 'Hamburg;12.0\nBulawayo;8.9\nHamburg;-3.4\n' > measurements.txt
//     nikaia run -- measurements.txt
//     {Bulawayo=8.9/8.9/8.9, Hamburg=-3.4/4.3/12.0}
//
// To parse on every core, add `user-parallelism = "yes"` under `[build]` in
// nikaia.toml; for timing, also `opt-level = 3` under `[build.x86_64-linux]`.
// Name the project `brc`, not `1brc`: a package name cannot start with a digit.
//
// Shows: a `grammar` that describes the whole file, a parallel `par_fold`,
// names that are slices of a memory-mapped file rather than copies, and a
// format spec inside an f-string.
//
// To run it, copy this file to `src/main.nika` of a project (guide/getting-started.md).

use std::fs
use std::cli
use std::text
use std::collections

// ---------------------------------------------------------------------------
// 1. A grammar for the file
// ---------------------------------------------------------------------------

// A grammar for the whole file, not for one line: it is called once and walks
// the bytes itself. There is no tokenizer, so station names may contain blanks,
// dots and non-ASCII text.
grammar Measurements {

    // An uppercase rule is *lexical*: nothing is skipped between its parts, so
    // "Hamburg " with a trailing blank is a station of its own.
    //
    // The `String` is a slice of the input, never a copy.
    //
    // The scan stops at the `;` or at `frame_end`, the line end that
    // MEASUREMENT below declares. A plain `until(";")` could run into the next
    // line, and the compiler rejects it here.
    rule NAME -> String = s:until(";" | frame_end) -> { s }

    // Temperatures always carry exactly one decimal, so they are parsed as
    // tenths and all the arithmetic stays in integers until printing.
    // `digit{1,2}` is the width the format allows (-99.9 to 99.9), so
    // "123.0" is a parse error. `dec[i32](…)` hands over the number those
    // digits write rather than their text, so the action has nothing left to
    // decode.
    rule TENTHS -> i32 =
        neg:"-"? whole:dec[i32](digit{1,2}) "." frac:digit
        -> {
            let value = whole * 10 + text::digit_value(frac)
            if neg != null { -value } else { value }
        }

    // `=>` is a commit point: once the `;` is there, the line is a measurement,
    // and a bad temperature after it is reported, not retried as something else.
    //
    // `@frame` says a measurement can be found from any offset by scanning to
    // the next "\n", so the file can be cut into pieces and parsed in parallel.
    // The compiler checks that no scan inside the frame can run past it (see NAME).
    @frame(boundary: "\n")
    rule MEASUREMENT -> Reading =
        name:NAME ";" => temp:TENTHS frame_end
        -> { Reading { name, temp } }

    // `MEASUREMENT*` would collect a billion readings into memory. A fold
    // instead passes each one to an accumulator, so the parse runs in
    // constant space. `par_fold` also says how two partial results merge;
    // with that and the frame, the compiler writes the splitting, one
    // accumulator per core and the final merge. The merge is exact because
    // everything is summed in integer tenths.
    entry rule file -> Summary =
        par_fold(MEASUREMENT, Summary, fn(acc, m) { acc.record(m) },
                 Summary::merge)
}

struct Reading {
    name: String,
    temp: i32,
}

// ---------------------------------------------------------------------------
// 2. Aggregation
// ---------------------------------------------------------------------------

struct Stats {
    min: i32,
    max: i32,
    sum: i64,
    count: i64,
}

impl Stats {
    fn(first: i32) -> Stats {
        return Stats { min: first, max: first, sum: first as i64, count: 1 }
    }

    fn add(ref mut self, temp: i32) {
        if temp < self.min { self.min = temp }
        if temp > self.max { self.max = temp }
        self.sum += temp as i64
        self.count += 1
    }

    fn merge(ref mut self, other: ref Stats) {
        if other.min < self.min { self.min = other.min }
        if other.max > self.max { self.max = other.max }
        self.sum += other.sum
        self.count += other.count
    }

    // The mean in tenths, rounded as the challenge rounds: a tie goes up,
    // toward +infinity (`Math.round` in the official Java baseline). Exact
    // integer arithmetic, no float: `div_euclid` divides rounding down, also
    // when the sum is negative.
    fn mean(ref self) -> i64 {
        return (2 * self.sum + self.count).div_euclid(2 * self.count)
    }
}

// The stations seen so far and their statistics. The keys are slices of the
// input too.
struct Summary {
    stations: collections::HashMap[String, Stats],
}

impl Summary {
    fn() -> Summary {
        return Summary { stations: collections::HashMap() }
    }

    fn record(ref mut self, m: Reading) {
        self.stations
            .entry(m.name)
            .and_modify fn (stats) { stats.add(m.temp) }
            .or_insert_with fn { Stats(m.temp) }
    }

    fn merge(ref mut self, other: Summary) {
        for (name, stats) in other.stations.drain() {
            self.stations
                .entry(name)
                .and_modify fn (into) { into.merge(ref stats) }
                .or_insert(stats)
        }
    }
}

// ---------------------------------------------------------------------------
// 3. Main: the I/O half
// ---------------------------------------------------------------------------

fn main() throws {
    let path = cli::args().nth(1) ?? "measurements.txt"

    // `fs::map` maps the file into memory instead of reading it, so a file of
    // many gigabytes is never copied, and every station name points into it.
    // No `async` or `await` is written: I/O does not block, and the code reads
    // straight through. `catch` runs its block when the call fails, with the
    // failure named `error`.
    let data = fs::map(ref path, fs::Root::Anywhere) catch {
        eprintln(f"cannot open {path}: {error}")
        return
    }

    // The whole parse. With no `catch`, a bad line leaves `main` through
    // `throws` and ends the program with an error that gives its line number
    // in the whole file, even when the file was parsed in pieces.
    let totals = Measurements::file(data)

    let mut names = totals.stations.keys().collect()
    names.sort()

    let body = names
        .map fn(name) {
            // A map lookup gives a `Stats?` that may be absent. `??` supplies
            // the other case; the key came from this map, so it cannot happen.
            let s = totals.stations[name] ?? panic(f"{name} was a key a moment ago")
            f"{name}={tenths(s.min as i64)}/{tenths(s.mean())}/{tenths(s.max as i64)}"
        }
        .join(", ")

    // In an f-string, `{{` and `}}` are literal braces.
    println(f"{{{body}}}")
}

// A value in tenths as the challenge prints it: one decimal, never `-0.0`.
fn tenths(t: i64) -> String {
    if t < 0 {
        let a = -t
        return f"-{a / 10}.{a % 10}"
    }
    return f"{t / 10}.{t % 10}"
}

//! AI variant 2: the same aggregation, written by an AI assistant over several
//! rounds, each measured with callgrind and improved. One vector compare finds
//! the `;` (SSE2, AVX2 where the CPU level has it), the name is hashed a word
//! at a time (crc32c where the CPU level has it, a multiply mixer otherwise),
//! the temperature is parsed without a branch, the stations live in an own
//! open-addressing table, and the end of the input is run from a zero-padded
//! copy, so that nothing reads past the end of the mapping. See
//! `benches/brc/README.md`.

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

/// The main loop only starts rows below `len - MARGIN`; every read of a row
/// (the scan, the name's words, the temperature's word) stays below the row's
/// start + MARGIN.
const MARGIN: usize = 256;

fn main() {
    let path = std::env::args()
        .nth(1)
        .expect("a measurements file is required");
    let file = std::fs::File::open(&path).expect("open the measurements");
    // SAFETY: the file is not changed while the program runs.
    let map = unsafe { memmap2::Mmap::map(&file).expect("map the measurements") };
    let data: &[u8] = ↦
    let cut = cut(data);
    let tail = tail(&data[cut..]);
    let mut stats = Table::new();
    run(data, cut, &mut stats);
    run(&tail, data.len() - cut, &mut stats);
    report(&stats);
}

/// Where the main loop ends: just after the last newline at or before
/// `len - MARGIN` (0 for a short file).
fn cut(data: &[u8]) -> usize {
    if data.len() <= MARGIN {
        return 0;
    }
    let mut i = data.len() - MARGIN;
    while i > 0 && data[i - 1] != b'\n' {
        i -= 1;
    }
    i
}

/// The rows after `cut`, copied into a buffer padded with zeros, and a `;`
/// right after them.
fn tail(rest: &[u8]) -> Vec<u8> {
    let mut v = vec![0u8; rest.len() + MARGIN];
    v[..rest.len()].copy_from_slice(rest);
    v[rest.len()] = b';';
    v
}

#[inline(never)]
fn run<'a>(buf: &'a [u8], stop: usize, stats: &mut Table<'a>) {
    let p = buf.as_ptr();
    let mut pos = 0usize;
    while pos < stop {
        // SAFETY: a row starts below `stop`, which is at least MARGIN bytes
        // before the end of `buf`, and every read stays within MARGIN of it.
        unsafe {
            let sep = find_semi(p, pos);
            let h = name_hash(p.add(pos), sep - pos);
            let (v, n) = temp_swar(p.add(sep + 1));
            stats.update(h, &buf[pos..sep], v);
            pos = sep + 1 + n;
        }
    }
}

// ---- the `;`: compare a vector of bytes, take the first set bit -------------

#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
const W: usize = 32;
#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))]
const W: usize = 16;

#[cfg(target_arch = "x86_64")]
#[inline(always)]
unsafe fn mask_semi(p: *const u8) -> u32 {
    #[cfg(target_feature = "avx2")]
    {
        let v = _mm256_loadu_si256(p as *const __m256i);
        _mm256_movemask_epi8(_mm256_cmpeq_epi8(v, _mm256_set1_epi8(b';' as i8))) as u32
    }
    #[cfg(not(target_feature = "avx2"))]
    {
        let v = _mm_loadu_si128(p as *const __m128i);
        _mm_movemask_epi8(_mm_cmpeq_epi8(v, _mm_set1_epi8(b';' as i8))) as u32
    }
}

/// The `;` at or after `pos`.
#[cfg(target_arch = "x86_64")]
#[inline(always)]
unsafe fn find_semi(p: *const u8, pos: usize) -> usize {
    let mut base = pos;
    loop {
        let m = mask_semi(p.add(base));
        if m != 0 {
            return base + m.trailing_zeros() as usize;
        }
        base += W;
    }
}

/// The `;` at or after `pos`, a word at a time where there is no x86 vector.
#[cfg(not(target_arch = "x86_64"))]
#[inline(always)]
unsafe fn find_semi(p: *const u8, pos: usize) -> usize {
    const ONES: u64 = 0x0101_0101_0101_0101;
    let mut base = pos;
    loop {
        let x = word(p.add(base)) ^ (ONES * b';' as u64);
        let m = x.wrapping_sub(ONES) & !x & (ONES << 7);
        if m != 0 {
            return base + (m.trailing_zeros() / 8) as usize;
        }
        base += 8;
    }
}

// ---- the name's hash, from its words ----------------------------------------

#[inline(always)]
unsafe fn word(p: *const u8) -> u64 {
    u64::from_le((p as *const u64).read_unaligned())
}

/// The low `n` bytes of `w` (all of it when `n` is 8 or more).
#[inline(always)]
fn low_bytes(w: u64, n: usize) -> u64 {
    #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
    {
        // SAFETY: the target has bmi2.
        unsafe { _bzhi_u64(w, (n * 8) as u32) }
    }
    #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
    {
        if n >= 8 {
            w
        } else {
            w & ((1u64 << (n * 8)) - 1)
        }
    }
}

/// A hash of the name's first 16 bytes and its length; the table compares the
/// whole name. crc32c where the CPU level has it, else a multiply mixer.
#[inline(always)]
unsafe fn name_hash(p: *const u8, len: usize) -> u64 {
    let a = low_bytes(word(p), len);
    let b = low_bytes(word(p.add(8)), len.saturating_sub(8));
    #[cfg(all(target_arch = "x86_64", target_feature = "sse4.2"))]
    {
        _mm_crc32_u64(_mm_crc32_u64(len as u64, a), b)
    }
    #[cfg(not(all(target_arch = "x86_64", target_feature = "sse4.2")))]
    {
        mix(mix(len as u64, a), b)
    }
}

#[cfg(not(all(target_arch = "x86_64", target_feature = "sse4.2")))]
#[inline(always)]
fn mix(h: u64, w: u64) -> u64 {
    (h.rotate_left(5) ^ w).wrapping_mul(0x51_7c_c1_b7_27_22_0a_95)
}

// ---- the temperature, without a branch --------------------------------------

/// The word after the `;`: (tenths, bytes up to and including the `\n`).
#[inline(always)]
unsafe fn temp_swar(p: *const u8) -> (i64, usize) {
    let w = word(p);
    let dot = (!w & 0x1010_1000).trailing_zeros();
    let signed = (((!w) << 59) as i64) >> 63;
    let keep = !((signed as u64) & 0xFF);
    let digits = ((w & keep) << (28 - dot)) & 0x0F_000F_0F00;
    let abs = (digits.wrapping_mul(0x640a_0001) >> 32) & 0x3FF;
    (((abs as i64) ^ signed) - signed, (dot >> 3) as usize + 3)
}

// ---- the stations: open addressing, linear probing ---------------------------

const BITS: u32 = 12;

#[derive(Clone, Copy)]
struct Slot<'a> {
    name: &'a [u8],
    h: u64,
    /// min, max, sum, count.
    s: [i64; 4],
}

struct Table<'a> {
    slots: Vec<Slot<'a>>,
}

impl<'a> Table<'a> {
    fn new() -> Self {
        let empty = Slot {
            name: &[],
            h: 0,
            s: [0; 4],
        };
        Table {
            slots: vec![empty; 1 << BITS],
        }
    }

    #[inline(always)]
    fn update(&mut self, h: u64, name: &'a [u8], v: i64) {
        let mask = (1usize << BITS) - 1;
        let mut i = (h.wrapping_mul(0x9E37_79B9_7F4A_7C15) >> (64 - BITS)) as usize;
        loop {
            // SAFETY: `i` is masked to the table's length.
            let s = unsafe { self.slots.get_unchecked_mut(i) };
            if s.h == h && s.name == name {
                let e = &mut s.s;
                if v < e[0] {
                    e[0] = v
                }
                if v > e[1] {
                    e[1] = v
                }
                e[2] += v;
                e[3] += 1;
                return;
            }
            if s.name.is_empty() {
                *s = Slot {
                    name,
                    h,
                    s: [v, v, v, 1],
                };
                return;
            }
            i = (i + 1) & mask;
        }
    }
}

// ---- the output --------------------------------------------------------------

fn report(stats: &Table) {
    let mut rows: Vec<_> = stats.slots.iter().filter(|s| !s.name.is_empty()).collect();
    rows.sort_unstable_by(|a, b| a.name.cmp(b.name));
    let body: Vec<String> = rows
        .into_iter()
        .map(|r| {
            let name = std::str::from_utf8(r.name).expect("a station name is text");
            let s = r.s;
            format!(
                "{}={}/{}/{}",
                name,
                tenths(s[0]),
                tenths((2 * s[2] + s[3]).div_euclid(2 * s[3])),
                tenths(s[1])
            )
        })
        .collect();
    println!("{{{}}}", body.join(", "));
}

fn tenths(t: i64) -> String {
    let a = t.unsigned_abs();
    format!("{}{}.{}", if t < 0 { "-" } else { "" }, a / 10, a % 10)
}

//! AI variant 1: written by an AI assistant at medium effort, in one request,
//! as an optimised solution.
//!
//! What a 1BRC entry looks like before it is parallelised: `mmap`, raw bytes
//! with no UTF-8 validation, one `memchr` searcher over the whole file, the
//! separator found by walking back from the end of the line, integer
//! temperatures and a word-at-a-time hash. The program `examples/1brc.nika` is
//! measured against; see `benches/brc/README.md`.

use std::collections::HashMap;
use std::hash::{BuildHasherDefault, Hasher};

fn main() {
    let path = std::env::args()
        .nth(1)
        .expect("a measurements file is required");
    ai1(&path);
}

/// `rustc-hash`'s mixer, written out rather than depended on.
///
/// The station name is the key, it is short, and the default `SipHash` costs
/// about a quarter of the whole program when it is the only thing between a
/// line and its slot. A 1BRC entry replaces it; so does this.
#[derive(Default)]
struct Fx(u64);

const SEED: u64 = 0x51_7c_c1_b7_27_22_0a_95;

impl Fx {
    #[inline]
    fn add(&mut self, w: u64) {
        self.0 = (self.0.rotate_left(5) ^ w).wrapping_mul(SEED);
    }
}

impl Hasher for Fx {
    fn finish(&self) -> u64 {
        self.0
    }

    #[inline]
    fn write(&mut self, mut bytes: &[u8]) {
        while bytes.len() >= 8 {
            self.add(u64::from_ne_bytes(
                bytes[..8].try_into().expect("eight bytes"),
            ));
            bytes = &bytes[8..];
        }
        if !bytes.is_empty() {
            let mut buf = [0u8; 8];
            buf[..bytes.len()].copy_from_slice(bytes);
            self.add(u64::from_ne_bytes(buf));
        }
    }
}

/// What a 1BRC entry looks like before it is parallelised.
fn ai1(path: &str) {
    let file = std::fs::File::open(path).expect("open the measurements");
    // SAFETY: the mapping is read-only and private, and the one thing no type
    // system rules out - another process rewriting the file underneath it - is
    // the documented caveat of every memory map. `std::fs::map` accepts it on
    // the same terms (ADR-016).
    let map = unsafe { memmap2::Mmap::map(&file).expect("map the measurements") };
    let data: &[u8] = ↦

    let mut stats: HashMap<&[u8], [i64; 4], BuildHasherDefault<Fx>> = HashMap::default();

    // One searcher for the whole file rather than one per line. The separator
    // is then found by walking *back* from the end of the line, because the
    // format fixes the temperature at no more than five characters - the same
    // fact the example's grammar writes down as `"-"? digit{1,2} "." digit`.
    let mut start = 0usize;
    for end in memchr::memchr_iter(b'\n', data) {
        let line = &data[start..end];
        start = end + 1;

        let mut sep = line.len() - 1;
        while line[sep] != b';' {
            sep -= 1;
        }
        let name = &line[..sep];

        let mut t = &line[sep + 1..];
        let neg = t[0] == b'-';
        if neg {
            t = &t[1..];
        }
        // `dd.d` or `d.d`: the two widths the format allows, branchlessly
        // enough that the compiler keeps both in registers.
        let v: i64 = if t.len() == 4 {
            (t[0] - b'0') as i64 * 100 + (t[1] - b'0') as i64 * 10 + (t[3] - b'0') as i64
        } else {
            (t[0] - b'0') as i64 * 10 + (t[2] - b'0') as i64
        };
        let v = if neg { -v } else { v };

        let e = stats.entry(name).or_insert([i64::MAX, i64::MIN, 0, 0]);
        if v < e[0] {
            e[0] = v
        }
        if v > e[1] {
            e[1] = v
        }
        e[2] += v;
        e[3] += 1;
    }

    let mut names: Vec<&[u8]> = stats.keys().copied().collect();
    names.sort_unstable();
    report(names.into_iter().map(|n| {
        let s = stats[n];
        // The only place this half looks at the text as text, and the only
        // place it could fail on input `fs::map` would have rejected outright.
        let name = std::str::from_utf8(n).expect("a station name is text");
        (name, s[0], s[1], s[2], s[3])
    }));
}

/// The one line the benchmark asks for. Shared, so that a difference in the
/// output is never a difference in how it was printed.
fn report<'a>(rows: impl Iterator<Item = (&'a str, i64, i64, i64, i64)>) {
    let body: Vec<String> = rows
        .map(|(name, min, max, sum, count)| {
            format!(
                "{}={}/{}/{}",
                name,
                tenths(min),
                tenths(mean_tenths(sum, count)),
                tenths(max)
            )
        })
        .collect();
    println!("{{{}}}", body.join(", "));
}

/// The official rounding (the Java baseline's `Math.round(x * 10.0) / 10.0`):
/// a tie goes toward +infinity. `sum` is in tenths, so the mean in tenths is
/// `floor((2 * sum + count) / (2 * count))`, exact and with no float. With
/// readings of at most 99.9 the numerator stays below 2^63 up to about 4.6e15
/// rows, far past the challenge's 1e9.
fn mean_tenths(sum: i64, count: i64) -> i64 {
    (2 * sum + count).div_euclid(2 * count)
}

/// A value in tenths as the challenge prints it: one decimal, never `-0.0`.
fn tenths(t: i64) -> String {
    let a = t.unsigned_abs();
    format!("{}{}.{}", if t < 0 { "-" } else { "" }, a / 10, a % 10)
}