<commit_before>package org.jetbrains.plugins.scala.util

import java.io._
import java.util.regex.Pattern

import com.intellij.openapi.actionSystem.{AnAction, AnActionEvent}
import com.intellij.openapi.application.PathManager
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.vfs.VirtualFileManager
import com.intellij.psi.impl.source.tree.LeafPsiElement
import com.intellij.psi._
import org.jetbrains.plugin.scala.util.MacroExpansion
import org.jetbrains.plugins.scala.lang.psi.ScalaPsiElement
import org.jetbrains.plugins.scala.lang.psi.api.{ScalaRecursiveElementVisitor, ScalaElementVisitor, ScalaFile}
import org.jetbrains.plugins.scala.lang.psi.api.expr.{ScAnnotation, ScMethodCall}
import org.jetbrains.plugins.scala.lang.psi.api.statements.ScFunction
import org.jetbrains.plugins.scala.lang.psi.api.statements.params.ScParameter
import org.jetbrains.plugins.scala.lang.psi.api.toplevel.typedef.{ScObject, ScClass}
import org.jetbrains.plugins.scala.lang.psi.impl.ScalaPsiElementFactory
import org.jetbrains.plugins.scala.extensions.{inWriteCommandAction, executeOnPooledThread}

import scala.annotation.tailrec


class MacroExpandAction extends AnAction {

  case class ResolvedMacroExpansion(expansion: MacroExpansion, psiElement: Option[SmartPsiElementPointer[PsiElement]])
  class UnresolvedExpansion extends Exception

  private val LOG = Logger.getInstance(getClass)

  override def actionPerformed(e: AnActionEvent): Unit = {
    implicit val currentEvent = e

    val sourceEditor = FileEditorManager.getInstance(e.getProject).getSelectedTextEditor
    val psiFile = PsiDocumentManager.getInstance(e.getProject).getPsiFile(sourceEditor.getDocument)
    val candidates = psiFile match {
      case file: ScalaFile => findCandidatesInFile(file)
      case _ => Seq.empty
    }

    val expansions = deserializeExpansions(e)
    val filtered = expansions.filter { exp =>
      psiFile.getVirtualFile.getPath == exp.place.sourceFile
    }
    val ensugared = filtered.map(e => MacroExpansion(e.place, ensugarExpansion(e.body)))
    val resolved = tryResolveExpansionPlaces(ensugared)

    // if macro is under cursor, expand it, otherwise expand all macros in current file
    resolved
      .find(_.expansion.place.line == sourceEditor.getCaretModel.getLogicalPosition.line+1)
      .map(expandMacroUnderCursor)
      .getOrElse(expandAllMacroInCurrentFile(resolved))
  }


  def expandMacroUnderCursor(expansion: ResolvedMacroExpansion)(implicit e: AnActionEvent) = {
    inWriteCommandAction(e.getProject) {
      applyExpansion(expansion)
      e.getProject
    }
  }

  def expandAllMacroInCurrentFile(expansions: Seq[ResolvedMacroExpansion])(implicit e: AnActionEvent)  = {
    inWriteCommandAction(e.getProject) {
      applyExpansions(expansions.toList)
      e.getProject
    }
  }

  def findCandidatesInFile(file: ScalaFile): Seq[ScalaPsiElement] = {
    val buffer = scala.collection.mutable.ListBuffer[ScalaPsiElement]()
    val visitor = new ScalaRecursiveElementVisitor {
      override def visitAnnotation(annotation: ScAnnotation) = {
        // TODO
      }
      override def visitMethodCallExpression(call: ScMethodCall) = {
        // TODO
      }
    }
    file.accept(visitor)
    buffer.toSeq
  }

  def applyExpansion(resolved: ResolvedMacroExpansion)(implicit e: AnActionEvent): Unit = {
    if (resolved.psiElement.isEmpty)
      throw new UnresolvedExpansion
    if (resolved.expansion.body.isEmpty) {
      LOG.warn(s"got empty expansion at ${resolved.expansion.place}, skipping")
      return
    }
    resolved.psiElement.get.getElement match {
      case (annot: ScAnnotation) => // TODO
        expandAnnotation(annot, resolved.expansion)
      case (mc: ScMethodCall) =>
        expandMacroCall(mc, resolved.expansion)
      case (other) => () // unreachable
    }
  }

  def applyExpansions(expansions: Seq[ResolvedMacroExpansion], triedResolving: Boolean = false)(implicit e: AnActionEvent): Unit = {
    expansions match {
      case x::xs =>
        try {
          applyExpansion(x)
          applyExpansions(xs)
        }
        catch {
          case exc: UnresolvedExpansion if !triedResolving =>
            applyExpansions(tryResolveExpansionPlace(x.expansion) :: xs, triedResolving = true)
          case exc: UnresolvedExpansion if triedResolving =>
            LOG.warn(s"unable to expand ${x.expansion.place}, cannot resolve place, skipping")
            applyExpansions(xs)
        }
      case Nil =>
    }
  }

  def expandAnnotation(place: ScAnnotation, expansion: MacroExpansion)(implicit e: AnActionEvent) = {
    // we can only macro-annotate scala code
    place.getParent.getParent match {
      case clazz: ScClass =>
        // FIXME: parse companion class as well(if present)
        val newClazz = ScalaPsiElementFactory.createBlockExpressionWithoutBracesFromText(expansion.body, PsiManager.getInstance(e.getProject))
        clazz.replace(newClazz)
      case obj: ScObject => // TODO
      case fun: ScFunction => // TODO
      case param: ScParameter => // TODO
      case other => LOG.warn(s"Unexpected annotated element: $other at ${other.getText}")
    }
  }

  def expandMacroCall(call: ScMethodCall, expansion: MacroExpansion)(implicit e: AnActionEvent) = {
    val blockImpl = ScalaPsiElementFactory.createBlockExpressionWithoutBracesFromText(s"${expansion.body}", PsiManager.getInstance(e.getProject))
    call.getParent.addAfter(blockImpl, call)
    call.delete()
  }

  def tryResolveExpansionPlace(expansion: MacroExpansion)(implicit e: AnActionEvent): ResolvedMacroExpansion = {
    ResolvedMacroExpansion(expansion, getRealOwner(expansion).map(new IdentitySmartPointer[PsiElement](_)))
  }

  def tryResolveExpansionPlaces(expansions: Seq[MacroExpansion])(implicit e: AnActionEvent): Seq[ResolvedMacroExpansion] = {
    expansions.map(tryResolveExpansionPlace)
  }
<commit_msg>macroexpand: macro annotated methods dont' expand #SCL-9244 fixed<commit_after>package org.jetbrains.plugins.scala.util

import java.io._
import java.util.regex.Pattern

import com.intellij.openapi.actionSystem.{AnAction, AnActionEvent}
import com.intellij.openapi.application.PathManager
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.vfs.VirtualFileManager
import com.intellij.psi.impl.source.tree.LeafPsiElement
import com.intellij.psi._
import org.jetbrains.plugin.scala.util.MacroExpansion
import org.jetbrains.plugins.scala.lang.psi.ScalaPsiElement
import org.jetbrains.plugins.scala.lang.psi.api.{ScalaRecursiveElementVisitor, ScalaElementVisitor, ScalaFile}
import org.jetbrains.plugins.scala.lang.psi.api.expr.{ScBlock, ScAnnotation, ScMethodCall}
import org.jetbrains.plugins.scala.lang.psi.api.statements.{ScAnnotationsHolder, ScFunction}
import org.jetbrains.plugins.scala.lang.psi.api.statements.params.ScParameter
import org.jetbrains.plugins.scala.lang.psi.api.toplevel.typedef.{ScObject, ScClass}
import org.jetbrains.plugins.scala.lang.psi.impl.ScalaPsiElementFactory
import org.jetbrains.plugins.scala.extensions.{inWriteCommandAction, executeOnPooledThread}

import scala.annotation.tailrec


class MacroExpandAction extends AnAction {

  case class ResolvedMacroExpansion(expansion: MacroExpansion, psiElement: Option[SmartPsiElementPointer[PsiElement]])
  class UnresolvedExpansion extends Exception

  private val LOG = Logger.getInstance(getClass)

  override def actionPerformed(e: AnActionEvent): Unit = {
    implicit val currentEvent = e

    val sourceEditor = FileEditorManager.getInstance(e.getProject).getSelectedTextEditor
    val psiFile = PsiDocumentManager.getInstance(e.getProject).getPsiFile(sourceEditor.getDocument)
    val candidates = psiFile match {
      case file: ScalaFile => findCandidatesInFile(file)
      case _ => Seq.empty
    }

    val expansions = deserializeExpansions(e)
    val filtered = expansions.filter { exp =>
      psiFile.getVirtualFile.getPath == exp.place.sourceFile
    }
    val ensugared = filtered.map(e => MacroExpansion(e.place, ensugarExpansion(e.body)))
    val resolved = tryResolveExpansionPlaces(ensugared)

    // if macro is under cursor, expand it, otherwise expand all macros in current file
    resolved
      .find(_.expansion.place.line == sourceEditor.getCaretModel.getLogicalPosition.line+1)
      .map(expandMacroUnderCursor)
      .getOrElse(expandAllMacroInCurrentFile(resolved))
  }


  def expandMacroUnderCursor(expansion: ResolvedMacroExpansion)(implicit e: AnActionEvent) = {
    inWriteCommandAction(e.getProject) {
      applyExpansion(expansion)
      e.getProject
    }
  }

  def expandAllMacroInCurrentFile(expansions: Seq[ResolvedMacroExpansion])(implicit e: AnActionEvent)  = {
    inWriteCommandAction(e.getProject) {
      applyExpansions(expansions.toList)
      e.getProject
    }
  }

  def findCandidatesInFile(file: ScalaFile): Seq[ScalaPsiElement] = {
    val buffer = scala.collection.mutable.ListBuffer[ScalaPsiElement]()
    val visitor = new ScalaRecursiveElementVisitor {
      override def visitAnnotation(annotation: ScAnnotation) = {
        // TODO
      }
      override def visitMethodCallExpression(call: ScMethodCall) = {
        // TODO
      }
    }
    file.accept(visitor)
    buffer.toSeq
  }

  def applyExpansion(resolved: ResolvedMacroExpansion)(implicit e: AnActionEvent): Unit = {
    if (resolved.psiElement.isEmpty)
      throw new UnresolvedExpansion
    if (resolved.expansion.body.isEmpty) {
      LOG.warn(s"got empty expansion at ${resolved.expansion.place}, skipping")
      return
    }
    resolved.psiElement.get.getElement match {
      case (annot: ScAnnotation) => // TODO
        expandAnnotation(annot, resolved.expansion)
      case (mc: ScMethodCall) =>
        expandMacroCall(mc, resolved.expansion)
      case (other) => () // unreachable
    }
  }

  def applyExpansions(expansions: Seq[ResolvedMacroExpansion], triedResolving: Boolean = false)(implicit e: AnActionEvent): Unit = {
    expansions match {
      case x::xs =>
        try {
          applyExpansion(x)
          applyExpansions(xs)
        }
        catch {
          case exc: UnresolvedExpansion if !triedResolving =>
            applyExpansions(tryResolveExpansionPlace(x.expansion) :: xs, triedResolving = true)
          case exc: UnresolvedExpansion if triedResolving =>
            LOG.warn(s"unable to expand ${x.expansion.place}, cannot resolve place, skipping")
            applyExpansions(xs)
        }
      case Nil =>
    }
  }

  def expandAnnotation(place: ScAnnotation, expansion: MacroExpansion)(implicit e: AnActionEvent) = {
    // we can only macro-annotate scala code
    place .getParent.getParent match {
      case holder: ScAnnotationsHolder =>
        val body = expansion.body
        val newPsi = ScalaPsiElementFactory.createBlockExpressionWithoutBracesFromText(body, PsiManager.getInstance(e.getProject))
        newPsi.firstChild match {
          case Some(block: ScBlock) => // insert content of block expression(annotation can generate >1 expression)
            val children = block.getChildren
            holder.getParent.addRangeAfter(children.tail.head, children.dropRight(1).last, holder)
            holder.delete()
          case Some(psi: PsiElement) => // defns/method bodies/etc...
            holder.replace(psi)
          case None => LOG.warn(s"Failed to parse expansion: $body")
        }
      case other =>  LOG.warn(s"Unexpected annotated element: $other at ${other.getText}")
    }
  }

  def expandMacroCall(call: ScMethodCall, expansion: MacroExpansion)(implicit e: AnActionEvent) = {
    val blockImpl = ScalaPsiElementFactory.createBlockExpressionWithoutBracesFromText(s"${expansion.body}", PsiManager.getInstance(e.getProject))
    call.getParent.addAfter(blockImpl, call)
    call.delete()
  }

  def tryResolveExpansionPlace(expansion: MacroExpansion)(implicit e: AnActionEvent): ResolvedMacroExpansion = {
    ResolvedMacroExpansion(expansion, getRealOwner(expansion).map(new IdentitySmartPointer[PsiElement](_)))
  }

  def tryResolveExpansionPlaces(expansions: Seq[MacroExpansion])(implicit e: AnActionEvent): Seq[ResolvedMacroExpansion] = {
    expansions.map(tryResolveExpansionPlace)
  }

{'fraction_non_alphanumeric': 0.09573344169036976, 'fraction_numerical': 0.000650142218610321, 'mean_word_length': 4.057953144266338, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6718766', 'n_tokens_mistral': 3742, 'n_tokens_neox': 3593, 'n_words': 673}