Add unit tests, fmt

src/lib/kclHelpers.test.ts

@@ -0,0 +1,44 @@
+import { getCalculatedKclExpressionValue } from './kclHelpers'
+
+describe('KCL expression calculations', () => {
+  it('calculates a simple expression', async () => {
+    const actual = await getCalculatedKclExpressionValue({
+      value: '1 + 2',
+      variables: [],
+    })
+    expect(actual?.valueAsString).toEqual('3')
+    expect(actual?.astNode).toBeDefined()
+  })
+  it('calculates a simple expression with a variable', async () => {
+    const actual = await getCalculatedKclExpressionValue({
+      value: '1 + x',
+      variables: [{ key: 'x', value: 2 }],
+    })
+    expect(actual?.valueAsString).toEqual('3')
+    expect(actual?.astNode).toBeDefined()
+  })
+  it('returns NAN for an invalid expression', async () => {
+    const actual = await getCalculatedKclExpressionValue({
+      value: '1 + x',
+      variables: [],
+    })
+    expect(actual?.valueAsString).toEqual('NAN')
+    expect(actual?.astNode).toBeDefined()
+  })
+  it('returns NAN for an expression with an invalid variable', async () => {
+    const actual = await getCalculatedKclExpressionValue({
+      value: '1 + x',
+      variables: [{ key: 'y', value: 2 }],
+    })
+    expect(actual?.valueAsString).toEqual('NAN')
+    expect(actual?.astNode).toBeDefined()
+  })
+  it('calculates a more complex expression with a variable', async () => {
+    const actual = await getCalculatedKclExpressionValue({
+      value: '(1 + x * x) * 2',
+      variables: [{ key: 'x', value: 2 }],
+    })
+    expect(actual?.valueAsString).toEqual('10')
+    expect(actual?.astNode).toBeDefined()
+  })
+})

src/lib/kclHelpers.ts

@@ -0,0 +1,74 @@
+import { err, trap } from './trap'
+import { engineCommandManager, kclManager } from 'lib/singletons'
+import { parse, ProgramMemory, resultIsOk } from 'lang/wasm'
+import { PrevVariable } from 'lang/queryAst'
+import { executeAst } from 'lang/langHelpers'
+
+const DUMMY_VARIABLE_NAME = '__result__'
+
+/**
+ * Calculate the value of the KCL expression,
+ * given the value and the variables that are available
+ */
+export async function getCalculatedKclExpressionValue({
+  value,
+  variables,
+}: {
+  value: string
+  variables: PrevVariable<string | number>[]
+}) {
+  // Create a one-line program that assigns the value to a variable
+  const dummyProgramCode = `const ${DUMMY_VARIABLE_NAME} = ${value}`
+  const pResult = parse(dummyProgramCode)
+  if (err(pResult) || !resultIsOk(pResult)) return
+  const ast = pResult.program
+
+  // Populate the program memory with the passed-in variables
+  const programMemoryOverride: ProgramMemory = ProgramMemory.empty()
+  for (const { key, value } of variables) {
+    const error = programMemoryOverride.set(
+      key,
+      typeof value === 'number'
+        ? {
+            type: 'Number',
+            value,
+            __meta: [],
+          }
+        : {
+            type: 'String',
+            value,
+            __meta: [],
+          }
+    )
+    if (trap(error, { suppress: true })) return
+  }
+
+  console.log(
+    'programMemoryOverride',
+    JSON.stringify(programMemoryOverride, null, 2)
+  )
+
+  // Execute the program without hitting the engine
+  const { execState } = await executeAst({
+    ast,
+    engineCommandManager,
+    programMemoryOverride,
+  })
+
+  // Find the variable declaration for the result
+  const resultDeclaration = ast.body.find(
+    (a) =>
+      a.type === 'VariableDeclaration' &&
+      a.declaration.id?.name === DUMMY_VARIABLE_NAME
+  )
+  const variableDeclaratorAstNode =
+    resultDeclaration?.type === 'VariableDeclaration' &&
+    resultDeclaration?.declaration.init
+  const resultRawValue = execState.memory?.get(DUMMY_VARIABLE_NAME)?.value
+
+  return {
+    astNode: variableDeclaratorAstNode,
+    valueAsString:
+      typeof resultRawValue === 'number' ? String(resultRawValue) : 'NAN',
+  }
+}

src/lib/useCalculateKclExpression.ts

@@ -1,12 +1,11 @@
 import { useModelingContext } from 'hooks/useModelingContext'
-import { kclManager, engineCommandManager } from 'lib/singletons'
+import { kclManager } from 'lib/singletons'
 import { useKclContext } from 'lang/KclProvider'
 import { findUniqueName } from 'lang/modifyAst'
 import { PrevVariable, findAllPreviousVariables } from 'lang/queryAst'
-import { ProgramMemory, Expr, parse, resultIsOk } from 'lang/wasm'
+import { Expr } from 'lang/wasm'
 import { useEffect, useRef, useState } from 'react'
-import { executeAst } from 'lang/langHelpers'
-import { err, trap } from 'lib/trap'
+import { getCalculatedKclExpressionValue } from './kclHelpers'
 
 const isValidVariableName = (name: string) =>
   /^[a-zA-Z_][a-zA-Z0-9_]*$/.test(name)
@@ -86,37 +85,12 @@ export function useCalculateKclExpression({
 
   useEffect(() => {
     const execAstAndSetResult = async () => {
-      const _code = `const __result__ = ${value}`
-      const pResult = parse(_code)
-      if (err(pResult) || !resultIsOk(pResult)) return
-      const ast = pResult.program
-
-      const _programMem: ProgramMemory = ProgramMemory.empty()
-      for (const { key, value } of availableVarInfo.variables) {
-        const error = _programMem.set(key, {
-          type: 'String',
-          value,
-          __meta: [],
-        })
-        if (trap(error, { suppress: true })) return
-      }
-      const { execState } = await executeAst({
-        ast,
-        engineCommandManager,
-        // We make sure to send an empty program memory to denote we mean mock mode.
-        programMemoryOverride: kclManager.programMemory.clone(),
+      const result = await getCalculatedKclExpressionValue({
+        value,
+        variables: availableVarInfo.variables,
       })
-      const resultDeclaration = ast.body.find(
-        (a) =>
-          a.type === 'VariableDeclaration' &&
-          a.declaration.id?.name === '__result__'
-      )
-      const init =
-        resultDeclaration?.type === 'VariableDeclaration' &&
-        resultDeclaration?.declaration.init
-      const result = execState.memory?.get('__result__')?.value
-      setCalcResult(typeof result === 'number' ? String(result) : 'NAN')
-      init && setValueNode(init)
+      setCalcResult(result?.valueAsString || 'NAN')
+      result?.astNode && setValueNode(result.astNode)
     }
     if (!value) return
     execAstAndSetResult().catch(() => {