Set the type of locals based on the expression they are initialized with

As a byproduct, the environment types get refined as well.

Author: vouillon

CHANGES.md

@@ -46,6 +46,7 @@
 * Compiler: rewrote inlining pass (#1935, #2018, #2027)
 * Compiler/wasm: optimize integer operations (#2032)
 * Compiler/wasm: use type analysis to remove some unnecessary uses of JavasScript strict equality (#2040)
+* Compiler/wasm: use more precise environment types (#2041)
 
 ## Bug fixes
 * Compiler: fix stack overflow issues with double translation (#1869)

compiler/lib-wasm/code_generation.ml

@@ -493,6 +493,11 @@ let tee ?typ x e =
     let* () = register_constant x e in
     return e
   else
+    let* typ =
+      match typ with
+      | Some _ -> return typ
+      | None -> expression_type e
+    in
     let* i = add_var ?typ x in
     return (W.LocalTee (i, e))
 
@@ -563,7 +568,14 @@ let rec store ?(always = false) ?typ x e =
               let* typ =
                 match typ with
                 | Some typ -> return typ
-                | None -> value_type
+                | None -> (
+                    if always
+                    then value_type
+                    else
+                      let* typ = expression_type e in
+                      match typ with
+                      | None -> value_type
+                      | Some typ -> return typ)
               in
               let* default, typ', cast = default_value typ in
               let* () =
@@ -577,6 +589,11 @@ let rec store ?(always = false) ?typ x e =
           in
           instr (GlobalSet (x, e))
         else
+          let* typ =
+            match typ with
+            | Some _ -> return typ
+            | None -> if always then return None else expression_type e
+          in
           let* i = add_var ?typ x in
           instr (LocalSet (i, e))