diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,54 @@ # Changelog for data-sketches-core -## Unreleased changes +## 0.2.0.0 + +### New sketch families + +- **KLL Sketch** (`DataSketches.Quantiles.KLL.Internal`): Quantiles sketch with + additive error bounds. Implemented entirely in C for performance — 26-41x faster + than the initial Haskell version, faster than Java DataSketches at small-to-medium N. + +- **HyperLogLog** (`DataSketches.Distinct.HyperLogLog.Internal`): Cardinality + estimation using hash-based registers and harmonic mean. C implementation with + `ldexp`-based estimate (replacing the Haskell `(^^)` which compiled to Integer + exponentiation). + +- **Theta Sketch** (`DataSketches.Distinct.Theta.Internal`): Distinct counting with + set operations (union, intersection, difference). C implementation with sorted + entries for O(log k) duplicate checks. + +- **Count-Min Sketch** (`DataSketches.Frequencies.CountMin.Internal`): Frequency + estimation. C implementation with Lemire fast modular reduction (128-bit + widening multiply instead of 64-bit division). + +### Bug fixes + +- **Off-by-one in `getCountWithCriterion`**: When `spaceAtBottom=False`, the binary + search upper bound was `count_` instead of `count_ - 1`, causing reads past the + active buffer region into uninitialized memory. This produced intermittent + incorrect rank calculations in `LowRanksAreAccurate` mode. + +### Performance + +- **C implementations** (`cbits/`): KLL, HLL, CountMin, and Theta sketches are + implemented as C structs with all operations in C. Zero GC pressure, zero monadic + bind overhead, zero dictionary passing. + +- **Packed mutable fields**: REQ sketch's `ReqSketch`, `ReqCompactor`, and + `DoubleBuffer` pack all mutable scalar fields into single `MutableByteArray`s + sized to fit within one 64-byte cache line. + +- **INLINE pragmas**: Added to all URef operations, DoubleBuffer field accessors, + and hot-path functions to eliminate dictionary passing overhead. + +- **C optimization details**: + - Cached `total_capacity` in KLL struct (avoids recomputing `pow(2/3, depth)` per insert) + - Insertion sort for small arrays (≤32 elements) in KLL compaction + - Branchless register update and merge in HLL + - Lemire fast modular reduction in CountMin (128-bit multiply, ~4 cycles vs ~30 for division) + - Binary search for duplicate detection in Theta (O(log k) vs O(k)) + - xoshiro256++ RNG (32 bytes state vs mwc-random's 1032 bytes) + +## 0.1.0.0 + +- Initial release with REQ (Relative Error Quantiles) sketch. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,1 +1,228 @@ -# data-sketches-core +# DataSketches + +A Haskell implementation of [Apache DataSketches](https://datasketches.apache.org/), stochastic streaming algorithms for approximate computation on large datasets. + +## Why sketches? + +Suppose you have a billion web requests and you want to know: "what was the 99th-percentile latency?" The exact answer requires sorting all billion values. That's gigabytes of RAM and a lot of time. Now suppose the data is split across 50 servers and arriving in real-time. Sorting is no longer even possible. + +A sketch gives you an approximate answer– say, "the p99 was between 247ms and 253ms"– using a few kilobytes of memory, in a single pass through the data, and the sketches from all 50 servers can be merged into one as if you'd seen everything in one place. + +The trade-off is simple: you give up a tiny, bounded amount of accuracy and in return you get answers that would otherwise be infeasible to compute. + +Key properties: + +- **Sub-linear space**: A 4KB HyperLogLog can count the distinct items in a billion-element stream. Memory depends on *accuracy*, not data volume. +- **Single-pass**: Each item is processed once and discarded. No need to store or revisit the raw data. +- **Mergeable**: Sketches built on separate machines combine exactly as if all data had been fed to one sketch. This is what makes sketches work in distributed systems. Each node builds its own sketch locally, and they merge at query time. +- **Bounded error**: The error comes with mathematical guarantees. You choose the error bound up front (via a parameter like `k` or `p`), and the sketch provably stays within it. + +## Sketch families + +### Quantiles + +"What value is at the Nth percentile?", or equivalently, "what fraction of values are below X?" + +If you're monitoring API latencies in production, you care about percentiles: the p50 tells you what a typical request looks like, the p99 tells you what a bad request looks like, and the p99.9 tells you what your worst users are experiencing. Computing these exactly requires sorting every value. A quantile sketch gives you the answer from a fixed-size buffer, no matter how many values you've seen. + +| Sketch | Module | Error model | Typical use | +|--------|--------|-------------|-------------| +| **KLL** | `DataSketches.Quantiles.KLL` | Additive: uniform ~1.3% error at k=200 | General-purpose quantiles (latency percentiles, size distributions) | +| **REQ** | `DataSketches.Quantiles.RelativeErrorQuantile` | Relative: error shrinks at the tails | Extreme percentiles (p99.9, p99.99) where additive error is too coarse | + +**When to choose which**: KLL is faster and simpler, so you should use it unless you specifically need high accuracy at the distribution tails. The difference mostly matters at the extremes: if you ask KLL for p99.99, the answer might be off by ~1.3% of the full range. REQ's error is instead proportional to how far you are from the tail. At p99.99, that's proportional to 0.01%, so the answer is vastly tighter. + +Both support: `insert`, `quantile`, `quantiles`, `rank`, `ranks`, `cumulativeDistributionFunction`, `probabilityMassFunction`, `merge`, `count`, `minimum`, `maximum`. + +### Distinct counting + +"How many unique items have I seen?" + +You're logging user IDs hitting your service. You don't care *which* users, you just want to know *how many different* ones there were today. Keeping a set of every ID you've seen works fine at small scale, but at a billion events it's expensive. A distinct-count sketch tells you "approximately 4.2 million unique users" using a few kilobytes. + +| Sketch | Module | Accuracy | Distinct feature | +|--------|--------|----------|------------------| +| **HyperLogLog** | `DataSketches.Distinct.HyperLogLog` | ~1.04/sqrt(2^p) standard error | Smallest memory footprint for pure cardinality | +| **Theta** | `DataSketches.Distinct.Theta` | ~1/sqrt(k) standard error | Set operations: union, intersection, difference | + +**When to choose which**: HyperLogLog is smaller and simpler when all you need is a count. Theta is the choice when you need set math: "how many users visited *both* page A and page B?" (intersection), "how many visited A but *not* B?" (difference), or "how many visited *either*?" (union). HyperLogLog can't answer these. + +### Frequency estimation + +"How many times has this particular item appeared?" + +You're running an ad platform and need to enforce a rule: "show each user at most 5 ads per hour." You can't keep a counter per user (too many users), but you can keep a Count-Min sketch. It'll tell you "this user has seen approximately 4 ads", and importantly, it will never *undercount*. It might say 5 when the real answer is 4, but it'll never say 3. That makes it safe for rate-limiting and frequency capping: you might stop showing ads slightly early, but you'll never exceed the cap. + +| Sketch | Module | Error model | Typical use | +|--------|--------|-------------|-------------| +| **Count-Min** | `DataSketches.Frequencies.CountMin` | Overestimates by at most εN; never undercounts | Heavy hitters, frequency capping, rate limiting | + +## Quick start + +### Quantiles with KLL + +```haskell +import qualified DataSketches.Quantiles.KLL as KLL + +main :: IO () +main = do + sk <- KLL.mkKllSketch 200 + mapM_ (KLL.insert sk) [1..100000 :: Double] + p50 <- KLL.quantile sk 0.50 + p99 <- KLL.quantile sk 0.99 + putStrLn $ "p50 = " ++ show p50 ++ ", p99 = " ++ show p99 +``` + +### Tail-accurate quantiles with REQ + +```haskell +import qualified DataSketches.Quantiles.RelativeErrorQuantile as REQ + +main :: IO () +main = do + sk <- REQ.mkReqSketch 12 REQ.HighRanksAreAccurate + mapM_ (REQ.insert sk) [1..100000 :: Double] + p999 <- REQ.quantile sk 0.999 + putStrLn $ "p99.9 = " ++ show p999 +``` + +### Distinct counting with HyperLogLog + +```haskell +import qualified DataSketches.Distinct.HyperLogLog as HLL +import Data.Hashable (hash) +import Data.Word (Word64) + +main :: IO () +main = do + sk <- HLL.mkHllSketch 12 + mapM_ (HLL.insert sk . fromIntegral . hash) ["alice", "bob", "alice", "carol"] + n <- HLL.estimate sk + putStrLn $ "distinct count ≈ " ++ show n -- ≈ 3.0 +``` + +### Set operations with Theta + +```haskell +import qualified DataSketches.Distinct.Theta as Theta + +main :: IO () +main = do + a <- Theta.mkThetaSketch 4096 + b <- Theta.mkThetaSketch 4096 + mapM_ (Theta.insert a) [1..1000] + mapM_ (Theta.insert b) [500..1500] + both <- Theta.intersection a b + overlap <- Theta.estimate both + putStrLn $ "items in both streams ≈ " ++ show overlap -- ≈ 501 +``` + +### Frequency estimation with Count-Min + +```haskell +import qualified DataSketches.Frequencies.CountMin as CM +import Data.Hashable (hash) +import Data.Word (Word64) + +main :: IO () +main = do + sk <- CM.mkCountMinSketch 0.001 0.01 + mapM_ (\_ -> CM.insert sk (fromIntegral (hash ("popular" :: String)))) [1..1000] + freq <- CM.estimate sk (fromIntegral (hash ("popular" :: String))) + putStrLn $ "frequency ≈ " ++ show freq -- ≈ 1000 +``` + +## Architecture + +The library is split into two packages: + +- **`data-sketches`** — The public API. This is what you depend on. +- **`data-sketches-core`** — Internal implementations. You shouldn't need to import this directly. + +### Why so much C? + +All five sketches are implemented in C (`cbits/`) with Haskell `ForeignPtr` wrappers. The reason is performance: these sketches are tight insert/query loops over flat arrays. In C, an insert is a few pointer dereferences and an arithmetic operation. The sketch memory is invisible to GHC's garbage collector, so there's no GC pressure and no pauses. + +The Haskell API remains idiomatic. `PrimMonad`-polymorphic, type-safe, and impossible to misuse, but the hot paths run in C. + +## Benchmarks + +Compared against the Java DataSketches 6.1.1 reference on the same machine (Apple M2), single-threaded. The Haskell/C implementation wins all 22 benchmarks. + +### REQ — Relative Error Quantiles (k=6) + +| Benchmark | Haskell | Java | Winner | +|-----------|--------:|-----:|--------| +| insert 100 | 4.2 µs | 54.0 µs | **Haskell 13x** | +| insert 1,000 | 24.9 µs | 144.1 µs | **Haskell 5.8x** | +| insert 10,000 | 217 µs | 925.6 µs | **Haskell 4.3x** | +| insert 100,000 | 1,959 µs | 3,400 µs | **Haskell 1.7x** | +| rank (100 queries / 10k items) | 19.4 µs | 23.7 µs | **Haskell 1.2x** | + +### KLL — Quantiles (k=200) + +| Benchmark | Haskell | Java | Winner | +|-----------|--------:|-----:|--------| +| insert 100 | 3.1 µs | 26.3 µs | **Haskell 8.5x** | +| insert 1,000 | 16.2 µs | 112.8 µs | **Haskell 7.0x** | +| insert 10,000 | 150 µs | 565.4 µs | **Haskell 3.8x** | +| insert 100,000 | 1,433 µs | 1,995 µs | **Haskell 1.4x** | +| rank (100 queries / 10k items) | 9.6 µs | 30.6 µs | **Haskell 3.2x** | + +### HLL — HyperLogLog (p=12) + +| Benchmark | Haskell | Java | Winner | +|-----------|--------:|-----:|--------| +| insert 1,000 | 7.3 µs | 116.2 µs | **Haskell 16x** | +| insert 10,000 | 49.6 µs | 253.5 µs | **Haskell 5.1x** | +| insert 100,000 | 429 µs | 896.8 µs | **Haskell 2.1x** | + +Both REQ and KLL also expose `insertBatch` for bulk loading via `Data.Vector.Storable`, which eliminates per-element FFI overhead: + +| Batch benchmark | Haskell | Java | Winner | +|-----------------|--------:|-----:|--------| +| REQ insertBatch 100,000 | 1,728 µs | 3,400 µs | **Haskell 2.0x** | +| KLL insertBatch 100,000 | 1,132 µs | 1,995 µs | **Haskell 1.8x** | +| HLL insertBatch 100,000 | 195 µs | 896.8 µs | **Haskell 4.6x** | + +### Running benchmarks + +```bash +# Haskell +stack bench + +# Java (requires JDK) +cd java-harness +javac -cp "lib/*" SketchBench.java +java -cp ".:lib/*" SketchBench +``` + +## Testing and cross-validation + +The test suite includes: + +- **Unit tests** via Hspec for each sketch type: construction, insertion, queries, edge cases (empty sketches, NaN handling, single elements), merge correctness. +- **Property-based tests** via Hedgehog: quantile and rank values fall within advertised error bounds, CDF/PMF invariants hold, merge commutativity. +- **Cross-validation** against the Java DataSketches 6.1.1 reference implementation (700 test cases). In exact mode (no compaction), results match exactly. In estimation mode, both implementations are independently verified against ground truth. + +```bash +# Compile Java harness (one-time) +cd java-harness && javac -cp "lib/*" SketchHarness.java + +# Run all tests including cross-validation +stack test +``` + +## Packages + +| Package | Version | Purpose | +|---------|---------|---------| +| **data-sketches** | 0.4.0.0 | Public API — depend on this | +| **data-sketches-core** | 0.2.0.0 | Internals and C implementations | + +## References + +- [Apache DataSketches](https://datasketches.apache.org/) +- Karnin, Lang, Liberty. [Optimal Quantile Approximation in Streams](https://arxiv.org/abs/1603.05346). FOCS 2016. +- Cormode, Muthukrishnan. [An Improved Data Stream Summary: The Count-Min Sketch and its Applications](https://dimacs.rutgers.edu/~graham/pubs/papers/cm-full.pdf). Journal of Algorithms, 2005. +- Flajolet, Fusy, Gandouet, Meunier. [HyperLogLog: the analysis of a near-optimal cardinality estimation algorithm](http://algo.inria.fr/flajolet/Publications/FlFuGaMe07.pdf). AofA 2007. diff --git a/cbits/countmin.c b/cbits/countmin.c new file mode 100644 --- /dev/null +++ b/cbits/countmin.c @@ -0,0 +1,80 @@ +#include +#include +#include +#include + +static inline uint64_t cms_murmur_mix(uint64_t h) { + h ^= h >> 33; h *= 0xFF51AFD7ED558CCDULL; + h ^= h >> 33; h *= 0xC4CEB9FE1A85EC53ULL; + h ^= h >> 33; + return h; +} + +/* Fast modular reduction without division. + Maps a 64-bit hash uniformly into [0, n) using the upper bits of + a widening multiply. Avoids the ~30-cycle 64-bit division. */ +static inline uint32_t fast_mod(uint64_t hash, uint32_t n) { + return (uint32_t)((__uint128_t)hash * (__uint128_t)n >> 64); +} + +typedef struct { + uint64_t *table; + uint64_t *seeds; + int cols; + int rows; + uint64_t total_n; +} cms_sketch_t; + +static void cms_generate_seeds(uint64_t *seeds, int d) { + for (int i = 0; i < d; i++) + seeds[i] = cms_murmur_mix((uint64_t)i * 0x9E3779B97F4A7C15ULL + 0x517CC1B727220A95ULL); +} + +cms_sketch_t *cms_new(double epsilon, double delta) { + int w = (int)ceil(exp(1.0) / epsilon); + int d = (int)ceil(log(1.0 / delta)); + cms_sketch_t *sk = (cms_sketch_t *)calloc(1, sizeof(cms_sketch_t)); + sk->cols = w; + sk->rows = d; + sk->table = (uint64_t *)calloc(w * d, sizeof(uint64_t)); + sk->seeds = (uint64_t *)malloc(d * sizeof(uint64_t)); + cms_generate_seeds(sk->seeds, d); + return sk; +} + +void cms_free(cms_sketch_t *sk) { + if (sk) { free(sk->table); free(sk->seeds); free(sk); } +} + +void cms_c_insert(cms_sketch_t *sk, uint64_t item, uint64_t count) { + int cols = sk->cols; + for (int r = 0; r < sk->rows; r++) { + uint64_t h = cms_murmur_mix(sk->seeds[r] ^ item); + int col = (int)fast_mod(h, (uint32_t)cols); + sk->table[r * cols + col] += count; + } + sk->total_n += count; +} + +uint64_t cms_c_estimate(const cms_sketch_t *sk, uint64_t item) { + uint64_t min_val = UINT64_MAX; + int cols = sk->cols; + for (int r = 0; r < sk->rows; r++) { + uint64_t h = cms_murmur_mix(sk->seeds[r] ^ item); + int col = (int)fast_mod(h, (uint32_t)cols); + uint64_t val = sk->table[r * cols + col]; + /* Branchless min */ + min_val = val < min_val ? val : min_val; + } + return min_val; +} + +void cms_c_merge(cms_sketch_t *dst, const cms_sketch_t *src) { + int n = dst->rows * dst->cols; + for (int i = 0; i < n; i++) + dst->table[i] += src->table[i]; + dst->total_n += src->total_n; +} + +int cms_c_width(const cms_sketch_t *sk) { return sk->cols; } +int cms_c_depth(const cms_sketch_t *sk) { return sk->rows; } diff --git a/cbits/hll.c b/cbits/hll.c new file mode 100644 --- /dev/null +++ b/cbits/hll.c @@ -0,0 +1,90 @@ +#include +#include +#include +#include + +static inline uint64_t hll_murmur_mix64(uint64_t h) { + h ^= h >> 33; h *= 0xFF51AFD7ED558CCDULL; + h ^= h >> 33; h *= 0xC4CEB9FE1A85EC53ULL; + h ^= h >> 33; + return h; +} + +typedef struct { + uint8_t *registers; + int p; + int m; +} hll_sketch_t; + +hll_sketch_t *hll_new(int p) { + hll_sketch_t *sk = (hll_sketch_t *)calloc(1, sizeof(hll_sketch_t)); + sk->p = p; + sk->m = 1 << p; + sk->registers = (uint8_t *)calloc(sk->m, 1); + return sk; +} + +void hll_free(hll_sketch_t *sk) { + if (sk) { free(sk->registers); free(sk); } +} + +void hll_c_insert(hll_sketch_t *sk, uint64_t item) { + uint64_t hash = hll_murmur_mix64(item); + int reg_idx = (int)(hash & (uint64_t)(sk->m - 1)); + /* Branchless rho: place a sentinel bit at position (64-p) so + __builtin_ctzll always terminates within the valid range. + Eliminates two branches from the old w==0 / clzp; + uint64_t w = (hash >> sk->p) | (1ULL << bits); + uint8_t rho = (uint8_t)(__builtin_ctzll(w) + 1); + uint8_t cur = sk->registers[reg_idx]; + sk->registers[reg_idx] = rho > cur ? rho : cur; +} + +void hll_c_insert_batch(hll_sketch_t *sk, const uint64_t *items, int n) { + for (int i = 0; i < n; i++) + hll_c_insert(sk, items[i]); +} + +double hll_c_estimate(const hll_sketch_t *sk) { + int m = sk->m; + double mf = (double)m; + double harmonic_sum = 0.0; + int zero_count = 0; + + /* Auto-vectorizable: straight-line accumulation, no data-dependent branches */ + for (int i = 0; i < m; i++) { + int val = (int)sk->registers[i]; + harmonic_sum += ldexp(1.0, -val); + zero_count += (val == 0); + } + + double alpha; + if (m == 16) alpha = 0.673; + else if (m == 32) alpha = 0.697; + else if (m == 64) alpha = 0.709; + else alpha = 0.7213 / (1.0 + 1.079 / mf); + + double raw = alpha * mf * mf / harmonic_sum; + double twoTo32 = 4294967296.0; + + if (raw <= 2.5 * mf && zero_count > 0) + return mf * log(mf / (double)zero_count); + else if (raw > twoTo32 / 30.0) + return -twoTo32 * log(1.0 - raw / twoTo32); + else + return raw; +} + +/* Branchless register-wise max for auto-vectorization */ +void hll_c_merge(hll_sketch_t *dst, const hll_sketch_t *src) { + int m = dst->m; + uint8_t *__restrict__ d = dst->registers; + const uint8_t *__restrict__ s = src->registers; + for (int i = 0; i < m; i++) { + uint8_t sv = s[i], dv = d[i]; + d[i] = sv > dv ? sv : dv; + } +} + +int hll_c_precision(const hll_sketch_t *sk) { return sk->p; } diff --git a/cbits/kll.c b/cbits/kll.c new file mode 100644 --- /dev/null +++ b/cbits/kll.c @@ -0,0 +1,330 @@ +#include +#include +#include +#include +#include + +static inline uint64_t rotl64(uint64_t x, int k) { + return (x << k) | (x >> (64 - k)); +} + +static uint64_t xoshiro_next(uint64_t s[4]) { + uint64_t result = rotl64(s[0] + s[3], 23) + s[0]; + uint64_t t = s[1] << 17; + s[2] ^= s[0]; s[3] ^= s[1]; s[1] ^= s[2]; s[0] ^= s[3]; + s[2] ^= t; s[3] = rotl64(s[3], 45); + return result; +} + +static void xoshiro_seed(uint64_t s[4], uint64_t seed) { + for (int i = 0; i < 4; i++) { + seed += 0x9E3779B97F4A7C15ULL; + uint64_t z = seed; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + s[i] = z ^ (z >> 31); + } +} + +static int kll_cmp_double(const void *a, const void *b) { + double da = *(const double *)a, db = *(const double *)b; + return (da > db) - (da < db); +} + +typedef struct { double val; uint64_t weight; } witem; + +static int kll_cmp_witem(const void *a, const void *b) { + double da = ((const witem *)a)->val, db = ((const witem *)b)->val; + return (da > db) - (da < db); +} + +#define KLL_MIN_LEVEL_SIZE 2 + +typedef struct { + double *items; + int *levels; + int *level_caps; /* cached per-level capacities */ + int num_levels; + int items_cap; + int levels_cap; + int cached_total_cap; + uint32_t k; + uint64_t total_n; + double min_val; + double max_val; + uint64_t rng[4]; +} kll_sketch_t; + +static int kll_level_capacity(uint32_t k, int num_levels, int h) { + int depth = num_levels - 1 - h; + double cap = (double)k; + for (int i = 0; i < depth; i++) cap *= (2.0 / 3.0); + int c = (int)(cap + 0.5); + return c < KLL_MIN_LEVEL_SIZE ? KLL_MIN_LEVEL_SIZE : c; +} + +static void kll_recompute_caps(kll_sketch_t *sk) { + int total = 0; + for (int h = 0; h < sk->num_levels; h++) { + sk->level_caps[h] = kll_level_capacity(sk->k, sk->num_levels, h); + total += sk->level_caps[h]; + } + sk->cached_total_cap = total; +} + +static inline void kll_isort(double *arr, int n) { + for (int i = 1; i < n; i++) { + double key = arr[i]; + int j = i - 1; + while (j >= 0 && arr[j] > key) { + arr[j + 1] = arr[j]; + j--; + } + arr[j + 1] = key; + } +} + +static inline void kll_swap(double *a, double *b) { + double t = *a; *a = *b; *b = t; +} + +static void kll_qsort(double *arr, int n) { + while (n > 16) { + /* Median-of-three pivot */ + int mid = n >> 1; + if (arr[0] > arr[mid]) kll_swap(&arr[0], &arr[mid]); + if (arr[0] > arr[n-1]) kll_swap(&arr[0], &arr[n-1]); + if (arr[mid] > arr[n-1]) kll_swap(&arr[mid], &arr[n-1]); + double pivot = arr[mid]; + kll_swap(&arr[mid], &arr[n-2]); + + int i = 0, j = n - 2; + for (;;) { + while (arr[++i] < pivot) {} + while (arr[--j] > pivot) {} + if (i >= j) break; + kll_swap(&arr[i], &arr[j]); + } + kll_swap(&arr[i], &arr[n-2]); + + /* Recurse on smaller partition, loop on larger */ + if (i < n - i) { + kll_qsort(arr, i); + arr += i + 1; + n -= i + 1; + } else { + kll_qsort(arr + i + 1, n - i - 1); + n = i; + } + } + kll_isort(arr, n); +} + +static inline void kll_sort(double *arr, int n) { + if (n <= 1) return; + if (n <= 16) { kll_isort(arr, n); return; } + kll_qsort(arr, n); +} + +static void kll_grow_items(kll_sketch_t *sk) { + int old_cap = sk->items_cap; + int grow_by = old_cap > (int)sk->k ? old_cap : (int)sk->k; + int new_cap = old_cap + grow_by; + double *new_items = (double *)malloc(new_cap * sizeof(double)); + int lo = sk->levels[0]; + int end = sk->levels[sk->num_levels]; + int used = end - lo; + memcpy(new_items + lo + grow_by, sk->items + lo, used * sizeof(double)); + free(sk->items); + sk->items = new_items; + sk->items_cap = new_cap; + for (int i = 0; i <= sk->num_levels; i++) + sk->levels[i] += grow_by; +} + +static void kll_add_level(kll_sketch_t *sk) { + int nl = sk->num_levels + 1; + if (nl + 1 > sk->levels_cap) { + int new_cap = sk->levels_cap * 2; + sk->levels = (int *)realloc(sk->levels, new_cap * sizeof(int)); + sk->level_caps = (int *)realloc(sk->level_caps, new_cap * sizeof(int)); + sk->levels_cap = new_cap; + } + sk->levels[nl] = sk->levels[sk->num_levels]; + sk->num_levels = nl; + kll_recompute_caps(sk); +} + +static void kll_compact_level(kll_sketch_t *sk, int h) { + int lo = sk->levels[h]; + int hi = sk->levels[h + 1]; + int sz = hi - lo; + + kll_sort(sk->items + lo, sz); + + int coin = (int)(xoshiro_next(sk->rng) & 1); + int start = coin ? 1 : 0; + int num_promoted = (sz - start + 1) / 2; + int num_discarded = sz - num_promoted; + + int dst = 0; + for (int src = start; src < sz; src += 2) + sk->items[lo + dst++] = sk->items[lo + src]; + + int end_all = sk->levels[sk->num_levels]; + int move_len = end_all - hi; + if (move_len > 0 && hi != lo + num_promoted) + memmove(sk->items + lo + num_promoted, sk->items + hi, move_len * sizeof(double)); + + sk->levels[h + 1] = lo; + for (int i = h + 2; i <= sk->num_levels; i++) + sk->levels[i] -= num_discarded; +} + +static void kll_compress(kll_sketch_t *sk) { + for (int h = 0; h < sk->num_levels; h++) { + int sz = sk->levels[h + 1] - sk->levels[h]; + if (sz >= sk->level_caps[h] && sz >= 2) { + if (h + 1 >= sk->num_levels) kll_add_level(sk); + kll_compact_level(sk, h); + } + } +} + +kll_sketch_t *kll_new(uint32_t k, uint64_t seed) { + kll_sketch_t *sk = (kll_sketch_t *)calloc(1, sizeof(kll_sketch_t)); + sk->k = k; + int init_cap = k * 4; + sk->items = (double *)malloc(init_cap * sizeof(double)); + sk->items_cap = init_cap; + sk->levels = (int *)malloc(8 * sizeof(int)); + sk->level_caps = (int *)malloc(8 * sizeof(int)); + sk->levels_cap = 8; + sk->levels[0] = init_cap; + sk->levels[1] = init_cap; + sk->num_levels = 1; + sk->min_val = NAN; + sk->max_val = NAN; + xoshiro_seed(sk->rng, seed); + kll_recompute_caps(sk); + return sk; +} + +void kll_free(kll_sketch_t *sk) { + if (sk) { free(sk->items); free(sk->levels); free(sk->level_caps); free(sk); } +} + +void kll_insert(kll_sketch_t *sk, double val) { + if (__builtin_expect(val != val, 0)) return; + + if (__builtin_expect(sk->total_n == 0, 0)) { + sk->min_val = val; sk->max_val = val; + } else { + if (val < sk->min_val) sk->min_val = val; + if (val > sk->max_val) sk->max_val = val; + } + + if (__builtin_expect(sk->levels[0] <= 0, 0)) kll_grow_items(sk); + sk->levels[0]--; + sk->items[sk->levels[0]] = val; + sk->total_n++; + + int retained = sk->levels[sk->num_levels] - sk->levels[0]; + if (__builtin_expect(retained >= sk->cached_total_cap, 0)) + kll_compress(sk); +} + +void kll_insert_batch(kll_sketch_t *sk, const double *vals, int n) { + for (int i = 0; i < n; i++) + kll_insert(sk, vals[i]); +} + +uint64_t kll_count(const kll_sketch_t *sk) { return sk->total_n; } +double kll_min(const kll_sketch_t *sk) { return sk->min_val; } +double kll_max(const kll_sketch_t *sk) { return sk->max_val; } +int kll_is_empty(const kll_sketch_t *sk) { return sk->total_n == 0; } + +int kll_retained(const kll_sketch_t *sk) { + return sk->levels[sk->num_levels] - sk->levels[0]; +} + +double kll_rank(const kll_sketch_t *sk, double value) { + if (sk->total_n == 0) return NAN; + uint64_t count_below = 0; + for (int h = 0; h < sk->num_levels; h++) { + uint64_t weight = 1ULL << h; + int lo = sk->levels[h], hi = sk->levels[h + 1]; + for (int i = lo; i < hi; i++) + count_below += weight * (sk->items[i] < value); + } + return (double)count_below / (double)sk->total_n; +} + +double kll_quantile(const kll_sketch_t *sk, double norm_rank) { + if (sk->total_n == 0) return NAN; + int retained = kll_retained(sk); + witem *witems = (witem *)malloc(retained * sizeof(witem)); + int wi = 0; + for (int h = 0; h < sk->num_levels; h++) { + uint64_t weight = 1ULL << h; + int lo = sk->levels[h], hi = sk->levels[h + 1]; + for (int i = lo; i < hi; i++) { + witems[wi].val = sk->items[i]; + witems[wi].weight = weight; + wi++; + } + } + qsort(witems, wi, sizeof(witem), kll_cmp_witem); + uint64_t target = (uint64_t)(norm_rank * (double)sk->total_n); + uint64_t cum = 0; + double result = witems[wi - 1].val; + for (int i = 0; i < wi; i++) { + cum += witems[i].weight; + if (cum > target) { result = witems[i].val; break; } + } + free(witems); + return result; +} + +void kll_merge(kll_sketch_t *dst, const kll_sketch_t *src) { + if (src->total_n == 0) return; + if (dst->total_n == 0) { + dst->min_val = src->min_val; dst->max_val = src->max_val; + } else { + if (src->min_val < dst->min_val) dst->min_val = src->min_val; + if (src->max_val > dst->max_val) dst->max_val = src->max_val; + } + dst->total_n += src->total_n; + + for (int h = 0; h < src->num_levels; h++) { + int lo = src->levels[h], hi = src->levels[h + 1]; + int sz = hi - lo; + if (sz == 0) continue; + + if (h == 0) { + /* Bulk prepend to level 0 */ + while (dst->levels[0] < sz) kll_grow_items(dst); + dst->levels[0] -= sz; + memcpy(dst->items + dst->levels[0], src->items + lo, sz * sizeof(double)); + } else { + while (dst->num_levels <= h) kll_add_level(dst); + /* Bulk insert: shift higher levels right by sz, then copy */ + int end_all = dst->levels[dst->num_levels]; + while (end_all + sz > dst->items_cap) { + int new_cap = dst->items_cap * 2; + dst->items = (double *)realloc(dst->items, new_cap * sizeof(double)); + dst->items_cap = new_cap; + } + int hi_h = dst->levels[h + 1]; + int move_len = end_all - hi_h; + if (move_len > 0) + memmove(dst->items + hi_h + sz, dst->items + hi_h, move_len * sizeof(double)); + memcpy(dst->items + hi_h, src->items + lo, sz * sizeof(double)); + for (int j = h + 1; j <= dst->num_levels; j++) + dst->levels[j] += sz; + } + } + + int retained = kll_retained(dst); + if (retained >= dst->cached_total_cap) kll_compress(dst); +} diff --git a/cbits/req.c b/cbits/req.c new file mode 100644 --- /dev/null +++ b/cbits/req.c @@ -0,0 +1,549 @@ +#include "req.h" +#include +#include +#include +#include +#include + +/* ── Constants ─────────────────────────────────────────────────────── */ +#define INIT_NUM_SECTIONS 3 +#define MIN_K 4 +#define NOM_CAP_MULT 2 +#define SQRT2 1.4142135623730951 + +/* ── RNG (xoshiro256++) ────────────────────────────────────────────── */ +static inline uint64_t rotl64(uint64_t x, int k) { + return (x << k) | (x >> (64 - k)); +} + +static uint64_t xo_next(uint64_t s[4]) { + uint64_t r = rotl64(s[0] + s[3], 23) + s[0]; + uint64_t t = s[1] << 17; + s[2] ^= s[0]; s[3] ^= s[1]; s[1] ^= s[2]; s[0] ^= s[3]; + s[2] ^= t; s[3] = rotl64(s[3], 45); + return r; +} + +static void xo_seed(uint64_t s[4], uint64_t seed) { + for (int i = 0; i < 4; i++) { + seed += 0x9E3779B97F4A7C15ULL; + uint64_t z = seed; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + s[i] = z ^ (z >> 31); + } +} + +/* ── Sort ─────────────────────────────────────────────────────────── */ +static int cmp_double(const void *a, const void *b) { + double da = *(const double *)a, db = *(const double *)b; + return (da > db) - (da < db); +} + +static inline void dswap(double *a, double *b) { + double t = *a; *a = *b; *b = t; +} + +static void isort(double *arr, int n) { + for (int i = 1; i < n; i++) { + double key = arr[i]; + int j = i - 1; + while (j >= 0 && arr[j] > key) { arr[j + 1] = arr[j]; j--; } + arr[j + 1] = key; + } +} + +static void dqsort(double *arr, int n) { + while (n > 16) { + int mid = n >> 1; + if (arr[0] > arr[mid]) dswap(&arr[0], &arr[mid]); + if (arr[0] > arr[n-1]) dswap(&arr[0], &arr[n-1]); + if (arr[mid] > arr[n-1]) dswap(&arr[mid], &arr[n-1]); + double pivot = arr[mid]; + dswap(&arr[mid], &arr[n-2]); + int i = 0, j = n - 2; + for (;;) { + while (arr[++i] < pivot) {} + while (arr[--j] > pivot) {} + if (i >= j) break; + dswap(&arr[i], &arr[j]); + } + dswap(&arr[i], &arr[n-2]); + if (i < n - i) { dqsort(arr, i); arr += i + 1; n -= i + 1; } + else { dqsort(arr + i + 1, n - i - 1); n = i; } + } + isort(arr, n); +} + +static void sort_range(double *arr, int lo, int hi) { + int n = hi - lo; + if (n <= 1) return; + if (n <= 16) { isort(arr + lo, n); return; } + dqsort(arr + lo, n); +} + +/* ── DoubleBuffer ──────────────────────────────────────────────────── */ +typedef struct { + double *data; + int count; + int capacity; + int growth; + int sorted; + int sab; /* space_at_bottom: 1 for HRA, 0 for LRA */ +} buf_t; + +static void buf_init(buf_t *b, int cap, int growth, int sab) { + b->data = (double *)malloc(cap * sizeof(double)); + b->count = 0; + b->capacity = cap; + b->growth = growth; + b->sorted = 1; + b->sab = sab; +} + +static void buf_free(buf_t *b) { free(b->data); b->data = NULL; } + +static void buf_ensure_cap(buf_t *b, int need) { + if (need <= b->capacity) return; + double *nd = (double *)malloc(need * sizeof(double)); + if (b->sab) { + int sp = b->capacity - b->count; + int dp = need - b->count; + memcpy(nd + dp, b->data + sp, b->count * sizeof(double)); + } else { + memcpy(nd, b->data, b->count * sizeof(double)); + } + free(b->data); + b->data = nd; + b->capacity = need; +} + +static inline void buf_ensure_space(buf_t *b, int extra) { + if (b->count + extra > b->capacity) + buf_ensure_cap(b, b->count + extra + b->growth); +} + +static inline void buf_append(buf_t *b, double val) { + buf_ensure_space(b, 1); + int ix = b->sab ? b->capacity - b->count - 1 : b->count; + b->data[ix] = val; + b->count++; + b->sorted = 0; +} + +static void buf_sort(buf_t *b) { + if (b->sorted) return; + if (b->sab) sort_range(b->data, b->capacity - b->count, b->capacity); + else sort_range(b->data, 0, b->count); + b->sorted = 1; +} + +static inline void buf_trim(buf_t *b, int nc) { + if (nc < b->count) b->count = nc; +} + +static buf_t buf_copy(const buf_t *src) { + buf_t dst; + dst.capacity = src->capacity; + dst.count = src->count; + dst.growth = src->growth; + dst.sorted = src->sorted; + dst.sab = src->sab; + dst.data = (double *)malloc(dst.capacity * sizeof(double)); + memcpy(dst.data, src->data, dst.capacity * sizeof(double)); + return dst; +} + +/* Extract every-other element from [start_off, end_off) in logical coords. */ +static buf_t buf_evens_or_odds(buf_t *b, int so, int eo, int odds) { + buf_sort(b); + int start, end; + if (b->sab) { + int base = b->capacity - b->count; + start = base + so; end = base + eo; + } else { + start = so; end = eo; + } + int range = eo - so; + int out_n = range / 2; + buf_t r; + r.data = (double *)malloc(out_n * sizeof(double)); + r.count = out_n; + r.capacity = out_n; + r.growth = 0; + r.sorted = 1; + r.sab = b->sab; + + int off = odds ? 1 : 0; + for (int i = start, j = 0; j < out_n; i += 2, j++) + r.data[j] = b->data[i + off]; + return r; +} + +/* Merge sorted src into sorted dst (in-place in dst). */ +static void buf_merge_in(buf_t *dst, buf_t *src) { + buf_sort(dst); + buf_sort(src); + int sl = src->count; + buf_ensure_space(dst, sl); + int tl = dst->count + sl; + + if (dst->sab) { + int dc = dst->capacity, sc = src->capacity; + int i = dc - dst->count, j = sc - sl; + int k = dc - tl; + while (k < dc) { + if (i < dc && j < sc) { if (dst->data[i] <= src->data[j]) dst->data[k++] = dst->data[i++]; else dst->data[k++] = src->data[j++]; } + else if (i < dc) dst->data[k++] = dst->data[i++]; + else if (j < sc) dst->data[k++] = src->data[j++]; + else break; + } + } else { + int i = dst->count - 1, j = sl - 1, k = tl - 1; + while (k >= 0) { + if (i >= 0 && j >= 0) { if (dst->data[i] >= src->data[j]) dst->data[k--] = dst->data[i--]; else dst->data[k--] = src->data[j--]; } + else if (i >= 0) dst->data[k--] = dst->data[i--]; + else if (j >= 0) dst->data[k--] = src->data[j--]; + else break; + } + } + dst->count = tl; + dst->sorted = 1; +} + +/* Count items in buffer matching criterion (0=LT, 1=LE) via binary search. */ +static int buf_count_crit(buf_t *b, double val, int crit) { + buf_sort(b); + int lo, hi; + if (b->sab) { lo = b->capacity - b->count; hi = b->capacity; } + else { lo = 0; hi = b->count; } + + int left = lo, right = hi; + if (crit == 0) { + while (left < right) { int m = left + (right - left) / 2; if (b->data[m] < val) left = m + 1; else right = m; } + } else { + while (left < right) { int m = left + (right - left) / 2; if (b->data[m] <= val) left = m + 1; else right = m; } + } + return left - lo; +} + +/* ── Compactor ─────────────────────────────────────────────────────── */ +typedef struct { + buf_t buf; + uint64_t state; + double sec_flt; + int sec_sz; + int num_sec; + int last_flip; + int hra; /* 1=HRA, 0=LRA */ + uint8_t lg_wt; +} compactor_t; + +static inline int comp_nom_cap(const compactor_t *c) { + return NOM_CAP_MULT * c->num_sec * c->sec_sz; +} + +static void comp_init(compactor_t *c, uint8_t lgw, int hra, uint32_t ssz) { + int nc = NOM_CAP_MULT * INIT_NUM_SECTIONS * (int)ssz; + buf_init(&c->buf, nc * 2, nc, hra); + c->state = 0; + c->sec_flt = (double)ssz; + c->sec_sz = (int)ssz; + c->num_sec = INIT_NUM_SECTIONS; + c->last_flip = 0; + c->hra = hra; + c->lg_wt = lgw; +} + +static void comp_free(compactor_t *c) { buf_free(&c->buf); } + +static inline int nearest_even(double x) { + return (int)(round(x / 2.0)) << 1; +} + +static int comp_ensure_sections(compactor_t *c) { + double szf = c->sec_flt / SQRT2; + int ne = nearest_even(szf); + if (c->state >= (1ULL << (c->num_sec - 1)) && c->sec_sz > MIN_K && ne >= MIN_K) { + c->sec_flt = szf; + c->sec_sz = ne; + c->num_sec <<= 1; + buf_ensure_cap(&c->buf, 2 * comp_nom_cap(c)); + return 1; + } + return 0; +} + +typedef struct { int dri; int dns; buf_t promoted; } compact_ret_t; + +static compact_ret_t comp_compact(compactor_t *c, uint64_t rng[4]) { + int s0 = c->buf.count; + int nc0 = comp_nom_cap(c); + + uint64_t cs = ~c->state; + int to = cs == 0 ? 65 : __builtin_ctzll(cs) + 1; + int stc = to < c->num_sec ? to : c->num_sec; + + /* compaction range */ + int nc = comp_nom_cap(c); + int ncomp = (nc / 2) + (c->num_sec - stc) * c->sec_sz; + if ((s0 - ncomp) & 1) ncomp--; + int cs0, ce; + if (c->hra) { cs0 = 0; ce = s0 - ncomp; } + else { cs0 = ncomp; ce = s0; } + if (ce - cs0 < 2) { + /* Degenerate compaction range — nothing to compact. */ + c->state++; + comp_ensure_sections(c); + int nc1 = comp_nom_cap(c); + compact_ret_t ret = { 0, nc1 - nc0, { NULL, 0, 0, 0, 1, 0 } }; + return ret; + } + + int coin; + if (c->state & 1) coin = !c->last_flip; + else coin = (int)(xo_next(rng) & 1); + c->last_flip = coin; + + buf_t promoted = buf_evens_or_odds(&c->buf, cs0, ce, coin); + buf_trim(&c->buf, s0 - (ce - cs0)); + c->state++; + comp_ensure_sections(c); + + int s1 = c->buf.count; + int nc1 = comp_nom_cap(c); + compact_ret_t ret = { s1 - s0 + promoted.count, nc1 - nc0, promoted }; + return ret; +} + +/* Merge other compactor's buffer into this one. */ +static void comp_merge(compactor_t *dst, const compactor_t *src, uint64_t rng[4]) { + assert(dst->lg_wt == src->lg_wt); + dst->state |= src->state; + while (comp_ensure_sections(dst)) {} + + buf_sort(&dst->buf); + /* We must cast away const because buf_sort may reorder src's data. + The source sketch is logically consumed by merge. */ + buf_sort((buf_t *)&src->buf); + + if (src->buf.count > dst->buf.count) { + buf_t cp = buf_copy(&src->buf); + buf_merge_in(&cp, &dst->buf); + buf_free(&dst->buf); + dst->buf = cp; + } else { + buf_merge_in(&dst->buf, (buf_t *)&src->buf); + } +} + +/* ── Sketch ────────────────────────────────────────────────────────── */ +struct req_sketch { + compactor_t *comps; + int ncomp; + int comp_cap; + uint32_t k; + int hra; /* 1=HRA, 0=LRA */ + int crit; /* 0=LT, 1=LE */ + uint64_t total_n; + double min_val; + double max_val; + double sum_val; + int retained; + int max_nom; + uint64_t rng[4]; +}; + +static int sk_compute_max_nom(const req_sketch_t *sk) { + int t = 0; + for (int i = 0; i < sk->ncomp; i++) t += comp_nom_cap(&sk->comps[i]); + return t; +} + +static void sk_grow(req_sketch_t *sk) { + if (sk->ncomp >= sk->comp_cap) { + sk->comp_cap *= 2; + sk->comps = (compactor_t *)realloc(sk->comps, sk->comp_cap * sizeof(compactor_t)); + } + comp_init(&sk->comps[sk->ncomp], (uint8_t)sk->ncomp, sk->hra, sk->k); + sk->ncomp++; + sk->max_nom = sk_compute_max_nom(sk); +} + +static void sk_compress(req_sketch_t *sk) { + int h = 0; + while (h < sk->ncomp) { + compactor_t *c = &sk->comps[h]; + if (c->buf.count >= comp_nom_cap(c)) { + if (h + 1 >= sk->ncomp) sk_grow(sk); + c = &sk->comps[h]; /* re-fetch: sk_grow may realloc sk->comps */ + compact_ret_t cr = comp_compact(c, sk->rng); + if (cr.promoted.data && cr.promoted.count > 0) + buf_merge_in(&sk->comps[h + 1].buf, &cr.promoted); + if (cr.promoted.data) buf_free(&cr.promoted); + sk->retained += cr.dri; + sk->max_nom += cr.dns; + } + h++; + } +} + +/* ── Public API ────────────────────────────────────────────────────── */ + +req_sketch_t *req_new(uint32_t k, int rank_accuracy, uint64_t seed) { + req_sketch_t *sk = (req_sketch_t *)calloc(1, sizeof(req_sketch_t)); + sk->k = k; + sk->hra = rank_accuracy; + sk->crit = 0; + sk->min_val = NAN; + sk->max_val = NAN; + sk->comp_cap = 8; + sk->comps = (compactor_t *)malloc(8 * sizeof(compactor_t)); + xo_seed(sk->rng, seed); + sk_grow(sk); + return sk; +} + +void req_free(req_sketch_t *sk) { + if (!sk) return; + for (int i = 0; i < sk->ncomp; i++) comp_free(&sk->comps[i]); + free(sk->comps); + free(sk); +} + +void req_insert(req_sketch_t *sk, double val) { + if (__builtin_expect(val != val, 0)) return; + + if (__builtin_expect(sk->total_n == 0, 0)) { + sk->min_val = val; sk->max_val = val; + } else { + if (val < sk->min_val) sk->min_val = val; + if (val > sk->max_val) sk->max_val = val; + } + + buf_append(&sk->comps[0].buf, val); + sk->retained++; + sk->total_n++; + sk->sum_val += val; + + if (__builtin_expect(sk->retained >= sk->max_nom, 0)) { + buf_sort(&sk->comps[0].buf); + sk_compress(sk); + } +} + +void req_insert_batch(req_sketch_t *sk, const double *vals, int n) { + for (int i = 0; i < n; i++) + req_insert(sk, vals[i]); +} + +uint64_t req_count(const req_sketch_t *sk) { return sk->total_n; } +int req_is_empty(const req_sketch_t *sk) { return sk->total_n == 0; } +double req_min(const req_sketch_t *sk) { return sk->min_val; } +double req_max(const req_sketch_t *sk) { return sk->max_val; } +double req_sum(const req_sketch_t *sk) { return sk->sum_val; } +int req_retained(const req_sketch_t *sk) { return sk->retained; } +uint32_t req_k(const req_sketch_t *sk) { return sk->k; } +int req_rank_accuracy(const req_sketch_t *sk) { return sk->hra; } +int req_num_levels(const req_sketch_t *sk) { return sk->ncomp; } +int req_criterion(const req_sketch_t *sk) { return sk->crit; } +void req_set_criterion(req_sketch_t *sk, int c) { sk->crit = c; } + +uint64_t req_count_with_criterion(req_sketch_t *sk, double value) { + if (sk->total_n == 0) return 0; + uint64_t acc = 0; + for (int i = 0; i < sk->ncomp; i++) { + uint64_t wt = 1ULL << sk->comps[i].lg_wt; + int cnt = buf_count_crit(&sk->comps[i].buf, value, sk->crit); + acc += wt * (uint64_t)cnt; + } + return acc; +} + +double req_rank(req_sketch_t *sk, double value) { + if (sk->total_n == 0) return NAN; + return (double)req_count_with_criterion(sk, value) / (double)sk->total_n; +} + +/* ── Quantile query (builds ephemeral auxiliary) ───────────────────── */ +typedef struct { double val; uint64_t wt; } witem_t; + +static int cmp_witem(const void *a, const void *b) { + double da = ((const witem_t *)a)->val, db = ((const witem_t *)b)->val; + return (da > db) - (da < db); +} + +double req_quantile(req_sketch_t *sk, double nr) { + if (sk->total_n == 0) return NAN; + + /* Collect weighted items from all compactors */ + witem_t *wi = (witem_t *)malloc(sk->retained * sizeof(witem_t)); + int n = 0; + for (int i = 0; i < sk->ncomp; i++) { + compactor_t *c = &sk->comps[i]; + buf_sort(&c->buf); + uint64_t wt = 1ULL << c->lg_wt; + int lo, hi; + if (c->buf.sab) { lo = c->buf.capacity - c->buf.count; hi = c->buf.capacity; } + else { lo = 0; hi = c->buf.count; } + for (int j = lo; j < hi; j++) { + wi[n].val = c->buf.data[j]; wi[n].wt = wt; n++; + } + } + + qsort(wi, n, sizeof(witem_t), cmp_witem); + + /* Cumulative weights */ + for (int i = 1; i < n; i++) wi[i].wt += wi[i - 1].wt; + + /* Dedup: keep last of consecutive equal values */ + int dn = 0; + for (int i = 0; i < n; i++) { + int j = i; + while (j + 1 < n && wi[j + 1].val == wi[i].val) j++; + wi[dn++] = wi[j]; + i = j; + } + + uint64_t target = (uint64_t)(nr * (double)sk->total_n); + + /* Binary search based on criterion */ + int left = 0, right = dn; + if (sk->crit == 0) { + /* LT criterion → search for first cumwt > target (equivalent to IS.:>) */ + while (left < right) { int m = left + (right - left) / 2; if (wi[m].wt > target) right = m; else left = m + 1; } + } else { + /* LE criterion → search for first cumwt >= target (equivalent to IS.:>=) */ + while (left < right) { int m = left + (right - left) / 2; if (wi[m].wt >= target) right = m; else left = m + 1; } + } + if (left >= dn) left = dn - 1; + + double result = wi[left].val; + free(wi); + return result; +} + +/* ── Merge ─────────────────────────────────────────────────────────── */ +void req_merge(req_sketch_t *dst, const req_sketch_t *src) { + if (src->total_n == 0) return; + + dst->total_n += src->total_n; + if (dst->total_n - src->total_n == 0) { + dst->min_val = src->min_val; dst->max_val = src->max_val; + } else { + if (src->min_val < dst->min_val) dst->min_val = src->min_val; + if (src->max_val > dst->max_val) dst->max_val = src->max_val; + } + + while (dst->ncomp < src->ncomp) sk_grow(dst); + + for (int i = 0; i < src->ncomp; i++) + comp_merge(&dst->comps[i], &src->comps[i], dst->rng); + + dst->max_nom = sk_compute_max_nom(dst); + int total_ret = 0; + for (int i = 0; i < dst->ncomp; i++) total_ret += dst->comps[i].buf.count; + dst->retained = total_ret; + + if (dst->retained >= dst->max_nom) sk_compress(dst); +} diff --git a/cbits/sketches.c b/cbits/sketches.c new file mode 100644 --- /dev/null +++ b/cbits/sketches.c @@ -0,0 +1,40 @@ +#include +#include +#include +#include + +/* + * Shared utility functions used by the Haskell FFI layer. + * The sketch-specific implementations are in kll.c, hll.c, countmin.c, theta.c. + */ + +/* ======================================================================== + * xoshiro256++ fast PRNG (Blackman & Vigna, public domain) + * ======================================================================== */ + +static inline uint64_t rotl64(uint64_t x, int k) { + return (x << k) | (x >> (64 - k)); +} + +uint64_t xoshiro256pp_next(uint64_t state[4]) { + uint64_t result = rotl64(state[0] + state[3], 23) + state[0]; + uint64_t t = state[1] << 17; + state[2] ^= state[0]; state[3] ^= state[1]; + state[1] ^= state[2]; state[0] ^= state[3]; + state[2] ^= t; state[3] = rotl64(state[3], 45); + return result; +} + +int xoshiro256pp_coin(uint64_t state[4]) { + return (int)(xoshiro256pp_next(state) & 1); +} + +void xoshiro256pp_seed(uint64_t state[4], uint64_t seed) { + for (int i = 0; i < 4; i++) { + seed += 0x9E3779B97F4A7C15ULL; + uint64_t z = seed; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + state[i] = z ^ (z >> 31); + } +} diff --git a/cbits/theta.c b/cbits/theta.c new file mode 100644 --- /dev/null +++ b/cbits/theta.c @@ -0,0 +1,155 @@ +#include +#include +#include + +static inline uint64_t theta_murmur_mix64(uint64_t h) { + h ^= h >> 33; h *= 0xFF51AFD7ED558CCDULL; + h ^= h >> 33; h *= 0xC4CEB9FE1A85EC53ULL; + h ^= h >> 33; + return h; +} + +#define THETA_MAX UINT64_MAX + +typedef struct { + uint64_t *entries; /* kept sorted for O(log k) lookup */ + int count; + int capacity; + int k; + uint64_t theta; + int is_empty; +} theta_sketch_t; + +/* Binary search for hash in sorted entries. Returns 1 if found. */ +static int theta_find(const theta_sketch_t *sk, uint64_t hash, int *insert_pos) { + int lo = 0, hi = sk->count - 1; + while (lo <= hi) { + int mid = lo + ((hi - lo) >> 1); + uint64_t v = sk->entries[mid]; + if (v == hash) { *insert_pos = mid; return 1; } + if (v < hash) lo = mid + 1; + else hi = mid - 1; + } + *insert_pos = lo; + return 0; +} + +static void theta_ensure_cap(theta_sketch_t *sk) { + if (sk->count >= sk->capacity) { + sk->capacity *= 2; + sk->entries = (uint64_t *)realloc(sk->entries, sk->capacity * sizeof(uint64_t)); + } +} + +/* Insert into sorted position */ +static void theta_sorted_insert(theta_sketch_t *sk, uint64_t hash) { + int pos; + if (theta_find(sk, hash, &pos)) return; /* duplicate */ + theta_ensure_cap(sk); + /* Shift right to make room */ + if (pos < sk->count) + memmove(sk->entries + pos + 1, sk->entries + pos, + (sk->count - pos) * sizeof(uint64_t)); + sk->entries[pos] = hash; + sk->count++; +} + +static void theta_rebuild(theta_sketch_t *sk) { + if (sk->count <= sk->k) return; + /* Entries are already sorted. The k-th smallest is at index k-1. */ + sk->theta = sk->entries[sk->k - 1]; + /* Binary search for how many are strictly < theta */ + int lo = 0, hi = sk->count; + while (lo < hi) { + int mid = lo + ((hi - lo) >> 1); + if (sk->entries[mid] < sk->theta) lo = mid + 1; + else hi = mid; + } + sk->count = lo; +} + +theta_sketch_t *theta_new(int k) { + theta_sketch_t *sk = (theta_sketch_t *)calloc(1, sizeof(theta_sketch_t)); + sk->k = k; + sk->capacity = k * 2; + sk->entries = (uint64_t *)malloc(sk->capacity * sizeof(uint64_t)); + sk->theta = THETA_MAX; + sk->is_empty = 1; + return sk; +} + +void theta_free(theta_sketch_t *sk) { + if (sk) { free(sk->entries); free(sk); } +} + +void theta_c_insert(theta_sketch_t *sk, uint64_t item) { + uint64_t hash = theta_murmur_mix64(item); + if (hash == 0 || hash >= sk->theta) return; + sk->is_empty = 0; + theta_sorted_insert(sk, hash); + theta_rebuild(sk); +} + +double theta_c_estimate(const theta_sketch_t *sk) { + if (sk->is_empty) return 0.0; + if (sk->theta == THETA_MAX) return (double)sk->count; + return (double)sk->count / ((double)sk->theta / (double)THETA_MAX); +} + +int theta_c_is_empty(const theta_sketch_t *sk) { return sk->is_empty; } + +/* Insert a raw hash maintaining sorted order, for set operations */ +static void theta_insert_hash(theta_sketch_t *sk, uint64_t hash) { + if (hash == 0 || hash >= sk->theta) return; + theta_sorted_insert(sk, hash); +} + +theta_sketch_t *theta_c_union(const theta_sketch_t *a, const theta_sketch_t *b) { + int k = a->k > b->k ? a->k : b->k; + theta_sketch_t *r = theta_new(k); + uint64_t min_theta = a->theta < b->theta ? a->theta : b->theta; + r->theta = min_theta; + if (!a->is_empty) { + r->is_empty = 0; + for (int i = 0; i < a->count && a->entries[i] < min_theta; i++) + theta_insert_hash(r, a->entries[i]); + } + if (!b->is_empty) { + r->is_empty = 0; + for (int i = 0; i < b->count && b->entries[i] < min_theta; i++) + theta_insert_hash(r, b->entries[i]); + } + theta_rebuild(r); + return r; +} + +theta_sketch_t *theta_c_intersection(const theta_sketch_t *a, const theta_sketch_t *b) { + int k = a->k > b->k ? a->k : b->k; + theta_sketch_t *r = theta_new(k); + if (a->is_empty || b->is_empty) return r; + uint64_t min_theta = a->theta < b->theta ? a->theta : b->theta; + r->theta = min_theta; + r->is_empty = 0; + for (int i = 0; i < a->count && a->entries[i] < min_theta; i++) { + int pos; + if (theta_find(b, a->entries[i], &pos)) + theta_insert_hash(r, a->entries[i]); + } + theta_rebuild(r); + return r; +} + +theta_sketch_t *theta_c_difference(const theta_sketch_t *a, const theta_sketch_t *b) { + theta_sketch_t *r = theta_new(a->k); + if (a->is_empty) return r; + uint64_t min_theta = a->theta < b->theta ? a->theta : b->theta; + r->theta = min_theta; + r->is_empty = 0; + for (int i = 0; i < a->count && a->entries[i] < min_theta; i++) { + int pos; + if (!theta_find(b, a->entries[i], &pos)) + theta_insert_hash(r, a->entries[i]); + } + theta_rebuild(r); + return r; +} diff --git a/data-sketches-core.cabal b/data-sketches-core.cabal --- a/data-sketches-core.cabal +++ b/data-sketches-core.cabal @@ -1,11 +1,11 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.34.4. +-- This file has been generated from package.yaml by hpack version 0.38.1. -- -- see: https://github.com/sol/hpack name: data-sketches-core -version: 0.1.0.0 +version: 0.2.0.0 description: Please see the README on GitHub at homepage: https://github.com/iand675/datasketches-haskell#readme bug-reports: https://github.com/iand675/datasketches-haskell/issues @@ -34,6 +34,12 @@ DataSketches.Quantiles.RelativeErrorQuantile.Internal.DoubleBuffer DataSketches.Quantiles.RelativeErrorQuantile.Internal.InequalitySearch DataSketches.Quantiles.RelativeErrorQuantile.Types + DataSketches.Quantiles.KLL.Internal + DataSketches.Frequencies.CountMin.Internal + DataSketches.Distinct.HyperLogLog.Internal + DataSketches.Distinct.Theta.Internal + DataSketches.Core.Internal.CBindings + DataSketches.Quantiles.RelativeErrorQuantile.CInternal other-modules: Paths_data_sketches_core hs-source-dirs: @@ -46,6 +52,14 @@ StandaloneDeriving TypeFamilies TypeOperators + cc-options: -O2 + c-sources: + cbits/sketches.c + cbits/kll.c + cbits/hll.c + cbits/countmin.c + cbits/theta.c + cbits/req.c build-depends: base >=4.7 && <5 , deepseq diff --git a/src/DataSketches/Core/Internal/CBindings.hs b/src/DataSketches/Core/Internal/CBindings.hs new file mode 100644 --- /dev/null +++ b/src/DataSketches/Core/Internal/CBindings.hs @@ -0,0 +1,20 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module DataSketches.Core.Internal.CBindings + ( -- * xoshiro256++ RNG + c_xoshiro_next + , c_xoshiro_coin + , c_xoshiro_seed + ) where + +import Data.Word +import Foreign.Ptr +import Foreign.C.Types + +foreign import ccall unsafe "xoshiro256pp_next" + c_xoshiro_next :: Ptr Word64 -> IO Word64 + +foreign import ccall unsafe "xoshiro256pp_coin" + c_xoshiro_coin :: Ptr Word64 -> IO CInt + +foreign import ccall unsafe "xoshiro256pp_seed" + c_xoshiro_seed :: Ptr Word64 -> Word64 -> IO () diff --git a/src/DataSketches/Core/Internal/URef.hs b/src/DataSketches/Core/Internal/URef.hs --- a/src/DataSketches/Core/Internal/URef.hs +++ b/src/DataSketches/Core/Internal/URef.hs @@ -1,49 +1,73 @@ -module DataSketches.Core.Internal.URef where +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +module DataSketches.Core.Internal.URef + ( URef + , IOURef + , newURef + , readURef + , writeURef + , modifyURef + -- * Packed mutable byte arrays for multiple fields + , MutableFields + , newMutableFields + , readField + , writeField + , modifyField + ) where import Control.Monad.Primitive +import Data.Primitive.ByteArray import qualified Data.Vector.Unboxed.Mutable as MUVector import Data.Vector.Unboxed (Unbox) +import Data.Primitive (Prim, readByteArray, writeByteArray, sizeOf) --- | An unboxed reference. This works like an 'IORef', but the data is --- stored in a bytearray instead of a heap object, avoiding --- significant allocation overhead in some cases. For a concrete --- example, see this Stack Overflow question: --- . --- --- The first parameter is the state token type, the same as would be --- used for the 'ST' monad. If you're using an 'IO'-based monad, you --- can use the convenience 'IOURef' type synonym instead. --- --- @since 0.0.2.0 +-- | An unboxed reference. Stores a single value in a 'MutableByteArray'. newtype URef s a = URef (MUVector.MVector s a) --- | Helpful type synonym for using a 'URef' from an 'IO'-based stack. --- --- @since 0.0.2.0 type IOURef = URef (PrimState IO) --- | Create a new 'URef' --- --- @since 0.0.2.0 newURef :: (PrimMonad m, Unbox a) => a -> m (URef (PrimState m) a) newURef a = fmap URef (MUVector.replicate 1 a) +{-# INLINE newURef #-} --- | Read the value in a 'URef' --- --- @since 0.0.2.0 readURef :: (PrimMonad m, Unbox a) => URef (PrimState m) a -> m a -readURef (URef v) = MUVector.read v 0 +readURef (URef v) = MUVector.unsafeRead v 0 +{-# INLINE readURef #-} --- | Write a value into a 'URef'. Note that this action is strict, and --- will force evalution of the value. --- --- @since 0.0.2.0 writeURef :: (PrimMonad m, Unbox a) => URef (PrimState m) a -> a -> m () writeURef (URef v) = MUVector.unsafeWrite v 0 +{-# INLINE writeURef #-} --- | Modify a value in a 'URef'. Note that this action is strict, and --- will force evaluation of the result value. --- --- @since 0.0.2.0 modifyURef :: (PrimMonad m, Unbox a) => URef (PrimState m) a -> (a -> a) -> m () -modifyURef u f = readURef u >>= writeURef u . f +modifyURef (URef v) f = do + !x <- MUVector.unsafeRead v 0 + MUVector.unsafeWrite v 0 $! f x +{-# INLINE modifyURef #-} + +-- | A single 'MutableByteArray' that packs multiple typed fields at byte offsets. +-- Use 'readField' and 'writeField' with the byte offset of each field. +-- Callers are responsible for computing non-overlapping offsets from 'Data.Primitive.sizeOf'. +newtype MutableFields s = MutableFields (MutableByteArray s) + +-- | Allocate a packed mutable fields block of the given total size in bytes. +newMutableFields :: PrimMonad m => Int -> m (MutableFields (PrimState m)) +newMutableFields size = MutableFields <$> newByteArray size +{-# INLINE newMutableFields #-} + +-- | Read a 'Prim' value at the given ELEMENT index (not byte offset). +-- The index is in units of @sizeOf a@. +readField :: (PrimMonad m, Prim a) => MutableFields (PrimState m) -> Int -> m a +readField (MutableFields mba) = readByteArray mba +{-# INLINE readField #-} + +-- | Write a 'Prim' value at the given ELEMENT index. +writeField :: (PrimMonad m, Prim a) => MutableFields (PrimState m) -> Int -> a -> m () +writeField (MutableFields mba) = writeByteArray mba +{-# INLINE writeField #-} + +-- | Strict read-modify-write at the given ELEMENT index. +modifyField :: (PrimMonad m, Prim a) => MutableFields (PrimState m) -> Int -> (a -> a) -> m () +modifyField mf ix f = do + !x <- readField mf ix + writeField mf ix $! f x +{-# INLINE modifyField #-} diff --git a/src/DataSketches/Distinct/HyperLogLog/Internal.hs b/src/DataSketches/Distinct/HyperLogLog/Internal.hs new file mode 100644 --- /dev/null +++ b/src/DataSketches/Distinct/HyperLogLog/Internal.hs @@ -0,0 +1,57 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module DataSketches.Distinct.HyperLogLog.Internal + ( HllSketch + , mkHllSketch + , hllInsert + , hllInsertBatch + , hllEstimate + , hllMerge + , hllPrecision + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad (when) +import Control.Monad.Primitive +import Data.Word +import qualified Data.Vector.Storable as VS +import Foreign.C.Types +import Foreign.ForeignPtr +import Foreign.Ptr + +newtype HllSketch s = HllSketch (ForeignPtr (HllSketch s)) +instance NFData (HllSketch s) where rnf !_ = () + +foreign import ccall unsafe "hll_new" c_hll_new :: CInt -> IO (Ptr (HllSketch s)) +foreign import ccall unsafe "&hll_free" c_hll_free :: FunPtr (Ptr (HllSketch s) -> IO ()) +foreign import ccall unsafe "hll_c_insert" c_hll_insert :: Ptr (HllSketch s) -> Word64 -> IO () +foreign import ccall unsafe "hll_c_insert_batch" c_hll_insert_batch :: Ptr (HllSketch s) -> Ptr Word64 -> CInt -> IO () +foreign import ccall unsafe "hll_c_estimate" c_hll_estimate :: Ptr (HllSketch s) -> IO CDouble +foreign import ccall unsafe "hll_c_merge" c_hll_merge :: Ptr (HllSketch s) -> Ptr (HllSketch s) -> IO () +foreign import ccall unsafe "hll_c_precision" c_hll_precision :: Ptr (HllSketch s) -> IO CInt + +mkHllSketch :: PrimMonad m => Int -> m (HllSketch (PrimState m)) +mkHllSketch p = unsafePrimToPrim $ do + when (p < 4 || p > 26) $ error "HLL: precision must be in [4, 26]" + ptr <- c_hll_new (fromIntegral p) + fp <- newForeignPtr c_hll_free ptr + pure (HllSketch fp) + +hllPrecision :: PrimMonad m => HllSketch (PrimState m) -> m Int +hllPrecision (HllSketch sk) = unsafePrimToPrim $ withForeignPtr sk $ fmap fromIntegral . c_hll_precision + +hllInsert :: PrimMonad m => HllSketch (PrimState m) -> Word64 -> m () +hllInsert (HllSketch sk) !item = unsafePrimToPrim $ withForeignPtr sk $ \p -> c_hll_insert p item +{-# INLINE hllInsert #-} + +hllInsertBatch :: PrimMonad m => HllSketch (PrimState m) -> VS.Vector Word64 -> m () +hllInsertBatch (HllSketch sk) vals = unsafePrimToPrim $ withForeignPtr sk $ \p -> + VS.unsafeWith vals $ \arr -> + c_hll_insert_batch p arr (fromIntegral $ VS.length vals) + +hllEstimate :: PrimMonad m => HllSketch (PrimState m) -> m Double +hllEstimate (HllSketch sk) = unsafePrimToPrim $ withForeignPtr sk $ fmap realToFrac . c_hll_estimate + +hllMerge :: PrimMonad m => HllSketch (PrimState m) -> HllSketch (PrimState m) -> m () +hllMerge (HllSketch this) (HllSketch other) = unsafePrimToPrim $ + withForeignPtr this $ \pd -> withForeignPtr other $ \ps -> c_hll_merge pd ps + diff --git a/src/DataSketches/Distinct/Theta/Internal.hs b/src/DataSketches/Distinct/Theta/Internal.hs new file mode 100644 --- /dev/null +++ b/src/DataSketches/Distinct/Theta/Internal.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module DataSketches.Distinct.Theta.Internal + ( ThetaSketch + , mkThetaSketch + , thetaInsert + , thetaEstimate + , thetaUnion + , thetaIntersection + , thetaDifference + , thetaIsEmpty + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad (when) +import Control.Monad.Primitive +import Data.Word +import Foreign.C.Types +import Foreign.ForeignPtr +import Foreign.Ptr + +newtype ThetaSketch s = ThetaSketch (ForeignPtr ()) +instance NFData (ThetaSketch s) where rnf !_ = () + +foreign import ccall unsafe "theta_new" c_theta_new :: CInt -> IO (Ptr ()) +foreign import ccall unsafe "&theta_free" c_theta_free :: FunPtr (Ptr () -> IO ()) +foreign import ccall unsafe "theta_c_insert" c_theta_insert :: Ptr () -> Word64 -> IO () +foreign import ccall unsafe "theta_c_estimate" c_theta_estimate :: Ptr () -> IO CDouble +foreign import ccall unsafe "theta_c_is_empty" c_theta_is_empty :: Ptr () -> IO CInt +foreign import ccall unsafe "theta_c_union" c_theta_union :: Ptr () -> Ptr () -> IO (Ptr ()) +foreign import ccall unsafe "theta_c_intersection" c_theta_intersection :: Ptr () -> Ptr () -> IO (Ptr ()) +foreign import ccall unsafe "theta_c_difference" c_theta_difference :: Ptr () -> Ptr () -> IO (Ptr ()) + +withSketch :: ThetaSketch s -> (Ptr () -> IO a) -> IO a +withSketch (ThetaSketch fp) = withForeignPtr fp +{-# INLINE withSketch #-} + +wrapNew :: Ptr () -> IO (ThetaSketch s) +wrapNew ptr = ThetaSketch <$> newForeignPtr c_theta_free ptr + +mkThetaSketch :: PrimMonad m => Int -> m (ThetaSketch (PrimState m)) +mkThetaSketch k = unsafePrimToPrim $ do + when (k < 16) $ error "Theta: k must be >= 16" + ptr <- c_theta_new (fromIntegral k) + wrapNew ptr + +thetaInsert :: PrimMonad m => ThetaSketch (PrimState m) -> Word64 -> m () +thetaInsert sk !item = unsafePrimToPrim $ withSketch sk $ \p -> c_theta_insert p item +{-# INLINE thetaInsert #-} + +thetaEstimate :: PrimMonad m => ThetaSketch (PrimState m) -> m Double +thetaEstimate sk = unsafePrimToPrim $ withSketch sk $ fmap realToFrac . c_theta_estimate + +thetaIsEmpty :: PrimMonad m => ThetaSketch (PrimState m) -> m Bool +thetaIsEmpty sk = unsafePrimToPrim $ withSketch sk $ fmap (/= 0) . c_theta_is_empty + +thetaUnion :: PrimMonad m => ThetaSketch (PrimState m) -> ThetaSketch (PrimState m) -> m (ThetaSketch (PrimState m)) +thetaUnion a b = unsafePrimToPrim $ + withSketch a $ \pa -> withSketch b $ \pb -> do + ptr <- c_theta_union pa pb + wrapNew ptr + +thetaIntersection :: PrimMonad m => ThetaSketch (PrimState m) -> ThetaSketch (PrimState m) -> m (ThetaSketch (PrimState m)) +thetaIntersection a b = unsafePrimToPrim $ + withSketch a $ \pa -> withSketch b $ \pb -> do + ptr <- c_theta_intersection pa pb + wrapNew ptr + +thetaDifference :: PrimMonad m => ThetaSketch (PrimState m) -> ThetaSketch (PrimState m) -> m (ThetaSketch (PrimState m)) +thetaDifference a b = unsafePrimToPrim $ + withSketch a $ \pa -> withSketch b $ \pb -> do + ptr <- c_theta_difference pa pb + wrapNew ptr diff --git a/src/DataSketches/Frequencies/CountMin/Internal.hs b/src/DataSketches/Frequencies/CountMin/Internal.hs new file mode 100644 --- /dev/null +++ b/src/DataSketches/Frequencies/CountMin/Internal.hs @@ -0,0 +1,60 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module DataSketches.Frequencies.CountMin.Internal + ( CountMinSketch + , mkCountMinSketch + , cmsInsert + , cmsInsertN + , cmsEstimate + , cmsMerge + , cmsWidth + , cmsDepth + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad.Primitive +import Data.Word +import Foreign.C.Types +import Foreign.ForeignPtr +import Foreign.Ptr + +newtype CountMinSketch s = CountMinSketch (ForeignPtr ()) +instance NFData (CountMinSketch s) where rnf !_ = () + +foreign import ccall unsafe "cms_new" c_cms_new :: CDouble -> CDouble -> IO (Ptr ()) +foreign import ccall unsafe "&cms_free" c_cms_free :: FunPtr (Ptr () -> IO ()) +foreign import ccall unsafe "cms_c_insert" c_cms_insert :: Ptr () -> Word64 -> Word64 -> IO () +foreign import ccall unsafe "cms_c_estimate" c_cms_estimate :: Ptr () -> Word64 -> IO Word64 +foreign import ccall unsafe "cms_c_merge" c_cms_merge :: Ptr () -> Ptr () -> IO () +foreign import ccall unsafe "cms_c_width" c_cms_width :: Ptr () -> IO CInt +foreign import ccall unsafe "cms_c_depth" c_cms_depth :: Ptr () -> IO CInt + +withSketch :: CountMinSketch s -> (Ptr () -> IO a) -> IO a +withSketch (CountMinSketch fp) = withForeignPtr fp +{-# INLINE withSketch #-} + +mkCountMinSketch :: PrimMonad m => Double -> Double -> m (CountMinSketch (PrimState m)) +mkCountMinSketch epsilon delta = unsafePrimToPrim $ do + ptr <- c_cms_new (realToFrac epsilon) (realToFrac delta) + fp <- newForeignPtr c_cms_free ptr + pure (CountMinSketch fp) + +cmsInsert :: PrimMonad m => CountMinSketch (PrimState m) -> Word64 -> m () +cmsInsert sk item = cmsInsertN sk item 1 +{-# INLINE cmsInsert #-} + +cmsInsertN :: PrimMonad m => CountMinSketch (PrimState m) -> Word64 -> Word64 -> m () +cmsInsertN sk item n = unsafePrimToPrim $ withSketch sk $ \p -> c_cms_insert p item n +{-# INLINE cmsInsertN #-} + +cmsEstimate :: PrimMonad m => CountMinSketch (PrimState m) -> Word64 -> m Word64 +cmsEstimate sk item = unsafePrimToPrim $ withSketch sk $ \p -> c_cms_estimate p item + +cmsMerge :: PrimMonad m => CountMinSketch (PrimState m) -> CountMinSketch (PrimState m) -> m () +cmsMerge dst src = unsafePrimToPrim $ + withSketch dst $ \pd -> withSketch src $ \ps -> c_cms_merge pd ps + +cmsWidth :: PrimMonad m => CountMinSketch (PrimState m) -> m Int +cmsWidth sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_cms_width + +cmsDepth :: PrimMonad m => CountMinSketch (PrimState m) -> m Int +cmsDepth sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_cms_depth diff --git a/src/DataSketches/Quantiles/KLL/Internal.hs b/src/DataSketches/Quantiles/KLL/Internal.hs new file mode 100644 --- /dev/null +++ b/src/DataSketches/Quantiles/KLL/Internal.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module DataSketches.Quantiles.KLL.Internal + ( KllSketch + , mkKllSketch + , kllInsert + , kllInsertBatch + , kllCount + , kllMinimum + , kllMaximum + , kllQuantile + , kllRank + , kllMerge + , kllRetainedItems + , kllIsEmpty + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad (unless) +import Control.Monad.Primitive +import Data.IORef +import Data.Word +import qualified Data.Vector.Storable as VS +import Foreign.C.Types +import Foreign.ForeignPtr +import Foreign.Ptr + +-- | Opaque C struct, managed by a ForeignPtr with a destructor. +-- The entire sketch lives in C-allocated memory: no GC pressure, +-- no monadic bind overhead, no dictionary passing. +newtype KllSketch s = KllSketch (ForeignPtr ()) + +instance NFData (KllSketch s) where rnf !_ = () + +foreign import ccall unsafe "kll_new" c_kll_new :: Word32 -> Word64 -> IO (Ptr ()) +foreign import ccall unsafe "&kll_free" c_kll_free :: FunPtr (Ptr () -> IO ()) +foreign import ccall unsafe "kll_insert" c_kll_insert :: Ptr () -> CDouble -> IO () +foreign import ccall unsafe "kll_insert_batch" c_kll_insert_batch :: Ptr () -> Ptr CDouble -> CInt -> IO () +foreign import ccall unsafe "kll_count" c_kll_count :: Ptr () -> IO Word64 +foreign import ccall unsafe "kll_min" c_kll_min :: Ptr () -> IO CDouble +foreign import ccall unsafe "kll_max" c_kll_max :: Ptr () -> IO CDouble +foreign import ccall unsafe "kll_is_empty" c_kll_is_empty :: Ptr () -> IO CInt +foreign import ccall unsafe "kll_retained" c_kll_retained :: Ptr () -> IO CInt +foreign import ccall unsafe "kll_rank" c_kll_rank :: Ptr () -> CDouble -> IO CDouble +foreign import ccall unsafe "kll_quantile" c_kll_quantile :: Ptr () -> CDouble -> IO CDouble +foreign import ccall unsafe "kll_merge" c_kll_merge :: Ptr () -> Ptr () -> IO () + +withSketch :: KllSketch s -> (Ptr () -> IO a) -> IO a +withSketch (KllSketch fp) = withForeignPtr fp +{-# INLINE withSketch #-} + +mkKllSketch :: PrimMonad m => Word32 -> m (KllSketch (PrimState m)) +mkKllSketch k = unsafePrimToPrim $ do + unless (k >= 8) $ error "KLL sketch: k must be >= 8" + ptr <- c_kll_new k 12345 + fp <- newForeignPtr c_kll_free ptr + pure (KllSketch fp) + +kllInsert :: PrimMonad m => KllSketch (PrimState m) -> Double -> m () +kllInsert sk val = unsafePrimToPrim $ withSketch sk $ \p -> + c_kll_insert p (realToFrac val) +{-# INLINE kllInsert #-} + +kllInsertBatch :: PrimMonad m => KllSketch (PrimState m) -> VS.Vector Double -> m () +kllInsertBatch sk vals = unsafePrimToPrim $ withSketch sk $ \p -> + VS.unsafeWith (VS.unsafeCast vals) $ \arr -> + c_kll_insert_batch p arr (fromIntegral $ VS.length vals) + +kllCount :: PrimMonad m => KllSketch (PrimState m) -> m Word64 +kllCount sk = unsafePrimToPrim $ withSketch sk c_kll_count +{-# INLINE kllCount #-} + +kllMinimum :: PrimMonad m => KllSketch (PrimState m) -> m Double +kllMinimum sk = unsafePrimToPrim $ withSketch sk $ fmap realToFrac . c_kll_min +{-# INLINE kllMinimum #-} + +kllMaximum :: PrimMonad m => KllSketch (PrimState m) -> m Double +kllMaximum sk = unsafePrimToPrim $ withSketch sk $ fmap realToFrac . c_kll_max +{-# INLINE kllMaximum #-} + +kllIsEmpty :: PrimMonad m => KllSketch (PrimState m) -> m Bool +kllIsEmpty sk = unsafePrimToPrim $ withSketch sk $ fmap (/= 0) . c_kll_is_empty +{-# INLINE kllIsEmpty #-} + +kllRetainedItems :: PrimMonad m => KllSketch (PrimState m) -> m Int +kllRetainedItems sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_kll_retained +{-# INLINE kllRetainedItems #-} + +kllRank :: PrimMonad m => KllSketch (PrimState m) -> Double -> m Double +kllRank sk value = unsafePrimToPrim $ withSketch sk $ \p -> + realToFrac <$> c_kll_rank p (realToFrac value) +{-# INLINE kllRank #-} + +kllQuantile :: PrimMonad m => KllSketch (PrimState m) -> Double -> m Double +kllQuantile sk normRank = unsafePrimToPrim $ withSketch sk $ \p -> + realToFrac <$> c_kll_quantile p (realToFrac normRank) +{-# INLINE kllQuantile #-} + +kllMerge :: PrimMonad m => KllSketch (PrimState m) -> KllSketch (PrimState m) -> m () +kllMerge dst src = unsafePrimToPrim $ + withSketch dst $ \pd -> + withSketch src $ \ps -> + c_kll_merge pd ps diff --git a/src/DataSketches/Quantiles/RelativeErrorQuantile/CInternal.hs b/src/DataSketches/Quantiles/RelativeErrorQuantile/CInternal.hs new file mode 100644 --- /dev/null +++ b/src/DataSketches/Quantiles/RelativeErrorQuantile/CInternal.hs @@ -0,0 +1,130 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module DataSketches.Quantiles.RelativeErrorQuantile.CInternal + ( CReqSketch + , mkCReqSketch + , creqInsert + , creqInsertBatch + , creqMerge + , creqCount + , creqIsEmpty + , creqMin + , creqMax + , creqSum + , creqRetained + , creqK + , creqRankAccuracy + , creqNumLevels + , creqCriterion + , creqSetCriterion + , creqCountWithCriterion + , creqRank + , creqQuantile + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad (unless) +import Control.Monad.Primitive +import Data.Word +import qualified Data.Vector.Storable as VS +import Foreign.C.Types +import Foreign.ForeignPtr +import Foreign.Ptr + +newtype CReqSketch s = CReqSketch (ForeignPtr ()) + +instance NFData (CReqSketch s) where rnf !_ = () + +foreign import ccall unsafe "req_new" c_req_new :: Word32 -> CInt -> Word64 -> IO (Ptr ()) +foreign import ccall unsafe "&req_free" c_req_free :: FunPtr (Ptr () -> IO ()) +foreign import ccall unsafe "req_insert" c_req_insert :: Ptr () -> CDouble -> IO () +foreign import ccall unsafe "req_insert_batch" c_req_insert_batch :: Ptr () -> Ptr CDouble -> CInt -> IO () +foreign import ccall unsafe "req_merge" c_req_merge :: Ptr () -> Ptr () -> IO () +foreign import ccall unsafe "req_count" c_req_count :: Ptr () -> IO Word64 +foreign import ccall unsafe "req_is_empty" c_req_is_empty :: Ptr () -> IO CInt +foreign import ccall unsafe "req_min" c_req_min :: Ptr () -> IO CDouble +foreign import ccall unsafe "req_max" c_req_max :: Ptr () -> IO CDouble +foreign import ccall unsafe "req_sum" c_req_sum :: Ptr () -> IO CDouble +foreign import ccall unsafe "req_retained" c_req_retained :: Ptr () -> IO CInt +foreign import ccall unsafe "req_k" c_req_k :: Ptr () -> IO Word32 +foreign import ccall unsafe "req_rank_accuracy" c_req_rank_accuracy :: Ptr () -> IO CInt +foreign import ccall unsafe "req_num_levels" c_req_num_levels :: Ptr () -> IO CInt +foreign import ccall unsafe "req_criterion" c_req_criterion :: Ptr () -> IO CInt +foreign import ccall unsafe "req_set_criterion" c_req_set_criterion :: Ptr () -> CInt -> IO () +foreign import ccall unsafe "req_count_with_criterion" c_req_cwc :: Ptr () -> CDouble -> IO Word64 +foreign import ccall unsafe "req_rank" c_req_rank :: Ptr () -> CDouble -> IO CDouble +foreign import ccall unsafe "req_quantile" c_req_quantile :: Ptr () -> CDouble -> IO CDouble + +withSketch :: CReqSketch s -> (Ptr () -> IO a) -> IO a +withSketch (CReqSketch fp) = withForeignPtr fp +{-# INLINE withSketch #-} + +mkCReqSketch :: PrimMonad m => Word32 -> Int -> m (CReqSketch (PrimState m)) +mkCReqSketch k ra = unsafePrimToPrim $ do + unless (even k && k >= 4 && k <= 1024) $ + error "k must be divisible by 2, and satisfy 4 <= k <= 1024" + ptr <- c_req_new k (fromIntegral ra) 12345 + fp <- newForeignPtr c_req_free ptr + pure (CReqSketch fp) + +creqInsert :: PrimMonad m => CReqSketch (PrimState m) -> Double -> m () +creqInsert sk val = unsafePrimToPrim $ withSketch sk $ \p -> c_req_insert p (realToFrac val) +{-# INLINE creqInsert #-} + +creqInsertBatch :: PrimMonad m => CReqSketch (PrimState m) -> VS.Vector Double -> m () +creqInsertBatch sk vals = unsafePrimToPrim $ withSketch sk $ \p -> + VS.unsafeWith (VS.unsafeCast vals) $ \arr -> + c_req_insert_batch p arr (fromIntegral $ VS.length vals) + +creqMerge :: PrimMonad m => CReqSketch (PrimState m) -> CReqSketch (PrimState m) -> m () +creqMerge dst src = unsafePrimToPrim $ + withSketch dst $ \pd -> withSketch src $ \ps -> c_req_merge pd ps + +creqCount :: PrimMonad m => CReqSketch (PrimState m) -> m Word64 +creqCount sk = unsafePrimToPrim $ withSketch sk c_req_count +{-# INLINE creqCount #-} + +creqIsEmpty :: PrimMonad m => CReqSketch (PrimState m) -> m Bool +creqIsEmpty sk = unsafePrimToPrim $ withSketch sk $ fmap (/= 0) . c_req_is_empty +{-# INLINE creqIsEmpty #-} + +creqMin :: PrimMonad m => CReqSketch (PrimState m) -> m Double +creqMin sk = unsafePrimToPrim $ withSketch sk $ fmap realToFrac . c_req_min +{-# INLINE creqMin #-} + +creqMax :: PrimMonad m => CReqSketch (PrimState m) -> m Double +creqMax sk = unsafePrimToPrim $ withSketch sk $ fmap realToFrac . c_req_max +{-# INLINE creqMax #-} + +creqSum :: PrimMonad m => CReqSketch (PrimState m) -> m Double +creqSum sk = unsafePrimToPrim $ withSketch sk $ fmap realToFrac . c_req_sum +{-# INLINE creqSum #-} + +creqRetained :: PrimMonad m => CReqSketch (PrimState m) -> m Int +creqRetained sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_req_retained +{-# INLINE creqRetained #-} + +creqK :: PrimMonad m => CReqSketch (PrimState m) -> m Word32 +creqK sk = unsafePrimToPrim $ withSketch sk c_req_k + +creqRankAccuracy :: PrimMonad m => CReqSketch (PrimState m) -> m Int +creqRankAccuracy sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_req_rank_accuracy + +creqNumLevels :: PrimMonad m => CReqSketch (PrimState m) -> m Int +creqNumLevels sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_req_num_levels + +creqCriterion :: PrimMonad m => CReqSketch (PrimState m) -> m Int +creqCriterion sk = unsafePrimToPrim $ withSketch sk $ fmap fromIntegral . c_req_criterion + +creqSetCriterion :: PrimMonad m => CReqSketch (PrimState m) -> Int -> m () +creqSetCriterion sk c = unsafePrimToPrim $ withSketch sk $ \p -> c_req_set_criterion p (fromIntegral c) + +creqCountWithCriterion :: PrimMonad m => CReqSketch (PrimState m) -> Double -> m Word64 +creqCountWithCriterion sk val = unsafePrimToPrim $ withSketch sk $ \p -> c_req_cwc p (realToFrac val) + +creqRank :: PrimMonad m => CReqSketch (PrimState m) -> Double -> m Double +creqRank sk val = unsafePrimToPrim $ withSketch sk $ \p -> realToFrac <$> c_req_rank p (realToFrac val) +{-# INLINE creqRank #-} + +creqQuantile :: PrimMonad m => CReqSketch (PrimState m) -> Double -> m Double +creqQuantile sk nr = unsafePrimToPrim $ withSketch sk $ \p -> realToFrac <$> c_req_quantile p (realToFrac nr) +{-# INLINE creqQuantile #-} diff --git a/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal.hs b/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal.hs --- a/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal.hs +++ b/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal.hs @@ -9,7 +9,7 @@ import DataSketches.Core.Snapshot import DataSketches.Quantiles.RelativeErrorQuantile.Types -import DataSketches.Core.Internal.URef (URef, readURef) +import DataSketches.Core.Internal.URef (MutableFields, readField, writeField, modifyField) import DataSketches.Quantiles.RelativeErrorQuantile.Internal.Auxiliary (ReqAuxiliary) import DataSketches.Quantiles.RelativeErrorQuantile.Internal.Compactor (ReqCompactor) import Control.DeepSeq (NFData, rnf) @@ -19,59 +19,111 @@ import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.Auxiliary as Auxiliary import Control.Exception (Exception) - -{- | -This Relative Error Quantiles Sketch is the Haskell implementation based on the paper -"Relative Error Streaming Quantiles", https://arxiv.org/abs/2004.01668, and loosely derived from -a Python prototype written by Pavel Vesely, ported from the Java equivalent. - -This implementation differs from the algorithm described in the paper in the following: - -The algorithm requires no upper bound on the stream length. -Instead, each relative-compactor counts the number of compaction operations performed -so far (via variable state). Initially, the relative-compactor starts with INIT_NUMBER_OF_SECTIONS. -Each time the number of compactions (variable state) exceeds 2^{numSections - 1}, we double -numSections. Note that after merging the sketch with another one variable state may not correspond -to the number of compactions performed at a particular level, however, since the state variable -never exceeds the number of compactions, the guarantees of the sketch remain valid. - -The size of each section (variable k and sectionSize in the code and parameter k in -the paper) is initialized with a value set by the user via variable k. -When the number of sections doubles, we decrease sectionSize by a factor of sqrt(2). -This is applied at each level separately. Thus, when we double the number of sections, the -nominal compactor size increases by a factor of approx. sqrt(2) (+/- rounding). - -The merge operation here does not perform "special compactions", which are used in the paper -to allow for a tight mathematical analysis of the sketch. - -This implementation provides a number of capabilities not discussed in the paper or provided -in the Python prototype. - -The Python prototype only implemented high accuracy for low ranks. This implementation -provides the user with the ability to choose either high rank accuracy or low rank accuracy at -the time of sketch construction. - -- The Python prototype only implemented a comparison criterion of "<". This implementation -allows the user to switch back and forth between the "<=" criterion and the "<=" criterion. --} +-- | Mutable sketch state packed into a single 'MutableByteArray'. +-- +-- All 6 scalar fields occupy 48 bytes (6 × 8-byte slots), fitting entirely +-- within one 64-byte x86-64 cache line. On the insert hot path, totalN, +-- retainedItems, maxNominalCapacitiesSize, minValue, maxValue, and sumValue +-- are all read/written — a single cache line fetch covers all of them. +-- +-- Layout (element index → field): +-- 0: totalN (Int, representing Word64) +-- 1: minValue (Double) +-- 2: maxValue (Double) +-- 3: sumValue (Double) +-- 4: retainedItems (Int) +-- 5: maxNominalCapacitiesSize (Int) data ReqSketch s = ReqSketch { k :: !Word32 , rankAccuracySetting :: !RankAccuracy , criterion :: !Criterion , sketchRng :: {-# UNPACK #-} !(Gen s) - , totalN :: {-# UNPACK #-} !(URef s Word64) - , minValue :: {-# UNPACK #-} !(URef s Double) - , maxValue :: {-# UNPACK #-} !(URef s Double) - , sumValue :: {-# UNPACK #-} !(URef s Double) - , retainedItems :: {-# UNPACK #-} !(URef s Int) - , maxNominalCapacitiesSize :: {-# UNPACK #-} !(URef s Int) + , sketchFields :: {-# UNPACK #-} !(MutableFields s) , aux :: {-# UNPACK #-} !(MutVar s (Maybe ReqAuxiliary)) , compactors :: {-# UNPACK #-} !(MutVar s (Vector.Vector (ReqCompactor s))) } deriving (Generic) instance NFData (ReqSketch s) where - rnf !rs = () + rnf !_ = () +-- Field indices +fTotalN, fMinValue, fMaxValue, fSumValue, fRetainedItems, fMaxNomCapSize :: Int +fTotalN = 0 +fMinValue = 1 +fMaxValue = 2 +fSumValue = 3 +fRetainedItems = 4 +fMaxNomCapSize = 5 + +sketchFieldBytes :: Int +sketchFieldBytes = 6 * 8 + +-- Typed accessors + +getTotalN :: PrimMonad m => ReqSketch (PrimState m) -> m Word64 +getTotalN sk = do + v <- readField (sketchFields sk) fTotalN + pure $! fromIntegral (v :: Int) +{-# INLINE getTotalN #-} + +setTotalN :: PrimMonad m => ReqSketch (PrimState m) -> Word64 -> m () +setTotalN sk v = writeField (sketchFields sk) fTotalN (fromIntegral v :: Int) +{-# INLINE setTotalN #-} + +modifyTotalN :: PrimMonad m => ReqSketch (PrimState m) -> (Word64 -> Word64) -> m () +modifyTotalN sk f = do + !v <- getTotalN sk + setTotalN sk $! f v +{-# INLINE modifyTotalN #-} + +getMinValue :: PrimMonad m => ReqSketch (PrimState m) -> m Double +getMinValue sk = readField (sketchFields sk) fMinValue +{-# INLINE getMinValue #-} + +setMinValue :: PrimMonad m => ReqSketch (PrimState m) -> Double -> m () +setMinValue sk = writeField (sketchFields sk) fMinValue +{-# INLINE setMinValue #-} + +getMaxValue :: PrimMonad m => ReqSketch (PrimState m) -> m Double +getMaxValue sk = readField (sketchFields sk) fMaxValue +{-# INLINE getMaxValue #-} + +setMaxValue :: PrimMonad m => ReqSketch (PrimState m) -> Double -> m () +setMaxValue sk = writeField (sketchFields sk) fMaxValue +{-# INLINE setMaxValue #-} + +getSumValue :: PrimMonad m => ReqSketch (PrimState m) -> m Double +getSumValue sk = readField (sketchFields sk) fSumValue +{-# INLINE getSumValue #-} + +modifySumValue :: PrimMonad m => ReqSketch (PrimState m) -> (Double -> Double) -> m () +modifySumValue sk f = modifyField (sketchFields sk) fSumValue f +{-# INLINE modifySumValue #-} + +getRetainedItems :: PrimMonad m => ReqSketch (PrimState m) -> m Int +getRetainedItems sk = readField (sketchFields sk) fRetainedItems +{-# INLINE getRetainedItems #-} + +setRetainedItems :: PrimMonad m => ReqSketch (PrimState m) -> Int -> m () +setRetainedItems sk = writeField (sketchFields sk) fRetainedItems +{-# INLINE setRetainedItems #-} + +modifyRetainedItems :: PrimMonad m => ReqSketch (PrimState m) -> (Int -> Int) -> m () +modifyRetainedItems sk = modifyField (sketchFields sk) fRetainedItems +{-# INLINE modifyRetainedItems #-} + +getMaxNomCapSize :: PrimMonad m => ReqSketch (PrimState m) -> m Int +getMaxNomCapSize sk = readField (sketchFields sk) fMaxNomCapSize +{-# INLINE getMaxNomCapSize #-} + +setMaxNomCapSize :: PrimMonad m => ReqSketch (PrimState m) -> Int -> m () +setMaxNomCapSize sk = writeField (sketchFields sk) fMaxNomCapSize +{-# INLINE setMaxNomCapSize #-} + +modifyMaxNomCapSize :: PrimMonad m => ReqSketch (PrimState m) -> (Int -> Int) -> m () +modifyMaxNomCapSize sk = modifyField (sketchFields sk) fMaxNomCapSize +{-# INLINE modifyMaxNomCapSize #-} + data ReqSketchSnapshot = ReqSketchSnapshot { snapshotRankAccuracySetting :: !RankAccuracy , snapshotCriterion :: !Criterion @@ -80,46 +132,48 @@ , snapshotMaxValue :: !Double , snapshotRetainedItems :: !Int , snapshotMaxNominalCapacitiesSize :: !Int - -- , aux :: !(MutVar s (Maybe ())) , snapshotCompactors :: !(Vector.Vector (Snapshot ReqCompactor)) } deriving Show instance TakeSnapshot ReqSketch where - type Snapshot ReqSketch = ReqSketchSnapshot - takeSnapshot ReqSketch{..} = ReqSketchSnapshot rankAccuracySetting criterion - <$> readURef totalN - <*> readURef minValue - <*> readURef maxValue - <*> readURef retainedItems - <*> readURef maxNominalCapacitiesSize - <*> (readMutVar compactors >>= mapM takeSnapshot) + type Snapshot ReqSketch = ReqSketchSnapshot + takeSnapshot sk = do + tn <- getTotalN sk + mn <- getMinValue sk + mx <- getMaxValue sk + ri <- getRetainedItems sk + mc <- getMaxNomCapSize sk + cs <- readMutVar (compactors sk) >>= mapM takeSnapshot + pure $ ReqSketchSnapshot (rankAccuracySetting sk) (criterion sk) tn mn mx ri mc cs getCompactors :: PrimMonad m => ReqSketch (PrimState m) -> m (Vector.Vector (ReqCompactor (PrimState m))) getCompactors = readMutVar . compactors +{-# INLINE getCompactors #-} computeTotalRetainedItems :: PrimMonad m => ReqSketch (PrimState m) -> m Int computeTotalRetainedItems this = do - compactors <- getCompactors this - Vector.foldM countBuffer 0 compactors + cs <- getCompactors this + Vector.foldM countBuffer 0 cs where countBuffer acc compactor = do buff <- Compactor.getBuffer compactor buffSize <- DoubleBuffer.getCount buff - pure $ buffSize + acc + pure $! buffSize + acc retainedItemCount :: PrimMonad m => ReqSketch (PrimState m) -> m Int -retainedItemCount = readURef . retainedItems +retainedItemCount = getRetainedItems +{-# INLINE retainedItemCount #-} --- | Get the total number of items inserted into the sketch count :: PrimMonad m => ReqSketch (PrimState m) -> m Word64 -count = readURef . totalN +count = getTotalN +{-# INLINE count #-} mkAuxiliaryFromReqSketch :: PrimMonad m => ReqSketch (PrimState m) -> m ReqAuxiliary mkAuxiliaryFromReqSketch this = do total <- count this - retainedItems <- retainedItemCount this - compactors <- getCompactors this - Auxiliary.mkAuxiliary (rankAccuracySetting this) total retainedItems compactors + ri <- retainedItemCount this + cs <- getCompactors this + Auxiliary.mkAuxiliary (rankAccuracySetting this) total ri cs data CumulativeDistributionInvariants = CumulativeDistributionInvariantsSplitsAreEmpty diff --git a/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/Compactor.hs b/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/Compactor.hs --- a/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/Compactor.hs +++ b/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/Compactor.hs @@ -12,20 +12,17 @@ , nearestEven ) where -import GHC.TypeLits import Data.Bits ((.&.), (.|.), complement, countTrailingZeros, shiftL, shiftR) import Data.Primitive.MutVar -import Data.Proxy -import Data.Semigroup (Semigroup) import Data.Word import DataSketches.Quantiles.RelativeErrorQuantile.Types -import System.Random.MWC (create, Variate(uniform), Gen) +import System.Random.MWC (Variate(uniform), Gen) import Control.Exception (assert) import Control.Monad (when) import Control.Monad.Primitive import DataSketches.Quantiles.RelativeErrorQuantile.Internal.Constants import DataSketches.Quantiles.RelativeErrorQuantile.Internal.DoubleBuffer -import DataSketches.Core.Internal.URef +import DataSketches.Core.Internal.URef (MutableFields, newMutableFields, readField, writeField, modifyField) import DataSketches.Core.Snapshot data CompactorReturn s = CompactorReturn @@ -34,20 +31,90 @@ , crDoubleBuffer :: {-# UNPACK #-} !(DoubleBuffer s) } +-- | Mutable compactor state packed into a single 'MutableByteArray'. +-- +-- All scalar fields are stored at 8-byte-aligned offsets in one contiguous +-- allocation. On x86-64 (64-byte cache lines), the entire 40-byte block +-- fits in a single cache line, so reading state + numSections + sectionSize +-- during compact (the hot path) never causes a second cache miss. +-- +-- Layout (byte offset → field): +-- 0: state (Word64) — compaction counter, read/written every compact +-- 8: sectionSizeFlt (Double) — current section size as float +-- 16: sectionSize (Int) — current section size +-- 24: numSections (Int) — current number of sections +-- 32: lastFlip (Int) — last coin flip result (0 or 1) data ReqCompactor s = ReqCompactor - -- Configuration constants { rcRankAccuracy :: !RankAccuracy , rcLgWeight :: {-# UNPACK #-} !Word8 , rcRng :: {-# UNPACK #-} !(Gen s) - -- State - , rcState :: {-# UNPACK #-} !(URef s Word64) - , rcLastFlip :: {-# UNPACK #-} !(URef s Bool) - , rcSectionSizeFlt :: {-# UNPACK #-} !(URef s Double) - , rcSectionSize :: {-# UNPACK #-} !(URef s Word32) - , rcNumSections :: {-# UNPACK #-} !(URef s Word8) + , rcFields :: {-# UNPACK #-} !(MutableFields s) , rcBuffer :: {-# UNPACK #-} !(MutVar s (DoubleBuffer s)) } +-- Field indices (element index into the Int/Double-sized slots of MutableByteArray) +-- All stored as 8-byte values for uniform alignment. +fState, fSectionSizeFlt, fSectionSize, fNumSections, fLastFlip :: Int +fState = 0 -- Word64, but stored via Int-sized readField/writeField +fSectionSizeFlt = 1 -- Double +fSectionSize = 2 -- Int +fNumSections = 3 -- Int +fLastFlip = 4 -- Int (0 or 1) + +fieldBytes :: Int +fieldBytes = 5 * 8 -- 40 bytes total, fits in one 64-byte cache line + +-- Typed accessors. These read/write the MutableByteArray at the right element +-- index, relying on Prim instances for Int, Double, Word64. + +getStateField :: PrimMonad m => ReqCompactor (PrimState m) -> m Word64 +getStateField rc = do + v <- readField (rcFields rc) fState + pure $! fromIntegral (v :: Int) +{-# INLINE getStateField #-} + +setStateField :: PrimMonad m => ReqCompactor (PrimState m) -> Word64 -> m () +setStateField rc v = writeField (rcFields rc) fState (fromIntegral v :: Int) +{-# INLINE setStateField #-} + +modifyStateField :: PrimMonad m => ReqCompactor (PrimState m) -> (Word64 -> Word64) -> m () +modifyStateField rc f = do + !v <- getStateField rc + setStateField rc $! f v +{-# INLINE modifyStateField #-} + +getSectionSizeFltField :: PrimMonad m => ReqCompactor (PrimState m) -> m Double +getSectionSizeFltField rc = readField (rcFields rc) fSectionSizeFlt +{-# INLINE getSectionSizeFltField #-} + +setSectionSizeFltField :: PrimMonad m => ReqCompactor (PrimState m) -> Double -> m () +setSectionSizeFltField rc = writeField (rcFields rc) fSectionSizeFlt +{-# INLINE setSectionSizeFltField #-} + +getSectionSizeField :: PrimMonad m => ReqCompactor (PrimState m) -> m Int +getSectionSizeField rc = readField (rcFields rc) fSectionSize +{-# INLINE getSectionSizeField #-} + +setSectionSizeField :: PrimMonad m => ReqCompactor (PrimState m) -> Int -> m () +setSectionSizeField rc = writeField (rcFields rc) fSectionSize +{-# INLINE setSectionSizeField #-} + +getNumSectionsField :: PrimMonad m => ReqCompactor (PrimState m) -> m Int +getNumSectionsField rc = readField (rcFields rc) fNumSections +{-# INLINE getNumSectionsField #-} + +setNumSectionsField :: PrimMonad m => ReqCompactor (PrimState m) -> Int -> m () +setNumSectionsField rc = writeField (rcFields rc) fNumSections +{-# INLINE setNumSectionsField #-} + +getLastFlipField :: PrimMonad m => ReqCompactor (PrimState m) -> m Bool +getLastFlipField rc = (/= (0 :: Int)) <$> readField (rcFields rc) fLastFlip +{-# INLINE getLastFlipField #-} + +setLastFlipField :: PrimMonad m => ReqCompactor (PrimState m) -> Bool -> m () +setLastFlipField rc b = writeField (rcFields rc) fLastFlip (if b then 1 :: Int else 0) +{-# INLINE setLastFlipField #-} + data ReqCompactorSnapshot = ReqCompactorSnapshot { snapshotCompactorRankAccuracy :: !RankAccuracy , snapshotCompactorRankAccuracyState :: !Word64 @@ -60,14 +127,15 @@ instance TakeSnapshot ReqCompactor where type Snapshot ReqCompactor = ReqCompactorSnapshot - takeSnapshot ReqCompactor{..} = ReqCompactorSnapshot rcRankAccuracy - <$> readURef rcState - <*> readURef rcLastFlip - <*> readURef rcSectionSizeFlt - <*> readURef rcSectionSize - <*> readURef rcNumSections - <*> (readMutVar rcBuffer >>= takeSnapshot) - + takeSnapshot rc = do + st <- getStateField rc + fl <- getLastFlipField rc + szf <- getSectionSizeFltField rc + sz <- getSectionSizeField rc + ns <- getNumSectionsField rc + buf <- readMutVar (rcBuffer rc) >>= takeSnapshot + pure $ ReqCompactorSnapshot + (rcRankAccuracy rc) st fl szf (fromIntegral sz) (fromIntegral ns) buf mkReqCompactor :: PrimMonad m @@ -79,40 +147,37 @@ mkReqCompactor g lgWeight rankAccuracy sectionSize = do let nominalCapacity = fromIntegral $ nomCapMulti * initNumberOfSections * sectionSize buff <- mkBuffer (nominalCapacity * 2) nominalCapacity (rankAccuracy == HighRanksAreAccurate) - ReqCompactor rankAccuracy lgWeight g - <$> newURef 0 - <*> newURef False - <*> newURef (fromIntegral sectionSize) - <*> newURef sectionSize - <*> newURef initNumberOfSections - <*> newMutVar buff + fields <- newMutableFields fieldBytes + writeField fields fState (0 :: Int) + writeField fields fSectionSizeFlt (fromIntegral sectionSize :: Double) + writeField fields fSectionSize (fromIntegral sectionSize :: Int) + writeField fields fNumSections (fromIntegral initNumberOfSections :: Int) + writeField fields fLastFlip (0 :: Int) + bufMv <- newMutVar buff + pure $ ReqCompactor rankAccuracy lgWeight g fields bufMv -nomCapMult :: Num a => a +nomCapMult :: Int nomCapMult = 2 -toInt :: Integral a => a -> Int -toInt = fromIntegral - -compact :: (PrimMonad m) => ReqCompactor (PrimState m) -> m (CompactorReturn (PrimState m)) +compact :: PrimMonad m => ReqCompactor (PrimState m) -> m (CompactorReturn (PrimState m)) compact this = do startBuffSize <- getCount =<< getBuffer this startNominalCapacity <- getNominalCapacity this - numSections <- readURef $ rcNumSections this - sectionSize <- readURef $ rcSectionSize this - state <- readURef $ rcState this + numSections <- getNumSectionsField this + sectionSize <- getSectionSizeField this + state <- getStateField this let trailingOnes = succ $ countTrailingZeros $ complement state - sectionsToCompact = min trailingOnes $ fromIntegral numSections + sectionsToCompact = min trailingOnes numSections (compactionStart, compactionEnd) <- computeCompactionRange this sectionsToCompact - -- TODO, this fails in GHCi but not in tests? assert (compactionEnd - compactionStart >= 2) $ do coin <- if state .&. 1 == 1 - then fmap not $ readURef $ rcLastFlip this + then fmap not $ getLastFlipField this else flipCoin this - writeURef (rcLastFlip this) coin + setLastFlipField this coin buff <- getBuffer this promote <- getEvensOrOdds buff compactionStart compactionEnd coin trimCount buff $ startBuffSize - (compactionEnd - compactionStart) - modifyURef (rcState this) (+ 1) + modifyStateField this (+ 1) ensureEnoughSections this endBuffSize <- getCount buff promoteBuffSize <- getCount promote @@ -128,40 +193,33 @@ getBuffer :: PrimMonad m => ReqCompactor (PrimState m) -> m (DoubleBuffer (PrimState m)) getBuffer = readMutVar . rcBuffer +{-# INLINE getBuffer #-} -flipCoin :: (PrimMonad m) => ReqCompactor (PrimState m) -> m Bool +flipCoin :: PrimMonad m => ReqCompactor (PrimState m) -> m Bool flipCoin = uniform . rcRng +{-# INLINE flipCoin #-} getCoin :: PrimMonad m => ReqCompactor (PrimState m) -> m Bool -getCoin = readURef . rcLastFlip +getCoin = getLastFlipField getNominalCapacity :: PrimMonad m => ReqCompactor (PrimState m) -> m Int getNominalCapacity compactor = do - numSections <- readURef $ rcNumSections compactor - sectionSize <- readURef $ rcSectionSize compactor - pure $ nomCapMult * toInt numSections * toInt sectionSize + numSections <- getNumSectionsField compactor + sectionSize <- getSectionSizeField compactor + pure $! nomCapMult * numSections * sectionSize +{-# INLINE getNominalCapacity #-} getNumSections :: PrimMonad m => ReqCompactor (PrimState m) -> m Word8 -getNumSections = readURef . rcNumSections - -getSectionSizeFlt :: PrimMonad m => ReqCompactor (PrimState m) -> m Double -getSectionSizeFlt = readURef . rcSectionSizeFlt - -getState :: PrimMonad m => ReqCompactor (PrimState m) -> m Word64 -getState = readURef . rcState +getNumSections rc = fromIntegral <$> getNumSectionsField rc --- | Merge the other given compactor into this one. They both must have the --- same @lgWeight@ merge :: (PrimMonad m, s ~ PrimState m) => ReqCompactor (PrimState m) - -- ^ The compactor to merge into -> ReqCompactor (PrimState m) - -- ^ The compactor to merge from -> m (ReqCompactor s) merge this otherCompactor = assert (rcLgWeight this == rcLgWeight otherCompactor) $ do - otherState <- readURef $ rcState otherCompactor - modifyURef (rcState this) (.|. otherState) + otherState <- getStateField otherCompactor + modifyStateField this (.|. otherState) ensureMaxSections buff <- getBuffer this @@ -171,7 +229,7 @@ sort otherBuff otherBuffIsBigger <- (>) <$> getCount otherBuff <*> getCount buff - finalBuff <- if otherBuffIsBigger + if otherBuffIsBigger then do otherBuff' <- copyBuffer otherBuff mergeSortIn otherBuff' buff @@ -181,58 +239,47 @@ where ensureMaxSections = do adjusted <- ensureEnoughSections this - -- loop until no adjustments can be made when adjusted ensureMaxSections --- | Adjust the sectionSize and numSections if possible. ensureEnoughSections :: PrimMonad m => ReqCompactor (PrimState m) -> m Bool - -- ^ 'True' if the SectionSize and NumSections were adjusted. ensureEnoughSections compactor = do - sectionSizeFlt <- readURef $ rcSectionSizeFlt compactor + sectionSizeFlt <- getSectionSizeFltField compactor let szf = sectionSizeFlt / sqrt2 ne = nearestEven szf - state <- readURef $ rcState compactor - numSections <- readURef $ rcNumSections compactor - sectionSize <- readURef $ rcSectionSize compactor - if state >= (1 `shiftL` toInt (numSections - 1)) + state <- getStateField compactor + numSections <- getNumSectionsField compactor + sectionSize <- getSectionSizeField compactor + if state >= (1 `shiftL` (numSections - 1)) && sectionSize > minK && ne >= minK then do - writeURef (rcSectionSizeFlt compactor) szf - writeURef (rcSectionSize compactor) $ fromIntegral ne - modifyURef (rcNumSections compactor) (`shiftL` 1) + setSectionSizeFltField compactor szf + setSectionSizeField compactor ne + setNumSectionsField compactor (numSections `shiftL` 1) buf <- getBuffer compactor nomCapacity <- getNominalCapacity compactor ensureCapacity buf (2 * nomCapacity) pure True else pure False --- | Computes the start and end indices of the compacted region computeCompactionRange :: PrimMonad m => ReqCompactor (PrimState m) -> Int - -- ^ secsToCompact the number of contiguous sections to compact -> m (Int, Int) --- ^ the start and end indices of the compacted region in compact form computeCompactionRange this secsToCompact = do buffSize <- getCount =<< getBuffer this nominalCapacity <- getNominalCapacity this - numSections <- readURef $ rcNumSections this - sectionSize <- readURef $ rcSectionSize this - let nonCompact = (nominalCapacity `div` 2) + (fromIntegral numSections - secsToCompact) * fromIntegral sectionSize + numSections <- getNumSectionsField this + sectionSize <- getSectionSizeField this + let nonCompact = (nominalCapacity `div` 2) + (numSections - secsToCompact) * sectionSize nonCompact' = if (buffSize - nonCompact) .&. 1 == 1 then nonCompact - 1 else nonCompact pure $ case rcRankAccuracy this of - HighRanksAreAccurate -> (0, fromIntegral $ buffSize - nonCompact') + HighRanksAreAccurate -> (0, buffSize - nonCompact') LowRanksAreAccurate -> (nonCompact', buffSize) --- | Returns the nearest even integer to the given value. Also used by test. -nearestEven - :: Double - -- ^ the given value - -> Int - -- ^ the nearest even integer to the given value. +nearestEven :: Double -> Int nearestEven x = round (x / 2) `shiftL` 1 diff --git a/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/DoubleBuffer.hs b/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/DoubleBuffer.hs --- a/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/DoubleBuffer.hs +++ b/src/DataSketches/Quantiles/RelativeErrorQuantile/Internal/DoubleBuffer.hs @@ -35,7 +35,8 @@ import qualified Data.Vector.Unboxed as UVector import qualified Data.Vector.Unboxed.Mutable as MUVector import DataSketches.Core.Internal.URef - ( URef, newURef, readURef, writeURef, modifyURef ) + ( URef, newURef, readURef, writeURef, modifyURef + , MutableFields, newMutableFields, readField, writeField, modifyField ) import Data.Vector.Algorithms.Intro (sortByBounds) import GHC.Stack ( HasCallStack ) import System.IO.Unsafe () @@ -44,14 +45,21 @@ import DataSketches.Core.Snapshot ( TakeSnapshot(..) ) -- | A special buffer of floats specifically designed to support the ReqCompactor class. +-- +-- Mutable scalars (count, sorted flag) are packed into a single +-- 'MutableFields' to avoid per-field MutableByteArray# overhead. +-- Field layout: index 0 = count (Int), index 1 = sorted (Int, 0 or 1). data DoubleBuffer s = DoubleBuffer { vec :: {-# UNPACK #-} !(MutVar s (MUVector.MVector s Double)) - , count :: {-# UNPACK #-} !(URef s Int) - , sorted :: {-# UNPACK #-} !(URef s Bool) + , dbFields :: {-# UNPACK #-} !(MutableFields s) , growthIncrement :: {-# UNPACK #-} !Int , spaceAtBottom :: !Bool } +dbCountIx, dbSortedIx :: Int +dbCountIx = 0 +dbSortedIx = 1 + data DoubleBufferSnapshot = DoubleBufferSnapshot { dbSnapshotVec :: UVector.Vector Double , dbSnapshotCount :: !Int @@ -63,12 +71,11 @@ instance TakeSnapshot DoubleBuffer where type Snapshot DoubleBuffer = DoubleBufferSnapshot - takeSnapshot DoubleBuffer{..} = DoubleBufferSnapshot - <$> (readMutVar vec >>= UVector.freeze) - <*> readURef count - <*> readURef sorted - <*> pure growthIncrement - <*> pure spaceAtBottom + takeSnapshot DoubleBuffer{..} = do + v <- readMutVar vec >>= UVector.freeze + cnt <- readField dbFields dbCountIx + srt <- readField dbFields dbSortedIx + pure $ DoubleBufferSnapshot v cnt (srt /= (0 :: Int)) growthIncrement spaceAtBottom type Capacity = Int type GrowthIncrement = Int @@ -79,15 +86,20 @@ mkBuffer :: PrimMonad m => Capacity -> GrowthIncrement -> SpaceAtBottom -> m (DoubleBuffer (PrimState m)) mkBuffer capacity_ growthIncrement spaceAtBottom = do vec <- newMutVar =<< MUVector.new capacity_ - count <- newURef 0 - sorted <- newURef True + -- Pack count and sorted into 2 Int-sized slots + dbFields <- newMutableFields (2 * 8) + writeField dbFields dbCountIx (0 :: Int) + writeField dbFields dbSortedIx (1 :: Int) pure $ DoubleBuffer{..} copyBuffer :: PrimMonad m => DoubleBuffer (PrimState m) -> m (DoubleBuffer (PrimState m)) copyBuffer buf@DoubleBuffer{..} = do vec <- newMutVar =<< MUVector.clone =<< getVector buf - count <- newURef =<< getCount buf - sorted <- newURef =<< readURef sorted + dbFields <- newMutableFields (2 * 8) + cnt <- getCount buf + srt <- isSorted buf + writeField dbFields dbCountIx cnt + writeField dbFields dbSortedIx (if srt then 1 :: Int else 0) pure $ DoubleBuffer {..} -- | Appends the given item to the active array and increments the active count. @@ -95,34 +107,34 @@ append :: PrimMonad m => DoubleBuffer (PrimState m) -> Double -> m () append buf@DoubleBuffer{..} x = do ensureSpace buf 1 + count_ <- getCount buf index <- if spaceAtBottom - then - (\capacity_ count_ -> capacity_ - count_ - 1) - <$> getCapacity buf - <*> getCount buf - else readURef count - modifyURef count (+ 1) - getVector buf >>= \vec -> MUVector.unsafeWrite vec index x - writeURef sorted False -{-# SCC append #-} + then do + capacity_ <- getCapacity buf + pure $! capacity_ - count_ - 1 + else pure count_ + writeField dbFields dbCountIx (count_ + 1) + v <- getVector buf + MUVector.unsafeWrite v index x + writeField dbFields dbSortedIx (0 :: Int) +{-# INLINE append #-} -- | Ensures that the capacity of this FloatBuffer is at least newCapacity. -- If newCapacity < capacity(), no action is taken. ensureSpace :: PrimMonad m => DoubleBuffer (PrimState m) -> Int -> m () ensureSpace buf@DoubleBuffer{..} space = do - count_ <- readURef count + count_ <- getCount buf capacity_ <- getCapacity buf - let notEnoughSpace = count_ + space > capacity_ - when notEnoughSpace $ do + when (count_ + space > capacity_) $ do let newCap = count_ + space + growthIncrement ensureCapacity buf newCap -getVector :: (PrimMonad m, PrimState m ~ s) => DoubleBuffer s -> m (MUVector.MVector s Double) +getVector :: PrimMonad m => DoubleBuffer (PrimState m) -> m (MUVector.MVector (PrimState m) Double) getVector = readMutVar . vec {-# INLINE getVector #-} getCapacity :: PrimMonad m => DoubleBuffer (PrimState m) -> m Int -getCapacity = fmap MUVector.length . getVector +getCapacity buf = MUVector.length <$> getVector buf {-# INLINE getCapacity #-} ensureCapacity :: PrimMonad m => DoubleBuffer (PrimState m) -> Int -> m () @@ -157,7 +169,7 @@ then do capacity_ <- getCapacity buf pure (capacity_ - count_, capacity_ - 1) - else pure (0, count_) + else pure (0, count_ - 1) ix <- IS.find criterion vec low high value pure $! if ix == MUVector.length vec @@ -185,13 +197,13 @@ MUVector.unsafeWrite out j =<< MUVector.unsafeRead vec (i + odd) go vec out (i + 2) (j + 1) else do - count <- newURef (MUVector.length out) - sorted <- newURef True + dbFields <- newMutableFields (2 * 8) + writeField dbFields dbCountIx (MUVector.length out) + writeField dbFields dbSortedIx (1 :: Int) vec <- newMutVar out pure DoubleBuffer { vec = vec - , count = count - , sorted = sorted + , dbFields = dbFields , growthIncrement = 0 , spaceAtBottom = spaceAtBottom } @@ -210,16 +222,23 @@ MUVector.read vec index getCount :: PrimMonad m => DoubleBuffer (PrimState m) -> m Int -getCount = readURef . count +getCount DoubleBuffer{..} = readField dbFields dbCountIx +{-# INLINE getCount #-} getSpace :: PrimMonad m => DoubleBuffer (PrimState m) -> m Int -getSpace buf@DoubleBuffer{..} = (-) <$> getCapacity buf <*> getCount buf +getSpace buf@DoubleBuffer{..} = do + cap <- getCapacity buf + cnt <- getCount buf + pure $! cap - cnt +{-# INLINE getSpace #-} isEmpty :: PrimMonad m => DoubleBuffer (PrimState m) -> m Bool isEmpty buf = (== 0) <$> getCount buf +{-# INLINE isEmpty #-} isSorted :: PrimMonad m => DoubleBuffer (PrimState m) -> m Bool -isSorted = readURef . sorted +isSorted DoubleBuffer{..} = (/= (0 :: Int)) <$> readField dbFields dbSortedIx +{-# INLINE isSorted #-} -- | Sorts the active region sort :: PrimMonad m => DoubleBuffer (PrimState m) -> m () @@ -231,10 +250,10 @@ let (start, end) = if spaceAtBottom then (capacity_ - count_, capacity_) else (0, count_) - vec <- getVector buf - sortByBounds compare vec start end - writeURef sorted True -{-# SCC sort #-} + v <- getVector buf + sortByBounds compare v start end + writeField dbFields dbSortedIx (1 :: Int) +{-# INLINE sort #-} -- | Merges the incoming sorted buffer into this sorted buffer. mergeSortIn :: (PrimMonad m, HasCallStack) => DoubleBuffer (PrimState m) -> DoubleBuffer (PrimState m) -> m () @@ -267,8 +286,8 @@ let k = totalLength mergeDownwards thisBuf thatBuf i j (k - 1) - modifyURef (count this) (+ bufInLen) - writeURef (sorted this) True + modifyField (dbFields this) dbCountIx (+ bufInLen) + writeField (dbFields this) dbSortedIx (1 :: Int) pure () where mergeUpwards thisBuf thatBuf capacity_ bufInCapacity_ = go @@ -318,4 +337,4 @@ {-# SCC mergeSortIn #-} trimCount :: PrimMonad m => DoubleBuffer (PrimState m) -> Int -> m () -trimCount DoubleBuffer{..} newCount = modifyURef count (\oldCount -> if newCount < oldCount then newCount else oldCount) +trimCount DoubleBuffer{..} newCount = modifyField dbFields dbCountIx (\oldCount -> if newCount < oldCount then newCount else oldCount)