SERVER-30468 Optimize errorcodes.py

* Only look in src/mongo rather than filtering after the fact
* Use re2 if available
* Lazily convert byte offsets to line numbers only when needed

Also I added a --quiet option to make it silent on success.

1 files changed, 45 insertions, 30 deletions

diff --git a/buildscripts/errorcodes.py b/buildscripts/errorcodes.py
index 471cb7fabda..7f51b7512b1 100755
--- a/buildscripts/errorcodes.py
+++ b/buildscripts/errorcodes.py
@@ -8,19 +8,24 @@ Optionally replaces zero codes in source code with new distinct values.
 
 import bisect
 import os
-import re
 import sys
 import utils
 from collections import defaultdict, namedtuple
 from optparse import OptionParser
 
+try:
+    import re2 as re
+except ImportError:
+    print("*** Run 'pip2 install --user re2' to speed up error code checking")
+    import re
+
 ASSERT_NAMES = [ "uassert" , "massert", "fassert", "fassertFailed" ]
 MINIMUM_CODE = 10000
 
 codes = []
 
 # Each AssertLocation identifies the C++ source location of an assertion
-AssertLocation = namedtuple( "AssertLocation", ['sourceFile', 'lineNum', 'lines', 'code'] )
+AssertLocation = namedtuple( "AssertLocation", ['sourceFile', 'byteOffset', 'lines', 'code'] )
 
 
 # Of historical interest only
@@ -59,11 +64,7 @@ def parseSourceFiles( callback ):
 
     bad = [ re.compile( r"^\s*assert *\(" ) ]
 
-    for sourceFile in utils.getAllSourceFiles():
-        if not sourceFile.find( "src/mongo/" ) >= 0:
-            # Skip generated sources
-            continue
-
+    for sourceFile in utils.getAllSourceFiles(prefix='src/mongo/'):
         with open(sourceFile) as f:
             text = f.read()
 
@@ -77,12 +78,6 @@ def parseSourceFiles( callback ):
                     print( "%s: %s" % (sourceFile, msg) )
                     raise Exception(msg)
 
-            splitlines = text.splitlines(True)
-
-            line_offsets = [0]
-            for line in splitlines:
-                line_offsets.append(line_offsets[-1] + len(line))
-
             matchiters = [p.finditer(text) for p in patterns]
             for matchiter in matchiters:
                 for match in matchiter:
@@ -90,15 +85,23 @@ def parseSourceFiles( callback ):
                     span = match.span()
 
                     thisLoc = AssertLocation(sourceFile,
-                                             getLineForPosition(line_offsets, span[1]),
+                                             span[1],
                                              text[span[0]:span[1]],
                                              code)
 
                     callback( thisLoc )
 
 # Converts an absolute position in a file into a line number.
-def getLineForPosition(line_offsets, position):
-    return bisect.bisect(line_offsets, position)
+def getLineForPosition(loc, _file_cache={}):
+    if loc.sourceFile not in _file_cache:
+        with open(loc.sourceFile) as f:
+            text = f.read()
+            line_offsets = [0]
+            for line in text.splitlines(True):
+                line_offsets.append(line_offsets[-1] + len(line))
+            _file_cache[loc.sourceFile] = line_offsets
+
+    return bisect.bisect(_file_cache[loc.sourceFile], loc.byteOffset)
 
 def isTerminated( lines ):
     """Given .cpp/.h source lines as text, determine if assert is terminated."""
@@ -159,12 +162,12 @@ def readErrorCodes():
         bad = seen[code]
         errors.append( bad )
         print( "ZERO_CODE:" )
-        print( "  %s:%d:%s" % (bad.sourceFile, bad.lineNum, bad.lines) )
+        print( "  %s:%d:%s" % (bad.sourceFile, getLineForPosition(bad), bad.lines) )
 
     for code, locations in dups.items():
         print( "DUPLICATE IDS: %s" % code )
         for loc in locations:
-            print( "  %s:%d:%s" % (loc.sourceFile, loc.lineNum, loc.lines) )
+            print( "  %s:%d:%s" % (loc.sourceFile, getLineForPosition(loc), loc.lines) )
 
     return (codes, errors)
 
@@ -181,25 +184,30 @@ def replaceBadCodes( errors, nextCode ):
     skip_errors = [e for e in errors if int(e.code) != 0]
 
     for loc in skip_errors:
-        print "SKIPPING NONZERO code=%s: %s:%s" % (loc.code, loc.sourceFile, loc.lineNum)
+        print ("SKIPPING NONZERO code=%s: %s:%s"
+                % (loc.code, loc.sourceFile, getLineForPosition(loc)))
 
-    for assertLoc in zero_errors:
-        (sourceFile, lineNum, lines, code) = assertLoc
+    # Dedupe, sort, and reverse so we don't have to update offsets as we go.
+    for assertLoc in reversed(sorted(set(zero_errors))):
+        (sourceFile, byteOffset, lines, code) = assertLoc
+        lineNum = getLineForPosition(assertLoc)
         print "UPDATING_FILE: %s:%s" % (sourceFile, lineNum)
 
         ln = lineNum - 1
 
         with open(sourceFile, 'r+') as f:
-            fileLines = f.readlines()
-            line = fileLines[ln]
-            print "LINE_%d_BEFORE:%s" % (lineNum, line.rstrip())
-            line = re.sub(r'(\( *)(\d+)',
-                          r'\g<1>' + str(nextCode),
-                          line)
-            print "LINE_%d_AFTER :%s" % (lineNum, line.rstrip())
-            fileLines[ln] = line
+            print "LINE_%d_BEFORE:%s" % (lineNum, f.readlines()[ln].rstrip())
+
+            f.seek(0)
+            text = f.read()
+            assert text[byteOffset-1] == '0'
+            f.seek(0)
+            f.write(text[:byteOffset-1])
+            f.write(str(nextCode))
+            f.write(text[byteOffset:])
             f.seek(0)
-            f.writelines(fileLines)
+
+            print "LINE_%d_AFTER :%s" % (lineNum, f.readlines()[ln].rstrip())
         nextCode += 1
 
 
@@ -241,10 +249,17 @@ def main():
     parser.add_option("--fix", dest="replace",
                       action="store_true", default=False,
                       help="Fix zero codes in source files [default: %default]")
+    parser.add_option("-q", "--quiet", dest="quiet",
+                      action="store_true", default=False,
+                      help="Suppress output on success [default: %default]")
     (options, args) = parser.parse_args()
 
     (codes, errors) = readErrorCodes()
     ok = len(errors) == 0
+
+    if ok and options.quiet:
+        return
+
     next = getNextCode()
 
     print("ok: %s" % ok)