From 4234d902a93f50b71974cc02bf7f51b59f47ad9c Mon Sep 17 00:00:00 2001
From: kstich <kevin@kstich.com>
Date: Thu, 16 Nov 2023 09:30:10 -0800
Subject: [PATCH] Fix removal of applied non-prelude meta traits

---
 .../UnreferencedTraitDefinitions.java         | 26 ++++++++++++-
 .../model/neighbor/unreferenced-test.json     | 39 ++++++++++++++++++-
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/UnreferencedTraitDefinitions.java b/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/UnreferencedTraitDefinitions.java
index 1860c599251..be216e006d2 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/UnreferencedTraitDefinitions.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/UnreferencedTraitDefinitions.java
@@ -15,6 +15,10 @@
 
 package software.amazon.smithy.model.neighbor;
 
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -23,9 +27,9 @@
 import software.amazon.smithy.model.loader.Prelude;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.utils.FunctionalUtils;
-import software.amazon.smithy.utils.OptionalUtils;
 
 /**
  * Finds trait definitions that are not connected to a service shape.
@@ -64,10 +68,28 @@ public Set<Shape> compute(Model model) {
                 .map(Shape::getAllTraits)
                 .flatMap(traits -> traits.keySet().stream())
                 .distinct()
-                .flatMap(traitId -> OptionalUtils.stream(model.getShape(traitId)))
+                .flatMap(traitId -> getTraitShapes(model, traitId).stream())
                 .filter(keepFilter)
                 .forEach(unused::remove);
 
         return unused;
     }
+
+    private Collection<Shape> getTraitShapes(Model model, ShapeId traitId) {
+        return getTraitShapes(model, traitId, new HashMap<>()).values();
+    }
+
+    private Map<ShapeId, Shape> getTraitShapes(Model model, ShapeId traitId, Map<ShapeId, Shape> traitShapes) {
+        Optional<Shape> initialTraitShapeOp = model.getShape(traitId);
+        if (initialTraitShapeOp.isPresent()) {
+            Shape initialTraitShape = initialTraitShapeOp.get();
+            traitShapes.put(traitId, initialTraitShape);
+            for (ShapeId metaTraitId : initialTraitShape.getAllTraits().keySet()) {
+                if (!metaTraitId.equals(TraitDefinition.ID) && !traitShapes.containsKey(metaTraitId)) {
+                    traitShapes.putAll(getTraitShapes(model, metaTraitId, traitShapes));
+                }
+            }
+        }
+        return traitShapes;
+    }
 }
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/neighbor/unreferenced-test.json b/smithy-model/src/test/resources/software/amazon/smithy/model/neighbor/unreferenced-test.json
index eb377d2e6dd..3226109ec47 100644
--- a/smithy-model/src/test/resources/software/amazon/smithy/model/neighbor/unreferenced-test.json
+++ b/smithy-model/src/test/resources/software/amazon/smithy/model/neighbor/unreferenced-test.json
@@ -9,7 +9,9 @@
                 }
             },
             "traits": {
-                "smithy.api#trait": {}
+                "smithy.api#trait": {},
+                "ns.foo#meta": {},
+                "ns.foo#threeMeta": {}
             }
         },
         "ns.foo#quux": {
@@ -23,6 +25,41 @@
                 "smithy.api#trait": {}
             }
         },
+        "ns.foo#meta": {
+            "type": "structure",
+            "members": {
+                "member": {
+                    "target": "smithy.api#String"
+                }
+            },
+            "traits": {
+                "smithy.api#trait": {},
+                "ns.foo#tooMeta": {},
+                "ns.foo#threeMeta": {}
+            }
+        },
+        "ns.foo#tooMeta": {
+            "type": "structure",
+            "members": {
+                "member": {
+                    "target": "smithy.api#String"
+                }
+            },
+            "traits": {
+                "smithy.api#trait": {}
+            }
+        },
+        "ns.foo#threeMeta": {
+            "type": "structure",
+            "members": {
+                "member": {
+                    "target": "smithy.api#String"
+                }
+            },
+            "traits": {
+                "smithy.api#trait": {}
+            }
+        },
         "ns.foo#MyService": {
             "type": "service",
             "version": "2017-01-19",