From 179a6002c20009ed07a93f5ec16d3481187d085e Mon Sep 17 00:00:00 2001
From: Viktor Lofgren <vlofgren@marginalia.nu>
Date: Fri, 12 Jul 2024 23:50:28 +0200
Subject: [PATCH] (coded-sequence) Add a callback for re-filling underlying
 buffer

---
 .../index/PrioIndexEntrySource.java           | 11 +++++-----
 .../nu/marginalia/sequence/io/BitReader.java  | 22 +++++++++++++++++--
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/code/index/index-reverse/java/nu/marginalia/index/PrioIndexEntrySource.java b/code/index/index-reverse/java/nu/marginalia/index/PrioIndexEntrySource.java
index 6346290f..e3b93d44 100644
--- a/code/index/index-reverse/java/nu/marginalia/index/PrioIndexEntrySource.java
+++ b/code/index/index-reverse/java/nu/marginalia/index/PrioIndexEntrySource.java
@@ -15,7 +15,7 @@ public class PrioIndexEntrySource implements EntrySource {
     private final String name;
 
     private final ByteBuffer readData = ByteBuffer.allocate(1024);
-    private final BitReader bitReader = new BitReader(readData);
+    private final BitReader bitReader = new BitReader(readData, this::fillReadBuffer);
 
     private final FileChannel docsFileChannel;
     private long dataOffsetStartB;
@@ -69,8 +69,6 @@ public class PrioIndexEntrySource implements EntrySource {
         outputBuffer.clear();
 
         while (outputBuffer.hasRemaining() && readItems++ < numItems) {
-            fillReadBuffer();
-
             int rank;
             int domainId;
             int docOrd;
@@ -119,8 +117,8 @@ public class PrioIndexEntrySource implements EntrySource {
         buffer.uniq();
     }
 
-    private void fillReadBuffer() throws IOException {
-        if (readData.remaining() < 8) {
+    private void fillReadBuffer() {
+        try {
             readData.compact();
             int rb = docsFileChannel.read(readData, dataOffsetStartB);
             if (rb > 0) {
@@ -128,6 +126,9 @@ public class PrioIndexEntrySource implements EntrySource {
             }
             readData.flip();
         }
+        catch (IOException ex) {
+            throw new IllegalStateException("Failed to read index data.", ex);
+        }
     }
 
     @Override
diff --git a/code/libraries/coded-sequence/java/nu/marginalia/sequence/io/BitReader.java b/code/libraries/coded-sequence/java/nu/marginalia/sequence/io/BitReader.java
index 930f67f0..524e2a6b 100644
--- a/code/libraries/coded-sequence/java/nu/marginalia/sequence/io/BitReader.java
+++ b/code/libraries/coded-sequence/java/nu/marginalia/sequence/io/BitReader.java
@@ -1,5 +1,8 @@
 package nu.marginalia.sequence.io;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.nio.ByteBuffer;
 
 /** A utility class for reading bits from a ByteBuffer
@@ -7,6 +10,9 @@ import java.nio.ByteBuffer;
  */
 public class BitReader {
     private final ByteBuffer underlying;
+    private final Runnable refillCallback;
+
+    private static final Logger logger = LoggerFactory.getLogger(BitReader.class);
 
     /** The current value being decoded */
     private long currentValue;
@@ -14,12 +20,23 @@ public class BitReader {
     /** Bit index in the current value */
     private int bitPosition;
 
-    public BitReader(ByteBuffer buffer) {
+
+    /** Create a new BitReader for the given buffer.  The supplied callback will be
+     * invoked when the underlying buffer is out of data.  The callback should
+     * refill the buffer with more data.
+     */
+    public BitReader(ByteBuffer buffer, Runnable refillCallback) {
         this.underlying = buffer;
+        this.refillCallback = refillCallback;
         this.bitPosition = 0;
         this.currentValue = 0;
     }
 
+    /** Create a new BitReader for the given buffer */
+    public BitReader(ByteBuffer buffer) {
+        this(buffer, () -> { throw new IllegalStateException("No more data to read and no re-fill callback provided"); });
+    }
+
     /** Read the next bit from the buffer */
     public boolean getBit() {
         if (bitPosition <= 0) {
@@ -132,7 +149,8 @@ public class BitReader {
             bitPosition = 8;
         }
         else { // There's no more data to read!
-            throw new ArrayIndexOutOfBoundsException("No more data to read");
+            refillCallback.run();
+            readNext();
         }
     }
 }