Add lexical scope and redefining variables in functions (#3015)

* Fix to allow variable shadowing inside functions

* Implement closures

* Fix KCL test code to not reference future tag definition

* Remove tag declarator from function parameters

This is an example where the scoping change revealed a subtle issue
with TagDeclarators.  You cannot bind a new tag using a function
parameter.

The issue is that evaluating a TagDeclarator like $foo binds an
identifier to its corresponding TagIdentifier, but returns the
TagDeclarator.  If you have a TagDeclarator passed in as a parameter
to a function, you can never get its corresponding TagIdentifier.

This seems like a case where TagDeclarator evaluation needs to be
revisited, especially now that we have scoped tags.

* Fix to query return, functions, and tag declarator AST nodes correctly

src/clientSideScene/sceneEntities.ts

@@ -47,7 +47,6 @@ import {
   PipeExpression,
   Program,
   ProgramMemory,
-  programMemoryInit,
   recast,
   SketchGroup,
   ExtrudeGroup,
@@ -130,7 +129,7 @@ export const HIDE_HOVER_SEGMENT_LENGTH = 60 // in pixels
 export class SceneEntities {
   engineCommandManager: EngineCommandManager
   scene: Scene
-  sceneProgramMemory: ProgramMemory = { root: {}, return: null }
+  sceneProgramMemory: ProgramMemory = ProgramMemory.empty()
   activeSegments: { [key: string]: Group } = {}
   intersectionPlane: Mesh | null = null
   axisGroup: Group | null = null
@@ -550,9 +549,9 @@ export class SceneEntities {
     const variableDeclarationName =
       _node1.node?.declarations?.[0]?.id?.name || ''
 
-    const sg = kclManager.programMemory.root[
+    const sg = kclManager.programMemory.get(
       variableDeclarationName
-    ] as SketchGroup
+    ) as SketchGroup
     const lastSeg = sg.value.slice(-1)[0] || sg.start
 
     const index = sg.value.length // because we've added a new segment that's not in the memory yet, no need for `-1`
@@ -768,9 +767,9 @@ export class SceneEntities {
           programMemoryOverride,
         })
         this.sceneProgramMemory = programMemory
-        const sketchGroup = programMemory.root[
+        const sketchGroup = programMemory.get(
           variableDeclarationName
-        ] as SketchGroup
+        ) as SketchGroup
         const sgPaths = sketchGroup.value
         const orthoFactor = orthoScale(sceneInfra.camControls.camera)
 
@@ -820,9 +819,9 @@ export class SceneEntities {
 
           // Prepare to update the THREEjs scene
           this.sceneProgramMemory = programMemory
-          const sketchGroup = programMemory.root[
+          const sketchGroup = programMemory.get(
             variableDeclarationName
-          ] as SketchGroup
+          ) as SketchGroup
           const sgPaths = sketchGroup.value
           const orthoFactor = orthoScale(sceneInfra.camControls.camera)
 
@@ -1081,9 +1080,9 @@ export class SceneEntities {
       })
       this.sceneProgramMemory = programMemory
 
-      const maybeSketchGroup = programMemory.root[variableDeclarationName]
+      const maybeSketchGroup = programMemory.get(variableDeclarationName)
       let sketchGroup = undefined
-      if (maybeSketchGroup.type === 'SketchGroup') {
+      if (maybeSketchGroup?.type === 'SketchGroup') {
         sketchGroup = maybeSketchGroup
       } else if ((maybeSketchGroup as ExtrudeGroup).sketchGroup) {
         sketchGroup = (maybeSketchGroup as ExtrudeGroup).sketchGroup
@@ -1773,7 +1772,7 @@ function prepareTruncatedMemoryAndAst(
   if (err(_node)) return _node
   const variableDeclarationName = _node.node?.declarations?.[0]?.id?.name || ''
   const lastSeg = (
-    programMemory.root[variableDeclarationName] as SketchGroup
+    programMemory.get(variableDeclarationName) as SketchGroup
   ).value.slice(-1)[0]
   if (draftSegment) {
     // truncatedAst needs to setup with another segment at the end
@@ -1824,33 +1823,27 @@ function prepareTruncatedMemoryAndAst(
     ..._ast,
     body: [JSON.parse(JSON.stringify(_ast.body[bodyIndex]))],
   }
-  const programMemoryOverride = programMemoryInit()
-  if (err(programMemoryOverride)) return programMemoryOverride
 
   // Grab all the TagDeclarators and TagIdentifiers from memory.
   let start = _node.node.start
-  for (const key in programMemory.root) {
+  const programMemoryOverride = programMemory.filterVariables(true, (value) => {
-    const value = programMemory.root[key]
-    if (!('__meta' in value)) {
-      continue
-    }
     if (
+      !('__meta' in value) ||
       value.__meta === undefined ||
       value.__meta.length === 0 ||
       value.__meta[0].sourceRange === undefined
     ) {
-      continue
+      return false
     }
 
     if (value.__meta[0].sourceRange[0] >= start) {
       // We only want things before our start point.
-      continue
+      return false
     }
 
-    if (value.type === 'TagIdentifier') {
+    return value.type === 'TagIdentifier'
-      programMemoryOverride.root[key] = JSON.parse(JSON.stringify(value))
+  })
-    }
+  if (err(programMemoryOverride)) return programMemoryOverride
-  }
 
   for (let i = 0; i < bodyIndex; i++) {
     const node = _ast.body[i]
@@ -1858,12 +1851,15 @@ function prepareTruncatedMemoryAndAst(
       continue
     }
     const name = node.declarations[0].id.name
-    // const memoryItem = kclManager.programMemory.root[name]
+    const memoryItem = programMemory.get(name)
-    const memoryItem = programMemory.root[name]
     if (!memoryItem) {
       continue
     }
-    programMemoryOverride.root[name] = JSON.parse(JSON.stringify(memoryItem))
+    const error = programMemoryOverride.set(
+      name,
+      JSON.parse(JSON.stringify(memoryItem))
+    )
+    if (err(error)) return error
   }
   return {
     truncatedAst,
@@ -1900,7 +1896,7 @@ export function sketchGroupFromPathToNode({
   )
   if (err(_varDec)) return _varDec
   const varDec = _varDec.node
-  const result = programMemory.root[varDec?.id?.name || '']
+  const result = programMemory.get(varDec?.id?.name || '')
   if (result?.type === 'ExtrudeGroup') {
     return result.sketchGroup
   }

src/components/AvailableVarsHelpers.tsx

@@ -1,5 +1,5 @@
 import { useEffect, useState, useRef } from 'react'
-import { parse, BinaryPart, Value } from '../lang/wasm'
+import { parse, BinaryPart, Value, ProgramMemory } from '../lang/wasm'
 import {
   createIdentifier,
   createLiteral,
@@ -120,8 +120,7 @@ export function useCalc({
   }, [])
 
   useEffect(() => {
-    const allVarNames = Object.keys(programMemory.root)
+    if (programMemory.has(newVariableName)) {
-    if (allVarNames.includes(newVariableName)) {
       setIsNewVariableNameUnique(false)
     } else {
       setIsNewVariableNameUnique(true)
@@ -143,17 +142,20 @@ export function useCalc({
       const code = `const __result__ = ${value}`
       const ast = parse(code)
       if (trap(ast)) return
-      const _programMem: any = { root: {}, return: null }
+      const _programMem: ProgramMemory = ProgramMemory.empty()
-      availableVarInfo.variables.forEach(({ key, value }) => {
+      for (const { key, value } of availableVarInfo.variables) {
-        _programMem.root[key] = { type: 'userVal', value, __meta: [] }
+        const error = _programMem.set(key, {
+          type: 'UserVal',
+          value,
+          __meta: [],
         })
+        if (trap(error)) return
+      }
       executeAst({
         ast,
         engineCommandManager,
         useFakeExecutor: true,
-        programMemoryOverride: JSON.parse(
+        programMemoryOverride: kclManager.programMemory.clone(),
-          JSON.stringify(kclManager.programMemory)
-        ),
       }).then(({ programMemory }) => {
         const resultDeclaration = ast.body.find(
           (a) =>
@@ -163,7 +165,7 @@ export function useCalc({
         const init =
           resultDeclaration?.type === 'VariableDeclaration' &&
           resultDeclaration?.declarations?.[0]?.init
-        const result = programMemory?.root?.__result__?.value
+        const result = programMemory?.get('__result__')?.value
         setCalcResult(typeof result === 'number' ? String(result) : 'NAN')
         init && setValueNode(init)
       })

src/components/ModelingSidebar/ModelingPanes/MemoryPane.test.tsx

@@ -1,6 +1,6 @@
 import { processMemory } from './MemoryPane'
 import { enginelessExecutor } from '../../../lib/testHelpers'
-import { initPromise, parse } from '../../../lang/wasm'
+import { initPromise, parse, ProgramMemory } from '../../../lang/wasm'
 
 beforeAll(async () => {
   await initPromise
@@ -29,10 +29,7 @@ describe('processMemory', () => {
     |> lineTo([2.15, 4.32], %)
     // |> rx(90, %)`
     const ast = parse(code)
-    const programMemory = await enginelessExecutor(ast, {
+    const programMemory = await enginelessExecutor(ast, ProgramMemory.empty())
-      root: {},
-      return: null,
-    })
     const output = processMemory(programMemory)
     expect(output.myVar).toEqual(5)
     expect(output.otherVar).toEqual(3)

src/components/ModelingSidebar/ModelingPanes/MemoryPane.tsx

@@ -82,8 +82,7 @@ export const MemoryPane = () => {
 
 export const processMemory = (programMemory: ProgramMemory) => {
   const processedMemory: any = {}
-  Object.keys(programMemory?.root || {}).forEach((key) => {
+  for (const [key, val] of programMemory?.visibleEntries()) {
-    const val = programMemory.root[key]
     if (typeof val.value !== 'function') {
       if (val.type === 'SketchGroup') {
         processedMemory[key] = val.value.map(({ __geoMeta, ...rest }: Path) => {
@@ -103,6 +102,6 @@ export const processMemory = (programMemory: ProgramMemory) => {
     } else if (key !== 'log') {
       processedMemory[key] = '__function__'
     }
-  })
+  }
   return processedMemory
 }

src/lang/KclSingleton.ts

@@ -14,9 +14,7 @@ import {
   Program,
   ProgramMemory,
   recast,
-  SketchGroup,
   SourceRange,
-  ExtrudeGroup,
 } from 'lang/wasm'
 import { getNodeFromPath } from './queryAst'
 import { codeManager, editorManager, sceneInfra } from 'lib/singletons'
@@ -33,10 +31,7 @@ export class KclManager {
     },
     digest: null,
   }
-  private _programMemory: ProgramMemory = {
+  private _programMemory: ProgramMemory = ProgramMemory.empty()
-    root: {},
-    return: null,
-  }
   private _logs: string[] = []
   private _kclErrors: KCLError[] = []
   private _isExecuting = false
@@ -505,10 +500,7 @@ function defaultSelectionFilter(
   programMemory: ProgramMemory,
   engineCommandManager: EngineCommandManager
 ) {
-  const firstSketchOrExtrudeGroup = Object.values(programMemory.root).find(
+  programMemory.hasSketchOrExtrudeGroup() &&
-    (node) => node.type === 'ExtrudeGroup' || node.type === 'SketchGroup'
-  ) as SketchGroup | ExtrudeGroup
-  firstSketchOrExtrudeGroup &&
     engineCommandManager.sendSceneCommand({
       type: 'modeling_cmd_req',
       cmd_id: uuidv4(),

src/lang/artifact.test.ts

@@ -16,7 +16,7 @@ const mySketch001 = startSketchOn('XY')
   // |> rx(45, %)`
     const programMemory = await enginelessExecutor(parse(code))
     // @ts-ignore
-    const sketch001 = programMemory?.root?.mySketch001
+    const sketch001 = programMemory?.get('mySketch001')
     expect(sketch001).toEqual({
       type: 'SketchGroup',
       on: expect.any(Object),
@@ -66,7 +66,7 @@ const mySketch001 = startSketchOn('XY')
   |> extrude(2, %)`
     const programMemory = await enginelessExecutor(parse(code))
     // @ts-ignore
-    const sketch001 = programMemory?.root?.mySketch001
+    const sketch001 = programMemory?.get('mySketch001')
     expect(sketch001).toEqual({
       type: 'ExtrudeGroup',
       id: expect.any(String),
@@ -146,7 +146,7 @@ const sk2 = startSketchOn('XY')
 `
     const programMemory = await enginelessExecutor(parse(code))
     // @ts-ignore
-    const geos = [programMemory?.root?.theExtrude, programMemory?.root?.sk2]
+    const geos = [programMemory?.get('theExtrude'), programMemory?.get('sk2')]
     expect(geos).toEqual([
       {
         type: 'ExtrudeGroup',

src/lang/executor.test.ts

@@ -12,25 +12,25 @@ describe('test executor', () => {
   it('test assigning two variables, the second summing with the first', async () => {
     const code = `const myVar = 5
 const newVar = myVar + 1`
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(5)
+    expect(mem.get('myVar')?.value).toBe(5)
-    expect(root.newVar.value).toBe(6)
+    expect(mem.get('newVar')?.value).toBe(6)
   })
   it('test assigning a var with a string', async () => {
     const code = `const myVar = "a str"`
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe('a str')
+    expect(mem.get('myVar')?.value).toBe('a str')
   })
   it('test assigning a var by cont concatenating two strings string execute', async () => {
     const code = fs.readFileSync(
       './src/lang/testExamples/variableDeclaration.cado',
       'utf-8'
     )
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe('a str another str')
+    expect(mem.get('myVar')?.value).toBe('a str another str')
   })
   it('fn funcN = () => {} execute', async () => {
-    const { root } = await exe(
+    const mem = await exe(
       [
         'fn funcN = (a, b) => {',
         '  return a + b',
@@ -39,8 +39,8 @@ const newVar = myVar + 1`
         'const magicNum = funcN(9, theVar)',
       ].join('\n')
     )
-    expect(root.theVar.value).toBe(60)
+    expect(mem.get('theVar')?.value).toBe(60)
-    expect(root.magicNum.value).toBe(69)
+    expect(mem.get('magicNum')?.value).toBe(69)
   })
   it('sketch declaration', async () => {
     let code = `const mySketch = startSketchOn('XY')
@@ -50,9 +50,9 @@ const newVar = myVar + 1`
   |> lineTo([5,-1], %, "rightPath")
   // |> close(%)
 `
-    const { root } = await exe(code)
+    const mem = await exe(code)
     // geo is three js buffer geometry and is very bloated to have in tests
-    const minusGeo = root.mySketch.value
+    const minusGeo = mem.get('mySketch')?.value
     expect(minusGeo).toEqual([
       {
         type: 'ToPoint',
@@ -104,8 +104,8 @@ const newVar = myVar + 1`
       'fn myFn = (a) => { return a + 1 }',
       'const myVar = 5 + 1 |> myFn(%)',
     ].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(7)
+    expect(mem.get('myVar')?.value).toBe(7)
   })
 
   // Enable rotations #152
@@ -117,16 +117,16 @@ const newVar = myVar + 1`
   //     '  |> lineTo([1, 1], %)',
   //     'const rotated = rx(90, mySk1)',
   //   ].join('\n')
-  //   const { root } = await exe(code)
+  //   const mem = await exe(code)
-  //   expect(root.mySk1.value).toHaveLength(3)
+  //   expect(mem.get('mySk1')?.value).toHaveLength(3)
-  //   expect(root?.rotated?.type).toBe('SketchGroup')
+  //   expect(mem.get('rotated')?.type).toBe('SketchGroup')
   //   if (
-  //     root?.mySk1?.type !== 'SketchGroup' ||
+  //     mem.get('mySk1')?.type !== 'SketchGroup' ||
-  //     root?.rotated?.type !== 'SketchGroup'
+  //     mem.get('rotated')?.type !== 'SketchGroup'
   //   )
   //     throw new Error('not a sketch group')
-  //   expect(root.mySk1.rotation).toEqual([0, 0, 0, 1])
+  //   expect(mem.get('mySk1')?.rotation).toEqual([0, 0, 0, 1])
-  //   expect(root.rotated.rotation.map((a) => a.toFixed(4))).toEqual([
+  //   expect(mem.get('rotated')?.rotation.map((a) => a.toFixed(4))).toEqual([
   //     '0.7071',
   //     '0.0000',
   //     '0.0000',
@@ -144,8 +144,8 @@ const newVar = myVar + 1`
       '  |> lineTo([1,1], %)',
       // '  |> rx(90, %)',
     ].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.mySk1).toEqual({
+    expect(mem.get('mySk1')).toEqual({
       type: 'SketchGroup',
       on: expect.any(Object),
       start: {
@@ -214,10 +214,9 @@ const newVar = myVar + 1`
     const code = ['const three = 3', "const yo = [1, '2', three, 4 + 5]"].join(
       '\n'
     )
-    const { root } = await exe(code)
+    const mem = await exe(code)
     // TODO path to node is probably wrong here, zero indexes are not correct
-    expect(root).toEqual({
+    expect(mem.get('three')).toEqual({
-      three: {
       type: 'UserVal',
       value: 3,
       __meta: [
@@ -225,8 +224,8 @@ const newVar = myVar + 1`
           sourceRange: [14, 15],
         },
       ],
-      },
+    })
-      yo: {
+    expect(mem.get('yo')).toEqual({
       type: 'UserVal',
       value: [1, '2', 3, 9],
       __meta: [
@@ -234,16 +233,18 @@ const newVar = myVar + 1`
           sourceRange: [27, 49],
         },
       ],
-      },
     })
+    // Check that there are no other variables or environments.
+    expect(mem.numEnvironments()).toBe(1)
+    expect(mem.numVariables(0)).toBe(2)
   })
   it('execute object expression', async () => {
     const code = [
       'const three = 3',
       "const yo = {aStr: 'str', anum: 2, identifier: three, binExp: 4 + 5}",
     ].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.yo).toEqual({
+    expect(mem.get('yo')).toEqual({
       type: 'UserVal',
       value: { aStr: 'str', anum: 2, identifier: 3, binExp: 9 },
       __meta: [
@@ -257,8 +258,8 @@ const newVar = myVar + 1`
     const code = ["const yo = {a: {b: '123'}}", "const myVar = yo.a['b']"].join(
       '\n'
     )
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar).toEqual({
+    expect(mem.get('myVar')).toEqual({
       type: 'UserVal',
       value: '123',
       __meta: [
@@ -273,81 +274,81 @@ const newVar = myVar + 1`
 describe('testing math operators', () => {
   it('can sum', async () => {
     const code = ['const myVar = 1 + 2'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(3)
+    expect(mem.get('myVar')?.value).toBe(3)
   })
   it('can subtract', async () => {
     const code = ['const myVar = 1 - 2'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(-1)
+    expect(mem.get('myVar')?.value).toBe(-1)
   })
   it('can multiply', async () => {
     const code = ['const myVar = 1 * 2'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(2)
+    expect(mem.get('myVar')?.value).toBe(2)
   })
   it('can divide', async () => {
     const code = ['const myVar = 1 / 2'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(0.5)
+    expect(mem.get('myVar')?.value).toBe(0.5)
   })
   it('can modulus', async () => {
     const code = ['const myVar = 5 % 2'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(1)
+    expect(mem.get('myVar')?.value).toBe(1)
   })
   it('can do multiple operations', async () => {
     const code = ['const myVar = 1 + 2 * 3'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(7)
+    expect(mem.get('myVar')?.value).toBe(7)
   })
   it('big example with parans', async () => {
     const code = ['const myVar = 1 + 2 * (3 - 4) / -5 + 6'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(7.4)
+    expect(mem.get('myVar')?.value).toBe(7.4)
   })
   it('with identifier', async () => {
     const code = ['const yo = 6', 'const myVar = yo / 2'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(3)
+    expect(mem.get('myVar')?.value).toBe(3)
   })
   it('with lots of testing', async () => {
     const code = ['const myVar = 2 * ((2 + 3 ) / 4 + 5)'].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(12.5)
+    expect(mem.get('myVar')?.value).toBe(12.5)
   })
   it('with callExpression at start', async () => {
     const code = 'const myVar = min(4, 100) + 2'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(6)
+    expect(mem.get('myVar')?.value).toBe(6)
   })
   it('with callExpression at end', async () => {
     const code = 'const myVar = 2 + min(4, 100)'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(6)
+    expect(mem.get('myVar')?.value).toBe(6)
   })
   it('with nested callExpression', async () => {
     const code = 'const myVar = 2 + min(100, legLen(5, 3))'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(6)
+    expect(mem.get('myVar')?.value).toBe(6)
   })
   it('with unaryExpression', async () => {
     const code = 'const myVar = -min(100, 3)'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(-3)
+    expect(mem.get('myVar')?.value).toBe(-3)
   })
   it('with unaryExpression in callExpression', async () => {
     const code = 'const myVar = min(-legLen(5, 4), 5)'
     const code2 = 'const myVar = min(5 , -legLen(5, 4))'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    const { root: root2 } = await exe(code2)
+    const mem2 = await exe(code2)
-    expect(root.myVar.value).toBe(-3)
+    expect(mem.get('myVar')?.value).toBe(-3)
-    expect(root.myVar.value).toBe(root2.myVar.value)
+    expect(mem.get('myVar')?.value).toBe(mem2.get('myVar')?.value)
   })
   it('with unaryExpression in ArrayExpression', async () => {
     const code = 'const myVar = [1,-legLen(5, 4)]'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toEqual([1, -3])
+    expect(mem.get('myVar')?.value).toEqual([1, -3])
   })
   it('with unaryExpression in ArrayExpression in CallExpression, checking nothing funny happens when used in a sketch', async () => {
     const code = [
@@ -355,8 +356,8 @@ describe('testing math operators', () => {
       '  |> startProfileAt([0, 0], %)',
       '|> line([-2.21, -legLen(5, min(3, 999))], %)',
     ].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    const sketch = root.part001
+    const sketch = mem.get('part001')
     // result of `-legLen(5, min(3, 999))` should be -4
     const yVal = (sketch as SketchGroup).value?.[0]?.to?.[1]
     expect(yVal).toBe(-4)
@@ -373,8 +374,8 @@ describe('testing math operators', () => {
       `], %)`,
       ``,
     ].join('\n')
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    const sketch = root.part001
+    const sketch = mem.get('part001')
     // expect -legLen(segLen('seg01', %), myVar) to equal -4 setting the y value back to 0
     expect((sketch as SketchGroup).value?.[1]?.from).toEqual([3, 4])
     expect((sketch as SketchGroup).value?.[1]?.to).toEqual([6, 0])
@@ -382,18 +383,18 @@ describe('testing math operators', () => {
       `-legLen(segLen('seg01', %), myVar)`,
       `legLen(segLen('seg01', %), myVar)`
     )
-    const { root: removedUnaryExpRoot } = await exe(removedUnaryExp)
+    const removedUnaryExpMem = await exe(removedUnaryExp)
-    const removedUnaryExpRootSketch = removedUnaryExpRoot.part001
+    const removedUnaryExpMemSketch = removedUnaryExpMem.get('part001')
 
     // without the minus sign, the y value should be 8
-    expect((removedUnaryExpRootSketch as SketchGroup).value?.[1]?.to).toEqual([
+    expect((removedUnaryExpMemSketch as SketchGroup).value?.[1]?.to).toEqual([
       6, 8,
     ])
   })
   it('with nested callExpression and binaryExpression', async () => {
     const code = 'const myVar = 2 + min(100, -1 + legLen(5, 3))'
-    const { root } = await exe(code)
+    const mem = await exe(code)
-    expect(root.myVar.value).toBe(5)
+    expect(mem.get('myVar')?.value).toBe(5)
   })
 })
 
@@ -421,7 +422,7 @@ const theExtrude = startSketchOn('XY')
 
 async function exe(
   code: string,
-  programMemory: ProgramMemory = { root: {}, return: null }
+  programMemory: ProgramMemory = ProgramMemory.empty()
 ) {
   const ast = parse(code)
 

src/lang/langHelpers.ts

@@ -79,20 +79,14 @@ export async function executeAst({
       return {
         errors: [e],
         logs: [],
-        programMemory: {
+        programMemory: ProgramMemory.empty(),
-          root: {},
-          return: null,
-        },
       }
     } else {
       console.log(e)
       return {
         logs: [e],
         errors: [],
-        programMemory: {
+        programMemory: ProgramMemory.empty(),
-          root: {},
-          return: null,
-        },
       }
     }
   }

src/lang/modifyAst.ts

@@ -983,7 +983,7 @@ export async function deleteFromSelection(
           if (err(parent)) {
             return
           }
-          const sketchToPreserve = programMemory.root[sketchName] as SketchGroup
+          const sketchToPreserve = programMemory.get(sketchName) as SketchGroup
           console.log('sketchName', sketchName)
           // Can't kick off multiple requests at once as getFaceDetails
           // is three engine calls in one and they conflict

src/lang/queryAst.ts

@@ -130,8 +130,14 @@ function moreNodePathFromSourceRange(
 
   const isInRange = _node.start <= start && _node.end >= end
 
-  if ((_node.type === 'Identifier' || _node.type === 'Literal') && isInRange)
+  if (
+    (_node.type === 'Identifier' ||
+      _node.type === 'Literal' ||
+      _node.type === 'TagDeclarator') &&
+    isInRange
+  ) {
     return path
+  }
 
   if (_node.type === 'CallExpression' && isInRange) {
     const { callee, arguments: args } = _node
@@ -277,6 +283,15 @@ function moreNodePathFromSourceRange(
         }
       }
     }
+    return path
+  }
+  if (_node.type === 'ReturnStatement' && isInRange) {
+    const { argument } = _node
+    if (argument.start <= start && argument.end >= end) {
+      path.push(['argument', 'ReturnStatement'])
+      return moreNodePathFromSourceRange(argument, sourceRange, path)
+    }
+    return path
   }
   if (_node.type === 'MemberExpression' && isInRange) {
     const { object, property } = _node
@@ -459,8 +474,8 @@ export function findAllPreviousVariablesPath(
   bodyItems?.forEach?.((item) => {
     if (item.type !== 'VariableDeclaration' || item.end > startRange) return
     const varName = item.declarations[0].id.name
-    const varValue = programMemory?.root[varName]
+    const varValue = programMemory?.get(varName)
-    if (typeof varValue?.value !== type) return
+    if (!varValue || typeof varValue?.value !== type) return
     variables.push({
       key: varName,
       value: varValue.value,
@@ -640,7 +655,7 @@ export function isLinesParallelAndConstrained(
     if (err(_varDec)) return _varDec
     const varDec = _varDec.node
     const varName = (varDec as VariableDeclaration)?.declarations[0]?.id?.name
-    const path = programMemory?.root[varName] as SketchGroup
+    const path = programMemory?.get(varName) as SketchGroup
     const _primarySegment = getSketchSegmentFromSourceRange(
       path,
       primaryLine.range
@@ -687,7 +702,7 @@ export function isLinesParallelAndConstrained(
       constraintType === 'angle' || constraintLevel === 'full'
 
     // get the previous segment
-    const prevSegment = (programMemory.root[varName] as SketchGroup).value[
+    const prevSegment = (programMemory.get(varName) as SketchGroup).value[
       secondaryIndex - 1
     ]
     const prevSourceRange = prevSegment.__geoMeta.sourceRange
@@ -757,7 +772,7 @@ export function hasExtrudeSketchGroup({
   const varDec = varDecMeta.node
   if (varDec.type !== 'VariableDeclaration') return false
   const varName = varDec.declarations[0].id.name
-  const varValue = programMemory?.root[varName]
+  const varValue = programMemory?.get(varName)
   return varValue?.type === 'ExtrudeGroup' || varValue?.type === 'SketchGroup'
 }
 

src/lang/std/sketch.ts

@@ -1009,8 +1009,8 @@ export const angledLineOfXLength: SketchLineHelper = {
     const { node: varDec } = nodeMeta2
 
     const variableName = varDec.id.name
-    const sketch = previousProgramMemory?.root?.[variableName]
+    const sketch = previousProgramMemory?.get(variableName)
-    if (sketch.type !== 'SketchGroup') {
+    if (!sketch || sketch.type !== 'SketchGroup') {
       return new Error('not a SketchGroup')
     }
     const angle = createLiteral(roundOff(getAngle(from, to), 0))
@@ -1105,8 +1105,8 @@ export const angledLineOfYLength: SketchLineHelper = {
     if (err(nodeMeta2)) return nodeMeta2
     const { node: varDec } = nodeMeta2
     const variableName = varDec.id.name
-    const sketch = previousProgramMemory?.root?.[variableName]
+    const sketch = previousProgramMemory?.get(variableName)
-    if (sketch.type !== 'SketchGroup') {
+    if (!sketch || sketch.type !== 'SketchGroup') {
       return new Error('not a SketchGroup')
     }
 
@@ -1443,7 +1443,7 @@ export const angledLineThatIntersects: SketchLineHelper = {
 
     const { node: varDec } = nodeMeta2
     const varName = varDec.declarations[0].id.name
-    const sketchGroup = previousProgramMemory.root[varName] as SketchGroup
+    const sketchGroup = previousProgramMemory.get(varName) as SketchGroup
     const intersectPath = sketchGroup.value.find(
       ({ tag }: Path) => tag && tag.value === intersectTagName
     )

src/lang/std/sketchConstraints.test.ts

@@ -363,7 +363,7 @@ const part001 = startSketchOn('XY')
     const programMemory = await enginelessExecutor(parse(code))
     const index = code.indexOf('// normal-segment') - 7
     const _segment = getSketchSegmentFromSourceRange(
-      programMemory.root['part001'] as SketchGroup,
+      programMemory.get('part001') as SketchGroup,
       [index, index]
     )
     if (err(_segment)) throw _segment
@@ -379,7 +379,7 @@ const part001 = startSketchOn('XY')
     const programMemory = await enginelessExecutor(parse(code))
     const index = code.indexOf('// segment-in-start') - 7
     const _segment = getSketchSegmentFromSourceRange(
-      programMemory.root['part001'] as SketchGroup,
+      programMemory.get('part001') as SketchGroup,
       [index, index]
     )
     if (err(_segment)) throw _segment

src/lang/std/sketchcombos.ts

@@ -1636,8 +1636,8 @@ export function transformAstSketchLines({
     })
 
     const varName = varDec.node.id.name
-    let sketchGroup = programMemory.root?.[varName]
+    let sketchGroup = programMemory.get(varName)
-    if (sketchGroup.type === 'ExtrudeGroup') {
+    if (sketchGroup?.type === 'ExtrudeGroup') {
       sketchGroup = sketchGroup.sketchGroup
     }
     if (!sketchGroup || sketchGroup.type !== 'SketchGroup')

src/lang/std/std.test.ts

@@ -17,9 +17,9 @@ describe('testing angledLineThatIntersects', () => {
   offset: ${offset},
 }, %, "yo2")
 const intersect = segEndX('yo2', part001)`
-    const { root } = await enginelessExecutor(parse(code('-1')))
+    const mem = await enginelessExecutor(parse(code('-1')))
-    expect(root.intersect.value).toBe(1 + Math.sqrt(2))
+    expect(mem.get('intersect')?.value).toBe(1 + Math.sqrt(2))
-    const { root: noOffset } = await enginelessExecutor(parse(code('0')))
+    const noOffset = await enginelessExecutor(parse(code('0')))
-    expect(noOffset.intersect.value).toBeCloseTo(1)
+    expect(noOffset.get('intersect')?.value).toBeCloseTo(1)
   })
 })

src/lang/wasm.ts

@@ -143,14 +143,200 @@ interface Memory {
   [key: string]: MemoryItem
 }
 
-export interface ProgramMemory {
+type EnvironmentRef = number
-  root: Memory
+
+const ROOT_ENVIRONMENT_REF: EnvironmentRef = 0
+
+interface Environment {
+  bindings: Memory
+  parent: EnvironmentRef | null
+}
+
+function emptyEnvironment(): Environment {
+  return { bindings: {}, parent: null }
+}
+
+interface RawProgramMemory {
+  environments: Environment[]
+  currentEnv: EnvironmentRef
   return: ProgramReturn | null
 }
 
+/**
+ * This duplicates logic in Rust.  The hope is to keep ProgramMemory internals
+ * isolated from the rest of the TypeScript code so that we can move it to Rust
+ * in the future.
+ */
+export class ProgramMemory {
+  private environments: Environment[]
+  private currentEnv: EnvironmentRef
+  private return: ProgramReturn | null
+
+  /**
+   * Empty memory doesn't include prelude definitions.
+   */
+  static empty(): ProgramMemory {
+    return new ProgramMemory()
+  }
+
+  static fromRaw(raw: RawProgramMemory): ProgramMemory {
+    return new ProgramMemory(raw.environments, raw.currentEnv, raw.return)
+  }
+
+  constructor(
+    environments: Environment[] = [emptyEnvironment()],
+    currentEnv: EnvironmentRef = ROOT_ENVIRONMENT_REF,
+    returnVal: ProgramReturn | null = null
+  ) {
+    this.environments = environments
+    this.currentEnv = currentEnv
+    this.return = returnVal
+  }
+
+  /**
+   * Returns a deep copy.
+   */
+  clone(): ProgramMemory {
+    return ProgramMemory.fromRaw(JSON.parse(JSON.stringify(this.toRaw())))
+  }
+
+  has(name: string): boolean {
+    let envRef = this.currentEnv
+    while (true) {
+      const env = this.environments[envRef]
+      if (env.bindings.hasOwnProperty(name)) {
+        return true
+      }
+      if (!env.parent) {
+        break
+      }
+      envRef = env.parent
+    }
+    return false
+  }
+
+  get(name: string): MemoryItem | null {
+    let envRef = this.currentEnv
+    while (true) {
+      const env = this.environments[envRef]
+      if (env.bindings.hasOwnProperty(name)) {
+        return env.bindings[name]
+      }
+      if (!env.parent) {
+        break
+      }
+      envRef = env.parent
+    }
+    return null
+  }
+
+  set(name: string, value: MemoryItem): Error | null {
+    if (this.environments.length === 0) {
+      return new Error('No environment to set memory in')
+    }
+    const env = this.environments[this.currentEnv]
+    env.bindings[name] = value
+    return null
+  }
+
+  /**
+   * Returns a new ProgramMemory with only `MemoryItem`s that pass the
+   * predicate.  Values are deep copied.
+   *
+   * Note: Return value of the returned ProgramMemory is always null.
+   */
+  filterVariables(
+    keepPrelude: boolean,
+    predicate: (value: MemoryItem) => boolean
+  ): ProgramMemory | Error {
+    const environments: Environment[] = []
+    for (const [i, env] of this.environments.entries()) {
+      let bindings: Memory
+      if (i === ROOT_ENVIRONMENT_REF && keepPrelude) {
+        // Get prelude definitions.  Create these first so that they're always
+        // first in iteration order.
+        const memoryOrError = programMemoryInit()
+        if (err(memoryOrError)) return memoryOrError
+        bindings = memoryOrError.environments[0].bindings
+      } else {
+        bindings = emptyEnvironment().bindings
+      }
+
+      for (const [name, value] of Object.entries(env.bindings)) {
+        // Check the predicate.
+        if (!predicate(value)) {
+          continue
+        }
+        // Deep copy.
+        bindings[name] = JSON.parse(JSON.stringify(value))
+      }
+      environments.push({ bindings, parent: env.parent })
+    }
+    return new ProgramMemory(environments, this.currentEnv, null)
+  }
+
+  numEnvironments(): number {
+    return this.environments.length
+  }
+
+  numVariables(envRef: EnvironmentRef): number {
+    return Object.keys(this.environments[envRef]).length
+  }
+
+  /**
+   * Returns all variable entries in memory that are visible, in a flat
+   * structure.  If variables are shadowed, they're not visible, and therefore,
+   * not included.
+   *
+   * This should only be used to display in the MemoryPane UI.
+   */
+  visibleEntries(): Map<string, MemoryItem> {
+    const map = new Map<string, MemoryItem>()
+    let envRef = this.currentEnv
+    while (true) {
+      const env = this.environments[envRef]
+      for (const [name, value] of Object.entries(env.bindings)) {
+        // Don't include shadowed variables.
+        if (!map.has(name)) {
+          map.set(name, value)
+        }
+      }
+      if (!env.parent) {
+        break
+      }
+      envRef = env.parent
+    }
+    return map
+  }
+
+  /**
+   * Returns true if any visible variables are a SketchGroup or ExtrudeGroup.
+   */
+  hasSketchOrExtrudeGroup(): boolean {
+    for (const node of this.visibleEntries().values()) {
+      if (node.type === 'ExtrudeGroup' || node.type === 'SketchGroup') {
+        return true
+      }
+    }
+    return false
+  }
+
+  /**
+   * Return the representation that can be serialized to JSON.  This should only
+   * be used within this module.
+   */
+  toRaw(): RawProgramMemory {
+    return {
+      environments: this.environments,
+      currentEnv: this.currentEnv,
+      return: this.return,
+    }
+  }
+}
+
 export const executor = async (
   node: Program,
-  programMemory: ProgramMemory | Error = { root: {}, return: null },
+  programMemory: ProgramMemory | Error = ProgramMemory.empty(),
   engineCommandManager: EngineCommandManager,
   isMock: boolean = false
 ): Promise<ProgramMemory> => {
@@ -171,7 +357,7 @@ export const executor = async (
 
 export const _executor = async (
   node: Program,
-  programMemory: ProgramMemory | Error = { root: {}, return: null },
+  programMemory: ProgramMemory | Error = ProgramMemory.empty(),
   engineCommandManager: EngineCommandManager,
   isMock: boolean
 ): Promise<ProgramMemory> => {
@@ -186,15 +372,15 @@ export const _executor = async (
       baseUnit =
         (await getSettingsState)()?.modeling.defaultUnit.current || 'mm'
     }
-    const memory: ProgramMemory = await execute_wasm(
+    const memory: RawProgramMemory = await execute_wasm(
       JSON.stringify(node),
-      JSON.stringify(programMemory),
+      JSON.stringify(programMemory.toRaw()),
       baseUnit,
       engineCommandManager,
       fileSystemManager,
       isMock
     )
-    return memory
+    return ProgramMemory.fromRaw(memory)
   } catch (e: any) {
     console.log(e)
     const parsed: RustKclError = JSON.parse(e.toString())
@@ -329,10 +515,17 @@ export function getTangentialArcToInfo({
   }
 }
 
+/**
+ * Returns new ProgramMemory with prelude definitions.
+ */
 export function programMemoryInit(): ProgramMemory | Error {
   try {
-    const memory: ProgramMemory = program_memory_init()
+    const memory: RawProgramMemory = program_memory_init()
-    return memory
+    return new ProgramMemory(
+      memory.environments,
+      memory.currentEnv,
+      memory.return
+    )
   } catch (e: any) {
     console.log(e)
     const parsed: RustKclError = JSON.parse(e.toString())

src/lib/testHelpers.ts

@@ -75,7 +75,7 @@ class MockEngineCommandManager {
 
 export async function enginelessExecutor(
   ast: Program | Error,
-  pm: ProgramMemory | Error = { root: {}, return: null }
+  pm: ProgramMemory | Error = ProgramMemory.empty()
 ): Promise<ProgramMemory> {
   if (err(ast)) return Promise.reject(ast)
   if (err(pm)) return Promise.reject(pm)
@@ -93,7 +93,7 @@ export async function enginelessExecutor(
 
 export async function executor(
   ast: Program,
-  pm: ProgramMemory = { root: {}, return: null }
+  pm: ProgramMemory = ProgramMemory.empty()
 ): Promise<ProgramMemory> {
   const engineCommandManager = new EngineCommandManager()
   engineCommandManager.start({

src/lib/useCalculateKclExpression.ts

@@ -3,7 +3,7 @@ import { kclManager, engineCommandManager } from 'lib/singletons'
 import { useKclContext } from 'lang/KclProvider'
 import { findUniqueName } from 'lang/modifyAst'
 import { PrevVariable, findAllPreviousVariables } from 'lang/queryAst'
-import { Value, parse } from 'lang/wasm'
+import { ProgramMemory, Value, parse } from 'lang/wasm'
 import { useEffect, useRef, useState } from 'react'
 import { executeAst } from 'lang/langHelpers'
 import { err, trap } from 'lib/trap'
@@ -60,9 +60,8 @@ export function useCalculateKclExpression({
   }, [])
 
   useEffect(() => {
-    const allVarNames = Object.keys(programMemory.root)
     if (
-      allVarNames.includes(newVariableName) ||
+      programMemory.has(newVariableName) ||
       newVariableName === '' ||
       !isValidVariableName(newVariableName)
     ) {
@@ -89,17 +88,20 @@ export function useCalculateKclExpression({
       if (err(ast)) return
       if (trap(ast, { suppress: true })) return
 
-      const _programMem: any = { root: {}, return: null }
+      const _programMem: ProgramMemory = ProgramMemory.empty()
-      availableVarInfo.variables.forEach(({ key, value }) => {
+      for (const { key, value } of availableVarInfo.variables) {
-        _programMem.root[key] = { type: 'userVal', value, __meta: [] }
+        const error = _programMem.set(key, {
+          type: 'UserVal',
+          value,
+          __meta: [],
         })
+        if (trap(error, { suppress: true })) return
+      }
       const { programMemory } = await executeAst({
         ast,
         engineCommandManager,
         useFakeExecutor: true,
-        programMemoryOverride: JSON.parse(
+        programMemoryOverride: kclManager.programMemory.clone(),
-          JSON.stringify(kclManager.programMemory)
-        ),
       })
       const resultDeclaration = ast.body.find(
         (a) =>
@@ -109,7 +111,7 @@ export function useCalculateKclExpression({
       const init =
         resultDeclaration?.type === 'VariableDeclaration' &&
         resultDeclaration?.declarations?.[0]?.init
-      const result = programMemory?.root?.__result__?.value
+      const result = programMemory?.get('__result__')?.value
       setCalcResult(typeof result === 'number' ? String(result) : 'NAN')
       init && setValueNode(init)
     }

src/machines/modelingMachine.ts

@@ -1139,8 +1139,8 @@ export const modelingMachine = createMachine(
         )
         if (err(varDecNode)) return
         const sketchVar = varDecNode.node.declarations[0].id.name
-        const sketchGroup = kclManager.programMemory.root[sketchVar]
+        const sketchGroup = kclManager.programMemory.get(sketchVar)
-        if (sketchGroup.type !== 'SketchGroup') return
+        if (sketchGroup?.type !== 'SketchGroup') return
         const idArtifact = engineCommandManager.artifactMap[sketchGroup.id]
         if (idArtifact.commandType !== 'start_path') return
         const extrusionArtifactId = (idArtifact as any)?.extrusions?.[0]

src/wasm-lib/kcl/src/ast/types.rs

@@ -1388,7 +1388,7 @@ impl CallExpression {
             }
             FunctionKind::UserDefined => {
                 let func = memory.get(&fn_name, self.into())?;
-                let result = func.call_fn(fn_args, memory.clone(), ctx.clone()).await.map_err(|e| {
+                let result = func.call_fn(fn_args, ctx.clone()).await.map_err(|e| {
                     // Add the call expression to the source ranges.
                     e.add_source_ranges(vec![self.into()])
                 })?;

src/wasm-lib/kcl/src/executor.rs

@@ -23,7 +23,8 @@ use crate::{
 #[ts(export)]
 #[serde(rename_all = "camelCase")]
 pub struct ProgramMemory {
-    pub root: HashMap<String, MemoryItem>,
+    pub environments: Vec<Environment>,
+    pub current_env: EnvironmentRef,
     #[serde(rename = "return")]
     pub return_: Option<ProgramReturn>,
 }
@@ -31,7 +32,105 @@ pub struct ProgramMemory {
 impl ProgramMemory {
     pub fn new() -> Self {
         Self {
-            root: HashMap::from([
+            environments: vec![Environment::root()],
+            current_env: EnvironmentRef::root(),
+            return_: None,
+        }
+    }
+
+    pub fn new_env_for_call(&mut self, parent: EnvironmentRef) -> EnvironmentRef {
+        let new_env_ref = EnvironmentRef(self.environments.len());
+        let new_env = Environment::new(parent);
+        self.environments.push(new_env);
+        new_env_ref
+    }
+
+    /// Add to the program memory in the current scope.
+    pub fn add(&mut self, key: &str, value: MemoryItem, source_range: SourceRange) -> Result<(), KclError> {
+        if self.environments[self.current_env.index()].contains_key(key) {
+            return Err(KclError::ValueAlreadyDefined(KclErrorDetails {
+                message: format!("Cannot redefine `{}`", key),
+                source_ranges: vec![source_range],
+            }));
+        }
+
+        self.environments[self.current_env.index()].insert(key.to_string(), value);
+
+        Ok(())
+    }
+
+    /// Get a value from the program memory.
+    /// Return Err if not found.
+    pub fn get(&self, var: &str, source_range: SourceRange) -> Result<&MemoryItem, KclError> {
+        let mut env_ref = self.current_env;
+        loop {
+            let env = &self.environments[env_ref.index()];
+            if let Some(item) = env.bindings.get(var) {
+                return Ok(item);
+            }
+            if let Some(parent) = env.parent {
+                env_ref = parent;
+            } else {
+                break;
+            }
+        }
+
+        Err(KclError::UndefinedValue(KclErrorDetails {
+            message: format!("memory item key `{}` is not defined", var),
+            source_ranges: vec![source_range],
+        }))
+    }
+
+    /// Find all extrude groups in the memory that are on a specific sketch group id.
+    /// This does not look inside closures.  But as long as we do not allow
+    /// mutation of variables in KCL, closure memory should be a subset of this.
+    pub fn find_extrude_groups_on_sketch_group(&self, sketch_group_id: uuid::Uuid) -> Vec<Box<ExtrudeGroup>> {
+        self.environments
+            .iter()
+            .flat_map(|env| {
+                env.bindings
+                    .values()
+                    .filter_map(|item| match item {
+                        MemoryItem::ExtrudeGroup(eg) if eg.sketch_group.id == sketch_group_id => Some(eg.clone()),
+                        _ => None,
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .collect()
+    }
+}
+
+impl Default for ProgramMemory {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// An index pointing to an environment.
+#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
+pub struct EnvironmentRef(usize);
+
+impl EnvironmentRef {
+    pub fn root() -> Self {
+        Self(0)
+    }
+
+    pub fn index(&self) -> usize {
+        self.0
+    }
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
+pub struct Environment {
+    bindings: HashMap<String, MemoryItem>,
+    parent: Option<EnvironmentRef>,
+}
+
+impl Environment {
+    pub fn root() -> Self {
+        Self {
+            // Prelude
+            bindings: HashMap::from([
                 (
                     "ZERO".to_string(),
                     MemoryItem::UserVal(UserVal {
@@ -61,28 +160,19 @@ impl ProgramMemory {
                     }),
                 ),
             ]),
-            return_: None,
+            parent: None,
         }
     }
 
-    /// Add to the program memory.
+    pub fn new(parent: EnvironmentRef) -> Self {
-    pub fn add(&mut self, key: &str, value: MemoryItem, source_range: SourceRange) -> Result<(), KclError> {
+        Self {
-        if self.root.contains_key(key) {
+            bindings: HashMap::new(),
-            return Err(KclError::ValueAlreadyDefined(KclErrorDetails {
+            parent: Some(parent),
-                message: format!("Cannot redefine `{}`", key),
+        }
-                source_ranges: vec![source_range],
-            }));
     }
 
-        self.root.insert(key.to_string(), value);
-
-        Ok(())
-    }
-
-    /// Get a value from the program memory.
-    /// Return Err if not found.
     pub fn get(&self, key: &str, source_range: SourceRange) -> Result<&MemoryItem, KclError> {
-        self.root.get(key).ok_or_else(|| {
+        self.bindings.get(key).ok_or_else(|| {
             KclError::UndefinedValue(KclErrorDetails {
                 message: format!("memory item key `{}` is not defined", key),
                 source_ranges: vec![source_range],
@@ -90,21 +180,12 @@ impl ProgramMemory {
         })
     }
 
-    /// Find all extrude groups in the memory that are on a specific sketch group id.
+    pub fn insert(&mut self, key: String, value: MemoryItem) {
-    pub fn find_extrude_groups_on_sketch_group(&self, sketch_group_id: uuid::Uuid) -> Vec<Box<ExtrudeGroup>> {
+        self.bindings.insert(key, value);
-        self.root
-            .values()
-            .filter_map(|item| match item {
-                MemoryItem::ExtrudeGroup(eg) if eg.sketch_group.id == sketch_group_id => Some(eg.clone()),
-                _ => None,
-            })
-            .collect()
     }
-}
 
-impl Default for ProgramMemory {
+    pub fn contains_key(&self, key: &str) -> bool {
-    fn default() -> Self {
+        self.bindings.contains_key(key)
-        Self::new()
     }
 }
 
@@ -161,6 +242,7 @@ pub enum MemoryItem {
         #[serde(skip)]
         func: Option<MemoryFunction>,
         expression: Box<FunctionExpression>,
+        memory: Box<ProgramMemory>,
         #[serde(rename = "__meta")]
         meta: Vec<Metadata>,
     },
@@ -646,6 +728,7 @@ impl MemoryItem {
         let MemoryItem::Function {
             func,
             expression,
+            memory,
             meta: _,
         } = &self
         else {
@@ -655,6 +738,7 @@ impl MemoryItem {
         Some(FnAsArg {
             func,
             expr: expression.to_owned(),
+            memory: memory.to_owned(),
         })
     }
 
@@ -728,10 +812,15 @@ impl MemoryItem {
     pub async fn call_fn(
         &self,
         args: Vec<MemoryItem>,
-        memory: ProgramMemory,
         ctx: ExecutorContext,
     ) -> Result<Option<ProgramReturn>, KclError> {
-        let MemoryItem::Function { func, expression, meta } = &self else {
+        let MemoryItem::Function {
+            func,
+            expression,
+            memory: closure_memory,
+            meta,
+        } = &self
+        else {
             return Err(KclError::Semantic(KclErrorDetails {
                 message: "not a in memory function".to_string(),
                 source_ranges: vec![],
@@ -743,7 +832,14 @@ impl MemoryItem {
                 source_ranges: vec![],
             }));
         };
-        func(args, memory, expression.clone(), meta.clone(), ctx).await
+        func(
+            args,
+            closure_memory.as_ref().clone(),
+            expression.clone(),
+            meta.clone(),
+            ctx,
+        )
+        .await
     }
 }
 
@@ -1552,16 +1648,13 @@ impl ExecutorContext {
                                 memory.return_ = result.return_;
                             }
                             FunctionKind::UserDefined => {
-                                if let Some(func) = memory.clone().root.get(&fn_name) {
+                                // TODO: Why do we change the source range to
-                                    let result = func.call_fn(args.clone(), memory.clone(), self.clone()).await?;
+                                // the call expression instead of keeping the
+                                // range of the callee?
+                                let func = memory.get(&fn_name, call_expr.into())?;
+                                let result = func.call_fn(args.clone(), self.clone()).await?;
 
                                 memory.return_ = result;
-                                } else {
-                                    return Err(KclError::Semantic(KclErrorDetails {
-                                        message: format!("No such name {} defined", fn_name),
-                                        source_ranges: vec![call_expr.into()],
-                                    }));
-                                }
                             }
                         }
                     }
@@ -1672,7 +1765,15 @@ impl ExecutorContext {
                      _metadata: Vec<Metadata>,
                      ctx: ExecutorContext| {
                         Box::pin(async move {
-                            let mut fn_memory = assign_args_to_params(&function_expression, args, memory.clone())?;
+                            // Create a new environment to execute the function
+                            // body in so that local variables shadow variables
+                            // in the parent scope.  The new environment's
+                            // parent should be the environment of the closure.
+                            let mut body_memory = memory.clone();
+                            let closure_env = memory.current_env;
+                            let body_env = body_memory.new_env_for_call(closure_env);
+                            body_memory.current_env = body_env;
+                            let mut fn_memory = assign_args_to_params(&function_expression, args, body_memory)?;
 
                             let result = ctx
                                 .inner_execute(&function_expression.body, &mut fn_memory, BodyType::Block)
@@ -1682,10 +1783,14 @@ impl ExecutorContext {
                         })
                     },
                 );
+                // Cloning memory here is crucial for semantics so that we close
+                // over variables.  Variables defined lexically later shouldn't
+                // be available to the function body.
                 MemoryItem::Function {
                     expression: function_expression.clone(),
                     meta: vec![metadata.to_owned()],
                     func: Some(mem_func),
+                    memory: Box::new(memory.clone()),
                 }
             }
             Value::CallExpression(call_expression) => call_expression.execute(memory, pipe_info, self).await?,
@@ -1788,7 +1893,8 @@ fn assign_args_to_params(
         return Err(err_wrong_number_args);
     }
 
-    // Add the arguments to the memory.
+    // Add the arguments to the memory.  A new call frame should have already
+    // been created.
     for (index, param) in function_expression.params.iter().enumerate() {
         if let Some(arg) = args.get(index) {
             // Argument was provided.
@@ -1854,11 +1960,19 @@ const newVar = myVar + 1"#;
         let memory = parse_execute(ast).await.unwrap();
         assert_eq!(
             serde_json::json!(5),
-            memory.root.get("myVar").unwrap().get_json_value().unwrap()
+            memory
+                .get("myVar", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
         assert_eq!(
             serde_json::json!(6.0),
-            memory.root.get("newVar").unwrap().get_json_value().unwrap()
+            memory
+                .get("newVar", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -1883,13 +1997,21 @@ const intersect = segEndX('yo2', part001)"#,
         let memory = parse_execute(&ast_fn("-1")).await.unwrap();
         assert_eq!(
             serde_json::json!(1.0 + 2.0f64.sqrt()),
-            memory.root.get("intersect").unwrap().get_json_value().unwrap()
+            memory
+                .get("intersect", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
 
         let memory = parse_execute(&ast_fn("0")).await.unwrap();
         assert_eq!(
             serde_json::json!(1.0000000000000002),
-            memory.root.get("intersect").unwrap().get_json_value().unwrap()
+            memory
+                .get("intersect", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -2207,13 +2329,201 @@ const thisBox = box([[0,0], 6, 10, 3])
         parse_execute(ast).await.unwrap();
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_function_cannot_access_future_definitions() {
+        let ast = r#"
+fn returnX = () => {
+  // x shouldn't be defined yet.
+  return x
+}
+
+const x = 5
+
+const answer = returnX()"#;
+
+        let result = parse_execute(ast).await;
+        let err = result.unwrap_err().downcast::<KclError>().unwrap();
+        assert_eq!(
+            err,
+            KclError::UndefinedValue(KclErrorDetails {
+                message: "memory item key `x` is not defined".to_owned(),
+                source_ranges: vec![SourceRange([64, 65]), SourceRange([97, 106])],
+            }),
+        );
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_pattern_transform_function_cannot_access_future_definitions() {
+        let ast = r#"
+fn transform = (replicaId) => {
+  // x shouldn't be defined yet.
+  let scale = x
+  return {
+    translate: [0, 0, replicaId * 10],
+    scale: [scale, 1, 0],
+  }
+}
+
+fn layer = () => {
+  return startSketchOn("XY")
+    |> circle([0, 0], 1, %, 'tag1')
+    |> extrude(10, %)
+}
+
+const x = 5
+
+// The 10 layers are replicas of each other, with a transform applied to each.
+let shape = layer() |> patternTransform(10, transform, %)
+"#;
+
+        let result = parse_execute(ast).await;
+        let err = result.unwrap_err().downcast::<KclError>().unwrap();
+        assert_eq!(
+            err,
+            KclError::UndefinedValue(KclErrorDetails {
+                message: "memory item key `x` is not defined".to_owned(),
+                source_ranges: vec![SourceRange([80, 81])],
+            }),
+        );
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_execute_function_with_parameter_redefined_outside() {
+        let ast = r#"
+fn myIdentity = (x) => {
+  return x
+}
+
+const x = 33
+
+const two = myIdentity(2)"#;
+
+        let memory = parse_execute(ast).await.unwrap();
+        assert_eq!(
+            serde_json::json!(2),
+            memory
+                .get("two", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+        assert_eq!(
+            serde_json::json!(33),
+            memory
+                .get("x", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_execute_function_referencing_variable_in_parent_scope() {
+        let ast = r#"
+const x = 22
+const y = 3
+
+fn add = (x) => {
+  return x + y
+}
+
+const answer = add(2)"#;
+
+        let memory = parse_execute(ast).await.unwrap();
+        assert_eq!(
+            serde_json::json!(5.0),
+            memory
+                .get("answer", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+        assert_eq!(
+            serde_json::json!(22),
+            memory
+                .get("x", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_execute_function_redefining_variable_in_parent_scope() {
+        let ast = r#"
+const x = 1
+
+fn foo = () => {
+  const x = 2
+  return x
+}
+
+const answer = foo()"#;
+
+        let memory = parse_execute(ast).await.unwrap();
+        assert_eq!(
+            serde_json::json!(2),
+            memory
+                .get("answer", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+        assert_eq!(
+            serde_json::json!(1),
+            memory
+                .get("x", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_execute_pattern_transform_function_redefining_variable_in_parent_scope() {
+        let ast = r#"
+const scale = 100
+fn transform = (replicaId) => {
+  // Redefine same variable as in parent scope.
+  const scale = 2
+  return {
+    translate: [0, 0, replicaId * 10],
+    scale: [scale, 1, 0],
+  }
+}
+
+fn layer = () => {
+  return startSketchOn("XY")
+    |> circle([0, 0], 1, %, 'tag1')
+    |> extrude(10, %)
+}
+
+// The 10 layers are replicas of each other, with a transform applied to each.
+let shape = layer() |> patternTransform(10, transform, %)"#;
+
+        let memory = parse_execute(ast).await.unwrap();
+        // TODO: Assert that scale 2 was used.
+        assert_eq!(
+            serde_json::json!(100),
+            memory
+                .get("scale", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
+        );
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_math_execute_with_functions() {
         let ast = r#"const myVar = 2 + min(100, -1 + legLen(5, 3))"#;
         let memory = parse_execute(ast).await.unwrap();
         assert_eq!(
             serde_json::json!(5.0),
-            memory.root.get("myVar").unwrap().get_json_value().unwrap()
+            memory
+                .get("myVar", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -2223,7 +2533,11 @@ const thisBox = box([[0,0], 6, 10, 3])
         let memory = parse_execute(ast).await.unwrap();
         assert_eq!(
             serde_json::json!(7.4),
-            memory.root.get("myVar").unwrap().get_json_value().unwrap()
+            memory
+                .get("myVar", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -2233,7 +2547,11 @@ const thisBox = box([[0,0], 6, 10, 3])
         let memory = parse_execute(ast).await.unwrap();
         assert_eq!(
             serde_json::json!(1.0),
-            memory.root.get("myVar").unwrap().get_json_value().unwrap()
+            memory
+                .get("myVar", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -2243,7 +2561,11 @@ const thisBox = box([[0,0], 6, 10, 3])
         let memory = parse_execute(ast).await.unwrap();
         assert_eq!(
             serde_json::json!(std::f64::consts::TAU),
-            memory.root.get("myVar").unwrap().get_json_value().unwrap()
+            memory
+                .get("myVar", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -2253,7 +2575,11 @@ const thisBox = box([[0,0], 6, 10, 3])
         let memory = parse_execute(ast).await.unwrap();
         assert_eq!(
             serde_json::json!(7.4),
-            memory.root.get("thing").unwrap().get_json_value().unwrap()
+            memory
+                .get("thing", SourceRange::default())
+                .unwrap()
+                .get_json_value()
+                .unwrap()
         );
     }
 
@@ -2383,7 +2709,9 @@ const bracket = startSketchOn('XY')
         fn additional_program_memory(items: &[(String, MemoryItem)]) -> ProgramMemory {
             let mut program_memory = ProgramMemory::new();
             for (name, item) in items {
-                program_memory.root.insert(name.to_string(), item.clone());
+                program_memory
+                    .add(name.as_str(), item.clone(), SourceRange::default())
+                    .unwrap();
             }
             program_memory
         }

src/wasm-lib/kcl/src/std/mod.rs

@@ -31,7 +31,7 @@ use crate::{
     ast::types::FunctionExpression,
     docs::StdLibFn,
     errors::KclError,
-    executor::{MemoryItem, SketchGroup, SketchSurface},
+    executor::{MemoryItem, ProgramMemory, SketchGroup, SketchSurface},
     std::kcl_stdlib::KclStdLibFn,
 };
 pub use args::Args;
@@ -281,6 +281,7 @@ pub enum Primitive {
 pub struct FnAsArg<'a> {
     pub func: &'a crate::executor::MemoryFunction,
     pub expr: Box<FunctionExpression>,
+    pub memory: Box<ProgramMemory>,
 }
 
 #[cfg(test)]

src/wasm-lib/kcl/src/std/patterns.rs

@@ -87,7 +87,7 @@ pub async fn pattern_transform(args: Args) -> Result<MemoryItem, KclError> {
             fn_expr: transform.expr,
             meta: vec![args.source_range.into()],
             ctx: args.ctx.clone(),
-            memory: args.current_program_memory.clone(),
+            memory: *transform.memory,
         },
         extr,
         &args,

src/wasm-lib/tests/executor/main.rs

@@ -1784,31 +1784,31 @@ const part002 = startSketchOn(part001, 'end')
 
 #[tokio::test(flavor = "multi_thread")]
 async fn serial_test_plumbus_fillets() {
-    let code = r#"fn make_circle = (ext, face, tag ,pos, radius) => {
+    let code = r#"fn make_circle = (ext, face, pos, radius) => {
   const sg = startSketchOn(ext, face)
   |> startProfileAt([pos[0] + radius, pos[1]], %)
   |> arc({
        angle_end: 360,
        angle_start: 0,
        radius: radius
-     }, %, tag)
+     }, %, $arc1)
   |> close(%)
 
   return sg
 }
 
-fn pentagon = (len, taga, tagb, tagc) => {
+fn pentagon = (len) => {
   const sg = startSketchOn('XY')
   |> startProfileAt([-len / 2, -len / 2], %)
-  |> angledLine({ angle: 0, length: len }, %,taga)
+  |> angledLine({ angle: 0, length: len }, %, $a)
   |> angledLine({
        angle: segAng(a, %) + 180 - 108,
        length: len
-     }, %, tagb)
+     }, %, $b)
   |> angledLine({
        angle: segAng(b, %) + 180 - 108,
        length: len
-     }, %,tagc)
+     }, %, $c)
   |> angledLine({
        angle: segAng(c, %) + 180 - 108,
        length: len
@@ -1821,21 +1821,23 @@ fn pentagon = (len, taga, tagb, tagc) => {
   return sg
 }
 
-const p = pentagon(32, $a, $b, $c)
+const p = pentagon(32)
   |> extrude(10, %)
 
-const plumbus0 = make_circle(p,a,  $arc_a, [0, 0], 2.5)
+const circle0 = make_circle(p, p.sketchGroup.tags.a, [0, 0], 2.5)
+const plumbus0 = circle0
   |> extrude(10, %)
   |> fillet({
        radius: 0.5,
-       tags: [arc_a, getOppositeEdge(arc_a, %)]
+       tags: [circle0.tags.arc1, getOppositeEdge(circle0.tags.arc1, %)]
      }, %)
 
-const plumbus1 = make_circle(p, b,$arc_b, [0, 0], 2.5)
+const circle1 = make_circle(p, p.sketchGroup.tags.b, [0, 0], 2.5)
+const plumbus1 = circle1
    |> extrude(10, %)
    |> fillet({
         radius: 0.5,
-        tags: [arc_b, getOppositeEdge(arc_b, %)]
+        tags: [circle1.tags.arc1, getOppositeEdge(circle1.tags.arc1, %)]
       }, %)
 "#;
 

src/wasm-lib/tests/modify/main.rs

@@ -39,7 +39,7 @@ async fn setup(code: &str, name: &str) -> Result<(ExecutorContext, Program, uuid
 
     // We need to get the sketch ID.
     // Get the sketch group ID from memory.
-    let MemoryItem::SketchGroup(sketch_group) = memory.root.get(name).unwrap() else {
+    let MemoryItem::SketchGroup(sketch_group) = memory.get(name, SourceRange::default()).unwrap() else {
         anyhow::bail!("part001 not found in memory: {:?}", memory);
     };
     let sketch_id = sketch_group.id;