Fixes lost dependencies for pruning during monomorphization.

Source/Core/AST/Absy.cs

@@ -1429,7 +1429,23 @@ public class Constant : Variable
     // from all other constants.
     public readonly bool Unique;
 
-    public IList<Axiom> DefinitionAxioms { get; }
+    public IList<Axiom> DefinitionAxioms => definitionAxioms;
+
+    private IList<Axiom> definitionAxioms = new List<Axiom>();
+
+    public bool RemoveDefinitionAxiom(Axiom axiom)
+    {
+      Contract.Requires(axiom != null);
+
+      return definitionAxioms.Remove(axiom);
+    }
+
+    public void AddDefinitionAxiom(Axiom axiom)
+    {
+      Contract.Requires(axiom != null);
+
+      DefinitionAxioms.Add(axiom);
+    }
 
     public Constant(IToken /*!*/ tok, TypedIdent /*!*/ typedIdent)
       : this(tok, typedIdent, true)
@@ -1452,7 +1468,7 @@ public Constant(IToken /*!*/ tok, TypedIdent /*!*/ typedIdent,
       Contract.Requires(typedIdent.Name != null && typedIdent.Name.Length > 0);
       Contract.Requires(typedIdent.WhereExpr == null);
       this.Unique = unique;
-      this.DefinitionAxioms = definitionAxioms ?? new List<Axiom>();
+      this.definitionAxioms = definitionAxioms ?? new List<Axiom>();
     }
 
     public override bool IsMutable => false;
@@ -2058,6 +2074,13 @@ public void AddOtherDefinitionAxiom(Axiom axiom)
       otherDefinitionAxioms.Add(axiom);
     }
 
+    public bool RemoveOtherDefinitionAxiom(Axiom axiom)
+    {
+      Contract.Requires(axiom != null);
+
+      return otherDefinitionAxioms.Remove(axiom);
+    }
+
     private bool neverTrigger;
     private bool neverTriggerComputed;
 

Source/Core/Monomorphization.cs

@@ -7,7 +7,7 @@
 
 namespace Microsoft.Boogie
 {
-  public class MonomorphismChecker : ReadOnlyVisitor
+  sealed public class MonomorphismChecker : ReadOnlyVisitor
   {
     // This visitor checks if the program is already monomorphic.
 
@@ -131,11 +131,13 @@ class MonomorphizableChecker : ReadOnlyVisitor
      * This visitor checks if the program is monomorphizable. It calculates one status value from
      * the enum MonomorphizableStatus for the program.
      */
-    public static MonomorphizableStatus IsMonomorphizable(Program program, out HashSet<Axiom> polymorphicFunctionAxioms)
+    public static MonomorphizableStatus IsMonomorphizable(Program program, out HashSet<Axiom> polymorphicFunctionAxioms,
+      out Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)> axiomsThatWillLoseQuantifiers)
     {
       var checker = new MonomorphizableChecker(program);
       checker.VisitProgram(program);
       polymorphicFunctionAxioms = checker.polymorphicFunctionAxioms;
+      axiomsThatWillLoseQuantifiers = checker.axiomsThatWillLoseQuantifiers;
       if (!checker.isMonomorphizable)
       {
         return MonomorphizableStatus.UnhandledPolymorphism;
@@ -148,7 +150,9 @@ public static MonomorphizableStatus IsMonomorphizable(Program program, out HashS
     }
 
     private bool isMonomorphizable;
+    private Axiom curAxiom;
     private HashSet<Axiom> polymorphicFunctionAxioms;
+    private Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)> axiomsThatWillLoseQuantifiers;
     private Graph<TypeVariable> typeVariableDependencyGraph; // (T,U) in this graph iff T flows to U
     private HashSet<Tuple<TypeVariable, TypeVariable>> strongDependencyEdges; // (T,U) in this set iff a type constructed from T flows into U
 
@@ -158,6 +162,8 @@ private MonomorphizableChecker(Program program)
       polymorphicFunctionAxioms = program.TopLevelDeclarations.OfType<Function>()
         .Where(f => f.TypeParameters.Count > 0 && f.DefinitionAxiom != null)
         .Select(f => f.DefinitionAxiom).ToHashSet();
+      axiomsThatWillLoseQuantifiers = new Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)>();
+      curAxiom = null;
       typeVariableDependencyGraph = new Graph<TypeVariable>();
       strongDependencyEdges = new HashSet<Tuple<TypeVariable, TypeVariable>>();
     }
@@ -179,7 +185,51 @@ private bool IsFinitelyInstantiable()
       }
       return true;
     }
-
+
+    public override Axiom VisitAxiom(Axiom node)
+    {
+      curAxiom = node;
+      var res = base.VisitAxiom(node);
+      curAxiom = null;
+      return res;
+    }
+
+    public override QuantifierExpr VisitQuantifierExpr(QuantifierExpr node)
+    {
+      if (curAxiom != null && node.TypeParameters.Any() && !node.Dummies.Any() &&
+          node.Triggers != null && node.Triggers.Tr.Any())
+      {
+        // This expression will be instantiated into non-quantified expressions.
+        // Since pruning (currently) depends on triggers and `uses` clauses,
+        // if the enclosing axiom is not already inside a uses clause, it can
+        // be incorrectly pruned away.
+        if (!axiomsThatWillLoseQuantifiers.ContainsKey(curAxiom))
+        {
+          axiomsThatWillLoseQuantifiers[curAxiom] = (new HashSet<Function>(), new HashSet<Constant>());
+        }
+
+        var fc = new MonomorphizationVisitor.FunctionAndConstantCollector();
+        node.Triggers.Tr.ForEach(expr => fc.VisitExpr(expr));
+        var triggerConstants = fc.GetConstants();
+        var triggerFunctions = fc.GetFunctions();
+
+        var functionsAndConstants = axiomsThatWillLoseQuantifiers[curAxiom];
+
+        if (triggerConstants.Any())
+        {
+          triggerConstants.ForEach(c => functionsAndConstants.Item2.Add(c));
+        }
+
+        if (triggerFunctions.Any())
+        {
+          triggerFunctions.ForEach(f => functionsAndConstants.Item1.Add(f));
+        }
+
+        axiomsThatWillLoseQuantifiers[curAxiom] = functionsAndConstants;
+      }
+      return base.VisitQuantifierExpr(node);
+    }
+
     public override Expr VisitNAryExpr(NAryExpr node)
     {
       if (node.Fun is FunctionCall functionCall)
@@ -343,7 +393,7 @@ abstract class BinderExprMonomorphizer
      */
 
     protected MonomorphizationVisitor monomorphizationVisitor;
-    protected Dictionary<List<Type>, Expr> instanceExprs;
+    public Dictionary<List<Type>, Expr> instanceExprs { get; }
 
     public BinderExprMonomorphizer(MonomorphizationVisitor monomorphizationVisitor)
     {
@@ -649,7 +699,11 @@ public override Axiom VisitAxiom(Axiom node)
         }
         else
         {
-          splitAxioms.Add(new Axiom(Token.NoToken, expr));
+          var newAxiom = new Axiom(Token.NoToken, expr);
+          if(monomorphizationVisitor.oldToNewAxioms.ContainsKey(axiom)) {
+            monomorphizationVisitor.oldToNewAxioms[axiom].Add(newAxiom);
+          }
+          splitAxioms.Add(newAxiom);
         }
       }
       return axiom;
@@ -1281,7 +1335,7 @@ class MonomorphizationVisitor : VarDeclOnceStandardVisitor
     private Dictionary<string, Procedure> nameToProcedure;
     private Dictionary<string, Implementation> nameToImplementation;
     private Dictionary<string, TypeCtorDecl> nameToTypeCtorDecl;
-    private Dictionary<Function, Dictionary<List<Type>, Function>> functionInstantiations;
+    public Dictionary<Function, Dictionary<List<Type>, Function>> functionInstantiations;
     private Dictionary<Procedure, Dictionary<List<Type>, Procedure>> procInstantiations;
     private Dictionary<Implementation, Dictionary<List<Type>, Implementation>> implInstantiations;
     private Dictionary<TypeCtorDecl, Dictionary<List<Type>, TypeCtorDecl>> typeInstantiations;
@@ -1293,11 +1347,19 @@ class MonomorphizationVisitor : VarDeclOnceStandardVisitor
     private HashSet<TypeCtorDecl> visitedTypeCtorDecls;
     private HashSet<Function> visitedFunctions;
     private Dictionary<Procedure, Implementation> procToImpl;
+    public readonly HashSet<Function> originalFunctions;
+    public readonly HashSet<Constant> originalConstants;
+    public readonly Dictionary<Axiom, HashSet<Axiom>> oldToNewAxioms; // TODO: optimize using this
+    public readonly Dictionary<Axiom, Axiom> oldAxiomToOriginalAxioms; //
 
     private MonomorphizationVisitor(CoreOptions options, Program program, HashSet<Axiom> polymorphicFunctionAxioms)
     {
       Options = options;
       this.program = program;
+      originalFunctions = program.Functions.ToHashSet();
+      originalConstants = program.Constants.ToHashSet();
+      oldToNewAxioms = program.Axioms.ToDictionary(ax => ax, _ => new HashSet<Axiom>());
+      oldAxiomToOriginalAxioms = program.Axioms.ToDictionary(ax => ax, ax => (Axiom) ax.Clone() );
       implInstantiations = new Dictionary<Implementation, Dictionary<List<Type>, Implementation>>();
       nameToImplementation = new Dictionary<string, Implementation>();
       program.TopLevelDeclarations.OfType<Implementation>().Where(impl => impl.TypeParameters.Count > 0).ForEach(
@@ -1349,9 +1411,11 @@ decl is TypeSynonymDecl typeSynonymDecl && typeSynonymDecl.TypeParameters.Count
         decl is Axiom axiom && polymorphicFunctionAxioms.Contains(axiom));
     }
 
-    public static MonomorphizationVisitor Initialize(CoreOptions options, Program program, HashSet<Axiom> polymorphicFunctionAxioms)
+    public static MonomorphizationVisitor Initialize(CoreOptions options, Program program, HashSet<Axiom> polymorphicFunctionAxioms,
+      Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)> axiomsThatWillLoseQuantifiers)
     {
-      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 +1442,139 @@ public static MonomorphizationVisitor Initialize(CoreOptions options, Program pr
         monomorphizationVisitor.polymorphicMapInfos.Values.Select(polymorphicMapInfo =>
           polymorphicMapInfo.CreateDatatypeTypeCtorDecl(polymorphicMapAndBinderSubstituter)).ToList();
       var splitAxioms = polymorphicMapAndBinderSubstituter.Substitute(program);
+      UpdateDependencies(splitAxioms, monomorphizationVisitor, axiomsThatWillLoseQuantifiers);
       program.RemoveTopLevelDeclarations(decl => decl is Axiom);
       program.AddTopLevelDeclarations(splitAxioms);
       program.AddTopLevelDeclarations(polymorphicMapDatatypeCtorDecls);
       Contract.Assert(MonomorphismChecker.IsMonomorphic(program));
       return monomorphizationVisitor;
     }
 
+    private static void UpdateDependencies(List<Axiom> splitAxioms, MonomorphizationVisitor monomorphizationVisitor,
+      Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)> axiomsThatWillLoseQuantifiers)
+    {
+      // 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.
+      // Furthermore, current pruning depends on symbols in triggers in quantified axioms to compute dependencies,
+      // which might have been lost during instantiation (i.e., axiom [forall | exists] <T> :: ...). These are tracked
+      // separately and added as dependencies to corresponding constants and functions
+
+      foreach (var (oldAxiom, newAxioms) in monomorphizationVisitor.oldToNewAxioms)
+      {
+        var oldAxiomFC = new FunctionAndConstantCollector();
+        var unmodifiedOldAxiom = monomorphizationVisitor.oldAxiomToOriginalAxioms[oldAxiom];
+        oldAxiomFC.Visit(unmodifiedOldAxiom);
+        var oldAxiomFunctions = oldAxiomFC.GetFunctions();
+
+        Dictionary<Axiom, HashSet<Function>> newAxiomFunctions = new Dictionary<Axiom, HashSet<Function>>();
+        Dictionary<Axiom, HashSet<Constant>> newAxiomConstants = new Dictionary<Axiom, HashSet<Constant>>();
+
+        foreach (var newAxiom in newAxioms)
+        {
+          var newAxiomFC = new FunctionAndConstantCollector();
+          newAxiomFC.Visit(newAxiom);
+          newAxiomFunctions.Add(newAxiom, newAxiomFC.GetFunctions());
+          newAxiomConstants.Add(newAxiom, newAxiomFC.GetConstants());
+        }
+
+        void UpdateFunctionDependencies(HashSet<Function> functions, bool forceDependency = false)
+        {
+          foreach (var function in functions)
+          {
+            if (function.DefinitionAxioms.Contains(oldAxiom) || forceDependency)
+            {
+              if (function.DefinitionAxiom == oldAxiom)
+              {
+                function.DefinitionAxiom = null; // is moving DefinitionAxiom to OtherDefinitionAxioms OK?
+              }
+              foreach (var newAxiom in newAxioms)
+              {
+                if (function.TypeParameters.Any()) // Polymorphic function
+                {
+                  var instancesToAddDependency = oldAxiomFunctions.Contains(function)
+                    ? newAxiomFunctions[newAxiom]
+                    : monomorphizationVisitor.functionInstantiations[function].Values.ToHashSet();
+                  foreach (var inst in instancesToAddDependency)
+                  {
+                    inst.AddOtherDefinitionAxiom(newAxiom);
+                  }
+                }
+                else
+                {
+                  // Non-monomorphized function
+                  if (!oldAxiomFunctions.Contains(function) || newAxiomFunctions[newAxiom].Contains(function))
+                  {
+                    function.AddOtherDefinitionAxiom(newAxiom);
+                  }
+                }
+              }
+              function.RemoveOtherDefinitionAxiom(oldAxiom);
+            }
+          }
+        }
+
+        void UpdateConstantDependencies(HashSet<Constant> constants)
+        {
+          foreach (var constant in constants)
+          {
+            if (constant.DefinitionAxioms.Contains(oldAxiom))
+            {
+              foreach (var newAxiom in newAxioms.Where(ax => newAxiomConstants[ax].Contains(constant)))
+              {
+                constant.AddDefinitionAxiom(newAxiom);
+              }
+              constant.RemoveDefinitionAxiom(oldAxiom);
+            }
+          }
+        }
+
+        UpdateConstantDependencies(monomorphizationVisitor.originalConstants);
+        UpdateFunctionDependencies(monomorphizationVisitor.originalFunctions);
+
+        if (axiomsThatWillLoseQuantifiers.TryGetValue(oldAxiom, out var pair))
+        {
+          UpdateFunctionDependencies(pair.Item1, true);
+          UpdateConstantDependencies(pair.Item2);
+        }
+      }
+    }
+
+    public class FunctionAndConstantCollector : ReadOnlyVisitor
+    {
+      private readonly HashSet<Function> functions = new HashSet<Function>();
+      private readonly HashSet<Constant> constants = new HashSet<Constant>();
+      private bool hasTriggers = false;
+
+      public override Expr VisitNAryExpr(NAryExpr node)
+      {
+        if (node.Fun is FunctionCall functionCall)
+        {
+          functions.Add(functionCall.Func);
+        }
+        return base.VisitNAryExpr(node);
+      }
+
+      public override Trigger VisitTrigger(Trigger node)
+      {
+        hasTriggers = true;
+        return base.VisitTrigger(node);
+      }
+
+      public override Expr VisitIdentifierExpr(IdentifierExpr node)
+      {
+        if (node.Decl is Constant c)
+        {
+          constants.Add(c);
+        }
+        return base.VisitIdentifierExpr(node);
+      }
+
+      public HashSet<Function> GetFunctions() => functions;
+      public HashSet<Constant> GetConstants() => constants;
+      public bool HasTriggers() => hasTriggers;
+    }
+
     /*
      * 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 +1682,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(oldToNewAxioms.ContainsKey(axiom)) {
+        oldToNewAxioms[axiom].Add(instantiatedAxiom);
+      }
       newInstantiatedDeclarations.Add(instantiatedAxiom);
       return instantiatedAxiom;
     }
@@ -1888,11 +2081,13 @@ public class Monomorphizer
 
     public static MonomorphizableStatus Monomorphize(CoreOptions options, Program program)
     {
-      var monomorphizableStatus = MonomorphizableChecker.IsMonomorphizable(program, out var polymorphicFunctionAxioms);
+      var monomorphizableStatus = MonomorphizableChecker.IsMonomorphizable(program,
+        out var polymorphicFunctionAxioms, out var axiomsThatWillLoseQuantifiers);
       if (monomorphizableStatus == MonomorphizableStatus.Monomorphizable)
       {
-        var monomorphizationVisitor = MonomorphizationVisitor.Initialize(options, program, polymorphicFunctionAxioms);
-        program.monomorphizer = new Monomorphizer(monomorphizationVisitor);
+        var monomorphizationVisitor =
+          MonomorphizationVisitor.Initialize(options, program, polymorphicFunctionAxioms, axiomsThatWillLoseQuantifiers);
+        program.monomorphizer = new Monomorphizer(monomorphizationVisitor, axiomsThatWillLoseQuantifiers);
       }
       return monomorphizableStatus;
     }
@@ -1936,9 +2131,12 @@ public List<Type> GetTypeInstantiation(TypeCtorDecl decl)
     }
 
     private MonomorphizationVisitor monomorphizationVisitor;
-    private Monomorphizer(MonomorphizationVisitor monomorphizationVisitor)
+    private Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)> axiomsThatWillLoseQuantifiers;
+    private Monomorphizer(MonomorphizationVisitor monomorphizationVisitor,
+      Dictionary<Axiom, (HashSet<Function>, HashSet<Constant>)> axiomsThatWillLoseQuantifiers)
     {
       this.monomorphizationVisitor = monomorphizationVisitor;
+      this.axiomsThatWillLoseQuantifiers = axiomsThatWillLoseQuantifiers;
     }
   }
 }

Test/prover/z3mutl.bpl.expect

@@ -1,7 +1,7 @@
 z3mutl.bpl(22,5): Error: this assertion could not be proved
 Execution trace:
     z3mutl.bpl(7,1): start
-    z3mutl.bpl(13,1): L2
+    z3mutl.bpl(16,1): L3
     z3mutl.bpl(22,1): L5
 
 Boogie program verifier finished with 0 verified, 1 error