WIP: Add Doc.align that aligns subsequent line breaks to the current line position

Author: DavidGregory084

src/main/java/com/opencastsoftware/prettier4j/Doc.java

@@ -446,10 +446,10 @@ public String text() {
         @Override
         public Doc append(Doc other) {
             // By string concat equivalency law
-            if (other instanceof Text) {
+            /*if (other instanceof Text) {
                 Text otherText = (Text) other;
                 return text(this.text() + otherText.text());
-            } else if (other instanceof Empty) {
+            } else*/ if (other instanceof Empty) {
                 // By left unit law
                 return this;
             }
@@ -849,6 +849,67 @@ public String toString() {
         }
     }
 
+    /**
+     * Represents an aligned {@link Doc}.
+     *
+     * Sets the indentation for line breaks within its inner {@link Doc} at the current line position.
+     */
+    public static class Align extends Doc {
+        private final Doc doc;
+
+        Align(Doc doc) {
+            this.doc = doc;
+        }
+
+        public Doc doc() {
+            return doc;
+        }
+
+        @Override
+        Doc flatten() {
+            return new Align(doc.flatten());
+        }
+
+        @Override
+        boolean hasParams() {
+            return doc.hasParams();
+        }
+
+        @Override
+        boolean hasLineSeparators() {
+            return doc.hasLineSeparators();
+        }
+
+        @Override
+        public Doc bind(String name, Doc value) {
+            return new Align(doc.bind(name, value));
+        }
+
+        @Override
+        public Doc bind(Map<String, Doc> bindings) {
+            return new Align(doc.bind(bindings));
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (o == null || getClass() != o.getClass()) return false;
+            Align align = (Align) o;
+            return Objects.equals(doc, align.doc);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hashCode(doc);
+        }
+
+        @Override
+        public String toString() {
+            return "Align[" +
+                    "doc=" + doc +
+                    ']';
+        }
+    }
+
     /**
      * Represents a line break which cannot be flattened into a more compact layout.
      */
@@ -1615,6 +1676,16 @@ public static Doc indent(int indent, Doc doc) {
         return new Indent(indent, doc);
     }
 
+    /**
+     * Align subsequent lines of the current {@link Doc} to the current position in the line.
+     *
+     * @param doc the input document
+     * @return the aligned document.
+     */
+    public static Doc align(Doc doc) {
+        return new Align(doc);
+    }
+
     /**
      * Apply the margin document {@code margin} to the current {@link Doc}, emitting the
      * margin at the start of every new line from the start of this document until the
@@ -2073,6 +2144,11 @@ private static int layoutEntry(RenderOptions options, Deque<Entry> inQueue, Queu
             Indent indentDoc = (Indent) entryDoc;
             int newIndent = entryIndent + indentDoc.indent();
             inQueue.addFirst(entry(newIndent, entryMargin, indentDoc.doc()));
+        } else if (entryDoc instanceof Align) {
+            // Eliminate Align
+            Align alignDoc = (Align) entryDoc;
+            int newIndent = Math.max(0, position - entryIndent);
+            inQueue.addFirst(entry(newIndent, entryMargin, alignDoc.doc()));
         } else if (entryDoc instanceof Margin) {
             // Eliminate Margin
             Margin marginDoc = (Margin) entryDoc;

src/test/java/com/opencastsoftware/prettier4j/DocTest.java

@@ -21,10 +21,10 @@
 import java.io.StringWriter;
 import java.io.Writer;
 import java.net.URI;
-import java.util.Arrays;
-import java.util.Collections;
+import java.util.*;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static com.opencastsoftware.prettier4j.Doc.*;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -113,6 +113,15 @@ void testAppendLineFlattening() {
         assertThat(actual, is(equalTo(expected)));
     }
 
+    @Test
+    void testAppendLineWithAlign() {
+        String expected = "one two\n    three";
+        String actual = text("one")
+                .appendSpace(group(align(text("two").appendLine(text("three")))))
+                .render(30);
+        assertThat(actual, is(equalTo(expected)));
+    }
+
     @Test
     void testAppendLineOrSpace() {
         String expected = "one two three";
@@ -123,6 +132,15 @@ void testAppendLineOrSpace() {
         assertThat(actual, is(equalTo(expected)));
     }
 
+    @Test
+    void testAppendLineOrSpaceWithAlign() {
+        String expected = "one two three";
+        String actual = text("one")
+                .appendSpace(group(align(text("two").appendLineOrSpace(text("three")))))
+                .render(30);
+        assertThat(actual, is(equalTo(expected)));
+    }
+
     @Test
     void testAppendLineOrSpaceFlattening() {
         String expected = "one\ntwo\nthree";
@@ -133,6 +151,15 @@ void testAppendLineOrSpaceFlattening() {
         assertThat(actual, is(equalTo(expected)));
     }
 
+    @Test
+    void testAppendLineOrSpaceWithAlignFlattening() {
+        String expected = "one two\n    three";
+        String actual = text("one")
+                .appendSpace(group(align(text("two").appendLineOrSpace(text("three")))))
+                .render(10);
+        assertThat(actual, is(equalTo(expected)));
+    }
+
     @Test
     void testAppendLineOrEmpty() {
         String expected = "onetwothree";
@@ -217,6 +244,38 @@ void testBracketFlattening() {
         assertThat(actual, is(equalTo(expected)));
     }
 
+    @Test
+    void testNestedBracketFlattening() {
+        String expectedWidth80 = "let x = functionCall(with, args, nestedFunctionCall(with, more, args))";
+        String expectedWidth40 = "let x = functionCall(\n  with,\n  args,\n  nestedFunctionCall(with, more, args)\n)";
+        String expectedWidth20 = "let x = functionCall(\n  with,\n  args,\n  nestedFunctionCall(\n    with,\n    more,\n    args\n  )\n)";
+
+        Doc inputDoc = text("let")
+            .appendSpace(text("x"))
+            .appendSpace(text("="))
+            .appendSpace(text("functionCall")
+                .append(
+                    intersperse(
+                        text(",").append(lineOrSpace()),
+                        Stream.concat(
+                            Stream.of("with", "args").map(Doc::text),
+                            Stream.of(text("nestedFunctionCall")
+                                .append(
+                                    intersperse(
+                                        text(",").append(lineOrSpace()),
+                                        Stream.of("with", "more", "args").map(Doc::text)
+                                    ).bracket(2, lineOrEmpty(), text("("), text(")"))
+                                ))
+                        )
+                    ).bracket(2, lineOrEmpty(), text("("), text(")"))
+                )
+            );
+
+        assertThat(inputDoc.render(80), is(equalTo(expectedWidth80)));
+        assertThat(inputDoc.render(40), is(equalTo(expectedWidth40)));
+        assertThat(inputDoc.render(20), is(equalTo(expectedWidth20)));
+    }
+
     @Test
     void testMarginWithLineSeparator() {
         assertThrows(IllegalArgumentException.class, () -> {