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
use crate::bwt::{AminoBwtBlock, NucleotideBwtBlock};
use crate::compressed_suffix_array::CompressedSuffixArray;
use crate::fm_index::FmIndex;
use crate::kmer_lookup_table::KmerLookupTable;
use crate::sequence_index::SequenceIndex;
use crate::simd_instructions::Vec256;
use crate::{alphabet::SymbolAlphabet, bwt::Bwt};
use std::convert::TryInto;
use std::io::{self, Read};
use std::slice;
use std::{
io::{Error, Write},
path::Path,
};
const FM_FILE_LABEL_STRING:&[u8;11] = b"AWRY-Index\n";
impl FmIndex {
/// Saves them FM-index to disk at the given file path
///
/// # Example
/// ```no_run
/// use awry::fm_index::{FmIndex, FmBuildArgs};
/// use awry::alphabet::SymbolAlphabet;
/// use std::path::Path;
///
///
/// let build_args = FmBuildArgs {
/// input_file_src: "test.fasta".into(),
/// suffix_array_output_src: None,
/// suffix_array_compression_ratio: Some(16),
/// lookup_table_kmer_len: None,
/// alphabet: SymbolAlphabet::Nucleotide,
/// max_query_len: None,
/// remove_intermediate_suffix_array_file: true,
/// };
/// let fm_index = FmIndex::new(&build_args).expect("unable to build fm index");
/// fm_index.save(&Path::new("test.awry")).expect("unable to save fm index to file");
/// ```
pub fn save(&self, file_output_src: &Path) -> Result<(), Error> {
let mut fm_index_file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true) // Overwrite if exists
.open(file_output_src)?;
//write the file label string so that the file type can be identified in a text editor
fm_index_file.write(FM_FILE_LABEL_STRING)?;
let header = self.generate_file_header();
for value in header {
fm_index_file.write_all(&value.to_le_bytes())?;
}
//write the BWT to file
match self.bwt() {
crate::bwt::Bwt::Nucleotide(vec) => {
for block in vec.iter() {
for bit_vector in block.bit_vectors() {
let bit_vector_values = bit_vector.data();
for bit_vector_value in bit_vector_values {
fm_index_file.write_all(&bit_vector_value.to_le_bytes())?;
}
}
for milestone in block.milestones() {
fm_index_file.write_all(&milestone.to_le_bytes())?;
}
}
}
crate::bwt::Bwt::Amino(vec) => {
for block in vec.iter() {
for bit_vector in block.bit_vectors() {
let bit_vector_values = bit_vector.data();
for bit_vector_value in bit_vector_values {
fm_index_file.write_all(&bit_vector_value.to_le_bytes())?;
}
}
for milestone in block.milestones() {
fm_index_file.write_all(&milestone.to_le_bytes())?;
}
}
}
}
//write the prefix sums
for prefix_sum in self.prefix_sums() {
fm_index_file.write_all(&prefix_sum.to_le_bytes())?;
}
//write the sampled suffix array
for suffix_array_value in self.sampled_suffix_array().data() {
fm_index_file.write_all(&suffix_array_value.to_le_bytes())?;
}
//write the kmer lookup table
debug_assert!(self.kmer_lookup_table().kmer_len() != 0);
fm_index_file.write(&self.kmer_lookup_table().kmer_len().to_le_bytes())?;
for range in self.kmer_lookup_table().table(){
fm_index_file.write_all(&range.start_ptr.to_le_bytes())?;
fm_index_file.write_all(&range.end_ptr.to_le_bytes())?;
}
self.sequence_index().serialize(&mut fm_index_file).expect("unable to serialize sequence index");
return Ok(());
}
///Loads the fm-index file from the given file path
///
/// # Example
/// ```no_run
/// use awry::fm_index::{FmIndex, FmBuildArgs};
/// use awry::alphabet::SymbolAlphabet;
/// use std::path::Path;
///
///
/// let build_args = FmBuildArgs {
/// input_file_src: "test.fasta".into(),
/// suffix_array_output_src: None,
/// suffix_array_compression_ratio: Some(16),
/// lookup_table_kmer_len: None,
/// alphabet: SymbolAlphabet::Nucleotide,
/// max_query_len: None,
/// remove_intermediate_suffix_array_file: true,
/// };
/// let fm_index = FmIndex::new(&build_args).expect("unable to build fm index");
/// fm_index.save(&Path::new("test.awry")).expect("unable to save fm index to file");
///
///
/// let loaded_fm_index = FmIndex::load(&Path::new("test.awry")).expect("unable to load fm index from file");
/// ```
pub fn load(fm_file_src: &Path) -> Result<FmIndex, Error> {
let mut fm_index_file = std::fs::OpenOptions::new()
.write(false)
.read(true)
.open(fm_file_src)?;
//read and check the file label
let mut file_label_buffer: [u8; FM_FILE_LABEL_STRING.len()] =
[0; FM_FILE_LABEL_STRING.len()];
let mut u64_buffer: [u8; 8] = [0; 8];
//read the label, and check to make sure it matches what we expect. if it doesn't, it's probably not an fm index file.
fm_index_file.read_exact(&mut file_label_buffer)?;
//compare the buffer to the expected file label
let file_label_validated = file_label_buffer.iter().zip(FM_FILE_LABEL_STRING.iter()).all(|(a,b)| a==b);
if !file_label_validated{
return Err(std::io::Error::new(io::ErrorKind::InvalidData,
"file provided did not start with expected label, and is probably not an fm index file"));
}
//read the header. this section may be rewritten and refactored if more than 1 version is supported
fm_index_file.read_exact(&mut u64_buffer)?;
let version_number = u64::from_le_bytes(u64_buffer);
return FmIndex::read_fm_index_by_version_number(& mut fm_index_file, version_number);
}
/// generates the file header from the data in the fm-index. This header
/// may differ on different versions of the index
fn generate_file_header(&self) -> [u64;4] {
match self.version_number() {
_ => {
let alphabet_idx: u64 = match self.bwt() {
Bwt::Nucleotide(_) => 0,
Bwt::Amino(_) => 1,
};
//Matches version 1
return [self.version_number(),
self.suffix_array_compression_ratio(),
self.bwt_len(),
alphabet_idx];
}
}
}
/// Reads the main contents of an fm index file, depending on the found version number.
fn read_fm_index_by_version_number(fm_index_file: &mut std::fs::File, version_number: u64) -> Result<FmIndex, Error> {
match version_number{
_=>{
let mut u64_buffer: [u8; 8] = [0; 8];
//currently only version 1 is supported.
fm_index_file.read_exact(&mut u64_buffer)?;
let suffix_array_compression_ratio = u64::from_le_bytes(u64_buffer);
fm_index_file.read_exact(&mut u64_buffer)?;
let bwt_len = u64::from_le_bytes(u64_buffer);
fm_index_file.read_exact(&mut u64_buffer)?;
let alphabet_idx = u64::from_le_bytes(u64_buffer);
let alphabet = match alphabet_idx{
0=>{SymbolAlphabet::Nucleotide},
1=>{SymbolAlphabet::Amino},
_=>panic!("invalid symbol alphabet , did not match any supported alphabet")
};
let suffix_array_num_words = CompressedSuffixArray::compressed_word_len(bwt_len as usize, suffix_array_compression_ratio);
let num_bwt_blocks = Bwt::num_blocks(bwt_len);
let bwt:Bwt = match alphabet{
SymbolAlphabet::Nucleotide => {
let mut bwt_block_list = vec![NucleotideBwtBlock::new();num_bwt_blocks];
for block_idx in 0..bwt_block_list.len(){
const BIT_VECTORS_PER_BLOCK:usize = NucleotideBwtBlock::NUM_BIT_VECTORS;
const BIT_VECTOR_BYTES_PER_BLOCK:usize = 32*BIT_VECTORS_PER_BLOCK;
const NUM_MILESTONES:usize = NucleotideBwtBlock::NUM_MILESTONES;
//read the bit vectors for this block
let mut bit_vector_buffer:[u8;BIT_VECTOR_BYTES_PER_BLOCK] = [0;BIT_VECTOR_BYTES_PER_BLOCK];
fm_index_file.read_exact(&mut bit_vector_buffer)?;
let buffer_ptr = bit_vector_buffer.as_ptr();
let vector_ptr = buffer_ptr as *const Vec256;
let bit_vector_slice = unsafe{
slice::from_raw_parts(vector_ptr, BIT_VECTORS_PER_BLOCK)
};
//read the milestones for this block
let mut milestones:[u64;NUM_MILESTONES] = [0; NUM_MILESTONES];
for milestone_idx in 0..milestones.len(){
fm_index_file.read_exact(&mut u64_buffer)?;
milestones[milestone_idx] = u64::from_le_bytes(u64_buffer);
}
bwt_block_list[block_idx] = NucleotideBwtBlock::from_data( bit_vector_slice.try_into().unwrap(), milestones);
}
Bwt::Nucleotide(bwt_block_list)
},
SymbolAlphabet::Amino =>{
let mut bwt_block_list = vec![AminoBwtBlock::new();num_bwt_blocks];
for block_idx in 0..bwt_block_list.len(){
const BIT_VECTORS_PER_BLOCK:usize = AminoBwtBlock::NUM_BIT_VECTORS;
//read the bit vectors for this block
let mut bit_vector_buffer:[u8;32*BIT_VECTORS_PER_BLOCK] = [0;32*BIT_VECTORS_PER_BLOCK];
fm_index_file.read_exact(&mut bit_vector_buffer)?;
let buffer_ptr = bit_vector_buffer.as_ptr();
let vector_ptr = buffer_ptr as *const Vec256;
let bit_vector_slice = unsafe{
slice::from_raw_parts(vector_ptr, BIT_VECTORS_PER_BLOCK)
};
//read the milestones for this block
let mut milestones:[u64;AminoBwtBlock::NUM_MILESTONES] = [0; AminoBwtBlock::NUM_MILESTONES];
for milestone_idx in 0..milestones.len(){
fm_index_file.read_exact(&mut u64_buffer)?;
milestones[milestone_idx] = u64::from_le_bytes(u64_buffer);
}
bwt_block_list[block_idx] = AminoBwtBlock::from_data( bit_vector_slice.try_into().unwrap(), milestones);
}
Bwt::Amino(bwt_block_list)
},
};
//write the prefix sums
let mut prefix_sums:Vec<u64> = vec![0; alphabet.cardinality() as usize+1];
for prefix_sum_idx in 0..prefix_sums.len() {
fm_index_file.read_exact(&mut u64_buffer)?;
let prefix_sum_value = u64::from_le_bytes(u64_buffer);
prefix_sums[prefix_sum_idx] = prefix_sum_value;
}
let mut sampled_suffix_array = CompressedSuffixArray::new(bwt_len as usize, suffix_array_compression_ratio);
//write the sampled suffix array
for suffix_array_idx in 0..suffix_array_num_words{
fm_index_file.read_exact(&mut u64_buffer)?;
let sa_value = u64::from_le_bytes(u64_buffer);
sampled_suffix_array.set_word(sa_value, suffix_array_idx);
}
let kmer_lookup_table = KmerLookupTable::from_file(fm_index_file, alphabet)?;
//read in the sequence index
let sequence_index = SequenceIndex::from_file(fm_index_file).expect("unable to read sequence index from file");
return Ok(FmIndex::from_elements(bwt, prefix_sums, sampled_suffix_array, kmer_lookup_table, bwt_len, version_number, sequence_index));
}
}
}
}