shrimp/vscode-extension/server/src/semanticTokens.ts

import { parser } from '../../../src/parser/shrimp'
import * as Terms from '../../../src/parser/shrimp.terms'
import { SyntaxNode } from '@lezer/common'
import { TextDocument } from 'vscode-languageserver-textdocument'
import {
  SemanticTokensBuilder,
  SemanticTokenTypes,
  SemanticTokenModifiers,
} from 'vscode-languageserver/node'
import { ScopeTracker } from './scopeTracker'

export const TOKEN_TYPES = [
  SemanticTokenTypes.function,
  SemanticTokenTypes.variable,
  SemanticTokenTypes.string,
  SemanticTokenTypes.number,
  SemanticTokenTypes.operator,
  SemanticTokenTypes.keyword,
  SemanticTokenTypes.parameter,
  SemanticTokenTypes.property,
  SemanticTokenTypes.regexp,
  SemanticTokenTypes.comment,
]

export const TOKEN_MODIFIERS = [
  SemanticTokenModifiers.declaration,
  SemanticTokenModifiers.modification,
  SemanticTokenModifiers.readonly,
]

export function buildSemanticTokens(document: TextDocument): number[] {
  const text = document.getText()
  const tree = parser.parse(text)
  const builder = new SemanticTokensBuilder()
  const scopeTracker = new ScopeTracker(document)

  walkTree(tree.topNode, document, builder, scopeTracker)

  return builder.build().data
}

// Emit split tokens for NamedArgPrefix (e.g., "color=" → "color" + "=")
function emitNamedArgPrefix(
  node: SyntaxNode,
  document: TextDocument,
  builder: SemanticTokensBuilder
) {
  const text = document.getText({
    start: document.positionAt(node.from),
    end: document.positionAt(node.to),
  })

  const nameLength = text.length - 1  // Everything except the =
  const start = document.positionAt(node.from)

  // Emit token for the name part (e.g., "color")
  builder.push(
    start.line,
    start.character,
    nameLength,
    TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
    0
  )

  // Emit token for the "=" part
  builder.push(
    start.line,
    start.character + nameLength,
    1,  // Just the = character
    TOKEN_TYPES.indexOf(SemanticTokenTypes.operator),
    0
  )
}

// Walk the tree and collect tokens
function walkTree(
  node: SyntaxNode,
  document: TextDocument,
  builder: SemanticTokensBuilder,
  scopeTracker: ScopeTracker
) {
  // Special handling for NamedArgPrefix to split "name=" into two tokens
  if (node.type.id === Terms.NamedArgPrefix) {
    emitNamedArgPrefix(node, document, builder)
  } else {
    const tokenInfo = getTokenType(node, document, scopeTracker)

    if (tokenInfo !== undefined) {
      const start = document.positionAt(node.from)
      const length = node.to - node.from
      builder.push(start.line, start.character, length, tokenInfo.type, tokenInfo.modifiers)
    }
  }

  let child = node.firstChild
  while (child) {
    walkTree(child, document, builder, scopeTracker)
    child = child.nextSibling
  }
}

// Map Lezer node IDs to semantic token type indices and modifiers
type TokenInfo = { type: number; modifiers: number } | undefined
function getTokenType(
  node: SyntaxNode,
  document: TextDocument,
  scopeTracker: ScopeTracker
): TokenInfo {
  const nodeTypeId = node.type.id
  const parentTypeId = node.parent?.type.id

  // Special case for now, eventually keywords will go away
  if (node.type.name === 'keyword') {
    return {
      type: TOKEN_TYPES.indexOf(SemanticTokenTypes.keyword),
      modifiers: 0,
    }
  }

  switch (nodeTypeId) {
    case Terms.Identifier:
      // Check parent to determine context
      if (parentTypeId === Terms.FunctionCall) {
        return {
          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
          modifiers: 0,
        }
      }
      if (parentTypeId === Terms.FunctionDef) {
        return {
          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
          modifiers: getModifierBits(SemanticTokenModifiers.declaration),
        }
      }
      if (parentTypeId === Terms.FunctionCallOrIdentifier) {
        return {
          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
          modifiers: 0,
        }
      }
      if (parentTypeId === Terms.Params) {
        return {
          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.parameter),
          modifiers: 0,
        }
      }
      if (parentTypeId === Terms.DotGet) {
        return {
          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
          modifiers: 0,
        }
      }

      // Special case: Identifier in PositionalArg or NamedArg- check scope
      if (parentTypeId === Terms.PositionalArg || parentTypeId === Terms.NamedArg) {
        const identifierText = document.getText({
          start: document.positionAt(node.from),
          end: document.positionAt(node.to),
        })

        // If not in scope, treat as string (like a Word)
        if (!scopeTracker.isInScope(identifierText, node)) {
          return {
            type: TOKEN_TYPES.indexOf(SemanticTokenTypes.string),
            modifiers: 0,
          }
        }
        // If in scope, fall through to treat as variable
      }

      // Otherwise it's a regular variable
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
        modifiers: 0,
      }

    case Terms.IdentifierBeforeDot:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
        modifiers: 0,
      }

    case Terms.AssignableIdentifier:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
        modifiers: getModifierBits(SemanticTokenModifiers.modification),
      }

    case Terms.String:
    case Terms.StringFragment:
    case Terms.Word:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.string),
        modifiers: 0,
      }

    case Terms.Number:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.number),
        modifiers: 0,
      }

    case Terms.Plus:
    case Terms.Minus:
    case Terms.Star:
    case Terms.Slash:
    case Terms.Eq:
    case Terms.EqEq:
    case Terms.Neq:
    case Terms.Lt:
    case Terms.Lte:
    case Terms.Gt:
    case Terms.Gte:
    case Terms.Modulo:
    case Terms.And:
    case Terms.Or:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.operator),
        modifiers: 0,
      }

    case Terms.Do:
    case Terms.colon:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.keyword),
        modifiers: 0,
      }

    case Terms.Regex:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.regexp),
        modifiers: 0,
      }

    case Terms.Comment:
      return {
        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.comment),
        modifiers: 0,
      }

    default:
      return undefined
  }
}

const getModifierBits = (...modifiers: SemanticTokenModifiers[]): number => {
  let bits = 0
  for (const modifier of modifiers) {
    const index = TOKEN_MODIFIERS.indexOf(modifier)
    if (index !== -1) bits |= 1 << index
  }
  return bits
}