From b9067ede232a4d81357286dd42a51e34f3d18ddd Mon Sep 17 00:00:00 2001
From: Andy Fingerhut <andy_fingerhut@alum.wustl.edu>
Date: Wed, 11 Sep 2024 03:20:03 -0400
Subject: [PATCH] Speed up STF tests by only compiling P4 program once for all
 test packets

Signed-off-by: Andy Fingerhut <andy_fingerhut@alum.wustl.edu>
---
 backends/bmv2/run-bmv2-test.py                | 41 ++++++++++++++++++-
 .../targets/bmv2/test/TestTemplate.cmake      | 12 +++++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/backends/bmv2/run-bmv2-test.py b/backends/bmv2/run-bmv2-test.py
index 94ec54d5dae..79c36f48c4b 100755
--- a/backends/bmv2/run-bmv2-test.py
+++ b/backends/bmv2/run-bmv2-test.py
@@ -98,6 +98,25 @@ def parse_args():
             " stf file in the same folder."
         ),
     )
+    parser.add_argument(
+        "-cf",
+        "--compiledP4File",
+        dest="compiledP4File",
+        help=(
+            "Provide a path for an already-compiled P4 program for this test. "
+            "If this option is not provided, the script will run the P4"
+            " compiler to create it."
+        ),
+    )
+    parser.add_argument(
+        "-wf",
+        "--writeCompiledP4File",
+        dest="writeCompiledP4File",
+        help=(
+            "Provide a path where a copy of the compiled P4 program should "
+            "be written, after this script compiles it."
+        ),
+    )
     return parser.parse_known_args()
 
 
@@ -109,6 +128,8 @@ def __init__(self):
         self.compilerSrcDir = ""  # path to compiler source tree
         self.compilerBuildDir = ""  # path to compiler build directory
         self.testFile = ""  # path to stf test file that is used
+        self.compiledP4File = None
+        self.writeCompiledP4File = None
         self.testName = None  # Name of the test
         self.verbose = False
         self.replace = False  # replace previous outputs
@@ -298,7 +319,15 @@ def process_file(options, argv):
     if options.runDebugger:
         args[0:0] = options.runDebugger.split()
         os.execvp(args[0], args)
-    result = run_timeout(options, args, timeout, stderr)
+    if options.compiledP4File:
+        shutil.copyfile(options.compiledP4File, jsonfile)
+        print("Copied compiled P4 file from " + options.compiledP4File)
+    if os.path.isfile(jsonfile):
+        print("Compiled P4 file already exists.  Skipping compile: " + jsonfile)
+        result = SUCCESS
+    else:
+        print("Compiled P4 file did not already exist, so compiling it: " + jsonfile)
+        result = run_timeout(options, args, timeout, stderr)
 
     if result != SUCCESS:
         print("Error compiling")
@@ -308,6 +337,10 @@ def process_file(options, argv):
             if "Compiler Bug" in stderr_file.read():
                 return FAILURE
 
+    if options.writeCompiledP4File:
+        shutil.copyfile(jsonfile, options.writeCompiledP4File)
+        print("Copied compiled P4 file to file " + options.writeCompiledP4File)
+
     expected_error = isError(options.p4filename)
     if expected_error:
         # invert result
@@ -346,6 +379,12 @@ def main(argv):
     # TODO: Convert these paths to pathlib's Path.
     options.p4filename = check_if_file(args.p4filename).as_posix()
     options.compilerSrcDir = check_if_dir(args.rootdir).as_posix()
+    if args.compiledP4File:
+        options.compiledP4File = check_if_file(args.compiledP4File).as_posix()
+    if args.writeCompiledP4File:
+        print("dbg jafinger '%s'" % (args.writeCompiledP4File))
+        p1 = Path(args.writeCompiledP4File)
+        options.writeCompiledP4File = Path(p1.absolute()).as_posix()
 
     # If no build directory is provided, use current working directory
     if args.builddir:
diff --git a/backends/p4tools/modules/testgen/targets/bmv2/test/TestTemplate.cmake b/backends/p4tools/modules/testgen/targets/bmv2/test/TestTemplate.cmake
index ab7a73a2ace..38f97312a09 100644
--- a/backends/p4tools/modules/testgen/targets/bmv2/test/TestTemplate.cmake
+++ b/backends/p4tools/modules/testgen/targets/bmv2/test/TestTemplate.cmake
@@ -10,10 +10,20 @@ function(check_with_bmv2 testfile testfolder p4test)
   set(__bmv2runner "${CMAKE_BINARY_DIR}/run-bmv2-test.py")
   # Find all the stf tests generated for this P4 file and test them with bmv2 model
   file(APPEND ${testfile} "stffiles=($(find ${testfolder} -name \"*.stf\"  | sort -n ))\n")
+  file(APPEND ${testfile} "first_stffile=1\n")
   file(APPEND ${testfile} "for item in \${stffiles[@]}\n")
   file(APPEND ${testfile} "do\n")
   file(APPEND ${testfile} "\techo \"Found \${item}\"\n")
-  file(APPEND ${testfile} "\tpython3 ${__bmv2runner} . -v -b -tf \${item} -bd ${__p4cbmv2path} ${p4test}\n")
+  file(APPEND ${testfile} "\tjsonfile=\"`basename ${p4test} .p4`.json\"\n")
+  file(APPEND ${testfile} "\tjsonfile=\"${testfolder}/\${jsonfile}\"\n")
+  file(APPEND ${testfile} "\tif [ \${first_stffile} == 1 ]\n")
+  file(APPEND ${testfile} "\tthen\n")
+  file(APPEND ${testfile} "\t\tfirst_stffile=0\n")
+  file(APPEND ${testfile} "\t\tmore_args=\"-wf \${jsonfile}\"\n")
+  file(APPEND ${testfile} "\telse\n")
+  file(APPEND ${testfile} "\t\tmore_args=\"-cf \${jsonfile}\"\n")
+  file(APPEND ${testfile} "\tfi\n")
+  file(APPEND ${testfile} "\tpython3 ${__bmv2runner} . -v -b -tf \${item} \${more_args} -bd ${__p4cbmv2path} ${p4test}\n")
   file(APPEND ${testfile} "done\n")
 endfunction(check_with_bmv2)