------------------------------------------------------------------------
-- The Agda standard library
--
-- Trie, basic type and operations
------------------------------------------------------------------------

-- See README.Data.Trie.NonDependent for an example of using a trie to
-- build a lexer.

{-# OPTIONS --cubical-compatible --sized-types #-}

open import Relation.Binary.Core using (Rel)
open import Relation.Binary.Bundles using (StrictTotalOrder)

module Data.Trie {k e r} (S : StrictTotalOrder k e r) where

open import Level
open import Size
open import Data.List.Base using (List; []; _∷_; _++_)
import Data.List.NonEmpty as List⁺
open import Data.Maybe.Base as Maybe using (Maybe; just; nothing; maybe′)
open import Data.Product.Base using (∃)
open import Data.These.Base as These using (These)
open import Function.Base using (_∘′_; const)
open import Relation.Unary using (IUniversal; _⇒_)

open StrictTotalOrder S
  using (module Eq)
  renaming (Carrier to Key)

open import Data.List.Relation.Binary.Equality.Setoid Eq.setoid
open import Data.Tree.AVL.Value ≋-setoid using (Value)

------------------------------------------------------------------------
-- Definition

-- Trie is defined in terms of Trie⁺, the type of non-empty trie. This
-- guarantees that the trie is minimal: each path in the tree leads to
-- either a value or a number of non-empty sub-tries.

open import Data.Trie.NonEmpty S as Trie⁺ public
  using (Trie⁺; Tries⁺; Word; eat)

Trie : ∀ {v} (V : Value v) → Size → Set (v ⊔ k ⊔ e ⊔ r)
Trie V i = Maybe (Trie⁺ V i)

------------------------------------------------------------------------
-- Operations

-- Functions acting on Trie are wrappers for functions acting on Tries.
-- Sometimes the empty case is handled in a special way (e.g. insertWith
-- calls singleton when faced with an empty Trie).

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

  private Val = Value.family V

------------------------------------------------------------------------
-- Lookup

  lookup : Trie V ∞ → ∀ ks → Maybe (These (Val ks) (Tries⁺ (eat V ks) ∞))
  lookup t ks = t Maybe.>>= λ ts → Trie⁺.lookup ts ks

  lookupValue : Trie V ∞ → ∀ ks → Maybe (Val ks)
  lookupValue t ks = t Maybe.>>= λ ts → Trie⁺.lookupValue ts ks

  lookupTries⁺ : Trie V ∞ → ∀ ks → Maybe (Tries⁺ (eat V ks) ∞)
  lookupTries⁺ t ks = t Maybe.>>= λ ts → Trie⁺.lookupTries⁺ ts ks

  lookupTrie : Trie V ∞ → ∀ k → Trie (eat V (k ∷ [])) ∞
  lookupTrie t k = t Maybe.>>= λ ts → Trie⁺.lookupTrie⁺ ts k

------------------------------------------------------------------------
-- Construction

  empty : Trie V ∞
  empty = nothing

  singleton : ∀ ks → Val ks → Trie V ∞
  singleton ks v = just (Trie⁺.singleton ks v)

  insertWith : ∀ ks → (Maybe (Val ks) → Val ks) → Trie V ∞ → Trie V ∞
  insertWith ks f (just t) = just (Trie⁺.insertWith ks f t)
  insertWith ks f nothing  = singleton ks (f nothing)

  insert : ∀ ks → Val ks → Trie V ∞ → Trie V ∞
  insert ks = insertWith ks ∘′ const

  fromList : List (∃ Val) → Trie V ∞
  fromList = Maybe.map Trie⁺.fromList⁺ ∘′ List⁺.fromList

  toList : Trie V ∞ → List (∃ Val)
  toList (just t) = List⁺.toList (Trie⁺.toList⁺ t)
  toList nothing  = []

------------------------------------------------------------------------
-- Modification

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

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

  map : ∀ {i} → ∀[ Val ⇒ Wal ] → Trie V i → Trie W i
  map = Maybe.map ∘′ Trie⁺.map V W

-- Deletion

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

  -- Use a function to decide how to modify the sub-Trie⁺ whose root is
  -- at the end of path ks.
  deleteWith : ∀ {i} (ks : Word) →
     (∀ {i} → Trie⁺ (eat V ks) i → Maybe (Trie⁺ (eat V ks) i)) →
     Trie V i → Trie V i
  deleteWith ks f t = t Maybe.>>= Trie⁺.deleteWith ks f

  -- Remove the whole node
  deleteTrie⁺ : ∀ {i} (ks : Word) → Trie V i → Trie V i
  deleteTrie⁺ ks t = t Maybe.>>= Trie⁺.deleteTrie⁺ ks

  -- Remove the value and keep the sub-Tries (if any)
  deleteValue : ∀ {i} (ks : Word) → Trie V i → Trie V i
  deleteValue ks t = t Maybe.>>= Trie⁺.deleteValue ks

  -- Remove the sub-Tries and keep the value (if any)
  deleteTries⁺ : ∀ {i} (ks : Word) → Trie V i → Trie V i
  deleteTries⁺ ks t = t Maybe.>>= Trie⁺.deleteTries⁺ ks