# AST Grep and Transform

> Learn how to parse source code using ast-grep and update nodes in the tree.

[Your browser does not support the audio element.](https://microsoft.github.io/genaiscript/llms-full.txt.mp3 "Generated by 🤖 AI with random voice personality. Content interpretation may vary.")

Generated by 🤖 AI

This page describes a strategy to build GenAI scripts that use Abstract Syntax Trees (AST) to parse and modify source code. When applicable, it provides an extremely flexible and stable way to apply large scale changes to source code. Interested? Let’s dive in!

## The strategy of AST-based code transformation

[Section titled “The strategy of AST-based code transformation”](#the-strategy-of-ast-based-code-transformation)

One of the challenge when creating GenAI scripts that update source code is to correctly locate and update source code. The slightest mistake in the location of the code to update can lead to a broken code. This is especially true when the code to update is not a simple string, but a complex structure like an object or a function call.

In some cases, you know “precisely” which part of the code you want to update. For example, you want to refresh the documentation of a function after a change. You know that the documentation is located just before the function definition at least in the sense of the programming language but the number of empty lines or space may vary.

math.ts

```ts
/** sums a and b */
function fn(a: number, b: number): number {
  return a - b; // oops outdated
}
```

In such scenario, you can use the Abstract Syntax Tree (AST) to locate the code to update. The AST is a tree representation of code.

![](<data:image/svg+xml,%3csvg aria-roledescription='flowchart-v2' role='graphics-document document' viewBox='0 0 189.21875 350' style='max-width: 189.21875px%3b' class='flowchart' xmlns:xlink='http://www.w3.org/1999/xlink' xmlns='http://www.w3.org/2000/svg' width='100%25' id='mermaid-0'%3e%3cstyle%3e%23mermaid-0%7bfont-family:arial%2csans-serif%3bfont-size:16px%3bfill:%23333%3b%7d%40keyframes edge-animation-frame%7bfrom%7bstroke-dashoffset:0%3b%7d%7d%40keyframes dash%7bto%7bstroke-dashoffset:0%3b%7d%7d%23mermaid-0 .edge-animation-slow%7bstroke-dasharray:9%2c5!important%3bstroke-dashoffset:900%3banimation:dash 50s linear infinite%3bstroke-linecap:round%3b%7d%23mermaid-0 .edge-animation-fast%7bstroke-dasharray:9%2c5!important%3bstroke-dashoffset:900%3banimation:dash 20s linear infinite%3bstroke-linecap:round%3b%7d%23mermaid-0 .error-icon%7bfill:%23552222%3b%7d%23mermaid-0 .error-text%7bfill:%23552222%3bstroke:%23552222%3b%7d%23mermaid-0 .edge-thickness-normal%7bstroke-width:1px%3b%7d%23mermaid-0 .edge-thickness-thick%7bstroke-width:3.5px%3b%7d%23mermaid-0 .edge-pattern-solid%7bstroke-dasharray:0%3b%7d%23mermaid-0 .edge-thickness-invisible%7bstroke-width:0%3bfill:none%3b%7d%23mermaid-0 .edge-pattern-dashed%7bstroke-dasharray:3%3b%7d%23mermaid-0 .edge-pattern-dotted%7bstroke-dasharray:2%3b%7d%23mermaid-0 .marker%7bfill:%23333333%3bstroke:%23333333%3b%7d%23mermaid-0 .marker.cross%7bstroke:%23333333%3b%7d%23mermaid-0 svg%7bfont-family:arial%2csans-serif%3bfont-size:16px%3b%7d%23mermaid-0 p%7bmargin:0%3b%7d%23mermaid-0 .label%7bfont-family:arial%2csans-serif%3bcolor:%23333%3b%7d%23mermaid-0 .cluster-label text%7bfill:%23333%3b%7d%23mermaid-0 .cluster-label span%7bcolor:%23333%3b%7d%23mermaid-0 .cluster-label span p%7bbackground-color:transparent%3b%7d%23mermaid-0 .label text%2c%23mermaid-0 span%7bfill:%23333%3bcolor:%23333%3b%7d%23mermaid-0 .node rect%2c%23mermaid-0 .node circle%2c%23mermaid-0 .node ellipse%2c%23mermaid-0 .node polygon%2c%23mermaid-0 .node path%7bfill:%23ECECFF%3bstroke:%239370DB%3bstroke-width:1px%3b%7d%23mermaid-0 .rough-node .label text%2c%23mermaid-0 .node .label text%2c%23mermaid-0 .image-shape .label%2c%23mermaid-0 .icon-shape .label%7btext-anchor:middle%3b%7d%23mermaid-0 .node .katex path%7bfill:black%3bstroke:black%3bstroke-width:1px%3b%7d%23mermaid-0 .rough-node .label%2c%23mermaid-0 .node .label%2c%23mermaid-0 .image-shape .label%2c%23mermaid-0 .icon-shape .label%7btext-align:center%3b%7d%23mermaid-0 .node.clickable%7bcursor:pointer%3b%7d%23mermaid-0 .root .anchor path%7bfill:%23333333!important%3bstroke-width:0%3bstroke:%23333333%3b%7d%23mermaid-0 .arrowheadPath%7bfill:%23333333%3b%7d%23mermaid-0 .edgePath .path%7bstroke:%23333333%3bstroke-width:2.0px%3b%7d%23mermaid-0 .flowchart-link%7bstroke:%23333333%3bfill:none%3b%7d%23mermaid-0 .edgeLabel%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3btext-align:center%3b%7d%23mermaid-0 .edgeLabel p%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3b%7d%23mermaid-0 .edgeLabel rect%7bopacity:0.5%3bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3bfill:rgba(232%2c232%2c232%2c 0.8)%3b%7d%23mermaid-0 .labelBkg%7bbackground-color:rgba(232%2c 232%2c 232%2c 0.5)%3b%7d%23mermaid-0 .cluster rect%7bfill:%23ffffde%3bstroke:%23aaaa33%3bstroke-width:1px%3b%7d%23mermaid-0 .cluster text%7bfill:%23333%3b%7d%23mermaid-0 .cluster span%7bcolor:%23333%3b%7d%23mermaid-0 div.mermaidTooltip%7bposition:absolute%3btext-align:center%3bmax-width:200px%3bpadding:2px%3bfont-family:arial%2csans-serif%3bfont-size:12px%3bbackground:hsl(80%2c 100%25%2c 96.2745098039%25)%3bborder:1px solid %23aaaa33%3bborder-radius:2px%3bpointer-events:none%3bz-index:100%3b%7d%23mermaid-0 .flowchartTitleText%7btext-anchor:middle%3bfont-size:18px%3bfill:%23333%3b%7d%23mermaid-0 rect.text%7bfill:none%3bstroke-width:0%3b%7d%23mermaid-0 .icon-shape%2c%23mermaid-0 .image-shape%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3btext-align:center%3b%7d%23mermaid-0 .icon-shape p%2c%23mermaid-0 .image-shape p%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3bpadding:2px%3b%7d%23mermaid-0 .icon-shape rect%2c%23mermaid-0 .image-shape rect%7bopacity:0.5%3bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3bfill:rgba(232%2c232%2c232%2c 0.8)%3b%7d%23mermaid-0 .label-icon%7bdisplay:inline-block%3bheight:1em%3boverflow:visible%3bvertical-align:-0.125em%3b%7d%23mermaid-0 .node .label-icon path%7bfill:currentColor%3bstroke:revert%3bstroke-width:revert%3b%7d%23mermaid-0 :root%7b--mermaid-font-family:arial%2csans-serif%3b%7d%3c/style%3e%3cg%3e%3cmarker orient='auto' markerHeight='8' markerWidth='8' markerUnits='userSpaceOnUse' refY='5' refX='5' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-0_flowchart-v2-pointEnd'%3e%3cpath style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 0 0 L 10 5 L 0 10 z'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='8' markerWidth='8' markerUnits='userSpaceOnUse' refY='5' refX='4.5' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-0_flowchart-v2-pointStart'%3e%3cpath style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 0 5 L 10 10 L 10 0 z'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5' refX='11' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-0_flowchart-v2-circleEnd'%3e%3ccircle style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' r='5' cy='5' cx='5'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5' refX='-1' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-0_flowchart-v2-circleStart'%3e%3ccircle style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' r='5' cy='5' cx='5'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5.2' refX='12' viewBox='0 0 11 11' class='marker cross flowchart-v2' id='mermaid-0_flowchart-v2-crossEnd'%3e%3cpath style='stroke-width: 2%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 1%2c1 l 9%2c9 M 10%2c1 l -9%2c9'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5.2' refX='-1' viewBox='0 0 11 11' class='marker cross flowchart-v2' id='mermaid-0_flowchart-v2-crossStart'%3e%3cpath style='stroke-width: 2%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 1%2c1 l 9%2c9 M 10%2c1 l -9%2c9'/%3e%3c/marker%3e%3cg class='root'%3e%3cg class='clusters'/%3e%3cg class='edgePaths'%3e%3cpath marker-end='url(%23mermaid-0_flowchart-v2-pointEnd)' style='' class='edge-thickness-normal edge-pattern-solid edge-thickness-normal edge-pattern-solid flowchart-link' id='L_A_B_0' d='M94.609%2c86L94.609%2c90.167C94.609%2c94.333%2c94.609%2c102.667%2c94.609%2c110.333C94.609%2c118%2c94.609%2c125%2c94.609%2c128.5L94.609%2c132'/%3e%3cpath marker-end='url(%23mermaid-0_flowchart-v2-pointEnd)' style='' class='edge-thickness-normal edge-pattern-solid edge-thickness-normal edge-pattern-solid flowchart-link' id='L_B_D_0' d='M94.609%2c214L94.609%2c218.167C94.609%2c222.333%2c94.609%2c230.667%2c94.609%2c238.333C94.609%2c246%2c94.609%2c253%2c94.609%2c256.5L94.609%2c260'/%3e%3c/g%3e%3cg class='edgeLabels'%3e%3cg class='edgeLabel'%3e%3cg transform='translate(0%2c 0)' class='label'%3e%3cforeignObject height='0' width='0'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' class='labelBkg' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='edgeLabel'%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3cg class='edgeLabel'%3e%3cg transform='translate(0%2c 0)' class='label'%3e%3cforeignObject height='0' width='0'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' class='labelBkg' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='edgeLabel'%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3c/g%3e%3cg class='nodes'%3e%3cg transform='translate(94.609375%2c 47)' id='flowchart-A-0' class='node default'%3e%3crect height='78' width='162.515625' y='-39' x='-81.2578125' ry='5' rx='5' style='' class='basic label-container'/%3e%3cg transform='translate(-66.2578125%2c -24)' style='' class='label'%3e%3crect/%3e%3cforeignObject height='48' width='132.515625'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='nodeLabel'%3e%3cp%3ecomment%3cbr /%3e/** sums a and b */%3c/p%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3cg transform='translate(94.609375%2c 175)' id='flowchart-B-1' class='node default'%3e%3crect height='78' width='173.21875' y='-39' x='-86.609375' ry='5' rx='5' style='' class='basic label-container'/%3e%3cg transform='translate(-71.609375%2c -24)' style='' class='label'%3e%3crect/%3e%3cforeignObject height='48' width='143.21875'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='nodeLabel'%3e%3cp%3efunction_declaration%3cbr /%3efn%3c/p%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3cg transform='translate(94.609375%2c 303)' id='flowchart-D-3' class='node default'%3e%3crect height='78' width='108.265625' y='-39' x='-54.1328125' ry='5' rx='5' style='' class='basic label-container'/%3e%3cg transform='translate(-39.1328125%2c -24)' style='' class='label'%3e%3crect/%3e%3cforeignObject height='48' width='78.265625'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='nodeLabel'%3e%3cp%3estatement%3cbr /%3ereturn a - b%3c/p%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3c/g%3e%3c/g%3e%3c/g%3e%3c/svg%3e>)

So instead of fighting spaces and new lines, you can just locate the `function_declaration` node that follows a `comment` node.

docs.genai.mts

```ts
const node = sg.search("functions without comments");
```

Once you’ve located the node to update, you could do any transformation you want, e.g. replace it with another text. In terms of GenAI script, this means that you can build a prompt that includes as much context as you need, generate a response.

docs.genai.mts

```ts
$`Update the documentation of the function 'fn' to reflect the new behavior of the function.`;
fence(node.text());
```

```txt
/* subs a and b */
```

Once the LLM responds with the new comment, you could insert it as the new content of the node in the AST.

docs.genai.mts

```ts
edits.replace(node.comment(), response);
```

![](<data:image/svg+xml,%3csvg aria-roledescription='flowchart-v2' role='graphics-document document' viewBox='0 0 189.21875 350' style='max-width: 189.21875px%3b' class='flowchart' xmlns:xlink='http://www.w3.org/1999/xlink' xmlns='http://www.w3.org/2000/svg' width='100%25' id='mermaid-1'%3e%3cstyle%3e%23mermaid-1%7bfont-family:arial%2csans-serif%3bfont-size:16px%3bfill:%23333%3b%7d%40keyframes edge-animation-frame%7bfrom%7bstroke-dashoffset:0%3b%7d%7d%40keyframes dash%7bto%7bstroke-dashoffset:0%3b%7d%7d%23mermaid-1 .edge-animation-slow%7bstroke-dasharray:9%2c5!important%3bstroke-dashoffset:900%3banimation:dash 50s linear infinite%3bstroke-linecap:round%3b%7d%23mermaid-1 .edge-animation-fast%7bstroke-dasharray:9%2c5!important%3bstroke-dashoffset:900%3banimation:dash 20s linear infinite%3bstroke-linecap:round%3b%7d%23mermaid-1 .error-icon%7bfill:%23552222%3b%7d%23mermaid-1 .error-text%7bfill:%23552222%3bstroke:%23552222%3b%7d%23mermaid-1 .edge-thickness-normal%7bstroke-width:1px%3b%7d%23mermaid-1 .edge-thickness-thick%7bstroke-width:3.5px%3b%7d%23mermaid-1 .edge-pattern-solid%7bstroke-dasharray:0%3b%7d%23mermaid-1 .edge-thickness-invisible%7bstroke-width:0%3bfill:none%3b%7d%23mermaid-1 .edge-pattern-dashed%7bstroke-dasharray:3%3b%7d%23mermaid-1 .edge-pattern-dotted%7bstroke-dasharray:2%3b%7d%23mermaid-1 .marker%7bfill:%23333333%3bstroke:%23333333%3b%7d%23mermaid-1 .marker.cross%7bstroke:%23333333%3b%7d%23mermaid-1 svg%7bfont-family:arial%2csans-serif%3bfont-size:16px%3b%7d%23mermaid-1 p%7bmargin:0%3b%7d%23mermaid-1 .label%7bfont-family:arial%2csans-serif%3bcolor:%23333%3b%7d%23mermaid-1 .cluster-label text%7bfill:%23333%3b%7d%23mermaid-1 .cluster-label span%7bcolor:%23333%3b%7d%23mermaid-1 .cluster-label span p%7bbackground-color:transparent%3b%7d%23mermaid-1 .label text%2c%23mermaid-1 span%7bfill:%23333%3bcolor:%23333%3b%7d%23mermaid-1 .node rect%2c%23mermaid-1 .node circle%2c%23mermaid-1 .node ellipse%2c%23mermaid-1 .node polygon%2c%23mermaid-1 .node path%7bfill:%23ECECFF%3bstroke:%239370DB%3bstroke-width:1px%3b%7d%23mermaid-1 .rough-node .label text%2c%23mermaid-1 .node .label text%2c%23mermaid-1 .image-shape .label%2c%23mermaid-1 .icon-shape .label%7btext-anchor:middle%3b%7d%23mermaid-1 .node .katex path%7bfill:black%3bstroke:black%3bstroke-width:1px%3b%7d%23mermaid-1 .rough-node .label%2c%23mermaid-1 .node .label%2c%23mermaid-1 .image-shape .label%2c%23mermaid-1 .icon-shape .label%7btext-align:center%3b%7d%23mermaid-1 .node.clickable%7bcursor:pointer%3b%7d%23mermaid-1 .root .anchor path%7bfill:%23333333!important%3bstroke-width:0%3bstroke:%23333333%3b%7d%23mermaid-1 .arrowheadPath%7bfill:%23333333%3b%7d%23mermaid-1 .edgePath .path%7bstroke:%23333333%3bstroke-width:2.0px%3b%7d%23mermaid-1 .flowchart-link%7bstroke:%23333333%3bfill:none%3b%7d%23mermaid-1 .edgeLabel%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3btext-align:center%3b%7d%23mermaid-1 .edgeLabel p%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3b%7d%23mermaid-1 .edgeLabel rect%7bopacity:0.5%3bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3bfill:rgba(232%2c232%2c232%2c 0.8)%3b%7d%23mermaid-1 .labelBkg%7bbackground-color:rgba(232%2c 232%2c 232%2c 0.5)%3b%7d%23mermaid-1 .cluster rect%7bfill:%23ffffde%3bstroke:%23aaaa33%3bstroke-width:1px%3b%7d%23mermaid-1 .cluster text%7bfill:%23333%3b%7d%23mermaid-1 .cluster span%7bcolor:%23333%3b%7d%23mermaid-1 div.mermaidTooltip%7bposition:absolute%3btext-align:center%3bmax-width:200px%3bpadding:2px%3bfont-family:arial%2csans-serif%3bfont-size:12px%3bbackground:hsl(80%2c 100%25%2c 96.2745098039%25)%3bborder:1px solid %23aaaa33%3bborder-radius:2px%3bpointer-events:none%3bz-index:100%3b%7d%23mermaid-1 .flowchartTitleText%7btext-anchor:middle%3bfont-size:18px%3bfill:%23333%3b%7d%23mermaid-1 rect.text%7bfill:none%3bstroke-width:0%3b%7d%23mermaid-1 .icon-shape%2c%23mermaid-1 .image-shape%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3btext-align:center%3b%7d%23mermaid-1 .icon-shape p%2c%23mermaid-1 .image-shape p%7bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3bpadding:2px%3b%7d%23mermaid-1 .icon-shape rect%2c%23mermaid-1 .image-shape rect%7bopacity:0.5%3bbackground-color:rgba(232%2c232%2c232%2c 0.8)%3bfill:rgba(232%2c232%2c232%2c 0.8)%3b%7d%23mermaid-1 .label-icon%7bdisplay:inline-block%3bheight:1em%3boverflow:visible%3bvertical-align:-0.125em%3b%7d%23mermaid-1 .node .label-icon path%7bfill:currentColor%3bstroke:revert%3bstroke-width:revert%3b%7d%23mermaid-1 :root%7b--mermaid-font-family:arial%2csans-serif%3b%7d%3c/style%3e%3cg%3e%3cmarker orient='auto' markerHeight='8' markerWidth='8' markerUnits='userSpaceOnUse' refY='5' refX='5' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-1_flowchart-v2-pointEnd'%3e%3cpath style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 0 0 L 10 5 L 0 10 z'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='8' markerWidth='8' markerUnits='userSpaceOnUse' refY='5' refX='4.5' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-1_flowchart-v2-pointStart'%3e%3cpath style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 0 5 L 10 10 L 10 0 z'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5' refX='11' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-1_flowchart-v2-circleEnd'%3e%3ccircle style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' r='5' cy='5' cx='5'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5' refX='-1' viewBox='0 0 10 10' class='marker flowchart-v2' id='mermaid-1_flowchart-v2-circleStart'%3e%3ccircle style='stroke-width: 1%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' r='5' cy='5' cx='5'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5.2' refX='12' viewBox='0 0 11 11' class='marker cross flowchart-v2' id='mermaid-1_flowchart-v2-crossEnd'%3e%3cpath style='stroke-width: 2%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 1%2c1 l 9%2c9 M 10%2c1 l -9%2c9'/%3e%3c/marker%3e%3cmarker orient='auto' markerHeight='11' markerWidth='11' markerUnits='userSpaceOnUse' refY='5.2' refX='-1' viewBox='0 0 11 11' class='marker cross flowchart-v2' id='mermaid-1_flowchart-v2-crossStart'%3e%3cpath style='stroke-width: 2%3b stroke-dasharray: 1%2c 0%3b' class='arrowMarkerPath' d='M 1%2c1 l 9%2c9 M 10%2c1 l -9%2c9'/%3e%3c/marker%3e%3cg class='root'%3e%3cg class='clusters'/%3e%3cg class='edgePaths'%3e%3cpath marker-end='url(%23mermaid-1_flowchart-v2-pointEnd)' style='' class='edge-thickness-normal edge-pattern-solid edge-thickness-normal edge-pattern-solid flowchart-link' id='L_A_B_0' d='M94.609%2c86L94.609%2c90.167C94.609%2c94.333%2c94.609%2c102.667%2c94.609%2c110.333C94.609%2c118%2c94.609%2c125%2c94.609%2c128.5L94.609%2c132'/%3e%3cpath marker-end='url(%23mermaid-1_flowchart-v2-pointEnd)' style='' class='edge-thickness-normal edge-pattern-solid edge-thickness-normal edge-pattern-solid flowchart-link' id='L_B_D_0' d='M94.609%2c214L94.609%2c218.167C94.609%2c222.333%2c94.609%2c230.667%2c94.609%2c238.333C94.609%2c246%2c94.609%2c253%2c94.609%2c256.5L94.609%2c260'/%3e%3c/g%3e%3cg class='edgeLabels'%3e%3cg class='edgeLabel'%3e%3cg transform='translate(0%2c 0)' class='label'%3e%3cforeignObject height='0' width='0'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' class='labelBkg' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='edgeLabel'%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3cg class='edgeLabel'%3e%3cg transform='translate(0%2c 0)' class='label'%3e%3cforeignObject height='0' width='0'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' class='labelBkg' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='edgeLabel'%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3c/g%3e%3cg class='nodes'%3e%3cg transform='translate(94.609375%2c 47)' id='flowchart-A-0' class='node default'%3e%3crect height='78' width='158.09375' y='-39' x='-79.046875' ry='5' rx='5' style='' class='basic label-container'/%3e%3cg transform='translate(-64.046875%2c -24)' style='' class='label'%3e%3crect/%3e%3cforeignObject height='48' width='128.09375'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='nodeLabel'%3e%3cp%3ecomment%3cbr /%3e/** subs a and b */%3c/p%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3cg transform='translate(94.609375%2c 175)' id='flowchart-B-1' class='node default'%3e%3crect height='78' width='173.21875' y='-39' x='-86.609375' ry='5' rx='5' style='' class='basic label-container'/%3e%3cg transform='translate(-71.609375%2c -24)' style='' class='label'%3e%3crect/%3e%3cforeignObject height='48' width='143.21875'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='nodeLabel'%3e%3cp%3efunction_declaration%3cbr /%3efn%3c/p%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3cg transform='translate(94.609375%2c 303)' id='flowchart-D-3' class='node default'%3e%3crect height='78' width='108.265625' y='-39' x='-54.1328125' ry='5' rx='5' style='' class='basic label-container'/%3e%3cg transform='translate(-39.1328125%2c -24)' style='' class='label'%3e%3crect/%3e%3cforeignObject height='48' width='78.265625'%3e%3cdiv style='display: table-cell%3b white-space: nowrap%3b line-height: 1.5%3b max-width: 200px%3b text-align: center%3b' xmlns='http://www.w3.org/1999/xhtml'%3e%3cspan class='nodeLabel'%3e%3cp%3estatement%3cbr /%3ereturn a - b%3c/p%3e%3c/span%3e%3c/div%3e%3c/foreignObject%3e%3c/g%3e%3c/g%3e%3c/g%3e%3c/g%3e%3c/g%3e%3c/svg%3e>)

Voila! You’ve only touched the part of the file you wanted to update!

math.ts

```ts
/** subs a and b */
function fn(a: number, b: number): number {
  return a - b;
}
```

To recap, this strategy is based on the following steps:

1. **search** Use the AST to locate the node to update.
2. **transform and replace** Use the LLM to generate the new content of the node.
3. **commit** Update the node in the AST with the new content.

## ast-grep

[Section titled “ast-grep”](#ast-grep)

![AST-Grep Logo](https://ast-grep.github.io/logo.svg)

[ast-grep(sg)](https://ast-grep.github.io/) is a fast and polyglot tool for code structural search, lint, rewriting at large scale. **sg** provides us the AST-search-and-replace capabilities we need to implement the strategy above.

GenAIScript benefits from the excellent [Node.JS integration](https://ast-grep.github.io/reference/api.html#napi), which is available through the [`astGrep()`](https://microsoft.github.io/genaiscript/reference/scripts/ast-grep) method.

### Searching

[Section titled “Searching”](#searching)

The `sg.search` method allows you to search for nodes in the AST. It takes the language, file glob and pattern syntax and returns a list of matches.

docs.genai.mts

```ts
// search
const { matches, replace } = await sg.search(
  "ts",
  "src/*fib*.ts",
  {
    rule: {
      kind: "function_declaration",
      not: {
        precedes: {
          kind: "comment",
          stopBy: "neighbor",
        },
      },
    },
  },
);
```

### Editing

[Section titled “Editing”](#editing)

The `sg.changeset` method creates a changeset that can be used to apply a set of edits to a set of files.

```ts
// transform
const edits = sg.changeset();
for (const match of matches) {
  const { text } =
    await prompt`Generate new docs for ${match.text()}`;
  // replace
  edits.replace(match.comment(), text); // it's somewhat more involved
}
// commit all edits to file
await workspace.writeFiles(edits.commit());
```

## Sample: Doc generator / updater[]()

[Section titled “Sample: Doc generator / updater ”](#sample-doc-generator--updater)

You will find a full write down of the making of the documentation generator/updater script below in the [documentation](/genaiscript/reference/scripts/ast-grep). I encourage you to read it to dig deeper.

The `docs` scripts is a documentation generator/updater.

* uses ast-grep to find and generate missing documentation for exported TypeScript function. A second LLM-as-Judge request is used to check that the generated documentation is correct.
* if the `diff` option is selected, it will filter out functions that do not intersect with the diff (this is rather naive but a good start…).
* it can also be used to update the documentation of a function that has changed.
* it works regardless of the file size or the number of files as most transformations are hyper-localized.

generate and refresh my docs plz

```sh
genaiscript run docs -- --diff
```

Here are some example of applications of the scripts (one-shot, no human edit, multi-edit per file):

* [TypeScript codebase](https://github.com/pelikhan/TypeScript/pull/1), I stopped the script after a while, it was humming.
* [GenAIScript codebase](https://github.com/microsoft/genaiscript/pull/1376/files), this change also contains the update to GenAIScript as I was building the feature.

Here it goes:

docs.genai.mts

```ts
import { classify } from "@genaiscript/runtime";
import { astGrep } from "@genaiscript/plugin-ast-grep";


script({
  title: "Generate TypeScript function documentation using AST insertion",
  group: "dev",
  description: `
## Docs!


This script generates and updates TypeScript function using an AST/LLM hybrid approach.
It uses ast-grep to look for undocumented and documented functions,
then uses a combination of LLM, and LLM-as-a-judge to generate and validate the documentation.
It also uses prettier to format the code before and after the generation.


By default,


- no edits are applied on disk. It is recommended to
run this script with \`--vars 'applyEdits=true'\` to apply the edits.
- if a diff is available, it will only process the files with changes.


    `,
  accept: ".ts",
  files: "src/cowsay.ts",
  parameters: {
    diff: {
      type: "boolean",
      default: false,
      description: "If true, the script will only process files with changes with respect to main.",
    },
    pretty: {
      type: "boolean",
      default: false,
      description: "If true, the script will prettify the files before analysis.",
    },
    applyEdits: {
      type: "boolean",
      default: false,
      description: "If true, the script will modify the files.",
    },
    missing: {
      type: "boolean",
      default: true,
      description: "Generate missing docs.",
    },
    update: {
      type: "boolean",
      default: true,
      description: "Update existing docs.",
    },
    maxFiles: {
      type: "integer",
      description: "Maximum number of files to process.",
    },
  },
});
const { output, dbg, vars } = env;
let { files } = env;
const { applyEdits, diff, pretty, missing, update, maxFiles } = vars;


dbg({ applyEdits, diff, pretty, missing, update, maxFiles });


if (!missing && !update) cancel(`not generating or updating docs, exiting...`);


if (!applyEdits) output.warn(`edit not applied, use --vars 'applyEdits=true' to apply the edits`);


// filter by diff
const gitDiff = diff ? await git.diff({ base: "dev" }) : undefined;
dbg(`diff: %s`, gitDiff);
const diffFiles = gitDiff ? DIFF.parse(gitDiff) : undefined;
if (diff && !diffFiles?.length) cancel(`no diff files found, exiting...`);
if (diffFiles?.length) {
  dbg(`diff files: ${diffFiles.map((f) => f.to)}`);
  files = files.filter(({ filename }) =>
    diffFiles.some((f) => path.resolve(f.to) === path.resolve(filename)),
  );
  dbg(`diff filtered files: ${files.length}`);
}
if (!files.length) cancel(`no files to process, exiting...`);


if (maxFiles && files.length > maxFiles) {
  dbg(`random slicing files to ${maxFiles}`);
  files = parsers.tidyData(files, {
    sliceSample: maxFiles,
  }) as WorkspaceFile[];
}


const sg = await astGrep();
const stats = [];
for (const file of files) {
  console.debug(file.filename);
  // normalize spacing
  if (pretty) await prettier(file);


  // generate updated docs
  if (update) {
    stats.push({
      filename: file.filename,
      kind: "update",
      gen: 0,
      genCost: 0,
      judge: 0,
      judgeCost: 0,
      edits: 0,
      updated: 0,
    });
    await updateDocs(file, stats.at(-1));
  }


  // generate missing docs
  if (missing) {
    stats.push({
      filename: file.filename,
      kind: "new",
      gen: 0,
      genCost: 0,
      judge: 0,
      judgeCost: 0,
      edits: 0,
      updated: 0,
      nits: 0,
    });
    await generateDocs(file, stats.at(-1));
  }
}


if (stats.length)
  output.table(
    stats.filter((row) => Object.values(row).some((d) => typeof d === "number" && d > 0)),
  );


async function generateDocs(file: WorkspaceFile, fileStats: any) {
  const { matches: missingDocs } = await sg.search(
    "ts",
    file.filename,
    {
      rule: {
        kind: "export_statement",
        not: {
          follows: {
            kind: "comment",
            stopBy: "neighbor",
          },
        },
        has: {
          kind: "function_declaration",
        },
      },
    },
    { diff: gitDiff, applyGitIgnore: false },
  );
  dbg(`found ${missingDocs.length} missing docs`);
  const edits = sg.changeset();
  // for each match, generate a docstring for functions not documented
  for (const missingDoc of missingDocs) {
    const res = await runPrompt(
      (_) => {
        _.def("FILE", missingDoc.getRoot().root().text());
        _.def("FUNCTION", missingDoc.text());
        // this needs more eval-ing
        _.$`Generate a TypeScript function documentation for <FUNCTION>.
                - Make sure parameters are documented.
                - Be concise. Use technical tone.
                - do NOT include types, this is for TypeScript.
                - Use docstring syntax. do not wrap in markdown code section.


                The full source of the file is in <FILE> for reference.`;
      },
      {
        model: "large",
        responseType: "text",
        label: missingDoc.text()?.slice(0, 20) + "...",
      },
    );
    // if generation is successful, insert the docs
    fileStats.gen += res.usage?.total || 0;
    fileStats.genCost += res.usage?.cost || 0;
    if (res.error) {
      output.warn(res.error.message);
      continue;
    }
    const docs = docify(res.text.trim());


    // sanity check
    const judge = await classify(
      (_) => {
        _.def("FUNCTION", missingDoc.text());
        _.def("DOCS", docs);
      },
      {
        ok: "The content in <DOCS> is an accurate documentation for the code in <FUNCTION>.",
        err: "The content in <DOCS> does not match with the code in <FUNCTION>.",
      },
      {
        model: "small",
        responseType: "text",
        temperature: 0.2,
        systemSafety: false,
        system: ["system.technical", "system.typescript"],
      },
    );
    fileStats.judge += judge.usage?.total || 0;
    fileStats.judgeCost += judge.usage?.cost || 0;
    if (judge.label !== "ok") {
      output.warn(judge.label);
      output.fence(judge.answer);
      continue;
    }
    const updated = `${docs}\n${missingDoc.text()}`;
    edits.replace(missingDoc, updated);
    fileStats.edits++;
    fileStats.nits++;
  }


  // apply all edits and write to the file
  const modifiedFiles = edits.commit();
  if (!modifiedFiles?.length) {
    dbg("no edits to apply");
    return;
  }
  fileStats.updated = 1;
  if (applyEdits) {
    await workspace.writeFiles(modifiedFiles);
    await prettier(file);
  } else {
    output.diff(file, modifiedFiles[0]);
  }
}


async function updateDocs(file: WorkspaceFile, fileStats: any) {
  const { matches } = await sg.search(
    "ts",
    file.filename,
    YAML`
rule:
  kind: "export_statement"
  follows:
    kind: "comment"
    stopBy: neighbor
  has:
      kind: "function_declaration"
`,
    { diff: gitDiff, applyGitIgnore: false },
  );
  dbg(`found ${matches.length} docs to update`);
  const edits = sg.changeset();
  // for each match, generate a docstring for functions not documented
  for (const match of matches) {
    const comment = match.prev();


    const res = await runPrompt(
      (_) => {
        _.def("FILE", match.getRoot().root().text(), { flex: 1 });
        _.def("DOCSTRING", comment.text(), { flex: 10 });
        _.def("FUNCTION", match.text(), { flex: 10 });
        // this needs more eval-ing
        _.$`Update the TypeScript docstring <DOCSTRING> to match the code in function <FUNCTION>.
                - If the docstring is up to date, return /NO/. It's ok to leave it as is.
                - do not rephrase an existing sentence if it is correct.
                - Make sure parameters are documented.
                - do NOT include types, this is for TypeScript.
                - Use docstring syntax. do not wrap in markdown code section.
                - Minimize updates to the existing docstring.


                The full source of the file is in <FILE> for reference.
                The source of the function is in <FUNCTION>.
                The current docstring is <DOCSTRING>.


                docstring:


                /**
                 * description
                 * @param param1 - description
                 * @param param2 - description
                 * @returns description
                 */
                `;
      },
      {
        model: "large",
        responseType: "text",
        flexTokens: 12000,
        label: match.text()?.slice(0, 20) + "...",
        temperature: 0.2,
        systemSafety: false,
        system: ["system.technical", "system.typescript"],
      },
    );
    fileStats.gen += res.usage?.total || 0;
    fileStats.genCost += res.usage?.cost || 0;
    // if generation is successful, insert the docs
    if (res.error) {
      output.warn(res.error.message);
      continue;
    }


    if (res.text.includes("/NO/")) continue;


    const docs = docify(res.text.trim());


    // ask LLM if change is worth it
    const judge = await classify(
      (_) => {
        _.def("FUNCTION", match.text());
        _.def("ORIGINAL_DOCS", comment.text());
        _.def("NEW_DOCS", docs);
        _.$`An LLM generated an updated docstring <NEW_DOCS> for function <FUNCTION>. The original docstring is <ORIGINAL_DOCS>.`;
      },
      {
        APPLY: "The <NEW_DOCS> is a significant improvement to <ORIGINAL_DOCS>.",
        NIT: "The <NEW_DOCS> contains nitpicks (minor adjustments) to <ORIGINAL_DOCS>.",
      },
      {
        model: "large",
        responseType: "text",
        temperature: 0.2,
        systemSafety: false,
        system: ["system.technical", "system.typescript"],
      },
    );


    fileStats.judge += judge.usage?.total || 0;
    fileStats.judgeCost += judge.usage?.cost || 0;
    if (judge.label === "NIT") {
      output.warn("LLM suggests minor adjustments, skipping");
      continue;
    }
    edits.replace(comment, docs);
    fileStats.edits++;
  }


  // apply all edits and write to the file
  const modifiedFiles = edits.commit();
  if (!modifiedFiles?.length) {
    dbg("no edits to apply");
    return;
  }
  fileStats.updated = 1;
  if (applyEdits) {
    await workspace.writeFiles(modifiedFiles);
    await prettier(file);
  } else {
    output.diff(file, modifiedFiles[0]);
  }
}


function docify(docs: string) {
  docs = parsers.unfence(docs, "*");
  if (!/^\/\*\*.*.*\*\/$/s.test(docs)) docs = `/**\n* ${docs.split(/\r?\n/g).join("\n* ")}\n*/`;
  return docs.replace(/\n+$/, "");
}


async function prettier(file: WorkspaceFile, options?: { curly?: boolean }) {
  dbg(file.filename);
  const args = ["--write"];
  if (options?.curly) args.push("--plugin=prettier-plugin-curly");
  // format
  const res = await host.exec("prettier", [...args, file.filename]);
  if (res.exitCode) {
    dbg(`error: %d\n%s`, res.exitCode, res.stderr);
    throw new Error(`${res.stdout} (${res.exitCode})`);
  }
}
```

* [original source](https://github.com/microsoft/genaiscript/blob/main/genaisrc/docs.genai.mts)