public PsiElement getContext(@Nullable String referenceName) {
    PsiClassStub<?> stub = getStub();
    if (stub == null) return getParent();

    // if AST is not loaded, then we only can need context to resolve class references, which specify the name
    if (referenceName == null) return super.getContext();
    
    // class names can be resolved by stubs unless this is a local/anonymous class referencing other local classes
    StubElement<?> parent = stub.getParentStub();
    while (parent != null && !(parent instanceof PsiClassStub) && !(parent instanceof PsiFileStub)) {
      PsiClass[] allLocalClasses = parent.getChildrenByType(JavaStubElementTypes.CLASS, PsiClass.ARRAY_FACTORY);
      if (allLocalClasses.length > 0 && ContainerUtil.exists(allLocalClasses, c -> referenceName.equals(c.getName()))) {
        return getParent();
      }
      if (parent instanceof PsiMethodStub) {
        return parent.getPsi();
      }
      parent = parent.getParentStub();
    }
    return parent != null ? parent.getPsi() : getParent();
  }
public PsiElement getContext(@Nullable String referenceName)