Refactor: search more intelligently within findUniqueName

src/lang/modifyAst.test.ts

@@ -1,4 +1,4 @@
-import { assertParse, recast, initPromise, Identifier } from './wasm'
+import { assertParse, recast, initPromise, VariableDeclaration } from './wasm'
 import {
   createLiteral,
   createIdentifier,
@@ -16,6 +16,7 @@ import {
   deleteSegmentFromPipeExpression,
   removeSingleConstraintInfo,
   deleteFromSelection,
+  createCallExpressionStdLib,
 } from './modifyAst'
 import { enginelessExecutor } from '../lib/testHelpers'
 import { findUsesOfTagInPipe, getNodePathFromSourceRange } from './queryAst'
@@ -100,23 +101,36 @@ describe('Testing createPipeExpression', () => {
 
 describe('Testing findUniqueName', () => {
   it('should find a unique name', () => {
-    const result = findUniqueName(
-      JSON.stringify([
-        { type: 'Identifier', name: 'yo01', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo02', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo03', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo04', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo05', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo06', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo07', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo08', start: 0, end: 0, moduleId: 0 },
-        { type: 'Identifier', name: 'yo09', start: 0, end: 0, moduleId: 0 },
-      ] satisfies Node<Identifier>[]),
-      'yo',
-      2
-    )
+    const sampleCode = `yo01 = 5
+    yo02 = 5
+    yo03 = 5
+    yo04 = 5
+    yo05 = 5
+    yo06 = startSketchOn('XY')
+      |> startProfileAt([0, 0], %)
+      |> line([0, 1], %, $yo07)
+    yo08 = 5
+    yo09 = 5
+    `
+    const ast = assertParse(sampleCode)
+    const result = findUniqueName(ast, 'yo', 2)
     expect(result).toBe('yo10')
   })
+
+  it(`should find a unique unique name even with user-defined functions`, () => {
+    const sampleCode = `fn yo01 = (x) => {
+      yo03 = 6
+      return x + yo03
+    }
+    yo02 = 5
+    sketch001 = startSketchOn('XY')
+      |> startProfileAt([0, 0], %)
+      |> line([0, yo01(4)], %, $yo04)
+    `
+    const ast = assertParse(sampleCode)
+    const result = findUniqueName(ast, 'yo', 2)
+    expect(result).toBe('yo05')
+  })
 })
 describe('Testing addSketchTo', () => {
   it('should add a sketch to a program', () => {

src/lang/modifyAst.ts

@@ -163,39 +163,51 @@ export function addSketchTo(
   }
 }
 
+/**
+ * Given a target name and an AST,
+ * return a unique name that is not already used as
+ * a constant or tag.
+ */
 export function findUniqueName(
-  ast: Program | string,
+  ast: Program,
   name: string,
   pad = 3,
   index = 1
 ): string {
-  let searchStr: string = typeof ast === 'string' ? ast : JSON.stringify(ast)
+  // Searches for partial TagDeclarator or VariableDeclarator strings
+  const searchStrings = (n: string) => [
+    `{"type":"TagDeclarator","value":"${n}"`,
+    `{"type":"VariableDeclarator","id":{"type":"Identifier","name":"${n}"`,
+  ]
+  const nameIsInAstIdentifier = (target: string) =>
+    searchStrings(target).reduce(
+      (prev, curr) => prev || JSON.stringify(ast).includes(curr),
+      false
+    )
   const indexStr = String(index).padStart(pad, '0')
 
   const endingDigitsMatcher = /\d+$/
   const nameEndsInDigits = name.match(endingDigitsMatcher)
-  let nameIsInString = searchStr.includes(`:"${name}"`)
 
   if (nameEndsInDigits !== null) {
     // base case: name is unique and ends in digits
-    if (!nameIsInString) return name
+    if (!nameIsInAstIdentifier(name)) return name
 
     // recursive case: name is not unique and ends in digits
     const newPad = nameEndsInDigits[1].length
     const newIndex = parseInt(nameEndsInDigits[1]) + 1
     const nameWithoutDigits = name.replace(endingDigitsMatcher, '')
 
-    return findUniqueName(searchStr, nameWithoutDigits, newPad, newIndex)
+    return findUniqueName(ast, nameWithoutDigits, newPad, newIndex)
   }
 
   const newName = `${name}${indexStr}`
-  nameIsInString = searchStr.includes(`:"${newName}"`)
 
-  // base case: name is unique and does not end in digits
-  if (!nameIsInString) return newName
-
-  // recursive case: name is not unique and does not end in digits
-  return findUniqueName(searchStr, name, pad, index + 1)
+  return !nameIsInAstIdentifier(newName)
+    ? // base case: name is unique and does not end in digits
+      newName
+    : // recursive case: name is not unique and does not end in digits
+      findUniqueName(ast, name, pad, index + 1)
 }
 
 export function mutateArrExp(node: Expr, updateWith: ArrayExpression): boolean {