1use std::collections::BTreeMap;
4
5use crate::error::Result;
6use crate::index::{
7 Index, IndexEntry, MODE_EXECUTABLE, MODE_GITLINK, MODE_REGULAR, MODE_SYMLINK, MODE_TREE,
8};
9use crate::objects::{serialize_tree, tree_entry_cmp, ObjectId, ObjectKind, TreeEntry};
10use crate::odb::Odb;
11
12pub fn write_tree_from_index(odb: &Odb, index: &Index, prefix: &str) -> Result<ObjectId> {
16 let prefix_bytes = prefix.as_bytes();
17 let mut entries: Vec<&IndexEntry> = index
18 .entries
19 .iter()
20 .filter(|entry| entry.stage() == 0 && entry.path.starts_with(prefix_bytes))
21 .collect();
22 entries.sort_by(|a, b| a.path.cmp(&b.path).then_with(|| a.stage().cmp(&b.stage())));
23 build_tree(odb, &entries, prefix_bytes)
24}
25
26fn build_tree(odb: &Odb, entries: &[&IndexEntry], dir_prefix: &[u8]) -> Result<ObjectId> {
27 let mut children: BTreeMap<Vec<u8>, ChildKind> = BTreeMap::new();
28
29 for entry in entries {
30 let path = &entry.path;
31 let rel = if dir_prefix.is_empty() {
32 path.as_slice()
33 } else {
34 path.strip_prefix(dir_prefix)
35 .and_then(|suffix| suffix.strip_prefix(b"/"))
36 .unwrap_or(path.as_slice())
37 };
38
39 if let Some(slash_pos) = rel.iter().position(|&byte| byte == b'/') {
40 let child_name = rel[..slash_pos].to_vec();
41 let sub_prefix = if dir_prefix.is_empty() {
42 child_name.clone()
43 } else {
44 let mut sub_prefix = dir_prefix.to_vec();
45 sub_prefix.push(b'/');
46 sub_prefix.extend_from_slice(&child_name);
47 sub_prefix
48 };
49 children
50 .entry(child_name)
51 .or_insert_with(|| ChildKind::Tree(sub_prefix, Vec::new()))
52 .push_entry(entry);
53 } else {
54 children
55 .entry(rel.to_vec())
56 .or_insert_with(|| ChildKind::Blob {
57 mode: canonicalize_blob_mode(entry.mode),
58 oid: entry.oid,
59 });
60 }
61 }
62
63 let mut tree_entries = Vec::with_capacity(children.len());
64 for (name, child) in children {
65 match child {
66 ChildKind::Blob { mode, oid } => tree_entries.push(TreeEntry { mode, name, oid }),
67 ChildKind::Tree(sub_prefix, sub_entries) => {
68 let sub_oid = build_tree(odb, &sub_entries, &sub_prefix)?;
69 tree_entries.push(TreeEntry {
70 mode: MODE_TREE,
71 name,
72 oid: sub_oid,
73 });
74 }
75 }
76 }
77
78 tree_entries.sort_by(|a, b| {
79 let a_tree = a.mode == MODE_TREE;
80 let b_tree = b.mode == MODE_TREE;
81 tree_entry_cmp(&a.name, a_tree, &b.name, b_tree)
82 });
83
84 let data = serialize_tree(&tree_entries);
85 odb.write(ObjectKind::Tree, &data)
86}
87
88fn canonicalize_blob_mode(mode: u32) -> u32 {
89 match mode & 0o170000 {
90 0o120000 => MODE_SYMLINK,
91 0o160000 => MODE_GITLINK,
92 0o100000 => {
93 if mode & 0o111 != 0 {
94 MODE_EXECUTABLE
95 } else {
96 MODE_REGULAR
97 }
98 }
99 _ => MODE_REGULAR,
100 }
101}
102
103enum ChildKind<'a> {
104 Blob { mode: u32, oid: ObjectId },
105 Tree(Vec<u8>, Vec<&'a IndexEntry>),
106}
107
108impl<'a> ChildKind<'a> {
109 fn push_entry(&mut self, entry: &'a IndexEntry) {
110 if let Self::Tree(_, entries) = self {
111 entries.push(entry);
112 }
113 }
114}
115
116#[cfg(test)]
117mod tests {
118 #![allow(clippy::expect_used, clippy::unwrap_used)]
119
120 use super::*;
121 use crate::index::{IndexEntry, MODE_EXECUTABLE, MODE_REGULAR, MODE_SYMLINK, MODE_TREE};
122 use crate::objects::parse_tree;
123 use tempfile::TempDir;
124
125 fn entry(path: &str, mode: u32, oid: ObjectId) -> IndexEntry {
126 IndexEntry {
127 ctime_sec: 0,
128 ctime_nsec: 0,
129 mtime_sec: 0,
130 mtime_nsec: 0,
131 dev: 0,
132 ino: 0,
133 mode,
134 uid: 0,
135 gid: 0,
136 size: 0,
137 oid,
138 flags: path.len().min(0xFFF) as u16,
139 flags_extended: None,
140 path: path.as_bytes().to_vec(),
141 }
142 }
143
144 #[test]
145 fn writes_sorted_tree_with_canonical_modes() {
146 let temp_dir = TempDir::new().unwrap();
147 let odb = Odb::new(temp_dir.path());
148
149 let oid_a = odb.write(ObjectKind::Blob, b"a").unwrap();
150 let oid_exec = odb.write(ObjectKind::Blob, b"exec").unwrap();
151 let oid_link = odb.write(ObjectKind::Blob, b"target").unwrap();
152
153 let mut index = Index::new();
154 index.add_or_replace(entry("bin/run.sh", 0o100777, oid_exec));
155 index.add_or_replace(entry("link", 0o120777, oid_link));
156 index.add_or_replace(entry("a.txt", 0o100664, oid_a));
157
158 let root_oid = write_tree_from_index(&odb, &index, "").unwrap();
159 let root_tree_obj = odb.read(&root_oid).unwrap();
160 let root_entries = parse_tree(&root_tree_obj.data).unwrap();
161
162 assert_eq!(root_entries.len(), 3);
163 assert_eq!(root_entries[0].name, b"a.txt");
164 assert_eq!(root_entries[0].mode, MODE_REGULAR);
165 assert_eq!(root_entries[1].name, b"bin");
166 assert_eq!(root_entries[1].mode, MODE_TREE);
167 assert_eq!(root_entries[2].name, b"link");
168 assert_eq!(root_entries[2].mode, MODE_SYMLINK);
169
170 let bin_tree_obj = odb.read(&root_entries[1].oid).unwrap();
171 let bin_entries = parse_tree(&bin_tree_obj.data).unwrap();
172 assert_eq!(bin_entries.len(), 1);
173 assert_eq!(bin_entries[0].name, b"run.sh");
174 assert_eq!(bin_entries[0].mode, MODE_EXECUTABLE);
175 }
176}