;; =============================================================================
;;
;;  pprint.ss - a pretty-printer for PLT Scheme
;;  Copyright (C) 2005 David Herman
;;
;;  Portions based on PPrint.hs - a pretty-printer for Haskell
;;  Copyright 2000, Daan Leijen. All rights reserved.
;;  See COPYING.HASKELL for accompanying license.
;;
;;  Portions based on pprint.m - a pretty-printer for Mercury
;;  Copyright (C) 2000-2002 The University of Melbourne
;;  Written by Ralph Becket
;;  See COPYING for accompanying license.
;;
;;  This library is free software; you can redistribute it and/or modify it
;;  under the terms of the GNU Lesser General Public License as published by
;;  the Free Software Foundation; either version 2.1 of the License, or (at
;;  your option) any later version.
;;
;;  This library is distributed in the hope that it will be useful, but WITHOUT
;;  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
;;  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
;;  License for more details.
;;
;;  You should have received a copy of the GNU Lesser General Public License
;;  along with this library; if not, write to the Free Software Foundation,
;;  Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
;;
;; =============================================================================

;; See doc.txt for information about using this library.
;; See design.txt for information about the design of this library.

(module pprint mzscheme
  (require (planet "datatype.ss" ("dherman" "struct.plt" 2 0)))
  (require (lib "contract.ss"))
  (require (lib "match.ss"))
  (require (all-except (lib "list.ss") empty))
  (require (lib "etc.ss"))

  ;; ===========================================================================
  ;; Primitives
  ;; ===========================================================================

  (define (write-doc v port write?)
    (fprintf port "#<struct:doc>"))

  (define-datatype (doc ([prop:custom-write write-doc]))
    [NIL     ()]
    [CAT     (left right)]
    [NEST    (depth doc)]
    [LABEL   (label doc)]
    [MARKUP  (f doc)]
    [TEXT    (text)]
    [LINE    (break?)]
    [GROUP   (doc)]
    [COLUMN  (f)]
    [NESTING (f)])

  (define (doc->string doc)
    (match doc
      [($ NIL) "NIL"]
      [($ CAT x y) (format "(CAT ~a ~a)" (doc->string x) (doc->string y))]
      [($ NEST n x) (format "(NEST ~a ~a)" n (doc->string x))]
      [($ LABEL l x) (format "(LABEL ~v ~a)" l (doc->string x))]
      [($ MARKUP f x) (format "(MARKUP ~a ~a)" f (doc->string x))]
      [($ LINE break?) (format "(LINE ~a)" break?)]
      [($ GROUP x) (format "(GROUP ~a)" (doc->string x))]
      [($ TEXT t) (format "~v" t)]
      [($ COLUMN f) (format "(COLUMN ~a)" f)]
      [($ NESTING f) (format "(NESTING ~a)" f)]))

  (define-datatype simple-doc
    [SEMPTY ()]
    [STEXT (text rest)]
    [SPUSH (f rest)]
    [SPOP (rest)]
    [SLINE (is rest)])

  (define empty         (make-NIL))
  (define (nest i x)    (make-NEST i x))
  (define (text s)      (make-TEXT s))
  (define (label l d)   (make-LABEL l d))
  (define (markup f d)  (make-MARKUP f d))
  (define (column f)    (make-COLUMN f))
  (define (nesting f)   (make-NESTING f))
  (define (group x)     (make-GROUP x))
  (define (char c)      (if (char=? c #\newline) line (text (string c))))

  (define line          (make-LINE #f))
  (define break         (make-LINE #t))
  (define soft-line     (group line))
  (define soft-break    (group break))

  ;; ===========================================================================
  ;; Semi-primitives
  ;; ===========================================================================

  (define (fill/break f x)
    (width x (lambda (w)
               (if (> w f)
                   (nest f break)
                   (text (spaces (- f w)))))))
  (define (fill f d)
    (width d (lambda (w)
               (if (>= w f)
                   empty
                   (text (spaces (- f w)))))))
  (define (width d f)
    (column (lambda (k1)
              (h-append d (column (lambda (k2)
                              (f (- k2 k1))))))))
  (define (indent i d)
    (hang i (h-append (text (spaces i)) d)))
  (define (hang i d)
    (align (nest i d)))
  (define (align d)
    (column (lambda (k)
              (nesting (lambda (i)
                         (nest (- k i) d))))))

  ;; ===========================================================================
  ;; High-level combinators
  ;; ===========================================================================

  (define comma     (char #\,))
  (define semi      (char #\;))
  (define colon     (char #\:))
  (define lparen    (char #\())
  (define rparen    (char #\)))
  (define lbracket  (char #\[))
  (define rbracket  (char #\]))
  (define langle    (char #\<))
  (define rangle    (char #\>))
  (define lbrace    (char #\{))
  (define rbrace    (char #\}))
  (define space     (char #\space))
  (define ellipsis  (text "..."))
  (define squote    (char #\'))
  (define dquote    (char #\"))
  (define dot       (char #\.))
  (define backslash (char #\\))
  (define equals    (char #\=))

  (define (foldr1 f xs)
    (match xs
      [(x) x]
      [(x . xs) (f x (foldr1 f xs))]))

  (define (fold f ds)
    (if (null? ds)
        empty
        (foldr1 f ds)))

  (define (cat-with sep)
    (letrec ([f (match-lambda
                  [()       empty]
                  [(x)      x]
                  [(x y)    (h-append x sep y)]
                  [(d . ds) (h-append d sep (f ds))])])
      (lambda ds
        (f ds))))

  (define h-append
    (letrec ([f (match-lambda
                  [()       empty]
                  [(x)      x]
                  [(x y)    (make-CAT x y)]
                  [(d . ds) (make-CAT d (f ds))])])
      (lambda ds
        (f ds))))

  (define hs-append     (cat-with space))
  (define v-append      (cat-with line))
  (define vs-append     (cat-with soft-line))
  (define vb-append     (cat-with break))
  (define vsb-append    (cat-with soft-break))

  (define hs-concat     (lambda (ds) (fold hs-append ds)))
  (define v-concat      (lambda (ds) (fold v-append ds)))
  (define vs-concat     (lambda (ds) (fold vs-append ds)))
  (define v-concat/s    (compose group v-concat))

  (define h-concat      (lambda (ds) (fold h-append ds)))
  (define vb-concat     (lambda (ds) (fold vb-append ds)))
  (define vsb-concat    (lambda (ds) (fold vsb-append ds)))
  (define vb-concat/s   (compose group vb-concat))

  (define (next-newline s i)
    (if (or (>= i (string-length s))
            (char=? (string-ref s i) #\newline))
        i
        (next-newline s (add1 i))))

  (define (split-newlines s)
    (let ([len (string-length s)])
      (let f ([start 0])
        (cond
          [(>= start len)
           null]
          [(char=? (string-ref s start) #\newline)
           (cons "\n" (f (add1 start)))]
          [else (let ([end (next-newline s start)])
                  (if (= end len)
                      (list (substring s start))
                      (cons (substring s start end)
                            (cons "\n" (f (add1 end))))))]))))

  (define (string->doc s)
    (foldr h-append
           empty
           (map (lambda (s)
                  (if (string=? "\n" s) line (text s)))
                (split-newlines s))))

  (define (value->doc x)
    (string->doc (format "~a" x)))

  ;; (punctuate p (list d1 d2 ... dn)) => (list (<> d1 p) (<> d2 p) ... dn)

  (define (apply-infix p ds)
    (match ds
      [() null]
      [(d) (list d)]
      [(d . ds) (cons (h-append d p) (apply-infix p ds))]))

  (define (spaces n)
    (build-string n (lambda (i) #\space)))

  (define (extend s n)
    (string-append s (spaces n)))

  ;; flatten : doc -> doc
  (define (flatten doc)
    (match doc
      [($ CAT x y) (make-CAT (flatten x) (flatten y))]
      [($ NEST n x) (flatten x)]
      [($ LABEL l x) (flatten x)]
      [($ MARKUP f x) (make-MARKUP f (flatten x))]
      [($ LINE #t) (make-NIL)]
      [($ LINE #f) (make-TEXT " ")]
      [($ GROUP x) (flatten x)]
      [($ COLUMN f) (make-COLUMN (compose flatten f))]
      [($ NESTING f) (make-NESTING (compose flatten f))]
      [_ doc]))

  ;; NOTE: If you want to add ribbon-width, you'll need an extra accumulator
  ;;       is0, which represents the current line's indentation (as opposed
  ;;       to the nesting level, which represents the indentation of future
  ;;       lines). I don't think I want to bother, though, because it seems
  ;;       like a pretty esoteric feature, and doesn't really make sense in
  ;;       the context of the `label' combinator.

  (define backtrack!
    (let-struct backtrack ()
      (make-backtrack)))

  (define (backtrack? x)
    (eq? x backtrack!))

  (define (too-big? text col width)
    (> (+ col (string-length text)) width))

  ;; layout : nat * doc -> simple-doc
  (define (layout width doc)
    ;; best : nat * (listof (cons string doc)) * boolean -> simple-doc
    (let best ([col 0] [docs (list (cons "" doc))] [alternate? #f])
      (match docs
        [() (make-SEMPTY)]
        [(#f . docs*)
         (make-SPOP (best col docs* alternate?))]
        [((is . ($ NIL)) . docs*)
         (best col docs* alternate?)]
        [((is . ($ CAT x y)) . docs*)
         (best col (cons (cons is x)
                         (cons (cons is y) docs*)) alternate?)]
        [((is . ($ NEST n x)) . docs*)
         (best col (cons (cons (extend is n) x) docs*) alternate?)]
        [((is . ($ LABEL l x)) . docs*)
         (best col (cons (cons (string-append is l) x) docs*) alternate?)]
        [((is . ($ MARKUP f x)) . docs*)
         (make-SPUSH f (best col (cons (cons is x) (cons #f docs*)) alternate?))]
        [((is . ($ LINE _)) . docs*)
         (make-SLINE is (best (string-length is) docs* alternate?))]
        [((is . ($ GROUP x)) . docs*)
         (with-handlers ([backtrack? (lambda (exn)
                                       (best col (cons (cons is x) docs*) alternate?))])
           (best col (cons (cons is (flatten x)) docs*) #t))]
        [((is . ($ TEXT t)) . docs*)
         (if (and alternate? (too-big? t col width))
             (raise backtrack!)
             (make-STEXT t (best (+ col (string-length t)) docs* alternate?)))]
        [((is . ($ COLUMN f)) . docs*)
         (best col (cons (cons is (f col)) docs*) alternate?)]
        [((is . ($ NESTING f)) . docs*)
         (best col (cons (cons is (f (string-length is))) docs*) alternate?)])))

  ;; ===========================================================================
  ;; Front-end and utilities
  ;; ===========================================================================

  (define current-page-width (make-parameter 80))

  (define pretty-print
    (opt-lambda (doc [port (current-output-port)] [width (current-page-width)])
      (let print ([sdoc (layout width doc)])
        (match sdoc
          [($ SEMPTY) (void)]
          [($ STEXT t rest)
           (display t port)
           (print rest)]
          ;; Discard markup for ordinary printing.
          [($ SPUSH f rest)
           (print rest)]
          [($ SPOP rest)
           (print rest)]
          [($ SLINE is rest)
           (newline port)
           (display is port)
           (print rest)]))))

  (define pretty-format
    (opt-lambda (doc [width (current-page-width)])
      (let ([out (open-output-string)])
        (pretty-print doc out width)
        (get-output-string out))))

  ;; This is sorta like the exception monad, except:
  ;;
  ;; a) it accumulates (tupling) return values instead of choosing (summation), and
  ;; b) it generalizes this accumulation to an unbounded stack of return values.
  ;;
  ;; So every time you reach a PUSH node, you expect one more return value,
  ;; and every time you reach a POP node, you return an extra value.
  ;;
  ;; The idea is that we fork a separate space in which to return pretty printed
  ;; text every time we reach a PUSH node, so that we can format just that portion,
  ;; but we still want the rest of the document, which we combine with the marked-up
  ;; text.

  ;; pretty-markup : (doc a) ((union string a) (union string a) -> (union string a)) [nat] -> (union string a)
  (define pretty-markup
    (opt-lambda (doc combine [width (current-page-width)])
      (car
       (let markup ([sdoc (layout width doc)])
         (match sdoc
           [($ SEMPTY) (list "")]
           [($ STEXT t rest)
            (let ([r (markup rest)])
              (cons (combine t (car r))
                    (cdr r)))]
           [($ SPUSH f rest)
            (let ([r (markup rest)])
              (cons (combine (f (car r)) (cadr r))
                    (cddr r)))]
           [($ SPOP rest)
            (cons "" (markup rest))]
           [($ SLINE is rest)
            (let ([r (markup rest)])
              (cons (combine (string-append "\n" is)
                             (car r))
                    (cdr r)))])))))


  ;; Front end:
  (provide/contract [pretty-print ((doc?) (output-port? natural-number/c) . opt-> . any)]
                    [pretty-format ((doc?) (natural-number/c) . opt-> . string?)]
                    [pretty-markup ((doc? (any/c any/c . -> . any)) (natural-number/c) . opt-> . any)]
                    [current-page-width parameter?])
  (provide/contract [doc? (any/c . -> . boolean?)]
                    [string->doc (string? . -> . doc?)]
                    [value->doc  (any/c . -> . doc?)])

  ;; Basic combinators:
  (provide/contract [empty doc?]
                    [char (char? . -> . doc?)]
                    [text (string? . -> . doc?)]
                    [nest (natural-number/c doc? . -> . doc?)]
                    [label (string? doc? . -> . doc?)]
                    [markup (procedure? doc? . -> . doc?)]
                    [group (doc? . -> . doc?)]
                    [line doc?]
                    [break doc?]
                    [soft-line doc?]
                    [soft-break doc?])

  ;; Alignment:
  (provide/contract [align (doc? . -> . doc?)]
                    [hang (natural-number/c doc? . -> . doc?)]
                    [indent (natural-number/c doc? . -> . doc?)])

  ;; Operators:
  (provide/contract [h-append   (() (listof doc?) . ->* . (doc?))]
                    [hs-append  (() (listof doc?) . ->* . (doc?))]
                    [vs-append  (() (listof doc?) . ->* . (doc?))]
                    [vsb-append (() (listof doc?) . ->* . (doc?))]
                    [v-append   (() (listof doc?) . ->* . (doc?))]
                    [vb-append  (() (listof doc?) . ->* . (doc?))])

  ;; List combinators:
  (provide/contract [v-concat/s   ((listof doc?) . -> . doc?)]
                    [vs-concat    ((listof doc?) . -> . doc?)]
                    [hs-concat    ((listof doc?) . -> . doc?)]
                    [v-concat     ((listof doc?) . -> . doc?)])
  (provide/contract [vb-concat/s  ((listof doc?) . -> . doc?)]
                    [vsb-concat   ((listof doc?) . -> . doc?)]
                    [h-concat     ((listof doc?) . -> . doc?)]
                    [vb-concat    ((listof doc?) . -> . doc?)])
  (provide/contract [apply-infix  (doc? (listof doc?) . -> . (listof doc?))])

  ;; Fillers:
  (provide/contract [fill (natural-number/c doc? . -> . doc?)]
                    [fill/break (natural-number/c doc? . -> . doc?)])

  ;; Constants:
  (provide/contract [lparen doc?]
                    [rparen doc?]
                    [lbrace doc?]
                    [rbrace doc?]
                    [lbracket doc?]
                    [rbracket doc?]
                    [langle doc?]
                    [rangle doc?]
                    [squote doc?]
                    [dquote doc?]
                    [semi doc?]
                    [colon doc?]
                    [comma doc?]
                    [space doc?]
                    [dot doc?]
                    [backslash doc?]
                    [equals doc?]
                    [ellipsis doc?]))