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-pretty
  12
  13 use std::str;
  14 use std::slice;
  15
  16 static TABLE: [u8, ..4] = [ 'A' as u8, 'C' as u8, 'G' as u8, 'T' as u8 ];
  17 static TABLE_SIZE: uint = 2 << 16;
  18
  19 static OCCURRENCES: [&'static str, ..5] = [
  20     "GGT",
  21     "GGTA",
  22     "GGTATT",
  23     "GGTATTTTAATT",
  24     "GGTATTTTAATTTATAGT",
  25 ];
  26
  27 // Code implementation
  28
  29 #[deriving(Eq, Ord, TotalOrd, TotalEq)]
  30 struct Code(u64);
  31
  32 impl Code {
  33     fn hash(&self) -> u64 {
  34         let Code(ret) = *self;
  35         return ret;
  36     }
  37
  38     fn push_char(&self, c: u8) -> Code {
  39         Code((self.hash() << 2) + (pack_symbol(c) as u64))
  40     }
  41
  42     fn rotate(&self, c: u8, frame: i32) -> Code {
  43         Code(self.push_char(c).hash() & ((1u64 << (2 * (frame as u64))) - 1))
  44     }
  45
  46     fn pack(string: &str) -> Code {
  47         string.bytes().fold(Code(0u64), |a, b| a.push_char(b))
  48     }
  49
  50     // FIXME: Inefficient.
  51     fn unpack(&self, frame: i32) -> ~str {
  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         str::from_utf8_owned(result.move_iter().collect()).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: i32,
  90     next: Option<~Entry>,
  91 }
  92
  93 struct Table {
  94     count: i32,
  95     items: Vec<Option<~Entry>> }
  96
  97 struct Items<'a> {
  98     cur: Option<&'a Entry>,
  99     items: slice::Items<'a, Option<~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 = ~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 = ~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' | 'A' => 0,
 194         'c' | 'C' => 1,
 195         'g' | 'G' => 2,
 196         't' | 'T' => 3,
 197         _ => fail!("{}", c as char),
 198     }
 199 }
 200
 201 fn unpack_symbol(c: u8) -> u8 {
 202     TABLE[c]
 203 }
 204
 205 fn next_char<'a>(mut buf: &'a [u8]) -> &'a [u8] {
 206     loop {
 207         buf = buf.slice(1, buf.len());
 208         if buf.len() == 0 {
 209             break;
 210         }
 211         if buf[0] != (' ' as u8) && buf[0] != ('\t' as u8) &&
 212                 buf[0] != ('\n' as u8) && buf[0] != 0 {
 213             break;
 214         }
 215     }
 216     buf
 217 }
 218
 219 fn generate_frequencies(frequencies: &mut Table,
 220                         mut input: &[u8],
 221                         frame: i32) {
 222     let mut code = Code(0);
 223
 224     // Pull first frame.
 225     for _ in range(0, frame) {
 226         code = code.push_char(input[0]);
 227         input = next_char(input);
 228     }
 229     frequencies.lookup(code, BumpCallback);
 230
 231     while input.len() != 0 && input[0] != ('>' as u8) {
 232         code = code.rotate(input[0], frame);
 233         frequencies.lookup(code, BumpCallback);
 234         input = next_char(input);
 235     }
 236 }
 237
 238 fn print_frequencies(frequencies: &Table, frame: i32) {
 239     let mut vector = Vec::new();
 240     for entry in frequencies.iter() {
 241         vector.push((entry.code, entry.count));
 242     }
 243     vector.as_mut_slice().sort();
 244
 245     let mut total_count = 0;
 246     for &(_, count) in vector.iter() {
 247         total_count += count;
 248     }
 249
 250     for &(key, count) in vector.iter() {
 251         println!("{} {:.3f}",
 252                  key.unpack(frame),
 253                  (count as f32 * 100.0) / (total_count as f32));
 254     }
 255 }
 256
 257 fn print_occurrences(frequencies: &mut Table, occurrence: &'static str) {
 258     frequencies.lookup(Code::pack(occurrence), PrintCallback(occurrence))
 259 }
 260
 261 fn main() {
 262     let input = include_str!("shootout-k-nucleotide.data");
 263     let pos = input.find_str(">THREE").unwrap();
 264     let pos2 = pos + input.slice_from(pos).find_str("\n").unwrap();
 265     let input = input.slice_from(pos2 + 1).as_bytes();
 266
 267     let mut frequencies = Table::new();
 268     generate_frequencies(&mut frequencies, input, 1);
 269     print_frequencies(&frequencies, 1);
 270
 271     frequencies = Table::new();
 272     generate_frequencies(&mut frequencies, input, 2);
 273     print_frequencies(&frequencies, 2);
 274
 275     for occurrence in OCCURRENCES.iter() {
 276         frequencies = Table::new();
 277         generate_frequencies(&mut frequencies,
 278                              input,
 279                              occurrence.len() as i32);
 280         print_occurrences(&mut frequencies, *occurrence);
 281     }
 282 }