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