3 files changed, 34 insertions, 0 deletions
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp
index 5032131ab..e028de526 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp
@@ -23,6 +23,20 @@ const int BufferWithExtendableBuffer::NEAR_BUFFER_LIMIT_THRESHOLD_PERCENTILE = 9
 // TODO: Needs to allocate larger memory corresponding to the current vector size.
 const size_t BufferWithExtendableBuffer::EXTEND_ADDITIONAL_BUFFER_SIZE_STEP = 128 * 1024;
 
+uint32_t BufferWithExtendableBuffer::readUintAndAdvancePosition(const int size,
+        int *const pos) const {
+    const bool readingPosIsInAdditionalBuffer = isInAdditionalBuffer(*pos);
+    if (readingPosIsInAdditionalBuffer) {
+        *pos -= mOriginalBufferSize;
+    }
+    const int value = ByteArrayUtils::readUintAndAdvancePosition(
+            getBuffer(readingPosIsInAdditionalBuffer), size, pos);
+    if (readingPosIsInAdditionalBuffer) {
+        *pos += mOriginalBufferSize;
+    }
+    return value;
+}
+
 bool BufferWithExtendableBuffer::writeUintAndAdvancePosition(const uint32_t data, const int size,
         int *const pos) {
     if (!(size >= 1 && size <= 4)) {
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h
index 1e27a1bec..2d89f71b1 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h
@@ -71,6 +71,8 @@ class BufferWithExtendableBuffer {
         }
     }
 
+    uint32_t readUintAndAdvancePosition(const int size, int *const pos) const;
+
     AK_FORCE_INLINE int getOriginalBufferSize() const {
         return mOriginalBufferSize;
     }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/byte_array_utils.h b/native/jni/src/suggest/policyimpl/dictionary/utils/byte_array_utils.h
index 0c1576818..1ca01b868 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/byte_array_utils.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/byte_array_utils.h
@@ -114,6 +114,24 @@ class ByteArrayUtils {
         return buffer[(*pos)++];
     }
 
+    static AK_FORCE_INLINE int readUintAndAdvancePosition(const uint8_t *const buffer,
+            const int size, int *const pos) {
+        // size must be in 1 to 4.
+        ASSERT(size >= 1 && size <= 4);
+        switch (size) {
+            case 1:
+                return ByteArrayUtils::readUint8AndAdvancePosition(buffer, pos);
+            case 2:
+                return ByteArrayUtils::readUint16AndAdvancePosition(buffer, pos);
+            case 3:
+                return ByteArrayUtils::readUint24AndAdvancePosition(buffer, pos);
+            case 4:
+                return ByteArrayUtils::readUint32AndAdvancePosition(buffer, pos);
+            default:
+                return 0;
+        }
+    }
+
     /**
      * Code Point Reading
      *