-- |
-- Module      :  Cryptol.REPL.Trie
-- 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 (Int -> Trie a -> ShowS
[Trie a] -> ShowS
Trie a -> String
(Int -> Trie a -> ShowS)
-> (Trie a -> String) -> ([Trie a] -> ShowS) -> Show (Trie a)
forall a. Show a => Int -> Trie a -> ShowS
forall a. Show a => [Trie a] -> ShowS
forall a. Show a => Trie a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Trie a] -> ShowS
$cshowList :: forall a. Show a => [Trie a] -> ShowS
show :: Trie a -> String
$cshow :: forall a. Show a => Trie a -> String
showsPrec :: Int -> Trie a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
Show)

emptyTrie :: Trie a
emptyTrie :: Trie a
emptyTrie  = Map Char (Trie a) -> Maybe a -> Trie a
forall a. Map Char (Trie a) -> Maybe a -> Trie a
Node Map Char (Trie a)
forall k a. Map k a
Map.empty Maybe a
forall a. Maybe a
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 :: String -> a -> Trie a -> Trie a
insertTrie String
k a
a = String -> Trie a -> Trie a
loop String
k
  where
  loop :: String -> Trie a -> Trie a
loop String
key (Node Map Char (Trie a)
m Maybe a
mb) = case String
key of
    Char
c:String
cs -> Map Char (Trie a) -> Maybe a -> Trie a
forall a. Map Char (Trie a) -> Maybe a -> Trie a
Node ((Maybe (Trie a) -> Maybe (Trie a))
-> Char -> Map Char (Trie a) -> Map Char (Trie a)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (Trie a -> Maybe (Trie a)
forall a. a -> Maybe a
Just (Trie a -> Maybe (Trie a))
-> (Maybe (Trie a) -> Trie a) -> Maybe (Trie a) -> Maybe (Trie a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Trie a -> Trie a
loop String
cs (Trie a -> Trie a)
-> (Maybe (Trie a) -> Trie a) -> Maybe (Trie a) -> Trie a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie a -> Maybe (Trie a) -> Trie a
forall a. a -> Maybe a -> a
fromMaybe Trie a
forall a. Trie a
emptyTrie) Char
c Map Char (Trie a)
m) Maybe a
mb
    []   -> case Maybe a
mb of
      Maybe a
Nothing -> Map Char (Trie a) -> Maybe a -> Trie a
forall a. Map Char (Trie a) -> Maybe a -> Trie a
Node Map Char (Trie a)
m (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
      Just a
_  -> String -> [String] -> Trie a
forall a. HasCallStack => String -> [String] -> a
panic String
"[REPL] Trie" [String
"key already exists:", String
"\t" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
k]

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

  Char
c:String
cs -> case Char -> Map Char (Trie a) -> Maybe (Trie a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Char
c Map Char (Trie a)
mp of
    Just Trie a
m' -> String -> Trie a -> [a]
forall a. String -> Trie a -> [a]
lookupTrie String
cs Trie a
m'
    Maybe (Trie a)
Nothing -> []

  [] -> Trie a -> [a]
forall a. Trie a -> [a]
leaves Trie a
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 :: String -> Trie a -> [a]
lookupTrieExact []     (Node Map Char (Trie a)
_ (Just a
x)) = a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
lookupTrieExact []     Trie a
t                 = Trie a -> [a]
forall a. Trie a -> [a]
leaves Trie a
t
lookupTrieExact (Char
c:String
cs) (Node Map Char (Trie a)
mp Maybe a
_)       =
  case Char -> Map Char (Trie a) -> Maybe (Trie a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Char
c Map Char (Trie a)
mp of
    Just Trie a
m' -> String -> Trie a -> [a]
forall a. String -> Trie a -> [a]
lookupTrieExact String
cs Trie a
m'
    Maybe (Trie a)
Nothing -> []

-- | Return all of the values from a Trie.
leaves :: Trie a -> [a]
leaves :: Trie a -> [a]
leaves (Node Map Char (Trie a)
mp Maybe a
mb) = Maybe a -> [a]
forall a. Maybe a -> [a]
maybeToList Maybe a
mb [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (Trie a -> [a]) -> [Trie a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Trie a -> [a]
forall a. Trie a -> [a]
leaves (Map Char (Trie a) -> [Trie a]
forall k a. Map k a -> [a]
Map.elems Map Char (Trie a)
mp)