src/test/bench/shootout-k-nucleotide.rs

   1 // Copyright 2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // ignore-android see #10393 #13206
  12
  13 use std::string::String;
  14 use std::slice;
  15 use std::sync::{Arc, Future};
  16
  17 static TABLE: [u8, ..4] = [ 'A' as u8, 'C' as u8, 'G' as u8, 'T' as u8 ];
  18 static TABLE_SIZE: uint = 2 << 16;
  19
  20 static OCCURRENCES: [&'static str, ..5] = [
  21     "GGT",
  22     "GGTA",
  23     "GGTATT",
  24     "GGTATTTTAATT",
  25     "GGTATTTTAATTTATAGT",
  26 ];
  27
  28 // Code implementation
  29
  30 #[deriving(PartialEq, PartialOrd, Ord, Eq)]
  31 struct Code(u64);
  32
  33 impl Code {
  34     fn hash(&self) -> u64 {
  35         let Code(ret) = *self;
  36         return ret;
  37     }
  38
  39     fn push_char(&self, c: u8) -> Code {
  40         Code((self.hash() << 2) + (pack_symbol(c) as u64))
  41     }
  42
  43     fn rotate(&self, c: u8, frame: uint) -> Code {
  44         Code(self.push_char(c).hash() & ((1u64 << (2 * frame)) - 1))
  45     }
  46
  47     fn pack(string: &str) -> Code {
  48         string.bytes().fold(Code(0u64), |a, b| a.push_char(b))
  49     }
  50
  51     fn unpack(&self, frame: uint) -> String {
  52         let mut key = self.hash();
  53         let mut result = Vec::new();
  54         for _ in range(0, frame) {
  55             result.push(unpack_symbol((key as u8) & 3));
  56             key >>= 2;
  57         }
  58
  59         result.reverse();
  60         String::from_utf8(result).unwrap()
  61     }
  62 }
  63
  64 // Hash table implementation
  65
  66 trait TableCallback {
  67     fn f(&self, entry: &mut Entry);
  68 }
  69
  70 struct BumpCallback;
  71
  72 impl TableCallback for BumpCallback {
  73     fn f(&self, entry: &mut Entry) {
  74         entry.count += 1;
  75     }
  76 }
  77
  78 struct PrintCallback(&'static str);
  79
  80 impl TableCallback for PrintCallback {
  81     fn f(&self, entry: &mut Entry) {
  82         let PrintCallback(s) = *self;
  83         println!("{}\t{}", entry.count as int, s);
  84     }
  85 }
  86
  87 struct Entry {
  88     code: Code,
  89     count: uint,
  90     next: Option<Box<Entry>>,
  91 }
  92
  93 struct Table {
  94     count: uint,
  95     items: Vec<Option<Box<Entry>>> }
  96
  97 struct Items<'a> {
  98     cur: Option<&'a Entry>,
  99     items: slice::Items<'a, Option<Box<Entry>>>,
 100 }
 101
 102 impl Table {
 103     fn new() -> Table {
 104         Table {
 105             count: 0,
 106             items: Vec::from_fn(TABLE_SIZE, |_| None),
 107         }
 108     }
 109
 110     fn search_remainder<C:TableCallback>(item: &mut Entry, key: Code, c: C) {
 111         match item.next {
 112             None => {
 113                 let mut entry = box Entry {
 114                     code: key,
 115                     count: 0,
 116                     next: None,
 117                 };
 118                 c.f(entry);
 119                 item.next = Some(entry);
 120             }
 121             Some(ref mut entry) => {
 122                 if entry.code == key {
 123                     c.f(*entry);
 124                     return;
 125                 }
 126
 127                 Table::search_remainder(*entry, key, c)
 128             }
 129         }
 130     }
 131
 132     fn lookup<C:TableCallback>(&mut self, key: Code, c: C) {
 133         let index = key.hash() % (TABLE_SIZE as u64);
 134
 135         {
 136             if self.items.get(index as uint).is_none() {
 137                 let mut entry = box Entry {
 138                     code: key,
 139                     count: 0,
 140                     next: None,
 141                 };
 142                 c.f(entry);
 143                 *self.items.get_mut(index as uint) = Some(entry);
 144                 return;
 145             }
 146         }
 147
 148         {
 149             let entry = &mut *self.items.get_mut(index as uint).get_mut_ref();
 150             if entry.code == key {
 151                 c.f(*entry);
 152                 return;
 153             }
 154
 155             Table::search_remainder(*entry, key, c)
 156         }
 157     }
 158
 159     fn iter<'a>(&'a self) -> Items<'a> {
 160         Items { cur: None, items: self.items.iter() }
 161     }
 162 }
 163
 164 impl<'a> Iterator<&'a Entry> for Items<'a> {
 165     fn next(&mut self) -> Option<&'a Entry> {
 166         let ret = match self.cur {
 167             None => {
 168                 let i;
 169                 loop {
 170                     match self.items.next() {
 171                         None => return None,
 172                         Some(&None) => {}
 173                         Some(&Some(ref a)) => { i = &**a; break }
 174                     }
 175                 }
 176                 self.cur = Some(&*i);
 177                 &*i
 178             }
 179             Some(c) => c
 180         };
 181         match ret.next {
 182             None => { self.cur = None; }
 183             Some(ref next) => { self.cur = Some(&**next); }
 184         }
 185         return Some(ret);
 186     }
 187 }
 188
 189 // Main program
 190
 191 fn pack_symbol(c: u8) -> u8 {
 192     match c as char {
 193         'A' => 0,
 194         'C' => 1,
 195         'G' => 2,
 196         'T' => 3,
 197         _ => fail!("{}", c as char),
 198     }
 199 }
 200
 201 fn unpack_symbol(c: u8) -> u8 {
 202     TABLE[c as uint]
 203 }
 204
 205 fn generate_frequencies(mut input: &[u8], frame: uint) -> Table {
 206     let mut frequencies = Table::new();
 207     if input.len() < frame { return frequencies; }
 208     let mut code = Code(0);
 209
 210     // Pull first frame.
 211     for _ in range(0, frame) {
 212         code = code.push_char(input[0]);
 213         input = input.slice_from(1);
 214     }
 215     frequencies.lookup(code, BumpCallback);
 216
 217     while input.len() != 0 && input[0] != ('>' as u8) {
 218         code = code.rotate(input[0], frame);
 219         frequencies.lookup(code, BumpCallback);
 220         input = input.slice_from(1);
 221     }
 222     frequencies
 223 }
 224
 225 fn print_frequencies(frequencies: &Table, frame: uint) {
 226     let mut vector = Vec::new();
 227     for entry in frequencies.iter() {
 228         vector.push((entry.count, entry.code));
 229     }
 230     vector.as_mut_slice().sort();
 231
 232     let mut total_count = 0;
 233     for &(count, _) in vector.iter() {
 234         total_count += count;
 235     }
 236
 237     for &(count, key) in vector.iter().rev() {
 238         println!("{} {:.3f}",
 239                  key.unpack(frame).as_slice(),
 240                  (count as f32 * 100.0) / (total_count as f32));
 241     }
 242     println!("");
 243 }
 244
 245 fn print_occurrences(frequencies: &mut Table, occurrence: &'static str) {
 246     frequencies.lookup(Code::pack(occurrence), PrintCallback(occurrence))
 247 }
 248
 249 fn get_sequence<R: Buffer>(r: &mut R, key: &str) -> Vec<u8> {
 250     let mut res = Vec::new();
 251     for l in r.lines().map(|l| l.ok().unwrap())
 252         .skip_while(|l| key != l.as_slice().slice_to(key.len())).skip(1)
 253     {
 254         res.push_all(l.as_slice().trim().as_bytes());
 255     }
 256     for b in res.mut_iter() {
 257         *b = b.to_ascii().to_upper().to_byte();
 258     }
 259     res
 260 }
 261
 262 fn main() {
 263     let input = if std::os::getenv("RUST_BENCH").is_some() {
 264         let fd = std::io::File::open(&Path::new("shootout-k-nucleotide.data"));
 265         get_sequence(&mut std::io::BufferedReader::new(fd), ">THREE")
 266     } else {
 267         get_sequence(&mut std::io::stdin(), ">THREE")
 268     };
 269     let input = Arc::new(input);
 270
 271     let nb_freqs: Vec<(uint, Future<Table>)> = range(1u, 3).map(|i| {
 272         let input = input.clone();
 273         (i, Future::spawn(proc() generate_frequencies(input.as_slice(), i)))
 274     }).collect();
 275     let occ_freqs: Vec<Future<Table>> = OCCURRENCES.iter().map(|&occ| {
 276         let input = input.clone();
 277         Future::spawn(proc() generate_frequencies(input.as_slice(), occ.len()))
 278     }).collect();
 279
 280     for (i, freq) in nb_freqs.move_iter() {
 281         print_frequencies(&freq.unwrap(), i);
 282     }
 283     for (&occ, freq) in OCCURRENCES.iter().zip(occ_freqs.move_iter()) {
 284         print_occurrences(&mut freq.unwrap(), occ);
 285     }
 286 }