[Experiments] Type representation

see #2348

src/ir_def/arrange_ir.ml

@@ -67,6 +67,7 @@ and prim = function
   | AwaitPrim         -> Atom "AwaitPrim"
   | AssertPrim        -> Atom "AssertPrim"
   | ThrowPrim         -> Atom "ThrowPrim"
+  | TypRep t          -> "TypRep" $$ [typ t]
   | ShowPrim t        -> "ShowPrim" $$ [typ t]
   | SerializePrim t   -> "SerializePrim" $$ List.map typ t
   | DeserializePrim t -> "DeserializePrim" $$ List.map typ t

src/ir_def/check_ir.ml

@@ -497,6 +497,9 @@ let rec check_exp env (exp:Ir.exp) : unit =
     | AssertPrim, [exp1] ->
       typ exp1 <: T.bool;
       T.unit <: t
+    | TypRep ot, [] ->
+      check (T.shared ot) "gen_typrep is not defined for operand type";
+      Construct.typRepT <: t
     | ShowPrim ot, [exp1] ->
       check env.flavor.has_show "show expression in non-show flavor";
       check (Show.can_show ot) "show is not defined for operand type";

src/ir_def/construct.ml

@@ -185,6 +185,14 @@ let optE e =
    at = no_region;
  }
 
+let arrayE t es =
+  let effs = List.map eff es in
+  let eff = List.fold_left max_eff T.Triv effs in
+  { it = PrimE (ArrayPrim (Const, t), es);
+    note = Note.{ def with typ = T.Array t; eff };
+    at = no_region;
+  }
+
 let tagE i e =
  { it = PrimE (TagPrim i, [e]);
    note = Note.{ def with typ = T.Variant [{T.lab = i; typ = typ e; depr = None}]; eff = eff e };
@@ -353,6 +361,7 @@ let tupE exps =
     note = Note.{ def with typ = T.Tup (List.map typ exps); eff };
   }
 
+(* TODO: Is it ok to share this value? *)
 let unitE = tupE []
 
 let breakE l exp =
@@ -505,16 +514,6 @@ let funcD ((id, typ) as f) x exp =
 let nary_funcD ((id, typ) as f) xs exp =
   letD f (nary_funcE id typ xs exp)
 
-(* Continuation types *)
-
-let answerT = T.unit
-
-let contT typ = T.Func (T.Local, T.Returns, [], T.as_seq typ, [])
-
-let err_contT =  T.Func (T.Local, T.Returns, [], [T.catch], [])
-
-let cpsT typ = T.Func (T.Local, T.Returns, [], [contT typ; err_contT], [])
-
 (* Sequence expressions *)
 
 let seqE es =
@@ -626,3 +625,32 @@ let unreachableE =
   (* Do we want a dedicated UnreachableE in the AST? *)
   loopE unitE
 
+(* Internal types *)
+
+(*
+These type definitions mirror the type definitions in prelude/internals.mo.
+It would be good to get rid of that duplication somehow, and use those
+in `prelude/internals.mo` directly.
+*)
+
+let answerT = T.unit
+
+let contT typ = T.Func (T.Local, T.Returns, [], T.as_seq typ, [])
+
+let err_contT =  T.Func (T.Local, T.Returns, [], [T.catch], [])
+
+let cpsT typ = T.Func (T.Local, T.Returns, [], [contT typ; err_contT], [])
+
+let typRepT =
+  let open T in
+  let c = Con.fresh "TypRep" (Abs ([], Pre)) in
+  let t = Con (c, []) in
+  let rhs =
+    Variant [
+      { lab = "bool"; typ = T.unit; depr = None };
+      { lab = "int"; typ = T.unit; depr = None };
+      { lab = "opt"; typ = t; depr = None };
+      { lab = "tup"; typ = T.Array t; depr = None };
+    ] in
+  set_kind c (Def ([], rhs));
+  t

src/ir_def/construct.mli

@@ -58,6 +58,7 @@ val ic_rejectE : exp -> exp
 val ic_callE : exp -> exp -> exp -> exp -> exp
 val projE : exp -> int -> exp
 val optE : exp -> exp
+val arrayE : typ -> exp list -> exp
 val tagE : id -> exp -> exp
 val blockE : dec list -> exp -> exp
 val textE : string -> exp
@@ -113,13 +114,6 @@ val nary_funcD : var -> var list -> exp -> dec
 
 val let_no_shadow : var -> exp -> dec list -> dec list
 
-(* Continuations *)
-
-val answerT : typ
-val contT : typ -> typ
-val err_contT : typ
-val cpsT : typ -> typ
-
 (* Sequence expressions *)
 
 val seqE : exp list -> exp
@@ -130,3 +124,19 @@ val (-->) : var -> exp -> exp
 val (-->*) : var list -> exp -> exp (* n-ary local *)
 val forall : typ_bind list -> exp -> exp (* generalization *)
 val (-*-) : exp -> exp -> exp       (* application *)
+
+(* Internal types *)
+
+(*
+These type definitions mirror the type definitions in prelude/internals.mo.
+It would be good to get rid of that duplication somehow, and use those
+in `prelude/internals.mo` directly.
+*)
+
+
+val answerT : typ
+val contT : typ -> typ
+val err_contT : typ
+val cpsT : typ -> typ
+val typRepT : typ
+

src/ir_def/ir.ml

@@ -119,6 +119,7 @@ and prim =
   | AwaitPrim                         (* await *)
   | AssertPrim                        (* assertion *)
   | ThrowPrim                         (* throw *)
+  | TypRep of Type.typ                (* Generates the type representation *)
   | ShowPrim of Type.typ              (* debug_show *)
   | SerializePrim of Type.typ list    (* Candid serialization prim *)
   | DeserializePrim of Type.typ list  (* Candid deserialization prim *)

src/ir_passes/typrep.ml

@@ -0,0 +1,179 @@
+(* Translates away calls to the `TypRep` prim *)
+open Ir_def
+open Mo_types
+
+open Source
+open Ir
+module T = Type
+open Construct
+open Typ_hash
+
+(* Environment *)
+
+(* We go through the file and collect all type arguments to `show`.
+   We store them in `params`, indexed by their `type_id`
+*)
+
+module M = Map.Make(String)
+type env =
+  { params : T.typ M.t ref
+  }
+
+let empty_env () : env = {
+  params = ref M.empty;
+  }
+
+let add_type env t : unit =
+  env.params := M.add (typ_hash t) t !(env.params)
+
+(* Definition names *)
+
+let name_for t =
+  "@typ_rep<" ^ typ_hash t ^ ">"
+
+let var_for t : Construct.var = var (name_for t) typRepT
+
+let exp_for t : exp = varE (var_for t)
+
+(* Synthesizing a single show function *)
+
+(* Returns the new declarations, as well as a list of further types it needs *)
+
+
+let decl_for : T.typ -> Ir.dec * T.typ list = fun t ->
+  match T.normalize t with
+  | T.(Prim Bool) ->
+    letD (var_for t) (tagE "bool" unitE),
+    []
+  | T.(Prim Int) ->
+    letD (var_for t) (tagE "int" unitE),
+    []
+  | T.Tup ts ->
+    letD (var_for t) (tagE "tup" (arrayE typRepT (List.map exp_for ts))),
+    ts
+  | T.Opt t ->
+    letD (var_for t) (tagE "opt" (exp_for t)),
+    [t]
+  | _ ->
+    raise (Invalid_argument ("Typrep.decl_for :" ^ T.string_of_typ_expand t))
+
+(* Synthesizing the types recursively. Hopefully well-founded. *)
+
+let decls : T.typ M.t -> Ir.dec list = fun roots ->
+  let seen = ref M.empty in
+
+  let rec go = function
+    | [] -> []
+    | t::todo when M.mem (typ_hash t) !seen ->
+      go todo
+    | t::todo ->
+      seen := M.add (typ_hash t) () !seen;
+      let (decl, deps) = decl_for t in
+      decl :: go (deps @ todo)
+  in go (List.map snd (M.bindings roots))
+
+(* The AST traversal *)
+
+(* Does two things:
+ - collects all uses of `debug_show` in the `env`
+ - for each actor, resets the environment, recurses,
+   and adds the show functions (this keeps closed actors closed)
+*)
+
+let rec t_exps env decs = List.map (t_exp env) decs
+
+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
+  | PrimE (TypRep ot, []) ->
+    let t' = T.normalize ot in
+    add_type env t';
+    (exp_for t').it
+  | PrimE (p, es) -> PrimE (p, t_exps env es)
+  | AssignE (lexp1, exp2) ->
+    AssignE (t_lexp env lexp1, t_exp env exp2)
+  | FuncE (s, c, id, typbinds, pat, typT, exp) ->
+    FuncE (s, c, id, typbinds, pat, typT, t_exp env exp)
+  | BlockE block -> BlockE (t_block env block)
+  | IfE (exp1, exp2, exp3) ->
+    IfE (t_exp env exp1, t_exp env exp2, t_exp env exp3)
+  | SwitchE (exp1, cases) ->
+    let cases' =
+      List.map
+        (fun {it = {pat;exp}; at; note} ->
+          {it = {pat = pat; exp = t_exp env exp}; at; note})
+        cases
+    in
+    SwitchE (t_exp env exp1, cases')
+  | TryE (exp1, cases) ->
+    let cases' =
+      List.map
+        (fun {it = {pat;exp}; at; note} ->
+          {it = {pat = pat; exp = t_exp env exp}; at; note})
+        cases
+    in
+    TryE (t_exp env exp1, cases')
+  | LoopE exp1 ->
+    LoopE (t_exp env exp1)
+  | LabelE (id, typ, exp1) ->
+    LabelE (id, typ, t_exp env exp1)
+  | AsyncE (tb, e, typ) -> AsyncE (tb, t_exp env e, typ)
+  | DeclareE (id, typ, exp1) ->
+    DeclareE (id, typ, t_exp env exp1)
+  | DefineE (id, mut ,exp1) ->
+    DefineE (id, mut, t_exp env exp1)
+  | NewObjE (sort, ids, t) ->
+    NewObjE (sort, ids, t)
+  | SelfCallE (ts, e1, e2, e3) ->
+    SelfCallE (ts, t_exp env e1, t_exp env e2, t_exp env e3)
+  | ActorE (ds, fields, {pre; post}, typ) ->
+    (* Until Actor expressions become their own units,
+       we repeat what we do in `comp_unit` below *)
+    let env1 = empty_env () in
+    let ds' = t_decs env1 ds in
+    let pre' = t_exp env1 pre in
+    let post' = t_exp env1 post in
+    let ds = decls !(env1.params) in
+    ActorE (ds @ ds', fields, {pre = pre'; post = post'}, typ)
+
+and t_lexp env (e : Ir.lexp) = { e with it = t_lexp' env e.it }
+and t_lexp' env = function
+  | VarLE id -> VarLE id
+  | IdxLE (exp1, exp2) ->
+    IdxLE (t_exp env exp1, t_exp env exp2)
+  | DotLE (exp1, n) ->
+    DotLE (t_exp env exp1, n)
+
+and t_dec env dec = { dec with it = t_dec' env dec.it }
+
+and t_dec' env dec' =
+  match dec' with
+  | LetD (pat,exp) -> LetD (pat,t_exp env exp)
+  | VarD (id, typ, exp) -> VarD (id, typ, t_exp env exp)
+
+and t_decs env decs = List.map (t_dec env) decs
+
+and t_block env (ds, exp) = (t_decs env ds, t_exp env exp)
+
+and t_comp_unit = function
+  | LibU _ -> raise (Invalid_argument "cannot compile library")
+  | ProgU ds ->
+    let env = empty_env () in
+    let ds' = t_decs env ds in
+    let ds = decls !(env.params) in
+    ProgU (ds @ ds')
+  | ActorU (as_opt, ds, fields, {pre; post}, typ) ->
+    let env = empty_env () in
+    let ds' = t_decs env ds in
+    let pre' = t_exp env pre in
+    let post' = t_exp env post in
+    let ds = decls !(env.params) in
+    ActorU (as_opt, ds @ ds', fields, {pre = pre'; post = post'}, typ)
+
+(* Entry point for the program transformation *)
+
+let transform (cu, flavor) =
+  (t_comp_unit cu, flavor)

src/ir_passes/typrep.mli

@@ -0,0 +1,3 @@
+open Ir_def
+
+val transform : Ir.prog -> Ir.prog

src/lowering/desugar.ml

@@ -138,6 +138,14 @@ and exp' at note = function
       I.PrimE (I.CastPrim (T.seq ts1, T.seq ts2), [exp e])
     | _ -> assert false
     end
+  | S.CallE ({it=S.AnnotE ({it=S.PrimE "gen_typrep";_}, _);note;_}, _, {it = S.LitE l; _}) ->
+    begin
+    if !l != NullLit then assert false;
+    match note.S.note_typ with
+    | T.Func (T.Local, T.Returns, [], [T.Opt t], _) ->
+      I.PrimE (I.TypRep t, [])
+    | _ -> assert false
+    end
   | S.CallE ({it=S.AnnotE ({it=S.PrimE "serialize";_}, _);note;_}, _, e) ->
     begin match note.S.note_typ with
     | T.Func (T.Local, T.Returns, [], ts1, ts2) ->

src/pipeline/pipeline.ml

@@ -560,6 +560,9 @@ let tailcall_optimization =
 let show_translation =
   transform_if "Translate show" Show.transform
 
+let typ_rep_translation =
+  transform_if "Translate type rep generation" Typrep.transform
+
 let eq_translation =
   transform_if "Translate polymorphic equality" Eq.transform
 
@@ -571,6 +574,7 @@ let analyze analysis_name analysis prog name =
 
 let ir_passes mode prog_ir name =
   (* translations that extend the progam and must be done before await/cps conversion *)
+  let prog_ir = typ_rep_translation true prog_ir name in
   let prog_ir = show_translation true prog_ir name in
   let prog_ir = eq_translation true prog_ir name in
   (* cps conversion and local transformations *)

test/run/ok/typrep-show.run-low.ok

@@ -0,0 +1 @@
+(unknown location): execution error, accessing identifier before its definition

test/run/ok/typrep-show.wasm-run.ok

@@ -0,0 +1,16 @@
+#bool
+pattern failed
+Error: failed to run main module `_out/typrep-show.wasm`
+
+Caused by:
+    0: failed to invoke command default
+    1: wasm trap: unreachable
+       wasm backtrace:
+         0: 4,word16:u,word32:u,word64:u,word8:u)>
+         1: @text_of_array
+         2: 5,word16:u,word32:u,word64:u,word8:u)>
+         3: @text_of_variant
+         4: 4,word16:u,word32:u,word64:u,word8:u)>
+         5: init
+         6: _start
+

test/run/ok/typrep-show.wasm-run.ret.ok

@@ -0,0 +1 @@
+Return code 134

test/run/typrep-show.mo

@@ -0,0 +1,19 @@
+
+// These tests use `debug_show` on the generated typrep to exercise the type
+// rep generation code. Only works for non-cyclic types of course.
+
+
+// Need to unlock prim to access `@TypRep` and the `gen_typrep` prim
+//MOC-ENV MOC_UNLOCK_PRIM=yesplease
+
+import Prim "mo:prim";
+
+Prim.debugPrint (debug_show ((prim "gen_typrep" : ?Bool -> @TypRep) null));
+Prim.debugPrint (debug_show ((prim "gen_typrep" : ?(Bool, Int) -> @TypRep) null));
+type T = ??T;
+Prim.debugPrint (debug_show ((prim "gen_typrep" : ?T -> @TypRep) null));
+
+
+//SKIP run
+//SKIP run-ir
+// NB: Run-low should work!