A fix for https://github.com/JetBrains-Research/Lama/issues/13

regression/orig/test108.log

@@ -0,0 +1 @@
+> 100

regression/test108.input

@@ -0,0 +1 @@
+100

regression/test108.lama

@@ -0,0 +1,8 @@
+fun foo (x) {
+    return x;
+    0
+}
+
+local n = read ();
+
+write (foo (n))

src/SM.ml

@@ -26,6 +26,7 @@ let show_scope = show(scope);;
 (* store a value into a reference            *) | STI
 (* store a value into array/sexp/string      *) | STA                                  
 (* a label                                   *) | LABEL   of string
+(* a forwarded label                         *) | FLABEL  of string
 (* a scope label                             *) | SLABEL  of string
 (* unconditional jump                        *) | JMP     of string
 (* conditional jump                          *) | CJMP    of string * string
@@ -150,7 +151,8 @@ let rec eval env (((cstack, stack, glob, loc, i, o) as conf) : config) = functio
                                  Value.update_elem x (Value.to_int j) v;
                                  eval env (cstack, v::stack', glob, loc, i, o) prg'
                                  
-    | SLABEL _ | LABEL  _     -> eval env conf prg'
+    | SLABEL _ | LABEL  _ | FLABEL _ -> eval env conf prg'
+                                      
     | JMP    l                -> eval env conf (env#labeled l)
     | CJMP  (c, l)            -> let x::stack' = stack in
                                  eval env (cstack, stack', glob, loc, i, o) (if (c = "z" && Value.to_int x = 0) || (c = "nz" && Value.to_int x <> 0) then env#labeled l else prg')
@@ -232,7 +234,8 @@ module M = Map.Make (String)
 class indexer prg =
   let rec make_env m = function
   | []               -> m
-  | (LABEL l) :: tl -> make_env (M.add l tl m) tl
+  | (LABEL l)  :: tl
+  | (FLABEL l) :: tl -> make_env (M.add l tl m) tl
   | _ :: tl          -> make_env m tl
   in 
   let m = make_env M.empty prg in
@@ -442,6 +445,7 @@ object (self : 'self)
   val decls        = []
   val funinfo      = new funinfo
   val line         = None
+  val end_label    = ""
 
   method show_funinfo = funinfo#show_funinfo
 
@@ -482,6 +486,11 @@ object (self : 'self)
   method global_scope = scope_index = 0
                       
   method get_label     = (label @@ string_of_int label_index), {< label_index = label_index + 1 >}
+  method get_end_label =    
+    let lab = label @@ string_of_int label_index in
+    lab, {< end_label = lab; label_index = label_index + 1 >}
+    
+  method end_label = end_label
     
   method nargs   = scope.arg_index
   method nlocals = scope.nlocals
@@ -837,7 +846,7 @@ let compile cmd ((imports, infixes), p) =
                                let cond, env = env#get_label in
                                let env, fe, se = compile_expr false lexp env c in
                                let env, _ , s  = compile_expr false cond env s in
-                               env, false, [JMP cond; LABEL loop] @ s @ [LABEL cond] @ se @ (if fe then [LABEL lexp] else []) @ [CJMP ("nz", loop)]
+                               env, false, [JMP cond; FLABEL loop] @ s @ [LABEL cond] @ se @ (if fe then [LABEL lexp] else []) @ [CJMP ("nz", loop)]
                                                                                                   
   | Expr.Repeat (s, c)      -> let lexp , env = env#get_label in
                                let loop , env = env#get_label in
@@ -847,9 +856,9 @@ let compile cmd ((imports, infixes), p) =
                                env, false, [LABEL loop] @ body @ (if flag then [LABEL check] else []) @ se @ (if fe then [LABEL lexp] else []) @ [CJMP ("z", loop)]
 
   | Expr.Return (Some e)    -> let lret, env = env#get_label in
-                               add_code (compile_expr true lret env e) lret false [RET]
+                               add_code (compile_expr true lret env e) lret false [JMP env#end_label] (* [RET] *)
                                
-  | Expr.Return None        -> env, false, [CONST 0; RET]
+  | Expr.Return None        -> env, false, [CONST 0; (*RET*) JMP env#end_label]
 
   | Expr.Leave              -> env, false, []
                                  
@@ -890,7 +899,7 @@ let compile cmd ((imports, infixes), p) =
     let env             = env#open_fun_scope blab elab fd in    
     (*Printf.eprintf "Lookup: %s\n%!" (try show(Value.designation) @@ snd (env#lookup "inner") with _ -> "no inner..."); *)
     let env             = List.fold_left (fun env arg -> env#add_arg arg) env args in 
-    let lend, env       = env#get_label in
+    let lend, env       = env#get_end_label in
     let env, flag, code = compile_expr true lend env stmt in
     let env, funcode    = compile_fundefs [] env in
     (*Printf.eprintf "Function: %s, closure: %s\n%!" name (show(list) (show(Value.designation)) env#closure);*)
@@ -900,8 +909,7 @@ let compile cmd ((imports, infixes), p) =
     let code =
       ([LABEL name; BEGIN (name, nargs, nlocals, closure, args, scopes); SLABEL blab] @ 
        code @
-       (if flag then [LABEL lend] else []) @
+       [LABEL lend; SLABEL elab; END]) :: funcode
-       [SLABEL elab; END]) :: funcode
     in
     env, code
   and compile_fundefs acc env =

src/X86.ml

@@ -229,7 +229,15 @@ let compile cmd env imports code =
     | [] -> env, []
     | instr :: scode' ->
         let stack = "" (* env#show_stack*) in
+        (* Printf.printf "insn=%s, stack=%s\n%!" (GT.show(insn) instr) (env#show_stack);  *)
         let env', code' =
+          if env#is_barrier
+          then match instr with
+               | LABEL  s -> if env#has_stack s then (env#drop_barrier)#retrieve_stack s, [Label s] else env, []
+               | FLABEL s -> env#drop_barrier, [Label s]
+               | SLABEL s -> env, [Label s]
+               | _        -> env, []
+          else
           match instr with
           | PUBLIC name -> env#register_public name, []
           | EXTERN name -> env#register_extern name, []
@@ -389,8 +397,8 @@ let compile cmd env imports code =
                  else [Binop (op, x, y); Or1 y]
              )
              
-          | LABEL s     -> (if env#is_barrier then (env#drop_barrier)#retrieve_stack s else env), [Label s]
+          | LABEL  s 
-
+          | FLABEL s
           | SLABEL s    -> env, [Label s]
 
 	  | JMP   l     -> (env#set_stack l)#set_barrier, [Jmp l]
@@ -622,18 +630,23 @@ class env prg =
     method is_barrier = barrier
 
     (* set barrier *)
-    method set_barrier = {< barrier = true >}
+    method set_barrier = {< stack = []; barrier = true >}
 
     (* drop barrier *)
     method drop_barrier = {< barrier = false >}
 
     (* associates a stack to a label *)
-    method set_stack l = (*Printf.printf "Setting stack for %s\n" l;*) {< stackmap = M.add l stack stackmap >}
+    method set_stack l = (*Printf.printf "Setting stack for %s\n" l;*)
+      {< stackmap = M.add l stack stackmap >}
 
     (* retrieves a stack for a label *)
     method retrieve_stack l = (*Printf.printf "Retrieving stack for %s\n" l;*)
       try {< stack = M.find l stackmap >} with Not_found -> self
 
+    (* checks if there is a stack for a label *)
+    method has_stack l = (*Printf.printf "Retrieving stack for %s\n" l;*)
+      M.mem l stackmap
+
     (* gets a name for a global variable *)
     method loc x =
       match x with

src/version.ml

@@ -1 +1 @@
-let version = "Version 1.00, 47d42aa4a, Wed Oct 21 07:40:50 2020 +0300"
+let version = "Version 1.00, 60e69ff31, Thu Oct 22 23:00:31 2020 +0300"