1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

//! Contains a [`Tree`] type that is used to process the dotted package names into
//! directory structured files.

use std::{
    collections::HashMap,
    ffi::{OsStr, OsString},
    fmt::{Debug, Display},
    iter::FromIterator,
    path::{Path, PathBuf},
};

use anyhow::{Context, Result};

#[derive(Default, Debug, PartialEq)]
pub struct Tree(pub(crate) HashMap<PathBuf, Tree>);

impl Extend<PathBuf> for Tree {
    fn extend<T: IntoIterator<Item = PathBuf>>(&mut self, iter: T) {
        for path in iter {
            self.insert_path(path)
        }
    }
}

impl FromIterator<PathBuf> for Tree {
    fn from_iter<T: IntoIterator<Item = PathBuf>>(iter: T) -> Self {
        let mut tree = Tree::default();
        tree.extend(iter);
        tree
    }
}

impl Tree {
    /// Given a file path that is `.` separated, it loads it into the tree.
    pub fn insert_path(mut self: &mut Self, path: PathBuf) {
        for comp in path.file_stem().unwrap().to_str().unwrap().split('.') {
            self = (*self).0.entry(PathBuf::from(comp)).or_default()
        }
    }

    /// Loop through the tree, determining where all the files should be
    /// and moving them there
    pub fn move_paths(&self, root: &Path, filename: OsString, output: PathBuf) -> Result<()> {
        if self.0.is_empty() {
            std::fs::create_dir_all(root.join(&output).parent().unwrap())
                .with_context(|| format!("could not create dir for file {}", output.display()))?;

            let from = root.join(filename.add("rs"));
            let to = root.join(output.with_extension("rs"));
            std::fs::rename(&from, &to).with_context(|| {
                format!("could not move {} to {}", from.display(), to.display())
            })?;
        } else {
            for (k, tree) in &self.0 {
                tree.move_paths(root, filename.add(k), output.join(k))?;
            }
        }
        Ok(())
    }
}

// private helper trait
trait OsStrExt {
    fn add(&self, add: impl AsRef<OsStr>) -> OsString;
}

impl OsStrExt for OsStr {
    /// Adds `add` to the [`OsStr`], returning a new [`OsString`]. If there already exists data in the string,
    /// this puts a `.` separator inbetween
    fn add(&self, add: impl AsRef<OsStr>) -> OsString {
        let mut _self = self.to_owned();
        if !_self.is_empty() {
            _self.push(".");
        }
        _self.push(add);
        _self
    }
}

impl Display for Tree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for (k, tree) in &self.0 {
            write!(f, "pub mod {}", k.display())?;
            if tree.0.is_empty() {
                write!(f, ";")?;
            } else {
                write!(f, "{{{}}}", tree)?;
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::path::PathBuf;

    use super::Tree;

    macro_rules! tree {
        ($($key:literal : $val:expr,)*) => {
            Tree(HashMap::from_iter([
                $(
                    (PathBuf::from($key), $val)
                ),*
            ]))
        };
    }

    #[test]
    fn build_tree() {
        let tree: Tree = [
            "grpc_build.client.helloworld.rs",
            "grpc_build.request.helloworld.rs",
            "grpc_build.response.helloworld.rs",
            "google.protobuf.foo.rs",
            "google.protobuf.bar.rs",
        ]
        .into_iter()
        .map(PathBuf::from)
        .collect();

        let expected = tree! {
            "grpc_build": tree! {
                "client": tree! {
                    "helloworld": tree!{},
                },
                "request": tree! {
                    "helloworld": tree!{},
                },
                "response": tree! {
                    "helloworld": tree!{},
                },
            },
            "google": tree! {
                "protobuf": tree! {
                    "foo": tree!{},
                    "bar": tree!{},
                },
            },
        };

        assert_eq!(tree, expected);
    }
}