From a102a7497e9ade1cf950f6450b8bfee5e50bcd8b Mon Sep 17 00:00:00 2001
From: Ayaz Hafiz <ayaz@rwx.com>
Date: Sat, 9 Apr 2022 16:21:41 -0400
Subject: [PATCH] Support box/unbox for types that are typically only on the
 stack

Closes #2786
---
 compiler/mono/src/layout.rs             | 14 +++++---
 compiler/test_gen/src/gen_primitives.rs | 44 +++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/compiler/mono/src/layout.rs b/compiler/mono/src/layout.rs
index 325aec6b0e..07cf05ea5e 100644
--- a/compiler/mono/src/layout.rs
+++ b/compiler/mono/src/layout.rs
@@ -468,7 +468,7 @@ impl<'a> UnionLayout<'a> {
 
     pub fn allocation_alignment_bytes(&self, target_info: TargetInfo) -> u32 {
         let allocation = match self {
-            UnionLayout::NonRecursive(_) => unreachable!("not heap-allocated"),
+            UnionLayout::NonRecursive(tags) => Self::tags_alignment_bytes(tags, target_info),
             UnionLayout::Recursive(tags) => Self::tags_alignment_bytes(tags, target_info),
             UnionLayout::NonNullableUnwrapped(field_layouts) => {
                 Layout::struct_no_name_order(field_layouts).alignment_bytes(target_info)
@@ -1150,9 +1150,11 @@ impl<'a> Layout<'a> {
     }
 
     pub fn allocation_alignment_bytes(&self, target_info: TargetInfo) -> u32 {
+        let ptr_width = target_info.ptr_width() as u32;
+
         match self {
             Layout::Builtin(builtin) => builtin.allocation_alignment_bytes(target_info),
-            Layout::Struct { .. } => unreachable!("not heap-allocated"),
+            Layout::Struct { .. } => self.alignment_bytes(target_info).max(ptr_width),
             Layout::Union(union_layout) => union_layout.allocation_alignment_bytes(target_info),
             Layout::LambdaSet(lambda_set) => lambda_set
                 .runtime_representation()
@@ -1545,9 +1547,6 @@ impl<'a> Builtin<'a> {
         let ptr_width = target_info.ptr_width() as u32;
 
         let allocation = match self {
-            Builtin::Int(_) | Builtin::Float(_) | Builtin::Bool | Builtin::Decimal => {
-                unreachable!("not heap-allocated")
-            }
             Builtin::Str => ptr_width,
             Builtin::Dict(k, v) => k
                 .alignment_bytes(target_info)
@@ -1555,6 +1554,11 @@ impl<'a> Builtin<'a> {
                 .max(ptr_width),
             Builtin::Set(k) => k.alignment_bytes(target_info).max(ptr_width),
             Builtin::List(e) => e.alignment_bytes(target_info).max(ptr_width),
+            // The following are usually not heap-allocated, but they might be when inside a Box.
+            Builtin::Int(int_width) => int_width.alignment_bytes(target_info).max(ptr_width),
+            Builtin::Float(float_width) => float_width.alignment_bytes(target_info).max(ptr_width),
+            Builtin::Bool => (core::mem::align_of::<bool>() as u32).max(ptr_width),
+            Builtin::Decimal => IntWidth::I128.alignment_bytes(target_info).max(ptr_width),
         };
 
         allocation.max(ptr_width)
diff --git a/compiler/test_gen/src/gen_primitives.rs b/compiler/test_gen/src/gen_primitives.rs
index be6ef612a5..4c4f020755 100644
--- a/compiler/test_gen/src/gen_primitives.rs
+++ b/compiler/test_gen/src/gen_primitives.rs
@@ -3275,3 +3275,47 @@ fn box_and_unbox_string() {
         RocStr
     )
 }
+
+#[test]
+#[cfg(any(feature = "gen-llvm"))]
+fn box_and_unbox_num() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            Box.unbox (Box.box (123u8))
+            "#
+        ),
+        123,
+        u8
+    )
+}
+
+#[test]
+#[cfg(any(feature = "gen-llvm"))]
+fn box_and_unbox_record() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            Box.unbox (Box.box { a: 15u8, b: 27u8 })
+            "#
+        ),
+        (15, 27),
+        (u8, u8)
+    )
+}
+
+#[test]
+#[cfg(any(feature = "gen-llvm"))]
+fn box_and_unbox_tag_union() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            v : [ A U8, B U8 ] # usually stack allocated
+            v = B 27u8
+            Box.unbox (Box.box v)
+            "#
+        ),
+        (27, 1),
+        (u8, u8)
+    )
+}