(module pretty-helper mzscheme
  (require (lib "class.ss")
           "partition.ss")
  (provide (all-defined))
  
  ;; Problem: If stx1 and stx2 are two distinguishable syntax objects, it
  ;; still may be the case that (syntax-e stx1) and (syntax-e stx2) are 
  ;; indistinguishable.
  
  ;; Solution: Rather than map stx to (syntax-e stx), in the cases where
  ;; (syntax-e stx) is confusable, map it to a different, unique, value.
  ;;   - stx is identifier : map it to an uninterned symbol w/ same rep
  ;;     (Symbols are useful: see pretty-print's style table)
  ;;   - else : map it to a syntax-dummy object

  ;; NOTE: Nulls are only wrapped when *not* list-terminators.  
  ;; If they were always wrapped, the pretty-printer would screw up
  ;; list printing (I think).
  
  (define-struct syntax-dummy (val))
  
  ;; A SuffixOption is one of
  ;; - 'never             -- never
  ;; - 'always            -- suffix > 0
  ;; - 'over-limit        -- suffix > limit
  ;; - 'all-if-over-limit -- suffix > 0 if any over limit
  
  ;; syntax->datum/tables : stx [partition% num SuffixOption]
  ;;                        -> (values s-expr hashtable hashtable)
  ;; When partition is not false, tracks the partititions that subterms belong to
  ;; When limit is a number, restarts processing with numbering? set to true
  ;; 
  ;; Returns three values:
  ;;   - an S-expression
  ;;   - a hashtable mapping S-expressions to syntax objects
  ;;   - a hashtable mapping syntax objects to S-expressions
  ;; Syntax objects which are eq? will map to same flat values
  (define syntax->datum/tables
    (case-lambda
      [(stx) (table stx #f #f 'never)]
      [(stx partition limit suffixopt) (table stx partition limit suffixopt)]))
  
  ;; table : syntax maybe-partition% maybe-num SuffixOption -> (values s-expr hashtable hashtable)
  (define (table stx partition limit suffixopt)
    (define (make-identifier-proxy id)
      (case suffixopt
        ((never) (unintern (syntax-e id)))
        ((always)
         (let ([n (send partition get-partition id)])
           (if (zero? n) (unintern (syntax-e id)) (suffix (syntax-e id) n))))
        ((over-limit)
         (let ([n (send partition get-partition id)])
           (if (<= n limit)
               (unintern (syntax-e id))
               (suffix (syntax-e id) n))))))
    
    (let/ec escape
      (let ([flat=>stx (make-hash-table)]
            [stx=>flat (make-hash-table)])
        (define (loop obj)
          (cond [(hash-table-get stx=>flat obj (lambda _ #f))
                 => (lambda (datum) datum)]
                [(and partition (identifier? obj))
                 (when (and (eq? suffixopt 'all-if-over-limit)
                            (> (send partition count) limit))
                   (call-with-values (lambda () (table stx partition #f 'always))
                                     escape))
                 (let ([lp-datum (make-identifier-proxy obj)])
                   (hash-table-put! flat=>stx lp-datum obj)
                   (hash-table-put! stx=>flat obj lp-datum)
                   lp-datum)]
                [(syntax? obj)
                 (void (send partition get-partition obj))
                 (let ([lp-datum (loop (syntax-e obj))])
                   (hash-table-put! flat=>stx lp-datum obj)
                   (hash-table-put! stx=>flat obj lp-datum)
                   lp-datum)]
                [(pair? obj)
                 (pairloop obj)]
                [(symbol? obj)
                 (unintern obj)]
                [(null? obj)
                 (make-syntax-dummy obj)]
                [(boolean? obj)
                 (make-syntax-dummy obj)]
                [(number? obj)
                 (make-syntax-dummy obj)]
                [(keyword? obj)
                 (make-syntax-dummy obj)]
                [(vector? obj) 
                 (list->vector (map loop (vector->list obj)))]
                [(box? obj)
                 (box (loop (unbox obj)))]
                [else obj]))
        (define (pairloop obj)
          (cond [(pair? obj)
                 (cons (loop (car obj))
                       (pairloop (cdr obj)))]
                [(null? obj)
                 null]
                [(and (syntax? obj) (null? (syntax-e obj)))
                 null]
                [else (loop obj)]))
        (values (loop stx)
                flat=>stx
                stx=>flat))))

  (define (unintern sym)
    (string->uninterned-symbol (symbol->string sym)))
  
  (define (suffix sym n)
    (string->uninterned-symbol (format "~a:~a" sym n)))

  )