Implement and use ModuleCell

Author: phcoder

grub_rust_core/src/lib.rs

@@ -8,20 +8,23 @@
 
 extern crate alloc;
 
+use alloc::boxed::Box;
 use alloc::vec::Vec;
 use alloc::vec;
 use alloc::string::String;
 use alloc::string::ToString;
+use alloc::ffi::CString;
+
 use core::alloc::{GlobalAlloc, Layout};
+use core::cell::Cell;
+use core::cell::RefCell;
 use core::ffi::c_char;
 use core::ffi::c_int;
 use core::ffi::c_void;
-use core::fmt::{self, Arguments, Write};
-
-use alloc::ffi::CString;
 use core::ffi::CStr;
-use core::panic::PanicInfo;
+use core::fmt::{self, Arguments, Write};
 use core::num::TryFromIntError;
+use core::panic::PanicInfo;
 
 #[unsafe(link_section = ".module_license")]
 pub static LICENSE: [u8; 15] = *b"LICENSE=GPLv3+\0";
@@ -363,11 +366,9 @@ pub struct Command {
     cmd: *mut GrubCommand,
 }
 
-static mut commands: Vec<Command> = vec![];
-
 impl Command {
     pub fn register(name: &str, cb: fn (argv: &[&str]) -> Result<(), GrubError>,
-		    summary: &str, description: &str) {
+		    summary: &str, description: &str) -> Self {
 	let mut ret = Command {
 	    name: CString::new(name).unwrap(),
 	    summary: CString::new(summary).unwrap(),
@@ -381,15 +382,15 @@ impl Command {
 						  ret.description.as_ptr(),
 						  0);
 	    (*ret.cmd).data = cb as *mut c_void;
-	    commands.push(ret);
 	}
+	return ret;
     }
+}
 
-    pub fn unregister_all() {
+impl Drop for Command {
+    fn drop(&mut self) {
 	unsafe {
-	    for cmd in commands.iter() {
-		grub_unregister_command(cmd.cmd);
-	    }
+	    grub_unregister_command(self.cmd);
 	}
     }
 }
@@ -470,3 +471,58 @@ impl Drop for File {
 	unsafe { grub_file_close(self.file); }
     }
 }
+
+pub trait ModuleFini {
+    fn module_fini(&self);
+}
+
+pub struct ModuleCell<T> {
+    inner: Cell<Option<T>>,
+}
+
+pub struct ModuleRefHolder {
+    val: Cell<*const c_void>,
+    destructors: RefCell<Vec<Box<dyn ModuleFini>>>,
+}
+
+impl<T> ModuleCell<T> {
+    pub fn set(&'static self, dl: &ModuleRefHolder, value: T) {
+	self.inner.set(Some(value));
+	dl.destructors.borrow_mut().push(Box::new(self));
+    }
+
+    pub const fn new() -> Self {
+	Self{inner: Cell::new(None)}
+    }
+}
+
+impl<T> ModuleFini for &ModuleCell<T> {
+    fn module_fini(&self) {
+	self.inner.set(None);
+    }
+}
+
+// SAFETY: We're single-threaded with disabled interrupts
+unsafe impl<T> Sync for ModuleCell<T> {
+}
+
+impl ModuleRefHolder {
+    pub fn init(&self, value: *const c_void) {
+	self.val.set(value);
+    }
+
+    pub fn fini(&self) {
+	for d in self.destructors.borrow().iter() {
+	    d.module_fini();
+	}
+	self.destructors.borrow_mut().clear();
+    }
+
+    pub const fn new() -> Self {
+	Self{val: Cell::new(core::ptr::null()), destructors: RefCell::new(vec![])}
+    }
+}
+
+// SAFETY: We're single-threaded with disabled interrupts
+unsafe impl Sync for ModuleRefHolder {
+}

grub_rust_example/src/lib.rs

@@ -11,6 +11,7 @@ use core::cmp::min;
 use core::convert::TryFrom;
 
 use alloc::format;
+use core::ffi::c_void;
 
 use grub::dprintln;
 use grub::println;
@@ -116,17 +117,23 @@ fn rust_hexdump (args: &[&str]) -> Result<(), grub::GrubError> {
     return Ok(());
 }
 
+static HELLO_CMD: grub::ModuleCell<grub::Command> = grub::ModuleCell::new();
+static ERR_CMD: grub::ModuleCell<grub::Command> = grub::ModuleCell::new();
+static HEXDUMP_CMD: grub::ModuleCell<grub::Command> = grub::ModuleCell::new();
+static MODULE: grub::ModuleRefHolder = grub::ModuleRefHolder::new();
+
 #[unsafe(no_mangle)]
-pub extern "C" fn grub_mod_init() {
-    grub::Command::register("rust_hello", rust_hello,
-				"Rust hello", "Say hello from Rust.");
-    grub::Command::register("rust_err", rust_err,
-				"Rust error", "Error out from Rust.");
-    grub::Command::register("rust_hexdump", rust_hexdump,
-				"Rust hexdump", "Hexdump a file from Rust.");
+pub extern "C" fn grub_mod_init(module: *const c_void) {
+    MODULE.init(module);
+    HELLO_CMD.set(&MODULE, grub::Command::register("rust_hello", rust_hello,
+				      "Rust hello", "Say hello from Rust."));
+    ERR_CMD.set(&MODULE, grub::Command::register("rust_err", rust_err,
+					"Rust error", "Error out from Rust."));
+    HEXDUMP_CMD.set(&MODULE, grub::Command::register("rust_hexdump", rust_hexdump,
+					    "Rust hexdump", "Hexdump a file from Rust."));
 }
 
 #[unsafe(no_mangle)]
 pub extern "C" fn grub_mod_fini() {
-    grub::Command::unregister_all();
+    MODULE.fini();
 }