<commit_before>/*
 * Use of this source code is governed by the MIT license that can be
 * found in the LICENSE file.
 */

package org.rust.lang.core.macros

import com.intellij.lang.LanguageParserDefinitions
import com.intellij.lang.PsiBuilder
import com.intellij.lang.PsiBuilderFactory
import com.intellij.lang.PsiBuilderUtil
import com.intellij.lang.parser.GeneratedParserUtilBase
import com.intellij.openapi.progress.ProgressManager
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiComment
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.impl.source.DummyHolderFactory
import com.intellij.psi.util.CachedValueProvider
import com.intellij.psi.util.CachedValuesManager
import com.intellij.psi.util.PsiModificationTracker
import org.rust.cargo.project.workspace.PackageOrigin
import org.rust.lang.RsLanguage
import org.rust.lang.core.parser.RustParser
import org.rust.lang.core.parser.RustParserUtil.collapsedTokenType
import org.rust.lang.core.psi.*
import org.rust.lang.core.psi.RsElementTypes.IDENTIFIER
import org.rust.lang.core.psi.ext.*
import org.rust.openapiext.Testmark
import org.rust.openapiext.forEachChild
import org.rust.stdext.joinToWithBuffer
import org.rust.stdext.mapNotNullToSet
import java.util.*
import java.util.Collections.singletonMap

private data class MacroSubstitution(
    /**
     * Maps meta variable names with the actual values, e.g. for this macro
     * ```rust
     * macro_rules! foo {
     *     ($ i:item) => ( $i )
     * }
     * foo!( mod a {} )
     * ```
     * we map "i" to "mod a {}"
     */
    val variables: Map<String, String>,

    /**
     * Contains macro groups values. E.g. for this macro
     * ```rust
     * macro_rules! foo {
     *     ($($ i:item),*; $($ e:expr),*) => {
     *         fn foo() { $($ e);*; }
     *         $($ i)*
     *     }
     * }
     * foo! { mod a {}, mod b {}; 1, 2 }
     * ```
     * It will contains
     * ```
     * [
     *     MacroGroup {
     *         substs: [
     *             {"i" => "mod a {}"},
     *             {"i" => "mod b {}"}
     *         ]
     *     },
     *     MacroGroup {
     *         substs: [
     *             {"e" => "1"},
     *             {"e" => "2"}
     *         ]
     *     }
     * ]
     * ```
     */
    val groups: List<MacroGroup>
)

private data class MacroGroup(
    val definition: RsMacroBindingGroup,
    val substs: List<MacroSubstitution>
)

private data class WithParent(
    private val subst: MacroSubstitution,
    private val parent: WithParent?
) {
    val groups: List<MacroGroup>
        get() = subst.groups

    fun getVar(name: String): String? =
        subst.variables[name] ?: parent?.getVar(name)
}

private val STD_MACRO_WHITELIST = setOf("write", "writeln")

class MacroExpander(val project: Project) {
    private val psiFactory = RsPsiFactory(project)

    fun expandMacro(def: RsMacro, call: RsMacroCall): List<RsExpandedElement>? {
        // All std macros contain the only `impl`s which are not supported for now, so ignoring them
        if (def.containingCargoTarget?.pkg?.origin == PackageOrigin.STDLIB && def.name !in STD_MACRO_WHITELIST) {
            return null
        }

        val expandedText = expandMacroAsText(def, call) ?: return null
        return when (call.context) {
            is RsMacroExpr -> listOfNotNull(psiFactory.tryCreateExpression(expandedText))
            else -> psiFactory.createFile(expandedText).childrenOfType()
        }
    }

    private fun expandMacroAsText(def: RsMacro, call: RsMacroCall): CharSequence? {
        val (case, subst) = findMatchingPattern(def, call) ?: return null
        val macroExpansion = case.macroExpansion ?: return null

        val crate = def.containingMod.crateRelativePath ?: ""
        val substWithGlobalVars = WithParent(
            subst,
            WithParent(
                MacroSubstitution(
                    singletonMap("crate", crate),
                    emptyList()
                ),
                null
            )
        )

        return substituteMacro(macroExpansion.macroExpansionContents, substWithGlobalVars)
    }

    private fun findMatchingPattern(
        def: RsMacro,
        call: RsMacroCall
    ): Pair<RsMacroCase, MacroSubstitution>? {
        val macroCallBody = createPsiBuilder(call, call.macroBody ?: return null)
        var start = macroCallBody.mark()
        return def.macroBodyStubbed?.macroCaseList?.asSequence()
            ?.mapNotNull { case ->
                val subst = case.pattern.match(macroCallBody)
                if (subst != null) {
                    case to subst
                } else {
                    start.rollbackTo()
                    start = macroCallBody.mark()
                    null
                }
            }
            ?.firstOrNull()
    }

    private fun createPsiBuilder(context: PsiElement, text: String): PsiBuilder {
        val holder = DummyHolderFactory.createHolder(PsiManager.getInstance(project), context).treeElement
        val parserDefinition = LanguageParserDefinitions.INSTANCE.forLanguage(RsLanguage)
            ?: error("No parser definition for language $RsLanguage")
        val lexer = parserDefinition.createLexer(project)
        return PsiBuilderFactory.getInstance().createBuilder(project, holder, lexer, RsLanguage, text)
    }

    private fun substituteMacro(root: PsiElement, subst: WithParent): CharSequence? =
        buildString { if (!substituteMacro(this, root, subst)) return null }

    private fun substituteMacro(sb: StringBuilder, root: PsiElement, subst: WithParent): Boolean {<commit_msg>Cleanup<commit_after>/*
 * Use of this source code is governed by the MIT license that can be
 * found in the LICENSE file.
 */

package org.rust.lang.core.macros

import com.intellij.lang.LanguageParserDefinitions
import com.intellij.lang.PsiBuilder
import com.intellij.lang.PsiBuilderFactory
import com.intellij.lang.PsiBuilderUtil
import com.intellij.lang.parser.GeneratedParserUtilBase
import com.intellij.openapi.progress.ProgressManager
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiComment
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.impl.source.DummyHolderFactory
import com.intellij.psi.util.CachedValueProvider
import com.intellij.psi.util.CachedValuesManager
import com.intellij.psi.util.PsiModificationTracker
import org.rust.cargo.project.workspace.PackageOrigin
import org.rust.lang.RsLanguage
import org.rust.lang.core.parser.RustParser
import org.rust.lang.core.parser.RustParserUtil.collapsedTokenType
import org.rust.lang.core.psi.*
import org.rust.lang.core.psi.RsElementTypes.IDENTIFIER
import org.rust.lang.core.psi.ext.*
import org.rust.openapiext.Testmark
import org.rust.openapiext.forEachChild
import org.rust.stdext.joinToWithBuffer
import org.rust.stdext.mapNotNullToSet
import java.util.*
import java.util.Collections.singletonMap

private data class MacroSubstitution(
    /**
     * Maps meta variable names with the actual values, e.g. for this macro
     * ```rust
     * macro_rules! foo {
     *     ($ i:item) => ( $i )
     * }
     * foo!( mod a {} )
     * ```
     * we map "i" to "mod a {}"
     */
    val variables: Map<String, String>,

    /**
     * Contains macro groups values. E.g. for this macro
     * ```rust
     * macro_rules! foo {
     *     ($($ i:item),*; $($ e:expr),*) => {
     *         fn foo() { $($ e);*; }
     *         $($ i)*
     *     }
     * }
     * foo! { mod a {}, mod b {}; 1, 2 }
     * ```
     * It will contains
     * ```
     * [
     *     MacroGroup {
     *         substs: [
     *             {"i" => "mod a {}"},
     *             {"i" => "mod b {}"}
     *         ]
     *     },
     *     MacroGroup {
     *         substs: [
     *             {"e" => "1"},
     *             {"e" => "2"}
     *         ]
     *     }
     * ]
     * ```
     */
    val groups: List<MacroGroup>
)

private data class MacroGroup(
    val definition: RsMacroBindingGroup,
    val substs: List<MacroSubstitution>
)

private data class WithParent(
    private val subst: MacroSubstitution,
    private val parent: WithParent?
) {
    val groups: List<MacroGroup>
        get() = subst.groups

    fun getVar(name: String): String? =
        subst.variables[name] ?: parent?.getVar(name)
}

private val STD_MACRO_WHITELIST = setOf("write", "writeln")

class MacroExpander(val project: Project) {
    private val psiFactory = RsPsiFactory(project)

    fun expandMacro(def: RsMacro, call: RsMacroCall): List<RsExpandedElement>? {
        // All std macros contain the only `impl`s which are not supported for now, so ignoring them
        if (def.containingCargoTarget?.pkg?.origin == PackageOrigin.STDLIB && def.name !in STD_MACRO_WHITELIST) {
            return null
        }

        val expandedText = expandMacroAsText(def, call) ?: return null
        return when (call.context) {
            is RsMacroExpr -> listOfNotNull(psiFactory.tryCreateExpression(expandedText))
            else -> psiFactory.createFile(expandedText).childrenOfType()
        }
    }

    private fun expandMacroAsText(def: RsMacro, call: RsMacroCall): CharSequence? {
        val (case, subst) = findMatchingPattern(def, call) ?: return null
        val macroExpansion = case.macroExpansion ?: return null

        val crate = def.containingMod.crateRelativePath ?: ""
        val substWithGlobalVars = WithParent(
            subst,
            WithParent(
                MacroSubstitution(
                    singletonMap("crate", crate),
                    emptyList()
                ),
                null
            )
        )

        return substituteMacro(macroExpansion.macroExpansionContents, substWithGlobalVars)
    }

    private fun findMatchingPattern(
        def: RsMacro,
        call: RsMacroCall
    ): Pair<RsMacroCase, MacroSubstitution>? {
        val macroCallBody = createPsiBuilder(call, call.macroBody ?: return null)
        var start = macroCallBody.mark()
        val macroCaseList = def.macroBodyStubbed?.macroCaseList ?: return null

        for (case in macroCaseList) {
            val subst = case.pattern.match(macroCallBody)
            if (subst != null) {
                return case to subst
            } else {
                start.rollbackTo()
                start = macroCallBody.mark()
            }
        }

        return null
    }

    private fun createPsiBuilder(context: PsiElement, text: String): PsiBuilder {
        val holder = DummyHolderFactory.createHolder(PsiManager.getInstance(project), context).treeElement
        val parserDefinition = LanguageParserDefinitions.INSTANCE.forLanguage(RsLanguage)
            ?: error("No parser definition for language $RsLanguage")
        val lexer = parserDefinition.createLexer(project)
        return PsiBuilderFactory.getInstance().createBuilder(project, holder, lexer, RsLanguage, text)
    }

    private fun substituteMacro(root: PsiElement, subst: WithParent): CharSequence? =
        buildString { if (!substituteMacro(this, root, subst)) return null }

    private fun substituteMacro(sb: StringBuilder, root: PsiElement, subst: WithParent): Boolean {
{'fraction_non_alphanumeric': 0.10793398640896655, 'fraction_numerical': 0.0007060277115876799, 'mean_word_length': 2.508359133126935, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 32, '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': '6125089', 'n_tokens_mistral': 3413, 'n_tokens_neox': 3345, 'n_words': 746}