enhance constraints to parse references

package.json

@@ -1,7 +1,7 @@
 {
   "name": "spex-parser",
   "repository": "github:Freelansys/spex-parser",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "A parser for Spex, a language designed for programming with LLMs",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

src/ast.ts

@@ -55,10 +55,27 @@ export type ExponentialObject = {
   exponent: ObjectExpression
 }
 
+export type ConstraintReference = {
+  kind: 'ConstraintReference'
+  name: string
+}
+
+export type ConstraintText = {
+  kind: 'ConstraintText'
+  text: string
+}
+
+export type ConstraintPart = ConstraintReference | ConstraintText
+
+export type Constraint = {
+  raw: string
+  parts: ConstraintPart[]
+}
+
 export type SubObject = {
   kind: 'SubObject'
   base: ObjectExpression
-  constraint: string
+  constraint: Constraint
 }
 
 export type ArrayObject = {

src/index.ts

@@ -10,4 +10,8 @@ export type {
   ProductObject,
   ExponentialObject,
   SubObject,
+  Constraint,
+  ConstraintPart,
+  ConstraintReference,
+  ConstraintText,
 } from './ast.js'

src/visitor.ts

@@ -11,13 +11,38 @@ import type {
   ProductObject,
   SubObject,
   ArrayObject,
+  Constraint,
+  ConstraintPart,
+  ConstraintReference,
 } from './ast.js'
 import { SpexLexer } from './lexer.js'
 import { SpexParser } from './parser.js'
 
 const parserInstance = new SpexParser()
 const BaseSpexVisitor = parserInstance.getBaseCstVisitorConstructor()
 
+export const REFERENCE_PATTERN = /@([a-zA-Z_][\w]*(?:\.[a-zA-Z_][\w]*)*)/g
+
+export function parseConstraint(raw: string): Constraint {
+  const parts: ConstraintPart[] = []
+  let lastIndex = 0
+  let match: RegExpExecArray | null
+
+  while ((match = REFERENCE_PATTERN.exec(raw)) !== null) {
+    if (match.index > lastIndex) {
+      parts.push({ kind: 'ConstraintText', text: raw.slice(lastIndex, match.index) })
+    }
+    parts.push({ kind: 'ConstraintReference', name: match[1]! })
+    lastIndex = match.index + match[0].length
+  }
+
+  if (lastIndex < raw.length) {
+    parts.push({ kind: 'ConstraintText', text: raw.slice(lastIndex) })
+  }
+
+  return { raw, parts }
+}
+
 export class SpexParserVisitor extends BaseSpexVisitor implements ICstVisitor<any, any> {
   constructor() {
     super()
@@ -112,11 +137,11 @@ export class SpexParserVisitor extends BaseSpexVisitor implements ICstVisitor<an
 
   subObject(ctx: any): SubObject {
     const rawText: string = ctx.SelectBlock[0].image
-    const constraint = rawText.slice(1, -1).trim()
+    const rawConstraint = rawText.slice(1, -1).trim()
     return {
       kind: 'SubObject',
       base: this.visit(ctx.base),
-      constraint,
+      constraint: parseConstraint(rawConstraint),
     }
   }
 

tests/visitor.test.ts

@@ -1,10 +1,11 @@
 import { describe, it, expect } from 'vitest'
-import { parseToAst } from '../src/visitor.js'
+import { parseToAst, parseConstraint } from '../src/visitor.js'
 import type {
   ObjectDeclaration,
   ImportDeclaration,
   ExportDeclaration,
   GenerateDeclaration,
+  Constraint,
 } from '../src/ast.js'
 
 describe('SpexParserVisitor', () => {
@@ -92,7 +93,10 @@ describe('SpexParserVisitor', () => {
             p: {
               kind: 'SubObject',
               base: { kind: 'NamedObject', name: 'Number' },
-              constraint: 'value is positive',
+              constraint: {
+                raw: 'value is positive',
+                parts: [{ kind: 'ConstraintText', text: 'value is positive' }],
+              },
             },
             n: { kind: 'NamedObject', name: 'Number' },
           },
@@ -178,12 +182,18 @@ describe('SpexParserVisitor', () => {
           exponent: {
             kind: 'SubObject',
             base: { kind: 'NamedObject', name: 'Number' },
-            constraint: 'value is positive',
+            constraint: {
+              raw: 'value is positive',
+              parts: [{ kind: 'ConstraintText', text: 'value is positive' }],
+            },
           },
           base: {
             kind: 'SubObject',
             base: { kind: 'NamedObject', name: 'Number' },
-            constraint: 'value is positive',
+            constraint: {
+              raw: 'value is positive',
+              parts: [{ kind: 'ConstraintText', text: 'value is positive' }],
+            },
           },
         },
       })
@@ -199,7 +209,10 @@ describe('SpexParserVisitor', () => {
         object: {
           kind: 'SubObject',
           base: { kind: 'NamedObject', name: 'Number' },
-          constraint: 'isPositive',
+          constraint: {
+            raw: 'isPositive',
+            parts: [{ kind: 'ConstraintText', text: 'isPositive' }],
+          },
         },
       })
     })
@@ -214,7 +227,10 @@ describe('SpexParserVisitor', () => {
         object: {
           kind: 'SubObject',
           base: { kind: 'NamedObject', name: 'Number' },
-          constraint: 'the number is positive',
+          constraint: {
+            raw: 'the number is positive',
+            parts: [{ kind: 'ConstraintText', text: 'the number is positive' }],
+          },
         },
       })
     })
@@ -236,7 +252,13 @@ describe('SpexParserVisitor', () => {
               s: { kind: 'NamedObject', name: 'String' },
             },
           },
-          constraint: '@n is positive',
+          constraint: {
+            raw: '@n is positive',
+            parts: [
+              { kind: 'ConstraintReference', name: 'n' },
+              { kind: 'ConstraintText', text: ' is positive' },
+            ],
+          },
         },
       })
     })
@@ -262,7 +284,10 @@ describe('SpexParserVisitor', () => {
             },
             base: { kind: 'NamedObject', name: 'Bool' },
           },
-          constraint: 'logs the given input',
+          constraint: {
+            raw: 'logs the given input',
+            parts: [{ kind: 'ConstraintText', text: 'logs the given input' }],
+          },
         },
       })
     })
@@ -280,9 +305,15 @@ describe('SpexParserVisitor', () => {
           base: {
             kind: 'SubObject',
             base: { kind: 'NamedObject', name: 'Number' },
-            constraint: 'value is positive',
+            constraint: {
+              raw: 'value is positive',
+              parts: [{ kind: 'ConstraintText', text: 'value is positive' }],
+            },
+          },
+          constraint: {
+            raw: 'value is odd',
+            parts: [{ kind: 'ConstraintText', text: 'value is odd' }],
           },
-          constraint: 'value is odd',
         },
       })
     })
@@ -460,4 +491,107 @@ describe('SpexParserVisitor', () => {
       })
     })
   })
+
+  describe('parseConstraint', () => {
+    it('should parse plain text constraint with no references', () => {
+      const result = parseConstraint('value is positive')
+      expect(result).toEqual<Constraint>({
+        raw: 'value is positive',
+        parts: [{ kind: 'ConstraintText', text: 'value is positive' }],
+      })
+    })
+
+    it('should parse constraint with single reference', () => {
+      const result = parseConstraint('@n is positive')
+      expect(result).toEqual<Constraint>({
+        raw: '@n is positive',
+        parts: [
+          { kind: 'ConstraintReference', name: 'n' },
+          { kind: 'ConstraintText', text: ' is positive' },
+        ],
+      })
+    })
+
+    it('should parse constraint with reference in the middle', () => {
+      const result = parseConstraint('use @path for storage')
+      expect(result).toEqual<Constraint>({
+        raw: 'use @path for storage',
+        parts: [
+          { kind: 'ConstraintText', text: 'use ' },
+          { kind: 'ConstraintReference', name: 'path' },
+          { kind: 'ConstraintText', text: ' for storage' },
+        ],
+      })
+    })
+
+    it('should parse constraint with dotted references', () => {
+      const result = parseConstraint('return @z.real^2 + @z.imag^2')
+      expect(result).toEqual<Constraint>({
+        raw: 'return @z.real^2 + @z.imag^2',
+        parts: [
+          { kind: 'ConstraintText', text: 'return ' },
+          { kind: 'ConstraintReference', name: 'z.real' },
+          { kind: 'ConstraintText', text: '^2 + ' },
+          { kind: 'ConstraintReference', name: 'z.imag' },
+          { kind: 'ConstraintText', text: '^2' },
+        ],
+      })
+    })
+
+    it('should parse constraint with multiple references', () => {
+      const result = parseConstraint('call @LoadTodos using @path')
+      expect(result).toEqual<Constraint>({
+        raw: 'call @LoadTodos using @path',
+        parts: [
+          { kind: 'ConstraintText', text: 'call ' },
+          { kind: 'ConstraintReference', name: 'LoadTodos' },
+          { kind: 'ConstraintText', text: ' using ' },
+          { kind: 'ConstraintReference', name: 'path' },
+        ],
+      })
+    })
+
+    it('should parse constraint with reference at start', () => {
+      const result = parseConstraint('@validate the input')
+      expect(result).toEqual<Constraint>({
+        raw: '@validate the input',
+        parts: [
+          { kind: 'ConstraintReference', name: 'validate' },
+          { kind: 'ConstraintText', text: ' the input' },
+        ],
+      })
+    })
+
+    it('should parse constraint with no @ symbols', () => {
+      const result = parseConstraint('the user is authenticated')
+      expect(result).toEqual<Constraint>({
+        raw: 'the user is authenticated',
+        parts: [{ kind: 'ConstraintText', text: 'the user is authenticated' }],
+      })
+    })
+
+    it('should parse constraint with underscore in reference', () => {
+      const result = parseConstraint('@todo_item is valid')
+      expect(result).toEqual<Constraint>({
+        raw: '@todo_item is valid',
+        parts: [
+          { kind: 'ConstraintReference', name: 'todo_item' },
+          { kind: 'ConstraintText', text: ' is valid' },
+        ],
+      })
+    })
+
+    it('should parse empty constraint', () => {
+      const result = parseConstraint('')
+      expect(result).toEqual<Constraint>({
+        raw: '',
+        parts: [],
+      })
+    })
+
+    it('should preserve raw text exactly', () => {
+      const result = parseConstraint('  @foo  ')
+      expect(result.raw).toBe('  @foo  ')
+    })
+  })
 })