diff --git a/app/gui2/shared/ast/index.ts b/app/gui2/shared/ast/index.ts
index e1088e3efb9f..f3b9be6bebdd 100644
--- a/app/gui2/shared/ast/index.ts
+++ b/app/gui2/shared/ast/index.ts
@@ -43,7 +43,7 @@ export function parentId(ast: Ast): AstId | undefined {
 export function subtrees(module: Module, ids: Iterable<AstId>) {
   return reachable(ids, (id) => {
     const parent = module.tryGet(id)?.parent()
-    return parent ? [parent.id] : []
+    return parent ? [id, parent.id] : [id]
   })
 }
 
diff --git a/app/gui2/src/stores/graph/graphDatabase.ts b/app/gui2/src/stores/graph/graphDatabase.ts
index b7890766a034..3bd7429d02ef 100644
--- a/app/gui2/src/stores/graph/graphDatabase.ts
+++ b/app/gui2/src/stores/graph/graphDatabase.ts
@@ -341,8 +341,6 @@ export class GraphDb {
     dirtyNodes: Set<AstId>,
   ) {
     const functionChanged = functionAst_.id !== this.currentFunction
-    // Note: `subtrees` returns a set that has the iteration order of all `Ast.ID`s in the order they appear in the
-    // module AST. This is important to ensure that nodes are updated in the correct order.
     const knownDirtySubtrees = functionChanged ? null : subtrees(functionAst_.module, dirtyNodes)
     const subtreeDirty = (id: AstId) => !knownDirtySubtrees || knownDirtySubtrees.has(id)
     this.currentFunction = functionAst_.id
diff --git a/app/gui2/src/util/ast/__tests__/abstract.test.ts b/app/gui2/src/util/ast/__tests__/abstract.test.ts
index 1d63c1a58163..a7f80983daee 100644
--- a/app/gui2/src/util/ast/__tests__/abstract.test.ts
+++ b/app/gui2/src/util/ast/__tests__/abstract.test.ts
@@ -4,8 +4,10 @@ import {
   MutableModule,
   TextLiteral,
   escapeTextLiteral,
+  findModuleMethod,
   substituteIdentifier,
   substituteQualifiedName,
+  subtrees,
   unescapeTextLiteral,
   type Identifier,
 } from '@/util/ast/abstract'
@@ -423,18 +425,33 @@ test('Insert new expression', () => {
 
 type SimpleModule = {
   root: Ast.BodyBlock
+  main: Ast.Function
+  mainBlock: Ast.BodyBlock
   assignment: Ast.Assignment
 }
 function simpleModule(): SimpleModule {
   const code = 'main =\n    text1 = "foo"\n'
   const root = Ast.parseBlock(code)
-  const main = Ast.functionBlock(root, 'main')!
-  expect(main).not.toBeNull()
-  const assignment: Ast.Assignment = main.statements().next().value
+  const main = findModuleMethod(root, 'main')!
+  const mainBlock = main.body instanceof Ast.BodyBlock ? main.body : null
+  assert(mainBlock != null)
+  expect(mainBlock).toBeInstanceOf(Ast.BodyBlock)
+  const assignment: Ast.Assignment = mainBlock.statements().next().value
   expect(assignment).toBeInstanceOf(Ast.Assignment)
-  return { root, assignment }
+  return { root, main, mainBlock, assignment }
 }
 
+test('Subtrees', () => {
+  const { root, main, mainBlock, assignment } = simpleModule()
+  const module = root.module
+  expect(subtrees(module, [assignment.id])).toEqual(
+    new Set([assignment.id, mainBlock.id, main.id, root.id]),
+  )
+  expect(subtrees(module, [mainBlock.id])).toEqual(new Set([mainBlock.id, main.id, root.id]))
+  expect(subtrees(module, [main.id])).toEqual(new Set([main.id, root.id]))
+  expect(subtrees(module, [root.id])).toEqual(new Set([root.id]))
+})
+
 test('Modify subexpression', () => {
   const { root, assignment } = simpleModule()
   expect(assignment.expression).not.toBeNull()