19073: Fix dup tag detection.
[lightning.git] / tilelib.go
index e0d0e1bc4f4a603b7ec8e3801a26d71031e32d9d..df6cb73b3be2b9e16418f6d5eb7c5135b7686efe 100644 (file)
@@ -1,3 +1,7 @@
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package lightning
 
 import (
@@ -57,16 +61,23 @@ type tileLibrary struct {
        retainNoCalls       bool
        skipOOO             bool
        retainTileSequences bool
+       useDups             bool
 
        taglib         *tagLibrary
        variant        [][][blake2b.Size256]byte
        refseqs        map[string]map[string][]tileLibRef
        compactGenomes map[string][]tileVariantID
-       // count [][]int
-       seq      map[[blake2b.Size256]byte][]byte
-       variants int64
+       seq2           map[[2]byte]map[[blake2b.Size256]byte][]byte
+       seq2lock       map[[2]byte]sync.Locker
+       variants       int64
        // if non-nil, write out any tile variants added while tiling
        encoder *gob.Encoder
+       // set Ref flag when writing new variants to encoder
+       encodeRef bool
+
+       onAddTileVariant func(libref tileLibRef, hash [blake2b.Size256]byte, seq []byte) error
+       onAddGenome      func(CompactGenome) error
+       onAddRefseq      func(CompactSequence) error
 
        mtx   sync.RWMutex
        vlock []sync.Locker
@@ -112,12 +123,12 @@ func (tilelib *tileLibrary) loadTileVariants(tvs []TileVariant, variantmap map[t
        for _, tv := range tvs {
                // Assign a new variant ID (unique across all inputs)
                // for each input variant.
-               variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+               variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence, tv.Ref).Variant
        }
        return nil
 }
 
-func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap map[tileLibRef]tileVariantID, onLoadGenome func(CompactGenome)) error {
+func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap map[tileLibRef]tileVariantID) error {
        log.Debugf("loadCompactGenomes: %d", len(cgs))
        var wg sync.WaitGroup
        errs := make(chan error, 1)
@@ -143,11 +154,18 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
                                        }
                                        return
                                }
-                               log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
+                               // log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
                                cg.Variants[i] = newvariant
                        }
-                       if onLoadGenome != nil {
-                               onLoadGenome(cg)
+                       if tilelib.onAddGenome != nil {
+                               err := tilelib.onAddGenome(cg)
+                               if err != nil {
+                                       select {
+                                       case errs <- err:
+                                       default:
+                                       }
+                                       return
+                               }
                        }
                        if tilelib.encoder != nil {
                                err := tilelib.encoder.Encode(LibraryEntry{
@@ -174,8 +192,9 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
 }
 
 func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, variantmap map[tileLibRef]tileVariantID) error {
-       log.Debugf("loadCompactSequences: %d", len(cseqs))
+       log.Infof("loadCompactSequences: %d todo", len(cseqs))
        for _, cseq := range cseqs {
+               log.Infof("loadCompactSequences: checking %s", cseq.Name)
                for _, tseq := range cseq.TileSequences {
                        for i, libref := range tseq {
                                if libref.Variant == 0 {
@@ -198,6 +217,13 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
                                return err
                        }
                }
+               if tilelib.onAddRefseq != nil {
+                       err := tilelib.onAddRefseq(cseq)
+                       if err != nil {
+                               return err
+                       }
+               }
+               log.Infof("loadCompactSequences: checking %s done", cseq.Name)
        }
        tilelib.mtx.Lock()
        defer tilelib.mtx.Unlock()
@@ -207,52 +233,56 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
        for _, cseq := range cseqs {
                tilelib.refseqs[cseq.Name] = cseq.TileSequences
        }
+       log.Info("loadCompactSequences: done")
        return nil
 }
 
-func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGenome func(CompactGenome)) error {
+func allFiles(path string, re *regexp.Regexp) ([]string, error) {
        var files []string
-       var walk func(string) error
-       walk = func(path string) error {
-               f, err := open(path)
-               if err != nil {
-                       return err
-               }
-               defer f.Close()
-               fis, err := f.Readdir(-1)
-               if err != nil {
-                       files = append(files, path)
-                       return nil
-               }
-               for _, fi := range fis {
-                       if fi.Name() == "." || fi.Name() == ".." {
-                               continue
-                       } else if child := path + "/" + fi.Name(); fi.IsDir() {
-                               err = walk(child)
-                               if err != nil {
-                                       return err
-                               }
-                       } else if strings.HasSuffix(child, ".gob") || strings.HasSuffix(child, ".gob.gz") {
-                               files = append(files, child)
+       f, err := open(path)
+       if err != nil {
+               return nil, err
+       }
+       defer f.Close()
+       fis, err := f.Readdir(-1)
+       if err != nil {
+               return []string{path}, nil
+       }
+       for _, fi := range fis {
+               if fi.Name() == "." || fi.Name() == ".." {
+                       continue
+               } else if child := path + "/" + fi.Name(); fi.IsDir() {
+                       add, err := allFiles(child, re)
+                       if err != nil {
+                               return nil, err
                        }
+                       files = append(files, add...)
+               } else if re == nil || re.MatchString(child) {
+                       files = append(files, child)
                }
-               return nil
        }
+       sort.Strings(files)
+       return files, nil
+}
+
+var matchGobFile = regexp.MustCompile(`\.gob(\.gz)?$`)
+
+func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string) error {
        log.Infof("LoadDir: walk dir %s", path)
-       err := walk(path)
+       files, err := allFiles(path, matchGobFile)
        if err != nil {
                return err
        }
        ctx, cancel := context.WithCancel(ctx)
        defer cancel()
        var mtx sync.Mutex
-       cgs := []CompactGenome{}
-       cseqs := []CompactSequence{}
-       variantmap := map[tileLibRef]tileVariantID{}
+       allcgs := make([][]CompactGenome, len(files))
+       allcseqs := make([][]CompactSequence, len(files))
+       allvariantmap := map[tileLibRef]tileVariantID{}
        errs := make(chan error, len(files))
        log.Infof("LoadDir: read %d files", len(files))
-       for _, path := range files {
-               path := path
+       for fileno, path := range files {
+               fileno, path := fileno, path
                go func() {
                        f, err := open(path)
                        if err != nil {
@@ -261,7 +291,11 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
                        }
                        defer f.Close()
                        defer log.Infof("LoadDir: finished reading %s", path)
-                       errs <- DecodeLibrary(f, strings.HasSuffix(path, ".gz"), func(ent *LibraryEntry) error {
+
+                       var variantmap = map[tileLibRef]tileVariantID{}
+                       var cgs []CompactGenome
+                       var cseqs []CompactSequence
+                       err = DecodeLibrary(f, strings.HasSuffix(path, ".gz"), func(ent *LibraryEntry) error {
                                if ctx.Err() != nil {
                                        return ctx.Err()
                                }
@@ -278,19 +312,21 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
                                        }
                                        mtx.Unlock()
                                }
-                               variantmapadd := map[tileLibRef]tileVariantID{}
                                for _, tv := range ent.TileVariants {
-                                       variantmapadd[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+                                       variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence, tv.Ref).Variant
                                }
-                               mtx.Lock()
                                cgs = append(cgs, ent.CompactGenomes...)
                                cseqs = append(cseqs, ent.CompactSequences...)
-                               for k, v := range variantmapadd {
-                                       variantmap[k] = v
-                               }
-                               mtx.Unlock()
                                return nil
                        })
+                       allcgs[fileno] = cgs
+                       allcseqs[fileno] = cseqs
+                       mtx.Lock()
+                       defer mtx.Unlock()
+                       for k, v := range variantmap {
+                               allvariantmap[k] = v
+                       }
+                       errs <- err
                }()
        }
        for range files {
@@ -299,22 +335,34 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
                        return err
                }
        }
+
        log.Info("LoadDir: loadCompactGenomes")
-       err = tilelib.loadCompactGenomes(cgs, variantmap, onLoadGenome)
+       var flatcgs []CompactGenome
+       for _, cgs := range allcgs {
+               flatcgs = append(flatcgs, cgs...)
+       }
+       err = tilelib.loadCompactGenomes(flatcgs, allvariantmap)
        if err != nil {
                return err
        }
+
        log.Info("LoadDir: loadCompactSequences")
-       err = tilelib.loadCompactSequences(cseqs, variantmap)
+       var flatcseqs []CompactSequence
+       for _, cseqs := range allcseqs {
+               flatcseqs = append(flatcseqs, cseqs...)
+       }
+       err = tilelib.loadCompactSequences(flatcseqs, allvariantmap)
        if err != nil {
                return err
        }
+
        log.Info("LoadDir done")
        return nil
 }
 
 func (tilelib *tileLibrary) WriteDir(dir string) error {
-       nfiles := 128
+       ntilefiles := 128
+       nfiles := ntilefiles + len(tilelib.refseqs)
        files := make([]*os.File, nfiles)
        for i := range files {
                f, err := os.OpenFile(fmt.Sprintf("%s/library.%04d.gob.gz", dir, i), os.O_CREATE|os.O_WRONLY, 0666)
@@ -338,6 +386,18 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
                encoders[i] = gob.NewEncoder(zws[i])
        }
 
+       cgnames := make([]string, 0, len(tilelib.compactGenomes))
+       for name := range tilelib.compactGenomes {
+               cgnames = append(cgnames, name)
+       }
+       sort.Strings(cgnames)
+
+       refnames := make([]string, 0, len(tilelib.refseqs))
+       for name := range tilelib.refseqs {
+               refnames = append(refnames, name)
+       }
+       sort.Strings(refnames)
+
        log.Infof("WriteDir: writing %d files", nfiles)
        ctx, cancel := context.WithCancel(context.Background())
        defer cancel()
@@ -350,39 +410,35 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
                                errs <- err
                                return
                        }
-                       if start == 0 {
-                               // For now, just write all the genomes and refs
-                               // to the first file
-                               for name, cg := range tilelib.compactGenomes {
-                                       err := encoders[start].Encode(LibraryEntry{CompactGenomes: []CompactGenome{{
-                                               Name:     name,
-                                               Variants: cg,
-                                       }}})
-                                       if err != nil {
-                                               errs <- err
-                                               return
-                                       }
-                               }
-                               for name, tseqs := range tilelib.refseqs {
-                                       err := encoders[start].Encode(LibraryEntry{CompactSequences: []CompactSequence{{
-                                               Name:          name,
-                                               TileSequences: tseqs,
-                                       }}})
-                                       if err != nil {
-                                               errs <- err
-                                               return
-                                       }
+                       if refidx := start - ntilefiles; refidx >= 0 {
+                               // write each ref to its own file
+                               // (they seem to load very slowly)
+                               name := refnames[refidx]
+                               errs <- encoders[start].Encode(LibraryEntry{CompactSequences: []CompactSequence{{
+                                       Name:          name,
+                                       TileSequences: tilelib.refseqs[name],
+                               }}})
+                               return
+                       }
+                       for i := start; i < len(cgnames); i += ntilefiles {
+                               err := encoders[start].Encode(LibraryEntry{CompactGenomes: []CompactGenome{{
+                                       Name:     cgnames[i],
+                                       Variants: tilelib.compactGenomes[cgnames[i]],
+                               }}})
+                               if err != nil {
+                                       errs <- err
+                                       return
                                }
                        }
                        tvs := []TileVariant{}
-                       for tag := start; tag < len(tilelib.variant) && ctx.Err() == nil; tag += nfiles {
+                       for tag := start; tag < len(tilelib.variant) && ctx.Err() == nil; tag += ntilefiles {
                                tvs = tvs[:0]
                                for idx, hash := range tilelib.variant[tag] {
                                        tvs = append(tvs, TileVariant{
                                                Tag:      tagID(tag),
                                                Variant:  tileVariantID(idx + 1),
                                                Blake2b:  hash,
-                                               Sequence: tilelib.seq[hash],
+                                               Sequence: tilelib.hashSequence(hash),
                                        })
                                }
                                err := encoders[start].Encode(LibraryEntry{TileVariants: tvs})
@@ -422,9 +478,7 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
 // Load library data from rdr. Tile variants might be renumbered in
 // the process; in that case, genomes variants will be renumbered to
 // match.
-//
-// If onLoadGenome is non-nil, call it on each CompactGenome entry.
-func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool, onLoadGenome func(CompactGenome)) error {
+func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool) error {
        cgs := []CompactGenome{}
        cseqs := []CompactSequence{}
        variantmap := map[tileLibRef]tileVariantID{}
@@ -448,7 +502,7 @@ func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool,
        if ctx.Err() != nil {
                return ctx.Err()
        }
-       err = tilelib.loadCompactGenomes(cgs, variantmap, onLoadGenome)
+       err = tilelib.loadCompactGenomes(cgs, variantmap)
        if err != nil {
                return err
        }
@@ -490,15 +544,16 @@ func (tilelib *tileLibrary) dump(out io.Writer) {
 }
 
 type importStats struct {
-       InputFile              string
-       InputLabel             string
-       InputLength            int
-       InputCoverage          int
-       PathLength             int
-       DroppedOutOfOrderTiles int
+       InputFile             string
+       InputLabel            string
+       InputLength           int
+       InputCoverage         int
+       PathLength            int
+       DroppedRepeatedTags   int
+       DroppedOutOfOrderTags int
 }
 
-func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChromosome *regexp.Regexp) (tileSeq, []importStats, error) {
+func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChromosome *regexp.Regexp, isRef bool) (tileSeq, []importStats, error) {
        ret := tileSeq{}
        type jobT struct {
                label string
@@ -506,6 +561,7 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
        }
        todo := make(chan jobT, 1)
        scanner := bufio.NewScanner(rdr)
+       scanner.Buffer(make([]byte, 256), 1<<29) // 512 MiB, in case fasta does not have line breaks
        go func() {
                defer close(todo)
                var fasta []byte
@@ -551,14 +607,33 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
                        log.Warnf("%s %s no tags found", filelabel, job.label)
                }
 
-               skipped := 0
+               droppedDup := 0
+               if !tilelib.useDups {
+                       // Remove any tags that appeared more than once
+                       dup := map[tagID]bool{}
+                       for _, ft := range found {
+                               _, dup[ft.tagid] = dup[ft.tagid]
+                       }
+                       dst := 0
+                       for _, ft := range found {
+                               if !dup[ft.tagid] {
+                                       found[dst] = ft
+                                       dst++
+                               }
+                       }
+                       droppedDup = len(found) - dst
+                       log.Infof("%s %s dropping %d non-unique tags", filelabel, job.label, droppedDup)
+                       found = found[:dst]
+               }
+
+               droppedOOO := 0
                if tilelib.skipOOO {
-                       log.Infof("%s %s keeping longest increasing subsequence", filelabel, job.label)
                        keep := longestIncreasingSubsequence(len(found), func(i int) int { return int(found[i].tagid) })
                        for i, x := range keep {
                                found[i] = found[x]
                        }
-                       skipped = len(found) - len(keep)
+                       droppedOOO = len(found) - len(keep)
+                       log.Infof("%s %s dropping %d out-of-order tags", filelabel, job.label, droppedOOO)
                        found = found[:len(keep)]
                }
 
@@ -582,7 +657,7 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
                                } else {
                                        endpos = found[i+1].pos + taglen
                                }
-                               path[i] = tilelib.getRef(f.tagid, job.fasta[startpos:endpos])
+                               path[i] = tilelib.getRef(f.tagid, job.fasta[startpos:endpos], isRef)
                                if countBases(job.fasta[startpos:endpos]) != endpos-startpos {
                                        atomic.AddInt64(&lowquality, 1)
                                }
@@ -597,14 +672,15 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
                ret[job.label] = pathcopy
 
                basesIn := countBases(job.fasta)
-               log.Infof("%s %s fasta in %d coverage in %d path len %d low-quality %d skipped-out-of-order %d", filelabel, job.label, len(job.fasta), basesIn, len(path), lowquality, skipped)
+               log.Infof("%s %s fasta in %d coverage in %d path len %d low-quality %d", filelabel, job.label, len(job.fasta), basesIn, len(path), lowquality)
                stats = append(stats, importStats{
-                       InputFile:              filelabel,
-                       InputLabel:             job.label,
-                       InputLength:            len(job.fasta),
-                       InputCoverage:          basesIn,
-                       PathLength:             len(path),
-                       DroppedOutOfOrderTiles: skipped,
+                       InputFile:             filelabel,
+                       InputLabel:            job.label,
+                       InputLength:           len(job.fasta),
+                       InputCoverage:         basesIn,
+                       PathLength:            len(path),
+                       DroppedOutOfOrderTags: droppedOOO,
+                       DroppedRepeatedTags:   droppedDup,
                })
 
                totalPathLen += len(path)
@@ -619,7 +695,7 @@ func (tilelib *tileLibrary) Len() int64 {
 
 // Return a tileLibRef for a tile with the given tag and sequence,
 // adding the sequence to the library if needed.
-func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
+func (tilelib *tileLibrary) getRef(tag tagID, seq []byte, usedByRef bool) tileLibRef {
        dropSeq := false
        if !tilelib.retainNoCalls {
                for _, b := range seq {
@@ -704,37 +780,66 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
        vlock.Unlock()
 
        if tilelib.retainTileSequences && !dropSeq {
-               tilelib.mtx.Lock()
-               if tilelib.seq == nil {
-                       tilelib.seq = map[[blake2b.Size256]byte][]byte{}
+               seqCopy := append([]byte(nil), seq...)
+               if tilelib.seq2 == nil {
+                       tilelib.mtx.Lock()
+                       if tilelib.seq2 == nil {
+                               tilelib.seq2lock = map[[2]byte]sync.Locker{}
+                               m := map[[2]byte]map[[blake2b.Size256]byte][]byte{}
+                               var k [2]byte
+                               for i := 0; i < 256; i++ {
+                                       k[0] = byte(i)
+                                       for j := 0; j < 256; j++ {
+                                               k[1] = byte(j)
+                                               m[k] = map[[blake2b.Size256]byte][]byte{}
+                                               tilelib.seq2lock[k] = &sync.Mutex{}
+                                       }
+                               }
+                               tilelib.seq2 = m
+                       }
+                       tilelib.mtx.Unlock()
                }
-               tilelib.seq[seqhash] = append([]byte(nil), seq...)
-               tilelib.mtx.Unlock()
+               var k [2]byte
+               copy(k[:], seqhash[:])
+               locker := tilelib.seq2lock[k]
+               locker.Lock()
+               tilelib.seq2[k][seqhash] = seqCopy
+               locker.Unlock()
        }
 
+       saveSeq := seq
+       if dropSeq {
+               // Save the hash, but not the sequence
+               saveSeq = nil
+       }
        if tilelib.encoder != nil {
-               saveSeq := seq
-               if dropSeq {
-                       // Save the hash, but not the sequence
-                       saveSeq = nil
-               }
                tilelib.encoder.Encode(LibraryEntry{
                        TileVariants: []TileVariant{{
                                Tag:      tag,
+                               Ref:      usedByRef,
                                Variant:  variant,
                                Blake2b:  seqhash,
                                Sequence: saveSeq,
                        }},
                })
        }
+       if tilelib.onAddTileVariant != nil {
+               tilelib.onAddTileVariant(tileLibRef{tag, variant}, seqhash, saveSeq)
+       }
        return tileLibRef{Tag: tag, Variant: variant}
 }
 
+func (tilelib *tileLibrary) hashSequence(hash [blake2b.Size256]byte) []byte {
+       var partition [2]byte
+       copy(partition[:], hash[:])
+       return tilelib.seq2[partition][hash]
+}
+
 func (tilelib *tileLibrary) TileVariantSequence(libref tileLibRef) []byte {
        if libref.Variant == 0 || len(tilelib.variant) <= int(libref.Tag) || len(tilelib.variant[libref.Tag]) < int(libref.Variant) {
                return nil
        }
-       return tilelib.seq[tilelib.variant[libref.Tag][libref.Variant-1]]
+       return tilelib.hashSequence(tilelib.variant[libref.Tag][libref.Variant-1])
 }
 
 // Tidy deletes unreferenced tile variants and renumbers variants so