rose-ash/lib/haskell/map.sx

;; map.sx — Phase 11 Data.Map: weight-balanced BST in pure SX.
;;
;; Algorithm: Adams's weight-balanced tree (the same family as Haskell's
;; Data.Map). Each node tracks its size; rotations maintain the invariant
;;
;;     size(small-side) * delta >= size(large-side)        (delta = 3)
;;
;; with single or double rotations chosen by the gamma ratio (gamma = 2).
;; The size field is an Int and is included so `size`, `lookup`, etc. are
;; O(log n) on both extremes of the tree.
;;
;; Representation:
;;   Empty → ("Map-Empty")
;;   Node  → ("Map-Node" key val left right size)
;;
;; All operations are pure SX — no mutation of nodes once constructed.
;; The user-facing Haskell layer (Phase 11 next iteration) wraps these
;; for `import Data.Map as Map`.

;; ── Constructors ────────────────────────────────────────────
(define hk-map-empty (list "Map-Empty"))

(define
  hk-map-node
  (fn
    (k v l r)
    (list "Map-Node" k v l r (+ 1 (+ (hk-map-size l) (hk-map-size r))))))

;; ── Predicates and accessors ────────────────────────────────
(define hk-map-empty? (fn (m) (and (list? m) (= (first m) "Map-Empty"))))

(define hk-map-node? (fn (m) (and (list? m) (= (first m) "Map-Node"))))

(define
  hk-map-size
  (fn (m) (cond ((hk-map-empty? m) 0) (:else (nth m 5)))))

(define hk-map-key (fn (m) (nth m 1)))
(define hk-map-val (fn (m) (nth m 2)))
(define hk-map-left (fn (m) (nth m 3)))
(define hk-map-right (fn (m) (nth m 4)))

;; ── Weight-balanced rotations ───────────────────────────────
;; delta and gamma per Adams 1992 / Haskell Data.Map.

(define hk-map-delta 3)
(define hk-map-gamma 2)

(define
  hk-map-single-l
  (fn
    (k v l r)
    (let
      ((rk (hk-map-key r))
        (rv (hk-map-val r))
        (rl (hk-map-left r))
        (rr (hk-map-right r)))
      (hk-map-node rk rv (hk-map-node k v l rl) rr))))

(define
  hk-map-single-r
  (fn
    (k v l r)
    (let
      ((lk (hk-map-key l))
        (lv (hk-map-val l))
        (ll (hk-map-left l))
        (lr (hk-map-right l)))
      (hk-map-node lk lv ll (hk-map-node k v lr r)))))

(define
  hk-map-double-l
  (fn
    (k v l r)
    (let
      ((rk (hk-map-key r))
        (rv (hk-map-val r))
        (rl (hk-map-left r))
        (rr (hk-map-right r))
        (rlk (hk-map-key (hk-map-left r)))
        (rlv (hk-map-val (hk-map-left r)))
        (rll (hk-map-left (hk-map-left r)))
        (rlr (hk-map-right (hk-map-left r))))
      (hk-map-node
        rlk
        rlv
        (hk-map-node k v l rll)
        (hk-map-node rk rv rlr rr)))))

(define
  hk-map-double-r
  (fn
    (k v l r)
    (let
      ((lk (hk-map-key l))
        (lv (hk-map-val l))
        (ll (hk-map-left l))
        (lr (hk-map-right l))
        (lrk (hk-map-key (hk-map-right l)))
        (lrv (hk-map-val (hk-map-right l)))
        (lrl (hk-map-left (hk-map-right l)))
        (lrr (hk-map-right (hk-map-right l))))
      (hk-map-node
        lrk
        lrv
        (hk-map-node lk lv ll lrl)
        (hk-map-node k v lrr r)))))

;; ── Balanced node constructor ──────────────────────────────
;; Use this in place of hk-map-node when one side may have grown
;; or shrunk by one and we need to restore the weight invariant.
(define
  hk-map-balance
  (fn
    (k v l r)
    (let
      ((sl (hk-map-size l)) (sr (hk-map-size r)))
      (cond
        ((<= (+ sl sr) 1) (hk-map-node k v l r))
        ((> sr (* hk-map-delta sl))
          (let
            ((rl (hk-map-left r)) (rr (hk-map-right r)))
            (cond
              ((< (hk-map-size rl) (* hk-map-gamma (hk-map-size rr)))
                (hk-map-single-l k v l r))
              (:else (hk-map-double-l k v l r)))))
        ((> sl (* hk-map-delta sr))
          (let
            ((ll (hk-map-left l)) (lr (hk-map-right l)))
            (cond
              ((< (hk-map-size lr) (* hk-map-gamma (hk-map-size ll)))
                (hk-map-single-r k v l r))
              (:else (hk-map-double-r k v l r)))))
        (:else (hk-map-node k v l r))))))

(define
  hk-map-singleton
  (fn (k v) (hk-map-node k v hk-map-empty hk-map-empty)))

(define
  hk-map-insert
  (fn
    (k v m)
    (cond
      ((hk-map-empty? m) (hk-map-singleton k v))
      (:else
        (let
          ((mk (hk-map-key m)))
          (cond
            ((< k mk)
              (hk-map-balance
                mk
                (hk-map-val m)
                (hk-map-insert k v (hk-map-left m))
                (hk-map-right m)))
            ((> k mk)
              (hk-map-balance
                mk
                (hk-map-val m)
                (hk-map-left m)
                (hk-map-insert k v (hk-map-right m))))
            (:else (hk-map-node k v (hk-map-left m) (hk-map-right m)))))))))

(define
  hk-map-lookup
  (fn
    (k m)
    (cond
      ((hk-map-empty? m) (list "Nothing"))
      (:else
        (let
          ((mk (hk-map-key m)))
          (cond
            ((< k mk) (hk-map-lookup k (hk-map-left m)))
            ((> k mk) (hk-map-lookup k (hk-map-right m)))
            (:else (list "Just" (hk-map-val m)))))))))

(define
  hk-map-member
  (fn
    (k m)
    (cond
      ((hk-map-empty? m) false)
      (:else
        (let
          ((mk (hk-map-key m)))
          (cond
            ((< k mk) (hk-map-member k (hk-map-left m)))
            ((> k mk) (hk-map-member k (hk-map-right m)))
            (:else true)))))))

(define hk-map-null hk-map-empty?)

(define
  hk-map-find-min
  (fn
    (m)
    (cond
      ((hk-map-empty? (hk-map-left m))
        (list (hk-map-key m) (hk-map-val m)))
      (:else (hk-map-find-min (hk-map-left m))))))

(define
  hk-map-delete-min
  (fn
    (m)
    (cond
      ((hk-map-empty? (hk-map-left m)) (hk-map-right m))
      (:else
        (hk-map-balance
          (hk-map-key m)
          (hk-map-val m)
          (hk-map-delete-min (hk-map-left m))
          (hk-map-right m))))))

(define
  hk-map-find-max
  (fn
    (m)
    (cond
      ((hk-map-empty? (hk-map-right m))
        (list (hk-map-key m) (hk-map-val m)))
      (:else (hk-map-find-max (hk-map-right m))))))

(define
  hk-map-delete-max
  (fn
    (m)
    (cond
      ((hk-map-empty? (hk-map-right m)) (hk-map-left m))
      (:else
        (hk-map-balance
          (hk-map-key m)
          (hk-map-val m)
          (hk-map-left m)
          (hk-map-delete-max (hk-map-right m)))))))

(define
  hk-map-glue
  (fn
    (l r)
    (cond
      ((hk-map-empty? l) r)
      ((hk-map-empty? r) l)
      ((> (hk-map-size l) (hk-map-size r))
        (let
          ((mp (hk-map-find-max l)))
          (hk-map-balance (first mp) (nth mp 1) (hk-map-delete-max l) r)))
      (:else
        (let
          ((mp (hk-map-find-min r)))
          (hk-map-balance (first mp) (nth mp 1) l (hk-map-delete-min r)))))))

(define
  hk-map-delete
  (fn
    (k m)
    (cond
      ((hk-map-empty? m) m)
      (:else
        (let
          ((mk (hk-map-key m)))
          (cond
            ((< k mk)
              (hk-map-balance
                mk
                (hk-map-val m)
                (hk-map-delete k (hk-map-left m))
                (hk-map-right m)))
            ((> k mk)
              (hk-map-balance
                mk
                (hk-map-val m)
                (hk-map-left m)
                (hk-map-delete k (hk-map-right m))))
            (:else (hk-map-glue (hk-map-left m) (hk-map-right m)))))))))

(define
  hk-map-from-list
  (fn
    (pairs)
    (reduce
      (fn (acc p) (hk-map-insert (first p) (nth p 1) acc))
      hk-map-empty
      pairs)))

(define
  hk-map-to-asc-list
  (fn
    (m)
    (cond
      ((hk-map-empty? m) (list))
      (:else
        (append
          (hk-map-to-asc-list (hk-map-left m))
          (cons
            (list (hk-map-key m) (hk-map-val m))
            (hk-map-to-asc-list (hk-map-right m))))))))

(define hk-map-to-list hk-map-to-asc-list)

(define
  hk-map-keys
  (fn
    (m)
    (cond
      ((hk-map-empty? m) (list))
      (:else
        (append
          (hk-map-keys (hk-map-left m))
          (cons (hk-map-key m) (hk-map-keys (hk-map-right m))))))))

(define
  hk-map-elems
  (fn
    (m)
    (cond
      ((hk-map-empty? m) (list))
      (:else
        (append
          (hk-map-elems (hk-map-left m))
          (cons (hk-map-val m) (hk-map-elems (hk-map-right m))))))))

(define
  hk-map-union-with
  (fn
    (f m1 m2)
    (reduce
      (fn
        (acc p)
        (let
          ((k (first p)) (v (nth p 1)))
          (let
            ((look (hk-map-lookup k acc)))
            (cond
              ((= (first look) "Just")
                (hk-map-insert k (f (nth look 1) v) acc))
              (:else (hk-map-insert k v acc))))))
      m1
      (hk-map-to-asc-list m2))))

(define
  hk-map-intersection-with
  (fn
    (f m1 m2)
    (reduce
      (fn
        (acc p)
        (let
          ((k (first p)) (v1 (nth p 1)))
          (let
            ((look (hk-map-lookup k m2)))
            (cond
              ((= (first look) "Just")
                (hk-map-insert k (f v1 (nth look 1)) acc))
              (:else acc)))))
      hk-map-empty
      (hk-map-to-asc-list m1))))

(define
  hk-map-difference
  (fn
    (m1 m2)
    (reduce
      (fn
        (acc p)
        (let
          ((k (first p)) (v (nth p 1)))
          (cond ((hk-map-member k m2) acc) (:else (hk-map-insert k v acc)))))
      hk-map-empty
      (hk-map-to-asc-list m1))))

(define
  hk-map-foldl-with-key
  (fn
    (f acc m)
    (cond
      ((hk-map-empty? m) acc)
      (:else
        (let
          ((acc1 (hk-map-foldl-with-key f acc (hk-map-left m))))
          (let
            ((acc2 (f acc1 (hk-map-key m) (hk-map-val m))))
            (hk-map-foldl-with-key f acc2 (hk-map-right m))))))))

(define
  hk-map-foldr-with-key
  (fn
    (f acc m)
    (cond
      ((hk-map-empty? m) acc)
      (:else
        (let
          ((acc1 (hk-map-foldr-with-key f acc (hk-map-right m))))
          (let
            ((acc2 (f (hk-map-key m) (hk-map-val m) acc1)))
            (hk-map-foldr-with-key f acc2 (hk-map-left m))))))))

(define
  hk-map-map-with-key
  (fn
    (f m)
    (cond
      ((hk-map-empty? m) m)
      (:else
        (list
          "Map-Node"
          (hk-map-key m)
          (f (hk-map-key m) (hk-map-val m))
          (hk-map-map-with-key f (hk-map-left m))
          (hk-map-map-with-key f (hk-map-right m))
          (hk-map-size m))))))

(define
  hk-map-filter-with-key
  (fn
    (p m)
    (hk-map-foldr-with-key
      (fn (k v acc) (cond ((p k v) (hk-map-insert k v acc)) (:else acc)))
      hk-map-empty
      m)))