taglib.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "fmt"
   7         "io"
   8 )
   9
  10 const tagmapKeySize = 32
  11
  12 type tagmapKey uint64
  13
  14 type tagID int32
  15
  16 type tagInfo struct {
  17         id     tagID // 0-based position in input tagset
  18         tagseq []byte
  19 }
  20
  21 type tagLibrary struct {
  22         tagmap  map[tagmapKey]tagInfo
  23         keylen  int
  24         keymask tagmapKey
  25 }
  26
  27 func (taglib *tagLibrary) Load(rdr io.Reader) error {
  28         var seqs [][]byte
  29         scanner := bufio.NewScanner(rdr)
  30         for scanner.Scan() {
  31                 data := scanner.Bytes()
  32                 if len(data) > 0 && data[0] == '>' {
  33                 } else {
  34                         seqs = append(seqs, append([]byte(nil), data...))
  35                 }
  36         }
  37         if err := scanner.Err(); err != nil {
  38                 return err
  39         }
  40         return taglib.setTags(seqs)
  41 }
  42
  43 func (taglib *tagLibrary) FindAll(buf []byte, fn func(id tagID, pos, taglen int)) {
  44         var key tagmapKey
  45         valid := 0 // if valid < taglib.keylen, key has "no data" zeroes that are otherwise indistinguishable from "A"
  46         for i, base := range buf {
  47                 if !isbase[int(base)] {
  48                         valid = 0
  49                         continue
  50                 }
  51                 key = ((key << 2) | twobit[int(base)]) & taglib.keymask
  52                 valid++
  53
  54                 if valid < taglib.keylen {
  55                         continue
  56                 } else if taginfo, ok := taglib.tagmap[key]; !ok {
  57                         continue
  58                 } else if tagstart := i - taglib.keylen + 1; len(taginfo.tagseq) > taglib.keylen && (len(buf) < i+len(taginfo.tagseq) || !bytes.Equal(taginfo.tagseq, buf[i:i+len(taginfo.tagseq)])) {
  59                         // key portion matches, but not the entire tag
  60                         continue
  61                 } else {
  62                         fn(taginfo.id, tagstart, len(taginfo.tagseq))
  63                         valid = 0 // don't try to match overlapping tags
  64                 }
  65         }
  66 }
  67
  68 func (taglib *tagLibrary) Len() int {
  69         return len(taglib.tagmap)
  70 }
  71
  72 var (
  73         twobit = func() []tagmapKey {
  74                 r := make([]tagmapKey, 256)
  75                 r[int('a')] = 0
  76                 r[int('A')] = 0
  77                 r[int('c')] = 1
  78                 r[int('C')] = 1
  79                 r[int('g')] = 2
  80                 r[int('G')] = 2
  81                 r[int('t')] = 3
  82                 r[int('T')] = 3
  83                 return r
  84         }()
  85         isbase = func() []bool {
  86                 r := make([]bool, 256)
  87                 r[int('a')] = true
  88                 r[int('A')] = true
  89                 r[int('c')] = true
  90                 r[int('C')] = true
  91                 r[int('g')] = true
  92                 r[int('G')] = true
  93                 r[int('t')] = true
  94                 r[int('T')] = true
  95                 return r
  96         }()
  97 )
  98
  99 func (taglib *tagLibrary) setTags(tags [][]byte) error {
 100         taglib.keylen = tagmapKeySize
 101         for _, t := range tags {
 102                 if l := len(t); taglib.keylen > l {
 103                         taglib.keylen = l
 104                 }
 105         }
 106         taglib.keymask = tagmapKey((1 << (taglib.keylen * 2)) - 1)
 107         taglib.tagmap = map[tagmapKey]tagInfo{}
 108         for i, tag := range tags {
 109                 var key tagmapKey
 110                 for _, b := range tag[:taglib.keylen] {
 111                         key = (key << 2) | twobit[int(b)]
 112                 }
 113                 if _, ok := taglib.tagmap[key]; ok {
 114                         return fmt.Errorf("first %d bytes of tag %d (%x) are not unique", taglib.keylen, i, key)
 115                 }
 116                 taglib.tagmap[key] = tagInfo{tagID(i), tag}
 117         }
 118         return nil
 119 }
 120
 121 func (taglib *tagLibrary) Tags() [][]byte {
 122         out := make([][]byte, len(taglib.tagmap))
 123         untwobit := []byte{'a', 'c', 'g', 't'}
 124         for key, info := range taglib.tagmap {
 125                 seq := make([]byte, taglib.keylen)
 126                 for i := len(seq) - 1; i >= 0; i-- {
 127                         seq[i] = untwobit[int(key)&3]
 128                         key = key >> 2
 129                 }
 130                 out[int(info.id)] = seq
 131         }
 132         return out
 133 }