tries-0.0.3: bench/Build.hs

module Build where



import Data.List.NonEmpty (NonEmpty (..))
import Data.Tree (Tree (..))
import Data.Trie.List              as L
import Data.Trie.Map               as M
import Data.Trie.HashMap           as HM
import Data.Tree.Knuth.Forest      as K
import Data.Trie.Knuth             as K
import qualified Data.Map          as Map
import qualified Data.HashMap.Lazy as HMap
import Control.Monad.State


buildMiddleQuery :: Int -> NonEmpty Int
buildMiddleQuery x = 0 :| go x
  where
    go n =
      let half = floor ((fromIntegral n)/2)
      in if half == 0
         then []
         else half : go (floor $ (fromIntegral n)/2)


genListTrie :: Int -> ListTrie Int Int
genListTrie n = ListTrie $ Node (0, Just 0) $ genTree n
  where
    genTree :: Int -> [Tree (Int, Maybe Int)]
    genTree n = flip evalState 1 $ replicateM n $ do
      i <- getSucc
      return $ Node (i, Just 0) $ genTree (floor $ (fromIntegral n)/2)


getSucc :: State Int Int
getSucc = do
  i <- get
  modify (+1)
  pure i


genMapTrie :: Int -> MapTrie Int Int
genMapTrie n = MapTrie $ MapStep $ Map.singleton 0
  (Just 0, Just $ MapTrie $ MapStep $ genMapTree n)
  where
    genMapTree :: Int -> Map.Map Int (Maybe Int, Maybe (MapTrie Int Int))
    genMapTree n =
      Map.unions . flip evalState 1 $
        replicateM n $ do
          i <- getSucc
          pure $ Map.singleton i ( Just 0
                                 , if n <= 1
                                   then Nothing
                                   else Just . MapTrie . MapStep $ genMapTree (floor $ (fromIntegral n)/2)
                                 )

genHashMapTrie :: Int -> HashMapTrie Int Int
genHashMapTrie n =
  HashMapTrie $ HashMapStep $ HMap.singleton 0
    (Just 0, Just . HashMapTrie . HashMapStep $ genHashMapTree n)
  where
    genHashMapTree :: Int -> HMap.HashMap Int (Maybe Int, Maybe (HashMapTrie Int Int))
    genHashMapTree n = HMap.unions $ flip evalState 1 $
      replicateM n $ do
        i <- getSucc
        return $ HMap.singleton i ( Just 0
                                  , if n <= 1
                                    then Nothing
                                    else Just . HashMapTrie . HashMapStep $ genHashMapTree (floor $ (fromIntegral n)/2)
                                  )


-- | This one is ordered largest first
genKnuthTrie :: Int -> KnuthTrie Int Int
genKnuthTrie n = KnuthTrie $ Fork (0, Just 0) (genKnuthForest n n) Nil
  where
    genKnuthForest :: Int -> Int -> KnuthForest (Int, Maybe Int)
    genKnuthForest 0 _ = Nil
    genKnuthForest n s = let cs = if n == 0 then Nil else genKnuthForest (floor $ (fromIntegral n)/2) (n-1)
                             ss = if s == 0 then Nil else genKnuthForest n (s-1)
                         in Fork (s, Just 0) cs ss