`fmt`-style pretty printing, but of syntax objects

syntax-format.rkt

#lang racket


;; This module defines a Racket syntax formatter that uses pretty-expressive as
;; the underlying formatting engine. The formatter takes *syntax objects* as
;; input and produces a pretty-expressie document as output. Formatting occurs
;; in steps, similar to macro expansion: each step analyzes a syntax object and
;; returns a *partially formatted* syntax object, which consists of strings
;; interleaved with unformatted syntax objects representing the subforms of the
;; input object. Formatting begins with the outermost syntax object and
;; recursively traverses into unformatted subforms.
;;
;; This doesn't have anywhere close to all of fmt's features. It's just a proof
;; of concept that explores an alternative formatting API that's driven by
;; syntax objects instead of a custom tokenizer and parser. Open questions:
;;
;; - How hard is it to support a wider variety of forms?
;; - What about fixed-width indentation of subforms, like in a function's body?
;; - Rhombus attaches "raw source" to syntax objects via a property, can we use
;;   that in order to recover the exact text of forms? See the "raw source"
;;   property in https://docs.racket-lang.org/rhombus@rhombus/stxobj.html for
;;   details. Matthew mentioned being open to adding this property to objects
;;   produced by the regular Racket reader as well.
;; - What about comments? Pretty important detail, that one.


(require guard
         pretty-expressive
         syntax/parse)

(module+ test
  (require racket/syntax
           rackunit))


;; An unformatted syntax object, to be formatted later by a subsequent pass.
(struct unformatted-syntax

  (;; A string to concatenate before the syntax object once it's formatted.
   prefix
   ;; The syntax object to format later.
   object
   ;; A string to concatenate after the syntax object once it's formatted.
   suffix
   ;; Either 'align, indicating this object needs to be wrapped with (align ...)
   ;; after it's formatted (and concatenated with the prefix and suffix), or
   ;; 'no-align indicating no alignment transformation needs to be made.
   alignment-mode)
  
  #:transparent
  #:guard (struct-guard/c string? syntax? string? (or/c 'align 'no-align)))


;; A partially formatted syntax object, consisting of a series of strings and
;; unformatted subforms. We use strings instead of doc? objects so that equals
;; works correctly on instances of this struct.
(struct partially-formatted-syntax (pieces)
  #:transparent
  #:guard (struct-guard/c (listof (or/c string? unformatted-syntax?))))


;; Fully formats a syntax object into a pretty printable document.
(define/contract (syntax-format stx)
  (-> syntax? doc?)
  (define choices (syntax-format-partially stx))
  (define formatted-choices
    (for/list ([choice (in-list choices)])
      (define doc-pieces
        (for/list ([piece (in-list (partially-formatted-syntax-pieces choice))])
          (match piece
            [(? string?) (string->doc piece)]
            [(unformatted-syntax prefix substx suffix 'no-align)
             (<> (string->doc prefix) (syntax-format substx) (string->doc suffix))]
            [(unformatted-syntax prefix substx suffix 'align)
             (align (<> (string->doc prefix) (syntax-format substx) (string->doc suffix)))])))
      (apply <> doc-pieces)))
  (apply alt formatted-choices))


;; Utility for turning a string that may contain newlines into a doc? consisting
;; of concatenated text and hard newlines.
(define/contract (string->doc s)
  (-> string? doc?)
  (v-concat (map text (string-split s "\n" #:trim? #false))))
             

;; The main dispatch table, which chooses what formatters to use based on the shape
;; of the input syntax object. Returns a list of possible formats to choose from,
;; to be combined later with `alt`.
(define/contract (syntax-format-partially stx)
  (-> syntax? (listof partially-formatted-syntax?))
  (syntax-parse stx
    [(~or atom:id atom:number) (list (format-atom stx))]
    [(form ...) (list (format-one-line-s-exp stx) (format-multi-line-s-exp stx))]
    [_ (list)]))


(define/contract (format-atom stx)
  (-> syntax? partially-formatted-syntax?)
  (syntax-parse stx
    [(~or atom:id atom:number atom:string atom:keyword)
     (partially-formatted-syntax (list (~s (syntax-e #'atom))))]))


(module+ test
  (test-case "format-atom"
    (check-equal? (format-atom #'hello) (partially-formatted-syntax (list "hello")))
    (check-equal? (format-atom #'"hello") (partially-formatted-syntax (list "\"hello\"")))
    (check-equal? (format-atom #'42) (partially-formatted-syntax (list "42")))
    (check-equal? (format-atom #'#:hello) (partially-formatted-syntax (list "#:hello")))))


(define/contract (format-one-line-s-exp stx)
  (-> syntax? partially-formatted-syntax?)
  (define forms (syntax-parse stx [(form ...) (attribute form)]))
  (define unformatted
    (for/list ([form (in-list forms)])
      (unformatted-syntax "" form "" 'no-align)))
  (partially-formatted-syntax
   (add-between unformatted (list " ")
                #:before-first (list "(")
                #:after-last (list ")")
                #:splice? #true)))


(module+ test
  (test-case "format-one-line-s-exp"
    (define stx #'(+ 1 2 3))
    (define/with-syntax (plus one two three) stx)
    (check-equal? (format-one-line-s-exp stx)
                  (partially-formatted-syntax
                   (list "("
                         (unformatted-syntax "" #'plus "" 'no-align)
                         " "
                         (unformatted-syntax "" #'one "" 'no-align)
                         " "
                         (unformatted-syntax "" #'two "" 'no-align)
                         " "
                         (unformatted-syntax "" #'three "" 'no-align)
                         ")")))))


(define/contract (format-multi-line-s-exp stx)
  (-> syntax? partially-formatted-syntax?)
  (guarded-block
    (define forms (syntax-parse stx [(form ...) (attribute form)]))
    (guard-match (cons head-form arg-forms) forms #:else
      (partially-formatted-syntax (list "()")))
    (define unformatted-head (unformatted-syntax "" head-form "" 'no-align))
    (define arg-count (length arg-forms))
    (guard (positive? arg-count) #:else
      (partially-formatted-syntax (list "(" unformatted-head ")")))
    (define unformatted-args
      (for/list ([arg (in-list arg-forms)]
                 [i (in-naturals 1)])
        (define arg-suffix (if (equal? i arg-count) "" "\n"))
        (unformatted-syntax "" arg arg-suffix 'align)))
    (partially-formatted-syntax
     (append (list "(" unformatted-head " ") unformatted-args (list ")")))))


(module+ test
  (test-case "format-multi-line-s-exp"
    (define stx #'(+ 1 2 3))
    (define/with-syntax (plus one two three) stx)
    (check-equal? (format-multi-line-s-exp stx)
                  (partially-formatted-syntax
                   (list "("
                         (unformatted-syntax "" #'plus "" 'no-align)
                         " "
                         (unformatted-syntax "" #'one "\n" 'align)
                         (unformatted-syntax "" #'two "\n" 'align)
                         (unformatted-syntax "" #'three "" 'align)
                         ")")))))


;; Basic integration test of the whole shebang.
(module+ test
  (test-case "syntax-format"
    (define stx #'(+ 1 (+ 2 3 4) 5))
    (check-equal? (pretty-format (syntax-format stx)) "(+ 1 (+ 2 3 4) 5)")
    (check-equal? (pretty-format (syntax-format stx) #:page-width 10)
                  #<<EOS
(+ 1
   (+ 2
      3
      4)
   5)
EOS
                  )
    (check-equal? (pretty-format (syntax-format stx) #:page-width 15)
                  #<<EOS
(+ 1
   (+ 2 3 4)
   5)
EOS
                  )))