Enable/disable "start sketch", "edit sketch" and "extrude" appropriately (#1449)

* test that fails for when to enable extrude and sketch features

* add fix to make test pass

* A snapshot a day keeps the bugs away! 📷🐛 (OS: ubuntu)

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

e2e/playwright/flow-tests.spec.ts

@@ -577,7 +577,7 @@ test('Selections work on fresh and edited sketch', async ({ page }) => {
   await page.waitForTimeout(200)
 
   // enter sketch again
-  await page.getByRole('button', { name: 'Start Sketch' }).click()
+  await page.getByRole('button', { name: 'Edit Sketch' }).click()
   await page.waitForTimeout(700) // wait for animation
 
   // hover again and check it works
@@ -848,3 +848,103 @@ test('ProgramMemory can be serialised', async ({ page, context }) => {
     })
   })
 })
+
+test('edit selections', async ({ page, context }) => {
+  const u = getUtils(page)
+  const selectionsSnippets = {
+    extrudeAndEditBlocked: '|> startProfileAt([10.81, 32.99], %)',
+    extrudeAndEditBlockedInFunction: '|> startProfileAt(pos, %)',
+    extrudeAndEditAllowed: '|> startProfileAt([15.72, 4.7], %)',
+    editOnly: '|> startProfileAt([15.79, -14.6], %)',
+  }
+  await context.addInitScript(
+    async ({
+      extrudeAndEditBlocked,
+      extrudeAndEditBlockedInFunction,
+      extrudeAndEditAllowed,
+      editOnly,
+    }: any) => {
+      localStorage.setItem(
+        'persistCode',
+        `const part001 = startSketchOn('-XZ')
+  ${extrudeAndEditBlocked}
+  |> line([25.96, 2.93], %)
+  |> line([5.25, -5.72], %)
+  |> line([-2.01, -10.35], %)
+  |> line([-27.65, -2.78], %)
+  |> close(%)
+  |> extrude(5, %)
+const part002 = startSketchOn('-XZ')
+  ${extrudeAndEditAllowed}
+  |> line([10.32, 6.47], %)
+  |> line([9.71, -6.16], %)
+  |> line([-3.08, -9.86], %)
+  |> line([-12.02, -1.54], %)
+  |> close(%)
+const part003 = startSketchOn('-XZ')
+  ${editOnly}
+  |> line([27.55, -1.65], %)
+  |> line([4.95, -8], %)
+  |> line([-20.38, -10.12], %)
+  |> line([-15.79, 17.08], %)
+
+fn yohey = (pos) => {
+  const part004 = startSketchOn('-XZ')
+  ${extrudeAndEditBlockedInFunction}
+  |> line([27.55, -1.65], %)
+  |> line([4.95, -10.53], %)
+  |> line([-20.38, -8], %)
+  |> line([-15.79, 17.08], %)
+  return ''
+}
+    
+    yohey([15.79, -34.6])
+`
+      )
+    },
+    selectionsSnippets
+  )
+  await page.setViewportSize({ width: 1200, height: 500 })
+  await page.goto('/')
+  await u.waitForAuthSkipAppStart()
+
+  // wait for execution done
+  await u.openDebugPanel()
+  await u.expectCmdLog('[data-message-type="execution-done"]')
+  await u.closeDebugPanel()
+
+  await page.getByText(selectionsSnippets.extrudeAndEditBlocked).click()
+  await expect(page.getByRole('button', { name: 'Extrude' })).toBeDisabled()
+  await expect(
+    page.getByRole('button', { name: 'Edit Sketch' })
+  ).not.toBeVisible()
+
+  await page.getByText(selectionsSnippets.extrudeAndEditAllowed).click()
+  await expect(page.getByRole('button', { name: 'Extrude' })).not.toBeDisabled()
+  await expect(
+    page.getByRole('button', { name: 'Edit Sketch' })
+  ).not.toBeDisabled()
+
+  await page.getByText(selectionsSnippets.editOnly).click()
+  await expect(page.getByRole('button', { name: 'Extrude' })).toBeDisabled()
+  await expect(
+    page.getByRole('button', { name: 'Edit Sketch' })
+  ).not.toBeDisabled()
+
+  await page
+    .getByText(selectionsSnippets.extrudeAndEditBlockedInFunction)
+    .click()
+  await expect(page.getByRole('button', { name: 'Extrude' })).toBeDisabled()
+  await expect(
+    page.getByRole('button', { name: 'Edit Sketch' })
+  ).not.toBeVisible()
+
+  // selecting an editable sketch but clicking "start sktech" should start a new sketch and not edit the existing one
+  await page.getByText(selectionsSnippets.extrudeAndEditAllowed).click()
+  await page.getByRole('button', { name: 'Start Sketch' }).click()
+  await page.mouse.click(700, 200)
+  // expect main content to contain `part005` i.e. started a new sketch
+  await expect(page.locator('.cm-content')).toHaveText(
+    /part005 = startSketchOn\('-XZ'\)/
+  )
+})

src/Toolbar.tsx

@@ -5,6 +5,8 @@ import { useModelingContext } from 'hooks/useModelingContext'
 import { useCommandsContext } from 'hooks/useCommandsContext'
 import { ActionButton } from 'components/ActionButton'
 import usePlatform from 'hooks/usePlatform'
+import { isSingleCursorInPipe } from 'lang/queryAst'
+import { kclManager } from 'lang/KclSingleton'
 
 export const Toolbar = () => {
   const platform = usePlatform()
@@ -13,14 +15,15 @@ export const Toolbar = () => {
   const toolbarButtonsRef = useRef<HTMLUListElement>(null)
   const bgClassName =
     'group-enabled:group-hover:bg-energy-10 group-pressed:bg-energy-10 dark:group-enabled:group-hover:bg-chalkboard-80 dark:group-pressed:bg-chalkboard-80'
-  const pathId = useMemo(
-    () =>
-      isCursorInSketchCommandRange(
+  const pathId = useMemo(() => {
+    if (!isSingleCursorInPipe(context.selectionRanges, kclManager.ast)) {
+      return false
+    }
+    return isCursorInSketchCommandRange(
       engineCommandManager.artifactMap,
       context.selectionRanges
-      ),
-    [engineCommandManager.artifactMap, context.selectionRanges]
     )
+  }, [engineCommandManager.artifactMap, context.selectionRanges])
 
   function handleToolbarButtonsWheelEvent(ev: WheelEvent<HTMLSpanElement>) {
     const span = toolbarButtonsRef.current
@@ -50,7 +53,9 @@ export const Toolbar = () => {
           <li className="contents">
             <ActionButton
               Element="button"
-              onClick={() => send({ type: 'Enter sketch' })}
+              onClick={() =>
+                send({ type: 'Enter sketch', data: { forceNewSketch: true } })
+              }
               icon={{
                 icon: 'sketch',
                 bgClassName,

src/components/ModelingMachineProvider.tsx

@@ -37,6 +37,7 @@ import { sceneInfra } from 'clientSideScene/sceneInfra'
 import { getSketchQuaternion } from 'clientSideScene/sceneEntities'
 import { startSketchOnDefault } from 'lang/modifyAst'
 import { Program } from 'lang/wasm'
+import { isSingleCursorInPipe } from 'lang/queryAst'
 
 type MachineContext<T extends AnyStateMachine> = {
   state: StateFrom<T>
@@ -182,7 +183,10 @@ export const ModelingMachineProvider = ({
 
           return canExtrudeSelection(selectionRanges)
         },
-        'Selection is one face': ({ selectionRanges }) => {
+        'Selection is on face': ({ selectionRanges }, { data }) => {
+          if (data?.forceNewSketch) return false
+          if (!isSingleCursorInPipe(selectionRanges, kclManager.ast))
+            return false
           return !!isCursorInSketchCommandRange(
             engineCommandManager.artifactMap,
             selectionRanges

src/lang/queryAst.ts

@@ -1,5 +1,5 @@
 import { ToolTip } from '../useStore'
-import { Selection } from 'lib/selections'
+import { Selection, Selections } from 'lib/selections'
 import {
   BinaryExpression,
   Program,
@@ -558,3 +558,24 @@ export function hasExtrudeSketchGroup({
   const varValue = programMemory?.root[varName]
   return varValue?.type === 'ExtrudeGroup' || varValue?.type === 'SketchGroup'
 }
+
+export function isSingleCursorInPipe(
+  selectionRanges: Selections,
+  ast: Program
+) {
+  if (selectionRanges.codeBasedSelections.length !== 1) return false
+  if (
+    doesPipeHaveCallExp({
+      ast,
+      selection: selectionRanges.codeBasedSelections[0],
+      calleeName: 'extrude',
+    })
+  )
+    return false
+  const selection = selectionRanges.codeBasedSelections[0]
+  const pathToNode = getNodePathFromSourceRange(ast, selection.range)
+  const nodeTypes = pathToNode.map(([, type]) => type)
+  if (nodeTypes.includes('FunctionExpression')) return false
+  if (nodeTypes.includes('PipeExpression')) return true
+  return false
+}

src/lib/selections.ts

@@ -9,7 +9,11 @@ import { SelectionRange } from '@uiw/react-codemirror'
 import { isOverlap } from 'lib/utils'
 import { isCursorInSketchCommandRange } from 'lang/util'
 import { Program } from 'lang/wasm'
-import { doesPipeHaveCallExp, getNodeFromPath } from 'lang/queryAst'
+import {
+  doesPipeHaveCallExp,
+  getNodeFromPath,
+  isSingleCursorInPipe,
+} from 'lang/queryAst'
 import { CommandArgument } from './commandTypes'
 import {
   STRAIGHT_SEGMENT,
@@ -455,6 +459,7 @@ function resetAndSetEngineEntitySelectionCmds(
 }
 
 export function isSketchPipe(selectionRanges: Selections) {
+  if (!isSingleCursorInPipe(selectionRanges, kclManager.ast)) return false
   return isCursorInSketchCommandRange(
     engineCommandManager.artifactMap,
     selectionRanges

src/machines/modelingMachine.ts

@@ -71,7 +71,12 @@ export type SetSelections =
     }
 
 export type ModelingMachineEvent =
-  | { type: 'Enter sketch' }
+  | {
+      type: 'Enter sketch'
+      data?: {
+        forceNewSketch?: boolean
+      }
+    }
   | {
       type: 'Select default plane'
       data: { plane: DefaultPlaneStr; normal: [number, number, number] }
@@ -153,7 +158,7 @@ export const modelingMachine = createMachine(
           'Enter sketch': [
             {
               target: 'animating to existing sketch',
-              cond: 'Selection is one face',
+              cond: 'Selection is on face',
               actions: ['set sketch metadata'],
             },
             'Sketch no face',