feedback

crates/compiler/mono/src/drop_specialization.rs

@@ -13,7 +13,7 @@ use std::iter::Iterator;
 use bumpalo::collections::vec::Vec;
 use bumpalo::collections::CollectIn;
 
-use roc_module::low_level::{LowLevel, LowLevelWrapperType};
+use roc_module::low_level::LowLevel;
 use roc_module::symbol::{IdentIds, ModuleId, Symbol};
 use roc_target::TargetInfo;
 
@@ -92,22 +92,17 @@ fn specialize_drops_stmt<'a, 'i>(
                     call_type,
                     arguments,
                 }) => {
-                    match call_type {
+                    match call_type.clone().replace_lowlevel_wrapper() {
-                        CallType::ByName { name, .. }
-                            if (LowLevelWrapperType::from_symbol(name.name())
-                                == LowLevelWrapperType::CanBeReplacedBy(
-                                    LowLevel::ListGetUnsafe,
-                                )) =>
-                        {
-                            environment.add_list_child(arguments[0], *binding, &arguments[1]);
-
-                            alloc_let_with_continuation!(environment)
-                        }
                         CallType::LowLevel {
                             op: LowLevel::ListGetUnsafe,
                             ..
                         } => {
-                            environment.add_list_child(arguments[0], *binding, &arguments[1]);
+                            let [structure, index] = match arguments {
+                                [structure, index] => [structure, index],
+                                _ => unreachable!("List get should have two arguments"),
+                            };
+
+                            environment.add_list_child(*structure, *binding, index);
 
                             alloc_let_with_continuation!(environment)
                         }
@@ -668,14 +663,12 @@ fn specialize_union<'a, 'i>(
                         continuation,
                     );
 
+                    type RCFun<'a> =
+                        Option<fn(arena: &'a Bump, Symbol, &'a Stmt<'a>) -> &'a Stmt<'a>>;
                     let refcount_fields = |layout_interner: &mut STLayoutInterner<'a>,
                                            ident_ids: &mut IdentIds,
-                                           rc_popped: Option<
+                                           rc_popped: RCFun<'a>,
-                        fn(arena: &'a Bump, Symbol, &'a Stmt<'a>) -> &'a Stmt<'a>,
+                                           rc_unpopped: RCFun<'a>,
-                    >,
-                                           rc_unpopped: Option<
-                        fn(arena: &'a Bump, Symbol, &'a Stmt<'a>) -> &'a Stmt<'a>,
-                    >,
                                            continuation: &'a Stmt<'a>|
                      -> &'a Stmt<'a> {
                         let mut new_continuation = continuation;
@@ -865,10 +858,9 @@ fn specialize_list<'a, 'i>(
     ) {
         (true, Some(length)) => {
             match environment.list_children.get(symbol) {
+                // Only specialize lists if all children are known.
+                // Otherwise we might have to insert an unbouned number of decrements.
                 Some(children) if children.len() as u64 == length => {
-                    // Only specialize lists if all children are known.
-                    // Otherwise we might have to insert an unbouned number of decrements.
-
                     // TODO perhaps this allocation can be avoided.
                     let children_clone = children.clone();
 
@@ -877,10 +869,7 @@ fn specialize_list<'a, 'i>(
                     let mut index_symbols = MutMap::default();
 
                     for index in 0..length {
-                        for (child, i) in children_clone
+                        for (child, i) in children_clone.iter().filter(|(_child, i)| *i == index) {
-                            .iter()
-                            .filter(|(_child, i)| *i == index as u64)
-                        {
                             debug_assert!(length > *i);
 
                             let removed = incremented_children.remove(child);
@@ -900,49 +889,25 @@ fn specialize_list<'a, 'i>(
                         continuation,
                     );
 
-                    |rc_popped: Option<
+                    let mut newer_continuation = arena.alloc(Stmt::Refcounting(
-                        fn(arena: &'a Bump, Symbol, &'a Stmt<'a>) -> &'a Stmt<'a>,
+                        ModifyRc::DecRef(*symbol),
-                    >,
+                        new_continuation,
-                     rc_unpopped: Option<
+                    ));
-                        fn(arena: &'a Bump, Symbol, &'a Stmt<'a>) -> &'a Stmt<'a>,
-                    >,
-                     continuation: &'a Stmt<'a>|
-                     -> &'a Stmt<'a> {
-                        let mut new_continuation = continuation;
 
-                        // Reversed to ensure that the generated code decrements the items in the correct order.
+                    // Reversed to ensure that the generated code decrements the items in the correct order.
-                        for i in (0..length).rev() {
+                    for i in (0..length).rev() {
-                            let (s, popped) = index_symbols.get(&i).unwrap();
+                        let (s, popped) = index_symbols.get(&i).unwrap();
-                            new_continuation = {
+
-                                if *popped {
+                        if !*popped {
-                                    // This symbol was popped, so we can skip the decrement.
+                            // Decrement the children that were not incremented before. And thus don't cancel out.
-                                    match rc_popped {
+                            newer_continuation = arena
-                                        Some(rc) => rc(arena, *s, new_continuation),
+                                .alloc(Stmt::Refcounting(ModifyRc::Dec(*s), newer_continuation));
-                                        None => new_continuation,
-                                    }
-                                } else {
-                                    // This symbol was indexed but not decremented, so we will decrement it.
-                                    match rc_unpopped {
-                                        Some(rc) => rc(arena, *s, new_continuation),
-                                        None => new_continuation,
-                                    }
-                                }
-                            };
                         }
 
-                        new_continuation
-                    }(
                         // Do nothing for the children that were incremented before, as the decrement will cancel out.
-                        None,
+                    }
-                        // Decrement the children that were not incremented before. And thus don't cancel out.
+
-                        Some(|arena, symbol, continuation| {
+                    newer_continuation
-                            arena.alloc(Stmt::Refcounting(ModifyRc::Dec(symbol), continuation))
-                        }),
-                        arena.alloc(Stmt::Refcounting(
-                            ModifyRc::DecRef(*symbol),
-                            new_continuation,
-                        )),
-                    )
                 }
                 _ => keep_original_decrement!(),
             }
@@ -1020,16 +985,20 @@ fn get_union_tag_layout(union_layout: UnionLayout<'_>, tag: Option<Tag>) -> Unio
 Branch on the uniqueness of a symbol.
 Using a joinpoint with the continuation as the body.
 */
-fn branch_uniqueness<'a, 'i>(
+fn branch_uniqueness<'a, 'i, F1, F2>(
     arena: &'a Bump,
     ident_ids: &'i mut IdentIds,
     layout_interner: &'i mut STLayoutInterner<'a>,
     environment: &DropSpecializationEnvironment<'a>,
     symbol: Symbol,
-    unique: impl FnOnce(&mut STLayoutInterner<'a>, &mut IdentIds, &'a Stmt<'a>) -> &'a Stmt<'a>,
+    unique: F1,
-    not_unique: impl FnOnce(&mut STLayoutInterner<'a>, &mut IdentIds, &'a Stmt<'a>) -> &'a Stmt<'a>,
+    not_unique: F2,
     continutation: &'a Stmt<'a>,
-) -> &'a Stmt<'a> {
+) -> &'a Stmt<'a>
+where
+    F1: FnOnce(&mut STLayoutInterner<'a>, &mut IdentIds, &'a Stmt<'a>) -> &'a Stmt<'a>,
+    F2: FnOnce(&mut STLayoutInterner<'a>, &mut IdentIds, &'a Stmt<'a>) -> &'a Stmt<'a>,
+{
     match continutation {
         // The continuation is a single stmt. So we can insert it inline and skip creating a joinpoint.
         Stmt::Ret(_) | Stmt::Jump(_, _) => {
@@ -1239,19 +1208,19 @@ impl<'a> DropSpecializationEnvironment<'a> {
         let mut res = Vec::new_in(self.arena);
 
         if let Some(children) = self.struct_children.get(parent) {
-            children.iter().for_each(|(child, _)| res.push(*child));
+            res.extend(children.iter().map(|(child, _)| child));
         }
 
         if let Some(children) = self.union_children.get(parent) {
-            children.iter().for_each(|(child, _, _)| res.push(*child));
+            res.extend(children.iter().map(|(child, _, _)| child));
         }
 
         if let Some(children) = self.box_children.get(parent) {
-            children.iter().for_each(|child| res.push(*child));
+            res.extend(children.iter());
         }
 
         if let Some(children) = self.list_children.get(parent) {
-            children.iter().for_each(|(child, _)| res.push(*child));
+            res.extend(children.iter().map(|(child, _)| child));
         }
 
         res

crates/compiler/mono/src/inc_dec.rs

@@ -888,35 +888,6 @@ fn insert_refcount_operations_binding<'a>(
         };
     }
 
-    macro_rules! refcount_listget {
-        ($arguments:expr) => {{
-            // TODO this index can be used to store the children of the list symbol.
-            // On the decrement of the list (if the list size is known and) the function can be specialized.
-            let [structure, _index] = match $arguments {
-                [structure, index] => Some([*structure, *index]),
-                _ => None,
-            }
-            .unwrap();
-
-            // All structures are alive at this point and don't have to be copied in order to take an index out/get tag id/copy values to the stack.
-            // But we do want to make sure to decrement this item if it is the last reference.
-            let new_stmt = dec_borrowed!([structure], stmt);
-
-            // Add an increment operation for the binding if it is reference counted and if the expression creates a new reference to a value.
-            let newer_stmt = if matches!(
-                environment.get_symbol_rc_type(binding),
-                VarRcType::ReferenceCounted
-            ) {
-                insert_inc_stmt(arena, *binding, 1, new_stmt)
-            } else {
-                // If the symbol is not reference counted, we don't need to increment it.
-                new_stmt
-            };
-
-            new_let!(newer_stmt)
-        }};
-    }
-
     match expr {
         Expr::Literal(_) | Expr::NullPointer | Expr::EmptyArray | Expr::RuntimeErrorFunction(_) => {
             // Literals, empty arrays, and runtime errors are not (and have nothing) reference counted.
@@ -933,24 +904,7 @@ fn insert_refcount_operations_binding<'a>(
 
             inc_owned!([*symbol], new_let)
         }
-        Expr::Call(Call {
+
-            arguments,
-            call_type:
-                CallType::LowLevel {
-                    op: LowLevel::ListGetUnsafe,
-                    ..
-                },
-        }) => {
-            refcount_listget!(arguments)
-        }
-        Expr::Call(Call {
-            arguments,
-            call_type: CallType::ByName { name, .. },
-        }) if (LowLevelWrapperType::from_symbol(name.name())
-            == LowLevelWrapperType::CanBeReplacedBy(LowLevel::ListGetUnsafe)) =>
-        {
-            refcount_listget!(arguments)
-        }
         Expr::GetTagId { structure, .. }
         | Expr::StructAtIndex { structure, .. }
         | Expr::UnionAtIndex { structure, .. }
@@ -1000,46 +954,13 @@ fn insert_refcount_operations_binding<'a>(
             arguments,
             call_type,
         }) => {
-            macro_rules! rc_lowlevel {
+            match call_type.clone().replace_lowlevel_wrapper() {
-                ($operator:expr) => {{
-                    let borrow_signature = lowlevel_borrow_signature(arena, $operator);
-                    let arguments_with_borrow_signature = arguments
-                        .iter()
-                        .copied()
-                        .zip(borrow_signature.iter().copied());
-                    let owned_arguments = arguments_with_borrow_signature
-                        .clone()
-                        .filter_map(|(symbol, ownership)| ownership.is_owned().then_some(symbol));
-                    let borrowed_arguments =
-                        arguments_with_borrow_signature.filter_map(|(symbol, ownership)| {
-                            ownership.is_borrowed().then_some(symbol)
-                        });
-
-                    let new_stmt = dec_borrowed!(borrowed_arguments, stmt);
-
-                    let new_let = new_let!(new_stmt);
-
-                    inc_owned!(owned_arguments, new_let)
-                }};
-            }
-            match call_type {
                 // A by name call refers to a normal function call.
                 // Normal functions take all their parameters as owned, so we can mark them all as such.
-                CallType::ByName { name, .. } => {
+                CallType::ByName { .. } => {
-                    // Lowlevels are wrapped in another function in order to add type signatures which help with inference.
+                    let new_let = new_let!(stmt);
-                    // But the reference counting algorithm inserts reference counting operations in the wrapper function.
-                    // But in a later stage, calls to the wrapper function were replaced by calls to the lowlevel function.
-                    // Effectively removing the inserted reference counting operations.
-                    // Thus to prevent that, we inline the operations here already.
-                    if let LowLevelWrapperType::CanBeReplacedBy(operator) =
-                        LowLevelWrapperType::from_symbol(name.name())
-                    {
-                        rc_lowlevel!(operator)
-                    } else {
-                        let new_let = new_let!(stmt);
 
-                        inc_owned!(arguments.iter().copied(), new_let)
+                    inc_owned!(arguments.iter().copied(), new_let)
-                    }
                 }
                 CallType::Foreign { .. } => {
                     // Foreign functions should be responsible for their own memory management.
@@ -1052,9 +973,44 @@ fn insert_refcount_operations_binding<'a>(
                 CallType::LowLevel {
                     op: operator,
                     update_mode: _,
-                } => {
+                } => match operator {
-                    rc_lowlevel!(*operator)
+                    // List get unsafe is a special case, because it returns a reference to the list element.
-                }
+                    // This means that we have to increment the reference count of this element.
+                    LowLevel::ListGetUnsafe => {
+                        let structure = match arguments {
+                            [structure, _index] => *structure,
+                            _ => unreachable!("List get should have two arguments"),
+                        };
+                        let new_stmt = dec_borrowed!([structure], stmt);
+                        let newer_stmt = if matches!(
+                            environment.get_symbol_rc_type(binding),
+                            VarRcType::ReferenceCounted
+                        ) {
+                            insert_inc_stmt(arena, *binding, 1, new_stmt)
+                        } else {
+                            new_stmt
+                        };
+                        new_let!(newer_stmt)
+                    }
+                    // Otherwise, perform regular reference counting using the lowlevel borrow signature.
+                    _ => {
+                        let borrow_signature = lowlevel_borrow_signature(arena, operator);
+                        let arguments_with_borrow_signature = arguments
+                            .iter()
+                            .copied()
+                            .zip(borrow_signature.iter().copied());
+                        let owned_arguments = arguments_with_borrow_signature.clone().filter_map(
+                            |(symbol, ownership)| ownership.is_owned().then_some(symbol),
+                        );
+                        let borrowed_arguments =
+                            arguments_with_borrow_signature.filter_map(|(symbol, ownership)| {
+                                ownership.is_borrowed().then_some(symbol)
+                            });
+                        let new_stmt = dec_borrowed!(borrowed_arguments, stmt);
+                        let new_let = new_let!(new_stmt);
+                        inc_owned!(owned_arguments, new_let)
+                    }
+                },
                 CallType::HigherOrder(HigherOrderLowLevel {
                     op: operator,
 

crates/compiler/mono/src/ir.rs

@@ -25,7 +25,7 @@ use roc_error_macros::{internal_error, todo_abilities};
 use roc_late_solve::storage::{ExternalModuleStorage, ExternalModuleStorageSnapshot};
 use roc_late_solve::{resolve_ability_specialization, AbilitiesView, Resolved, UnificationFailed};
 use roc_module::ident::{ForeignSymbol, Lowercase, TagName};
-use roc_module::low_level::LowLevel;
+use roc_module::low_level::{LowLevel, LowLevelWrapperType};
 use roc_module::symbol::{IdentIds, ModuleId, Symbol};
 use roc_problem::can::{RuntimeError, ShadowKind};
 use roc_region::all::{Loc, Region};
@@ -1790,6 +1790,24 @@ pub enum CallType<'a> {
     HigherOrder(&'a HigherOrderLowLevel<'a>),
 }
 
+impl<'a> CallType<'a> {
+    /**
+    Replace calls to wrappers of lowlevel functions with the lowlevel function itself
+    */
+    pub fn replace_lowlevel_wrapper(self) -> Self {
+        match self {
+            CallType::ByName { name, .. } => match LowLevelWrapperType::from_symbol(name.name()) {
+                LowLevelWrapperType::CanBeReplacedBy(lowlevel) => CallType::LowLevel {
+                    op: lowlevel,
+                    update_mode: UpdateModeId::BACKEND_DUMMY,
+                },
+                LowLevelWrapperType::NotALowLevelWrapper => self,
+            },
+            _ => self,
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct PassedFunction<'a> {
     /// name of the top-level function that is passed as an argument