Fix issues with AbstractCodeWriter state stacks

Code writer context inheritance wasn't working in the previous
implementation of the stack based inheritance system. We were able to
resolve context from the current state and parent state (because of
eagerly copying states and flattening them per/state), but we weren't
able to look at parent states that used a CodeSection. Methods of a
CodeSection state can be accessed in named template labels. However,
when another state is pushed after setting a CodeSection value for a
state, code writer wasn't crawling the state stack to see if previous
CodeSections could provide a value for a named label that wasn't
explicitly set as parameter in the context map. This was because
instead of looking up values lazily by iterating over the state stack,
we were making eager copies of different pieces of states. This
matters when using things like for loops in templates that push an
implicit state to capture the current loop iteration state.

This commit changes our approach to now actually use the stack itself
and iterate over states when attempting to resolve context keys,
formatters, iterators, and CodeSection values.

Some tests were added to ensure the behavior of the writer is the
same. For example, when using inheritance based lookups, removing a
context key now means we need to set a context key to an explicit null
value to emulate the previous behavior. If we didn't do this, and we
just called Map#remove(), the context key may have been set by a parent
state, making the remove call do nothing (whereas the previous
implementation would actually take effect).

Author: mtdowling

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

@@ -17,9 +17,12 @@
 
 import java.lang.reflect.Method;
 import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -574,6 +577,12 @@
  */
 public abstract class AbstractCodeWriter<T extends AbstractCodeWriter<T>> {
 
+    // Valid formatter characters that can be registered. Must be sorted for binary search to work.
+    static final char[] VALID_FORMATTER_CHARS = {
+            '!', '%', '&', '*', '+', ',', '-', '.', ';', '=', '@',
+            'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S',
+            'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '_', '`'};
+
     private static final Pattern LINES = Pattern.compile("\\r?\\n");
     private static final Map<Character, BiFunction<Object, String, String>> DEFAULT_FORMATTERS = MapUtils.of(
             'L', (s, i) -> formatLiteral(s),
@@ -603,22 +612,19 @@ public AbstractCodeWriter() {
     /**
      * Copies settings from the given AbstractCodeWriter into this AbstractCodeWriter.
      *
-     * <p>The settings of the {@code other} AbstractCodeWriter will overwrite
-     * both global and state-based settings of this AbstractCodeWriter. Formatters of
-     * the {@code other} AbstractCodeWriter will be merged with the formatters of this
-     * AbstractCodeWriter, and in the case of conflicts, the formatters of the
-     * {@code other} will take precedence.
+     * <p>The settings of the {@code other} AbstractCodeWriter will overwrite both global and state-based settings
+     * of this AbstractCodeWriter.
      *
-     * <p>Stateful settings of the {@code other} AbstractCodeWriter are copied into
-     * the <em>current</em> state of this AbstractCodeWriter. Only the settings of
-     * the top-most state is copied. Other states, and the contents of the
-     * top-most state are not copied.
+     * <p>Stateful settings of the {@code other} AbstractCodeWriter like formatters, interceptors, and context are
+     * flattened and then copied into the <em>current</em> state of this AbstractCodeWriter. Any conflicts between
+     * formatters, interceptors, or context of the current writer are overwritten by the other writer. The stack of
+     * states and the contents written to {@code other} are not copied.
      *
      * <pre>{@code
-     * SimpleCodeWritera = new SimpleCodeWriter();
+     * SimpleCodeWriter a = new SimpleCodeWriter();
      * a.setExpressionStart('#');
      *
-     * SimpleCodeWriterb = new SimpleCodeWriter();
+     * SimpleCodeWriter b = new SimpleCodeWriter();
      * b.copySettingsFrom(a);
      *
      * assert(b.getExpressionStart() == '#');
@@ -634,6 +640,16 @@ public void copySettingsFrom(AbstractCodeWriter<T> other) {
 
         // Copy the current state settings of other into the current state.
         currentState.copyStateFrom(other.currentState);
+
+        // Flatten containers into the current state. This is done in reverse order to ensure that more recent
+        // state changes supersede earlier changes.
+        Iterator<State> reverseOtherStates = other.states.descendingIterator();
+        while (reverseOtherStates.hasNext()) {
+            State otherState = reverseOtherStates.next();
+            currentState.interceptors.addAll(otherState.interceptors);
+            currentState.formatters.putAll(otherState.formatters);
+            currentState.context.putAll(otherState.context);
+        }
     }
 
     /**
@@ -686,7 +702,7 @@ public static String formatLiteral(Object value) {
      */
     @SuppressWarnings("unchecked")
     public T putFormatter(char identifier, BiFunction<Object, String, String> formatFunction) {
-        this.currentState.formatters.get().putFormatter(identifier, formatFunction);
+        this.currentState.putFormatter(identifier, formatFunction);
         return (T) this;
     }
 
@@ -962,9 +978,15 @@ public T popState() {
 
             // Don't attempt to intercept anonymous sections.
             if (!(sectionValue instanceof AnonymousCodeSection)) {
-                for (CodeInterceptor<CodeSection, T> interceptor : popped.interceptors.peek().get(sectionValue)) {
-                    result = interceptSection(popped, interceptor, result);
+                // Ensure the remaining parent interceptors are applied in the order they were inserted.
+                // This is the reverse order used when normally iterating over the states deque.
+                Iterator<State> insertionOrderedStates = states.descendingIterator();
+                while (insertionOrderedStates.hasNext()) {
+                    State state = insertionOrderedStates.next();
+                    result = applyPoppedInterceptors(popped, state, sectionValue, result);
                 }
+                // Now ensure the popped state's interceptors are applied.
+                result = applyPoppedInterceptors(popped, popped, sectionValue, result);
             }
 
             if (popped.isInline) {
@@ -986,6 +1008,13 @@ public T popState() {
         return (T) this;
     }
 
+    private String applyPoppedInterceptors(State popped, State state, CodeSection sectionValue, String result) {
+        for (CodeInterceptor<CodeSection, T> interceptor : state.getInterceptors(sectionValue)) {
+            result = interceptSection(popped, interceptor, result);
+        }
+        return result;
+    }
+
     // This method exists because inlining in popSection is impossible due to needing to mutate a result variable.
     @SuppressWarnings("unchecked")
     private String interceptSection(State popped, CodeInterceptor<CodeSection, T> interceptor, String previous) {
@@ -1070,7 +1099,7 @@ private String interceptSection(State popped, CodeInterceptor<CodeSection, T> in
      */
     @SuppressWarnings("unchecked")
     public T onSection(String sectionName, Consumer<Object> interceptor) {
-        currentState.interceptors.get().putInterceptor(CodeInterceptor.forName(sectionName, (w, p) -> {
+        currentState.putInterceptor(CodeInterceptor.forName(sectionName, (w, p) -> {
             String trimmedContent = removeTrailingNewline(p);
             interceptor.accept(trimmedContent);
         }));
@@ -1098,7 +1127,7 @@ public T onSection(String sectionName, Consumer<Object> interceptor) {
      */
     @SuppressWarnings("unchecked")
     public <S extends CodeSection> T onSection(CodeInterceptor<S, T> interceptor) {
-        currentState.interceptors.get().putInterceptor(interceptor);
+        currentState.putInterceptor(interceptor);
         return (T) this;
     }
 
@@ -1857,7 +1886,7 @@ public T unwrite(Object content, Object... args) {
      */
     @SuppressWarnings("unchecked")
     public T putContext(String key, Object value) {
-        currentState.context.get().put(key, value);
+        currentState.context.put(key, value);
         return (T) this;
     }
 
@@ -1879,40 +1908,45 @@ public T putContext(Map<String, Object> mappings) {
     /**
      * Removes a named key-value pair from the context of the current state.
      *
+     * <p>This method has no effect if the parent state defines the context key value pair.
+     *
      * @param key Key to add to remove from the current context.
      * @return Returns self.
      */
     @SuppressWarnings("unchecked")
     public T removeContext(String key) {
-        if (currentState.context.peek().containsKey(key)) {
-            currentState.context.get().remove(key);
+        if (currentState.context.containsKey(key)) {
+            currentState.context.remove(key);
+        } else {
+            // Parent states might have a value for this context key, so explicitly set it to null in this context.
+            currentState.context.put(key, null);
         }
         return (T) this;
     }
 
     /**
-     * Gets a named contextual key-value pair from the current state.
+     * Gets a named contextual key-value pair from the current state or any parent states.
      *
      * @param key Key to retrieve.
      * @return Returns the associated value or null if not present.
      */
     public Object getContext(String key) {
-        CodeSection section = currentState.sectionValue;
-        Map<String, Object> currentContext = currentState.context.peek();
-        if (currentContext.containsKey(key)) {
-            return currentContext.get(key);
-        } else if (section != null) {
-            Method method = findContextMethod(section, key);
-            if (method != null) {
-                try {
-                    return method.invoke(section);
-                } catch (ReflectiveOperationException e) {
-                    String message = String.format(
-                            "Unable to get context '%s' from a matching method of the current CodeSection: %s %s",
-                            key,
-                            e.getCause() != null ? e.getCause().getMessage() : e.getMessage(),
-                            getDebugInfo());
-                    throw new RuntimeException(message, e);
+        for (State state : states) {
+            if (state.context.containsKey(key)) {
+                return state.context.get(key);
+            } else if (state.sectionValue != null) {
+                Method method = findContextMethod(state.sectionValue, key);
+                if (method != null) {
+                    try {
+                        return method.invoke(state.sectionValue);
+                    } catch (ReflectiveOperationException e) {
+                        String message = String.format(
+                                "Unable to get context '%s' from a matching method of the current CodeSection: %s %s",
+                                key,
+                                e.getCause() != null ? e.getCause().getMessage() : e.getMessage(),
+                                getDebugInfo());
+                        throw new RuntimeException(message, e);
+                    }
                 }
             }
         }
@@ -1984,7 +2018,7 @@ String expandSection(CodeSection section, String previousContent, Consumer<Strin
     // Used only by CodeFormatter to apply formatters.
     @SuppressWarnings("unchecked")
     String applyFormatter(char identifier, Object value) {
-        BiFunction<Object, String, String> f = currentState.formatters.peek().getFormatter(identifier);
+        BiFunction<Object, String, String> f = resolveFormatter(identifier);
         if (f != null) {
             return f.apply(value, getIndentText());
         } else if (identifier == 'C') {
@@ -2009,14 +2043,24 @@ String applyFormatter(char identifier, Object value) {
                 throw new ClassCastException(String.format(
                         "Expected value for 'C' formatter to be an instance of %s or %s, but found %s %s",
                         Runnable.class.getName(), Consumer.class.getName(),
-                        value.getClass().getName(), getDebugInfo()));
+                        value == null ? "null" : value.getClass().getName(), getDebugInfo()));
             }
         } else {
             // Return null if no formatter was found.
             return null;
         }
     }
 
+    BiFunction<Object, String, String> resolveFormatter(char identifier) {
+        for (State state : states) {
+            BiFunction<Object, String, String> result = state.getFormatter(identifier);
+            if (result != null) {
+                return result;
+            }
+        }
+        return null;
+    }
+
     private final class State {
         private final boolean isRoot;
         private String indentText = "    ";
@@ -2030,9 +2074,9 @@ private final class State {
         private boolean needsIndentation;
 
         private CodeSection sectionValue;
-        private CopyOnWriteRef<Map<String, Object>> context;
-        private CopyOnWriteRef<CodeWriterFormatterContainer> formatters;
-        private CopyOnWriteRef<CodeInterceptorContainer<T>> interceptors;
+        private final Map<String, Object> context = new HashMap<>();
+        private final Map<Character, BiFunction<Object, String, String>> formatters = new HashMap<>();
+        private final List<CodeInterceptor<CodeSection, T>> interceptors = new ArrayList<>();
 
         private StringBuilder builder;
 
@@ -2046,11 +2090,7 @@ private final class State {
         State() {
             builder = new StringBuilder();
             isRoot = true;
-            CodeWriterFormatterContainer formatterContainer = new CodeWriterFormatterContainer();
-            DEFAULT_FORMATTERS.forEach(formatterContainer::putFormatter);
-            this.formatters = CopyOnWriteRef.fromOwned(formatterContainer);
-            this.context = CopyOnWriteRef.fromOwned(new HashMap<>());
-            this.interceptors = CopyOnWriteRef.fromOwned(new CodeInterceptorContainer<>());
+            DEFAULT_FORMATTERS.forEach(this::putFormatter);
         }
 
         @SuppressWarnings("CopyConstructorMissesField")
@@ -2060,20 +2100,18 @@ private final class State {
             this.builder = copy.builder;
         }
 
+        // This does not copy context, interceptors, or formatters.
+        // State inheritance relies on stacks of States in an AbstractCodeWriter.
         private void copyStateFrom(State copy) {
             this.newline = copy.newline;
             this.expressionStart = copy.expressionStart;
-            this.context = copy.context;
             this.indentText = copy.indentText;
             this.leadingIndentString = copy.leadingIndentString;
             this.indentation = copy.indentation;
             this.newlinePrefix = copy.newlinePrefix;
             this.trimTrailingSpaces = copy.trimTrailingSpaces;
             this.disableNewline = copy.disableNewline;
             this.needsIndentation = copy.needsIndentation;
-            this.context = CopyOnWriteRef.fromBorrowed(copy.context.peek(), HashMap::new);
-            this.formatters = CopyOnWriteRef.fromBorrowed(copy.formatters.peek(), CodeWriterFormatterContainer::new);
-            this.interceptors = CopyOnWriteRef.fromBorrowed(copy.interceptors.peek(), CodeInterceptorContainer::new);
         }
 
         @Override
@@ -2200,6 +2238,48 @@ private String getSectionName() {
                 return sectionValue.sectionName();
             }
         }
+
+        void putFormatter(Character identifier, BiFunction<Object, String, String> formatFunction) {
+            if (Arrays.binarySearch(VALID_FORMATTER_CHARS, identifier) < 0) {
+                throw new IllegalArgumentException("Invalid formatter identifier: " + identifier);
+            }
+            formatters.put(identifier, formatFunction);
+        }
+
+        BiFunction<Object, String, String> getFormatter(char identifier) {
+            return formatters.get(identifier);
+        }
+
+        @SuppressWarnings("unchecked")
+        void putInterceptor(CodeInterceptor<? extends CodeSection, T> interceptor) {
+            interceptors.add((CodeInterceptor<CodeSection, T>) interceptor);
+        }
+
+        /**
+         * Gets a list of interceptors that match the given type and for which the
+         * result of {@link CodeInterceptor#isIntercepted(CodeSection)} returns true
+         * when given {@code forSection}.
+         *
+         * @param forSection The section that is being intercepted.
+         * @param <S> The type of section being intercepted.
+         * @return Returns the list of matching interceptors.
+         */
+        <S extends CodeSection> List<CodeInterceptor<CodeSection, T>> getInterceptors(S forSection) {
+            // Add in parent interceptors.
+            List<CodeInterceptor<CodeSection, T>> result = new ArrayList<>();
+            // Merge in local interceptors.
+            for (CodeInterceptor<CodeSection, T> interceptor : interceptors) {
+                // Add the interceptor only if it's the right type.
+                if (interceptor.sectionType().isInstance(forSection)) {
+                    // Only add if the filter passes.
+                    if (interceptor.isIntercepted(forSection)) {
+                        result.add(interceptor);
+                    }
+                }
+            }
+
+            return result;
+        }
     }
 
     String removeTrailingNewline(String value) {

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

@@ -635,7 +635,7 @@ private Operation parseNormalArgument() {
             // Parse the formatter and apply it.
             int line = parser.line();
             int column = parser.column();
-            char identifier = parser.expect(CodeWriterFormatterContainer.VALID_FORMATTER_CHARS);
+            char identifier = parser.expect(AbstractCodeWriter.VALID_FORMATTER_CHARS);
 
             // The error message needs to be created here and given to the operation in way that it can
             // throw with an appropriate message.