optimisation: annotations to track (reads from) mutable declarations …

…and adapt `Ir.VarE` (#4637)

By making reads from mutable variables distinguishable, we restore the referential transparency of `VarE (Const, id)`. This was the reason why #4611 failed (and had to be reverted in #4623).

This opens up a bunch of optimisation possibilities
- `VarE` path compression
- dead `LetD` elimination
- `BlockE`/`NewObjE` simplification
- nested `BlockE` merging (given the defs are disjoint and outer uses don't reference inner ones; use `freevars.ml`)
- etc.

In the long run this may allow us to extend the IR effect system with {`Load`, `Store`} effects, and thus pave the way for transactional memory.

We now have a note on `Syntax.VarE`'s identifier, being immutable by default. When performing type checking, this field can change to signify a mutable variable source after looking up the identifier in the environment. Then desugaring uses mutability info to create either `IR.VarE (Const, _)` (for immutable) or a`IR.VarE (Var, _) (for mutable). The IR type checker then asserts that the right kind of access is used.

The Wasm output should be unchanged.

_NOTE_: with this PR `git revert 13d9d57` can be done and it passes. So it seems to work 😀

src/codegen/compile.ml

@@ -11911,7 +11911,7 @@ and compile_exp_with_hint (env : E.t) ae sr_hint exp =
 
   | PrimE (p, es) ->
     compile_prim_invocation (env : E.t) ae p es exp.at
-  | VarE var ->
+  | VarE (_, var) ->
     Var.get_val env ae var
   | AssignE (e1,e2) ->
     SR.unit,
@@ -12442,7 +12442,7 @@ and compile_const_exp env pre_ae exp : Const.t * (E.t -> VarEnv.t -> unit) =
       | Type.Local, Type.Returns, [], PrimE (prim, prim_args) when
           inlineable_prim prim &&
           List.length args = List.length prim_args &&
-          List.for_all2 (fun p a -> a.it = VarE p.it) args prim_args ->
+          List.for_all2 (fun p a -> a.it = VarE (Const, p.it)) args prim_args ->
         Const.PrimWrapper prim
       | _, _, _, _ -> Const.Complicated
     in
@@ -12465,7 +12465,7 @@ and compile_const_exp env pre_ae exp : Const.t * (E.t -> VarEnv.t -> unit) =
       let ae' = extend ae in
       fill1 env ae';
       fill2 env ae')
-  | VarE v ->
+  | VarE (_, v) ->
     let c =
       match VarEnv.lookup_var pre_ae v with
       | Some (VarEnv.Const c) -> c

src/ir_def/arrange_ir.ml

@@ -13,7 +13,7 @@ let prim_ty p = typ (Type.Prim p)
 let kind k = Atom (Type.string_of_kind k)
 
 let rec exp e = match e.it with
-  | VarE i              -> "VarE"    $$ [id i]
+  | VarE (_, i)         -> "VarE"    $$ [id i] (* FIXME: EXPOSE *)
   | LitE l              -> "LitE"    $$ [lit l]
   | PrimE (p, es)       -> "PrimE"   $$ [prim p] @ List.map exp es
   | AssignE (le1, e2)   -> "AssignE" $$ [lexp le1; exp e2]

src/ir_def/check_ir.ml

@@ -387,12 +387,19 @@ let rec check_exp env (exp:Ir.exp) : unit =
     "inferred effect not a subtype of expected effect";
   (* check typing *)
   begin match exp.it with
-  | VarE id ->
-    let { typ; loc_known; const } =
+  | VarE (m, id) ->
+    let { typ; _ } =
       try T.Env.find id env.vals
       with Not_found -> error env exp.at "unbound variable %s" id
     in
-    T.as_immut typ <: t
+    begin match m with
+    | Const ->
+       assert (not (T.is_mut typ));
+       typ <: t
+    | Var ->
+       assert (T.is_mut typ);
+       T.as_immut typ <: t
+    end
   | LitE lit ->
     T.Prim (type_lit env lit exp.at) <: t
   | PrimE (p, es) ->
@@ -851,7 +858,7 @@ let rec check_exp env (exp:Ir.exp) : unit =
   if exp.note.Note.const
   then begin
     match exp.it with
-    | VarE id -> check_var "VarE" id
+    | VarE (Const, id) -> check_var "VarE" id
     | FuncE (x, s, c, tp, as_ , ts, body) ->
       check (s = T.Local) "constant FuncE cannot be of shared sort";
       if env.lvl = NotTopLvl then

src/ir_def/construct.ml

@@ -76,7 +76,9 @@ let wildP =
 (* Primitives *)
 
 let varE (id, typ) =
-  { it = VarE id; at = no_region; note = Note.{ def with typ = T.as_immut typ } }
+  { it = VarE ((if T.is_mut typ then Var else Const), id)
+  ; at = no_region
+  ; note = Note.{ def with typ = T.as_immut typ } }
 
 let varLE (id, typ) =
   { it = VarLE id; at = no_region; note = typ }

src/ir_def/freevars.ml

@@ -103,7 +103,7 @@ let fields fs = unions (fun f ->
 ) fs
 
 let rec exp e : f = match e.it with
-  | VarE i              -> id i
+  | VarE (_, i)         -> id i
   | LitE l              -> M.empty
   | PrimE (_, es)       -> exps es
   | AssignE (e1, e2)    -> lexp e1 ++ exp e2

src/ir_def/ir.ml

@@ -58,7 +58,7 @@ type arg = (string, Type.typ) Source.annotated_phrase
 type exp = exp' phrase
 and exp' =
   | PrimE of (prim * exp list)                 (* primitive *)
-  | VarE of id                                 (* variable *)
+  | VarE of mut * id                           (* variable *)
   | LitE of lit                                (* literal *)
   | AssignE of lexp * exp                      (* assignment *)
   | BlockE of (dec list * exp)                 (* block *)

src/ir_def/rename.ml

@@ -30,7 +30,7 @@ let rec prim rho p =
 
 and exp rho e  =  {e with it = exp' rho e.it}
 and exp' rho = function
-  | VarE i              -> VarE (id rho i)
+  | VarE (m, i)         -> VarE (m, id rho i)
   | LitE _ as e         -> e
   | PrimE (p, es)       -> PrimE (prim rho p, List.map (exp rho) es)
   | ActorE (ds, fs, { meta; preupgrade; postupgrade; heartbeat; timer; inspect }, t) ->

src/ir_interpreter/interpret_ir.ml

@@ -300,7 +300,7 @@ and interpret_exp_mut env exp (k : V.value V.cont) =
   last_env := env;
   Profiler.bump_region exp.at ;
   match exp.it with
-  | VarE id ->
+  | VarE (_, id) ->
     (match Lib.Promise.value_opt (find id env.vals) with
     | Some v -> k v
     | None -> trap exp.at "accessing identifier before its definition"

src/ir_passes/async.ml

@@ -130,7 +130,7 @@ let new_nary_async_reply ts =
 
 let let_eta e scope =
   match e.it with
-  | VarE _ -> scope e (* pure, so reduce *)
+  | VarE (Const, _) -> scope e (* pure, so reduce *)
   | _  ->
     let f = fresh_var "x" (typ e) in
     letD f e :: (scope (varE f)) (* maybe impure; sequence *)
@@ -162,7 +162,7 @@ let let_seq ts e d_of_vs =
 (* name e in f unless named already *)
 let ensureNamed e f =
   match e.it with
-  | VarE v -> f (var v (typ e))
+  | VarE (Const, v) -> f (var v (typ e))
   | _ ->
     let v = fresh_var "v" (typ e) in
     blockE [letD v e] (f v)
@@ -251,8 +251,8 @@ let transform prog =
   and t_exp' (exp:exp) =
     let exp' = exp.it in
     match exp' with
-    | LitE _ -> exp'
-    | VarE id -> exp'
+    | LitE _
+    | VarE (_, _) -> exp'
     | AssignE (exp1, exp2) ->
       AssignE (t_lexp exp1, t_exp exp2)
     | PrimE (CPSAwait (Fut, cont_typ), [a; krb]) ->

src/ir_passes/await.ml

@@ -48,7 +48,7 @@ let ( -@- ) k exp2 =
     varE v -*- exp2
   | MetaCont (typ0, k) ->
     match exp2.it with
-    | VarE v -> k (var v (typ exp2))
+    | VarE (Const, v) -> k (var v (typ exp2))
     | _ ->
       let u = fresh_var "u" typ0 in
       letE u exp2 (k u)

src/ir_passes/const.ml

@@ -100,7 +100,7 @@ let set_lazy_const e lb =
 let rec exp lvl (env : env) e : Lbool.t =
   let lb =
     match e.it with
-    | VarE v -> (find v env).const
+    | VarE (_, v) -> (find v env).const (*FIXME: use the mutability marker?*)
     | FuncE (x, s, c, tp, as_ , ts, body) ->
       exp_ NotTopLvl (args NotTopLvl env as_) body;
       begin match s, lvl with

src/ir_passes/eq.ml

@@ -202,8 +202,7 @@ and t_exp env (e : Ir.exp) =
   { e with it = t_exp' env e.it }
 
 and t_exp' env = function
-  | LitE l -> LitE l
-  | VarE id -> VarE id
+  | (LitE _ | VarE _) as e -> e
   | PrimE (RelPrim (ot, Operator.EqOp), [exp1; exp2]) when T.singleton ot ->
     (* Handle singletons here, but beware of side-effects *)
     let e1 = t_exp env exp1 in

src/ir_passes/erase_typ_field.ml

@@ -100,8 +100,8 @@ let transform prog =
   and t_exp' (exp : exp) =
     let exp' = exp.it in
     match exp' with
-    | LitE _ -> exp'
-    | VarE id -> exp'
+    | LitE _
+    | VarE _ -> exp'
     | AssignE (exp1, exp2) ->
       AssignE (t_lexp exp1, t_exp exp2)
     | PrimE (p, exps) ->

src/ir_passes/show.ml

@@ -249,8 +249,7 @@ and t_exp env (e : Ir.exp) =
   { e with it = t_exp' env e.it }
 
 and t_exp' env = function
-  | LitE l -> LitE l
-  | VarE id -> VarE id
+  | (LitE _ | VarE _) as e -> e
   | PrimE (ShowPrim ot, [exp1]) ->
     let t' = T.normalize ot in
     add_type env t';

src/ir_passes/tailcall.ml

@@ -92,12 +92,12 @@ and assignEs vars exp : dec list =
     List.mapi (fun i v -> expD (assignE v (projE (varE v) i))) vars
 
 and exp' env e  : exp' = match e.it with
-  | (VarE _ | LitE _) as it -> it
+  | (VarE (_, _) | LitE _) as it -> it
   | AssignE (e1, e2)    -> AssignE (lexp env e1, exp env e2)
   | PrimE (CallPrim insts, [e1; e2])  ->
     begin match e1.it, env with
-    | VarE f1, { tail_pos = true;
-                 info = Some { func; typ_binds; temps; label; tail_called } }
+    | VarE (_, f1), { tail_pos = true;
+                      info = Some { func; typ_binds; temps; label; tail_called } }
          when f1 = func && are_generic_insts typ_binds insts  ->
       tail_called := true;
       (blockE (assignEs temps (exp env e2)) (breakE label (unitE ()))).it

src/lang_utils/source.ml

@@ -3,7 +3,8 @@ type region = {left : pos; right : pos}
 type ('a, 'b) annotated_phrase = {at : region; it : 'a; mutable note: 'b}
 type 'a phrase = ('a, unit) annotated_phrase
 
-let (@@) it at = {it; at; note = ()}
+let annotate note it at = {it; at; note}
+let (@@) it at = annotate () it at
 
 (* Positions and regions *)
 
@@ -27,4 +28,3 @@ let string_of_region r =
 (* generic parse error *)
 
 exception ParseError of region * string
-

src/lang_utils/source.mli

@@ -12,6 +12,7 @@ val string_of_region : region -> string
 val span : region -> region -> region
 val between : region -> region -> region
 
+val annotate : 'b -> 'a -> region -> ('a, 'b) annotated_phrase
 val (@@) : 'a -> region -> 'a phrase
 
 exception ParseError of region * string

src/lowering/desugar.ml

@@ -35,7 +35,7 @@ let apply_sign op l = Syntax.(match op, l with
 let phrase f x = { x with it = f x.it }
 
 let typ_note : S.typ_note -> Note.t =
-  fun S.{ note_typ; note_eff } -> Note.{ def with typ = note_typ; eff = note_eff }
+  fun S.{ note_typ; note_eff; _ } -> Note.{ def with typ = note_typ; eff = note_eff }
 
 let phrase' f x =
   { x with it = f x.at x.note x.it }
@@ -54,7 +54,7 @@ and exp e =
     | _ -> typed_phrase' exp' e
 
 and exp' at note = function
-  | S.VarE i -> I.VarE i.it
+  | S.VarE i -> I.VarE ((match i.note with Var -> I.Var | Const -> I.Const), i.it)
   | S.ActorUrlE e ->
     I.(PrimE (ActorOfIdBlob note.Note.typ, [url e at]))
   | S.LitE l -> I.LitE (lit !l)

src/mo_def/compUnit.ml

@@ -2,6 +2,7 @@ open Mo_types
 
 open Syntax
 
+let (@~) it at = Source.annotate Const it at
 
 (* Compilation unit detection *)
 
@@ -91,7 +92,7 @@ let obj_decs obj_sort at note id_opt fields =
         None);
       at; note
     };
-    { it = ExpD { it = VarE id; at; note };
+    { it = ExpD { it = VarE (id.it @~ id.at); at; note };
       at; note }
     ]
 
@@ -106,10 +107,10 @@ let decs_of_lib (cu : comp_unit) =
       pat,
       { it = ImportE (fp, ri);
         at;
-        note = { note_typ = note; note_eff = Type.Triv} },
+        note = { empty_typ_note with note_typ = note } },
       None);
       at;
-      note = { note_typ = note; note_eff = Type.Triv } }) imports
+      note = { empty_typ_note with note_typ = note } }) imports
   in
   import_decs,
   match cub.it with

src/mo_def/syntax.ml

@@ -22,6 +22,7 @@ type resolved_import =
 (* Identifiers *)
 
 type id = string Source.phrase
+(* type id_ref, see below *)
 type typ_id = (string, Type.con option) Source.annotated_phrase
 
 
@@ -32,6 +33,7 @@ type func_sort = Type.func_sort Source.phrase
 
 type mut = mut' Source.phrase
 and mut' = Const | Var
+and id_ref = (string, mut') Source.annotated_phrase
 
 and path = (path', Type.typ) Source.annotated_phrase
 and path' =
@@ -149,7 +151,7 @@ type sugar = bool (* Is the source of a function body a block `<block>`,
 type exp = (exp', typ_note) Source.annotated_phrase
 and exp' =
   | PrimE of string                            (* primitive *)
-  | VarE of id                                 (* variable *)
+  | VarE of id_ref                             (* variable *)
   | LitE of lit ref                            (* literal *)
   | ActorUrlE of exp                           (* actor reference *)
   | UnE of op_typ * unop * exp                 (* unary operator *)
@@ -260,15 +262,15 @@ type lib = comp_unit
 (* Helpers *)
 
 let (@@) = Source.(@@)
-let (@?) it at = Source.({it; at; note = empty_typ_note})
-let (@!) it at = Source.({it; at; note = Type.Pre})
-let (@=) it at = Source.({it; at; note = None})
-
+let (@~) it at = Source.annotate Const it at
+let (@?) it at = Source.annotate empty_typ_note it at
+let (@!) it at = Source.annotate Type.Pre it at
+let (@=) it at = Source.annotate None it at
 
 (* NB: This function is currently unused *)
 let string_of_lit = function
   | BoolLit false -> "false"
-  | BoolLit true  ->  "true"
+  | BoolLit true  -> "true"
   | IntLit n
   | NatLit n      -> Numerics.Int.to_pretty_string n
   | Int8Lit n     -> Numerics.Int_8.to_pretty_string n

src/mo_frontend/parser.mly

@@ -54,7 +54,7 @@ let funcT (sort, tbs, t1, t2) =
   | _ -> FuncT(sort, tbs, t1, t2)
 
 
-let dup_var x = VarE (x.it @@ x.at) @? x.at
+let dup_var x = VarE (x.it @~ x.at) @? x.at
 
 let name_exp e =
   match e.it with
@@ -99,7 +99,7 @@ let rec normalize_let p e =
 
 let let_or_exp named x e' at =
   if named
-  then LetD(VarP(x) @! at, e' @? at, None) @? at
+  then LetD(VarP x @! at, e' @? at, None) @? at
        (* If you change the above regions,
           modify is_sugared_func_or_module to match *)
   else ExpD(e' @? at) @? at
@@ -587,11 +587,11 @@ exp_nullary(B) :
   | e=exp_plain
     { e }
   | x=id
-    { VarE(x) @? at $sloc }
+    { VarE (x.it @~ x.at) @? at $sloc }
   | PRIM s=TEXT
     { PrimE(s) @? at $sloc }
   | UNDERSCORE
-    { VarE ("_" @@ at $sloc) @? at $sloc }
+    { VarE ("_" @~ at $sloc) @? at $sloc }
 
 exp_post(B) :
   | e=exp_nullary(B)
@@ -767,7 +767,7 @@ catch :
 
 exp_field :
   | m=var_opt x=id t=annot_opt
-    { let e = VarE(x.it @@ x.at) @? x.at in
+    { let e = VarE (x.it @~ x.at) @? x.at in
       { mut = m; id = x; exp = annot_exp e t; } @@ at $sloc }
   | m=var_opt x=id t=annot_opt EQ e=exp(ob)
     { { mut = m; id = x; exp = annot_exp e t; } @@ at $sloc }