refactor: deduplicate two functions (#4083)

No functionality changes. Also some cosmetics in `compile.ml`.
dfinity · Jun 28, 2023 · c09b751 · c09b751
1 parent 6c64da4
commit c09b751
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 37 deletions.
diff --git a/src/codegen/compile.ml b/src/codegen/compile.ml
@@ -40,7 +40,7 @@ let ptr_unskew = 1l
 let prim_fun_name p stem = Printf.sprintf "%s<%s>" stem (Type.string_of_prim p)
 
 (* Helper functions to produce annotated terms (Wasm.AST) *)
-let nr x = { Wasm.Source.it = x; Wasm.Source.at = Wasm.Source.no_region }
+let nr x = Wasm.Source.{ it = x; at = no_region }
 
 let todo fn se x = Printf.eprintf "%s: %s" fn (Wasm.Sexpr.to_string 80 se); x
 
@@ -1650,7 +1650,7 @@ module Tagged = struct
         load_forwarding_pointer env ^^
         get_object ^^
         G.i (Compare (Wasm.Values.I32 I32Op.Eq)) ^^
-        E.else_trap_with env "missing object forwarding"  ^^
+        E.else_trap_with env "missing object forwarding" ^^
         get_object ^^
         (if unskewed then
           compile_unboxed_const ptr_unskew ^^
@@ -2401,7 +2401,7 @@ module TaggedSmallWord = struct
         get_exp ^^
         compile_unboxed_const 0l ^^
         G.i (Compare (Wasm.Values.I32 I32Op.GeS)) ^^
-        E.else_trap_with env "negative power"  ^^
+        E.else_trap_with env "negative power" ^^
         get_n ^^ get_exp ^^ compile_nat_power env ty
       )
 
@@ -4893,7 +4893,7 @@ module StableMem = struct
           get_offset ^^
           compile_const_64 (Int64.of_int page_size_bits) ^^
           G.i (Binary (Wasm.Values.I64 I64Op.ShrU)) ^^
-          get_mem_size env  ^^
+          get_mem_size env ^^
           G.i (Compare (Wasm.Values.I64 I64Op.LtU)) ^^
           E.else_trap_with env "StableMemory offset out of bounds")
     | _ -> assert false
@@ -7218,7 +7218,7 @@ module Stabilization = struct
         get_N ^^
         extend64 get_len ^^
         compile_add64_const 16L ^^
-        StableMem.ensure env  ^^
+        StableMem.ensure env ^^
 
         get_N ^^
         get_len ^^
@@ -8156,7 +8156,7 @@ module FuncDec = struct
       get_meth_pair ^^ Arr.load_field env 1l ^^ Blob.as_ptr_len env ^^
       (* The reply and reject callback *)
       push_continuations ^^
-      set_cb_index  ^^ get_cb_index ^^
+      set_cb_index ^^ get_cb_index ^^
       (* initiate call *)
       IC.system_call env "call_new" ^^
       cleanup_callback env ^^ get_cb_index ^^
@@ -9661,7 +9661,7 @@ and compile_prim_invocation (env : E.t) ae p es at =
     compile_exp_vanilla env ae e ^^
     Serialization.buffer_size env t ^^
     G.i Drop ^^
-    compile_add_const tydesc_len  ^^
+    compile_add_const tydesc_len ^^
     G.i (Convert (Wasm.Values.I64 I64Op.ExtendUI32))
 
   (* Other prims, unary *)

diff --git a/src/linking/linkModule.ml b/src/linking/linkModule.ml
@@ -541,7 +541,8 @@ let set_table_size new_size : module_' -> module_' = fun m ->
     }
   | _ -> raise (LinkError "Expect one table in first module")
 
-let fill_memory_base_import new_base : module_' -> module_' = fun m ->
+
+let fill_item_import module_name item_name new_base (m : module_') : module_' =
   (* We need to find the right import,
      replace all uses of get_global of that import with the constant,
      and finally rename all globals
@@ -551,8 +552,8 @@ let fill_memory_base_import new_base : module_' -> module_' = fun m ->
       | [] -> assert false
       | imp::is -> match imp.it.idesc.it with
         | GlobalImport _ty
-          when imp.it.module_name = Lib.Utf8.decode "env" &&
-               imp.it.item_name = Lib.Utf8.decode "__memory_base" ->
+          when imp.it.module_name = Lib.Utf8.decode module_name &&
+               imp.it.item_name = Lib.Utf8.decode item_name ->
           Int32.of_int i
         | GlobalImport _ ->
           go (i + 1) is
@@ -561,39 +562,18 @@ let fill_memory_base_import new_base : module_' -> module_' = fun m ->
     in go 0 m.imports in
 
     m |> fill_global base_global new_base
-      |> remove_imports is_global_import [(base_global, base_global)]
+      |> remove_imports is_global_import [base_global, base_global]
       |> rename_globals Int32.(fun i ->
           if i < base_global then i
           else if i = base_global then assert false
-          else sub i 1l
+          else sub i one
         )
 
-let fill_table_base_import new_base : module_' -> module_' = fun m ->
-  (* We need to find the right import,
-     replace all uses of get_global of that import with the constant,
-     and finally rename all globals
-  *)
-  let base_global =
-    let rec go i = function
-      | [] -> assert false
-      | imp::is -> match imp.it.idesc.it with
-        | GlobalImport _ty
-          when imp.it.module_name = Lib.Utf8.decode "env" &&
-               imp.it.item_name = Lib.Utf8.decode "__table_base" ->
-          Int32.of_int i
-        | GlobalImport _ ->
-          go (i + 1) is
-        | _ ->
-          go i is
-    in go 0 m.imports in
+let fill_memory_base_import new_base : module_' -> module_' =
+  fill_item_import "env" "__memory_base" new_base
 
-    m |> fill_global base_global new_base
-      |> remove_imports is_global_import [(base_global, base_global)]
-      |> rename_globals Int32.(fun i ->
-          if i < base_global then i
-          else if i = base_global then assert false
-          else sub i 1l
-        )
+let fill_table_base_import new_base : module_' -> module_' =
+  fill_item_import "env" "__table_base" new_base
 
 
 (* Concatenation of modules *)