added mt comression to merge

[samtools.git] / bam_sort.c

diff --git a/bam_sort.c b/bam_sort.c
index e3a6c472cb1b50ddb0eab7273433a438c6dabede..f436fb9d80bcbe9227bddf14f75dc5b1660898bf 100644 (file)
--- a/bam_sort.c
+++ b/bam_sort.c
@@ -10,24 +10,28 @@
 
 static int g_is_by_qname = 0;
 
-static inline int strnum_cmp(const char *a, const char *b)
+static int strnum_cmp(const char *_a, const char *_b)
 {
-       char *pa, *pb;
-       pa = (char*)a; pb = (char*)b;
+       const unsigned char *a = (const unsigned char*)_a, *b = (const unsigned char*)_b;
+       const unsigned char *pa = a, *pb = b;
        while (*pa && *pb) {
                if (isdigit(*pa) && isdigit(*pb)) {
-                       long ai, bi;
-                       ai = strtol(pa, &pa, 10);
-                       bi = strtol(pb, &pb, 10);
-                       if (ai != bi) return ai<bi? -1 : ai>bi? 1 : 0;
+                       while (*pa == '0') ++pa;
+                       while (*pb == '0') ++pb;
+                       while (isdigit(*pa) && isdigit(*pb) && *pa == *pb) ++pa, ++pb;
+                       if (isdigit(*pa) && isdigit(*pb)) {
+                               int i = 0;
+                               while (isdigit(pa[i]) && isdigit(pb[i])) ++i;
+                               return isdigit(pa[i])? 1 : isdigit(pb[i])? -1 : (int)*pa - (int)*pb;
+                       } else if (isdigit(*pa)) return 1;
+                       else if (isdigit(*pb)) return -1;
+                       else if (pa - a != pb - b) return pa - a < pb - b? 1 : -1;
                } else {
-                       if (*pa != *pb) break;
+                       if (*pa != *pb) return (int)*pa - (int)*pb;
                        ++pa; ++pb;
                }
        }
-       if (*pa == *pb)
-               return (pa-a) < (pb-b)? -1 : (pa-a) > (pb-b)? 1 : 0;
-       return *pa<*pb? -1 : *pa>*pb? 1 : 0;
+       return *pa? 1 : *pb? -1 : 0;
 }
 
 #define HEAP_EMPTY 0xffffffffffffffffull
@@ -85,8 +89,7 @@ static void swap_header_text(bam_header_t *h1, bam_header_t *h2)
   @discussion Padding information may NOT correctly maintained. This
   function is NOT thread safe.
  */
-int bam_merge_core(int by_qname, const char *out, const char *headers, int n, char * const *fn,
-                                       int flag, const char *reg)
+int bam_merge_core2(int by_qname, const char *out, const char *headers, int n, char * const *fn, int flag, const char *reg, int n_threads)
 {
        bamFile fpout, *fp;
        heap1_t *heap;
@@ -218,6 +221,7 @@ int bam_merge_core(int by_qname, const char *out, const char *headers, int n, ch
        }
        bam_header_write(fpout, hout);
        bam_header_destroy(hout);
+       if (!(flag & MERGE_UNCOMP)) bgzf_mt(fpout, n_threads, 256);
 
        ks_heapmake(heap, n, heap);
        while (heap->pos != HEAP_EMPTY) {
@@ -252,12 +256,17 @@ int bam_merge_core(int by_qname, const char *out, const char *headers, int n, ch
        return 0;
 }
 
+int bam_merge_core(int by_qname, const char *out, const char *headers, int n, char * const *fn, int flag, const char *reg)
+{
+       return bam_merge_core2(by_qname, out, headers, n, fn, flag, reg, 0);
+}
+
 int bam_merge(int argc, char *argv[])
 {
-       int c, is_by_qname = 0, flag = 0, ret = 0;
+       int c, is_by_qname = 0, flag = 0, ret = 0, n_threads = 0;
        char *fn_headers = NULL, *reg = 0;
 
-       while ((c = getopt(argc, argv, "h:nru1R:f")) >= 0) {
+       while ((c = getopt(argc, argv, "h:nru1R:f@:")) >= 0) {
                switch (c) {
                case 'r': flag |= MERGE_RG; break;
                case 'f': flag |= MERGE_FORCE; break;
@@ -266,6 +275,7 @@ int bam_merge(int argc, char *argv[])
                case '1': flag |= MERGE_LEVEL1; break;
                case 'u': flag |= MERGE_UNCOMP; break;
                case 'R': reg = strdup(optarg); break;
+               case '@': n_threads = atoi(optarg); break;
                }
        }
        if (optind + 2 >= argc) {
@@ -276,6 +286,7 @@ int bam_merge(int argc, char *argv[])
                fprintf(stderr, "         -u       uncompressed BAM output\n");
                fprintf(stderr, "         -f       overwrite the output BAM if exist\n");
                fprintf(stderr, "         -1       compress level 1\n");
+               fprintf(stderr, "         -@ INT   number of BAM compression threads [0]\n");
                fprintf(stderr, "         -R STR   merge file in the specified region STR [all]\n");
                fprintf(stderr, "         -h FILE  copy the header in FILE to <out.bam> [in1.bam]\n\n");
                fprintf(stderr, "Note: Samtools' merge does not reconstruct the @RG dictionary in the header. Users\n");
@@ -291,7 +302,7 @@ int bam_merge(int argc, char *argv[])
                        return 1;
                }
        }
-       if (bam_merge_core(is_by_qname, argv[optind], fn_headers, argc - optind - 1, argv + optind + 1, flag, reg) < 0) ret = 1;
+       if (bam_merge_core2(is_by_qname, argv[optind], fn_headers, argc - optind - 1, argv + optind + 1, flag, reg, n_threads) < 0) ret = 1;
        free(reg);
        free(fn_headers);
        return ret;