grit_lib/
write_tree.rs

1//! Build tree objects from index entries (`git write-tree` core logic).
2
3use 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::{parse_tree, serialize_tree, tree_entry_cmp, ObjectId, ObjectKind, TreeEntry};
10use crate::odb::Odb;
11
12fn ensure_empty_blob_for_intent_to_add(odb: &Odb, index: &Index) -> Result<()> {
13    if index
14        .entries
15        .iter()
16        .any(|e| e.stage() == 0 && e.intent_to_add())
17    {
18        let _ = odb.write(ObjectKind::Blob, b"")?;
19    }
20    Ok(())
21}
22
23/// Build and write tree object(s) from index entries and return the tree OID.
24///
25/// The `prefix` argument optionally limits the write to a subtree path.
26/// Like [`write_tree_from_index`], but only index entries whose path is listed in `paths`
27/// (repository-relative, as stored in the index) are included in the tree.
28pub fn write_tree_from_index_subset(
29    odb: &Odb,
30    index: &Index,
31    paths: &std::collections::HashSet<Vec<u8>>,
32) -> Result<ObjectId> {
33    ensure_empty_blob_for_intent_to_add(odb, index)?;
34
35    let mut entries: Vec<&IndexEntry> = index
36        .entries
37        .iter()
38        .filter(|entry| {
39            entry.stage() == 0
40                && !entry.intent_to_add()
41                && entry.mode != MODE_TREE
42                && paths.contains(&entry.path)
43        })
44        .collect();
45    entries.sort_by(|a, b| a.path.cmp(&b.path).then_with(|| a.stage().cmp(&b.stage())));
46    build_tree(odb, &entries, b"")
47}
48
49/// Build and write tree object(s) from index entries and return the tree OID.
50pub fn write_tree_from_index(odb: &Odb, index: &Index, prefix: &str) -> Result<ObjectId> {
51    ensure_empty_blob_for_intent_to_add(odb, index)?;
52
53    let prefix_bytes = prefix.as_bytes();
54    let mut entries: Vec<&IndexEntry> = index
55        .entries
56        .iter()
57        .filter(|entry| {
58            entry.stage() == 0
59                && !entry.intent_to_add()
60                && entry.mode != MODE_TREE
61                && entry.path.starts_with(prefix_bytes)
62        })
63        .collect();
64    entries.sort_by(|a, b| a.path.cmp(&b.path).then_with(|| a.stage().cmp(&b.stage())));
65    build_tree(odb, &entries, prefix_bytes)
66}
67
68fn build_tree(odb: &Odb, entries: &[&IndexEntry], dir_prefix: &[u8]) -> Result<ObjectId> {
69    let mut children: BTreeMap<Vec<u8>, ChildKind> = BTreeMap::new();
70
71    for entry in entries {
72        let path = &entry.path;
73        let rel = if dir_prefix.is_empty() {
74            path.as_slice()
75        } else {
76            path.strip_prefix(dir_prefix)
77                .and_then(|suffix| suffix.strip_prefix(b"/"))
78                .unwrap_or(path.as_slice())
79        };
80
81        if let Some(slash_pos) = rel.iter().position(|&byte| byte == b'/') {
82            let child_name = rel[..slash_pos].to_vec();
83            let sub_prefix = if dir_prefix.is_empty() {
84                child_name.clone()
85            } else {
86                let mut sub_prefix = dir_prefix.to_vec();
87                sub_prefix.push(b'/');
88                sub_prefix.extend_from_slice(&child_name);
89                sub_prefix
90            };
91            children
92                .entry(child_name)
93                .or_insert_with(|| ChildKind::Tree(sub_prefix, Vec::new()))
94                .push_entry(entry);
95        } else {
96            children
97                .entry(rel.to_vec())
98                .or_insert_with(|| ChildKind::Blob {
99                    mode: canonicalize_blob_mode(entry.mode),
100                    oid: entry.oid,
101                });
102        }
103    }
104
105    let mut tree_entries = Vec::with_capacity(children.len());
106    for (name, child) in children {
107        match child {
108            ChildKind::Blob { mode, oid } => tree_entries.push(TreeEntry { mode, name, oid }),
109            ChildKind::Tree(sub_prefix, sub_entries) => {
110                let sub_oid = build_tree(odb, &sub_entries, &sub_prefix)?;
111                tree_entries.push(TreeEntry {
112                    mode: MODE_TREE,
113                    name,
114                    oid: sub_oid,
115                });
116            }
117        }
118    }
119
120    tree_entries.sort_by(|a, b| {
121        let a_tree = a.mode == MODE_TREE;
122        let b_tree = b.mode == MODE_TREE;
123        tree_entry_cmp(&a.name, a_tree, &b.name, b_tree)
124    });
125
126    let data = serialize_tree(&tree_entries);
127    odb.write(ObjectKind::Tree, &data)
128}
129
130/// Build a tree for a **partial** commit: paths listed in `paths_from_index` (repository-relative,
131/// UTF-8 path bytes) are taken from `index`; every other path is copied from `base_tree_oid`
132/// (typically `HEAD^{tree}`).
133///
134/// This matches Git's behaviour when committing with pathspecs while the index contains additional
135/// staged paths: the commit tree merges `HEAD` with only the pathspec-selected index updates.
136pub fn write_tree_partial_from_index(
137    odb: &Odb,
138    index: &Index,
139    base_tree_oid: &ObjectId,
140    paths_from_index: &std::collections::HashSet<Vec<u8>>,
141) -> Result<ObjectId> {
142    let _ = odb.write(ObjectKind::Blob, b"");
143
144    fn full_path(prefix: &[u8], name: &[u8]) -> Vec<u8> {
145        if prefix.is_empty() {
146            name.to_vec()
147        } else {
148            let mut p = prefix.to_vec();
149            p.push(b'/');
150            p.extend_from_slice(name);
151            p
152        }
153    }
154
155    fn subtree_affected(paths_from_index: &std::collections::HashSet<Vec<u8>>, dir: &[u8]) -> bool {
156        paths_from_index
157            .iter()
158            .any(|p| p == dir || (p.starts_with(dir) && p.get(dir.len()) == Some(&b'/')))
159    }
160
161    fn merge_level(
162        odb: &Odb,
163        index: &Index,
164        base_tree_oid: &ObjectId,
165        prefix: &[u8],
166        paths_from_index: &std::collections::HashSet<Vec<u8>>,
167    ) -> Result<ObjectId> {
168        let base_obj = odb.read(base_tree_oid)?;
169        let base_entries = parse_tree(&base_obj.data)?;
170
171        let mut by_name: BTreeMap<Vec<u8>, TreeEntry> = BTreeMap::new();
172        for te in base_entries {
173            let fp = full_path(prefix, &te.name);
174            if !subtree_affected(paths_from_index, &fp) {
175                by_name.insert(te.name.clone(), te);
176            } else if te.mode == MODE_TREE {
177                let sub_oid = merge_level(odb, index, &te.oid, &fp, paths_from_index)?;
178                by_name.insert(
179                    te.name.clone(),
180                    TreeEntry {
181                        mode: MODE_TREE,
182                        name: te.name,
183                        oid: sub_oid,
184                    },
185                );
186            } else if paths_from_index.contains(&fp) {
187                if let Some(ie) = index.entries.iter().find(|e| {
188                    e.stage() == 0 && !e.intent_to_add() && e.mode != MODE_TREE && e.path == fp
189                }) {
190                    by_name.insert(
191                        te.name.clone(),
192                        TreeEntry {
193                            mode: canonicalize_blob_mode(ie.mode),
194                            name: te.name,
195                            oid: ie.oid,
196                        },
197                    );
198                }
199                // No index entry: path was removed — omit from the merged tree.
200            } else {
201                by_name.insert(te.name.clone(), te);
202            }
203        }
204
205        for ie in &index.entries {
206            if ie.stage() != 0 || ie.intent_to_add() || ie.mode == MODE_TREE {
207                continue;
208            }
209            if !paths_from_index.contains(&ie.path) {
210                continue;
211            }
212            let rel = if prefix.is_empty() {
213                ie.path.as_slice()
214            } else if ie.path.starts_with(prefix) && ie.path.get(prefix.len()) == Some(&b'/') {
215                &ie.path[prefix.len() + 1..]
216            } else {
217                continue;
218            };
219            if rel.is_empty() {
220                continue;
221            }
222            if let Some(slash) = rel.iter().position(|&b| b == b'/') {
223                let dir_name = rel[..slash].to_vec();
224                if by_name.contains_key(&dir_name) {
225                    continue;
226                }
227                let sub_prefix = full_path(prefix, &dir_name);
228                let sub_oid =
229                    write_tree_from_index(odb, index, &String::from_utf8_lossy(&sub_prefix))?;
230                by_name.insert(
231                    dir_name.clone(),
232                    TreeEntry {
233                        mode: MODE_TREE,
234                        name: dir_name,
235                        oid: sub_oid,
236                    },
237                );
238            } else {
239                let name = rel.to_vec();
240                if !by_name.contains_key(&name) {
241                    by_name.insert(
242                        name.clone(),
243                        TreeEntry {
244                            mode: canonicalize_blob_mode(ie.mode),
245                            name,
246                            oid: ie.oid,
247                        },
248                    );
249                }
250            }
251        }
252
253        let mut out: Vec<TreeEntry> = by_name.into_values().collect();
254        out.sort_by(|a, b| {
255            let a_tree = a.mode == MODE_TREE;
256            let b_tree = b.mode == MODE_TREE;
257            tree_entry_cmp(&a.name, a_tree, &b.name, b_tree)
258        });
259        let data = serialize_tree(&out);
260        odb.write(ObjectKind::Tree, &data)
261    }
262
263    merge_level(odb, index, base_tree_oid, b"", paths_from_index)
264}
265
266fn canonicalize_blob_mode(mode: u32) -> u32 {
267    match mode & 0o170000 {
268        0o120000 => MODE_SYMLINK,
269        0o160000 => MODE_GITLINK,
270        0o100000 => {
271            if mode & 0o111 != 0 {
272                MODE_EXECUTABLE
273            } else {
274                MODE_REGULAR
275            }
276        }
277        _ => MODE_REGULAR,
278    }
279}
280
281enum ChildKind<'a> {
282    Blob { mode: u32, oid: ObjectId },
283    Tree(Vec<u8>, Vec<&'a IndexEntry>),
284}
285
286impl<'a> ChildKind<'a> {
287    fn push_entry(&mut self, entry: &'a IndexEntry) {
288        if let Self::Tree(_, entries) = self {
289            entries.push(entry);
290        }
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    #![allow(clippy::expect_used, clippy::unwrap_used)]
297
298    use super::*;
299    use crate::index::{IndexEntry, MODE_EXECUTABLE, MODE_REGULAR, MODE_SYMLINK, MODE_TREE};
300    use crate::objects::parse_tree;
301    use tempfile::TempDir;
302
303    fn entry(path: &str, mode: u32, oid: ObjectId) -> IndexEntry {
304        IndexEntry {
305            ctime_sec: 0,
306            ctime_nsec: 0,
307            mtime_sec: 0,
308            mtime_nsec: 0,
309            dev: 0,
310            ino: 0,
311            mode,
312            uid: 0,
313            gid: 0,
314            size: 0,
315            oid,
316            flags: path.len().min(0xFFF) as u16,
317            flags_extended: None,
318            path: path.as_bytes().to_vec(),
319        }
320    }
321
322    #[test]
323    fn writes_sorted_tree_with_canonical_modes() {
324        let temp_dir = TempDir::new().unwrap();
325        let odb = Odb::new(temp_dir.path());
326
327        let oid_a = odb.write(ObjectKind::Blob, b"a").unwrap();
328        let oid_exec = odb.write(ObjectKind::Blob, b"exec").unwrap();
329        let oid_link = odb.write(ObjectKind::Blob, b"target").unwrap();
330
331        let mut index = Index::new();
332        index.add_or_replace(entry("bin/run.sh", 0o100777, oid_exec));
333        index.add_or_replace(entry("link", 0o120777, oid_link));
334        index.add_or_replace(entry("a.txt", 0o100664, oid_a));
335
336        let root_oid = write_tree_from_index(&odb, &index, "").unwrap();
337        let root_tree_obj = odb.read(&root_oid).unwrap();
338        let root_entries = parse_tree(&root_tree_obj.data).unwrap();
339
340        assert_eq!(root_entries.len(), 3);
341        assert_eq!(root_entries[0].name, b"a.txt");
342        assert_eq!(root_entries[0].mode, MODE_REGULAR);
343        assert_eq!(root_entries[1].name, b"bin");
344        assert_eq!(root_entries[1].mode, MODE_TREE);
345        assert_eq!(root_entries[2].name, b"link");
346        assert_eq!(root_entries[2].mode, MODE_SYMLINK);
347
348        let bin_tree_obj = odb.read(&root_entries[1].oid).unwrap();
349        let bin_entries = parse_tree(&bin_tree_obj.data).unwrap();
350        assert_eq!(bin_entries.len(), 1);
351        assert_eq!(bin_entries[0].name, b"run.sh");
352        assert_eq!(bin_entries[0].mode, MODE_EXECUTABLE);
353    }
354}
grit_lib/write_tree.rs

grit_lib/
write_tree.rs
