From 5498a180cf4e2ace45d82ae16d4d486dcf7a3760 Mon Sep 17 00:00:00 2001
From: Fahad Zubair <fahadzubair@gmail.com>
Date: Fri, 28 Jun 2024 12:32:19 +0100
Subject: [PATCH] Define a function to return operation shapes that need a
 `ValidationException` (#3720)

Refactor and define a separate function that returns a set of operation
shapes that must have a supported validation exception shape in their
associated errors list.

This helps identify which type of `ValidationException` has been added
to the operation shape's errors list.

Closes: [3722](https://github.com/smithy-lang/smithy-rs/issues/3722)

---------

Co-authored-by: Fahad Zubair <fahadzub@amazon.com>
Co-authored-by: david-perez <d@vidp.dev>
---
 .../smithy/ValidateUnsupportedConstraints.kt  | 29 ++++++++++++++-----
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ValidateUnsupportedConstraints.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ValidateUnsupportedConstraints.kt
index 0d3f50ddca..6fb4379c39 100644
--- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ValidateUnsupportedConstraints.kt
+++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ValidateUnsupportedConstraints.kt
@@ -175,6 +175,26 @@ data class LogMessage(val level: Level, val message: String)
 data class ValidationResult(val shouldAbort: Boolean, val messages: List<LogMessage>) :
     Throwable(message = messages.joinToString("\n") { it.message })
 
+/*
+ * Returns the set of operation shapes that must have a supported validation exception shape
+ * in their associated errors list.
+ */
+fun operationShapesThatMustHaveValidationException(
+    model: Model,
+    service: ServiceShape,
+): Set<OperationShape> {
+    val walker = DirectedWalker(model)
+    return walker.walkShapes(service)
+        .filterIsInstance<OperationShape>()
+        .asSequence()
+        .filter { operationShape ->
+            // Walk the shapes reachable via this operation input.
+            walker.walkShapes(operationShape.inputShape(model))
+                .any { it is SetShape || it is EnumShape || it.hasConstraintTrait() }
+        }
+        .toSet()
+}
+
 /**
  * Validate that all constrained operations have the shape [validationExceptionShapeId] shape attached to their errors.
  */
@@ -189,14 +209,7 @@ fun validateOperationsWithConstrainedInputHaveValidationExceptionAttached(
     //  `disableDefaultValidation` set to `true`, allowing service owners to map from constraint violations to operation errors.
     val walker = DirectedWalker(model)
     val operationsWithConstrainedInputWithoutValidationExceptionSet =
-        walker.walkShapes(service)
-            .filterIsInstance<OperationShape>()
-            .asSequence()
-            .filter { operationShape ->
-                // Walk the shapes reachable via this operation input.
-                walker.walkShapes(operationShape.inputShape(model))
-                    .any { it is SetShape || it is EnumShape || it.hasConstraintTrait() }
-            }
+        operationShapesThatMustHaveValidationException(model, service)
             .filter { !it.errors.contains(validationExceptionShapeId) }
             .map { OperationWithConstrainedInputWithoutValidationException(it) }
             .toSet()