From 5741eb9ad71875b45d37f2bc36a939ce1b20952a Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Mon, 14 Sep 2020 12:45:59 -0700
Subject: [PATCH] Expanded benchmarking script and added one size opt to the
 encoder.

---
 benchmark.py | 44 +++++++++++++++++++++++++++++++++++++++-----
 upb/encode.c | 10 +++-------
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/benchmark.py b/benchmark.py
index 98c008864d..4a6e9f839b 100755
--- a/benchmark.py
+++ b/benchmark.py
@@ -1,16 +1,36 @@
 #!/usr/bin/env python3
+"""Benchmarks the current working directory against a given baseline.
 
+This script benchmarks both size and speed. Sample output:
+"""
+
+import contextlib
 import json
-import subprocess
+import os
 import re
+import subprocess
+import sys
+import tempfile
+
+@contextlib.contextmanager
+def GitWorktree(commit):
+  tmpdir = tempfile.mkdtemp()
+  subprocess.run(['git', 'worktree', 'add', '-q', tmpdir, commit], check=True)
+  cwd = os.getcwd()
+  os.chdir(tmpdir)
+  try:
+    yield tmpdir
+  finally:
+    os.chdir(cwd)
+    subprocess.run(['git', 'worktree', 'remove', tmpdir], check=True)
 
 def Run(cmd):
   subprocess.check_call(cmd, shell=True)
 
-def RunAgainstBranch(branch, outbase, runs=12):
+def Benchmark(outbase, runs=12):
   tmpfile = "/tmp/bench-output.json"
   Run("rm -rf {}".format(tmpfile))
-  Run("git checkout {}".format(branch))
+  Run("bazel test :all")
   Run("bazel build -c opt :benchmark")
 
   Run("./bazel-bin/benchmark --benchmark_out_format=json --benchmark_out={} --benchmark_repetitions={}".format(tmpfile, runs))
@@ -21,6 +41,7 @@ def RunAgainstBranch(branch, outbase, runs=12):
   with open(tmpfile) as f:
     bench_json = json.load(f)
 
+  # Translate into the format expected by benchstat.
   with open(outbase + ".txt", "w") as f:
     for run in bench_json["benchmarks"]:
       name = re.sub(r'^BM_', 'Benchmark', run["name"])
@@ -29,8 +50,21 @@ def RunAgainstBranch(branch, outbase, runs=12):
       values = (name, run["iterations"], run["cpu_time"])
       print("{} {} {} ns/op".format(*values), file=f)
 
-RunAgainstBranch("6e140c267cc9bf6a4ef89d3f9a842209d7537599", "/tmp/old")
-RunAgainstBranch("encoder", "/tmp/new")
+baseline = "master"
+
+if len(sys.argv) > 1:
+  baseline = sys.argv[1]
+
+  # Quickly verify that the baseline exists.
+  with GitWorktree(baseline):
+    pass
+
+# Benchmark our current directory first, since it's more likely to be broken.
+Benchmark("/tmp/new")
+
+# Benchmark the baseline.
+with GitWorktree(baseline):
+  Benchmark("/tmp/old")
 
 print()
 print()
diff --git a/upb/encode.c b/upb/encode.c
index 6bb876270f..299db1cbcc 100644
--- a/upb/encode.c
+++ b/upb/encode.c
@@ -146,7 +146,6 @@ static void encode_fixedarray(upb_encstate *e, const upb_array *arr,
     }
   } else {
     encode_bytes(e, data, bytes);
-    encode_varint(e, bytes);
   }
 }
 
@@ -242,6 +241,7 @@ static void encode_array(upb_encstate *e, const char *field_mem,
                          const upb_msglayout *m, const upb_msglayout_field *f) {
   const upb_array *arr = *(const upb_array**)field_mem;
   bool packed = f->label == _UPB_LABEL_PACKED;
+  size_t pre_len = e->limit - e->ptr;
 
   if (arr == NULL || arr->len == 0) {
     return;
@@ -251,19 +251,14 @@ static void encode_array(upb_encstate *e, const char *field_mem,
   {                                                                      \
     const ctype *start = _upb_array_constptr(arr);                       \
     const ctype *ptr = start + arr->len;                                 \
-    size_t pre_len = e->limit - e->ptr;                                  \
     uint32_t tag = packed ? 0 : (f->number << 3) | UPB_WIRE_TYPE_VARINT; \
     do {                                                                 \
       ptr--;                                                             \
       encode_varint(e, encode);                                          \
       if (tag) encode_varint(e, tag);                                    \
     } while (ptr != start);                                              \
-    if (!tag) encode_varint(e, e->limit - e->ptr - pre_len);             \
   }                                                                      \
-  break;                                                                 \
-  do {                                                                   \
-    ;                                                                    \
-  } while (0)
+  break;
 
 #define TAG(wire_type) (packed ? 0 : (f->number << 3 | wire_type))
 
@@ -338,6 +333,7 @@ static void encode_array(upb_encstate *e, const char *field_mem,
 #undef VARINT_CASE
 
   if (packed) {
+    encode_varint(e, e->limit - e->ptr - pre_len);
     encode_tag(e, f->number, UPB_WIRE_TYPE_DELIMITED);
   }
 }