RedBear-OS/local/recipes/drivers/linux-kpi/source/src/rust_impl/idr.rs

use std::collections::HashMap;
use std::ffi::c_void;
use std::ptr;

const EINVAL: i32 = 22;
const ENOSPC: i32 = 28;

#[repr(C)]
pub struct Idr {
    map: HashMap<u32, usize>,
    next_id: u32,
}

#[no_mangle]
pub extern "C" fn rust_idr_init(internal: *mut *mut c_void) {
    if internal.is_null() {
        return;
    }

    unsafe {
        let boxed = Box::new(Idr {
            map: HashMap::new(),
            next_id: 0,
        });
        *internal = Box::into_raw(boxed).cast();
    }
}

fn normalize_id(value: i32) -> Option<u32> {
    if value < 0 {
        None
    } else {
        Some(value as u32)
    }
}

unsafe fn as_idr_mut<'a>(internal: *mut c_void) -> Option<&'a mut Idr> {
    if internal.is_null() {
        None
    } else {
        Some(&mut *(internal.cast::<Idr>()))
    }
}

unsafe fn as_idr_ref<'a>(internal: *mut c_void) -> Option<&'a Idr> {
    if internal.is_null() {
        None
    } else {
        Some(&*(internal.cast::<Idr>()))
    }
}

#[no_mangle]
pub extern "C" fn rust_idr_alloc(
    internal: *mut c_void,
    ptr: *mut c_void,
    start: i32,
    end: i32,
    _gfp: u32,
) -> i32 {
    let idr_ref = match unsafe { as_idr_mut(internal) } {
        Some(idr_ref) => idr_ref,
        None => return -EINVAL,
    };

    let start = match normalize_id(start) {
        Some(start) => start,
        None => return -EINVAL,
    };
    let end = match end {
        0 => None,
        value if value > 0 => Some(value as u32),
        _ => return -EINVAL,
    };

    if let Some(end) = end {
        if start >= end {
            return -EINVAL;
        }
    }

    let initial = idr_ref.next_id.max(start);

    if let Some(end) = end {
        for candidate in initial..end {
            if let std::collections::hash_map::Entry::Vacant(entry) = idr_ref.map.entry(candidate) {
                entry.insert(ptr as usize);
                idr_ref.next_id = candidate.saturating_add(1);
                if idr_ref.next_id >= end {
                    idr_ref.next_id = start;
                }
                return candidate as i32;
            }
        }

        for candidate in start..initial {
            if let std::collections::hash_map::Entry::Vacant(entry) = idr_ref.map.entry(candidate) {
                entry.insert(ptr as usize);
                idr_ref.next_id = candidate.saturating_add(1);
                if idr_ref.next_id >= end {
                    idr_ref.next_id = start;
                }
                return candidate as i32;
            }
        }

        return -ENOSPC;
    }

    for candidate in initial..=u32::MAX {
        if let std::collections::hash_map::Entry::Vacant(entry) = idr_ref.map.entry(candidate) {
            entry.insert(ptr as usize);
            idr_ref.next_id = if candidate == u32::MAX {
                start
            } else {
                candidate.saturating_add(1).max(start)
            };
            return candidate as i32;
        }
    }

    for candidate in start..initial {
        if let std::collections::hash_map::Entry::Vacant(entry) = idr_ref.map.entry(candidate) {
            entry.insert(ptr as usize);
            idr_ref.next_id = if candidate == u32::MAX {
                start
            } else {
                candidate.saturating_add(1).max(start)
            };
            return candidate as i32;
        }
    }

    -ENOSPC
}

#[no_mangle]
pub extern "C" fn rust_idr_find(internal: *mut c_void, id: i32) -> *mut c_void {
    let idr_ref = match unsafe { as_idr_ref(internal) } {
        Some(idr_ref) => idr_ref,
        None => return ptr::null_mut(),
    };

    let id = match normalize_id(id) {
        Some(id) => id,
        None => return ptr::null_mut(),
    };

    match idr_ref.map.get(&id) {
        Some(value) => *value as *mut c_void,
        None => ptr::null_mut(),
    }
}

#[no_mangle]
pub extern "C" fn rust_idr_remove(internal: *mut c_void, id: i32) {
    let idr_ref = match unsafe { as_idr_mut(internal) } {
        Some(idr_ref) => idr_ref,
        None => return,
    };

    let id = match normalize_id(id) {
        Some(id) => id,
        None => return,
    };

    idr_ref.map.remove(&id);
    if id < idr_ref.next_id {
        idr_ref.next_id = id;
    }
}

#[no_mangle]
pub extern "C" fn rust_idr_destroy(internal: *mut *mut c_void) {
    if internal.is_null() {
        return;
    }

    unsafe {
        let raw = *internal;
        if raw.is_null() {
            return;
        }

        drop(Box::from_raw(raw.cast::<Idr>()));
        *internal = ptr::null_mut();
    }
}

#[no_mangle]
pub extern "C" fn rust_idr_for_each_entry(internal: *mut c_void, id: *mut i32) -> *mut c_void {
    if id.is_null() {
        return ptr::null_mut();
    }

    let start = match normalize_id(unsafe { *id }) {
        Some(start) => start,
        None => return ptr::null_mut(),
    };

    let idr_ref = match unsafe { as_idr_ref(internal) } {
        Some(idr_ref) => idr_ref,
        None => return ptr::null_mut(),
    };

    let mut found_id: Option<u32> = None;
    let mut found_ptr: Option<usize> = None;
    for (&key, &value) in idr_ref.map.iter() {
        if key < start {
            continue;
        }
        if found_id.is_none() || key < found_id.unwrap() {
            found_id = Some(key);
            found_ptr = Some(value);
        }
    }

    match (found_id, found_ptr) {
        (Some(key), Some(ptr)) => {
            unsafe { *id = key as i32 };
            ptr as *mut c_void
        }
        _ => ptr::null_mut(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn idr_alloc_and_find_round_trip() {
        let mut internal = ptr::null_mut();
        rust_idr_init(&mut internal);

        let ptr1: *mut c_void = 0x1000 as *mut c_void;
        let id1 = rust_idr_alloc(internal, ptr1, 1, 0, 0);
        assert!(id1 >= 1, "allocated ID should be >= start");

        assert_eq!(rust_idr_find(internal, id1), ptr1);
        assert_eq!(rust_idr_find(internal, 9999), std::ptr::null_mut());

        rust_idr_destroy(&mut internal);
    }

    #[test]
    fn idr_remove_frees_slot() {
        let mut internal = ptr::null_mut();
        rust_idr_init(&mut internal);

        let ptr1: *mut c_void = 0x2000 as *mut c_void;
        let id1 = rust_idr_alloc(internal, ptr1, 10, 0, 0);
        assert!(id1 >= 10);

        rust_idr_remove(internal, id1);
        assert_eq!(rust_idr_find(internal, id1), std::ptr::null_mut());

        rust_idr_destroy(&mut internal);
    }

    #[test]
    fn idr_alloc_with_bounded_range() {
        let mut internal = ptr::null_mut();
        rust_idr_init(&mut internal);

        let ptr1: *mut c_void = 0x3000 as *mut c_void;
        let id1 = rust_idr_alloc(internal, ptr1, 5, 8, 0);
        assert!(id1 >= 5 && id1 < 8, "ID should be in [5, 8)");

        rust_idr_destroy(&mut internal);
    }

    #[test]
    fn idr_alloc_returns_enospc_when_full() {
        let mut internal = ptr::null_mut();
        rust_idr_init(&mut internal);

        let ptr1: *mut c_void = 0x4000 as *mut c_void;
        let id1 = rust_idr_alloc(internal, ptr1, 1, 2, 0);
        assert_eq!(id1, 1);

        let ptr2: *mut c_void = 0x4001 as *mut c_void;
        let id2 = rust_idr_alloc(internal, ptr2, 1, 2, 0);
        assert_eq!(id2, -ENOSPC);

        rust_idr_destroy(&mut internal);
    }

    #[test]
    fn idr_null_pointers_are_safe() {
        assert_eq!(
            rust_idr_alloc(std::ptr::null_mut(), std::ptr::null_mut(), 1, 0, 0),
            -EINVAL
        );
        assert_eq!(rust_idr_find(std::ptr::null_mut(), 1), std::ptr::null_mut());
        rust_idr_remove(std::ptr::null_mut(), 1);
        rust_idr_destroy(std::ptr::null_mut());
        rust_idr_init(std::ptr::null_mut());
    }

    #[test]
    fn idr_alloc_reuses_removed_id() {
        let mut internal = ptr::null_mut();
        rust_idr_init(&mut internal);

        let ptr1: *mut c_void = 0x5000 as *mut c_void;
        let id1 = rust_idr_alloc(internal, ptr1, 1, 0, 0);

        rust_idr_remove(internal, id1);

        let ptr2: *mut c_void = 0x5001 as *mut c_void;
        let id2 = rust_idr_alloc(internal, ptr2, 1, 0, 0);
        assert_eq!(id2, id1, "should reuse removed ID");

        rust_idr_destroy(&mut internal);
    }
}