[compiler] Make use of i31 literals (#1208)

Make the optimized form of variant actually i31 without assuming the
representation during generics specialization. Then we do the `i * 2 +
1` conversion during LIR printing and wasm lowering. Eventually, the
WASM one will be replaced by the actual WASM i31.

crates/samlang-ast/src/lir.rs

@@ -113,8 +113,10 @@ impl Expression {
     str_table: &HashMap<PStr, usize>,
   ) {
     match self {
-      Expression::Int32Literal(i) | Expression::Int31Literal(i) => {
-        collector.push_str(&i.to_string())
+      Expression::Int32Literal(i) => collector.push_str(&i.to_string()),
+      Expression::Int31Literal(i) => {
+        let i32_form = i * 2 + 1;
+        collector.push_str(&i32_form.to_string())
       }
       Expression::Variable(n, _) => collector.push_str(n.as_str(heap)),
       Expression::StringName(n) => {
@@ -574,6 +576,7 @@ pub fn ts_prolog() -> String {
   let table = &SymbolTable::new();
   let mut collector = String::new();
 
+  collector.push_str("type i31 = number;\n");
   collector.push_str("const ");
   FunctionName::STR_CONCAT.write_encoded(&mut collector, heap, table);
   collector.push_str(" = ([, a]: _Str, [, b]: _Str): _Str => [1, a + b];\n");

crates/samlang-ast/src/lir_tests.rs

@@ -161,7 +161,7 @@ mod tests {
                 heap.alloc_str_for_test("dd"),
                 BinaryOperator::MINUS,
                 ZERO,
-                Expression::Int31Literal(-2147483648),
+                Expression::Int31Literal(-21478),
               ),
               Statement::binary(heap.alloc_str_for_test("dd"), BinaryOperator::MUL, ZERO, ZERO),
               Statement::binary(heap.alloc_str_for_test("dd"), BinaryOperator::DIV, ZERO, ZERO),
@@ -291,7 +291,7 @@ function __$f(v1: (t0: number) => number): number {{
     let dd = typeof 0 === 'object';
     let dd = 0 + 0;
     let dd = 0 + 0;
-    let dd = 0 - -2147483648;
+    let dd = 0 - -42955;
     let dd = 0 * 0;
     let dd = Math.floor(0 / 0);
     let dd = 0 % 0;

crates/samlang-ast/src/mir.rs

@@ -240,7 +240,7 @@ impl ClosureTypeDefinition {
 pub enum EnumTypeDefinition {
   Boxed(Vec<Type>),
   Unboxed(TypeNameId),
-  Int,
+  Int31,
 }
 
 impl EnumTypeDefinition {
@@ -250,7 +250,7 @@ impl EnumTypeDefinition {
         format!("Boxed({})", types.iter().map(|it| it.pretty_print(heap, table)).join(", "))
       }
       EnumTypeDefinition::Unboxed(t) => format!("Unboxed({})", t.encoded_for_test(heap, table)),
-      EnumTypeDefinition::Int => "int".to_string(),
+      EnumTypeDefinition::Int31 => "i31".to_string(),
     }
   }
 }
@@ -423,7 +423,8 @@ impl Expression {
 
   pub fn debug_print(&self, heap: &Heap, table: &SymbolTable) -> String {
     match self {
-      Expression::Int32Literal(i) | Expression::Int31Literal(i) => i.to_string(),
+      Expression::Int32Literal(i) => i.to_string(),
+      Expression::Int31Literal(i) => format!("{} as i31", i),
       Expression::StringName(n) => format!("\"{}\"", n.as_str(heap)),
       Expression::Variable(v) => v.debug_print(heap, table),
     }

crates/samlang-ast/src/mir_tests.rs

@@ -148,7 +148,7 @@ mod tests {
     };
     let ed1 = EnumTypeDefinition::Unboxed(table.create_type_name_for_test(PStr::UPPER_D));
     let ed2 = EnumTypeDefinition::Boxed(vec![INT_32_TYPE, INT_32_TYPE]);
-    let ed3 = EnumTypeDefinition::Int;
+    let ed3 = EnumTypeDefinition::Int31;
     assert!(ed1.eq(&ed1));
     assert!(ed2.eq(&ed2));
     assert!(ed3.eq(&ed3));
@@ -158,7 +158,7 @@ mod tests {
     };
     assert_eq!("object type _A = [int, int]", d1.pretty_print(heap, table));
     assert_eq!(
-      "variant type _B = [Unboxed(_D), Boxed(int, int), int]",
+      "variant type _B = [Unboxed(_D), Boxed(int, int), i31]",
       d2.pretty_print(heap, table)
     );
     format!("{d1:?} {d2:?}");
@@ -335,7 +335,7 @@ if 0 {
   let dd = 0 is _Str;
   let dd = 0 + 0;
   let dd = 0 + 0;
-  let dd = 0 - -2147483648;
+  let dd = 0 - -2147483648 as i31;
   let dd = 0 * 0;
   let dd = 0 / 0;
   let dd = 0 % 0;

crates/samlang-compiler/src/lir_lowering.rs

@@ -435,7 +435,7 @@ fn generate_inc_ref_fn() -> lir::Function {
             is_zero,
             hir::BinaryOperator::EQ,
             lir::Expression::Variable(old_ref_count, lir::INT_32_TYPE),
-            lir::Expression::Int31Literal(0),
+            lir::Expression::Int32Literal(0),
           ),
           lir::Statement::SingleIf {
             condition: lir::Expression::Variable(is_zero, lir::INT_32_TYPE),
@@ -711,7 +711,7 @@ pub fn compile_mir_to_lir(heap: &mut Heap, sources: mir::Sources) -> lir::Source
       mir::TypeDefinitionMappings::Enum(variants) => {
         for (i, variant) in variants.iter().enumerate() {
           match variant {
-            mir::EnumTypeDefinition::Unboxed(_) | mir::EnumTypeDefinition::Int => {}
+            mir::EnumTypeDefinition::Unboxed(_) | mir::EnumTypeDefinition::Int31 => {}
             mir::EnumTypeDefinition::Boxed(types) => {
               let name = symbol_table.derived_type_name_with_subtype_tag(type_def.name, i as u32);
               let mut mappings = Vec::with_capacity(types.len() + 1);
@@ -816,7 +816,7 @@ mod tests {
         TypeDefinition {
           name: table.create_type_name_for_test(heap.alloc_str_for_test("Variant")),
           mappings: TypeDefinitionMappings::Enum(vec![
-            EnumTypeDefinition::Int,
+            EnumTypeDefinition::Int31,
             EnumTypeDefinition::Unboxed(TypeNameId::STR),
             EnumTypeDefinition::Boxed(vec![INT_32_TYPE, INT_31_TYPE]),
           ]),
@@ -1068,7 +1068,7 @@ type _Variant = [number, number];
 function __$cc(): i31 {{
   let _t1: (t0: any, t1: number) => number = cc[1];
   let _t2: any = cc[2];
-  _t1(_t2, 0);
+  _t1(_t2, 1);
   let v1: number = a[1];
   let v2: number = b[1];
   let v3: number = b[2];

crates/samlang-compiler/src/mir_constant_param_elimination.rs

@@ -660,7 +660,7 @@ function __$str_const(): int {
 function __$func_with_consts(c: _C, e: _E): int {
   let _ = !0;
   let _ = 0 + (c: int);
-  let _ = 0 + 0;
+  let _ = 0 + 0 as i31;
   let _: int = 0[0];
   let _: __ = Closure { fun: (__$otherwise_optimizable: () -> int), context: 0 };
   (_: int)(0);

crates/samlang-compiler/src/mir_generics_specialization.rs

@@ -195,21 +195,21 @@ impl Rewriter {
               ),
             });
           }
-          mir::EnumTypeDefinition::Int => {
+          mir::EnumTypeDefinition::Int31 => {
             let casted_collector = heap.alloc_temp_str();
             let comparison_temp = heap.alloc_temp_str();
             // We cast the test expression to int to perform the test.
             // Once we have i31 type, the test should be performed on i31.
             collector.push(mir::Statement::Cast {
               name: casted_collector,
-              type_: mir::INT_32_TYPE,
+              type_: mir::INT_31_TYPE,
               assigned_expression: test_expr,
             });
             collector.push(mir::Statement::binary(
               comparison_temp,
               hir::BinaryOperator::EQ,
               mir::Expression::var_name(casted_collector, mir::INT_32_TYPE),
-              mir::Expression::i32(i32::try_from(*tag * 2 + 1).unwrap()),
+              mir::Expression::Int31Literal(i32::try_from(*tag).unwrap()),
             ));
             collector.push(mir::Statement::IfElse {
               condition: mir::Expression::var_name(comparison_temp, mir::INT_32_TYPE),
@@ -310,13 +310,13 @@ impl Rewriter {
               ),
             });
           }
-          mir::EnumTypeDefinition::Int => {
+          mir::EnumTypeDefinition::Int31 => {
             debug_assert!(associated_data_list.is_empty());
             // Cast from more specific subtype to the general enum type.
             collector.push(mir::Statement::Cast {
               name: *enum_variable_name,
               type_: mir::Type::Id(enum_type),
-              assigned_expression: mir::Expression::i32(i32::try_from(*tag * 2 + 1).unwrap()),
+              assigned_expression: mir::Expression::Int31Literal(i32::try_from(*tag).unwrap()),
             });
           }
         }
@@ -532,7 +532,7 @@ impl Rewriter {
             let mut already_unused_boxed_optimization = None;
             for (tag, (_, types)) in hir_variants.iter().enumerate() {
               if types.is_empty() {
-                mir_variants.push(mir::EnumTypeDefinition::Int);
+                mir_variants.push(mir::EnumTypeDefinition::Int31);
               } else {
                 if let Some((i, t)) = already_unused_boxed_optimization {
                   mir_variants[i] = mir::EnumTypeDefinition::Boxed(vec![mir::INT_32_TYPE, t]);
@@ -595,7 +595,9 @@ impl Rewriter {
                 mir::EnumTypeDefinition::Boxed(_) => {}
                 // Deopt if enum can be int or unboxed.
                 // This is essential to correctly pattern match on Some(Some(Some(_)))
-                mir::EnumTypeDefinition::Unboxed(_) | mir::EnumTypeDefinition::Int => return false,
+                mir::EnumTypeDefinition::Unboxed(_) | mir::EnumTypeDefinition::Int31 => {
+                  return false
+                }
               }
             }
             true
@@ -1294,11 +1296,11 @@ closure type DUMMY_CC___Str__Str = (_Str) -> _Str
 closure type DUMMY_CC__int__Str = (int) -> _Str
 variant type _Str = []
 object type DUMMY_J = [int]
-variant type DUMMY_I__int__Str = [int, int]
-variant type DUMMY_I___Str__Str = [int, int]
-variant type DUMMY_Enum = [Unboxed(DUMMY_J), int]
-variant type DUMMY_Enum3 = [Boxed(int, DUMMY_J), Boxed(int, DUMMY_J), int]
-variant type DUMMY_Enum2 = [Boxed(int, int), int]
+variant type DUMMY_I__int__Str = [i31, i31]
+variant type DUMMY_I___Str__Str = [i31, i31]
+variant type DUMMY_Enum = [Unboxed(DUMMY_J), i31]
+variant type DUMMY_Enum3 = [Boxed(int, DUMMY_J), Boxed(int, DUMMY_J), i31]
+variant type DUMMY_Enum2 = [Boxed(int, int), i31]
 function _DUMMY_I$main(): int {
   let finalV: int;
   if 1 {
@@ -1328,8 +1330,8 @@ function _DUMMY_I$main(): int {
     let a = (b: DUMMY_Enum) as DUMMY_J;
   } else {
   }
-  let _t3 = (b: DUMMY_Enum) as int;
-  let _t4 = (_t3: int) == 3;
+  let _t3 = (b: DUMMY_Enum) as i31;
+  let _t4 = (_t3: int) == 1 as i31;
   if (_t4: int) {
   } else {
   }
@@ -1350,12 +1352,12 @@ function _DUMMY_I$main(): int {
     let _t12 = (b: DUMMY_Enum2) as DUMMY_Enum2$_Sub0;
   } else {
   }
-  let _t13 = (b: DUMMY_Enum2) as int;
-  let _t14 = (_t13: int) == 3;
+  let _t13 = (b: DUMMY_Enum2) as i31;
+  let _t14 = (_t13: int) == 1 as i31;
   if (_t14: int) {
   } else {
   }
-  let b = 3 as DUMMY_Enum2;
+  let b = 1 as i31 as DUMMY_Enum2;
   return 0;
 }
 
@@ -1512,8 +1514,8 @@ sources.mains = [_DUMMY_I$main]"#,
       heap,
       r#"
 object type DUMMY_J = [DUMMY_I__int_int]
-variant type DUMMY_I__int_int = [int, int]
-variant type DUMMY_StrOption = [Unboxed(DUMMY_K), int]
+variant type DUMMY_I__int_int = [i31, i31]
+variant type DUMMY_StrOption = [Unboxed(DUMMY_K), i31]
 variant type DUMMY_K = [Boxed(int, int), Boxed(int, int)]
 function _DUMMY_I$creatorJ(): DUMMY_J {
   let v1: DUMMY_I__int_int = [];

crates/samlang-compiler/src/mir_type_deduplication.rs

@@ -229,7 +229,7 @@ pub(super) fn deduplicate(
                 EnumTypeDefinition::Unboxed(t) => {
                   EnumTypeDefinition::Unboxed(rewrite_id_type_name(&state, t))
                 }
-                EnumTypeDefinition::Int => EnumTypeDefinition::Int,
+                EnumTypeDefinition::Int31 => EnumTypeDefinition::Int31,
               })
               .collect(),
           ),
@@ -344,7 +344,7 @@ mod tests {
           mappings: TypeDefinitionMappings::Enum(vec![
             EnumTypeDefinition::Boxed(vec![INT_32_TYPE]),
             EnumTypeDefinition::Unboxed(TypeNameId::STR),
-            EnumTypeDefinition::Int,
+            EnumTypeDefinition::Int31,
           ]),
         },
       ],
@@ -417,7 +417,7 @@ mod tests {
       r#"closure type _A = () -> int
 closure type _C = () -> _C
 object type _C = [int, _Str]
-variant type _E = [Boxed(int), Unboxed(_Str), int]
+variant type _E = [Boxed(int), Unboxed(_Str), i31]
 function __$main(): int {
   let _: int;
   if 1 {

crates/samlang-compiler/src/wasm_lowering.rs

@@ -243,9 +243,8 @@ impl<'a> LoweringManager<'a> {
 
   fn lower_expr(&mut self, e: &lir::Expression) -> wasm::InlineInstruction {
     match e {
-      lir::Expression::Int32Literal(v) | lir::Expression::Int31Literal(v) => {
-        wasm::InlineInstruction::Const(*v)
-      }
+      lir::Expression::Int32Literal(v) => wasm::InlineInstruction::Const(*v),
+      lir::Expression::Int31Literal(v) => wasm::InlineInstruction::Const(*v * 2 + 1),
       lir::Expression::Variable(n, _) => self.get(n),
       lir::Expression::StringName(n) => {
         let index = self.global_variables_to_pointer_mapping.get(n).unwrap();

crates/samlang-optimization/src/unused_name_elimination.rs

@@ -186,7 +186,7 @@ fn analyze_all_used_names(
             EnumTypeDefinition::Unboxed(t) => {
               type_set.insert(*t);
             }
-            EnumTypeDefinition::Int => {}
+            EnumTypeDefinition::Int31 => {}
           }
         }
       }
@@ -308,7 +308,7 @@ mod tests {
         TypeDefinition {
           name: table.create_type_name_for_test(heap.alloc_str_for_test("Baz")),
           mappings: TypeDefinitionMappings::Enum(vec![
-            EnumTypeDefinition::Int,
+            EnumTypeDefinition::Int31,
             EnumTypeDefinition::Unboxed(TypeNameId::STR),
             EnumTypeDefinition::Boxed(vec![INT_32_TYPE]),
           ]),