<commit_before>package com.intellij.psi.impl.source.tree.injected;

import com.intellij.lang.ASTNode;
import com.intellij.lang.Language;
import com.intellij.lang.LanguageDialect;
import com.intellij.lang.ParserDefinition;
import com.intellij.lexer.Lexer;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.editor.RangeMarker;
import com.intellij.openapi.editor.SelectionModel;
import com.intellij.openapi.editor.ex.DocumentEx;
import com.intellij.openapi.editor.impl.DocumentImpl;
import com.intellij.openapi.editor.impl.EditorImpl;
import com.intellij.openapi.editor.impl.injected.DocumentRange;
import com.intellij.openapi.editor.impl.injected.EditorDelegate;
import com.intellij.openapi.editor.impl.injected.VirtualFileDelegate;
import com.intellij.openapi.extensions.Extensions;
import com.intellij.openapi.fileEditor.FileEditorManager;
import com.intellij.openapi.fileEditor.OpenFileDescriptor;
import com.intellij.openapi.fileEditor.impl.FileDocumentManagerImpl;
import com.intellij.openapi.fileTypes.SyntaxHighlighter;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.*;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.*;
import com.intellij.psi.impl.CachedValueImpl;
import com.intellij.psi.impl.PsiDocumentManagerImpl;
import com.intellij.psi.impl.PsiManagerEx;
import com.intellij.psi.impl.PsiManagerImpl;
import com.intellij.psi.impl.source.PsiFileImpl;
import com.intellij.psi.impl.source.resolve.ResolveUtil;
import com.intellij.psi.impl.source.tree.*;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.util.CachedValue;
import com.intellij.psi.util.CachedValueProvider;
import com.intellij.psi.util.PsiModificationTracker;
import com.intellij.util.SmartList;
import gnu.trove.THashMap;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @author cdr
 */
public class InjectedLanguageUtil {
  private static final Logger LOG = Logger.getInstance("com.intellij.psi.impl.source.tree.injected.InjectedLanguageUtill");
  private static final Key<CachedValue<List<Pair<PsiElement, TextRange>>>> INJECTED_PSI = Key.create("injectedPsi");
  private static final Key<List<Pair<IElementType, TextRange>>> HIGHLIGHT_TOKENS = Key.create("HIGHLIGHT_TOKENS");

  @Nullable
  public static <T extends PsiLanguageInjectionHost> List<Pair<PsiElement, TextRange>> getInjectedPsiFiles(@NotNull T host, @Nullable LiteralTextEscaper<T> textEscaper) {
    if (!host.isPhysical()) {
      CachedValueProvider.Result<List<Pair<PsiElement, TextRange>>> result = new InjectedPsiProvider<T>(host, textEscaper).compute();
      return result == null ? null : result.getValue();
    }

    CachedValue<List<Pair<PsiElement, TextRange>>> cachedPsi = host.getUserData(INJECTED_PSI);
    if (cachedPsi == null) {
      InjectedPsiProvider<T> provider = new InjectedPsiProvider<T>(host, textEscaper);
      CachedValueProvider.Result<List<Pair<PsiElement, TextRange>>> result = provider.compute();
      if (result == null) return null;
      cachedPsi = host.getManager().getCachedValuesManager().createCachedValue(provider, false);
      ((CachedValueImpl<List<Pair<PsiElement, TextRange>>>)cachedPsi).setValue(result);
      ((UserDataHolderEx)host).putUserDataIfAbsent(INJECTED_PSI, cachedPsi);
      return result.getValue();
    }
    return cachedPsi.getValue();
  }

  private static <T extends PsiLanguageInjectionHost> PsiElement parseInjectedPsiFile(@NotNull final T host,
                                                                                      @NotNull final TextRange rangeInsideHost,
                                                                                      @NotNull final Language language,
                                                                                      @NotNull final VirtualFile hostVirtualFile,
                                                                                      @NotNull final TextRange hostRange,
                                                                                      @NotNull final DocumentEx hostDocument,
                                                                                      LiteralTextEscaper<T> textEscaper,
                                                                                      @NotNull String prefix,
                                                                                      @NotNull String suffix) {
    final ParserDefinition parserDefinition = language.getParserDefinition();
    if (parserDefinition == null) return null;
    PsiFile hostFile = host.getContainingFile();
    if (hostFile == null) return null;
    PsiManagerEx psiManager = (PsiManagerEx)host.getManager();
    final Project project = psiManager.getProject();

    final String hostText = host.getText();
    StringBuilder outChars = new StringBuilder(hostText.length());
    outChars.append(prefix);
    if (textEscaper == null) {
      outChars.append(hostText, rangeInsideHost.getStartOffset(), rangeInsideHost.getEndOffset());
    }
    else {
      boolean result = textEscaper.decode(host, rangeInsideHost, outChars);
      if (!result) return null;
    }
    outChars.append(suffix);

    TextRange documentWindow = hostRange.cutOut(rangeInsideHost);
    DocumentRange documentRange = new DocumentRange(hostDocument, documentWindow,prefix,suffix);
    VirtualFileDelegate virtualFile = InjectedManagerImpl.getInstance().createVirtualFile(language, hostVirtualFile, documentRange, outChars, project);

    DocumentImpl decodedDocument = new DocumentImpl(outChars);
    FileDocumentManagerImpl.registerDocument(decodedDocument, virtualFile);

    SingleRootFileViewProvider viewProvider = new MyFileViewProvider(project, virtualFile);
    PsiFile psiFile = parserDefinition.createFile(viewProvider);
    assert psiFile.getViewProvider() instanceof MyFileViewProvider : psiFile.getViewProvider();

    SmartPsiElementPointer<T> pointer = createHostSmartPointer(host);
    psiFile.putUserData(ResolveUtil.INJECTED_IN_ELEMENT, pointer);
    psiFile.putUserData(PsiManagerImpl.LANGUAGE_DIALECT,language instanceof LanguageDialect ? (LanguageDialect)language:null);

    final ASTNode parsedNode = psiFile.getNode();
    if (!(parsedNode instanceof FileElement)) return null;

    List<Pair<IElementType, TextRange>> tokens = obtainHighlightTokensFromLexer(language, outChars, textEscaper, virtualFile, project, rangeInsideHost, documentRange);

    String sourceRawText = prefix + hostText.substring(rangeInsideHost.getStartOffset(), rangeInsideHost.getEndOffset()) + suffix;
    if (textEscaper != null) {
      patchLeafs(parsedNode, textEscaper, rangeInsideHost, hostText, prefix.length(), suffix.length());
    }
    String parsedText = parsedNode.getText();
    LOG.assertTrue(parsedText.equals(sourceRawText));

    parsedNode.putUserData(TreeElement.MANAGER_KEY, psiManager);
    virtualFile.setContent(null, documentRange.getText(), false);
    FileDocumentManagerImpl.registerDocument(documentRange, virtualFile);
    synchronized (PsiLock.LOCK) {
      psiFile = registerDocumentRange(documentRange, psiFile);
    }

    ((MyFileViewProvider)psiFile.getViewProvider()).setVirtualFile(virtualFile);
    ((SingleRootFileViewProvider)psiFile.getViewProvider()).forceCachedPsi(psiFile);
    psiFile.putUserData(HIGHLIGHT_TOKENS, tokens);

    PsiDocumentManagerImpl.checkConsistency(psiFile, documentRange);
    PsiDocumentManagerImpl.checkConsistency(hostFile, documentRange.getDelegate());

    return psiFile;
  }

  public static List<Pair<IElementType, TextRange>> getHighlightTokens(PsiFile file) {
    return file.getUserData(HIGHLIGHT_TOKENS);
  }

  private static <T extends PsiLanguageInjectionHost> List<Pair<IElementType, TextRange>> obtainHighlightTokensFromLexer(final Language language,
                                                                                                                         final StringBuilder outChars,
                                                                                                                         final LiteralTextEscaper<T> textEscaper,
                                                                                                                         final VirtualFileDelegate virtualFile,
                                                                                                                         final Project project,
                                                                                                                         final TextRange rangeInsideHost,
                                                                                                                         final DocumentRange documentRange) {<commit_msg>IDEADEV-18538<commit_after>package com.intellij.psi.impl.source.tree.injected;

import com.intellij.lang.ASTNode;
import com.intellij.lang.Language;
import com.intellij.lang.LanguageDialect;
import com.intellij.lang.ParserDefinition;
import com.intellij.lexer.Lexer;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.editor.RangeMarker;
import com.intellij.openapi.editor.SelectionModel;
import com.intellij.openapi.editor.ex.DocumentEx;
import com.intellij.openapi.editor.impl.DocumentImpl;
import com.intellij.openapi.editor.impl.EditorImpl;
import com.intellij.openapi.editor.impl.injected.DocumentRange;
import com.intellij.openapi.editor.impl.injected.EditorDelegate;
import com.intellij.openapi.editor.impl.injected.VirtualFileDelegate;
import com.intellij.openapi.extensions.Extensions;
import com.intellij.openapi.fileEditor.FileEditorManager;
import com.intellij.openapi.fileEditor.OpenFileDescriptor;
import com.intellij.openapi.fileEditor.impl.FileDocumentManagerImpl;
import com.intellij.openapi.fileTypes.SyntaxHighlighter;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.*;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.*;
import com.intellij.psi.impl.CachedValueImpl;
import com.intellij.psi.impl.PsiDocumentManagerImpl;
import com.intellij.psi.impl.PsiManagerEx;
import com.intellij.psi.impl.PsiManagerImpl;
import com.intellij.psi.impl.source.PsiFileImpl;
import com.intellij.psi.impl.source.resolve.ResolveUtil;
import com.intellij.psi.impl.source.tree.*;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.util.CachedValue;
import com.intellij.psi.util.CachedValueProvider;
import com.intellij.psi.util.PsiModificationTracker;
import com.intellij.util.SmartList;
import gnu.trove.THashMap;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @author cdr
 */
public class InjectedLanguageUtil {
  private static final Logger LOG = Logger.getInstance("com.intellij.psi.impl.source.tree.injected.InjectedLanguageUtill");
  private static final Key<CachedValue<List<Pair<PsiElement, TextRange>>>> INJECTED_PSI = Key.create("injectedPsi");
  private static final Key<List<Pair<IElementType, TextRange>>> HIGHLIGHT_TOKENS = Key.create("HIGHLIGHT_TOKENS");

  @Nullable
  public static <T extends PsiLanguageInjectionHost> List<Pair<PsiElement, TextRange>> getInjectedPsiFiles(@NotNull T host, @Nullable LiteralTextEscaper<T> textEscaper) {
    if (!host.isPhysical()) {
      CachedValueProvider.Result<List<Pair<PsiElement, TextRange>>> result = new InjectedPsiProvider<T>(host, textEscaper).compute();
      return result == null ? null : result.getValue();
    }

    CachedValue<List<Pair<PsiElement, TextRange>>> cachedPsi = host.getUserData(INJECTED_PSI);
    if (cachedPsi == null) {
      InjectedPsiProvider<T> provider = new InjectedPsiProvider<T>(host, textEscaper);
      CachedValueProvider.Result<List<Pair<PsiElement, TextRange>>> result = provider.compute();
      if (result == null) return null;
      cachedPsi = host.getManager().getCachedValuesManager().createCachedValue(provider, false);
      ((CachedValueImpl<List<Pair<PsiElement, TextRange>>>)cachedPsi).setValue(result);
      ((UserDataHolderEx)host).putUserDataIfAbsent(INJECTED_PSI, cachedPsi);
      return result.getValue();
    }
    return cachedPsi.getValue();
  }

  private static <T extends PsiLanguageInjectionHost> PsiElement parseInjectedPsiFile(@NotNull final T host,
                                                                                      @NotNull final TextRange rangeInsideHost,
                                                                                      @NotNull final Language language,
                                                                                      @NotNull final VirtualFile hostVirtualFile,
                                                                                      @NotNull final TextRange hostRange,
                                                                                      @NotNull final DocumentEx hostDocument,
                                                                                      LiteralTextEscaper<T> textEscaper,
                                                                                      @NotNull String prefix,
                                                                                      @NotNull String suffix) {
    final ParserDefinition parserDefinition = language.getParserDefinition();
    if (parserDefinition == null) return null;
    PsiFile hostFile = host.getContainingFile();
    if (hostFile == null) return null;
    PsiManagerEx psiManager = (PsiManagerEx)host.getManager();
    final Project project = psiManager.getProject();

    final String hostText = host.getText();
    StringBuilder outChars = new StringBuilder(hostText.length());
    outChars.append(prefix);
    if (textEscaper == null) {
      outChars.append(hostText, rangeInsideHost.getStartOffset(), rangeInsideHost.getEndOffset());
    }
    else {
      boolean result = textEscaper.decode(host, rangeInsideHost, outChars);
      if (!result) return null;
    }
    outChars.append(suffix);

    TextRange documentWindow = hostRange.cutOut(rangeInsideHost);
    DocumentRange documentRange = new DocumentRange(hostDocument, documentWindow,prefix,suffix);
    VirtualFileDelegate virtualFile = InjectedManagerImpl.getInstance().createVirtualFile(language, hostVirtualFile, documentRange, outChars, project);

    DocumentImpl decodedDocument = new DocumentImpl(outChars);
    FileDocumentManagerImpl.registerDocument(decodedDocument, virtualFile);

    SingleRootFileViewProvider viewProvider = new MyFileViewProvider(project, virtualFile);
    PsiFile psiFile = parserDefinition.createFile(viewProvider);
    assert psiFile.getViewProvider() instanceof MyFileViewProvider : psiFile.getViewProvider();

    SmartPsiElementPointer<T> pointer = createHostSmartPointer(host);
    psiFile.putUserData(ResolveUtil.INJECTED_IN_ELEMENT, pointer);
    psiFile.putUserData(PsiManagerImpl.LANGUAGE_DIALECT,language instanceof LanguageDialect ? (LanguageDialect)language:null);

    final ASTNode parsedNode = psiFile.getNode();
    if (!(parsedNode instanceof FileElement)) return null;

    List<Pair<IElementType, TextRange>> tokens = obtainHighlightTokensFromLexer(language, outChars, textEscaper, virtualFile, project, rangeInsideHost, documentRange);

    String sourceRawText = prefix + hostText.substring(rangeInsideHost.getStartOffset(), rangeInsideHost.getEndOffset()) + suffix;
    if (textEscaper != null) {
      patchLeafs(parsedNode, textEscaper, rangeInsideHost, hostText, prefix.length(), suffix.length());
    }
    String parsedText = parsedNode.getText();
    LOG.assertTrue(parsedText.equals(sourceRawText));

    parsedNode.putUserData(TreeElement.MANAGER_KEY, psiManager);
    virtualFile.setContent(null, documentRange.getText(), false);
    FileDocumentManagerImpl.registerDocument(documentRange, virtualFile);
    synchronized (PsiLock.LOCK) {
      psiFile = registerDocumentRange(documentRange, psiFile);
      MyFileViewProvider myFileViewProvider = (MyFileViewProvider)psiFile.getViewProvider();
      myFileViewProvider.setVirtualFile(virtualFile);
      myFileViewProvider.forceCachedPsi(psiFile);
    }

    psiFile.putUserData(HIGHLIGHT_TOKENS, tokens);

    PsiDocumentManagerImpl.checkConsistency(psiFile, documentRange);
    PsiDocumentManagerImpl.checkConsistency(hostFile, documentRange.getDelegate());

    return psiFile;
  }

  public static List<Pair<IElementType, TextRange>> getHighlightTokens(PsiFile file) {
    return file.getUserData(HIGHLIGHT_TOKENS);
  }

  private static <T extends PsiLanguageInjectionHost> List<Pair<IElementType, TextRange>> obtainHighlightTokensFromLexer(final Language language,
                                                                                                                         final StringBuilder outChars,
                                                                                                                         final LiteralTextEscaper<T> textEscaper,
                                                                                                                         final VirtualFileDelegate virtualFile,
                                                                                                                         final Project project,
                                                                                                                         final TextRange rangeInsideHost,
                                                                                                                         final DocumentRange documentRange) {
{'fraction_non_alphanumeric': 0.08599167822468794, 'fraction_numerical': 0.00027739251040221914, 'mean_word_length': 2.6467732146469753, 'pattern_counts': {'":': 0, '<': 77, '<?xml version=': 0, '>': 77, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 124, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4232662', 'n_tokens_mistral': 4588, 'n_tokens_neox': 4456, 'n_words': 807}