K2: Extract PersistentCheckerContext::copy

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/PersistentCheckerContext.kt

@@ -48,99 +48,34 @@ class PersistentCheckerContext private constructor(
         allErrorsSuppressed = false
     )
 
-    override fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext {
+    override fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext =
-        return PersistentCheckerContext(
+        copy(implicitReceiverStack = implicitReceiverStack.add(name, value))
-            implicitReceiverStack.add(name, value),
-            containingDeclarations,
-            qualifiedAccessOrAssignmentsOrAnnotationCalls,
-            getClassCalls,
-            annotationContainers,
-            isContractBody,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics,
-            allInfosSuppressed,
-            allWarningsSuppressed,
-            allErrorsSuppressed
-        )
-    }
 
-    override fun addDeclaration(declaration: FirDeclaration): PersistentCheckerContext {
+    override fun addDeclaration(declaration: FirDeclaration): PersistentCheckerContext =
-        return PersistentCheckerContext(
+        copy(containingDeclarations = containingDeclarations.add(declaration))
-            implicitReceiverStack,
-            containingDeclarations.add(declaration),
-            qualifiedAccessOrAssignmentsOrAnnotationCalls,
-            getClassCalls,
-            annotationContainers,
-            isContractBody,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics,
-            allInfosSuppressed,
-            allWarningsSuppressed,
-            allErrorsSuppressed
-        )
-    }
 
     override fun dropDeclaration() {
     }
 
-    override fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): PersistentCheckerContext {
+    override fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): PersistentCheckerContext =
-        return PersistentCheckerContext(
+        copy(
-            implicitReceiverStack,
+            qualifiedAccessOrAssignmentsOrAnnotationCalls =
-            containingDeclarations,
+            qualifiedAccessOrAssignmentsOrAnnotationCalls.add(qualifiedAccessOrAnnotationCall)
-            this.qualifiedAccessOrAssignmentsOrAnnotationCalls.add(qualifiedAccessOrAnnotationCall),
-            getClassCalls,
-            annotationContainers,
-            isContractBody,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics,
-            allInfosSuppressed,
-            allWarningsSuppressed,
-            allErrorsSuppressed
         )
-    }
 
     override fun dropQualifiedAccessOrAnnotationCall() {
     }
 
-    override fun addGetClassCall(getClassCall: FirGetClassCall): PersistentCheckerContext {
+    override fun addGetClassCall(getClassCall: FirGetClassCall): PersistentCheckerContext =
-        return PersistentCheckerContext(
+        copy(
-            implicitReceiverStack,
+            getClassCalls = getClassCalls.add(getClassCall),
-            containingDeclarations,
-            qualifiedAccessOrAssignmentsOrAnnotationCalls,
-            getClassCalls.add(getClassCall),
-            annotationContainers,
-            isContractBody,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics,
-            allInfosSuppressed,
-            allWarningsSuppressed,
-            allErrorsSuppressed
         )
-    }
 
     override fun dropGetClassCall() {
     }
 
-    override fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): PersistentCheckerContext {
+    override fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): PersistentCheckerContext =
-        return PersistentCheckerContext(
+        copy(annotationContainers = annotationContainers.add(annotationContainer))
-            implicitReceiverStack,
-            containingDeclarations,
-            qualifiedAccessOrAssignmentsOrAnnotationCalls,
-            getClassCalls,
-            annotationContainers.add(annotationContainer),
-            isContractBody,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics,
-            allInfosSuppressed,
-            allWarningsSuppressed,
-            allErrorsSuppressed
-        )
-    }
 
     override fun dropAnnotationContainer() {
     }
@@ -152,46 +87,40 @@ class PersistentCheckerContext private constructor(
         allErrorsSuppressed: Boolean
     ): CheckerContextForProvider {
         if (diagnosticNames.isEmpty()) return this
+        return copy(
+            suppressedDiagnostics = suppressedDiagnostics.addAll(diagnosticNames),
+            allInfosSuppressed = this.allInfosSuppressed || allInfosSuppressed,
+            allWarningsSuppressed = this.allWarningsSuppressed || allWarningsSuppressed,
+            allErrorsSuppressed = this.allErrorsSuppressed || allErrorsSuppressed
+        )
+    }
+
+    private fun copy(
+        implicitReceiverStack: PersistentImplicitReceiverStack = this.implicitReceiverStack,
+        qualifiedAccessOrAssignmentsOrAnnotationCalls: PersistentList<FirStatement> = this.qualifiedAccessOrAssignmentsOrAnnotationCalls,
+        getClassCalls: PersistentList<FirGetClassCall> = this.getClassCalls,
+        annotationContainers: PersistentList<FirAnnotationContainer> = this.annotationContainers,
+        containingDeclarations: PersistentList<FirDeclaration> = this.containingDeclarations,
+        isContractBody: Boolean = this.isContractBody,
+        allInfosSuppressed: Boolean = this.allInfosSuppressed,
+        allWarningsSuppressed: Boolean = this.allWarningsSuppressed,
+        allErrorsSuppressed: Boolean = this.allErrorsSuppressed,
+        suppressedDiagnostics: PersistentSet<String> = this.suppressedDiagnostics,
+    ): PersistentCheckerContext {
         return PersistentCheckerContext(
-            implicitReceiverStack,
+            implicitReceiverStack, containingDeclarations, qualifiedAccessOrAssignmentsOrAnnotationCalls,
-            containingDeclarations,
+            getClassCalls, annotationContainers, isContractBody, sessionHolder, returnTypeCalculator, suppressedDiagnostics,
-            qualifiedAccessOrAssignmentsOrAnnotationCalls,
+            allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed,
-            getClassCalls,
-            annotationContainers,
-            isContractBody,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics.addAll(diagnosticNames),
-            this.allInfosSuppressed || allInfosSuppressed,
-            this.allWarningsSuppressed || allWarningsSuppressed,
-            this.allErrorsSuppressed || allErrorsSuppressed
         )
     }
 
     private fun toggleContractBody(newValue: Boolean): CheckerContextForProvider {
         check(isContractBody != newValue)
 
-        return PersistentCheckerContext(
+        return copy(isContractBody = newValue)
-            implicitReceiverStack,
-            containingDeclarations,
-            qualifiedAccessOrAssignmentsOrAnnotationCalls,
-            getClassCalls,
-            annotationContainers,
-            isContractBody = newValue,
-            sessionHolder,
-            returnTypeCalculator,
-            suppressedDiagnostics,
-            allInfosSuppressed,
-            allWarningsSuppressed,
-            allErrorsSuppressed
-        )
     }
 
-    override fun enterContractBody(): CheckerContextForProvider {
+    override fun enterContractBody(): CheckerContextForProvider = toggleContractBody(newValue = true)
-        return toggleContractBody(newValue = true)
-    }
 
-    override fun exitContractBody(): CheckerContextForProvider {
+    override fun exitContractBody(): CheckerContextForProvider = toggleContractBody(newValue = false)
-        return toggleContractBody(newValue = false)
-    }
 }