PsiElementVisitor (@NotNull ProblemsHolder holder, boolean isOnTheFly) { if (!HighlightingFeature.ENHANCED_SWITCH.isAvailable(holder.getFile())) { return PsiElementVisitor.EMPTY_VISITOR; } return new JavaElementVisitor() { @Override public void visitSwitchLabeledRuleStatement(@NotNull PsiSwitchLabeledRuleStatement statement) { super.visitSwitchLabeledRuleStatement(statement); PsiStatement body = statement.getBody(); if (body instanceof PsiBlockStatement) { PsiCodeBlock codeBlock = ((PsiBlockStatement)body).getCodeBlock(); PsiStatement bodyStatement = getSingleStatement(codeBlock); if (bodyStatement instanceof PsiYieldStatement yieldStatement) { if (yieldStatement.getExpression() != null && PsiTreeUtil.getParentOfType(yieldStatement, PsiSwitchBlock.class) instanceof PsiSwitchExpression) { registerProblem(bodyStatement.getFirstChild()); } } else if (bodyStatement instanceof PsiThrowStatement || bodyStatement instanceof PsiExpressionStatement) { registerProblem(codeBlock.getLBrace()); if (isOnTheFly) registerProblem(codeBlock.getRBrace()); } } } private void registerProblem(@Nullable PsiElement element) { if (element != null) { holder.registerProblem(element, message("inspection.labeled.switch.rule.redundant.code.block.message"), new UnwrapCodeBlockFix()); } } }; }