compile.js

/* eslint @stylistic/indent: 0 */

import {readFileSync} from 'fs';

const version = JSON.parse(readFileSync(new URL('package.json', import.meta.url))).version;

export function compile(proto) {
    return new Function(`const exports = {};\n${compileRaw(proto, {legacy: true})}\nreturn exports;`)();
}

export function compileRaw(proto, options = {}) {
    const context = buildDefaults(buildContext(proto, null), proto.syntax);
    return `${options.dev ? '' : `// code generated by pbf v${version}\n`}${writeContext(context, options)}`;
}

function writeContext(ctx, options) {
    let code = '';
    if (ctx._proto.fields) code += writeMessage(ctx, options);
    if (ctx._proto.values) code += writeEnum(ctx, options);

    for (const child of ctx._children) code += writeContext(child, options);
    return code;
}

function writeMessage(ctx, options) {
    const fields = ctx._proto.fields;

    let code = '\n';

    if (!options.noRead) {
        const readName = `read${ctx._name}`;

        code += `${writeFunctionExport(options, readName)}function ${readName}(pbf, end) {\n`;
        if (fields.length === 0) {
            code += '    pbf.pos = end;\n    return {};\n';
        } else {
            code += `    const obj = ${compileDest(ctx)};\n`;
            code += '    let field;\n';
            code += '    while ((field = pbf.nextField(end))) {\n';

            for (let i = 0; i < fields.length; i++) {
                code += `        ${i ? 'else ' : ''}if (field === ${fields[i].tag}) ${compileFieldRead(ctx, fields[i])}\n`;
            }
            code += '    }\n    return obj;\n';
        }
        code += '}\n';
    }

    if (!options.noWrite) {
        const writeName = `write${ctx._name}`;
        code += `${writeFunctionExport(options, writeName)}function ${writeName}(obj, pbf) {`;
        if (fields.length === 0) {
            code += '}\n';
        } else {
            code += '\n';
            for (const field of fields) code += compileFieldWriteLine(ctx, field);
            code += '}\n';
        }
    }
    return code;
}

function compileFieldWriteLine(ctx, field) {
    const v = `obj.${field.name}`;
    let body;
    if (field.repeated && !isPacked(field)) {
        body = `for (const item of ${v}) ${compileFieldWrite(ctx, field, 'item')}`;
    } else if (field.type === 'map') {
        body = `for (const key of Object.keys(${v})) ${compileFieldWrite(ctx, field, `{key, value: ${v}[key]}`)}`;
    } else {
        body = compileFieldWrite(ctx, field, v);
    }
    return `${getDefaultWriteTest(ctx, field)}${body};\n`;
}

function writeFunctionExport({legacy}, name) {
    return legacy ? `exports.${name} = ${name};\n` : 'export ';
}

function getEnumValues(ctx) {
    const enums = {};
    for (const [name, {value}] of Object.entries(ctx._proto.values)) enums[name] = value;
    return enums;
}

function writeEnum(ctx, {legacy}) {
    const enums = JSON.stringify(getEnumValues(ctx), null, 4);
    const name = ctx._name;
    return `\n${legacy ? `const ${name} = exports.${name}` : `export const ${name}`} = ${enums};\n`;
}

function compileDest(ctx) {
    const props = new Set();
    for (const {name, oneof} of ctx._proto.fields) {
        props.add(`${name}: ${JSON.stringify(ctx._defaults[name])}`);
        if (oneof) props.add(`${oneof}: undefined`);
    }
    return `{${[...props].join(', ')}}`;
}

function isEnum(type) {
    return type && type._proto.values;
}

function getType(ctx, field) {
    if (field.type === 'map') return ctx._mapEntries[field.name];
    return field.type.split('.').reduce((ctx, name) => ctx && ctx[name], ctx);
}

function capitalize(s) {
    return s.charAt(0).toUpperCase() + s.slice(1);
}

function getSubMessage(ctx, field) {
    const type = getType(ctx, field);
    return type && type._proto.fields ? type : null;
}

// Map of scalar protobuf type → suffix used in read/write/readPacked method names,
// plus per-type flags. `packable` matches the protobuf spec list of packable types.
// `signed` is for varint types that take a signedness flag.
const TYPES = {
    string: {m: 'String'},
    float: {m: 'Float', packable: true},
    double: {m: 'Double', packable: true},
    bool: {m: 'Boolean', packable: true},
    uint32: {m: 'Varint', packable: true},
    uint64: {m: 'Varint', packable: true},
    int32: {m: 'Varint', packable: true, signed: true},
    int64: {m: 'Varint', packable: true, signed: true},
    sint32: {m: 'SVarint', packable: true},
    sint64: {m: 'SVarint', packable: true},
    fixed32: {m: 'Fixed32', packable: true},
    fixed64: {m: 'Fixed64', packable: true},
    sfixed32: {m: 'SFixed32', packable: true},
    sfixed64: {m: 'SFixed64', packable: true},
    bytes: {m: 'Bytes'},
    enum: {m: 'Varint', packable: true},
};

function getMethod(ctx, field) {
    if (isEnum(getType(ctx, field))) return TYPES.enum;
    const t = TYPES[field.type];
    if (!t) throw new Error(`Unexpected type: ${field.type}`);
    return t;
}

// JS_STRING serializes numeric scalars as strings — applies to every TYPES entry except the non-numeric ones.
function fieldShouldUseStringAsNumber(field) {
    if (field.options.jstype !== 'JS_STRING') return false;
    const t = TYPES[field.type];
    return !!t && t.m !== 'String' && t.m !== 'Boolean' && t.m !== 'Bytes';
}

function compileScalarFieldRead(ctx, field) {
    const sub = getSubMessage(ctx, field);
    if (sub) return `read${sub._name}(pbf, pbf.readVarint() + pbf.pos)`;

    const {m, signed} = getMethod(ctx, field);
    let suffix = `(${signed ? 'true' : ''})`;
    if (fieldShouldUseStringAsNumber(field)) suffix += '.toString()';
    return `pbf.read${m}${suffix}`;
}

function compilePackedRead(ctx, field) {
    const {m, signed} = getMethod(ctx, field);
    return `pbf.readPacked${m}(obj.${field.name}${signed ? ', true' : ''})`;
}

function compileFieldRead(ctx, field) {
    const {type, name, repeated, oneof} = field;
    let body;
    if (repeated && willSupportPacked(ctx, field)) {
        body = compilePackedRead(ctx, field);
    } else {
        const scalar = compileScalarFieldRead(ctx, field);
        if (type === 'map') body = `const {key, value} = ${scalar}; obj.${name}[key] = value`;
        else if (repeated) body = `obj.${name}.push(${scalar})`;
        else body = `obj.${name} = ${scalar}`;
    }
    if (oneof) body += `; obj.${oneof} = ${JSON.stringify(name)}`;
    return type === 'map' || oneof ? `{ ${body}; }` : `${body};`;
}

function compileFieldWrite(ctx, field, name) {
    const sub = getSubMessage(ctx, field);
    if (sub) return `pbf.writeMessage(${field.tag}, write${sub._name}, ${name})`;

    if (fieldShouldUseStringAsNumber(field)) {
        name = field.type === 'float' || field.type === 'double' ?
            `parseFloat(${name})` : `parseInt(${name}, 10)`;
    }
    const {m} = getMethod(ctx, field);
    const fn = isPacked(field) ? `writePacked${m}` : `write${m}Field`;
    return `pbf.${fn}(${field.tag}, ${name})`;
}

function getMapMessage(field) {
    const f = (name, type, tag) => ({name, type, tag, oneof: null, repeated: false, options: {}});
    return {
        name: `${capitalize(field.name)}Entry`,
        fields: [f('key', field.map.from, 1), f('value', field.map.to, 2)],
    };
}

// Protobuf identifier per spec: starts with a letter or underscore, followed by letters, digits, or underscores.
// Validated to prevent code injection via untrusted .proto schemas, since names are interpolated into generated JS.
const identifierRegex = /^[A-Za-z_][A-Za-z0-9_]*$/;

function validateIdentifier(name) {
    if (typeof name !== 'string' || !identifierRegex.test(name)) {
        throw new Error(`Invalid protobuf identifier: ${JSON.stringify(name)}`);
    }
}

function buildContext(proto, parent, mapEntryField) {
    const obj = Object.create(parent);
    obj._proto = proto;
    obj._children = [];
    obj._defaults = {};
    obj._mapEntries = {};

    if (parent) {
        obj._name = (parent._name ?? '') + proto.name;
        if (mapEntryField) {
            parent._mapEntries[mapEntryField.name] = obj;
        } else {
            validateIdentifier(proto.name);
            parent[proto.name] = obj;
        }
    }

    for (const field of proto.fields ?? []) {
        validateIdentifier(field.name);
        if (field.oneof) validateIdentifier(field.oneof);
    }
    for (const valueName of Object.keys(proto.values ?? {})) validateIdentifier(valueName);

    for (const e of proto.enums ?? []) obj._children.push(buildContext(e, obj));
    for (const m of proto.messages ?? []) obj._children.push(buildContext(m, obj));
    for (const f of proto.fields ?? []) {
        if (f.type !== 'map') continue;
        const entryProto = getMapMessage(f);
        // Disambiguate against sibling messages/enums with a `$` suffix — `$` is a valid
        // JS identifier char but disallowed in protobuf identifiers, so it can't re-collide.
        if (obj._children.some(c => c._proto.name === entryProto.name)) entryProto.name += '$';
        obj._children.push(buildContext(entryProto, obj, f));
    }

    return obj;
}

function getDefaultValue(field, value) {
    if (field.repeated) return []; // defaults not supported for repeated fields
    if (field.type === 'map') return {};

    const t = TYPES[field.type];
    if (!t || t.m === 'Bytes') return undefined; // unknown / message / bytes
    if (t.m === 'String') return value || '';
    if (t.m === 'Boolean') return value === 'true';

    const isFloat = t.m === 'Float' || t.m === 'Double';
    const num = value ? (isFloat ? parseFloat(value) : parseInt(value, 10)) : 0;
    return fieldShouldUseStringAsNumber(field) ? num.toString() : num;
}

function willSupportPacked(ctx, field) {
    if (!field.repeated || getSubMessage(ctx, field)) return false;
    return !!getMethod(ctx, field).packable;
}

function setPackedOption(ctx, field, syntax) {
    // proto3 packs eligible repeated fields by default; older syntax requires explicit [packed=true].
    if (syntax >= 3 && field.options.packed === undefined && willSupportPacked(ctx, field)) {
        field.options.packed = 'true';
    }
}

function setDefaultValue(ctx, field, syntax) {
    const type = getType(ctx, field);
    // Proto3 does not support overriding defaults
    const explicitDefault = syntax < 3 ? field.options.default : undefined;

    ctx._defaults[field.name] = isEnum(type) && !field.repeated ?
        (getEnumValues(type)[explicitDefault] || 0) :
        getDefaultValue(field, explicitDefault);
}

function buildDefaults(ctx, syntax) {
    for (const child of ctx._children) buildDefaults(child, syntax);
    for (const field of ctx._proto.fields ?? []) {
        setPackedOption(ctx, field, syntax);
        setDefaultValue(ctx, field, syntax);
    }
    return ctx;
}

function getDefaultWriteTest(ctx, field) {
    const def = ctx._defaults[field.name];
    let code = `    if (obj.${field.name}`;

    if (!field.repeated && !getSubMessage(ctx, field)) {
        if (def === undefined || def || field.oneof) code += ' != null';
        if (def) code += ` && obj.${field.name} !== ${JSON.stringify(def)}`;
    }
    return `${code}) `;
}

function isPacked(field) {
    return field.options.packed === 'true';
}