feat: implement chunking for FASTA records and enhance CLI options

Author: ericksamera

cmd/radigest/main.go

@@ -1,35 +1,105 @@
 package main
 
 import (
+	"encoding/json"
 	"flag"
 	"fmt"
 	"log"
+	"os"
 	"runtime"
+	"sort"
 	"strings"
 	"sync"
-	"encoding/json"
-	"os"
 
 	"radigest/internal/collector"
 	"radigest/internal/digest"
 	"radigest/internal/enzyme"
 	"radigest/internal/fasta"
 )
 
+var (
+	version = "dev"
+	commit  = "none"
+	date    = "unknown"
+)
+
+func produceChunks(faCh <-chan fasta.Record, jobs chan<- fasta.Record, chunkSz int) {
+	defer close(jobs)
+	for rec := range faCh {
+		seq := rec.Seq
+		n := len(seq)
+		for from := 0; from < n; from += chunkSz {
+			to := from + chunkSz
+			if to > n {
+				to = n
+			}
+			jobs <- fasta.Record{
+				ID:  rec.ID,
+				Seq: seq[from:to],
+			}
+		}
+	}
+}
+
 func main() {
 	// ---- CLI flags ----------------------------------------------------------
 	fastaPath := flag.String("fasta", "", "reference FASTA file (required)")
-	enzFlag   := flag.String("enzymes", "", "comma-separated enzyme names (≥2, first two form the AB pair)")
-	minLen    := flag.Int("min", 0, "minimum fragment length")
-	maxLen    := flag.Int("max", 1<<30, "maximum fragment length")
-	gffPath   := flag.String("gff",  "fragments.gff3", "output GFF3")
-	jsonPath  := flag.String("json", "", "write run summary JSON")
-	chunkSz   := flag.Int("chunk", 8<<20, "chunk size (bp) sent to each worker")
-	threads   := flag.Int("threads", runtime.NumCPU(), "worker goroutines")
+	enzFlag := flag.String("enzymes", "", "comma-separated enzyme names (≥2, first two form the AB pair)")
+	minLen := flag.Int("min", 0, "minimum fragment length (bp)")
+	maxLen := flag.Int("max", 1<<30, "maximum fragment length (bp)")
+	gffPath := flag.String("gff", "fragments.gff3", "output GFF3 file")
+	jsonPath := flag.String("json", "", "optional: write run summary JSON here")
+	chunkSz := flag.Int("chunk", 8<<20, "chunk size (bp) sent to each worker")
+	threads := flag.Int("threads", runtime.NumCPU(), "number of worker goroutines")
+	showVer := flag.Bool("version", false, "print version and exit")
+	listEns := flag.Bool("list-enzymes", false, "list available enzyme names and exit")
+
+	flag.Usage = func() {
+		b := &strings.Builder{}
+		fmt.Fprintln(b, "radigest — in-silico double-digest and GFF3 fragment export")
+		fmt.Fprintln(b)
+		fmt.Fprintln(b, "Usage:")
+		fmt.Fprintln(b, "  radigest -fasta <ref.fa> -enzymes <E1,E2[,E3...]> [options]")
+		fmt.Fprintln(b)
+		fmt.Fprintln(b, "Required flags:")
+		fmt.Fprintln(b, "  -fasta, -enzymes")
+		fmt.Fprintln(b)
+		fmt.Fprintln(b, "Options:")
+		flag.CommandLine.SetOutput(b)
+		flag.PrintDefaults()
+		flag.CommandLine.SetOutput(os.Stderr)
+		fmt.Fprintln(b)
+		fmt.Fprintln(b, "Examples:")
+		fmt.Fprintln(b, "  # Basic EcoRI/MseI digest to GFF3")
+		fmt.Fprintln(b, "  radigest -fasta ref.fa -enzymes EcoRI,MseI -gff out.gff3")
+		fmt.Fprintln(b, "  # Restrict fragment size and emit JSON summary")
+		fmt.Fprintln(b, "  radigest -fasta ref.fa -enzymes EcoRI,MseI -min 100 -max 800 -json run.json")
+		fmt.Fprintln(b, "  # See supported enzymes")
+		fmt.Fprintln(b, "  radigest -list-enzymes")
+		fmt.Fprintln(os.Stderr, b.String())
+	}
+
 	flag.Parse()
 
+	if *showVer {
+		fmt.Printf("radigest %s (commit %s, %s)\n", version, commit, date)
+		return
+	}
+	if *listEns {
+		names := make([]string, 0, len(enzyme.DB))
+		for name := range enzyme.DB {
+			names = append(names, name)
+		}
+		sort.Strings(names)
+		for _, n := range names {
+			fmt.Println(n)
+		}
+		return
+	}
 	if *fastaPath == "" || *enzFlag == "" {
-		log.Fatal("flags --fasta and --enzymes are required")
+		fmt.Fprintln(os.Stderr, "error: flags -fasta and -enzymes are required\n")
+		flag.Usage()
+		os.Exit(2)
 	}
 
 	// ---- build enzyme slice -------------------------------------------------
@@ -68,53 +138,44 @@ func main() {
 	}
 
 	// ---- stream FASTA into jobs --------------------------------------------
-	faCh := make(chan fasta.Record, 2)
+	faCh := make(chan fasta.Record)
 	go func() {
 		if err := fasta.Stream(*fastaPath, faCh); err != nil {
 			log.Fatalf("fasta stream: %v", err)
 		}
+		// NOTE: assume fasta.Stream closes faCh when it returns.
 	}()
 
-	for rec := range faCh {
-		// split sequence into windows of *chunkSz bases
-		for from := 0; from < len(rec.Seq); from += *chunkSz {
-			to := from + *chunkSz
-			if to > len(rec.Seq) { to = len(rec.Seq) }
-			jobs <- fasta.Record{
-				ID:  rec.ID,
-				Seq: rec.Seq[from:to],
-			}
-		}
-	}
-	
-	for rec := range faCh {
-		jobs <- rec
-	}
-	close(jobs) // no more work
-	wg.Wait()   // workers done
-	close(cIn)  // tell collector to finish
+	// single consumer / producer path
+	go produceChunks(faCh, jobs, *chunkSz)
+
+	// wait for workers, finish collector
+	wg.Wait()  // jobs closed by produceChunks
+	close(cIn) // tell collector to finish
 
 	// ---- summary ------------------------------------------------------------
 	stats := <-done
 	fmt.Printf("Fragments kept: %d\nBases covered: %d\nChromosomes: %d\n",
-	    stats.TotalFragments, stats.TotalBases, len(stats.PerChr))	
+		stats.TotalFragments, stats.TotalBases, len(stats.PerChr))
 	if *jsonPath != "" {
-	    out := struct {
-	        Enzymes     []string               `json:"enzymes"`
-	        MinLength   int                    `json:"min_length"`
-	        MaxLength   int                    `json:"max_length"`
-	        collector.Stats
-	    }{
-	        Enzymes:   strings.Split(*enzFlag, ","),
-	        MinLength: *minLen,
-	        MaxLength: *maxLen,
-	        Stats:     stats,
-	    }
-	    f, err := os.Create(*jsonPath)
-	    if err != nil { log.Fatalf("write json: %v", err) }
-	    if err := json.NewEncoder(f).Encode(out); err != nil {
-	        log.Fatalf("encode json: %v", err)
-	    }
-	    f.Close()
+		out := struct {
+			Enzymes   []string `json:"enzymes"`
+			MinLength int      `json:"min_length"`
+			MaxLength int      `json:"max_length"`
+			collector.Stats
+		}{
+			Enzymes:   strings.Split(*enzFlag, ","),
+			MinLength: *minLen,
+			MaxLength: *maxLen,
+			Stats:     stats,
+		}
+		f, err := os.Create(*jsonPath)
+		if err != nil {
+			log.Fatalf("write json: %v", err)
+		}
+		if err := json.NewEncoder(f).Encode(out); err != nil {
+			log.Fatalf("encode json: %v", err)
+		}
+		_ = f.Close()
 	}
-}
+}