;; Swap-trees
;; Implemented using streams
;; (basic implementation included)
;;
;; (C) 2006 Michel Salim <msalim@cs.indiana.edu>
;;

(define-syntax delay
  (syntax-rules ()
    [(_ e)
     (lambda () e)]))

(define-syntax force
  (syntax-rules ()
    [(_ th)
     (th)]))

(define-syntax stream-cons
  (syntax-rules ()
    [(_ a d)
     (cons a (delay d))]))

(define stream-car car)
(define-syntax stream-cdr
  (syntax-rules ()
    [(_ s)
     (force (cdr s))]))

(define stream-ref
  (lambda (s n)
    (if (zero? n)
        (stream-car s)
        (stream-ref (stream-cdr s) (sub1 n)))))

(define empty-stream '())
(define empty-stream? null?)

(define stream-append
  (lambda (s1 s2)
    (if (empty-stream? s1)
        s2
        (stream-cons (stream-car s1)
                     (stream-append (stream-cdr s1)
                                    s2)))))

(define gen-stream-equal?
  (lambda (eqp)
    (letrec ([stream-equal?
              (lambda (s1 s2)
                (cond
                  [(empty-stream? s1)
                   (empty-stream? s2)]
                  [(empty-stream? s2) #f]
                  [else (and (eqp (stream-car s1)
                                  (stream-car s2))
                             (stream-equal?
                              (stream-cdr s1)
                              (stream-cdr s2)))]))])
      stream-equal?)))

(define stream-equal?
  (gen-stream-equal? equal?))

(define stream-=
  (gen-stream-equal? =))

(define treewalk
  (lambda (t)
    (if (not (pair? t))
        (stream-cons `(t . (lambda (v)
                             (error "Cannot modify singleton")))
                     empty-stream)
        (stream-append (treewalk^ (car t) #t t)
                       (treewalk^ (cdr t) #f t)))))

(define treewalk^
  (lambda (t left? parent)
    (if (not (pair? t))
        (stream-cons `(,t . ,(if left?
                               (lambda (v)
                                 (set-car! parent v))
                               (lambda (v)
                                 (set-cdr! parent v))))
                     empty-stream)
        (stream-append (treewalk^ (car t) #t t)
                       (treewalk^ (cdr t) #f t)))))
(define seq
  (gen-stream-equal?
   (lambda (p1 p2)
     (printf "Comparing ~s and ~s\n" (car p1) (car p2))
     (printf "Result: ~s\n" (= (car p1) (car p2)))
     (equal? (car p1) (car p2)))))

(define same-fringe?
  (lambda (t1 t2)
    (let ([nodes1 (treewalk t1)]
          [nodes2 (treewalk t2)])
      (seq nodes1 nodes2))))

(define swap-leaves
  (lambda (t1 t2)
    (let loop ([nodes1 (treewalk t1)]
               [nodes2 (treewalk t2)])
      (cond
        [(or (empty-stream? nodes1)
             (empty-stream? nodes2))
         'done]
        [else (let ([np1 (stream-car nodes1)]
                    [np2 (stream-car nodes2)])
                ((cdr np1) (car np2))
                ((cdr np2) (car np1))
                (loop (stream-cdr nodes1)
                      (stream-cdr nodes2)))]))))


(define t1
  '(1 . (2 . 3)))
(define t2
  '((1 . 2) . 3))
(define t3
  '(1 . ((2 . 3) . 4)))
(define t4
  '((a . b) . c))
(define s1
  (treewalk t1))
(define s2
  (treewalk t2))

;; Also copies the tree lazily
;; The closure associated with each node modifies its copy,
;; not the original node

(define treewalk-copy
  (lambda (t)
    (if (not (pair? t))
        (cons (stream-cons
               `(t . (lambda (v)
                       (error "Cannot modify singleton")))
               empty-stream)
              t)
        (let ([copy (cons (car t) (cdr t))])
          (cons
           (stream-append (treewalk-copy^ (car t) #t copy)
                          (treewalk-copy^ (cdr t) #f copy))
           copy)))))

(define treewalk-copy^
  (lambda (t left? copy-parent)
    (if (not (pair? t))
        (stream-cons `(,t . ,(if left?
                               (lambda (v)
                                 (set-car! copy-parent v))
                               (lambda (v)
                                 (set-cdr! copy-parent v))))
                     empty-stream)
        (let ([copy (cons (car t) (cdr t))])
               (if left?
                   (set-car! copy-parent copy)
                   (set-cdr! copy-parent copy))
          (stream-append (treewalk-copy^ (car t) #t copy)
                         (treewalk-copy^ (cdr t) #f copy))))))

;; Modify the copy of the tree
(define swap-leaves-copy
  (lambda (t1 t2)
    (let* ([res1 (treewalk-copy t1)]
           [res2 (treewalk-copy t2)]
           [copy1 (cdr res1)]
           [copy2 (cdr res2)])
      (let loop ([nodes1 (car res1)]
                 [nodes2 (car res2)])
        (cond
          [(or (empty-stream? nodes1)
               (empty-stream? nodes2))
           (list copy1 copy2)]
        [else (let ([np1 (stream-car nodes1)]
                    [np2 (stream-car nodes2)])
                ((cdr np1) (car np2))
                ((cdr np2) (car np1))
                (loop (stream-cdr nodes1)
                      (stream-cdr nodes2)))])))))