clawhdf5/crates/clawhdf5-format/src/writer_tree.rs

//! The group hierarchy `FileWriter` writes: builders flattened into a tree
//! of groups, datasets and links, with path names expanded into
//! intermediate groups, hard links resolved to objects, reference counts
//! counted, and everything put in layout order.

#[cfg(not(feature = "std"))]
use alloc::{
    collections::BTreeMap,
    format,
    string::{String, ToString},
    vec,
    vec::Vec,
};
#[cfg(feature = "std")]
use std::collections::BTreeMap;

use crate::error::FormatError;
use crate::type_builders::{AttrValue, DatasetBuilder, GroupBuilder, GroupItem};

/// Depth of the chain of unresolved hard links followed while resolving one
/// hard-link target path (a bound on recursion; cycles are found exactly).
const MAX_LINK_DEPTH: usize = 64;

fn err(msg: String) -> FormatError {
    FormatError::SerializationError(msg)
}

/// A link name must be one path component: not empty, not ".", and without
/// '/' (a '/' separates components, so it cannot be part of a name).
fn check_link_name(name: &str, path: &str) -> Result<(), FormatError> {
    if name.is_empty() || name == "." || name.contains('/') {
        return Err(err(format!(
            "invalid object name {path:?}: every path component must be a \
             non-empty name other than \".\""
        )));
    }
    Ok(())
}

/// What a link in the final tree points at.
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum LinkTo {
    /// A group, by index into [`Tree::groups`] (layout order).
    Group(usize),
    /// A dataset, by index into [`Tree::datasets`] (layout order).
    Dataset(usize),
    Soft(String),
    External {
        file: String,
        path: String,
    },
}

pub(crate) struct Link {
    pub(crate) name: String,
    pub(crate) to: LinkTo,
    /// Set when the group tracks creation order.
    pub(crate) creation_order: Option<u64>,
}

pub(crate) struct Group {
    pub(crate) attrs: Vec<(String, AttrValue)>,
    /// Links in the order they are written.
    pub(crate) links: Vec<Link>,
    pub(crate) track_order: bool,
    /// Number of hard links to this group (the root counts one for the
    /// superblock's reference).
    pub(crate) refcount: u32,
}

/// The flattened file: groups (root first) and datasets, both in the order
/// they are laid out in the file.
pub(crate) struct Tree {
    pub(crate) groups: Vec<Group>,
    pub(crate) datasets: Vec<(DatasetBuilder, u32)>,
}

// ---- construction ----

enum Target {
    Group(usize),
    Dataset(usize),
    Soft(String),
    Hard(String),
    External { file: String, path: String },
}

struct BuildGroup {
    /// Full path, for messages.
    path: String,
    attrs: Vec<(String, AttrValue)>,
    links: Vec<(String, Target)>,
    by_name: BTreeMap<String, usize>,
    track_order: Option<bool>,
}

struct Builder {
    groups: Vec<BuildGroup>,
    datasets: Vec<DatasetBuilder>,
}

fn join(parent: &str, name: &str) -> String {
    if parent == "/" {
        format!("/{name}")
    } else {
        format!("{parent}/{name}")
    }
}

impl Builder {
    fn new_group(&mut self, path: String) -> usize {
        self.groups.push(BuildGroup {
            path,
            attrs: Vec::new(),
            links: Vec::new(),
            by_name: BTreeMap::new(),
            track_order: None,
        });
        self.groups.len() - 1
    }

    /// Split `path` (relative to group `g`) into the group holding its last
    /// component, creating missing intermediate groups, and that component.
    fn parent_of<'p>(&mut self, g: usize, path: &'p str) -> Result<(usize, &'p str), FormatError> {
        // An absolute path is accepted at the root only.
        let rel = match path.strip_prefix('/') {
            Some(rest) if g == 0 => rest,
            Some(_) => {
                return Err(err(format!(
                    "invalid object name {path:?} in {}: absolute paths are accepted \
                     only at the root",
                    self.groups[g].path
                )));
            }
            None => path,
        };
        let mut comps: Vec<&str> = rel.split('/').collect();
        let last = comps.pop().unwrap_or("");
        check_link_name(last, path)?;
        let mut cur = g;
        for c in comps {
            check_link_name(c, path)?;
            cur = match self.groups[cur].by_name.get(c).copied() {
                Some(i) => match self.groups[cur].links[i].1 {
                    Target::Group(child) => child,
                    _ => {
                        return Err(err(format!(
                            "cannot create {path:?} in {}: {c:?} exists and is not a group",
                            self.groups[g].path
                        )));
                    }
                },
                None => {
                    let child = self.new_group(join(&self.groups[cur].path, c));
                    self.push_link(cur, c, Target::Group(child))?;
                    child
                }
            };
        }
        Ok((cur, last))
    }

    fn push_link(&mut self, g: usize, name: &str, to: Target) -> Result<(), FormatError> {
        let grp = &mut self.groups[g];
        if grp.by_name.contains_key(name) {
            return Err(err(format!("{:?} already exists", join(&grp.path, name))));
        }
        grp.by_name.insert(name.to_string(), grp.links.len());
        grp.links.push((name.to_string(), to));
        Ok(())
    }

    /// Add `item` to group `g`.
    fn add_item(&mut self, g: usize, item: GroupItem) -> Result<(), FormatError> {
        match item {
            GroupItem::Dataset(db) => {
                let (parent, name) = self.parent_of(g, &db.name)?;
                let name = name.to_string();
                self.push_link(parent, &name, Target::Dataset(self.datasets.len()))?;
                self.datasets.push(*db);
            }
            GroupItem::Group(gb) => self.add_group(g, gb)?,
            GroupItem::Soft { name, target } => {
                if target.is_empty() {
                    return Err(err(format!("soft link {name:?} has an empty target")));
                }
                let (parent, last) = self.parent_of(g, &name)?;
                self.push_link(parent, last, Target::Soft(target))?;
            }
            GroupItem::Hard { name, target } => {
                let (parent, last) = self.parent_of(g, &name)?;
                self.push_link(parent, last, Target::Hard(target))?;
            }
            GroupItem::External { name, file, path } => {
                if file.is_empty() || path.is_empty() {
                    return Err(err(format!(
                        "external link {name:?} needs a file name and an object path"
                    )));
                }
                let (parent, last) = self.parent_of(g, &name)?;
                self.push_link(parent, last, Target::External { file, path })?;
            }
        }
        Ok(())
    }

    /// Add the group `gb` (named by a path relative to group `g`), merging it
    /// into a group already at that path.
    fn add_group(&mut self, g: usize, gb: GroupBuilder) -> Result<(), FormatError> {
        let (parent, last) = self.parent_of(g, &gb.name)?;
        let idx = match self.groups[parent].by_name.get(last).copied() {
            Some(i) => match self.groups[parent].links[i].1 {
                Target::Group(child) => child,
                _ => {
                    return Err(err(format!(
                        "{:?} already exists and is not a group",
                        join(&self.groups[parent].path, last)
                    )));
                }
            },
            None => {
                let child = self.new_group(join(&self.groups[parent].path, last));
                self.push_link(parent, last, Target::Group(child))?;
                child
            }
        };
        self.merge_into(idx, gb)
    }

    /// Merge a builder's attributes, setting and items into group `idx`.
    fn merge_into(&mut self, idx: usize, gb: GroupBuilder) -> Result<(), FormatError> {
        // An attribute set again (by this builder or a merged one) takes the
        // new value, as assigning `attrs[name]` in h5py does.
        for (name, value) in gb.attrs {
            let attrs = &mut self.groups[idx].attrs;
            match attrs.iter_mut().find(|(n, _)| *n == name) {
                Some(slot) => slot.1 = value,
                None => attrs.push((name, value)),
            }
        }
        if let Some(t) = gb.track_order {
            match self.groups[idx].track_order {
                Some(old) if old != t => {
                    return Err(err(format!(
                        "conflicting track_order settings for {}",
                        self.groups[idx].path
                    )));
                }
                _ => self.groups[idx].track_order = Some(t),
            }
        }
        for item in gb.items {
            self.add_item(idx, item)?;
        }
        Ok(())
    }

    /// The object a hard link's `target` path names, from group `from`.
    ///
    /// Hard links met on the way are resolved once and remembered in
    /// `memo` (by group and link index), so a target that goes through
    /// other hard links costs time linear in the links, not exponential; a
    /// hard link met again while it is being resolved is a cycle.
    fn resolve(
        &self,
        memo: &mut [Vec<Resolution>],
        from: usize,
        target: &str,
        depth: usize,
    ) -> Result<Obj, FormatError> {
        if depth > MAX_LINK_DEPTH {
            return Err(err(format!(
                "hard link target {target:?}: more than {MAX_LINK_DEPTH} hard links \
                 to follow"
            )));
        }
        let (mut cur, rest) = match target.strip_prefix('/') {
            Some(rest) => (0, rest),
            None => (from, target),
        };
        if target.is_empty() {
            return Err(err("a hard link needs a target path".to_string()));
        }
        let comps: Vec<&str> = rest
            .split('/')
            .filter(|c| !c.is_empty() && *c != ".")
            .collect();
        let mut obj = Obj::Group(cur);
        for (i, c) in comps.iter().enumerate() {
            let Obj::Group(g) = obj else {
                return Err(err(format!(
                    "hard link target {target:?}: {:?} is not a group",
                    comps[..i].join("/")
                )));
            };
            cur = g;
            let grp = &self.groups[cur];
            let Some(&li) = grp.by_name.get(*c) else {
                return Err(err(format!(
                    "hard link target {target:?} does not exist in the file"
                )));
            };
            obj = match &grp.links[li].1 {
                Target::Group(child) => Obj::Group(*child),
                Target::Dataset(d) => Obj::Dataset(*d),
                Target::Hard(p) => match memo[cur][li] {
                    Resolution::Done(o) => o,
                    Resolution::InProgress => {
                        return Err(err(format!(
                            "hard link target {target:?}: the hard link {:?} leads \
                             back to itself (a cycle)",
                            join(&grp.path, c)
                        )));
                    }
                    Resolution::Todo => {
                        memo[cur][li] = Resolution::InProgress;
                        let o = self.resolve(memo, cur, p, depth + 1)?;
                        memo[cur][li] = Resolution::Done(o);
                        o
                    }
                },
                Target::Soft(_) | Target::External { .. } => {
                    return Err(err(format!(
                        "hard link target {target:?} goes through a soft or external \
                         link ({:?}); name the object by its hard-link path",
                        join(&grp.path, c)
                    )));
                }
            };
        }
        Ok(obj)
    }
}

/// Where resolving one hard link has got to.
#[derive(Clone, Copy)]
enum Resolution {
    Todo,
    InProgress,
    Done(Obj),
}

#[derive(Clone, Copy)]
enum Obj {
    Group(usize),
    Dataset(usize),
}

/// Flatten the root group builder into a [`Tree`]. `default_track_order`
/// applies to every group that does not set its own.
pub(crate) fn build(root: GroupBuilder, default_track_order: bool) -> Result<Tree, FormatError> {
    let mut b = Builder {
        groups: Vec::new(),
        datasets: Vec::new(),
    };
    b.new_group("/".to_string());
    b.merge_into(0, root)?;

    // Resolve hard links and count references.
    let mut group_refs = vec![0u32; b.groups.len()];
    let mut ds_refs = vec![0u32; b.datasets.len()];
    group_refs[0] = 1; // the superblock's reference to the root
    let mut memo: Vec<Vec<Resolution>> = b
        .groups
        .iter()
        .map(|g| vec![Resolution::Todo; g.links.len()])
        .collect();
    let mut resolved: Vec<Vec<Option<Obj>>> = Vec::with_capacity(b.groups.len());
    for (gi, g) in b.groups.iter().enumerate() {
        let mut row = Vec::with_capacity(g.links.len());
        for (li, (_, t)) in g.links.iter().enumerate() {
            let obj = match t {
                Target::Group(i) => Some(Obj::Group(*i)),
                Target::Dataset(d) => Some(Obj::Dataset(*d)),
                Target::Hard(p) => Some(match memo[gi][li] {
                    Resolution::Done(o) => o,
                    _ => {
                        memo[gi][li] = Resolution::InProgress;
                        let o = b.resolve(&mut memo, gi, p, 0)?;
                        memo[gi][li] = Resolution::Done(o);
                        o
                    }
                }),
                Target::Soft(_) | Target::External { .. } => None,
            };
            match obj {
                Some(Obj::Group(i)) => group_refs[i] += 1,
                Some(Obj::Dataset(d)) => ds_refs[d] += 1,
                None => {}
            }
            row.push(obj);
        }
        resolved.push(row);
    }

    // The order each group's links are written in: creation order when
    // tracked; otherwise datasets, then groups, then other links (the order
    // earlier versions wrote, so one-level files keep their layout).
    let tracked: Vec<bool> = b
        .groups
        .iter()
        .map(|g| g.track_order.unwrap_or(default_track_order))
        .collect();
    let link_order: Vec<Vec<usize>> = b
        .groups
        .iter()
        .enumerate()
        .map(|(gi, g)| {
            let mut idx: Vec<usize> = (0..g.links.len()).collect();
            if !tracked[gi] {
                idx.sort_by_key(|&i| match g.links[i].1 {
                    Target::Dataset(_) => 0,
                    Target::Group(_) => 1,
                    _ => 2,
                });
            }
            idx
        })
        .collect();

    // Layout order: groups depth-first from the root, following the links
    // that created them; datasets group by group in that order.
    let mut group_order = Vec::with_capacity(b.groups.len());
    let mut stack = vec![0usize];
    while let Some(g) = stack.pop() {
        group_order.push(g);
        let children: Vec<usize> = link_order[g]
            .iter()
            .filter_map(|&i| match b.groups[g].links[i].1 {
                Target::Group(c) => Some(c),
                _ => None,
            })
            .collect();
        stack.extend(children.into_iter().rev());
    }
    let mut ds_order = Vec::with_capacity(b.datasets.len());
    for &g in &group_order {
        for &i in &link_order[g] {
            if let Target::Dataset(d) = b.groups[g].links[i].1 {
                ds_order.push(d);
            }
        }
    }
    let mut group_pos = vec![0usize; b.groups.len()];
    for (pos, &g) in group_order.iter().enumerate() {
        group_pos[g] = pos;
    }
    let mut ds_pos = vec![0usize; b.datasets.len()];
    for (pos, &d) in ds_order.iter().enumerate() {
        ds_pos[d] = pos;
    }

    let mut groups_by_id: Vec<Option<BuildGroup>> = b.groups.into_iter().map(Some).collect();
    let mut groups = Vec::with_capacity(group_order.len());
    for &g in &group_order {
        let bg = groups_by_id[g].take().expect("each group is laid out once");
        let mut targets: Vec<Option<(String, Target)>> = bg.links.into_iter().map(Some).collect();
        let links = link_order[g]
            .iter()
            .map(|&i| {
                let (name, t) = targets[i].take().expect("each link is written once");
                let to = match (resolved[g][i], t) {
                    (Some(Obj::Group(c)), _) => LinkTo::Group(group_pos[c]),
                    (Some(Obj::Dataset(d)), _) => LinkTo::Dataset(ds_pos[d]),
                    (None, Target::Soft(s)) => LinkTo::Soft(s),
                    (None, Target::External { file, path }) => LinkTo::External { file, path },
                    (None, _) => unreachable!("hard links are resolved"),
                };
                Link {
                    name,
                    to,
                    creation_order: tracked[g].then_some(i as u64),
                }
            })
            .collect();
        groups.push(Group {
            attrs: bg.attrs,
            links,
            track_order: tracked[g],
            refcount: group_refs[g],
        });
    }
    let mut ds_by_id: Vec<Option<DatasetBuilder>> = b.datasets.into_iter().map(Some).collect();
    let datasets = ds_order
        .iter()
        .map(|&d| {
            (
                ds_by_id[d].take().expect("each dataset is laid out once"),
                ds_refs[d],
            )
        })
        .collect();
    Ok(Tree { groups, datasets })
}