feat: Add newtype support for map and imap

.gitignore

@@ -80,4 +80,4 @@ Source/IntegrationTests/TestFiles/LitTests/LitTest/server/*.bvd
 /Source/IntegrationTests/TestFiles/LitTests/LitTest/cli/log.smt2
 /Source/IntegrationTests/TestFiles/LitTests/LitTest/git-issues/model
 Source/IntegrationTests/TestFiles/LitTests/LitTest/logger/*.html
-Source/IntegrationTests/TestFiles/LitTests/LitTest/comp/*.html
+Source/IntegrationTests/TestFiles/LitTests/LitTest/comp/**/*.html

Source/DafnyCore/Resolver/PreType/PreTypeConstraints.cs

@@ -51,8 +51,28 @@ public DPreType ApproximateReceiverType(PreType preType, [CanBeNull] string memb
       }
       var proxy = (PreTypeProxy)preType;
 
+      var approximateReceiverType = ApproximateReceiverTypeViaBounds(proxy, memberName, out var subProxies);
+      if (approximateReceiverType != null) {
+        return approximateReceiverType;
+      }
+
+      // The bounds didn't give any results, but perhaps one of the proxies visited (in the sub direction) has
+      // associated default advice.
+      foreach (var subProxy in subProxies) {
+        TryApplyDefaultAdviceFor(subProxy);
+        if (proxy.Normalize() is DPreType defaultType) {
+          return defaultType;
+        }
+      }
+
+      // Try once more, in case the application of default advice changed the situation
+      return ApproximateReceiverTypeViaBounds(proxy, memberName, out _);
+    }
+
+    [CanBeNull]
+    private DPreType ApproximateReceiverTypeViaBounds(PreTypeProxy proxy, [CanBeNull] string memberName, out HashSet<PreTypeProxy> subProxies) {
       // If there is a subtype constraint "proxy :> sub<X>", then (if the program is legal at all, then) "sub" must have the member "memberName".
-      var subProxies = new HashSet<PreTypeProxy>();
+      subProxies = new HashSet<PreTypeProxy>();
       foreach (var sub in AllSubBounds(proxy, subProxies)) {
         return sub;
       }
@@ -68,16 +88,7 @@ public DPreType ApproximateReceiverType(PreType preType, [CanBeNull] string memb
         }
       }
 
-      // The bounds didn't give any results, but perhaps one of the proxies visited (in the sub direction) has
-      // associated default advice.
-      foreach (var subProxy in subProxies) {
-        TryApplyDefaultAdviceFor(subProxy);
-        if (proxy.Normalize() is DPreType defaultType) {
-          return defaultType;
-        }
-      }
-
-      return null; // could not be determined
+      return null;
     }
 
     /// <summary>
@@ -536,10 +547,12 @@ public enum CommonConfirmationBag {
       InRealFamily,
       InBoolFamily,
       InCharFamily,
-      InSeqFamily,
       InSetFamily,
       InIsetFamily,
       InMultisetFamily,
+      InSeqFamily,
+      InMapFamily,
+      InImapFamily,
       IsNullableRefType,
       IsBitvector,
       IntLikeOrBitvector,
@@ -590,6 +603,10 @@ private bool ConfirmConstraint(CommonConfirmationBag check, PreType preType, DPr
           return familyDeclName == PreType.TypeNameMultiset;
         case CommonConfirmationBag.InSeqFamily:
           return familyDeclName == PreType.TypeNameSeq;
+        case CommonConfirmationBag.InMapFamily:
+          return familyDeclName == PreType.TypeNameMap;
+        case CommonConfirmationBag.InImapFamily:
+          return familyDeclName == PreType.TypeNameImap;
         case CommonConfirmationBag.IsNullableRefType:
           return DPreType.IsReferenceTypeDecl(pt.Decl);
         case CommonConfirmationBag.IsBitvector:
@@ -695,7 +712,7 @@ private bool ConfirmConstraint(CommonConfirmationBag check, PreType preType, DPr
         case CommonConfirmationBag.IsNewtypeBaseTypeLegacy:
           return pt.Decl is NewtypeDecl || pt.Decl.Name is PreType.TypeNameInt or PreType.TypeNameReal;
         case CommonConfirmationBag.IsNewtypeBaseTypeGeneral:
-          if (familyDeclName is PreType.TypeNameMap or PreType.TypeNameImap || pt.Decl is DatatypeDecl) {
+          if (pt.Decl is DatatypeDecl) {
             // These base types are not yet supported, but they will be soon.
             return false;
           }

Source/DafnyCore/Resolver/PreType/PreTypeResolve.Expressions.cs

@@ -153,8 +153,8 @@ public void ResolveExpression(Expression expr, ResolutionContext resolutionConte
           AddSubtypeConstraint(rangePreType, p.B.PreType, p.B.tok,
             "All elements of display must have some common supertype (got {1}, but needed type or type of previous elements is {0})");
         }
-        var argTypes = new List<PreType>() { domainPreType, rangePreType };
-        expr.PreType = new DPreType(BuiltInTypeDecl(PreType.MapTypeName(e.Finite)), argTypes);
+
+        ResolveMapProducingExpr(e.Finite, "display", expr, domainPreType, rangePreType);
 
       } else if (expr is NameSegment) {
         var e = (NameSegment)expr;
@@ -664,9 +664,6 @@ resolutionContext.CodeContext is ConstantField ||
         foreach (BoundVar v in e.BoundVars) {
           resolver.ResolveType(v.tok, v.Type, resolutionContext, ResolveTypeOptionEnum.InferTypeProxies, null);
           ScopePushAndReport(v, "bound-variable", true);
-          if (v.Type is InferredTypeProxy inferredProxy) {
-            Contract.Assert(!inferredProxy.KeepConstraints);  // in general, this proxy is inferred to be a base type
-          }
         }
         ResolveExpression(e.Range, resolutionContext);
         ConstrainTypeExprBool(e.Range, "range of comprehension must be of type bool (instead got {0})");
@@ -677,8 +674,8 @@ resolutionContext.CodeContext is ConstantField ||
 
         ResolveAttributes(e, resolutionContext, false);
         scope.PopMarker();
-        expr.PreType = new DPreType(BuiltInTypeDecl(PreType.MapTypeName(e.Finite)),
-          new List<PreType>() { e.TermLeft != null ? e.TermLeft.PreType : e.BoundVars[0].PreType, e.Term.PreType });
+
+        ResolveMapProducingExpr(e.Finite, "comprehension", expr, e.TermLeft?.PreType ?? e.BoundVars[0].PreType, e.Term.PreType);
 
       } else if (expr is LambdaExpr) {
         var e = (LambdaExpr)expr;
@@ -754,19 +751,35 @@ private void ResolveCollectionProducingExpr(string typeName, string exprKindSuff
       AddConfirmation(confirmationFamily, expr.PreType, expr.tok, $"{exprKind} used as if it had type {{0}}");
     }
 
-    private void SetupCollectionProducingExpr(string typeName, string exprKind, Expression expr, PreType elementPreType) {
+    private void ResolveMapProducingExpr(bool finite, string exprKindSuffix, Expression expr, PreType keyPreType, PreType valuePreType) {
+      var typeName = PreType.MapTypeName(finite);
+      PreTypeConstraints.CommonConfirmationBag confirmationFamily =
+        finite ? PreTypeConstraints.CommonConfirmationBag.InMapFamily : PreTypeConstraints.CommonConfirmationBag.InImapFamily;
+      var exprKind = $"{typeName} {exprKindSuffix}";
+
+      SetupCollectionProducingExpr(typeName, exprKind, expr, keyPreType, valuePreType);
+      AddConfirmation(confirmationFamily, expr.PreType, expr.tok, $"{exprKind} used as if it had type {{0}}");
+    }
+
+    private void SetupCollectionProducingExpr(string typeName, string exprKind, Expression expr, PreType elementPreType, PreType valuePreType = null) {
       expr.PreType = CreatePreTypeProxy($"{exprKind}");
 
-      var defaultType = new DPreType(BuiltInTypeDecl(typeName), new List<PreType>() { elementPreType });
+      var arguments = valuePreType == null ? new List<PreType>() { elementPreType } : new List<PreType>() { elementPreType, valuePreType };
+      var defaultType = new DPreType(BuiltInTypeDecl(typeName), arguments);
       Constraints.AddDefaultAdvice(expr.PreType, defaultType);
 
       Constraints.AddGuardedConstraint(() => {
         if (expr.PreType.UrAncestor(this) is DPreType dPreType) {
-          if (dPreType.Decl.Name == typeName) {
+          if (dPreType.Decl.Name != typeName) {
+            ReportError(expr, $"{exprKind} used as if it had type {{0}}", expr.PreType);
+          } else if (valuePreType == null) {
             AddSubtypeConstraint(dPreType.Arguments[0], elementPreType, expr.tok,
               $"element type of {exprKind} expected to be {{0}} (got {{1}})");
           } else {
-            ReportError(expr, $"{exprKind} used as if it had type {{0}}", expr.PreType);
+            AddSubtypeConstraint(dPreType.Arguments[0], elementPreType, expr.tok,
+              $"key type of {exprKind} expected to be {{0}} (got {{1}})");
+            AddSubtypeConstraint(dPreType.Arguments[1], valuePreType, expr.tok,
+              $"value type of {exprKind} expected to be {{0}} (got {{1}})");
           }
           return true;
         }
@@ -893,21 +906,33 @@ private PreType ResolveBinaryExpr(IToken tok, BinaryExpr.Opcode opcode, Expressi
             // that only the expected types are allowed.
             var a0 = e0.PreType;
             var a1 = e1.PreType;
-            var left = a0.NormalizeWrtScope() as DPreType;
-            var right = a1.NormalizeWrtScope() as DPreType;
             var familyDeclNameLeft = AncestorName(a0);
             var familyDeclNameRight = AncestorName(a1);
             if (familyDeclNameLeft is PreType.TypeNameMap or PreType.TypeNameImap) {
+              var left = (DPreType)a0.UrAncestor(this);
               Contract.Assert(left.Arguments.Count == 2);
               var st = new DPreType(BuiltInTypeDecl(PreType.TypeNameSet), new List<PreType>() { left.Arguments[0] });
               Constraints.DebugPrint($"    DEBUG: guard applies: Minusable {a0} {a1}, converting to {st} :> {a1}");
-              AddSubtypeConstraint(st, a1, tok,
-                "map subtraction expects right-hand operand to have type {0} (instead got {1})");
+              Constraints.AddDefaultAdvice(a1, st);
+
+              var messageFormat = $"map subtraction expects right-hand operand to have type {st} (instead got {{0}})";
+              Constraints.AddGuardedConstraint(() => {
+                if (a1.UrAncestor(this) is DPreType dPreType) {
+                  if (dPreType.Decl.Name != PreType.TypeNameSet) {
+                    ReportError(e1, messageFormat, a1);
+                  } else {
+                    AddSubtypeConstraint(dPreType.Arguments[0], left.Arguments[0], e1.tok,
+                      $"element type of {PreType.TypeNameSet} expected to be {{0}} (got {{1}})");
+                  }
+                  return true;
+                }
+                return false;
+              });
+              AddConfirmation(PreTypeConstraints.CommonConfirmationBag.InSetFamily, a1, e1.tok, messageFormat);
               return true;
             } else if (familyDeclNameLeft != null || (familyDeclNameRight != null && familyDeclNameRight != PreType.TypeNameSet)) {
               Constraints.DebugPrint($"    DEBUG: guard applies: Minusable {a0} {a1}, converting to {a0} :> {a1}");
-              AddSubtypeConstraint(a0, a1, tok,
-                "type of right argument to - ({0}) must agree with the result type ({1})");
+              AddSubtypeConstraint(a0, a1, tok, "type of right argument to - ({0}) must agree with the result type ({1})");
               return true;
             }
             return false;

Source/IntegrationTests/TestFiles/LitTests/LitTest/dafny0/GeneralNewtypeCollections.dfy

@@ -9,6 +9,8 @@ method Main() {
   Multiset.Test();
   Seq.Test();
   SeqToMultisetConversion.Test();
+  Map.Test();
+  Imap.Test();
 }
 
 module Set {
@@ -407,3 +409,105 @@ module SeqToMultisetConversion {
     print m0' == m1', " ", |m0'|, "\n"; // true 4
   }
 }
+
+module Map {
+  newtype IntMap = m: map<int, real> | true
+  newtype IntSet = s: set<int> | true
+
+  method Test() {
+    var m: IntMap;
+    m := map[];
+    print m, " "; // map[]
+    m := map[6 := 2.0, 19 := 2.1];
+    var n: IntMap := m;
+    print |m|, " "; // 2
+    print 7 in m, " ", 6 in m, "\n"; // false true
+    print m[6], " ", m[19], "\n"; // 2.0 2.1
+    m := m[17 := 3.0][1800 := 0.0][6 := 1.0];
+    print m[6], " ", m[19], " ", m[17], "\n"; // 1.0 2.1 3.0
+
+    m := m + n;
+    // m is now:  map[6 := 2.0, 17 := 3.0, 19 := 2.1, 1800 := 0.0]
+    print |m|, " ", m[6], "\n"; // 4 2.0
+
+    var sf: set<int> := {3, 4, 5, 17, 20};
+    m := m - sf;
+    var nf: IntSet := {3, 1800, 4, 5};
+    m := m - nf;
+    // m is now:  map[6 := 2.0, 19 := 2.1]
+    print |m|, " ", m[6] + m[19], "\n"; // 2 4.1
+
+    var u: map<int, real>;
+    u := m as map<int, real>;
+    m := u as IntMap;
+    var b0 := m is map<int, real>;
+    var b1 := u is IntMap;
+
+    print |m|, " ", |u|, " ", b0, " ", b1, "\n"; // 2 2 true true
+
+    b0 := m == n;
+    b1 := m != n;
+    print b0, " ", b1, "\n"; // true false
+
+    TestComprehensions();
+  }
+
+  method TestComprehensions() {
+    var m: IntMap := map x: int | 2 <= x < 5 :: (2 * x) as real;
+    var n: IntMap := map x: int | 2 <= x < 5 :: 2 * x := 10.0 - x as real;
+    expect |m| == |n| == 3;
+    print 2 in m, " ", 2 in n, " ", 4 in m, " ", 4 in n, "\n"; // true false true true
+    assert 2 * 3 in n;
+    print m[3], " ", n[6], "\n"; // 6.0 7.0
+  }
+}
+
+module Imap {
+  newtype IntImap = m: imap<int, real> | true
+  newtype IntSet = s: set<int> | true
+
+  method Test() {
+    var m: IntImap;
+    m := imap[];
+    print m, " "; // imap[]
+    m := imap[6 := 2.0, 19 := 2.1];
+    var n: IntImap := m;
+    print 7 in m, " ", 6 in m, "\n"; // false true
+    print m[6], " ", m[19], "\n"; // 2.0 2.1
+    m := m[17 := 3.0][1800 := 0.0][6 := 1.0];
+    print m[6], " ", m[19], " ", m[17], "\n"; // 1.0 2.1 3.0
+
+    m := m + n;
+    // m is now:  imap[6 := 2.0, 17 := 3.0, 19 := 2.1, 1800 := 0.0]
+    print m[6], "\n"; // 2.0
+
+    var sf: set<int> := {3, 4, 5, 17, 20};
+    m := m - sf;
+    var nf: IntSet := {3, 1800, 4, 5};
+    m := m - nf;
+    // m is now:  imap[6 := 2.0, 19 := 2.1]
+    print m[6] + m[19], "\n"; // 4.1
+
+    var u: imap<int, real>;
+    u := m as imap<int, real>;
+    m := u as IntImap;
+    var b0 := m is imap<int, real>;
+    var b1 := u is IntImap;
+
+    print b0, " ", b1, "\n"; // true true
+
+    b0 := m == n;
+    b1 := m != n;
+    print b0, " ", b1, "\n"; // true false
+
+    TestComprehensions();
+  }
+
+  method TestComprehensions() {
+    var m: IntImap := imap x: int | 2 <= x < 5 :: (2 * x) as real;
+    var n: IntImap := imap x: int | 2 <= x < 5 :: 2 * x := 10.0 - x as real;
+    print 2 in m, " ", 2 in n, " ", 4 in m, " ", 4 in n, "\n"; // true false true true
+    assert 2 * 3 in n;
+    print m[3], " ", n[6], "\n"; // 6.0 7.0
+  }
+}

Source/IntegrationTests/TestFiles/LitTests/LitTest/dafny0/GeneralNewtypeCollections.dfy.expect

@@ -32,3 +32,21 @@ true 3
 true 3
 true 4
 true 4
+map[] 2 false true
+2.0 2.1
+1.0 2.1 3.0
+4 2.0
+2 4.1
+2 2 true true
+true false
+true false true true
+6.0 7.0
+map[] false true
+2.0 2.1
+1.0 2.1 3.0
+2.0
+4.1
+true true
+true false
+true false true true
+6.0 7.0

docs/dev/news/5175.feat

@@ -0,0 +1 @@
+Support newtypes based on map and imap