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
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".to_owned(),
/// 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 milestone in block.milestones() {
fm_index_file.write_all(&milestone.to_le_bytes())?;
}
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())?;
}
}
}
}
crate::bwt::Bwt::Amino(vec) => {
for block in vec.iter() {
for milestone in block.milestones() {
fm_index_file.write_all(&milestone.to_le_bytes())?;
}
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())?;
}
}
}
}
}
//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
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".to_owned(),
/// 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) -> Vec<u64> {
match self.version_number() {
_ => {
let alphabet_idx: u64 = match self.bwt() {
Bwt::Nucleotide(_) => 0,
Bwt::Amino(_) => 1,
};
//Matches version 1
let mut header: Vec<u64> = vec![0; 8];
header[0] = self.version_number();
header[1] = self.suffix_array_compression_ratio();
header[2] = self.bwt_len();
header[3] = alphabet_idx;
//the remaining 32 bytes are left empty for now
header
}
}
}
/// 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 compressed_suffix_array_len = (bwt_len / suffix_array_compression_ratio) as usize;
let num_bwt_blocks = (bwt_len as usize).div_ceil(Bwt::NUM_SYMBOLS_PER_BLOCK as usize);
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(){
//read the milestones for this block
let mut milestones:[u64;NucleotideBwtBlock::NUM_MILESTONES] = [0; NucleotideBwtBlock::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);
}
//read the bit vectors for this block
let mut bit_vector_buffer:[u8;32*NucleotideBwtBlock::NUM_BIT_VECTORS] = [0;32*NucleotideBwtBlock::NUM_BIT_VECTORS];
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, NucleotideBwtBlock::NUM_BIT_VECTORS)
};
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(){
//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);
}
//read the bit vectors for this block
let mut bit_vector_buffer:[u8;32*AminoBwtBlock::NUM_BIT_VECTORS] = [0;32*AminoBwtBlock::NUM_BIT_VECTORS];
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, AminoBwtBlock::NUM_BIT_VECTORS)
};
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)?;
prefix_sums[prefix_sum_idx] = u64::from_le_bytes(u64_buffer);
}
let mut sampled_suffix_array = CompressedSuffixArray::new(compressed_suffix_array_len, suffix_array_compression_ratio);
//write the sampled suffix array
for suffix_array_idx in 0..compressed_suffix_array_len{
fm_index_file.read_exact(&mut u64_buffer)?;
sampled_suffix_array.set_value(u64::from_le_bytes(u64_buffer), 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));
}
}
}
}