Add a DOC type which constructs Doc.

Following Wadler.

prototype/Doubt/Diff.swift

@@ -20,15 +20,13 @@ public enum Diff: Comparable, CustomDebugStringConvertible, CustomDocConvertible
 		}
 	}
 
-	public var doc: Doc {
+	public var doc: DOC {
 		switch self {
 		case .Empty:
 			return .Empty
 		case let .Patch(a, b):
-			return .Horizontal([
-				.Wrap(.Text("{-"), Doc(a), .Text("-}")),
-				.Wrap(.Text("{+"), Doc(b), .Text("+}"))
-			])
+			return .bracket("{-", DOC(a), "-}")
+				<> .bracket("{+", DOC(b), "+}")
 		case let .Copy(a):
 			return a.doc
 		}

prototype/Doubt/Doc.swift

@@ -1,40 +1,148 @@
 public enum Doc: CustomStringConvertible, Equatable {
 	case Empty
-	case Text(String)
-	case Horizontal([Doc])
-	case Vertical([Doc])
-	indirect case Wrap(Doc, Doc, Doc)
-	indirect case Join(Doc, [Doc])
-
-	public init<T>(_ value: T) {
-		self = (value as? CustomDocConvertible)?.doc ?? .Text(String(value))
-	}
+	indirect case Text(String, Doc)
+	indirect case Line(Int, Doc)
 
 	public var description: String {
 		switch self {
 		case .Empty:
 			return ""
-		case let .Text(s):
-			return s
-		case let .Horizontal(a):
-			return a.lazy.map { String($0) }.joinWithSeparator("")
-		case let .Vertical(a):
-			return a.lazy.map { String($0) }.joinWithSeparator("\n")
-		case let .Wrap(open, body, close):
-			return "\(String(open))\(String(body))\(String(close))"
-		case let .Join(separator, elements):
-			return elements.map { String($0) }.joinWithSeparator(String(separator))
+		case let .Text(s, doc):
+			return s + doc.description
+		case let .Line(n, doc):
+			return "\n" + String(count: n, repeatedValue: " " as Character) + doc.description
 		}
 	}
 }
 
+public enum DOC {
+	case Empty
+	indirect case Concat(DOC, DOC)
+	indirect case Union(DOC, DOC)
+	case Text(String)
+	indirect case Nest(Int, DOC)
+	case Line
+
+	public init<T>(_ value: T) {
+		self = (value as? CustomDocConvertible)?.doc ?? .Text(String(value))
+	}
+
+	public static func group(doc: DOC) -> DOC {
+		return Union(doc.flattened, doc)
+	}
+
+	public static func bracket(l: String, _ x: DOC, _ r: String) -> DOC {
+		return group(Concat(Text(l), Concat(Nest(2, Concat(Line, x)), Concat(Line, Text(r)))))
+	}
+
+	public static func folddoc<C: CollectionType where C.Generator.Element == DOC>(docs: C, combine: (DOC, DOC) -> DOC) -> DOC {
+		func fold(docs: Stream<DOC>) -> DOC {
+			switch docs {
+			case .Nil:
+				return .Empty
+			case let .Cons(x, rest) where rest.value.isEmpty:
+				return x
+			case let .Cons(x, rest):
+				return combine(x, fold(rest.value))
+			}
+		}
+		return fold(Stream(sequence: docs))
+	}
+
+	public static func spread<C: CollectionType where C.Generator.Element == DOC>(docs: C) -> DOC {
+		return folddoc(docs, combine: <+>)
+	}
+
+	public static func stack<C: CollectionType where C.Generator.Element == DOC>(docs: C) -> DOC {
+		return folddoc(docs, combine: </>)
+	}
+
+	public static func join<C: CollectionType where C.Generator.Element == DOC>(separator: String, _ docs: C) -> DOC {
+		return folddoc(docs) {
+			$0 <> Text(separator) <+> $1
+		}
+	}
+
+	public var flattened: DOC {
+		switch self {
+		case .Empty, .Text:
+			return self
+		case let .Concat(a, b):
+			return .Concat(a.flattened, b.flattened)
+		case let .Union(a, _):
+			return a.flattened
+		case let .Nest(_, doc):
+			return doc.flattened
+		case .Line:
+			return .Text(" ")
+		}
+	}
+}
+
+func best(w: Int, _ k: Int, _ x: DOC) -> Doc {
+	return be(w, k, Stream<(Int, DOC)>.pure((0, x)))
+}
+
+func be(w: Int, _ k: Int, _ z: Stream<(Int, DOC)>) -> Doc {
+	switch z {
+	case .Nil:
+		return .Empty
+	case let .Cons((_, .Empty), z):
+		return be(w, k, z.value)
+	case let .Cons((i, .Concat(x, y)), z):
+		return be(w, k, .Cons((i, x), Memo(evaluated: .Cons((i, y), z))))
+	case let .Cons((i, .Nest(j, x)), z):
+		return be(w, k, .Cons((i + j, x), z))
+	case let .Cons((_, .Text(s)), z):
+		return .Text(s, be(w, k + Int(s.characters.count), z.value))
+	case let .Cons((i, .Line), z):
+		return .Line(i, be(w, i, z.value))
+	case let .Cons((i, .Union(x, y)), z):
+		return better(w, k, be(w, k, .Cons((i, x), z)), be(w, k, .Cons((i, y), z)))
+	}
+}
+
+func better(w: Int, _ k: Int, _ x: Doc, _ y: Doc) -> Doc {
+	return fits(w - k, x)
+		? x
+		: y
+}
+
+func fits(w: Int, _ x: Doc) -> Bool {
+	guard w >= 0 else { return false }
+	switch x {
+	case .Empty, .Line:
+		return true
+	case let .Text(s, x):
+		return fits(w - Int(s.characters.count), x)
+	}
+}
+
 
 public protocol CustomDocConvertible: CustomStringConvertible {
-	var doc: Doc { get }
+	var doc: DOC { get }
 }
 
 extension CustomDocConvertible {
+	private func pretty(w: Int) -> String {
+		return best(w, 0, doc).description
+	}
+
 	public var description: String {
-		return doc.description
+		return pretty(70)
 	}
 }
+
+
+public func <> (left: DOC, right: DOC) -> DOC {
+	return .Concat(left, right)
+}
+
+
+public func <+> (left: DOC, right: DOC) -> DOC {
+	return left <> .Text(" ") <> right
+}
+
+public func </> (left: DOC, right: DOC) -> DOC {
+	return left <> .Line <> right
+}

prototype/Doubt/Equatable.swift

@@ -46,16 +46,10 @@ public func == (left: Doc, right: Doc) -> Bool {
 	switch (left, right) {
 	case (.Empty, .Empty):
 		return true
-	case let (.Text(a), .Text(b)):
-		return a == b
-	case let (.Horizontal(a), .Horizontal(b)):
-		return a == b
-	case let (.Vertical(a), .Vertical(b)):
-		return a == b
-	case let (.Wrap(a1, b1, c1), .Wrap(a2, b2, c2)):
-		return a1 == a2 && b1 == b2 && c1 == c2
-	case let (.Join(s1, e1), .Join(s2, e2)):
-		return s1 == s2 && e1 == e2
+	case let (.Text(a, x), .Text(b, y)):
+		return a == b && x == y
+	case let (.Line(i, x), .Line(j, y)):
+		return i == j && x == y
 	default:
 		return false
 	}

prototype/Doubt/Syntax.swift

@@ -15,7 +15,7 @@ public enum Term: CustomDebugStringConvertible, CustomDocConvertible, CustomStri
 		}
 	}
 
-	public var doc: Doc {
+	public var doc: DOC {
 		switch self {
 		case .Empty:
 			return .Empty
@@ -182,31 +182,31 @@ public enum Syntax<Payload>: CustomDebugStringConvertible, CustomDocConvertible
 		}
 	}
 
-	public var doc: Doc {
+	public var doc: DOC {
 		switch self {
 		case let .Apply(f, vs):
-			return .Horizontal([
-				Doc(f),
-				.Wrap(.Text("("), .Join(.Text(", "), vs.map(Doc.init)), .Text(")"))
+			return .spread([
+				DOC(f),
+				.bracket("(", .join(", ", vs.map(DOC.init)), ")")
 			])
 		case let .Abstract(parameters, body):
-			return .Horizontal([
+			return .folddoc([
 				.Text("λ"),
-				.Join(.Text(", "), parameters.map(Doc.init)),
+				.join(", ", parameters.map(DOC.init)),
 				.Text("."),
-				.Vertical(body.map(Doc.init))
-			])
+				.stack(body.map(DOC.init))
+			], combine: <>)
 		case let .Assign(n, v):
-			return .Horizontal([ .Text(n), .Text("="), Doc(v) ])
+			return .spread([ .Text(n), .Text("="), DOC(v) ])
 		case let .Variable(n):
 			return .Text(n)
 		case let .Literal(s):
 			return .Text(s)
 		case let .Group(n, vs):
-			return .Horizontal([
-				Doc(n),
-				.Wrap(.Text("{"), .Vertical(vs.map(Doc.init)), .Text("}"))
-			])
+			return .folddoc([
+				DOC(n),
+				.bracket("{", .stack(vs.map(DOC.init)), "}")
+			], combine: <>)
 		}
 	}
 }