@needle-tools/engine/plugins/vite/treeshake.js

Needle Engine is a web-based runtime for 3D apps. It runs on your machine for development with great integrations into editors like Unity or Blender - and can be deployed onto any device! It is flexible, extensible and networking and XR are built-in.

import { existsSync, readFileSync, readdirSync, statSync } from 'fs';
import path from 'path';
import { tryLoadProjectConfig } from './config.js';
import { needleLog } from './logging.js';

/**
 * Scans local GLB/glTF files to determine which engine components are actually used,
 * then rewrites the engine's `codegen/register_types.ts` so that unused components
 * are registered lazily via `TypeStore.addLazy()` instead of eagerly imported.
 *
 * @param {"build" | "serve"} command
 * @param {import('../types/needleConfig').needleMeta | null | undefined} config
 * @param {import('../types').userSettings} userSettings
 * @returns {import('vite').Plugin | undefined}
 */
export function needleTreeshake(command, config, userSettings) {

    if (command !== 'build') return;

    // Temporarily disabled: the treeshake plugin's dynamic imports are ineffective
    // because codegen/components.ts statically re-exports all components via the
    // `export *` chain from needle-engine.ts. The dynamic import() calls in
    // register_types cannot create split points for already-statically-imported modules.
    // TODO: Find a way to also transform components.ts without breaking the public API.
    return;

    /** @type {Set<string>} component names found in local GLB/glTF files */
    const usedComponents = new Set();

    return {
        name: 'needle:treeshake',
        apply: 'build',

        buildStart() {
            const projectConfig = tryLoadProjectConfig();
            const assetsDir = projectConfig?.assetsDirectory || 'assets';
            const assetsPath = path.resolve(process.cwd(), assetsDir);

            if (!existsSync(assetsPath)) {
                needleLog('needle-treeshake', `Assets directory "${assetsPath}" not found — skipping component scanning`);
                return;
            }

            const files = findGltfFiles(assetsPath);
            needleLog('needle-treeshake', `Found ${files.length} GLB/glTF file(s) in "${assetsDir}"`);

            for (const file of files) {
                try {
                    const names = extractComponentNames(file);
                    for (const name of names) usedComponents.add(name);
                }
                catch (err) {
                    needleLog('needle-treeshake', `Failed to parse "${file}": ${err.message}`, 'warn');
                }
            }

            if (usedComponents.size > 0) {
                needleLog('needle-treeshake', `Components used in local scenes: ${[...usedComponents].sort().join(', ')}`);
            }
            else {
                needleLog('needle-treeshake', 'No components found in local scenes — all components will be loaded eagerly');
            }
        },

        /**
         * @param {string} src
         * @param {string} id
         */
        transform(src, id) {
            if (usedComponents.size === 0) return;

            if (id.includes('engine/codegen/register_types')) {
                return { code: rewriteRegisterTypes(src, usedComponents), map: null };
            }
        }
    };
}


/**
 * Recursively find all .glb and .gltf files in a directory
 * @param {string} dir
 * @returns {string[]}
 */
function findGltfFiles(dir) {
    /** @type {string[]} */
    const results = [];
    try {
        for (const entry of readdirSync(dir)) {
            const full = path.join(dir, entry);
            try {
                const stat = statSync(full);
                if (stat.isDirectory()) {
                    results.push(...findGltfFiles(full));
                }
                else if (entry.endsWith('.glb') || entry.endsWith('.gltf')) {
                    results.push(full);
                }
            }
            catch { /* skip inaccessible entries */ }
        }
    }
    catch { /* skip inaccessible directories */ }
    return results;
}


/**
 * Extract all NEEDLE_components names from a GLB or glTF file
 * @param {string} filePath
 * @returns {Set<string>}
 */
function extractComponentNames(filePath) {
    /** @type {Set<string>} */
    const names = new Set();

    let json;
    if (filePath.endsWith('.glb')) {
        const buf = readFileSync(filePath);
        if (buf.length < 20) return names;
        const jsonLen = buf.readUInt32LE(12);
        json = JSON.parse(buf.slice(20, 20 + jsonLen).toString('utf8'));
    }
    else {
        json = JSON.parse(readFileSync(filePath, 'utf8'));
    }

    if (json.nodes) {
        for (const node of json.nodes) {
            const comps = node?.extensions?.NEEDLE_components?.builtin_components;
            if (comps) {
                for (const comp of comps) {
                    if (comp?.name) names.add(comp.name);
                }
            }
        }
    }

    return names;
}


/**
 * Rewrite the engine's register_types.ts source:
 * - Eager components: keep static imports + TypeStore.add()
 * - Lazy components: per-file dynamic import() + TypeStore.addLazy()
 *
 * @param {string} src
 * @param {Set<string>} usedComponents
 * @returns {string}
 */
function rewriteRegisterTypes(src, usedComponents) {
    /** @type {Map<string, { importPath: string, importLine: string }>} */
    const importMap = new Map();
    const importRegex = /^import\s*\{\s*(\w+)\s*\}\s*from\s*"([^"]+)"\s*;?\s*$/gm;
    let match;
    while ((match = importRegex.exec(src)) !== null) {
        importMap.set(match[1], { importPath: match[2], importLine: match[0] });
    }

    /** @type {Map<string, string>} */
    const addMap = new Map();
    const addRegex = /TypeStore\.add\("(\w+)",\s*(\w+)\)/g;
    while ((match = addRegex.exec(src)) !== null) {
        addMap.set(match[1], match[2]);
    }

    /** @type {{ name: string, varName: string, importPath: string, importLine: string }[]} */
    const eager = [];
    /** @type {{ name: string, varName: string, importPath: string, importLine: string }[]} */
    const lazy = [];

    for (const [compName, varName] of addMap) {
        const imp = importMap.get(varName);
        if (!imp) continue;
        if (usedComponents.has(compName)) {
            eager.push({ name: compName, varName, ...imp });
        }
        else {
            lazy.push({ name: compName, varName, ...imp });
        }
    }

    needleLog('needle-treeshake', `Eager: ${eager.length} components, Lazy: ${lazy.length} components`);

    const lines = [];
    lines.push('/* eslint-disable */');
    lines.push('import { TypeStore } from "./../engine_typestore.js"');
    lines.push('');

    const eagerPaths = new Set(eager.map(e => e.importLine));
    lines.push('// Eagerly loaded components (used in local scenes)');
    for (const importLine of eagerPaths) {
        lines.push(importLine);
    }

    lines.push('');
    lines.push('// Register types');
    lines.push('export function initBuiltinTypes() {');

    for (const e of eager) {
        lines.push(`\tTypeStore.add("${e.name}", ${e.varName});`);
    }

    if (lazy.length > 0) {
        lines.push('');
        lines.push(`\t// Lazy-loaded components (${lazy.length} components, loaded on demand)`);
        for (const comp of lazy) {
            lines.push(`\tTypeStore.addLazy("${comp.name}", () => import("${comp.importPath}").then(m => m.${comp.varName}));`);
        }
    }

    lines.push('}');

    return lines.join('\n');
}