    -- |
    -- Module      :  $Header$
    -- Copyright   :  (c) 2013-2016 Galois, Inc.
    -- License     :  BSD3
    -- Maintainer  :  cryptol@galois.com
    -- Stability   :  provisional
    -- Portability :  portable

    module Cryptol.REPL.Trie where

    import           Cryptol.Utils.Panic (panic)
    import qualified Data.Map as Map
    import           Data.Maybe (fromMaybe,maybeToList)

    -- | Maps string names to values, allowing for partial key matches and querying.
    data Trie a = Node (Map.Map Char (Trie a)) (Maybe a)
        deriving (Show)

    emptyTrie :: Trie a
    emptyTrie  = Node Map.empty Nothing

    -- | Insert a value into the Trie.  Will call `panic` if a value already exists
    -- with that key.
    insertTrie :: String -> a -> Trie a -> Trie a
    insertTrie k a = loop k
      where
      loop key (Node m mb) = case key of
        c:cs -> Node (Map.alter (Just . loop cs . fromMaybe emptyTrie) c m) mb
        []   -> case mb of
          Nothing -> Node m (Just a)
          Just _  -> panic "[REPL] Trie" ["key already exists:", "\t" ++ k]

    -- | Return all matches with the given prefix.
    lookupTrie :: String -> Trie a -> [a]
    lookupTrie key t@(Node mp _) = case key of

      c:cs -> case Map.lookup c mp of
        Just m' -> lookupTrie cs m'
        Nothing -> []

      [] -> leaves t

    -- | Given a key, return either an exact match for that key, or all
    -- matches with the given prefix.
    lookupTrieExact :: String -> Trie a -> [a]
    lookupTrieExact []     (Node _ (Just x)) = return x
    lookupTrieExact []     t                 = leaves t
    lookupTrieExact (c:cs) (Node mp _)       =
      case Map.lookup c mp of
        Just m' -> lookupTrieExact cs m'
        Nothing -> []

    -- | Return all of the values from a Trie.
    leaves :: Trie a -> [a]
    leaves (Node mp mb) = maybeToList mb ++ concatMap leaves (Map.elems mp)