32 bit count in reference counting (#1050)

* Add sync bit

* Change count type

* Revert "Add sync bit"

This reverts commit 8342715.

* Fix record type

* Fix counter pointer

* Rename count type

* Refactor

* Fix build

* Fix header

* Update snapshots

* Fix FFI

* Bump version

* Update lock file

* Update Cargo.lock

* Fix versions

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

cmd/test/Cargo.toml

@@ -5,4 +5,4 @@ publish = false
 edition = "2021"
 
 [dependencies]
-ffi = { package = "pen-ffi", version = "0.3" }
+ffi = { package = "pen-ffi", version = "0.4" }

lib/ffi/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "pen-ffi"
 description = "FFI library for Pen programming language"
-version = "0.3.70"
+version = "0.4.0"
 publish = true
 edition = "2021"
 license = "MIT"

lib/ffi/src/arc/arc_block.rs

@@ -2,20 +2,28 @@ use alloc::alloc::{alloc, dealloc};
 use core::{
     alloc::Layout,
     ptr::{drop_in_place, null},
-    sync::atomic::{fence, AtomicUsize, Ordering},
+    sync::atomic::{fence, AtomicU32, Ordering},
 };
 
-const INITIAL_COUNT: usize = 0;
+const INITIAL_COUNT: u32 = 0;
+const EMPTY_TAG: u32 = 0;
 
 #[derive(Debug)]
 #[repr(C)]
 pub struct ArcBlock {
     pointer: *const u8,
 }
 
+#[derive(Debug)]
+#[repr(C)]
+struct ArcHeader {
+    count: AtomicU32,
+    tag: u32,
+}
+
 #[repr(C)]
 struct ArcInner {
-    count: AtomicUsize,
+    header: ArcHeader,
     payload: (),
 }
 
@@ -26,7 +34,10 @@ impl ArcBlock {
         } else {
             let pointer = unsafe { &mut *(alloc(Self::inner_layout(layout)) as *mut ArcInner) };
 
-            pointer.count = AtomicUsize::new(INITIAL_COUNT);
+            pointer.header = ArcHeader {
+                count: AtomicU32::new(INITIAL_COUNT),
+                tag: EMPTY_TAG,
+            };
 
             Self {
                 pointer: &pointer.payload as *const () as *const u8,
@@ -43,7 +54,7 @@ impl ArcBlock {
     }
 
     pub fn get_mut(&mut self) -> Option<*mut u8> {
-        if !self.is_static() && self.inner().count.load(Ordering::Acquire) == INITIAL_COUNT {
+        if !self.is_static() && self.inner().header.count.load(Ordering::Acquire) == INITIAL_COUNT {
             Some(self.ptr_mut())
         } else {
             None
@@ -69,7 +80,7 @@ impl ArcBlock {
     }
 
     fn inner_layout(layout: Layout) -> Layout {
-        Layout::new::<AtomicUsize>()
+        Layout::new::<ArcHeader>()
             .extend(layout)
             .unwrap()
             .0
@@ -78,7 +89,7 @@ impl ArcBlock {
 
     pub fn clone(&self) -> Self {
         if !self.is_static() {
-            self.inner().count.fetch_add(1, Ordering::Relaxed);
+            self.inner().header.count.fetch_add(1, Ordering::Relaxed);
         }
 
         Self {
@@ -91,7 +102,7 @@ impl ArcBlock {
             return;
         }
 
-        if self.inner().count.fetch_sub(1, Ordering::Release) == INITIAL_COUNT {
+        if self.inner().header.count.fetch_sub(1, Ordering::Release) == INITIAL_COUNT {
             fence(Ordering::Acquire);
 
             unsafe {

lib/mir-fmm/src/reference_count.rs

@@ -1,3 +1,4 @@
+mod count;
 mod expression;
 pub mod function;
 pub mod heap;

lib/mir-fmm/src/reference_count/count.rs

@@ -0,0 +1,13 @@
+const INITIAL_COUNT: isize = 0;
+
+pub const fn compile(count: isize) -> fmm::ir::Primitive {
+    fmm::ir::Primitive::Integer32(count as u32)
+}
+
+pub const fn compile_initial() -> fmm::ir::Primitive {
+    compile(INITIAL_COUNT)
+}
+
+pub const fn compile_type() -> fmm::types::Primitive {
+    fmm::types::Primitive::Integer32
+}

lib/mir-fmm/src/reference_count/heap.rs

@@ -1,7 +1,11 @@
-use super::super::error::CompileError;
+use super::{super::error::CompileError, count};
+use once_cell::sync::Lazy;
 
-pub(super) const COUNT_TYPE: fmm::types::Primitive = fmm::types::Primitive::PointerInteger;
-pub(super) const INITIAL_COUNT: usize = 0;
+const TAG_TYPE: fmm::types::Primitive = fmm::types::Primitive::Integer32;
+const EMPTY_TAG: fmm::ir::Primitive = fmm::ir::Primitive::Integer32(0);
+
+static HEADER_TYPE: Lazy<fmm::types::Record> =
+    Lazy::new(|| fmm::types::Record::new(vec![count::compile_type().into(), TAG_TYPE.into()]));
 
 pub fn allocate(
     builder: &fmm::build::InstructionBuilder,
@@ -10,17 +14,17 @@ pub fn allocate(
     let type_ = type_.into();
     let pointer = fmm::build::bit_cast(
         fmm::types::Pointer::new(fmm::types::Record::new(vec![
-            COUNT_TYPE.into(),
+            HEADER_TYPE.clone().into(),
             type_.clone(),
         ])),
         builder.allocate_heap(fmm::build::size_of(fmm::types::Record::new(vec![
-            COUNT_TYPE.into(),
+            HEADER_TYPE.clone().into(),
             type_,
         ]))),
     );
 
     builder.store(
-        fmm::ir::Primitive::PointerInteger(INITIAL_COUNT as i64),
+        fmm::build::record(vec![count::compile_initial().into(), EMPTY_TAG.into()]),
         fmm::build::record_address(pointer.clone(), 0)?,
     );
 
@@ -34,13 +38,26 @@ pub fn free(
     builder.free_heap(fmm::build::bit_cast(
         fmm::types::GENERIC_POINTER_TYPE.clone(),
         fmm::build::pointer_address(
-            fmm::build::bit_cast(
-                fmm::types::Pointer::new(fmm::types::Primitive::PointerInteger),
-                pointer,
-            ),
+            fmm::build::bit_cast(fmm::types::Pointer::new(HEADER_TYPE.clone()), pointer),
             fmm::ir::Primitive::PointerInteger(-1),
         )?,
     ));
 
     Ok(())
 }
+
+pub fn get_counter_pointer(
+    heap_pointer: &fmm::build::TypedExpression,
+) -> Result<fmm::build::TypedExpression, fmm::build::BuildError> {
+    Ok(fmm::build::record_address(
+        fmm::build::pointer_address(
+            fmm::build::bit_cast(
+                fmm::types::Pointer::new(HEADER_TYPE.clone()),
+                heap_pointer.clone(),
+            ),
+            fmm::ir::Primitive::PointerInteger(-1),
+        )?,
+        0,
+    )?
+    .into())
+}

lib/mir-fmm/src/reference_count/pointer.rs

@@ -1,4 +1,4 @@
-use super::{super::error::CompileError, heap};
+use super::{super::error::CompileError, count, heap};
 
 pub fn clone(
     builder: &fmm::build::InstructionBuilder,
@@ -88,12 +88,11 @@ pub fn is_owned(
         |builder| {
             Ok(builder.branch(fmm::build::comparison_operation(
                 fmm::ir::ComparisonOperator::Equal,
-                // TODO Does this actually need to be atomic?
                 builder.atomic_load(
-                    get_counter_pointer(pointer)?,
+                    heap::get_counter_pointer(pointer)?,
                     fmm::ir::AtomicOrdering::Relaxed,
                 )?,
-                fmm::ir::Primitive::PointerInteger(heap::INITIAL_COUNT as i64),
+                count::compile_initial(),
             )?))
         },
         |builder| Ok(builder.branch(fmm::ir::Primitive::Boolean(false))),
@@ -107,8 +106,8 @@ fn increment_count(
 ) -> Result<(), CompileError> {
     builder.atomic_operation(
         fmm::ir::AtomicOperator::Add,
-        get_counter_pointer(pointer)?,
-        fmm::ir::Primitive::PointerInteger(1),
+        heap::get_counter_pointer(pointer)?,
+        count::compile(1),
         ordering,
     )?;
 
@@ -123,11 +122,11 @@ fn decrement_and_compare_count(
         fmm::ir::ComparisonOperator::Equal,
         builder.atomic_operation(
             fmm::ir::AtomicOperator::Subtract,
-            get_counter_pointer(pointer)?,
-            fmm::ir::Primitive::PointerInteger(1),
+            heap::get_counter_pointer(pointer)?,
+            count::compile(1),
             fmm::ir::AtomicOrdering::Release,
         )?,
-        fmm::ir::Primitive::PointerInteger(heap::INITIAL_COUNT as i64),
+        count::compile_initial(),
     )?
     .into())
 }
@@ -194,16 +193,3 @@ fn is_heap_pointer(
     )?
     .into())
 }
-
-fn get_counter_pointer(
-    heap_pointer: &fmm::build::TypedExpression,
-) -> Result<fmm::build::TypedExpression, fmm::build::BuildError> {
-    Ok(fmm::build::pointer_address(
-        fmm::build::bit_cast(
-            fmm::types::Pointer::new(heap::COUNT_TYPE),
-            heap_pointer.clone(),
-        ),
-        fmm::ir::Primitive::PointerInteger(-1),
-    )?
-    .into())
-}

lib/mir-fmm/src/snapshots/mir_fmm__entry_function__tests__do_not_overwrite_global_functions_in_variables.snap

@@ -15,7 +15,7 @@ expression: "fmm::analysis::format_module(&context.module_builder().as_module())
   (function _fmm_0 _closure
     (block
       (allocate-heap (size-of record) _fmm_1)
-      (store 0 (record-address (bit-cast _fmm_1) 0))
+      (store (record 0 0) (record-address (bit-cast _fmm_1) 0))
       (store (record _fmm_2 _fmm_5 (record)) (record-address (bit-cast _fmm_1) 1))
       (atomic-load (record-address (bit-cast (bit& (bit-cast (bit-cast (record-address (bit-cast _fmm_1) 1))) (bit! 1))) 0) _fmm_7)
       (call _fmm_7 (bit-cast (bit-cast (record-address (bit-cast _fmm_1) 1))) _fmm_8)

packages/core/ffi/Cargo.toml

@@ -8,5 +8,5 @@ edition = "2021"
 crate-type = ["staticlib"]
 
 [dependencies]
-ffi = { package = "pen-ffi", version = "0.3" }
+ffi = { package = "pen-ffi", version = "0.4" }
 siphasher = "0.3"

packages/http/ffi/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 crate-type = ["staticlib"]
 
 [dependencies]
-ffi = { package = "pen-ffi", version = "0.3", features = ["runtime"] }
+ffi = { package = "pen-ffi", version = "0.4", features = ["runtime"] }
 futures = "0.3"
 hyper = { version = "0.14", features = [
   "client",

packages/json/ffi/Cargo.toml

@@ -8,4 +8,4 @@ edition = "2021"
 crate-type = ["staticlib"]
 
 [dependencies]
-ffi = { package = "pen-ffi", version = "0.3" }
+ffi = { package = "pen-ffi", version = "0.4" }