use call and return for internal function call

CHANGELOG.md

@@ -43,6 +43,10 @@
 
 ## New features
 
+- Local functions with return address on the stack use usual `CALL`
+  and `RET` x86 instructions instead of (direct & computed) `JMP`
+  ([PR #194](https://github.com/jasmin-lang/jasmin/pull/194)).
+
 - Fill an array with “random” data using `p = #randombytes(p);`
 
 - provide access to mmx registers:

compiler/src/stackAlloc.ml

@@ -127,6 +127,7 @@ let memory_analysis pp_err ~debug tbl up =
     let sao = Stack_alloc.{
         sao_align  = align;
         sao_size   = Conv.cz_of_int size;
+        sao_ioff   = Z0;
         sao_extra_size = Z0;
         sao_max_size = Z0;
         sao_params = List.map (omap conv_pi) sao.sao_params;
@@ -209,10 +210,14 @@ let memory_analysis pp_err ~debug tbl up =
     let ra = V.mk "RA" (Stack Direct) (tu Arch.reg_size) L._dummy [] in
     let extra =
       let extra = to_save in
-      let extra = if rastack then ra :: extra else extra in
       if has_stack && ro.ro_rsp = None then extra @ [rsp]
       else extra in
     let extra_size, align, extrapos = Varalloc.extend_sao sao extra in
+    let extrapos = if rastack then (ra, 0):: extrapos else extrapos in
+    let align =
+      if rastack && wsize_lt align Arch.reg_size then Arch.reg_size
+      else align in
+
     let align, max_stk = 
       Sf.fold (fun fn (align, max_stk) ->
           let sao = get_sao fn in
@@ -251,10 +256,10 @@ let memory_analysis pp_err ~debug tbl up =
     let convert_to_save m =
       m |> List.rev_map conv_to_save |> List.sort compare_to_save |> List.rev_map (fun (x, n) -> x, Conv.cz_of_int n)
     in
-
     let csao =
       Stack_alloc.{ csao with
         sao_align = align;
+        sao_ioff = Conv.cz_of_int (if rastack then size_of_ws Arch.reg_size else 0);
         sao_extra_size = Conv.cz_of_int extra_size;
         sao_max_size = Conv.cz_of_z max_size;
         sao_to_save = convert_to_save ro.ro_to_save;

compiler/src/varalloc.ml

@@ -369,7 +369,16 @@ let alloc_stack_fd pd is_move_op get_info gtbl fd =
 
   let sao_params, atbl = all_alignment pd ctbl alias sao_return fd.f_args lalloc in
 
-  let sao_align, sao_slots, sao_size = alloc_local_stack 0 (Sv.elements slots) atbl in
+  let ra_on_stack =
+    match fd.f_annot.retaddr_kind with
+    | Some OnStack -> true
+    | _ -> false in
+
+  let sao_align, sao_slots, sao_size =
+    alloc_local_stack
+      (* if ra_on_stack we add some space for the return address,
+         the alignment will be automatically corrected *)
+      (if ra_on_stack then size_of_ws pd else 0) (Sv.elements slots) atbl in
 
   (* FIXME: 1- make this U128 arch (call-conv) dependent; 2- make it a semantic requirement. *)
   let sao_align = if has_syscall fd.f_body && wsize_lt sao_align U128 then U128 else sao_align in

compiler/tests/success/subroutines/exec.jazz

@@ -0,0 +1,32 @@
+#returnaddress=stack
+fn add(reg u128 x, reg u128 z) -> reg u128 {
+  reg u128 y;
+  stack u128[2] t;
+  y = z;
+  t[0] = x;
+  t[1] = y;
+  x ^= t[0];
+  y ^= t[1];
+  x ^= y;
+  return x;
+}
+
+u128 c3 = 3;
+u128 c7 = 7;
+
+export fn main() -> reg u128 {
+
+  reg u128 x y z;
+  stack u128 s;
+  x = c3;
+  y = c7;
+  z = #set0_128();
+  s = z;
+  x = add(x, y);
+  x ^= s;
+  x = add(x, y);
+  x ^= s;
+  x = add(x, y);
+  x ^= s;
+  return x;
+}

compiler/tests/success/subroutines/rastack.jazz

@@ -1,4 +1,4 @@
-#[returnaddress=stack, stacksize=8] 
+#[returnaddress=stack, stacksize=16]
 fn copy(reg u64 x) -> reg u64 {
 reg u64 r;
 stack u64 t;
@@ -48,7 +48,8 @@ result = copy32_32(arg);
 return result;
 }
 
-#[returnaddress=stack, stacksize=20, stackalign=u128]
+// Remark: The compiler can reduce the stack space to 32, if it try to fill the hole with q.
+#[returnaddress=stack, stacksize=36, stackalign=u128]
 fn copy128(reg u128 x) -> reg u128 {
 reg u128 r w;
 stack u128 t;

proofs/compiler/linearization.v

@@ -257,6 +257,13 @@ Definition stack_frame_allocation_size (e: stk_fun_extra) : Z :=
      is_align (wrepr Uptr ofs) ws (* Stack slot is aligned *)
     ].
 
+  Let check_stack_ofs_internal_call e ofs ws : bool :=
+    [&&
+     ofs == 0%Z,
+     wsize_size ws == sf_stk_ioff e &
+     (ws ≤ sf_align e)%CMP (* Stack frame is aligned for storing words of size ws *)
+    ].
+
   Definition all_disjoint_aligned_between (lo hi: Z) (al: wsize) A (m: seq A) (slot: A → cexec (Z * wsize)) : cexec unit :=
     Let last := foldM (λ a base,
                        Let: (ofs, ws) := slot a in
@@ -346,7 +353,7 @@ Definition check_fd (fn: funname) (fd:sfundef) :=
   Let _ := assert match sf_return_address e with
                   | RAnone => true
                   | RAreg ra => vtype ra == sword Uptr
-                  | RAstack ofs => check_stack_ofs e ofs Uptr
+                  | RAstack ofs => check_stack_ofs_internal_call e ofs Uptr
                   end
                   (E.error "bad return-address") in
   let ok_save_stack :=
@@ -387,10 +394,12 @@ Definition check_prog :=
   Let _ := map_cfprog_name check_fd (p_funcs p) in
   ok tt.
 
-Definition allocate_stack_frame (free: bool) (ii: instr_info) (sz: Z) : lcmd :=
+Definition allocate_stack_frame (free: bool) (ii: instr_info) (sz: Z) (rastack: bool) : lcmd :=
   if sz == 0%Z
   then [::]
-  else let args := if free
+  else
+    let sz := if rastack then (sz - wsize_size Uptr)%Z else sz in
+    let args := if free
                    then (lip_allocate_stack_frame liparams) rspi sz
                    else (lip_free_stack_frame liparams) rspi sz
        in [:: MkLI ii (Lopn args.1.1 args.1.2 args.2) ].
@@ -513,9 +522,10 @@ Fixpoint linear_i (i:instr) (lbl:label) (lc:lcmd) :=
       let ra := sf_return_address e in
       if ra == RAnone then (lbl, lc)
       else
+        let rastack := if ra is RAstack _ then true else false in
         let sz := stack_frame_allocation_size e in
-        let before := allocate_stack_frame false ii sz in
-        let after := allocate_stack_frame true ii sz in
+        let before := allocate_stack_frame false ii sz rastack in
+        let after := allocate_stack_frame true ii sz rastack in
         let lret := lbl in
         let lbl := next_lbl lbl in
         let lcall := (fn', if fn' == fn
@@ -543,19 +553,15 @@ Fixpoint linear_i (i:instr) (lbl:label) (lc:lcmd) :=
         | RAstack z =>
           (* Save return address to the stack with an offset.
            * 1. Allocate stack frame.
-           * 2. Store return label in ra.
-           * 3. Push ra to stack.
-           * 4. Insert jump to callee.
-           * 5. Insert return label (callee will jump back here).
-           * 6. Free stack frame.
-           * 7. Continue.
+           * 2. Call callee.
+           * 3. Insert return label (callee will jump back here).
+           * 4. Free stack frame.
+           * 5. Continue.
            *)
           if extra_free_registers ii is Some ra
-          then let glob_ra := Gvar (VarI ra dummy_var_info) Slocal in
+          then
                (lbl, before
-                       ++ MkLI ii (LstoreLabel ra lret)
-                       :: of_olinstr_r ii (lstore rspi z Uptr glob_ra)
-                       :: MkLI ii (Lgoto lcall)
+                       ++ MkLI ii (Lcall lcall)
                        :: MkLI ii (Llabel lret)
                        :: after
                        ++ lc
@@ -576,7 +582,7 @@ Definition linear_body (e: stk_fun_extra) (body: cmd) : lcmd :=
        , 2%positive
        )
      | RAstack z =>
-       ( [:: MkLI dummy_instr_info (Ligoto (Pload Uptr rspi (cast_const z))) ]
+       ( [:: MkLI dummy_instr_info Lret ]
        , [:: MkLI dummy_instr_info (Llabel 1) ]
        , 2%positive
        )
@@ -592,7 +598,7 @@ Definition linear_body (e: stk_fun_extra) (body: cmd) : lcmd :=
          let r := VarI x dummy_var_info in
          ( [:: of_olinstr_r dummy_instr_info (lmove rspi Uptr (Gvar r Slocal)) ]
          , of_olinstr_r dummy_instr_info (lmove r Uptr rspg)
-             :: allocate_stack_frame false dummy_instr_info (sf_stk_sz e + sf_stk_extra_sz e)
+             :: allocate_stack_frame false dummy_instr_info (sf_stk_sz e + sf_stk_extra_sz e) false
              ++ [:: ensure_rsp_alignment dummy_instr_info e.(sf_align) ]
          , 1%positive
          )
@@ -607,7 +613,7 @@ Definition linear_body (e: stk_fun_extra) (body: cmd) : lcmd :=
          ( pop_to_save dummy_instr_info e.(sf_to_save)
              ++ [:: of_olinstr_r dummy_instr_info (lload rspi Uptr rspi ofs) ]
          , of_olinstr_r dummy_instr_info (lmove var_tmpi Uptr rspg)
-             :: allocate_stack_frame false dummy_instr_info (sf_stk_sz e + sf_stk_extra_sz e)
+             :: allocate_stack_frame false dummy_instr_info (sf_stk_sz e + sf_stk_extra_sz e) false
              ++ ensure_rsp_alignment dummy_instr_info e.(sf_align)
              :: of_olinstr_r dummy_instr_info (lstore rspi ofs Uptr var_tmpg)
              :: push_to_save dummy_instr_info e.(sf_to_save)