improve documentation truncation behavior

Author: kuhe

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/documentation/StructureExampleGenerator.java

@@ -18,6 +18,7 @@
 import software.amazon.smithy.model.shapes.ListShape;
 import software.amazon.smithy.model.shapes.MapShape;
 import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.SetShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.shapes.UnionShape;
@@ -104,9 +105,10 @@ private static void shape(Shape shape,
             target = shape;
         }
 
-        shapeTracker.mark(shape, indentation);
-        if (shapeTracker.getOccurrenceDepths(shape) > 2) {
-            append(indentation, buffer, "\"<" + shape.getId().getName() + ">\",\n");
+        shapeTracker.mark(target, indentation);
+        if (shapeTracker.shouldTruncate(target)) {
+            append(indentation, buffer, "\"<" + target.getId().getName() + ">\",");
+            checkRequired(indentation, buffer, shape);
         } else {
             switch (target.getType()) {
                 case BIG_DECIMAL:
@@ -261,23 +263,41 @@ private static void append(int indentation, StringBuilder buffer, String tail) {
      * This handles the case of recursive shapes.
      */
     private static class ShapeTracker {
-        private Map<Shape, Set<Integer>> data = new HashMap<Shape, Set<Integer>>();
+        private Map<Shape, Set<Integer>> depths = new HashMap<Shape, Set<Integer>>();
+        private Map<Shape, Integer> occurrences = new HashMap<Shape, Integer>();
 
         /**
          * Mark that a shape is observed at depth.
          */
         public void mark(Shape shape, int depth) {
-            if (!data.containsKey(shape)) {
-                data.put(shape, new HashSet<>());
+            if (!depths.containsKey(shape)) {
+                depths.put(shape, new HashSet<>());
             }
-            data.get(shape).add(depth);
+            depths.get(shape).add(depth);
+            occurrences.put(shape, occurrences.getOrDefault(shape, 0) + 1);
+        }
+
+        /**
+         * @return whether the shape should be truncated.
+         */
+        public boolean shouldTruncate(Shape shape) {
+            return (shape instanceof MapShape || shape instanceof UnionShape || shape instanceof StructureShape
+                    || shape instanceof ListShape || shape instanceof SetShape)
+                    && (getOccurrenceCount(shape) > 2 || getOccurrenceDepths(shape) > 2);
         }
 
         /**
          * @return the number of distinct depths in which the shape appears.
          */
         public int getOccurrenceDepths(Shape shape) {
-            return data.getOrDefault(shape, Collections.emptySet()).size();
+            return depths.getOrDefault(shape, Collections.emptySet()).size();
+        }
+
+        /**
+         * @return total appearances of the shape.
+         */
+        public int getOccurrenceCount(Shape shape) {
+            return occurrences.getOrDefault(shape, 0);
         }
     }
 }