Print Lezer Trees

print-lezer-tree.ts

// MIT License
// 
// Copyright (c) 2021 Matthijs Steen
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

import { Text } from "@codemirror/text"
import { Input, NodeType, SyntaxNode, Tree, TreeCursor } from "@lezer/common"

class StringInput implements Input {
  constructor(private readonly input: string) {}

  get length() {
    return this.input.length
  }

  chunk(from: number): string {
    return this.input.slice(from)
  }

  lineChunks = false

  read(from: number, to: number): string {
    return this.input.slice(from, to)
  }
}

export function sliceType(cursor: TreeCursor, input: Input, type: number): string | null {
  if (cursor.type.id === type) {
    const s = input.read(cursor.from, cursor.to)
    cursor.nextSibling()
    return s
  }
  return null
}

export function isType(cursor: TreeCursor, type: number): boolean {
  const cond = cursor.type.id === type
  if (cond) cursor.nextSibling()
  return cond
}

export type CursorNode = { type: NodeType; from: number; to: number; isLeaf: boolean }

function cursorNode({ type, from, to }: TreeCursor, isLeaf = false): CursorNode {
  return { type, from, to, isLeaf }
}

export type TreeTraversal = {
  beforeEnter?: (cursor: TreeCursor) => void
  onEnter: (node: CursorNode) => false | void
  onLeave?: (node: CursorNode) => false | void
}

type TreeTraversalOptions = {
  from?: number
  to?: number
  includeParents?: boolean
} & TreeTraversal

export function traverseTree(
  cursor: TreeCursor | Tree | SyntaxNode,
  {
    from = -Infinity,
    to = Infinity,
    includeParents = false,
    beforeEnter,
    onEnter,
    onLeave,
  }: TreeTraversalOptions,
): void {
  if (!(cursor instanceof TreeCursor)) cursor = cursor instanceof Tree ? cursor.cursor() : cursor.cursor
  for (;;) {
    let node = cursorNode(cursor)
    let leave = false
    if (node.from <= to && node.to >= from) {
      const enter = !node.type.isAnonymous && (includeParents || (node.from >= from && node.to <= to))
      if (enter && beforeEnter) beforeEnter(cursor)
      node.isLeaf = !cursor.firstChild()
      if (enter) {
        leave = true
        if (onEnter(node) === false) return
      }
      if (!node.isLeaf) continue
    }
    for (;;) {
      node = cursorNode(cursor, node.isLeaf)
      if (leave && onLeave) if (onLeave(node) === false) return
      leave = cursor.type.isAnonymous
      node.isLeaf = false
      if (cursor.nextSibling()) break
      if (!cursor.parent()) return
      leave = true
    }
  }
}

function isChildOf(child: CursorNode, parent: CursorNode): boolean {
  return (
    child.from >= parent.from && child.from <= parent.to && child.to <= parent.to && child.to >= parent.from
  )
}

export function validatorTraversal(
  input: Input | string,
  { fullMatch = true }: { fullMatch?: boolean } = {},
) {
  if (typeof input === "string") input = new StringInput(input)
  const state = {
    valid: true,
    parentNodes: [] as CursorNode[],
    lastLeafTo: 0,
  }
  return {
    state,
    traversal: {
      onEnter(node) {
        state.valid = true
        if (!node.isLeaf) state.parentNodes.unshift(node)
        if (node.from > node.to || node.from < state.lastLeafTo) {
          state.valid = false
        } else if (node.isLeaf) {
          if (state.parentNodes.length && !isChildOf(node, state.parentNodes[0])) state.valid = false
          state.lastLeafTo = node.to
        } else {
          if (state.parentNodes.length) {
            if (!isChildOf(node, state.parentNodes[0])) state.valid = false
          } else if (fullMatch && (node.from !== 0 || node.to !== input.length)) {
            state.valid = false
          }
        }
      },
      onLeave(node) {
        if (!node.isLeaf) state.parentNodes.shift()
      },
    } as TreeTraversal,
  }
}

export function validateTree(
  tree: TreeCursor | Tree | SyntaxNode,
  input: Input | string,
  options?: { fullMatch?: boolean },
): boolean {
  const { state, traversal } = validatorTraversal(input, options)
  traverseTree(tree, traversal)
  return state.valid
}

enum Color {
  Red = 31,
  Green = 32,
  Yellow = 33,
}

function colorize(value: any, color: number): string {
  return "\u001b[" + color + "m" + String(value) + "\u001b[39m"
}

type PrintTreeOptions = { from?: number; to?: number; start?: number; includeParents?: boolean }

export function printTree(
  cursor: TreeCursor | Tree | SyntaxNode,
  input: Input | string,
  { from, to, start = 0, includeParents }: PrintTreeOptions = {},
): string {
  const inp = typeof input === "string" ? new StringInput(input) : input
  const text = Text.of(inp.read(0, inp.length).split("\n"))
  const state = {
    output: "",
    prefixes: [] as string[],
    hasNextSibling: false,
  }
  const validator = validatorTraversal(inp)
  traverseTree(cursor, {
    from,
    to,
    includeParents,
    beforeEnter(cursor) {
      state.hasNextSibling = cursor.nextSibling() && cursor.prevSibling()
    },
    onEnter(node) {
      validator.traversal.onEnter(node)
      const isTop = state.output === ""
      const hasPrefix = !isTop || node.from > 0
      if (hasPrefix) {
        state.output += (!isTop ? "\n" : "") + state.prefixes.join("")
        if (state.hasNextSibling) {
          state.output += " ├─ "
          state.prefixes.push(" │  ")
        } else {
          state.output += " └─ "
          state.prefixes.push("    ")
        }
      }
      const hasRange = node.from !== node.to
      state.output +=
        (node.type.isError || !validator.state.valid ? colorize(node.type.name, Color.Red) : node.type.name) +
        " " +
        (hasRange
          ? "[" +
            colorize(locAt(text, start + node.from), Color.Yellow) +
            ".." +
            colorize(locAt(text, start + node.to), Color.Yellow) +
            "]"
          : colorize(locAt(text, start + node.from), Color.Yellow))
      if (hasRange && node.isLeaf) {
        state.output += ": " + colorize(JSON.stringify(inp.read(node.from, node.to)), Color.Green)
      }
    },
    onLeave(node) {
      validator.traversal.onLeave!(node)
      state.prefixes.pop()
    },
  })
  return state.output
}

function locAt(text: Text, pos: number): string {
  const line = text.lineAt(pos)
  return line.number + ":" + (pos - line.from)
}

export function logTree(
  tree: TreeCursor | Tree | SyntaxNode,
  input: string,
  options?: PrintTreeOptions,
): void {
  console.log(printTree(tree, input, options))
}

Thanks for this! I managed to get this working with:

logTree(tree, cm.state.doc.toString());

Is this how this was intended to be used?

I made this for Lezer specifically, so I haven't looked into how to use it with CodeMirror, but what you have seems alright to me. The Lezer tree being the first argument, and the input string that produced the tree, being the second argument.

@msteen, I packaged your code in https://github.com/lezer-parser-unofficial/printer. Do you mind?

@deemp I don't mind you packaging it at all, I am glad to see it be of use to others, if anything. When I was making ANTLR parsers I made a similar tool in Java at the time (not public), and found it a simple but must have tool for debugging, which is why I remade it while trying to make a Lezer parser.

Yes, thanks for this! This pretty printer was very helpful.

Great utility!
Found this as well while starting working on a grammar: https://lezer-playground.vercel.app
Hope that can help others!