Skip to content

Commit 20cea29

Browse files
streamichstreamich
authored
feat: move capacity estimation codegen into a separate module
2 parents 465cc7d + ec46d1c commit 20cea29

File tree

16 files changed

+308
-163
lines changed

16 files changed

+308
-163
lines changed

‎src/codegen/capacity/__tests__/CapacityEstimatorCodegenContext.spec.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,3 +215,46 @@ describe('or', () => {
215215
expect(estimator([1, 2, 3])).toBe(maxEncodingCapacity([1, 2, 3]));
216216
});
217217
});
218+
219+
describe('standalone codegen function', () => {
220+
test('generates capacity estimator equivalent to compileCapacityEstimator', () => {
221+
const system = new TypeSystem();
222+
const type = system.t.Array(system.t.str);
223+
224+
// Compare standalone codegen function with the class method
225+
const {codegen} = require('../estimators');
226+
const standaloneEstimator = codegen(type, {});
227+
const classEstimator = type.compileCapacityEstimator({});
228+
229+
const testData = ['hello', 'world', 'test'];
230+
expect(standaloneEstimator(testData)).toBe(classEstimator(testData));
231+
expect(standaloneEstimator(testData)).toBe(maxEncodingCapacity(testData));
232+
});
233+
234+
test('works with complex nested types', () => {
235+
const system = new TypeSystem();
236+
const type = system.t.Object(
237+
system.t.prop('name', system.t.str),
238+
system.t.prop('items', system.t.Array(system.t.num)),
239+
);
240+
241+
const {codegen} = require('../estimators');
242+
const standaloneEstimator = codegen(type, {});
243+
const classEstimator = type.compileCapacityEstimator({});
244+
245+
const testData = {name: 'test', items: [1, 2, 3, 4, 5]};
246+
expect(standaloneEstimator(testData)).toBe(classEstimator(testData));
247+
expect(standaloneEstimator(testData)).toBe(maxEncodingCapacity(testData));
248+
});
249+
250+
test('works with const types', () => {
251+
const system = new TypeSystem();
252+
const type = system.t.Const('hello world');
253+
254+
const {codegen} = require('../estimators');
255+
const standaloneEstimator = codegen(type, {});
256+
257+
// For const types, the value doesn't matter
258+
expect(standaloneEstimator(null)).toBe(maxEncodingCapacity('hello world'));
259+
});
260+
});

‎src/codegen/capacity/estimators.ts‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,255 @@
1+
import {JsExpression} from '@jsonjoy.com/util/lib/codegen/util/JsExpression';
2+
import {MaxEncodingOverhead, maxEncodingCapacity} from '@jsonjoy.com/util/lib/json-size';
3+
import {CapacityEstimatorCodegenContext} from './CapacityEstimatorCodegenContext';
4+
import type {
5+
CapacityEstimatorCodegenContextOptions,
6+
CompiledCapacityEstimator,
7+
} from './CapacityEstimatorCodegenContext';
8+
import type {Type} from '../../type';
9+
10+
type EstimatorFunction = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type) => void;
11+
12+
const normalizeAccessor = (key: string): string => {
13+
// Simple property access for valid identifiers, bracket notation otherwise
14+
if (/^[a-zA-Z_$][a-zA-Z0-9_$]*$/.test(key)) {
15+
return `.${key}`;
16+
}
17+
return `[${JSON.stringify(key)}]`;
18+
};
19+
20+
export const any = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
21+
const codegen = ctx.codegen;
22+
codegen.link('Value');
23+
const r = codegen.var(value.use());
24+
codegen.if(
25+
`${r} instanceof Value`,
26+
() => {
27+
codegen.if(
28+
`${r}.type`,
29+
() => {
30+
ctx.codegen.js(`size += ${r}.type.capacityEstimator()(${r}.data);`);
31+
},
32+
() => {
33+
ctx.codegen.js(`size += maxEncodingCapacity(${r}.data);`);
34+
},
35+
);
36+
},
37+
() => {
38+
ctx.codegen.js(`size += maxEncodingCapacity(${r});`);
39+
},
40+
);
41+
};
42+
43+
export const bool = (ctx: CapacityEstimatorCodegenContext, value: JsExpression): void => {
44+
ctx.inc(MaxEncodingOverhead.Boolean);
45+
};
46+
47+
export const num = (ctx: CapacityEstimatorCodegenContext, value: JsExpression): void => {
48+
ctx.inc(MaxEncodingOverhead.Number);
49+
};
50+
51+
export const str = (ctx: CapacityEstimatorCodegenContext, value: JsExpression): void => {
52+
ctx.inc(MaxEncodingOverhead.String);
53+
ctx.codegen.js(`size += ${MaxEncodingOverhead.StringLengthMultiplier} * ${value.use()}.length;`);
54+
};
55+
56+
export const bin = (ctx: CapacityEstimatorCodegenContext, value: JsExpression): void => {
57+
ctx.inc(MaxEncodingOverhead.Binary);
58+
ctx.codegen.js(`size += ${MaxEncodingOverhead.BinaryLengthMultiplier} * ${value.use()}.length;`);
59+
};
60+
61+
export const const_ = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
62+
const constType = type as any; // ConstType
63+
ctx.inc(maxEncodingCapacity(constType.value()));
64+
};
65+
66+
export const arr = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
67+
const codegen = ctx.codegen;
68+
ctx.inc(MaxEncodingOverhead.Array);
69+
const rLen = codegen.var(`${value.use()}.length`);
70+
const arrayType = type as any; // ArrayType
71+
const elementType = arrayType.type;
72+
codegen.js(`size += ${MaxEncodingOverhead.ArrayElement} * ${rLen}`);
73+
const fn = elementType.compileCapacityEstimator({
74+
system: ctx.options.system,
75+
name: ctx.options.name,
76+
});
77+
const isConstantSizeType = ['const', 'bool', 'num'].includes(elementType.getTypeName());
78+
if (isConstantSizeType) {
79+
const rFn = codegen.linkDependency(fn);
80+
codegen.js(`size += ${rLen} * ${rFn}(${elementType.random()});`);
81+
} else {
82+
const rFn = codegen.linkDependency(fn);
83+
const ri = codegen.var('0');
84+
codegen.js(`for (; ${ri} < ${rLen}; ${ri}++) {`);
85+
codegen.js(`size += ${rFn}(${value.use()}[${ri}]);`);
86+
codegen.js(`}`);
87+
}
88+
};
89+
90+
export const tup = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
91+
const codegen = ctx.codegen;
92+
const r = codegen.var(value.use());
93+
const tupleType = type as any; // TupleType
94+
const types = tupleType.types;
95+
const overhead = MaxEncodingOverhead.Array + MaxEncodingOverhead.ArrayElement * types.length;
96+
ctx.inc(overhead);
97+
for (let i = 0; i < types.length; i++) {
98+
const elementType = types[i];
99+
const fn = elementType.compileCapacityEstimator({
100+
system: ctx.options.system,
101+
name: ctx.options.name,
102+
});
103+
const rFn = codegen.linkDependency(fn);
104+
codegen.js(`size += ${rFn}(${r}[${i}]);`);
105+
}
106+
};
107+
108+
export const obj = (
109+
ctx: CapacityEstimatorCodegenContext,
110+
value: JsExpression,
111+
type: Type,
112+
estimateCapacityFn: EstimatorFunction,
113+
): void => {
114+
const codegen = ctx.codegen;
115+
const r = codegen.var(value.use());
116+
const objectType = type as any; // ObjectType
117+
const encodeUnknownFields = !!objectType.schema.encodeUnknownFields;
118+
if (encodeUnknownFields) {
119+
codegen.js(`size += maxEncodingCapacity(${r});`);
120+
return;
121+
}
122+
const fields = objectType.fields;
123+
const overhead = MaxEncodingOverhead.Object + fields.length * MaxEncodingOverhead.ObjectElement;
124+
ctx.inc(overhead);
125+
for (const field of fields) {
126+
ctx.inc(maxEncodingCapacity(field.key));
127+
const accessor = normalizeAccessor(field.key);
128+
const isOptional = field.optional || field.constructor?.name === 'ObjectOptionalFieldType';
129+
const block = () => estimateCapacityFn(ctx, new JsExpression(() => `${r}${accessor}`), field.value);
130+
if (isOptional) {
131+
codegen.if(`${r}${accessor} !== undefined`, block);
132+
} else block();
133+
}
134+
};
135+
136+
export const map = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
137+
const codegen = ctx.codegen;
138+
ctx.inc(MaxEncodingOverhead.Object);
139+
const r = codegen.var(value.use());
140+
const rKeys = codegen.var(`Object.keys(${r})`);
141+
const rKey = codegen.var();
142+
const rLen = codegen.var(`${rKeys}.length`);
143+
codegen.js(`size += ${MaxEncodingOverhead.ObjectElement} * ${rLen}`);
144+
const mapType = type as any; // MapType
145+
const valueType = mapType.type;
146+
const fn = valueType.compileCapacityEstimator({
147+
system: ctx.options.system,
148+
name: ctx.options.name,
149+
});
150+
const rFn = codegen.linkDependency(fn);
151+
const ri = codegen.var('0');
152+
codegen.js(`for (; ${ri} < ${rLen}; ${ri}++) {`);
153+
codegen.js(`${rKey} = ${rKeys}[${ri}];`);
154+
codegen.js(`size += maxEncodingCapacity(${rKey}) + ${rFn}(${r}[${rKey}]);`);
155+
codegen.js(`}`);
156+
};
157+
158+
export const ref = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
159+
const refType = type as any; // RefType
160+
const system = ctx.options.system || refType.system;
161+
if (!system) throw new Error('NO_SYSTEM');
162+
const estimator = system.resolve(refType.schema.ref).type.capacityEstimator();
163+
const d = ctx.codegen.linkDependency(estimator);
164+
ctx.codegen.js(`size += ${d}(${value.use()});`);
165+
};
166+
167+
export const or = (
168+
ctx: CapacityEstimatorCodegenContext,
169+
value: JsExpression,
170+
type: Type,
171+
estimateCapacityFn: EstimatorFunction,
172+
): void => {
173+
const codegen = ctx.codegen;
174+
const orType = type as any; // OrType
175+
const discriminator = orType.discriminator();
176+
const d = codegen.linkDependency(discriminator);
177+
const types = orType.types;
178+
codegen.switch(
179+
`${d}(${value.use()})`,
180+
types.map((childType: Type, index: number) => [
181+
index,
182+
() => {
183+
estimateCapacityFn(ctx, value, childType);
184+
},
185+
]),
186+
);
187+
};
188+
189+
/**
190+
* Main router function that dispatches capacity estimation to the appropriate
191+
* estimator function based on the type's kind.
192+
*/
193+
export const generate = (ctx: CapacityEstimatorCodegenContext, value: JsExpression, type: Type): void => {
194+
const kind = type.getTypeName();
195+
196+
switch (kind) {
197+
case 'any':
198+
any(ctx, value, type);
199+
break;
200+
case 'bool':
201+
bool(ctx, value);
202+
break;
203+
case 'num':
204+
num(ctx, value);
205+
break;
206+
case 'str':
207+
str(ctx, value);
208+
break;
209+
case 'bin':
210+
bin(ctx, value);
211+
break;
212+
case 'const':
213+
const_(ctx, value, type);
214+
break;
215+
case 'arr':
216+
arr(ctx, value, type);
217+
break;
218+
case 'tup':
219+
tup(ctx, value, type);
220+
break;
221+
case 'obj':
222+
obj(ctx, value, type, generate);
223+
break;
224+
case 'map':
225+
map(ctx, value, type);
226+
break;
227+
case 'ref':
228+
ref(ctx, value, type);
229+
break;
230+
case 'or':
231+
or(ctx, value, type, generate);
232+
break;
233+
default:
234+
throw new Error(`${kind} type capacity estimation not implemented`);
235+
}
236+
};
237+
238+
/**
239+
* Standalone function to generate a capacity estimator for a given type.
240+
*/
241+
export const codegen = (
242+
type: Type,
243+
options: Omit<CapacityEstimatorCodegenContextOptions, 'type'>,
244+
): CompiledCapacityEstimator => {
245+
const ctx = new CapacityEstimatorCodegenContext({
246+
system: type.system,
247+
...options,
248+
type: type as any,
249+
});
250+
const r = ctx.codegen.options.args[0];
251+
const value = new JsExpression(() => r);
252+
// Use the centralized router instead of the abstract method
253+
generate(ctx, value, type);
254+
return ctx.compile();
255+
};

‎src/codegen/capacity/index.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
export * from './estimators';
2+
3+
export {CapacityEstimatorCodegenContext} from './CapacityEstimatorCodegenContext';
4+
export type {
5+
CapacityEstimatorCodegenContextOptions,
6+
CompiledCapacityEstimator,
7+
} from './CapacityEstimatorCodegenContext';

‎src/type/classes/AbstractType.ts‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import {
3636
type CapacityEstimatorCodegenContextOptions,
3737
type CompiledCapacityEstimator,
3838
} from '../../codegen/capacity/CapacityEstimatorCodegenContext';
39+
import {generate} from '../../codegen/capacity/estimators';
3940
import type {JsonValueCodec} from '@jsonjoy.com/json-pack/lib/codecs/types';
4041
import type * as jsonSchema from '../../json-schema';
4142
import type {BaseType} from '../types';
@@ -265,14 +266,11 @@ export abstract class AbstractType<S extends schema.Schema> implements BaseType<
265266
});
266267
const r = ctx.codegen.options.args[0];
267268
const value = new JsExpression(() => r);
268-
this.codegenCapacityEstimator(ctx, value);
269+
// Use the centralized router instead of the abstract method
270+
generate(ctx, value, this as any);
269271
return ctx.compile();
270272
}
271273

272-
public codegenCapacityEstimator(ctx: CapacityEstimatorCodegenContext, value: JsExpression): void {
273-
throw new Error(`${this.toStringName()}.codegenCapacityEstimator() not implemented`);
274-
}
275-
276274
private __capacityEstimator: CompiledCapacityEstimator | undefined;
277275
public capacityEstimator(): CompiledCapacityEstimator {
278276
return (

‎src/type/classes/AnyType.ts‎

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -82,29 +82,6 @@ export class AnyType extends AbstractType<schema.AnySchema> {
8282
this.codegenBinaryEncoder(ctx, value);
8383
}
8484

85-
public codegenCapacityEstimator(ctx: CapacityEstimatorCodegenContext, value: JsExpression): void {
86-
const codegen = ctx.codegen;
87-
codegen.link('Value');
88-
const r = codegen.var(value.use());
89-
codegen.if(
90-
`${r} instanceof Value`,
91-
() => {
92-
codegen.if(
93-
`${r}.type`,
94-
() => {
95-
ctx.codegen.js(`size += ${r}.type.capacityEstimator()(${r}.data);`);
96-
},
97-
() => {
98-
ctx.codegen.js(`size += maxEncodingCapacity(${r}.data);`);
99-
},
100-
);
101-
},
102-
() => {
103-
ctx.codegen.js(`size += maxEncodingCapacity(${r});`);
104-
},
105-
);
106-
}
107-
10885
public random(): unknown {
10986
return RandomJson.generate({nodeCount: 5});
11087
}

0 commit comments

Comments
 (0)