annotate method

TreeSitterNode annotate(TreeSitterNode root, String source)

Implementation

TreeSitterNode annotate(TreeSitterNode root, String source) {
  final captures = query.captures(root, source: source);
  final parents = <TreeSitterNode, TreeSitterNode?>{};
  void index(TreeSitterNode node, TreeSitterNode? parent) {
    parents[node] = parent;
    for (final child in node.children) {
      index(child, node);
    }
  }

  index(root, null);
  final scopes = <TreeSitterNode, bool>{root: scopeInherits};
  final definitions = <TreeSitterNode>[];
  final references = <TreeSitterNode>[];
  final ignored = {
    for (final capture in captures)
      if (capture.name == 'ignore') capture.node,
  };
  for (final capture in captures) {
    if (ignored.contains(capture.node)) continue;
    if (capture.name == 'local.scope') {
      scopes[capture.node] =
          capture.setProperties['local.scope-inherits'] != 'false';
    } else if (capture.name.startsWith('local.definition') &&
        capture.name != 'local.definition-value') {
      definitions.add(capture.node);
    } else if (capture.name.startsWith('local.reference')) {
      references.add(capture.node);
    }
  }

  TreeSitterNode enclosingScope(TreeSitterNode node) {
    var cursor = parents[node];
    while (cursor != null) {
      if (scopes.containsKey(cursor)) return cursor;
      cursor = parents[cursor];
    }
    return root;
  }

  String? textOf(TreeSitterNode node) {
    final start = node.startOffset.value;
    final end = node.endOffset.value;
    if (start < 0 || end < start || end > source.length) return null;
    return source.substring(start, end);
  }

  final definitionsByScope =
      <TreeSitterNode, Map<String, List<TreeSitterNode>>>{};
  final localNodes = <TreeSitterNode>{};
  for (final definition in definitions) {
    final text = textOf(definition);
    if (text == null) continue;
    final scope = enclosingScope(definition);
    definitionsByScope
        .putIfAbsent(scope, () => <String, List<TreeSitterNode>>{})
        .putIfAbsent(text, () => <TreeSitterNode>[])
        .add(definition);
    localNodes.add(definition);
  }

  bool resolves(TreeSitterNode reference, String text) {
    var scope = enclosingScope(reference);
    while (true) {
      final candidates = definitionsByScope[scope]?[text];
      if (candidates != null &&
          candidates.any(
            (definition) => definition.startByte <= reference.startByte,
          )) {
        return true;
      }
      if (!(scopes[scope] ?? scopeInherits) || identical(scope, root)) {
        return false;
      }
      final parent = parents[scope];
      if (parent == null) return false;
      scope = enclosingScope(scope);
    }
  }

  for (final reference in references) {
    final text = textOf(reference);
    if (text != null && resolves(reference, text)) localNodes.add(reference);
  }

  TreeSitterNode clone(TreeSitterNode node) => TreeSitterNode(
    kind: node.kind,
    startByte: node.startByte,
    endByte: node.endByte,
    startOffset: node.startOffset,
    endOffset: node.endOffset,
    children: [for (final child in node.children) clone(child)],
    isNamed: node.isNamed,
    isMissing: node.isMissing,
    isExtra: node.isExtra,
    isError: node.isError,
    symbolId: node.symbolId,
    grammarSymbolId: node.grammarId,
    grammarKind: node.grammarName,
    language: node.language,
    parseState: node.parseState,
    nodeIdentity: node._nodeIdentity ?? node,
    paddingBytes: node.paddingBytes,
    hasChanges: node.hasChanges,
    fieldNames: node.fieldNames,
    supertypeKinds: node.supertypeKinds,
    properties: {...node.properties, if (localNodes.contains(node)) 'local'},
    querySiblingGroup: node._querySiblingGroup,
  );

  return clone(root);
}