How to implement PsiReferenceContributor (code completion) for a Custom Language to reference function declaration and function call?

I'm following the official documentation to create a psiReferenceContributor for my Firebase rules plugin. But unfortunately, the autocomplete is not working as show in the documenation, all I get is the “No Suggestions” popup. 

----

I have 2 functions declared example1() and example2() . When I ctrl + space after let variable = example  I would like to show a code completion popup with all the list of available functions that match the name - example. Fairly straightforward as way functions behave in most programming languages (Like javascript, java, etc.).

grammar bnf: of my function call and function definition:

functionBlock ::= FUNCTION functionName '(' functionParameterList? ')' functionBody {
 pin = 1
 mixin="com.github.dinbtechit.firebase.language.referenceContributor.function.FirestoreFunctionDeclarationNamedElementImpl"
 implements="com.github.dinbtechit.firebase.language.referenceContributor.function.FirestoreFunctionDeclarationNamedElement"
}

...

functionCall ::= functionName '(' functionArgumentList? ')' {
 mixin = "com.github.dinbtechit.firebase.language.referenceContributor.function.FirestoreFunctionCallNamedElementImpl"
 implements = "com.github.dinbtechit.firebase.language.referenceContributor.function.FirestoreFunctionCallNamedElement"
}

Here is my code:

FunctionDeclarationNamedElement.kt

interface FirestoreFunctionDeclarationNamedElement : PsiNameIdentifierOwner, NavigatablePsiElement

open class FirestoreFunctionDeclarationNamedElementImpl(node: ASTNode): ASTWrapperPsiElement(node),
    FirestoreFunctionDeclarationNamedElement {
    override fun getName(): String? {
        return RulesPsiImplUtil.getFunctionName(this)
    }

    override fun getTextOffset(): Int {
        val element = RulesPsiImplUtil.getNameIdentifier(this)
        return element?.startOffset ?: super.getTextOffset()
    }

    override fun setName(p0: String): PsiElement {
        return RulesPsiImplUtil.setName(this, p0)
    }

    override fun getNameIdentifier(): PsiElement? {
        return RulesPsiImplUtil.getNameIdentifier(this)
    }
}

FunctionCallNamedElement.kt

interface FirestoreFunctionCallNamedElement : PsiNamedElement, NavigatablePsiElement

open class FirestoreFunctionCallNamedElementImpl(node: ASTNode) : ASTWrapperPsiElement(node),
    FirestoreFunctionCallNamedElement {

    override fun getName(): String? {
        return RulesPsiImplUtil.getFunctionName(this)
    }

    override fun setName(p0: String): PsiElement {
        TODO("Not yet implemented")
    }

    override fun getReferences(): Array<PsiReference> {
        val functionCalls =
            PsiTreeUtil.collectElementsOfType(this.containingFile, PsiFirestoreRulesFunctionBlock::class.java)
        val references = mutableListOf<PsiReference>()

        for (functionCall in functionCalls) {
            val textRange = RulesPsiImplUtil.getNameIdentifier(functionCall)?.textRange
            if (textRange != null) {
                val ref = FirestoreFunctionReference(functionCall, textRange)
                references.add(ref)
            }
        }

        return if (references.isNotEmpty()) references.toTypedArray() else PsiReference.EMPTY_ARRAY
        //return PsiReferenceService.getService().getContributedReferences(this) - I get stack overflow for this if this line is uncommented.
    }
}

FirestoreFunctionReference.kt

package com.github.dinbtechit.firebase.language.referenceContributor.function

import com.github.dinbtechit.firebase.psi.RulesPsiImplUtil
import com.github.dinbtechit.firebase.psi.RulesUtil
import com.intellij.codeInsight.lookup.LookupElementBuilder
import com.intellij.icons.AllIcons
import com.intellij.openapi.util.TextRange
import com.intellij.psi.*

class FirestoreFunctionReference(element: PsiElement, textRange: TextRange) :
    PsiReferenceBase<PsiElement>(element, textRange), PsiPolyVariantReference {

    //private var key: String = element.text.substring(textRange.startOffset, textRange.endOffset)

    private var key = ""

    init {
        key = RulesPsiImplUtil.getFunctionName(element) ?: ""
    }

    override fun resolve(): PsiElement? {
        val resolveResult = multiResolve(false)
        return resolveResult.map { it.element }.firstOrNull()
    }

    override fun multiResolve(incompleteCode: Boolean): Array<ResolveResult> {
        val functions = RulesUtil.findFunctionsByElementParent(myElement, key)
        val results = mutableListOf<ResolveResult>()
        functions.forEach {
            results.add(PsiElementResolveResult(it))
        }
        return if (results.isNotEmpty())
            results.toTypedArray()
        else ResolveResult.EMPTY_ARRAY
    }


    override fun getVariants(): Array<Any> {
        val project = myElement.project
        val functions = RulesUtil.findFunctionsByElementParent(myElement)
        val variants = mutableListOf<Any>()
        functions.forEach {
                variants.add(
                    LookupElementBuilder.create(it.functionName?.text!!)
                        .withIcon(AllIcons.Nodes.Function)
                )

        }
        return variants.toTypedArray()
    }
}

RuleFunctionReferenceProvider.kt

class RuleFunctionReferenceProvider : PsiReferenceProvider() {
    override fun getReferencesByElement(element: PsiElement, p1: ProcessingContext): Array<out PsiReference> {
        val functionCalls =
            PsiTreeUtil.collectElementsOfType(element.containingFile, PsiFirestoreRulesFunctionBlock::class.java)
        val references = mutableListOf<PsiReference>()
            for (functionCall in functionCalls) {
                val textRange = RulesPsiImplUtil.getNameIdentifier(functionCall)?.textRange
                if (textRange != null) {
                    val ref = FirestoreFunctionReference(element, textRange)
                    references.add(ref)
                }
            }

        return if (references.isNotEmpty()) references.toTypedArray() else PsiReference.EMPTY_ARRAY
    }
}

and lastly the RuleFunctionReferenceContributor.kt - I have kept the ElementPattern pretty open for now. But still it does not show the code completion popup.

class RuleFunctionReferenceContributor : PsiReferenceContributor() {
    override fun registerReferenceProviders(registrar: PsiReferenceRegistrar) {
        registrar.registerReferenceProvider(
            psiElement(),
            RuleFunctionReferenceProvider()
        )
    }
}

plugin.xml.

<psi.referenceContributor language="FirestoreRules"
                                  implementation="com.github.dinbtechit.firebase.language.referenceContributor.RuleFunctionReferenceContributor"/>

Observations:

- The breakpoint does not hit RuleFunctionReferenceProvider. But when implemented the  getReferecenes() method in  FirestoreFunctionDeclarationNamedElementImpl the breakpoint hit the reference provider however no popup is shown.

- The breakpoints in getVariants() does not get hit at all.

I feel like the whole implementation might be wrong. Any advice on how to apply the mixin and implementation would be very helpful. 

Let me know if you would like to see the full code.