Refactoring before text edits (#8956)

Changes in preparation for #8238 features. # Important Notes Changed edit APIs: - **`graph.astModule` is deprecated.** It will be removed in my next PR. - Prefer `graph.edit` to start and commit an edit. - Use `graph.startEdit` / `graph.commitEdit` if the edit can't be confined to one scope.
2024-11-27 02:55:06 +03:00 · 2024-02-06 06:43:17 -08:00 · 2024-02-06 06:43:17 -08:00 · 06f18864f4
commit 06f18864f4
parent 6517384bbb
15 changed files with 492 additions and 364 deletions
--- a/app/gui2/shared/ast/parse.ts
+++ b/app/gui2/shared/ast/parse.ts
@ -1,6 +1,14 @@
 import * as map from 'lib0/map'
-import type { AstId, NodeChild, Owned } from '.'
-import { Token, asOwned, parentId, subtreeRoots } from '.'
+import {
+  Token,
+  asOwned,
+  isTokenId,
+  parentId,
+  subtreeRoots,
+  type AstId,
+  type NodeChild,
+  type Owned,
+} from '.'
 import { assert, assertDefined, assertEqual } from '../util/assert'
 import type { SourceRange, SourceRangeKey } from '../yjsModel'
 import { IdMap, isUuid, sourceRangeFromKey, sourceRangeKey } from '../yjsModel'
@ -52,42 +60,31 @@ export function abstract(
  tree: RawAst.Tree,
  code: string,
 ): { root: Owned; spans: SpanMap; toRaw: Map<AstId, RawAst.Tree> } {
-  const tokens = new Map()
-  const nodes = new Map()
-  const toRaw = new Map()
-  const root = abstractTree(module, tree, code, nodes, tokens, toRaw).node
-  const spans = { tokens, nodes }
-  return { root, spans, toRaw }
+  const abstractor = new Abstractor(module, code)
+  const root = abstractor.abstractTree(tree).node
+  const spans = { tokens: abstractor.tokens, nodes: abstractor.nodes }
+  return { root, spans, toRaw: abstractor.toRaw }
 }

-function abstractTree(
-  module: MutableModule,
-  tree: RawAst.Tree,
-  code: string,
-  nodesOut: NodeSpanMap,
-  tokensOut: TokenSpanMap,
-  toRaw: Map<AstId, RawAst.Tree>,
-): { whitespace: string | undefined; node: Owned } {
-  const recurseTree = (tree: RawAst.Tree) =>
-    abstractTree(module, tree, code, nodesOut, tokensOut, toRaw)
-  const recurseToken = (token: RawAst.Token.Token) => abstractToken(token, code, tokensOut)
-  const visitChildren = (tree: LazyObject) => {
-    const children: NodeChild<Owned | Token>[] = []
-    const visitor = (child: LazyObject) => {
-      if (RawAst.Tree.isInstance(child)) {
-        children.push(recurseTree(child))
-      } else if (RawAst.Token.isInstance(child)) {
-        children.push(recurseToken(child))
-      } else {
-        child.visitChildren(visitor)
-      }
-    }
-    tree.visitChildren(visitor)
-    return children
+class Abstractor {
+  private readonly module: MutableModule
+  private readonly code: string
+  readonly nodes: NodeSpanMap
+  readonly tokens: TokenSpanMap
+  readonly toRaw: Map<AstId, RawAst.Tree>
+
+  constructor(module: MutableModule, code: string) {
+    this.module = module
+    this.code = code
+    this.nodes = new Map()
+    this.tokens = new Map()
+    this.toRaw = new Map()
  }
+
+  abstractTree(tree: RawAst.Tree): { whitespace: string | undefined; node: Owned } {
    const whitespaceStart = tree.whitespaceStartInCodeParsed
    const whitespaceEnd = whitespaceStart + tree.whitespaceLengthInCodeParsed
-  const whitespace = code.substring(whitespaceStart, whitespaceEnd)
+    const whitespace = this.code.substring(whitespaceStart, whitespaceEnd)
    const codeStart = whitespaceEnd
    const codeEnd = codeStart + tree.childrenLengthInCodeParsed
    const spanKey = nodeKey(codeStart, codeEnd - codeStart)
@ -95,90 +92,90 @@ function abstractTree(
    switch (tree.type) {
      case RawAst.Tree.Type.BodyBlock: {
        const lines = Array.from(tree.statements, (line) => {
-        const newline = recurseToken(line.newline)
-        const expression = line.expression ? recurseTree(line.expression) : undefined
+          const newline = this.abstractToken(line.newline)
+          const expression = line.expression ? this.abstractTree(line.expression) : undefined
          return { newline, expression }
        })
-      node = BodyBlock.concrete(module, lines)
+        node = BodyBlock.concrete(this.module, lines)
        break
      }
      case RawAst.Tree.Type.Function: {
-      const name = recurseTree(tree.name)
-      const argumentDefinitions = Array.from(tree.args, (arg) => visitChildren(arg))
-      const equals = recurseToken(tree.equals)
-      const body = tree.body !== undefined ? recurseTree(tree.body) : undefined
-      node = Function.concrete(module, name, argumentDefinitions, equals, body)
+        const name = this.abstractTree(tree.name)
+        const argumentDefinitions = Array.from(tree.args, (arg) => this.abstractChildren(arg))
+        const equals = this.abstractToken(tree.equals)
+        const body = tree.body !== undefined ? this.abstractTree(tree.body) : undefined
+        node = Function.concrete(this.module, name, argumentDefinitions, equals, body)
        break
      }
      case RawAst.Tree.Type.Ident: {
-      const token = recurseToken(tree.token)
-      node = Ident.concrete(module, token)
+        const token = this.abstractToken(tree.token)
+        node = Ident.concrete(this.module, token)
        break
      }
      case RawAst.Tree.Type.Assignment: {
-      const pattern = recurseTree(tree.pattern)
-      const equals = recurseToken(tree.equals)
-      const value = recurseTree(tree.expr)
-      node = Assignment.concrete(module, pattern, equals, value)
+        const pattern = this.abstractTree(tree.pattern)
+        const equals = this.abstractToken(tree.equals)
+        const value = this.abstractTree(tree.expr)
+        node = Assignment.concrete(this.module, pattern, equals, value)
        break
      }
      case RawAst.Tree.Type.App: {
-      const func = recurseTree(tree.func)
-      const arg = recurseTree(tree.arg)
-      node = App.concrete(module, func, undefined, undefined, arg)
+        const func = this.abstractTree(tree.func)
+        const arg = this.abstractTree(tree.arg)
+        node = App.concrete(this.module, func, undefined, undefined, arg)
        break
      }
      case RawAst.Tree.Type.NamedApp: {
-      const func = recurseTree(tree.func)
-      const open = tree.open ? recurseToken(tree.open) : undefined
-      const name = recurseToken(tree.name)
-      const equals = recurseToken(tree.equals)
-      const arg = recurseTree(tree.arg)
-      const close = tree.close ? recurseToken(tree.close) : undefined
+        const func = this.abstractTree(tree.func)
+        const open = tree.open ? this.abstractToken(tree.open) : undefined
+        const name = this.abstractToken(tree.name)
+        const equals = this.abstractToken(tree.equals)
+        const arg = this.abstractTree(tree.arg)
+        const close = tree.close ? this.abstractToken(tree.close) : undefined
        const parens = open && close ? { open, close } : undefined
        const nameSpecification = { name, equals }
-      node = App.concrete(module, func, parens, nameSpecification, arg)
+        node = App.concrete(this.module, func, parens, nameSpecification, arg)
        break
      }
      case RawAst.Tree.Type.UnaryOprApp: {
-      const opr = recurseToken(tree.opr)
-      const arg = tree.rhs ? recurseTree(tree.rhs) : undefined
+        const opr = this.abstractToken(tree.opr)
+        const arg = tree.rhs ? this.abstractTree(tree.rhs) : undefined
        if (arg && opr.node.code() === '-') {
-        node = NegationApp.concrete(module, opr, arg)
+          node = NegationApp.concrete(this.module, opr, arg)
        } else {
-        node = UnaryOprApp.concrete(module, opr, arg)
+          node = UnaryOprApp.concrete(this.module, opr, arg)
        }
        break
      }
      case RawAst.Tree.Type.OprApp: {
-      const lhs = tree.lhs ? recurseTree(tree.lhs) : undefined
+        const lhs = tree.lhs ? this.abstractTree(tree.lhs) : undefined
        const opr = tree.opr.ok
-        ? [recurseToken(tree.opr.value)]
-        : Array.from(tree.opr.error.payload.operators, recurseToken)
-      const rhs = tree.rhs ? recurseTree(tree.rhs) : undefined
+          ? [this.abstractToken(tree.opr.value)]
+          : Array.from(tree.opr.error.payload.operators, this.abstractToken.bind(this))
+        const rhs = tree.rhs ? this.abstractTree(tree.rhs) : undefined
        if (opr.length === 1 && opr[0]?.node.code() === '.' && rhs?.node instanceof MutableIdent) {
          // Propagate type.
          const rhs_ = { ...rhs, node: rhs.node }
-        node = PropertyAccess.concrete(module, lhs, opr[0], rhs_)
+          node = PropertyAccess.concrete(this.module, lhs, opr[0], rhs_)
        } else {
-        node = OprApp.concrete(module, lhs, opr, rhs)
+          node = OprApp.concrete(this.module, lhs, opr, rhs)
        }
        break
      }
      case RawAst.Tree.Type.Number: {
        const tokens = []
-      if (tree.base) tokens.push(recurseToken(tree.base))
-      if (tree.integer) tokens.push(recurseToken(tree.integer))
+        if (tree.base) tokens.push(this.abstractToken(tree.base))
+        if (tree.integer) tokens.push(this.abstractToken(tree.integer))
        if (tree.fractionalDigits) {
-        tokens.push(recurseToken(tree.fractionalDigits.dot))
-        tokens.push(recurseToken(tree.fractionalDigits.digits))
+          tokens.push(this.abstractToken(tree.fractionalDigits.dot))
+          tokens.push(this.abstractToken(tree.fractionalDigits.digits))
        }
-      node = NumericLiteral.concrete(module, tokens)
+        node = NumericLiteral.concrete(this.module, tokens)
        break
      }
      case RawAst.Tree.Type.Wildcard: {
-      const token = recurseToken(tree.token)
-      node = Wildcard.concrete(module, token)
+        const token = this.abstractToken(tree.token)
+        node = Wildcard.concrete(this.module, token)
        break
      }
      // These expression types are (or will be) used for backend analysis.
@ -186,86 +183,98 @@ function abstractTree(
      // (which makes it impossible to give them unique IDs in the current IdMap format).
      case RawAst.Tree.Type.OprSectionBoundary:
      case RawAst.Tree.Type.TemplateFunction:
-      return { whitespace, node: recurseTree(tree.ast).node }
+        return { whitespace, node: this.abstractTree(tree.ast).node }
      case RawAst.Tree.Type.Invalid: {
-      const expression = recurseTree(tree.ast)
-      node = Invalid.concrete(module, expression)
+        const expression = this.abstractTree(tree.ast)
+        node = Invalid.concrete(this.module, expression)
        break
      }
      case RawAst.Tree.Type.Group: {
-      const open = tree.open ? recurseToken(tree.open) : undefined
-      const expression = tree.body ? recurseTree(tree.body) : undefined
-      const close = tree.close ? recurseToken(tree.close) : undefined
-      node = Group.concrete(module, open, expression, close)
+        const open = tree.open ? this.abstractToken(tree.open) : undefined
+        const expression = tree.body ? this.abstractTree(tree.body) : undefined
+        const close = tree.close ? this.abstractToken(tree.close) : undefined
+        node = Group.concrete(this.module, open, expression, close)
        break
      }
      case RawAst.Tree.Type.TextLiteral: {
-      const open = tree.open ? recurseToken(tree.open) : undefined
-      const newline = tree.newline ? recurseToken(tree.newline) : undefined
+        const open = tree.open ? this.abstractToken(tree.open) : undefined
+        const newline = tree.newline ? this.abstractToken(tree.newline) : undefined
        const elements = []
        for (const e of tree.elements) {
-        elements.push(...visitChildren(e))
+          elements.push(...this.abstractChildren(e))
        }
-      const close = tree.close ? recurseToken(tree.close) : undefined
-      node = TextLiteral.concrete(module, open, newline, elements, close)
+        const close = tree.close ? this.abstractToken(tree.close) : undefined
+        node = TextLiteral.concrete(this.module, open, newline, elements, close)
        break
      }
      case RawAst.Tree.Type.Documented: {
-      const open = recurseToken(tree.documentation.open)
+        const open = this.abstractToken(tree.documentation.open)
        const elements = []
        for (const e of tree.documentation.elements) {
-        elements.push(...visitChildren(e))
+          elements.push(...this.abstractChildren(e))
        }
-      const newlines = Array.from(tree.documentation.newlines, recurseToken)
-      const expression = tree.expression ? recurseTree(tree.expression) : undefined
-      node = Documented.concrete(module, open, elements, newlines, expression)
+        const newlines = Array.from(tree.documentation.newlines, this.abstractToken.bind(this))
+        const expression = tree.expression ? this.abstractTree(tree.expression) : undefined
+        node = Documented.concrete(this.module, open, elements, newlines, expression)
        break
      }
      case RawAst.Tree.Type.Import: {
        const recurseBody = (tree: RawAst.Tree) => {
-        const body = recurseTree(tree)
+          const body = this.abstractTree(tree)
          if (body.node instanceof Invalid && body.node.code() === '') return undefined
          return body
        }
        const recurseSegment = (segment: RawAst.MultiSegmentAppSegment) => ({
-        header: recurseToken(segment.header),
+          header: this.abstractToken(segment.header),
          body: segment.body ? recurseBody(segment.body) : undefined,
        })
        const polyglot = tree.polyglot ? recurseSegment(tree.polyglot) : undefined
        const from = tree.from ? recurseSegment(tree.from) : undefined
        const import_ = recurseSegment(tree.import)
-      const all = tree.all ? recurseToken(tree.all) : undefined
+        const all = tree.all ? this.abstractToken(tree.all) : undefined
        const as = tree.as ? recurseSegment(tree.as) : undefined
        const hiding = tree.hiding ? recurseSegment(tree.hiding) : undefined
-      node = Import.concrete(module, polyglot, from, import_, all, as, hiding)
+        node = Import.concrete(this.module, polyglot, from, import_, all, as, hiding)
        break
      }
      default: {
-      node = Generic.concrete(module, visitChildren(tree))
+        node = Generic.concrete(this.module, this.abstractChildren(tree))
      }
    }
-  toRaw.set(node.id, tree)
-  map.setIfUndefined(nodesOut, spanKey, (): Ast[] => []).unshift(node)
+    this.toRaw.set(node.id, tree)
+    map.setIfUndefined(this.nodes, spanKey, (): Ast[] => []).unshift(node)
    return { node, whitespace }
  }

-function abstractToken(
-  token: RawAst.Token,
-  code: string,
-  tokensOut: TokenSpanMap,
-): { whitespace: string; node: Token } {
+  private abstractToken(token: RawAst.Token): { whitespace: string; node: Token } {
    const whitespaceStart = token.whitespaceStartInCodeBuffer
    const whitespaceEnd = whitespaceStart + token.whitespaceLengthInCodeBuffer
-  const whitespace = code.substring(whitespaceStart, whitespaceEnd)
+    const whitespace = this.code.substring(whitespaceStart, whitespaceEnd)
    const codeStart = token.startInCodeBuffer
    const codeEnd = codeStart + token.lengthInCodeBuffer
-  const tokenCode = code.substring(codeStart, codeEnd)
+    const tokenCode = this.code.substring(codeStart, codeEnd)
    const key = tokenKey(codeStart, codeEnd - codeStart)
    const node = Token.new(tokenCode, token.type)
-  tokensOut.set(key, node)
+    this.tokens.set(key, node)
    return { whitespace, node }
  }

+  private abstractChildren(tree: LazyObject) {
+    const children: NodeChild<Owned | Token>[] = []
+    const visitor = (child: LazyObject) => {
+      if (RawAst.Tree.isInstance(child)) {
+        children.push(this.abstractTree(child))
+      } else if (RawAst.Token.isInstance(child)) {
+        children.push(this.abstractToken(child))
+      } else {
+        child.visitChildren(visitor)
+      }
+    }
+    tree.visitChildren(visitor)
+    return children
+  }
+}
+
 declare const nodeKeyBrand: unique symbol
 export type NodeKey = SourceRangeKey & { [nodeKeyBrand]: never }
 declare const tokenKeyBrand: unique symbol
@ -325,6 +334,87 @@ export function print(ast: Ast): PrintedSource {
  return { info, code }
 }

+/** @internal Used by `Ast.printSubtree`. Note that some AST types have overrides. */
+export function printAst(
+  ast: Ast,
+  info: SpanMap,
+  offset: number,
+  parentIndent: string | undefined,
+  verbatim?: boolean,
+): string {
+  let code = ''
+  for (const child of ast.concreteChildren(verbatim)) {
+    if (!isTokenId(child.node) && ast.module.checkedGet(child.node) === undefined) continue
+    if (child.whitespace != null) {
+      code += child.whitespace
+    } else if (code.length != 0) {
+      code += ' '
+    }
+    if (isTokenId(child.node)) {
+      const tokenStart = offset + code.length
+      const token = ast.module.getToken(child.node)
+      const span = tokenKey(tokenStart, token.code().length)
+      info.tokens.set(span, token)
+      code += token.code()
+    } else {
+      const childNode = ast.module.checkedGet(child.node)
+      assert(childNode != null)
+      code += childNode.printSubtree(info, offset + code.length, parentIndent, verbatim)
+      // Extra structural validation.
+      assertEqual(childNode.id, child.node)
+      if (parentId(childNode) !== ast.id) {
+        console.error(`Inconsistent parent pointer (expected ${ast.id})`, childNode)
+      }
+      assertEqual(parentId(childNode), ast.id)
+    }
+  }
+  const span = nodeKey(offset, code.length)
+  map.setIfUndefined(info.nodes, span, (): Ast[] => []).unshift(ast)
+  return code
+}
+
+/** @internal Use `Ast.code()' to stringify. */
+export function printBlock(
+  block: BodyBlock,
+  info: SpanMap,
+  offset: number,
+  parentIndent: string | undefined,
+  verbatim?: boolean,
+): string {
+  let blockIndent: string | undefined
+  let code = ''
+  for (const line of block.fields.get('lines')) {
+    code += line.newline.whitespace ?? ''
+    const newlineCode = block.module.getToken(line.newline.node).code()
+    // Only print a newline if this isn't the first line in the output, or it's a comment.
+    if (offset || code || newlineCode.startsWith('#')) {
+      // If this isn't the first line in the output, but there is a concrete newline token:
+      // if it's a zero-length newline, ignore it and print a normal newline.
+      code += newlineCode || '\n'
+    }
+    if (line.expression) {
+      if (blockIndent === undefined) {
+        if ((line.expression.whitespace?.length ?? 0) > (parentIndent?.length ?? 0)) {
+          blockIndent = line.expression.whitespace!
+        } else if (parentIndent !== undefined) {
+          blockIndent = parentIndent + '    '
+        } else {
+          blockIndent = ''
+        }
+      }
+      const validIndent = (line.expression.whitespace?.length ?? 0) > (parentIndent?.length ?? 0)
+      code += validIndent ? line.expression.whitespace : blockIndent
+      const lineNode = block.module.checkedGet(line.expression.node)
+      assertEqual(lineNode.id, line.expression.node)
+      assertEqual(parentId(lineNode), block.id)
+      code += lineNode.printSubtree(info, offset + code.length, blockIndent, verbatim)
+    }
+  }
+  const span = nodeKey(offset, code.length)
+  map.setIfUndefined(info.nodes, span, (): Ast[] => []).unshift(block)
+  return code
+}
+
 /** Parse the input as a block. */
 export function parseBlock(code: string, inModule?: MutableModule) {
  return parseBlockWithSpans(code, inModule).root
--- a/app/gui2/shared/ast/tree.ts
+++ b/app/gui2/shared/ast/tree.ts
@ -1,4 +1,3 @@
-import * as map from 'lib0/map'
 import type {
  Identifier,
  IdentifierOrOperatorIdentifier,
@ -18,12 +17,12 @@ import {
  isToken,
  isTokenId,
  newExternalId,
-  nodeKey,
  parentId,
  parse,
  parseBlock,
  print,
-  tokenKey,
+  printAst,
+  printBlock,
 } from '.'
 import { assert, assertDefined, assertEqual, bail } from '../util/assert'
 import type { Result } from '../util/data/result'
@ -123,35 +122,7 @@ export abstract class Ast {
    parentIndent: string | undefined,
    verbatim?: boolean,
  ): string {
-    let code = ''
-    for (const child of this.concreteChildren(verbatim)) {
-      if (!isTokenId(child.node) && this.module.checkedGet(child.node) === undefined) continue
-      if (child.whitespace != null) {
-        code += child.whitespace
-      } else if (code.length != 0) {
-        code += ' '
-      }
-      if (isTokenId(child.node)) {
-        const tokenStart = offset + code.length
-        const token = this.module.getToken(child.node)
-        const span = tokenKey(tokenStart, token.code().length)
-        info.tokens.set(span, token)
-        code += token.code()
-      } else {
-        const childNode = this.module.checkedGet(child.node)
-        assert(childNode != null)
-        code += childNode.printSubtree(info, offset + code.length, parentIndent, verbatim)
-        // Extra structural validation.
-        assertEqual(childNode.id, child.node)
-        if (parentId(childNode) !== this.id) {
-          console.error(`Inconsistent parent pointer (expected ${this.id})`, childNode)
-        }
-        assertEqual(parentId(childNode), this.id)
-      }
-    }
-    const span = nodeKey(offset, code.length)
-    map.setIfUndefined(info.nodes, span, (): Ast[] => []).unshift(this)
-    return code
+    return printAst(this, info, offset, parentIndent, verbatim)
  }

  /** Returns child subtrees, without information about the whitespace between them. */
@ -237,6 +208,12 @@ export abstract class MutableAst extends Ast {
    return old
  }

+  replaceValueChecked<T extends MutableAst>(replacement: Owned<T>): Owned<typeof this> {
+    const parentId = this.fields.get('parent')
+    assertDefined(parentId)
+    return this.replaceValue(replacement)
+  }
+
  /** Replace the parent of this object with a reference to a new placeholder object.
   *  Returns the object, now parentless, and the placeholder. */
  takeToReplace(): Removed<this> {
@ -354,11 +331,15 @@ interface AppFields {
 }
 export class App extends Ast {
  declare fields: FixedMap<AstFields & AppFields>
-
  constructor(module: Module, fields: FixedMapView<AstFields & AppFields>) {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableApp> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableApp) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    func: NodeChild<Owned>,
@ -486,6 +467,11 @@ export class UnaryOprApp extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableUnaryOprApp> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableUnaryOprApp) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    operator: NodeChild<Token>,
@ -549,6 +535,11 @@ export class NegationApp extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableNegationApp> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableNegationApp) return parsed
+  }
+
  static concrete(module: MutableModule, operator: NodeChild<Token>, argument: NodeChild<Owned>) {
    const base = module.baseObject('NegationApp')
    const id_ = base.get('id')
@ -606,6 +597,11 @@ export class OprApp extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableOprApp> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableOprApp) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    lhs: NodeChild<Owned> | undefined,
@ -693,6 +689,14 @@ export class PropertyAccess extends Ast {
    super(module, fields)
  }

+  static tryParse(
+    source: string,
+    module?: MutableModule,
+  ): Owned<MutablePropertyAccess> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutablePropertyAccess) return parsed
+  }
+
  static new(module: MutableModule, lhs: Owned, rhs: IdentLike) {
    const dot = unspaced(Token.new('.', RawAst.Token.Type.Operator))
    return this.concrete(
@ -894,6 +898,11 @@ export class Import extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableImport> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableImport) return parsed
+  }
+
  get polyglot(): Ast | undefined {
    return this.module.checkedGet(this.fields.get('polyglot')?.body?.node)
  }
@ -1022,17 +1031,17 @@ export class MutableImport extends Import implements MutableAst {

  replaceChild<T extends MutableAst>(target: AstId, replacement: Owned<T>) {
    const { polyglot, from, import: import_, as, hiding } = getAll(this.fields)
-    ;(polyglot?.body?.node === target
-      ? this.setPolyglot
+    polyglot?.body?.node === target
+      ? this.setPolyglot(replacement)
      : from?.body?.node === target
-      ? this.setFrom
+      ? this.setFrom(replacement)
      : import_.body?.node === target
-      ? this.setImport
+      ? this.setImport(replacement)
      : as?.body?.node === target
-      ? this.setAs
+      ? this.setAs(replacement)
      : hiding?.body?.node === target
-      ? this.setHiding
-      : bail(`Failed to find child ${target} in node ${this.externalId}.`))(replacement)
+      ? this.setHiding(replacement)
+      : bail(`Failed to find child ${target} in node ${this.externalId}.`)
  }
 }
 export interface MutableImport extends Import, MutableAst {
@ -1074,6 +1083,11 @@ export class TextLiteral extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableTextLiteral> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableTextLiteral) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    open: NodeChild<Token> | undefined,
@ -1134,6 +1148,11 @@ export class Documented extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableDocumented> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableDocumented) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    open: NodeChild<Token> | undefined,
@ -1259,6 +1278,11 @@ export class Group extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableGroup> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableGroup) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    open: NodeChild<Token> | undefined,
@ -1315,6 +1339,14 @@ export class NumericLiteral extends Ast {
    super(module, fields)
  }

+  static tryParse(
+    source: string,
+    module?: MutableModule,
+  ): Owned<MutableNumericLiteral> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableNumericLiteral) return parsed
+  }
+
  static concrete(module: MutableModule, tokens: NodeChild<Token>[]) {
    const base = module.baseObject('NumericLiteral')
    const fields = setAll(base, { tokens })
@ -1365,6 +1397,11 @@ export class Function extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableFunction> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableFunction) return parsed
+  }
+
  get name(): Ast {
    return this.module.checkedGet(this.fields.get('name').node)
  }
@ -1506,6 +1543,11 @@ export class Assignment extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableAssignment> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableAssignment) return parsed
+  }
+
  static concrete(
    module: MutableModule,
    pattern: NodeChild<Owned>,
@ -1580,6 +1622,11 @@ export class BodyBlock extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableBodyBlock> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableBodyBlock) return parsed
+  }
+
  static concrete(module: MutableModule, lines: OwnedBlockLine[]) {
    const base = module.baseObject('BodyBlock')
    const id_ = base.get('id')
@ -1616,38 +1663,7 @@ export class BodyBlock extends Ast {
    parentIndent: string | undefined,
    verbatim?: boolean,
  ): string {
-    let blockIndent: string | undefined
-    let code = ''
-    for (const line of this.fields.get('lines')) {
-      code += line.newline.whitespace ?? ''
-      const newlineCode = this.module.getToken(line.newline.node).code()
-      // Only print a newline if this isn't the first line in the output, or it's a comment.
-      if (offset || code || newlineCode.startsWith('#')) {
-        // If this isn't the first line in the output, but there is a concrete newline token:
-        // if it's a zero-length newline, ignore it and print a normal newline.
-        code += newlineCode || '\n'
-      }
-      if (line.expression) {
-        if (blockIndent === undefined) {
-          if ((line.expression.whitespace?.length ?? 0) > (parentIndent?.length ?? 0)) {
-            blockIndent = line.expression.whitespace!
-          } else if (parentIndent !== undefined) {
-            blockIndent = parentIndent + '    '
-          } else {
-            blockIndent = ''
-          }
-        }
-        const validIndent = (line.expression.whitespace?.length ?? 0) > (parentIndent?.length ?? 0)
-        code += validIndent ? line.expression.whitespace : blockIndent
-        const lineNode = this.module.checkedGet(line.expression.node)
-        assertEqual(lineNode.id, line.expression.node)
-        assertEqual(parentId(lineNode), this.id)
-        code += lineNode.printSubtree(info, offset + code.length, blockIndent, verbatim)
-      }
-    }
-    const span = nodeKey(offset, code.length)
-    map.setIfUndefined(info.nodes, span, (): Ast[] => []).unshift(this)
-    return code
+    return printBlock(this, info, offset, parentIndent, verbatim)
  }
 }
 export class MutableBodyBlock extends BodyBlock implements MutableAst {
@ -1772,6 +1788,11 @@ export class Ident extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableIdent> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableIdent) return parsed
+  }
+
  get token(): IdentifierToken {
    return this.module.getToken(this.fields.get('token').node) as IdentifierToken
  }
@ -1825,6 +1846,11 @@ export class Wildcard extends Ast {
    super(module, fields)
  }

+  static tryParse(source: string, module?: MutableModule): Owned<MutableWildcard> | undefined {
+    const parsed = parse(source, module)
+    if (parsed instanceof MutableWildcard) return parsed
+  }
+
  get token(): Token {
    return this.module.getToken(this.fields.get('token').node)
  }
--- a/app/gui2/src/components/GraphEditor.vue
+++ b/app/gui2/src/components/GraphEditor.vue
@ -265,12 +265,12 @@ const graphBindingsHandler = graphBindings.handler({
        bail(`Cannot get the method name for the current execution stack item. ${currentMethod}`)
      }
      const currentFunctionEnv = environmentForNodes(selected.values())
-      const module = graphStore.astModule
      const topLevel = graphStore.topLevel
      if (!topLevel) {
        bail('BUG: no top level, collapsing not possible.')
      }
-      const edit = module.edit()
+      const { position } = collapsedNodePlacement(DEFAULT_NODE_SIZE, currentFunctionEnv)
+      graphStore.edit((edit) => {
        const { refactoredNodeId, collapsedNodeIds, outputNodeId } = performCollapse(
          info,
          edit.getVersion(topLevel),
@ -280,19 +280,21 @@ const graphBindingsHandler = graphBindings.handler({
        const collapsedFunctionEnv = environmentForNodes(collapsedNodeIds.values())
        // For collapsed function, only selected nodes would affect placement of the output node.
        collapsedFunctionEnv.nodeRects = collapsedFunctionEnv.selectedNodeRects
-      const { position } = collapsedNodePlacement(DEFAULT_NODE_SIZE, currentFunctionEnv)
        edit
          .checkedGet(refactoredNodeId)
          .mutableNodeMetadata()
          .set('position', { x: position.x, y: position.y })
        if (outputNodeId != null) {
-        const { position } = previousNodeDictatedPlacement(DEFAULT_NODE_SIZE, collapsedFunctionEnv)
+          const { position } = previousNodeDictatedPlacement(
+            DEFAULT_NODE_SIZE,
+            collapsedFunctionEnv,
+          )
          edit
            .checkedGet(outputNodeId)
            .mutableNodeMetadata()
            .set('position', { x: position.x, y: position.y })
        }
-      graphStore.commitEdit(edit)
+      })
    } catch (err) {
      console.log('Error while collapsing, this is not normal.', err)
    }
--- a/app/gui2/src/components/GraphEditor/GraphEdges.vue
+++ b/app/gui2/src/components/GraphEditor/GraphEdges.vue
@ -56,7 +56,7 @@ const editingEdge: Interaction = {
 interaction.setWhen(() => graph.unconnectedEdge != null, editingEdge)

 function disconnectEdge(target: PortId) {
-  graph.commitDirect((edit) => {
+  graph.edit((edit) => {
    if (!graph.updatePortValue(edit, target, undefined)) {
      if (isAstId(target)) {
        console.warn(`Failed to disconnect edge from port ${target}, falling back to direct edit.`)
@ -79,7 +79,7 @@ function createEdge(source: AstId, target: PortId) {
    return console.error(`Failed to connect edge, source or target node not found.`)
  }

-  const edit = graph.astModule.edit()
+  const edit = graph.startEdit()
  const reorderResult = graph.ensureCorrectNodeOrder(edit, sourceNode, targetNode)
  if (reorderResult === 'circular') {
    // Creating this edge would create a circular dependency. Prevent that and display error.
--- a/app/gui2/src/components/GraphEditor/NodeWidgetTree.vue
+++ b/app/gui2/src/components/GraphEditor/NodeWidgetTree.vue
@ -31,11 +31,9 @@ const observedLayoutTransitions = new Set([
 ])

 function handleWidgetUpdates(update: WidgetUpdate) {
+  const edit = update.edit ?? graph.startEdit()
  if (update.portUpdate) {
-    const {
-      edit,
-      portUpdate: { value, origin },
-    } = update
+    const { value, origin } = update.portUpdate
    if (Ast.isAstId(origin)) {
      const ast =
        value instanceof Ast.Ast
@ -48,7 +46,7 @@ function handleWidgetUpdates(update: WidgetUpdate) {
      console.error(`[UPDATE ${origin}] Invalid top-level origin. Expected expression ID.`)
    }
  }
-  graph.commitEdit(update.edit)
+  graph.commitEdit(edit)
  // This handler is guaranteed to be the last handler in the chain.
  return true
 }
--- a/app/gui2/src/components/GraphEditor/widgets/WidgetCheckbox.vue
+++ b/app/gui2/src/components/GraphEditor/widgets/WidgetCheckbox.vue
@ -17,7 +17,7 @@ const value = computed({
    return WidgetInput.valueRepr(props.input)?.endsWith('True') ?? false
  },
  set(value) {
-    const edit = graph.astModule.edit()
+    const edit = graph.startEdit()
    if (props.input.value instanceof Ast.Ast) {
      setBoolNode(
        edit.getVersion(props.input.value),
--- a/app/gui2/src/components/GraphEditor/widgets/WidgetFunction.vue
+++ b/app/gui2/src/components/GraphEditor/widgets/WidgetFunction.vue
@ -162,10 +162,8 @@ const widgetConfiguration = computed(() => {
 function handleArgUpdate(update: WidgetUpdate): boolean {
  const app = application.value
  if (update.portUpdate && app instanceof ArgumentApplication) {
-    const {
-      edit,
-      portUpdate: { value, origin },
-    } = update
+    const { value, origin } = update.portUpdate
+    const edit = update.edit ?? graph.startEdit()
    // Find the updated argument by matching origin port/expression with the appropriate argument.
    // We are interested only in updates at the top level of the argument AST. Updates from nested
    // widgets do not need to be processed at the function application level.
--- a/app/gui2/src/components/GraphEditor/widgets/WidgetNumber.vue
+++ b/app/gui2/src/components/GraphEditor/widgets/WidgetNumber.vue
@ -16,7 +16,6 @@ const value = computed({
  },
  set(value) {
    props.onUpdate({
-      edit: graph.astModule.edit(),
      portUpdate: { value: value.toString(), origin: asNot<TokenId>(props.input.portId) },
    })
  },
--- a/app/gui2/src/components/GraphEditor/widgets/WidgetSelection.vue
+++ b/app/gui2/src/components/GraphEditor/widgets/WidgetSelection.vue
@ -10,6 +10,7 @@ import { useGraphStore } from '@/stores/graph'
 import { requiredImports, type RequiredImport } from '@/stores/graph/imports.ts'
 import { useSuggestionDbStore } from '@/stores/suggestionDatabase'
 import { type SuggestionEntry } from '@/stores/suggestionDatabase/entry.ts'
+import { Ast } from '@/util/ast'
 import type { TokenId } from '@/util/ast/abstract.ts'
 import { ArgumentInfoKey } from '@/util/callTree'
 import { asNot } from '@/util/data/types.ts'
@ -117,9 +118,11 @@ function toggleDropdownWidget() {

 // When the selected index changes, we update the expression content.
 watch(selectedIndex, (_index) => {
-  const edit = graph.astModule.edit()
-  if (selectedTag.value?.requiredImports)
+  let edit: Ast.MutableModule | undefined
+  if (selectedTag.value?.requiredImports) {
+    edit = graph.startEdit()
    graph.addMissingImports(edit, selectedTag.value.requiredImports)
+  }
  props.onUpdate({
    edit,
    portUpdate: {
--- a/app/gui2/src/components/GraphEditor/widgets/WidgetText.vue
+++ b/app/gui2/src/components/GraphEditor/widgets/WidgetText.vue
@ -16,7 +16,6 @@ const value = computed({
  },
  set(value) {
    props.onUpdate({
-      edit: graph.astModule.edit(),
      portUpdate: { value: value.toString(), origin: asNot<TokenId>(props.input.portId) },
    })
  },
--- a/app/gui2/src/components/GraphEditor/widgets/WidgetVector.vue
+++ b/app/gui2/src/components/GraphEditor/widgets/WidgetVector.vue
@ -36,9 +36,7 @@ const value = computed({
  set(value) {
    // TODO[ao]: here we re-create AST. It would be better to reuse existing AST nodes.
    const newCode = `[${value.map((item) => item.code()).join(', ')}]`
-    const edit = graph.astModule.edit()
    props.onUpdate({
-      edit,
      portUpdate: { value: newCode, origin: asNot<TokenId>(props.input.portId) },
    })
  },
--- a/app/gui2/src/composables/stackNavigator.ts
+++ b/app/gui2/src/composables/stackNavigator.ts
@ -55,8 +55,12 @@ export function useStackNavigator() {
  }

  function enterNode(id: AstId) {
-    const node = graphStore.astModule.get(id)!
-    const expressionInfo = graphStore.db.getExpressionInfo(id)
+    const externalId = graphStore.db.idToExternal(id)
+    if (externalId == null) {
+      console.debug("Cannot enter node that hasn't been committed yet.")
+      return
+    }
+    const expressionInfo = graphStore.db.getExpressionInfo(externalId)
    if (expressionInfo == null || expressionInfo.methodCall == null) {
      console.debug('Cannot enter node that has no method call.')
      return
@ -71,7 +75,7 @@ export function useStackNavigator() {
      console.debug('Cannot enter node that is not defined on current module.')
      return
    }
-    projectStore.executionContext.push(node.externalId)
+    projectStore.executionContext.push(externalId)
    graphStore.updateState()
    breadcrumbs.value = projectStore.executionContext.desiredStack.slice()
  }
--- a/app/gui2/src/providers/widgetRegistry.ts
+++ b/app/gui2/src/providers/widgetRegistry.ts
@ -141,7 +141,7 @@ export interface WidgetProps<T> {
 * is committed in {@link NodeWidgetTree}.
 */
 export interface WidgetUpdate {
-  edit: MutableModule
+  edit?: MutableModule | undefined
  portUpdate?: {
    value: Ast.Owned | string | undefined
    origin: PortId
--- a/app/gui2/src/stores/graph/imports.ts
+++ b/app/gui2/src/stores/graph/imports.ts
@ -537,8 +537,8 @@ if (import.meta.vitest) {
  })

  const parseImport = (code: string): Import | null => {
-    const ast: Ast.Ast = Ast.parse(code)
-    return ast instanceof Ast.Import ? recognizeImport(ast) : null
+    const ast = Ast.Import.tryParse(code)
+    return ast ? recognizeImport(ast) : null
  }

  test.each([
--- a/app/gui2/src/stores/graph/index.ts
+++ b/app/gui2/src/stores/graph/index.ts
@ -212,22 +212,21 @@ export const useGraphStore = defineStore('graph', () => {
    const mod = proj.module
    if (!mod) return
    const ident = generateUniqueIdent()
-    const edit = astModule.value.edit()
-    if (withImports) addMissingImports(edit, withImports)
    const currentFunc = 'main'
    const method = Ast.findModuleMethod(topLevel.value!, currentFunc)
    if (!method) {
      console.error(`BUG: Cannot add node: No current function.`)
      return
    }
-    const functionBlock = edit.getVersion(method).bodyAsBlock()
-    const rhs = Ast.parse(expression, edit)
    metadata.position = { x: position.x, y: position.y }
+    return edit((edit) => {
+      if (withImports) addMissingImports(edit, withImports)
+      const rhs = Ast.parse(expression, edit)
      rhs.setNodeMetadata(metadata)
      const assignment = Ast.Assignment.new(edit, ident, rhs)
-    functionBlock.push(assignment)
-    commitEdit(edit)
+      edit.getVersion(method).bodyAsBlock().push(assignment)
      return asNodeId(rhs.id)
+    })
  }

  function addMissingImports(edit: MutableModule, newImports: RequiredImport[]) {
@ -244,21 +243,25 @@ export const useGraphStore = defineStore('graph', () => {
  }

  function deleteNodes(ids: NodeId[]) {
-    commitDirect((edit) => {
+    edit(
+      (edit) => {
        for (const id of ids) {
          const node = db.nodeIdToNode.get(id)
-        if (!node) return
+          if (!node) continue
          const outerExpr = edit.get(node.outerExprId)
          if (outerExpr) Ast.deleteFromParentBlock(outerExpr)
          nodeRects.delete(id)
        }
-    }, true)
+      },
+      true,
+      true,
+    )
  }

  function setNodeContent(id: NodeId, content: string) {
    const node = db.nodeIdToNode.get(id)
    if (!node) return
-    commitDirect((edit) => {
+    edit((edit) => {
      edit.getVersion(node.rootSpan).replaceValue(Ast.parse(content, edit))
    })
  }
@ -276,12 +279,9 @@ export const useGraphStore = defineStore('graph', () => {
    if (!nodeAst) return
    const oldPos = nodeAst.nodeMetadata.get('position')
    if (oldPos?.x !== position.x || oldPos?.y !== position.y) {
-      commitDirect((edit) => {
-        edit
-          .getVersion(nodeAst)
-          .mutableNodeMetadata()
-          .set('position', { x: position.x, y: position.y })
-      }, true)
+      editNodeMetadata(nodeAst, (metadata) =>
+        metadata.set('position', { x: position.x, y: position.y }),
+      )
    }
  }

@ -297,25 +297,23 @@ export const useGraphStore = defineStore('graph', () => {
  function setNodeVisualizationId(nodeId: NodeId, vis: Opt<VisualizationIdentifier>) {
    const nodeAst = astModule.value?.get(nodeId)
    if (!nodeAst) return
-    commitDirect((edit) => {
-      const metadata = edit.getVersion(nodeAst).mutableNodeMetadata()
+    editNodeMetadata(nodeAst, (metadata) =>
      metadata.set(
        'visualization',
        normalizeVisMetadata(vis, metadata.get('visualization')?.visible),
+      ),
    )
-    }, true)
  }

  function setNodeVisualizationVisible(nodeId: NodeId, visible: boolean) {
    const nodeAst = astModule.value?.get(nodeId)
    if (!nodeAst) return
-    commitDirect((edit) => {
-      const metadata = edit.getVersion(nodeAst).mutableNodeMetadata()
+    editNodeMetadata(nodeAst, (metadata) =>
      metadata.set(
        'visualization',
        normalizeVisMetadata(metadata.get('visualization')?.identifier, visible),
+      ),
    )
-    }, true)
  }

  function updateNodeRect(nodeId: NodeId, rect: Rect) {
@ -329,12 +327,9 @@ export const useGraphStore = defineStore('graph', () => {
        screenBounds: Rect.Zero,
        mousePosition: Vec2.Zero,
      })
-      commitDirect((edit) => {
-        edit
-          .getVersion(nodeAst)
-          .mutableNodeMetadata()
-          .set('position', { x: position.x, y: position.y })
-      }, true)
+      editNodeMetadata(nodeAst, (metadata) =>
+        metadata.set('position', { x: position.x, y: position.y }),
+      )
      nodeRects.set(nodeId, new Rect(position, rect.size))
    } else {
      nodeRects.set(nodeId, rect)
@ -405,6 +400,10 @@ export const useGraphStore = defineStore('graph', () => {
    return true
  }

+  function startEdit(): MutableModule {
+    return astModule.value.edit()
+  }
+
  /** Apply the given `edit` to the state.
   *
   *  @param skipTreeRepair - If the edit is known not to require any parenthesis insertion, this may be set to `true`
@ -422,23 +421,34 @@ export const useGraphStore = defineStore('graph', () => {
    Y.applyUpdateV2(syncModule.value!.ydoc, Y.encodeStateAsUpdateV2(edit.ydoc), 'local')
  }

-  /** Run the given callback with direct access to the document module. Any edits to the module will be committed
-   *  unconditionally; use with caution to avoid committing partial edits.
+  /** Edit the AST module.
   *
-   *  @param skipTreeRepair - If the edit is known not to require any parenthesis insertion, this may be set to `true`
-   *  for better performance.
+   *  Optimization options: These are safe to use for metadata-only edits; otherwise, they require extreme caution.
+   *
+   *  @param skipTreeRepair - If the edit is certain not to produce incorrect or non-canonical syntax, this may be set
+   *  to `true` for better performance.
+   *  @param direct - Apply all changes directly to the synchronized module; they will be committed even if the callback
+   *  exits by throwing an exception.
   */
-  function commitDirect(f: (edit: MutableModule) => void, skipTreeRepair?: boolean) {
-    const edit = syncModule.value
+  function edit<T>(f: (edit: MutableModule) => T, skipTreeRepair?: boolean, direct?: boolean): T {
+    const edit = direct ? syncModule.value : syncModule.value?.edit()
    assert(edit != null)
+    let result
    edit.ydoc.transact(() => {
-      f(edit)
+      result = f(edit)
      if (!skipTreeRepair) {
        const root = edit.root()
        assert(root instanceof Ast.BodyBlock)
        Ast.repair(root, edit)
      }
    }, 'local')
+    if (!direct)
+      Y.applyUpdateV2(syncModule.value!.ydoc, Y.encodeStateAsUpdateV2(edit.ydoc), 'local')
+    return result!
+  }
+
+  function editNodeMetadata(ast: Ast.Ast, f: (metadata: Ast.MutableNodeMetadata) => void) {
+    edit((edit) => f(edit.getVersion(ast).mutableNodeMetadata()), true, true)
  }

  function mockExpressionUpdate(
@ -564,8 +574,9 @@ export const useGraphStore = defineStore('graph', () => {
    updatePortValue,
    setEditedNode,
    updateState,
+    startEdit,
    commitEdit,
-    commitDirect,
+    edit,
    addMissingImports,
  }
 })