diff --git a/harness.py b/harness.py index f7b5ce8..aca7735 100644 --- a/harness.py +++ b/harness.py @@ -558,8 +558,8 @@ class Harness: self.format_document(lines, doc.resolve(), indent) case wadler.Cons(): - self.format_document(lines, doc.left, indent) - self.format_document(lines, doc.right, indent) + for child in doc.docs: + self.format_document(lines, child, indent) case wadler.Marker(): append("Marker") diff --git a/parser/wadler.py b/parser/wadler.py index 9278deb..925c116 100644 --- a/parser/wadler.py +++ b/parser/wadler.py @@ -13,8 +13,7 @@ from . import runtime @dataclasses.dataclass(frozen=True) class Cons: - left: "Document" - right: "Document" + docs: list["Document"] @dataclasses.dataclass(frozen=True) @@ -73,14 +72,19 @@ Document = None | Text | Literal | NewLine | ForceBreak | Cons | Indent | Group def cons(*documents: Document) -> Document: - result = None - for doc in documents: - if result is None: - result = doc - elif doc is not None: - result = Cons(result, doc) + if len(documents) == 0: + return None - return result + result = [] + for document in documents: + if isinstance(document, Cons): + result.extend(document.docs) + elif document is not None: + result.append(document) + + if len(result) == 0: + return None + return Cons(result) def group(document: Document) -> Document: @@ -174,9 +178,8 @@ def layout_document(doc: Document, width: int, indent: str) -> DocumentLayout: # line and yes, whatever you were asking about has fit. return not chunk.flat - case Cons(left, right): - stack.append(chunk.with_document(right)) - stack.append(chunk.with_document(left)) + case Cons(docs): + stack.extend(chunk.with_document(doc) for doc in reversed(docs)) case Lazy(): stack.append(chunk.with_document(chunk.doc.resolve())) @@ -236,9 +239,8 @@ def layout_document(doc: Document, width: int, indent: str) -> DocumentLayout: output.append("\n" + (chunk.indent * indent)) column = chunk.indent * len(indent) - case Cons(left, right): - chunks.append(chunk.with_document(right)) - chunks.append(chunk.with_document(left)) + case Cons(docs): + chunks.extend(chunk.with_document(doc) for doc in reversed(docs)) case Indent(amount, doc): chunks.append(chunk.with_document(doc, and_indent=amount)) @@ -264,20 +266,25 @@ def layout_document(doc: Document, width: int, indent: str) -> DocumentLayout: def resolve_document(doc: Document) -> Document: match doc: - case Cons(left, right): - lr = resolve_document(left) - rr = resolve_document(right) - if lr is not left or rr is not right: - return cons(lr, rr) - else: - return doc + case Cons(docs): + docs = [resolve_document(d) for d in docs] + return cons(*docs) case Lazy(_): return resolve_document(doc.resolve()) - case _: + case Group(doc): + return group(resolve_document(doc)) + + case Marker(child, meta): + return Marker(resolve_document(child), meta) + + case Text() | Literal() | NewLine() | ForceBreak() | Indent() | None: return doc + case _: + typing.assert_never(doc) + def child_to_name(child: runtime.Tree | runtime.TokenValue) -> str: if isinstance(child, runtime.Tree): diff --git a/tests/test_wadler.py b/tests/test_wadler.py index a5081de..bfd8d90 100644 --- a/tests/test_wadler.py +++ b/tests/test_wadler.py @@ -11,8 +11,6 @@ from parser.parser import ( alt, indent, seq, - Rule, - Assoc, sp, nl, br, @@ -129,9 +127,14 @@ def flatten_document(doc: wadler.Document, src: str) -> list: case wadler.Lazy(): return flatten_document(doc.resolve(), src) case wadler.Cons(): - return flatten_document(doc.left, src) + flatten_document(doc.right, src) + result = [] + for d in doc.docs: + result += flatten_document(d, src) + return result case None: return [] + case wadler.Marker(): + return [f"", flatten_document(doc.child, src)] case _: typing.assert_never(doc)