Use 2bit-packed uint64 as tagmapKey.

Arvados-DCO-1.1-Signed-off-by: Tom Clegg <tom@tomclegg.ca>

diff --git a/taglib.go b/taglib.go
index 91d4d4df299f0b6fe4ff8dfe2d6576cccfa5aed9..0de37089a791d1914ce560717dfe950fc156e0c4 100644 (file)
--- a/taglib.go
+++ b/taglib.go
@@ -7,9 +7,9 @@ import (
        "io"
 )
 
-const tagmapKeySize = 24
+const tagmapKeySize = 32
 
-type tagmapKey [tagmapKeySize]byte
+type tagmapKey uint64
 
 type tagID int32
 
@@ -19,8 +19,9 @@ type tagInfo struct {
 }
 
 type tagLibrary struct {
-       tagmap map[tagmapKey]tagInfo
-       keylen int
+       tagmap  map[tagmapKey]tagInfo
+       keylen  int
+       keymask tagmapKey
 }
 
 func (taglib *tagLibrary) Load(rdr io.Reader) error {
@@ -41,16 +42,25 @@ func (taglib *tagLibrary) Load(rdr io.Reader) error {
 
 func (taglib *tagLibrary) FindAll(buf []byte, fn func(id tagID, pos, taglen int)) {
        var key tagmapKey
-       for i := 0; i <= len(buf)-taglib.keylen; i++ {
-               copy(key[:taglib.keylen], buf[i:])
-               if taginfo, ok := taglib.tagmap[key]; !ok {
+       valid := 0 // if valid < taglib.keylen, key has "no data" zeroes that are otherwise indistinguishable from "A"
+       for i, base := range buf {
+               if !isbase[int(base)] {
+                       valid = 0
                        continue
-               } else if len(taginfo.tagseq) > taglib.keylen && (len(buf) < i+len(taginfo.tagseq) || !bytes.Equal(taginfo.tagseq, buf[i:i+len(taginfo.tagseq)])) {
+               } else {
+                       valid++
+               }
+               key = ((key << 2) | twobit[int(base)]) & taglib.keymask
+               if valid < taglib.keylen {
+                       continue
+               } else if taginfo, ok := taglib.tagmap[key]; !ok {
+                       continue
+               } 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)])) {
                        // key portion matches, but not the entire tag
                        continue
                } else {
-                       fn(taginfo.id, i, len(taginfo.tagseq))
-                       i += len(taginfo.tagseq) - 1 // don't try to match overlapping tags
+                       fn(taginfo.id, tagstart, len(taginfo.tagseq))
+                       valid = 0 // don't try to match overlapping tags
                }
        }
 }
@@ -59,6 +69,33 @@ func (taglib *tagLibrary) Len() int {
        return len(taglib.tagmap)
 }
 
+var (
+       twobit = func() []tagmapKey {
+               r := make([]tagmapKey, 256)
+               r[int('a')] = 0
+               r[int('A')] = 0
+               r[int('c')] = 1
+               r[int('C')] = 1
+               r[int('g')] = 2
+               r[int('G')] = 2
+               r[int('t')] = 3
+               r[int('T')] = 3
+               return r
+       }()
+       isbase = func() []bool {
+               r := make([]bool, 256)
+               r[int('a')] = true
+               r[int('A')] = true
+               r[int('c')] = true
+               r[int('C')] = true
+               r[int('g')] = true
+               r[int('G')] = true
+               r[int('t')] = true
+               r[int('T')] = true
+               return r
+       }()
+)
+
 func (taglib *tagLibrary) setTags(tags [][]byte) error {
        taglib.keylen = tagmapKeySize
        for _, t := range tags {
@@ -66,15 +103,17 @@ func (taglib *tagLibrary) setTags(tags [][]byte) error {
                        taglib.keylen = l
                }
        }
+       taglib.keymask = tagmapKey((1 << (taglib.keylen * 2)) - 1)
        taglib.tagmap = map[tagmapKey]tagInfo{}
-       for i, t := range tags {
-               t = bytes.ToLower(t)
+       for i, tag := range tags {
                var key tagmapKey
-               copy(key[:], t[:taglib.keylen])
+               for _, b := range tag[:taglib.keylen] {
+                       key = (key << 2) | twobit[int(b)]
+               }
                if _, ok := taglib.tagmap[key]; ok {
-                       return fmt.Errorf("first %d bytes of tag %d (%s) are not unique", taglib.keylen, i, key)
+                       return fmt.Errorf("first %d bytes of tag %d (%x) are not unique", taglib.keylen, i, key)
                }
-               taglib.tagmap[key] = tagInfo{tagID(i), t}
+               taglib.tagmap[key] = tagInfo{tagID(i), tag}
        }
        return nil
 }