Address Andreas' feedback

Author: dhil

interpreter/exec/eval.ml

@@ -73,7 +73,7 @@ and admin_instr' =
   | Frame of int * frame * code
   | Handler of int * catch list * code
   | Handle of handle_table * code
-  | Suspending of tag_inst * value stack * ref_ option * ctxt
+  | Suspending of tag_inst * value stack * (int32 * ref_) option * ctxt
 
 and ctxt = code -> code
 and handle_table = (tag_inst * idx) list * tag_inst list
@@ -104,15 +104,6 @@ let is_jumping e =
 
 let compose (vs1, es1) (vs2, es2) = vs1 @ vs2, es1 @ es2
 
-let switcher_contref_arity contarg =
-  match contarg with
-  | RefT (_, DefHT def) ->
-     (match as_cont_str_type (expand_def_type def) with
-      | ContT (DefHT def) ->
-         let FuncT (ts, _) = as_func_str_type (expand_def_type def) in
-         Lib.List32.length ts
-      | _ -> assert false)
-  | _ -> assert false
 
 (* Configurations *)
 
@@ -422,14 +413,12 @@ let rec step (c : config) : config =
 
       | Switch (x, y), Ref (ContRef ({contents = Some (n, ctxt)} as cont)) :: vs ->
          let FuncT (ts, _) = func_type_of_cont_type c.frame.inst (cont_type c.frame.inst x) in
-         let arity = switcher_contref_arity (Lib.List.last ts) in
+         let FuncT (ts', _) = as_cont_func_ref_type (Lib.List.last ts) in
+         let arity = Lib.List32.length ts' in
          let tagt = tag c.frame.inst y in
          let args, vs' = i32_split (Int32.sub n 1l) vs e.at in
-         (* NOTE(dhil): Slight hack here. We update the arity of given
-            contref to be the arity of the current continuation
-            generated by the handler for switch. *)
          cont := Some (arity, ctxt);
-         vs', [Suspending (tagt, args, Some (ContRef cont), fun code -> code) @@ e.at]
+         vs', [Suspending (tagt, args, Some (arity, ContRef cont), fun code -> code) @@ e.at]
 
       | ReturnCall x, vs ->
         (match (step {c with code = (vs, [Plain (Call x) @@ e.at])}).code with
@@ -1277,10 +1266,10 @@ let rec step (c : config) : config =
       [Ref (ContRef (ref (Some (Lib.List32.length ts, ctxt'))))] @ vs1 @ vs,
       [Plain (Br (List.assq tagt hs)) @@ e.at]
 
-    | Handle ((_, hs) as hso, (vs', {it = Suspending (tagt, vs1, Some (ContRef ({contents = Some (n, ctxt)} as cont)), ctxt'); at} :: es')), vs
+    | Handle ((_, hs) as hso, (vs', {it = Suspending (tagt, vs1, Some (ar, ContRef ({contents = Some (_, ctxt)} as cont)), ctxt'); at} :: es')), vs
        when List.memq tagt hs ->
        let ctxt'' code = compose (ctxt' code) (vs', es') in
-       let cont' = Ref (ContRef (ref (Some (n, ctxt'')))) in
+       let cont' = Ref (ContRef (ref (Some (ar, ctxt'')))) in
        let args = cont' :: vs1 in
        cont := None;
        vs' @ vs, [Handle (hso, ctxt (args, [])) @@ e.at]

interpreter/syntax/types.ml

@@ -277,6 +277,11 @@ let unpacked_storage_type = function
 
 let unpacked_field_type (FieldT (_mut, t)) = unpacked_storage_type t
 
+let as_def_heap_type (ht : heap_type) : def_type =
+  match ht with
+  | DefHT def -> def
+  | _ -> assert false
+
 let as_func_str_type (st : str_type) : func_type =
   match st with
   | DefFuncT ft -> ft
@@ -297,6 +302,15 @@ let as_array_str_type (st : str_type) : array_type =
   | DefArrayT at -> at
   | _ -> assert false
 
+let as_cont_func_heap_type (ht : heap_type) : func_type =
+  let ContT ht' = as_cont_str_type (expand_def_type (as_def_heap_type ht)) in
+  as_func_str_type (expand_def_type (as_def_heap_type ht'))
+
+let as_cont_func_ref_type (rt : val_type) : func_type =
+  match rt with
+  | RefT (_, ht) -> as_cont_func_heap_type ht
+  | _ -> assert false
+
 let extern_type_of_import_type (ImportT (et, _, _)) = et
 let extern_type_of_export_type (ExportT (et, _)) = et