feat(compiler)!: Stack-allocated Chars (#1103)

Author: ospencer

compiler/src/codegen/compcore.re

@@ -1679,59 +1679,11 @@ let allocate_bytes_uninitialized = (wasm_mod, env, size) => {
   );
 };
 
-let allocate_char = (wasm_mod, env, char) => {
-  // Copy bytes into a fresh buffer so we can guarantee a copy of a full word
-  let bytes = Bytes.make(4, Char.chr(0));
-  // OCaml String#length is byte length, not Unicode character length
-  // Guaranteed not to be longer than 4 bytes by well-formedness
-  Bytes.blit_string(char, 0, bytes, 0, String.length(char));
-  let value = Bytes.get_int32_le(bytes, 0);
-
-  let get_swap = () => get_swap(wasm_mod, env, 0);
-  let tee_swap = tee_swap(wasm_mod, env, 0);
-  Expression.Block.make(
-    wasm_mod,
-    gensym_label("allocate_char"),
-    [
-      store(
-        ~offset=0,
-        wasm_mod,
-        tee_swap(heap_allocate(wasm_mod, env, 2)),
-        Expression.Const.make(
-          wasm_mod,
-          const_int32(tag_val_of_heap_tag_type(CharType)),
-        ),
-      ),
-      store(
-        ~offset=4,
-        wasm_mod,
-        get_swap(),
-        Expression.Const.make(wasm_mod, wrap_int32(value)),
-      ),
-      get_swap(),
-    ],
-  );
-};
-
-let allocate_char_uninitialized = (wasm_mod, env) => {
-  let get_swap = () => get_swap(wasm_mod, env, 0);
-  let tee_swap = tee_swap(wasm_mod, env, 0);
-  Expression.Block.make(
-    wasm_mod,
-    gensym_label("allocate_char_uninitialized"),
-    [
-      store(
-        ~offset=0,
-        wasm_mod,
-        tee_swap(heap_allocate(wasm_mod, env, 2)),
-        Expression.Const.make(
-          wasm_mod,
-          const_int32(tag_val_of_heap_tag_type(CharType)),
-        ),
-      ),
-      get_swap(),
-    ],
-  );
+let create_char = (wasm_mod, env, char) => {
+  let uchar = List.hd @@ Utf8.decodeUtf8String(char);
+  let uchar_int: int = Utf8__Uchar.toInt(uchar);
+  let grain_char = uchar_int lsl 3 lor 0b010;
+  Expression.Const.make(wasm_mod, const_int32(grain_char));
 };
 
 let allocate_closure =
@@ -2261,7 +2213,6 @@ let allocate_rational = (wasm_mod, env, n, d) => {
 
 let compile_prim0 = (wasm_mod, env, p0): Expression.t => {
   switch (p0) {
-  | AllocateChar => allocate_char_uninitialized(wasm_mod, env)
   | AllocateInt32 => allocate_number_uninitialized(wasm_mod, env, BoxedInt32)
   | AllocateInt64 => allocate_number_uninitialized(wasm_mod, env, BoxedInt64)
   | AllocateFloat32 =>
@@ -2307,6 +2258,25 @@ let compile_prim1 = (wasm_mod, env, p1, arg, loc): Expression.t => {
       compiled_arg,
       Expression.Const.make(wasm_mod, const_int32(0x1)),
     )
+  | TagChar =>
+    Expression.Binary.make(
+      wasm_mod,
+      Op.xor_int32,
+      Expression.Binary.make(
+        wasm_mod,
+        Op.shl_int32,
+        compiled_arg,
+        Expression.Const.make(wasm_mod, const_int32(0x3)),
+      ),
+      Expression.Const.make(wasm_mod, const_int32(0b10)),
+    )
+  | UntagChar =>
+    Expression.Binary.make(
+      wasm_mod,
+      Op.shr_s_int32,
+      compiled_arg,
+      Expression.Const.make(wasm_mod, const_int32(0x3)),
+    )
   | Not =>
     /* Flip the first bit */
     Expression.Binary.make(
@@ -2691,7 +2661,7 @@ let compile_allocation = (wasm_mod, env, alloc_type) =>
   | MRecord(ttag, elts) => allocate_record(wasm_mod, env, ttag, elts)
   | MBytes(bytes) => allocate_bytes(wasm_mod, env, bytes)
   | MString(str) => allocate_string(wasm_mod, env, str)
-  | MChar(char) => allocate_char(wasm_mod, env, char)
+  | MChar(char) => create_char(wasm_mod, env, char)
   | MADT(ttag, vtag, elts) => allocate_adt(wasm_mod, env, ttag, vtag, elts)
   | MInt32(i) =>
     allocate_int32(

compiler/src/codegen/mashtree.re

@@ -172,7 +172,6 @@ type wasm_op =
 
 type prim0 =
   Parsetree.prim0 =
-    | AllocateChar
     | AllocateInt32
     | AllocateInt64
     | AllocateFloat32
@@ -194,6 +193,8 @@ type prim1 =
     | BytesSize
     | TagSimpleNumber
     | UntagSimpleNumber
+    | TagChar
+    | UntagChar
     | Not
     | Box
     | Unbox

compiler/src/codegen/value_tags.re

@@ -4,7 +4,6 @@ open Sexplib.Conv;
 [@deriving sexp]
 type heap_tag_type =
   | StringType
-  | CharType
   | ADTType
   | RecordType
   | ArrayType
@@ -16,26 +15,24 @@ type heap_tag_type =
 let tag_val_of_heap_tag_type =
   fun
   | StringType => 1
-  | CharType => 2
-  | ADTType => 3
-  | RecordType => 4
-  | ArrayType => 5
-  | BoxedNumberType => 6
-  | LambdaType => 7
-  | TupleType => 8
-  | BytesType => 9;
+  | ADTType => 2
+  | RecordType => 3
+  | ArrayType => 4
+  | BoxedNumberType => 5
+  | LambdaType => 6
+  | TupleType => 7
+  | BytesType => 8;
 
 let heap_tag_type_of_tag_val =
   fun
   | 1 => StringType
-  | 2 => CharType
-  | 3 => ADTType
-  | 4 => RecordType
-  | 5 => ArrayType
-  | 6 => BoxedNumberType
-  | 7 => LambdaType
-  | 8 => TupleType
-  | 9 => BytesType
+  | 2 => ADTType
+  | 3 => RecordType
+  | 4 => ArrayType
+  | 5 => BoxedNumberType
+  | 6 => LambdaType
+  | 7 => TupleType
+  | 8 => BytesType
   | x => failwith(Printf.sprintf("Unknown tag type: %d", x));
 
 [@deriving sexp]
@@ -67,18 +64,5 @@ let boxed_number_tag_type_of_tag_val =
 type tag_type =
   | NumberTagType
   | ConstTagType
+  | CharTagType
   | GenericHeapType(option(heap_tag_type));
-
-let and_mask_of_tag_type =
-  fun
-  | NumberTagType => 0b0001
-  | ConstTagType => 0b0111
-  | GenericHeapType(_) => 0b0111;
-
-let tag_val_of_tag_type =
-  fun
-  | NumberTagType => 0b0001
-  | ConstTagType => 0b0110
-  | GenericHeapType(_) => 0b0000;
-
-let shift_amount_of_tag_type = tt => 31 - tag_val_of_tag_type(tt);

compiler/src/middle_end/analyze_purity.re

@@ -44,8 +44,7 @@ let rec analyze_comp_expression =
     | CImmExpr({imm_desc: ImmTrap}) => false
     | CImmExpr(_) => true
     | CPrim0(
-        AllocateChar | AllocateInt32 | AllocateInt64 | AllocateFloat32 |
-        AllocateFloat64 |
+        AllocateInt32 | AllocateInt64 | AllocateFloat32 | AllocateFloat64 |
         AllocateRational,
       ) =>
       true
@@ -60,6 +59,8 @@ let rec analyze_comp_expression =
         BytesSize |
         TagSimpleNumber |
         UntagSimpleNumber |
+        TagChar |
+        UntagChar |
         Not |
         Box |
         Unbox |

compiler/src/middle_end/anftree.re

@@ -158,7 +158,6 @@ type wasm_op =
 
 type prim0 =
   Parsetree.prim0 =
-    | AllocateChar
     | AllocateInt32
     | AllocateInt64
     | AllocateFloat32
@@ -180,6 +179,8 @@ type prim1 =
     | BytesSize
     | TagSimpleNumber
     | UntagSimpleNumber
+    | TagChar
+    | UntagChar
     | Not
     | Box
     | Unbox

compiler/src/middle_end/anftree.rei

@@ -159,7 +159,6 @@ type wasm_op =
 
 type prim0 =
   Parsetree.prim0 =
-    | AllocateChar
     | AllocateInt32
     | AllocateInt64
     | AllocateFloat32
@@ -181,6 +180,8 @@ type prim1 =
     | BytesSize
     | TagSimpleNumber
     | UntagSimpleNumber
+    | TagChar
+    | UntagChar
     | Not
     | Box
     | Unbox

compiler/src/parsing/parsetree.re

@@ -303,7 +303,6 @@ type wasm_op =
 /** Zero-argument operators */
 [@deriving (sexp, yojson)]
 type prim0 =
-  | AllocateChar
   | AllocateInt32
   | AllocateInt64
   | AllocateFloat32
@@ -326,6 +325,8 @@ type prim1 =
   | BytesSize
   | TagSimpleNumber
   | UntagSimpleNumber
+  | TagChar
+  | UntagChar
   | Not
   | Box
   | Unbox

compiler/src/typed/builtin_types.re

@@ -174,7 +174,7 @@ let initial_env = (add_type, empty_env) =>
   |> add_type(ident_bool, decl_bool)
   |> add_type(ident_box, decl_box)
   |> add_type(ident_string, decl_abstr(path_string))
-  |> add_type(ident_char, decl_abstr(path_char))
+  |> add_type(ident_char, decl_abstr_imm(WasmI32, path_char))
   |> add_type(ident_void, decl_void)
   |> add_type(ident_array, decl_array)
   |> add_type(ident_fd, decl_abstr(path_fd))

compiler/src/typed/translprim.re

@@ -44,7 +44,6 @@ let prim_map =
       ("@heap.base", PrimitiveConstant(HeapBase)),
       ("@heap.start", PrimitiveConstant(HeapStart)),
       ("@heap.type_metadata", PrimitiveConstant(HeapTypeMetadata)),
-      ("@allocate.char", Primitive0(AllocateChar)),
       ("@allocate.int32", Primitive0(AllocateInt32)),
       ("@allocate.int64", Primitive0(AllocateInt64)),
       ("@allocate.float32", Primitive0(AllocateFloat32)),
@@ -63,6 +62,8 @@ let prim_map =
       ("@bytes.size", Primitive1(BytesSize)),
       ("@tag.simple_number", Primitive1(TagSimpleNumber)),
       ("@untag.simple_number", Primitive1(UntagSimpleNumber)),
+      ("@tag.char", Primitive1(TagChar)),
+      ("@untag.char", Primitive1(UntagChar)),
       ("@not", Primitive1(Not)),
       ("@box", Primitive1(Box)),
       ("@unbox", Primitive1(Unbox)),
@@ -1499,8 +1500,7 @@ let transl_prim = (env, desc) => {
       (Exp.constant(~loc, ~attributes=disable_gc, value), attrs);
     | Primitive0(
         (
-          AllocateChar | AllocateInt32 | AllocateInt64 | AllocateFloat32 |
-          AllocateFloat64 |
+          AllocateInt32 | AllocateInt64 | AllocateFloat32 | AllocateFloat64 |
           AllocateRational
         ) as p,
       ) => (

compiler/src/typed/typecore.re

@@ -100,7 +100,6 @@ let grain_type_of_wasm_prim_type =
 
 let prim0_type =
   fun
-  | AllocateChar
   | AllocateInt32
   | AllocateInt64
   | AllocateFloat32
@@ -125,6 +124,8 @@ let prim1_type =
       Builtin_types.type_number,
       Builtin_types.type_wasmi32,
     )
+  | TagChar => (Builtin_types.type_wasmi32, Builtin_types.type_char)
+  | UntagChar => (Builtin_types.type_char, Builtin_types.type_wasmi32)
   | Not => (Builtin_types.type_bool, Builtin_types.type_bool)
   | Box
   | BoxBind => {