[ast][ir] Introduce i31 values in MIR and LIR (#1199)

SamChou19815 · May 19, 2024 · e30c687 · e30c687
1 parent fa11029
commit e30c687
Show file tree

Hide file tree

Showing 13 changed files with 66 additions and 24 deletions.
diff --git a/crates/samlang-ast/src/lir.rs b/crates/samlang-ast/src/lir.rs
@@ -100,6 +100,7 @@ pub const ANY_TYPE: Type = Type::Primitive(PrimitiveType::Any);
 #[derive(Debug, Clone, EnumAsInner)]
 pub enum Expression {
   Int32Literal(i32),
+  Int31Literal(i32),
   StringName(PStr),
   Variable(PStr, Type),
   FnName(FunctionName, Type),
@@ -112,7 +113,7 @@ impl Expression {
 
   pub fn ref_countable(&self) -> bool {
     match self {
-      Expression::Int32Literal(_) | Expression::FnName(_, _) => false,
+      Expression::Int32Literal(_) | Expression::Int31Literal(_) | Expression::FnName(_, _) => false,
       Expression::StringName(_) => true,
       Expression::Variable(_, t) => t.as_id().is_some(),
     }
@@ -126,7 +127,9 @@ impl Expression {
     str_table: &HashMap<PStr, usize>,
   ) {
     match self {
-      Expression::Int32Literal(i) => collector.push_str(&i.to_string()),
+      Expression::Int32Literal(i) | Expression::Int31Literal(i) => {
+        collector.push_str(&i.to_string())
+      }
       Expression::Variable(n, _) => collector.push_str(n.as_str(heap)),
       Expression::StringName(n) => {
         collector.push_str("GLOBAL_STRING_");

diff --git a/crates/samlang-ast/src/lir_tests.rs b/crates/samlang-ast/src/lir_tests.rs
@@ -172,7 +172,7 @@ mod tests {
                 heap.alloc_str_for_test("dd"),
                 BinaryOperator::MINUS,
                 ZERO,
-                Expression::int32(-2147483648),
+                Expression::Int31Literal(-2147483648),
               ),
               Statement::binary(heap.alloc_str_for_test("dd"), BinaryOperator::MUL, ZERO, ZERO),
               Statement::binary(heap.alloc_str_for_test("dd"), BinaryOperator::DIV, ZERO, ZERO),

diff --git a/crates/samlang-ast/src/mir.rs b/crates/samlang-ast/src/mir.rs
@@ -353,6 +353,7 @@ pub struct FunctionNameExpression {
 #[derive(Debug, Clone, Copy, EnumAsInner)]
 pub enum Expression {
   Int32Literal(i32),
+  Int31Literal(i32),
   StringName(PStr),
   Variable(VariableName),
 }
@@ -364,8 +365,13 @@ impl Ord for Expression {
         Expression::Int32Literal(i2) => i1.cmp(i2),
         _ => Ordering::Less,
       },
-      Expression::StringName(n1) => match other {
+      Expression::Int31Literal(i1) => match other {
         Expression::Int32Literal(_) => Ordering::Greater,
+        Expression::Int31Literal(i2) => i1.cmp(i2),
+        _ => Ordering::Less,
+      },
+      Expression::StringName(n1) => match other {
+        Expression::Int32Literal(_) | Expression::Int31Literal(_) => Ordering::Greater,
         Expression::StringName(n2) => n1.cmp(n2),
         Expression::Variable(_) => Ordering::Less,
       },
@@ -395,7 +401,7 @@ impl Hash for Expression {
   fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
     core::mem::discriminant(self).hash(state);
     match self {
-      Expression::Int32Literal(i) => i.hash(state),
+      Expression::Int32Literal(i) | Expression::Int31Literal(i) => i.hash(state),
       Expression::StringName(n) => n.hash(state),
       Expression::Variable(v) => v.hash(state),
     }
@@ -413,15 +419,15 @@ impl Expression {
 
   pub fn debug_print(&self, heap: &Heap, table: &SymbolTable) -> String {
     match self {
-      Expression::Int32Literal(i) => i.to_string(),
+      Expression::Int32Literal(i) | Expression::Int31Literal(i) => i.to_string(),
       Expression::StringName(n) => format!("\"{}\"", n.as_str(heap)),
       Expression::Variable(v) => v.debug_print(heap, table),
     }
   }
 
   pub fn convert_to_callee(self) -> Option<Callee> {
     match self {
-      Expression::Int32Literal(_) | Expression::StringName(_) => None,
+      Expression::Int32Literal(_) | Expression::Int31Literal(_) | Expression::StringName(_) => None,
       Expression::Variable(v) => Some(Callee::Variable(v)),
     }
   }

diff --git a/crates/samlang-ast/src/mir_tests.rs b/crates/samlang-ast/src/mir_tests.rs
@@ -15,9 +15,14 @@ mod tests {
     format!("{:?}", INT_32_TYPE.cmp(&INT_32_TYPE));
     assert!(INT_31_TYPE <= INT_31_TYPE);
     format!("{:?}", INT_31_TYPE.cmp(&INT_31_TYPE));
+    assert!(Expression::i32(100) <= Expression::Int31Literal(1));
+    assert!(Expression::Int31Literal(1) >= Expression::i32(100));
+    assert!(Expression::Int31Literal(1) <= Expression::Int31Literal(2));
+    assert!(Expression::Int31Literal(1) <= Expression::StringName(PStr::EMPTY));
     assert!(ZERO.as_int32_literal().is_some());
     Type::Id(table.create_type_name_for_test(PStr::UPPER_A)).as_id();
     Type::Id(table.create_type_name_for_test(PStr::UPPER_A)).is_int32();
+    Type::Id(table.create_type_name_for_test(PStr::UPPER_A)).is_int31();
     assert!(INT_32_TYPE.is_int32());
     assert!(INT_32_TYPE.as_id().is_none());
     assert!(INT_31_TYPE.is_int31());
@@ -82,6 +87,7 @@ mod tests {
     );
     let mut hasher = DefaultHasher::new();
     ZERO.hash(&mut hasher);
+    Expression::Int31Literal(1).hash(&mut hasher);
     Expression::StringName(PStr::LOWER_A).hash(&mut hasher);
     Expression::var_name(PStr::LOWER_A, INT_32_TYPE).hash(&mut hasher);
     Expression::var_name(PStr::LOWER_A, Type::Id(table.create_type_name_for_test(PStr::UPPER_A)))
@@ -244,7 +250,7 @@ mod tests {
           heap.alloc_str_for_test("dd"),
           BinaryOperator::MINUS,
           ZERO,
-          Expression::i32(-2147483648),
+          Expression::Int31Literal(-2147483648),
         ),
         Statement::binary(heap.alloc_str_for_test("dd"), BinaryOperator::MUL, ZERO, ZERO),
         Statement::binary(heap.alloc_str_for_test("dd"), BinaryOperator::DIV, ZERO, ZERO),

diff --git a/crates/samlang-compiler/src/lir_lowering.rs b/crates/samlang-compiler/src/lir_lowering.rs
@@ -28,6 +28,7 @@ fn unknown_member_destructor_type() -> lir::FunctionType {
 fn lower_expression(expr: mir::Expression) -> lir::Expression {
   match expr {
     mir::Expression::Int32Literal(i) => lir::Expression::Int32Literal(i),
+    mir::Expression::Int31Literal(i) => lir::Expression::Int31Literal(i),
     mir::Expression::StringName(n) => lir::Expression::StringName(n),
     mir::Expression::Variable(mir::VariableName { name, type_ }) => {
       lir::Expression::Variable(name, lower_type(type_))
@@ -430,7 +431,7 @@ fn generate_inc_ref_fn() -> lir::Function {
             is_zero,
             hir::BinaryOperator::EQ,
             lir::Expression::Variable(old_ref_count, lir::INT_32_TYPE),
-            lir::ZERO,
+            lir::Expression::Int31Literal(0),
           ),
           lir::Statement::SingleIf {
             condition: lir::Expression::Variable(is_zero, lir::INT_32_TYPE),
@@ -839,7 +840,7 @@ mod tests {
                 heap.alloc_str_for_test("cc"),
                 closure_type.clone(),
               )),
-              arguments: vec![ZERO],
+              arguments: vec![Expression::Int31Literal(0)],
               return_type: INT_32_TYPE,
               return_collector: None,
             },

diff --git a/crates/samlang-compiler/src/lir_unused_name_elimination.rs b/crates/samlang-compiler/src/lir_unused_name_elimination.rs
@@ -19,7 +19,7 @@ fn collect_used_names_from_expression(
   expression: &Expression,
 ) {
   match expression {
-    Expression::Int32Literal(_) => {}
+    Expression::Int32Literal(_) | Expression::Int31Literal(_) => {}
     Expression::Variable(_, t) => collect_for_type_set(t, type_set),
     Expression::StringName(n) => {
       str_name_set.insert(*n);

diff --git a/crates/samlang-compiler/src/mir_constant_param_elimination.rs b/crates/samlang-compiler/src/mir_constant_param_elimination.rs
@@ -11,6 +11,7 @@ enum ParamUsageAnalysisState {
   Unused,
   Referenced,
   Int32Constant(i32),
+  Int31Constant(i32),
   StrConstant(PStr),
   Unoptimizable,
 }
@@ -202,6 +203,7 @@ fn collect_global_usages_stmt(
             param_states[i],
             match arg {
               Expression::Int32Literal(n) => ParamUsageAnalysisState::Int32Constant(*n),
+              Expression::Int31Literal(n) => ParamUsageAnalysisState::Int31Constant(*n),
               Expression::StringName(p) => ParamUsageAnalysisState::StrConstant(*p),
               Expression::Variable(_) => ParamUsageAnalysisState::Unoptimizable,
             },
@@ -236,7 +238,8 @@ fn collect_all_usages(sources: &Sources) -> HashMap<FunctionName, FunctionAnalys
 }
 
 enum VariableRewriteInstruction {
-  IntConstant(i32),
+  Int32(i32),
+  Int31(i32),
   StrConstant(PStr),
 }
 
@@ -247,10 +250,11 @@ struct RewriteState<'a> {
 
 fn rewrite_expr(state: &RewriteState, expr: &mut Expression) {
   match &expr {
-    Expression::Int32Literal(_) | Expression::StringName(_) => {}
+    Expression::Int32Literal(_) | Expression::Int31Literal(_) | Expression::StringName(_) => {}
     Expression::Variable(v) => match state.local_rewrite.get(&v.name) {
       None => {}
-      Some(VariableRewriteInstruction::IntConstant(n)) => *expr = Expression::Int32Literal(*n),
+      Some(VariableRewriteInstruction::Int32(n)) => *expr = Expression::Int32Literal(*n),
+      Some(VariableRewriteInstruction::Int31(n)) => *expr = Expression::Int31Literal(*n),
       Some(VariableRewriteInstruction::StrConstant(s)) => *expr = Expression::StringName(*s),
     },
   }
@@ -373,7 +377,11 @@ pub(super) fn rewrite_sources(mut sources: Sources) -> Sources {
         current_index += 1;
         match state {
           ParamUsageAnalysisState::Int32Constant(i) => {
-            local_rewrite.insert(*name, VariableRewriteInstruction::IntConstant(i));
+            local_rewrite.insert(*name, VariableRewriteInstruction::Int32(i));
+            false
+          }
+          ParamUsageAnalysisState::Int31Constant(i) => {
+            local_rewrite.insert(*name, VariableRewriteInstruction::Int31(i));
             false
           }
           ParamUsageAnalysisState::StrConstant(s) => {
@@ -464,6 +472,7 @@ mod tests {
             PStr::LOWER_C,
             PStr::LOWER_D,
             PStr::LOWER_E,
+            PStr::LOWER_F,
           ],
           type_: FunctionType {
             argument_types: vec![
@@ -472,6 +481,7 @@ mod tests {
               Type::Id(table.create_type_name_for_test(PStr::UPPER_C)),
               Type::Id(table.create_type_name_for_test(PStr::UPPER_D)),
               Type::Id(table.create_type_name_for_test(PStr::UPPER_E)),
+              Type::Id(table.create_type_name_for_test(PStr::UPPER_F)),
             ],
             return_type: Box::new(INT_32_TYPE),
           },
@@ -481,6 +491,12 @@ mod tests {
               operator: UnaryOperator::Not,
               operand: Expression::var_name(PStr::LOWER_A, INT_32_TYPE),
             },
+            Statement::binary(
+              dummy_name,
+              BinaryOperator::PLUS,
+              Expression::var_name(PStr::LOWER_A, INT_32_TYPE),
+              Expression::var_name(PStr::LOWER_C, INT_32_TYPE),
+            ),
             Statement::binary(
               dummy_name,
               BinaryOperator::PLUS,
@@ -518,6 +534,7 @@ mod tests {
                     INT_32_TYPE,
                     INT_32_TYPE,
                     INT_32_TYPE,
+                    INT_32_TYPE,
                   ],
                   return_type: Box::new(INT_32_TYPE),
                 },
@@ -529,10 +546,11 @@ mod tests {
               // e: recursive param used
               arguments: vec![
                 ZERO,
+                Expression::Int31Literal(0),
                 Expression::i32(1),
                 Expression::i32(2),
-                Expression::var_name(PStr::LOWER_D, INT_32_TYPE),
                 Expression::var_name(PStr::LOWER_E, INT_32_TYPE),
+                Expression::var_name(PStr::LOWER_F, INT_32_TYPE),
               ],
               return_type: INT_32_TYPE,
               return_collector: None,
@@ -547,16 +565,18 @@ mod tests {
                     INT_32_TYPE,
                     INT_32_TYPE,
                     INT_32_TYPE,
+                    INT_32_TYPE,
                   ],
                   return_type: Box::new(INT_32_TYPE),
                 },
               }),
               arguments: vec![
                 ZERO,
+                Expression::Int31Literal(0),
                 Expression::i32(3),
                 Expression::i32(3),
-                Expression::var_name(PStr::LOWER_D, INT_32_TYPE),
                 Expression::var_name(PStr::LOWER_E, INT_32_TYPE),
+                Expression::var_name(PStr::LOWER_F, INT_32_TYPE),
               ],
               return_type: INT_32_TYPE,
               return_collector: None,
@@ -636,9 +656,10 @@ function __$str_const(): int {
   return 0;
 }
 
-function __$func_with_consts(b: _B, e: _E): int {
+function __$func_with_consts(c: _C, e: _E): int {
   let _ = !0;
-  let _ = 0 + (b: int);
+  let _ = 0 + (c: int);
+  let _ = 0 + 0;
   let _: int = 0[0];
   let _: __ = Closure { fun: (__$otherwise_optimizable: () -> int), context: 0 };
   (_: int)(0);

diff --git a/crates/samlang-compiler/src/mir_tail_recursion_rewrite.rs b/crates/samlang-compiler/src/mir_tail_recursion_rewrite.rs
@@ -158,7 +158,7 @@ fn optimize_function_by_tailrec_rewrite_aux(
   function: Function,
 ) -> Result<Function, Function> {
   let expected_return_collector = match &function.return_value {
-    Expression::Int32Literal(_) => None,
+    Expression::Int32Literal(_) | Expression::Int31Literal(_) => None,
     Expression::Variable(v) => Some(v.name),
     Expression::StringName(_) => return Err(function),
   };

diff --git a/crates/samlang-compiler/src/mir_type_deduplication.rs b/crates/samlang-compiler/src/mir_type_deduplication.rs
@@ -47,7 +47,7 @@ fn rewrite_fn_name(
 
 fn rewrite_expr(state: &State, expr: Expression) -> Expression {
   match expr {
-    Expression::Int32Literal(_) | Expression::StringName(_) => expr,
+    Expression::Int32Literal(_) | Expression::Int31Literal(_) | Expression::StringName(_) => expr,
     Expression::Variable(n) => Expression::Variable(rewrite_var_name(state, n)),
   }
 }

diff --git a/crates/samlang-compiler/src/wasm_lowering.rs b/crates/samlang-compiler/src/wasm_lowering.rs
@@ -243,7 +243,9 @@ impl<'a> LoweringManager<'a> {
 
   fn lower_expr(&mut self, e: &lir::Expression) -> wasm::InlineInstruction {
     match e {
-      lir::Expression::Int32Literal(v) => wasm::InlineInstruction::Const(*v),
+      lir::Expression::Int32Literal(v) | lir::Expression::Int31Literal(v) => {
+        wasm::InlineInstruction::Const(*v)
+      }
       lir::Expression::Variable(n, _) => self.get(n),
       lir::Expression::StringName(n) => {
         let index = self.global_variables_to_pointer_mapping.get(n).unwrap();

diff --git a/crates/samlang-optimization/src/conditional_constant_propagation.rs b/crates/samlang-optimization/src/conditional_constant_propagation.rs
@@ -114,7 +114,7 @@ fn optimize_variable_name(
 
 fn optimize_expr(value_cx: &mut LocalValueContextForOptimization, e: &Expression) -> Expression {
   match e {
-    Expression::Int32Literal(_) | Expression::StringName(_) => *e,
+    Expression::Int32Literal(_) | Expression::Int31Literal(_) | Expression::StringName(_) => *e,
     Expression::Variable(v) => optimize_variable_name(value_cx, v),
   }
 }

diff --git a/crates/samlang-optimization/src/loop_induction_analysis.rs b/crates/samlang-optimization/src/loop_induction_analysis.rs
@@ -323,6 +323,9 @@ fn get_loop_invariant_expression_opt(
 ) -> Option<PotentialLoopInvariantExpression> {
   match expression {
     Expression::Int32Literal(i) => Some(PotentialLoopInvariantExpression::Int(*i)),
+    // i31 currently represents the optimized form of variants.
+    // We should just treat it as opaque, and shouldn't do anything funny with it.
+    Expression::Int31Literal(_) |
     // We are doing algebraic operations here. Name is hopeless.
     Expression::StringName(_) => None,
     Expression::Variable(v) => {

diff --git a/crates/samlang-optimization/src/unused_name_elimination.rs b/crates/samlang-optimization/src/unused_name_elimination.rs
@@ -25,7 +25,7 @@ fn collect_used_names_from_expression(
   expression: &Expression,
 ) {
   match expression {
-    Expression::Int32Literal(_) => {}
+    Expression::Int32Literal(_) | Expression::Int31Literal(_) => {}
     Expression::Variable(v) => {
       collect_for_type_set(&v.type_, type_set);
     }