Add stack trace comment support to code writer

smithy-utils/src/main/java/software/amazon/smithy/utils/AbstractCodeWriter.java

@@ -574,6 +574,7 @@ public abstract class AbstractCodeWriter<T extends AbstractCodeWriter<T>> {
     private State currentState;
     private boolean trailingNewline = true;
     private int trimBlankLines = -1;
+    private boolean enableStackTraceComments;
 
     /**
      * Creates a new SimpleCodeWriter that uses "\n" for a newline, four spaces
@@ -620,6 +621,7 @@ public void copySettingsFrom(AbstractCodeWriter<T> other) {
         // Copy global settings.
         trailingNewline = other.trailingNewline;
         trimBlankLines = other.trimBlankLines;
+        enableStackTraceComments = other.enableStackTraceComments;
 
         // Copy the current state settings of other into the current state.
         currentState.copyStateFrom(other.currentState);
@@ -1540,28 +1542,94 @@ public T closeBlock(String textAfterNewline, Object... args) {
      * <p>Indentation and the newline prefix is only prepended if the writer's
      * cursor is at the beginning of a newline.
      *
+     * <p>Stack trace comments are written along with the given content if
+     * {@link #enableStackTraceComments(boolean)} was called with {@code true}.
+     *
      * @param content Content to write.
      * @return Returns self.
      */
     @SuppressWarnings("unchecked")
     public T writeWithNoFormatting(Object content) {
-        currentState.writeLine(content.toString());
+        currentState.writeLine(findAndFormatStackTraceElement(content.toString(), false));
         return (T) this;
     }
 
+    private String findAndFormatStackTraceElement(String content, boolean inline) {
+        if (enableStackTraceComments) {
+            for (StackTraceElement e : Thread.currentThread().getStackTrace()) {
+                if (isStackTraceRelevant(e)) {
+                    return formatWithStackTraceElement(content, e, inline);
+                }
+            }
+        }
+
+        return content;
+    }
+
+    /**
+     * Tests if the given {@code StackTraceElement} is relevant for a comment
+     * used when writing debug information before calls to write.
+     *
+     * <p>The default implementation filters out all methods in "java.*",
+     * AbstractCodeWriter, software.amazon.smithy.utils.SymbolWriter,
+     * SimpleCodeWriter, and methods of the implementing subclass of
+     * AbstractCodeWriter. This method can be overridden to further filter
+     * stack frames as needed.
+     *
+     * @param e StackTraceElement to test.
+     * @return Returns true if this element should be in a comment.
+     */
+    protected boolean isStackTraceRelevant(StackTraceElement e) {
+        String normalized = e.getClassName().replace("$", ".");
+        return !normalized.startsWith("java.")
+               // Ignore writes made by AbstractCodeWriter or AbstractCodeWriter$State.
+               && !normalized.startsWith(AbstractCodeWriter.class.getCanonicalName())
+               // Ignore writes made by subclasses of this class.
+               && !normalized.startsWith(getClass().getCanonicalName())
+               // Ignore writes made by SimpleCodeWriter.
+               && !normalized.equals(SimpleCodeWriter.class.getCanonicalName())
+               // Ignore any writes made by the well-known SymbolWriter from smithy-codegen-core.
+               && !normalized.equals("software.amazon.smithy.utils.SymbolWriter");
+    }
+
+    /**
+     * Formats content for the given stack frame.
+     *
+     * <p>Subclasses are expected to override this method as needed to handle
+     * language-specific comment requirements. By default, this class will use
+     * C/Java style "traditional" comments that come on the same line before
+     * both calls to writeInline and calls to write with a newline
+     * {@see https://docs.oracle.com/javase/specs/jls/se18/html/jls-3.html#jls-3.7}.
+     *
+     * <p>Programming languages that do not support inline comments should return
+     * the given {@code content} string as-is when {@code writingInline} is set
+     * to {@code true}.
+     *
+     * @param content The content about to be written.
+     * @param element The {@code StackFrameElement} to format.
+     * @param inline Set to true when this is a comment intended to appear before inline content.
+     * @return Returns the formatted content that includes a leading comment.
+     */
+    protected String formatWithStackTraceElement(String content, StackTraceElement element, boolean inline) {
+        return "/* " + element + " */ " + content;
+    }
+
     /**
      * Writes inline text to the AbstractCodeWriter with no formatting.
      *
      * <p>The provided text does not use any kind of expression formatting.
      * Indentation and the newline prefix is only prepended if the writer's
      * cursor is at the beginning of a newline.
      *
+     * <p>Stack trace comments are written along with the given content if
+     * {@link #enableStackTraceComments(boolean)} was called with {@code true}.
+     *
      * @param content Inline content to write.
      * @return Returns self.
      */
     @SuppressWarnings("unchecked")
     public final T writeInlineWithNoFormatting(Object content) {
-        currentState.write(content.toString());
+        currentState.write(findAndFormatStackTraceElement(content.toString(), true));
         return (T) this;
     }
 
@@ -1571,7 +1639,7 @@ public final T writeInlineWithNoFormatting(Object content) {
      *
      * <p>Important: if the formatters that are executed while formatting the
      * given {@code content} string mutate the AbstractCodeWriter, it could leave the
-     * SimpleCodeWriterin an inconsistent state. For example, some AbstractCodeWriter
+     * SimpleCodeWriter in an inconsistent state. For example, some AbstractCodeWriter
      * implementations manage imports and dependencies automatically based on
      * code that is referenced by formatters. If such an expression is used
      * with this format method but the returned String is never written to the
@@ -1643,15 +1711,16 @@ public T call(Runnable task) {
      * <p>Indentation and the newline prefix is only prepended if the writer's
      * cursor is at the beginning of a newline.
      *
+     * <p>If a subclass overrides this method, it <em>should</em> first
+     * perform formatting and then delegate to {@link #writeWithNoFormatting}
+     * to perform the actual write.
+     *
      * @param content Content to write.
      * @param args String arguments to use for formatting.
      * @return Returns self.
      */
-    @SuppressWarnings("unchecked")
     public T write(Object content, Object... args) {
-        String value = format(content, args);
-        currentState.writeLine(value);
-        return (T) this;
+        return writeWithNoFormatting(format(content, args));
     }
 
     /**
@@ -1665,15 +1734,16 @@ public T write(Object content, Object... args) {
      *
      * <p>If newlines are present in the given string, each of those lines will receive proper indentation.
      *
+     * <p>If a subclass overrides this method, it <em>should</em> first
+     * perform formatting and then delegate to {@link #writeInlineWithNoFormatting}
+     * to perform the actual write.
+     *
      * @param content Content to write.
      * @param args String arguments to use for formatting.
      * @return Returns self.
      */
-    @SuppressWarnings("unchecked")
     public T writeInline(Object content, Object... args) {
-        String value = format(content, args);
-        currentState.write(value);
-        return (T) this;
+        return writeInlineWithNoFormatting(format(content, args));
     }
 
     /**
@@ -1844,6 +1914,25 @@ public <C> C getContext(String key, Class<C> type) {
         }
     }
 
+    /**
+     * Enable or disable writing stack trace comments before each call to
+     * {@link #write}, {@link #writeWithNoFormatting}, {@link #writeInline},
+     * and {@link #writeInlineWithNoFormatting}.
+     *
+     * <p>It's sometimes useful to know where in a code generator a line of code
+     * generated text came from. Enabling stack trace comments will output
+     * the last relevant stack trace information caused text to appear in the
+     * code writer's output.
+     *
+     * @param enableStackTraceComments Set to true to enable stack trace comments.
+     * @return Returns self.
+     */
+    @SuppressWarnings("unchecked")
+    public final T enableStackTraceComments(boolean enableStackTraceComments) {
+        this.enableStackTraceComments = enableStackTraceComments;
+        return (T) this;
+    }
+
     String expandSection(CodeSection section, String previousContent, Consumer<String> consumer) {
         StringBuilder buffer = new StringBuilder();
         pushState(section);

smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java

@@ -17,7 +17,9 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.endsWith;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.startsWith;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -1358,4 +1360,91 @@ public void handlesNestedBlockAlignment() {
         assertThat(actual, equalTo("Hello: . * Hi\n"
                                    + "         * There"));
     }
+
+    @Test
+    public void writesStackTraceInfo() {
+        MyCustomWriter writer = new MyCustomWriter();
+        writer.enableStackTraceComments(true);
+
+        writer.writeWithNoFormatting("Hello 1");
+        writer.write("Hello 2");
+        writer.indent().write("Hello 3");
+        writer.writeInline("Hello 4");
+        writer.writeInline("Hello 5");
+        String[] result = writer.toString().split("\n");
+
+        assertThat(result[0], startsWith("/* software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[0], endsWith("*/ Hello 1"));
+        assertThat(result[1], startsWith("/* software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[1], endsWith("*/ Hello 2"));
+        assertThat(result[2], startsWith("    /* software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[2], endsWith("*/ Hello 3"));
+        assertThat(result[3], startsWith("    /* software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[3], containsString("Hello 4"));
+        assertThat(result[3], endsWith("*/ Hello 5"));
+    }
+
+    @Test
+    public void writesStackTraceInfoIgnoringInlineWrites() {
+        // This writer ignores inline writes and puts the comment on the line before.
+        MyCustomPythonWriter writer = new MyCustomPythonWriter();
+        writer.enableStackTraceComments(true);
+
+        writer.writeWithNoFormatting("Hello 1");
+        writer.write("Hello 2");
+        writer.indent().write("Hello 3");
+        writer.writeInline("Hello 4|");
+        writer.writeInline("Hello 5");
+        String[] result = writer.toString().split("\n");
+
+        assertThat(result[0], startsWith("# software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[1], equalTo("Hello 1"));
+        assertThat(result[2], startsWith("# software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[3], equalTo("Hello 2"));
+        assertThat(result[4], startsWith("    # software.amazon.smithy.utils.CodeWriterTest.writesStackTraceInfo"));
+        assertThat(result[5], equalTo("    Hello 3"));
+        assertThat(result[6], equalTo("    Hello 4|Hello 5"));
+    }
+
+    @Test
+    public void filteringAllStackFramesEmitsNoStackComment() {
+        // This writer ignores inline writes and puts the comment on the line before.
+        MyCustomFilteredWriter writer = new MyCustomFilteredWriter();
+        writer.enableStackTraceComments(true);
+        writer.write("Hello");
+
+        assertThat(writer.toString(), equalTo("Hello\n"));
+    }
+
+    private static final class MyCustomWriter extends AbstractCodeWriter<MyCustomWriter> {
+        // Ensure that subclass methods are automatically filtered out as irrelevant frames.
+        @Override
+        public MyCustomWriter write(Object content, Object... args) {
+            return super.write(content, args);
+        }
+
+        // Ensure that subclass methods are automatically filtered out as irrelevant frames.
+        @Override
+        public MyCustomWriter writeInline(Object content, Object... args) {
+            return super.writeInline(content, args);
+        }
+    }
+
+    private static final class MyCustomPythonWriter extends AbstractCodeWriter<MyCustomPythonWriter> {
+        @Override
+        protected String formatWithStackTraceElement(String content, StackTraceElement element, boolean inline) {
+            if (inline) {
+                return content;
+            } else {
+                return "# " + element + getNewline() + content;
+            }
+        }
+    }
+
+    private static final class MyCustomFilteredWriter extends AbstractCodeWriter<MyCustomFilteredWriter> {
+        @Override
+        protected boolean isStackTraceRelevant(StackTraceElement e) {
+            return false;
+        }
+    }
 }