taglib.go

   1 package lightning
   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 func (taglib *tagLibrary) TagLen() int {
  73         return taglib.keylen
  74 }
  75
  76 var (
  77         twobit = func() []tagmapKey {
  78                 r := make([]tagmapKey, 256)
  79                 r[int('a')] = 0
  80                 r[int('A')] = 0
  81                 r[int('c')] = 1
  82                 r[int('C')] = 1
  83                 r[int('g')] = 2
  84                 r[int('G')] = 2
  85                 r[int('t')] = 3
  86                 r[int('T')] = 3
  87                 return r
  88         }()
  89         isbase = func() []bool {
  90                 r := make([]bool, 256)
  91                 r[int('a')] = true
  92                 r[int('A')] = true
  93                 r[int('c')] = true
  94                 r[int('C')] = true
  95                 r[int('g')] = true
  96                 r[int('G')] = true
  97                 r[int('t')] = true
  98                 r[int('T')] = true
  99                 return r
 100         }()
 101 )
 102
 103 func (taglib *tagLibrary) setTags(tags [][]byte) error {
 104         taglib.keylen = tagmapKeySize
 105         for _, t := range tags {
 106                 if l := len(t); taglib.keylen > l {
 107                         taglib.keylen = l
 108                 }
 109         }
 110         taglib.keymask = tagmapKey((1 << (taglib.keylen * 2)) - 1)
 111         taglib.tagmap = map[tagmapKey]tagInfo{}
 112         for i, tag := range tags {
 113                 var key tagmapKey
 114                 for _, b := range tag[:taglib.keylen] {
 115                         key = (key << 2) | twobit[int(b)]
 116                 }
 117                 if _, ok := taglib.tagmap[key]; ok {
 118                         return fmt.Errorf("first %d bytes of tag %d (%x) are not unique", taglib.keylen, i, key)
 119                 }
 120                 taglib.tagmap[key] = tagInfo{tagID(i), tag}
 121         }
 122         return nil
 123 }
 124
 125 func (taglib *tagLibrary) Tags() [][]byte {
 126         out := make([][]byte, len(taglib.tagmap))
 127         untwobit := []byte{'a', 'c', 'g', 't'}
 128         for key, info := range taglib.tagmap {
 129                 seq := make([]byte, taglib.keylen)
 130                 for i := len(seq) - 1; i >= 0; i-- {
 131                         seq[i] = untwobit[int(key)&3]
 132                         key = key >> 2
 133                 }
 134                 out[int(info.id)] = seq
 135         }
 136         return out
 137 }