From de738e153deb88d9f366a11701090b52497cbb1b Mon Sep 17 00:00:00 2001
From: Florent Kermarrec <florent@enjoy-digital.fr>
Date: Tue, 28 Sep 2021 09:58:43 +0200
Subject: [PATCH] tools/litex_sim: Avoid double build iteration with
 pre_run_callback function.

---
 litex/build/sim/verilator.py | 35 +++++++++++++++++++----------------
 litex/tools/litex_sim.py     | 34 +++++++++++++++-------------------
 2 files changed, 34 insertions(+), 35 deletions(-)

diff --git a/litex/build/sim/verilator.py b/litex/build/sim/verilator.py
index 83d774c7d..ea58cc82f 100644
--- a/litex/build/sim/verilator.py
+++ b/litex/build/sim/verilator.py
@@ -175,22 +175,23 @@ def _run_sim(build_name, as_root=False, interactive=True):
 
 class SimVerilatorToolchain:
     def build(self, platform, fragment,
-            build_dir    = "build",
-            build_name   = "sim",
-            serial       = "console",
-            build        = True,
-            run          = True,
-            threads      = 1,
-            verbose      = True,
-            sim_config   = None,
-            coverage     = False,
-            opt_level    = "O0",
-            trace        = False,
-            trace_fst    = False,
-            trace_start  = 0,
-            trace_end    = -1,
-            regular_comb = False,
-            interactive  = True):
+            build_dir        = "build",
+            build_name       = "sim",
+            serial           = "console",
+            build            = True,
+            run              = True,
+            threads          = 1,
+            verbose          = True,
+            sim_config       = None,
+            coverage         = False,
+            opt_level        = "O0",
+            trace            = False,
+            trace_fst        = False,
+            trace_start      = 0,
+            trace_end        = -1,
+            regular_comb     = False,
+            interactive      = True,
+            pre_run_callback = None):
 
         # Create build directory
         os.makedirs(build_dir, exist_ok=True)
@@ -228,6 +229,8 @@ class SimVerilatorToolchain:
 
         # Run
         if run:
+            if pre_run_callback is not None:
+                pre_run_callback(v_output.ns)
             if which("verilator") is None:
                 msg = "Unable to find Verilator toolchain, please either:\n"
                 msg += "- Install Verilator.\n"
diff --git a/litex/tools/litex_sim.py b/litex/tools/litex_sim.py
index e38bae4b4..94dcde053 100755
--- a/litex/tools/litex_sim.py
+++ b/litex/tools/litex_sim.py
@@ -454,27 +454,23 @@ def main():
             soc.add_constant("REMOTEIP{}".format(i+1), int(args.remote_ip.split(".")[i]))
 
     # Build/Run ------------------------------------------------------------------------------------
+    def pre_run_callback(vns):
+        generate_gtkw_savefile(builder, vns, args.trace_fst)
+
     builder_kwargs["csr_csv"] = "csr.csv"
     builder = Builder(soc, **builder_kwargs)
-    for i in range(2):
-        build = (i == 0)
-        run   = (i == 1)
-        vns = builder.build(
-            build       = build,
-            run         = run,
-            threads     = args.threads,
-            sim_config  = sim_config,
-            opt_level   = args.opt_level,
-            trace       = args.trace,
-            trace_fst   = args.trace_fst,
-            trace_start = trace_start,
-            trace_end   = trace_end,
-            interactive = not args.non_interactive
-        )
-        if args.with_analyzer:
-            soc.analyzer.export_csv(vns, "analyzer.csv")
-        if args.gtkwave_savefile:
-            generate_gtkw_savefile(builder, vns, args.trace_fst)
+    soc.platform.toolchain.pre_run_callback = pre_run_callback
+    builder.build(
+        threads          = args.threads,
+        sim_config       = sim_config,
+        opt_level        = args.opt_level,
+        trace            = args.trace,
+        trace_fst        = args.trace_fst,
+        trace_start      = trace_start,
+        trace_end        = trace_end,
+        interactive      = not args.non_interactive,
+        pre_run_callback = pre_run_callback
+    )
 
 if __name__ == "__main__":
     main()