shrimp/src/parser/node.ts

import { type Token, TokenType } from './tokenizer2'

export type NodeType =
  | 'Program'
  | 'Block'
  | 'FunctionCall'
  | 'FunctionCallOrIdentifier'
  | 'FunctionCallWithBlock'
  | 'PositionalArg'
  | 'NamedArg'
  | 'NamedArgPrefix'
  | 'FunctionDef'
  | 'Params'
  | 'NamedParam'
  | 'Null'
  | 'Boolean'
  | 'Number'
  | 'String'
  | 'StringFragment'
  | 'CurlyString'
  | 'DoubleQuote'
  | 'EscapeSeq'
  | 'Interpolation'
  | 'Regex'
  | 'Identifier'
  | 'AssignableIdentifier'
  | 'IdentifierBeforeDot'
  | 'Word'
  | 'Array'
  | 'Dict'
  | 'Comment'
  | 'BinOp'
  | 'ConditionalOp'
  | 'ParenExpr'
  | 'Assign'
  | 'CompoundAssign'
  | 'DotGet'
  | 'PipeExpr'
  | 'IfExpr'
  | 'ElseIfExpr'
  | 'ElseExpr'
  | 'WhileExpr'
  | 'TryExpr'
  | 'CatchExpr'
  | 'FinallyExpr'
  | 'Throw'
  | 'Not'
  | 'Eq'
  | 'Modulo'
  | 'Plus'
  | 'Star'
  | 'Slash'
  | 'Import'
  | 'Do'
  | 'Underscore'
  | 'colon'
  | 'keyword'
  | 'operator'

// TODO: remove this when we switch from lezer
export const operators: Record<string, any> = {
  // Logic
  and: 'And',
  or: 'Or',

  // Bitwise
  band: 'Band',
  bor: 'Bor',
  bxor: 'Bxor',
  '>>>': 'Ushr',
  '>>': 'Shr',
  '<<': 'Shl',

  // Comparison
  '>=': 'Gte',
  '<=': 'Lte',
  '>': 'Gt',
  '<': 'Lt',
  '!=': 'Neq',
  '==': 'EqEq',

  // Compound assignment operators
  '??=': 'NullishEq',
  '+=': 'PlusEq',
  '-=': 'MinusEq',
  '*=': 'StarEq',
  '/=': 'SlashEq',
  '%=': 'ModuloEq',

  // Nullish coalescing
  '??': 'NullishCoalesce',

  // Math
  '*': 'Star',
  '**': 'StarStar',
  '=': 'Eq',
  '/': 'Slash',
  '+': 'Plus',
  '-': 'Minus',
  '%': 'Modulo',

  // Dotget
  '.': 'Dot',

  // Pipe
  '|': 'operator',
}

export class Tree {
  constructor(public topNode: SyntaxNode) {}

  get length(): number {
    return this.topNode.to
  }

  cursor() {
    return {
      type: this.topNode.type,
      from: this.topNode.from,
      to: this.topNode.to,
      node: this.topNode,
    }
  }

  iterate(options: { enter: (node: SyntaxNode) => void }) {
    const iter = (node: SyntaxNode) => {
      for (const n of node.children) iter(n)
      options.enter(node)
    }

    iter(this.topNode)
  }
}

export class SyntaxNode {
  #type: NodeType
  #isError = false
  from: number
  to: number
  parent: SyntaxNode | null
  children: SyntaxNode[] = []

  constructor(type: NodeType, from: number, to: number, parent: SyntaxNode | null = null) {
    this.#type = type
    this.from = from
    this.to = to
    this.parent = parent
  }

  static from(token: Token, parent?: SyntaxNode): SyntaxNode {
    return new SyntaxNode(TokenType[token.type] as NodeType, token.from, token.to, parent ?? null)
  }

  get type(): {
    type: NodeType
    name: NodeType
    isError: boolean
    is: (other: NodeType) => boolean
  } {
    return {
      type: this.#type,
      name: this.#type,
      isError: this.#isError,
      is: (other: NodeType) => other === this.#type,
    }
  }

  set type(name: NodeType) {
    this.#type = name
  }

  get name(): string {
    return this.type.name
  }

  get isError(): boolean {
    return this.#isError
  }

  set isError(err: boolean) {
    this.#isError = err
  }

  get firstChild(): SyntaxNode | null {
    return this.children[0] ?? null
  }

  get lastChild(): SyntaxNode | null {
    return this.children.at(-1) ?? null
  }

  get nextSibling(): SyntaxNode | null {
    if (!this.parent) return null
    const siblings = this.parent.children
    const index = siblings.indexOf(this)
    return index >= 0 && index < siblings.length - 1 ? siblings[index + 1]! : null
  }

  get prevSibling(): SyntaxNode | null {
    if (!this.parent) return null
    const siblings = this.parent.children
    const index = siblings.indexOf(this)
    return index > 0 ? siblings[index - 1]! : null
  }

  add(node: SyntaxNode) {
    node.parent = this
    this.children.push(node)
  }

  push(...nodes: SyntaxNode[]): SyntaxNode {
    nodes.forEach((child) => (child.parent = this))
    this.children.push(...nodes)
    return this
  }

  toString(): string {
    return this.type.name
  }
}

// Operator precedence (binding power) - higher = tighter binding
export const precedence: Record<string, number> = {
  // Logical
  or: 10,
  and: 20,

  // Comparison
  '==': 30,
  '!=': 30,
  '<': 30,
  '>': 30,
  '<=': 30,
  '>=': 30,

  // Nullish coalescing
  '??': 35,

  // Bitwise shifts (lower precedence than addition)
  '<<': 37,
  '>>': 37,
  '>>>': 37,

  // Addition/Subtraction
  '+': 40,
  '-': 40,

  // Bitwise AND/OR/XOR (higher precedence than addition)
  band: 45,
  bor: 45,
  bxor: 45,

  // Multiplication/Division/Modulo
  '*': 50,
  '/': 50,
  '%': 50,

  // Exponentiation (right-associative)
  '**': 60,
}

export const conditionals = new Set(['==', '!=', '<', '>', '<=', '>=', '??', 'and', 'or'])

export const compounds = ['??=', '+=', '-=', '*=', '/=', '%=']