How to init a Struct within a Enum in unsafe codes

ghz 8months ago ⋅ 63 views

I'm trying to initialize an enum variant, which is a struct in my case, using Box::new_uninit_in and ptr::addr_of_mut!. However, I'm struggling to access the right field. Can someone help me with this problem?

#![feature(allocator_api)]
use std::alloc::Allocator;
use std::mem::MaybeUninit;
use core::ptr::{self, NonNull};

fn main() {}

enum Node<K, V> {
    LeafNode {
        size: u16,
        keys: [MaybeUninit<K>; 10],
        vals: [MaybeUninit<V>; 10],
        prev: Option<NonNull<Node<K, V>>>,
        next: Option<NonNull<Node<K, V>>>,
    },
    InternalNode {
        size: u16,
        keys: [MaybeUninit<K>; 10],
        vals: [MaybeUninit<NonNull<Node<K, V>>>; 11],
    },
}

impl<K, V> Node<K, V> {
    unsafe fn init(this: *mut Self) {
        unsafe {
            // 1. How do I access the fields of the struct within the enum?
            // 2. How can I initialize the enum as either the LeafNode variant or the InternalNode variant?
            ptr::addr_of_mut!((*this).size).write(0);
        }
    } 
    fn new<A: Allocator + Clone>(alloc: A) -> Box<Self, A> {
        unsafe {
            let mut node = Box::new_uninit_in(alloc);
            Node::init(node.as_mut_ptr());
            node.assume_init()
        }
    }
}

https://play.rust-lang.org/?version=nightly&mode=debug&edition=2024&gist=a6f780637a195cce4a5c4aabff131bfc

I tried to cast the node to Node::LeafNode using node as Node::LeafNode but rust doesn't allow me to do it. Maybe I did it in a wrong.

Answers

You're on the right track, but the syntax for accessing enum variants is a bit different. Here's how you can access the fields of the enum variants and initialize them:

#![feature(allocator_api)]

use std::alloc::Allocator;
use std::mem::MaybeUninit;
use core::ptr::{self, NonNull};

fn main() {}

enum Node<K, V> {
    LeafNode {
        size: u16,
        keys: [MaybeUninit<K>; 10],
        vals: [MaybeUninit<V>; 10],
        prev: Option<NonNull<Node<K, V>>>,
        next: Option<NonNull<Node<K, V>>>,
    },
    InternalNode {
        size: u16,
        keys: [MaybeUninit<K>; 10],
        vals: [MaybeUninit<NonNull<Node<K, V>>>; 11],
    },
}

impl<K, V> Node<K, V> {
    unsafe fn init(this: *mut Self) {
        match &mut *this {
            Node::LeafNode { size, keys, vals, prev, next } => {
                size.write(0);
                // Initialize other fields of LeafNode
            },
            Node::InternalNode { size, keys, vals } => {
                size.write(0);
                // Initialize other fields of InternalNode
            }
        }
    } 
    fn new<A: Allocator + Clone>(alloc: A) -> Box<Self, A> {
        unsafe {
            let mut node = Box::new_uninit_in(alloc);
            Node::init(node.as_mut_ptr());
            node.assume_init()
        }
    }
}

In the init method, you can use pattern matching (match) to access the fields of the enum variants. Depending on which variant this points to, you can initialize its fields accordingly. This approach allows you to initialize the enum as either the LeafNode variant or the InternalNode variant based on the pointer passed to the init method.