From cacde2524f5d5d068e48dc41fcae64b1e9b23db9 Mon Sep 17 00:00:00 2001
From: Zhicheng HUI <huizhicheng123@gmail.com>
Date: Thu, 22 Aug 2024 13:11:58 +0200
Subject: [PATCH] simplify code

---
 src/ast/code_generator.ml | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/ast/code_generator.ml b/src/ast/code_generator.ml
index 025e8a09c..6fa7009c8 100644
--- a/src/ast/code_generator.ml
+++ b/src/ast/code_generator.ml
@@ -12,6 +12,7 @@ type type_env =
   ; global_types : binary val_type list
   ; result_types : binary val_type list
   ; temp : binary indice
+  ; result : int -> binary indice
   ; owi_i32 : binary indice
   ; owi_i64 : binary indice
   ; owi_f32 : binary indice
@@ -34,6 +35,7 @@ let build_type_env (m : modul)
   in
   let result_types = snd func_ty in
   let temp = Raw (List.length param_types) in
+  let result i = Raw (List.length param_types + i + 1) in
   let owi_i32 =
     match
       Array.find_index
@@ -88,6 +90,7 @@ let build_type_env (m : modul)
   ; global_types
   ; result_types
   ; temp
+  ; result
   ; owi_i32
   ; owi_i64
   ; owi_f32
@@ -222,12 +225,11 @@ let rec term_generate tenv (term : binary term) :
     binop_generate b expr1 ty1 expr2 ty2
   | Result (Some i) -> (
     match List.nth_opt tenv.result_types i with
-    | Some t ->
-      Ok ([ Local_get (Raw (List.length tenv.param_types + i + 1)) ], t)
+    | Some t -> Ok ([ Local_get (tenv.result i) ], t)
     | None -> Error (`Spec_type_error Fmt.(str "%a" pp_term term)) )
   | Result None -> (
     match List.nth_opt tenv.result_types 0 with
-    | Some t -> Ok ([ Local_get (Raw (List.length tenv.param_types + 1)) ], t)
+    | Some t -> Ok ([ Local_get (tenv.result 0) ], t)
     | None -> Error (`Spec_type_error Fmt.(str "%a" pp_term term)) )
 
 let binpred_generate (b : binpred) (expr1 : binary expr) (ty1 : binary val_type)
@@ -451,11 +453,11 @@ let contract_generate (add_assert_to_assume : bool)
     List.init (List.length tenv.param_types) (fun i -> Local_get (Raw i))
     @ [ Call (Raw old_index) ]
     @ List.init (List.length tenv.result_types) (fun i ->
-          Local_set (Raw (List.length tenv.param_types + i + 1)) )
+          Local_set (tenv.result i) )
   in
   let return =
     List.init (List.length tenv.result_types) (fun i ->
-        Local_get (Raw (List.length tenv.param_types + i + 1)) )
+        Local_get (tenv.result i) )
   in
   let* precond_checker =
     list_concat_map (prop_generate add_assert_to_assume tenv) preconditions