cli: added directory sizes

1 files changed, 24 insertions, 2 deletions

diff --git a/src/main.rs b/src/main.rs
index ebbc817..68b82ab 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -531,6 +531,26 @@ impl Default for Fs {
     }
 }
 
+pub fn compute_directory_sizes(fs: &mut Fs) {
+    fn compute(fs: &mut Fs, node_id: FsNodeId) {
+        let node = &fs.nodes[node_id];
+        if node.kind == FsNodeKind::File {
+            return;
+        }
+
+        let children = node.children.values().copied().collect::<Vec<_>>();
+        children.into_iter().for_each(|id| compute(fs, id));
+
+        let mut size = 0;
+        for child in fs.nodes[node_id].children.values() {
+            let child = &fs.nodes[*child];
+            size += child.size;
+        }
+        fs.nodes[node_id].size = size;
+    }
+    compute(fs, fs.root);
+}
+
 pub fn apply(fs: &mut Fs, op: &Operation) -> Result<(), FsError> {
     match &op.data {
         OperationData::CreateFile(data) => apply_create_file(fs, &op.header, &data),
@@ -951,6 +971,8 @@ fn cmd_ls(args: LsArgs) {
     let mut fs = Fs::default();
     let ops = common_read_log_file(&args.common);
     ops.iter().for_each(|op| apply(&mut fs, op).unwrap());
+    compute_directory_sizes(&mut fs);
+
     let node_id = match &args.path {
         Some(path) => find_node(&fs, path),
         None => Some(fs.root),
@@ -999,8 +1021,8 @@ fn write_node(
             if !recurse && (filter_type == None || filter_type == Some(FsNodeKind::Directory)) {
                 writeln!(
                     writer,
-                    "{}\tdir\t{}\t-\t-\t{}",
-                    path, node.lastmod, node.author
+                    "{}\tdir\t{}\t-\t{}\t{}",
+                    path, node.lastmod, node.size, node.author
                 )
                 .unwrap();
             }