1 #include "byte_offsets.h"
2 #include <arpa/inet.h>
3 
NewByteOffsets()4 RSByteOffsets *NewByteOffsets() {
5   RSByteOffsets *ret = rm_calloc(1, sizeof(*ret));
6   return ret;
7 }
8 
RSByteOffsets_Free(RSByteOffsets * offsets)9 void RSByteOffsets_Free(RSByteOffsets *offsets) {
10   rm_free(offsets->offsets.data);
11   rm_free(offsets->fields);
12   rm_free(offsets);
13 }
14 
RSByteOffsets_ReserveFields(RSByteOffsets * offsets,size_t numFields)15 void RSByteOffsets_ReserveFields(RSByteOffsets *offsets, size_t numFields) {
16   offsets->fields = rm_realloc(offsets->fields, sizeof(*offsets->fields) * numFields);
17 }
18 
RSByteOffsets_AddField(RSByteOffsets * offsets,uint32_t fieldId,uint32_t startPos)19 RSByteOffsetField *RSByteOffsets_AddField(RSByteOffsets *offsets, uint32_t fieldId,
20                                           uint32_t startPos) {
21   RSByteOffsetField *field = &(offsets->fields[offsets->numFields++]);
22   field->fieldId = fieldId;
23   field->firstTokPos = startPos;
24   return field;
25 }
26 
ByteOffsetWriter_Move(ByteOffsetWriter * w,RSByteOffsets * offsets)27 void ByteOffsetWriter_Move(ByteOffsetWriter *w, RSByteOffsets *offsets) {
28   offsets->offsets.data = w->buf.data;
29   offsets->offsets.len = w->buf.offset;
30   memset(&w->buf, 0, sizeof w->buf);
31 }
32 
RSByteOffsets_Serialize(const RSByteOffsets * offsets,Buffer * b)33 void RSByteOffsets_Serialize(const RSByteOffsets *offsets, Buffer *b) {
34   BufferWriter w = NewBufferWriter(b);
35 
36   Buffer_WriteU8(&w, offsets->numFields);
37 
38   for (size_t ii = 0; ii < offsets->numFields; ++ii) {
39     Buffer_WriteU8(&w, offsets->fields[ii].fieldId);
40     Buffer_WriteU32(&w, offsets->fields[ii].firstTokPos);
41     Buffer_WriteU32(&w, offsets->fields[ii].lastTokPos);
42   }
43 
44   Buffer_WriteU32(&w, offsets->offsets.len);
45   Buffer_Write(&w, offsets->offsets.data, offsets->offsets.len);
46 }
47 
LoadByteOffsets(Buffer * buf)48 RSByteOffsets *LoadByteOffsets(Buffer *buf) {
49   BufferReader r = NewBufferReader(buf);
50 
51   RSByteOffsets *offsets = NewByteOffsets();
52   uint8_t numFields = Buffer_ReadU8(&r);
53   RSByteOffsets_ReserveFields(offsets, numFields);
54 
55   for (size_t ii = 0; ii < numFields; ++ii) {
56     uint8_t fieldId = Buffer_ReadU8(&r);
57     uint32_t firstTok = Buffer_ReadU32(&r);
58     uint32_t lastTok = Buffer_ReadU32(&r);
59     RSByteOffsetField *fieldInfo = RSByteOffsets_AddField(offsets, fieldId, firstTok);
60     fieldInfo->lastTokPos = lastTok;
61   }
62 
63   uint32_t offsetsLen = Buffer_ReadU32(&r);
64   offsets->offsets.len = offsetsLen;
65   if (offsetsLen) {
66     offsets->offsets.data = rm_malloc(offsetsLen);
67     Buffer_Read(&r, offsets->offsets.data, offsetsLen);
68   } else {
69     offsets->offsets.data = NULL;
70   }
71 
72   return offsets;
73 }
74 
RSByteOffset_Iterate(const RSByteOffsets * offsets,uint32_t fieldId,RSByteOffsetIterator * iter)75 int RSByteOffset_Iterate(const RSByteOffsets *offsets, uint32_t fieldId,
76                          RSByteOffsetIterator *iter) {
77   const RSByteOffsetField *offField = NULL;
78   for (size_t ii = 0; ii < offsets->numFields; ++ii) {
79     if (offsets->fields[ii].fieldId == fieldId) {
80       offField = offsets->fields + ii;
81       break;
82     }
83   }
84   if (!offField) {
85     return REDISMODULE_ERR;
86   }
87 
88   // printf("Generating iterator for fieldId=%lu. BeginPos=%lu. EndPos=%lu\n", fieldId,
89   //        offField->firstTokPos, offField->lastTokPos);
90 
91   iter->buf.cap = 0;
92   iter->buf.data = offsets->offsets.data;
93   iter->buf.offset = offsets->offsets.len;
94   iter->rdr = NewBufferReader(&iter->buf);
95   iter->curPos = 1;
96   iter->endPos = offField->lastTokPos;
97 
98   iter->lastValue = 0;
99 
100   while (iter->curPos < offField->firstTokPos && !BufferReader_AtEnd(&iter->rdr)) {
101     // printf("Seeking & incrementing\n");
102     iter->lastValue = ReadVarint(&iter->rdr) + iter->lastValue;
103     iter->curPos++;
104   }
105 
106   // printf("Iterator is now at %lu\n", iter->curPos);
107   iter->curPos--;
108   return REDISMODULE_OK;
109 }
110 
RSByteOffsetIterator_Next(RSByteOffsetIterator * iter)111 uint32_t RSByteOffsetIterator_Next(RSByteOffsetIterator *iter) {
112   if (BufferReader_AtEnd(&iter->rdr) || ++iter->curPos > iter->endPos) {
113     return RSBYTEOFFSET_EOF;
114   }
115 
116   iter->lastValue = ReadVarint(&iter->rdr) + iter->lastValue;
117   return iter->lastValue;
118 }