Source code on Github
------------------------------------------------------------------------
-- The Agda standard library
--
-- AVL trees where the stored values may depend on their key
------------------------------------------------------------------------

{-# OPTIONS --cubical-compatible --safe #-}

open import Relation.Binary using (StrictTotalOrder)

module Data.Tree.AVL.Indexed
  {a ℓ₁ ℓ₂} (strictTotalOrder : StrictTotalOrder a ℓ₁ ℓ₂) where

open import Level using (Level; _⊔_)
open import Data.Nat.Base using (; zero; suc; _+_)
open import Data.Product using (Σ; ; _×_; _,_; proj₁)
open import Data.Maybe.Base using (Maybe; just; nothing)
open import Data.List.Base as List using (List)
open import Data.DifferenceList using (DiffList; []; _∷_; _++_)
open import Function.Base as F hiding (const)
open import Relation.Unary
open import Relation.Binary using (_Respects_; Tri; tri<; tri≈; tri>)
open import Relation.Binary.PropositionalEquality using (_≡_; refl)

private
  variable
    l v : Level
    A : Set l

open StrictTotalOrder strictTotalOrder renaming (Carrier to Key)

------------------------------------------------------------------------
-- Re-export core definitions publicly

open import Data.Tree.AVL.Key strictTotalOrder public
open import Data.Tree.AVL.Value Eq.setoid public
open import Data.Tree.AVL.Height public

------------------------------------------------------------------------
-- Definitions of the tree

-- The trees have three parameters/indices: a lower bound on the
-- keys, an upper bound, and a height.
--
-- (The bal argument is the balance factor.)

data Tree {v} (V : Value v) (l u : Key⁺) :   Set (a  v  ℓ₂) where
  leaf : (l<u : l <⁺ u)  Tree V l u 0
  node :  {  h}
         (kv : K& V)
         (lk : Tree V l [ kv .key ] )
         (ku : Tree V [ kv .key ] u )
         (bal :     h) 
         Tree V l u (suc h)

module _ {v} {V : Value v} where

  ordered :  {l u n}  Tree V l u n  l <⁺ u
  ordered (leaf l<u)          = l<u
  ordered (node kv lk ku bal) = trans⁺ _ (ordered lk) (ordered ku)

  private
    Val = Value.family V
    V≈  = Value.respects V

  leaf-injective :  {l u} {p q : l <⁺ u}  (Tree V l u 0  leaf p)  leaf q  p  q
  leaf-injective refl = refl

  node-injective-key :
     {hˡ₁ hˡ₂ hʳ₁ hʳ₂ h l u k₁ k₂}
      {lk₁ : Tree V l [ k₁ .key ] hˡ₁} {lk₂ : Tree V l [ k₂ .key ] hˡ₂}
      {ku₁ : Tree V [ k₁ .key ] u hʳ₁} {ku₂ : Tree V [ k₂ .key ] u hʳ₂}
      {bal₁ : hˡ₁  hʳ₁  h} {bal₂ : hˡ₂  hʳ₂  h} 
    node k₁ lk₁ ku₁ bal₁  node k₂ lk₂ ku₂ bal₂  k₁  k₂
  node-injective-key refl = refl

  -- See also node-injectiveˡ, node-injectiveʳ, and node-injective-bal
  -- in Data.Tree.AVL.Indexed.WithK.

  -- Cast operations. Logarithmic in the size of the tree, if we don't
  -- count the time needed to construct the new proofs in the leaf
  -- cases. (The same kind of caveat applies to other operations
  -- below.)
  --
  -- Perhaps it would be worthwhile changing the data structure so
  -- that the casts could be implemented in constant time (excluding
  -- proof manipulation). However, note that this would not change the
  -- worst-case time complexity of the operations below (up to Θ).

  castˡ :  {l m u h}  l <⁺ m  Tree V m u h  Tree V l u h
  castˡ {l} l<m (leaf m<u)         = leaf (trans⁺ l l<m m<u)
  castˡ     l<m (node k mk ku bal) = node k (castˡ l<m mk) ku bal

  castʳ :  {l m u h}  Tree V l m h  m <⁺ u  Tree V l u h
  castʳ {l} (leaf l<m)         m<u = leaf (trans⁺ l l<m m<u)
  castʳ     (node k lk km bal) m<u = node k lk (castʳ km m<u) bal

  -- Various constant-time functions which construct trees out of
  -- smaller pieces, sometimes using rotation.

  pattern node⁺ k₁ t₁ k₂ t₂ t₃ bal = node k₁ t₁ (node k₂ t₂ t₃ bal) ∼+

  joinˡ⁺ :  {l u   h} 
           (k : K& V) 
           ( λ i  Tree V l [ k .key ] (i  )) 
           Tree V [ k .key ] u  
           (bal :     h) 
            λ i  Tree V l u (i  (1 + h))
  joinˡ⁺ k₂ (0# , t₁)                t₃ bal = (0# , node k₂ t₁ t₃ bal)
  joinˡ⁺ k₂ (1# , t₁)                t₃ ∼0  = (1# , node k₂ t₁ t₃ ∼-)
  joinˡ⁺ k₂ (1# , t₁)                t₃ ∼+  = (0# , node k₂ t₁ t₃ ∼0)
  joinˡ⁺ k₄ (1# , node  k₂ t₁ t₃ ∼-) t₅ ∼-  = (0# , node k₂ t₁ (node k₄ t₃ t₅ ∼0) ∼0)
  joinˡ⁺ k₄ (1# , node  k₂ t₁ t₃ ∼0) t₅ ∼-  = (1# , node k₂ t₁ (node k₄ t₃ t₅ ∼-) ∼+)
  joinˡ⁺ k₆ (1# , node⁺ k₂ t₁ k₄ t₃ t₅ bal) t₇ ∼-
    = (0# , node k₄ (node k₂ t₁ t₃ (max∼ bal))
                    (node k₆ t₅ t₇ (∼max bal))
                    ∼0)

  pattern node⁻ k₁ k₂ t₁ t₂ bal t₃ = node k₁ (node k₂ t₁ t₂ bal) t₃ ∼-

  joinʳ⁺ :  {l u   h} 
           (k : K& V) 
           Tree V l [ k .key ]  
           ( λ i  Tree V [ k .key ] u (i  )) 
           (bal :     h) 
            λ i  Tree V l u (i  (1 + h))
  joinʳ⁺ k₂ t₁ (0# , t₃)               bal = (0# , node k₂ t₁ t₃ bal)
  joinʳ⁺ k₂ t₁ (1# , t₃)               ∼0  = (1# , node k₂ t₁ t₃ ∼+)
  joinʳ⁺ k₂ t₁ (1# , t₃)               ∼-  = (0# , node k₂ t₁ t₃ ∼0)
  joinʳ⁺ k₂ t₁ (1# , node k₄ t₃ t₅ ∼+) ∼+  = (0# , node k₄ (node k₂ t₁ t₃ ∼0) t₅ ∼0)
  joinʳ⁺ k₂ t₁ (1# , node k₄ t₃ t₅ ∼0) ∼+  = (1# , node k₄ (node k₂ t₁ t₃ ∼+) t₅ ∼-)
  joinʳ⁺ k₂ t₁ (1# , node⁻ k₆ k₄ t₃ t₅ bal t₇) ∼+
    = (0# , node k₄ (node k₂ t₁ t₃ (max∼ bal))
                    (node k₆ t₅ t₇ (∼max bal))
                    ∼0)

  joinˡ⁻ :  {l u}  { h} 
           (k : K& V) 
           ( λ i  Tree V l [ k .key ] pred[ i   ]) 
           Tree V [ k .key ] u  
           (bal :     h) 
            λ i  Tree V l u (i  h)
  joinˡ⁻ zero    k₂ (0# , t₁) t₃ bal = (1# , node k₂ t₁ t₃ bal)
  joinˡ⁻ zero    k₂ (1# , t₁) t₃ bal = (1# , node k₂ t₁ t₃ bal)
  joinˡ⁻ (suc _) k₂ (0# , t₁) t₃ ∼+  = joinʳ⁺ k₂ t₁ (1# , t₃) ∼+
  joinˡ⁻ (suc _) k₂ (0# , t₁) t₃ ∼0  = (1# , node k₂ t₁ t₃ ∼+)
  joinˡ⁻ (suc _) k₂ (0# , t₁) t₃ ∼-  = (0# , node k₂ t₁ t₃ ∼0)
  joinˡ⁻ (suc _) k₂ (1# , t₁) t₃ bal = (1# , node k₂ t₁ t₃ bal)

  joinʳ⁻ :  {l u }  {h} 
           (k : K& V) 
           Tree V l [ k .key ]  
           ( λ i  Tree V [ k .key ] u pred[ i   ]) 
           (bal :     h) 
            λ i  Tree V l u (i  h)
  joinʳ⁻ zero    k₂ t₁ (0# , t₃) bal = (1# , node k₂ t₁ t₃ bal)
  joinʳ⁻ zero    k₂ t₁ (1# , t₃) bal = (1# , node k₂ t₁ t₃ bal)
  joinʳ⁻ (suc _) k₂ t₁ (0# , t₃) ∼-  = joinˡ⁺ k₂ (1# , t₁) t₃ ∼-
  joinʳ⁻ (suc _) k₂ t₁ (0# , t₃) ∼0  = (1# , node k₂ t₁ t₃ ∼-)
  joinʳ⁻ (suc _) k₂ t₁ (0# , t₃) ∼+  = (0# , node k₂ t₁ t₃ ∼0)
  joinʳ⁻ (suc _) k₂ t₁ (1# , t₃) bal = (1# , node k₂ t₁ t₃ bal)

  -- Extracts the smallest element from the tree, plus the rest.
  -- Logarithmic in the size of the tree.

  headTail :  {l u h}  Tree V l u (1 + h) 
              λ (k : K& V)  l <⁺ [ k .key ] ×
                             λ i  Tree V [ k .key ] u (i  h)
  headTail (node k₁ (leaf l<k₁) t₂ ∼+) = (k₁ , l<k₁ , 0# , t₂)
  headTail (node k₁ (leaf l<k₁) t₂ ∼0) = (k₁ , l<k₁ , 0# , t₂)
  headTail (node { = suc _} k₃ t₁₂ t₄ bal) with headTail t₁₂
  ... | (k₁ , l<k₁ , t₂) = (k₁ , l<k₁ , joinˡ⁻ _ k₃ t₂ t₄ bal)

  -- Extracts the largest element from the tree, plus the rest.
  -- Logarithmic in the size of the tree.

  initLast :  {l u h}  Tree V l u (1 + h) 
              λ (k : K& V)  [ k .key ] <⁺ u ×
                             λ i  Tree V l [ k .key ] (i  h)
  initLast (node k₂ t₁ (leaf k₂<u) ∼-) = (k₂ , k₂<u , (0# , t₁))
  initLast (node k₂ t₁ (leaf k₂<u) ∼0) = (k₂ , k₂<u , (0# , t₁))
  initLast (node { = suc _} k₂ t₁ t₃₄ bal) with initLast t₃₄
  ... | (k₄ , k₄<u , t₃) = (k₄ , k₄<u , joinʳ⁻ _ k₂ t₁ t₃ bal)

  -- Another joining function. Logarithmic in the size of either of
  -- the input trees (which need to have almost equal heights).

  join :  {l m u   h} 
         Tree V l m   Tree V m u   (bal :     h) 
          λ i  Tree V l u (i  h)
  join t₁ (leaf m<u) ∼0 = (0# , castʳ t₁ m<u)
  join t₁ (leaf m<u) ∼- = (0# , castʳ t₁ m<u)
  join { = suc _} t₁ t₂₃ bal with headTail t₂₃
  ... | (k₂ , m<k₂ , t₃) = joinʳ⁻ _ k₂ (castʳ t₁ m<k₂) t₃ bal

  -- An empty tree.

  empty :  {l u}  l <⁺ u  Tree V l u 0
  empty = leaf

  -- A singleton tree.

  singleton :  {l u} (k : Key)  Val k  l < k < u  Tree V l u 1
  singleton k v (l<k , k<u) = node (k , v) (leaf l<k) (leaf k<u) ∼0

  -- Inserts a key into the tree, using a function to combine any
  -- existing value with the new value. Logarithmic in the size of the
  -- tree (assuming constant-time comparisons and a constant-time
  -- combining function).

  insertWith :  {l u h} (k : Key)  (Maybe (Val k)  Val k)   -- Maybe old → result.
               Tree V l u h  l < k < u 
                λ i  Tree V l u (i  h)
  insertWith k f (leaf l<u) l<k<u = (1# , singleton k (f nothing) l<k<u)
  insertWith k f (node (k′ , v′) lp pu bal) (l<k , k<u) with compare k k′
  ... | tri< k<k′ _ _ = joinˡ⁺ (k′ , v′) (insertWith k f lp (l<k , [ k<k′ ]ᴿ)) pu bal
  ... | tri> _ _ k′<k = joinʳ⁺ (k′ , v′) lp (insertWith k f pu ([ k′<k ]ᴿ , k<u)) bal
  ... | tri≈ _ k≈k′ _ = (0# , node (k′ , V≈ k≈k′ (f (just (V≈ (Eq.sym k≈k′) v′)))) lp pu bal)

  -- Inserts a key into the tree. If the key already exists, then it
  -- is replaced. Logarithmic in the size of the tree (assuming
  -- constant-time comparisons).

  insert :  {l u h}  (k : Key)  Val k  Tree V l u h  l < k < u 
            λ i  Tree V l u (i  h)
  insert k v = insertWith k (F.const v)

  -- Deletes the key/value pair containing the given key, if any.
  -- Logarithmic in the size of the tree (assuming constant-time
  -- comparisons).

  delete :  {l u h} (k : Key)  Tree V l u h  l < k < u 
            λ i  Tree V l u pred[ i  h ]
  delete k (leaf l<u) l<k<u = (0# , leaf l<u)
  delete k (node p@(k′ , v) lp pu bal) (l<k , k<u) with compare k′ k
  ... | tri< k′<k _ _ = joinʳ⁻ _ p lp (delete k pu ([ k′<k ]ᴿ , k<u)) bal
  ... | tri> _ _ k′>k = joinˡ⁻ _ p (delete k lp (l<k , [ k′>k ]ᴿ)) pu bal
  ... | tri≈ _ k′≡k _ = join lp pu bal

  -- Looks up a key. Logarithmic in the size of the tree (assuming
  -- constant-time comparisons).

  lookup :  {l u h} (k : Key)  Tree V l u h  l < k < u  Maybe (Val k)
  lookup k (leaf _) l<k<u = nothing
  lookup k (node (k′ , v) lk′ k′u _) (l<k , k<u) with compare k′ k
  ... | tri< k′<k _ _ = lookup k k′u ([ k′<k ]ᴿ , k<u)
  ... | tri> _ _ k′>k = lookup k lk′ (l<k , [ k′>k ]ᴿ)
  ... | tri≈ _ k′≡k _ = just (V≈ k′≡k v)

  -- Converts the tree to an ordered list. Linear in the size of the
  -- tree.

  foldr :  {l u h}  (∀ {k}  Val k  A  A)  A  Tree V l u h  A
  foldr cons nil (leaf l<u)             = nil
  foldr cons nil (node (_ , v) l r bal) = foldr cons (cons v (foldr cons nil r)) l

  toDiffList :  {l u h}  Tree V l u h  DiffList (K& V)
  toDiffList (leaf _)       = []
  toDiffList (node k l r _) = toDiffList l ++ k  toDiffList r

  toList :  {l u h}  Tree V l u h  List (K& V)
  toList t = toDiffList t List.[]

  size :  {l u h}  Tree V l u h  
  size = List.length ∘′ toList

module _ {v w} {V : Value v} {W : Value w} where

  private
    Val = Value.family V
    Wal = Value.family W

  -- Maps a function over all values in the tree.

  map : ∀[ Val  Wal ]   {l u h}  Tree V l u h  Tree W l u h
  map f (leaf l<u)             = leaf l<u
  map f (node (k , v) l r bal) = node (k , f v) (map f l) (map f r) bal