incbwt/rlcsa_grep.cpp

   1 #include <iostream>
   2
   3 #include "rlcsa.h"
   4 #include "misc/utils.h"
   5
   6 using namespace CSA;
   7
   8
   9 enum mode_type { COUNT, DISPLAY, CONTEXT };
  10
  11
  12 void printUsage()
  13 {
  14   std::cout << "Usage: rlcsa_grep [-c] pattern base_name" << std::endl;
  15   std::cout << "  -c    print the number of matching lines" << std::endl;
  16   std::cout << "  -NUM  display NUM characters of leading and trailing context instead of" << std::endl;
  17   std::cout << "        the entire line" << std::endl;
  18 }
  19
  20
  21 int main(int argc, char** argv)
  22 {
  23   int base_arg = 2, pattern_arg = 1;
  24   mode_type mode = DISPLAY;
  25   usint context = 0;
  26
  27   if(argc < base_arg + 1)
  28   {
  29     printUsage();
  30     return 1;
  31   }
  32
  33   if(argv[1][0] == '-')
  34   {
  35     base_arg++; pattern_arg++;
  36     if(std::string("-c").compare(argv[1]) == 0)
  37     {
  38       mode = COUNT;
  39     }
  40     else
  41     {
  42       mode = CONTEXT;
  43       context = atoi(&(argv[1][1]));
  44     }
  45     if(argc < base_arg + 1)
  46     {
  47       printUsage();
  48       return 2;
  49     }
  50   }
  51
  52   RLCSA rlcsa(argv[base_arg], false);
  53   if(!rlcsa.isOk())
  54   {
  55     return 3;
  56   }
  57
  58   usint len = std::string(argv[pattern_arg]).length();
  59   pair_type result_range = rlcsa.count(argv[pattern_arg]);
  60   usint occurrences = length(result_range);
  61   if(occurrences == 0)
  62   {
  63     if(mode == COUNT)
  64     {
  65       std::cout << 0 << std::endl;
  66     }
  67     return 0;
  68   }
  69
  70   usint last_row = 0;
  71   usint* results = rlcsa.locate(result_range);
  72   if(mode != CONTEXT)
  73   {
  74     rlcsa.getSequenceForPosition(results, occurrences);
  75   }
  76   std::sort(results, results + occurrences);
  77   if(mode != CONTEXT)
  78   {
  79     for(usint i = 1; i < occurrences; i++)
  80     {
  81       if(results[i] != results[last_row])
  82       {
  83         last_row++; results[last_row] = results[i];
  84       }
  85     }
  86   }
  87
  88   if(mode == COUNT)
  89   {
  90     std::cout << (last_row + 1) << std::endl;
  91   }
  92   else if(mode == DISPLAY)
  93   {
  94     for(usint i = 0; i <= last_row; i++)
  95     {
  96       uchar* row = rlcsa.display(results[i]);
  97       std::cout.write((char*)row, length(rlcsa.getSequenceRange(results[i])));
  98       std::cout << std::endl;
  99       delete[] row;
 100     }
 101   }
 102   else if(mode == CONTEXT)
 103   {
 104     usint result_length = 0;
 105     for(usint i = 0; i < occurrences; i++)
 106     {
 107       uchar* text = rlcsa.display(results[i], len, context, result_length);
 108       std::cout.write((char*)text, result_length);
 109       std::cout << std::endl;
 110       delete[] text;
 111     }
 112   }
 113
 114   delete[] results;
 115   return 0;
 116 }