Merge pull request #992 from candy-lang/more-efficient-if-else

More efficient if/else

compiler/cli/src/debug.rs

@@ -110,7 +110,7 @@ pub struct PathAndExecutionTargetAndTracing {
     #[arg(
         long,
         default_value("off"),
-        default_missing_value("only-potentially-panicking"),
+        default_missing_value("only-for-panic-traces"),
         num_args(0..=1),
         require_equals(true)
     )]

compiler/frontend/src/lir/expression.rs

@@ -48,6 +48,15 @@ pub enum Expression {
         responsible: Id,
     },
 
+    IfElse {
+        condition: Id,
+        then_body_id: BodyId,
+        then_captured: Vec<Id>,
+        else_body_id: BodyId,
+        else_captured: Vec<Id>,
+        responsible: Id,
+    },
+
     Panic {
         reason: Id,
         responsible: Id,
@@ -128,6 +137,23 @@ impl Expression {
                 }
                 *responsible = replacer(*responsible);
             }
+            Self::IfElse {
+                condition,
+                then_body_id: _,
+                then_captured,
+                else_body_id: _,
+                else_captured,
+                responsible,
+            } => {
+                *condition = replacer(*condition);
+                for captured in then_captured {
+                    *captured = replacer(*captured);
+                }
+                for captured in else_captured {
+                    *captured = replacer(*captured);
+                }
+                *responsible = replacer(*responsible);
+            }
             Self::Panic {
                 reason,
                 responsible,
@@ -266,6 +292,40 @@ impl Expression {
                 responsible.build_rich_ir_with_constants(builder, constants, body);
                 builder.push(" is responsible)", None, EnumSet::empty());
             }
+            Self::IfElse {
+                condition,
+                then_body_id,
+                then_captured,
+                else_body_id,
+                else_captured,
+                responsible,
+            } => {
+                builder.push("if ", None, EnumSet::empty());
+                condition.build_rich_ir_with_constants(builder, constants, body);
+                builder.push(" then call ", None, EnumSet::empty());
+                then_body_id.build_rich_ir(builder);
+                if !then_captured.is_empty() {
+                    builder.push(" capturing ", None, EnumSet::empty());
+                    builder.push_children_custom(
+                        then_captured,
+                        |builder, it| it.build_rich_ir_with_constants(builder, constants, body),
+                        ", ",
+                    );
+                }
+                builder.push(" else call ", None, EnumSet::empty());
+                else_body_id.build_rich_ir(builder);
+                if !else_captured.is_empty() {
+                    builder.push(" capturing ", None, EnumSet::empty());
+                    builder.push_children_custom(
+                        else_captured,
+                        |builder, it| it.build_rich_ir_with_constants(builder, constants, body),
+                        ", ",
+                    );
+                }
+                builder.push(" (", None, EnumSet::empty());
+                responsible.build_rich_ir_with_constants(builder, constants, body);
+                builder.push(" is responsible)", None, EnumSet::empty());
+            }
             Self::Panic {
                 reason,
                 responsible,

compiler/frontend/src/mir_to_lir.rs

@@ -1,10 +1,11 @@
 use crate::{
+    builtin_functions::BuiltinFunction,
     error::CompilerError,
-    hir::{self},
+    hir,
     hir_to_mir::ExecutionTarget,
     id::CountableId,
     lir::{self, Lir},
-    mir::{self},
+    mir,
     mir_optimize::OptimizeMir,
     string_to_rcst::ModuleError,
     utils::{HashMapExtension, HashSetExtension},
@@ -38,11 +39,18 @@ fn lir(db: &dyn MirToLir, target: ExecutionTarget, tracing: TracingConfig) -> Li
     Ok((Arc::new(lir), errors))
 }
 
-#[derive(Clone, Debug, Default)]
+#[derive(Debug, Default)]
 struct LoweringContext {
     constants: lir::Constants,
     constant_mapping: FxHashMap<mir::Id, lir::ConstantId>,
     bodies: lir::Bodies,
+    potential_if_else_bodies: FxHashMap<mir::Id, IfElseBody>,
+    delegating_if_else_body_id: Option<lir::BodyId>,
+}
+#[derive(Clone, Debug)]
+struct IfElseBody {
+    body_id: lir::BodyId,
+    captured: Vec<lir::Id>,
 }
 impl LoweringContext {
     fn constant_for(&self, id: mir::Id) -> Option<lir::ConstantId> {
@@ -67,6 +75,41 @@ impl LoweringContext {
         );
         self.bodies.push(body)
     }
+
+    /// Creates a body capturing a function and then executing it.
+    ///
+    /// It's used for ifElse calls for which at least one of the branches'
+    /// functions is not exclusively used for ifElse bodies.
+    #[must_use]
+    fn delegating_if_else_body_id(&mut self) -> lir::BodyId {
+        self.delegating_if_else_body_id.unwrap_or_else(|| {
+            let target_function_id = mir::Id::from_usize(0);
+            let responsible_parameter_id = mir::Id::from_usize(1);
+            let call_id = mir::Id::from_usize(2);
+
+            let mut current_body = CurrentBody::new(
+                FxHashSet::default(),
+                &[target_function_id],
+                &[],
+                responsible_parameter_id,
+            );
+            let function = *current_body.id_mapping.get(&target_function_id).unwrap();
+            let responsible = *current_body
+                .id_mapping
+                .get(&responsible_parameter_id)
+                .unwrap();
+            current_body.push(
+                call_id,
+                lir::Expression::Call {
+                    function,
+                    arguments: vec![],
+                    responsible,
+                },
+            );
+            let body = current_body.finish(&self.constant_mapping);
+            self.bodies.push(body)
+        })
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -86,9 +129,9 @@ impl CurrentBody {
         body: &mir::Body,
     ) -> lir::Body {
         let mut lir_body = Self::new(original_hirs, captured, parameters, responsible_parameter);
-        for (id, expression) in body.iter() {
+        for (index, (id, expression)) in body.iter().enumerate() {
             lir_body.current_constant = None;
-            lir_body.compile_expression(context, id, expression);
+            lir_body.compile_expression(context, id, expression, &body.expressions[index + 1..]);
         }
         lir_body.finish(&context.constant_mapping)
     }
@@ -140,6 +183,7 @@ impl CurrentBody {
         context: &mut LoweringContext,
         id: mir::Id,
         expression: &mir::Expression,
+        remaining_expressions: &[(mir::Id, mir::Expression)],
     ) {
         match expression {
             mir::Expression::Int(int) => self.push_constant(context, id, int.clone()),
@@ -247,10 +291,43 @@ impl CurrentBody {
                     *responsible_parameter,
                     body,
                 );
-                if captured.is_empty() {
+
+                let captured = self.ids_for(context, &captured);
+
+                if parameters.is_empty() {
+                    context.potential_if_else_bodies.force_insert(
+                        id,
+                        IfElseBody {
+                            body_id,
+                            captured: captured.clone(),
+                        },
+                    );
+                }
+
+                if parameters.is_empty()
+                    && remaining_expressions.iter().all(|(_, expression)| {
+                        if let mir::Expression::Call {
+                            function,
+                            arguments,
+                            ..
+                        } = expression
+                            && Self::is_builtin_if_else(context, *function)
+                            && arguments[0] != id
+                        {
+                            // Call to the ifElse builtin function
+                            true
+                        } else {
+                            // The expression doesn't reference the current body
+                            !expression.referenced_ids().contains(&id)
+                        }
+                    })
+                {
+                    // If the function takes no parameters and is only used in
+                    // calls to the ifElse builtin function, we refer to it in
+                    // the special ifElse LIR expression.
+                } else if captured.is_empty() {
                     self.push_constant(context, id, body_id);
                 } else {
-                    let captured = self.ids_for(context, &captured);
                     self.push(id, lir::Expression::CreateFunction { captured, body_id });
                 }
             }
@@ -260,6 +337,44 @@ impl CurrentBody {
                 arguments,
                 responsible,
             } => {
+                if Self::is_builtin_if_else(context, *function) {
+                    let [condition, then_body, else_body] = arguments.as_slice() else {
+                        unreachable!()
+                    };
+
+                    let condition = self.id_for(context, *condition);
+                    let then_body = context
+                        .potential_if_else_bodies
+                        .get(then_body)
+                        .cloned()
+                        .unwrap_or_else(|| IfElseBody {
+                            body_id: context.delegating_if_else_body_id(),
+                            captured: vec![self.id_for(context, *then_body)],
+                        });
+                    let else_body = context
+                        .potential_if_else_bodies
+                        .get(else_body)
+                        .cloned()
+                        .unwrap_or_else(|| IfElseBody {
+                            body_id: context.delegating_if_else_body_id(),
+                            captured: vec![self.id_for(context, *else_body)],
+                        });
+                    let responsible = self.id_for_without_dup(context, *responsible);
+
+                    self.push(
+                        id,
+                        lir::Expression::IfElse {
+                            condition,
+                            then_body_id: then_body.body_id,
+                            then_captured: then_body.captured,
+                            else_body_id: else_body.body_id,
+                            else_captured: else_body.captured,
+                            responsible,
+                        },
+                    );
+                    return;
+                }
+
                 let function = self.id_for(context, *function);
                 let arguments = self.ids_for(context, arguments);
                 let responsible = self.id_for_without_dup(context, *responsible);
@@ -381,6 +496,13 @@ impl CurrentBody {
 
         self.push(id, context.constant_for(id).unwrap())
     }
+    #[must_use]
+    fn is_builtin_if_else(context: &LoweringContext, id: mir::Id) -> bool {
+        matches!(
+            context.constant_for(id),
+            Some(constant_id) if context.constants.get(constant_id) == &lir::Constant::Builtin(BuiltinFunction::IfElse),
+        )
+    }
 
     fn push_constant(
         &mut self,

compiler/vm/src/byte_code.rs

@@ -103,6 +103,17 @@ pub enum Instruction {
     /// called.
     Return,
 
+    /// Conditionally calls either `then_target` or `else_target` depending on
+    /// the `condition`.
+    ///
+    /// a, condition, responsible -> a
+    IfElse {
+        then_target: InstructionPointer,
+        then_captured: Vec<StackOffset>,
+        else_target: InstructionPointer,
+        else_captured: Vec<StackOffset>,
+    },
+
     /// Panics. Because the panic instruction only occurs inside the generated
     /// needs function, the reason is already guaranteed to be a text.
     ///
@@ -185,6 +196,11 @@ impl Instruction {
                 // Only modifies the call stack and the instruction pointer.
                 // Leaves the return value untouched on the stack.
             }
+            Self::IfElse { .. } => {
+                stack.pop(); // responsible
+                stack.pop(); // condition
+                stack.push(result); // return value
+            }
             Self::Panic => {
                 stack.pop(); // responsible
                 stack.pop(); // reason
@@ -392,6 +408,30 @@ impl Instruction {
                 );
             }
             Self::Return => {}
+            Self::IfElse {
+                then_target,
+                then_captured,
+                else_target,
+                else_captured,
+            } => {
+                builder.push(
+                    format!(
+                        " then call {then_target:?} capturing {} else call {else_target:?} capturing {}",
+                        if then_captured.is_empty() {
+                            "nothing".to_string()
+                        } else {
+                            then_captured.iter().join(", ")
+                        },
+                        if else_captured.is_empty() {
+                            "nothing".to_string()
+                        } else {
+                            else_captured.iter().join(", ")
+                        },
+                    ),
+                    None,
+                     EnumSet::empty(),
+                );
+            }
             Self::Panic => {}
             Self::TraceCallStarts { num_args } | Self::TraceTailCall { num_args } => {
                 builder.push(

compiler/vm/src/instructions.rs

@@ -23,6 +23,7 @@ impl MachineState {
         tracer: &mut impl Tracer,
     ) -> InstructionResult {
         if TRACE {
+            trace!("");
             trace!("Running instruction: {instruction:?}");
             trace!("Instruction pointer: {:?}", self.next_instruction.unwrap());
             trace!(
@@ -147,6 +148,39 @@ impl MachineState {
                 self.next_instruction = self.call_stack.pop();
                 InstructionResult::Done
             }
+            Instruction::IfElse {
+                then_target,
+                then_captured,
+                else_target,
+                else_captured,
+            } => {
+                let responsible = self.pop_from_data_stack();
+                let condition = self.pop_from_data_stack();
+                let condition = Tag::try_from(condition)
+                    .unwrap()
+                    .try_into_bool(heap)
+                    .unwrap();
+                let (target, captured) = if condition {
+                    (*then_target, then_captured)
+                } else {
+                    (*else_target, else_captured)
+                };
+
+                if let Some(next_instruction) = self.next_instruction {
+                    self.call_stack.push(next_instruction);
+                }
+
+                // Initially, we need to adjust the offset because we already
+                // popped two values from the data stack. Afterwards, increment
+                // it for each value.
+                for (index, offset) in captured.iter().enumerate() {
+                    let captured = self.get_from_data_stack(*offset - 2 + index);
+                    self.data_stack.push(captured);
+                }
+                self.push_to_data_stack(responsible);
+                self.next_instruction = Some(target);
+                InstructionResult::Done
+            }
             Instruction::Panic => {
                 let responsible_for_panic = self.pop_from_data_stack();
                 let reason = self.pop_from_data_stack();