Actually fix stack for return

2 files changed, 11 insertions, 12 deletions

diff --git a/src/lem_interp/interp.lem b/src/lem_interp/interp.lem
index 89ffc315..79b1e21a 100644
--- a/src/lem_interp/interp.lem
+++ b/src/lem_interp/interp.lem
@@ -854,13 +854,13 @@ let rec clear_stack_state stack =
     | Thunk_frame e t_level env lmem s -> Thunk_frame e t_level env lmem (clear_stack_state s)
 end
 
-let remove_top_stack_frame stack =
+let rec remove_top_stack_frame stack =
   match stack with
-  | Top -> Nothing
-  | Hole_frame _ _ _ _ _ Top -> Nothing
-  | Thunk_frame _ _ _ _ Top -> Nothing
-  | Hole_frame _ _ _ _ _ stack -> Just stack
-  | Thunk_frame _ _ _ _ stack -> Just stack
+  | Top -> Top
+  | Hole_frame _ _ _ _ _ Top -> Top
+  | Thunk_frame _ _ _ _ Top -> Top
+  | Hole_frame id e t_level env lmem stack -> Hole_frame id e t_level env lmem (remove_top_stack_frame stack)
+  | Thunk_frame e t_level env lmem stack -> Thunk_frame e t_level env lmem (remove_top_stack_frame stack)
 end
 
 (*functions for converting in progress evaluation back into expression for building current continuation*)
@@ -2194,8 +2194,7 @@ and interp_main mode t_level l_env l_mem (E_aux exp (l,annot)) =
   | E_return exp ->
     resolve_outcome
       (interp_main mode t_level l_env l_mem exp)
-      (fun v lm le ->
-         (Action (Return v) Top, l_mem, l_env))
+      (fun v lm le -> (Action (Return v) Top, l_mem, l_env))
       (fun a -> update_stack a (add_to_top_frame (fun e env -> (E_aux (E_return e) (l,annot), env)))) 
   | E_assert cond msg ->
     resolve_outcome 
diff --git a/src/lem_interp/interp_inter_imp.lem b/src/lem_interp/interp_inter_imp.lem
index 04af1f94..4824ecf7 100644
--- a/src/lem_interp/interp_inter_imp.lem
+++ b/src/lem_interp/interp_inter_imp.lem
@@ -336,7 +336,7 @@ let rec interp_to_value_helper arg ivh_mode err_str instr direction registers ev
     | (Interp.Error l msg,_,_) -> (Ivh_error (Internal_error msg), events_out)
     | (Interp.Action (Interp.Return value) stack,_,_) ->
       (match Interp.remove_top_stack_frame stack with
-       | Nothing ->
+       | Top ->
          if exn_seen
          then (Ivh_value_after_exn value, events_out)
          else (match ivh_mode with
@@ -351,7 +351,7 @@ let rec interp_to_value_helper arg ivh_mode err_str instr direction registers ev
                | Ivh_unsupported -> (Ivh_error (Unsupported_instruction_error instr), events_out)
                | Ivh_translate -> Assert_extra.failwith "Reached unreachable pattern" end)
             | _ -> (Ivh_error (Internal_error ("Value not an option for " ^ errk_str)), events_out)end) end)
-       | Just stack -> interp_to_value_helper arg ivh_mode err_str instr direction registers events exn_seen
+       | stack -> interp_to_value_helper arg ivh_mode err_str instr direction registers events exn_seen
                                               (fun _ -> Interp.resume mode stack (Just value)) end)
     | (Interp.Action (Interp.Call_extern i value) stack,_,_) ->
       match List.lookup i (Interp_lib.library_functions direction) with
@@ -664,8 +664,8 @@ let rec interp_to_outcome mode context thunk =
             end)
          | Interp.Return value ->
            (match Interp.remove_top_stack_frame next_state with
-            | Nothing -> (Done, lm)
-            | Just stack ->
+            | Top -> (Done, lm)
+            | stack ->
               interp_to_outcome mode context (fun _ -> Interp.resume mode.internal_mode stack (Just value)) end)
          | Interp.Step l Nothing Nothing -> (Internal Nothing Nothing (IState next_state context), lm)
          | Interp.Step l (Just name) Nothing -> (Internal (Just name) Nothing (IState next_state context), lm)