util/dbuffer.lua

changeset 0
89e39cd5a7cd
equal deleted inserted replaced
-1:000000000000 0:89e39cd5a7cd
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 };

mercurial