Fixes lost dependencies for pruning during monomorphization. (#767)

This PR attempts to fix more axioms being pruned with the `/prune` flag
than intended with monomorphization.

Monomorphization introduces new axioms that are monomorphized, and tries
splitting all axioms (polymorphic or not) into their conjuncts. This PR
adds logic to update functions and constants to point to the new axioms
as dependencies.

---------

Co-authored-by: Remy Willems <[email protected]>

Source/Core/AST/Absy.cs

@@ -1429,7 +1429,7 @@ public class Constant : Variable
     // from all other constants.
     public readonly bool Unique;
 
-    public IList<Axiom> DefinitionAxioms { get; }
+    public IList<Axiom> DefinitionAxioms { get; set; }
 
     public Constant(IToken /*!*/ tok, TypedIdent /*!*/ typedIdent)
       : this(tok, typedIdent, true)
@@ -2045,18 +2045,9 @@ public class Function : DeclWithFormals
     public NAryExpr DefinitionBody; // Only set if the function is declared with {:define}
     public Axiom DefinitionAxiom;
 
-    public IList<Axiom> otherDefinitionAxioms = new List<Axiom>();
+    public IList<Axiom> OtherDefinitionAxioms = new List<Axiom>();
     public IEnumerable<Axiom> DefinitionAxioms => 
-      (DefinitionAxiom == null ? Enumerable.Empty<Axiom>() : new[]{ DefinitionAxiom }).Concat(otherDefinitionAxioms);
-
-    public IEnumerable<Axiom> OtherDefinitionAxioms => otherDefinitionAxioms;
-
-    public void AddOtherDefinitionAxiom(Axiom axiom)
-    {
-      Contract.Requires(axiom != null);
-
-      otherDefinitionAxioms.Add(axiom);
-    }
+      (DefinitionAxiom == null ? Enumerable.Empty<Axiom>() : new[]{ DefinitionAxiom }).Concat(OtherDefinitionAxioms);
 
     private bool neverTrigger;
     private bool neverTriggerComputed;

Source/Core/BoogiePL.atg

@@ -477,7 +477,7 @@ Function<.out List<Declaration>/*!*/ ds.>
                                       new Formal(retTyd.tok, retTyd, false, argKv), null, kv);
     foreach(var axiom in axioms) {
         ds.Add(axiom);
-        func.AddOtherDefinitionAxiom(axiom);
+        func.OtherDefinitionAxioms.Add(axiom);
     }
 
     Contract.Assert(func != null);

Source/Core/Monomorphization.cs

@@ -179,7 +179,7 @@ private bool IsFinitelyInstantiable()
       }
       return true;
     }
-    
+
     public override Expr VisitNAryExpr(NAryExpr node)
     {
       if (node.Fun is FunctionCall functionCall)
@@ -343,9 +343,9 @@ abstract class BinderExprMonomorphizer
      */
 
     protected MonomorphizationVisitor monomorphizationVisitor;
-    protected Dictionary<List<Type>, Expr> instanceExprs;
+    protected Dictionary<List<Type>, Expr> instanceExprs { get; }
 
-    public BinderExprMonomorphizer(MonomorphizationVisitor monomorphizationVisitor)
+    protected BinderExprMonomorphizer(MonomorphizationVisitor monomorphizationVisitor)
     {
       this.monomorphizationVisitor = monomorphizationVisitor;
       instanceExprs = new Dictionary<List<Type>, Expr>(new ListComparer<Type>());
@@ -649,7 +649,11 @@ public override Axiom VisitAxiom(Axiom node)
         }
         else
         {
-          splitAxioms.Add(new Axiom(Token.NoToken, expr));
+          var newAxiom = new Axiom(Token.NoToken, expr);
+          if (monomorphizationVisitor.originalAxiomToSplitAxioms.ContainsKey(axiom)) {
+            monomorphizationVisitor.originalAxiomToSplitAxioms[axiom].Add(newAxiom);
+          }
+          splitAxioms.Add(newAxiom);
         }
       }
       return axiom;
@@ -1274,7 +1278,7 @@ class MonomorphizationVisitor : VarDeclOnceStandardVisitor
      * children has been finalized.
      */
 
-    public CoreOptions Options { get; }
+    private CoreOptions Options { get; }
 
     private Program program;
     private Dictionary<string, Function> nameToFunction;
@@ -1293,11 +1297,20 @@ class MonomorphizationVisitor : VarDeclOnceStandardVisitor
     private HashSet<TypeCtorDecl> visitedTypeCtorDecls;
     private HashSet<Function> visitedFunctions;
     private Dictionary<Procedure, Implementation> procToImpl;
+    private readonly HashSet<Function> originalFunctions;
+    private readonly HashSet<Constant> originalConstants;
+    // Note that original axioms refer to axioms of the original program which might have been updated in-place during monomorphization.
+    public readonly Dictionary<Axiom, HashSet<Axiom>> originalAxiomToSplitAxioms;
+    private readonly Dictionary<Axiom, FunctionAndConstantCollector> originalAxiomToOriginalSymbols;
 
     private MonomorphizationVisitor(CoreOptions options, Program program, HashSet<Axiom> polymorphicFunctionAxioms)
     {
       Options = options;
       this.program = program;
+      originalFunctions = program.Functions.ToHashSet();
+      originalConstants = program.Constants.ToHashSet();
+      originalAxiomToSplitAxioms = program.Axioms.ToDictionary(ax => ax, _ => new HashSet<Axiom>());
+      originalAxiomToOriginalSymbols = program.Axioms.ToDictionary(ax => ax, ax => new FunctionAndConstantCollector(ax));
       implInstantiations = new Dictionary<Implementation, Dictionary<List<Type>, Implementation>>();
       nameToImplementation = new Dictionary<string, Implementation>();
       program.TopLevelDeclarations.OfType<Implementation>().Where(impl => impl.TypeParameters.Count > 0).ForEach(
@@ -1351,7 +1364,8 @@ decl is TypeSynonymDecl typeSynonymDecl && typeSynonymDecl.TypeParameters.Count
 
     public static MonomorphizationVisitor Initialize(CoreOptions options, Program program, HashSet<Axiom> polymorphicFunctionAxioms)
     {
-      var monomorphizationVisitor = new MonomorphizationVisitor(options, program, polymorphicFunctionAxioms);
+      var monomorphizationVisitor =
+        new MonomorphizationVisitor(options, program, polymorphicFunctionAxioms);
       // ctorTypes contains all the uninterpreted types created for monomorphizing top-level polymorphic implementations 
       // that must be verified. The types in ctorTypes are reused across different implementations.
       var ctorTypes = new List<Type>();
@@ -1378,13 +1392,118 @@ public static MonomorphizationVisitor Initialize(CoreOptions options, Program pr
         monomorphizationVisitor.polymorphicMapInfos.Values.Select(polymorphicMapInfo =>
           polymorphicMapInfo.CreateDatatypeTypeCtorDecl(polymorphicMapAndBinderSubstituter)).ToList();
       var splitAxioms = polymorphicMapAndBinderSubstituter.Substitute(program);
+      UpdateDependencies(monomorphizationVisitor);
       program.RemoveTopLevelDeclarations(decl => decl is Axiom);
       program.AddTopLevelDeclarations(splitAxioms);
       program.AddTopLevelDeclarations(polymorphicMapDatatypeCtorDecls);
       Contract.Assert(MonomorphismChecker.IsMonomorphic(program));
       return monomorphizationVisitor;
     }
 
+    /// <summary>
+    /// Original program axioms were replaced with new splitAxioms. Program constants and functions that had
+    /// those original axioms as dependencies need their references updated to the new axioms.
+    /// Axioms / functions could both have been instantiated, which adds some complexity.
+    /// </summary>
+    private static void UpdateDependencies(MonomorphizationVisitor monomorphizationVisitor)
+    {
+      foreach (var (originalAxiom, newAxioms) in monomorphizationVisitor.originalAxiomToSplitAxioms)
+      {
+        var originalAxiomFc = monomorphizationVisitor.originalAxiomToOriginalSymbols[originalAxiom];
+        var originalAxiomFunctions = originalAxiomFc.Functions;
+
+        var newAxiomFunctions = new Dictionary<Axiom, HashSet<Function>>();
+        var newAxiomConstants = new Dictionary<Axiom, HashSet<Constant>>();
+
+        foreach (var newAxiom in newAxioms)
+        {
+          var newAxiomFc = new FunctionAndConstantCollector(newAxiom);
+          newAxiomFunctions.Add(newAxiom, newAxiomFc.Functions);
+          newAxiomConstants.Add(newAxiom, newAxiomFc.Constants);
+        }
+
+        // Update function dependencies.
+        foreach (var function in monomorphizationVisitor.originalFunctions)
+        {
+          if (!function.DefinitionAxioms.Contains(originalAxiom))
+          {
+            continue;
+          }
+          if (function.DefinitionAxiom == originalAxiom)
+          {
+            function.DefinitionAxiom = null;
+          }
+          foreach (var newAxiom in newAxioms)
+          {
+            if (function.TypeParameters.Any()) // Polymorphic function
+            {
+              var functionInstantiations =
+                monomorphizationVisitor.functionInstantiations[function].Values.ToHashSet();
+              var instancesToAddDependency = originalAxiomFunctions.Contains(function)
+                ? newAxiomFunctions[newAxiom].Intersect(functionInstantiations)
+                : functionInstantiations;
+              foreach (var inst in instancesToAddDependency)
+              {
+                inst.OtherDefinitionAxioms.Add(newAxiom);
+              }
+            }
+            else // Non-monomorphized function
+            {
+              if (!originalAxiomFunctions.Contains(function) || newAxiomFunctions[newAxiom].Contains(function))
+              {
+                function.OtherDefinitionAxioms.Add(newAxiom);
+              }
+            }
+          }
+          function.OtherDefinitionAxioms.Remove(originalAxiom);
+        }
+
+        // Update constant dependencies.
+        foreach (var constant in monomorphizationVisitor.originalConstants)
+        {
+          if (constant.DefinitionAxioms.Contains(originalAxiom))
+          {
+            foreach (var newAxiom in newAxioms.Where(ax => newAxiomConstants[ax].Contains(constant)))
+            {
+              constant.DefinitionAxioms.Add(newAxiom);
+            }
+            constant.DefinitionAxioms.Remove(originalAxiom);
+          }
+        }
+      }
+    }
+
+    private sealed class FunctionAndConstantCollector : ReadOnlyVisitor
+    {
+      public HashSet<Function> Functions { get; }
+      public HashSet<Constant> Constants { get; }
+
+      public FunctionAndConstantCollector(Axiom axiom)
+      {
+        Functions = new HashSet<Function>();
+        Constants = new HashSet<Constant>();
+        VisitAxiom(axiom);
+      }
+
+      public override Expr VisitNAryExpr(NAryExpr node)
+      {
+        if (node.Fun is FunctionCall functionCall)
+        {
+          Functions.Add(functionCall.Func);
+        }
+        return base.VisitNAryExpr(node);
+      }
+
+      public override Expr VisitIdentifierExpr(IdentifierExpr node)
+      {
+        if (node.Decl is Constant c)
+        {
+          Constants.Add(c);
+        }
+        return base.VisitIdentifierExpr(node);
+      }
+    }
+
     /*
      * Binder expressions may be nested inside other binder expressions. The instantiation of known binder
      * expressions using known type and function instantiations may discover more binder expressions and
@@ -1492,6 +1611,9 @@ private Axiom InstantiateAxiom(Axiom axiom, List<Type> actualTypeParams)
     {
       var axiomExpr = InstantiateBinderExpr((BinderExpr)axiom.Expr, actualTypeParams);
       var instantiatedAxiom = new Axiom(axiom.tok, axiomExpr, axiom.Comment, axiom.Attributes);
+      if (originalAxiomToSplitAxioms.ContainsKey(axiom)) {
+        originalAxiomToSplitAxioms[axiom].Add(instantiatedAxiom);
+      }
       newInstantiatedDeclarations.Add(instantiatedAxiom);
       return instantiatedAxiom;
     }

Source/Core/Parser.cs

@@ -451,7 +451,7 @@ void Function(out List<Declaration>/*!*/ ds) {
 		                                 new Formal(retTyd.tok, retTyd, false, argKv), null, kv);
 		foreach(var axiom in axioms) {
 		   ds.Add(axiom);
-		   func.AddOtherDefinitionAxiom(axiom);
+		   func.OtherDefinitionAxioms.Add(axiom);
 		}
 
 		Contract.Assert(func != null);

Source/ExecutionEngine/VerificationResultCache.cs

@@ -507,7 +507,7 @@ public override Expr VisitNAryExpr(NAryExpr node)
         {
           // We found a function call within a trigger of a quantifier expression, or the function does not take any
           // arguments so we don't expect it ever to sit inside a quantifier.
-          funCall.Func.AddOtherDefinitionAxiom(currentAxiom);
+          funCall.Func.OtherDefinitionAxioms.Add(currentAxiom);
         }
       }
 

Source/VCGeneration/Split.cs

@@ -151,9 +151,16 @@ private void PrintTopLevelDeclarationsForPruning(Program program, Implementation
         }
 
         using var writer = new TokenTextWriter(
-          $"{options.PrintPrunedFile}-{suffix}-{Util.EscapeFilename(implementation.Name)}", false,
+          $"{options.PrintPrunedFile}-{suffix}-{Util.EscapeFilename(implementation.Name)}.bpl", false,
           options.PrettyPrint, options);
-        foreach (var declaration in TopLevelDeclarations ?? program.TopLevelDeclarations) {
+
+        var functionAxioms = 
+          program.Functions.Where(f => f.DefinitionAxioms.Any()).SelectMany(f => f.DefinitionAxioms);
+        var constantAxioms = 
+          program.Constants.Where(f => f.DefinitionAxioms.Any()).SelectMany(c => c.DefinitionAxioms);
+
+        foreach (var declaration in (TopLevelDeclarations ?? program.TopLevelDeclarations).
+                 Except(functionAxioms.Concat(constantAxioms)).ToList()) {
           declaration.Emit(writer, 0);
         }
 

Test/pruning/MonomorphicSplit.bpl

@@ -0,0 +1,39 @@
+// RUN: %parallel-boogie /prune /errorTrace:0 /printPruned:"%t" "%s" > "%t"
+// RUN: %OutputCheck "%s" --file-to-check="%t-after-monomorphicSplit.bpl"
+
+// The following checks are a bit simplistic, but this is
+// on purpose to reduce brittleness. We assume there would now be two uses clauses
+// with one axiom each, and those axioms should not be a conjunction of
+// the instantiations of the original axiom.
+// 
+// Last CHECK-NOT is for ensuring definition axioms are not printed outside
+// uses clauses when using /printPruned.
+
+// CHECK-L: uses
+// CHECK-NEXT-L: axiom
+// CHECK-NOT-L: &&
+// CHECK-L: }
+// CHECK-L: uses
+// CHECK-NEXT-L: axiom
+// CHECK-NOT-L: &&
+// CHECK-L: }
+// CHECK-NOT: axiom
+
+// Related PR #767.
+
+function f1 <T> (x: T) : int uses {
+  axiom (forall <T> x: T :: {f1(x)} f1(x) == 42);
+}
+
+// Both f1 and the axiom will be monomorphized into bool and int instances.
+// Automatic edge inference would already ensure that only the monomorphized instances
+// are incoming, however we want to test here how the new uses clauses are determined.
+//
+// After instantiation, each instance of f1 for some T should *only* have 
+// the T-instantiated axiom as outgoing. 
+
+procedure monomorphicSplit()
+  ensures f1(true) == 42;
+  ensures f1(3) == 42;
+{
+}

Test/pruning/NonmonomorphicSplit.bpl

@@ -0,0 +1,37 @@
+// RUN: %parallel-boogie /prune /trace /errorTrace:0 "%s" > "%t"
+// RUN: %OutputCheck "%s" --file-to-check="%t"
+
+// Related PR #767.
+
+function f1 (x: int) : int;
+function f2 (x: int) : int uses
+{
+  axiom f1(0) == 1 && f2(0) == 2;
+}
+
+// Above axiom will be split into two during monomorphization into
+// axiom f1(0) == 1 and axiom f2(0) == 2.
+// Current implementation ensures that f1(0) == 1 is kept as a dependency
+// of f1, and f2(0) is moved to be a dependency of f2. If there would be
+// other symbols s_0,...,s_n in the axiom, any split axiom containing some
+// symbol s_i (with i = 0,...,n) would be added as a dependency to s_i.
+
+procedure nonMonomorphicSplitPass()
+  ensures f1(0) == 1 && f2(0) == 2;
+{
+}
+// CHECK-L: 1 proof obligation]  verified
+
+function f3 (x: int) : int;
+function f4 (x: int) : int;
+
+axiom f3(0) == 1 && f4(0) == 2;
+
+// This one is expected to fail. We do not want to preserve axioms not inside
+// uses clauses automatically, as this weakens pruning.
+
+procedure nonMonomorphicSplitFail()
+  ensures f3(0) == 1 && f4(0) == 2;
+{
+}
+// CHECK-L: 1 proof obligation]  error