36 Commits
Author SHA1 Message Date
LeoVasanko 665966ede2 Update README.md 2026-01-03 04:46:50 +00:00
LeoVasanko 9a3f3516c2 Fix old references to chacha20.c/h that were renamed to charandom.*. 2025-02-14 05:51:28 +00:00
LeoVasanko 5fb8af478d Position seek/tell. 2023-12-01 02:40:29 +00:00
LeoVasanko e9fa61177b Remove buffer check that fails with Numpy arrays 2023-11-30 22:55:38 +00:00
LeoVasanko 160eee95b1 Update README 2023-11-29 23:43:01 +00:00
LeoVasanko fcc6ad6d1d Update README 2023-11-29 23:36:21 +00:00
LeoVasanko 372ea4a18d Update README 2023-11-29 23:31:12 +00:00
LeoVasanko 46af28f19c Fix Neon implementation. 2023-11-29 22:55:48 +00:00
LeoVasanko e063648ecf Fixes, tests passing. 2023-11-29 21:56:55 +00:00
LeoVasanko 0b01458a98 Add SI and binary units k/M/G/T for sizes specified via -b on CLI. 2023-11-29 19:32:16 +00:00
LeoVasanko ed5ba3379f Native Apple Silicon SIMD code, cleanup. 2023-11-29 19:07:44 +00:00
LeoVasanko e081d19b2f Missing files 2023-11-29 17:41:58 +00:00
LeoVasanko 00f504f2c5 Updated Python module, Numpy module and randquik CLI for new core. 2023-11-29 17:32:06 +00:00
LeoVasanko 236f7eedf6 Rewrite, cleanup, configurable ChaCha rounds. 2023-11-28 23:39:38 +00:00
LeoVasanko 998560f5a0 Fix compilation, still needs more work. 2023-11-22 05:43:27 +00:00
LeoVasanko d102fa1d28 Implement Numpy BitGenerator. 2023-11-22 04:03:48 +00:00
LeoVasanko c573505826 More Pythonic Python binding 2023-11-21 23:52:40 +00:00
LeoVasanko 749f7c227c Add missing file for Apple compatibility. 2023-11-21 23:46:05 +00:00
LeoVasanko d2a0433344 Benchmark for cha random too. 2023-11-21 23:45:49 +00:00
LeoVasanko 570b071d36 Make it compile and run (fairly quick) on Apple M1/M2. 2023-11-21 23:45:10 +00:00
User e465d099c7 Make it compile and run (fairly quick) on Apple M1/M2. 2023-11-21 23:39:05 +00:00
LeoVasanko 4f42fb3447 Added Python convenience for different nonce sizes. Statistical tests of output randomness. 2023-10-29 12:17:05 +00:00
LeoVasanko baf3ef31a5 Test our implementation against Cryptography module's 2023-10-29 10:59:32 +00:00
LeoVasanko 8e38985bf8 Hunting a bug that causes test failure 2023-10-28 23:02:13 +00:00
LeoVasanko a0b55bc406 Docs 2023-10-28 18:44:59 +00:00
LeoVasanko 8acebc571e More cleanup 2023-10-28 18:42:46 +00:00
LeoVasanko b4fd8c29ff Python module with Numpy BitGenerator to integrate with np.random, and other Python use. 2023-10-28 18:42:34 +00:00
LeoVasanko a4b75a3bfe Renamed files. 2023-10-28 16:26:14 +00:00
LeoVasanko fe80552d6b Refactor 2023-10-28 16:24:18 +00:00
LeoVasanko ab67daf050 Fixed cha vectorized algorithms 2023-10-28 16:22:51 +00:00
LeoVasanko a6daac4ce6 Arrayized the variables, no change in performance but a lot less code. 2023-10-28 15:23:30 +00:00
LeoVasanko 771212029f Code formatting, refactoring and fixes. 2023-10-28 15:15:41 +00:00
LeoVasanko 44551920fa README 2023-10-28 11:58:06 +00:00
LeoVasanko 720f60ed71 README 2023-10-28 11:22:12 +00:00
LeoVasanko 6972c0a8b2 README 2023-10-28 11:21:22 +00:00
LeoVasanko 3deea0c962 Initial commit 2023-10-28 10:51:13 +00:00
31 changed files with 1310 additions and 2276 deletions
-7
View File
@@ -1,7 +0,0 @@
.*
*.lock
*.egg-info
/dist
/build
__pycache__
!.gitignore
+65 -91
View File
@@ -1,117 +1,91 @@
# The World's Fastest Random Generator v2
# The World's Fastest Random Generator
![RandQuik v2 Benchmark Stats](https://github.com/LeoVasanko/RandQuik/raw/main/docs/v2stat.webp)
*RandQuik running at full speed on a Macbook laptop, writing /dev/null.*
I was disappointed with the sad state of random number generators. Many languages don't ship anything useful and some are stuck with whatever the OS provides. All existing implementations are slow, typically maxing out at a few hundred megabytes per second.
I was disappointed with the sad state of random number generators. Many languages don't ship anything useful and some are stuck with whatever the OS provides. Most existing implementations are slow, often maxing out at a few hundred megabytes per second, which becomes a real bottleneck in high-throughput systems.
To give some perspective, my tool reaches 37.8 GB/s, writing /dev/null. Several gigabytes per second also on actual SSDs. And this is using full 20 rounds of shuffling.
Secondly, [flaws have been found](https://numpy.org/doc/stable/reference/random/upgrading-pcg64.html) in popular non-cryptographic algorithms such as Mersenne Twister and PCG-style generators. These issues range from detectable structure to repeating sequences, making them unsuitable for serious or long-running workloads.
Secondly, [flaws have been found](https://numpy.org/doc/stable/reference/random/upgrading-pcg64.html) in non-cryptographic algorithms such as the popular Mersenne Twister and PCG64 algorithms.
The cryptographic alternative is simply better, being free of such issues as sequences repeating, and additionally being, well, cryptographically secure. Surprisingly, they appear even to be faster now, so there really ought to be no reason to stick with the old.
The cryptographic alternative is simply better. Proper CSPRNGs avoid these pitfalls entirely and, on modern hardware, can now be *faster* than legacy non-cryptographic designs. There is little reason left to accept weaker guarantees for worse performance.
We use the widely used encryption algorithm **ChaCha20** as an extremely fast Cryptographically Secure Pseudorandom Number Generator (CSPRNG).
This version uses **AEGIS**, a modern authenticated encryption primitive that leverages AES hardware acceleration. AEGIS provides extremely high throughput while retaining strong cryptographic properties, significantly outperforming our legacy implementation and serving as an ideal foundation for a high-performance CSPRNG.
## CLI: randquik
## Quick start
A simple shell tool that simply produces randomness to a file or pipe. It uses ChaCha20 encryption algorithm to produce a random stream that cannot be predicted unless one knows the key - the seed - being used. Keeping always the same seed may be useful for researchers and such who need repeatable results.
Install [UV](https://docs.astral.sh/uv/getting-started/installation/) and install the CLI tool with it:
<img src="https://github.com/LeoVasanko/RandQuik/blob/legacy/docs/random.webp?raw=true" width="800" alt="Screenshot">
```sh
uv tool install randquik
## Installation
Clone the repository, use Meson build:
```
meson setup build
cd build
ninja
sudo ninja install # optionally
randquik > /dev/null
```
You can try how it performs on your machine and find the optimal parameters:
```sh
randquik --benchmark
Alternatively you may compile it by hand with `gcc *.c -o randquik -O3 -pthread`.
## Python module
Python module fills any buffers with random data very quickly,
```python
from randquik import Cha, generate
import secrets
key = secrets.token_bytes(32)
# Allocate bytearray and fill with random
data = generate(1_000_000, key)
# Or into an existing buffer
generate_into(data, key)
```
Wipe an entire file without altering its size:
```sh
randquik -o sensitive.dat
Given the same key, the generate functions will on each call produce the same sequence. For incremental updates, create a generator object and extract as many non-identical bytes from it as needed. Re-initializing with the same key of course once again repeats the requence.
```python
rng = Cha(key)
# Fill some buffer with next bytes iteratively
rng(data)
rng(data)
...
```
Piping and redirection:
```sh
randquik --quiet | hexdump -C | head
## Numpy module
Numpy.Random BitGenerator is also provided for use with Numpy distributions. We do not recommend using it for filling byte buffers, where the Python module and the C code are far faster, but it will still provide better quality random numbers and faster than Numpy's own PCG64.
```python
import numpy as np
from nprand import Cha
gen = np.random.Generator(Cha()) # System random seeding by default
gen.normal(size=10)
```
## Features
This module needs to be built by hand with Cython for now. Will be packaged properly in randquik module eventually.
- Blazing-fast CSPRNG built on AEGIS single or multithreaded
- Deterministic seeding: same seed, same byte stream
- Seekable random stream and file output
- Flexible I/O: piping, files, and `mmap`
- Built-in benchmarking and dry-run modes
- Full screen console graphics for speed display
## Tests
Below are the most important features with example commands.
Some pytest tests are provided for verifying ChaCha20 implementation correctness against the Cryptography module. Run by `pytest` in the main folder after installation.
### Performance
## Performance
You'll be looking at up to 100 GB/s raw generation speed, making this some orders of magnitude faster than your traditional random number generation. Single-threaded performance still is 10-20 times faster than other options that don't have threading.
ChaCha20 as its name implies uses 20 "rounds" of shuffling for each output block. A lower number of rounds can be used for extra performance, where 8 is the minimum that is considered secure, and 12 provides a balanced option, while 20 has comfortable headroom to stay cryptographically secure.
Your output, e.g. writing a file, will always be the bottle neck, not your random generator, but modern SSDs allow up to 10 GB/s write speeds already.
All functions and constructors of this module take `rounds` kwarg for adjusting this. On CLI the equivalent option is `-r`. By default 20 rounds are used.
### Size units
The CLI uses a configurable number of threads for extremely high performance, while the Python and Numpy modules don't - for now at least.
Many options such as `--len`, `--seek` accept human-readable units. The table uses conventional, capitalized forms (e.g. `KiB`, `MB`), but you may write them in lower case and with or without the trailing `B` (for example `1m` or `5gi`).
The implementation is optimized for Apple Silicon SIMD (Neon) and x86 CPUs using AVX2 where available, falling back to SSSE3 and ultimately plain C on other platforms. The implementation is loosely based on code from libsodium but runs faster than the library can.
| SI unit | Binary unit | Meaning | Bytes factor |
|--------:|------------:|------------------------|-----------------------|
| 100 | — | 100 bytes | 1 |
| 1kB | 1KiB | Kilobyte / kibibyte | 1_000 / 1_024 |
| 1MB | 1MiB | Megabyte / mebibyte | 1_000_000 / 1_048_576 |
| 1GB | 1GiB | Gigabyte / gibibyte | 1_000_000_000 / 1_073_741_824 |
| 1TB | 1TiB | Terabyte / tebibyte | 1_000_000_000_000 / 1_099_511_627_776 |
| 1PB | 1PiB | Petabyte / pebibyte | 1_000_... / 1_125_... |
| — | 1sect | Sectors of output device | 512 (typical), 4096 (rarely) |
## Seekability
### Seeding for repeatable output
You can provide an explicit seed string, always providing the same output, which can be useful e.g. for memory/disk testing where the data needs to be read back and verified.
```sh
randquik -l 64MiB -s my-seed-string -o chunk.bin
```
If no seed is provided, a secure new random one is created and printed on console (unless hidden by `-q`).
### Seekable random stream and output file
It is possible to seek to any byte position in the stream without delay.
- `--iseek`: seek the input random stream
- `--oseek`: seek in the output file
- `--seek`: set both input and output to the same position
Example: resume as if 5 terabytes had already been written, and continue writing to `out.dat`. The seed from the prior invocation should be included:
```sh
randquik --seek 5T --len 1G -s a5Z8Ew1Hfc2VfEtY -o out.dat
```
Bytes prior to seek position are kept as they were while the file is expanded to fit all the data starting at five terabytes mark (using sparse allocation so it doesn't actually consume 5 terabytes).
Wipe a specific range of a disk or USB drive (using sector numbers e.g. from gdisk):
```sh
randquik -oseek 2048sect --len 100MiB -o /dev/sde
```
### Benchmark and dry-run modes
Benchmark different modes and thread counts. Prints the options that perform the best on your system:
```sh
randquik --benchmark
```
To do a single run without actually writing anywhere, use `--dry`:
```sh
randquik --len 50GiB -t8 --dry
```
## Legacy
The original implementation is preserved in the [legacy](https://github.com/LeoVasanko/RandQuik/tree/legacy) git branch.
That version was once the fastest CSPRNG available, built around ChaCha20 with SIMD Assembly and C code written by me, making it faster than traditional algorithms without such optimizations and faster than the Linux kernel that also uses ChaCha20 to make random numbers. While historically significant, it has been greatly surpassed by the current AEGIS-based design in performance and features.
The legacy branch remains available for reference and benchmarking, but version 2 is the recommended implementation.
It is possible to seek ChaCha to any byte position in the stream without delay. This is implemented in C API only for now, and is not exposed via Numpy, Python or CLI interfaces.
+55
View File
@@ -0,0 +1,55 @@
#include <openssl/evp.h>
#include <stdio.h>
#include <string.h>
int main() {
// Key and IV should be appropriately sized for ChaCha20
unsigned char key[] = {
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f
};
unsigned char iv[] = {
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00
};
// Initialize context
EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
if (!ctx) {
perror("EVP_CIPHER_CTX_new failed");
return 1;
}
// Initialize the ChaCha20 cipher
if (!EVP_EncryptInit_ex(ctx, EVP_chacha20(), NULL, key, iv)) {
perror("EVP_EncryptInit_ex failed");
EVP_CIPHER_CTX_free(ctx);
return 1;
}
// Buffer for the keystream
unsigned char keystream[1000000];
memset(keystream, 0, sizeof(keystream));
for (int i = 0; i < 1000; i++) {
// Generate keystream
int len;
if (!EVP_EncryptUpdate(ctx, keystream, &len, keystream, sizeof keystream)) {
perror("EVP_EncryptUpdate failed");
EVP_CIPHER_CTX_free(ctx);
return 1;
}
}
// Clean up
EVP_CIPHER_CTX_free(ctx);
// Print the generated keystream
for (int i = 0; i < 64; i++) {
printf("%02x", keystream[i]);
}
printf("\n");
return 0;
}
+20
View File
@@ -0,0 +1,20 @@
#include "chacha20.h"
#include <stdio.h>
#include <stdlib.h>
int main(void) {
uint64_t N = 1000000;
uint8_t* buf = malloc(N);
const uint8_t key[32] = {0};
const uint8_t nonce[16] = {0};
for (uint64_t i = 0; i < 1000; ++i) {
cha_generate(buf, N, key, nonce);
}
for (unsigned i = 0; i < 16; ++i)
printf("%02X ", buf[i]);
for (unsigned i = 0; i < 16; ++i)
printf(" %02X", buf[1024 + i]);
puts("");
free(buf);
}
+12
View File
@@ -0,0 +1,12 @@
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
int main(void)
{
printf("RAND_MAX = %d (%.1lf bit)\n", RAND_MAX, log2(RAND_MAX));
const uint64_t rounds = 1000000000ull * 8 / (unsigned)log2(RAND_MAX);
for (uint64_t i = rounds; i-->0;) {
rand();
}
}
BIN
View File
Binary file not shown.

After

Width:  |  Height:  |  Size: 66 KiB

BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 10 KiB

+15
View File
@@ -0,0 +1,15 @@
project('randquik', 'c')
executable(
'randquik',
'src/cli.c',
c_args: ['-Wall', '-O3', '-march=native'],
install: true,
)
dependency('threads')
library(
'randquik-chacha20',
'src/charandom.c',
build_by_default: true,
c_args: ['-Wall', '-O3', '-march=native'],
)
+57 -36
View File
@@ -1,42 +1,63 @@
[project]
name = "randquik"
dynamic = ["version"]
description = "CLI tool for extremely fast random bytes."
readme = "README.md"
requires-python = ">=3.13"
authors = [
{ name = "Leo Vasanko" }
]
keywords = ["random", "CSPRNG", "AEGIS", "shred", "benchmark"]
classifiers = [
"Development Status :: 5 - Production/Stable",
"Environment :: Console",
"Intended Audience :: Developers",
"Intended Audience :: System Administrators",
"Topic :: Security :: Cryptography",
"Topic :: Utilities",
]
dependencies = [
"aeg>=0.4.3",
]
[project.scripts]
randquik = "randquik.cli:main"
[project.urls]
Homepage = "https://git.zi.fi/LeoVasanko/randquik"
Repository = "https://github.com/LeoVasanko/randquik"
[build-system]
requires = ["hatchling", "hatch-vcs"]
requires = ["hatchling", "hatch-vcs", "wheel", "cffi"]
build-backend = "hatchling.build"
[tool.hatch.version]
source = "vcs"
[project]
name = "randquik"
version = "0.1.0"
description = "Extremely fast and cryptographically secure random number generator."
readme = "README.md"
license.text = "Public Domain"
authors = [{ name = "Vasanko" }]
classifiers = [
"Operating System :: POSIX",
"Operating System :: Unix",
"Topic :: Security :: Cryptography",
"Topic :: Security",
"Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = ["cffi>=1.0.1", "numpy"]
requires-python = ">=3.10"
keywords = [
"random",
"generator",
"fast",
"secure",
"cryptographic",
"randomness",
]
[project.urls]
[project.optional-dependencies]
dev = ["pytest", "ruff", "cryptography", "scipy"]
[tool.hatchling]
[tool.ruff]
line-length = 100
target-version = "py313"
extend-select = ["I", "W", "UP", "C4", "ISC", "S"]
# Worth selecting but still too broken: ASYNC, B, DTZ, FA
ignore = [
"D100",
"D101",
"D102",
"D103",
"E402",
"E741",
"F811",
"F821",
# ruff format complains about these:
"ISC001",
"S101",
"S102",
"S104",
"S311",
"S603",
"S607",
"W191",
]
show-source = true
show-fixes = true
[tool.ruff.lint.isort]
known-first-party = ["randquik"]
[tool.pytest.ini_options]
pythonpath = ["."]
+2 -24
View File
@@ -1,25 +1,3 @@
"""RandQuik - High-performance cryptographic random data generator.
from randquik.cha import Cha, generate, generate_into
This package provides fast random data generation using AEGIS ciphers,
with support for multi-threaded generation and various I/O modes.
"""
from randquik.crypto import derive_key, generate_random_seed
from randquik.progress import ProgressDisplay
from randquik.stats import format_size, format_time
from randquik.utils import parse_size
try:
from randquik._version import __version__
except ImportError:
__version__ = "0.0.0.dev0"
__all__ = [
"ProgressDisplay",
"__version__",
"derive_key",
"format_size",
"format_time",
"generate_random_seed",
"parse_size",
]
__all__ = ["Cha", "generate_into", "generate"]
-6
View File
@@ -1,6 +0,0 @@
"""Entry point for running as `python -m randquik`."""
from randquik.cli import main
if __name__ == "__main__":
main()
-142
View File
@@ -1,142 +0,0 @@
"""Benchmark functions for measuring performance."""
import contextlib
import os
import pathlib
import re
import subprocess
import sys
import time
from randquik.utils import sparse_range
__all__ = ["bench_mode", "run_benchmark"]
def bench_mode(
tcounts: list[int],
io_mode: str,
length: str,
alg: str | None,
bench_file: pathlib.Path | None,
) -> list[tuple[int, float, list[str]]]:
max_repeats = 5
max_time = 0.5 # Quit early if more than 500ms has passed
results = []
# Use a file in current folder for file modes
if "file" in io_mode:
iocmd = ["-o", str(bench_file)]
elif "dry" in io_mode:
iocmd = ["--dry"]
elif "null" in io_mode:
iocmd = ["-o", os.devnull]
else:
raise ValueError(f"Unknown io_mode: {io_mode}")
# Print iocmd at start of row
sys.stdout.write(f"{' '.join(iocmd)[:20]:<20}")
sys.stdout.flush()
for workers in tcounts:
speeds = []
worker_start = time.perf_counter()
for rep in range(max_repeats):
if rep > 0 and (time.perf_counter() - worker_start) > max_time:
break
cmd = [
sys.executable,
"-m",
"randquik",
f"-l{length}",
f"-t{workers}",
*([f"-a{alg}"] if alg else []),
*iocmd,
]
try:
proc = subprocess.run(
cmd,
stdout=subprocess.DEVNULL,
stderr=subprocess.PIPE,
text=False,
)
except KeyboardInterrupt:
sys.stderr.write(
f" Interrupted\n\n>>> {' '.join(cmd)}\n"
) # 6 not 8 to account for ^C
sys.exit(1)
stderr = proc.stderr.decode(errors="ignore")
if proc.returncode != 0:
sys.stderr.write(f"{'ERROR':>8}\n\n>>> {' '.join(cmd)}\n{stderr}")
sys.exit(1)
m2 = re.findall(r"([0-9]+\.[0-9]+)\s+GB/s", stderr)
if m2:
speeds.append(float(m2[-1]))
if speeds:
sorted_speeds = sorted(speeds)
median = sorted_speeds[len(speeds) // 2]
sys.stdout.write(f"{median:>8.2f}")
sys.stdout.flush()
results.append((workers, median, iocmd))
else:
sys.stdout.write(f"{'---':>8}")
sys.stdout.flush()
sys.stdout.write("\n") # newline after row
sys.stdout.flush()
# Cleanup bench file
if bench_file:
with contextlib.suppress(OSError):
bench_file.unlink()
return results
def run_benchmark(args):
"""Run comprehensive benchmark across all I/O modes."""
# Check if output file already exists
bench_file = pathlib.Path(args.output or "test.dat")
if bench_file.exists():
if not args.output:
raise ValueError(
f"File test.dat already exists. Use -o {bench_file} to benchmark over it or choose another name."
)
bench_file.unlink()
try:
length = "1G" if args.len is None else args.len
max_threads = args.threads if args.threads is not None else os.cpu_count()
all_results = {}
tcounts = sparse_range(max_threads)
# Print header
header = f"{'randquik':<20}"
for w in tcounts:
header += f"{'-t' + str(w):>8}"
header += "\n" + "-" * (20 + 8 * len(tcounts)) + "\n"
sys.stdout.write(header)
for io_mode in ["dry", "null", "file"]:
results = bench_mode(tcounts, io_mode, length, alg=args.alg, bench_file=bench_file)
all_results[io_mode] = results
sys.stdout.write("-" * (20 + 8 * len(tcounts)) + "\n")
# Find fastest configuration and RNG speed
gen_speed = max(r[1] for res in all_results.values() for r in res)
best_speed, best_threads, best_iocmd = max(
[(sp, w, iocmd) for w, sp, iocmd in all_results["file"]],
)
threads = f" -t{best_threads}" if best_threads != 1 else ""
sys.stderr.write(
f"\n>>> Fastest wrote {best_speed:.2f} GB/s, plain RNG {gen_speed:.0f} GB/s\n"
f"randquik {' '.join(best_iocmd)}{threads}\n"
)
finally:
# Cleanup bench file even if interrupted
with contextlib.suppress(OSError):
bench_file.unlink()
+104
View File
@@ -0,0 +1,104 @@
import sys
from pathlib import Path
from typing import Any
import cffi
src = Path(__file__).parent.parent / "src"
if not src.is_dir():
raise RuntimeError("Unable to find RandQuik C sources in {src}")
ffi = cffi.FFI()
ffi.cdef(
"""
typedef uint64_t (*genfunc)(
uint8_t* out, size_t outsize, uint32_t state[16], unsigned rounds
);
typedef struct cha_ctx {
uint32_t input[16];
uint8_t unconsumed[512];
uint32_t offset, end;
unsigned rounds;
genfunc gen;
} cha_ctx;
int cha_generate(uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16], unsigned rounds);
void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds);
void cha_wipe(cha_ctx* ctx);
int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen);
"""
)
libname = "librandquik-chacha20.so"
if sys.platform == "darwin":
libname = "librandquik-chacha20.dylib"
elif sys.platform == "win32":
libname = "randquik-chacha20.dll"
lib = ffi.dlopen((Path(__file__).parent.parent / f"build/{libname}").as_posix())
def _processKeys(key, iv):
key = memoryview(key)
iv = memoryview(iv)
if key.nbytes != 32:
raise ValueError("key must be 32 bytes")
if iv.nbytes != 16:
# Allow original and IETF nonces with zero counter
if iv.nbytes == 8:
iv = bytes(8) + iv
elif iv.nbytes == 12:
iv = bytes(4) + iv
else:
raise ValueError(
"iv lenth must be 8 (ChaCha20 original), 12 (IETF) or 16 (counter in initial 8 bytes)"
)
return ffi.from_buffer(key), ffi.from_buffer(iv)
def _processBuffer(out):
try:
outlen = out.nbytes
except AttributeError:
out = memoryview(out)
outlen = out.nbytes
if getattr(out, "readonly", None):
raise ValueError("The output buffer must be writable, not e.g. `bytes`")
return ffi.from_buffer(out), outlen
class Cha:
def __init__(self, key: bytes | Any, iv: bytes | Any, *, rounds=20):
"""Construct a generator that holds its internal state, moving forward on each call."""
key, iv = _processKeys(key, iv)
self.ctx = ffi.new("cha_ctx*")
lib.cha_init(self.ctx, key, iv, rounds)
def __del__(self):
lib.cha_wipe(self.ctx)
def __call__(self, out: bytearray | Any):
"""Fill the parameter with random bytes"""
outbuf, outlen = _processBuffer(out)
lib.cha_update(self.ctx, outbuf, outlen)
return out
def generate_into(
out: bytearray | memoryview | Any,
key: bytes | Any,
iv: bytes | Any = bytes(16),
*,
rounds=20,
):
"""Fill in random bytes into an existing array (buffer interface)"""
key, iv = _processKeys(key, iv)
outbuf, outlen = _processBuffer(out)
lib.cha_generate(outbuf, outlen, key, iv, rounds)
return out
def generate(outlen: int, key: bytes | Any, iv: bytes | Any = bytes(16), *, rounds=20):
"""Return a bytearray of random bytes"""
assert outlen >= 0
return generate_into(bytearray(outlen), key, iv, rounds=rounds)
-206
View File
@@ -1,206 +0,0 @@
"""Command-line interface for RandQuik."""
import argparse
import gc
import sys
import aeg
from randquik.benchmark import run_benchmark
from randquik.crypto import derive_key, generate_random_seed
from randquik.utils import parse_size
from randquik.workers import run
__all__ = ["main"]
# Disable GC for performance
gc.disable()
DEFAULT_ALG = "AEGIS-128X2"
def prepare_seed(args):
"""Prepare seed and determine if it was generated."""
generated_seed = args.seed is None
seed = generate_random_seed() if generated_seed else args.seed
return seed, generated_seed
def prepare_key(seed, keybytes):
"""Derive key from seed."""
key = derive_key(seed, keybytes)
return key
def parse_seeks(args):
"""Parse seek values from arguments."""
try:
iseek = parse_size(args.iseek, args.output) or 0
oseek = parse_size(args.oseek, args.output) or 0
if args.seek:
iseek = oseek = parse_size(args.seek, args.output) or 0
except ValueError as e:
raise ValueError(f"Error parsing seek: {e}") from None
return iseek, oseek
def _main():
"""Internal main function that may raise exceptions."""
parser = argparse.ArgumentParser(description="Generate random bytes using AEGIS ciphers")
parser.add_argument("-s", "--seed", help="Alphanumeric seed string", type=str)
parser.add_argument(
"-l",
"--len",
help="Length to generate (e.g. 1g, 100mi, 1000sect)",
type=str,
default=None,
)
parser.add_argument("-o", "--output", help="Output file (default: stdout)", type=str)
parser.add_argument(
"-t",
"--threads",
help="Number of worker threads (benchmark: upper limit)",
type=int,
default=None,
)
parser.add_argument(
"-a",
"--alg",
help=f"Cipher algorithm (default: {DEFAULT_ALG})",
type=str,
default=None,
)
parser.add_argument(
"--benchmark",
action="store_true",
help="Run benchmark (generates 1GB and reports speed)",
)
parser.add_argument(
"--dry",
action="store_true",
help="Dry run: open output but skip writes (for benchmarking)",
)
parser.add_argument(
"--seek",
type=str,
default=None,
help="Seek both input stream and output to position (e.g. 1g, 100mi)",
)
parser.add_argument(
"--iseek",
type=str,
default=None,
help="Seek input random stream to position (overrides --seek)",
)
parser.add_argument(
"--oseek",
type=str,
default=None,
help="Seek output file to position (overrides --seek)",
)
parser.add_argument(
"-q",
"--quiet",
action="store_true",
help="Quiet mode: suppress all output except errors",
)
parser.add_argument(
"-v",
"--verbose",
action="count",
default=0,
help="Verbose mode: -v for I/O mode, -vv for worker statistics",
)
args = parser.parse_args()
# Normalize "-" output to None (stdout)
if args.output == "-":
args.output = None
ciph = aeg.cipher(args.alg or DEFAULT_ALG)
# Validate and process args
seed, generated_seed = prepare_seed(args)
key = prepare_key(seed, ciph.KEYBYTES)
iseek, oseek = parse_seeks(args)
total_bytes = parse_size(args.len) # None if not specified, 0 if -l0
# Always track the seed for commands, but only show repeat for generated seeds
seed_for_display = seed
if args.benchmark:
if args.seed is not None:
raise ValueError("Cannot specify seed in benchmark mode")
if iseek or oseek:
raise ValueError("Cannot use seek options in benchmark mode")
run_benchmark(args)
return
# Build continue command for interruption
action = "generated" if args.dry else "wrote"
continue_cmd = None
repeat_cmd = None
if args.output and seed_for_display:
# Will be updated with actual written bytes after run
continue_cmd = f"randquik -s {seed_for_display} --seek {{seek}} -o {args.output}"
if args.len:
continue_cmd += f" -l {args.len}"
# Build repeat command for randomly generated seeds (so user can reproduce)
if generated_seed and not args.quiet:
repeat_cmd = f"randquik -s {seed_for_display}"
if args.len:
repeat_cmd += f" -l {args.len}"
if args.output:
repeat_cmd += f" -o {args.output}"
# Run generation
workers = args.threads if args.threads is not None else 1
result = run(
output=args.output,
total_bytes=total_bytes,
iseek=iseek,
oseek=oseek,
key=key,
ciph=ciph,
workers=workers,
dry=args.dry,
quiet=args.quiet,
action=action,
)
# Set repeat command for generated seeds
result.repeat_cmd = repeat_cmd
# Update continue command with actual written bytes
if result.interrupted and result.written > 0 and args.output:
new_iseek = iseek + result.written
new_oseek = oseek + result.written
if new_iseek == new_oseek:
result.continue_cmd = (
f"randquik -s {seed_for_display} --seek {new_iseek} -o {args.output}"
)
else:
result.continue_cmd = f"randquik -s {seed_for_display} --iseek {new_iseek} --oseek {new_oseek} -o {args.output}"
if args.len:
result.continue_cmd += f" -l {args.len}"
# Print summary
show_summary = not args.quiet or args.verbose >= 1 or result.interrupted
if show_summary and (total_bytes is not None or result.interrupted):
result.print_summary(verbose=args.verbose)
if args.verbose >= 2:
result.print_detailed_stats()
if result.interrupted:
sys.exit(1)
def main():
"""Main entry point for the CLI with exception handling."""
try:
_main()
except (KeyboardInterrupt, BrokenPipeError):
sys.exit(1)
except ValueError as e:
sys.stderr.write(f"Error: {e}\n")
sys.exit(1)
-22
View File
@@ -1,22 +0,0 @@
"""Cryptographic functions for key derivation and seed generation."""
import hashlib
import secrets
import string
__all__ = [
"derive_key",
"generate_random_seed",
]
def generate_random_seed() -> str:
"""Generate a random alphanumeric seed string."""
chars = string.ascii_letters + string.digits
return "".join(secrets.choice(chars) for _ in range(16))
def derive_key(seed: str, key_bytes: int) -> bytes:
"""Derive a key from a seed string using SHA-512."""
assert 16 <= key_bytes <= 64, "Only 128-512 bits supported"
return hashlib.sha512(seed.encode()).digest()[:key_bytes]
-115
View File
@@ -1,115 +0,0 @@
"""File I/O helpers for output handling."""
import contextlib
import errno
import os
import pathlib
import sys
from collections.abc import Generator
__all__ = [
"open_fd",
"open_memoryview",
]
def _open_output(
output_path: str,
total_bytes: int | None,
oseek: int = 0,
) -> tuple[int, bool]:
"""Open output file descriptor, preallocate and apply platform hints.
Returns:
Tuple of (file descriptor, whether we created the file)
"""
created = False
path = pathlib.Path(output_path)
created = not path.exists()
flags = os.O_WRONLY | os.O_CREAT
fd = os.open(str(path), flags, 0o644)
required_size = oseek + (total_bytes if total_bytes is not None else 0)
current_size = os.fstat(fd).st_size
if required_size > current_size:
with contextlib.suppress(OSError):
os.ftruncate(fd, required_size)
# Seek to output position
if oseek > 0:
try:
os.lseek(fd, oseek, os.SEEK_SET)
except OSError as e:
raise ValueError(
f"Cannot oseek in {output_path or 'stdout'}. Use only --iseek or specify a seekable file."
) from e
# macOS: try to bypass unified buffer cache (F_NOCACHE)
if sys.platform == "darwin":
try:
import fcntl
fcntl.fcntl(fd, fcntl.F_NOCACHE, 1)
except (OSError, AttributeError, ImportError):
pass
return fd, created
@contextlib.contextmanager
def open_fd(
output_path: str | None,
total_bytes: int | None,
dry: bool = False,
oseek: int = 0,
) -> Generator[int]:
"""Context manager for output file descriptor.
Args:
output_path: Path to output file, or None for stdout
total_bytes: Total bytes to write
dry: If True, skip truncation/preallocation and tty check
oseek: Seek position for output
Yields:
Integer file descriptor
"""
if dry:
yield -1
return
if not output_path:
if sys.stdout.isatty():
raise ValueError("Refusing to write binary data to terminal. Use -o to specify a file.")
yield sys.stdout.fileno()
return
fd, created = _open_output(output_path, total_bytes, oseek)
try:
yield fd
except OSError as e:
if e.errno == errno.ENOSPC:
# Clean up file we created on disk full
if created:
with contextlib.suppress(Exception):
os.unlink(output_path)
raise ValueError(f"No space left on device: {output_path}") from None
raise
finally:
with contextlib.suppress(Exception):
os.close(fd)
@contextlib.contextmanager
def open_memoryview(buf) -> Generator[memoryview]:
"""Context manager for memoryview.
Args:
buf: Buffer to create memoryview from
Yields:
memoryview object
"""
view = memoryview(buf)
try:
yield view
finally:
view.release()
-844
View File
@@ -1,844 +0,0 @@
"""Progress display with speed graph at bottom of terminal."""
import math
import os
import sys
import threading
import time
from randquik.stats import format_size, format_time
__all__ = ["ProgressDisplay"]
# Unicode block characters for graph (8 levels per cell)
GRAPH_BLOCKS = " ▁▂▃▄▅▆▇█"
# Maximum height of progress display in terminal rows
MAX_HEIGHT = 10
class ProgressDisplay:
"""Progress display with speed graph at bottom of terminal, updated every 100ms.
Only active when stderr is a tty. Reads progress from a shared state dict
with a single 'written' key. All display logic is encapsulated here.
Uses the bottom portion of the terminal with a scrolling region preserved
at the top, allowing normal output to scroll above the progress display.
The graph fills from left to right as progress advances, doubling as both
a progress bar and a speed-over-time visualization.
"""
def __init__(
self,
total_bytes: int | None,
start_time: float,
state: dict,
infinite: bool | None = None,
output_name: str | None = None,
oseek: int = 0,
):
self.total_bytes = total_bytes
self.start_time = start_time
self.state = state # Must have 'written' key
self.infinite = infinite if infinite is not None else total_bytes is None
self.output_name = output_name or "<stdout>"
self.oseek = oseek
self.active = sys.stderr.isatty()
self._stop = threading.Event()
self._thread = None
self._last_written = 0
self._last_time = start_time
# Speed history for graph - fixed size, filled from left as progress advances
self._graph_width = 80 # Will be updated on first render
self._speed_history: list[float] = [] # Stores GB/s values, one per column
self._max_speed: float = 0.01 # Start with small value to avoid div by zero
# For infinite mode: track time of each speed sample
self._time_history: list[float] = []
# X-axis scale smoothing (hysteresis for estimated total time)
self._smoothed_scale_time: float | None = None
# Terminal handling
self._current_scroll_bottom: int | None = None
self._hidden_cursor = False
self._first_draw = True
def start(self):
if not self.active:
return
self._setup_terminal_state()
self._thread = threading.Thread(target=self._run, daemon=True)
self._thread.start()
def stop(self):
if not self.active or self._thread is None:
return
self._stop.set()
self._thread.join(timeout=0.5)
# Final render so the finished state stays on screen
if self.active:
cols, rows, lines, overlay = self._render_frame()
self._draw_frame(cols, rows, lines, overlay)
self._restore_terminal_state()
def _setup_terminal_state(self):
"""Prepare terminal: hide cursor and start using bottom reserved block."""
sys.stderr.write("\x1b[?25l") # Hide cursor to reduce flicker
sys.stderr.flush()
self._hidden_cursor = True
def _get_smoothed_speed(self, window_secs: float = 1.0) -> float:
"""Calculate average speed over the last window_secs seconds.
Returns speed in bytes/sec, averaged from recent samples in _speed_history.
Falls back to the most recent sample if not enough history.
"""
if not self._speed_history or not self._time_history:
return 0.0
current_time = self._time_history[-1]
cutoff_time = current_time - window_secs
# Find samples within the window
total_speed = 0.0
count = 0
for idx in range(len(self._time_history) - 1, -1, -1):
if self._time_history[idx] < cutoff_time:
break
total_speed += self._speed_history[idx]
count += 1
if count == 0:
return self._speed_history[-1] * 1_000_000_000
# Return average in bytes/sec (speed_history stores GB/s)
return (total_speed / count) * 1_000_000_000
def _restore_terminal_state(self):
"""Restore terminal scrolling and cursor after progress is done."""
# Reset scrolling region to full screen
sys.stderr.write("\x1b[r")
self._current_scroll_bottom = None
# Move cursor to a fresh line under the progress block
cols, rows = self._get_terminal_size()
sys.stderr.write(f"\x1b[{rows};1H\n")
if self._hidden_cursor:
sys.stderr.write("\x1b[?25h")
self._hidden_cursor = False
sys.stderr.flush()
def _get_terminal_size(self) -> tuple[int, int]:
"""Return (columns, rows)."""
try:
size = os.get_terminal_size(sys.stderr.fileno())
return size.columns, size.lines
except (OSError, ValueError):
return 80, 24
def _render_graph_row(
self,
values: list[float],
max_val: float,
row: int,
total_rows: int,
width: int,
avg_speed: float = 0,
) -> str:
"""Render one row of the graph using Unicode blocks.
Row 0 is top, total_rows-1 is bottom. Each cell can show 8 levels.
Values list may be shorter than width (unfilled area shown as dim bar at avg_speed).
"""
filled_chars = []
unfilled_chars = []
filled_cols = len(values)
# Calculate the row threshold for average speed
avg_normalized = (avg_speed / max_val) * total_rows * 8 if max_val > 0 else 0
row_bottom = (total_rows - row - 1) * 8
row_top = row_bottom + 8
# Build filled portion
for i in range(filled_cols):
v = values[i]
# Normalize value to 0..total_rows*8 range
normalized = (v / max_val) * total_rows * 8 if max_val > 0 else 0
if normalized <= row_bottom:
filled_chars.append(" ")
elif normalized >= row_top:
filled_chars.append("")
else:
level = math.ceil(normalized - row_bottom)
filled_chars.append(GRAPH_BLOCKS[min(level, 8)])
# Build unfilled portion (dim grey at avg_speed level)
unfilled_width = width - filled_cols
if unfilled_width > 0:
if avg_normalized <= row_bottom:
unfilled_char = " "
elif avg_normalized >= row_top:
unfilled_char = ""
else:
level = int(avg_normalized - row_bottom)
unfilled_char = GRAPH_BLOCKS[min(level, 8)]
unfilled_chars = [unfilled_char] * unfilled_width
# Combine with color codes only at transitions
filled_str = "".join(filled_chars)
unfilled_str = "".join(unfilled_chars)
if unfilled_str:
return f"{filled_str}\x1b[0m\x1b[38;5;235m{unfilled_str}\x1b[0m\x1b[33m"
return filled_str
def _build_header(
self,
cols: int,
written: int,
speed: float,
elapsed: float,
eta: float | None = None,
total_bytes: int | None = None,
) -> str:
"""Build the header line with stats, output name, and position.
Args:
cols: Terminal width
written: Bytes written so far
speed: Current speed in bytes/sec
elapsed: Elapsed time in seconds
eta: Estimated time remaining (None for infinite mode)
total_bytes: Total bytes to write (None for infinite mode)
"""
spinner = "\u25d0\u25d3\u25d1\u25d2"[int(elapsed * 4) % 4]
written_gb = written / 1_000_000_000
speed_gbs = speed / 1_000_000_000
# Build the fixed stats portion
if total_bytes is not None:
# Finite mode: show progress and ETA
total_gb = total_bytes / 1_000_000_000
eta_str = format_time(eta) if eta is not None else "--"
stats = (
f"\x1b[1;36mRandQuik {spinner}\x1b[0m "
f"{written_gb:6.2f}\x1b[2m/\x1b[0m{total_gb:.2f} GB "
)
stats += (
f"\x1b[2m@\x1b[0m {speed_gbs:5.2f} GB/s \x1b[2mest.\x1b[0m {eta_str:<8}"
if written < total_bytes
else f"\x1b[2m{'done':>27}\x1b[0m"
)
# Visible: "RandQuik X " (14) + "XXXX.XX/XXXX.XX GB " (20) + "@ XX.XX GB/s " (15) + "est. XXXXXXXX" (13) = 62
stats_len = 62
else:
# Infinite mode: show written and elapsed
stats = (
f" \x1b[1;36mRandQuik {spinner}\x1b[0m "
f"{written_gb:6.2f} GB \x1b[2m\u221e\x1b[0m "
f"\x1b[2m@\x1b[0m {speed_gbs:5.2f} GB/s "
f"\x1b[2m\u2502\x1b[0m {format_time(elapsed):>8}"
)
# Visible: " RandQuik X " (16) + "XXXX.XX GB \u221e " (14) + "@ XX.XX GB/s " (15) + "\u2502 XXXXXXXX" (11) = 56
stats_len = 56
# Build position suffix if oseek was used
if self.oseek > 0:
file_pos = self.oseek + written
pos_str = format_size(file_pos).replace(" ", "")
if total_bytes is not None:
pos_suffix = f" \x1b[2m[\x1b[0m{pos_str}\x1b[2m]\x1b[0m"
pos_suffix_len = 3 + len(pos_str) # " []" + size
else:
pos_suffix = f" \x1b[2m@\x1b[0m{pos_str}"
pos_suffix_len = 2 + len(pos_str) # " @" + size
else:
pos_suffix = ""
pos_suffix_len = 0
# Calculate available space for filename
# Format: {stats} > {name}{pos_suffix}
available = cols - stats_len - 4 - pos_suffix_len # 4 for " > "
name = self.output_name
if len(name) > available > 3:
name = "\u2026" + name[-(available - 1) :]
elif available <= 3:
name = ""
if name:
return f"{stats} \x1b[2m>\x1b[0m {name}{pos_suffix}"
return stats
def _render_progress_block(
self, cols: int, rows: int, max_height: int
) -> tuple[list[str], tuple[int, int, str] | None]:
"""Render the progress block constrained to max_height lines."""
if self.infinite:
return self._render_infinite_block(cols, rows, max_height), None
return self._render_finite_block(cols, rows, max_height)
def _render_infinite_block(self, cols: int, rows: int, max_height: int) -> list[str]:
"""Render progress block for infinite mode (no known total)."""
written = self.state.get("written", 0)
now = time.perf_counter()
elapsed = now - self.start_time
# Calculate graph width (leave room for Y-axis labels)
graph_width = max(10, cols - 8)
self._graph_width = graph_width
# Calculate speeds
overall_speed = written / elapsed if elapsed > 0 else 0
dt = now - self._last_time
instant_speed = (written - self._last_written) / dt if dt > 0 else 0
self._last_written = written
self._last_time = now
# Update speed history
speed_gbs = instant_speed / 1_000_000_000
self._speed_history.append(speed_gbs)
self._time_history.append(elapsed)
# Update max speed
if speed_gbs > self._max_speed:
self._max_speed = speed_gbs
scale_max = self._nice_scale(self._max_speed)
# Use smoothed speed for header display
display_speed = self._get_smoothed_speed()
# Build output
lines: list[str] = []
header = self._build_header(cols, written, display_speed, elapsed)
lines.append(header)
# Calculate graph dimensions based on remaining height
# Layout: [header][GB/s label][graph rows][time axis]
remaining_after_label = max_height - len(lines) - 2
if remaining_after_label < 1:
# Terminal is too short; fall back to header-only view
return lines
graph_rows = max(1, remaining_after_label)
# Downsample speed history for display - average samples within each column's time range
min_scale_time = 10.0
scale_time = max(elapsed, min_scale_time)
col_width_time = scale_time / (graph_width - 1) if graph_width > 1 else scale_time
display_values = []
for col in range(graph_width):
col_time = col / (graph_width - 1) * scale_time if graph_width > 1 else 0
if col_time > elapsed:
break
# Find all samples within this column's time range
col_start = col_time - col_width_time / 2
col_end = col_time + col_width_time / 2
samples = [
self._speed_history[idx]
for idx, t in enumerate(self._time_history)
if col_start <= t <= col_end
]
if samples:
display_values.append(sum(samples) / len(samples))
elif self._speed_history:
# Fallback to closest if no samples in range
best_idx = 0
best_diff = float("inf")
for idx, t in enumerate(self._time_history):
diff = abs(t - col_time)
if diff < best_diff:
best_diff = diff
best_idx = idx
display_values.append(self._speed_history[best_idx])
avg_speed_gbs = overall_speed / 1_000_000_000
# Use MB/s scale if max speed < 1 GB/s
use_mb = scale_max < 1
unit_label = "MB/s" if use_mb else "GB/s"
lines.append(f" \x1b[36m{unit_label}\x1b[0m")
# Compute nice Y-axis tick values and map each to its best row
nice_ticks = self._nice_y_ticks(scale_max, graph_rows)
row_labels = self._assign_ticks_to_rows(nice_ticks, scale_max, graph_rows, use_mb)
for row in range(graph_rows):
graph_line = self._render_graph_row(
display_values,
scale_max,
row,
graph_rows,
len(display_values),
avg_speed_gbs,
)
graph_line = graph_line.ljust(graph_width)
# Y-axis label from pre-computed mapping
label = row_labels.get(row, " ")
lines.append(f" \x1b[36m{label}\x1b[0m \x1b[33m{graph_line}\x1b[0m")
# Time axis
time_axis = self._build_infinite_time_axis(graph_width, scale_time)
lines.append(f" {''.join(time_axis)}")
return lines
def _nice_scale(self, max_speed: float) -> float:
"""Round up to next nice number for scale."""
if max_speed <= 0.01:
return 0.01
log_val = math.log10(max_speed)
power = math.floor(log_val)
mantissa = max_speed / (10**power)
nice_mantissa = math.ceil(mantissa)
if nice_mantissa > 9:
nice_mantissa = 1
power += 1
return nice_mantissa * (10**power)
def _format_label(self, val: float, use_mb: bool = False) -> str:
"""Format Y-axis label.
Args:
val: Value in GB/s (will be converted to MB/s if use_mb is True)
use_mb: If True, multiply by 1000 and format as MB/s values
"""
if use_mb:
val = val * 1000 # Convert GB/s to MB/s
if val == 0:
return "0"
elif val >= 1:
return f"{val:.0f}"
else:
return f"{val:.1f}"
def _nice_y_ticks(self, scale_max: float, graph_rows: int = 10) -> list[float]:
"""Return nice Y-axis tick values from 0 to scale_max.
Chooses a nice interval (1, 2, 5 × 10^N) that gives labels with
sufficient spacing (at least 3 rows between labels).
"""
if scale_max <= 0:
return [0]
# We want at least 5 empty rows between labels for readability
min_row_spacing = 5
max_ticks = max(2, graph_rows // min_row_spacing)
# Nice intervals: 0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, ...
nice_bases = [1, 2, 5]
best_interval = scale_max
for exp in range(-1, 10):
for base in nice_bases:
interval = base * (10**exp)
num_ticks = scale_max / interval
if 2 <= num_ticks <= max_ticks:
best_interval = interval
break
else:
continue
break
# Generate ticks from 0 to scale_max at best_interval
ticks = []
val = 0.0
while val <= scale_max + 1e-9:
ticks.append(val)
val += best_interval
return ticks
def _assign_ticks_to_rows(
self, ticks: list[float], scale_max: float, graph_rows: int, use_mb: bool = False
) -> dict[int, str]:
"""Assign each tick to the row closest to its value.
Returns a dict mapping row index to formatted label string.
Each tick is assigned to exactly one row.
Args:
ticks: List of tick values in GB/s
scale_max: Maximum scale value in GB/s
graph_rows: Number of rows in the graph
use_mb: If True, format labels as MB/s instead of GB/s
"""
row_labels: dict[int, str] = {}
if graph_rows <= 1 or scale_max <= 0:
return {0: f"{self._format_label(ticks[0] if ticks else 0, use_mb):>4}"}
for tick in ticks:
# Calculate which row this tick value corresponds to
# Row 0 is top (scale_max), row graph_rows-1 is bottom (0)
exact_row = (1 - tick / scale_max) * (graph_rows - 1)
best_row = round(exact_row)
best_row = max(0, min(graph_rows - 1, best_row))
# Only assign if row is not already taken (first tick wins)
if best_row not in row_labels:
row_labels[best_row] = f"{self._format_label(tick, use_mb):>4}"
return row_labels
def _build_infinite_time_axis(self, graph_width: int, scale_time: float) -> list[str]:
"""Build time axis for infinite mode with nice interval labels.
Shows from 0 to scale_time with nice interval markers.
"""
time_axis = [" "] * graph_width
# Nice time intervals
nice_intervals = [
1,
2,
5,
10,
15,
30,
60,
120,
300,
600,
900,
1800,
3600,
7200,
18000,
36000,
]
def format_time_short(secs):
"""Format time for axis label."""
if secs == 0:
return "0"
elif secs < 120:
return f"{int(secs)}s"
elif secs < 3600:
m = int(secs // 60)
s = int(secs % 60)
if s == 0:
return f"{m}m"
return f"{m}m{s}s"
else:
h = int(secs // 3600)
m = int((secs % 3600) // 60)
if m == 0:
return f"{h}h"
return f"{h}h{m}m"
# Find a nice interval that gives us ~4-8 labels
interval = nice_intervals[-1]
for ni in nice_intervals:
if scale_time / ni <= 8:
interval = ni
break
# Place labels at nice intervals starting from 0
t = 0
while t <= scale_time:
col = int(t / scale_time * (graph_width - 1)) if scale_time > 0 else 0
if 0 <= col < graph_width:
label = format_time_short(t)
label_start = max(0, col - len(label) // 2)
label_end = min(graph_width, label_start + len(label))
if all(c == " " for c in time_axis[label_start:label_end]):
for i, ch in enumerate(label):
if label_start + i < graph_width:
time_axis[label_start + i] = ch
t += interval
return time_axis
def _render_finite_block(
self, cols: int, rows: int, max_height: int
) -> tuple[list[str], tuple[int, int, str] | None]:
"""Render the progress block for finite progress."""
written = self.state.get("written", 0)
now = time.perf_counter()
elapsed = now - self.start_time
# Calculate graph width (leave room for Y-axis labels)
graph_width = max(10, cols - 8)
self._graph_width = graph_width
# Calculate speeds
overall_speed = written / elapsed if elapsed > 0 else 0
dt = now - self._last_time
instant_speed = (written - self._last_written) / dt if dt > 0 else 0
self._last_written = written
self._last_time = now
# Collect time-based speed samples (like infinite mode)
speed_gbs = instant_speed / 1_000_000_000
self._speed_history.append(speed_gbs)
self._time_history.append(elapsed)
# Use smoothed speed for header display and ETA
display_speed = self._get_smoothed_speed()
# ETA and total estimated time
remaining = self.total_bytes - written
# ETA based on smoothed speed for stability
eta = remaining / display_speed if display_speed > 0 else -1
# Graph X-axis scaling based on overall average speed for stability
avg_eta = remaining / overall_speed if overall_speed > 0 else -1
estimated_total_time = elapsed + avg_eta if avg_eta > 0 else elapsed
# Apply hysteresis to scale_time to prevent jumping
# Only update if change is significant (>20%) or if new estimate is larger
raw_scale_time = max(estimated_total_time, 1.0)
if self._smoothed_scale_time is None:
self._smoothed_scale_time = raw_scale_time
else:
# Always grow immediately, shrink only gradually
if raw_scale_time > self._smoothed_scale_time:
self._smoothed_scale_time = raw_scale_time
else:
# Shrink slowly: blend 90% old, 10% new
self._smoothed_scale_time = 0.9 * self._smoothed_scale_time + 0.1 * raw_scale_time
scale_time = self._smoothed_scale_time
# Progress percentage (for display)
pct = min(100, written * 100 / self.total_bytes) if self.total_bytes > 0 else 0
# Update max speed
if speed_gbs > self._max_speed:
self._max_speed = speed_gbs
scale_max = self._nice_scale(self._max_speed)
# Build output
lines: list[str] = []
header = self._build_header(
cols, written, display_speed, elapsed, eta=eta, total_bytes=self.total_bytes
)
lines.append(header)
# Calculate graph dimensions based on remaining height
# Layout: [header][GB/s label][graph rows][time axis]
remaining_after_label = max_height - len(lines) - 2
if remaining_after_label < 1:
# Terminal is too short; fall back to a compact header-only view
return lines, None
graph_rows = max(1, remaining_after_label)
# Downsample speed history for display - average samples within each column's time range
# Graph fills based on elapsed / scale_time (smoothed)
col_width_time = scale_time / (graph_width - 1) if graph_width > 1 else scale_time
# Calculate how many columns should be filled based on elapsed time
if scale_time > 0:
time_pct = min(100, elapsed * 100 / scale_time)
else:
time_pct = 100
target_cols = min(graph_width, int(graph_width * time_pct / 100) + 1) if time_pct > 0 else 0
display_values = []
for col in range(target_cols):
col_time = col / (graph_width - 1) * scale_time if graph_width > 1 else 0
if col_time > elapsed:
break
# Find all samples within this column's time range
col_start = col_time - col_width_time / 2
col_end = col_time + col_width_time / 2
samples = [
self._speed_history[idx]
for idx, t in enumerate(self._time_history)
if col_start <= t <= col_end
]
if samples:
display_values.append(sum(samples) / len(samples))
elif self._speed_history:
# Fallback to closest if no samples in range
best_idx = 0
best_diff = float("inf")
for idx, t in enumerate(self._time_history):
diff = abs(t - col_time)
if diff < best_diff:
best_diff = diff
best_idx = idx
display_values.append(self._speed_history[best_idx])
# Render graph rows with Y-axis
avg_speed_gbs = overall_speed / 1_000_000_000
# Use MB/s scale if max speed < 1 GB/s
use_mb = scale_max < 1
unit_label = "MB/s" if use_mb else "GB/s"
lines.append(f" \x1b[36m{unit_label}\x1b[0m")
# Compute nice Y-axis tick values and map each to its best row
nice_ticks = self._nice_y_ticks(scale_max, graph_rows)
row_labels = self._assign_ticks_to_rows(nice_ticks, scale_max, graph_rows, use_mb)
for row in range(graph_rows):
graph_line = self._render_graph_row(
display_values,
scale_max,
row,
graph_rows,
graph_width,
avg_speed_gbs,
)
graph_line = graph_line.ljust(graph_width)
# Y-axis label from pre-computed mapping
label = row_labels.get(row, " ")
lines.append(f" \x1b[36m{label}\x1b[0m \x1b[33m{graph_line}\x1b[0m")
# Time labels on X-axis with nice intervals
time_axis = self._build_time_axis(graph_width, scale_time)
lines.append(f" {''.join(time_axis)}")
# Position percentage at top of bar at current progress point
# Find the height of the bar at current progress (last value in display_values)
current_speed_gbs = display_values[-1] if display_values else 0
# Normalize to find which row the top of the bar is at
# Bar fills from bottom up; row 0 is top, graph_rows-1 is bottom
if scale_max > 0 and current_speed_gbs > 0:
# How many rows from bottom does the bar fill?
bar_height_fraction = current_speed_gbs / scale_max
# The top of the bar is at this row (0 = top, graph_rows-1 = bottom)
bar_top_row = int(graph_rows * (1 - bar_height_fraction))
bar_top_row = max(0, min(graph_rows - 1, bar_top_row))
else:
bar_top_row = graph_rows - 1 # At bottom if no speed
# Overlay position inside the progress block
progress_col = int(pct / 100 * (graph_width - 1)) + 7
pct_label = f"{int(pct)}%"
pct_col = max(7, min(cols, progress_col - len(pct_label) // 2))
first_graph_row = len(lines) - (graph_rows + 2)
pct_row_offset = first_graph_row + bar_top_row
pct_position = (pct_row_offset, pct_col, f"\x1b[1;37m{pct_label}\x1b[0m")
return lines, pct_position
def _build_time_axis(self, graph_width: int, estimated_total_time: float) -> list[str]:
"""Build time axis with nice interval labels."""
def nice_time_interval(total_secs):
"""Return a nice interval for time axis labels."""
nice_intervals = [
1,
2,
5,
10,
15,
30,
60,
120,
300,
600,
900,
1800,
3600,
7200,
18000,
36000,
]
for interval in nice_intervals:
if total_secs / interval <= 8:
return interval
return 36000
def format_time_short(secs):
"""Format time for axis label."""
if secs == 0:
return "0s"
elif secs < 120:
return f"{int(secs)}s"
elif secs < 3600:
m = int(secs // 60)
s = int(secs % 60)
if s == 0:
return f"{m}m"
return f"{m}m{s}s"
else:
h = int(secs // 3600)
m = int((secs % 3600) // 60)
if m == 0:
return f"{h}h"
return f"{h}h{m}m"
time_axis = [" "] * graph_width
if estimated_total_time > 0:
interval = nice_time_interval(estimated_total_time)
t = 0
while t <= estimated_total_time:
col = (
int(t / estimated_total_time * (graph_width - 1))
if estimated_total_time > 0
else 0
)
if col < graph_width:
label = format_time_short(t)
start = max(0, col - len(label) // 2)
end = min(graph_width, start + len(label))
if all(c == " " for c in time_axis[start:end]):
for i, ch in enumerate(label):
if start + i < graph_width:
time_axis[start + i] = ch
t += interval
else:
time_axis = list(f"{'0s':<{graph_width}}")
return time_axis
def _render_frame(self) -> tuple[int, int, list[str], tuple[int, int, str] | None]:
cols, rows = self._get_terminal_size()
max_height = min(MAX_HEIGHT, rows) if rows > 0 else MAX_HEIGHT
lines, overlay = self._render_progress_block(cols, rows, max_height)
# Guard against empty renders
if not lines:
lines = [""]
return cols, rows, lines, overlay
def _draw_frame(
self, cols: int, rows: int, lines: list[str], overlay: tuple[int, int, str] | None
):
"""Draw the progress block at the bottom of the terminal."""
height = min(len(lines), max(1, rows))
progress_top = max(1, rows - height + 1)
# Build entire frame as a single string
buf: list[str] = []
# On first draw, scroll terminal up to make room for progress block
if self._first_draw:
self._first_draw = False
buf.append("\n" * height)
# Update scrolling region so other output scrolls above the progress block
top = 1
bottom = max(1, rows - height)
if self._current_scroll_bottom != bottom:
buf.append(f"\x1b[{top};{bottom}r")
self._current_scroll_bottom = bottom
# Paint each progress line
for idx in range(height):
row = progress_top + idx
line = lines[idx]
buf.append(f"\x1b[{row};1H\x1b[2K{line}")
# Overlay (e.g., percent marker)
if overlay:
row_offset, col, text = overlay
abs_row = progress_top + row_offset
abs_col = max(1, min(cols, col))
buf.append(f"\x1b[{abs_row};{abs_col}H{text}")
# Place cursor back at the bottom of the scrolling region
anchor_row = max(1, progress_top - 1)
buf.append(f"\x1b[{anchor_row};1H")
# Single atomic write
sys.stderr.write("".join(buf))
sys.stderr.flush()
def _run(self):
"""Background thread: update display every 100ms."""
while not self._stop.wait(0.1):
cols, rows, lines, overlay = self._render_frame()
self._draw_frame(cols, rows, lines, overlay)
-269
View File
@@ -1,269 +0,0 @@
"""Statistics collection and formatting for workers."""
import re
import sys
import time
from dataclasses import dataclass, field
__all__ = [
"ConsumerStats",
"RunResult",
"SingleThreadedStats",
"WorkerStats",
"format_size",
"format_time",
"format_worker_stats_report",
"stopwatch",
]
def stopwatch():
"""Generator that yields elapsed time since last yield."""
t = time.perf_counter()
while True:
now = time.perf_counter()
yield now - t
t = now
def format_size(size: float) -> str:
"""Format bytes as human-readable size."""
for unit in ["B", "kB", "MB", "GB", "TB"]:
if abs(size) < 1000:
return f"{size:.0f} {unit}"
size /= 1000
return f"{size:.0f} PB"
def format_time(seconds: float) -> str:
"""Format seconds as human-readable time."""
if seconds < 0:
return "--"
if seconds < 1:
return f"{seconds * 1000:.0f}ms"
if seconds < 120:
return f"{int(seconds)}s"
elif seconds < 3600:
m = int(seconds // 60)
s = int(seconds % 60)
if s == 0:
return f"{m}m"
return f"{m}m{s}s"
elif seconds < 172800: # 48 hours
h = int(seconds // 3600)
m = int((seconds % 3600) // 60)
if m == 0:
return f"{h}h"
return f"{h}h{m}m"
else:
d = int(seconds // 86400)
h = int((seconds % 86400) // 3600)
if h == 0:
return f"{d}d"
return f"{d}d{h}h"
@dataclass
class WorkerStats:
"""Timing statistics for a worker thread."""
worker_id: int = -1
# Lock timing breakdown
lock_acquire_time: float = 0.0 # Time to acquire the lock (contention)
lock_wait_space_time: float = 0.0 # Time waiting for has_space condition
lock_claim_time: float = 0.0 # Time inside lock claiming block number
lock_notify_time: float = 0.0 # Time inside lock marking ready + notify
# Work timing
crypto_time: float = 0.0
# Counters
blocks_processed: int = 0
bytes_generated: int = 0
wait_cycles: int = 0 # How many times we had to wait for space
def total_time(self) -> float:
"""Total measured time."""
return (
self.lock_acquire_time
+ self.lock_wait_space_time
+ self.lock_claim_time
+ self.lock_notify_time
+ self.crypto_time
)
@dataclass
class ConsumerStats:
"""Timing statistics for the consumer thread."""
wait_time: float = 0.0
write_time: float = 0.0
def total_time(self) -> float:
return self.wait_time + self.write_time
@dataclass
class SingleThreadedStats:
"""Timing statistics for single-threaded mode."""
crypto_time: float = 0.0
write_time: float = 0.0
def total_time(self) -> float:
return self.crypto_time + self.write_time
@dataclass
class RunResult:
"""Result from running a worker."""
written: int
elapsed: float
interrupted: bool
action: str = "wrote"
# Raw stats objects (optional, for verbose output)
consumer_stats: ConsumerStats | None = None
singlethreaded_stats: SingleThreadedStats | None = None
worker_stats: list[WorkerStats] = field(default_factory=list)
# For continue command on interrupt
continue_cmd: str | None = None
# For repeat command (when seed was randomly generated)
repeat_cmd: str | None = None
def _format_io_stats(self) -> str | None:
"""Format I/O timing stats for the summary line."""
if self.singlethreaded_stats is not None:
st = self.singlethreaded_stats
tt = st.total_time()
if tt <= 0:
return None
return f"crypto {st.crypto_time / tt:.0%} — write {st.write_time / tt:.0%}"
elif self.consumer_stats is not None:
cs = self.consumer_stats
tt = cs.total_time()
if tt <= 0:
return None
return f"wait {cs.wait_time / tt:.0%} — write {cs.write_time / tt:.0%}"
return None
def print_summary(self, verbose: int = 0):
"""Print a nice one-liner summary with optional colors."""
speed_gbs = (self.written / 1_000_000_000) / self.elapsed if self.elapsed > 0 else 0
size_str = format_size(self.written)
time_str = format_time(self.elapsed)
# I/O stats for verbose mode
io_stats = self._format_io_stats() if verbose >= 1 else None
stats_fmt = f"\033[0;32m • {io_stats}" if io_stats else ""
status_fmt = " \033[31m(interrupted)\033[0m" if self.interrupted else ""
# Continue or repeat command
cmd_line = ""
if self.interrupted and self.continue_cmd:
cmd_line = f"\n\033[2mContinue >>>\033[0;34m {self.continue_cmd}\033[0m"
elif not self.interrupted and self.repeat_cmd:
cmd_line = f"\n\033[2mRepeat >>>\033[0;34m {self.repeat_cmd}\033[0m"
msg = (
f"\033[36m[RandQuik]\033[32m {self.action} \033[1m{size_str}\033[0;32m in "
f"\033[1m{time_str}\033[0;32m @ \033[1;32m{speed_gbs:.2f} GB/s{stats_fmt}\033[0m"
f"{status_fmt}\033[1m{cmd_line}\n"
)
if not sys.stderr.isatty():
msg = re.sub(r"\033\[[0-9;]*m", "", msg)
sys.stderr.write(msg)
def print_detailed_stats(self):
"""Print detailed worker statistics table (for -vv)."""
if self.worker_stats and self.consumer_stats:
report = format_worker_stats_report(self.worker_stats, self.consumer_stats)
sys.stderr.write(report + "\n")
def format_worker_stats_report(
worker_stats: list[WorkerStats], consumer_stats: ConsumerStats
) -> str:
"""Format a complete stats report for all workers as a table."""
if not worker_stats:
return "No worker stats available"
def ms(val: float) -> str:
return f"{val * 1000:.0f} ms"
# Column width for worker data
col_w = 8
pct_w = 6 # Width for percentage column
# Total time across all workers for percentage calculation
total_all = sum(s.total_time() for s in worker_stats)
def pct(val: float) -> str:
return f"{100 * val / total_all:.0f}%" if total_all > 0 else "--"
# Header row with worker numbers
header = (
"Worker stats"
+ f"{'%':>{pct_w}}"
+ "".join(f"{'W' + str(s.worker_id):>{col_w}}" for s in worker_stats)
)
sep = "-" * len(header)
# Non-timing rows (counters)
total_blocks = sum(s.blocks_processed for s in worker_stats)
total_cycles = sum(s.wait_cycles for s in worker_stats)
def cycles_pct() -> str:
return f"{100 * total_cycles / total_blocks:.0f}%" if total_blocks > 0 else "--"
counter_rows = [
("1MiB blocks", "", [str(s.blocks_processed) for s in worker_stats]),
("wait_cycles", cycles_pct(), [str(s.wait_cycles) for s in worker_stats]),
]
# Timing rows with summed values for percentage
timing_rows = [
(
"crypto",
sum(s.crypto_time for s in worker_stats),
[ms(s.crypto_time) for s in worker_stats],
),
(
"lock_acq",
sum(s.lock_acquire_time for s in worker_stats),
[ms(s.lock_acquire_time) for s in worker_stats],
),
(
"wait_sp",
sum(s.lock_wait_space_time for s in worker_stats),
[ms(s.lock_wait_space_time) for s in worker_stats],
),
(
"claim",
sum(s.lock_claim_time for s in worker_stats),
[ms(s.lock_claim_time) for s in worker_stats],
),
(
"notify",
sum(s.lock_notify_time for s in worker_stats),
[ms(s.lock_notify_time) for s in worker_stats],
),
]
timing_rows.append(("total", total_all, [ms(s.total_time()) for s in worker_stats]))
lines = [header, sep]
for label, pct_val, values in counter_rows:
row = f"{label:<12}" + f"{pct_val:>{pct_w}}" + "".join(f"{v:>{col_w}}" for v in values)
lines.append(row)
lines.append(sep)
for label, total_val, values in timing_rows:
row = (
f"{label:<12}" + f"{pct(total_val):>{pct_w}}" + "".join(f"{v:>{col_w}}" for v in values)
)
lines.append(row)
# Consumer stats
lines.append(sep)
return "\n".join(lines)
-182
View File
@@ -1,182 +0,0 @@
"""Utility functions for formatting and parsing."""
import pathlib
import re
import sys
__all__ = [
"get_output_size",
"get_sector_size",
"parse_size",
"sparse_range",
]
# Cache for sector size lookup (path -> size)
_sector_size_cache: dict[str, int] = {}
def get_sector_size(path: str | pathlib.Path) -> int:
"""Get sector size for a block device, or 512 as fallback."""
import os
import stat
try:
st = pathlib.Path(path).stat()
if not stat.S_ISBLK(st.st_mode):
return 512
except OSError:
return 512
try:
fd = os.open(str(path), os.O_RDONLY)
try:
import fcntl
import struct
if sys.platform == "darwin":
# macOS: DKIOCGETBLOCKSIZE = 0x40046418
DKIOCGETBLOCKSIZE = 0x40046418
buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4)
return struct.unpack("I", buf)[0]
else:
# Linux: BLKSSZGET = 0x1268
BLKSSZGET = 0x1268
buf = fcntl.ioctl(fd, BLKSSZGET, b"\x00" * 4)
return struct.unpack("i", buf)[0]
finally:
os.close(fd)
except (OSError, ImportError, Exception):
return 512
def get_output_size(path: str | pathlib.Path | None) -> int | None:
"""Get the size of output file or block device.
Returns None for stdout or non-existent files.
Returns the size in bytes for existing files or block devices.
"""
import os
import stat
if not path:
return None # stdout
try:
st = pathlib.Path(path).stat()
except OSError:
return None # doesn't exist yet
if stat.S_ISBLK(st.st_mode):
# Block device - get size via ioctl
try:
fd = os.open(str(path), os.O_RDONLY)
try:
import fcntl
import struct
if sys.platform == "darwin":
# macOS: DKIOCGETBLOCKCOUNT and DKIOCGETBLOCKSIZE
DKIOCGETBLOCKCOUNT = 0x40086419
DKIOCGETBLOCKSIZE = 0x40046418
count_buf = fcntl.ioctl(fd, DKIOCGETBLOCKCOUNT, b"\x00" * 8)
size_buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4)
block_count = struct.unpack("Q", count_buf)[0]
block_size = struct.unpack("I", size_buf)[0]
return block_count * block_size
else:
# Linux: BLKGETSIZE64 = 0x80081272
BLKGETSIZE64 = 0x80081272
buf = fcntl.ioctl(fd, BLKGETSIZE64, b"\x00" * 8)
return struct.unpack("Q", buf)[0]
finally:
os.close(fd)
except (OSError, ImportError, Exception):
return None
elif stat.S_ISREG(st.st_mode):
return st.st_size
else:
return None
def parse_size(length: str | None, output_path: str | pathlib.Path | None = None) -> int | None:
"""Parse size string with SI/IEC prefixes.
Supports:
- Plain numbers: 1000, 1_000_000
- SI prefixes: k, m, g, t, p (powers of 1000)
- IEC prefixes: ki, mi, gi, ti, pi (powers of 1024)
- Optional 'b' suffix: kb, kib, mb, mib, etc.
- Special unit 'sect' = device sector size (detected, fallback 512)
- Case insensitive
Examples: 1k, 1ki, 1kb, 1kib, 100m, 100mi, 1g, 1gi, 10sect
"""
if length is None:
return None
s = length.strip().lower().replace("_", "")
# Handle sect unit
m = re.match(r"^(\d+)\s*sect?s?$", s)
if m:
if output_path:
path_str = str(output_path)
if path_str not in _sector_size_cache:
_sector_size_cache[path_str] = get_sector_size(output_path)
sector_size = _sector_size_cache.get(path_str, 512)
else:
sector_size = 512
return int(m.group(1)) * sector_size
# SI/IEC prefixes
si_prefixes = {"k": 1000, "m": 1000**2, "g": 1000**3, "t": 1000**4, "p": 1000**5}
iec_prefixes = {
"ki": 1024,
"mi": 1024**2,
"gi": 1024**3,
"ti": 1024**4,
"pi": 1024**5,
}
# Try IEC first (ki, mi, etc.) - must check before SI
m = re.match(r"^(\d+(?:\.\d+)?)\s*(ki|mi|gi|ti|pi)b?$", s)
if m:
num, prefix = m.groups()
return int(float(num) * iec_prefixes[prefix])
# Try SI (k, m, g, etc.)
m = re.match(r"^(\d+(?:\.\d+)?)\s*([kmgtp])b?$", s)
if m:
num, prefix = m.groups()
return int(float(num) * si_prefixes[prefix])
# Plain number
m = re.match(r"^(\d+)$", s)
if m:
return int(m.group(1))
raise ValueError(f"Invalid size format: {length}")
def sparse_range(n: int, max_items: int = 9) -> list[int]:
"""Generate a sparse range from 1 to N for benchmarking thread counts."""
if n < 1:
return [1]
if n <= max_items - 1:
return list(range(n + 1))
keep = 3 # dense prefix: 1,2,3
out = list(range(keep + 1))
remaining = max_items - keep
step = max(1, n // (remaining - 1))
for k in range(1, remaining):
v = k * step
if v > out[-1]:
out.append(v)
if out[-1] != n:
out[-1] = n
return out
-332
View File
@@ -1,332 +0,0 @@
"""Worker threads and ring buffer management for parallel generation."""
import os
import sys
import threading
import time
from randquik.io import open_fd
from randquik.progress import ProgressDisplay
from randquik.stats import (
ConsumerStats,
RunResult,
SingleThreadedStats,
WorkerStats,
stopwatch,
)
__all__ = [
"BLOCK_SIZE",
"RunResult",
"run",
]
BLOCK_SIZE = 1 << 20
class _FdProducer:
"""Multi-threaded producer with ring buffer for sequential file output.
Uses efficient synchronization:
- Single lock with two conditions (has_data, has_space)
- Workers wait on has_space, notify has_data when block is ready
- Consumer waits on has_data, notifies has_space when block is consumed
- Crypto runs outside the lock
"""
def __init__(
self,
workers: int,
key: bytes,
ciph,
total_bytes: int | None,
fd: int,
dry: bool = False,
iseek: int = 0,
block_size: int = BLOCK_SIZE,
):
self.workers = workers
self.key = key
self.ciph = ciph
self.total_bytes = total_bytes
self.fd = fd
self.dry = dry
self.iseek = iseek
self.block_size = block_size
# iseek handling: which block to start at, and offset within first block
self.start_block = iseek // block_size
self.start_offset = iseek % block_size
self.num_slots = workers + 2 # Tested optimal (+1 for I/O and +1 to avoid congestion)
self._buf = bytearray(self.num_slots * block_size)
# Separate conditions for producers and consumer
self._lock = threading.Lock()
self.has_data = threading.Condition(self._lock) # Consumer waits, workers notify
self.has_space = threading.Condition(self._lock) # Workers wait, consumer notifies
self.lock_blkno = threading.Lock()
self.blkno = self.start_block # next block to generate
self.ready = [False] * self.num_slots # which block number is there ready
self.quit = False
self.threads: list[threading.Thread] = []
self.written = 0
self.consumer_stats = ConsumerStats()
# Per-worker stats, collected after threads finish
self._worker_stats: list[WorkerStats] = []
self._stats_lock = threading.Lock()
def start(self):
self.threads = [
threading.Thread(target=self.worker, args=(i,)) for i in range(self.workers)
]
for t in self.threads:
t.start()
def worker(self, worker_id: int):
assert self.num_slots >= self.workers, "Ring buffer quarantee broken"
view = memoryview(self._buf)
slots = [
view[i * self.block_size : (i + 1) * self.block_size] for i in range(self.num_slots)
]
# Profiling setup
stats = WorkerStats(worker_id=worker_id)
timer = stopwatch()
try:
while True:
# Claim next block number
with self.lock_blkno:
blkno = self.blkno
self.blkno += 1
stats.lock_claim_time += next(timer)
with self._lock:
stats.lock_acquire_time += next(timer)
# Wait for the NEXT slot to be free
slot = blkno % self.num_slots
while self.ready[slot] and not self.quit:
stats.wait_cycles += 1
self.has_space.wait()
stats.lock_wait_space_time += next(timer)
if self.quit:
return
# Generate block
self.ciph.stream(
self.key,
blkno.to_bytes(self.ciph.NONCEBYTES, "little"),
into=slots[slot],
)
stats.crypto_time += next(timer)
stats.blocks_processed += 1
stats.bytes_generated += self.block_size
# Commit block (mark ready + notify consumer)
with self._lock:
self.ready[slot] = True
self.has_data.notify()
finally:
view.release()
with self._stats_lock:
self._worker_stats.append(stats)
def consumer(self, progress_state: dict | None = None):
"""Consume blocks and write to fd. Call start() first."""
view = memoryview(self._buf)
try:
slots = [
view[i * self.block_size : (i + 1) * self.block_size] for i in range(self.num_slots)
]
blkno = self.start_block
slot = blkno % self.num_slots
total = sys.maxsize if self.total_bytes is None else self.total_bytes
# Handle the first block: skip start_offset bytes (note: this is purposefully left out of stats)
with self.has_data:
while not self.ready[slot] and not self.quit:
self.has_data.wait()
if self.quit:
return
buf = slots[slot][self.start_offset : self.start_offset + total]
if not self.dry:
os.write(self.fd, buf)
self.written += len(buf)
# Other blocks
timer = stopwatch()
while self.written < total:
# Take/wait for expected slot
with self._lock:
# Release previous slot and notify workers
self.ready[slot] = False
self.has_space.notify_all()
# Wait for the next buffer to be ready
blkno += 1
slot = blkno % self.num_slots
while not self.ready[slot] and not self.quit:
self.has_data.wait()
if self.quit:
return
self.consumer_stats.wait_time += next(timer)
buf = slots[slot]
# Last block? Trim to remaining size
if self.written + len(buf) > total:
buf = buf[: total - self.written]
if not self.dry:
os.write(self.fd, buf)
self.consumer_stats.write_time += next(timer)
self.written += len(buf)
if progress_state is not None:
progress_state["written"] = self.written
finally:
self.stop()
view.release()
def stop(self):
"""Signal workers to stop and wait for them."""
with self._lock:
self.quit = True
self.has_data.notify_all()
self.has_space.notify_all()
for t in self.threads:
t.join()
def get_worker_stats(self) -> list[WorkerStats]:
"""Get stats for each worker. Call after stop()."""
with self._stats_lock:
# Sort by worker_id for consistent output
return sorted(self._worker_stats, key=lambda s: s.worker_id)
def run(self, progress_state: dict | None = None):
"""Run multi-threaded generation."""
self.start()
try:
self.consumer(progress_state)
finally:
self.stop()
class _SingleThreadedProducer:
"""Single-threaded producer for infinite output or workers=0 mode."""
def __init__(
self,
key: bytes,
ciph,
total_bytes: int | None,
fd: int,
dry: bool = False,
block_size: int = BLOCK_SIZE,
):
self.key = key
self.ciph = ciph
self.total_bytes = total_bytes
self.fd = fd
self.dry = dry
self.block_size = block_size
self.written = 0
self.stats = SingleThreadedStats()
def run(self, progress_state: dict | None = None):
"""Generate and write blocks sequentially."""
buf = bytearray(self.block_size)
view = memoryview(buf)
nonce = bytearray(self.ciph.NONCEBYTES)
total = sys.maxsize if self.total_bytes is None else self.total_bytes
timer = stopwatch()
try:
while self.written < total:
size = min(self.block_size, total - self.written)
chunk = view[:size]
self.ciph.stream(self.key, nonce, size, into=chunk)
self.stats.crypto_time += next(timer)
if not self.dry:
os.write(self.fd, chunk)
self.stats.write_time += next(timer)
self.ciph.nonce_increment(nonce)
self.written += size
if progress_state is not None:
progress_state["written"] = self.written
finally:
view.release()
def run(
output: str | None,
total_bytes: int | None,
iseek: int,
oseek: int,
key: bytes,
ciph,
workers: int = 1,
dry: bool = False,
quiet: bool = False,
action: str = "wrote",
continue_cmd: str | None = None,
) -> RunResult:
"""Run random generation with specified number of workers. Returns RunResult.
Args:
workers: Number of worker threads. 0 for single-threaded mode.
"""
start_time = time.perf_counter()
infinite = total_bytes is None
fd_size = 0 if infinite else total_bytes
with open_fd(output, fd_size, dry=dry, oseek=oseek) as fd:
if workers == 0:
producer = _SingleThreadedProducer(key, ciph, total_bytes, fd, dry=dry)
else:
producer = _FdProducer(workers, key, ciph, total_bytes, fd, dry=dry, iseek=iseek)
progress_state = {"written": 0}
progress = ProgressDisplay(
total_bytes,
start_time,
progress_state,
infinite=infinite,
output_name=output,
oseek=oseek,
)
if not quiet:
progress.start()
interrupted = False
try:
producer.run(progress_state)
except (KeyboardInterrupt, BrokenPipeError):
interrupted = True
finally:
progress.stop()
elapsed = time.perf_counter() - start_time
# Build result with raw stats
if workers == 0:
return RunResult(
written=producer.written,
elapsed=elapsed,
interrupted=interrupted,
action=action,
singlethreaded_stats=producer.stats,
continue_cmd=continue_cmd,
)
else:
return RunResult(
written=producer.written,
elapsed=elapsed,
interrupted=interrupted,
action=action,
consumer_stats=producer.consumer_stats,
worker_stats=producer.get_worker_stats(),
continue_cmd=continue_cmd,
)
+17
View File
@@ -0,0 +1,17 @@
import os
from distutils.core import setup
import numpy
from Cython.Build import cythonize
from setuptools import Extension
os.environ["CFLAGS"] = "-O3 -march=native -Wall -Wextra"
extensions = [
Extension(
"nprand",
["src/nprand.pyx"],
include_dirs=[numpy.get_include()],
define_macros=[("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")],
)
]
setup(ext_modules=cythonize(extensions))
+29
View File
@@ -0,0 +1,29 @@
#include <stdint.h>
#include <string.h>
// clang-format off
#define QUARTERSTEP(a, b, c, n) a += b; c ^= a; c = (c << n) | (c >> (32 - n))
#define QUARTERROUND(a, b, c, d) {\
QUARTERSTEP(a, b, d, 16); QUARTERSTEP(c, d, b, 12); \
QUARTERSTEP(a, b, d, 8); QUARTERSTEP(c, d, b, 7); }
static inline uint64_t _cha_block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
unsigned blocks = bufsize / 64;
uint32_t* out = (uint32_t*)buf;
uint32_t x[16];
for (unsigned b = blocks; b-->0;) {
for (unsigned i = 0; i < 16; ++i) x[i] = state[i]; // Faster than memcpy
for (unsigned i = rounds / 2; i-->0;) {
// Mix columns, then diagonals
for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + j], x[8 + j], x[12 + j]);
for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + (j+1)%4], x[8 + (j+2)%4], x[12 + (j+3)%4]);
}
for (unsigned i = 0; i < 16; ++i) *out++ = x[i] + state[i];
++*(uint64_t*)(state + 12); // Increment counter
}
memset(x, 0, sizeof x);
return blocks * CHA_BLOCK_SIZE;
}
#undef QUARTERROUND
#undef QUARTERSTEP
+106
View File
@@ -0,0 +1,106 @@
#include <arm_neon.h>
#include <stdint.h>
#include <stdlib.h>
// clang-format off
#define VEC4_ROT(A, IMM) \
vreinterpretq_u32_u8(vorrq_u8(vshlq_n_u32(A, IMM), vshrq_n_u32(A, 32 - IMM)))
/* same, but replace 2 of the shift/shift/or "rotation" by byte shuffles (8 &
* 16) (better) */
#define VEC4_QUARTERROUND(A, B, C, D) \
x[A] = vaddq_u32(x[A], x[B]); \
x[D] = vqtbl1q_u8(veorq_u32(x[D], x[A]), rot16); \
x[C] = vaddq_u32(x[C], x[D]); \
x[B] = VEC4_ROT(veorq_u32(x[B], x[C]), 12); \
x[A] = vaddq_u32(x[A], x[B]); \
x[D] = vqtbl1q_u8(veorq_u32(x[D], x[A]), rot8); \
x[C] = vaddq_u32(x[C], x[D]); \
x[B] = VEC4_ROT(veorq_u32(x[B], x[C]), 7)
#define ONEQUAD(A, B, C, D, OUT) \
{ \
/* Add original block */ \
x[A] = vaddq_u32(x[A], orig[A]); \
x[B] = vaddq_u32(x[B], orig[B]); \
x[C] = vaddq_u32(x[C], orig[C]); \
x[D] = vaddq_u32(x[D], orig[D]); \
/* Transpose */ \
uint32x4x2_t ab = vtrnq_u32(x[A], x[B]); \
uint32x4x2_t cd = vtrnq_u32(x[C], x[D]); \
x[A] = vcombine_u32(vget_low_u32(ab.val[0]), vget_low_u32(cd.val[0])); \
x[B] = vcombine_u32(vget_low_u32(ab.val[1]), vget_low_u32(cd.val[1])); \
x[C] = vcombine_u32(vget_high_u32(ab.val[0]), vget_high_u32(cd.val[0])); \
x[D] = vcombine_u32(vget_high_u32(ab.val[1]), vget_high_u32(cd.val[1])); \
/* Write out 1/4 of each block */ \
vst1q_u32((uint32_t*)(OUT), x[A]); \
vst1q_u32((uint32_t*)(OUT + 64), x[B]); \
vst1q_u32((uint32_t*)(OUT + 128), x[C]); \
vst1q_u32((uint32_t*)(OUT + 192), x[D]); \
}
#define COUNTER_INCREMENT(addv) \
{ \
orig[12] = vaddq_u32(orig[12], addv); \
orig[13] = vaddq_u32(orig[13], vshrq_n_u32(vcltq_u32(orig[12], addv), 31)); \
}
static inline uint64_t
_cha_4block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
/* constant for shuffling bytes (replacing multiple-of-8 rotates) */
const uint8x16_t rot16 = {
2, 3, 0, 1,
6, 7, 4, 5,
10, 11, 8, 9,
14, 15, 12, 13
};
const uint8x16_t rot8= {
3, 0, 1, 2,
7, 4, 5, 6,
11, 8, 9, 10,
15, 12, 13, 14
};
// Load state to vectors, duplicate four times, only different counters
uint32x4_t orig[16];
for (unsigned i = 0; i < 16; ++i) orig[i] = vdupq_n_u32(state[i]);
uint32x4_t addv = { 0, 1, 2, 3 };
COUNTER_INCREMENT(addv);
addv = vdupq_n_u32(4);
const unsigned batches = bufsize / 256;
for (unsigned b = batches; b-->0;) {
uint32x4_t x[16];
for (unsigned i = 0; i < 16; ++i) x[i] = orig[i];
for (unsigned r = rounds / 2; r-->0;) {
// Mix columns
VEC4_QUARTERROUND(0, 4, 8, 12);
VEC4_QUARTERROUND(1, 5, 9, 13);
VEC4_QUARTERROUND(2, 6, 10, 14);
VEC4_QUARTERROUND(3, 7, 11, 15);
// Mix diagonals
VEC4_QUARTERROUND(0, 5, 10, 15);
VEC4_QUARTERROUND(1, 6, 11, 12);
VEC4_QUARTERROUND(2, 7, 8, 13);
VEC4_QUARTERROUND(3, 4, 9, 14);
}
// Add original block, unpack output
ONEQUAD(0, 1, 2, 3, buf);
ONEQUAD(4, 5, 6, 7, buf + 16);
ONEQUAD(8, 9, 10, 11, buf + 32);
ONEQUAD(12, 13, 14, 15, buf + 48);
COUNTER_INCREMENT(addv);
buf += 256;
}
// Store counter
state[12] = vgetq_lane_u32(orig[12], 0);
state[13] = vgetq_lane_u32(orig[13], 0);
return batches * 256;
}
#undef COUNTER_INCREMENT
#undef ONEQUAD
#undef ONEQUAD_TRANSPOSE
#undef VEC4_ROT
#undef VEC4_QUARTERROUND
#undef VEC4_QUARTERROUND_SHUFFLE
+104
View File
@@ -0,0 +1,104 @@
#if defined(__x86_64__)
#include <emmintrin.h> // SSE2
#include <tmmintrin.h> // SSSE3
#elif defined(__aarch64__)
#include "sse2neon.h"
#endif
#include <stdio.h>
// clang-format off
#define VEC4_ROT(A, IMM) \
_mm_or_si128(_mm_slli_epi32(A, IMM), _mm_srli_epi32(A, (32 - IMM)))
/* same, but replace 2 of the shift/shift/or "rotation" by byte shuffles (8 &
* 16) (better) */
#define VEC4_QUARTERROUND(A, B, C, D) \
x[A] = _mm_add_epi32(x[A], x[B]); \
x[D] = _mm_shuffle_epi8(_mm_xor_si128(x[D], x[A]), rot16); \
x[C] = _mm_add_epi32(x[C], x[D]); \
x[B] = VEC4_ROT(_mm_xor_si128(x[B], x[C]), 12); \
x[A] = _mm_add_epi32(x[A], x[B]); \
x[D] = _mm_shuffle_epi8(_mm_xor_si128(x[D], x[A]), rot8); \
x[C] = _mm_add_epi32(x[C], x[D]); \
x[B] = VEC4_ROT(_mm_xor_si128(x[B], x[C]), 7)
#define ONEQUAD(A, B, C, D, OUT) \
{ \
/* Add original block */ \
x[A] = _mm_add_epi32(x[A], orig[A]); \
x[B] = _mm_add_epi32(x[B], orig[B]); \
x[C] = _mm_add_epi32(x[C], orig[C]); \
x[D] = _mm_add_epi32(x[D], orig[D]); \
/* Transpose */ \
__m128i abl = _mm_unpacklo_epi32(x[A], x[B]); \
__m128i cdl = _mm_unpacklo_epi32(x[C], x[D]); \
__m128i abh = _mm_unpackhi_epi32(x[A], x[B]); \
__m128i cdh = _mm_unpackhi_epi32(x[C], x[D]); \
x[A] = _mm_unpacklo_epi64(abl, cdl); /* a0 b0 c0 d0 */ \
x[B] = _mm_unpackhi_epi64(abl, cdl); /* a1 b1 c1 d1 */ \
x[C] = _mm_unpacklo_epi64(abh, cdh); /* a2 b2 c2 d2 */ \
x[D] = _mm_unpackhi_epi64(abh, cdh); /* a3 b3 c3 d3 */ \
/* Write out 1/4 of each block */ \
_mm_storeu_si128((__m128i*)(OUT), x[A]); \
_mm_storeu_si128((__m128i*)(OUT + 64), x[B]); \
_mm_storeu_si128((__m128i*)(OUT + 128), x[C]); \
_mm_storeu_si128((__m128i*)(OUT + 192), x[D]); \
}
#define COUNTER_INCREMENT(addv) \
{ \
__m128i carry = orig[12]; \
orig[12] = _mm_add_epi32(orig[12], addv); \
carry = _mm_srli_epi32(_mm_and_si128(_mm_xor_si128(orig[12], carry), carry), 31); \
orig[13] = _mm_add_epi32(orig[13], carry); \
}
static inline uint64_t
_cha_4block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
/* constant for shuffling bytes (replacing multiple-of-8 rotates) */
const __m128i rot16 =
_mm_set_epi8(13, 12, 15, 14, 9, 8, 11, 10, 5, 4, 7, 6, 1, 0, 3, 2);
const __m128i rot8 =
_mm_set_epi8(14, 13, 12, 15, 10, 9, 8, 11, 6, 5, 4, 7, 2, 1, 0, 3);
// Load state to vectors, duplicate four times, only different counters
__m128i orig[16];
for (unsigned i = 0; i < 16; ++i) orig[i] = _mm_set1_epi32(state[i]);
__m128i addv = _mm_set_epi32(3, 2, 1, 0);
COUNTER_INCREMENT(addv);
addv = _mm_set1_epi32(4);
const unsigned batches = bufsize / 256;
for (unsigned b = batches; b-->0;) {
__m128i x[16];
for (unsigned i = 0; i < 16; ++i) x[i] = orig[i];
for (unsigned r = rounds / 2; r-->0;) {
// Mix columns
VEC4_QUARTERROUND(0, 4, 8, 12);
VEC4_QUARTERROUND(1, 5, 9, 13);
VEC4_QUARTERROUND(2, 6, 10, 14);
VEC4_QUARTERROUND(3, 7, 11, 15);
// Mix diagonals
VEC4_QUARTERROUND(0, 5, 10, 15);
VEC4_QUARTERROUND(1, 6, 11, 12);
VEC4_QUARTERROUND(2, 7, 8, 13);
VEC4_QUARTERROUND(3, 4, 9, 14);
}
// Add original block, unpack output
ONEQUAD(0, 1, 2, 3, buf);
ONEQUAD(4, 5, 6, 7, buf + 16);
ONEQUAD(8, 9, 10, 11, buf + 32);
ONEQUAD(12, 13, 14, 15, buf + 48);
COUNTER_INCREMENT(addv);
buf += 256;
}
// Store counter
state[12] = _mm_cvtsi128_si32(orig[12]);
state[13] = _mm_cvtsi128_si32(orig[13]);
return batches * 256;
}
#undef COUNTER_INCREMENT
#undef ONEQUAD
#undef ONEQUAD_TRANSPOSE
#undef VEC4_ROT
#undef VEC4_QUARTERROUND
#undef VEC4_QUARTERROUND_SHUFFLE
+130
View File
@@ -0,0 +1,130 @@
#include <immintrin.h> // AVX2
// clang-format off
#define VEC8_ROT(A, IMM) \
_mm256_or_si256(_mm256_slli_epi32(A, IMM), _mm256_srli_epi32(A, (32 - IMM)))
#define VEC8_LINE1(A, B, C, D) \
x[A] = _mm256_add_epi32(x[A], x[B]); \
x[D] = _mm256_shuffle_epi8(_mm256_xor_si256(x[D], x[A]), rot16)
#define VEC8_LINE2(A, B, C, D) \
x[C] = _mm256_add_epi32(x[C], x[D]); \
x[B] = VEC8_ROT(_mm256_xor_si256(x[B], x[C]), 12)
#define VEC8_LINE3(A, B, C, D) \
x[A] = _mm256_add_epi32(x[A], x[B]); \
x[D] = _mm256_shuffle_epi8(_mm256_xor_si256(x[D], x[A]), rot8)
#define VEC8_LINE4(A, B, C, D) \
x[C] = _mm256_add_epi32(x[C], x[D]); \
x[B] = VEC8_ROT(_mm256_xor_si256(x[B], x[C]), 7)
#define VEC8_ROUND( \
A1, B1, C1, D1, A2, B2, C2, D2, A3, B3, C3, D3, A4, B4, C4, D4 \
) \
VEC8_LINE1(A1, B1, C1, D1); \
VEC8_LINE1(A2, B2, C2, D2); \
VEC8_LINE1(A3, B3, C3, D3); \
VEC8_LINE1(A4, B4, C4, D4); \
VEC8_LINE2(A1, B1, C1, D1); \
VEC8_LINE2(A2, B2, C2, D2); \
VEC8_LINE2(A3, B3, C3, D3); \
VEC8_LINE2(A4, B4, C4, D4); \
VEC8_LINE3(A1, B1, C1, D1); \
VEC8_LINE3(A2, B2, C2, D2); \
VEC8_LINE3(A3, B3, C3, D3); \
VEC8_LINE3(A4, B4, C4, D4); \
VEC8_LINE4(A1, B1, C1, D1); \
VEC8_LINE4(A2, B2, C2, D2); \
VEC8_LINE4(A3, B3, C3, D3); \
VEC8_LINE4(A4, B4, C4, D4)
#define TRANSPOSE(A, B, C, D) \
{ \
const __m256i t0 = _mm256_unpacklo_epi32(x[A], x[B]), \
t1 = _mm256_unpacklo_epi32(x[C], x[D]), \
t2 = _mm256_unpackhi_epi32(x[A], x[B]), \
t3 = _mm256_unpackhi_epi32(x[C], x[D]); \
x[A] = _mm256_unpacklo_epi64(t0, t1); \
x[B] = _mm256_unpackhi_epi64(t0, t1); \
x[C] = _mm256_unpacklo_epi64(t2, t3); \
x[D] = _mm256_unpackhi_epi64(t2, t3); \
}
#define ONEOCTO(A, B, C, D, A2, B2, C2, D2, c) \
{ \
TRANSPOSE(A, B, C, D); \
TRANSPOSE(A2, B2, C2, D2); \
_mm256_storeu_si256((__m256i*)(c), _mm256_permute2x128_si256(x[A], x[A2], 0x20)); \
_mm256_storeu_si256((__m256i*)(c + 64), _mm256_permute2x128_si256(x[B], x[B2], 0x20)); \
_mm256_storeu_si256((__m256i*)(c + 128), _mm256_permute2x128_si256(x[C], x[C2], 0x20)); \
_mm256_storeu_si256((__m256i*)(c + 192), _mm256_permute2x128_si256(x[D], x[D2], 0x20)); \
_mm256_storeu_si256((__m256i*)(c + 256), _mm256_permute2x128_si256(x[A], x[A2], 0x31)); \
_mm256_storeu_si256((__m256i*)(c + 320), _mm256_permute2x128_si256(x[B], x[B2], 0x31)); \
_mm256_storeu_si256((__m256i*)(c + 384), _mm256_permute2x128_si256(x[C], x[C2], 0x31)); \
_mm256_storeu_si256((__m256i*)(c + 448), _mm256_permute2x128_si256(x[D], x[D2], 0x31)); \
}
#define COUNTER_INCREMENT(addv) \
{ \
__m256i carry = orig[12]; \
orig[12] = _mm256_add_epi32(orig[12], addv); \
carry = _mm256_srli_epi32(_mm256_and_si256(_mm256_xor_si256(orig[12], carry), carry), 31); \
orig[13] = _mm256_add_epi32(orig[13], carry); \
}
static inline uint64_t
_cha_8block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
unsigned batches = bufsize / 512;
/* constant for shuffling bytes (replacing multiple-of-8 rotates) */
const __m256i rot16 = _mm256_set_epi8(
13, 12, 15, 14,
9, 8, 11, 10,
5, 4, 7, 6,
1, 0, 3, 2,
13, 12, 15, 14,
9, 8, 11, 10,
5, 4, 7, 6,
1, 0, 3, 2
);
const __m256i rot8 = _mm256_set_epi8(
14, 13, 12, 15,
10, 9, 8, 11,
6, 5, 4, 7,
2, 1, 0, 3,
14, 13, 12, 15,
10, 9, 8, 11,
6, 5, 4, 7,
2, 1, 0, 3
);
__m256i orig[16];
for (int i = 0; i < 16; ++i)
orig[i] = _mm256_set1_epi32(state[i]);
COUNTER_INCREMENT(_mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0));
for (unsigned b = batches; b-->0;) {
__m256i x[16];
for (int i = 0; i < 16; ++i) x[i] = orig[i];
for (unsigned r = rounds / 2; r-->0;) {
VEC8_ROUND(0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15);
VEC8_ROUND(0, 5, 10, 15, 1, 6, 11, 12, 2, 7, 8, 13, 3, 4, 9, 14);
}
for (unsigned i = 0; i < 16; ++i) x[i] = _mm256_add_epi32(x[i], orig[i]);
ONEOCTO(0, 1, 2, 3, 4, 5, 6, 7, buf);
ONEOCTO(8, 9, 10, 11, 12, 13, 14, 15, buf + 32);
COUNTER_INCREMENT(_mm256_set1_epi32(8));
buf += 512;
}
state[12] = _mm256_extract_epi32(orig[12], 0);
state[13] = _mm256_extract_epi32(orig[13], 0);
return batches * 512;
}
#undef COUNTER_INCREMENT
#undef ONEOCTO
#undef TRANSPOSE
#undef VEC8_ROT
#undef VEC8_LINE1
#undef VEC8_LINE2
#undef VEC8_LINE3
#undef VEC8_LINE4
#undef VEC8_ROUND
+18
View File
@@ -0,0 +1,18 @@
#include "charandom.h"
static uint64_t cha_uint64(void* st) {
cha_ctx* ctx = (cha_ctx*)st;
if (ctx->offset + sizeof(uint64_t) > ctx->end) {
ctx->offset -= ctx->end;
ctx->end =
ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
}
uint64_t ret = *(uint64_t*)(ctx->unconsumed + ctx->offset);
ctx->offset += sizeof(uint64_t);
return ret;
}
static uint32_t cha_uint32(void* st) { return cha_uint64(st); }
static double cha_double(void* st) {
// Fast uint64_to_double conversion from numpy/random/_common.pxd
return (cha_uint64(st) >> 11) * (1.0 / 9007199254740992.0);
}
+3
View File
@@ -0,0 +1,3 @@
#include "charandom.h"
// Library build for Python CFFI to use
+152
View File
@@ -0,0 +1,152 @@
#include <stdbool.h>
#include <stdint.h>
#define CHA_BLOCK_SIZE 64
#define BATCH_BLOCKS 8
#define BATCH_SIZE (BATCH_BLOCKS * CHA_BLOCK_SIZE)
#if defined(__x86_64__)
#ifdef __GNUC__
#pragma GCC target("sse2")
#pragma GCC target("ssse3")
#pragma GCC target("avx2")
#endif
#include "cha4ssse3.h"
#include "cha8avx2.h"
#elif defined(__aarch64__)
#include "cha4neon.h"
#endif
#include "cha1c.h"
#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <pthread.h>
#include <time.h>
#include <unistd.h>
typedef uint64_t (*genfunc)(
uint8_t* out, size_t outsize, uint32_t state[16], unsigned rounds
);
typedef struct cha_ctx {
uint32_t state[16];
uint8_t unconsumed[BATCH_SIZE];
uint32_t offset, end;
unsigned rounds;
genfunc gen;
} cha_ctx;
/// @brief Initialize cha_ctx
/// @param ctx holds ChaCha20 state
/// @param key 32 byte key
/// @param iv 16 bytes, usually the first 4-8 bytes are zeroes, the rest nonce
/// @param rounds ChaCha iteration count: 8=fast, 12=balanced, 20=secure
void cha_init(
cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds
) {
ctx->state[0] = 0x61707865;
ctx->state[1] = 0x3320646e;
ctx->state[2] = 0x79622d32;
ctx->state[3] = 0x6b206574;
memcpy(ctx->state + 4, key, 32);
memcpy(ctx->state + 12, iv, 16);
memset(ctx->unconsumed, 0, sizeof ctx->unconsumed);
ctx->offset = ctx->end = 0;
ctx->rounds = rounds;
#if defined(__x86_64__)
if (__builtin_cpu_supports("avx2"))
ctx->gen = _cha_8block;
else if (__builtin_cpu_supports("ssse3"))
ctx->gen = _cha_4block;
else
ctx->gen = _cha_block;
#elif defined(__aarch64__)
ctx->gen = _cha_4block;
#else
ctx->gen = _cha_block;
#endif
}
/// Dispose of sensitive data within the context
void cha_wipe(cha_ctx* ctx) { memset(ctx, 0, sizeof(cha_ctx)); }
/// @brief Advance or rewind the stream to any arbitrary location
/// Keeps byte offset inside batch untouched but clears the buffer
/// @param ctx ChaCha context
/// @param offset Offset in blocks of 64 bytes (counter change)
void cha_seek_blocks(cha_ctx* ctx, int64_t offset) {
*(uint64_t*)(ctx->state + 12) += offset - (int64_t)ctx->end / 64;
ctx->end = 0;
}
/// @brief Seek a number of bytes forward or backward in stream
/// Supports seeking backwards and forwards, even beyond start, but only up to
/// 64 bit distance which does not cover the whole counter range.
/// @param ctx ChaCha context
/// @param offset Positive or negative offset from current byte position
void cha_seek(cha_ctx* ctx, int64_t offset) {
offset += (int64_t)ctx->offset;
ctx->offset = ((offset % 64) + 64) % 64;
cha_seek_blocks(ctx, (offset - (int64_t)ctx->offset) / 64);
}
/// @brief Tell current byte position in stream (assuming initial counter 0).
/// Result is truncated to int64 range (positive or negative)
/// @param ctx ChaCha context
int64_t cha_tell(cha_ctx* ctx) {
int64_t counter = *(int64_t*)(ctx->state + 12);
return counter * CHA_BLOCK_SIZE + ctx->offset - ctx->end;
}
/// @brief Incremental generation, keeps state between calls
/// @param ctx ChaCha context
/// @param out output buffer
/// @param outlen output buffer length
void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) {
// The included header will mess with these variables
uint8_t* end = out + outlen;
if (ctx->offset) {
// Need to generate stored buffer?
if (ctx->end == 0)
ctx->end =
ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
// Deliver stored bytes first
uint64_t N = ctx->end - ctx->offset;
if (N > outlen)
N = outlen;
memcpy(out, ctx->unconsumed + ctx->offset, N);
ctx->offset += N;
out += N;
if (ctx->offset == ctx->end)
ctx->offset = ctx->end = 0;
if (out == end)
return;
}
out += ctx->gen(out, end - out, ctx->state, ctx->rounds);
const uint32_t N = end - out;
if (N) {
ctx->end =
ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
memcpy(out, ctx->unconsumed, N);
ctx->offset = N;
}
}
/// @brief Produce a requested number of random bytes, single shot.
/// @param out output buffer
/// @param outlen output buffer length
/// @param key 32 byte key
/// @param iv 16 bytes, where normally initial 4-8 bytes are 0 (counter)
void cha_generate(
uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16],
unsigned rounds
) {
cha_ctx ctx;
cha_init(&ctx, key, iv, rounds);
cha_update(&ctx, out, outlen);
cha_wipe(&ctx);
}
+295
View File
@@ -0,0 +1,295 @@
#include <errno.h>
#include <inttypes.h>
#include <pthread.h>
#include <signal.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include "charandom.h"
static volatile bool quit = false;
void signal_handler(int sig) {
quit = true;
signal(SIGINT, SIG_DFL);
signal(SIGTERM, SIG_DFL);
}
#define BLOCK_SIZE (1 << 21) // 2 MiB seems optimal for speed
static const unsigned char default_iv[16] = "\0\0\0\0\0\0\0\0RandQuik";
typedef struct thread_args {
int index;
int done;
unsigned char* buf;
unsigned char key[32];
unsigned workers;
unsigned rounds;
pthread_mutex_t lock;
pthread_cond_t cond;
pthread_t thread;
} thread_args;
void* producer_thread(void* a) {
thread_args* args = (thread_args*)a;
const uint64_t ivstep = args->workers * BATCH_BLOCKS;
cha_ctx ctx;
cha_init(&ctx, args->key, default_iv, args->rounds);
cha_seek_blocks(&ctx, args->index * BLOCK_SIZE / 64);
while (!quit) {
pthread_mutex_lock(&args->lock);
while (args->done) {
pthread_cond_wait(&args->cond, &args->lock);
}
cha_update(&ctx, args->buf, BLOCK_SIZE);
cha_seek_blocks(&ctx, ivstep);
args->done = 1;
pthread_cond_signal(&args->cond);
pthread_mutex_unlock(&args->lock);
}
cha_wipe(&ctx);
return NULL;
}
void print_status(
uint64_t bytes, uint64_t max_bytes, struct timespec start_time
) {
struct timespec end_time;
clock_gettime(CLOCK_MONOTONIC, &end_time);
double t = (end_time.tv_sec - start_time.tv_sec) +
1e-9 * (end_time.tv_nsec - start_time.tv_nsec);
char buf[64] = {};
double speed = bytes / t;
char const* unit = "MB";
double m = 1e-6;
if (speed > 0.5e9) {
unit = "GB";
m = 1e-9;
}
if (max_bytes) {
snprintf(buf, sizeof buf - 1, " of %'.0lf", m * max_bytes);
}
fprintf(
stderr, "\r%5.0lf%s %s written, %.2lf %s/s.\e[K", m * bytes, buf, unit,
m * speed, unit
);
}
int fast(
FILE* f, unsigned workers, uint64_t max_bytes, unsigned char const key[32],
unsigned char const iv[16], unsigned rounds
) {
thread_args args[workers];
memset(args, 0, sizeof args);
for (int i = 0; i < workers; ++i) {
args[i].index = i;
args[i].buf = malloc(BLOCK_SIZE);
args[i].workers = workers;
args[i].rounds = rounds;
memcpy(args[i].key, key, 32);
pthread_mutex_init(&args[i].lock, NULL);
pthread_cond_init(&args[i].cond, NULL);
pthread_create(&args[i].thread, NULL, producer_thread, &args[i]);
}
struct timespec start_time;
clock_gettime(CLOCK_MONOTONIC, &start_time);
int i = -1;
uint64_t bytes = 0;
while (!quit) {
i = (i + 1) % workers;
pthread_mutex_lock(&args[i].lock);
while (!args[i].done) {
pthread_cond_wait(&args[i].cond, &args[i].lock);
}
if (bytes % (1 << 30) == 0 || bytes + BLOCK_SIZE >= max_bytes) {
print_status(bytes, max_bytes, start_time);
}
uint64_t sz = BLOCK_SIZE;
if (max_bytes && bytes + sz >= max_bytes) {
fprintf(stderr, "\r\e[KMax reached\n");
sz = max_bytes - bytes;
quit = true;
}
if (fwrite(args[i].buf, sz, 1, f) != 1) {
quit = true;
fprintf(stderr, "\r\e[KWrite failed: %s\n", strerror(errno));
}
bytes += sz;
args[i].done = 0;
pthread_cond_signal(&args[i].cond);
pthread_mutex_unlock(&args[i].lock);
}
print_status(bytes, max_bytes, start_time);
for (int i = 0; i < workers; ++i) {
args[i].done = 0;
pthread_cancel(args[i].thread);
pthread_join(args[i].thread, NULL);
pthread_mutex_destroy(&args[i].lock);
pthread_cond_destroy(&args[i].cond);
free(args[i].buf);
}
fprintf(stderr, "\nRandQuik wrote %" PRIu64 " bytes!\n\n", bytes);
return 0;
}
bool parse_hex(char* str, unsigned char* buf, size_t len) {
for (size_t i = 0; i < len; ++i) {
int sz = 0;
if (sscanf(str, "%2hhx%n", buf + i, &sz) != 1) {
if (*str) {
fprintf(stderr, "Unable to read seed at `%s`\n\n", str);
return false;
}
return true; // Shorter than key length is OK
}
str += sz;
}
return true;
}
void print_hex(unsigned char* buf, size_t len) {
for (size_t i = 0; i < len; ++i) {
fprintf(stderr, "%02hhx", buf[i]);
}
}
void help(char** argv) {
fprintf(
stderr,
"Usage: %s [-t #threads] [-s hexseed] [-b #bytes] [-r #rounds] [-o "
"outputfile]\n\n",
argv[0]
);
}
int main(int argc, char** argv) {
unsigned char key[32] = {};
unsigned char iv[16] = {};
unsigned int workers = 8;
unsigned int rounds = 20;
char* output = NULL;
uint64_t max_bytes = 0;
bool seeded = false;
for (char opt; (opt = getopt(argc, argv, "bostr")) != -1;) {
if (opt == 't') {
if (optind >= argc || sscanf(argv[optind++], "%u", &workers) != 1) {
fprintf(
stderr, "Expected the number of worker threads after -t\n"
);
return 1;
}
continue;
}
if (opt == 'r') {
if (optind >= argc || sscanf(argv[optind++], "%u", &rounds) != 1) {
fprintf(
stderr,
"Expected the number ChaCha rounds (8, 12 or 20) after -r\n"
);
return 1;
}
continue;
}
if (opt == 's') {
if (optind >= argc || !parse_hex(argv[optind++], key, 32)) {
fprintf(stderr, "Expected a hex seed string after -s\n");
return 1;
}
seeded = true;
continue;
}
if (opt == 'o') {
if (optind >= argc) {
fprintf(stderr, "Expected output filename after -s\n");
return 1;
}
if (strcmp(argv[optind], "-") != 0) {
output = argv[optind++];
}
continue;
}
if (opt == 'b') {
char unit[16] = {};
if (optind >= argc || sscanf(argv[optind++], "%" SCNu64 "%15s", &max_bytes, unit) < 1) {
fprintf(
stderr,
"Expected a maximum number of bytes to read after -b\n"
);
return 1;
}
if (strcasecmp(unit, "k") == 0 || strcasecmp(unit, "kb") == 0)
max_bytes *= 1000ull;
else if (strcasecmp(unit, "m") == 0 || strcasecmp(unit, "mb") == 0)
max_bytes *= 1000000ull;
else if (strcasecmp(unit, "g") == 0 || strcasecmp(unit, "gb") == 0)
max_bytes *= 1000000000ull;
else if (strcasecmp(unit, "t") == 0 || strcasecmp(unit, "tb") == 0)
max_bytes *= 1000000000000ull;
else if (strcasecmp(unit, "ki") == 0 || strcasecmp(unit, "kib") == 0)
max_bytes <<= 10;
else if (strcasecmp(unit, "mi") == 0 || strcasecmp(unit, "mib") == 0)
max_bytes <<= 20;
else if (strcasecmp(unit, "gi") == 0 || strcasecmp(unit, "gib") == 0)
max_bytes <<= 30;
else if (strcasecmp(unit, "ti") == 0 || strcasecmp(unit, "tib") == 0)
max_bytes <<= 40;
continue;
}
help(argv);
return 1;
}
FILE* f = stdout;
if (output) {
f = fopen(output, "wb");
if (!f) {
fprintf(stderr, "Failed to open %s for writing.\n", output);
return 1;
}
} else if (isatty(1)) {
fprintf(
stderr,
"Won't print random on console. Pipe me to another program or "
"file instead.\n\n"
);
help(argv);
return 1;
}
if (!seeded) {
FILE* urand = fopen("/dev/urandom", "rb");
if (!urand || fread(key, 32, 1, urand) != 1) {
fprintf(
stderr, "Failed to seed from /dev/urandom. Use -s hexstring for "
"manual seeding.\n"
);
fclose(urand);
return 1;
}
fclose(urand);
fprintf(
stderr,
"Random seed generated. This sequence may be repeated by:\n%s ",
argv[0]
);
if (rounds != 20)
fprintf(stderr, "-r %u -s ", rounds);
else
fprintf(stderr, "-s ");
print_hex(key, 32);
fprintf(stderr, "\n\n");
}
signal(SIGINT, signal_handler);
signal(SIGTERM, signal_handler);
int ret = fast(f, workers, max_bytes, key, iv, rounds);
fclose(f);
return ret;
}
+51
View File
@@ -0,0 +1,51 @@
#cython: language_level=3
from libc.stdint cimport int64_t, uint32_t, uint8_t, uint64_t
from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
import numpy as np
cimport numpy as np
cimport cython
import secrets
from numpy.random cimport BitGenerator
np.import_array()
cdef extern from "chanumpy.h":
struct cha_ctx:
uint32_t state[16]
uint8_t unconsumed[512]
uint32_t offset, end;
unsigned rounds;
void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds) nogil
void cha_seek(cha_ctx* ctx, int64_t offset)
int64_t cha_tell(cha_ctx* ctx)
uint64_t cha_uint64(void *state) nogil
uint32_t cha_uint32(void *state) nogil
double cha_double(void *state) nogil
cdef class Cha(BitGenerator):
cdef cha_ctx rng_state
def __init__(self, seed=None, *, rounds=20):
BitGenerator.__init__(self, seed)
self._bitgen.state = <void *>&self.rng_state
self._bitgen.next_uint64 = &cha_uint64
self._bitgen.next_uint32 = &cha_uint32
self._bitgen.next_double = &cha_double
self._bitgen.next_raw = &cha_uint64
# Generated state is ChaCha20 key
key = self._seed_seq.generate_state(4, np.uint64)
cha_init(&self.rng_state, <uint8_t *>np.PyArray_DATA(key), bytes(16) + b"NumpRand", rounds)
def advance(self, delta):
cha_seek(&self.rng_state, delta << 3)
def tell(self):
return cha_tell(&self.rng_state) >> 3;
def state(self):
return self.rng_state.state[12], self.rng_state.state[13]
+75
View File
@@ -0,0 +1,75 @@
from secrets import randbelow, token_bytes
import pytest
from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20
from randquik import cha
def test_cipherstreams_fullblocks():
"""Requests in multiple of ChaCha20 block size 64 bytes"""
key = token_bytes(32)
iv = token_bytes(16)
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
c1 = cha.Cha(key, iv)
for i in range(2048):
N = 64 * (1 + randbelow(2048))
ct0 = c0.update(bytes(N))
ct1 = c1(bytearray(N))
assert len(ct0) == len(ct1)
assert ct0.hex() == ct1.hex(), f"{i=} {N=}"
@pytest.mark.parametrize(
"counter",
[
b"\x00\x00\x00\x00\x00\x00\x00\x00",
b"\xFF\xFF\xFF\xFF\x00\x00\x00\x00",
b"\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
b"\xFF\xFF\xFF\x7F\xFF\xFF\xFF\xFF",
b"\x00\x00\x00\x80\xFF\xFF\xFF\xFF",
],
)
def test_counter_wrap(counter):
"""Tests carry handling of counter increments"""
key = bytes(32)
iv = counter + b"--------"
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
c1 = cha.Cha(key, iv)
N = 128
ct0 = c0.update(bytes(N))
ct1 = c1(bytearray(N))
assert len(ct0) == len(ct1)
assert ct0.hex() == ct1.hex()
def test_cipherstreams_partial_updates():
"""Odd-sized requests that retain leftover buffers"""
key = token_bytes(32)
iv = token_bytes(16)
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
c1 = cha.Cha(key, iv)
for i in range(2048):
N = 1 + randbelow(2048)
ct0 = c0.update(bytes(N))
ct1 = c1(bytearray(N))
assert len(ct0) == len(ct1)
assert ct0.hex() == ct1.hex(), f"{i=} {N=}"
def test_cipherstreams_32leftover():
"""Test the special case where the second response comes entirely from the leftover buffer"""
key = token_bytes(32)
nonce = token_bytes(12) # IETF nonce size
iv = bytes(4) + nonce
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
c1 = cha.Cha(key, nonce)
for N in [512 - 32, 32]:
ct0 = c0.update(bytes(N))
ct1 = c1(bytearray(N))
assert len(ct0) == len(ct1)
assert ct0.hex() == ct1.hex()