ProgramMemory refactor - eliminate copies of memory (#5273)

* refactor program memory for encapsulation and remove special treatment of return values

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Refactor ProgramMemory to isolate mutation

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Refactor ProgramMemory

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Generated output

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Cache the result of executing modules for their items

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Compress envs when popped

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Remove the last traces of geometry from the memory module

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* docs

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Fixups

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Frontend

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Improve Environment GC

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Fixup mock execution frontend and interpreter, misc bug fixes

Signed-off-by: Nick Cameron <nrc@ncameron.org>

---------

Signed-off-by: Nick Cameron <nrc@ncameron.org>

src/components/AvailableVarsHelpers.tsx

@@ -1,11 +1,5 @@
 import { useEffect, useState, useRef } from 'react'
-import {
-  parse,
-  BinaryPart,
-  Expr,
-  ProgramMemory,
-  resultIsOk,
-} from '../lang/wasm'
+import { parse, BinaryPart, Expr, resultIsOk, VariableMap } from '../lang/wasm'
 import {
   createIdentifier,
   createLiteral,
@@ -100,7 +94,7 @@ export function useCalc({
   newVariableInsertIndex: number
   setNewVariableName: (a: string) => void
 } {
-  const { programMemory } = useKclContext()
+  const { variables } = useKclContext()
   const { context } = useModelingContext()
   const selectionRange =
     context.selectionRanges?.graphSelections[0]?.codeRef?.range
@@ -127,7 +121,7 @@ export function useCalc({
   }, [])
 
   useEffect(() => {
-    if (programMemory.has(newVariableName)) {
+    if (variables[newVariableName]) {
       setIsNewVariableNameUnique(false)
     } else {
       setIsNewVariableNameUnique(true)
@@ -135,14 +129,14 @@ export function useCalc({
   }, [newVariableName])
 
   useEffect(() => {
-    if (!programMemory || !selectionRange) return
+    if (!variables || !selectionRange) return
     const varInfo = findAllPreviousVariables(
       kclManager.ast,
-      kclManager.programMemory,
+      kclManager.variables,
       selectionRange
     )
     setAvailableVarInfo(varInfo)
-  }, [kclManager.ast, kclManager.programMemory, selectionRange])
+  }, [kclManager.ast, kclManager.variables, selectionRange])
 
   useEffect(() => {
     try {
@@ -150,9 +144,9 @@ export function useCalc({
       const pResult = parse(code)
       if (trap(pResult) || !resultIsOk(pResult)) return
       const ast = pResult.program
-      const _programMem: ProgramMemory = ProgramMemory.empty()
+      const _variables: VariableMap = {}
       for (const { key, value } of availableVarInfo.variables) {
-        const error = _programMem.set(key, {
+        const error = (_variables[key] = {
           type: 'String',
           value,
           __meta: [],
@@ -163,8 +157,8 @@ export function useCalc({
       executeAst({
         ast,
         engineCommandManager,
-        // We make sure to send an empty program memory to denote we mean mock mode.
-        programMemoryOverride: kclManager.programMemory.clone(),
+        isMock: true,
+        variables,
       }).then(({ execState }) => {
         const resultDeclaration = ast.body.find(
           (a) =>
@@ -174,7 +168,7 @@ export function useCalc({
         const init =
           resultDeclaration?.type === 'VariableDeclaration' &&
           resultDeclaration?.declaration.init
-        const result = execState.memory?.get('__result__')?.value
+        const result = execState.variables['__result__']?.value
         setCalcResult(typeof result === 'number' ? String(result) : 'NAN')
         init && setValueNode(init)
       })

src/components/LspProvider.tsx

@@ -157,8 +157,6 @@ export const LspProvider = ({ children }: { children: React.ReactNode }) => {
                 // eslint-disable-next-line @typescript-eslint/no-floating-promises
                 plugin.requestSemanticTokens()
                 break
-              case 'kcl/memoryUpdated':
-                break
             }
           } catch (error) {
             console.error(error)

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

@@ -1,6 +1,6 @@
 import { processMemory } from './MemoryPane'
 import { enginelessExecutor } from '../../../lib/testHelpers'
-import { assertParse, initPromise, ProgramMemory } from '../../../lang/wasm'
+import { assertParse, initPromise } from '../../../lang/wasm'
 
 beforeAll(async () => {
   await initPromise
@@ -29,15 +29,11 @@ describe('processMemory', () => {
     |> line(endAbsolute = [2.15, 4.32])
     // |> rx(90, %)`
     const ast = assertParse(code)
-    const execState = await enginelessExecutor(ast, ProgramMemory.empty())
-    const output = processMemory(execState.memory)
+    const execState = await enginelessExecutor(ast)
+    const output = processMemory(execState.variables)
     expect(output.myVar).toEqual(5)
     expect(output.otherVar).toEqual(3)
     expect(output).toEqual({
-      HALF_TURN: 180,
-      QUARTER_TURN: 90,
-      THREE_QUARTER_TURN: 270,
-      ZERO: 0,
       myVar: 5,
       myFn: '__function(a)__',
       otherVar: 3,

src/components/ModelingSidebar/ModelingPanes/MemoryPane.tsx

@@ -2,10 +2,10 @@ import toast from 'react-hot-toast'
 import ReactJson from 'react-json-view'
 import { useMemo } from 'react'
 import {
-  ProgramMemory,
   Path,
   ExtrudeSurface,
   sketchFromKclValueOptional,
+  VariableMap,
 } from 'lang/wasm'
 import { useKclContext } from 'lang/KclProvider'
 import { useResolvedTheme } from 'hooks/useResolvedTheme'
@@ -15,12 +15,12 @@ import Tooltip from 'components/Tooltip'
 import { useModelingContext } from 'hooks/useModelingContext'
 
 export const MemoryPaneMenu = () => {
-  const { programMemory } = useKclContext()
+  const { variables } = useKclContext()
 
   function copyProgramMemoryToClipboard() {
     if (globalThis && 'navigator' in globalThis) {
       navigator.clipboard
-        .writeText(JSON.stringify(programMemory))
+        .writeText(JSON.stringify(variables))
         .then(() => toast.success('Program memory copied to clipboard'))
         .catch((e) =>
           trap(new Error('Failed to copy program memory to clipboard'))
@@ -50,12 +50,9 @@ export const MemoryPaneMenu = () => {
 
 export const MemoryPane = () => {
   const theme = useResolvedTheme()
-  const { programMemory } = useKclContext()
+  const { variables } = useKclContext()
   const { state } = useModelingContext()
-  const ProcessedMemory = useMemo(
-    () => processMemory(programMemory),
-    [programMemory]
-  )
+  const ProcessedMemory = useMemo(() => processMemory(variables), [variables])
   return (
     <div className="h-full relative">
       <div className="absolute inset-0 p-2 flex flex-col items-start">
@@ -85,9 +82,10 @@ export const MemoryPane = () => {
   )
 }
 
-export const processMemory = (programMemory: ProgramMemory) => {
+export const processMemory = (variables: VariableMap) => {
   const processedMemory: any = {}
-  for (const [key, val] of programMemory?.visibleEntries()) {
+  for (const [key, val] of Object.entries(variables)) {
+    if (val === undefined) continue
     if (
       val.type === 'Sketch' ||
       // @ts-ignore

src/components/Toolbar/EqualAngle.tsx

@@ -95,7 +95,7 @@ export function applyConstraintEqualAngle({
     ast: kclManager.ast,
     selectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
   })
   if (err(transform)) return transform
   const { modifiedAst, pathToNodeMap } = transform

src/components/Toolbar/EqualLength.tsx

@@ -93,7 +93,7 @@ export function applyConstraintEqualLength({
     ast: kclManager.ast,
     selectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
   })
   if (err(transform)) return transform
   const { modifiedAst, pathToNodeMap } = transform

src/components/Toolbar/HorzVert.tsx

@@ -1,6 +1,6 @@
 import { toolTips } from 'lang/langHelpers'
 import { Selections } from 'lib/selections'
-import { Program, ProgramMemory, Expr } from '../../lang/wasm'
+import { Program, Expr, VariableMap } from '../../lang/wasm'
 import { getNodeFromPath } from '../../lang/queryAst'
 import {
   PathToNodeMap,
@@ -51,7 +51,7 @@ export function applyConstraintHorzVert(
   selectionRanges: Selections,
   horOrVert: 'vertical' | 'horizontal',
   ast: Node<Program>,
-  programMemory: ProgramMemory
+  memVars: VariableMap
 ):
   | {
       modifiedAst: Node<Program>
@@ -66,7 +66,7 @@ export function applyConstraintHorzVert(
     ast,
     selectionRanges,
     transformInfos,
-    programMemory,
+    memVars,
     referenceSegName: '',
   })
 }

src/components/Toolbar/Intersect.tsx

@@ -46,7 +46,7 @@ export function intersectInfo({
     isLinesParallelAndConstrained(
       kclManager.ast,
       engineCommandManager.artifactGraph,
-      kclManager.programMemory,
+      kclManager.variables,
       selectionRanges.graphSelections[0],
       selectionRanges.graphSelections[1]
     )
@@ -147,7 +147,7 @@ export async function applyConstraintIntersect({
     ast: structuredClone(kclManager.ast),
     selectionRanges: forcedSelectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
   })
   if (err(transform1)) return Promise.reject(transform1)
   const { modifiedAst, tagInfo, valueUsedInTransform, pathToNodeMap } =
@@ -184,7 +184,7 @@ export async function applyConstraintIntersect({
     ast: kclManager.ast,
     selectionRanges: forcedSelectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     forceSegName: segName,
     forceValueUsedInTransform: finalValue,
   })

src/components/Toolbar/RemoveConstrainingValues.tsx

@@ -88,7 +88,7 @@ export function applyRemoveConstrainingValues({
     ast: kclManager.ast,
     selectionRanges: updatedSelectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
   })
 }

src/components/Toolbar/SetAbsDistance.tsx

@@ -104,7 +104,7 @@ export async function applyConstraintAbsDistance({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
   })
   if (err(transform1)) return Promise.reject(transform1)
@@ -124,7 +124,7 @@ export async function applyConstraintAbsDistance({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
     forceValueUsedInTransform: finalValue,
   })
@@ -172,7 +172,7 @@ export function applyConstraintAxisAlign({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
     forceValueUsedInTransform: finalValue,
   })

src/components/Toolbar/SetAngleBetween.tsx

@@ -95,7 +95,7 @@ export async function applyConstraintAngleBetween({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
   })
   if (err(transformed1)) return Promise.reject(transformed1)
   const { modifiedAst, tagInfo, valueUsedInTransform, pathToNodeMap } =
@@ -133,7 +133,7 @@ export async function applyConstraintAngleBetween({
     ast: kclManager.ast,
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     forceSegName: segName,
     forceValueUsedInTransform: finalValue,
   })

src/components/Toolbar/SetHorzVertDistance.tsx

@@ -107,7 +107,7 @@ export async function applyConstraintHorzVertDistance({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
   })
   if (err(transformed)) return Promise.reject(transformed)
   const { modifiedAst, tagInfo, valueUsedInTransform, pathToNodeMap } =
@@ -145,7 +145,7 @@ export async function applyConstraintHorzVertDistance({
       ast: kclManager.ast,
       selectionRanges,
       transformInfos,
-      programMemory: kclManager.programMemory,
+      memVars: kclManager.variables,
       forceSegName: segName,
       forceValueUsedInTransform: finalValue,
     })
@@ -195,7 +195,7 @@ export function applyConstraintHorzVertAlign({
     ast: kclManager.ast,
     selectionRanges,
     transformInfos,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     forceValueUsedInTransform: finalValue,
   })
   if (err(retval)) return retval

src/components/Toolbar/setAngleLength.tsx

@@ -105,7 +105,7 @@ export async function applyConstraintLength({
     ast,
     selectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
     forceValueUsedInTransform: distanceExpression,
   })
@@ -137,7 +137,7 @@ export async function applyConstraintAngleLength({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
   })
   if (err(sketched)) return Promise.reject(sketched)
@@ -189,7 +189,7 @@ export async function applyConstraintAngleLength({
     ast: structuredClone(kclManager.ast),
     selectionRanges,
     transformInfos: transforms,
-    programMemory: kclManager.programMemory,
+    memVars: kclManager.variables,
     referenceSegName: '',
     forceValueUsedInTransform: finalValue,
   })