<commit_before>package uk.co.reecedunn.intellij.plugin.xquery.resolve.reference;

import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiReferenceBase;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import uk.co.reecedunn.intellij.plugin.xquery.ast.xquery.*;
import uk.co.reecedunn.intellij.plugin.xquery.psi.PsiNavigation;
import uk.co.reecedunn.intellij.plugin.xquery.psi.XQueryModuleProvider;
import uk.co.reecedunn.intellij.plugin.xquery.psi.XQueryNamespace;

import java.util.List;

public class XQueryFunctionNameReference extends PsiReferenceBase<XQueryEQName> {
    public XQueryFunctionNameReference(XQueryEQName element, TextRange range) {
        super(element, range);
    }

    @Nullable
    @Override
    public PsiElement resolve() {
        XQueryNamespace ns = getElement().resolvePrefixNamespace();
        if (ns == null) {
            return null;
        }

        if (ns.getDeclaration() instanceof XQueryModuleProvider) {
            XQueryModuleProvider provider = (XQueryModuleProvider)ns.getDeclaration();
            PsiElement prolog = provider.getReferencedProlog();
            if (prolog == null) {
                return null;
            }

            CharSequence localName = getElement().getLocalName().getText();
            PsiElement annotation = prolog.getFirstChild();
            while (annotation != null) {
                if (annotation instanceof XQueryAnnotatedDecl) {
                    XQueryFunctionDecl functionDecl = PsiNavigation.findChildByClass(annotation, XQueryFunctionDecl.class);
                    if (functionDecl != null) {
                        XQueryEQName functionName = PsiNavigation.findChildByClass(functionDecl, XQueryEQName.class);
                        if (functionName != null && functionName.getLocalName().getText().equals(localName)) {
                            // TODO: Check function arity.
                            return functionName;
                        }
                    }<commit_msg>Check function arity when resolving function references.<commit_after>package uk.co.reecedunn.intellij.plugin.xquery.resolve.reference;

import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiReferenceBase;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import uk.co.reecedunn.intellij.plugin.xquery.ast.xquery.*;
import uk.co.reecedunn.intellij.plugin.xquery.psi.PsiNavigation;
import uk.co.reecedunn.intellij.plugin.xquery.psi.XQueryModuleProvider;
import uk.co.reecedunn.intellij.plugin.xquery.psi.XQueryNamespace;

import java.util.List;

public class XQueryFunctionNameReference extends PsiReferenceBase<XQueryEQName> {
    public XQueryFunctionNameReference(XQueryEQName element, TextRange range) {
        super(element, range);
    }

    @Nullable
    @Override
    public PsiElement resolve() {
        XQueryNamespace ns = getElement().resolvePrefixNamespace();
        if (ns == null) {
            return null;
        }

        if (ns.getDeclaration() instanceof XQueryModuleProvider) {
            XQueryModuleProvider provider = (XQueryModuleProvider)ns.getDeclaration();
            PsiElement prolog = provider.getReferencedProlog();
            if (prolog == null) {
                return null;
            }

            PsiElement parent = getElement().getParent();
            int arity = 0;
            if (parent instanceof XQueryFunctionCall) {
                arity = ((XQueryFunctionCall)parent).getArity();
            } else if (parent instanceof XQueryNamedFunctionRef) {
                arity = ((XQueryNamedFunctionRef)parent).getArity();
            }

            CharSequence localName = getElement().getLocalName().getText();
            PsiElement annotation = prolog.getFirstChild();
            while (annotation != null) {
                if (annotation instanceof XQueryAnnotatedDecl) {
                    XQueryFunctionDecl functionDecl = PsiNavigation.findChildByClass(annotation, XQueryFunctionDecl.class);
                    if (functionDecl != null) {
                        XQueryEQName functionName = PsiNavigation.findChildByClass(functionDecl, XQueryEQName.class);
                        if (functionName != null && functionName.getLocalName().getText().equals(localName)) {
                            if (functionDecl.getArity() == arity) {
                                return functionName;
                            }
                        }
                    }
{'fraction_non_alphanumeric': 0.09162132752992383, 'fraction_numerical': 0.0002176278563656148, 'mean_word_length': 2.492401215805471, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '425300', 'n_tokens_mistral': 1218, 'n_tokens_neox': 1173, 'n_words': 216}