swaymux/tree/pstree.cpp

//
// Created by grimmauld on 25.02.24.
//

#include "pstree.h"
#include "../util/StringUtil.h"

#include <cstring>
#include <set>
#include <string>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <map>
#include <memory>

const char *procFileDelim = ":";

void populateTree(ProcessTreeNode *node, std::map<pid_t, std::set<ProcessRecord>> &ppid_map, std::set<ProcessTreeNode*>& changedParents) {
    if (node == nullptr)
        return;

    for (int i =  (int) node->childCount() - 1; i >= 0; i--) { // backwards to avoid changing stuff that might change later
        std::set<ProcessRecord> children = ppid_map[node->proc.ppid];
        auto* c = node->child(i);
        if (c->proc.pid == 0)
            continue;

        if (children.find(c->proc) == children.end()) {
            // node->removeChild(i); // fixme: segfaults
            changedParents.insert(node);
        }
    }

    for (const auto &child: ppid_map[node->proc.pid]) {
        // check child exists, if so no need to update
        auto* existingChildNode = node->findChild(ProcessTreeNode::matchingRecord(child));
        if (existingChildNode != nullptr) {
            populateTree(existingChildNode, ppid_map, changedParents);
            continue;
        }

        // else add the new child node, keep track in notifier
        changedParents.insert(node);
        auto childNode = std::make_unique<ProcessTreeNode>(child, node);
        populateTree(childNode.get(), ppid_map, changedParents);
        node->appendChild(childNode);
    }
}

ProcessTreeNode * get_process_records() {
    auto* root = new ProcessTreeNode(); // rootPrc
    update_process_records(root);
    return root;
}

std::set<ProcessTreeNode*> update_process_records(ProcessTreeNode *node) {
    std::map<pid_t, ProcessRecord> pid_map;
    std::map<pid_t, std::set<ProcessRecord>> ppid_map;
    std::set<ProcessTreeNode*> changedParents{};

    const auto proc_fs = fs::path(PROCFS_ROOT);
    auto content = fs::directory_iterator(proc_fs);
    for (const auto &proc_dir: content) {
        if (fs::is_directory(proc_dir)) {
            insert_process_record(proc_dir.path() / "status", pid_map);
        }
    }

    pid_t rootPid = INT32_MAX;
    for (auto &pair: pid_map) {
        auto proc = pair.second;
        if (proc.pid < rootPid) {
            rootPid = proc.pid;
        }
        ppid_map[proc.ppid].insert(proc);
    }

    // auto rootPrc = pid_map[rootPid];
    // ProcessTreeNode root = ProcessTreeNode(); // rootPrc
    populateTree(node, ppid_map, changedParents);
    return changedParents;
}

void insert_process_record(const std::filesystem::path &status_path, std::map<pid_t, ProcessRecord> &buff) {
    if (!fs::exists(status_path) || !fs::is_regular_file(status_path))
        return;

    std::string proc_name;
    pid_t pid, ppid;
    bool foundPid, foundPPid, foundName;

    auto reader = std::ifstream{status_path, std::ios::in};
    std::string line;
    while (std::getline(reader, line)) {
        auto data = line.data();
        auto name = strtok(data, procFileDelim);
        auto val = strtok(nullptr, procFileDelim);
        if (name == nullptr || val == nullptr)
            continue;

        if (strcmp(name, "Name") == 0) {
            proc_name = std::string(val);
            trim(proc_name);
            foundName = true;
        } else if (strcmp(name, "Pid") == 0) {
            pid = std::stoi(val);
            foundPid = true;
        } else if (strcmp(name, "PPid") == 0) {
            ppid = std::stoi(val);
            foundPPid = true;
        }
    }

    if (foundPPid && foundName && foundPid) {
        buff.insert({pid, ProcessRecord(proc_name, pid, ppid)});
    }
}

void printTree(const ProcessTreeNode &node, const std::string &prefix) {
    std::cout << prefix << node.proc.name << ": " << node.proc.pid << "\n";
    for (const auto &child: node.children) {
        printTree(*child, prefix + "\t");
    }
}

void printTree() {
    auto* tree = get_process_records();
    printTree(*tree, "");
    delete tree;
}

QVariant ProcessTreeNode::data(int column) const {
    switch (column) {
        case 0:
            return QString::fromStdString(std::to_string(proc.pid));
        case 1:
            return QString::fromStdString(proc.name);
        default:
            throw std::exception();
    }
}

QVariant ProcessTreeNode::headerData(int column) const {
    switch (column) {
        case 0:
            return "PID";
        case 1:
            return "Process Name";
        default:
            return QVariant{};
    }
}