repomix/lib/core/file/fileTreeGenerate.js

A tool to pack repository contents to single file for AI consumption

import nodepath from 'node:path';
const childLookupCache = new WeakMap();
const createTreeNode = (name, isDirectory) => ({ name, children: [], isDirectory });
const getOrCreateChildMap = (node) => {
    let map = childLookupCache.get(node);
    if (!map) {
        map = new Map();
        childLookupCache.set(node, map);
    }
    return map;
};
export const generateFileTree = (files, emptyDirPaths = []) => {
    const root = createTreeNode('root', true);
    for (const file of files) {
        addPathToTree(root, file, false);
    }
    for (const dir of emptyDirPaths) {
        addPathToTree(root, dir, true);
    }
    return root;
};
const addPathToTree = (root, path, isDirectory) => {
    const parts = path.split(nodepath.sep);
    let currentNode = root;
    for (let i = 0; i < parts.length; i++) {
        const part = parts[i];
        const isLastPart = i === parts.length - 1;
        const childMap = getOrCreateChildMap(currentNode);
        let child = childMap.get(part);
        if (!child) {
            child = createTreeNode(part, !isLastPart || isDirectory);
            currentNode.children.push(child);
            childMap.set(part, child);
        }
        currentNode = child;
    }
};
const sortTreeNodes = (node) => {
    node.children.sort((a, b) => {
        if (a.isDirectory === b.isDirectory) {
            return a.name.localeCompare(b.name);
        }
        return a.isDirectory ? -1 : 1;
    });
    for (const child of node.children) {
        sortTreeNodes(child);
    }
};
export const treeToString = (node, prefix = '', _isRoot = true) => {
    if (_isRoot) {
        sortTreeNodes(node);
    }
    let result = '';
    for (const child of node.children) {
        result += `${prefix}${child.name}${child.isDirectory ? '/' : ''}\n`;
        if (child.isDirectory) {
            result += treeToString(child, `${prefix}  `, false);
        }
    }
    return result;
};
export const treeToStringWithLineCounts = (node, lineCounts, prefix = '', currentPath = '', _isRoot = true) => {
    if (_isRoot) {
        sortTreeNodes(node);
    }
    let result = '';
    for (const child of node.children) {
        const childPath = currentPath ? `${currentPath}/${child.name}` : child.name;
        if (child.isDirectory) {
            result += `${prefix}${child.name}/\n`;
            result += treeToStringWithLineCounts(child, lineCounts, `${prefix}  `, childPath, false);
        }
        else {
            const lineCount = lineCounts[childPath];
            const lineCountSuffix = lineCount !== undefined ? ` (${lineCount} lines)` : '';
            result += `${prefix}${child.name}${lineCountSuffix}\n`;
        }
    }
    return result;
};
export const generateTreeString = (files, emptyDirPaths = []) => {
    const tree = generateFileTree(files, emptyDirPaths);
    return treeToString(tree).trim();
};
export const generateTreeStringWithLineCounts = (files, lineCounts, emptyDirPaths = []) => {
    const tree = generateFileTree(files, emptyDirPaths);
    return treeToStringWithLineCounts(tree, lineCounts).trim();
};
const generateMultiRootSections = (filesByRoot, treeToStringFn) => {
    const sections = [];
    for (const { rootLabel, files } of filesByRoot) {
        if (files.length === 0) {
            continue;
        }
        const tree = generateFileTree(files);
        const treeContent = treeToStringFn(tree, '  ');
        sections.push(`[${rootLabel}]/\n${treeContent}`);
    }
    return sections.join('\n').trim();
};
export const generateTreeStringWithRoots = (filesByRoot, emptyDirPaths = []) => {
    if (filesByRoot.length === 1) {
        return generateTreeString(filesByRoot[0].files, emptyDirPaths);
    }
    return generateMultiRootSections(filesByRoot, (tree, prefix) => treeToString(tree, prefix));
};
export const generateTreeStringWithRootsAndLineCounts = (filesByRoot, lineCounts, emptyDirPaths = []) => {
    if (filesByRoot.length === 1) {
        return generateTreeStringWithLineCounts(filesByRoot[0].files, lineCounts, emptyDirPaths);
    }
    return generateMultiRootSections(filesByRoot, (tree, prefix) => treeToStringWithLineCounts(tree, lineCounts, prefix));
};