; Scheme 9 from Empty Space, Function Library
; By Nils M Holm, 2009-2012
; Placed in the Public Domain
;
; (draw-tree object)  ==>  unspecific
;
; Print a tree structure resembling a Scheme datum. Each cons
; cell is represented by [o|o] with lines leading to their car
; and cdr parts. Conses with a cdr value of () are represented
; by [o|/].
;
; (Example): (draw-tree '((a) (b . c) (d e)))  ==>  unspecific
;
;            Output:  [o|o]---[o|o]---[o|/]
;                      |       |       |
;                     [o|/]    |      [o|o]---[o|/]
;                      |       |       |       |
;                      a       |       d       e
;                              |
;                             [o|o]--- c
;                              |
;                              b

#| ;treat vectors as "cells"
;NOTE min/max-width of 3 & 7
#((X ...) (Y) 1 (Z ...) 2)
      =>
[   o       o       1        o       2     ]
    |       |                |
    |       |               [o|o]---...
    |       |                |
    |       |                Z
    |       |
    |      [o|/]
    |       |
    |       Y
    |
   [o|o]---...
    |
    X
|#

(module (s9fes draw-tree)

(;export
  ;atom->string/for
  draw-tree)

(import scheme utf8) ;->string
(import (chicken base))
(import (chicken fixnum))
(import (chicken type))
(import (srfi 1))
(import utf8-srfi-13)

(define-type atom (not pair))

(: atom->string/for ((* -> string) --> (atom -> string)))

(: draw-tree (* #!optional fixnum (atom -> string) -> void))

;;

(define-constant MIN-WIDTH 3)

(define NOTHING-TAG '(N))
(define VISITED-TAG '(V))

(define-inline (empty? x) (eq? x NOTHING-TAG))
(define-inline (visited? x) (eq? (car x) VISITED-TAG))

(define-inline (mark-visited x) (cons VISITED-TAG x))
(define-inline (mark-empty x) (cons NOTHING-TAG x))

(define-inline (members-of x) (cdr x))

(define-inline (done? x)
  (and (pair? x)
       (visited? x)
       (null? (cdr x))))

;;

(define atom->string (atom->string/for ->string))

(define ((atom->string/for to-string) n)
  (cond ((string? n)
          (string-append "\"" n "\""))
        ((char? n)
          (string-append "#\\" (string n)))
        (else
          (to-string n))))

(define (draw-tree n #!optional (max-width 7) (atom->string atom->string))

  (define *branch* (string-pad-right "[o|o]" (add1 max-width) #\-))

  (define (draw-fixed-string s)
    (let* ((k (string-length s))
           (s (cond ((fx> k max-width) (substring s 0 max-width))
                    ((fx< k MIN-WIDTH) (string-pad s MIN-WIDTH #\space))
                    (else s))) )
      (display (string-pad-right s (fx+ max-width 1) #\space))) )

  (define (draw-atom n)
    (draw-fixed-string (atom->string n)))

  (define (draw-conses n)
    (let draw-conses ((n n) (r '()))
      (cond ((not (pair? n))
              (draw-atom n)
              (reverse! r))
            ((null? (cdr n))
              (display "[o|/]")
              (reverse! (cons (car n) r)))
            (else
              (display *branch*)
              (draw-conses (cdr n) (cons (car n) r))))))

  (define (draw-bars n)
    (let draw-bars ((n (members-of n)))
      (cond ((not (pair? n)))
            ((empty? (car n))
              (draw-fixed-string "")
              (draw-bars (cdr n)))
            ((and (pair? (car n))
                  (visited? (car n)))
              (draw-bars (members-of (car n)))
              (draw-bars (cdr n)))
            (else
              (draw-fixed-string "|")
              (draw-bars (cdr n))))))

  (define (skip-empty n)
    (if (and (pair? n)
             (or (empty? (car n))
                 (done? (car n))))
        (skip-empty (cdr n))
        n))

  (define (remove-trailing-nothing n)
    (reverse (skip-empty (reverse n))))

  (define (all-vertical? n)
    (or (not (pair? n))
        (and (null? (cdr n))
             (all-vertical? (car n)))))

  (define (draw-members n)
    (let draw-members ((n (members-of n))
                       (r '()))
      (cond ((not (pair? n))
              (mark-visited
                (remove-trailing-nothing
                  (reverse r))))
            ((empty? (car n))
              (draw-fixed-string "")
              (draw-members (cdr n)
                            (mark-empty r)))
            ((not (pair? (car n)))
              (draw-atom (car n))
              (draw-members (cdr n)
                            (mark-empty r)))
            ((null? (cdr n))
              (draw-members (cdr n)
                            (cons (draw-final (car n)) r)))
            ((all-vertical? (car n))
              (draw-fixed-string "[o|/]")
              (draw-members (cdr n)
                            (cons (caar n) r)))
            (else
              (draw-fixed-string "|")
              (draw-members (cdr n)
                            (cons (car n) r))))))

  (define (draw-final n)
    (cond ((not (pair? n))
            (draw-atom n)
            NOTHING-TAG)
          ((visited? n)
            (draw-members n))
          (else
            (mark-visited (draw-conses n)))))

  ;
  (assert (fixnum? max-width))
  (assert (fx< MIN-WIDTH max-width))
  (assert (procedure? atom->string))

  ;
  (if (not (pair? n)) (draw-atom n)
      (let draw-tree ((n (mark-visited (draw-conses n))))
        (unless (done? n)
          (newline)
          (draw-bars n)
          (newline)
          (draw-tree (draw-members n)))))
  (newline))

) ;module (s9fes draw-tree)