taglib.go

   1 // Copyright (C) The Lightning Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lightning
   6
   7 import (
   8         "bufio"
   9         "fmt"
  10         "io"
  11 )
  12
  13 const tagmapKeySize = 32
  14
  15 type tagmapKey uint64
  16
  17 type tagID int32
  18
  19 type tagInfo struct {
  20         id     tagID // 0-based position in input tagset
  21         tagseq []byte
  22 }
  23
  24 type tagLibrary struct {
  25         tagmap  map[tagmapKey]tagInfo
  26         keylen  int
  27         keymask tagmapKey
  28 }
  29
  30 func (taglib *tagLibrary) Load(rdr io.Reader) error {
  31         var seqs [][]byte
  32         scanner := bufio.NewScanner(rdr)
  33         for scanner.Scan() {
  34                 data := scanner.Bytes()
  35                 if len(data) > 0 && data[0] == '>' {
  36                 } else {
  37                         seqs = append(seqs, append([]byte(nil), data...))
  38                 }
  39         }
  40         if err := scanner.Err(); err != nil {
  41                 return err
  42         }
  43         return taglib.setTags(seqs)
  44 }
  45
  46 func (taglib *tagLibrary) FindAll(in *bufio.Reader, passthrough io.Writer, fn func(id tagID, pos, taglen int)) error {
  47         var window = make([]byte, 0, taglib.keylen*1000)
  48         var key tagmapKey
  49         for offset := 0; ; {
  50                 base, err := in.ReadByte()
  51                 if err == io.EOF {
  52                         return nil
  53                 } else if err != nil {
  54                         return err
  55                 } else if base == '\r' || base == '\n' {
  56                         if buf, err := in.Peek(1); err == nil && len(buf) > 0 && buf[0] == '>' {
  57                                 return nil
  58                         } else if err == io.EOF {
  59                                 return nil
  60                         }
  61                         continue
  62                 } else if base == '>' || base == ' ' {
  63                         return fmt.Errorf("unexpected char %q at offset %d in fasta data", base, offset)
  64                 }
  65
  66                 if passthrough != nil {
  67                         if base >= 'A' && base <= 'Z' {
  68                                 // lowercase for passthrough
  69                                 base += 'a' - 'A'
  70                         }
  71                         _, err = passthrough.Write([]byte{base})
  72                         if err != nil {
  73                                 return err
  74                         }
  75                 }
  76                 if !isbase[int(base)] {
  77                         // 'N' or various other chars meaning exact
  78                         // base not known
  79                         window = window[:0]
  80                         continue
  81                 }
  82                 offset++
  83                 window = append(window, base)
  84                 if len(window) == cap(window) {
  85                         copy(window, window[len(window)-taglib.keylen:])
  86                         window = window[:taglib.keylen]
  87                 }
  88                 key = ((key << 2) | twobit[int(base)]) & taglib.keymask
  89
  90                 if len(window) < taglib.keylen {
  91                         continue
  92                 } else if taginfo, ok := taglib.tagmap[key]; !ok {
  93                         continue
  94                 } else if len(taginfo.tagseq) != taglib.keylen {
  95                         return fmt.Errorf("assertion failed: len(%q) != keylen %d", taginfo.tagseq, taglib.keylen)
  96                 } else {
  97                         fn(taginfo.id, offset-taglib.keylen, len(taginfo.tagseq))
  98                         window = window[:0] // don't try to match overlapping tags
  99                 }
 100         }
 101         return nil
 102 }
 103
 104 func (taglib *tagLibrary) Len() int {
 105         return len(taglib.tagmap)
 106 }
 107
 108 func (taglib *tagLibrary) TagLen() int {
 109         return taglib.keylen
 110 }
 111
 112 var (
 113         twobit = func() []tagmapKey {
 114                 r := make([]tagmapKey, 256)
 115                 r[int('a')] = 0
 116                 r[int('A')] = 0
 117                 r[int('c')] = 1
 118                 r[int('C')] = 1
 119                 r[int('g')] = 2
 120                 r[int('G')] = 2
 121                 r[int('t')] = 3
 122                 r[int('T')] = 3
 123                 return r
 124         }()
 125         isbase = func() []bool {
 126                 r := make([]bool, 256)
 127                 r[int('a')] = true
 128                 r[int('A')] = true
 129                 r[int('c')] = true
 130                 r[int('C')] = true
 131                 r[int('g')] = true
 132                 r[int('G')] = true
 133                 r[int('t')] = true
 134                 r[int('T')] = true
 135                 return r
 136         }()
 137 )
 138
 139 func (taglib *tagLibrary) setTags(tags [][]byte) error {
 140         taglib.keylen = tagmapKeySize
 141         for _, t := range tags {
 142                 if l := len(t); taglib.keylen > l {
 143                         taglib.keylen = l
 144                 }
 145         }
 146         taglib.keymask = tagmapKey((1 << (taglib.keylen * 2)) - 1)
 147         taglib.tagmap = map[tagmapKey]tagInfo{}
 148         for i, tag := range tags {
 149                 var key tagmapKey
 150                 for _, b := range tag[:taglib.keylen] {
 151                         key = (key << 2) | twobit[int(b)]
 152                 }
 153                 if _, ok := taglib.tagmap[key]; ok {
 154                         return fmt.Errorf("first %d bytes of tag %d (%x) are not unique", taglib.keylen, i, key)
 155                 }
 156                 taglib.tagmap[key] = tagInfo{tagID(i), tag}
 157         }
 158         return nil
 159 }
 160
 161 func (taglib *tagLibrary) Tags() [][]byte {
 162         out := make([][]byte, len(taglib.tagmap))
 163         untwobit := []byte{'a', 'c', 'g', 't'}
 164         for key, info := range taglib.tagmap {
 165                 seq := make([]byte, taglib.keylen)
 166                 for i := len(seq) - 1; i >= 0; i-- {
 167                         seq[i] = untwobit[int(key)&3]
 168                         key = key >> 2
 169                 }
 170                 out[int(info.id)] = seq
 171         }
 172         return out
 173 }