Faster code generation

2 files changed, 48 insertions, 23 deletions

diff --git a/tools/buildgen/generate_projects.py b/tools/buildgen/generate_projects.py
index 965dd292af..0602d93e56 100755
--- a/tools/buildgen/generate_projects.py
+++ b/tools/buildgen/generate_projects.py
@@ -45,12 +45,12 @@ import jobset
 os.chdir(os.path.join(os.path.dirname(sys.argv[0]), '..', '..'))
 
 argp = argparse.ArgumentParser()
-argp.add_argument('json', nargs='+')
+argp.add_argument('build_files', nargs='+', default=[])
 argp.add_argument('--templates', nargs='+', default=[])
 argp.add_argument('--jobs', '-j', default=multiprocessing.cpu_count(), type=int)
 args = argp.parse_args()
 
-json = args.json
+json = args.build_files
 
 test = {} if 'TEST' in os.environ else None
 
@@ -62,21 +62,31 @@ if not templates:
     for f in files:
       templates.append(os.path.join(root, f))
 
+pre_jobs = []
+base_cmd = ['python2.7', 'tools/buildgen/mako_renderer.py']
+cmd = base_cmd[:]
+for plugin in plugins:
+  cmd.append('-p')
+  cmd.append(plugin)
+for js in json:
+  cmd.append('-d')
+  cmd.append(js)
+cmd.append('-w')
+preprocessed_build = '.preprocessed_build'
+cmd.append(preprocessed_build)
+pre_jobs.append(jobset.JobSpec(cmd, shortname='preprocess', timeout_seconds=None))
+
 jobs = []
-for template in templates:
+for template in reversed(sorted(templates)):
   root, f = os.path.split(template)
   if os.path.splitext(f)[1] == '.template':
     out_dir = '.' + root[len('templates'):]
     out = out_dir + '/' + os.path.splitext(f)[0]
     if not os.path.exists(out_dir):
       os.makedirs(out_dir)
-    cmd = ['python2.7', 'tools/buildgen/mako_renderer.py']
-    for plugin in plugins:
-      cmd.append('-p')
-      cmd.append(plugin)
-    for js in json:
-      cmd.append('-d')
-      cmd.append(js)
+    cmd = base_cmd[:]
+    cmd.append('-P')
+    cmd.append(preprocessed_build)
     cmd.append('-o')
     if test is None:
       cmd.append(out)
@@ -88,6 +98,7 @@ for template in templates:
     cmd.append(root + '/' + f)
     jobs.append(jobset.JobSpec(cmd, shortname=out, timeout_seconds=None))
 
+jobset.run(pre_jobs, maxjobs=args.jobs)
 jobset.run(jobs, maxjobs=args.jobs)
 
 if test is not None:
diff --git a/tools/buildgen/mako_renderer.py b/tools/buildgen/mako_renderer.py
index f1b28d352e..f629e68eb9 100755
--- a/tools/buildgen/mako_renderer.py
+++ b/tools/buildgen/mako_renderer.py
@@ -38,6 +38,7 @@ Just a wrapper around the mako rendering library.
 import getopt
 import imp
 import os
+import cPickle as pickle
 import shutil
 import sys
 
@@ -66,21 +67,23 @@ def out(msg):
 
 
 def showhelp():
-  out('mako-renderer.py [-o out] [-m cache] [-d dict] [-d dict...] template')
+  out('mako-renderer.py [-o out] [-m cache] [-P preprocessed_input] [-d dict] [-d dict...]'
+      ' [-t template] [-w preprocessed_output]')
 
 
 def main(argv):
   got_input = False
   module_directory = None
+  preprocessed_output = None
   dictionary = {}
   json_dict = {}
   got_output = False
-  output_file = sys.stdout
   plugins = []
   output_name = None
+  got_preprocessed_input = False
 
   try:
-    opts, args = getopt.getopt(argv, 'hm:d:o:p:')
+    opts, args = getopt.getopt(argv, 'hm:d:o:p:t:P:w:')
   except getopt.GetoptError:
     out('Unknown option')
     showhelp()
@@ -104,18 +107,31 @@ def main(argv):
         showhelp()
         sys.exit(4)
       module_directory = arg
+    elif opt == '-P':
+      assert not got_preprocessed_input
+      assert json_dict == {}
+      sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), 'plugins')))
+      with open(arg, 'r') as dict_file:
+        dictionary = pickle.load(dict_file)
+      got_preprocessed_input = True
     elif opt == '-d':
-      dict_file = open(arg, 'r')
-      bunch.merge_json(json_dict, yaml.load(dict_file.read()))
-      dict_file.close()
+      assert not got_preprocessed_input
+      with open(arg, 'r') as dict_file:
+        bunch.merge_json(json_dict, yaml.load(dict_file.read()))
     elif opt == '-p':
       plugins.append(import_plugin(arg))
+    elif opt == '-w':
+      preprocessed_output = arg
 
-  for plugin in plugins:
-    plugin.mako_plugin(json_dict)
+  if not got_preprocessed_input:
+    for plugin in plugins:
+      plugin.mako_plugin(json_dict)
+    for k, v in json_dict.items():
+      dictionary[k] = bunch.to_bunch(v)
 
-  for k, v in json_dict.items():
-    dictionary[k] = bunch.to_bunch(v)
+  if preprocessed_output:
+    with open(preprocessed_output, 'w') as dict_file:
+      pickle.dump(dictionary, dict_file)
 
   cleared_dir = False
   for arg in args:
@@ -168,11 +184,9 @@ def main(argv):
           with open(item_output_name, 'w') as output_file:
             template.render_context(Context(output_file, **args))
 
-  if not got_input:
+  if not got_input and not preprocessed_output:
     out('Got nothing to do')
     showhelp()
 
-  output_file.close()
-
 if __name__ == '__main__':
   main(sys.argv[1:])