Branch-like structured control flow

ml-proto/host/lexer.mll

@@ -129,12 +129,15 @@ rule token = parse
 
   | "nop" { NOP }
   | "block" { BLOCK }
-  | "if" { IF }
   | "loop" { LOOP }
-  | "label" { LABEL }
-  | "break" { BREAK }
+  | "br" { BR }
+  | "br_if" { BRIF }
+  | "br_unless" { BRUNLESS }
+  | "if" { IF }
+  | "forever" { FOREVER }
   | "case" { CASE }
   | "fallthrough" { FALLTHROUGH }
+  | "break" { BREAK }
   | "call" { CALL }
   | "call_import" { CALLIMPORT }
   | "call_indirect" { CALLINDIRECT }
@@ -157,6 +160,7 @@ rule token = parse
   | (ixx as t)".store"(mem_size as sz)"/"(align as a)
     { STOREWRAP (wrapop t sz a) }
 
+  | (nxx as t)".br_switch" { BRSWITCH (value_type t) }
   | (nxx as t)".switch" { SWITCH (value_type t) }
   | (nxx as t)".const" { CONST (value_type t) }
 

ml-proto/host/parser.mly

@@ -86,6 +86,12 @@ let bind_label c x =
     Error.error x.at ("duplicate label " ^ x.it);
   {c with labels = VarMap.add x.it 0 (VarMap.map ((+) 1) c.labels)}
 
+let sugar_label c =
+  (label c ("(syntax sugar)" @@ at())) @@ at()
+
+let bind_sugar_label c =
+  {c with labels = VarMap.add "(syntax sugar)" 0 (VarMap.map ((+) 1) c.labels)}
+
 let anon space n = space.count <- space.count + n
 
 let anon_func c = anon c.funcs 1
@@ -95,13 +101,14 @@ let anon_label c = {c with labels = VarMap.map ((+) 1) c.labels}
 %}
 
 %token INT FLOAT TEXT VAR TYPE LPAR RPAR
-%token NOP BLOCK IF LOOP LABEL BREAK SWITCH CASE FALLTHROUGH
+%token NOP BLOCK LOOP BR BRIF BRUNLESS BRSWITCH
 %token CALL CALLIMPORT CALLINDIRECT RETURN
 %token GETLOCAL SETLOCAL LOAD STORE
 %token CONST UNARY BINARY COMPARE CONVERT
 %token FUNC PARAM RESULT LOCAL MODULE MEMORY SEGMENT IMPORT EXPORT TABLE
 %token PAGESIZE MEMORYSIZE RESIZEMEMORY
 %token ASSERTINVALID ASSERTRETURN ASSERTRETURNNAN ASSERTTRAP INVOKE
+%token IF FOREVER SWITCH CASE FALLTHROUGH BREAK
 %token EOF
 
 %token<string> INT
@@ -110,6 +117,7 @@ let anon_label c = {c with labels = VarMap.map ((+) 1) c.labels}
 %token<string> VAR
 %token<Types.value_type> TYPE
 %token<Types.value_type> CONST
+%token<Types.value_type> BRSWITCH
 %token<Types.value_type> SWITCH
 %token<Ast.unop> UNARY
 %token<Ast.binop> BINARY
@@ -160,20 +168,50 @@ expr :
 ;
 oper :
   | NOP { fun c -> Nop }
-  | BLOCK expr expr_list { fun c -> Block ($2 c :: $3 c) }
-  | IF expr expr expr { fun c -> If ($2 c, $3 c, $4 c) }
-  | IF expr expr  /* Sugar */
-    { let at1 = ati 1 in fun c -> If ($2 c, $3 c, Nop @@ at1) }
-  | LOOP expr_block { fun c -> Loop ($2 c) }
-  | LABEL expr_block { fun c -> Label ($2 (anon_label c)) }
-  | LABEL bind_var expr_block  /* Sugar */
-    { fun c -> Label ($3 (bind_label c $2)) }
-  | BREAK var expr_opt { fun c -> Break ($2 c label, $3 c) }
-  | BREAK { let at = at() in fun c -> Break (0 @@ at, None) }  /* Sugar */
+  | BLOCK expr_list bind_var { fun c -> Block ($2 (bind_label c $3)) }
+  /* Sugar: block with no label */
+  | BLOCK expr_list { fun c -> Block ($2 (anon_label c)) }
+  | LOOP bind_var expr_list { fun c -> Loop ($3 (bind_label c $2)) }
+  /* Sugar: loop with additional label at end */
+  | LOOP bind_var expr_list bind_var
+    { fun c -> let at = at() in
+               Loop ((fun c -> [Block ($3 (bind_label c $4)) @@ at])
+                     (bind_label c $2)) }
+  | BR var expr_opt { fun c -> Br ($2 c label, $3 c) }
+  | BRIF var expr expr_opt { fun c -> BrIf ($2 c label, $3 c, $4 c) }
+  | BRUNLESS var expr expr_opt { fun c -> BrUnless ($2 c label, $3 c, $4 c) }
+  | BRSWITCH expr var br_switch_arms expr_opt
+    { let at = at() in
+      fun c -> BrSwitch ($1 @@ at, $2 c, $3 c label,
+                         List.map (fun (x,y) -> ((literal at x $1).it, y c label)) $4,
+                         $5 c) }
+  /* Sugar: if statements */
+  | IF expr expr
+    { fun c -> let at = at() in
+               let c' = (anon_label c) in
+               Block [BrUnless (0 @@ at, $2 c', None) @@ at; $3 c'] }
+  /* Sugar: if-else statements */
+  | IF expr expr expr
+    { fun c -> let at = at() in
+               let c' = (anon_label c) in
+               Block [Block [BrUnless (0 @@ at, $2 c', None) @@ at;
+                             Br (1 @@ at, Some ($3 (anon_label c'))) @@ at] @@ at;
+                     $4 c'] }
+  /* Sugar: forever-loop statements */
+  | FOREVER expr_list
+    { fun c -> let at = at() in
+               let c' = (bind_sugar_label c) in
+               let c'' = (anon_label c') in
+               Block [Loop (List.append ($2 c'') [Br (0 @@ at, None) @@ at]) @@ at]
+    }
+  /* Sugar: switch statements */
   | SWITCH expr arms
     { let at1 = ati 1 in
-      fun c -> let x, y = $3 c in
-        Switch ($1 @@ at1, $2 c, List.map (fun a -> a $1) x, y) }
+      fun c -> let c' = (bind_sugar_label c) in
+               let x, y = $3 c' in
+               Block [Switch ($1 @@ at1, $2 c', List.map (fun a -> a $1) x, y) @@ at1] }
+  /* Sugar: break statements */
+  | BREAK expr_opt { fun c -> Br (sugar_label c, $2 c) }
   | CALL var expr_list { fun c -> Call ($2 c func, $3 c) }
   | CALLIMPORT var expr_list { fun c -> CallImport ($2 c import, $3 c) }
   | CALLINDIRECT var expr expr_list
@@ -202,12 +240,12 @@ expr_list :
   | /* empty */ { fun c -> [] }
   | expr expr_list { fun c -> $1 c :: $2 c }
 ;
+
 expr_block :
   | expr { $1 }
   | expr expr expr_list  /* Sugar */
     { let at = at() in fun c -> Block ($1 c :: $2 c :: $3 c) @@ at }
 ;
-
 fallthrough :
   | /* empty */ { false }
   | FALLTHROUGH { true }
@@ -225,17 +263,24 @@ arm :
 arms :
   | expr { fun c -> [], $1 c }
   | arm arms { fun c -> let x, y = $2 c in $1 c :: x, y }
-;
 
+br_switch_arm :
+  | INT var { ($1, $2) }
+;
+br_switch_arms :
+  | /* empty */ { [] }
+  | br_switch_arm br_switch_arms { $1 :: $2 }
+;
 
 /* Functions */
 
 func_fields :
   | /* empty */  /* Sugar */
     { let at = at() in
       fun c -> {params = []; result = None; locals = []; body = Nop @@ at} }
-  | expr_block
-    { fun c -> {params = []; result = None; locals = []; body = $1 c} }
+  | expr_list
+    { let at = at() in
+      fun c -> {params = []; result = None; locals = []; body = Block ($1 c) @@ at} }
   | LPAR PARAM value_type_list RPAR func_fields
     { fun c -> anon_locals c $3; let f = $5 c in
       {f with params = $3 @ f.params} }

ml-proto/spec/ast.ml

@@ -75,12 +75,14 @@ type literal = value Source.phrase
 type expr = expr' Source.phrase
 and expr' =
   | Nop                                           (* do nothing *)
-  | Block of expr list                            (* execute in sequence *)
-  | If of expr * expr * expr                      (* conditional *)
-  | Loop of expr                                  (* infinite loop *)
-  | Label of expr                                 (* labelled expression *)
-  | Break of var * expr option                    (* break to n-th surrounding label *)
+  | Block of expr list                            (* execute in sequence, label at end *)
+  | Loop of expr list                             (* execute in sequence, label at beginning *)
+  | Br of var * expr option                       (* branch to label *)
+  | BrIf of var * expr * expr option              (* branch to label if expr is true *)
+  | BrUnless of var * expr * expr option          (* branch to label if expr is false *)
+  | BrSwitch of value_type * expr * var * (value * var) list * expr option
   | Switch of value_type * expr * arm list * expr (* switch, latter expr is default *)
+                                                  (* branch to label selected by expr *)
   | Call of var * expr list                       (* call function *)
   | CallImport of var * expr list                 (* call imported function *)
   | CallIndirect of var * expr * expr list        (* call function through table *)

ml-proto/spec/check.ml

@@ -118,26 +118,34 @@ let rec check_expr c et e =
     check_type None et e.at
 
   | Block es ->
-    require (es <> []) e.at "invalid block";
+    let c' = {c with labels = et :: c.labels} in
     let es', eN = Lib.List.split_last es in
-    List.iter (check_expr c None) es';
-    check_expr c et eN
+    List.iter (check_expr c' None) es';
+    check_expr c' et eN;
 
-  | If (e1, e2, e3) ->
-    check_expr c (Some Int32Type) e1;
-    check_expr c et e2;
-    check_expr c et e3
+  | Loop es ->
+    let c' = {c with labels = et :: c.labels} in
+    let es', eN = Lib.List.split_last es in
+    List.iter (check_expr c' None) es';
+    check_expr c' et eN;
 
-  | Loop e1 ->
-    check_expr c None e1
+  | Br (x, eo) ->
+    check_expr_option c (label c x) eo e.at
 
-  | Label e1 ->
-    let c' = {c with labels = et :: c.labels} in
-    check_expr c' et e1
+  | BrIf (x, ec, eo) ->
+    check_expr c (Some Int32Type) ec;
+    check_expr_option c (label c x) eo e.at
 
-  | Break (x, eo) ->
+  | BrUnless (x, ec, eo) ->
+    check_expr c (Some Int32Type) ec;
     check_expr_option c (label c x) eo e.at
 
+  | BrSwitch (t, ec, default, labels, eo) ->
+    (* TODO: Check that cases are unique. *)
+    require (t.it = Int32Type || t.it = Int64Type) t.at "invalid br_switch type";
+    check_expr c (Some t.it) ec;
+    check_expr_option c (label c default) eo e.at
+
   | Switch (t, e1, arms, e2) ->
     require (t.it = Int32Type || t.it = Int64Type) t.at "invalid switch type";
     (* TODO: Check that cases are unique. *)

ml-proto/spec/eval.ml

@@ -126,26 +126,48 @@ let rec eval_expr (c : config) (e : expr) =
     None
 
   | Block es ->
-    let es', eN = Lib.List.split_last es in
-    List.iter (fun eI -> ignore (eval_expr c eI)) es';
-    eval_expr c eN
-
-  | If (e1, e2, e3) ->
-    let i = int32 (eval_expr c e1) e1.at in
-    eval_expr c (if i <> Int32.zero then e2 else e3)
-
-  | Loop e1 ->
-    ignore (eval_expr c e1);
-    eval_expr c e
+    let module L = MakeLabel () in
+    let c' = {c with labels = L.label :: c.labels} in
+    (try
+      (let es', eN = Lib.List.split_last es in
+       List.iter (fun eI -> ignore (eval_expr c' eI)) es';
+       eval_expr c' eN)
+     with L.Label vo -> vo)
 
-  | Label e1 ->
+  | Loop es ->
     let module L = MakeLabel () in
     let c' = {c with labels = L.label :: c.labels} in
-    (try eval_expr c' e1 with L.Label vo -> vo)
+    (try
+      (let es', eN = Lib.List.split_last es in
+       List.iter (fun eI -> ignore (eval_expr c' eI)) es';
+       eval_expr c' eN)
+     with L.Label _ -> eval_expr c e)
 
-  | Break (x, eo) ->
+  | Br (x, eo) ->
     raise (label c x (eval_expr_option c eo))
 
+  | BrIf (x, ec, eo) ->
+    let i = int32 (eval_expr c ec) ec.at in
+    if i <> Int32.zero then
+      raise (label c x (eval_expr_option c eo))
+    else
+      None
+
+  | BrUnless (x, ec, eo) ->
+    let i = int32 (eval_expr c ec) ec.at in
+    if i = Int32.zero then
+      raise (label c x (eval_expr_option c eo))
+    else
+      None
+
+  | BrSwitch (_t, ec, default, labels, eo) ->
+    let e = some (eval_expr c ec) ec.at in
+    raise (label c
+           (try
+              let i, l = List.find (fun (i, l) -> i = e) labels in l
+            with Not_found -> default)
+           (eval_expr_option c eo))
+
   | Switch (_t, e1, arms, e2) ->
     let vo = some (eval_expr c e1) e1.at in
     (match List.fold_left (eval_arm c vo) `Seek arms with

ml-proto/test/br_switch.wase

@@ -0,0 +1,60 @@
+(module
+  ;; Statement br_switch
+  (func $stmt (param $i i32) (result i32)
+    (local $j i32)
+    (set_local $j (i32.const 100))
+    (block (block (block (block (block (block (block (block (block
+
+    (i32.br_switch (get_local $i)
+                   $default 0 $case0 1 $case1 2 $case2 3 $case3 4 $case4 5 $case5 6 $case6)
+    $case0) (return (get_local $i))
+    $case1)
+    $case2)
+    $case3) (set_local $j (i32.sub (i32.const 0) (get_local $i))) (br $end)
+    $case4) (br $end)
+    $case5) (set_local $j (i32.const 101)) (br $end)
+    $case6) (set_local $j (i32.const 101))
+    $default) (set_local $j (i32.const 102))
+    $end)
+    (return (get_local $j))
+  )
+
+  ;; Expression br_switch
+  (func $expr (param $i i64) (result i64)
+    (local $j i64)
+    (set_local $j (i64.const 100))
+    (return
+      (block (block (block (block (block (block (block
+      (i64.br_switch (get_local $i)
+                     $default 0 $case0 1 $case1 2 $case2 3 $case3 6 $case6)
+      $case0) (return (get_local $i))
+      $case1)
+      $case2)
+      $case3) (br $exit (i64.sub (i64.const 0) (get_local $i)))
+      $case6) (set_local $j (i64.const 101))
+      $default) (get_local $j)
+      $exit)
+    )
+  )
+
+  (export "stmt" $stmt)
+  (export "expr" $expr)
+)
+
+(assert_eq (invoke "stmt" (i32.const 0)) (i32.const 0))
+(assert_eq (invoke "stmt" (i32.const 1)) (i32.const -1))
+(assert_eq (invoke "stmt" (i32.const 2)) (i32.const -2))
+(assert_eq (invoke "stmt" (i32.const 3)) (i32.const -3))
+(assert_eq (invoke "stmt" (i32.const 4)) (i32.const 100))
+(assert_eq (invoke "stmt" (i32.const 5)) (i32.const 101))
+(assert_eq (invoke "stmt" (i32.const 6)) (i32.const 102))
+(assert_eq (invoke "stmt" (i32.const 7)) (i32.const 102))
+(assert_eq (invoke "stmt" (i32.const -10)) (i32.const 102))
+
+(assert_eq (invoke "expr" (i64.const 0)) (i64.const 0))
+(assert_eq (invoke "expr" (i64.const 1)) (i64.const -1))
+(assert_eq (invoke "expr" (i64.const 2)) (i64.const -2))
+(assert_eq (invoke "expr" (i64.const 3)) (i64.const -3))
+(assert_eq (invoke "expr" (i64.const 6)) (i64.const 101))
+(assert_eq (invoke "expr" (i64.const 7)) (i64.const 100))
+(assert_eq (invoke "expr" (i64.const -10)) (i64.const 100))