1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
//! # Indexed Line Reader
//! `IndexedLineReader` implements the `Seek` trait to allow seeking to specific lines
//!
//! ## Index Granularity
//! `IndexedLineReader` just stores a byte count every `n` lines to allow fast seeking by line.
//! The granularity can be configured by passing the valut to the `IndexedLineReader` constructor.
//!
//! There is a tradeoff to make between memory occupied by the index and the seek speed,
//! in general lower granularity means more line indexed and higher memory consumption, but fast seek time,
//! while a higher granularity slows down the seek time but less indexes are kept in memory.
//!
//! As an example, if a file has 100,000,000 lines, usually a granularity of 100,000 gives a good tradeoff
//! between performance and memory consumption.
//!
//! ## Example
//!
//! ```no_run
//! extern crate indexed_line_reader;
//! # fn main() {
//!
//! use indexed_line_reader::*;
//! use std::fs::*;
//! use std::io::{BufRead, BufReader, Seek, SeekFrom, Write};
//!
//! /* Creates an IndexedLineReader for a file with index granularity of 100 */
//! let file_reader = OpenOptions::new().read(true).open("file.txt").expect("Unable to open file reader");
//! let mut indexed_line_reader = &mut IndexedLineReader::new(BufReader::new(file_reader), 100);
//!
//! /* Seeks to line 100 from the start of the file */
//! indexed_line_reader.seek(SeekFrom::Start(100));
//!
//! /* Seeks forward by 50 lines from the current position on the file */
//! indexed_line_reader.seek(SeekFrom::Current(50));
//!
//! /* Seeks backward by 50 lines from the current position on the file */
//! indexed_line_reader.seek(SeekFrom::Current(-50));
//!
//! /* Seeks to the 100th line from the end on the file */
//! indexed_line_reader.seek(SeekFrom::End(100));
//!
//! # }
//! ```

use std::collections::BTreeMap;
use std::io::{BufRead, Error, Read, Seek, SeekFrom};

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LinesIndex {
    index: BTreeMap<u64, u64>,
    granularity: u64,
    line_count: u64,
    byte_count: u64
}

impl LinesIndex {
    pub fn new(granularity: u64) -> LinesIndex {
        LinesIndex {
            index: BTreeMap::new(),
            granularity: granularity,
            line_count: 0,
            byte_count: 0
        }
    }

    pub fn insert(&mut self, pos: u64, byte_count: u64) -> Option<u64> {
        self.index.insert(pos, byte_count)
    }

    pub fn clear(&mut self) {
        self.index.clear();
    }

    pub fn line_count(&self) -> u64 {
        self.line_count
    }

    pub fn byte_count(&self) -> u64 {
        self.byte_count
    }

    pub fn byte_count_at_pos(&self, pos: &u64) -> Option<u64> {
        self.index.get(pos).map(|&x| x)
    }

    pub fn last_indexed_pos(&self) -> Option<u64> {
        self.index.keys().map(|&x| x).max()
    }

    pub fn compute<T: BufRead + Seek>(&mut self, mut reader: &mut T) -> Result<u64, Error> {
        let initial_pos = self.last_indexed_pos().unwrap_or(0);
        let mut line_count = initial_pos;
        let mut byte_count = self.byte_count_at_pos(&line_count).unwrap_or(0);
        try!(reader.seek(SeekFrom::Start(byte_count)));
        if byte_count > 0 {
            reader.lines().next();
        }
        for (pos, line) in reader.lines().enumerate() {
            match line {
                Ok(line) => {
                    byte_count += line.as_bytes().len() as u64 + 1;
                    if (pos as u64 + 1) % self.granularity == 0 {
                        self.insert(initial_pos + pos as u64 + 1, byte_count);
                    }
                    line_count += 1;
                },
                Err(err) => return Err(err)
            }
        }
        self.line_count = line_count;
        self.byte_count = byte_count;
        Ok(line_count)
    }

    pub fn get_ref(&self) -> &BTreeMap<u64, u64> {
        &self.index
    }

    pub fn get_mut(&mut self) -> &mut BTreeMap<u64, u64> {
        &mut self.index
    }

    pub fn into_inner(self) -> BTreeMap<u64, u64> {
        self.index
    }
}

#[derive(Debug)]
pub struct IndexedLineReader<T> {
    index: LinesIndex,
    pos: u64,
    line_count: u64,
    reader: T
}

impl<T: BufRead + Seek> IndexedLineReader<T> {
    pub fn new(reader: T, index_granularity: u64) -> IndexedLineReader<T> {
        IndexedLineReader {
            index: LinesIndex::new(index_granularity),
            pos: 0,
            line_count: 0,
            reader: reader
        }
    }

    pub fn get_index(&self) -> &LinesIndex {
        &self.index
    }

    pub fn restore_index(&mut self, index: LinesIndex) {
        self.index = index;
    }

    pub fn compute_index(&mut self) -> Result<u64, Error> {
        self.index.compute(&mut self.reader).and_then(|line_count| {
            self.line_count = line_count;
            Ok(line_count)
        })
    }

    pub fn clear_index(&mut self) {
        self.index.clear()
    }

    pub fn get_current_pos(&self) -> u64 {
        self.pos
    }

    pub fn byte_count(&mut self) -> Result<u64, Error> {
        self.reader.seek(SeekFrom::End(0))
    }

    fn seek_to_index(&mut self, indexed_pos: u64) -> Result<u64, Error> {
        self.pos = indexed_pos;
        let byte_count = self.index.byte_count_at_pos(&indexed_pos).unwrap_or(0);
        self.reader.seek(SeekFrom::Start(byte_count))
    }

    fn seek_to_closest_index(&mut self, pos: SeekFrom) -> Result<u64, Error> {
        match pos {
            SeekFrom::Start(pos) => {
                let extra_lines = pos % self.index.granularity;
                let closest_index = pos - extra_lines;
                self.seek_to_index(closest_index)
            },
            SeekFrom::Current(pos) => {
                let extra_lines = pos as u64 % self.index.granularity;
                let extra_lines_from_current_pos = self.pos % self.index.granularity;
                let previous_closest_index = self.pos - extra_lines_from_current_pos;
                let closest_index = previous_closest_index + pos as u64 - extra_lines;
                self.seek_to_index(closest_index)
            },
            SeekFrom::End(pos) => {
                let pos = self.line_count - pos.abs() as u64;
                self.seek_to_closest_index(SeekFrom::Start(pos))
            }
        }
    }

    fn seek_forward(&mut self, lines: u64) -> Result<u64, Error> {
        let mut lines_left = lines;
        let mut extra_byte_count: u64 = 0;
        for line in (&mut self.reader).lines() {
            match line {
                Ok(line) => {
                    lines_left -= 1;
                    self.pos += 1;
                    extra_byte_count += line.as_bytes().len() as u64 + 1;
                    if lines_left == 0 { break }
                },
                Err(err) => return Err(err)
            }
        }
        Ok(extra_byte_count)
    }
}

impl<T: Read> Read for IndexedLineReader<T> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
        self.reader.read(buf)
    }
}

impl<T: BufRead> BufRead for IndexedLineReader<T> {
    fn fill_buf(&mut self) -> Result<&[u8], Error> {
        self.reader.fill_buf()
    }
    fn consume(&mut self, amt: usize) {
        self.reader.consume(amt)
    }
}

impl<T: BufRead + Seek> Seek for IndexedLineReader<T> {
    fn seek(&mut self, pos: SeekFrom) -> Result<u64, Error> {
        self.compute_index().and_then(|_| {
            match pos {
                SeekFrom::Start(pos) => {
                    let extra_lines = pos as u64 % self.index.granularity;
                    self.seek_to_closest_index(SeekFrom::Start(pos)).and_then(|new_pos| {
                        if extra_lines > 0 {
                            self.seek(SeekFrom::Current(extra_lines as i64))
                        } else {
                            Ok(new_pos)
                        }
                    })
                },
                SeekFrom::Current(pos) => {
                    if pos >= 0 {
                        let extra_lines = pos as u64 % self.index.granularity;
                        let extra_lines_from_current_pos = self.pos % self.index.granularity;
                        self.seek_to_closest_index(SeekFrom::Current(pos)).and_then(|new_pos| {
                            if extra_lines + extra_lines_from_current_pos > 0 {
                                self.seek_forward(extra_lines + extra_lines_from_current_pos)
                            } else {
                                Ok(new_pos)
                            }
                        })
                    } else {
                        let pos = self.pos - pos.abs() as u64;
                        self.seek(SeekFrom::Start(pos))
                    }
                },
                SeekFrom::End(pos) => {
                    let pos = self.line_count - pos.abs() as u64;
                    self.seek(SeekFrom::Start(pos))
                }
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::fs::*;
    use std::io::{BufRead, BufReader, Seek, SeekFrom, Write};

    fn seek_and_assert_line_number(mut reader: &mut IndexedLineReader<BufReader<File>>,
                                    seek_from: SeekFrom, expected_line_number: u64) {
        reader.seek(seek_from).expect(&format!("Unable to seek from {:?}", seek_from));
        let line = (&mut reader).lines().next().unwrap().unwrap();
        let line_number: u64 = line.parse().expect("Unable to deserialize line number");
        assert_eq!(line_number, expected_line_number);
    }

    #[test]
    fn test_seek() {
        let log_name = "indexed-line-reader.log";
        let mut file_writer = OpenOptions::new().create(true).write(true).append(true).open(log_name).expect("Unable to open file writer");

        for i in 0..10000 {
            assert!(write!(file_writer, "{}\n", i).is_ok());
        }

        let file_reader = OpenOptions::new().read(true).open(log_name).expect("Unable to open file reader");
        let mut line_reader = &mut IndexedLineReader::new(BufReader::new(file_reader), 100);

        line_reader.compute_index().expect("Unable to compute index");

        seek_and_assert_line_number(line_reader, SeekFrom::Start(1234), 1234);
        seek_and_assert_line_number(line_reader, SeekFrom::Start(2468), 2468);
        seek_and_assert_line_number(line_reader, SeekFrom::Current(1000), 3468);
        seek_and_assert_line_number(line_reader, SeekFrom::Current(1032), 4500);
        seek_and_assert_line_number(line_reader, SeekFrom::Current(-450), 4050);
        seek_and_assert_line_number(line_reader, SeekFrom::End(1234), 8766);
        seek_and_assert_line_number(line_reader, SeekFrom::End(-1234), 8766);

        remove_file(log_name).expect("Unable to delete log");
    }
}