Merge pull request #157 from mulimoen/feature/iteration

Add more iterators on Group

Author: aldanor

CHANGELOG.md

@@ -49,6 +49,8 @@
   but must read the entire attribute at once).
 - Added support in `hdf5-sys` for the new functions in `hdf5` `1.10.6` and `1.10.7`.
 - Added support for creating external links on a `Group` with `link_external`.
+- Added `Location` methods: `get_info`, `get_info_by_name`, `loc_type`, and `open_by_token`.
+- Added `Group` methods: `iter_visit`, `iter_visit_default`, `get_all_of_type`, `datasets`, `groups`, and `named_datatypes`.
 
  ### Changed
 

hdf5-sys/src/h5o.rs

@@ -9,6 +9,12 @@ pub use {
 };
 #[cfg(hdf5_1_12_0)]
 pub use {H5O_info2_t as H5O_info_t, H5O_iterate2_t as H5O_iterate_t};
+#[cfg(not(hdf5_1_10_3))]
+pub use {
+    H5Oget_info1 as H5Oget_info, H5Oget_info_by_idx1 as H5Oget_info_by_idx,
+    H5Oget_info_by_name1 as H5Oget_info_by_name, H5Ovisit1 as H5Ovisit,
+    H5Ovisit_by_name1 as H5Ovisit_by_name,
+};
 
 use crate::internal_prelude::*;
 
@@ -201,19 +207,24 @@ pub type H5O_mcdt_search_cb_t =
 
 #[cfg(not(hdf5_1_10_3))]
 extern "C" {
-    pub fn H5Oget_info(loc_id: hid_t, oinfo: *mut H5O_info1_t) -> herr_t;
-    pub fn H5Oget_info_by_name(
+    #[link_name = "H5Oget_info"]
+    pub fn H5Oget_info1(loc_id: hid_t, oinfo: *mut H5O_info1_t) -> herr_t;
+    #[link_name = "H5Oget_info_by_name"]
+    pub fn H5Oget_info_by_name1(
         loc_id: hid_t, name: *const c_char, oinfo: *mut H5O_info1_t, lapl_id: hid_t,
     ) -> herr_t;
-    pub fn H5Oget_info_by_idx(
+    #[link_name = "H5Oget_info_by_idx"]
+    pub fn H5Oget_info_by_idx1(
         loc_id: hid_t, group_name: *const c_char, idx_type: H5_index_t, order: H5_iter_order_t,
         n: hsize_t, oinfo: *mut H5O_info1_t, lapl_id: hid_t,
     ) -> herr_t;
-    pub fn H5Ovisit(
+    #[link_name = "H5Ovisit"]
+    pub fn H5Ovisit1(
         obj_id: hid_t, idx_type: H5_index_t, order: H5_iter_order_t, op: H5O_iterate1_t,
         op_data: *mut c_void,
     ) -> herr_t;
-    pub fn H5Ovisit_by_name(
+    #[link_name = "H5Ovisit_by_name"]
+    pub fn H5Ovisit_by_name1(
         loc_id: hid_t, obj_name: *const c_char, idx_type: H5_index_t, order: H5_iter_order_t,
         op: H5O_iterate1_t, op_data: *mut c_void, lapl_id: hid_t,
     ) -> herr_t;
@@ -353,7 +364,7 @@ extern "C" {
 pub const H5O_MAX_TOKEN_SIZE: usize = 16;
 
 #[repr(C)]
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 #[cfg(hdf5_1_12_0)]
 pub struct H5O_token_t {
     __data: [u8; H5O_MAX_TOKEN_SIZE],
@@ -370,15 +381,15 @@ impl Default for H5O_token_t {
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct H5O_info2_t {
-    fileno: c_ulong,
-    token: H5O_token_t,
-    type_: H5O_type_t,
-    rc: c_uint,
-    atime: time_t,
-    mtime: time_t,
-    ctime: time_t,
-    btime: time_t,
-    num_attrs: hsize_t,
+    pub fileno: c_ulong,
+    pub token: H5O_token_t,
+    pub type_: H5O_type_t,
+    pub rc: c_uint,
+    pub atime: time_t,
+    pub mtime: time_t,
+    pub ctime: time_t,
+    pub btime: time_t,
+    pub num_attrs: hsize_t,
 }
 
 #[cfg(hdf5_1_12_0)]

src/hl.rs

@@ -24,8 +24,8 @@ pub use self::{
     dataspace::Dataspace,
     datatype::{Conversion, Datatype},
     file::{File, FileBuilder, OpenMode},
-    group::Group,
-    location::Location,
+    group::{Group, LinkInfo, LinkType},
+    location::{Location, LocationInfo, LocationToken, LocationType},
     object::Object,
     plist::PropertyList,
 };

src/hl/group.rs

@@ -7,14 +7,16 @@ use hdf5_sys::{
     h5d::H5Dopen2,
     h5g::{H5G_info_t, H5Gcreate2, H5Gget_info, H5Gopen2},
     h5l::{
-        H5L_info_t, H5L_iterate_t, H5Lcreate_external, H5Lcreate_hard, H5Lcreate_soft, H5Ldelete,
-        H5Lexists, H5Literate, H5Lmove, H5L_SAME_LOC,
+        H5L_info_t, H5L_iterate_t, H5L_type_t, H5Lcreate_external, H5Lcreate_hard, H5Lcreate_soft,
+        H5Ldelete, H5Lexists, H5Literate, H5Lmove, H5L_SAME_LOC,
     },
     h5p::{H5Pcreate, H5Pset_create_intermediate_group},
+    h5t::H5T_cset_t,
 };
 
 use crate::globals::H5P_LINK_CREATE;
 use crate::internal_prelude::*;
+use crate::{Location, LocationType};
 
 /// Represents the HDF5 group object.
 #[repr(transparent)]
@@ -212,37 +214,193 @@ impl Group {
         let name = to_cstring(name)?;
         Dataset::from_id(h5try!(H5Dopen2(self.id(), name.as_ptr(), H5P_DEFAULT)))
     }
+}
 
-    /// Returns names of all the members in the group, non-recursively.
-    pub fn member_names(&self) -> Result<Vec<String>> {
-        extern "C" fn members_callback(
-            _id: hid_t, name: *const c_char, _info: *const H5L_info_t, op_data: *mut c_void,
-        ) -> herr_t {
-            panic::catch_unwind(|| {
-                let other_data: &mut Vec<String> = unsafe { &mut *(op_data.cast::<Vec<String>>()) };
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum TraversalOrder {
+    Name,
+    Creation,
+}
 
-                other_data.push(string_from_cstr(name));
+impl Default for TraversalOrder {
+    fn default() -> Self {
+        Self::Name
+    }
+}
+
+impl From<TraversalOrder> for H5_index_t {
+    fn from(v: TraversalOrder) -> Self {
+        match v {
+            TraversalOrder::Name => Self::H5_INDEX_NAME,
+            TraversalOrder::Creation => Self::H5_INDEX_CRT_ORDER,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum IterationOrder {
+    Increasing,
+    Decreasing,
+    Native,
+}
+
+impl Default for IterationOrder {
+    fn default() -> Self {
+        Self::Native
+    }
+}
 
-                0 // Continue iteration
+impl From<IterationOrder> for H5_iter_order_t {
+    fn from(v: IterationOrder) -> Self {
+        match v {
+            IterationOrder::Increasing => Self::H5_ITER_INC,
+            IterationOrder::Decreasing => Self::H5_ITER_DEC,
+            IterationOrder::Native => Self::H5_ITER_NATIVE,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum LinkType {
+    Hard,
+    Soft,
+    External,
+}
+
+impl From<H5L_type_t> for LinkType {
+    fn from(link_type: H5L_type_t) -> Self {
+        match link_type {
+            H5L_type_t::H5L_TYPE_HARD => Self::Hard,
+            H5L_type_t::H5L_TYPE_SOFT => Self::Soft,
+            _ => Self::External,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct LinkInfo {
+    pub link_type: LinkType,
+    pub creation_order: Option<i64>,
+    pub is_utf8: bool,
+}
+
+impl From<&H5L_info_t> for LinkInfo {
+    fn from(link: &H5L_info_t) -> Self {
+        let link_type = link.type_.into();
+        let creation_order = if link.corder_valid == 1 { Some(link.corder) } else { None };
+        let is_utf8 = link.cset == H5T_cset_t::H5T_CSET_UTF8;
+        Self { link_type, creation_order, is_utf8 }
+    }
+}
+
+/// Iteration methods
+impl Group {
+    /// Visits all objects in the group
+    pub fn iter_visit<F, G>(
+        &self, iteration_order: IterationOrder, traversal_order: TraversalOrder, mut val: G,
+        mut op: F,
+    ) -> Result<G>
+    where
+        F: Fn(&Self, &str, LinkInfo, &mut G) -> bool,
+    {
+        /// Struct used to pass a tuple
+        struct Vtable<'a, F, D> {
+            f: &'a mut F,
+            d: &'a mut D,
+        }
+        // Maps a closure to a C callback
+        //
+        // This function will be called multiple times, but never concurrently
+        extern "C" fn callback<F, G>(
+            id: hid_t, name: *const c_char, info: *const H5L_info_t, op_data: *mut c_void,
+        ) -> herr_t
+        where
+            F: FnMut(&Group, &str, LinkInfo, &mut G) -> bool,
+        {
+            panic::catch_unwind(|| {
+                let vtable = op_data.cast::<Vtable<F, G>>();
+                let vtable = unsafe { vtable.as_mut().expect("iter_visit: null op_data ptr") };
+                unsafe { name.as_ref().expect("iter_visit: null name ptr") };
+                let name = unsafe { std::ffi::CStr::from_ptr(name) };
+                let info = unsafe { info.as_ref().expect("iter_vist: null info ptr") };
+                let handle = Handle::try_new(id).expect("iter_visit: unable to create a handle");
+                handle.incref();
+                let group = Group::from_handle(handle);
+                if (vtable.f)(&group, name.to_string_lossy().as_ref(), info.into(), vtable.d) {
+                    0
+                } else {
+                    1
+                }
             })
             .unwrap_or(-1)
         }
 
-        let callback_fn: H5L_iterate_t = Some(members_callback);
-        let iteration_position: *mut hsize_t = &mut { 0_u64 };
-        let mut result: Vec<String> = Vec::new();
-        let other_data: *mut c_void = (&mut result as *mut Vec<String>).cast::<c_void>();
+        let callback_fn: H5L_iterate_t = Some(callback::<F, G>);
+        let iter_pos: *mut hsize_t = &mut 0_u64;
+
+        // Store our references on the heap
+        let mut vtable = Vtable { f: &mut op, d: &mut val };
+        let other_data = (&mut vtable as *mut Vtable<_, _>).cast::<c_void>();
 
         h5call!(H5Literate(
             self.id(),
-            H5_index_t::H5_INDEX_NAME,
-            H5_iter_order_t::H5_ITER_INC,
-            iteration_position,
+            traversal_order.into(),
+            iteration_order.into(),
+            iter_pos,
             callback_fn,
             other_data
-        ))?;
+        ))
+        .map(|_| val)
+    }
+
+    /// Visits all objects in the group using default iteration/traversal order.
+    pub fn iter_visit_default<F, G>(&self, val: G, op: F) -> Result<G>
+    where
+        F: Fn(&Self, &str, LinkInfo, &mut G) -> bool,
+    {
+        self.iter_visit(IterationOrder::default(), TraversalOrder::default(), val, op)
+    }
+
+    fn get_all_of_type(&self, loc_type: LocationType) -> Result<Vec<Location>> {
+        self.iter_visit_default(vec![], |group, name, _info, objects| {
+            if let Ok(info) = group.get_info_by_name(name) {
+                if info.loc_type == loc_type {
+                    if let Ok(loc) = group.open_by_token(info.token) {
+                        objects.push(loc);
+                        return true; // ok, object extracted and pushed
+                    }
+                } else {
+                    return true; // ok, object is of another type, skipped
+                }
+            }
+            false // an error occured somewhere along the way
+        })
+    }
+
+    /// Returns all groups in the group, non-recursively
+    pub fn groups(&self) -> Result<Vec<Self>> {
+        self.get_all_of_type(LocationType::Group)
+            .map(|vec| vec.into_iter().map(|obj| unsafe { obj.cast() }).collect())
+    }
 
-        Ok(result)
+    /// Returns all datasets in the group, non-recursively
+    pub fn datasets(&self) -> Result<Vec<Dataset>> {
+        self.get_all_of_type(LocationType::Dataset)
+            .map(|vec| vec.into_iter().map(|obj| unsafe { obj.cast() }).collect())
+    }
+
+    /// Returns all named types in the group, non-recursively
+    pub fn named_datatypes(&self) -> Result<Vec<Datatype>> {
+        self.get_all_of_type(LocationType::NamedDatatype)
+            .map(|vec| vec.into_iter().map(|obj| unsafe { obj.cast() }).collect())
+    }
+
+    /// Returns the names of all objects in the group, non-recursively.
+    pub fn member_names(&self) -> Result<Vec<String>> {
+        self.iter_visit_default(vec![], |_, name, _, names| {
+            names.push(name.to_owned());
+            true
+        })
     }
 }
 
@@ -471,4 +629,32 @@ pub mod tests {
             assert_eq!(dset2.read_scalar::<i32>().unwrap(), 13);
         })
     }
+
+    #[test]
+    pub fn test_iterators() {
+        with_tmp_file(|file| {
+            file.create_group("a").unwrap();
+            file.create_group("b").unwrap();
+            let group_a = file.group("a").unwrap();
+            let _group_b = file.group("b").unwrap();
+            file.new_dataset::<u32>().shape((10, 20)).create("a/foo").unwrap();
+            file.new_dataset::<u32>().shape((10, 20)).create("a/123").unwrap();
+            file.new_dataset::<u32>().shape((10, 20)).create("a/bar").unwrap();
+
+            let groups = file.groups().unwrap();
+            assert_eq!(groups.len(), 2);
+            for group in groups {
+                assert!(matches!(group.name().as_ref(), "/a" | "/b"));
+            }
+
+            let datasets = file.datasets().unwrap();
+            assert_eq!(datasets.len(), 0);
+
+            let datasets = group_a.datasets().unwrap();
+            assert_eq!(datasets.len(), 3);
+            for dataset in datasets {
+                assert!(matches!(dataset.name().as_ref(), "/a/foo" | "/a/123" | "/a/bar"));
+            }
+        })
+    }
 }