From c44bd6dbc141f4ec71d5d6ba6653f7279f153d78 Mon Sep 17 00:00:00 2001
From: Linus Dierheimer <Linus@Dierheimer.de>
Date: Sun, 28 Mar 2021 15:38:52 +0200
Subject: [PATCH] better disk output

---
 completions/bash    |  1 +
 src/common.c        |  4 +++
 src/fastfetch.c     |  6 +++-
 src/fastfetch.h     |  2 ++
 src/modules/cpu.c   |  1 -
 src/modules/disk.c  | 82 +++++++++++++++++++++++++++++++++++++--------
 src/util/FFstrbuf.c | 10 ++++++
 src/util/FFstrbuf.h |  1 +
 8 files changed, 91 insertions(+), 16 deletions(-)

diff --git a/completions/bash b/completions/bash
index 0c99f83e8..f21c14ffc 100644
--- a/completions/bash
+++ b/completions/bash
@@ -146,6 +146,7 @@ __fastfetch_completion()
         "--disk-format"
         "--battery-format"
         "--locale-format"
+        "--disk-folders"
     )
 
     local FF_OPTIONS_PATH=(
diff --git a/src/common.c b/src/common.c
index 987c48512..62cbc1919 100644
--- a/src/common.c
+++ b/src/common.c
@@ -56,6 +56,8 @@ void ffDefaultConfig(FFconfig* config)
     ffStrbufInitA(&config->libPCI, 256);
     ffStrbufInitA(&config->libX11, 256);
     ffStrbufInitA(&config->libXrandr, 256);
+
+    ffStrbufInitA(&config->diskFolders, 128);
 }
 
 static void ffCleanup(FFinstance* instance)
@@ -88,6 +90,8 @@ static void ffCleanup(FFinstance* instance)
     ffStrbufDestroy(&instance->config.libX11);
     ffStrbufDestroy(&instance->config.libXrandr);
 
+    ffStrbufDestroy(&instance->config.diskFolders);
+
     ffStrbufDestroy(&instance->state.terminal.value);
 }
 
diff --git a/src/fastfetch.c b/src/fastfetch.c
index 48666c036..a63a75f94 100644
--- a/src/fastfetch.c
+++ b/src/fastfetch.c
@@ -84,6 +84,8 @@ static inline void printHelp()
         "   --lib-PCI <path>\n"
         "   --lib-X11 <path>\n"
         "   --lib-Xrandr <path>\n"
+        "Module specific options:\n"
+        "   --disk-folders <folders>: A colon seperated list of folder paths for the disk output. Default is \"/:/home\"\n"
         "\n"
         "If an value starts with an ?, it is optional. \"true\" will be used if not set.\n"
         "An (+) at the end indicates that more help can be printed with --help <option>\n"
@@ -238,7 +240,7 @@ static inline void printCommandHelp(const char* command)
     else if(strcasecmp(command, "memory-format") == 0)
         constructAndPrintCommandHelpFormat("memory", "{}MiB / {}MiB ({}%)", 3, "Used memory", "Total memory", "Used memory percentage");
     else if(strcasecmp(command, "disk-format") == 0)
-        constructAndPrintCommandHelpFormat("disk", "{}GB / {}GB ({}%)", 3, "Used disk space", "Total disk space", "Used disk space percentage");
+        constructAndPrintCommandHelpFormat("disk", "{}GB / {}GB ({4}%)", 4, "Used disk space", "Total disk space", "Number of files", "Used disk space percentage");
     else if(strcasecmp(command, "battery-format") == 0)
         constructAndPrintCommandHelpFormat("battery", "{} {} ({}) [{}%; {}]", 5, "Battery manufactor", "Battery model", "Battery technology", "Battery capacity", "Battery status");
     else if(strcasecmp(command, "locale-format") == 0)
@@ -422,6 +424,8 @@ static void parseOption(FFinstance* instance, FFdata* data, const char* key, con
         optionParseString(key, value, &instance->config.libX11);
     else if(strcasecmp(key, "--lib-Xrandr") == 0)
         optionParseString(key, value, &instance->config.libXrandr);
+    else if(strcasecmp(key, "--disk-folders") == 0)
+        optionParseString(key, value, &instance->config.diskFolders);
     else
     {
         printf("Error: unknown option: %s\n", key);
diff --git a/src/fastfetch.h b/src/fastfetch.h
index 014ab9fb6..dcd0ec319 100644
--- a/src/fastfetch.h
+++ b/src/fastfetch.h
@@ -71,6 +71,8 @@ typedef struct FFconfig
     FFstrbuf libX11;
     FFstrbuf libXrandr;
 
+    FFstrbuf diskFolders;
+
 } FFconfig;
 
 typedef struct FFvalue
diff --git a/src/modules/cpu.c b/src/modules/cpu.c
index 2063ede5a..e80b7e351 100644
--- a/src/modules/cpu.c
+++ b/src/modules/cpu.c
@@ -87,7 +87,6 @@ void ffPrintCPU(FFinstance* instance)
         if(ghz == 0)
             ghz = infoMinFreq;
 
-
         if(namePretty.length > 0)
             ffStrbufAppend(&cpu, &namePretty);
         else if(name[0] != '\0')
diff --git a/src/modules/disk.c b/src/modules/disk.c
index 89949a3fe..3d74cc4bd 100644
--- a/src/modules/disk.c
+++ b/src/modules/disk.c
@@ -2,34 +2,88 @@
 
 #include <sys/statvfs.h>
 
-void ffPrintDisk(FFinstance* instance)
+static void printStatvfs(FFinstance* instance, const char* key, struct statvfs* fs)
 {
-    struct statvfs fs;
-    int ret = statvfs("/", &fs);
-    if(ret != 0)
-    {
-        ffPrintError(instance, "Disk (/)", "statvfs(\"/\", &fs) != 0");
-        return;
-    }
+    const uint32_t GB = 1024 * 1024 * 1024;
 
-    const uint32_t GB = (1024 * 1024) * 1024;
-
-    uint32_t total     = (fs.f_blocks * fs.f_frsize) / GB;
-    uint32_t available = (fs.f_bfree  * fs.f_frsize) / GB;
+    uint32_t total     = (fs->f_blocks * fs->f_frsize) / GB;
+    uint32_t available = (fs->f_bfree  * fs->f_frsize) / GB;
     uint32_t used      = total - available;
     uint8_t percentage = (used / (double) total) * 100.0;
 
+    uint32_t files = fs->f_files - fs->f_ffree;
+
     if(instance->config.diskFormat.length == 0)
     {
-        ffPrintLogoAndKey(instance, "Disk (/)");
+        ffPrintLogoAndKey(instance, key);
         printf("%uGB / %uGB (%u%%)\n", used, total, percentage);
     }
     else
     {
-        ffPrintFormatString(instance, "Disk (/)", &instance->config.diskFormat, 3,
+        ffPrintFormatString(instance, key, &instance->config.diskFormat, 4,
             (FFformatarg){FF_FORMAT_ARG_TYPE_UINT, &used},
             (FFformatarg){FF_FORMAT_ARG_TYPE_UINT, &total},
+            (FFformatarg){FF_FORMAT_ARG_TYPE_UINT, &files},
             (FFformatarg){FF_FORMAT_ARG_TYPE_UINT8, &percentage}
         );
     }
 }
+
+static void printFolder(FFinstance* instance, const char* folderPath)
+{
+    FF_STRBUF_CREATE(key);
+    ffStrbufAppendS(&key, "Disk (");
+    ffStrbufAppendS(&key, folderPath);
+    ffStrbufAppendC(&key, ')');
+
+    struct statvfs fs;
+    int ret = statvfs(folderPath, &fs);
+    if(ret != 0 && instance->config.diskFormat.length == 0)
+    {
+        ffPrintError(instance, key.chars, "statvfs(\"%s\", &fs) != 0", folderPath);
+        ffStrbufDestroy(&key);
+        return;
+    }
+
+    printStatvfs(instance, key.chars, &fs);
+
+    ffStrbufDestroy(&key);
+}
+
+void ffPrintDisk(FFinstance* instance)
+{
+    if(instance->config.diskFolders.length == 0)
+    {
+        struct statvfs fsRoot;
+        int rootRet = statvfs("/", &fsRoot);
+
+        struct statvfs fsHome;
+        int homeRet = statvfs("/home", &fsHome);
+
+        if(rootRet != 0 && homeRet != 0)
+        {
+            ffPrintError(instance, "Disk", "statvfs failed for both / and /home");
+            return;
+        }
+
+        if(rootRet == 0)
+            printStatvfs(instance, "Disk (/)", &fsRoot);
+
+        if(homeRet == 0 && (rootRet != 0 || fsRoot.f_fsid != fsHome.f_fsid))
+            printStatvfs(instance, "Disk (/home)", &fsHome);
+    }
+    else
+    {
+        uint32_t lastIndex = 0;
+        while (lastIndex < instance->config.diskFolders.length)
+        {
+            uint32_t colonIndex = ffStrbufFirstIndexC(&instance->config.diskFolders, ':');
+
+            if(colonIndex < instance->config.diskFolders.length)
+                instance->config.diskFolders.chars[colonIndex] = '\0';
+
+            printFolder(instance, instance->config.diskFolders.chars + lastIndex);
+            lastIndex = colonIndex + 1;
+        }
+    }
+}
diff --git a/src/util/FFstrbuf.c b/src/util/FFstrbuf.c
index e35a20897..274477ae4 100644
--- a/src/util/FFstrbuf.c
+++ b/src/util/FFstrbuf.c
@@ -333,6 +333,16 @@ void ffStrbufRemoveStrings(FFstrbuf* strbuf, uint32_t numStrings, ...)
     }
 }
 
+uint32_t ffStrbufFirstIndexC(FFstrbuf* strbuf, const char c)
+{
+    for(uint32_t i = 0; i < strbuf->length; i++)
+    {
+        if(strbuf->chars[i] == c)
+            return i;
+    }
+    return strbuf->length;
+}
+
 uint32_t ffStrbufLastIndexC(FFstrbuf* strbuf, const char c)
 {
     //We need to loop one higher than the actual index, because uint32_t is guranteed to be >= 0, so this statement would always be true
diff --git a/src/util/FFstrbuf.h b/src/util/FFstrbuf.h
index f3cf2c2f5..66e4dab44 100644
--- a/src/util/FFstrbuf.h
+++ b/src/util/FFstrbuf.h
@@ -68,6 +68,7 @@ void ffStrbufTrim(FFstrbuf* strbuf, char c);
 
 void ffStrbufRemoveStrings(FFstrbuf* strbuf, uint32_t numStrings, ...);
 
+uint32_t ffStrbufFirstIndexC(FFstrbuf* strbuf, const char c);
 uint32_t ffStrbufLastIndexC(FFstrbuf* strbuf, const char c);
 
 void ffStrbufLimitLength(FFstrbuf* strbuf, uint32_t length);