From 6b1e9c75c8c2a5c3fa972bedf6985c4d2539a2dd Mon Sep 17 00:00:00 2001
From: Brian Carroll <brian.carroll.ireland@gmail.com>
Date: Sun, 3 Jul 2022 23:31:49 +0100
Subject: [PATCH] mono: generate refcounting helper functions for Boxed layout

---
 crates/compiler/gen_wasm/src/wasm32_result.rs |  2 +-
 .../mono/src/code_gen_help/refcount.rs        | 69 ++++++++++++++++++-
 crates/compiler/test_gen/src/gen_refcount.rs  | 43 ++++++++++++
 3 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/crates/compiler/gen_wasm/src/wasm32_result.rs b/crates/compiler/gen_wasm/src/wasm32_result.rs
index e9fd6358e0..1863bff7e6 100644
--- a/crates/compiler/gen_wasm/src/wasm32_result.rs
+++ b/crates/compiler/gen_wasm/src/wasm32_result.rs
@@ -76,7 +76,7 @@ pub fn insert_wrapper_for_layout<'a>(
             bool::insert_wrapper(arena, module, wrapper_name, main_fn_index);
         }
         Layout::Union(UnionLayout::NonRecursive(_)) => stack_data_structure(),
-        Layout::Union(_) => {
+        Layout::Union(_) | Layout::Boxed(_) => {
             i32::insert_wrapper(arena, module, wrapper_name, main_fn_index);
         }
         _ => stack_data_structure(),
diff --git a/crates/compiler/mono/src/code_gen_help/refcount.rs b/crates/compiler/mono/src/code_gen_help/refcount.rs
index e3a6282ce0..24f1b9c9ba 100644
--- a/crates/compiler/mono/src/code_gen_help/refcount.rs
+++ b/crates/compiler/mono/src/code_gen_help/refcount.rs
@@ -129,7 +129,9 @@ pub fn refcount_generic<'a>(
         Layout::RecursivePointer => unreachable!(
             "We should never call a refcounting helper on a RecursivePointer layout directly"
         ),
-        Layout::Boxed(_) => rc_todo(),
+        Layout::Boxed(inner_layout) => {
+            refcount_boxed(root, ident_ids, ctx, &layout, inner_layout, structure)
+        }
     }
 }
 
@@ -343,7 +345,7 @@ pub fn is_rc_implemented_yet(layout: &Layout) -> bool {
             is_rc_implemented_yet(&lambda_set.runtime_representation())
         }
         Layout::RecursivePointer => true,
-        Layout::Boxed(_) => false,
+        Layout::Boxed(_) => true,
     }
 }
 
@@ -1465,3 +1467,66 @@ fn refcount_tag_fields<'a>(
 
     stmt
 }
+
+fn refcount_boxed<'a>(
+    root: &mut CodeGenHelp<'a>,
+    ident_ids: &mut IdentIds,
+    ctx: &mut Context<'a>,
+    layout: &Layout,
+    inner_layout: &'a Layout,
+    outer: Symbol,
+) -> Stmt<'a> {
+    let arena = root.arena;
+
+    //
+    // modify refcount of the inner and outer structures
+    // RC on inner first, to avoid use-after-free for Dec
+    // We're defining statements in reverse, so define outer first
+    //
+
+    let rc_ptr = root.create_symbol(ident_ids, "rc_ptr");
+    let alignment = layout.alignment_bytes(root.target_info);
+    let ret_stmt = rc_return_stmt(root, ident_ids, ctx);
+    let modify_outer = modify_refcount(
+        root,
+        ident_ids,
+        ctx,
+        rc_ptr,
+        alignment,
+        arena.alloc(ret_stmt),
+    );
+
+    let get_rc_and_modify_outer = rc_ptr_from_data_ptr(
+        root,
+        ident_ids,
+        outer,
+        rc_ptr,
+        false,
+        arena.alloc(modify_outer),
+    );
+
+    if inner_layout.is_refcounted() && !ctx.op.is_decref() {
+        let inner = root.create_symbol(ident_ids, "inner");
+        let inner_expr = Expr::ExprUnbox { symbol: outer };
+
+        let mod_inner_unit = root.create_symbol(ident_ids, "mod_inner_unit");
+        let mod_inner_args = refcount_args(root, ctx, inner);
+        let mod_inner_expr = root
+            .call_specialized_op(ident_ids, ctx, *inner_layout, mod_inner_args)
+            .unwrap();
+
+        Stmt::Let(
+            inner,
+            inner_expr,
+            *inner_layout,
+            arena.alloc(Stmt::Let(
+                mod_inner_unit,
+                mod_inner_expr,
+                LAYOUT_UNIT,
+                arena.alloc(get_rc_and_modify_outer),
+            )),
+        )
+    } else {
+        get_rc_and_modify_outer
+    }
+}
diff --git a/crates/compiler/test_gen/src/gen_refcount.rs b/crates/compiler/test_gen/src/gen_refcount.rs
index e707495919..9024d96a44 100644
--- a/crates/compiler/test_gen/src/gen_refcount.rs
+++ b/crates/compiler/test_gen/src/gen_refcount.rs
@@ -431,3 +431,46 @@ fn union_linked_list_long_dec() {
         &[Deallocated; 1_000]
     );
 }
+
+#[test]
+#[cfg(any(feature = "gen-wasm"))]
+fn boxed_str_inc() {
+    assert_refcounts!(
+        indoc!(
+            r#"
+                s = Str.concat "A long enough string " "to be heap-allocated"
+                b = Box.box s
+
+                Tuple b b
+            "#
+        ),
+        (Pointer, Pointer),
+        &[
+            Live(2), // s
+            Live(2), // b
+        ]
+    );
+}
+
+#[test]
+#[cfg(any(feature = "gen-wasm"))]
+fn boxed_str_dec() {
+    assert_refcounts!(
+        indoc!(
+            r#"
+                s = Str.concat "A long enough string " "to be heap-allocated"
+                b = Box.box s
+
+                if False then
+                    ReturnTheBox b
+                else
+                    DeallocateEverything
+            "#
+        ),
+        (i32, i32),
+        &[
+            Deallocated, // s
+            Deallocated, // b
+        ]
+    );
+}