From 5fb8af478dc50094fda4ddcdc46ccd432a2ca762 Mon Sep 17 00:00:00 2001
From: Leo Vasanko <git@vasanko.com>
Date: Fri, 1 Dec 2023 02:40:29 +0000
Subject: [PATCH] Position seek/tell.

---
 src/{npbitgen.h => chanumpy.h}  |  9 +++++----
 src/{chacha20.c => charandom.c} |  0
 src/{chacha20.h => charandom.h} | 24 ++++++++++++++++++++++--
 src/cli.c                       |  2 +-
 src/nprand.pyx                  | 23 +++++++++++++++++------
 5 files changed, 45 insertions(+), 13 deletions(-)
 rename src/{npbitgen.h => chanumpy.h} (61%)
 rename src/{chacha20.c => charandom.c} (100%)
 rename src/{chacha20.h => charandom.h} (79%)

diff --git a/src/npbitgen.h b/src/chanumpy.h
similarity index 61%
rename from src/npbitgen.h
rename to src/chanumpy.h
index e6ba65f..abc0a24 100644
--- a/src/npbitgen.h
+++ b/src/chanumpy.h
@@ -1,9 +1,9 @@
-#include "chacha20.h"
+#include "charandom.h"
 
 static uint64_t cha_uint64(void* st) {
     cha_ctx* ctx = (cha_ctx*)st;
-    if (ctx->offset == ctx->end) {
-        ctx->offset = 0;
+    if (ctx->offset + sizeof(uint64_t) > ctx->end) {
+        ctx->offset -= ctx->end;
         ctx->end =
           ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
     }
@@ -13,5 +13,6 @@ static uint64_t cha_uint64(void* st) {
 }
 static uint32_t cha_uint32(void* st) { return cha_uint64(st); }
 static double cha_double(void* st) {
-    return cha_uint64(st) / (UINT64_MAX + 1.0);
+    // Fast uint64_to_double conversion from numpy/random/_common.pxd
+    return (cha_uint64(st) >> 11) * (1.0 / 9007199254740992.0);
 }
diff --git a/src/chacha20.c b/src/charandom.c
similarity index 100%
rename from src/chacha20.c
rename to src/charandom.c
diff --git a/src/chacha20.h b/src/charandom.h
similarity index 79%
rename from src/chacha20.h
rename to src/charandom.h
index ea1db7d..5934b3f 100644
--- a/src/chacha20.h
+++ b/src/charandom.h
@@ -75,11 +75,31 @@ void cha_init(
 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;
-    ctx->offset = ctx->end = 0;
+    *(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
diff --git a/src/cli.c b/src/cli.c
index 0e132f2..5282c08 100644
--- a/src/cli.c
+++ b/src/cli.c
@@ -10,7 +10,7 @@
 #include <time.h>
 #include <unistd.h>
 
-#include "chacha20.h"
+#include "charandom.h"
 
 static volatile bool quit = false;
 
diff --git a/src/nprand.pyx b/src/nprand.pyx
index aad27e9..1204861 100644
--- a/src/nprand.pyx
+++ b/src/nprand.pyx
@@ -1,8 +1,7 @@
 #cython: language_level=3
 
-from libc.stdint cimport uint32_t, uint8_t, uint64_t
+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
@@ -12,7 +11,7 @@ from numpy.random cimport BitGenerator
 
 np.import_array()
 
-cdef extern from "npbitgen.h":
+cdef extern from "chanumpy.h":
     struct cha_ctx:
         uint32_t state[16]
         uint8_t unconsumed[512]
@@ -20,15 +19,18 @@ cdef extern from "npbitgen.h":
         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 class Cha(BitGenerator):
     cdef cha_ctx rng_state
 
-    def __init__(self, seed=None, rounds=8):
+    def __init__(self, seed=None, *, rounds=20):
         BitGenerator.__init__(self, seed)
         self._bitgen.state = <void *>&self.rng_state
         self._bitgen.next_uint64 = &cha_uint64
@@ -37,4 +39,13 @@ cdef class Cha(BitGenerator):
         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), b"NumpRand", rounds)
+        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]