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                         _, err = passthrough.Write([]byte{base})
  68                         if err != nil {
  69                                 return err
  70                         }
  71                 }
  72                 if !isbase[int(base)] {
  73                         // 'N' or various other chars meaning exact
  74                         // base not known
  75                         window = window[:0]
  76                         continue
  77                 }
  78                 offset++
  79                 window = append(window, base)
  80                 if len(window) == cap(window) {
  81                         copy(window, window[len(window)-taglib.keylen:])
  82                         window = window[:taglib.keylen]
  83                 }
  84                 key = ((key << 2) | twobit[int(base)]) & taglib.keymask
  85
  86                 if len(window) < taglib.keylen {
  87                         continue
  88                 } else if taginfo, ok := taglib.tagmap[key]; !ok {
  89                         continue
  90                 } else if len(taginfo.tagseq) != taglib.keylen {
  91                         return fmt.Errorf("assertion failed: len(%q) != keylen %d", taginfo.tagseq, taglib.keylen)
  92                 } else {
  93                         fn(taginfo.id, offset-taglib.keylen, len(taginfo.tagseq))
  94                         window = window[:0] // don't try to match overlapping tags
  95                 }
  96         }
  97         return nil
  98 }
  99
 100 func (taglib *tagLibrary) Len() int {
 101         return len(taglib.tagmap)
 102 }
 103
 104 func (taglib *tagLibrary) TagLen() int {
 105         return taglib.keylen
 106 }
 107
 108 var (
 109         twobit = func() []tagmapKey {
 110                 r := make([]tagmapKey, 256)
 111                 r[int('a')] = 0
 112                 r[int('A')] = 0
 113                 r[int('c')] = 1
 114                 r[int('C')] = 1
 115                 r[int('g')] = 2
 116                 r[int('G')] = 2
 117                 r[int('t')] = 3
 118                 r[int('T')] = 3
 119                 return r
 120         }()
 121         isbase = func() []bool {
 122                 r := make([]bool, 256)
 123                 r[int('a')] = true
 124                 r[int('A')] = true
 125                 r[int('c')] = true
 126                 r[int('C')] = true
 127                 r[int('g')] = true
 128                 r[int('G')] = true
 129                 r[int('t')] = true
 130                 r[int('T')] = true
 131                 return r
 132         }()
 133 )
 134
 135 func (taglib *tagLibrary) setTags(tags [][]byte) error {
 136         taglib.keylen = tagmapKeySize
 137         for _, t := range tags {
 138                 if l := len(t); taglib.keylen > l {
 139                         taglib.keylen = l
 140                 }
 141         }
 142         taglib.keymask = tagmapKey((1 << (taglib.keylen * 2)) - 1)
 143         taglib.tagmap = map[tagmapKey]tagInfo{}
 144         for i, tag := range tags {
 145                 var key tagmapKey
 146                 for _, b := range tag[:taglib.keylen] {
 147                         key = (key << 2) | twobit[int(b)]
 148                 }
 149                 if _, ok := taglib.tagmap[key]; ok {
 150                         return fmt.Errorf("first %d bytes of tag %d (%x) are not unique", taglib.keylen, i, key)
 151                 }
 152                 taglib.tagmap[key] = tagInfo{tagID(i), tag}
 153         }
 154         return nil
 155 }
 156
 157 func (taglib *tagLibrary) Tags() [][]byte {
 158         out := make([][]byte, len(taglib.tagmap))
 159         untwobit := []byte{'a', 'c', 'g', 't'}
 160         for key, info := range taglib.tagmap {
 161                 seq := make([]byte, taglib.keylen)
 162                 for i := len(seq) - 1; i >= 0; i-- {
 163                         seq[i] = untwobit[int(key)&3]
 164                         key = key >> 2
 165                 }
 166                 out[int(info.id)] = seq
 167         }
 168         return out
 169 }