main.cpp

#include <ncurses.h>
#include <iostream>
#include <fstream>
#include <filesystem>
#include <vector>
#include <mutex>
#include <condition_variable>
#include <thread>
#include <algorithm>
#include <chrono>
#include <format>
#include <atomic>
#include <string_view>
#include <cstring>
#include <unordered_set>
#include <unordered_map>
#include <limits>

#ifdef __APPLE__
#include <sys/attr.h>
#include <sys/vnode.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#endif

struct file_size
{
    std::filesystem::path path;
    size_t size;
};

template <typename T>
struct thread_safe_vector
{
    void add(T value)
    {
        {
            std::scoped_lock scope_lock(lock);
            vec.push_back(std::move(value));
        }
        available.notify_one();
    }

    void append(std::vector<T>&& others)
    {
        {
            std::scoped_lock scope_lock(lock);
            vec.insert(vec.end(),
                std::make_move_iterator(others.begin()),
                std::make_move_iterator(others.end()));
        }
        available.notify_all();
    }

    // Blocks until items are available or cancelled() returns true
    template <typename Cancelled>
    void wait_transfer(const size_t count, std::vector<T>& destination, Cancelled cancelled)
    {
        std::unique_lock scope_lock(lock);
        available.wait(scope_lock, [&] { return !vec.empty() || cancelled(); });

        if (!vec.empty())
        {
            const size_t take_count = std::min<size_t>(count, vec.size());

            // Take from the back: no O(n) shuffle of the remaining queue, and the paths
            // are moved out rather than copied
            destination.resize(take_count);
            std::move(vec.end() - take_count, vec.end(), destination.begin());

            vec.erase(vec.end() - take_count, vec.end());
        }
    }

    void wake_all_waiters()
    {
        available.notify_all();
    }

private:
    std::vector<T> vec;
    std::mutex lock;
    std::condition_variable available;
};

// Tracks the N largest files seen so far. Workers read `threshold` without locking to
// discard files that cannot make the list, so the mutex is only taken for candidates.
struct largest_files_tracker
{
    static constexpr size_t max_tracked = 100;

    // Size of the smallest tracked file once the list is full; files below this can't qualify
    std::atomic<size_t> threshold = 0;

    void merge(std::vector<file_size>& candidates)
    {
        std::scoped_lock scope_lock(lock);

        vec.insert(vec.end(),
            std::make_move_iterator(candidates.begin()),
            std::make_move_iterator(candidates.end()));

        std::stable_sort(vec.begin(), vec.end(), [](const file_size& lhs, const file_size& rhs)
        {
            return lhs.size > rhs.size;
        });

        if (vec.size() >= max_tracked)
        {
            vec.resize(max_tracked);
            threshold.store(vec.back().size, std::memory_order_relaxed);
        }
    }

    std::vector<file_size> get_copy()
    {
        std::scoped_lock scope_lock(lock);
        return vec;
    }

private:
    std::vector<file_size> vec;
    std::mutex lock;
};

// Accumulates the recursive size of every directory scanned so far. Workers credit a
// directory's direct file bytes to the directory and all of its ancestors, so an entry
// is the live total of everything scanned beneath it.
struct directory_size_tracker
{
    void add(const std::unordered_map<std::string, size_t>& deltas)
    {
        std::scoped_lock scope_lock(lock);

        for (const auto& [path, bytes] : deltas)
        {
            const auto [it, inserted] = index_of.try_emplace(path, static_cast<uint32_t>(sizes.size()));
            if (inserted)
            {
                sizes.push_back(bytes);
                names.push_back(&it->first);
            }
            else
            {
                sizes[it->second] += bytes;
            }
        }
    }

    // Returns the `count` largest directories seen so far. Costs one pass over every
    // tracked directory, but is only paid while the directories view is being rendered.
    std::vector<file_size> get_top(const size_t count)
    {
        std::scoped_lock scope_lock(lock);

        scratch.resize(sizes.size());
        for (uint32_t i = 0; i < scratch.size(); ++i)
        {
            scratch[i] = i;
        }

        const size_t take_count = std::min(count, scratch.size());
        std::partial_sort(scratch.begin(), scratch.begin() + take_count, scratch.end(),
            [this](const uint32_t lhs, const uint32_t rhs)
        {
            return sizes[lhs] > sizes[rhs];
        });

        std::vector<file_size> top;
        top.reserve(take_count);
        for (size_t i = 0; i < take_count; ++i)
        {
            top.emplace_back(*names[scratch[i]], sizes[scratch[i]]);
        }
        return top;
    }

private:
    // Sizes live in a dense vector so get_top scans contiguous memory; the map only
    // resolves paths to slots. Key pointers are stable across rehashes.
    std::unordered_map<std::string, uint32_t> index_of;
    std::vector<size_t> sizes;
    std::vector<const std::string*> names;
    std::vector<uint32_t> scratch;
    std::mutex lock;
};

// Identifies a directory by device + inode so the same directory reached through
// different paths (e.g. macOS firmlinks like /Users vs /System/Volumes/Data/Users)
// is only scanned once
struct dir_id
{
    uint64_t device;
    uint64_t inode;

    bool operator==(const dir_id&) const = default;
};

struct dir_id_hash
{
    size_t operator()(const dir_id& id) const
    {
        return std::hash<uint64_t>()(id.inode) ^ (std::hash<uint64_t>()(id.device) << 1);
    }
};

// Fallback for platforms/volumes where no real id is available; never collides,
// so deduplication becomes a no-op for these entries
dir_id unique_dir_id()
{
    static std::atomic<uint64_t> counter = 1;
    return dir_id{ std::numeric_limits<uint64_t>::max(), counter++ };
}

struct directory_dedupe
{
    void mark_visited(const dir_id id)
    {
        std::scoped_lock scope_lock(lock);
        visited.insert(id);
    }

    // Erases paths whose id has already been visited and marks the rest as visited.
    // `ids` must be parallel to `paths`.
    void filter(std::vector<std::filesystem::path>& paths, const std::vector<dir_id>& ids)
    {
        std::scoped_lock scope_lock(lock);

        size_t keep = 0;
        for (size_t i = 0; i < paths.size(); ++i)
        {
            if (visited.insert(ids[i]).second)
            {
                if (keep != i)
                {
                    paths[keep] = std::move(paths[i]);
                }
                ++keep;
            }
        }
        paths.resize(keep);
    }

private:
    std::unordered_set<dir_id, dir_id_hash> visited;
    std::mutex lock;
};

thread_safe_vector<std::filesystem::path> directory_queue;
largest_files_tracker largest_files;
directory_size_tracker directory_sizes;
directory_dedupe visited_directories;
std::filesystem::path scan_root;

std::atomic<bool> stop_flag(false);
std::atomic<int> failed_reads = 0;
std::atomic<int> threads_working = 0;
std::atomic<size_t> files_scanned = 0;
std::atomic<size_t> bytes_scanned = 0;
std::atomic<size_t> directories_scanned = 0;
// Directories queued or in-flight; the scan is complete when this reaches 0
std::atomic<size_t> outstanding_directories = 0;

// Space the file actually occupies on disk, so sparse files (e.g. VM disk images)
// aren't reported at their logical size
size_t file_size_on_disk(const std::filesystem::directory_entry& file, std::error_code& ec)
{
#ifdef __APPLE__
    struct stat file_stat;
    if (lstat(file.path().c_str(), &file_stat) != 0)
    {
        ec = std::error_code(errno, std::generic_category());
        return 0;
    }
    return static_cast<size_t>(file_stat.st_blocks) * 512;
#else
    return file.file_size(ec);
#endif
}

dir_id directory_id_of(const std::filesystem::path& path)
{
#ifdef __APPLE__
    struct stat dir_stat;
    if (lstat(path.c_str(), &dir_stat) == 0)
    {
        return dir_id{ static_cast<uint64_t>(dir_stat.st_dev), static_cast<uint64_t>(dir_stat.st_ino) };
    }
#endif
    return unique_dir_id();
}

void scan_directory_portable(const std::filesystem::path& path,
                             std::vector<file_size>& candidate_files,
                             std::vector<std::filesystem::path>& directories,
                             std::vector<dir_id>& directory_ids,
                             size_t& files_seen,
                             size_t& bytes_seen)
{
    try
    {
        const auto options = std::filesystem::directory_options::skip_permission_denied;
        for (const auto& file : std::filesystem::directory_iterator(path, options))
        {
            // symlink_status does not follow symlinks, so linked files aren't
            // double-counted and linked directories can't create scan cycles
            std::error_code ec;
            const auto type = file.symlink_status(ec).type();
            if (ec)
            {
                continue;
            }

            if (type == std::filesystem::file_type::regular)
            {
                const size_t size = file_size_on_disk(file, ec);
                if (!ec)
                {
                    ++files_seen;
                    bytes_seen += size;
                    if (size >= largest_files.threshold.load(std::memory_order_relaxed))
                    {
                        candidate_files.emplace_back(file.path(), size);
                    }
                }
            }
            else if (type == std::filesystem::file_type::directory)
            {
                directory_ids.push_back(directory_id_of(file.path()));
                directories.push_back(file.path());
            }
        }
    }
    catch (const std::exception& e)
    {
        ++failed_reads;
        // Cannot read this directory; skip it and continue with the rest of the batch
    }
}

#ifdef __APPLE__
// Reads name, type and size for whole batches of entries per syscall, instead of
// one stat() per file. Returns false if the volume doesn't support bulk enumeration
// and nothing was consumed, in which case the caller should use the portable scan.
bool scan_directory_bulk(const std::filesystem::path& path,
                         std::vector<file_size>& candidate_files,
                         std::vector<std::filesystem::path>& directories,
                         std::vector<dir_id>& directory_ids,
                         size_t& files_seen,
                         size_t& bytes_seen)
{
    const int dir_fd = open(path.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC);
    if (dir_fd < 0)
    {
        ++failed_reads;
        return true;
    }

    attrlist request{};
    request.bitmapcount = ATTR_BIT_MAP_COUNT;
    request.commonattr = ATTR_CMN_RETURNED_ATTRS | ATTR_CMN_NAME | ATTR_CMN_DEVID | ATTR_CMN_OBJTYPE | ATTR_CMN_FILEID;
    request.fileattr = ATTR_FILE_ALLOCSIZE;

    thread_local std::vector<char> buffer(256 * 1024);
    bool any_entries_consumed = false;

    for (;;)
    {
        const int entry_count = getattrlistbulk(dir_fd, &request, buffer.data(), buffer.size(), 0);
        if (entry_count < 0)
        {
            close(dir_fd);
            if (any_entries_consumed)
            {
                ++failed_reads;
                return true;
            }
            return false;
        }
        if (entry_count == 0)
        {
            break;
        }
        any_entries_consumed = true;

        const char* entry = buffer.data();
        for (int i = 0; i < entry_count; ++i)
        {
            const char* field = entry;

            uint32_t entry_length;
            std::memcpy(&entry_length, field, sizeof(entry_length));
            field += sizeof(uint32_t);

            attribute_set_t returned;
            std::memcpy(&returned, field, sizeof(returned));
            field += sizeof(attribute_set_t);

            std::string_view name;
            if (returned.commonattr & ATTR_CMN_NAME)
            {
                attrreference_t name_info;
                std::memcpy(&name_info, field, sizeof(name_info));
                // attr_length includes the null terminator
                name = std::string_view(field + name_info.attr_dataoffset, name_info.attr_length - 1);
                field += sizeof(attrreference_t);
            }

            dev_t device = 0;
            if (returned.commonattr & ATTR_CMN_DEVID)
            {
                std::memcpy(&device, field, sizeof(device));
                field += sizeof(dev_t);
            }

            fsobj_type_t obj_type = VNON;
            if (returned.commonattr & ATTR_CMN_OBJTYPE)
            {
                std::memcpy(&obj_type, field, sizeof(obj_type));
                field += sizeof(fsobj_type_t);
            }

            uint64_t file_id = 0;
            bool has_file_id = false;
            if (returned.commonattr & ATTR_CMN_FILEID)
            {
                std::memcpy(&file_id, field, sizeof(file_id));
                field += sizeof(uint64_t);
                has_file_id = (returned.commonattr & ATTR_CMN_DEVID) != 0;
            }

            if (obj_type == VREG && (returned.fileattr & ATTR_FILE_ALLOCSIZE) && !name.empty())
            {
                off_t size;
                std::memcpy(&size, field, sizeof(size));

                ++files_seen;
                bytes_seen += static_cast<size_t>(size);
                if (static_cast<size_t>(size) >= largest_files.threshold.load(std::memory_order_relaxed))
                {
                    candidate_files.emplace_back(path / name, static_cast<size_t>(size));
                }
            }
            else if (obj_type == VDIR && !name.empty())
            {
                directory_ids.push_back(has_file_id
                    ? dir_id{ static_cast<uint64_t>(device), file_id }
                    : unique_dir_id());
                directories.emplace_back(path / name);
            }

            entry += entry_length;
        }
    }

    close(dir_fd);
    return true;
}
#endif

// Credits a directory's direct file bytes to itself and every ancestor up to the scan root
void add_to_ancestors(std::unordered_map<std::string, size_t>& totals, const std::filesystem::path& path, const size_t bytes)
{
    for (std::filesystem::path current = path;;)
    {
        totals[current.string()] += bytes;

        if (current == scan_root)
        {
            break;
        }

        std::filesystem::path parent = current.parent_path();
        if (parent.empty() || parent == current)
        {
            break;
        }
        current = std::move(parent);
    }
}

void process_directories(const std::vector<std::filesystem::path>& paths)
{
    ++threads_working;

    std::vector<file_size> candidate_files;
    std::vector<std::filesystem::path> directories;
    std::vector<dir_id> directory_ids;
    std::unordered_map<std::string, size_t> directory_bytes;
    size_t batch_files_scanned = 0;
    size_t batch_bytes_scanned = 0;

    for (const auto& path : paths)
    {
        size_t bytes_in_directory = 0;
#ifdef __APPLE__
        if (!scan_directory_bulk(path, candidate_files, directories, directory_ids, batch_files_scanned, bytes_in_directory))
#endif
        {
            scan_directory_portable(path, candidate_files, directories, directory_ids, batch_files_scanned, bytes_in_directory);
        }

        batch_bytes_scanned += bytes_in_directory;
        if (bytes_in_directory > 0)
        {
            add_to_ancestors(directory_bytes, path, bytes_in_directory);
        }
    }

    visited_directories.filter(directories, directory_ids);

    const size_t new_directory_count = directories.size();

    // Count children before releasing our claim on the parents so the total never dips to 0 early
    outstanding_directories += new_directory_count;
    directory_queue.append(std::move(directories));

    if (!candidate_files.empty())
    {
        largest_files.merge(candidate_files);
    }

    if (!directory_bytes.empty())
    {
        directory_sizes.add(directory_bytes);
    }

    files_scanned += batch_files_scanned;
    bytes_scanned += batch_bytes_scanned;
    directories_scanned += new_directory_count;

    outstanding_directories -= paths.size();

    --threads_working;
}

void thread_job()
{
    std::vector<std::filesystem::path> directories;

    while (!stop_flag.load())
    {
        directory_queue.wait_transfer(100, directories, [] { return stop_flag.load(); });
        if (!directories.empty())
        {
            process_directories(directories);
        }
        directories.clear();
    }
}

struct terminal_context
{
    int max_x;
    int max_y;

    terminal_context()
        : max_x(0)
        , max_y(0)
    {
    }
};

std::string pretty_bytes(const size_t size)
{
    if (size < 1'000)
    {
        return std::format("{} B", size);
    }
    if (size < 1'000'000)
    {
        return std::format("{:.2f} KB", static_cast<double>(size) / 1000);
    }
    if (size < 1'000'000'000)
    {
        return std::format("{:.2f} MB", static_cast<double>(size) / 1'000'000); 
    }
    return std::format("{:.2f} GB", static_cast<double>(size) / 1'000'000'000); 
}

enum class view_mode
{
    files,
    directories,
    combined,
};

constexpr short directory_color_pair = 1;

view_mode next_view(const view_mode view)
{
    return static_cast<view_mode>((static_cast<int>(view) + 1) % 3);
}

void render_frame(const terminal_context& terminal_ctx, const std::chrono::steady_clock::duration elapsed, const view_mode view)
{
    std::vector<file_size> local_largest;
    const char* label = "";

    switch (view)
    {
    case view_mode::files:
        local_largest = largest_files.get_copy();
        label = "Largest files";
        break;
    case view_mode::directories:
        local_largest = directory_sizes.get_top(largest_files_tracker::max_tracked);
        label = "Largest directories";
        break;
    case view_mode::combined:
    {
        local_largest = largest_files.get_copy();

        // Directories get a trailing slash so they can be told apart from files
        std::vector<file_size> directories = directory_sizes.get_top(largest_files_tracker::max_tracked);
        for (auto& directory : directories)
        {
            directory.path += "/";
        }

        local_largest.insert(local_largest.end(),
            std::make_move_iterator(directories.begin()),
            std::make_move_iterator(directories.end()));

        std::stable_sort(local_largest.begin(), local_largest.end(), [](const file_size& lhs, const file_size& rhs)
        {
            return lhs.size > rhs.size;
        });

        if (local_largest.size() > largest_files_tracker::max_tracked)
        {
            local_largest.resize(largest_files_tracker::max_tracked);
        }

        label = "Largest files & directories";
        break;
    }
    }

    mvprintw(1, 5, label);

    for (int i = 0; i < std::min<int>(terminal_ctx.max_y - 4, local_largest.size()); i++)
    {
        const auto&[path, size] = local_largest[i];

        // Combined-view directories carry a trailing slash, so they have no filename
        const bool is_directory = view == view_mode::directories || !path.has_filename();
        if (is_directory)
        {
            attron(COLOR_PAIR(directory_color_pair));
        }

        mvprintw(i + 2, 5, pretty_bytes(size).c_str());
        mvprintw(i + 2, 20, path.c_str());

        if (is_directory)
        {
            attroff(COLOR_PAIR(directory_color_pair));
        }
    }

    const double elapsed_seconds = std::chrono::duration<double>(elapsed).count();
    const size_t total_bytes = bytes_scanned.load();
    const size_t bytes_per_second = elapsed_seconds > 0 ? static_cast<size_t>(total_bytes / elapsed_seconds) : 0;
    mvprintw(0, 0, "FILE SNIPER | %d files scanned | %s scanned | %d directories scanned | %d threads active | %d failed directories | %.1f s | %s/s | %s", files_scanned.load(), pretty_bytes(total_bytes).c_str(), directories_scanned.load(), threads_working.load(), failed_reads.load(), elapsed_seconds, pretty_bytes(bytes_per_second).c_str(), stop_flag.load() ? "Complete" : "Searching...");
    const char* next_label = "";
    switch (next_view(view))
    {
    case view_mode::files:       next_label = "Show largest files"; break;
    case view_mode::directories: next_label = "Show largest directories"; break;
    case view_mode::combined:    next_label = "Show combined"; break;
    }
    mvprintw(terminal_ctx.max_y - 1, 0, "(Q) - Quit    (TAB) - %s", next_label);
}

int main(int argc, char* argv[])
{
    const bool benchmark = (argc == 3 && std::string_view(argv[2]) == "--benchmark");
    const bool headless = (argc == 4 && std::string_view(argv[2]) == "--output");
    if (argc != 2 && !benchmark && !headless)
    {
        std::cerr << "Usage: filesniper <directory> [--benchmark | --output <file>]\n";
        return 1;
    }

    scan_root = std::filesystem::path(argv[1]);
    if (!scan_root.has_filename())
    {
        // Drop a trailing slash so ancestor propagation stops exactly at the root
        scan_root = scan_root.parent_path();
    }
    if (scan_root.empty())
    {
        scan_root = std::filesystem::path(argv[1]);
    }

    outstanding_directories = 1;
    visited_directories.mark_visited(directory_id_of(scan_root));
    directory_queue.add(scan_root);

    const auto scan_start = std::chrono::steady_clock::now();

    std::vector<std::thread> workers;
    workers.reserve(std::thread::hardware_concurrency());

    for (int i = 0; i < std::thread::hardware_concurrency(); ++i)
    {
        workers.emplace_back(thread_job);
    }

    if (benchmark || headless)
    {
        using namespace std::chrono_literals;
        while (outstanding_directories.load() != 0)
        {
            std::this_thread::sleep_for(1ms);
        }
        const auto elapsed = std::chrono::steady_clock::now() - scan_start;

        stop_flag = true;
        directory_queue.wake_all_waiters();
        for (auto& worker : workers)
        {
            worker.join();
        }

        const double seconds = std::chrono::duration<double>(elapsed).count();
        const auto top_files = largest_files.get_copy();
        const auto top_directories = directory_sizes.get_top(largest_files_tracker::max_tracked);

        if (benchmark)
        {
            std::cout << std::format("elapsed: {:.3f} s\n", seconds);
            std::cout << std::format("files scanned: {}\n", files_scanned.load());
            std::cout << std::format("directories scanned: {}\n", directories_scanned.load());
            std::cout << std::format("failed directories: {}\n", failed_reads.load());
            std::cout << std::format("throughput: {:.0f} files/s\n", files_scanned.load() / seconds);

            for (size_t i = 0; i < std::min<size_t>(5, top_files.size()); ++i)
            {
                std::cout << std::format("{:>12} {}\n", pretty_bytes(top_files[i].size), top_files[i].path.string());
            }
            std::cout << '\n';
            for (size_t i = 0; i < std::min<size_t>(5, top_directories.size()); ++i)
            {
                std::cout << std::format("{:>12} {}\n", pretty_bytes(top_directories[i].size), top_directories[i].path.string());
            }
            return 0;
        }

        std::ofstream output(argv[3]);
        if (!output)
        {
            std::cerr << std::format("Failed to open '{}' for writing\n", argv[3]);
            return 1;
        }

        output << std::format("scanned '{}' in {:.3f} s | {} files | {} directories | {} failed directories\n\n",
            argv[1], seconds, files_scanned.load(), directories_scanned.load(), failed_reads.load());

        output << "largest files:\n";
        for (const auto& [path, size] : top_files)
        {
            output << std::format("{:>12} {}\n", pretty_bytes(size), path.string());
        }

        output << "\nlargest directories:\n";
        for (const auto& [path, size] : top_directories)
        {
            output << std::format("{:>12} {}\n", pretty_bytes(size), path.string());
        }
        return 0;
    }

    initscr();
    noecho();
    cbreak();
    nodelay(stdscr, TRUE);

    if (has_colors())
    {
        start_color();
        // Keep the terminal's default background instead of ncurses' black
        use_default_colors();
        init_pair(directory_color_pair, COLOR_YELLOW, -1);
    }

    // Set once when the search completes so the timer freezes at the final scan time
    std::chrono::steady_clock::time_point scan_end;
    bool scan_ended = false;
    view_mode view = view_mode::files;

    while (true)
    {
        const int ch = getch();
        if (ch == 'q')
        {
            break;
        }
        if (ch == '\t')
        {
            view = next_view(view);
        }

        if (outstanding_directories.load() == 0 && !scan_ended)
        {
            scan_end = std::chrono::steady_clock::now();
            scan_ended = true;
            stop_flag = true;
        }

        // render
        {
            clear();

            terminal_context terminal_ctx;
            getmaxyx(stdscr, terminal_ctx.max_y, terminal_ctx.max_x);

            const auto elapsed = (scan_ended ? scan_end : std::chrono::steady_clock::now()) - scan_start;
            render_frame(terminal_ctx, elapsed, view);
            refresh();
        }

        using namespace std::chrono_literals;
        std::this_thread::sleep_for(100ms);
    }

    stop_flag = true;
    directory_queue.wake_all_waiters();
    for (auto& worker : workers)
    {
        worker.join();
    }

    nocbreak();
    endwin();

    return 0;
}