Merge pull request roc-lang#7334 from roc-lang/specialize-expr-unit

Do not discard empty types in `specialize_types`

crates/compiler/specialize_types/src/mono_expr.rs

@@ -84,7 +84,7 @@ impl<'a, 'c, 'd, 'i, 's, 't, P: Push<Problem>> Env<'a, 'c, 'd, 'i, 's, 't, P> {
         }
     }
 
-    pub fn to_mono_expr(&mut self, can_expr: &Expr) -> Option<MonoExpr> {
+    pub fn to_mono_expr(&mut self, can_expr: &Expr) -> MonoExpr {
         let problems = &mut self.problems;
         let mono_types = &mut self.mono_types;
         let mut mono_from_var = |var| {
@@ -103,7 +103,7 @@ impl<'a, 'c, 'd, 'i, 's, 't, P: Push<Problem>> Env<'a, 'c, 'd, 'i, 's, 't, P> {
         macro_rules! compiler_bug {
             ($problem:expr) => {{
                 problems.push($problem);
-                Some(MonoExpr::CompilerBug($problem))
+                MonoExpr::CompilerBug($problem)
             }};
         }
 
@@ -112,56 +112,48 @@ impl<'a, 'c, 'd, 'i, 's, 't, P: Push<Problem>> Env<'a, 'c, 'd, 'i, 's, 't, P> {
                 match self.subs.get_content_without_compacting(*var) {
                     Content::FlexVar(_) => {
                         // Plain decimal number literals like `4.2` can still have an unbound var.
-                        Some(MonoExpr::Number(Number::Dec(*val)))
+                        MonoExpr::Number(Number::Dec(*val))
                     }
-                    _ => match mono_from_var(*var) {
-                        Some(mono_id) => match mono_types.get(mono_id) {
-                            MonoType::Primitive(primitive) => {
-                                Some(to_frac(*primitive, *val, problems))
-                            }
+                    _ => {
+                        let mono_type = mono_from_var(*var);
+                        match mono_types.get(mono_type) {
+                            MonoType::Primitive(primitive) => to_frac(*primitive, *val, problems),
                             other => {
                                 compiler_bug!(Problem::NumSpecializedToWrongType(Some(*other)))
                             }
-                        },
-                        None => {
-                            compiler_bug!(Problem::NumSpecializedToWrongType(None))
                         }
-                    },
+                    }
                 }
             }
             Expr::Num(var, _, int_value, _) | Expr::Int(var, _, _, int_value, _) => {
+                let mono_type = mono_from_var(*var);
+
                 // Number literals and int literals both specify integer numbers, so to_num() can work on both.
-                match mono_from_var(*var) {
-                    Some(mono_id) => match mono_types.get(mono_id) {
-                        MonoType::Primitive(primitive) => {
-                            Some(to_num(*primitive, *int_value, problems))
-                        }
-                        other => compiler_bug!(Problem::NumSpecializedToWrongType(Some(*other))),
-                    },
-                    None => compiler_bug!(Problem::NumSpecializedToWrongType(None)),
+                match mono_types.get(mono_type) {
+                    MonoType::Primitive(primitive) => to_num(*primitive, *int_value, problems),
+                    other => compiler_bug!(Problem::NumSpecializedToWrongType(Some(*other))),
                 }
             }
-            Expr::SingleQuote(var, _, char, _) => match mono_from_var(*var) {
+            Expr::SingleQuote(var, _, char, _) => {
+                let mono_type = mono_from_var(*var);
+
                 // Single-quote characters monomorphize to an integer.
-                // TODO if we store these using the same representation as other ints (e.g. Expr::Int,
+                // TODO [mono2] if we store these using the same representation as other ints (e.g. Expr::Int,
                 // or keeping a separate value but storing an IntValue instead of a char), then
                 // even though we verify them differently, we can combine this branch with Num and Int.
-                Some(mono_id) => match mono_types.get(mono_id) {
-                    MonoType::Primitive(primitive) => {
-                        Some(char_to_int(*primitive, *char, problems))
-                    }
+                match mono_types.get(mono_type) {
+                    MonoType::Primitive(primitive) => char_to_int(*primitive, *char, problems),
                     other => compiler_bug!(Problem::CharSpecializedToWrongType(Some(*other))),
-                },
-                None => compiler_bug!(Problem::CharSpecializedToWrongType(None)),
-            },
-            Expr::Str(contents) => Some(MonoExpr::Str(self.string_interns.get_id(
+                }
+            }
+            Expr::Str(contents) => MonoExpr::Str(self.string_interns.get_id(
                 self.arena,
                 // TODO should be able to remove this alloc_str() once canonical Expr stores an arena-allocated string.
                 self.arena.alloc_str(contents),
-            ))),
+            )),
             Expr::EmptyRecord => {
                 // Empty records are zero-sized and should be discarded.
-                None
+                MonoExpr::Unit
             }
             Expr::Record {
                 record_var: _,
@@ -172,10 +164,10 @@ impl<'a, 'c, 'd, 'i, 's, 't, P: Push<Problem>> Env<'a, 'c, 'd, 'i, 's, 't, P> {
                 // Check for records with 0-1 fields before sorting or reserving a slice of IDs (which might be unnecessary).
                 // We'll check again after discarding zero-sized fields, because we might end up with 0 or 1 fields remaining.
                 if fields.len() <= 1 {
-                    return fields
-                        .into_iter()
-                        .next()
-                        .and_then(|(_, field)| self.to_mono_expr(&field.loc_expr.value));
+                    return match fields.into_iter().next() {
+                        Some((_, field)) => self.to_mono_expr(&field.loc_expr.value),
+                        None => MonoExpr::Unit,
+                    };
                 }
 
                 // Sort the fields alphabetically by name.
@@ -189,21 +181,18 @@ impl<'a, 'c, 'd, 'i, 's, 't, P: Push<Problem>> Env<'a, 'c, 'd, 'i, 's, 't, P> {
                 let mut buf: Vec<(MonoExpr, Region)> =
                     Vec::with_capacity_in(fields.len(), self.arena);
 
-                buf.extend(
-                    // flat_map these so we discard all the fields that monomorphized to None
-                    fields.into_iter().flat_map(|(_name, field)| {
-                        self.to_mono_expr(&field.loc_expr.value)
-                            .map(|mono_expr| (mono_expr, field.loc_expr.region))
-                    }),
-                );
+                buf.extend(fields.into_iter().map(|(_name, field)| {
+                    (
+                        self.to_mono_expr(&field.loc_expr.value),
+                        field.loc_expr.region,
+                    )
+                }));
 
-                // If we ended up with exactly 1 field, return it unwrapped.
-                if buf.len() == 1 {
-                    return buf.pop().map(|(expr, _region)| expr);
-                }
+                let slice = unsafe {
+                    NonEmptySlice::from_slice_unchecked(self.mono_exprs.extend(buf.into_iter()))
+                };
 
-                NonEmptySlice::from_slice(self.mono_exprs.extend(buf.iter().copied()))
-                    .map(MonoExpr::Struct)
+                MonoExpr::Struct(slice)
             }
             // Expr::Call((fn_var, fn_expr, capture_var, ret_var), args, called_via) => {
             //     let opt_ret_type = mono_from_var(*var);

crates/compiler/specialize_types/src/mono_ir.rs

@@ -129,14 +129,18 @@ impl MonoExprs {
         })
     }
 
-    pub fn extend(
-        &mut self,
-        exprs: impl Iterator<Item = (MonoExpr, Region)> + Clone,
-    ) -> Slice<MonoExpr> {
+    pub fn extend(&mut self, exprs: impl Iterator<Item = (MonoExpr, Region)>) -> Slice<MonoExpr> {
         let start = self.exprs.len();
 
-        self.exprs.extend(exprs.clone().map(|(expr, _region)| expr));
-        self.regions.extend(exprs.map(|(_expr, region)| region));
+        let (size_hint, _) = exprs.size_hint();
+
+        self.exprs.reserve(size_hint);
+        self.regions.reserve(size_hint);
+
+        for (expr, region) in exprs {
+            self.exprs.push(expr);
+            self.regions.push(region);
+        }
 
         let len = self.exprs.len() - start;
 
@@ -279,6 +283,8 @@ pub enum MonoExpr {
         recursive: Recursive,
     },
 
+    Unit,
+
     /// A record literal or a tuple literal.
     /// These have already been sorted alphabetically.
     Struct(NonEmptySlice<MonoExpr>),

crates/compiler/specialize_types/src/mono_type.rs

@@ -129,37 +129,21 @@ impl MonoTypes {
 
     pub(crate) fn add_function(
         &mut self,
-        ret: Option<MonoTypeId>,
+        ret: MonoTypeId,
         args: impl IntoIterator<Item = MonoTypeId>,
     ) -> MonoTypeId {
-        let mono_type = match ret {
-            Some(ret) => {
-                let ret_then_args = {
-                    let start = self.ids.len();
-                    self.ids.push(ret);
-                    self.ids.extend(args);
-                    // Safety: we definitely have at least 2 elements in here, even if the iterator is empty.
-                    let length =
-                        unsafe { NonZeroU16::new_unchecked((self.ids.len() - start) as u16) };
-
-                    NonEmptySlice::new(start as u32, length)
-                };
-
-                MonoType::Func { ret_then_args }
-            }
-            None => {
-                let args = {
-                    let start = self.ids.len();
-                    self.ids.extend(args);
-                    let length = (self.ids.len() - start) as u16;
-
-                    Slice::new(start as u32, length)
-                };
-
-                MonoType::VoidFunc { args }
-            }
+        let ret_then_args = {
+            let start = self.ids.len();
+            self.ids.push(ret);
+            self.ids.extend(args);
+            // Safety: we definitely have at least 2 elements in here, even if the iterator is empty.
+            let length = unsafe { NonZeroU16::new_unchecked((self.ids.len() - start) as u16) };
+
+            NonEmptySlice::new(start as u32, length)
         };
 
+        let mono_type = MonoType::Func { ret_then_args };
+
         let index = self.entries.len();
         self.entries.push(mono_type);
         MonoTypeId::new(Index::new(index as u32))
@@ -275,16 +259,6 @@ pub enum MonoType {
     Func {
         ret_then_args: NonEmptySlice<MonoTypeId>,
     },
-
-    /// A function that does not have a return value (e.g. the return value was {}, which
-    /// got eliminated), and which has 0 or more arguments.
-    /// This has to be its own variant because the other function variant uses one slice to
-    /// store the return value followed by the arguments. Without this separate variant,
-    /// the other one couldn't distinguish between a function with a return value and 0 arguments,
-    /// and a function with 1 argument but no return value.
-    VoidFunc {
-        args: Slice<MonoTypeId>,
-    },
     // This last slot is tentatively reserved for Dict, because in the past we've discussed wanting to
     // implement Dict in Zig (for performance) instead of on top of List, like it is as of this writing.
     //

crates/compiler/specialize_types/src/specialize_type.rs

@@ -71,7 +71,7 @@ impl MonoTypeCache {
         problems: &mut impl Push<Problem>,
         debug_info: &mut Option<DebugInfo>,
         var: Variable,
-    ) -> Option<MonoTypeId> {
+    ) -> MonoTypeId {
         let mut env = Env {
             arena,
             cache: self,
@@ -135,23 +135,23 @@ impl<'a, 'c, 'd, 'e, 'f, 'm, 'p, P: Push<Problem>> Env<'a, 'c, 'd, 'e, 'f, 'm, '
         mono_args.extend(
             arg_vars
                 .into_iter()
-                .flat_map(|var_index| self.lower_var(subs, subs[var_index])),
+                .map(|var_index| self.lower_var(subs, subs[var_index])),
         );
 
         // TODO [mono2] populate debuginfo as appropriate
 
         self.mono_types.add_function(func, mono_args)
     }
 
-    fn lower_var(&mut self, subs: &Subs, var: Variable) -> Option<MonoTypeId> {
+    fn lower_var(&mut self, subs: &Subs, var: Variable) -> MonoTypeId {
         let root_var = subs.get_root_key_without_compacting(var);
 
         // TODO: we could replace this cache by having Subs store a Content::Monomorphic(MonoTypeId)
         // and then overwrite it rather than having a separate cache. That memory is already in cache
         // for sure, and the lookups should be faster because they're O(1) but don't require hashing.
         // Kinda creates a cyclic dep though.
         if let Some(mono_id) = self.cache.inner.get(&root_var) {
-            return Some(*mono_id);
+            return *mono_id;
         }
 
         // Convert the Content to a MonoType, often by passing an iterator. None of these iterators introduce allocations.
@@ -291,7 +291,7 @@ impl<'a, 'c, 'd, 'e, 'f, 'm, 'p, P: Push<Problem>> Env<'a, 'c, 'd, 'e, 'f, 'm, '
                     }
                     _ => {
                         // TODO [mono2] record in debuginfo the alias name for whatever we're lowering.
-                        self.lower_var(subs, real)?
+                        self.lower_var(subs, real)
                     }
                 }
             }
@@ -313,10 +313,7 @@ impl<'a, 'c, 'd, 'e, 'f, 'm, 'p, P: Push<Problem>> Env<'a, 'c, 'd, 'e, 'f, 'm, '
         };
 
         // This var is now known to be monomorphic, so we don't repeat this work again later.
-        // (We don't insert entries for Unit values.)
-        self.cache.inner.insert(root_var, mono_id);
-
-        Some(mono_id)
+        mono_id
     }
 }
 

crates/compiler/specialize_types/tests/helpers/mod.rs

@@ -86,9 +86,8 @@ fn specialize_expr<'a>(
     assert_eq!(0, home_decls.expressions.len());
 
     let region = Region::zero();
-    let mono_expr_id = env
-        .to_mono_expr(&main_expr)
-        .map(|mono_expr| mono_exprs.add(mono_expr, region));
+    let mono_expr = env.to_mono_expr(&main_expr);
+    let mono_expr_id = mono_exprs.add(mono_expr, region);
 
     SpecializedExprOut {
         mono_expr_id,
@@ -100,13 +99,13 @@ fn specialize_expr<'a>(
 }
 
 #[track_caller]
-pub fn expect_no_expr(input: impl AsRef<str>) {
+pub fn expect_unit(input: impl AsRef<str>) {
     let arena = Bump::new();
     let mut interns = Interns::new();
     let out = specialize_expr(&arena, input.as_ref(), &mut interns);
-    let actual = out.mono_expr_id.map(|id| out.mono_exprs.get_expr(id));
+    let actual = out.mono_exprs.get_expr(out.mono_expr_id);
 
-    assert_eq!(None, actual, "This input expr should have specialized to being dicarded as zero-sized, but it didn't: {:?}", input.as_ref());
+    assert_eq!(MonoExpr::Unit, *actual, "This input expr should have specialized to being dicarded as zero-sized, but it didn't: {:?}", input.as_ref());
 }
 
 #[track_caller]
@@ -165,6 +164,9 @@ fn dbg_mono_expr_help<'a>(
 
             write!(buf, "])").unwrap();
         }
+        MonoExpr::Unit => {
+            write!(buf, "{{}}").unwrap();
+        }
         // MonoExpr::List { elem_type, elems } => todo!(),
         // MonoExpr::Lookup(symbol, mono_type_id) => todo!(),
         // MonoExpr::ParameterizedLookup {
@@ -270,11 +272,8 @@ pub fn expect_mono_expr_custom<T: PartialEq + core::fmt::Debug>(
     let arena = Bump::new();
     let mut string_interns = Interns::new();
     let out = specialize_expr(&arena, input.as_ref(), &mut string_interns);
-    let mono_expr_id = out
-            .mono_expr_id
-            .expect("This input expr should not have been discarded as zero-sized, but it was discarded: {input:?}");
 
-    let actual_expr = out.mono_exprs.get_expr(mono_expr_id); // Must run first, to populate string interns!
+    let actual_expr = out.mono_exprs.get_expr(out.mono_expr_id); // Must run first, to populate string interns!
     let actual = to_actual(&arena, &out.mono_exprs, &string_interns, actual_expr);
     let expected = to_expected(&arena, &out.mono_exprs, &string_interns);