Adjust slice memory.
[lightning.git] / tilelib.go
index 26b2fd2d8aea62796b28fbb0c9c99d94f7f3285f..289f8895e5ba420b00d0d449f9a916cfcaf5d4f5 100644 (file)
@@ -1,3 +1,7 @@
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package lightning
 
 import (
@@ -67,6 +71,12 @@ type tileLibrary struct {
        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 +122,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)
@@ -146,8 +156,15 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
                                // 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{
@@ -199,6 +216,12 @@ 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()
@@ -213,36 +236,36 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
        return nil
 }
 
-func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGenome func(CompactGenome)) error {
+func allGobFiles(path string) ([]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 := allGobFiles(child)
+                       if err != nil {
+                               return nil, err
                        }
+                       files = append(files, add...)
+               } else if strings.HasSuffix(child, ".gob") || strings.HasSuffix(child, ".gob.gz") {
+                       files = append(files, child)
                }
-               return nil
        }
+       return files, nil
+}
+
+func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string) error {
        log.Infof("LoadDir: walk dir %s", path)
-       err := walk(path)
+       files, err := allGobFiles(path)
        if err != nil {
                return err
        }
@@ -286,7 +309,7 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
                                        mtx.Unlock()
                                }
                                for _, tv := range ent.TileVariants {
-                                       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
                                }
                                cgs = append(cgs, ent.CompactGenomes...)
                                cseqs = append(cseqs, ent.CompactSequences...)
@@ -314,7 +337,7 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
        for _, cgs := range allcgs {
                flatcgs = append(flatcgs, cgs...)
        }
-       err = tilelib.loadCompactGenomes(flatcgs, allvariantmap, onLoadGenome)
+       err = tilelib.loadCompactGenomes(flatcgs, allvariantmap)
        if err != nil {
                return err
        }
@@ -334,7 +357,8 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 }
 
 func (tilelib *tileLibrary) WriteDir(dir string) error {
-       nfiles := 128 + len(tilelib.refseqs)
+       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)
@@ -377,7 +401,12 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
        for start := range files {
                start := start
                go func() {
-                       if refidx := start - (nfiles - len(tilelib.refseqs)); refidx >= 0 {
+                       err := encoders[start].Encode(LibraryEntry{TagSet: tilelib.taglib.Tags()})
+                       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]
@@ -387,12 +416,7 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
                                }}})
                                return
                        }
-                       err := encoders[start].Encode(LibraryEntry{TagSet: tilelib.taglib.Tags()})
-                       if err != nil {
-                               errs <- err
-                               return
-                       }
-                       for i := start; i < len(cgnames); i += nfiles {
+                       for i := start; i < len(cgnames); i += ntilefiles {
                                err := encoders[start].Encode(LibraryEntry{CompactGenomes: []CompactGenome{{
                                        Name:     cgnames[i],
                                        Variants: tilelib.compactGenomes[cgnames[i]],
@@ -403,7 +427,7 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
                                }
                        }
                        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{
@@ -450,9 +474,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{}
@@ -476,7 +498,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
        }
@@ -526,7 +548,7 @@ type importStats struct {
        DroppedOutOfOrderTiles 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
@@ -610,7 +632,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)
                                }
@@ -647,7 +669,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 {
@@ -759,21 +781,25 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
                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}
 }