Fix arity of generated switcher continuation. (#105)

This patch fixes a bug with the arity of the continuation reference
generated by a `switch` (aka the current continuation). Its arity was
mistakenly derived from the switch-tag's codomain. The arity of the
current continuation can be obtained at the `switch` point by
deconstructing the type annotation on it.

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
@@ -413,10 +413,13 @@ let rec step (c : config) : config =
       | Switch (x, y), Ref (ContRef {contents = None}) :: vs ->
          vs, [Trapping "continuation already consumed" @@ e.at]
 
-      | Switch (x, y), Ref (ContRef {contents = Some (n, ctxt)} as cont) :: vs ->
+      | 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 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
-         vs', [Suspending (tagt, args, Some 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
@@ -1292,11 +1295,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 (_, 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 FuncT (_, ts) = func_type_of_tag_type c.frame.inst (Tag.type_of tagt) in
        let ctxt'' code = compose (ctxt' code) (vs', es') in
-       let cont' = Ref (ContRef (ref (Some (Int32.add (Lib.List32.length ts) 1l, 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

@@ -114,44 +114,6 @@ let defaultable = function
   | BotT -> assert false
 
 
-(* Conversions & Projections *)
-
-let num_type_of_addr_type = function
-  | I32AT -> I32T
-  | I64AT -> I64T
-
-let addr_type_of_num_type = function
-  | I32T -> I32AT
-  | I64T -> I64AT
-  | _ -> assert false
-
-
-let unpacked_storage_type = function
-  | ValStorageT t -> t
-  | PackStorageT _ -> NumT I32T
-
-let unpacked_field_type (FieldT (_mut, t)) = unpacked_storage_type t
-
-
-let as_func_str_type (st : str_type) : func_type =
-  match st with
-  | DefFuncT ft -> ft
-  | _ -> assert false
-
-let as_struct_str_type (st : str_type) : struct_type =
-  match st with
-  | DefStructT st -> st
-  | _ -> assert false
-
-let as_array_str_type (st : str_type) : array_type =
-  match st with
-  | DefArrayT at -> at
-  | _ -> assert false
-
-let extern_type_of_import_type (ImportT (et, _, _)) = et
-let extern_type_of_export_type (ExportT (et, _)) = et
-
-
 (* Filters *)
 
 let funcs = List.filter_map (function ExternFuncT ft -> Some ft | _ -> None)
@@ -310,17 +272,59 @@ let expand_def_type (dt : def_type) : str_type =
   st
 
 
-(* Projections *)
+(* Conversions & Projections *)
+
+let num_type_of_addr_type = function
+  | I32AT -> I32T
+  | I64AT -> I64T
+
+let addr_type_of_num_type = function
+  | I32T -> I32AT
+  | I64T -> I64AT
+  | _ -> assert false
 
 let unpacked_storage_type = function
   | ValStorageT t -> t
   | PackStorageT _ -> NumT I32T
 
+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
+  | _ -> assert false
+
 let as_cont_str_type (dt : str_type) : cont_type =
   match dt with
   | DefContT ct -> ct
   | _ -> assert false
 
+let as_struct_str_type (st : str_type) : struct_type =
+  match st with
+  | DefStructT st -> st
+  | _ -> assert false
+
+let as_array_str_type (st : str_type) : array_type =
+  match st with
+  | 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
 
 (* String conversion *)
 

test/core/stack-switching/cont.wast

@@ -927,6 +927,38 @@
 )
 (assert_return (invoke "main") (i32.const 10))
 
+(module
+  (type $f1 (func (result i32)))
+  (type $c1 (cont $f1))
+  (type $f2 (func (param (ref null $c1)) (result i32)))
+  (type $c2 (cont $f2))
+  (type $f3 (func (param (ref null $c2)) (result i32)))
+  (type $c3 (cont $f3))
+  (tag $e (result i32))
+
+  (func $fn_1 (param (ref null $c2)) (result i32)
+    (local.get 0)
+    (switch $c2 $e)
+    (i32.const 24)
+  )
+  (elem declare func $fn_1)
+
+  (func $fn_2 (result i32)
+    (cont.new $c3 (ref.func $fn_1))
+    (switch $c3 $e)
+    (drop)
+    (i32.const -1)
+  )
+  (elem declare func $fn_2)
+
+  (func (export "main") (result i32)
+    (cont.new $c1 (ref.func $fn_2))
+    (resume $c1 (on $e switch))
+  )
+)
+
+(assert_return (invoke "main") (i32.const -1))
+
 ;; Syntax: check unfolded forms
 (module
   (type $ft (func))