Thu, 07 Apr 2022 18:11:33 +0100
Initial commit
0 | 1 | local queue = require "util.queue"; |
2 | ||
3 | local s_byte, s_sub = string.byte, string.sub; | |
4 | local dbuffer_methods = {}; | |
5 | local dynamic_buffer_mt = { __name = "dbuffer", __index = dbuffer_methods }; | |
6 | ||
7 | function dbuffer_methods:write(data) | |
8 | if self.max_size and #data + self._length > self.max_size then | |
9 | return nil; | |
10 | end | |
11 | local ok = self.items:push(data); | |
12 | if not ok then | |
13 | self:collapse(); | |
14 | ok = self.items:push(data); | |
15 | end | |
16 | if not ok then | |
17 | return nil; | |
18 | end | |
19 | self._length = self._length + #data; | |
20 | return true; | |
21 | end | |
22 | ||
23 | function dbuffer_methods:read_chunk(requested_bytes) | |
24 | local chunk, consumed = self.items:peek(), self.front_consumed; | |
25 | if not chunk then return; end | |
26 | local chunk_length = #chunk; | |
27 | local remaining_chunk_length = chunk_length - consumed; | |
28 | if not requested_bytes then | |
29 | requested_bytes = remaining_chunk_length; | |
30 | end | |
31 | if remaining_chunk_length <= requested_bytes then | |
32 | self.front_consumed = 0; | |
33 | self._length = self._length - remaining_chunk_length; | |
34 | self.items:pop(); | |
35 | assert(#chunk:sub(consumed + 1, -1) == remaining_chunk_length); | |
36 | return chunk:sub(consumed + 1, -1), remaining_chunk_length; | |
37 | end | |
38 | local end_pos = consumed + requested_bytes; | |
39 | self.front_consumed = end_pos; | |
40 | self._length = self._length - requested_bytes; | |
41 | assert(#chunk:sub(consumed + 1, end_pos) == requested_bytes); | |
42 | return chunk:sub(consumed + 1, end_pos), requested_bytes; | |
43 | end | |
44 | ||
45 | function dbuffer_methods:read(requested_bytes) | |
46 | local chunks; | |
47 | ||
48 | if requested_bytes and requested_bytes > self._length then | |
49 | return nil; | |
50 | end | |
51 | ||
52 | local chunk, read_bytes = self:read_chunk(requested_bytes); | |
53 | if not requested_bytes then | |
54 | return chunk; | |
55 | elseif chunk then | |
56 | requested_bytes = requested_bytes - read_bytes; | |
57 | if requested_bytes == 0 then -- Already read everything we need | |
58 | return chunk; | |
59 | end | |
60 | chunks = {}; | |
61 | else | |
62 | return nil; | |
63 | end | |
64 | ||
65 | -- Need to keep reading more chunks | |
66 | while chunk do | |
67 | table.insert(chunks, chunk); | |
68 | if requested_bytes > 0 then | |
69 | chunk, read_bytes = self:read_chunk(requested_bytes); | |
70 | requested_bytes = requested_bytes - read_bytes; | |
71 | else | |
72 | break; | |
73 | end | |
74 | end | |
75 | ||
76 | return table.concat(chunks); | |
77 | end | |
78 | ||
79 | -- Read to, and including, the specified character sequence (return nil if not found) | |
80 | function dbuffer_methods:read_until(char) | |
81 | local buffer_pos = 0; | |
82 | for i, chunk in self.items:items() do | |
83 | local start = 1 + ((i == 1) and self.front_consumed or 0); | |
84 | local char_pos = chunk:find(char, start, true); | |
85 | if char_pos then | |
86 | return self:read(1 + buffer_pos + char_pos - start); | |
87 | end | |
88 | buffer_pos = buffer_pos + #chunk - (start - 1); | |
89 | end | |
90 | return nil; | |
91 | end | |
92 | ||
93 | function dbuffer_methods:discard(requested_bytes) | |
94 | if requested_bytes > self._length then | |
95 | return nil; | |
96 | end | |
97 | ||
98 | local chunk, read_bytes = self:read_chunk(requested_bytes); | |
99 | if chunk then | |
100 | requested_bytes = requested_bytes - read_bytes; | |
101 | if requested_bytes == 0 then -- Already read everything we need | |
102 | return true; | |
103 | end | |
104 | else | |
105 | return nil; | |
106 | end | |
107 | ||
108 | while chunk do | |
109 | if requested_bytes > 0 then | |
110 | chunk, read_bytes = self:read_chunk(requested_bytes); | |
111 | requested_bytes = requested_bytes - read_bytes; | |
112 | else | |
113 | break; | |
114 | end | |
115 | end | |
116 | return true; | |
117 | end | |
118 | ||
119 | -- Normalize i, j into absolute offsets within the | |
120 | -- front chunk (accounting for front_consumed), and | |
121 | -- ensure there is enough data in the first chunk | |
122 | -- to cover any subsequent :sub() or :byte() operation | |
123 | function dbuffer_methods:_prep_sub(i, j) | |
124 | if j == nil then | |
125 | j = -1; | |
126 | end | |
127 | if j < 0 then | |
128 | j = self._length + (j+1); | |
129 | end | |
130 | if i < 0 then | |
131 | i = self._length + (i+1); | |
132 | end | |
133 | if i < 1 then | |
134 | i = 1; | |
135 | end | |
136 | if j > self._length then | |
137 | j = self._length; | |
138 | end | |
139 | if i > j then | |
140 | return nil; | |
141 | end | |
142 | ||
143 | self:collapse(j); | |
144 | ||
145 | if self.front_consumed > 0 then | |
146 | i = i + self.front_consumed; | |
147 | j = j + self.front_consumed; | |
148 | end | |
149 | ||
150 | return i, j; | |
151 | end | |
152 | ||
153 | function dbuffer_methods:sub(i, j) | |
154 | i, j = self:_prep_sub(i, j); | |
155 | if not i then | |
156 | return ""; | |
157 | end | |
158 | return s_sub(self.items:peek(), i, j); | |
159 | end | |
160 | ||
161 | function dbuffer_methods:byte(i, j) | |
162 | i = i or 1; | |
163 | j = j or i; | |
164 | i, j = self:_prep_sub(i, j); | |
165 | if not i then | |
166 | return; | |
167 | end | |
168 | return s_byte(self.items:peek(), i, j); | |
169 | end | |
170 | ||
171 | function dbuffer_methods:length() | |
172 | return self._length; | |
173 | end | |
174 | dbuffer_methods.len = dbuffer_methods.length; -- strings have :len() | |
175 | dynamic_buffer_mt.__len = dbuffer_methods.length; -- support # operator | |
176 | ||
177 | function dbuffer_methods:collapse(bytes) | |
178 | bytes = bytes or self._length; | |
179 | ||
180 | local front_chunk = self.items:peek(); | |
181 | ||
182 | if not front_chunk or #front_chunk - self.front_consumed >= bytes then | |
183 | return; | |
184 | end | |
185 | ||
186 | local front_chunks = { front_chunk:sub(self.front_consumed+1) }; | |
187 | local front_bytes = #front_chunks[1]; | |
188 | ||
189 | while front_bytes < bytes do | |
190 | self.items:pop(); | |
191 | local chunk = self.items:peek(); | |
192 | front_bytes = front_bytes + #chunk; | |
193 | table.insert(front_chunks, chunk); | |
194 | end | |
195 | self.items:replace(table.concat(front_chunks)); | |
196 | self.front_consumed = 0; | |
197 | end | |
198 | ||
199 | local function new(max_size, max_chunks) | |
200 | if max_size and max_size <= 0 then | |
201 | return nil; | |
202 | end | |
203 | return setmetatable({ | |
204 | front_consumed = 0; | |
205 | _length = 0; | |
206 | max_size = max_size; | |
207 | items = queue.new(max_chunks or 32); | |
208 | }, dynamic_buffer_mt); | |
209 | end | |
210 | ||
211 | return { | |
212 | new = new; | |
213 | }; |