[ReuseIR] add nullable (#22)

* adjust runtime * [ReuseIR] add nullable type
SchrodingerZhu · Aug 22, 2024 · 2badbb8 · 2badbb8
1 parent f6d485d
commit 2badbb8
Show file tree

Hide file tree

Showing 5 changed files with 102 additions and 58 deletions.
diff --git a/reuse-mlir/include/ReuseIR/IR/ReuseIRTypes.td b/reuse-mlir/include/ReuseIR/IR/ReuseIRTypes.td
@@ -68,6 +68,21 @@ def ReuseIR_RegionCtxType : ReuseIR_Type<"RegionCtx", "region_ctx",
   }];
 }
 
+def ReuseIR_NullableType : ReuseIR_Type<"Nullable", "nullable",
+    [DeclareTypeInterfaceMethods<DataLayoutTypeInterface>]> {
+  let summary = "ReuseIR Nullable Type";
+  let description = [{
+    `reuse_ir.nullable` is a nullable type. Rc pointers on default at nonnull in ReuseIR.
+    This is just a type wrapper to indicate that the pointer can be null.
+  }];
+  let parameters = (ins
+    "::mlir::reuse_ir::RcType":$pointer
+  );
+  let assemblyFormat = [{
+    `<` $pointer `>`
+  }];
+}
+
 def ReuseIR_RcBoxType : ReuseIR_Type<"RcBox", "rcbox",
     [DeclareTypeInterfaceMethods<DataLayoutTypeInterface>,DeclareTypeInterfaceMethods<ReuseIRCompositeLayoutInterface>]> {
   let summary = "ReuseIR Control Box";

diff --git a/reuse-mlir/src/cxx/IR/ReuseIROps.cpp b/reuse-mlir/src/cxx/IR/ReuseIROps.cpp
@@ -84,8 +84,10 @@ mlir::reuse_ir::LogicalResult LoadOp::verify() {
     if (input.getFreezingKind().getValue() == FreezingKind::nonfreezing)
       return emitOpError(
           "cannot load a mutable RC pointer through a nonfreezing reference");
-    targetType = RcType::get(getContext(), mref.getPointee(),
-                             mref.getAtomicKind(), input.getFreezingKind());
+    targetType = NullableType::get(getContext(),
+                                   RcType::get(getContext(), mref.getPointee(),
+                                               mref.getAtomicKind(),
+                                               input.getFreezingKind()));
   } else
     targetType = input.getPointee();
   if (targetType != getType())

diff --git a/reuse-mlir/src/cxx/IR/ReuseIRTypes.cpp b/reuse-mlir/src/cxx/IR/ReuseIRTypes.cpp
@@ -95,6 +95,7 @@ GENERATE_POINTER_ALIKE_LAYOUT(RefType)
 GENERATE_POINTER_ALIKE_LAYOUT(TokenType)
 GENERATE_POINTER_ALIKE_LAYOUT(MRefType)
 GENERATE_POINTER_ALIKE_LAYOUT(RegionCtxType)
+GENERATE_POINTER_ALIKE_LAYOUT(NullableType)
 #pragma pop_macro("GENERATE_POINTER_ALIKE_LAYOUT")
 
 // RcBox DataLayoutInterface:
@@ -343,6 +344,9 @@ void populateLLVMTypeConverter(CompositeLayoutCache &cache,
   converter.addConversion([](RegionCtxType type) -> Type {
     return mlir::LLVM::LLVMPointerType::get(type.getContext());
   });
+  converter.addConversion([](NullableType type) -> Type {
+    return mlir::LLVM::LLVMPointerType::get(type.getContext());
+  });
   converter.addConversion(
       [&converter, &cache](ReuseIRCompositeLayoutInterface type) -> Type {
         return cache.get(type).getLLVMType(converter);

diff --git a/reuse-mlir/test/integration/lowering/load_op_lowering.mlir b/reuse-mlir/test/integration/lowering/load_op_lowering.mlir
@@ -1,27 +1,29 @@
 // RUN: %reuse-opt %s -convert-reuse-ir-to-llvm | %FileCheck %s
-!test = !reuse_ir.composite<!reuse_ir.composite<i32, i32, f128>, i32>
+!mref = !reuse_ir.mref<f128, nonatomic>
+!inner = !reuse_ir.composite<i32, i32, !mref>
+!test = !reuse_ir.composite<!inner, i32>
 !array = !reuse_ir.array<!test, 16, 16>
 module @test attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<f80, dense<128> : vector<2xi64>>, #dlti.dl_entry<i128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<f64, dense<64> : vector<2xi64>>, #dlti.dl_entry<f16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>, #dlti.dl_entry<"dlti.endianness", "little">>, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"} 
 {
-    func.func @array_projection(%0: !reuse_ir.rc<!array, nonatomic, nonfreezing>) -> f128 {
+    func.func @array_projection(%0: !reuse_ir.rc<!array, nonatomic, frozen>) -> !reuse_ir.nullable<<f128, nonatomic, frozen>> {
         %1 = reuse_ir.borrow %0 : 
-            !reuse_ir.rc<!array, nonatomic, nonfreezing>
-            -> !reuse_ir.ref<!array, nonfreezing>
-        // CHECK: %[[REG0:[0-9a-z]+]] = llvm.getelementptr %{{[0-9a-z]+}}[13] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<16 x struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>>
+            !reuse_ir.rc<!array, nonatomic, frozen>
+            -> !reuse_ir.ref<!array, frozen>
+        // CHECK: %[[REG0:[0-9a-z]+]] = llvm.getelementptr %{{[0-9a-z]+}}[13] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<16 x struct<(struct<(i32, i32, ptr)>, i32, array<4 x i8>)>>
         %2 = reuse_ir.proj %1[13] : 
-            !reuse_ir.ref<!array, nonfreezing> -> !reuse_ir.ref<!reuse_ir.array<!test, 16>, nonfreezing>
-        // CHECK: %[[REG1:[0-9a-z]+]] = llvm.getelementptr %[[REG0]][7] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>
+            !reuse_ir.ref<!array, frozen> -> !reuse_ir.ref<!reuse_ir.array<!test, 16>, frozen>
+        // CHECK: %[[REG1:[0-9a-z]+]] = llvm.getelementptr %[[REG0]][7] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, ptr)>, i32, array<4 x i8>)>
         %3 = reuse_ir.proj %2[7] : 
-            !reuse_ir.ref<!reuse_ir.array<!test, 16>, nonfreezing> -> !reuse_ir.ref<!test, nonfreezing>
-        // CHECK: %[[REG2:[0-9a-z]+]] = llvm.getelementptr %[[REG1]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>
+            !reuse_ir.ref<!reuse_ir.array<!test, 16>, frozen> -> !reuse_ir.ref<!test, frozen>
+        // CHECK: %[[REG2:[0-9a-z]+]] = llvm.getelementptr %[[REG1]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, ptr)>, i32, array<4 x i8>)>
         %4 = reuse_ir.proj %3[0] : 
-            !reuse_ir.ref<!test, nonfreezing> -> !reuse_ir.ref<!reuse_ir.composite<i32, i32, f128>, nonfreezing>
-        // CHECK: %[[REG3:[0-9a-z]+]] = llvm.getelementptr %[[REG2]][0, 3] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, i32, array<8 x i8>, f128)>
+            !reuse_ir.ref<!test, frozen> -> !reuse_ir.ref<!inner, frozen>
+        // CHECK: %[[REG3:[0-9a-z]+]] = llvm.getelementptr %[[REG2]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, i32, ptr)>
         %5 = reuse_ir.proj %4[2] : 
-            !reuse_ir.ref<!reuse_ir.composite<i32, i32, f128>, nonfreezing> -> !reuse_ir.ref<f128, nonfreezing>
-        // CHECK: %{{[0-9a-z]+}} = llvm.load %[[REG3]] {alignment = 16 : i64} : !llvm.ptr -> f128
+            !reuse_ir.ref<!inner, frozen> -> !reuse_ir.ref<!mref, frozen>
+        // CHECK: %{{[0-9a-z]+}} = llvm.load %[[REG3]] {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
         %6 = reuse_ir.load %5 : 
-            !reuse_ir.ref<f128, nonfreezing> -> f128
-        return %6 : f128
+            !reuse_ir.ref<!mref, frozen> -> !reuse_ir.nullable<<f128, nonatomic, frozen>>
+        return %6 : !reuse_ir.nullable<<f128, nonatomic, frozen>>
     }
 }
diff --git a/reuse-runtime/src/freezable.rs b/reuse-runtime/src/freezable.rs
@@ -1,5 +1,7 @@
 #![allow(clippy::missing_safety_doc)]
 
+use std::ptr::NonNull;
+
 use smallvec::SmallVec;
 
 const STATUS_COUNTER_PADDING_BITS: usize = 2;
@@ -43,8 +45,8 @@ impl FreezingStatus {
     pub fn get_value(self) -> usize {
         unsafe { (self.scalar & (isize::MAX as usize)) >> 2 }
     }
-    pub fn get_pointer(self) -> *mut FreezableRcBoxHeader {
-        unsafe { self.ptr }
+    pub unsafe fn get_pointer_unchecked(self) -> NonNull<FreezableRcBoxHeader> {
+        unsafe { NonNull::new_unchecked(self.ptr) }
     }
     pub fn unmarked() -> Self {
         Self { scalar: 0 }
@@ -90,76 +92,81 @@ pub struct FreezableRcBoxHeader {
 
 struct FieldIterator {
     slice_iter: std::slice::Iter<'static, usize>,
-    base_pointer: *mut u8,
+    base_pointer: NonNull<u8>,
 }
 
 impl FieldIterator {
-    unsafe fn new(base_pointer: *mut FreezableRcBoxHeader) -> Self {
-        let vtable = (*base_pointer).vtable;
+    unsafe fn new(base_pointer: NonNull<FreezableRcBoxHeader>) -> Self {
+        let vtable = base_pointer.as_ref().vtable;
         let slice =
             std::slice::from_raw_parts((*vtable).scan_offset.as_ptr(), (*vtable).scan_count);
         Self {
             slice_iter: slice.iter(),
-            base_pointer: base_pointer as _,
+            base_pointer: base_pointer.cast(),
         }
     }
 }
 
 impl Iterator for FieldIterator {
-    type Item = *mut FreezableRcBoxHeader;
+    type Item = Option<NonNull<FreezableRcBoxHeader>>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.slice_iter.next().map(|offset| unsafe {
-            let pointer = self.base_pointer.byte_add(*offset) as *mut *mut FreezableRcBoxHeader;
-            *pointer
+            let pointer = self
+                .base_pointer
+                .byte_add(*offset)
+                .cast::<*mut FreezableRcBoxHeader>();
+            NonNull::new(pointer.read())
         })
     }
 }
 
-unsafe fn increase_refcnt(object: *mut FreezableRcBoxHeader, count: usize) {
-    (*object).status.scalar += count << STATUS_COUNTER_PADDING_BITS;
+unsafe fn increase_refcnt(mut object: NonNull<FreezableRcBoxHeader>, count: usize) {
+    object.as_mut().status.scalar += count << STATUS_COUNTER_PADDING_BITS;
 }
 
-unsafe fn decrease_refcnt(object: *mut FreezableRcBoxHeader, count: usize) -> bool {
-    (*object).status.scalar -= count << STATUS_COUNTER_PADDING_BITS;
-    (*object).status.scalar <= (isize::MIN as usize | 0b11)
+unsafe fn decrease_refcnt(mut object: NonNull<FreezableRcBoxHeader>, count: usize) -> bool {
+    object.as_mut().status.scalar -= count << STATUS_COUNTER_PADDING_BITS;
+    object.as_mut().status.scalar <= (isize::MIN as usize | 0b11)
 }
 
-unsafe fn find_representative(mut object: *mut FreezableRcBoxHeader) -> *mut FreezableRcBoxHeader {
+unsafe fn find_representative(
+    mut object: NonNull<FreezableRcBoxHeader>,
+) -> NonNull<FreezableRcBoxHeader> {
     let mut root = object;
-    while (*root).status.get_kind() == StatusKind::Representative {
-        root = (*root).status.get_pointer();
+    while root.as_ref().status.get_kind() == StatusKind::Representative {
+        root = root.as_ref().status.get_pointer_unchecked();
     }
-    while (*object).status.get_kind() == StatusKind::Representative {
-        let next = (*object).status.get_pointer();
-        (*object).status = FreezingStatus::representative(root);
+    while object.as_ref().status.get_kind() == StatusKind::Representative {
+        let next = object.as_ref().status.get_pointer_unchecked();
+        object.as_mut().status = FreezingStatus::representative(root.as_ptr());
         object = next;
     }
     root
 }
 
-unsafe fn union(x: *mut FreezableRcBoxHeader, y: *mut FreezableRcBoxHeader) -> bool {
+unsafe fn union(x: NonNull<FreezableRcBoxHeader>, y: NonNull<FreezableRcBoxHeader>) -> bool {
     let mut rep_x = find_representative(x);
     let mut rep_y = find_representative(y);
     if rep_x == rep_y {
         return false;
     }
-    if (*rep_x).status.get_value() < (*rep_y).status.get_value() {
+    if rep_x.as_ref().status.get_value() < rep_y.as_ref().status.get_value() {
         std::mem::swap(&mut rep_x, &mut rep_y);
     }
-    if (*rep_x).status.get_value() == (*rep_y).status.get_value() {
-        (*rep_x).status = FreezingStatus::rank((*rep_x).status.get_value() + 1);
+    if rep_x.as_ref().status.get_value() == rep_y.as_ref().status.get_value() {
+        rep_x.as_mut().status = FreezingStatus::rank(rep_x.as_ref().status.get_value() + 1);
     }
-    (*rep_y).status = FreezingStatus::representative(rep_x);
+    rep_y.as_mut().status = FreezingStatus::representative(rep_x.as_ptr());
     true
 }
 
 #[cold]
-unsafe fn dispose(object: *mut FreezableRcBoxHeader) {
-    type Stack = SmallVec<[*mut FreezableRcBoxHeader; 16]>;
-    unsafe fn add_stack(stack: &mut Stack, object: *mut FreezableRcBoxHeader) {
+unsafe fn dispose(object: NonNull<FreezableRcBoxHeader>) {
+    type Stack = SmallVec<[NonNull<FreezableRcBoxHeader>; 16]>;
+    unsafe fn add_stack(stack: &mut Stack, mut object: NonNull<FreezableRcBoxHeader>) {
         stack.push(object);
-        (*object).status = FreezingStatus::disposing();
+        object.as_mut().status = FreezingStatus::disposing();
     }
     let mut dfs = Stack::new();
     let mut scc = Stack::new();
@@ -169,9 +176,9 @@ unsafe fn dispose(object: *mut FreezableRcBoxHeader) {
         scc.push(obj);
         while let Some(obj) = scc.pop() {
             recycle.push(obj);
-            for field in FieldIterator::new(obj) {
+            for field in FieldIterator::new(obj).flatten() {
                 let next = find_representative(field);
-                match (*next).status.get_kind() {
+                match next.as_ref().status.get_kind() {
                     StatusKind::Disposing if field != next => add_stack(&mut scc, field),
                     StatusKind::Rc if decrease_refcnt(next, 1) => add_stack(&mut dfs, next),
                     _ => continue,
@@ -181,21 +188,28 @@ unsafe fn dispose(object: *mut FreezableRcBoxHeader) {
     }
     stacker::maybe_grow(16 * 1024, 1024 * 1024, || {
         while let Some(obj) = recycle.pop() {
-            let vtable = (*obj).vtable;
+            let vtable = obj.as_ref().vtable;
             if let Some(dtor) = (*vtable).drop {
-                dtor(obj);
+                dtor(obj.as_ptr());
             }
-            crate::allocator::__reuse_ir_dealloc(obj as _, (*vtable).size, (*vtable).alignment);
+            crate::allocator::__reuse_ir_dealloc(
+                obj.cast().as_ptr(),
+                (*vtable).size,
+                (*vtable).alignment,
+            );
         }
     });
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn __reuse_ir_freeze(object: *mut FreezableRcBoxHeader) {
-    type PendingList = SmallVec<[*mut FreezableRcBoxHeader; 32]>;
+    let Some(object) = NonNull::new(object) else {
+        panic!("attempt to freeze null object");
+    };
+    type PendingList = SmallVec<[NonNull<FreezableRcBoxHeader>; 32]>;
     let mut pending = PendingList::new();
-    unsafe fn freeze(object: *mut FreezableRcBoxHeader, pending: &mut PendingList) {
-        match (*object).status.get_kind() {
+    unsafe fn freeze(mut object: NonNull<FreezableRcBoxHeader>, pending: &mut PendingList) {
+        match object.as_ref().status.get_kind() {
             StatusKind::Rc => {
                 let root = find_representative(object);
                 increase_refcnt(root, 1);
@@ -209,11 +223,11 @@ pub unsafe extern "C" fn __reuse_ir_freeze(object: *mut FreezableRcBoxHeader) {
                 }
             },
             StatusKind::Unmarked => {
-                (*object).status = FreezingStatus::rc(1);
+                object.as_mut().status = FreezingStatus::rc(1);
                 pending.push(object);
                 stacker::maybe_grow(16 * 1024, 1024 * 1024, || {
                     let iter = FieldIterator::new(object);
-                    for field in iter {
+                    for field in iter.flatten() {
                         freeze(field, pending);
                     }
                 });
@@ -234,6 +248,9 @@ pub unsafe extern "C" fn __reuse_ir_acquire_freezable(
     object: *mut FreezableRcBoxHeader,
     count: usize,
 ) {
+    let Some(object) = NonNull::new(object) else {
+        panic!("attempt to acquire null object");
+    };
     let root = find_representative(object);
     increase_refcnt(root, count);
 }
@@ -243,8 +260,12 @@ pub unsafe extern "C" fn __reuse_ir_release_freezable(
     object: *mut FreezableRcBoxHeader,
     count: usize,
 ) {
-    if decrease_refcnt(object, count) {
-        dispose(object)
+    if let Some(object) = NonNull::new(object) {
+        if decrease_refcnt(object, count) {
+            dispose(object);
+        }
+    } else {
+        panic!("attempt to release null object");
     }
 }