From 045b5b999e03a84968ef18d167bc149299c2e389 Mon Sep 17 00:00:00 2001
From: Remy Willems <rwillems@amazon.com>
Date: Fri, 9 Aug 2024 11:43:16 +0200
Subject: [PATCH] Add extra cancellation

---
 Source/BoogieDriver/BoogieDriver.cs       | 14 +++---
 Source/ExecutionEngine/ExecutionEngine.cs | 56 ++++++++++++++---------
 2 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/Source/BoogieDriver/BoogieDriver.cs b/Source/BoogieDriver/BoogieDriver.cs
index 79b86bdd6..379c00385 100644
--- a/Source/BoogieDriver/BoogieDriver.cs
+++ b/Source/BoogieDriver/BoogieDriver.cs
@@ -11,7 +11,7 @@ public class OnlyBoogie
     public static int Main(string[] args)
     {
       Contract.Requires(cce.NonNullElements(args));
-
+      
       var options = new CommandLineOptions(Console.Out, new ConsolePrinter())
       {
         RunningBoogieFromCommandLine = true
@@ -21,6 +21,12 @@ public static int Main(string[] args)
       {
         return 1;
       }
+      var source = new CancellationTokenSource();
+      if (options.ProcessTimeLimit != 0)
+      {
+        source.CancelAfter(TimeSpan.FromSeconds(options.ProcessTimeLimit));
+      }
+      
       using var executionEngine = ExecutionEngine.CreateWithoutSharedCache(options);
       
       if (options.ProcessInfoFlags())
@@ -64,12 +70,6 @@ public static int Main(string[] args)
 
       Helpers.ExtraTraceInformation(options, "Becoming sentient");
 
-      var source = new CancellationTokenSource();
-      if (options.ProcessTimeLimit != 0)
-      {
-        source.CancelAfter(TimeSpan.FromSeconds(options.ProcessTimeLimit));
-      }
-
       var success = executionEngine.ProcessFiles(Console.Out, fileList, cancellationToken: source.Token).Result;
       if (options.XmlSink != null)
       {
diff --git a/Source/ExecutionEngine/ExecutionEngine.cs b/Source/ExecutionEngine/ExecutionEngine.cs
index a1e026246..bf3ece69d 100644
--- a/Source/ExecutionEngine/ExecutionEngine.cs
+++ b/Source/ExecutionEngine/ExecutionEngine.cs
@@ -74,50 +74,62 @@ public static ExecutionEngine CreateWithoutSharedCache(ExecutionEngineOptions op
 
     static DateTime FirstRequestStart;
 
-    static readonly ConcurrentDictionary<string, TimeSpan>
-      TimePerRequest = new ConcurrentDictionary<string, TimeSpan>();
+    static readonly ConcurrentDictionary<string, TimeSpan> TimePerRequest = new();
 
-    static readonly ConcurrentDictionary<string, PipelineStatistics> StatisticsPerRequest =
-      new ConcurrentDictionary<string, PipelineStatistics>();
+    static readonly ConcurrentDictionary<string, PipelineStatistics> StatisticsPerRequest = new();
 
-    static readonly ConcurrentDictionary<string, CancellationTokenSource> ImplIdToCancellationTokenSource =
-      new ConcurrentDictionary<string, CancellationTokenSource>();
+    static readonly ConcurrentDictionary<string, CancellationTokenSource> ImplIdToCancellationTokenSource = new();
 
     private readonly TaskScheduler largeThreadScheduler;
     private readonly bool disposeScheduler;
 
-    public async Task<bool> ProcessFiles(TextWriter output, IList<string> fileNames, bool lookForSnapshots = true,
+    public async Task<bool> ProcessFiles(TextWriter output, IList<string> fileNames, 
       string programId = null, CancellationToken cancellationToken = default) {
       Contract.Requires(cce.NonNullElements(fileNames));
 
-      if (Options.VerifySeparately && 1 < fileNames.Count) {
-        var success = true;
-        foreach (var fileName in fileNames) {
-          success &= await ProcessFiles(output, new List<string> { fileName }, lookForSnapshots, fileName, cancellationToken: cancellationToken);
-        }
-        return success;
-      }
-
-      if (0 <= Options.VerifySnapshots && lookForSnapshots) {
+      if (0 <= Options.VerifySnapshots) {
         var snapshotsByVersion = LookForSnapshots(fileNames);
         var success = true;
         foreach (var snapshots in snapshotsByVersion) {
+          cancellationToken.ThrowIfCancellationRequested();
           // BUG: Reusing checkers during snapshots doesn't work, even though it should. We create a new engine (and thus checker pool) to workaround this.
           using var engine = new ExecutionEngine(Options, Cache,
             CustomStackSizePoolTaskScheduler.Create(StackSize, Options.VcsCores), true);
-          success &= await engine.ProcessFiles(output, new List<string>(snapshots), false, programId, cancellationToken: cancellationToken);
+          success &= await engine.ProcessBoogieFiles(output, snapshots, programId, cancellationToken: cancellationToken);
+        }
+        return success;
+      }
+
+      return await ProcessBoogieFiles(output, fileNames, programId, cancellationToken);
+    }
+
+    private async Task<bool> ProcessBoogieFiles(TextWriter output, IList<string> fileNames, string programId,
+      CancellationToken cancellationToken)
+    {
+      if (Options.VerifySeparately) {
+        var success = true;
+        foreach (var fileName in fileNames) {
+          cancellationToken.ThrowIfCancellationRequested();
+          success &= await ProcessBoogieFilesTogether(output, new List<string> { fileName }, fileName, cancellationToken: cancellationToken);
         }
         return success;
       }
 
+      return await ProcessBoogieFilesTogether(output, fileNames, programId, cancellationToken);
+    }
+
+    private Task<bool> ProcessBoogieFilesTogether(TextWriter output, IList<string> fileNames, string programId,
+      CancellationToken cancellationToken)
+    {
       using XmlFileScope xf = new XmlFileScope(Options.XmlSink, fileNames[^1]);
-      Program program = ParseBoogieProgram(fileNames, false);
-      var bplFileName = fileNames[^1];
-      if (program == null) {
-        return true;
+      var program = ParseBoogieProgram(fileNames, false);
+      if (program == null)
+      {
+        return Task.FromResult(true);
       }
 
-      return await ProcessProgram(output, program, bplFileName, programId, cancellationToken: cancellationToken);
+      var bplFileName = fileNames[^1];
+      return ProcessProgram(output, program, bplFileName, programId, cancellationToken: cancellationToken);
     }
 
     [Obsolete("Please inline this method call")]