Trie in Haskell, using lens

Trie.hs

{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Data.Trie (
    Trie,
    singleton,
    insert,
    delete,
    lookup,
    unionWith,
    union,
    fromList,
    toList)
where

import           Control.Applicative ((<|>))
import           Control.Lens        hiding ((<.>))
import           Data.Foldable       (foldl')
import           Data.Function       (fix)
import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import           Data.Hashable       (Hashable)
import           Prelude             hiding (lookup)

data Trie k v
    = Trie{
        trieLabel :: !(Maybe v),
        trieNext  :: !(HashMap k (Trie k v))}
    deriving (Show, Eq, Functor, Traversable)

instance Foldable (Trie k) where
    foldMap f = fix $ \go Trie{..} -> foldMap f trieLabel <> foldMap go trieNext
    null Trie{..} = null trieLabel && null trieNext

instance AsEmpty (Trie k v) where
    _Empty = nearly empty null

type instance Index (Trie k v) = [k]
type instance IxValue (Trie k v) = v
instance (Hashable k) => Ixed (Trie k v)
instance (Hashable k) => At (Trie k v) where
    at segs = subTrie segs . label

label :: Lens' (Trie k v) (Maybe v)
label = lens trieLabel $ \trie newLabel -> trie{trieLabel = newLabel}

next :: Lens' (Trie k v) (HashMap k (Trie k v))
next = lens trieNext $ \trie newNext -> trie{trieNext = newNext}

prune :: Trie k v -> Maybe (Trie k v)
prune = review $ non' _Empty

empty :: Trie k v
empty = Trie Nothing HM.empty

subTrie :: (Hashable k) => [k] -> Lens' (Trie k v) (Trie k v)
subTrie []           = id
subTrie (seg : segs) = next . at seg . l where
    l f Nothing     = f empty <&> fmap tele . prune
    l f (Just trie) = subTrie segs f trie <&> prune
    tele trie = foldr (\seg' r -> Trie Nothing $ HM.singleton seg' r) trie segs

singleton :: (Hashable k) => [k] -> v -> Trie k v
singleton segs x = insert segs x empty

insert :: (Hashable k) => [k] -> v -> Trie k v -> Trie k v
insert segs x = at segs ?~ x

lookup :: (Hashable k) => [k] -> Trie k v -> Maybe v
lookup = view . at

delete :: (Hashable k) => [k] -> Trie k v -> Trie k v
delete segs = at segs .~ Nothing

unionWith :: (Hashable k) => (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith f = fix $ \go (Trie label0 next0) (Trie label1 next1) -> Trie{
    trieLabel = f <$> label0 <*> label1 <|> label0 <|> label1,
    trieNext  = HM.unionWith go next0 next1}

union :: (Hashable k) => Trie k v -> Trie k v -> Trie k v
union = unionWith const

fromList :: (Hashable k) => [([k], v)] -> Trie k v
fromList = foldl' (\trie (segs, x) -> insert segs x trie) empty

toList :: Trie k v -> [([k], v)]
toList Trie{..} = [([], x) | x <- toListOf _Just trieLabel] ++ do
    (seg, trie) <- HM.toList trieNext
    (segs, x) <- toList trie
    return (seg : segs, x)