avl.tbs

;; Invariants are baked into the data definition. Each Node construction
;; verifies the AVL invariant locally; the smart constructor `mk-node`
;; both verifies its inputs and certifies its output.

(deftype AVL [A]
  (Leaf)
  (Node :key    A
        :left   (AVL A)
        :right  (refine right (AVL A)
                  (<= (abs (- (avl-height left) (avl-height right))) 1))
        :height (refine h int
                  (= h (+ 1 (max (avl-height left) (avl-height right)))))))

(defmeasure avl-height :: (-> (AVL int) int)
  (fn [t] (if (Leaf? t) 0 (Node-height t))))

;; ---- Type aliases for readability ----

(deftype-alias AVLN [A H]
  (refine v (AVL A) (= (avl-height v) H)))

;; Takes a value + balanced subtrees and produces an AVL with the
;; correctly-computed height. The right-arg refinement enforces the
;; balance precondition; the output refinement guarantees the result's
;; height matches `1 + max(height left, height right)`.

(define mk-node :: (-> int (AVL int)
                       (refine right (AVL int)
                         (<= (abs (- (avl-height left) (avl-height right))) 1))
                       (refine v (AVL int)
                         (= (avl-height v)
                            (+ 1 (max (avl-height left) (avl-height right))))))
  (fn [k left right]
    (Node k left right (+ 1 (max (avl-height left) (avl-height right))))))

;; ---- Building balanced trees ----
;;
;; Each intermediate is typed with `AVLN int H` so its height is part
;; of the type and propagates into subsequent calls. mk-node's output
;; refinement is what makes this typecheck — it certifies the result's
;; height equals `1 + max(left, right)`.

(define t1 :: (AVLN int 1) (mk-node 5 (Leaf) (Leaf)))
(define t2 :: (AVLN int 2) (mk-node 10 t1 (Leaf)))
(define t3 :: (AVLN int 2) (mk-node 7 t1 t1))

;; ---- Demo of rejection ----
;;
;; (mk-node 1 t3 (Leaf))  ;; rejected: t3 has height 2, Leaf height 0, |2-0|=2 > 1
;;
;; Uncommenting the line above yields:
;;   refinement VIOLATED: mk-node arg 3 does not satisfy
;;     (<= (abs (- (avl-height t3) (avl-height right))) 1)