front_end/ui/components/text_editor/AiCodeGenerationParser.ts

// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

import * as CodeMirror from '../../../third_party/codemirror.next/codemirror.next.js';

const LINE_COMMENT_PATTERN = /^(?:\/\/|#)\s*/gm;
const BLOCK_COMMENT_START_PATTERN = /^\/\*+\s*/;
const BLOCK_COMMENT_END_PATTERN = /\s*\*+\/$/;
const BLOCK_COMMENT_LINE_PREFIX_PATTERN = /^\s*\*\s?/;

export interface CommentNodeInfo {
  text: string;
  to: number;
}

function findLastNonWhitespacePos(state: CodeMirror.EditorState, cursorPosition: number): number {
  const line = state.doc.lineAt(cursorPosition);
  const textBefore = line.text.substring(0, cursorPosition - line.from);
  const effectiveEnd = line.from + textBefore.trimEnd().length;
  return effectiveEnd;
}

function resolveCommentNode(state: CodeMirror.EditorState, cursorPosition: number): CodeMirror.SyntaxNode|undefined {
  const tree = CodeMirror.syntaxTree(state);
  const lookupPos = findLastNonWhitespacePos(state, cursorPosition);
  // Find the innermost syntax node at the last non-whitespace character position.
  // The bias of -1 makes it check the character to the left of the position.
  const node = tree.resolveInner(lookupPos, -1);
  const nodeType = node.type.name;
  // Check if the node type is a comment
  if (nodeType.includes('Comment')) {
    if (!nodeType.includes('BlockComment')) {
      return node;
    }
    // An unclosed block comment can result in the parser inserting an error.
    let hasInternalError = false;
    tree.iterate({
      from: node.from,
      to: node.to,
      enter: n => {
        if (n.type.isError) {
          hasInternalError = true;
          return false;
        }
        return true;
      },
    });
    return hasInternalError ? undefined : node;
  }
  return;
}

function extractBlockCommentText(rawText: string): string|undefined {
  // Remove /* and */, whitespace, and common leading asterisks on new lines
  if (!rawText.match(BLOCK_COMMENT_START_PATTERN)) {
    return;
  }
  let cleaned = rawText.replace(BLOCK_COMMENT_START_PATTERN, '');
  if (!cleaned.match(BLOCK_COMMENT_END_PATTERN)) {
    return;
  }
  cleaned = cleaned.replace(BLOCK_COMMENT_END_PATTERN, '');
  // Remove leading " * " from multi-line block comments
  cleaned = cleaned.split('\n').map(line => line.replace(BLOCK_COMMENT_LINE_PREFIX_PATTERN, '')).join('\n').trim();
  return cleaned;
}

function extractLineComment(node: CodeMirror.SyntaxNode, state: CodeMirror.EditorState): CommentNodeInfo|undefined {
  let firstNode = node;
  let lastNode = node;

  let prev = node.prevSibling;
  while (prev?.type.name.includes('LineComment')) {
    firstNode = prev;
    prev = prev.prevSibling;
  }

  let next = node.nextSibling;
  while (next?.type.name.includes('LineComment')) {
    lastNode = next;
    next = next.nextSibling;
  }

  // Extract all lines between the first and last identified node
  const fullRawText = state.doc.sliceString(firstNode.from, lastNode.to);

  // Process each line to remove prefixes (// or #)
  const concatenatedText = fullRawText.replaceAll(LINE_COMMENT_PATTERN, '').replace(/\n\s*\n/g, '\n').trim();

  return concatenatedText ? {text: concatenatedText, to: lastNode.to} : undefined;
}

export class AiCodeGenerationParser {
  static extractCommentNodeInfo(state: CodeMirror.EditorState, cursorPosition: number): CommentNodeInfo|undefined {
    const node = resolveCommentNode(state, cursorPosition);
    if (!node) {
      return;
    }
    const nodeType = node.type.name;
    const rawText = state.doc.sliceString(node.from, node.to);
    let text = '';
    if (nodeType.includes('LineComment')) {
      return extractLineComment(node, state);
    }
    if (nodeType.includes('BlockComment')) {
      text = extractBlockCommentText(rawText) ?? '';
    } else {
      text = rawText;
    }

    if (!Boolean(text)) {
      return;
    }

    return {text, to: node.to};
  }
}