Type annotated getters/setters

This commit introduces the following syntax:

  - `m.[:(ty) k]` for `(m.[k] :~ ty)`
  - `m.[:(ty) k <- v]` for `m.[k <- (v :~ ty)]`
  - `m.[:(ty) k] <- k` for `m.[k] <- (v :~ ty)`

This helps disambiguate the selection of getters and setters using
the type of the map's values.

Author: strub

src/ecParser.mly

@@ -45,11 +45,30 @@
   let pfapp_symb loc s ti es =
     PFapp(mk_pfid_symb loc s ti, es)
 
-  let pfget loc ti e1 e2    =
-    pfapp_symb loc EcCoreLib.s_get ti [e1;e2]
-
-  let pfset loc ti e1 e2 e3 =
-    pfapp_symb loc EcCoreLib.s_set ti [e1;e2;e3]
+  let pfget
+    (lc    : EcLocation.t)
+    (ti    : ptyannot option)
+    (codom : pty option)
+    (map   : pformula)
+    (k     : pformula)
+  =
+    ofold
+        (fun codom map -> PFcast (mk_loc lc map, codom))
+        (pfapp_symb lc EcCoreLib.s_get ti [map; k])
+        codom
+
+  let pfset
+    (lc    : EcLocation.t)
+    (ti    : ptyannot option)
+    (codom : pty option)
+    (map   : pformula)
+    (k     : pformula)
+    (v     : pformula)
+  =
+    let v =
+        ofold (fun codom v -> mk_loc (loc v) (PFcast (v, codom)))
+        v codom
+    in pfapp_symb lc EcCoreLib.s_set ti [map; k; v]
 
   let pf_nil loc ti =
     mk_pfid_symb loc EcCoreLib.s_nil ti
@@ -1053,15 +1072,19 @@ sform_u(P):
 | x=mident
    { PFmem x }
 
-| se=sform_r(P) DLBRACKET ti=tvars_app? e=loc(plist1(form_r(P), COMMA)) RBRACKET
+| se=sform_r(P) DLBRACKET
+    ti=tvars_app? codom=prefix(COLON, paren(loc(type_exp)))?
+    e=loc(plist1(form_r(P), COMMA))
+  RBRACKET
    { let e = List.reduce1 (fun _ -> lmap (fun x -> PFtuple x) e) (unloc e) in
-     pfget (EcLocation.make $startpos $endpos) ti se e }
+     pfget (EcLocation.make $startpos $endpos) ti codom se e }
 
 | se=sform_r(P) DLBRACKET
-    ti=tvars_app? e1=loc(plist1(form_r(P), COMMA))LARROW e2=form_r(P)
+    ti=tvars_app? codom=prefix(COLON, paren(loc(type_exp)))?
+    e1=loc(plist1(form_r(P), COMMA)) LARROW e2=form_r(P)
   RBRACKET
    { let e1 = List.reduce1 (fun _ -> lmap (fun x -> PFtuple x) e1) (unloc e1) in
-     pfset (EcLocation.make $startpos $endpos) ti se e1 e2 }
+     pfset (EcLocation.make $startpos $endpos) ti codom se e1 e2 }
 
 | x=sform_r(P) s=pside_force
    { PFside (x, s) }
@@ -1296,8 +1319,11 @@ lvalue_u:
 | LPAREN p=plist2(qident, COMMA) RPAREN
    { PLvTuple p }
 
-| x=lvalue_var DLBRACKET ti=tvars_app? e=plist1(expr, COMMA) RBRACKET
-   { PLvMap (x, ti, e) }
+| x=lvalue_var DLBRACKET
+    ti=tvars_app? codom=prefix(COLON, paren(loc(type_exp)))?
+    e=plist1(expr, COMMA)
+  RBRACKET
+   { PLvMap (x, ti, codom, e) }
 
 %inline lvalue:
 | x=loc(lvalue_u) { x }

src/ecParsetree.ml

@@ -148,7 +148,7 @@ type glob_or_var =
 type plvalue_r =
   | PLvSymbol of pqsymbol
   | PLvTuple  of pqsymbol list
-  | PLvMap    of pqsymbol * ptyannot option * pexpr list
+  | PLvMap    of pqsymbol * ptyannot option * pty option * pexpr list
 
 and plvalue = plvalue_r located
 

src/ecTyping.ml

@@ -2779,14 +2779,17 @@ and translvalue ue (env : EcEnv.env) lvalue =
       let ty = ttuple (List.map snd xs) in
       Lval (LvTuple xs), ty
 
-  | PLvMap (x, tvi, es) ->
+  | PLvMap (x, tvi, codom, es) ->
       let tvi = tvi |> omap (transtvi env ue) in
-      let codomty = UE.fresh ue in
+      let codom = Option.map (transty tp_relax env ue) codom in
+      let codom = ofdfl (fun () -> UE.fresh ue) codom in
       let pv, xty = trans_pv env x in
       let e, ety = List.split (List.map (transexp env `InProc ue) es) in
+
       let e, ety = e_tuple e, ttuple ety in
+
       let name = ([], EcCoreLib.s_set) in
-      let esig = [xty; ety; codomty] in
+      let esig = [xty; ety; codom] in
       let ops = select_exp_op env `InProc None name ue tvi (esig, None) in
 
       match ops with
@@ -2801,7 +2804,7 @@ and translvalue ue (env : EcEnv.env) lvalue =
           let esig = Tuni.subst_dom uidmap esig in
           let esig = toarrow esig xty in
           unify_or_fail env ue lvalue.pl_loc ~expct:esig opty;
-          LvMap ((p, tys), pv, e, xty), codomty
+          LvMap ((p, tys), pv, e, xty), codom
 
       | [_] ->
           let uidmap = UE.assubst ue in
@@ -3063,7 +3066,7 @@ and trans_form_or_pattern env mode ?mv ?ps ue pf tt =
 
     | PFcast (pf, pty) ->
         let ty = transty tp_relax env ue pty in
-        let aout = transf env pf in
+        let aout = transf env ~tt:ty pf in
         unify_or_fail env ue pf.pl_loc ~expct:ty aout.f_ty; aout
 
     | PFmem _ -> tyerror f.pl_loc env MemNotAllowed
@@ -3554,8 +3557,8 @@ and trans_memtype env ue (pmemtype : pmemtype) : memtype =
   List.fold_left add_decl mt pmemtype
 
 (* -------------------------------------------------------------------- *)
-and transexp env mode ue { pl_desc = Expr e; pl_loc = loc; } =
-  let f = trans_form_or_pattern env (`Expr mode) ue e None in
+and transexp env ?tt mode ue { pl_desc = Expr e; pl_loc = loc; } =
+  let f = trans_form_or_pattern env (`Expr mode) ue e tt in
   let m = Option.value ~default:mhr (EcEnv.Memory.get_active env) in
   let e =
     try

src/ecTyping.mli

@@ -213,7 +213,7 @@ val trans_gbinding : env -> EcUnify.unienv -> pgtybindings ->
 
 (* -------------------------------------------------------------------- *)
 val transexp :
-  env -> [`InProc|`InOp] -> EcUnify.unienv -> pexpr -> expr * ty
+  env -> ?tt:ty -> [`InProc|`InOp] -> EcUnify.unienv -> pexpr -> expr * ty
 
 val transexpcast :
   env -> [`InProc|`InOp] -> EcUnify.unienv -> ty -> pexpr -> expr

tests/annotated-getter-setter.ec

@@ -0,0 +1,39 @@
+(* -------------------------------------------------------------------- *)
+require import AllCore.
+
+type u32, u64.
+
+type key.
+type map.
+
+theory GS32.
+  op "_.[_]" : map -> key -> u32.
+  op "_.[_<-_]" : map -> key -> u32 -> map.
+end GS32.
+
+theory GS64.
+  op "_.[_]" : map -> key -> u64.
+  op "_.[_<-_]" : map -> key -> u64 -> map.
+end GS64.
+
+import GS32 GS64.
+
+op myget32 (m : map) k = m.[:(u32) k].
+
+op myset32 (m : map) k v = m.[:(u32) k <- v].
+
+op myget64 (m : map) k = m.[:(u64) k].
+
+op myset64 (m : map) k v = m.[:(u64) k <- v].
+
+module M = {
+  proc f32(m : map, k, v) = {
+    m.[:(u32) k] <- v;
+    return m;
+  }
+
+  proc f64(m : map, k, v) = {
+    m.[:(u64) k] <- v;
+    return m;
+  }
+}.