Improve codegen typing

.changeset/eleven-knives-complain.md

@@ -0,0 +1,5 @@
+---
+"@atproto/lex-cli": patch
+---
+
+Improve typing of isX and validateX return values

.changeset/fast-points-walk.md

@@ -0,0 +1,5 @@
+---
+"@atproto/syntax": patch
+---
+
+Improve performance of isValidTid

.changeset/popular-shirts-rescue.md

@@ -0,0 +1,5 @@
+---
+"@atproto/lexicon": patch
+---
+
+Various performance improvements

.changeset/silly-starfishes-pay.md

@@ -0,0 +1,5 @@
+---
+"@atproto/lexicon": patch
+---
+
+Add generic parameter to ValidationResult

packages/lex-cli/src/codegen/client.ts

@@ -452,15 +452,24 @@ const lexiconTs = (project, lexicons: Lexicons, lexiconDoc: LexiconDoc) =>
           moduleSpecifier: '@atproto/lexicon',
         })
         .addNamedImports([{ name: 'ValidationResult' }, { name: 'BlobRef' }])
-      //= import {isObj, hasProp} from '../../util.ts'
+
+      //= import {CID} from 'multiformats/cid'
+      file
+        .addImportDeclaration({
+          moduleSpecifier: 'multiformats/cid',
+        })
+        .addNamedImports([{ name: 'CID' }])
+
+      //= import { $Type, is$typed } from '../../util.ts'
       file
         .addImportDeclaration({
           moduleSpecifier: `${lexiconDoc.id
             .split('.')
             .map((_str) => '..')
             .join('/')}/util`,
         })
-        .addNamedImports([{ name: 'isObj' }, { name: 'hasProp' }])
+        .addNamedImports([{ name: '$Type' }, { name: 'is$typed' }])
+
       //= import {lexicons} from '../../lexicons.ts'
       file
         .addImportDeclaration({
@@ -470,12 +479,15 @@ const lexiconTs = (project, lexicons: Lexicons, lexiconDoc: LexiconDoc) =>
             .join('/')}/lexicons`,
         })
         .addNamedImports([{ name: 'lexicons' }])
-      //= import {CID} from 'multiformats/cid'
-      file
-        .addImportDeclaration({
-          moduleSpecifier: 'multiformats/cid',
-        })
-        .addNamedImports([{ name: 'CID' }])
+
+      //= const id = "{lexiconDoc.id}"
+      file.addVariableStatement({
+        isExported: false, // Do not export to allow tree-shaking
+        declarationKind: VariableDeclarationKind.Const,
+        declarations: [
+          { name: 'id', initializer: JSON.stringify(lexiconDoc.id) },
+        ],
+      })
 
       for (const defId in lexiconDoc.defs) {
         const def = lexiconDoc.defs[defId]

packages/lex-cli/src/codegen/common.ts

@@ -14,15 +14,41 @@ const PRETTIER_OPTS = {
 export const utilTs = (project) =>
   gen(project, '/util.ts', async (file) => {
     file.replaceWithText(`
-export function isObj(v: unknown): v is Record<string, unknown> {
-  return typeof v === 'object' && v !== null
+export type $Typed<V> = V & { $type: string }
+
+export type $Type<Id extends string, Hash extends string> = Hash extends 'main'
+  ? Id | \`\${Id}#\${Hash}\`
+  : \`\${Id}#\${Hash}\`
+
+function has$type<V>(v: V): v is $Typed<V & object> {
+  return (
+    v != null &&
+    typeof v === 'object' &&
+    '$type' in v &&
+    typeof v.$type === 'string'
+  )
+}
+
+function check$type<Id extends string, Hash extends string>(
+  $type: string,
+  id: Id,
+  hash: Hash,
+): $type is $Type<Id, Hash> {
+  return $type === id
+    ? hash === 'main'
+    : // $type === \`\${id}#\${hash}\`
+      $type.length === id.length + 1 + hash.length &&
+        $type.charCodeAt(id.length) === 35 /* '#' */ &&
+        $type.startsWith(id) &&
+        $type.endsWith(hash)
 }
 
-export function hasProp<K extends PropertyKey>(
-  data: object,
-  prop: K,
-): data is Record<K, unknown> {
-  return prop in data
+export function is$typed<V, Id extends string, Hash extends string>(
+  v: V,
+  id: Id,
+  hash: Hash,
+): v is V & object & { $type: $Type<Id, Hash> } {
+  return has$type(v) && check$type(v.$type, id, hash)
 }
 `)
   })

packages/lex-cli/src/codegen/lex-gen.ts

@@ -187,14 +187,16 @@ export function genObject(
 }
 
 export function genToken(file: SourceFile, lexUri: string, def: LexToken) {
+  //= /** <comment> */
+  //= export const <TOKEN> = `${id}#<token>`
   genComment(
     file.addVariableStatement({
       isExported: true,
       declarationKind: VariableDeclarationKind.Const,
       declarations: [
         {
           name: toScreamingSnakeCase(getHash(lexUri)),
-          initializer: `"${stripScheme(lexUri)}"`,
+          initializer: `\`\${id}#${getHash(lexUri)}\``,
         },
       ],
     }),
@@ -398,35 +400,30 @@ export function genXrpcOutput(
 export function genObjHelpers(
   file: SourceFile,
   lexUri: string,
-  ifaceName?: string,
+  ifaceName: string = toTitleCase(getHash(lexUri)),
 ) {
   const hash = getHash(lexUri)
 
-  //= export function is{X}(v: unknown): v is X {...}
+  //= export function is{X}(v: unknown): v is X & { $type: NS } {...}
   file
     .addFunction({
-      name: toCamelCase(`is-${ifaceName || hash}`),
+      name: toCamelCase(`is-${ifaceName}`),
       parameters: [{ name: `v`, type: `unknown` }],
-      returnType: `v is ${ifaceName || toTitleCase(hash)}`,
+      returnType: `v is ${ifaceName} & { $type: $Type<'${stripHash(lexUri)}', '${hash}'> }`,
       isExported: true,
     })
-    .setBodyText(
-      hash === 'main'
-        ? `return isObj(v) && hasProp(v, '$type') && (v.$type === "${lexUri}" || v.$type === "${stripHash(
-            lexUri,
-          )}")`
-        : `return isObj(v) && hasProp(v, '$type') && v.$type === "${lexUri}"`,
-    )
+    .setBodyText(`return is$typed(v, id, ${JSON.stringify(hash)})`)
 
-  //= export function validate{X}(v: unknown): ValidationResult {...}
+  //= export function validate{X}(v: unknown): ValidationResult<X> {...}
   file
     .addFunction({
-      name: toCamelCase(`validate-${ifaceName || hash}`),
+      name: toCamelCase(`validate-${ifaceName}`),
       parameters: [{ name: `v`, type: `unknown` }],
-      returnType: `ValidationResult`,
       isExported: true,
     })
-    .setBodyText(`return lexicons.validate("${lexUri}", v)`)
+    .setBodyText(
+      `return lexicons.validate(\`\${id}#${hash}\`, v) as ValidationResult<${ifaceName}>`,
+    )
 }
 
 export function stripScheme(uri: string): string {

packages/lex-cli/src/codegen/server.ts

@@ -348,6 +348,14 @@ const lexiconTs = (project, lexicons: Lexicons, lexiconDoc: LexiconDoc) =>
           moduleSpecifier: '@atproto/lexicon',
         })
         .addNamedImports([{ name: 'ValidationResult' }, { name: 'BlobRef' }])
+
+      //= import {CID} from 'multiformats/cid'
+      file
+        .addImportDeclaration({
+          moduleSpecifier: 'multiformats/cid',
+        })
+        .addNamedImports([{ name: 'CID' }])
+
       //= import {lexicons} from '../../lexicons.ts'
       file
         .addImportDeclaration({
@@ -357,21 +365,25 @@ const lexiconTs = (project, lexicons: Lexicons, lexiconDoc: LexiconDoc) =>
             .join('/')}/lexicons`,
         })
         .addNamedImports([{ name: 'lexicons' }])
-      //= import {isObj, hasProp} from '../../util.ts'
+
+      //= import { $Type, is$typed } from '../../util.ts'
       file
         .addImportDeclaration({
           moduleSpecifier: `${lexiconDoc.id
             .split('.')
             .map((_str) => '..')
             .join('/')}/util`,
         })
-        .addNamedImports([{ name: 'isObj' }, { name: 'hasProp' }])
-      //= import {CID} from 'multiformats/cid'
-      file
-        .addImportDeclaration({
-          moduleSpecifier: 'multiformats/cid',
-        })
-        .addNamedImports([{ name: 'CID' }])
+        .addNamedImports([{ name: '$Type' }, { name: 'is$typed' }])
+
+      //= const id = "{lexiconDoc.id}"
+      file.addVariableStatement({
+        isExported: false, // Do not export to allow tree-shaking
+        declarationKind: VariableDeclarationKind.Const,
+        declarations: [
+          { name: 'id', initializer: JSON.stringify(lexiconDoc.id) },
+        ],
+      })
 
       for (const defId in lexiconDoc.defs) {
         const def = lexiconDoc.defs[defId]

packages/lexicon/src/lexicons.ts

@@ -7,7 +7,6 @@ import {
   ValidationResult,
   ValidationError,
   isObj,
-  hasProp,
 } from './types'
 import {
   assertValidRecord,
@@ -17,7 +16,7 @@ import {
   assertValidXrpcMessage,
 } from './validation'
 import { toLexUri } from './util'
-import * as ComplexValidators from './validators/complex'
+import { object as validateObject } from './validators/complex'
 
 /**
  * A collection of compiled lexicons.
@@ -127,15 +126,17 @@ export class Lexicons implements Iterable<LexiconDoc> {
    * Validate a record or object.
    */
   validate(lexUri: string, value: unknown): ValidationResult {
-    lexUri = toLexUri(lexUri)
-    const def = this.getDefOrThrow(lexUri, ['record', 'object'])
     if (!isObj(value)) {
       throw new ValidationError(`Value must be an object`)
     }
+
+    const lexUriNormalized = toLexUri(lexUri)
+    const def = this.getDefOrThrow(lexUriNormalized, ['record', 'object'])
+
     if (def.type === 'record') {
-      return ComplexValidators.object(this, 'Record', def.record, value)
+      return validateObject(this, 'Record', def.record, value)
     } else if (def.type === 'object') {
-      return ComplexValidators.object(this, 'Object', def, value)
+      return validateObject(this, 'Object', def, value)
     } else {
       // shouldn't happen
       throw new InvalidLexiconError('Definition must be a record or object')
@@ -146,20 +147,25 @@ export class Lexicons implements Iterable<LexiconDoc> {
    * Validate a record and throw on any error.
    */
   assertValidRecord(lexUri: string, value: unknown) {
-    lexUri = toLexUri(lexUri)
-    const def = this.getDefOrThrow(lexUri, ['record'])
     if (!isObj(value)) {
       throw new ValidationError(`Record must be an object`)
     }
-    if (!hasProp(value, '$type') || typeof value.$type !== 'string') {
+    if (!('$type' in value)) {
       throw new ValidationError(`Record/$type must be a string`)
     }
-    const $type = (value as Record<string, string>).$type || ''
-    if (toLexUri($type) !== lexUri) {
+    const { $type } = value
+    if (typeof $type !== 'string') {
+      throw new ValidationError(`Record/$type must be a string`)
+    }
+
+    const lexUriNormalized = toLexUri(lexUri)
+    if (toLexUri($type) !== lexUriNormalized) {
       throw new ValidationError(
-        `Invalid $type: must be ${lexUri}, got ${$type}`,
+        `Invalid $type: must be ${lexUriNormalized}, got ${$type}`,
       )
     }
+
+    const def = this.getDefOrThrow(lexUriNormalized, ['record'])
     return assertValidRecord(this, def as LexRecord, value)
   }
 

packages/lexicon/src/types.ts

@@ -447,34 +447,24 @@ export function isValidLexiconDoc(v: unknown): v is LexiconDoc {
   return lexiconDoc.safeParse(v).success
 }
 
-export function isObj(obj: unknown): obj is Record<string, unknown> {
-  return obj !== null && typeof obj === 'object'
+export function isObj<V>(v: V): v is V & object {
+  return v != null && typeof v === 'object'
 }
 
-export function hasProp<K extends PropertyKey>(
-  data: object,
-  prop: K,
-): data is Record<K, unknown> {
-  return prop in data
-}
-
-export const discriminatedObject = z.object({ $type: z.string() })
-export type DiscriminatedObject = z.infer<typeof discriminatedObject>
-export function isDiscriminatedObject(
-  value: unknown,
-): value is DiscriminatedObject {
-  return discriminatedObject.safeParse(value).success
+export type DiscriminatedObject = { $type: string }
+export function isDiscriminatedObject(v: unknown): v is DiscriminatedObject {
+  return isObj(v) && '$type' in v && typeof v.$type === 'string'
 }
 
 export function parseLexiconDoc(v: unknown): LexiconDoc {
   lexiconDoc.parse(v)
   return v as LexiconDoc
 }
 
-export type ValidationResult =
+export type ValidationResult<V = unknown> =
   | {
       success: true
-      value: unknown
+      value: V
     }
   | {
       success: false

packages/lexicon/src/util.ts

@@ -1,6 +1,4 @@
 import { z } from 'zod'
-import { Lexicons } from './lexicons'
-import { LexRefVariant, LexUserType } from './types'
 
 export function toLexUri(str: string, baseUri?: string): string {
   if (str.split('#').length > 2) {
@@ -19,19 +17,6 @@ export function toLexUri(str: string, baseUri?: string): string {
   return `lex:${str}`
 }
 
-export function toConcreteTypes(
-  lexicons: Lexicons,
-  def: LexRefVariant | LexUserType,
-): LexUserType[] {
-  if (def.type === 'ref') {
-    return [lexicons.getDefOrThrow(def.ref)]
-  } else if (def.type === 'union') {
-    return def.refs.map((ref) => lexicons.getDefOrThrow(ref)).flat()
-  } else {
-    return [def]
-  }
-}
-
 export function requiredPropertiesRefinement<
   ObjectType extends {
     required?: string[]