#lang plait
(define-type BoxTree
  [mt]
  [node (b : (Boxof (Optionof Number))) (l : BoxTree) (r : BoxTree)])

#;(define (unify-trees! t1 t2)
  (cond
    [(and (mt? t1) (mt? t2)) (void)]
    [(or (mt? t1) (mt? t2)) (error 'unify-trees! "shape mismatch")]
    [else
     (let* ([t1-l (node-l t1)]
            [t1-r (node-r t1)]
            [t2-l (node-l t2)]
            [t2-r (node-r t2)])
       (begin
         (type-case BoxTree t1-l
           [(node b l r) (type-case (Optionof Number) (unbox b)
                           [(some b) (set-box! (node-b t1-r) (unbox (node-b t2-r)))]
                           [(none) (set-box! (node-b t1-l) (unbox (node-b t2-l)))])]
           [else (void)])
         (type-case BoxTree t2-l
           [(node b l r) (type-case (Optionof Number) (unbox b)
                           [(some b) (set-box! (node-b t2-r) (unbox (node-b t1-r)))]
                           [(none) (set-box! (node-b t2-l) (unbox (node-b t1-l)))])]
           [else (unify-trees! )])))]))

#;(define (unify-trees! t1 t2)
  (cond
    [(and (mt? t1) (mt? t2)) (void)]
    [(or (mt? t1) (mt? t2)) (error 'unify-trees! "shape mismatch")]
    [else
     (let* ([t1-l (node-l t1)]
            [t1-r (node-r t1)]
            [t2-l (node-l t2)]
            [t2-r (node-r t2)])
       (begin
         (type-case BoxTree t1-l
           [(node b l r) (type-case (Optionof Number) (unbox b)
                           [(some b) (set-box! (node-b t1-r) (unbox (node-b t2-r)))]
                           [(none) (set-box! (node-b t1-l) (unbox (node-b t2-l)))])]
           [else (unify-trees! t1-l t1-r)])
         (type-case BoxTree t2-l
           [(node b l r) (type-case (Optionof Number) (unbox b)
                           [(some b) (set-box! (node-b t2-r) (unbox (node-b t1-r)))]
                           [(none) (set-box! (node-b t2-l) (unbox (node-b t1-l)))])]
           [else (unify-trees! t2-l t2-r)])))]))

; should only call set box once per branch
; should probably call unify trees after calling set box
; 

(define (unify-trees! t1 t2)
  (cond
    [(and (mt? t1) (mt? t2)) (void)]
    [(or (mt? t1) (mt? t2)) (error 'unify-trees! "shape mismatch")]
    [else
     (begin
       (type-case (Optionof Number) (unbox (node-b t1))
         [(some v1)
          (type-case (Optionof Number) (unbox (node-b t2))
            [(some v2) (if (eq? v1 v2)
                           (type-case BoxTree (node-l t1)
                             [(node b l r) (begin
                                             (set-box! (node-b (node-r t1)) (unbox (node-b (node-r t2))))
                                             (set-box! (node-b (node-l t2)) (unbox (node-b (node-l t1)))))]
                             [(mt) (begin
                                             (set-box! (node-b (node-l t1)) (unbox (node-b (node-l t2))))
                                             (set-box! (node-b (node-r t2)) (unbox (node-b (node-r t1)))))])
                       (error 'unify-trees! "value mismatch"))]
            [(none) (begin
                      (set-box! (node-b (node-l t1)) (unbox (node-b (node-l t2))))
                      (set-box! (node-b (node-r t2)) (unbox (node-b (node-r t1)))))])]
         [(none) (begin
                   (unify-trees! (node-l t1) (node-r t1))
                   (unify-trees! (node-l t2) (node-r t2)))]))]))
     



(module+ test
  ;; Same shape, one "variable" in each
  (let ([t1 (node (box (some 0))
                  (node (box (some 1)) (mt) (mt))
                  (node (box (none)) (mt) (mt)))]
        [t2 (node (box (some 0))
                  (node (box (none)) (mt) (mt))
                  (node (box (some 2)) (mt) (mt)))]
        [result (node (box (some 0))
                      (node (box (some 1)) (mt) (mt))
                      (node (box (some 2)) (mt) (mt)))])
    (begin
      (unify-trees! t1 t2)
      (test t1 result)
      (test t2 result)))

  (test/exn (unify-trees! (node (box (some 1))
                                (node (box (none)) (mt) (mt))
                                (mt))
                          (node (box (none)) (mt) (mt)))
            "shape mismatch")
  (test/exn (unify-trees! (node (box (some 1)) (mt) (mt))
                          (node (box (some 2)) (mt) (mt))) "value mismatch")

  ;; Only variables
  (let ([t1 (node (box (none)) (mt) (mt))]
        [t2 (node (box (none)) (mt) (mt))]
        [result (node (box (none)) (mt) (mt))])
    (begin
      (unify-trees! t1 t2)
      (test t1 result)
      (test t2 result)))

  ;; compatible elements, different shape
  (test/exn
   (unify-trees!
    (node (box (some 0))
          (node (box (some 1)) (mt) (node (box (none)) (mt) (mt)))
          (mt))
    (node (box (some 0))
          (mt)
          (node (box (some 1)) (node (box (none)) (mt) (mt)) (mt))))
   "shape mismatch")

  ;; deeper trees
  (let ([result (node
                 (box (some 0))
                 (node
                  (box (some 1))
                  (node (box (some 2))
                        (node (box (some 3))
                              (node (box (some 4)) (mt) (mt))
                              (mt))
                        (mt))
                  (mt))
                 (mt))]
        [t1  (node
              (box (none))
              (node
               (box (some 1))
               (node (box (none))
                     (node (box (some 3))
                           (node (box (some 4)) (mt) (mt))
                           (mt))
                     (mt))
               (mt))
              (mt))]
        [t2  (node
              (box (some 0))
              (node
               (box (none))
               (node (box (some 2))
                     (node (box (none))
                           (node (box (some 4)) (mt) (mt))
                           (mt))
                     (mt))
               (mt))
              (mt))])
    (begin
      (unify-trees! t1 t2)
      (test t1 t2)
      (test t1 result)
      (test t2 result)))
  )