Position seek/tell.

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);
 }

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

src/cli.c

@@ -10,7 +10,7 @@
 #include <time.h>
 #include <unistd.h>
 
-#include "chacha20.h"
+#include "charandom.h"
 
 static volatile bool quit = false;
 

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]