1+ import { combineComments } from './comments' ;
12import { isPrimitive , type Primitive , type Schema } from './ir' ;
23
34export type OptimizeFn = ( schema : Schema ) => Schema ;
@@ -173,6 +174,15 @@ export function filterUndefinedUnion(schema: Schema): [boolean, Schema] {
173174 }
174175}
175176
177+ // This function is a helper that adds back any comments that were removed during optimization
178+ function withComment ( newSchema : Schema , oldSchema : Schema ) : Schema {
179+ if ( oldSchema . comment ) {
180+ newSchema . comment = combineComments ( oldSchema ) ;
181+ }
182+
183+ return newSchema ;
184+ }
185+
176186export function optimize ( schema : Schema ) : Schema {
177187 if ( schema . type === 'object' ) {
178188 const properties : Record < string , Schema > = { } ;
@@ -184,8 +194,8 @@ export function optimize(schema: Schema): Schema {
184194 }
185195 const [ isOptional , filteredSchema ] = filterUndefinedUnion ( optimized ) ;
186196
187- if ( prop . comment ) {
188- filteredSchema . comment = prop . comment ;
197+ if ( optimized . comment ) {
198+ filteredSchema . comment = optimized . comment ;
189199 }
190200
191201 properties [ key ] = filteredSchema ;
@@ -197,48 +207,31 @@ export function optimize(schema: Schema): Schema {
197207
198208 const schemaObject : Schema = { type : 'object' , properties, required } ;
199209
200- // only add comment field if there is a comment
201- if ( schema . comment ) {
202- return { ...schemaObject , comment : schema . comment } ;
203- }
204-
205- return schemaObject ;
210+ return withComment ( schemaObject , schema ) ;
206211 } else if ( schema . type === 'intersection' ) {
207212 const newSchema = foldIntersection ( schema , optimize ) ;
208- if ( schema . comment ) {
209- return { ...newSchema , comment : schema . comment } ;
210- }
211- return newSchema ;
213+
214+ return withComment ( newSchema , schema ) ;
212215 } else if ( schema . type === 'union' ) {
213216 const consolidated = consolidateUnion ( schema ) ;
214217 const simplified = simplifyUnion ( consolidated , optimize ) ;
215218 const merged = mergeUnions ( simplified ) ;
216219
217- if ( schema . comment ) {
218- return { ...merged , comment : schema . comment } ;
219- }
220-
221- return merged ;
220+ return withComment ( merged , schema ) ;
222221 } else if ( schema . type === 'array' ) {
223222 const optimized = optimize ( schema . items ) ;
224- if ( schema . comment ) {
225- return { type : 'array' , items : optimized , comment : schema . comment } ;
226- }
227- return { type : 'array' , items : optimized } ;
223+
224+ return withComment ( { type : 'array' , items : optimized } , schema ) ;
228225 } else if ( schema . type === 'record' ) {
229- return {
226+ return withComment ( {
230227 type : 'record' ,
231228 ...( schema . domain ? { domain : optimize ( schema . domain ) } : { } ) ,
232229 codomain : optimize ( schema . codomain ) ,
233- ...( schema . comment ? { comment : schema . comment } : { } ) ,
234- } ;
230+ } , schema )
235231 } else if ( schema . type === 'tuple' ) {
236232 const schemas = schema . schemas . map ( optimize ) ;
237- return { type : 'tuple' , schemas } ;
233+ return withComment ( { type : 'tuple' , schemas } , schema ) ;
238234 } else if ( schema . type === 'ref' ) {
239- if ( schema . comment ) {
240- return { ...schema , comment : schema . comment } ;
241- }
242235 return schema ;
243236 } else {
244237 return schema ;
0 commit comments